token-weather

1// Token Weather: a live forecast of the context window, above the prompt,
2// and how long the prompt cache stays warm after the last request.
3import { atom, read, update } from 'claude-code'
4import type { EngineInterface, Register, Timer } from 'claude-code'
5
6import type { CacheTtl, TokenWeatherReading } from '../types'
7import { bandsFile, parseOrder, stackBands } from './bands'
8
9const HISTORY = 12
10const BARS = '▁▂▃▄▅▆▇█'
11const FORECAST = [
12  { upTo: 25, icon: '☀', word: 'Clear', color: 'yellow' },
13  { upTo: 50, icon: '☁', word: 'Cloudy', color: 'cyan' },
14  { upTo: 75, icon: '☂', word: 'Showers', color: 'blue' },
15  { upTo: 90, icon: '☇', word: 'Storm', color: 'magenta' },
16  { upTo: Infinity, icon: '↯', word: 'Compact soon', color: 'red' },
17] as const
18
19// The cache countdown ticks every second only in its last two minutes;
20// before that it moves in 15 s steps (30 s with over ten minutes left),
21// so the band is not redrawn every second.
22const TTL_MS = { '5m': 5 * 60_000, '1h': 60 * 60_000 } as const
23const FAST_BELOW_MS = 2 * 60_000
24const SLOW_STEP_MS = 15_000
25const SLOWER_ABOVE_MS = 10 * 60_000
26const SLOWER_STEP_MS = 30_000
27const WARN_BELOW_MS = 60_000
28// How much of the transcript's end to read for the last response's usage.
29const TAIL_BYTES = 1024 * 1024
30
31// Held by the host, so the history survives a hot reload of this file.
32const readings = atom({ plugin: 'token-weather', key: 'readings' } as const, [] as TokenWeatherReading[])
33const lastRequestAt = atom({ plugin: 'token-weather', key: 'lastRequestAt' } as const, null as number | null)
34// The cache lifetime Claude's responses last showed; null until one wrote to the cache.
35const detectedTtl = atom({ plugin: 'token-weather', key: 'detectedTtl' } as const, null as CacheTtl | null)
36// Whether the cache-clock mod draws the cache on the status line, from Claude Code's own figures.
37const cacheClockOn = atom({ plugin: 'token-weather', key: 'cacheClockOn' } as const, false)
38
39let bandOrder: Promise<string[]> | undefined // read once per session: restart to apply
40
41// The band order file, as the claude-code-mods checklist saves it; [] on any error.
42async function readBandOrder($: EngineInterface): Promise<string[]> {
43  const file = bandsFile(await $.env.get('CLAUDE_CONFIG_DIR'), await $.env.get('HOME'))
44  return file && (await $.fs.exists(file)) ? parseOrder(await $.fs.read(file)) : []
45}
46
47// Forgets the cached band order, for tests.
48export function resetBandOrder() {
49  bandOrder = undefined
50}
51
52export const register: Register = (on, options) => {
53  // `5m` or `1h` set by hand wins; `auto` (the default) goes by what was detected.
54  const override: CacheTtl | null = options.cacheTtl === '5m' || options.cacheTtl === '1h' ? options.cacheTtl : null
55
56  on('session.start', async ($, e, next) => {
57    const result = await next(e)
58    await takeReading($)
59    await checkCacheClock($)
60    if ((await read($, detectedTtl)) === null) {
61      const known = await $.store.get('detectedTtl').catch(() => undefined) // the last session's, until this one's first response says
62      if (known === '5m' || known === '1h') await update($, detectedTtl, () => known)
63    }
64    const since = await read($, lastRequestAt) // kept across a hot reload: pick the countdown up again
65    if (since !== null) restart($, since, await $.clock.now(), await ttlMsFor($, override))
66    return result
67  })
68
69
70  // A /clear or a resume leaves the old conversation, and its cache, behind.
71  on('session.end', async ($, e, next) => {
72    if (e.reason === 'clear' || e.reason === 'resume') {
73      stop()
74      await update($, lastRequestAt, () => null)
75    }
76    return next(e)
77  })
78
79  // One model request of a turn. The cache's clock restarts when each one on the
80  // main thread finishes, tool-loop requests inside a turn included.
81  on('turn.step', async function* ($, e, next) {
82    const result = yield* next(e)
83    if (!e.agentId && result.stopReason !== null) {
84      const now = await $.clock.now()
85      await update($, lastRequestAt, () => now)
86      restart($, now, now, await ttlMsFor($, override))
87    }
88    return result
89  })
90
91  on('turn.complete', async ($, e, next) => {
92    const result = await next(e)
93    if (!e.agentId) {
94      await takeReading($) // main-loop turns only, not subagents
95      await checkCacheClock($) // set up or removed mid-session
96      if (override === null) {
97        // The last response's usage in the transcript says which lifetime it wrote to the cache with.
98        const found = ttlFromTranscript(await readTail($))
99        if (found !== null) await detect($, found)
100      }
101    }
102    return result
103  })
104
105  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
106    const rest = await next(e) // what other mods and Claude Code draw here stays
107    const history = await read($, readings)
108    const now = history.at(-1)
109    if (e.props.hasSurvey || !now) {
110      return rest
111    }
112
113    const { Box, Text } = $.ui.resolve(e)
114    const f = FORECAST.find(b => now.percent < b.upTo) ?? FORECAST[4]
115    const isWide = e.props.bodyColumns >= 60
116
117    // The cache part goes first when the line runs out of room.
118    const since = await read($, lastRequestAt)
119    const ttlMs = await ttlMsFor($, override)
120    let cache: ReturnType<typeof cachePart> | null = null
121    if (since !== null && !(await read($, cacheClockOn))) {
122      const clockNow = await $.clock.now()
123      if (armedFor !== since) restart($, since, clockNow, ttlMs) // its timer was lost (a reload): pick it up
124      cache = cachePart(since + ttlMs - clockNow)
125    }
126    const lineWidth =
127      2 + // paddingX
128      `${f.icon} ${f.word} ${now.percent}% of context ${short(now.tokens)} / ${short(now.window)}`.length +
129      (isWide ? `  last turns ${sparkline(history)}${history.length > 1 ? trend(history) : ''}`.length : 0)
130    const showCache = isWide && cache !== null && lineWidth + cache.text.length <= e.props.bodyColumns
131
132    const mine = (
133      <Box flexDirection="row" paddingX={1}>
134        <Text color={f.color} bold>{`${f.icon} ${f.word}`}</Text>
135        <Text>{` ${now.percent}% of context`}</Text>
136        <Text dimColor>{` ${short(now.tokens)} / ${short(now.window)}`}</Text>
137        {isWide && <Text dimColor>{'  last turns '}</Text>}
138        {isWide && <Text color={f.color}>{sparkline(history)}</Text>}
139        {isWide && history.length > 1 && <Text dimColor>{trend(history)}</Text>}
140        {showCache && cache && <Text color={cache.color} dimColor={cache.isCold}>{cache.text}</Text>}
141      </Box>
142    )
143    bandOrder ??= readBandOrder($).catch(() => [])
144    return stackBands(rest, 'token-weather', mine, 'above', await bandOrder) // above what other mods draw, unless the band order says otherwise
145  })
146}
147
148// cache-clock's status line shows the exact countdown, so this one steps aside.
149async function checkCacheClock($: EngineInterface) {
150  const settings = await $.settings.read().catch(() => ({}) as Record<string, unknown>)
151  const line = settings.statusLine as { command?: unknown } | undefined
152  const isOn = typeof line?.command === 'string' && line.command.includes('cache-clock/cache-clock.mjs')
153  if ((await read($, cacheClockOn)) !== isOn) await update($, cacheClockOn, () => isOn)
154}
155
156async function ttlMsFor($: EngineInterface, override: CacheTtl | null) {
157  return TTL_MS[override ?? (await read($, detectedTtl)) ?? '5m']
158}
159
160// Keeps a newly detected lifetime, here and for the next session, and moves the countdown
161// onto it. Only called with no override set.
162async function detect($: EngineInterface, ttl: CacheTtl) {
163  if ((await read($, detectedTtl)) === ttl) return
164  await update($, detectedTtl, () => ttl)
165  await $.store.set('detectedTtl', ttl).catch(() => {}) // remembered for the next session when the store allows
166  const now = await $.clock.now()
167  const since = await read($, lastRequestAt) // read last, so a request that just finished is the one re-armed
168  if (since !== null) restart($, since, now, TTL_MS[ttl])
169}
170
171// The end of this session's transcript, where the last response is: whole if small,
172// else its last bytes. '' when it cannot be found or read.
173async function readTail($: EngineInterface) {
174  try {
175    const path = await transcriptPath($)
176    if (path === null) return ''
177    const { size } = await $.fs.stat(path)
178    if (size <= TAIL_BYTES) return await $.fs.read(path)
179    const { exitCode, stdout } = await $.process.run(['tail', '-c', String(TAIL_BYTES), path])
180    return exitCode === 0 ? stdout : ''
181  } catch {
182    return '' // unreadable: keep the last detected value
183  }
184}
185
186// <config dir>/projects/<project root, each other character a dash>/<session id>.jsonl
187// A very long project path, or a root moved during the session, gives a path that does
188// not exist: then nothing is detected and the last value stays.
189async function transcriptPath($: EngineInterface) {
190  const home = await $.env.get('HOME')
191  const configDir = (await $.env.get('CLAUDE_CONFIG_DIR')) ?? (home ? `${home}/.claude` : undefined)
192  if (!configDir) return null
193  const project = (await $.session.root()).replace(/[^a-zA-Z0-9]/g, '-')
194  return `${configDir}/projects/${project}/${await $.session.id()}.jsonl`
195}
196
197// The lifetime the latest main-thread response wrote to the cache with; null when it
198// wrote nothing (a pure cache hit) or no response was found.
199function ttlFromTranscript(text: string): CacheTtl | null {
200  const lines = text.split('\n')
201  for (let i = lines.length - 1; i >= 0; i--) {
202    let entry: any
203    try {
204      entry = JSON.parse(lines[i] ?? '')
205    } catch {
206      continue // blank, or cut by the tail
207    }
208    if (entry?.type !== 'assistant' || entry.isSidechain || !entry.message?.usage) continue
209    const written = entry.message.usage.cache_creation
210    if (Number(written?.ephemeral_1h_input_tokens) > 0) return '1h'
211    if (Number(written?.ephemeral_5m_input_tokens) > 0) return '5m'
212    return null
213  }
214  return null
215}
216
217// The countdown's redraw timer, one at a time; a reload of this file starts it over.
218let tick: Timer | undefined
219let armedFor: number | null = null
220
221function stop() {
222  tick?.cancel()
223  tick = undefined
224  armedFor = null
225}
226
227function restart($: EngineInterface, since: number, now: number, ttlMs: number) {
228  armedFor = since
229  arm($, since, now, ttlMs)
230}
231
232// Redraws the band when the countdown's text next changes, until it reads cold.
233function arm($: EngineInterface, since: number, now: number, ttlMs: number) {
234  if (armedFor !== since) return // a newer request took over
235  tick?.cancel()
236  const left = since + ttlMs - now
237  if (left <= 0) return
238  const isFast = left <= FAST_BELOW_MS
239  const step = isFast ? 1000 : left > SLOWER_ABOVE_MS ? SLOWER_STEP_MS : SLOW_STEP_MS
240  const wait = Math.min(left % step || step, isFast ? step : left - (left > SLOWER_ABOVE_MS ? SLOWER_ABOVE_MS : FAST_BELOW_MS))
241  tick = $.clock.after(wait, () => {
242    $.ui.invalidate('ui.render')
243    $.clock.now().then(t => arm($, since, t, ttlMs)).catch(() => {})
244  })
245}
246
247async function takeReading($: EngineInterface) {
248  const { context } = await $.session.usage()
249  if (!context?.window) return
250  const tokens = context.tokens ?? 0
251  const percent = context.percent ?? Math.round((tokens / context.window) * 100)
252  await update($, readings, history => [...history, { tokens, window: context.window, percent }].slice(-HISTORY))
253}
254
255// What the cache part says with `leftMs` until the cache expires.
256function cachePart(leftMs: number) {
257  if (leftMs <= 0) {
258    return { text: '  ❄ cache cold', color: 'red', isCold: true }
259  }
260  const seconds = Math.ceil(leftMs / 1000)
261  const clock = `${Math.floor(seconds / 60)}:${String(seconds % 60).padStart(2, '0')}`
262  return { text: `  ❄ cache ${clock}`, color: seconds * 1000 < WARN_BELOW_MS ? 'yellow' : undefined, isCold: false }
263}
264
265function sparkline(history: TokenWeatherReading[]) {
266  const top = Math.max(...history.map(r => r.tokens), 1)
267  return history.map(r => BARS[Math.floor((r.tokens / top) * (BARS.length - 1))]).join('')
268}
269
270function trend(history: TokenWeatherReading[]) {
271  const delta = (history.at(-1)?.tokens ?? 0) - (history.at(-2)?.tokens ?? 0)
272  if (delta === 0) return ' steady'
273  return delta > 0 ? ` ▲ +${short(delta)} last turn` : ` ▼ ${short(-delta)} last turn`
274}
275
276function short(n: number) {
277  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
278  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
279  return String(n)
280}
281

