context-gauge

1// Context Gauge: one line under the prompt with repo/branch, model, effort and
2// how much context is left. Replaces the old status line script.
3//
4//   my-project/main · Opus 5.5 1M · high  ████████████░░░░░░░░ 62% · 620K free  ·  5h 72% · 2:14:05  ·  wk 60%  ·  68 t/s
5//
6// `5h` is the plan's 5-hour usage window: what is left, draining from 100% to
7// 0, and a countdown to its reset; `wk` is what is left of the weekly window.
8// Both only on a subscription. `t/s` is how fast Claude writes: output tokens
9// per second of the last few main-loop responses, timed from the first token.
10//
11// `/gauge` swaps that line for a detailed band above the prompt: one bar split
12// by what fills the window (/context's categories and colours), and a key.
13// `/gauge on` and `/gauge off` set it outright; the choice is kept.
14
15import { atom, read, update } from 'claude-code'
16import type { EngineInterface, Register } from 'claude-code'
17
18import type { GaugeBreakdown, GaugeCategory, GaugeLimit, GaugeReading } from '../types'
19
20const EMPTY: GaugeReading = { model: '', effort: null, tokens: null, window: null, percent: null, place: null }
21
22// Held by the host, so the reading survives a hot reload of this file.
23const reading = atom({ plugin: 'context-gauge', key: 'reading' } as const, EMPTY)
24const detail = atom({ plugin: 'context-gauge', key: 'detail' } as const, false)
25const breakdown = atom({ plugin: 'context-gauge', key: 'breakdown' } as const, null as GaugeBreakdown | null)
26const limits = atom({ plugin: 'context-gauge', key: 'limits' } as const, [] as GaugeLimit[])
27const clock = atom({ plugin: 'context-gauge', key: 'now' } as const, 0)
28const speed = atom({ plugin: 'context-gauge', key: 'speed' } as const, [] as number[])
29
30// The countdown's ticker; a reload drops it and session.start starts another.
31let lastTick = ''
32
33// A breakdown is a local estimate, but a long turn calls tools often.
34const BREAKDOWN_EVERY_MS = 3000
35
36// A response shorter than this says little about speed (a lone tool call).
37const SPEED_MIN_TOKENS = 50
38// Responses averaged into the t/s figure.
39const SPEED_SAMPLES = 3
40
41export const register: Register = on => {
42  on('session.start', async ($, e, next) => {
43    const result = await next(e)
44    const settings = (await $.settings.read()) as Record<string, unknown>
45    const effort = typeof settings.effortLevel === 'string' ? settings.effortLevel : null
46    const model = await $.session.model()
47    const place = await whereAmI($)
48    await update($, reading, r => ({ ...r, model: r.model || model, effort: r.effort ?? effort, place }))
49    await takeUsage($)
50    await update($, clock, () => Date.now())
51    lastTick = ''
52    // Moves the clock only when the countdown's text would change.
53    $.clock.every(1000, async () => {
54      const now = Date.now()
55      const label = countdown(await read($, limits), now)
56      if (label === lastTick) return
57      lastTick = label
58      await update($, clock, () => now)
59    })
60    await $.command.register({
61      name: 'gauge',
62      description: 'Switch the context gauge between one line and the detailed breakdown',
63      argumentHint: '[on|off]',
64      immediate: true,
65    })
66    const isDetailed = (await $.store.get('detail')) === true
67    await update($, detail, () => isDetailed)
68    if (isDetailed) await takeBreakdown($, true)
69
70    return result
71  })
72
73  on('command.run', { command: 'gauge' }, async ($, e) => {
74    const word = e.args.trim().toLowerCase()
75    if (word && !['on', 'off'].includes(word)) {
76      return { text: 'Usage: /gauge [on|off]. Bare /gauge switches views.' }
77    }
78    const isDetailed = word ? word === 'on' : !(await read($, detail))
79    await $.store.set('detail', isDetailed)
80    await update($, detail, () => isDetailed)
81    if (isDetailed) await takeBreakdown($, true)
82
83    return {
84      text: isDetailed
85        ? 'Context gauge: detailed, above the prompt. /gauge off for the one line.'
86        : 'Context gauge: one line under the prompt.',
87    }
88  })
89
90  // The model and effort each main-loop request actually goes out with, and
91  // how fast its response streams: output tokens over the time from the first
92  // token (text, thinking or a tool call) to the end, so the wait before
93  // Claude starts answering does not count against its speed.
94  on('turn.step', async function* ($, e, next) {
95    if (e.agentId) return yield* next(e)
96    const effort = e.effort === undefined ? null : String(e.effort)
97    await update($, reading, r => ({ ...r, model: e.model, effort }))
98
99    const stream = next(e)
100    let firstAt: number | undefined
101    for (;;) {
102      const n = await stream.next()
103      if (n.done) {
104        const out = n.value.usage?.output_tokens ?? 0
105        if (firstAt !== undefined && out >= SPEED_MIN_TOKENS) {
106          const seconds = ((await $.clock.now()) - firstAt) / 1000
107          if (seconds > 0) await update($, speed, s => [...s, out / seconds].slice(-SPEED_SAMPLES))
108        }
109        return n.value
110      }
111      // the engine's own chunks (the envelope) come before any token
112      if (firstAt === undefined && n.value.kind !== 'engine' && n.value.kind !== 'stop') firstAt = await $.clock.now()
113      yield n.value
114    }
115  })
116
117  // Pushed after every main-thread turn.
118  on('session.measure', async ($, e, next) => {
119    await setContext($, e.context)
120    await setLimits($, e.rateLimits)
121    await takeBreakdown($, true)
122
123    return next(e)
124  })
125
126  // Keeps the gauge moving during a long turn, not only at its end.
127  on('tool.call', async ($, e, next) => {
128    const result = await next(e)
129    if (!e.agentId) {
130      await takeUsage($)
131      await takeBreakdown($, false)
132    }
133
134    return result
135  })
136
137  on('turn.complete', async ($, e, next) => {
138    const result = await next(e)
139    if (!e.agentId) {
140      const place = await whereAmI($)
141      await update($, reading, r => ({ ...r, place }))
142    }
143
144    return result
145  })
146
147  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
148    const engine = await next(e)
149    if (await read($, detail)) {
150      // The detailed view's key sits under the input, its bar right above it.
151      const b = await read($, breakdown)
152      if (!b) return engine
153      const { Box, Text } = $.ui.resolve(e)
154
155      return (
156        <Box flexDirection="column">
157          {legend(Box, Text, b, e.viewport?.columns ?? 120)}
158          {engine}
159        </Box>
160      )
161    }
162    const r = await read($, reading)
163    if (!r.model && r.percent === null) {
164      return engine
165    }
166    const { Box, Text } = $.ui.resolve(e)
167    const columns = e.viewport?.columns ?? 120
168    const windows = await read($, limits)
169    const five = windows.find(l => l.kind === 'five_hour')
170    const week = windows.find(l => l.kind === 'seven_day')
171    const now = await read($, clock)
172    const tps = average(await read($, speed))
173
174    return (
175      <Box flexDirection="column">
176        <Box flexDirection="row">
177          {gauge(Text, r, columns)}
178          {five ? usage(Text, five, now, columns) : null}
179          {week ? weekly(Text, week, now) : null}
180          {tps !== null ? <Text dimColor>{`  ·  ${Math.round(tps)} t/s`}</Text> : null}
181        </Box>
182        {engine}
183      </Box>
184    )
185  })
186
187  // Composes: another mod (where-we-are) may draw in this band too. Theirs
188  // goes first so the bar sits right on the input.
189  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
190    const others = await next(e)
191    if (e.props.hasSurvey || !(await read($, detail))) return others
192    const r = await read($, reading)
193    const b = await read($, breakdown)
194    const { Box, Text } = $.ui.resolve(e)
195
196    return (
197      <Box flexDirection="column">
198        {others}
199        {band(Box, Text, r, b, e.props.bodyColumns)}
200      </Box>
201    )
202  })
203}
204
205async function takeUsage($: EngineInterface) {
206  const { context, rateLimits } = await $.session.usage()
207  await setContext($, context)
208  await setLimits($, rateLimits)
209}
210
211async function setLimits($: EngineInterface, windows: readonly { kind: string; percentUsed: number; resetsAt?: string }[]) {
212  const next: GaugeLimit[] = windows.map(w => {
213    const at = w.resetsAt ? Date.parse(w.resetsAt) : NaN
214    return { kind: w.kind, percentUsed: w.percentUsed, resetsAt: Number.isFinite(at) ? at : null }
215  })
216  const was = await read($, limits)
217  if (JSON.stringify(was) !== JSON.stringify(next)) await update($, limits, () => next)
218}
219
220// The 5-hour window's countdown text, or '' when there is none to show.
221function countdown(windows: GaugeLimit[], now: number): string {
222  const five = windows.find(l => l.kind === 'five_hour')
223  if (!five?.resetsAt || five.resetsAt <= now) return ''
224  return clockText(five.resetsAt - now)
225}
226
227function clockText(ms: number): string {
228  const s = Math.ceil(ms / 1000)
229  const h = Math.floor(s / 3600)
230  const m = Math.floor((s % 3600) / 60)
231  const sec = s % 60
232  const mm = String(m).padStart(2, '0')
233  const ss = String(sec).padStart(2, '0')
234  return h > 0 ? `${h}:${mm}:${ss}` : `${m}:${ss}`
235}
236
237// What is left of the weekly window, short: `wk 60%`.
238function weekly(Text: any, week: GaugeLimit, now: number) {
239  const left = leftOf(week, now)
240  return [<Text dimColor>{'  ·  wk '}</Text>, <Text color={leftColor(left)}>{`${left}%`}</Text>]
241}
242
243// Percent of a window left, 0 to 100; a window past its reset reads full
244// until the next response says more.
245function leftOf(w: GaugeLimit, now: number): number {
246  if (w.resetsAt !== null && w.resetsAt <= now) return 100
247  return Math.max(0, Math.min(100, Math.round(100 - w.percentUsed)))
248}
249
250function leftColor(left: number): string {
251  return left > 20 ? 'green' : left >= 10 ? 'yellow' : 'red'
252}
253
254function average(values: readonly number[]): number | null {
255  return values.length ? values.reduce((sum, v) => sum + v, 0) / values.length : null
256}
257
258// What is left of the 5-hour window, draining from 100% to 0, and the time to
259// its reset. Past the reset it reads full until the next response says more.
260function usage(Text: any, five: GaugeLimit, now: number, columns: number) {
261  const left = leftOf(five, now)
262  const parts = [<Text dimColor>{'  ·  5h '}</Text>]
263  parts.push(<Text color={leftColor(left)}>{`${left}%`}</Text>)
264  const label = countdown([five], now)
265  if (label) parts.push(<Text dimColor>{` · ${columns >= 100 ? label : label.replace(/^(\d+:\d\d):\d\d$/, '$1')}`}</Text>)
266
267  return parts
268}
269
270async function setContext(
271  $: EngineInterface,
272  context: { tokens?: number; window: number; percent?: number } | undefined,
273) {
274  if (!context?.window) return
275  const tokens = context.tokens ?? null
276  const percent =
277    context.percent ?? (tokens === null ? null : Math.round((tokens / context.window) * 100))
278  await update($, reading, r => ({ ...r, tokens, window: context.window, percent }))
279}
280
281// Only while the detailed band shows; `force` skips the throttle.
282async function takeBreakdown($: EngineInterface, force: boolean) {
283  if (!(await read($, detail))) return
284  const now = await $.clock.now()
285  const last = await read($, breakdown)
286  if (!force && last && now - last.takenAt < BREAKDOWN_EVERY_MS) return
287  let usage
288  try {
289    usage = await $.session.usage({ breakdown: 'summary' })
290  } catch {
291    return
292  }
293  const found = usage.context.breakdown
294  if (!found) return
295  const categories: GaugeCategory[] = []
296  for (const c of found.categories) {
297    if (c.kind === 'deferred' || c.isDeferred) continue
298    categories.push({ name: c.name, tokens: c.tokens, color: c.color, kind: c.kind })
299  }
300  await update($, breakdown, () => ({
301    categories,
302    totalTokens: found.totalTokens,
303    rawMaxTokens: found.rawMaxTokens,
304    percentage: found.percentage,
305    takenAt: now,
306  }))
307}
308
309async function whereAmI($: EngineInterface): Promise<string | null> {
310  const top = await $.process.run(['git', 'rev-parse', '--show-toplevel'], { timeoutMs: 3000 })
311  if (top.exitCode !== 0) return null
312  const repo = top.stdout.trim().split('/').pop() ?? ''
313  const branch = await $.process.run(['git', 'branch', '--show-current'], { timeoutMs: 3000 })
314  const name = branch.stdout.trim() || 'detached'
315
316  return `${repo}/${name}`
317}
318
319function gauge(Text: any, r: GaugeReading, columns: number) {
320  const parts = []
321  if (r.place && columns >= 100) parts.push(<Text dimColor>{`${r.place} · `}</Text>)
322  if (r.model) parts.push(<Text bold>{modelLabel(r.model)}</Text>)
323  if (r.window) parts.push(<Text dimColor>{` ${short(r.window)}`}</Text>)
324  if (r.effort) parts.push(<Text dimColor>{` · ${r.effort}`}</Text>)
325  if (r.percent !== null && r.window) {
326    const left = Math.max(0, 100 - r.percent)
327    const free = Math.max(0, r.window - (r.tokens ?? 0))
328    const color = left > 20 ? 'green' : left >= 10 ? 'yellow' : 'red'
329    const width = columns >= 100 ? 20 : 10
330    const filled = Math.min(width, Math.round((left * width) / 100))
331    parts.push(<Text dimColor>{'  '}</Text>)
332    parts.push(<Text color={color}>{'█'.repeat(filled)}</Text>)
333    parts.push(<Text dimColor>{'░'.repeat(width - filled)}</Text>)
334    parts.push(<Text color={color}>{` ${left}%`}</Text>)
335    parts.push(<Text dimColor>{` · ${short(free)} free`}</Text>)
336  }
337
338  return parts
339}
340
341// The detailed band above the input: who and where, then the window as one
342// bar split by category. Its key is `legend`, drawn under the input.
343//
344//   my-project/main · Opus 5.5 1M · high              38% used · 620K free
345//   ███▓▓▓▓████████████████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▒▒▒▒
346//   > prompt
347//   ■ sys 4K  ■ tools 18K  ■ msgs 340K  ▒ buffer 33K
348function band(Box: any, Text: any, r: GaugeReading, b: GaugeBreakdown | null, columns: number) {
349  const who = []
350  if (r.place && columns >= 70) who.push(<Text dimColor>{`${r.place} · `}</Text>)
351  if (r.model) who.push(<Text bold>{modelLabel(r.model)}</Text>)
352  if (r.window) who.push(<Text dimColor>{` ${short(r.window)}`}</Text>)
353  if (r.effort) who.push(<Text dimColor>{` · ${r.effort}`}</Text>)
354
355  if (!b) {
356    return [
357      <Box flexDirection="row">
358        {who}
359        <Text dimColor>{'  measuring context…'}</Text>
360      </Box>,
361    ]
362  }
363
364  const used = Math.min(100, Math.max(0, Math.round(b.percentage)))
365  const left = 100 - used
366  const color = left > 20 ? 'green' : left >= 10 ? 'yellow' : 'red'
367  const free = b.categories.find(c => c.kind === 'free')?.tokens ?? Math.max(0, b.rawMaxTokens - b.totalTokens)
368  const width = Math.max(10, Math.min(columns, 120))
369  const cells = split(b.categories, width)
370
371  return [
372    <Box flexDirection="row" justifyContent="space-between" width={columns}>
373      <Box flexDirection="row">{who}</Box>
374      <Box flexDirection="row">
375        <Text color={color} bold>{`${used}% used`}</Text>
376        <Text dimColor>{` · ${short(free)} free`}</Text>
377      </Box>
378    </Box>,
379    <Box flexDirection="row">
380      {b.categories.map((c, i) =>
381        cells[i] ? (
382          c.kind === 'free' ? (
383            <Text dimColor>{'░'.repeat(cells[i]!)}</Text>
384          ) : (
385            <Text color={c.color}>{(c.kind === 'buffer' ? '▒' : '█').repeat(cells[i]!)}</Text>
386          )
387        ) : null,
388      )}
389    </Box>,
390  ]
391}
392
393// The bar's key, short labels wrapped to the width: `■ msgs 340K  ■ sys 4K`.
394function legend(Box: any, Text: any, b: GaugeBreakdown, columns: number) {
395  const shown = b.categories.filter(c => c.kind !== 'free' && c.tokens > 0)
396  const lines: GaugeCategory[][] = [[]]
397  let room = columns
398  for (const c of shown) {
399    const size = 2 + label(c).length + 1 + short(c.tokens).length + 2
400    const line = lines[lines.length - 1]!
401    if (line.length > 0 && size > room) {
402      lines.push([c])
403      room = columns - size
404    } else {
405      line.push(c)
406      room -= size
407    }
408  }
409
410  return lines.map(line => (
411    <Box flexDirection="row">
412      {line.flatMap(c => [
413        <Text color={c.color}>{c.kind === 'buffer' ? '▒ ' : '■ '}</Text>,
414        <Text>{label(c)}</Text>,
415        <Text dimColor>{` ${short(c.tokens)}  `}</Text>,
416      ])}
417    </Box>
418  ))
419}
420
421// Short display names for /context's rows; one it does not know keeps its
422// first word. Display only: nothing branches on a row's name.
423const LABELS: Record<string, string> = {
424  'system prompt': 'sys',
425  'system tools': 'tools',
426  'mcp tools': 'mcp',
427  'custom agents': 'agents',
428  'memory files': 'memory',
429  skills: 'skills',
430  'slash commands': 'cmds',
431  messages: 'msgs',
432  'autocompact buffer': 'buffer',
433}
434
435function label(c: GaugeCategory): string {
436  const name = c.name.toLowerCase()
437  if (c.kind === 'buffer') return 'buffer'
438
439  return LABELS[name] ?? name.split(/\s+/)[0] ?? name
440}
441
442// Cells per category that add up to exactly `width`: largest remainder, then
443// a cell for any used category too small to round to one, taken from free.
444function split(categories: GaugeCategory[], width: number): number[] {
445  const total = categories.reduce((sum, c) => sum + c.tokens, 0)
446  if (total <= 0) return categories.map(c => (c.kind === 'free' ? width : 0))
447  const exact = categories.map(c => (c.tokens * width) / total)
448  const cells = exact.map(Math.floor)
449  let spare = width - cells.reduce((sum, n) => sum + n, 0)
450  const order = exact.map((x, i) => [x - Math.floor(x), i] as const).sort((a, b) => b[0] - a[0])
451  for (const [, i] of order) {
452    if (spare <= 0) break
453    cells[i]! += 1
454    spare -= 1
455  }
456  const freeAt = categories.findIndex(c => c.kind === 'free')
457  categories.forEach((c, i) => {
458    if (c.kind === 'used' && c.tokens > 0 && cells[i] === 0 && freeAt >= 0 && cells[freeAt]! > 1) {
459      cells[i] = 1
460      cells[freeAt]! -= 1
461    }
462  })
463
464  return cells
465}
466
467// Shows whatever id the engine reports; only recognized families get a short name.
468function modelLabel(id: string): string {
469  const lower = id.toLowerCase()
470  const versioned = lower.match(/(opus|sonnet|haiku|fable)-(\d+)(?:-(\d{1,2}))?(?!\d)/)
471  if (versioned) {
472    const [, name = '', major = '', minor] = versioned
473    const family = capitalize(name)
474    return minor ? `${family} ${major}.${minor}` : `${family} ${major}`
475  }
476  const alias = lower.match(/^(opus|sonnet|haiku|fable)\b/)
477  if (alias) return capitalize(alias[1] ?? id)
478
479  return id
480}
481
482function capitalize(word: string): string {
483  return word.charAt(0).toUpperCase() + word.slice(1)
484}
485
486function short(n: number): string {
487  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
488  if (n >= 1_000) return `${Math.round(n / 1_000)}K`
489  return String(n)
490}
491

