usage-band

1import { atom, read, update } from 'claude-code'
2import type { Register, SessionUsage } from 'claude-code'
3
4import type { TokenTotals, UsageSnapshot, UsageWindow } from '../types'
5
6const ZERO: TokenTotals = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }
7
8const snapshot = atom({ plugin: 'usage-band', key: 'snapshot' } as const, null)
9const tokens = atom({ plugin: 'usage-band', key: 'tokens' } as const, ZERO)
10const tick = atom({ plugin: 'usage-band', key: 'tick' } as const, 0)
11const isHidden = atom({ plugin: 'usage-band', key: 'isHidden' } as const, false)
12
13const HOUR = 3_600_000
14const WINDOW_MS: Record<string, number> = { five_hour: 5 * HOUR, seven_day: 168 * HOUR }
15const LABEL: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
16
17function toSnapshot(u: Pick<SessionUsage, 'context' | 'rateLimits' | 'cost'>): UsageSnapshot {
18  return {
19    windows: u.rateLimits.map(r => ({ kind: r.kind, percentUsed: r.percentUsed, resetsAt: r.resetsAt })),
20    costUsd: u.cost?.usd,
21    contextPercent: u.context.percent,
22  }
23}
24
25// ---------- formatting ----------
26
27function fmtTokens(n: number): string {
28  if (n < 1000) return String(n)
29  if (n < 1_000_000) return `${(n / 1000).toFixed(1)}k`
30  return `${(n / 1_000_000).toFixed(2)}M`
31}
32
33// The reading has one decimal; below 10% that decimal is most of the story (1.4% is not 1%).
34function fmtPct(pct: number): string {
35  return pct < 10 ? `${Number(pct.toFixed(1))}%` : `${Math.round(pct)}%`
36}
37
38function fmtCountdown(ms: number): string {
39  if (ms <= 0) return 'now'
40  const mins = Math.floor(ms / 60_000)
41  const d = Math.floor(mins / 1440)
42  const h = Math.floor((mins % 1440) / 60)
43  const m = mins % 60
44  if (d > 0) return `${d}d ${h}h`
45  if (h > 0) return `${h}h ${m}m`
46  return `${m}m`
47}
48
49type Window = {
50  label: string
51  pct: number
52  remaining?: string
53  pace?: number // fraction of the window already elapsed, 0..1
54  level: 'ok' | 'warm' | 'hot'
55}
56
57function describe(w: UsageWindow, now: number): Window {
58  const resetAt = w.resetsAt ? Date.parse(w.resetsAt) : NaN
59  const left = Number.isFinite(resetAt) ? resetAt - now : undefined
60  const span = WINDOW_MS[w.kind]
61  const pace =
62    left !== undefined && left > 0 && span ? Math.min(1, Math.max(0, 1 - left / span)) : undefined
63  // Past its reset the window starts over; the last reading is stale until the next response.
64  const hasReset = left !== undefined && left <= 0
65  const pct = hasReset ? 0 : w.percentUsed
66  const isAhead = pace !== undefined && pct > pace * 100 + 15
67  const level = pct >= 90 ? 'hot' : pct >= 75 || isAhead ? 'warm' : 'ok'
68  return {
69    label: LABEL[w.kind] ?? w.kind,
70    pct,
71    remaining: left !== undefined ? fmtCountdown(left) : undefined,
72    pace,
73    level,
74  }
75}
76
77// ---------- desktop: an instrument strip drawn as SVG ----------
78// Monochrome ink; amber or red appear only on a meter that needs attention.
79
80type Level = Window['level']
81
82const H = 20
83const MID = H / 2
84const FONT = 11
85const CHAR = FONT * 0.602 // monospace advance
86const LABEL_FONT = 9
87const LABEL_TRACK = 0.7
88const INK: Record<Level, string> = { ok: 'ink', warm: 'warm', hot: 'hot' }
89
90const STYLE =
91  `text{font-family:ui-monospace,SFMono-Regular,Menlo,monospace;font-size:${FONT}px;dominant-baseline:central}` +
92  `.ink{fill:#1d1d1b}.mut{fill:#8c8a85}.lbl{fill:#9b9993;font-size:${LABEL_FONT}px;letter-spacing:${LABEL_TRACK}px}` +
93  `.val{fill:#1d1d1b;font-weight:700}.trk{fill:#e4e2dd}.pace{fill:#bdbab3}` +
94  `.div{fill:#e4e2dd}.warm{fill:#c98a2e}.hot{fill:#c8553d}` +
95  `@media (prefers-color-scheme:dark){.ink,.val{fill:#ecebe8}.mut{fill:#8f8d88}.lbl{fill:#7d7b76}` +
96  `.trk,.div{fill:#3a3936}.pace{fill:#6a6863}.warm{fill:#e0a64b}.hot{fill:#e2745c}}`
97
98type Part = { w: number; draw: (x: number) => string }
99
100const esc = (s: string) => s.replace(/&/g, '&amp;').replace(/</g, '&lt;')
101
102const gap = (w: number): Part => ({ w, draw: () => '' })
103
104const text = (s: string, cls: string): Part => ({
105  w: s.length * CHAR,
106  draw: x => `<text x="${x}" y="${MID}" class="${cls}">${esc(s)}</text>`,
107})
108
109const label = (s: string): Part => ({
110  w: s.length * LABEL_FONT * 0.602 + (s.length - 1) * LABEL_TRACK,
111  draw: x => `<text x="${x}" y="${MID}" class="lbl">${esc(s.toUpperCase())}</text>`,
112})
113
114const slot = (p: Part, w: number, align: 'start' | 'end'): Part => ({
115  w,
116  draw: x => p.draw(align === 'end' ? x + w - p.w : x),
117})
118
119const divider: Part = {
120  w: 1,
121  draw: x => `<rect x="${x}" y="${MID - 6}" width="1" height="12" class="div"/>`,
122}
123
124const paceTick = (x: number) =>
125  `<rect x="${x - 0.75}" y="${MID - 5}" width="1.5" height="10" rx=".75" class="pace"/>`
126
127const clamp = (pct: number) => Math.min(100, Math.max(0, pct)) / 100
128
129// A solid rounded bar: the 5h window.
130function bar(pct: number, level: Level, pace?: number): Part {
131  const W = 34
132  return {
133    w: W,
134    draw: x =>
135      `<rect x="${x}" y="${MID - 2.5}" width="${W}" height="5" rx="2.5" class="trk"/>` +
136      `<rect x="${x}" y="${MID - 2.5}" width="${pct > 0 ? Math.max(5, W * clamp(pct)) : 0}" height="5" rx="2.5" class="${INK[level]}"/>` +
137      (pace !== undefined ? paceTick(x + W * pace) : ''),
138  }
139}
140
141// One dot per day: the 7d window.
142function dots(pct: number, level: Level, pace?: number): Part {
143  const N = 7
144  const STEP = 6
145  const W = STEP * N - 2
146  const filled = Math.round(clamp(pct) * N)
147  return {
148    w: W,
149    draw: x =>
150      Array.from({ length: N }, (_, i) =>
151        `<circle cx="${x + 2 + i * STEP}" cy="${MID}" r="2" class="${i < filled ? INK[level] : 'trk'}"/>`,
152      ).join('') + (pace !== undefined ? paceTick(x + W * pace) : ''),
153  }
154}
155
156// Input against output: how much of the traffic went in.
157function ratio(input: number, output: number): Part {
158  const W = 24
159  const share = input + output > 0 ? input / (input + output) : 0
160  return {
161    w: W,
162    draw: x =>
163      `<rect x="${x}" y="${MID - 1.25}" width="${W}" height="2.5" rx="1.25" class="trk"/>` +
164      `<rect x="${x}" y="${MID - 1.25}" width="${W * share}" height="2.5" rx="1.25" class="ink"/>`,
165  }
166}
167
168function strip(parts: Part[]): string {
169  const width = Math.ceil(parts.reduce((sum, p) => sum + p.w, 0))
170  let x = 0
171  const body = parts.map(p => {
172    const out = p.draw(x)
173    x += p.w
174    return out
175  })
176  return (
177    `<svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${H}" viewBox="0 0 ${width} ${H}">` +
178    `<style>${STYLE}</style>${body.join('')}</svg>`
179  )
180}
181
182function windowCell(w: Window, i: number): Part[] {
183  const meter = i === 0 ? bar(w.pct, w.level, w.pace) : dots(w.pct, w.level, w.pace)
184  return [
185    label(w.label),
186    gap(6),
187    meter,
188    gap(6),
189    text(fmtPct(w.pct), 'val'),
190  ]
191}
192
193// ---------- the band ----------
194
195export const register: Register = on => {
196  on('session.start', async ($, e, next) => {
197    const usage = await $.session.usage()
198    await update($, snapshot, () => toSnapshot(usage))
199    await $.command.register({
200      name: 'usage-band',
201      description: 'Show or hide the usage band above the prompt',
202    })
203    // Redraw so the reset countdowns and pace markers stay current.
204    $.clock.every(30_000, () => void update($, tick, n => n + 1))
205
206    return next(e)
207  })
208
209  on('command.run', { command: 'usage-band' }, async $ => {
210    const hidden = await update($, isHidden, v => !v)
211
212    return { text: hidden ? 'Usage band hidden.' : 'Usage band shown.' }
213  })
214
215  on('session.measure', async ($, e, next) => {
216    await update($, snapshot, () => toSnapshot(e))
217
218    return next(e)
219  })
220
221  on('turn.step', async function* ($, e, next) {
222    const result = yield* next(e)
223    const u = result.usage
224    if (u) {
225      await update($, tokens, t => ({
226        input: t.input + u.input_tokens,
227        output: t.output + u.output_tokens,
228        cacheRead: t.cacheRead + u.cache_read_input_tokens,
229        cacheWrite: t.cacheWrite + u.cache_creation_input_tokens,
230      }))
231    }
232    // Cost and limits move with every request, not only when the turn ends.
233    const usage = await $.session.usage()
234    await update($, snapshot, () => toSnapshot(usage))
235
236    return result
237  })
238
239  on('session.end', async ($, e, next) => {
240    if (e.reason === 'clear') {
241      await update($, tokens, () => ZERO)
242    }
243
244    return next(e)
245  })
246
247  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
248    if (e.props.hasSurvey || (await read($, isHidden))) {
249      return next(e)
250    }
251
252    await read($, tick)
253    const snap = await read($, snapshot)
254    const t = await read($, tokens)
255    const now = await $.clock.now()
256    const windows = (snap?.windows ?? []).map(w => describe(w, now))
257    const sent = t.input + t.cacheWrite
258    const cost = snap?.costUsd !== undefined ? `$${snap.costUsd.toFixed(2)}` : undefined
259    const limitsAlt = windows
260      .map(w => `${w.label} window ${w.pct}% used${w.remaining ? `, resets in ${w.remaining}` : ''}`)
261      .join('; ')
262    const ioAlt = `${fmtTokens(sent)} input tokens, ${fmtTokens(t.output)} output tokens`
263
264    // Limits left, tokens centred, cost right: the two sides share the room equally.
265    if (e.surface === 'desktop') {
266      const { Box, Svg } = $.ui.resolve({ ...e, surface: 'desktop' as const })
267      const limits: Part[] = []
268      windows.forEach((w, i) => {
269        if (i > 0) limits.push(gap(9), divider, gap(9))
270        limits.push(...windowCell(w, i))
271      })
272      const up = text(`↑${fmtTokens(sent)}`, 'val')
273      const down = text(`↓${fmtTokens(t.output)}`, 'val')
274      const side = Math.max(up.w, down.w)
275      const io = [slot(up, side, 'end'), gap(6), ratio(sent, t.output), gap(6), slot(down, side, 'start')]
276
277      return (
278        <Box flexDirection="row" flexWrap="nowrap" alignItems="center" width="100%">
279          <Box key="limits" flexGrow={1} width={0} flexDirection="row" justifyContent="flex-start">
280            {limits.length > 0 && <Svg key="svg" source={strip(limits)} alt={limitsAlt} />}
281          </Box>
282          <Box key="io" flexShrink={0}>
283            <Svg key="svg" source={strip(io)} alt={ioAlt} />
284          </Box>
285          <Box key="cost" flexGrow={1} width={0} flexDirection="row" justifyContent="flex-end">
286            {cost && <Svg key="svg" source={strip([text(cost, 'val')])} alt={`Session cost ${cost}`} />}
287          </Box>
288        </Box>
289      )
290    }
291
292    // Terminal: the same strip as one line, coloured only where it needs attention.
293    const { Box, Text } = $.ui.resolve(e)
294    const color = { ok: undefined, warm: 'yellow', hot: 'red' }
295    const meter = (pct: number, n: number) => {
296      const filled = Math.round((Math.min(100, Math.max(0, pct)) / 100) * n)
297      return '▰'.repeat(filled) + '▱'.repeat(n - filled)
298    }
299
300    return (
301      <Box flexDirection="row" flexWrap="nowrap" width="100%">
302        <Box key="limits" flexGrow={1} width={0} flexDirection="row" columnGap={2}>
303          {windows.map((w, i) => (
304            <Text key={`w${i}`}>
305              <Text dimColor>{w.label.toUpperCase()} </Text>
306              <Text color={color[w.level]}>{meter(w.pct, i === 0 ? 6 : 7)}</Text>
307              <Text bold> {fmtPct(w.pct)}</Text>
308            </Text>
309          ))}
310        </Box>
311        <Text key="io">
312          <Text bold>↑{fmtTokens(sent)}</Text>
313          <Text bold> ↓{fmtTokens(t.output)}</Text>
314        </Text>
315        <Box key="cost" flexGrow={1} width={0} flexDirection="row" justifyContent="flex-end">
316          {cost && <Text bold>{cost}</Text>}
317        </Box>
318      </Box>
319    )
320  })
321}
322

1export type UsageWindow = { kind: string; percentUsed: number; resetsAt?: string }
2
3export type UsageSnapshot = {
4  windows: UsageWindow[]
5  costUsd?: number
6  contextPercent?: number
7}
8
9export type TokenTotals = {
10  input: number
11  output: number
12  cacheRead: number
13  cacheWrite: number
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    'usage-band': {
19      snapshot: UsageSnapshot | null
20      tokens: TokenTotals
21      tick: number
22      isHidden: boolean
23    }
24  }
25}
26