token-weather

1// Token Weather: a live forecast of the context window, above the prompt,
2// with a bar and legend showing which elements fill it, in /context's colours.
3import { atom, read, update } from 'claude-code'
4import type { EngineInterface, Register } from 'claude-code'
5
6import { added, barCells, barWidth, forecast, freeTokens, isBuffer, push, bufferText, freeText, segmentText, share, short, toReading, toSegments } from './weather'
7
8// Held by the host, so the history survives a hot reload of this file.
9// A theme key, so the track follows dark and light themes; it is no element's /context colour.
10const FREE_TRACK = 'userMessageBackground'
11
12const readings = atom({ plugin: 'token-weather', key: 'readings' } as const, [])
13const segments = atom({ plugin: 'token-weather', key: 'segments' } as const, [])
14
15// One measurement: the window's fill plus /context's breakdown by element.
16async function sample($: EngineInterface) {
17  const { context } = await $.session.usage({ breakdown: 'summary' })
18  const r = toReading(context)
19  if (r) await update($, readings, h => push(h, r))
20  const rows = context.breakdown?.categories
21  if (rows) await update($, segments, () => toSegments(rows))
22}
23
24export const register: Register = on => {
25  on('session.start', async ($, e, next) => {
26    const result = await next(e)
27    await sample($)
28    return result
29  })
30
31  // Pushed after each main-thread turn, so subagent turns never land here.
32  on('session.measure', async ($, e, next) => {
33    const result = await next(e)
34    await sample($)
35    return result
36  })
37
38  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
39    if (e.props.hasSurvey) return next(e)
40
41    let history = await read($, readings)
42    if (history.length === 0) {
43      // Nothing stored yet (first render after a reload): draw from a live reading.
44      const r = toReading((await $.session.usage()).context)
45      if (!r) return next(e)
46      history = [r]
47    }
48
49    const now = history[history.length - 1]!
50    const f = forecast(now.percent)
51    const columns = e.props.bodyColumns
52    const wide = columns >= 60
53    const grew = added(history)
54    const { Box, Text } = $.ui.resolve(e)
55
56    const forecastRow = [
57      <Text color={f.color} bold>{`${f.icon}  ${f.word}`}</Text>,
58      <Text>{`  ${now.percent}% of context`}</Text>,
59      <Text dimColor>{`  ${short(now.tokens)} / ${short(now.window)}`}</Text>,
60    ]
61    if (wide && grew !== null && grew > 0) {
62      forecastRow.push(<Text dimColor>{`  +${short(grew)} last turn`}</Text>)
63    }
64
65    const stored = await read($, segments)
66    const parts = stored.filter(s => !isBuffer(s))
67    const reserve = stored.filter(isBuffer)
68    if (parts.length === 0) return <Box>{forecastRow}</Box>
69
70    // The bar spans the band's width (one cell short, so it never wraps), laid
71    // out as /context lays it out: each element's share of the window in its
72    // colour, then free space, then the autocompact buffer at the far end.
73    // Each segment carries its percentage where it fits.
74    const width = barWidth(columns)
75    const cells = barCells([...parts, ...reserve], now.window, width)
76    const filled = cells.reduce((a, b) => a + b, 0)
77    const reserveTokens = reserve.reduce((a, s) => a + s.tokens, 0)
78    const free = freeTokens(now.window, now.tokens, reserveTokens)
79    const bar = [
80      ...parts.map((p, i) => (
81        // inverse: the element's /context colour fills the cells and the label takes
82        // the terminal's own background colour, which contrasts in dark and light themes.
83        <Text color={p.color} inverse bold>
84          {segmentText(cells[i] ?? 0, share(p.tokens, now.window))}
85        </Text>
86      )),
87      // Free space: an unbroken track in the theme's user-message background, the room left
88      // before autocompact written into it as the elements' percentages are. A dim inverse fill
89      // drew the same grey as grey elements (System tools) and the buffer; a pattern or no fill
90      // left the label cutting the line.
91      <Text backgroundColor={FREE_TRACK}>{freeText(Math.max(0, width - filled), free, now.window)}</Text>,
92      ...reserve.map((b, i) => (
93        <Text color={b.color} inverse>
94          {bufferText(cells[parts.length + i] ?? 0)}
95        </Text>
96      )),
97    ]
98
99    // The legend: largest elements first; fewer of them on a narrow band.
100    const shown = wide ? parts : parts.slice(0, 3)
101    const legend = shown.map(p => (
102      <Text color={p.color}>{`■ ${p.name} ${short(p.tokens)}  `}</Text>
103    ))
104
105    return (
106      <Box flexDirection="column">
107        <Box>{forecastRow}</Box>
108        <Box>{bar}</Box>
109        <Box flexWrap="wrap">{legend}</Box>
110      </Box>
111    )
112  })
113}
114

