context-gauge

1import { atom, read, update } from 'claude-code'
2import type { Register, SessionContextUsage } from 'claude-code'
3
4import type { Reading } from '../types'
5
6const reading = atom({ plugin: 'context-gauge', key: 'reading' } as const, null)
7
8const CELLS = 6
9const GAP = '  '
10
11/** 124000 → "124k", 1000000 → "1M". */
12function short(n: number): string {
13  if (n >= 999_500) return `${+(n / 1_000_000).toFixed(1)}M`
14  if (n >= 1000) return `${Math.round(n / 1000)}k`
15  return String(n)
16}
17
18function gauge({ tokens, window }: Reading): { text: string; percent: number } {
19  const percent = Math.round((tokens / window) * 100)
20  const filled = Math.min(CELLS, Math.round((percent / 100) * CELLS))
21  const bar = '▰'.repeat(filled) + '▱'.repeat(CELLS - filled)
22
23  return { text: `ctx ${bar} ${percent}% · ${short(tokens)}/${short(window)}`, percent }
24}
25
26function colorOf(percent: number): 'warning' | 'error' | undefined {
27  if (percent > 85) return 'error'
28
29  return percent >= 60 ? 'warning' : undefined
30}
31
32/** `tokens` is absent until the live window's first response (a fresh or just compacted session). */
33function readingOf({ tokens, window }: SessionContextUsage, fallbackWindow: number): Reading | null {
34  return tokens === undefined ? null : { tokens, window: window > 0 ? window : fallbackWindow }
35}
36
37export const register: Register = (on, options) => {
38  const fallbackWindow =
39    typeof options.windowTokens === 'number' && options.windowTokens > 0 ? options.windowTokens : 200_000
40
41  // Also raised on an enable or a reload mid-session: the gauge shows at once.
42  on('session.start', async ($, e, next) => {
43    const started = await next(e)
44    const { context } = await $.session.usage()
45    await update($, reading, () => readingOf(context, fallbackWindow))
46
47    return started
48  })
49
50  // The status line's figures, pushed after each main-thread turn.
51  on('session.measure', async ($, e, next) => {
52    if (e.changed.includes('context')) await update($, reading, () => readingOf(e.context, fallbackWindow))
53
54    return next(e)
55  })
56
57  // A main-thread compaction that went through makes the last reading stale until the next turn.
58  // An observer: whatever fails here, the compaction goes on.
59  on('session.compact', async ($, e, next) => {
60    const done = await next(e)
61    if (done.skip === undefined && e.trigger !== 'precompute' && e.agentId === undefined) {
62      await update($, reading, () => null)
63    }
64
65    return done
66  }).catch(($, e, next) => next(e))
67
68  // Only a /clear or a resume goes on in this process, and no session.start follows them.
69  // On an exit the write would only spend the chain's one short budget.
70  on('session.end', async ($, e, next) => {
71    if (e.reason === 'clear' || e.reason === 'resume') await update($, reading, () => null)
72
73    return next(e)
74  })
75
76  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
77    const now = await read($, reading)
78    if (now === null) return next(e)
79
80    const { text, percent } = gauge(now)
81    const color = colorOf(percent)
82
83    // The terminal draws `tail` after its own line and keeps its pills live, but only dim:
84    // a color, and the desktop (which draws no tail), take a tree of our own.
85    if (e.surface === 'terminal' && color === undefined) {
86      return next({ ...e, props: { ...e.props, tail: `${e.props.tail ?? ''}${GAP}${text}` } })
87    }
88
89    const { Box, Text } = $.ui.resolve(e)
90    const line = [e.props.hint, e.props.tail?.trim()].filter(Boolean).join(GAP)
91
92    return (
93      <Box flexDirection="row" columnGap={2}>
94        {line !== '' && (
95          <Text dimColor wrap="truncate-end">
96            {line}
97          </Text>
98        )}
99        <Box flexShrink={0}>
100          <Text color={color} dimColor={color === undefined}>
101            {text}
102          </Text>
103        </Box>
104      </Box>
105    )
106  })
107}
108

1/** The live context window's fill, as the last measurement reported it. */
2export type Reading = { tokens: number; window: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'context-gauge': { reading: Reading | null }
7  }
8}
9