context-usage

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextUsage } from 'claude-code'
3
4import type { Part, Usage } from '../types'
5
6const usage = atom({ plugin: 'context-usage', key: 'usage' } as const, null)
7const isExpanded = atom({ plugin: 'context-usage', key: 'isExpanded' } as const, false)
8
9const COMMAND = 'context-bar'
10const STORE_KEY = 'isExpanded'
11const SMALL_BAR = 12
12const LARGE_BAR = 20
13
14function formatTokens(n: number): string {
15  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
16  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
17  return `${n}`
18}
19
20function colorFor(percent: number): 'success' | 'warning' | 'error' {
21  if (percent >= 85) return 'error'
22  if (percent >= 60) return 'warning'
23  return 'success'
24}
25
26// With a breakdown, every figure comes from it, so the bar and the number
27// agree with /context. Without one, the status line's figures stand alone.
28function fromContext(context: SessionContextUsage): Usage {
29  const b = context.breakdown
30  if (b) {
31    const parts = b.categories
32      .filter(c => c.kind === 'used' && c.tokens > 0)
33      .map(({ name, color, tokens }) => ({ name, color, tokens }))
34    return { tokens: b.totalTokens, window: b.rawMaxTokens, percent: b.percentage, parts }
35  }
36  const out: Usage = { window: context.window }
37  if (context.tokens !== undefined) out.tokens = context.tokens
38  if (context.percent !== undefined) out.percent = context.percent
39  return out
40}
41
42// A local estimate by category; it sends no token-count requests.
43async function measure($: EngineInterface): Promise<Usage | null> {
44  try {
45    const { context } = await $.session.usage({ breakdown: 'summary' })
46    return fromContext(context)
47  } catch {
48    return null
49  }
50}
51
52async function toggle($: EngineInterface) {
53  const value = await update($, isExpanded, v => !v)
54  await $.store.set(STORE_KEY, value)
55}
56
57// Split `filled` cells across the parts by token share, the sum kept exact.
58function cellsFor(parts: Part[], filled: number): number[] {
59  const total = parts.reduce((sum, p) => sum + p.tokens, 0)
60  if (total === 0) return parts.map(() => 0)
61  const exact = parts.map(p => (p.tokens / total) * filled)
62  const cells = exact.map(Math.floor)
63  let left = filled - cells.reduce((sum, n) => sum + n, 0)
64  const byRemainder = exact
65    .map((x, i) => ({ i, r: x - Math.floor(x) }))
66    .sort((a, b) => b.r - a.r)
67  for (const { i } of byRemainder) {
68    if (left <= 0) break
69    cells[i] = (cells[i] ?? 0) + 1
70    left -= 1
71  }
72  return cells
73}
74
75export const register: Register = on => {
76  on('session.start', async ($, e, next) => {
77    await $.command.register({
78      name: COMMAND,
79      description: 'Expand or minimize the context usage bar',
80      immediate: true,
81    })
82    const stored = await $.store.get(STORE_KEY)
83    await update($, isExpanded, () => stored === true)
84    const current = await measure($)
85    if (current) await update($, usage, () => current)
86
87    return next(e)
88  })
89
90  // Pushed by the engine after each main-thread turn.
91  on('session.measure', async ($, e, next) => {
92    if (e.changed.includes('context')) {
93      const current = (await measure($)) ?? fromContext(e.context)
94      await update($, usage, () => current)
95    }
96
97    return next(e)
98  })
99
100  // The fill drops after a compaction; read it again once it is done.
101  on('session.compact', async ($, e, next) => {
102    const result = await next(e)
103    // Never fail the compaction itself; the next measurement catches up.
104    try {
105      const current = await measure($)
106      if (current) await update($, usage, () => current)
107    } catch {}
108
109    return result
110  })
111
112  // Answers with no text, so the toggle adds nothing to the transcript.
113  on('command.run', { command: COMMAND }, async $ => {
114    await toggle($)
115
116    return {}
117  }).catch(() => ({}))
118
119  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
120    // What other plugins drew in the band: drawn above this bar, never replaced.
121    const below = await next(e)
122    if (e.props.hasSurvey) {
123      return below
124    }
125
126    const { Box, Button, Text } = $.ui.resolve(e)
127    const current = await read($, usage)
128    const expanded = await read($, isExpanded)
129
130    const size = current?.window
131    const percent = current?.percent
132    const tokens = current?.tokens
133    const parts = current?.parts ?? []
134
135    const figure =
136      percent === undefined
137        ? ' --%'
138        : expanded && tokens !== undefined && size
139          ? ` ${percent}% · ${formatTokens(tokens)} / ${formatTokens(size)}`
140          : ` ${percent}%`
141
142    // The toggle's label and the space after it take six cells.
143    const room = e.props.bodyColumns - 6 - figure.length
144    const barWidth = Math.max(0, Math.min(expanded ? LARGE_BAR : SMALL_BAR, room))
145    const fill = Math.min(100, Math.max(0, percent ?? 0))
146    const filled = Math.round((fill / 100) * barWidth)
147    const cells = cellsFor(parts, filled)
148    const empty = Math.max(0, barWidth - filled)
149
150    const bar = (
151      <Box flexDirection="column">
152        <Box>
153          <Button
154            key="toggle"
155            plain
156            dimColor
157            label={expanded ? '▾ ctx' : '▸ ctx'}
158            onPress={() => toggle($)}
159          />
160          <Text> </Text>
161          {parts.length > 0
162            ? parts.map(
163                (part, i) =>
164                  (cells[i] ?? 0) > 0 && (
165                    <Text color={part.color}>{'█'.repeat(cells[i] ?? 0)}</Text>
166                  ),
167              )
168            : filled > 0 && <Text color={colorFor(fill)}>{'█'.repeat(filled)}</Text>}
169          {empty > 0 && <Text dimColor>{'░'.repeat(empty)}</Text>}
170          <Text color={percent === undefined ? 'inactive' : colorFor(fill)} wrap="truncate">
171            {figure}
172          </Text>
173        </Box>
174        {expanded && parts.length > 0 && (
175          <Box flexDirection="row" flexWrap="wrap" columnGap={2}>
176            {parts.map(part => (
177              <Box>
178                <Text color={part.color}>■ </Text>
179                <Text dimColor>
180                  {part.name} {formatTokens(part.tokens)}
181                </Text>
182              </Box>
183            ))}
184          </Box>
185        )}
186      </Box>
187    )
188
189    return below.type === 'engine' ? bar : (
190      <Box flexDirection="column">
191        {below}
192        {bar}
193      </Box>
194    )
195  })
196}
197

1export type Part = { name: string; color: string; tokens: number }
2
3export type Usage = {
4  tokens?: number
5  window: number
6  percent?: number
7  parts?: Part[]
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'context-usage': { usage: Usage | null; isExpanded: boolean }
13  }
14}
15