context-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Row, Snapshot } from '../types'
5
6const snapshot = atom({ plugin: 'context-bar', key: 'snapshot' } as const, null)
7
8// Legend order and colors, matched against /context's category names.
9const CATEGORIES = [
10  { label: 'system prompt', match: 'system prompt', color: '#7095d1' },
11  { label: 'tools', match: 'system tools', color: '#5cc8c8' },
12  { label: 'mcp tools', match: 'mcp tools', color: '#a78bfa' },
13  { label: 'agents', match: 'agents', color: '#8fd17f' },
14  { label: 'memory files', match: 'memory', color: '#e5c07b' },
15  { label: 'skills', match: 'skills', color: '#f0a6c8' },
16  { label: 'messages', match: 'messages', color: '#e2714e' },
17]
18const FREE = '#3a3a4c'
19const MARKER = '#f5b942'
20const ACCENT = '#d97757'
21const BADGE = '#8fd17f'
22
23export const fmt = (n: number): string => {
24  const trim = (x: number) => x.toFixed(1).replace(/\.0$/, '')
25  if (n >= 1e6) return `${trim(n / 1e6)}M`
26  if (n >= 1e4) return `${Math.round(n / 1000)}k`
27  if (n >= 1000) return `${trim(n / 1000)}k`
28  return String(n)
29}
30
31export const pct = (n: number, of: number): number =>
32  n > 0 ? Math.max(1, Math.round((n / of) * 100)) : 0
33
34async function refresh($: EngineInterface) {
35  try {
36    const { context } = await $.session.usage({ breakdown: 'summary' })
37    const b = context.breakdown
38    if (!b) return
39    const used = b.categories.filter(c => c.kind === 'used')
40    const rows: Row[] = CATEGORIES.map(({ label, match, color }) => ({
41      label,
42      color,
43      tokens: used
44        .filter(c => c.name.toLowerCase().includes(match))
45        .reduce((sum, c) => sum + c.tokens, 0),
46    }))
47    const total = b.totalTokens
48    const free =
49      b.categories.find(c => c.kind === 'free')?.tokens ?? Math.max(0, context.window - total)
50    const snap: Snapshot = {
51      used: total,
52      window: context.window,
53      compactAt: b.isAutoCompactEnabled ? (b.autoCompactThreshold ?? null) : null,
54      rows,
55      free,
56    }
57    await update($, snapshot, () => snap)
58  } catch {
59    // ponytail: usage can fail before a session is bound; the next event retries
60  }
61}
62
63// One bar cell per slot: each category gets at least one cell when non-empty.
64export const barCells = (s: Snapshot, width: number): string[] => {
65  const cells: string[] = []
66  for (const r of s.rows) {
67    const n = r.tokens > 0 ? Math.max(1, Math.round((r.tokens / s.window) * width)) : 0
68    for (let i = 0; i < n && cells.length < width; i++) cells.push(r.color)
69  }
70  while (cells.length < width) cells.push(FREE)
71  if (s.compactAt !== null) {
72    cells[Math.min(width - 1, Math.floor((s.compactAt / s.window) * width))] = MARKER
73  }
74  return cells
75}
76
77export const register: Register = on => {
78  on('session.start', async ($, e, next) => {
79    const r = await next(e)
80    await refresh($)
81    return r
82  })
83
84  on('turn.complete', async ($, e, next) => {
85    const r = await next(e)
86    if (e.agentId === undefined) await refresh($)
87    return r
88  })
89
90  on('tool.call', async ($, e, next) => {
91    const r = await next(e)
92    void refresh($)
93    return r
94  })
95
96  on('session.compact', async ($, e, next) => {
97    const r = await next(e)
98    await refresh($)
99    return r
100  })
101
102  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
103    const s = await read($, snapshot)
104    if (e.props.hasSurvey || s === null) return next(e)
105
106    const { Box, Text } = $.ui.resolve(e)
107    const width = Math.min(e.props.bodyColumns, 56)
108    const inner = width - 4
109    const half = Math.floor(inner / 2)
110
111    // Merge equal neighbours into runs so the bar is a handful of Texts.
112    const runs: { color: string; n: number }[] = []
113    for (const c of barCells(s, inner)) {
114      const last = runs.at(-1)
115      if (last && last.color === c) last.n++
116      else runs.push({ color: c, n: 1 })
117    }
118
119    const legend = [
120      ...s.rows,
121      { label: 'free', tokens: s.free, color: FREE },
122    ]
123    const pairs = []
124    for (let i = 0; i < legend.length; i += 2) pairs.push(legend.slice(i, i + 2))
125
126    return (
127      <Box flexDirection="column" borderStyle="round" borderColor="#3b3b4f" paddingX={1} width={width}>
128        <Box justifyContent="space-between">
129          <Text wrap="truncate">
130            <Text color={ACCENT}>◆ </Text>
131            <Text bold>context window {fmt(s.used)}</Text>
132            <Text dimColor>
133              {' '}of {fmt(s.window)}
134              {s.compactAt !== null ? ` · compacts at ${fmt(s.compactAt)}` : ''}
135            </Text>
136          </Text>
137          <Text backgroundColor={BADGE} color="#1a1b26" bold>
138            {` ${pct(s.used, s.window)}% `}
139          </Text>
140        </Box>
141        <Text>
142          {runs.map((r, i) => (
143            <Text key={`b${i}`} color={r.color} backgroundColor={r.color === MARKER ? FREE : undefined}>
144              {(r.color === MARKER ? '▕' : '█').repeat(r.n)}
145            </Text>
146          ))}
147        </Text>
148        {pairs.map((pair, i) => (
149          <Box key={`l${i}`}>
150            {pair.map(r => (
151              <Box key={r.label} width={half}>
152                <Text wrap="truncate">
153                  <Text color={r.color}>■ </Text>
154                  {r.label} <Text bold>{fmt(r.tokens)}</Text>
155                  {r.label === 'free' ? '' : <Text dimColor> {pct(r.tokens, s.window)}%</Text>}
156                </Text>
157              </Box>
158            ))}
159          </Box>
160        ))}
161      </Box>
162    )
163  })
164}
165

1export type Row = { label: string; tokens: number; color: string }
2export type Snapshot = { used: number; window: number; compactAt: number | null; rows: Row[]; free: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'context-bar': { snapshot: Snapshot | null }
7  }
8}
9