context-gauge

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextUsage } from 'claude-code'
3
4import type { Reading } from '../types'
5import { HINT, barWidth, colorFor, crossed, formatTokens, label, segments, toastText, warnLevel } from './gauge'
6
7const reading = atom({ plugin: 'context-gauge', key: 'reading' } as const, null as Reading | null)
8const isHidden = atom({ plugin: 'context-gauge', key: 'isHidden' } as const, false)
9
10const PREFIX = 'context '
11
12/** Turns the engine's context figures into a reading; null until a response of the live window reported one. */
13function toReading(c: SessionContextUsage): Reading | null {
14  if (c.tokens === undefined || c.percent === undefined || !(c.window > 0)) return null
15  return { tokens: c.tokens, window: c.window, percent: Math.round(c.percent) }
16}
17
18/** Stores a new reading and toasts once per threshold crossed on the way up. */
19async function record($: EngineInterface, c: SessionContextUsage, warnAt: number) {
20  const next = toReading(c)
21  const prev = await read($, reading)
22  await update($, reading, () => next)
23  if (next === null) return
24  const hit = crossed(prev?.percent ?? null, next.percent, warnAt)
25  if (hit !== null) $.ui.toast(toastText(next.percent, hit), { timeoutMs: 8000 })
26}
27
28/** Reads the live window and records it; a failed read leaves the last reading in place. */
29async function refresh($: EngineInterface, warnAt: number) {
30  try {
31    const usage = await $.session.usage()
32    await record($, usage.context, warnAt)
33  } catch {
34    // No session bound yet, or the read was refused: keep what we had.
35  }
36}
37
38async function setHidden($: EngineInterface, hidden: boolean) {
39  await update($, isHidden, () => hidden)
40  await $.store.set('isHidden', hidden)
41}
42
43function describe(r: Reading | null, hidden: boolean, warnAt: number): string {
44  const state = hidden ? ' (band hidden; /gauge show)' : ''
45  if (r === null) return `Context gauge: no reading yet (one arrives after the next response)${state}.`
46  const hint = r.percent >= warnAt ? ' Consider /compact.' : ''
47  return `Context gauge: ${r.percent}% used · ${formatTokens(r.tokens)} of ${formatTokens(r.window)} tokens (${r.tokens.toLocaleString('en-US')} / ${r.window.toLocaleString('en-US')}).${hint}${state}`
48}
49
50export const register: Register = (on, options) => {
51  const warnAt = warnLevel((options as { warnAt?: unknown } | undefined)?.warnAt)
52
53  on('session.start', async ($, e, next) => {
54    await $.command.register({
55      name: 'gauge',
56      description: 'Context gauge: show context-window usage, or hide | show | toggle the band above the prompt',
57      argumentHint: '[hide | show | toggle]',
58      immediate: true,
59    })
60    const saved = await $.store.get('isHidden')
61    if (typeof saved === 'boolean') await update($, isHidden, () => saved)
62    await refresh($, warnAt)
63
64    return next(e)
65  }).catch(($, e, next) => next(e))
66
67  // Pushed by the engine after each main-thread turn: the cheapest source.
68  on('session.measure', async ($, e, next) => {
69    if (e.changed.includes('context')) await record($, e.context, warnAt)
70    return next(e)
71  }).catch(($, e, next) => next(e))
72
73  // Each main-loop step answered: refresh mid-turn so the bar moves during long turns.
74  on('turn.step', async function* ($, e, next) {
75    const result = yield* next(e)
76    if (e.agentId === undefined) await refresh($, warnAt)
77    return result
78  })
79
80  on('turn.complete', async ($, e, next) => {
81    const result = await next(e)
82    if (e.agentId === undefined) await refresh($, warnAt)
83    return result
84  }).catch(($, e, next) => next(e))
85
86  on('command.run', { command: 'gauge' }, async ($, e) => {
87    const verb = e.args.trim().toLowerCase()
88    const hidden = await read($, isHidden)
89
90    if (verb === 'hide' || verb === 'show' || verb === 'toggle') {
91      const nextHidden = verb === 'toggle' ? !hidden : verb === 'hide'
92      await setHidden($, nextHidden)
93      return { text: `Context gauge: band ${nextHidden ? 'hidden' : 'shown'}.` }
94    }
95
96    await refresh($, warnAt)
97    return { text: describe(await read($, reading), hidden, warnAt) }
98  })
99
100  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
101    if (e.props.hasSurvey || (await read($, isHidden))) return next(e)
102    const r = await read($, reading)
103    if (r === null) return next(e)
104
105    const { Box, Text } = $.ui.resolve(e)
106    const isWarn = r.percent >= warnAt
107    const tail = label(r.percent, r.tokens, r.window)
108    const hint = isWarn ? HINT : ''
109    const width = barWidth(e.props.bodyColumns, PREFIX.length + tail.length + hint.length)
110
111    return (
112      <Box>
113        <Text dimColor>{PREFIX}</Text>
114        {width > 0 && <Text dimColor>▕</Text>}
115        {segments(r.percent, width, warnAt).map(s =>
116          s.color === null ? <Text dimColor>{s.text}</Text> : <Text color={s.color}>{s.text}</Text>,
117        )}
118        {width > 0 && <Text dimColor>▏</Text>}
119        <Text bold color={colorFor(r.percent, warnAt)}>{` ${r.percent}%`}</Text>
120        <Text dimColor>{tail.slice(` ${r.percent}%`.length)}</Text>
121        {isWarn && <Text color={colorFor(r.percent, warnAt)}>{hint}</Text>}
122      </Box>
123    )
124  })
125}
126

