token-watch

1// Token Watch: plan usage (the 5-hour and weekly windows, and any other window
2// Claude Code reports) and the context window as thin bars above the prompt, one colour per category the way /context
3// breaks it down. A click anywhere on the band collapses it to one row;
4// /context-bar shows or hides it.
5//
6// session.start: register /context-bar, restore the saved choices and the last
7//   usage windows seen (a new session has none until its first response), take a
8//   reading, and tick a clock every 30s so the reset countdowns stay current.
9// session.measure: after each main-thread turn, or when a usage window moves,
10//   take a reading.
11// command.run (context-bar): show or hide the band.
12// ui.render (AbovePrompt): the band, drawn by the Client in ./band.tsx, which
13//   lays itself out to the room it is given and reports clicks.
14// ui.message: the band was clicked; flip collapsed.
15//
16// A reading is $.session.usage({ breakdown: 'summary' }): the /context rows,
17// estimated locally, so it costs no token-count request, and the plan's
18// rate-limit windows as the status line has them (none off a subscription).
19
20import { atom, read, update } from 'claude-code'
21import type { Register } from 'claude-code'
22
23import type { BandProps, Limit, Reading, Slice } from '../types'
24
25const reading = atom({ plugin: 'token-watch', key: 'reading' } as const, null)
26const isShown = atom({ plugin: 'token-watch', key: 'isShown' } as const, true)
27const isCollapsed = atom({ plugin: 'token-watch', key: 'isCollapsed' } as const, false)
28const limits = atom({ plugin: 'token-watch', key: 'limits' } as const, [])
29const nowMs = atom({ plugin: 'token-watch', key: 'nowMs' } as const, 0)
30
31// /context's row names, shortened for the legend.
32const LABELS: Record<string, string> = {
33  'system tools': 'tools',
34  'custom agents': 'agents',
35  'free space': 'free',
36  'autocompact buffer': 'buffer',
37}
38
39const COLORS: Record<string, string> = {
40  'system prompt': '#7b9cd8',
41  tools: '#7ecfc4',
42  'mcp tools': '#a78bfa',
43  agents: '#8fd18f',
44  'memory files': '#e8c66a',
45  skills: '#e89bb8',
46  messages: '#d97757',
47}
48
49export const register: Register = on => {
50  on('session.start', async ($, e, next) => {
51    const result = await next(e)
52    await $.command.register({
53      name: 'context-bar',
54      description: 'Show or hide the Token Watch band (plan usage and context) above the prompt',
55    })
56    const shown = await $.store.get('isShown')
57    const collapsed = await $.store.get('isCollapsed')
58    await update($, isShown, () => shown !== false)
59    await update($, isCollapsed, () => collapsed === true)
60    await tick($)
61    // A new session has no usage windows until its first response, so start
62    // from the last ones seen; any whose reset time has passed is drawn at 0%,
63    // and the first response replaces them.
64    const saved = await $.store.get('limits')
65    if (Array.isArray(saved)) {
66      await update($, limits, () => saved as Limit[])
67    }
68    await takeReading($)
69    $.clock.every(30_000, () => void tick($))
70
71    return result
72  })
73
74  on('session.measure', async ($, e, next) => {
75    const result = await next(e)
76    if (await read($, isShown)) {
77      await tick($)
78      await takeReading($)
79    }
80
81    return result
82  })
83
84  on('command.run', { command: 'context-bar' }, async $ => {
85    const shown = !(await read($, isShown))
86    await update($, isShown, () => shown)
87    await $.store.set('isShown', shown)
88    if (shown) {
89      await takeReading($)
90    }
91
92    return { text: shown ? 'Token Watch shown.' : 'Token Watch hidden.' }
93  })
94
95  on('ui.message', async ($, e, next) => {
96    if ((e.data as { toggle?: boolean } | null)?.toggle !== true) {
97      return next(e)
98    }
99    const collapsed = !(await read($, isCollapsed))
100    await update($, isCollapsed, () => collapsed)
101    await $.store.set('isCollapsed', collapsed)
102
103    return {}
104  })
105
106  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
107    const now = await read($, reading)
108    if (e.props.hasSurvey || now === null || !(await read($, isShown))) {
109      return next(e)
110    }
111
112    // The band holds one drawing, so whatever the plugins beneath draw there
113    // (another mod's hint) stacks above this band instead of being replaced.
114    const below = await next(e)
115    const { Box, Client } = $.ui.resolve(e)
116    const clock = await read($, nowMs)
117    const props: BandProps = {
118      reading: now,
119      limits: current(await read($, limits), clock),
120      nowMs: clock,
121      isCollapsed: await read($, isCollapsed),
122      surface: e.surface,
123    }
124
125    return (
126      <Box flexDirection="column">
127        {below}
128        <Client key="band" module="./band.tsx" props={props} width="100%" />
129      </Box>
130    )
131  })
132}
133
134async function takeReading($) {
135  try {
136    const { context, rateLimits } = await $.session.usage({ breakdown: 'summary' })
137    if (rateLimits.length > 0) {
138      const reported: Limit[] = rateLimits.map(({ kind, percentUsed, resetsAt }) => ({ kind, percentUsed, resetsAt }))
139      // A window the reading leaves out keeps its last values until its reset
140      // time; one still left out after that is no longer part of the plan.
141      const now = await $.clock.now()
142      const kept = (await read($, limits)).filter(
143        w => !reported.some(r => r.kind === w.kind) && !!w.resetsAt && Date.parse(w.resetsAt) > now,
144      )
145      const windows = [...reported, ...kept]
146      await update($, limits, () => windows)
147      await $.store.set('limits', windows)
148    }
149    const breakdown = context.breakdown
150    if (!breakdown || !breakdown.rawMaxTokens) {
151      return
152    }
153    const slices: Slice[] = breakdown.categories
154      .filter(row => row.kind !== 'deferred' && row.tokens > 0)
155      .map(row => {
156        const name = row.name.toLowerCase()
157        const label = LABELS[name] ?? name
158        return {
159          label,
160          tokens: row.tokens,
161          color: COLORS[label] ?? row.color,
162          kind: row.kind as Slice['kind'],
163        }
164      })
165    const next: Reading = {
166      slices,
167      total: breakdown.totalTokens,
168      window: breakdown.rawMaxTokens,
169      compactsAt: breakdown.isAutoCompactEnabled ? (breakdown.autoCompactThreshold ?? null) : null,
170    }
171    await update($, reading, () => next)
172  } catch {
173    // No reading this time; the band keeps the last one.
174  }
175}
176
177async function tick($) {
178  const now = await $.clock.now()
179  await update($, nowMs, () => now)
180}
181
182// The windows as they stand now: one whose reset time has passed has started
183// over, so it is drawn at 0% until a response reports it again.
184function current(windows: Limit[], nowMs: number) {
185  return windows.map(w =>
186    !w.resetsAt || !nowMs || Date.parse(w.resetsAt) > nowMs ? w : { kind: w.kind, percentUsed: 0 },
187  )
188}
189

