context-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Reading } from '../types'
5import { layout, reduce } from './layout'
6
7const USAGE = 'Usage: /context-bar [on|off]'
8
9// Mirrors $.store's isOn (R15), so the render hook reads state only.
10const isOn = atom({ plugin: 'context-bar', key: 'isOn' } as const, false)
11// Change the tag when Reading changes: it declines what an earlier version
12// stored, which outlives a reload.
13const reading = atom({ plugin: 'context-bar', key: 'reading' } as const, null, {
14  shape: 'reading-2',
15})
16
17// R13: starting a refresh, a compaction's clear and turning the bar off
18// each move the epoch; every write to the reading goes through `queue`.
19// Module variables, so a reload starts them over (spec concern 5).
20let epoch = 0
21let queue: Promise<void> = Promise.resolve()
22
23// One write at a time. A refresh's write (`mine` set) checks its epoch at
24// the head of the queue, immediately before the write, and is dropped if
25// the epoch has moved; nothing else writes the reading in between.
26function publish($: EngineInterface, value: Reading | null, mine?: number): Promise<void> {
27  queue = queue.then(async () => {
28    if (mine !== undefined && mine !== epoch) return
29    try {
30      await update($, reading, () => value)
31    } catch (err) {
32      $.ui.log(`context-bar: state write failed: ${err}`, { to: 'debug' })
33    }
34  })
35  return queue
36}
37
38// Queued, never awaited: a held write delays no compaction or command.
39function clear($: EngineInterface): void {
40  epoch++
41  void publish($, null)
42}
43
44// R10, R12: started with `void`, so nothing may escape it.
45async function refresh($: EngineInterface): Promise<void> {
46  const mine = ++epoch
47  let next: Reading | null = null
48  try {
49    next = reduce(await $.session.usage({ breakdown: 'summary' }))
50  } catch (err) {
51    $.ui.log(`context-bar: refresh failed: ${err}`, { to: 'debug' })
52  }
53  await publish($, next, mine)
54}
55
56export const register: Register = on => {
57  on('session.start', async ($, e, next) => {
58    const result = await next(e)
59    await $.command.register({
60      name: 'context-bar',
61      description: 'Show or hide the context breakdown above the prompt: [on|off]',
62    })
63    const stored = (await $.store.get('isOn')) === true
64    await update($, isOn, () => stored)
65    if (stored) void refresh($)
66    return result
67  })
68
69  on('turn.complete', async ($, e, next) => {
70    const result = await next(e)
71    if (e.agentId === undefined && (await read($, isOn))) void refresh($)
72    return result
73  })
74
75  // R9, R11: a main-loop compaction clears first and refreshes after, even
76  // when it was vetoed with { skip }.
77  on('session.compact', async ($, e, next) => {
78    if (e.agentId !== undefined || e.trigger === 'precompute') return next(e)
79    clear($)
80    const result = await next(e)
81    if (await read($, isOn)) void refresh($)
82    return result
83  })
84
85  // R14, R15.
86  on('command.run', { command: 'context-bar' }, async ($, e) => {
87    const arg = e.args.trim()
88    if (arg !== '' && arg !== 'on' && arg !== 'off') return { text: USAGE }
89    const isTurnedOn = arg === '' ? !(await read($, isOn)) : arg === 'on'
90    await $.store.set('isOn', isTurnedOn)
91    await update($, isOn, () => isTurnedOn)
92    if (isTurnedOn) void refresh($)
93    else clear($)
94    return { text: isTurnedOn ? 'Context bar on.' : 'Context bar off.' }
95  })
96
97  // R8: step aside for a survey, while off, with no reading, or under 20
98  // columns; otherwise draw the band layout gives, framed from 60 (R7).
99  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
100    if (e.props.hasSurvey || !(await read($, isOn))) return next(e)
101    const value = await read($, reading)
102    const band = value && layout(value, e.props.bodyColumns)
103    if (!band) return next(e)
104
105    const { Box, Text } = $.ui.resolve(e)
106    const frame = band.isFramed ? { borderStyle: 'round', borderColor: 'inactive', paddingX: 1 } : {}
107    return (
108      <Box flexDirection="column" {...frame}>
109        {band.rows.map(row => (
110          <Text>
111            {row.map(s => (
112              <Text color={s.color} dimColor={s.dim} bold={s.bold} backgroundColor={s.background}>
113                {s.text}
114              </Text>
115            ))}
116          </Text>
117        ))}
118      </Box>
119    )
120  })
121}
122

