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 isVisible = atom({ plugin: 'context-bar', key: 'isVisible' } as const, true)
7const snapshot = atom({ plugin: 'context-bar', key: 'snapshot' } as const, null)
8
9const GLYPH: Record<Segment['kind'], string> = { used: '█', buffer: '▒', free: '░' }
10
11export const formatTokens = (n: number): string => {
12  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
13  if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`
14  return String(n)
15}
16
17// Largest-remainder rounding, so the cells always sum to exactly `width`.
18export const allocateCells = (tokens: number[], width: number): number[] => {
19  const total = tokens.reduce((sum, t) => sum + t, 0)
20  if (total <= 0 || width <= 0) return tokens.map(() => 0)
21
22  const exact = tokens.map(t => (t / total) * width)
23  const cells = exact.map(Math.floor)
24  let left = width - cells.reduce((sum, c) => sum + c, 0)
25  const byRemainder = exact
26    .map((x, i) => ({ i, rem: x - Math.floor(x) }))
27    .sort((a, b) => b.rem - a.rem)
28
29  for (const { i } of byRemainder) {
30    if (left <= 0) break
31    cells[i] = (cells[i] ?? 0) + 1
32    left -= 1
33  }
34
35  return cells
36}
37
38const refresh = async ($: EngineInterface) => {
39  const usage = await $.session.usage({ breakdown: 'summary' })
40  const breakdown = usage.context.breakdown
41  if (!breakdown) return
42
43  const segments: Segment[] = breakdown.categories
44    .filter(row => row.kind !== 'deferred' && row.tokens > 0)
45    .map(row => ({
46      name: row.name,
47      tokens: row.tokens,
48      color: row.color,
49      kind: row.kind as Segment['kind'],
50    }))
51
52  const next: Snapshot = {
53    segments,
54    usedTokens: breakdown.totalTokens,
55    maxTokens: breakdown.rawMaxTokens,
56    percent: breakdown.percentage,
57  }
58  await update($, snapshot, () => next)
59}
60
61let isRefreshing = false
62
63const refreshQuietly = async ($: EngineInterface) => {
64  if (isRefreshing) return
65  isRefreshing = true
66  try {
67    await refresh($)
68  } catch {
69    // A failed estimate leaves the last bar up; the next measure retries.
70  } finally {
71    isRefreshing = false
72  }
73}
74
75export const register: Register = on => {
76  on('session.start', async ($, e, next) => {
77    await $.command.register({
78      name: 'context-bar',
79      description: 'Toggle the context window bar above the prompt',
80    })
81    if (await read($, isVisible)) void refreshQuietly($)
82
83    return next(e)
84  })
85
86  on('session.measure', async ($, e, next) => {
87    if (e.changed.includes('context') && (await read($, isVisible))) {
88      await refreshQuietly($)
89    }
90
91    return next(e)
92  })
93
94  on('command.run', { command: 'context-bar' }, async $ => {
95    try {
96      const shown = await update($, isVisible, value => !value)
97      if (shown) await refreshQuietly($)
98
99      return { text: shown ? 'Context bar on.' : 'Context bar off.' }
100    } catch {
101      return { text: 'Context bar could not switch just now. Try again.' }
102    }
103  })
104
105  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
106    if (e.props.hasSurvey || !(await read($, isVisible))) return next(e)
107
108    const data = await read($, snapshot)
109    const { Box, Text } = $.ui.resolve(e)
110
111    if (!data || data.segments.length === 0) {
112      return (
113        <Box>
114          <Text dimColor>Context bar: waiting for the first response…</Text>
115        </Box>
116      )
117    }
118
119    const label = ` ${data.percent}% · ${formatTokens(data.usedTokens)}/${formatTokens(data.maxTokens)}`
120    const barWidth = Math.max(10, e.props.bodyColumns - label.length)
121    const cells = allocateCells(
122      data.segments.map(s => s.tokens),
123      barWidth,
124    )
125    const used = data.segments.filter(s => s.kind === 'used')
126
127    return (
128      <Box flexDirection="column">
129        <Text wrap="truncate-end">
130          {data.segments.map((s, i) => {
131            const n = cells[i] ?? 0
132            return n > 0 ? (
133              <Text color={s.color}>
134                {GLYPH[s.kind].repeat(n)}
135              </Text>
136            ) : null
137          })}
138          <Text dimColor>{label}</Text>
139        </Text>
140        <Text wrap="truncate-end">
141          {used.map((s, i) => (
142            <Text>
143              <Text color={s.color}>{GLYPH.used}</Text>
144              <Text dimColor>
145                {' '}
146                {s.name} {formatTokens(s.tokens)}
147                {i < used.length - 1 ? '   ' : ''}
148              </Text>
149            </Text>
150          ))}
151        </Text>
152      </Box>
153    )
154  })
155}
156

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