usage-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, TurnUsage } from 'claude-code'
3
4import type { Pricing, RateLimit, Totals, WindowTokens } from '../types'
5
6const limitsAtom = atom({ plugin: 'usage-band', key: 'limits' } as const, [] as RateLimit[])
7const totalsAtom = atom({ plugin: 'usage-band', key: 'totals' } as const, null as Totals | null)
8const sessionAtom = atom({ plugin: 'usage-band', key: 'session' } as const, null as WindowTokens | null)
9const lastTurnAtom = atom({ plugin: 'usage-band', key: 'lastTurn' } as const, null as WindowTokens | null)
10
11const HOUR = 3600_000
12const WINDOWS = { fiveHour: { kind: 'five_hour', ms: 5 * HOUR }, week: { kind: 'seven_day', ms: 168 * HOUR } }
13
14/** Per-minute buckets from scan.py: minute -> [input, output, cacheWrite, cacheRead, cost]. */
15let buckets: Record<string, number[]> = {}
16/** Per-session totals from scan.py: session id -> [tokens, cacheRead, cost]. */
17let sessions: Record<string, number[]> = {}
18let pricing: Pricing | null = null
19
20export const sumSince = (from: number, b: Record<string, number[]>): WindowTokens => {
21  const fromMinute = Math.floor(from / 60_000)
22  let tokens = 0
23  let cacheRead = 0
24  let cost = 0
25  for (const [minute, [i = 0, o = 0, cw = 0, cr = 0, usd = 0]] of Object.entries(b)) {
26    if (Number(minute) < fromMinute) continue
27    tokens += i + o + cw
28    cacheRead += cr
29    cost += usd
30  }
31  return { tokens, cacheRead, cost }
32}
33
34/**
35 * USD for a turn's summed usage. A turn's usage has no cache-write split, so
36 * writes are priced at the 1-hour rate Claude Code uses; 0 for an unknown model.
37 */
38export const turnCost = (p: Pricing, u: TurnUsage) => {
39  const base = Object.keys(p.models)
40    .sort((a, b) => b.length - a.length)
41    .find(m => u.model.startsWith(m))
42  const rates = base === undefined ? undefined : p.models[base]
43  if (!rates) return 0
44  const [input, output, read] = rates
45  return (
46    (u.input_tokens * input +
47      u.output_tokens * output +
48      u.cache_creation_input_tokens * input * p.cacheWrite1h +
49      u.cache_read_input_tokens * read) /
50    1e6
51  )
52}
53
54export const fmtCost = (usd: number) =>
55  usd >= 100 ? `$${usd.toFixed(0)}` : usd >= 0.01 || usd === 0 ? `$${usd.toFixed(2)}` : '<$0.01'
56
57/** Where a window began: its reset time minus its length, else a rolling window. */
58const windowStart = (limits: RateLimit[], kind: string, ms: number, now: number) => {
59  const resetsAt = limits.find(l => l.kind === kind)?.resetsAt
60  const reset = resetsAt ? Date.parse(resetsAt) : NaN
61  return Number.isFinite(reset) && reset > now ? reset - ms : now - ms
62}
63
64export const fmtTokens = (n: number) =>
65  n >= 1e9 ? `${(n / 1e9).toFixed(2)}B` : n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${(n / 1e3).toFixed(1)}K` : `${n}`
66
67const pad = (n: number) => String(n).padStart(2, '0')
68const fmtReset = (iso: string | undefined, withDate: boolean) => {
69  if (!iso) return ''
70  const d = new Date(iso)
71  if (Number.isNaN(d.getTime())) return ''
72  const hm = `${pad(d.getHours())}:${pad(d.getMinutes())}`
73  return withDate ? `${d.getMonth() + 1}/${d.getDate()} ${hm}` : hm
74}
75
76async function recompute($: EngineInterface) {
77  const limits = await read($, limitsAtom)
78  const now = await $.clock.now()
79  const totals: Totals = {
80    fiveHour: sumSince(windowStart(limits, WINDOWS.fiveHour.kind, WINDOWS.fiveHour.ms, now), buckets),
81    week: sumSince(windowStart(limits, WINDOWS.week.kind, WINDOWS.week.ms, now), buckets),
82  }
83  await update($, totalsAtom, () => totals)
84
85  const [tokens = 0, cacheRead = 0, cost = 0] = sessions[await $.session.id()] ?? []
86  await update($, sessionAtom, () => ({ tokens, cacheRead, cost }))
87}
88
89export const register: Register = on => {
90  on('session.start', async ($, e, next) => {
91    const started = await next(e)
92    pricing = await $.fs
93      .read(`${$.plugin.root}/hooks/pricing.json`)
94      .then(text => JSON.parse(text) as Pricing)
95      .catch(() => null)
96    const usage = await $.session.usage().catch(() => undefined)
97    if (usage) await update($, limitsAtom, () => usage.rateLimits)
98
99    // A scanner for the session's life: one JSON line of buckets a minute.
100    void (async () => {
101      const scan = $.process.spawn({ argv: ['python3', `${$.plugin.root}/hooks/scan.py`, '60'] })
102      let pending = ''
103      for await (const chunk of scan) {
104        if (!('stream' in chunk) || chunk.stream !== 'stdout') continue
105        pending += chunk.text
106        let nl
107        while ((nl = pending.indexOf('\n')) >= 0) {
108          const line = pending.slice(0, nl)
109          pending = pending.slice(nl + 1)
110          try {
111            const parsed = JSON.parse(line)
112            buckets = parsed.b ?? {}
113            sessions = parsed.s ?? {}
114            await recompute($)
115          } catch {}
116        }
117      }
118    })()
119
120    return started
121  })
122
123  on('session.measure', async ($, e, next) => {
124    if (e.changed.includes('rateLimits')) {
125      await update($, limitsAtom, () => e.rateLimits)
126      await recompute($)
127    }
128    return next(e)
129  })
130
131  on('turn.complete', async ($, e, next) => {
132    const result = await next(e)
133    const u = result.usage
134    if (u && e.agentId === undefined) {
135      const turn = {
136        tokens: u.input_tokens + u.output_tokens + u.cache_creation_input_tokens,
137        cacheRead: u.cache_read_input_tokens,
138        cost: pricing ? turnCost(pricing, u) : 0,
139      }
140      await update($, lastTurnAtom, () => turn)
141    }
142    return result
143  })
144
145  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
146    if (e.props.hasSurvey) return next(e)
147    const totals = await read($, totalsAtom)
148    const limits = await read($, limitsAtom)
149    const session = await read($, sessionAtom)
150    const lastTurn = await read($, lastTurnAtom)
151    if (!totals && limits.length === 0) return next(e)
152
153    const { Box, Text } = $.ui.resolve(e)
154    const part = (label: string, w: WindowTokens | undefined, kind: string, withDate: boolean) => {
155      const limit = limits.find(l => l.kind === kind)
156      const pct = limit ? limit.percentUsed : undefined
157      const color = pct === undefined ? undefined : pct >= 90 ? 'red' : pct >= 70 ? 'yellow' : 'green'
158      const reset = fmtReset(limit?.resetsAt, withDate)
159      return (
160        <Text>
161          <Text bold>{label} </Text>
162          <Text>{w ? `${fmtTokens(w.tokens)} tok` : '…'}</Text>
163          {w ? <Text color="cyan"> {fmtCost(w.cost)}</Text> : null}
164          {w && w.cacheRead > 0 ? <Text dimColor> (+{fmtTokens(w.cacheRead)} cache read)</Text> : null}
165          {pct !== undefined ? <Text color={color}> {pct}%</Text> : null}
166          {reset ? <Text dimColor> ↻{reset}</Text> : null}
167        </Text>
168      )
169    }
170
171    const tok = (label: string, w: WindowTokens) => (
172      <Text>
173        <Text bold>{label} </Text>
174        <Text>{fmtTokens(w.tokens)} tok</Text>
175        <Text color="cyan"> {fmtCost(w.cost)}</Text>
176        {w.cacheRead > 0 ? <Text dimColor> (+{fmtTokens(w.cacheRead)} cache read)</Text> : null}
177      </Text>
178    )
179
180    return (
181      <Box flexDirection="column">
182        <Text>
183          {part('5h', totals?.fiveHour, 'five_hour', false)}
184          <Text dimColor>  │  </Text>
185          {part('Week', totals?.week, 'seven_day', true)}
186        </Text>
187        {session || lastTurn ? (
188          <Text>
189            {session ? tok('Session', session) : null}
190            {session && lastTurn ? <Text dimColor>  │  </Text> : null}
191            {lastTurn ? tok('Last turn', lastTurn) : null}
192          </Text>
193        ) : null}
194      </Box>
195    )
196  })
197}
198

1/** One rate-limit window as `$.session.usage()` reports it. */
2export type RateLimit = { kind: string; percentUsed: number; resetsAt?: string }
3
4/**
5 * Token totals of one window: input + output + cache writes, cache reads apart,
6 * and what all four cost in USD at API rates.
7 */
8export type WindowTokens = { tokens: number; cacheRead: number; cost: number }
9
10/** hooks/pricing.json: USD per million tokens, [input, output, cache read]. */
11export type Pricing = {
12  cacheWrite5m: number
13  cacheWrite1h: number
14  fast: Record<string, number>
15  models: Record<string, [number, number, number]>
16}
17
18export type Totals = { fiveHour: WindowTokens; week: WindowTokens }
19
20declare module 'claude-code' {
21  interface PluginState {
22    'usage-band': {
23      limits: RateLimit[]
24      totals: Totals | null
25      /** This conversation's tokens, its subagents included. */
26      session: WindowTokens | null
27      /** The last main-loop turn's tokens. */
28      lastTurn: WindowTokens | null
29    }
30  }
31}
32