1// Band order: every band mod ships this same file (canonical copy: shared/bands.ts in the repo).
2//   The order file, <configDir>/claude-code-mods/bands.json, lists mod names top to bottom
3//   (the last sits right above the prompt). Each band mod wraps its band in a Box keyed
4//   `band:<name>` and returns one column keyed `bands`; the next band mod unpacks that
5//   column, adds its own band and sorts again, so the last one to draw sorts them all.
6//   Edit shared/bands.ts, then copy it over every mod's hooks/bands.ts (a repo test checks).
7//   Reading the file needs $, which Claude Code follows only within one file, so each mod
8//   reads it in its own register file with bandsFile and parseOrder, once per session.
9import type { RenderElement, RenderNode } from 'claude-code'
10
11const COLUMN = 'bands'
12const PREFIX = 'band:'
13
14// Where the checklist saves the order: CLAUDE_CONFIG_DIR, or ~/.claude. Undefined with neither.
15export function bandsFile(configDir: string | undefined, home: string | undefined): string | undefined {
16  const dir = configDir ? configDir.replace(/\/+$/, '') : home ? `${home}/.claude` : ''
17  return dir ? `${dir}/claude-code-mods/bands.json` : undefined
18}
19
20// The names in an order file; [] when it is no JSON or not { order: string[] }.
21export function parseOrder(text: string): string[] {
22  try {
23    const order = (JSON.parse(text) as { order?: unknown } | null)?.order
24    return Array.isArray(order) && order.every(n => typeof n === 'string') ? order : []
25  } catch {
26    return []
27  }
28}
29
30/**
31 * One column of bands: `mine` wrapped in a Box keyed `band:<name>`, placed
32 * `at` (above or below) what `next(e)` returned, then sorted by `order`.
33 * With an order, items that are no band come first, then bands not in the
34 * order (in their current order), then the listed ones in file order. With
35 * no order nothing moves, so the screen looks as it did before.
36 */
37export function stackBands(rest: RenderElement | null | undefined, name: string, mine: RenderNode, at: 'above' | 'below', order: string[]): RenderElement {
38  const own: RenderElement = { type: 'Box', props: { key: PREFIX + name, flexDirection: 'column' }, children: [mine] }
39  const others = unpack(rest)
40  const items = at === 'above' ? [own, ...others] : [...others, own]
41  return { type: 'Box', props: { key: COLUMN, flexDirection: 'column' }, children: order.length ? sorted(items, order) : items }
42}
43
44// Another band mod's column gives up its items; anything else is one item.
45function unpack(rest: RenderElement | null | undefined): RenderNode[] {
46  if (!rest) return []
47  if (!isColumn(rest)) return [rest]
48  const kids = rest.children ?? []
49  const owner = rest.part?.plugin
50  // The engine stamped the column with its owner: that stamp belongs on the owner's own band.
51  return kids.map(k => (owner && typeof k !== 'string' && bandName(k) === owner && !('part' in k && k.part) ? { ...k, part: rest.part } : k))
52}
53
54function isColumn(el: RenderElement): el is RenderElement & { type: 'Box' } {
55  return el.type === 'Box' && el.props?.key === COLUMN
56}
57
58function bandName(node: RenderNode): string | undefined {
59  if (typeof node === 'string' || node.type !== 'Box') return undefined
60  const key = node.props?.key
61  return typeof key === 'string' && key.startsWith(PREFIX) ? key.slice(PREFIX.length) : undefined
62}
63
64function sorted(items: RenderNode[], order: string[]): RenderNode[] {
65  const rank = (n: RenderNode) => {
66    const name = bandName(n)
67    if (name === undefined) return -2 // not a band: first
68    const i = order.indexOf(name)
69    return i < 0 ? -1 : i // a band not in the order: on top of the listed ones
70  }
71  return items.map((n, i) => ({ n, i, r: rank(n) })).sort((a, b) => a.r - b.r || a.i - b.i).map(x => x.n)
72}
73

1export type CacheTtl = '5m' | '1h'
2export type TokenWeatherReading = { tokens: number; window: number; percent: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'token-weather': {
7      readings: TokenWeatherReading[]
8      // When the last main-thread model request finished, in ms since the epoch; null before the first.
9      lastRequestAt: number | null
10      // The cache lifetime Claude's responses last showed; null until one wrote to the cache.
11      detectedTtl: CacheTtl | null
12      // Whether the cache-clock mod's status line is set up; its countdown replaces this one.
13      cacheClockOn: boolean
14    }
15  }
16}
17