1export type GaugeReading = {
2  model: string
3  effort: string | null
4  tokens: number | null
5  window: number | null
6  percent: number | null
7  place: string | null
8}
9
10// One plan usage window as the last response reported it (`five_hour`,
11// `seven_day`): how much is used and when it resets (epoch ms, or null).
12export type GaugeLimit = {
13  kind: string
14  percentUsed: number
15  resetsAt: number | null
16}
17
18// One row of /context's breakdown, as the detailed band draws it.
19export type GaugeCategory = {
20  name: string
21  tokens: number
22  // A theme key (/context's own colour for the row) or a raw colour.
23  color: string
24  kind: 'used' | 'free' | 'buffer'
25}
26
27export type GaugeBreakdown = {
28  categories: GaugeCategory[]
29  totalTokens: number
30  rawMaxTokens: number
31  percentage: number
32  takenAt: number
33}
34
35declare module 'claude-code' {
36  interface PluginState {
37    'context-gauge': {
38      reading: GaugeReading
39      // On: the detailed band above the prompt; off: the one line under it.
40      detail: boolean
41      breakdown: GaugeBreakdown | null
42      // Plan usage windows; empty off a subscription (API key, cloud provider).
43      limits: GaugeLimit[]
44      // The clock the usage countdown reads, moved once a second while it shows.
45      now: number
46      // Output tokens per second of the last few main-loop responses, newest last.
47      speed: number[]
48    }
49  }
50}
51