context-bar

1// Context bar: the context window as one stacked bar above the prompt, one
2// color per /context category, plus how much room is left before
3// auto-compact (the number that actually matters with
4// CLAUDE_AUTOCOMPACT_PCT_OVERRIDE set below 100).
5//
6// session.start: register /context-bar, restore the on/off choice, read once.
7// session.measure: the engine pushes this after each main-thread turn; take a
8//   `summary` breakdown, which is estimated locally and sends no API request.
9// ui.render (AbovePrompt): draw the bar and a legend sized to the band.
10
11import { atom, read, update } from 'claude-code'
12import type { EngineInterface, Register } from 'claude-code'
13
14import type { Reading, Segment } from '../types'
15
16const reading = atom({ plugin: 'context-bar', key: 'reading' } as const, null)
17const isOn = atom({ plugin: 'context-bar', key: 'isOn' } as const, true)
18
19const STORE_KEY = 'isOn'
20const MIN_BAR = 10
21const MAX_BAR = 48
22
23export const register: Register = on => {
24  on('session.start', async ($, e, next) => {
25    const result = await next(e)
26    await $.command.register({
27      name: 'context-bar',
28      description: 'Toggle the stacked context-window bar above the prompt',
29    })
30    const stored = await $.store.get(STORE_KEY)
31    await update($, isOn, () => stored !== false)
32    await takeReading($)
33
34    return result
35  })
36
37  on('session.measure', async ($, e, next) => {
38    const result = await next(e)
39    await takeReading($)
40
41    return result
42  })
43
44  on('command.run', { command: 'context-bar' }, async $ => {
45    const now = !(await read($, isOn))
46    await update($, isOn, () => now)
47    await $.store.set(STORE_KEY, now)
48
49    return { text: now ? 'Context bar on.' : 'Context bar off.' }
50  })
51
52  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
53    const r = await read($, reading)
54    if (e.props.hasSurvey || r === null || !(await read($, isOn))) {
55      return next(e)
56    }
57
58    const { Box, Text } = $.ui.resolve(e)
59    const columns = e.props.bodyColumns
60    const headline = headlineFor(r)
61    const width = Math.max(MIN_BAR, Math.min(MAX_BAR, columns - headline.text.length - 3))
62    const cells = allocate(r.segments, width)
63    const legend = legendFor(r.segments, columns)
64
65    return (
66      <Box flexDirection="column" paddingX={1}>
67        <Box flexDirection="row">
68          {r.segments.map((s, i) => {
69            const n = cells[i] ?? 0
70            return n > 0 ? (
71              <Text color={s.kind === 'free' ? 'subtle' : s.color} dimColor={s.kind !== 'used'}>
72                {(s.kind === 'used' ? '█' : s.kind === 'buffer' ? '░' : '·').repeat(n)}
73              </Text>
74            ) : null
75          })}
76          <Text color={headline.color}> {headline.text}</Text>
77        </Box>
78        {legend.length > 0 && (
79          <Box flexDirection="row">
80            {legend.map(s => (
81              <Text>
82                <Text color={s.color}>■</Text>
83                <Text dimColor> {s.name} {short(s.tokens)}  </Text>
84              </Text>
85            ))}
86          </Box>
87        )}
88      </Box>
89    )
90  })
91}
92
93async function takeReading($: EngineInterface) {
94  try {
95    const { context } = await $.session.usage({ breakdown: 'summary' })
96    const b = context.breakdown
97    let next: Reading | null = null
98
99    if (b && b.rawMaxTokens > 0) {
100      const segments: Segment[] = b.categories
101        .filter(c => c.kind !== 'deferred' && c.tokens > 0)
102        .map(c => ({ name: c.name, color: c.color, kind: c.kind as Segment['kind'], tokens: c.tokens }))
103      next = {
104        segments,
105        used: b.totalTokens,
106        window: b.rawMaxTokens,
107        compactAt: b.isAutoCompactEnabled ? (b.autoCompactThreshold ?? null) : null,
108      }
109    } else if (context.window > 0 && context.tokens) {
110      next = {
111        segments: [
112          { name: 'Used', color: 'claude', kind: 'used', tokens: context.tokens },
113          { name: 'Free', color: 'subtle', kind: 'free', tokens: Math.max(0, context.window - context.tokens) },
114        ],
115        used: context.tokens,
116        window: context.window,
117        compactAt: null,
118      }
119    }
120
121    if (next) {
122      const value = next
123      await update($, reading, () => value)
124    }
125  } catch {
126    // Keep the last reading; the next measure tries again.
127  }
128}
129
130// Split `width` cells across segments by largest remainder, so the bar is
131// always exactly `width` wide; every used segment gets at least one cell.
132export function allocate(segments: Segment[], width: number): number[] {
133  const total = segments.reduce((n, s) => n + s.tokens, 0)
134  if (total <= 0) return segments.map(() => 0)
135
136  const exact = segments.map(s => (s.tokens / total) * width)
137  const cells = exact.map(Math.floor)
138  let left = width - cells.reduce((n, c) => n + c, 0)
139  const order = exact.map((x, i) => ({ i, rest: x - Math.floor(x) })).sort((a, b) => b.rest - a.rest)
140  for (const { i } of order) {
141    if (left <= 0) break
142    cells[i] = (cells[i] ?? 0) + 1
143    left -= 1
144  }
145
146  segments.forEach((s, i) => {
147    if (s.kind !== 'used' || (cells[i] ?? 0) > 0) return
148    const donor = cells.reduce((best, c, j) => (c > (cells[best] ?? 0) ? j : best), 0)
149    if ((cells[donor] ?? 0) > 1) {
150      cells[donor] = (cells[donor] ?? 0) - 1
151      cells[i] = 1
152    }
153  })
154
155  return cells
156}
157
158export function headlineFor(r: Reading): { text: string; color: string } {
159  const pct = Math.round((r.used / r.window) * 100)
160  const base = `${short(r.used)}/${short(r.window)} ${pct}%`
161  if (r.compactAt === null) {
162    return { text: base, color: pct >= 90 ? 'error' : pct >= 70 ? 'warning' : 'text' }
163  }
164
165  const room = r.compactAt - r.used
166  if (room <= 0) return { text: `${base} · compacting`, color: 'error' }
167  const share = room / r.compactAt
168  return {
169    text: `${base} · ${short(room)} to compact`,
170    color: share < 0.1 ? 'error' : share < 0.25 ? 'warning' : 'text',
171  }
172}
173
174// The biggest used categories that fit on one row, biggest first.
175export function legendFor(segments: Segment[], columns: number): Segment[] {
176  const used = segments.filter(s => s.kind === 'used').sort((a, b) => b.tokens - a.tokens)
177  const shown: Segment[] = []
178  let room = columns - 2
179  for (const s of used) {
180    const cost = s.name.length + short(s.tokens).length + 5
181    if (cost > room) break
182    shown.push(s)
183    room -= cost
184  }
185
186  return shown
187}
188
189export function short(n: number): string {
190  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1).replace(/\.0$/, '')}M`
191  if (n >= 1_000) return `${(n / 1_000).toFixed(n >= 100_000 ? 0 : 1).replace(/\.0$/, '')}k`
192  return String(n)
193}
194

1export type Segment = { name: string; color: string; kind: 'used' | 'free' | 'buffer'; tokens: number }
2
3export type Reading = {
4  segments: Segment[]
5  used: number
6  window: number
7  compactAt: number | null
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'context-bar': { reading: Reading | null; isOn: boolean }
13  }
14}
15