token-weather

1// Token Weather: how full the context window is, as a weather forecast,
2// one line above the prompt. Updates after every main-loop turn.
3
4import { atom, read, update } from 'claude-code'
5import type { EngineInterface, Register } from 'claude-code'
6
7import type { TokenWeatherReading } from '../types'
8
9const KEEP = 12
10const BARS = '▁▂▃▄▅▆▇█'
11
12type Sky = { below: number; icon: string; word: string; color: string }
13
14// Single-width symbols only, so the band lines up in every terminal font.
15const SKIES: readonly Sky[] = [
16  { below: 25, icon: '☀', word: 'Clear', color: 'yellow' },
17  { below: 50, icon: '☁', word: 'Cloudy', color: 'cyan' },
18  { below: 75, icon: '☂', word: 'Showers', color: 'blue' },
19  { below: 90, icon: '☇', word: 'Storm', color: 'magenta' },
20  { below: Infinity, icon: '↯', word: 'Compact soon', color: 'red' },
21]
22
23const readings = atom({ plugin: 'token-weather', key: 'readings' } as const, [])
24
25export const register: Register = on => {
26  on('session.start', async ($, e, next) => {
27    const result = await next(e)
28    await measure($)
29    return result
30  })
31
32  on('turn.complete', async ($, e, next) => {
33    const result = await next(e)
34    if (e.agentId === undefined) await measure($)
35    return result
36  })
37
38  on('session.end', async ($, e, next) => {
39    if (e.reason === 'clear') await update($, readings, () => [])
40    return next(e)
41  })
42
43  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
44    const history = await read($, readings)
45    const below = await next(e)
46    if (e.props.hasSurvey || history.length === 0) return below
47
48    const { Box, Text } = $.ui.resolve(e)
49    const now = history[history.length - 1]!
50    const sky = skyFor(now.percent)
51    const width = e.props.bodyColumns
52
53    const line = (
54      <Box key="token-weather" flexDirection="row" paddingX={1}>
55        <Text color={sky.color} bold>
56          {sky.icon} {sky.word}
57        </Text>
58        <Text> {now.percent}% of context</Text>
59        <Text dimColor>
60          {'  '}
61          {short(now.tokens)} / {short(now.window)}
62        </Text>
63        {width >= 64 && history.length > 1 && (
64          <Text color={sky.color}>{'  ' + sparkline(history)}</Text>
65        )}
66        {width >= 80 && history.length > 1 && <Text dimColor>{'  ' + trend(history)}</Text>}
67        {width >= 96 && now.usd !== undefined && <Text dimColor>{'  $' + now.usd.toFixed(2)}</Text>}
68      </Box>
69    )
70
71    return (
72      <Box flexDirection="column">
73        {below}
74        {line}
75      </Box>
76    )
77  })
78}
79
80async function measure($: EngineInterface): Promise<void> {
81  const usage = await $.session.usage()
82  const { context } = usage
83  if (!context.window) return
84  const tokens = context.tokens ?? 0
85  const percent = context.percent ?? Math.round((tokens / context.window) * 100)
86  const reading: TokenWeatherReading = { tokens, window: context.window, percent }
87  if (usage.cost !== undefined) reading.usd = usage.cost.usd
88  await update($, readings, list => {
89    const last = list[list.length - 1]
90    // A session.start after a hot reload reads the same figures again: skip the repeat.
91    if (last && last.tokens === reading.tokens && last.window === reading.window) {
92      return [...list.slice(0, -1), reading]
93    }
94    return [...list, reading].slice(-KEEP)
95  })
96}
97
98export function skyFor(percent: number): Sky {
99  return SKIES.find(s => percent < s.below) ?? SKIES[SKIES.length - 1]!
100}
101
102function sparkline(history: readonly TokenWeatherReading[]): string {
103  const top = Math.max(1, ...history.map(r => r.tokens))
104  return history
105    .map(r => BARS[Math.min(BARS.length - 1, Math.floor((r.tokens / top) * (BARS.length - 1)))])
106    .join('')
107}
108
109function trend(history: readonly TokenWeatherReading[]): string {
110  const a = history[history.length - 2]!
111  const b = history[history.length - 1]!
112  const delta = b.tokens - a.tokens
113  if (delta === 0) return 'steady'
114  return delta > 0 ? `▲ +${short(delta)} last turn` : `▼ ${short(-delta)} last turn`
115}
116
117export function short(n: number): string {
118  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
119  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
120  return String(n)
121}
122

1export type TokenWeatherReading = {
2  tokens: number
3  window: number
4  percent: number
5  /** Session cost in US dollars when the host keeps a ledger. */
6  usd?: number
7}
8
9declare module 'claude-code' {
10  interface PluginState {
11    'token-weather': { readings: TokenWeatherReading[] }
12  }
13}
14