1// Pure helpers for the band: no `$`, so every rule is tested without the
2// engine. Rule numbers are the spec's (intent/show-context-breakdown/spec.md).
3
4import type { ContextCategory, ContextCategoryKind, SessionUsage } from 'claude-code'
5
6import type { Part, Reading } from '../types'
7
8// One run of text in a row; the render hook draws each as a <Text>.
9export type Span = { text: string; color?: string; dim?: boolean; bold?: boolean; background?: string }
10export type Row = Span[]
11
12const SEP = '  '
13
14// R5: under 1000 as is, then k to one decimal (none from 10k), then M to
15// two decimals, trailing zeros dropped. Integer rounding, so 9950 is 10k.
16export function shortTokens(n: number): string {
17  if (n < 1000) return String(Math.round(n))
18  if (Math.round(n / 1000) < 1000) {
19    const tenths = Math.round(n / 100)
20    return tenths < 100 ? `${trim((tenths / 10).toFixed(1))}k` : `${Math.round(n / 1000)}k`
21  }
22  return `${trim((Math.round(n / 10_000) / 100).toFixed(2))}M`
23}
24
25const trim = (s: string) => s.replace(/\.?0+$/, '')
26
27// R5: whole percent of the window; any non-zero share is at least 1%.
28export function percent(tokens: number, window: number): string {
29  if (tokens <= 0 || window <= 0) return '0%'
30  return `${Math.max(1, Math.round((tokens / window) * 100))}%`
31}
32
33const LABELS: Record<string, string> = {
34  'System prompt': 'system prompt',
35  'System tools': 'tools',
36  'MCP tools': 'mcp tools',
37  'Custom agents': 'agents',
38  'Memory files': 'memory files',
39  Skills: 'skills',
40  Messages: 'messages',
41}
42
43// R5: by kind first; a used row through the map, else its name lowercased.
44export function labelFor(row: { name: string; kind: ContextCategoryKind }): string {
45  if (row.kind === 'free') return 'free'
46  if (row.kind === 'buffer') return 'reserved'
47  return LABELS[row.name] ?? row.name.toLowerCase()
48}
49
50// R3: share `width` cells among parts in bar order. When no part has tokens
51// the whole bar goes to `fill` (the free part).
52export function apportion(tokens: number[], width: number, fill: number): number[] {
53  const d = tokens.reduce((a, b) => a + b, 0)
54  if (d <= 0) return tokens.map((_, i) => (i === fill ? width : 0))
55
56  const parts = tokens.map((t, i) => {
57    const exact = (t * width) / d
58    return { i, t, cells: 0, floor: Math.floor(exact), remainder: exact - Math.floor(exact) }
59  })
60  const nonZero = parts.filter(p => p.t > 0)
61
62  // Step 3, crowded: the `width` largest parts (ties: earlier) get one each.
63  if (nonZero.length > width) {
64    const chosen = [...nonZero].sort((a, b) => b.t - a.t || a.i - b.i).slice(0, width)
65    for (const p of chosen) p.cells = 1
66    return parts.map(p => p.cells)
67  }
68
69  // Step 2: largest remainder; remainder ties go to the earlier part.
70  for (const p of parts) p.cells = p.floor
71  let left = width - parts.reduce((a, p) => a + p.cells, 0)
72  for (const p of [...nonZero].sort((a, b) => b.remainder - a.remainder || a.i - b.i)) {
73    if (left === 0) break
74    p.cells++
75    left--
76  }
77
78  // Step 3: a part with tokens and no cell takes one from the part with the
79  // most cells (ties: the later), which has at least 2 since k ≤ width.
80  for (const p of nonZero) {
81    if (p.cells > 0) continue
82    const donor = parts.reduce((most, q) => (q.cells >= most.cells ? q : most))
83    donor.cells--
84    p.cells = 1
85  }
86  return parts.map(p => p.cells)
87}
88
89// R2, R10: the breakdown reduced to what the layout needs, by kind; null
90// when there is no breakdown or no used row.
91export function reduce(usage: SessionUsage): Reading | null {
92  const b = usage.context.breakdown
93  if (!b) return null
94  const part = (c: ContextCategory): Part => ({ label: labelFor(c), tokens: c.tokens, color: c.color })
95  const used = b.categories.filter(c => c.kind === 'used').map(part)
96  if (used.length === 0) return null
97  const free = b.categories.find(c => c.kind === 'free')
98  const buffer = b.categories.find(c => c.kind === 'buffer')
99  const window = b.rawMaxTokens
100  return {
101    used,
102    free: free ? part(free) : { label: 'free', tokens: 0, color: 'inactive' },
103    reserved: buffer ? part(buffer) : null,
104    totalTokens: b.totalTokens,
105    window,
106    compactsAt: b.isAutoCompactEnabled ? (b.autoCompactThreshold ?? null) : null,
107  }
108}
109
110// R7: cut text to `n` columns, ending in … when anything was cut.
111function cut(s: string, n: number): string {
112  if (s.length <= n) return s
113  if (n <= 0) return ''
114  return s.slice(0, n - 1) + '…'
115}
116
117// R7: left and right parts with at least one space between; when both do
118// not fit, the right part wins and the left is cut.
119export function spread(left: Row, right: Row, width: number): Row {
120  const r = cutRow(right, width)
121  const l = cutRow(left, width - widthOf(r) - 1)
122  const gap = width - widthOf(l) - widthOf(r)
123  return [...l, ...(gap > 0 ? [{ text: ' '.repeat(gap) }] : []), ...r]
124}
125
126// R6: the share of the window in use, on green, yellow from 70%, red from 90%.
127function badge(totalTokens: number, window: number): Span {
128  const pct = percent(totalTokens, window)
129  const n = Number.parseInt(pct, 10)
130  const background = n >= 90 ? 'error' : n >= 70 ? 'warning' : 'success'
131  return { text: ` ${pct} `, color: 'inverseText', background }
132}
133
134// The band: its rows, and whether they sit in the frame (R7).
135export type Band = { rows: Row[]; isFramed: boolean }
136
137// The frame's border and padding take two columns each side.
138const FRAME = 4
139
140// R1, R4–R7: the band at `width` columns, or null to step aside.
141export function layout(reading: Reading, width: number): Band | null {
142  if (width < 20) return null
143  const { used, window, totalTokens, compactsAt } = reading
144  // R2: over the window, free and reserved show 0.
145  const isOver = totalTokens > window
146  const free: Part = { ...reading.free, tokens: isOver ? 0 : reading.free.tokens }
147  const reserved: Part | null = reading.reserved && {
148    ...reading.reserved,
149    tokens: isOver ? 0 : reading.reserved.tokens,
150  }
151  const u = shortTokens(totalTokens)
152  const w = shortTokens(window)
153
154  // R3, R4: one glyph per cell, each part in its row's colour; the reserve
155  // opens with the compaction marker.
156  const bar = (cells: number): Row => {
157    const parts = [...used, free, ...(reserved ? [reserved] : [])]
158    const counts = apportion(parts.map(p => p.tokens), cells, used.length)
159    const spans: Row = []
160    parts.forEach((p, i) => {
161      const n = counts[i] ?? 0
162      if (n === 0) return
163      if (i < used.length) spans.push({ text: '█'.repeat(n), color: p.color })
164      else if (p === free) spans.push({ text: '░'.repeat(n), color: p.color })
165      else {
166        spans.push({ text: '│', color: 'warning' })
167        if (n > 1) spans.push({ text: '░'.repeat(n - 1), color: p.color, dim: true })
168      }
169    })
170    return spans
171  }
172
173  const mark: Span = { text: '◆', color: 'claude' }
174  if (width < 40) {
175    const prefix = ` ${u}/${w} `
176    return { rows: [[mark, { text: prefix }, ...bar(width - 1 - prefix.length)]], isFramed: false }
177  }
178  const left: Row = [mark, { text: ' context', bold: true }]
179  const share = badge(totalTokens, window)
180  if (width < 60) {
181    const right: Row = [{ text: `${u}/${w}`, bold: true }, { text: ' ' }, share]
182    return { rows: [spread(left, right, width), bar(width)], isFramed: false }
183  }
184
185  const room = width - FRAME
186  const context = ` of ${w}` + (compactsAt === null ? '' : ` · compacts at ${shortTokens(compactsAt)}`)
187  const right: Row = [{ text: u, bold: true }, { text: context, dim: true }, { text: ' ' }, share]
188  const entries: Row[] = [...used.map(p => entry(p, percent(p.tokens, window))), entry(free, null)]
189  return { rows: [spread(left, right, room), bar(room), ...wrap(entries, room)], isFramed: true }
190}
191
192// R5: `■ <label> <tokens> <pct>`, the swatch coloured, the tokens bold and
193// the percent dim; the free entry is dim throughout, with no percent.
194function entry(p: Part, pct: string | null): Row {
195  const swatch: Span = { text: '■', color: p.color }
196  if (pct === null) return [swatch, { text: ` ${p.label} ${shortTokens(p.tokens)}`, dim: true }]
197  return [
198    swatch,
199    { text: ` ${p.label} ` },
200    { text: shortTokens(p.tokens), bold: true },
201    { text: ' ' },
202    { text: pct, dim: true },
203  ]
204}
205
206const widthOf = (r: Row) => r.reduce((a, s) => a + s.text.length, 0)
207
208// Cut a row to `n` columns ending in …, which takes the last column even
209// when the cut falls between two spans.
210function cutRow(r: Row, n: number): Row {
211  if (widthOf(r) <= n) return r
212  if (n <= 0) return []
213  const out: Row = []
214  let room = n - 1
215  for (const s of r) {
216    if (room <= 0) break
217    const text = s.text.slice(0, room)
218    out.push({ ...s, text })
219    room -= text.length
220  }
221  const last = out.at(-1)
222  if (last) last.text += '…'
223  else out.push({ text: '…' })
224  return out
225}
226
227// R5, R7: pack whole entries into rows of at most `width`, two spaces apart.
228function wrap(entries: Row[], width: number): Row[] {
229  const rows: Row[] = []
230  let current: Row = []
231  for (const e of entries.map(e => cutRow(e, width))) {
232    if (current.length > 0 && widthOf(current) + SEP.length + widthOf(e) <= width) {
233      current.push({ text: SEP }, ...e)
234    } else {
235      if (current.length > 0) rows.push(current)
236      current = [...e]
237    }
238  }
239  if (current.length > 0) rows.push(current)
240  return rows
241}
242

1// One segment of the bar and one legend entry: a breakdown row reduced to
2// what the layout needs. `color` is the row's theme key.
3export type Part = { label: string; tokens: number; color: string }
4
5// The breakdown reduced (spec "Reading"): every figure shares `window`.
6export type Reading = {
7  // kind 'used', in /context order
8  used: Part[]
9  // kind 'free', or 0 tokens when the row is missing
10  free: Part
11  // kind 'buffer'; null when absent
12  reserved: Part | null
13  totalTokens: number
14  // rawMaxTokens
15  window: number
16  // autoCompactThreshold in tokens, or null when auto-compact is off
17  compactsAt: number | null
18}
19
20declare module 'claude-code' {
21  interface PluginState {
22    'context-bar': {
23      // Mirrors $.store's isOn, which is the source of truth across sessions.
24      isOn: boolean
25      reading: Shaped<Reading | null>
26    }
27  }
28}
29