context-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextBreakdown } from 'claude-code'
3
4import type { Fill, Segment } from '../types'
5
6const fill = atom({ plugin: 'context-bar', key: 'fill' } as const, null)
7const isHidden = atom({ plugin: 'context-bar', key: 'isHidden' } as const, false)
8
9// Deferred tool schemas sit outside the window, so /context leaves them off its grid too.
10export function toFill(breakdown: SessionContextBreakdown): Fill {
11  const segments: Segment[] = breakdown.categories
12    .filter(row => row.kind !== 'deferred' && row.tokens > 0)
13    .map(row => ({ name: row.name, tokens: row.tokens, color: row.color, isFree: row.kind === 'free' }))
14
15  return {
16    segments,
17    totalTokens: breakdown.totalTokens,
18    maxTokens: breakdown.rawMaxTokens,
19    percentage: breakdown.percentage,
20  }
21}
22
23// Splits `width` cells across the segments by tokens (largest remainder), so the widths always sum to `width`.
24export function cellWidths(segments: Segment[], width: number): number[] {
25  const total = segments.reduce((sum, s) => sum + s.tokens, 0)
26  if (total === 0 || width <= 0) return segments.map(() => 0)
27
28  const exact = segments.map(s => (s.tokens / total) * width)
29  const widths = exact.map(Math.floor)
30  let left = width - widths.reduce((sum, w) => sum + w, 0)
31  const byRemainder = exact.map((x, i) => ({ i, rest: x - Math.floor(x) })).sort((a, b) => b.rest - a.rest)
32  for (const { i } of byRemainder) {
33    if (left === 0) break
34    widths[i] = (widths[i] ?? 0) + 1
35    left -= 1
36  }
37
38  return widths
39}
40
41export function formatTokens(tokens: number): string {
42  if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`
43  if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1)}k`
44  return String(tokens)
45}
46
47async function refresh($: EngineInterface) {
48  // `summary` estimates locally and sends no token-count requests.
49  const { context } = await $.session.usage({ breakdown: 'summary' })
50  if (context.breakdown) await update($, fill, () => toFill(context.breakdown!))
51}
52
53export const register: Register = on => {
54  on('session.start', async ($, e, next) => {
55    await $.command.register({
56      name: 'context-bar',
57      description: 'Show or hide the context window bar above the prompt',
58    })
59    const ran = await next(e)
60    await refresh($)
61
62    return ran
63  })
64
65  on('session.measure', async ($, e, next) => {
66    const ran = await next(e)
67    if (e.changed.includes('context')) await refresh($)
68
69    return ran
70  })
71
72  on('command.run', { command: 'context-bar' }, async $ => {
73    const wasHidden = await read($, isHidden)
74    await update($, isHidden, () => !wasHidden)
75    if (wasHidden) await refresh($)
76
77    return { text: wasHidden ? 'Context bar shown.' : 'Context bar hidden.' }
78  })
79
80  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
81    const current = await read($, fill)
82    if (e.props.hasSurvey || current === null || (await read($, isHidden))) return next(e)
83
84    // The band holds one tree. When a mod beneath draws there (Blast Radius's Proceed/Cancel
85    // in a narrow terminal, say), it gets the band and the bar steps aside.
86    const below = await next(e)
87    if (below.type !== 'engine') return below
88
89    const { Box, Text } = $.ui.resolve(e)
90    const widths = cellWidths(current.segments, e.props.bodyColumns)
91
92    return (
93      <Box flexDirection="column">
94        <Text wrap="truncate">
95          {current.segments.map((segment, i) => (
96            <Text key={segment.name} color={segment.color} dimColor={segment.isFree}>
97              {(segment.isFree ? '░' : '█').repeat(widths[i] ?? 0)}
98            </Text>
99          ))}
100        </Text>
101        <Text wrap="truncate">
102          <Text dimColor>
103            {formatTokens(current.totalTokens)}/{formatTokens(current.maxTokens)} ({current.percentage}%){'  '}
104          </Text>
105          {current.segments
106            .filter(segment => !segment.isFree)
107            .map(segment => (
108              <Text key={segment.name}>
109                <Text color={segment.color}>■</Text>
110                <Text dimColor> {segment.name} {formatTokens(segment.tokens)}  </Text>
111              </Text>
112            ))}
113        </Text>
114      </Box>
115    )
116  })
117}
118

1export type Segment = { name: string; tokens: number; color: string; isFree: boolean }
2export type Fill = { segments: Segment[]; totalTokens: number; maxTokens: number; percentage: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'context-bar': { fill: Fill | null; isHidden: boolean }
7  }
8}
9