token-weather

1// Token Weather: a live forecast of the context window, above the prompt.
2import { atom, read, update } from 'claude-code'
3import type { EngineInterface, Register, SessionContextUsage, SessionRateLimit, Timer } from 'claude-code'
4
5import type { TokenWeatherLimit, TokenWeatherReading } from '../types'
6import { DAY, WEEK_HEIGHT, WEEK_WIDTH, dial, elapsed, glyph, trend, trendWidth, week } from './sky'
7
8// Readings kept: one more than the bars, which are the gains between them.
9const HISTORY = 7
10const BARS = '▁▂▃▄▅▆▇█'
11// `color` is the theme's name for the text; `hex` the same hue in the band's drawings.
12const FORECAST = [
13  { upTo: 25, icon: '☀', sky: 'sun', word: 'Clear', color: 'yellow', hex: '#f6d909' },
14  { upTo: 50, icon: '☁', sky: 'cloud', word: 'Cloudy', color: 'cyan', hex: '#08a2c3' },
15  { upTo: 75, icon: '☂', sky: 'rain', word: 'Showers', color: 'blue', hex: '#5b8def' },
16  { upTo: 90, icon: '☇', sky: 'storm', word: 'Storm', color: 'magenta', hex: '#b77ce0' },
17  { upTo: Infinity, icon: '↯', sky: 'funnel', word: 'Compact soon', color: 'red', hex: '#e5484d' },
18] as const
19const WEEKDAYS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat']
20const ORANGE = '#f0883e'
21// The plan's rate-limit windows to show, in order, with their labels.
22const LIMIT_LABELS: Record<string, string> = { five_hour: '5h', seven_day: '7d' }
23// Windows whose reset countdown is drawn after the percentage.
24const SHOW_RESET = new Set(['five_hour'])
25// How often the band redraws so the countdown keeps time between turns.
26const TICK_MS = 30_000
27// Width of one text cell on the desktop, to fit the band to `bodyColumns`.
28const CELL_PX = 8
29
30// The band is one line of groups; one that does not fit leaves a dot in its color.
31const GAP = '  '
32const DOT = '•'
33// A terminal group is a run of spans, its first span's color its dot's.
34type Span = { text: string; color?: string; bold?: boolean; dimColor?: boolean }
35// A desktop group is a drawing and the text beside it.
36type Piece = { text: string; style?: { color?: string; bold?: boolean; dimColor?: boolean } }
37type Group = {
38  key: string
39  color: string
40  svg: { source: string; alt: string; width: number; height: number }
41  pieces: Piece[]
42}
43
44// Held by the host, so they survive a hot reload of this file: one reading a turn, the
45// quota windows, and the context fill as of the latest response.
46const readings = atom({ plugin: 'token-weather', key: 'readings' } as const, [] as TokenWeatherReading[])
47const limits = atom({ plugin: 'token-weather', key: 'limits' } as const, [] as TokenWeatherLimit[])
48const live = atom({ plugin: 'token-weather', key: 'live' } as const, null as TokenWeatherReading | null)
49
50export const register: Register = on => {
51  let tick: Timer | undefined
52
53  on('session.start', async ($, e, next) => {
54    const result = await next(e)
55    await startReading($)
56    tick?.cancel()
57    tick = $.clock.every(TICK_MS, () => $.ui.invalidate('ui.render'))
58
59    return result
60  })
61
62  // Every response of the turn, not just its end, moves the figures.
63  on('session.measure', async ($, e, next) => {
64    const result = await next(e)
65    await measure($, e.context, e.rateLimits)
66
67    return result
68  })
69
70  on('turn.complete', async ($, e, next) => {
71    const result = await next(e)
72    // Main-loop turns only, not subagents.
73    if (e.agentId === undefined) {
74      await takeReading($)
75    }
76
77    return result
78  })
79
80  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
81    const history = await read($, readings)
82    const now = (await read($, live)) ?? history.at(-1)
83    if (e.props.hasSurvey || !now) {
84      return next(e)
85    }
86
87    const f = FORECAST.find(b => now.percent < b.upTo) ?? FORECAST[4]
88    const quota = await read($, limits)
89    const at = await $.clock.now()
90
91    if (e.surface !== 'terminal') {
92      const { Box, Text, Svg } = $.ui.resolve(e)
93      const groups: Group[] = [
94        {
95          key: 'sky',
96          color: f.color,
97          svg: { source: glyph(f.sky, f.hex), alt: f.word, width: 16, height: 16 },
98          pieces: [
99            { text: f.word, style: { color: f.color, bold: true } },
100            { text: `${now.percent}%` },
101            { text: `${now.isEstimate ? '~' : ''}${short(now.tokens)} / ${short(now.window)}`, style: { dimColor: true } },
102          ],
103        },
104        ...quota.map(l => quotaGroup(l, at)),
105      ]
106      const { bars, label } = turns(history, now)
107      if (bars.length > 0) {
108        const alt = bars.length === 1 ? 'Context added in the last turn' : `Context added in each of the last ${bars.length} turns`
109        groups.push({
110          key: 'trend',
111          color: f.color,
112          svg: { source: trend(bars, f.hex), alt, width: trendWidth(bars.length), height: 14 },
113          pieces: [{ text: label, style: { dimColor: true } }],
114        })
115      }
116      const k = fit(groups.map(g => iconCells(g) + 1 + textCells(g)), e.props.bodyColumns - 2)
117      const hidden = groups.slice(k)
118
119      return (
120        <Box flexDirection="row" flexGrow={1} alignItems="center" justifyContent="space-between" paddingX={1} gap={2}>
121          {groups.slice(0, k).map(g => (
122            <Box key={g.key} flexDirection="row" alignItems="center" gap={1}>
123              <Svg {...g.svg} />
124              {g.pieces.map((p, i) => (
125                <Text key={String(i)} {...p.style}>
126                  {p.text}
127                </Text>
128              ))}
129            </Box>
130          ))}
131          {hidden.length > 0 && (
132            <Box key="dots" flexDirection="row">
133              {hidden.map(g => (
134                <Text key={g.key} color={g.color}>
135                  {DOT}
136                </Text>
137              ))}
138            </Box>
139          )}
140        </Box>
141      )
142    }
143
144    const { Box, Text } = $.ui.resolve(e)
145    const spans: Span[][] = [
146      [
147        { text: `${f.icon}  ${f.word}`, color: f.color, bold: true },
148        { text: ` ${now.percent}%` },
149        { text: ` ${now.isEstimate ? '~' : ''}${short(now.tokens)} / ${short(now.window)}`, dimColor: true },
150      ],
151      ...quota.map(l => {
152        const label: Span = { text: `${LIMIT_LABELS[l.kind]} ${l.percentUsed}%`, color: limitColor(l.percentUsed) }
153
154        return SHOW_RESET.has(l.kind) && l.resetsAt
155          ? [label, { text: ` Resets ${countdown(l.resetsAt, at)}`, dimColor: true }]
156          : [label]
157      }),
158    ]
159    const { bars, label } = turns(history, now)
160    if (bars.length > 0) {
161      spans.push([
162        { text: sparkline(bars), color: f.color },
163        { text: ` ${label}`, dimColor: true },
164      ])
165    }
166    const k = fit(
167      spans.map(group => group.reduce((n, span) => n + [...span.text].length, 0)),
168      e.props.bodyColumns - 2,
169    )
170    const shown = spans.slice(0, k)
171    const hidden = spans.slice(k)
172
173    return (
174      <Box flexDirection="row" paddingX={1}>
175        <Text wrap="truncate-end">
176          {shown.flatMap((group, g) => [
177            ...(g > 0 ? [<Text key={`gap-${g}`}>{GAP}</Text>] : []),
178            ...group.map((span, i) => (
179              <Text key={`${g}-${i}`} color={span.color} bold={span.bold} dimColor={span.dimColor}>
180                {span.text}
181              </Text>
182            )),
183          ])}
184          {hidden.length > 0 && <Text key="gap-dots">{GAP}</Text>}
185          {hidden.map((group, g) => (
186            <Text key={`dot-${g}`} color={group[0]?.color}>
187              {DOT}
188            </Text>
189          ))}
190        </Text>
191      </Box>
192    )
193  })
194}
195
196// At the start, or a reload: before the session's first response the context is
197// estimated locally; the quota waits for that response.
198async function startReading($: EngineInterface) {
199  const { context, rateLimits } = await $.session.usage()
200  const reading = (await measure($, context, rateLimits)) ?? (context?.window ? await estimate($, context.window) : null)
201  if (reading?.isEstimate) await update($, live, () => reading)
202  if (reading) await addReading($, reading)
203}
204
205// After a turn: the figures as it left them, and the turn's reading. A turn answered
206// with no quota at all is an API-key session, which shows none.
207async function takeReading($: EngineInterface) {
208  const { context, rateLimits } = await $.session.usage()
209  if (pickLimits(rateLimits).length === 0) await update($, limits, () => [])
210  const reading = await measure($, context, rateLimits)
211  if (reading) await addReading($, reading)
212}
213
214// Takes what a response reported: the quota and the context.
215async function measure($: EngineInterface, context: SessionContextUsage | undefined, rateLimits: SessionRateLimit[] | undefined) {
216  const shown = pickLimits(rateLimits)
217  if (shown.length > 0) await update($, limits, () => shown)
218  if (!context?.window || context.tokens === undefined) return null
219  const reading: TokenWeatherReading = {
220    tokens: context.tokens,
221    window: context.window,
222    percent: context.percent ?? Math.round((context.tokens / context.window) * 100),
223  }
224  await update($, live, () => reading)
225
226  return reading
227}
228
229// The windows the band shows, as the API reported them.
230function pickLimits(rateLimits: SessionRateLimit[] | undefined): TokenWeatherLimit[] {
231  return Object.keys(LIMIT_LABELS).flatMap(kind => {
232    const found = rateLimits?.find(l => l.kind === kind)
233
234    return found ? [{ kind, percentUsed: Math.round(found.percentUsed), resetsAt: found.resetsAt }] : []
235  })
236}
237
238// The context before any response, counted locally (free): system prompt, tools, memory.
239async function estimate($: EngineInterface, window: number): Promise<TokenWeatherReading | null> {
240  try {
241    const tokens = (await $.session.usage({ breakdown: 'summary' })).context?.breakdown?.totalTokens
242    if (!tokens) return null
243
244    return { tokens, window, percent: Math.round((tokens / window) * 100), isEstimate: true }
245  } catch {
246    return null
247  }
248}
249
250// One reading a turn; one the same as the last (a reload, a resume) is no turn of its own.
251async function addReading($: EngineInterface, reading: TokenWeatherReading) {
252  await update($, readings, history => {
253    const last = history.at(-1)
254    if (last && last.tokens === reading.tokens && last.window === reading.window) return history
255
256    return [...history, reading].slice(-HISTORY)
257  })
258}
259
260function quotaGroup(l: TokenWeatherLimit, at: number): Group {
261  const color = limitHex(l.percentUsed)
262  const label: Piece = { text: `${LIMIT_LABELS[l.kind]} ${l.percentUsed}%`, style: { color: limitColor(l.percentUsed) } }
263  if (l.kind === 'five_hour') {
264    const left = l.resetsAt ? countdown(l.resetsAt, at) : undefined
265    const alt =
266      `5-hour limit: ${l.percentUsed}% used` +
267      (left === undefined ? '' : left === 'now' ? ', resets now' : `, resets in ${left}`)
268
269    return {
270      key: l.kind,
271      color: label.style?.color ?? color,
272      svg: { source: dial({ used: l.percentUsed, now: at, color }), alt, width: 20, height: 20 },
273      pieces: left === undefined ? [label] : [label, { text: `Resets ${left}`, style: { dimColor: true } }],
274    }
275  }
276  const alt =
277    `Weekly limit: ${l.percentUsed}% used` +
278    (l.resetsAt
279      ? `, day ${Math.min(Math.floor(elapsed(l.resetsAt, 7 * DAY, at) * 7), 6) + 1} of 7, resets ${weekday(l.resetsAt)} ${clockTime(l.resetsAt)}`
280      : '')
281
282  return {
283    key: l.kind,
284    color: label.style?.color ?? color,
285    svg: { source: week({ used: l.percentUsed, color }), alt, width: WEEK_WIDTH, height: WEEK_HEIGHT },
286    pieces: [label],
287  }
288}
289
290// How many groups, from the left, fit one line of `columns` cells: GAP between groups,
291// and after the last a GAP and a dot for each group left out. Never fewer than one.
292function fit(widths: number[], columns: number) {
293  for (let k = widths.length; k > 1; k--) {
294    const hidden = widths.length - k
295    const cells =
296      widths.slice(0, k).reduce((n, w) => n + w, 0) + GAP.length * (k - 1) + (hidden > 0 ? GAP.length + hidden : 0)
297    if (cells <= columns) return k
298  }
299
300  return 1
301}
302
303function iconCells(g: Group) {
304  return Math.ceil(g.svg.width / CELL_PX)
305}
306
307function textCells(g: Group) {
308  return g.pieces.reduce((n, p) => n + p.text.length + 1, -1)
309}
310
311// Time left until an ISO timestamp: 2h 41m, 41m, <1m, or now once it has passed. Whole
312// minutes are floored, as the app's usage panel counts them.
313function countdown(iso: string, now: number) {
314  const ms = Date.parse(iso) - now
315  if (ms <= 0) return 'now'
316  const minutes = Math.floor(ms / 60_000)
317  if (minutes === 0) return '<1m'
318  const h = Math.floor(minutes / 60)
319  const m = minutes % 60
320  if (h === 0) return `${m}m`
321
322  return m === 0 ? `${h}h` : `${h}h ${m}m`
323}
324
325// An ISO timestamp as local HH:MM.
326function clockTime(iso: string) {
327  const d = new Date(iso)
328
329  return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}`
330}
331
332// An ISO timestamp's local weekday, as Tue.
333function weekday(iso: string) {
334  return WEEKDAYS[new Date(iso).getDay()]
335}
336
337// Green below 60%, orange below 90%, red from there. The theme has no orange name,
338// so that one is a hex on every surface.
339function limitColor(percentUsed: number) {
340  if (percentUsed >= 90) return 'red'
341  if (percentUsed >= 60) return ORANGE
342
343  return 'green'
344}
345
346// limitColor's hues as the desktop draws them, for the drawings.
347function limitHex(percentUsed: number) {
348  if (percentUsed >= 90) return '#e5484d'
349  if (percentUsed >= 60) return ORANGE
350
351  return '#46a759'
352}
353
354// What each turn added to the context; a compaction counts as nothing added.
355function gains(history: TokenWeatherReading[]) {
356  return history.slice(1).map((r, i) => Math.max(r.tokens - (history[i]?.tokens ?? r.tokens), 0))
357}
358
359function sparkline(values: number[]) {
360  const top = Math.max(...values, 1)
361
362  return values.map(v => BARS[Math.floor((v / top) * (BARS.length - 1))]).join('')
363}
364
365// The bars and their label. While a turn runs (the latest response is past the last
366// turn's reading) its gain so far is the newest bar: This turn; once it ends, Last turn.
367function turns(history: TokenWeatherReading[], now: TokenWeatherReading) {
368  const last = history.at(-1)
369  if (last && now.tokens !== last.tokens) {
370    const d = now.tokens - last.tokens
371
372    return { bars: [...gains(history), Math.max(d, 0)].slice(1 - HISTORY), label: `This turn ${arrow(d)}` }
373  }
374  const d = (history.at(-1)?.tokens ?? 0) - (history.at(-2)?.tokens ?? 0)
375
376  return { bars: gains(history), label: `Last turn ${arrow(d)}` }
377}
378
379function arrow(d: number) {
380  if (d === 0) return '0'
381
382  return d > 0 ? `↑${short(d)}` : `↓${short(-d)}`
383}
384
385function short(n: number) {
386  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
387  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
388
389  return String(n)
390}
391

1// The drawings the surfaces with Svg show: plain SVG markup built from the band's figures.
2
3export type Sky = 'sun' | 'cloud' | 'rain' | 'storm' | 'funnel'
4
5export const HOUR = 3_600_000
6export const DAY = 24 * HOUR
7export const WEEK_WIDTH = 54
8export const WEEK_HEIGHT = 9
9
10// Mid-tones that read on the desktop's dark card and on a light one.
11const TRACK = 'rgba(137,135,129,0.3)'
12const PAST = 'rgba(137,135,129,0.6)'
13const MARK = '#898781'
14// Deeper takes on the band's hues, for a light card where the scheme says so.
15const ON_LIGHT: Record<string, string> = {
16  '#f6d909': '#b39b00',
17  '#46a759': '#2f8a43',
18  '#e5484d': '#c9372f',
19  '#f0883e': '#c4661a',
20  '#08a2c3': '#06819b',
21  '#5b8def': '#3a6fd8',
22  '#b77ce0': '#9150c4',
23}
24
25const CLOUD = 'M17.5 19H9a7 7 0 1 1 6.71-9h1.79a4.5 4.5 0 1 1 0 9z'
26const RAISED = `<path d="${CLOUD}" transform="translate(0 -3)"/>`
27const SKY: Record<Sky, string> = {
28  sun:
29    '<circle cx="12" cy="12" r="4"/>' +
30    [0, 45, 90, 135, 180, 225, 270, 315].map(a => `<path d="M${at(a, 7.5)}L${at(a, 10)}"/>`).join(''),
31  cloud: `<path d="${CLOUD}"/>`,
32  rain: `${RAISED}<path d="M8 18.5l-1 3M12 18.5l-1 3M16 18.5l-1 3"/>`,
33  storm: `${RAISED}<path d="M13 16.5l-2.5 3.5h3.5l-2.5 3.5"/>`,
34  funnel: '<path d="M3 5h18M5 9h13M8 13h9M10 17h5M12 21h1.5"/>',
35}
36
37// The weather icon, stroked in the forecast's color.
38export function glyph(sky: Sky, color: string) {
39  return svg(
40    24,
41    24,
42    `class="k" fill="none" stroke="${color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"`,
43    light(color, 'stroke') + SKY[sky],
44  )
45}
46
47// A small clock: the rim fills clockwise from twelve as the quota is used, and the
48// two hands show the time and keep turning.
49export function dial(o: { used: number; now: number; color: string }) {
50  const minute = (new Date(o.now).getMinutes() + new Date(o.now).getSeconds() / 60) * 6
51
52  return svg(
53    24,
54    24,
55    'fill="none" stroke-linecap="round"',
56    light(o.color, 'stroke') +
57      `<circle cx="12" cy="12" r="10" stroke="${TRACK}" stroke-width="2.5"/>` +
58      rim(share(o.used) * 359.99, o.color) +
59      hand(face(o.now), 4.5, 43_200) +
60      hand(minute, 6.5, 3_600),
61  )
62}
63
64// Seven day cells filled as far as the quota is used.
65export function week(o: { used: number; color: string }) {
66  const filled = share(o.used) * 7
67  let cells = ''
68  for (let i = 0; i < 7; i++) {
69    const x = i * 8
70    const part = Math.min(Math.max(filled - i, 0), 1)
71    cells += `<rect x="${x}" y="0" width="6" height="9" rx="2" fill="${TRACK}"/>`
72    if (part > 0) {
73      cells += `<rect class="k" x="${x}" y="0" width="${+(6 * part).toFixed(2)}" height="9" rx="2" fill="${o.color}"/>`
74    }
75  }
76
77  return svg(WEEK_WIDTH, WEEK_HEIGHT, '', light(o.color, 'fill') + cells)
78}
79
80// What each of the last turns added to the context, as bars; the newest in the forecast's color.
81export function trend(values: number[], color: string) {
82  const top = Math.max(...values, 1)
83  const bars = values
84    .map((v, i) => {
85      const h = Math.max((v / top) * 14, 1.5)
86      const paint = i === values.length - 1 ? `class="k" fill="${color}"` : `fill="${PAST}"`
87
88      return `<rect x="${i * 5}" y="${+(14 - h).toFixed(2)}" width="3" height="${+h.toFixed(2)}" rx="1" ${paint}/>`
89    })
90    .join('')
91
92  return svg(trendWidth(values.length), 14, '', light(color, 'fill') + bars)
93}
94
95export function trendWidth(count: number) {
96  return Math.max(count * 5 - 2, 1)
97}
98
99// How much of a window ending at `resetsAt` has passed at `now`, 0 to 1.
100export function elapsed(resetsAt: string, span: number, now: number) {
101  return Math.min(Math.max((now - (Date.parse(resetsAt) - span)) / span, 0), 1)
102}
103
104// Repaints the `k` marks in the deeper hue when the scheme is light.
105function light(color: string, prop: 'fill' | 'stroke') {
106  const deeper = ON_LIGHT[color]
107
108  return deeper ? `<style>@media (prefers-color-scheme: light){.k{${prop}:${deeper}}}</style>` : ''
109}
110
111function share(percent: number) {
112  return Math.min(Math.max(percent, 0), 100) / 100
113}
114
115// Degrees on a 12-hour face, clockwise from twelve, for a moment in local time.
116function face(ms: number) {
117  const d = new Date(ms)
118
119  return ((d.getHours() % 12) * 60 + d.getMinutes() + d.getSeconds() / 60) / 2
120}
121
122function at(deg: number, r: number) {
123  const a = (deg * Math.PI) / 180
124
125  return `${+(12 + r * Math.sin(a)).toFixed(2)} ${+(12 - r * Math.cos(a)).toFixed(2)}`
126}
127
128// The colored part of the rim, clockwise from twelve.
129function rim(sweep: number, color: string) {
130  if (sweep <= 0) return ''
131
132  return `<path class="k" d="M${at(0, 10)}A10 10 0 ${sweep > 180 ? 1 : 0} 1 ${at(sweep, 10)}" stroke="${color}" stroke-width="2.5"/>`
133}
134
135// A hand at `deg` that goes once round the face every `period` seconds.
136function hand(deg: number, length: number, period: number) {
137  return (
138    `<g transform="rotate(${deg} 12 12)"><path d="M12 12V${12 - length}" stroke="${MARK}" stroke-width="1.8"/>` +
139    `<animateTransform attributeName="transform" type="rotate" from="${deg} 12 12" to="${deg + 360} 12 12" dur="${period}s" repeatCount="indefinite"/></g>`
140  )
141}
142
143function svg(w: number, h: number, attrs: string, body: string) {
144  return `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 ${w} ${h}" width="${w}" height="${h}" ${attrs}>${body}</svg>`
145}
146

1// `isEstimate`: counted locally before the session's first response, not reported by the API.
2export type TokenWeatherReading = { tokens: number; window: number; percent: number; isEstimate?: boolean }
3export type TokenWeatherLimit = { kind: string; percentUsed: number; resetsAt?: string }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'token-weather': { readings: TokenWeatherReading[]; limits: TokenWeatherLimit[]; live: TokenWeatherReading | null }
8  }
9}
10