context-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextBreakdown } from 'claude-code'
3
4import type { Item, Row, Snapshot } from '../types'
5
6const snapshot = atom({ plugin: 'context-band', key: 'snapshot' } as const, null)
7const isHidden = atom({ plugin: 'context-band', key: 'isHidden' } as const, false)
8const isDetailed = atom({ plugin: 'context-band', key: 'isDetailed' } as const, false)
9
10const GLYPH: Record<Row['kind'], string> = { used: '█', buffer: '▒', free: '░', deferred: '' }
11
12// 1234 -> "1.2k", 156000 -> "156k", 1000000 -> "1.0M"
13function fmt(n: number): string {
14  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
15  if (n >= 100_000) return `${Math.round(n / 1000)}k`
16  if (n >= 1000) return `${(n / 1000).toFixed(1)}k`
17  return String(n)
18}
19
20function pct(tokens: number, of: number): string {
21  if (of <= 0) return '0%'
22  const p = (tokens / of) * 100
23  return p > 0 && p < 1 ? '<1%' : `${Math.round(p)}%`
24}
25
26function top(items: Item[], n: number): Item[] {
27  return [...items].sort((a, b) => b.tokens - a.tokens).slice(0, n)
28}
29
30function basename(path: string): string {
31  const parts = path.split('/')
32  return parts.slice(-2).join('/')
33}
34
35function toSnapshot(b: SessionContextBreakdown): Snapshot {
36  const servers = new Map<string, number>()
37  for (const t of b.mcpTools) {
38    if (t.isLoaded) servers.set(t.serverName, (servers.get(t.serverName) ?? 0) + t.tokens)
39  }
40  return {
41    totalTokens: b.totalTokens,
42    maxTokens: b.rawMaxTokens,
43    percentage: b.percentage,
44    model: b.model,
45    autoCompactAt: b.isAutoCompactEnabled ? (b.autoCompactThreshold ?? null) : null,
46    rows: b.categories.map(c => ({ name: c.name, tokens: c.tokens, color: c.color, kind: c.kind })),
47    mcpServers: [...servers].map(([label, tokens]) => ({ label, tokens })),
48    memoryFiles: b.memoryFiles.map(f => ({ label: `${basename(f.path)} (${f.type})`, tokens: f.tokens })),
49    skills: b.skills
50      ? { listed: b.skills.includedSkills, total: b.skills.totalSkills, tokens: b.skills.tokens }
51      : null,
52    agents: { count: b.agents.length, tokens: b.agents.reduce((s, a) => s + a.tokens, 0) },
53  }
54}
55
56// Split `width` cells across the rows in proportion to their tokens
57// (largest remainder), so the bar always fills exactly `width`.
58function cells(rows: Row[], width: number): number[] {
59  const sum = rows.reduce((s, r) => s + r.tokens, 0)
60  if (sum <= 0) return rows.map(() => 0)
61  const exact = rows.map(r => (r.tokens / sum) * width)
62  // A row that holds any tokens gets at least one cell, so small ones stay visible.
63  const out = exact.map((x, i) => Math.max(Math.floor(x), (rows[i]?.tokens ?? 0) > 0 ? 1 : 0))
64  let left = width - out.reduce((s, c) => s + c, 0)
65  const order = exact
66    .map((x, i) => ({ i, rest: x - Math.floor(x) }))
67    .sort((a, b) => b.rest - a.rest)
68  for (const { i } of order) {
69    if (left <= 0) break
70    out[i] = (out[i] ?? 0) + 1
71    left -= 1
72  }
73  // Too many minimum cells: take back from the widest rows.
74  while (left < 0) {
75    const widest = out.indexOf(Math.max(...out))
76    out[widest] = (out[widest] ?? 1) - 1
77    left += 1
78  }
79  return out
80}
81
82// One refresh at a time; calls that arrive meanwhile fold into one more pass.
83let isBusy = false
84let isAgain = false
85
86async function refresh($: EngineInterface): Promise<void> {
87  if (isBusy) {
88    isAgain = true
89    return
90  }
91  isBusy = true
92  try {
93    do {
94      isAgain = false
95      const usage = await $.session.usage({ breakdown: 'summary' })
96      const b = usage.context.breakdown
97      if (b) await update($, snapshot, () => toSnapshot(b))
98    } while (isAgain)
99  } finally {
100    isBusy = false
101  }
102}
103
104function background($: EngineInterface): void {
105  refresh($).catch(err => $.ui.log(`context-band: ${String(err)}`, { to: 'debug' }))
106}
107
108export const register: Register = on => {
109  on('session.start', async ($, e, next) => {
110    await $.command.register({
111      name: 'context-band',
112      description: 'Show the context usage band again and refresh it',
113    })
114    const started = await next(e)
115    background($)
116    return started
117  })
118
119  on('command.run', { command: 'context-band' }, async $ => {
120    await update($, isHidden, () => false)
121    await refresh($)
122    return { text: 'Context band shown and refreshed.' }
123  })
124
125  // Each main-thread model response moves the Messages row: refresh after it.
126  on('turn.step', async function* ($, e, next) {
127    const r = yield* next(e)
128    if (e.agentId === undefined) background($)
129    return r
130  })
131
132  // Tool results land in the window as Claude works.
133  on('tool.call', async ($, e, next) => {
134    const r = await next(e)
135    background($)
136    return r
137  }).catch(($, e, next) => next(e))
138
139  on('session.measure', ($, e, next) => {
140    if (e.changed.includes('context')) background($)
141    return next(e)
142  })
143
144  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
145    const snap = await read($, snapshot)
146    if (e.props.hasSurvey || snap === null || (await read($, isHidden))) {
147      return next(e)
148    }
149    const detailed = await read($, isDetailed)
150    const { Box, Button, Text } = $.ui.resolve(e)
151
152    const inBar = snap.rows.filter(r => r.kind !== 'deferred')
153    const scale = Math.max(snap.maxTokens, snap.totalTokens)
154    // "Context " + bar + " 100% 999k/1.0M " + buttons
155    const width = Math.max(10, Math.min(100, e.props.bodyColumns - 44))
156    const widths = cells(inBar, width)
157    const level = snap.percentage >= 85 ? 'error' : snap.percentage >= 70 ? 'warning' : 'success'
158
159    const legend = snap.rows.filter(r => r.tokens > 0 || r.kind === 'free')
160    const mcp = top(snap.mcpServers, 4)
161    const memory = top(snap.memoryFiles, 3)
162
163    const mine = (
164      <Box flexDirection="column">
165        <Box>
166          <Text bold>Context </Text>
167          {inBar.map((r, i) => {
168            const n = widths[i] ?? 0
169            return n > 0 ? <Text color={r.color}>{GLYPH[r.kind].repeat(n)}</Text> : null
170          })}
171          <Text color={level} bold> {snap.percentage}%</Text>
172          <Text dimColor> {fmt(snap.totalTokens)}/{fmt(snap.maxTokens)} </Text>
173          <Button
174            key="details"
175            label={detailed ? 'Less' : 'Details'}
176            onPress={() => update($, isDetailed, v => !v)}
177          />
178          <Text> </Text>
179          <Button key="hide" label="Hide" onPress={() => update($, isHidden, () => true)} />
180        </Box>
181        <Box flexWrap="wrap" columnGap={2}>
182          {legend.map(r => (
183            <Box>
184              <Text color={r.color}>{r.kind === 'deferred' ? '◌' : '■'} </Text>
185              <Text dimColor={r.kind !== 'used'}>
186                {r.name} {fmt(r.tokens)}
187                {r.kind === 'deferred' ? ' (not loaded)' : ` ${pct(r.tokens, scale)}`}
188              </Text>
189            </Box>
190          ))}
191        </Box>
192        {detailed && (
193          <Box flexDirection="column">
194            <Text dimColor wrap="truncate-end">
195              {snap.model}
196              {snap.autoCompactAt !== null
197                ? ` · autocompact at ${fmt(snap.autoCompactAt)} (${fmt(Math.max(0, snap.autoCompactAt - snap.totalTokens))} left)`
198                : ' · autocompact off'}
199            </Text>
200            {mcp.length > 0 && (
201              <Text wrap="truncate-end">
202                <Text bold>MCP </Text>
203                {mcp.map(s => `${s.label} ${fmt(s.tokens)}`).join(' · ')}
204              </Text>
205            )}
206            {memory.length > 0 && (
207              <Text wrap="truncate-end">
208                <Text bold>Memory </Text>
209                {memory.map(m => `${m.label} ${fmt(m.tokens)}`).join(' · ')}
210              </Text>
211            )}
212            <Text wrap="truncate-end">
213              {snap.skills !== null && (
214                <Text>
215                  <Text bold>Skills </Text>
216                  {snap.skills.listed}/{snap.skills.total} listed {fmt(snap.skills.tokens)}
217                  {'  '}
218                </Text>
219              )}
220              <Text bold>Agents </Text>
221              {snap.agents.count} {fmt(snap.agents.tokens)}
222            </Text>
223          </Box>
224        )}
225      </Box>
226    )
227
228    // The band has one instance: keep whatever another mod draws there below ours.
229    const below = await next(e)
230    return (
231      <Box flexDirection="column">
232        {mine}
233        {below}
234      </Box>
235    )
236  })
237}
238

1export type Row = {
2  name: string
3  tokens: number
4  color: string
5  kind: 'used' | 'free' | 'buffer' | 'deferred'
6}
7
8export type Item = { label: string; tokens: number }
9
10export type Snapshot = {
11  totalTokens: number
12  maxTokens: number
13  percentage: number
14  model: string
15  autoCompactAt: number | null
16  rows: Row[]
17  mcpServers: Item[]
18  memoryFiles: Item[]
19  skills: { listed: number; total: number; tokens: number } | null
20  agents: { count: number; tokens: number }
21}
22
23declare module 'claude-code' {
24  interface PluginState {
25    'context-band': { snapshot: Snapshot | null; isHidden: boolean; isDetailed: boolean }
26  }
27}
28