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

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