1// Pure forecasting helpers: no `$`, so the tests can call them directly.
2import type { Reading, Segment } from '../types'
3
4/**
5 * The breakdown's `used` rows, largest first, then its `buffer` rows, each
6 * keeping the theme colour /context draws it in. Free space is derived, not
7 * stored, and deferred schemas sit outside the window, so both are left out.
8 */
9export function toSegments(
10  categories: readonly { name: string; tokens: number; kind: string; color: string }[],
11): Segment[] {
12  const rows = (kind: 'used' | 'buffer') =>
13    categories
14      .filter(c => c.kind === kind && c.tokens > 0)
15      .map(c => ({ name: c.name, tokens: c.tokens, color: c.color, kind }))
16  return [...rows('used').sort((a, b) => b.tokens - a.tokens), ...rows('buffer')]
17}
18
19export function isBuffer(s: Segment): boolean {
20  return s.kind === 'buffer'
21}
22
23/** An element's share of the window as a whole percentage; "<1%" for a sliver. */
24export function share(tokens: number, window: number): string {
25  if (window <= 0) return '0%'
26  const p = (tokens / window) * 100
27  return p > 0 && p < 1 ? '<1%' : `${Math.round(p)}%`
28}
29
30/**
31 * One segment of the bar, exactly `cells` wide. It is drawn as a filled
32 * background, so the text is spaces with the percentage centred in it when it
33 * fits with a cell of margin either side, and spaces alone when it does not.
34 */
35export function segmentText(cells: number, percent: string): string {
36  const n = Math.max(0, cells)
37  if (n < percent.length + 2) return ' '.repeat(n)
38  const left = Math.floor((n - percent.length) / 2)
39  return ' '.repeat(left) + percent + ' '.repeat(n - left - percent.length)
40}
41
42/**
43 * A segment's text: the first candidate that fits with a cell of margin either
44 * side, centred; blank when none fits. Always exactly `cells` wide.
45 */
46export function fitText(cells: number, candidates: readonly string[]): string {
47  const n = Math.max(0, cells)
48  return segmentText(n, candidates.find(c => n >= c.length + 2) ?? '')
49}
50
51/**
52 * The free segment: the room left before autocompact fires, longest wording
53 * that fits ("19% left before autocompact", "19% left", "19%"), blank otherwise.
54 */
55export function freeText(cells: number, tokens: number, window: number): string {
56  const pct = share(Math.max(0, tokens), window)
57  return fitText(cells, [`${pct} left before autocompact`, `${pct} left`, pct])
58}
59
60/**
61 * The autocompact buffer: a fixed reserve, so its size says nothing about the
62 * session. It is named, not measured, and left blank when the name does not fit.
63 */
64export function bufferText(cells: number): string {
65  return fitText(cells, ['autocompact'])
66}
67
68/** Window left once the used content and the autocompact buffer are taken out. */
69export function freeTokens(window: number, used: number, buffer: number): number {
70  return Math.max(0, window - used - buffer)
71}
72
73/** The bar's cells for a band `columns` wide: the full width less one, at least 1. */
74export function barWidth(columns: number): number {
75  return Math.max(1, Math.floor(columns) - 1)
76}
77
78/**
79 * Cells of a `width`-wide bar each segment fills, scaled to the window. Every
80 * non-empty segment gets at least one cell while room remains, so a small
81 * element still shows its colour; the rest of the bar is free space.
82 */
83export function barCells(segments: readonly Segment[], window: number, width: number): number[] {
84  if (window <= 0 || width <= 0) return segments.map(() => 0)
85  let left = width
86  return segments.map(s => {
87    const want = Math.max(1, Math.round((s.tokens / window) * width))
88    const cells = Math.min(left, want)
89    left -= cells
90    return cells
91  })
92}
93
94export const HISTORY = 12
95
96export type Forecast = { upTo: number; icon: string; word: string; color: string }
97
98export const FORECAST: readonly Forecast[] = [
99  { upTo: 25, icon: '☀', word: 'Clear', color: 'yellow' },
100  { upTo: 50, icon: '☁', word: 'Cloudy', color: 'cyan' },
101  { upTo: 75, icon: '☂', word: 'Showers', color: 'blue' },
102  { upTo: 90, icon: '☇', word: 'Storm', color: 'magenta' },
103  { upTo: Infinity, icon: '↯', word: 'Compact soon', color: 'red' },
104]
105
106export function forecast(percent: number): Forecast {
107  return FORECAST.find(f => percent < f.upTo) ?? FORECAST[FORECAST.length - 1]!
108}
109
110/** A reading from the engine's context figures, or null before the window is known. */
111export function toReading(context: { tokens?: number; window: number; percent?: number }): Reading | null {
112  if (!context.window) return null
113  const tokens = context.tokens ?? 0
114  const percent = context.percent ?? Math.round((tokens / context.window) * 100)
115  return { tokens, window: context.window, percent }
116}
117
118/** Append a reading, skipping an unchanged repeat, keeping the last HISTORY. */
119export function push(history: readonly Reading[], r: Reading): Reading[] {
120  const last = history[history.length - 1]
121  if (last && last.tokens === r.tokens && last.window === r.window) return [...history]
122  return [...history, r].slice(-HISTORY)
123}
124
125/** Tokens the last turn added, or null with fewer than two readings. */
126export function added(history: readonly Reading[]): number | null {
127  const now = history[history.length - 1]
128  const before = history[history.length - 2]
129  return now && before ? now.tokens - before.tokens : null
130}
131
132export function short(n: number): string {
133  if (Math.abs(n) >= 1_000_000) return `${(n / 1_000_000).toFixed(1).replace(/\.0$/, '')}M`
134  if (Math.abs(n) >= 1_000) return `${Math.round(n / 1_000)}k`
135  return String(n)
136}
137

1/** One reading of the context window, taken after a main-thread turn. */
2export type Reading = { tokens: number; window: number; percent: number }
3
4/**
5 * One row of /context's breakdown that the bar draws: an element occupying the
6 * window (`used`) or the autocompact reserve (`buffer`), with the theme colour
7 * /context draws it in. A segment stored before `kind` existed reads as `used`.
8 */
9export type Segment = { name: string; tokens: number; color: string; kind?: 'used' | 'buffer' }
10
11declare module 'claude-code' {
12  interface PluginState {
13    'token-weather': { readings: Reading[]; segments: Segment[] }
14  }
15}
16