token-weather-usage

1import { atom, read, update } from 'claude-code'
2import type { Elements, EngineInterface, Register, RenderChildren } from 'claude-code'
3
4import type { Limit, View } from '../types'
5
6// Two rows of colored pills above the prompt: context tokens with weather, prompt cache time left,
7// session cost and running subagents; under them, the 5h/7d limits against the time elapsed.
8// Reads $.session / $.agent / $.clock and four caching env vars; no network, files or processes.
9// ponytail: plain Unicode glyphs only (no Nerd Font, no emoji) so cell widths hold in any terminal.
10
11const MIN = 60_000
12const HOUR = 60 * MIN
13const SPAN: Record<string, number> = { five_hour: 5 * HOUR, seven_day: 7 * 24 * HOUR }
14const LABEL: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: '$' }
15const ORDER = ['five_hour', 'seven_day', 'spend_limit']
16const BARS = '▁▂▃▄▅▆▇█'
17const PIE = '○◔◑◕●'
18const SPARK = 6
19const GAUGE = 5
20const HISTORY = 12
21const COMPACT_AT = 100_000
22const CACHE_SOON = 10 * MIN
23
24// Pill accents; neutral parts use theme keys (subtle, inactive) to follow light and dark.
25const HUE: Record<string, string> = {
26  five_hour: '#3fb97f',
27  seven_day: '#a07cf0',
28  spend_limit: '#e0b040',
29  cache: '#4fb3e8',
30  cost: '#e0b040',
31  agents: '#ec6fa7',
32  ok: '#3fb97f',
33  warn: '#f0a030',
34  hot: '#ef5350',
35}
36
37const WEATHER = [
38  { upTo: 25, icon: '☀', color: '#f5c000' },
39  { upTo: 50, icon: '☁', color: '#94a3b8' },
40  { upTo: 75, icon: '☂', color: '#4c7cf0' },
41  { upTo: Infinity, icon: '↯', color: '#a855f7' },
42]
43
44type UI = Pick<Elements['terminal'], 'Box' | 'Text'>
45type Data = Omit<View, 'now'>
46
47const view = atom({ plugin: 'token-weather-usage', key: 'view' } as const, null)
48
49const blank = (): Data => ({
50  tokens: [],
51  window: 0,
52  limits: [],
53  cache: null,
54  cacheOff: false,
55  cacheTtl: '',
56  cost: null,
57  lastPrompt: null,
58  agents: 0,
59})
60
61// ponytail: module variables, rebuilt from $.session.usage() on start; the band reads the snapshot in $.state.
62let s = blank()
63let promptBase: number | null = null
64
65export const register: Register = on => {
66  let ticker: { cancel: () => void } | undefined
67
68  on('session.start', async ($, e, next) => {
69    ticker?.cancel()
70    s = { ...blank(), window: s.window }
71    const flag = (v: string | undefined) => /^(1|true|yes|on)$/i.test((v ?? '').trim())
72    s.cacheOff = flag(await $.env.get('DISABLE_PROMPT_CACHING'))
73    s.cacheTtl = flag(await $.env.get('FORCE_PROMPT_CACHING_5M'))
74      ? '5m'
75      : (await $.env.get('CLAUDE_CODE_PROMPT_CACHE_TTL')) ||
76        (flag(await $.env.get('ENABLE_PROMPT_CACHING_1H')) ? '1h' : '')
77    apply(await $.session.usage())
78    promptBase = s.cost
79    await refreshAgents($)
80    // Elapsed time, the cache countdown and agents change between turns: redraw every 15 s.
81    ticker = $.clock.every(15_000, async () => {
82      await refreshAgents($)
83      await show($)
84    })
85    await show($)
86    return next(e)
87  })
88
89  on('session.end', async ($, e, next) => {
90    // /clear, /resume and disconnect keep the ticker.
91    if (e.reason === 'prompt_input_exit' || e.reason === 'other') ticker?.cancel()
92    return next(e)
93  })
94
95  on('session.measure', async ($, e, next) => {
96    apply(e)
97    await show($)
98    return next(e)
99  })
100
101  // Main-loop requests only: how much of the prompt the cache served.
102  on('turn.step', async function* ($, e, next) {
103    if (e.agentId) return yield* next(e)
104    const at = await $.clock.now()
105    const result = yield* next(e)
106    if (result?.usage) {
107      s.cache = {
108        at,
109        read: result.usage.cache_read_input_tokens ?? 0,
110        write: result.usage.cache_creation_input_tokens ?? 0,
111        fresh: result.usage.input_tokens ?? 0,
112      }
113      await refreshAgents($)
114      await show($)
115    }
116    return result
117  })
118
119  on('turn.complete', async ($, e, next) => {
120    const result = await next(e)
121    try {
122      if (e.agentId) {
123        await refreshAgents($)
124      } else {
125        apply(await $.session.usage())
126        if (s.cost !== null && promptBase !== null && s.cost >= promptBase) s.lastPrompt = s.cost - promptBase
127        promptBase = s.cost
128      }
129      await show($)
130    } catch {
131      // No reading this turn: the band keeps the previous one.
132    }
133    return result
134  })
135
136  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
137    const v = await read($, view)
138    if (e.props.hasSurvey || !v || (v.tokens.length === 0 && v.limits.length === 0)) return next(e)
139    const ui = $.ui.resolve(e)
140    const { Box } = ui
141    return (
142      <Box width={e.props.bodyColumns} flexDirection="column" rowGap={1} marginTop={1}>
143        {rows(ui, v)}
144      </Box>
145    )
146  })
147}
148
149const rank = (k: string) => (ORDER.includes(k) ? ORDER.indexOf(k) : ORDER.length)
150
151function apply(u: { context: { tokens?: number; window: number }; rateLimits: Limit[]; cost?: { usd: number } }) {
152  const t = u.context.tokens ?? 0
153  let tokens = s.tokens.filter(n => n > 0)
154  // A reopened session reads the same context again: no duplicate reading, no false empty bar.
155  if (u.context.window && !(t > 0 && tokens[tokens.length - 1] === t)) tokens.push(t)
156  s.tokens = tokens.slice(-HISTORY)
157  s.window = u.context.window || s.window
158  if (u.rateLimits.length > 0) s.limits = [...u.rateLimits].sort((a, b) => rank(a.kind) - rank(b.kind))
159  if (u.cost) s.cost = u.cost.usd
160}
161
162async function refreshAgents($: EngineInterface) {
163  try {
164    s.agents = ((await $.agent.list()) ?? []).filter(a => a.status === 'running').length
165  } catch {
166    // Keep the previous count.
167  }
168}
169
170async function show($: EngineInterface) {
171  const now = await $.clock.now()
172  await update($, view, () => ({ ...s, tokens: [...s.tokens], limits: [...s.limits], now }))
173}
174
175function rows(ui: UI, v: View) {
176  const session: RenderChildren[] = []
177  if (v.tokens.length > 0) session.push(contextPill(ui, v))
178  session.push(cachePill(ui, v))
179  if (v.cost !== null && v.cost >= 0.005) session.push(costPill(ui, v))
180  if (v.agents > 0) session.push(agentsPill(ui, v.agents))
181  const limits: RenderChildren[] = []
182  for (const l of v.limits) {
183    if (Date.parse(l.resetsAt ?? '') <= v.now) continue // reset already: no valid reading
184    limits.push(limitPill(ui, l, v.now))
185  }
186  return [row(ui, 'session', session), row(ui, 'limits', limits)]
187}
188
189// One left-aligned row, a grey " | " between pills; an empty row draws nothing.
190function row({ Box, Text }: UI, key: string, items: RenderChildren[]) {
191  const shown = items.filter(Boolean)
192  if (shown.length === 0) return null
193  return (
194    <Box key={key} flexDirection="row" flexWrap="wrap">
195      {shown.flatMap((p, i) => (i > 0 ? [<Text key={`${key}-sep-${i}`} color="inactive"> | </Text>, p] : [p]))}
196    </Box>
197  )
198}
199
200// ponytail: no border, so a pill is one row instead of three; the separators in row() mark where each ends.
201function pill({ Box, Text }: UI, key: string, ...parts: RenderChildren[]) {
202  return (
203    <Box key={key} flexShrink={0}>
204      <Text wrap="truncate">{parts}</Text>
205    </Box>
206  )
207}
208
209// ☀ 48k ▂▅▁▃▇ ▲ +5.3k: weather by share of the window, a bar per turn's growth, the last step.
210function contextPill(ui: UI, v: View) {
211  const { Text } = ui
212  const cur = v.tokens[v.tokens.length - 1] ?? 0
213  const pct = v.window > 0 ? (cur / v.window) * 100 : 0
214  const w = WEATHER.find(b => pct < b.upTo) ?? WEATHER[WEATHER.length - 1]!
215  const deltas = v.tokens.slice(1).map((n, i) => Math.max(0, n - (v.tokens[i] ?? 0))).slice(-SPARK)
216  const top = Math.max(...deltas, 1)
217  const bars = deltas.map(d => BARS[Math.round((d / top) * (BARS.length - 1))]).join('')
218  const diff = v.tokens.length >= 2 ? cur - (v.tokens[v.tokens.length - 2] ?? 0) : 0
219  return pill(
220    ui,
221    'context',
222    <Text color={w.color}>{w.icon}</Text>,
223    ' ',
224    <Text bold>{short(cur)}</Text>,
225    bars ? [' ', <Text color="inactive">{bars.slice(0, -1)}</Text>, <Text color={w.color}>{bars.slice(-1)}</Text>] : null,
226    diff !== 0 ? <Text dimColor>{` ${diff > 0 ? '▲ +' : '▼ −'}${short(Math.abs(diff))}`}</Text> : null,
227  )
228}
229
230// ◔ 5h ━╍━━━ 22% ↻ 3h00 → 20:40
231function limitPill(ui: UI, l: Limit, now: number) {
232  const { Text } = ui
233  const used = Math.max(0, l.percentUsed)
234  const resetMs = l.resetsAt ? Date.parse(l.resetsAt) : NaN
235  const span = SPAN[l.kind]
236  const left = Number.isFinite(resetMs) ? Math.max(0, resetMs - now) : null
237  const elapsed = span && left !== null ? clamp(((span - left) / span) * 100) : null
238  const pace = elapsed === null ? 0 : used - elapsed
239  const hot = used >= 90 || pace > 15
240  const level = hot ? HUE.hot! : used >= 75 || pace > 5 ? HUE.warn! : HUE.ok!
241  const hue = HUE[l.kind] ?? HUE.spend_limit!
242  const icon = l.kind === 'five_hour' ? PIE[Math.round(clamp(used) / 25)] : l.kind === 'seven_day' ? '▦' : '¤'
243  const when = left === null ? '' : l.kind === 'five_hour' ? `${duration(left)} → ${clock(resetMs)}` : duration(left)
244  return pill(
245    ui,
246    l.kind,
247    <Text color={hue}>{icon}</Text>,
248    ` ${LABEL[l.kind] ?? l.kind} `,
249    gauge(ui, used, elapsed, level),
250    ' ',
251    <Text bold color={hot ? level : undefined}>{`${Math.round(used)}%`}</Text>,
252    when ? [' ', <Text dimColor>{`↻ ${when}`}</Text>] : null,
253  )
254}
255
256// Solid: the share used, up to where the clock stands. Dashed: the gap between usage and time
257// elapsed, grey while usage trails the clock, in the warning color once it runs ahead.
258function gauge({ Text }: UI, used: number, elapsed: number | null, color: string) {
259  const u = Math.round((clamp(used) / 100) * GAUGE)
260  const t = elapsed === null ? u : Math.round((elapsed / 100) * GAUGE)
261  const solid = Math.min(u, t)
262  const dashed = Math.abs(t - u)
263  const rest = GAUGE - Math.max(u, t)
264  return [
265    solid > 0 ? <Text color={color}>{'━'.repeat(solid)}</Text> : null,
266    dashed > 0 ? <Text color={u > t ? color : 'inactive'}>{'╍'.repeat(dashed)}</Text> : null,
267    rest > 0 ? <Text color="subtle">{'━'.repeat(rest)}</Text> : null,
268  ]
269}
270
271// Lifetime is inferred, as the host gives token counts not the TTL: env override, then plan.
272function cacheTtl(v: View) {
273  if (v.cacheTtl === '5m') return 5 * MIN
274  if (v.cacheTtl === '1h') return HOUR
275  // A subscription within its plan gets 1 hour; usage credits or an API key, 5 minutes.
276  const plan = v.limits.filter(l => l.kind === 'five_hour' || l.kind === 'seven_day')
277  return plan.length > 0 && plan.every(l => l.percentUsed < 100) ? HOUR : 5 * MIN
278}
279
280// ϟ cache 98% 54 min, amber under 10 min, red once expired.
281function cachePill(ui: UI, v: View) {
282  if (v.cacheOff) return null
283  const { Text } = ui
284  const head = (hue: string) => [<Text color={hue}>ϟ</Text>, ' cache ']
285  if (!v.cache) return pill(ui, 'cache', head(HUE.cache!), <Text dimColor>—</Text>)
286  const total = v.cache.read + v.cache.write + v.cache.fresh
287  const hit = total > 0 ? Math.round((v.cache.read / total) * 100) : 0
288  const left = v.cache.at + cacheTtl(v) - v.now
289  if (left <= 0) {
290    const big = (v.tokens[v.tokens.length - 1] ?? total) >= COMPACT_AT
291    return pill(
292      ui,
293      'cache',
294      head(HUE.hot!),
295      <Text bold color={HUE.hot}>expired</Text>,
296      big ? [' ', <Text dimColor>/compact</Text>] : null,
297    )
298  }
299  if (hit < 50 && v.cache.write > 1_000) {
300    return pill(ui, 'cache', head(HUE.warn!), <Text bold>{`${hit}%`}</Text>, ' ', <Text color={HUE.warn}>missed</Text>)
301  }
302  const soon = left < CACHE_SOON
303  const hue = soon ? HUE.warn! : HUE.cache!
304  const time = left < MIN ? '< 1 min' : duration(left)
305  return pill(
306    ui,
307    'cache',
308    head(hue),
309    <Text bold>{`${hit}%`}</Text>,
310    ' ',
311    soon ? <Text bold color={HUE.warn}>{time}</Text> : <Text dimColor>{time}</Text>,
312  )
313}
314
315// ¤ ≈ $18.42 ❯ +$2.31: the session so far, then what the last prompt added.
316function costPill(ui: UI, v: View) {
317  const { Text } = ui
318  const last = v.lastPrompt
319  return pill(
320    ui,
321    'cost',
322    <Text color={HUE.cost}>¤</Text>,
323    ' ≈ ',
324    <Text bold>{`$${(v.cost ?? 0).toFixed(2)}`}</Text>,
325    last !== null && last >= 0.005 ? [' ', <Text dimColor>{`❯ +$${last.toFixed(2)}`}</Text>] : null,
326  )
327}
328
329function agentsPill(ui: UI, n: number) {
330  const { Text } = ui
331  return pill(ui, 'agents', <Text color={HUE.agents}>✻</Text>, ' ', <Text bold>{String(n)}</Text>, n === 1 ? ' agent' : ' agents')
332}
333
334// 3h02, 42 min, 2d23h.
335function duration(ms: number) {
336  const m = Math.round(ms / MIN)
337  if (m < 60) return `${m} min`
338  const d = Math.floor(m / 1440)
339  const hours = Math.floor((m % 1440) / 60)
340  return d > 0 ? `${d}d${String(hours).padStart(2, '0')}h` : `${hours}h${String(m % 60).padStart(2, '0')}`
341}
342
343function clock(ms: number) {
344  const d = new Date(ms)
345  return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}`
346}
347
348const clamp = (n: number) => Math.min(100, Math.max(0, n))
349
350// 1M, 1.2M, 107k, 98.3k, 950: one decimal only when it matters.
351function short(n: number) {
352  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
353  if (n >= 100_000) return `${Math.round(n / 1_000)}k`
354  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
355  return String(n)
356}
357

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2export type Cache = { at: number; read: number; write: number; fresh: number }
3
4// Everything the band draws from, stamped with the time it was taken.
5export type View = {
6  now: number
7  tokens: number[]
8  window: number
9  limits: Limit[]
10  cache: Cache | null
11  cacheOff: boolean
12  cacheTtl: string
13  cost: number | null
14  lastPrompt: number | null
15  agents: number
16}
17
18declare module 'claude-code' {
19  interface PluginState {
20    'token-weather-usage': { view: View | null }
21  }
22}
23