token-weather

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Reading } from '../types'
5import { CHART_COLUMNS, HISTORY, formatTokens, sparkline, switchFor, trend, weatherFor } from './forecast'
6
7const COMMAND = 'token-weather'
8const USAGE = `Usage: /${COMMAND} [on|off|status]  (bare toggles)`
9
10// Held by the host, so the history survives a hot reload of this file.
11const readings = atom({ plugin: 'token-weather', key: 'readings' } as const, [] as Reading[])
12
13// The switch lives in $.store (kept between sessions); this mirrors it so the band redraws the moment it flips.
14const isOn = atom({ plugin: 'token-weather', key: 'isOn' } as const, true)
15
16// With no reading yet (loaded or switched on mid-session), take the live fill so the band shows at once.
17const seed = async ($: EngineInterface) => {
18  if ((await read($, readings)).length > 0) return
19  const { tokens, window, percent } = (await $.session.usage()).context
20  if (tokens === undefined || window <= 0) return
21  const reading: Reading = { tokens, window, percent: percent ?? Math.round((tokens / window) * 100) }
22  await update($, readings, list => (list.length > 0 ? list : [reading]))
23}
24
25export const register: Register = on => {
26  // Fires at session start, and again whenever the mod is enabled or reloaded mid-session.
27  on('session.start', async ($, e, next) => {
28    await $.command.register({
29      name: COMMAND,
30      description: 'Turn the context forecast above the prompt on or off (remembered across sessions)',
31      argumentHint: '[on|off|status]',
32      immediate: true,
33    })
34    const wasLeftOn = (await $.store.get('isOn')) !== false
35    await update($, isOn, () => wasLeftOn)
36    if (wasLeftOn) await seed($)
37    return next(e)
38  })
39
40  on('command.run', { command: COMMAND }, async ($, e) => {
41    const wasOn = await read($, isOn)
42    const choice = switchFor(e.args, wasOn)
43
44    if (choice === null) return { text: USAGE }
45    if (choice === 'status') {
46      return { text: `Token weather is ${wasOn ? 'on' : 'off'}. ${USAGE}` }
47    }
48
49    await $.store.set('isOn', choice)
50    await update($, isOn, () => choice)
51    if (!choice) return { text: 'Token weather off. Stays off in new sessions until /token-weather on.' }
52
53    await seed($)
54    const hasReading = (await read($, readings)).length > 0
55    return {
56      text: `Token weather on${hasReading ? '' : ': the forecast shows after the next turn'}. Stays on in new sessions.`,
57    }
58  })
59
60  // After each main-thread turn the engine measures the session; take the context's fill.
61  // Kept while off too, so switching back on shows the whole chart.
62  on('session.measure', async ($, e, next) => {
63    const { tokens, window, percent } = e.context
64
65    if (e.changed.includes('context') && tokens !== undefined && window > 0) {
66      const reading: Reading = {
67        tokens,
68        window,
69        percent: percent ?? Math.round((tokens / window) * 100),
70      }
71      await update($, readings, list => [...list, reading].slice(-HISTORY))
72    }
73
74    return next(e)
75  })
76
77  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
78    if (!(await read($, isOn))) return next(e)
79
80    const history = await read($, readings)
81    const now = history[history.length - 1]
82
83    if (e.props.hasSurvey || now === undefined) {
84      return next(e)
85    }
86
87    const { Box, Text } = $.ui.resolve(e)
88    const weather = weatherFor(now.percent)
89    const change = trend(history)
90    const isWide = e.props.bodyColumns >= CHART_COLUMNS
91
92    return (
93      <Box flexDirection="row" paddingX={1}>
94        <Text color={weather.color} bold>
95          {`${weather.icon}  ${weather.word}`}
96        </Text>
97        <Text>{`  ${now.percent}% of context`}</Text>
98        <Text dimColor>{`  ${formatTokens(now.tokens)} / ${formatTokens(now.window)}`}</Text>
99        {isWide && <Text dimColor>{'   last turns '}</Text>}
100        {isWide && <Text color={weather.color}>{sparkline(history)}</Text>}
101        {isWide && change !== null && <Text dimColor>{change}</Text>}
102      </Box>
103    )
104  })
105}
106

1import type { Reading } from '../types'
2
3export type Weather = { icon: string; word: string; color: string }
4
5const BARS = '▁▂▃▄▅▆▇█'
6
7/** The readings kept, one per turn: the chart draws them all. */
8export const HISTORY = 12
9
10/** Below this the band drops the chart and the trend. */
11export const CHART_COLUMNS = 60
12
13export const weatherFor = (percent: number): Weather => {
14  if (percent < 25) return { icon: '☀', word: 'Clear', color: 'yellow' }
15  if (percent < 50) return { icon: '☁', word: 'Cloudy', color: 'cyan' }
16  if (percent < 75) return { icon: '☂', word: 'Showers', color: 'blue' }
17  if (percent < 90) return { icon: '☇', word: 'Storm', color: 'magenta' }
18  return { icon: '↯', word: 'Compact soon', color: 'red' }
19}
20
21/** 134400 → "134.4k", 200000 → "200k", 1000000 → "1M". */
22export const formatTokens = (n: number): string => {
23  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
24  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
25  return String(n)
26}
27
28/** One bar per turn, scaled to the fullest turn kept (▁ … █). */
29export const sparkline = (readings: readonly Reading[]): string => {
30  const top = Math.max(...readings.map(r => r.tokens), 1)
31  return readings.map(r => BARS[Math.floor((r.tokens / top) * (BARS.length - 1))]).join('')
32}
33
34/** "  ▲ +98.3k last turn", "  ▼ 120k last turn" after a compaction, "  steady"; null before a second turn. */
35export const trend = (readings: readonly Reading[]): string | null => {
36  const [before, after] = readings.slice(-2)
37  if (before === undefined || after === undefined) return null
38  const delta = after.tokens - before.tokens
39  if (delta === 0) return '  steady'
40  return delta > 0 ? `  ▲ +${formatTokens(delta)} last turn` : `  ▼ ${formatTokens(-delta)} last turn`
41}
42
43/** `/token-weather` with its args: bare or "toggle" flips the switch; "status" reads it; null for anything else. */
44export const switchFor = (args: string, isOn: boolean): boolean | 'status' | null => {
45  const word = args.trim().toLowerCase()
46  if (word === '' || word === 'toggle') return !isOn
47  if (word === 'on') return true
48  if (word === 'off') return false
49  if (word === 'status') return 'status'
50  return null
51}
52

1export type Reading = { tokens: number; window: number; percent: number }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'token-weather': { readings: Reading[]; isOn: boolean }
6  }
7}
8