usage-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionUsage } from 'claude-code'
3
4import type { Usage } from '../types'
5
6const usage = atom({ plugin: 'usage-band', key: 'usage' } as const, null)
7
8const tokens = atom({ plugin: 'usage-band', key: 'tokens' } as const, 0)
9
10const LABELS: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
11
12export const slim = (u: Pick<SessionUsage, 'rateLimits' | 'cost'>): Usage => ({
13  limits: u.rateLimits.map(({ kind, percentUsed, resetsAt }) => ({ kind, percentUsed, resetsAt })),
14  usd: u.cost?.usd,
15})
16
17const minsUntil = (iso: string | undefined, now: number) =>
18  iso ? Math.max(0, Math.round((Date.parse(iso) - now) / 60000)) : undefined
19
20const until = (iso: string | undefined, now: number) => {
21  const mins = minsUntil(iso, now)
22  if (mins === undefined) return ''
23  if (mins >= 1440) return ` (${Math.floor(mins / 1440)}d ${Math.floor((mins % 1440) / 60)}h)`
24  return mins >= 60 ? ` (${Math.floor(mins / 60)}h ${mins % 60}m)` : ` (${mins}m)`
25}
26
27const WIDTH = 7
28
29// Solid blocks, the empty part drawn dim.
30export const bar = (pct: number) => {
31  const filled = Math.min(WIDTH, Math.max(0, Math.round((pct / 100) * WIDTH)))
32  return { filled: '█'.repeat(filled), empty: '░'.repeat(WIDTH - filled) }
33}
34
35const MOODS: [number, string][] = [
36  [1000, '🙀'],
37  [500, '🤑'],
38  [100, '😛'],
39  [50, '🫪'],
40  [25, '😮‍💨'],
41  [10, '🥱'],
42  [5, '😒'],
43  [0, '💩'],
44]
45
46export const mood = (usd: number) => MOODS.find(([over]) => usd > over)?.[1] ?? '💩'
47
48export const compact = (n: number) =>
49  n >= 1e9 ? `${(n / 1e9).toFixed(1)}B` : n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${(n / 1e3).toFixed(1)}k` : `${n}`
50
51// Pacing against the 5-hour window: unused capacity is lost at reset, so spend
52// freely near a reset with headroom; otherwise compare use with time elapsed.
53export const pace = (pct: number, minsLeft: number) => {
54  const expected = ((300 - minsLeft) / 300) * 100
55  if (pct >= 90) return { text: '🛑 nearly out', color: 'error' }
56  if (minsLeft <= 60 && pct <= 70) return { text: '🚀 spend freely', color: 'success' }
57  if (pct > expected + 15) return { text: '🐢 slow down', color: 'warning' }
58  return { text: '👌 on pace', color: 'subtle' }
59}
60
61const tone = (pct: number) => (pct >= 90 ? 'error' : pct >= 70 ? 'warning' : 'success')
62
63export const items = (u: Usage, now: number, spent = 0) => [
64  ...u.limits.map(l => ({
65    key: l.kind,
66    label: LABELS[l.kind] ?? l.kind,
67    bar: bar(l.percentUsed),
68    color: tone(l.percentUsed),
69    text: `${l.percentUsed}%${until(l.resetsAt, now)}`,
70  })),
71  ...(u.usd === undefined ? [] : [{ key: 'cost', bar: undefined, color: 'subtle', text: `$${u.usd.toFixed(2)}   ${mood(u.usd)}` }]),
72  { key: 'tokens', bar: undefined, color: 'subtle', text: `${compact(spent)} tokens` },
73  ...u.limits
74    .filter(l => l.kind === 'five_hour')
75    .flatMap(l => {
76      const left = minsUntil(l.resetsAt, now)
77      return left === undefined ? [] : [{ key: 'pace', bar: undefined, ...pace(l.percentUsed, left) }]
78    }),
79]
80
81// Sums the whole session's transcripts, earlier turns and subagents included.
82const recount = async ($: EngineInterface) => {
83  const ran = await $.process.run(['python3', '-I', `${$.plugin.root}/hooks/tokens.py`, await $.session.id()])
84  const n = Number(ran.stdout.trim())
85  if (ran.exitCode === 0 && Number.isFinite(n)) await update($, tokens, () => n)
86}
87
88export const register: Register = on => {
89  on('session.start', async ($, e, next) => {
90    const ran = await next(e)
91    const u = await $.session.usage()
92    await update($, usage, () => slim(u))
93    // Cost moves with every API response mid-turn; poll it so the emoji flips at each milestone.
94    let lastUsd = u.cost?.usd
95    $.clock.every(3000, async () => {
96      const now = await $.session.usage()
97      if (now.cost?.usd === lastUsd) return
98      lastUsd = now.cost?.usd
99      await update($, usage, () => slim(now))
100    })
101    await recount($)
102    return ran
103  })
104
105  on('session.measure', async ($, e, next) => {
106    await update($, usage, () => slim(e))
107    await recount($)
108    return next(e)
109  })
110
111  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
112    const u = await read($, usage)
113    if (u === null || e.props.hasSurvey) return next(e)
114
115    const { Box, Text } = $.ui.resolve(e)
116    const now = await $.clock.now()
117    return (
118      <Box flexDirection="row" columnGap={5} paddingX={1}>
119        {items(u, now, await read($, tokens)).map(i => (
120          <Box key={i.key} flexDirection="row" columnGap={1}>
121            {'label' in i ? <Text dimColor>{i.label}</Text> : null}
122            {i.bar ? (
123              <Text>
124                <Text color={i.color}>{i.bar.filled}</Text>
125                <Text dimColor>{i.bar.empty}</Text>
126              </Text>
127            ) : null}
128            <Text dimColor>{i.text}</Text>
129          </Box>
130        ))}
131      </Box>
132    )
133  })
134}
135

1export type Usage = {
2  limits: { kind: string; percentUsed: number; resetsAt?: string }[]
3  usd?: number
4}
5
6declare module 'claude-code' {
7  interface PluginState {
8    'usage-band': { usage: Usage | null; tokens: number }
9  }
10}
11