context-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Segment, Snapshot } from '../types'
5
6const snapshot = atom({ plugin: 'context-bar', key: 'snapshot' } as const, null)
7const isVisible = atom({ plugin: 'context-bar', key: 'isVisible' } as const, true)
8
9const STORE_KEY = 'isVisible'
10const FREE_COLOR = 'subtle'
11const MARKER_COLOR = 'warning'
12
13// 950 -> "950", 3_412 -> "3.4k", 12_300 -> "12k", 1_000_000 -> "1M"
14const formatTokens = (n: number) => {
15  const short = (x: number, unit: string) => `${x < 10 ? Number(x.toFixed(1)) : Math.round(x)}${unit}`
16  if (n >= 1_000_000) return short(n / 1_000_000, 'M')
17  if (n >= 1000) return short(n / 1000, 'k')
18  return `${n}`
19}
20
21// 0.27 -> "0.3%", 1.2 -> "1%", 18.6 -> "19%"
22const formatPercent = (p: number) => `${p < 1 ? Number(p.toFixed(1)) : Math.round(p)}%`
23
24// Splits `width` cells across the segments in proportion to their tokens,
25// giving every non-empty used segment at least one cell so it stays visible.
26export const allocate = (segments: Segment[], width: number): number[] => {
27  const total = segments.reduce((sum, s) => sum + s.tokens, 0)
28  if (total <= 0 || width <= 0) return segments.map(() => 0)
29
30  const exact = segments.map(s => (s.tokens / total) * width)
31  const cells = segments.map((s, i) => {
32    const floor = Math.floor(exact[i] ?? 0)
33    return s.kind === 'used' && s.tokens > 0 ? Math.max(1, floor) : floor
34  })
35  let left = width - cells.reduce((a, b) => a + b, 0)
36
37  const byRemainder = exact
38    .map((x, i) => ({ i, r: x - Math.floor(x) }))
39    .sort((a, b) => b.r - a.r)
40  for (let k = 0; left > 0 && byRemainder.length > 0; k++, left--) {
41    const { i } = byRemainder[k % byRemainder.length]!
42    cells[i] = (cells[i] ?? 0) + 1
43  }
44  // Minimums can overshoot: take the excess from the widest non-used segments.
45  while (left < 0) {
46    const order = cells
47      .map((c, i) => ({ i, c, isUsed: segments[i]?.kind === 'used' }))
48      .sort((a, b) => Number(a.isUsed) - Number(b.isUsed) || b.c - a.c)
49    const victim = order.find(o => o.c > (o.isUsed ? 1 : 0))
50    if (!victim) break
51    cells[victim.i] = victim.c - 1
52    left += 1
53  }
54
55  return cells
56}
57
58// Used categories first, then free space, then the compaction buffer at the end.
59const ORDER: Record<Segment['kind'], number> = { used: 0, free: 1, buffer: 2 }
60
61// /context draws some categories in greys that blur into each other and into
62// free space; those (and any repeated colour) get a colourful theme key instead.
63const GREYS = new Set(['inactive', 'subtle', 'promptBorder', 'text', FREE_COLOR])
64const SPARE = ['suggestion', 'bashBorder', 'ide', 'planMode', 'merged', 'autoAccept', 'remember', 'claude']
65
66export const distinctColors = (segments: Segment[]): Segment[] => {
67  const taken = new Set(
68    segments.filter(s => s.kind === 'used' && !GREYS.has(s.color)).map(s => s.color),
69  )
70  const spare = SPARE.filter(c => !taken.has(c))
71  const seen = new Set<string>()
72
73  return segments.map(s => {
74    if (s.kind !== 'used') return s
75    const isClash = GREYS.has(s.color) || seen.has(s.color)
76    const color = isClash ? (spare.shift() ?? s.color) : s.color
77    seen.add(color)
78    return color === s.color ? s : { ...s, color }
79  })
80}
81
82const refresh = async ($: EngineInterface) => {
83  const usage = await $.session.usage({ breakdown: 'summary' })
84  const breakdown = usage.context.breakdown
85  if (!breakdown) return
86
87  const segments: Segment[] = distinctColors(breakdown.categories.flatMap(c =>
88    c.kind === 'deferred' || c.tokens <= 0
89      ? []
90      : [{ name: c.name, color: c.color, tokens: c.tokens, kind: c.kind }],
91  ).sort((a, b) => ORDER[a.kind] - ORDER[b.kind]))
92  const next: Snapshot = {
93    segments,
94    totalTokens: breakdown.totalTokens,
95    maxTokens: breakdown.rawMaxTokens,
96    percentage: breakdown.percentage,
97    compactAt: breakdown.isAutoCompactEnabled ? breakdown.autoCompactThreshold : undefined,
98  }
99  await update($, snapshot, () => next)
100}
101
102export const register: Register = on => {
103  on('session.start', async ($, e, next) => {
104    await $.command.register({
105      name: 'context-bar',
106      description: 'Toggle the context window bar above the prompt',
107    })
108    const stored = await $.store.get(STORE_KEY)
109    if (typeof stored === 'boolean') await update($, isVisible, () => stored)
110    void refresh($).catch(() => {})
111
112    return next(e)
113  })
114
115  on('session.measure', async ($, e, next) => {
116    if (e.changed.includes('context')) await refresh($).catch(() => {})
117
118    return next(e)
119  })
120
121  on('command.run', { command: 'context-bar' }, async $ => {
122    const shown = !(await read($, isVisible))
123    await update($, isVisible, () => shown)
124    await $.store.set(STORE_KEY, shown)
125    if (shown) await refresh($).catch(() => {})
126
127    return { text: shown ? 'Context bar shown.' : 'Context bar hidden.' }
128  })
129
130  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
131    if (e.props.hasSurvey || !(await read($, isVisible))) return next(e)
132    const snap = await read($, snapshot)
133    if (!snap || snap.segments.length === 0) return next(e)
134
135    const { Box, Text } = $.ui.resolve(e)
136    // Border (2) and horizontal padding (2) around the bar.
137    const cells = allocate(snap.segments, Math.max(10, e.props.bodyColumns - 4))
138    const share = (tokens: number) => formatPercent((tokens / snap.maxTokens) * 100)
139    const legend = snap.segments.filter(s => s.kind !== 'buffer')
140
141    return (
142      <Box flexDirection="column" borderStyle="round" borderColor="subtle" paddingX={1}>
143        <Box flexDirection="row" justifyContent="space-between">
144          <Text>
145            <Text color="claude">◆ </Text>
146            <Text bold>context</Text>
147          </Text>
148          <Text wrap="truncate-start">
149            <Text bold>{formatTokens(snap.totalTokens)}</Text>
150            <Text dimColor> of </Text>
151            <Text>{formatTokens(snap.maxTokens)}</Text>
152            {snap.compactAt !== undefined && (
153              <Text dimColor> · compacts at {formatTokens(snap.compactAt)}</Text>
154            )}{' '}
155            <Text bold color="inverseText" backgroundColor={snap.percentage >= 80 ? 'warning' : 'success'}>
156              {` ${snap.percentage}% `}
157            </Text>
158          </Text>
159        </Box>
160        <Text wrap="truncate">
161          {snap.segments.map((s, i) => {
162            const n = cells[i] ?? 0
163            if (n <= 0) return null
164            if (s.kind === 'buffer') {
165              return (
166                <Text>
167                  <Text color={MARKER_COLOR}>▏</Text>
168                  {' '.repeat(n - 1)}
169                </Text>
170              )
171            }
172            return (
173              <Text color={s.kind === 'free' ? FREE_COLOR : s.color} dimColor={s.kind === 'free'}>
174                {'█'.repeat(n)}
175              </Text>
176            )
177          })}
178        </Text>
179        <Box flexDirection="row" flexWrap="wrap" columnGap={3}>
180          {legend.map(s =>
181            s.kind === 'free' ? (
182              <Text>
183                <Text color={FREE_COLOR} dimColor>▌</Text>
184                <Text dimColor> free </Text>
185                <Text>{formatTokens(s.tokens)}</Text>
186              </Text>
187            ) : (
188              <Text>
189                <Text color={s.color}>▌</Text>
190                <Text> {s.name.toLowerCase()} {formatTokens(s.tokens)} </Text>
191                <Text dimColor>{share(s.tokens)}</Text>
192              </Text>
193            ),
194          )}
195        </Box>
196      </Box>
197    )
198  })
199}
200

1export type Segment = {
2  name: string
3  color: string
4  tokens: number
5  kind: 'used' | 'free' | 'buffer'
6}
7
8export type Snapshot = {
9  segments: Segment[]
10  totalTokens: number
11  maxTokens: number
12  percentage: number
13  compactAt?: number
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    'context-bar': { snapshot: Snapshot | null; isVisible: boolean }
19  }
20}
21