1// The band itself, drawn as a Client so the whole thing is one click target.
2//
3// Every bar is a row of Boxes that grow in proportion to what they measure, so
4// the layout engine fits it to whatever width the band has, on the terminal
5// and the desktop alike, instead of the module counting characters.
6//
7// One layout and one palette on every surface; only the leaf that paints a
8// coloured piece depends on the surface (see `pieces`). The desktop paints a
9// Box background gaplessly, where any run of glyphs would leave a sliver of
10// remainder per piece in its proportional font, so there each piece is an
11// empty Box 60% as tall as its line. A terminal can only shade whole cells,
12// and only a glyph shape gives it less than a cell of height, so there each
13// piece is a run of lower blocks clipped to its Box. Either way the bars are
14// thinner than their line and stacked bars keep a gap between them.
15//
16// Expanded: one aligned table in a rounded box. A row per usage window (5H, WK
17//   and any other window the API sends, such as a model's own allowance), then
18//   CTX, each as label, bar, percentage and a dim detail (the time to reset,
19//   or the tokens of the window). Every bar starts and ends in the same
20//   columns, over a track in one quiet colour for what is still available.
21//   On the context bar the reserve past the auto-compact point is a darker
22//   shade than the free space before it. Pieces are pure proportion, with no
23//   one-cell minimum, so a bar's fill matches the printed percentage. Under
24//   CTX, the legend, largest category first, with 2-cell swatches, wrapping
25//   onto more lines when it does not fit.
26// Collapsed: one line in the table's language. Each usage window as label, a
27//   10-cell meter and percentage; then CTX as label, the context bar filling
28//   the rest, percentage and "212k / 1M". Counted in cells against the width:
29//   the meters are drawn only while the context bar keeps 20 cells, the
30//   detail only while it keeps 8, and below that the bar takes what is left,
31//   so the line never wraps.
32// A left click anywhere posts { toggle: true } to the hooks module.
33
34import type { ClientModule } from 'claude-code'
35
36import type { BandProps, Limit, Reading } from '../types'
37
38type Local = { isHovered: boolean }
39
40type Row = { label: string; bar: unknown[]; percent: number; detail: string }
41
42// One palette for both views and both surfaces: what is still available (the
43// usage tracks, the context's free space), and the reserve past the
44// auto-compact point a step darker.
45const TRACK = '#3d4250'
46const RESERVE = '#2e3139'
47
48// What a terminal draws a piece with: a lower block, the one glyph that is
49// shorter than its cell.
50const GLYPH = '▆'
51
52// Usage windows by kind, in the order they are drawn; any other window the
53// API sends is drawn after them.
54const WINDOWS: Record<string, string> = {
55  five_hour: '5H',
56  seven_day: 'WK',
57  spend_limit: '$$',
58}
59const ORDER = Object.keys(WINDOWS)
60
61// Windows Claude Code does not list today but the API may send, such as a
62// model's own allowance: named by the model when the kind mentions one.
63const MODELS: [RegExp, string][] = [
64  [/fable/i, 'FB'],
65  [/opus/i, 'OP'],
66  [/sonnet/i, 'SN'],
67  [/haiku/i, 'HK'],
68]
69
70// The legend's shorter names for the longer category labels.
71const LEGEND: Record<string, string> = {
72  'system prompt': 'system',
73  'memory files': 'memory',
74  'mcp tools': 'mcp',
75}
76
77// Below this many columns the expanded table leaves out the detail column.
78const DETAIL_MIN_COLUMNS = 50
79
80// Cells of slack in each fixed column: the desktop draws text in a
81// proportional font, where "CTX" or "100%" can be wider than its characters.
82const SLACK = 2
83
84// The percentage column: "100%" and its slack.
85const PERCENT_WIDTH = 6
86
87// The collapsed line: the gap between its cells, the width of a usage
88// window's meter, the room the context bar keeps before the line gives up
89// the detail, and the room it keeps before the meters are drawn.
90const CELL_GAP = 3
91const METER_WIDTH = 10
92const BAR_MIN = 8
93const METERS_BAR_MIN = 20
94
95const Band: ClientModule<BandProps, Local> = (props, surface) => {
96  const { Box, Text } = surface.elements
97  const { reading: now, isCollapsed, nowMs } = props
98  const limits = sorted(props.limits)
99  const isHovered = surface.state?.isHovered ?? false
100
101  surface.onPointer(event => {
102    if (event.type === 'up' && event.button === 'left') {
103      surface.post({ toggle: true })
104    } else if (event.type === 'enter' && !isHovered) {
105      surface.setState({ isHovered: true })
106    } else if (event.type === 'leave' && isHovered) {
107      surface.setState({ isHovered: false })
108    }
109  })
110
111  const percent = Math.round((now.total / now.window) * 100)
112
113  const draw = pieces(Box, Text, props.surface, surface.columns)
114
115  if (isCollapsed) {
116    const detail = `${short(now.total)} / ${short(now.window)}`
117    const { showDetail, showMeters } = fit(surface.columns, limits, percent, detail)
118    return (
119      <Box flexDirection="row" alignItems="center" columnGap={CELL_GAP} paddingX={1}>
120        {limits.map(limit => (
121          <Box key={limit.kind} flexDirection="row" alignItems="center" columnGap={1} flexShrink={0}>
122            <Text dimColor={!isHovered} wrap="truncate">
123              {label(limit)}
124            </Text>
125            {showMeters && (
126              <Box width={METER_WIDTH} height={1} flexDirection="row" alignItems="center" flexShrink={0}>
127                {meter(draw, limit.percentUsed)}
128              </Box>
129            )}
130            <Text wrap="truncate" {...warning(limit.percentUsed)}>{`${Math.round(limit.percentUsed)}%`}</Text>
131          </Box>
132        ))}
133        <Box flexDirection="row" alignItems="center" columnGap={1} flexGrow={1}>
134          <Text dimColor={!isHovered} wrap="truncate">
135            CTX
136          </Text>
137          <Box flexDirection="row" flexGrow={1} minWidth={4} height={1} alignItems="center">
138            {bar(draw, now, TRACK, RESERVE)}
139          </Box>
140          <Text wrap="truncate" {...warning(percent)}>{`${percent}%`}</Text>
141          {showDetail && (
142            <Box marginLeft={1} flexShrink={0}>
143              <Text dimColor wrap="truncate">
144                {detail}
145              </Text>
146            </Box>
147          )}
148        </Box>
149      </Box>
150    )
151  }
152
153  const rows: Row[] = [
154    ...limits.map(limit => ({
155      label: label(limit),
156      bar: meter(draw, limit.percentUsed),
157      percent: limit.percentUsed,
158      detail: countdown(limit.resetsAt, nowMs),
159    })),
160    {
161      label: 'CTX',
162      bar: bar(draw, now, TRACK, RESERVE),
163      percent,
164      detail: `${short(now.total)} / ${short(now.window)}`,
165    },
166  ]
167  const labelWidth = Math.max(3, ...rows.map(row => row.label.length)) + SLACK
168  const showDetail = surface.columns >= DETAIL_MIN_COLUMNS
169  const detailWidth = Math.max(...rows.map(row => row.detail.length)) + SLACK
170  const legend = now.slices.filter(s => s.kind === 'used').sort((a, b) => b.tokens - a.tokens)
171  // A legend swatch: a 2-cell bar, drawn the way the bars are.
172  const swatch = (color: string) => (
173    <Box width={2} height={1} flexDirection="row" alignItems="center" flexShrink={0}>
174      {draw(color, { grow: 1 })}
175    </Box>
176  )
177
178  return (
179    <Box flexDirection="column" borderStyle="round" borderDimColor={!isHovered} paddingX={1}>
180      {rows.map(row => (
181        <Box key={row.label} flexDirection="row" alignItems="center" columnGap={2}>
182          <Box width={labelWidth} flexShrink={0}>
183            <Text dimColor wrap="truncate">
184              {row.label}
185            </Text>
186          </Box>
187          <Box flexDirection="row" flexGrow={1} minWidth={4} height={1} alignItems="center">
188            {row.bar}
189          </Box>
190          <Box width={PERCENT_WIDTH} flexShrink={0} justifyContent="flex-end">
191            <Text wrap="truncate" {...warning(row.percent)}>{`${Math.round(row.percent)}%`}</Text>
192          </Box>
193          {showDetail && (
194            <Box width={detailWidth} flexShrink={0}>
195              <Text dimColor wrap="truncate">
196                {row.detail}
197              </Text>
198            </Box>
199          )}
200        </Box>
201      ))}
202      <Box flexDirection="row" flexWrap="wrap" alignItems="center" columnGap={3} paddingLeft={labelWidth + 2}>
203        {legend.map(slice => (
204          <Box key={slice.label} flexDirection="row" alignItems="center" columnGap={1} flexShrink={0}>
205            {swatch(slice.color)}
206            <Text dimColor>{`${LEGEND[slice.label] ?? slice.label} ${short(slice.tokens)}`}</Text>
207          </Box>
208        ))}
209      </Box>
210    </Box>
211  )
212}
213
214export default Band
215
216// How a bar draws one piece: a colour over a share of the bar (`grow`) or a
217// fixed number of cells (`width`). No minimum size: the share is exact, so a
218// bar's fill matches the percentage printed beside it.
219type Size = { grow?: number; width?: number }
220type Draw = (color: string, size: Size) => unknown
221
222const box = (size: Size) => ({
223  width: size.width ?? 0,
224  flexGrow: size.grow ?? 0,
225  flexShrink: size.width ? 0 : 1,
226  minWidth: 0,
227})
228
229// A piece shorter than its line, for both views. The surface chooses
230// the leaf, and only here: on the desktop an empty Box 60% as tall as its row
231// with a background, which is gapless; in a terminal a Box one cell tall
232// holding a run of GLYPH, which wraps exactly at the Box's width and whose
233// overflow is clipped. (The run is longer than any piece can be, so it fills
234// the Box whatever share it has.)
235function pieces(Box, Text, surface: BandProps['surface'], columns: number): Draw {
236  if (surface === 'desktop') {
237    return (color, size) => <Box {...box(size)} height="60%" backgroundColor={color} />
238  }
239
240  return (color, size) => (
241    <Box {...box(size)} height={1} overflow="hidden">
242      <Text color={color} wrap="wrap">
243        {GLYPH.repeat(size.width ?? Math.max(columns, 80) + 2)}
244      </Text>
245    </Box>
246  )
247}
248
249// The context bar: a piece per category, each growing by its tokens from a
250// zero base, so the widths stay proportional at any size. Free space is drawn
251// in `free`, the reserve past the auto-compact point in `reserveColor`.
252function bar(draw: Draw, now: Reading, free: string, reserveColor: string) {
253  // Largest category first, as the legend lists them, then the free space.
254  const used = now.slices.filter(s => s.kind === 'used').sort((a, b) => b.tokens - a.tokens)
255  const segments = [...used, ...now.slices.filter(s => s.kind === 'free')].map(s =>
256    draw(s.kind === 'free' ? free : s.color, { grow: grow(s.tokens, now.window) }),
257  )
258  const reserve = now.compactsAt !== null ? Math.max(0, now.window - now.compactsAt) : 0
259  if (reserve > 0) {
260    segments.push(draw(reserveColor, { grow: grow(reserve, now.window) }))
261  }
262
263  return segments
264}
265
266// A usage window's meter: the used share in its level colour and what is
267// still available as track, both growing from a zero base like the context
268// bar.
269function meter(draw: Draw, percentUsed: number) {
270  const used = Math.max(0, Math.min(100, percentUsed))
271  return [draw(levelColor(used), { grow: used * 100 }), draw(TRACK, { grow: (100 - used) * 100 })]
272}
273
274// Which of the collapsed line's extras fit: the detail, then the usage
275// meters. Counted in cells, the bar's room being what is left after the
276// padding, labels, percentages and gaps; the desktop's proportional text
277// measures a little wider, which the bar absorbs.
278function fit(columns: number, limits: Limit[], percent: number, detail: string) {
279  const usage = limits.reduce(
280    (n, limit) => n + label(limit).length + 1 + `${Math.round(limit.percentUsed)}%`.length + CELL_GAP,
281    0,
282  )
283  const base = 2 + usage + 'CTX'.length + 1 + 1 + `${percent}%`.length
284  const withDetail = base + 2 + detail.length
285  const withMeters = withDetail + limits.length * (METER_WIDTH + 1)
286  const showDetail = columns - withDetail >= BAR_MIN
287  const showMeters = showDetail && columns - withMeters >= METERS_BAR_MIN
288  return { showDetail, showMeters }
289}
290
291function sorted(list: Limit[]) {
292  const rank = (kind: string) => (ORDER.includes(kind) ? ORDER.indexOf(kind) : ORDER.length)
293  return [...list].sort((a, b) => rank(a.kind) - rank(b.kind))
294}
295
296function label(limit: Limit) {
297  return WINDOWS[limit.kind] ?? MODELS.find(([pattern]) => pattern.test(limit.kind))?.[1] ?? limit.kind
298}
299
300// Time left until a window resets: "3d 4h", "2h 14m", "9m".
301function countdown(resetsAt: string | undefined, nowMs: number) {
302  if (!resetsAt || !nowMs) return ''
303  const ms = Date.parse(resetsAt) - nowMs
304  if (Number.isNaN(ms)) return ''
305  if (ms <= 0) return 'now'
306  const minutes = Math.floor(ms / 60_000)
307  const days = Math.floor(minutes / 1_440)
308  const hours = Math.floor((minutes % 1_440) / 60)
309  if (days > 0) return `${days}d ${hours}h`
310  if (hours > 0) return `${hours}h ${String(minutes % 60).padStart(2, '0')}m`
311  return `${Math.max(1, minutes)}m`
312}
313
314// A share of the window in parts per 10,000, the most flexGrow takes.
315function grow(tokens: number, window: number) {
316  return Math.min(10_000, Math.round((tokens / window) * 10_000 * 100) / 100)
317}
318
319function levelColor(percent: number) {
320  if (percent < 50) return '#8fd18f'
321  if (percent < 80) return '#e8c66a'
322  return '#e06c6c'
323}
324
325// The percentages in both views: plain text, yellow from 50%, red from 80%.
326function warning(percent: number) {
327  return percent < 50 ? {} : { color: levelColor(percent) }
328}
329
330function short(n: number) {
331  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
332  if (n >= 10_000) return `${Math.round(n / 1_000)}k`
333  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
334  return String(n)
335}
336

1// One category of the context window, as the bar and the legend draw it.
2export type Slice = {
3  label: string
4  tokens: number
5  color: string
6  kind: 'used' | 'free' | 'buffer'
7}
8
9// The last reading of the window, broken down the way /context breaks it down.
10export type Reading = {
11  slices: Slice[]
12  total: number
13  window: number
14  compactsAt: number | null
15}
16
17// One plan usage window (five_hour, seven_day), as the status line has it.
18export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
19
20// What the hooks module hands the band's Client.
21export type BandProps = {
22  reading: Reading
23  limits: Limit[]
24  nowMs: number
25  isCollapsed: boolean
26  surface: 'terminal' | 'desktop' | 'mobile' | 'vscode'
27}
28
29declare module 'claude-code' {
30  interface PluginState {
31    'token-watch': {
32      reading: Reading | null
33      limits: Limit[]
34      nowMs: number
35      isShown: boolean
36      isCollapsed: boolean
37    }
38  }
39}
40