1// Pure helpers for the context gauge: no engine calls, so tests can hit them directly.
2
3export const GREEN = '#22c55e'
4export const ORANGE = '#f59e0b'
5export const RED = '#ef4444'
6
7/** Where the bar starts turning orange. */
8export const ORANGE_AT = 50
9/** The second, louder toast. */
10export const CRITICAL_AT = 90
11
12const MIN_BAR = 4
13const MAX_BAR = 40
14
15/** The configured warning percentage, defaulting to 80 and kept within 1..100. */
16export function warnLevel(raw: unknown): number {
17  const n = Number(raw)
18  if (!Number.isFinite(n) || n <= 0) return 80
19  return Math.min(100, Math.round(n))
20}
21
22/** green below 50%, orange below `warnAt`, red from it. */
23export function colorFor(percent: number, warnAt: number): string {
24  if (percent >= warnAt) return RED
25  if (percent >= Math.min(ORANGE_AT, warnAt)) return ORANGE
26  return GREEN
27}
28
29/** 620000 → "620k", 1000000 → "1M", 1500000 → "1.5M", 950 → "950". */
30export function formatTokens(n: number): string {
31  const trim = (x: number) => (Math.round(x * 10) / 10).toString()
32  if (n >= 1_000_000) return `${trim(n / 1_000_000)}M`
33  if (n >= 1_000) return `${n >= 100_000 ? Math.round(n / 1_000) : trim(n / 1_000)}k`
34  return String(Math.round(n))
35}
36
37/** The text after the bar: " 62% · 620k / 1M", plus " · /compact soon" from `warnAt`. */
38export function label(percent: number, tokens: number, window: number): string {
39  return ` ${percent}% · ${formatTokens(tokens)} / ${formatTokens(window)}`
40}
41
42export const HINT = ' · /compact soon'
43
44/** Cells the bar gets: what is left of `bodyColumns` after the prefix, brackets and label; 0 when under 4. */
45export function barWidth(bodyColumns: number, textColumns: number): number {
46  const room = Math.min(MAX_BAR, bodyColumns - textColumns - 2)
47  return room < MIN_BAR ? 0 : room
48}
49
50export type Segment = { text: string; color: string | null }
51
52/**
53 * The bar as runs of same-coloured cells: each filled cell is shaded by its own
54 * position (so a fuller bar reads green → orange → red), empty cells are `null`.
55 */
56export function segments(percent: number, width: number, warnAt: number): Segment[] {
57  if (width <= 0) return []
58  const filled = Math.max(0, Math.min(width, Math.round((percent / 100) * width)))
59  const out: Segment[] = []
60  for (let i = 0; i < width; i++) {
61    const isFilled = i < filled
62    const color = isFilled ? colorFor(((i + 0.5) / width) * 100, warnAt) : null
63    const last = out[out.length - 1]
64    const ch = isFilled ? '█' : '░'
65    if (last && last.color === color) last.text += ch
66    else out.push({ text: ch, color })
67  }
68  return out
69}
70
71/** The thresholds that warn: `warnAt` and 90, deduplicated and ascending. */
72export function thresholds(warnAt: number): number[] {
73  return [...new Set([warnAt, CRITICAL_AT])].filter(t => t > 0 && t <= 100).sort((a, b) => a - b)
74}
75
76/**
77 * The highest threshold the fill crossed going up from `prev` to `now`, or null.
78 * A drop (a compaction) re-arms the thresholds it falls below.
79 */
80export function crossed(prev: number | null, now: number, warnAt: number): number | null {
81  const before = prev ?? 0
82  const hit = thresholds(warnAt).filter(t => before < t && now >= t)
83  return hit.length ? hit[hit.length - 1]! : null
84}
85
86export function toastText(percent: number, threshold: number): string {
87  return threshold >= CRITICAL_AT
88    ? `Context ${percent}% full: run /compact now or the session will auto-compact soon`
89    : `Context ${percent}% full: consider /compact soon`
90}
91

1export type Reading = {
2  /** Input tokens the last main-loop response was answered over. */
3  tokens: number
4  /** The session model's context window, in tokens. */
5  window: number
6  /** tokens / window as a whole percentage. */
7  percent: number
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'context-gauge': {
13      /** The latest context reading; null before the first response or right after a compaction. */
14      reading: Reading | null
15      /** Whether the band is hidden (/gauge hide). */
16      isHidden: boolean
17    }
18  }
19}
20