usage-band

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Limit, Tokens } from '../types'
5
6const limits = atom({ plugin: 'usage-band', key: 'limits' } as const, [] as Limit[])
7const tokens = atom(
8  { plugin: 'usage-band', key: 'tokens' } as const,
9  { input: 0, output: 0, cached: 0 } as Tokens
10)
11
12const BAR = 16
13
14const compact = (n: number): string =>
15  n >= 1_000_000 ? `${(n / 1_000_000).toFixed(1)}M` : n >= 1_000 ? `${(n / 1_000).toFixed(1)}K` : `${n}`
16
17const untilReset = (iso: string | undefined, now: number): string => {
18  if (!iso) return ''
19  const minutes = Math.max(0, Math.round((Date.parse(iso) - now) / 60_000))
20  const days = Math.floor(minutes / 1440)
21  const hours = Math.floor((minutes % 1440) / 60)
22
23  if (days > 0) return `RESET ${days}D ${hours}H`
24  if (hours > 0) return `RESET ${hours}H ${minutes % 60}M`
25
26  return `RESET ${minutes}M`
27}
28
29// Above 80: red; above 70: yellow; otherwise green.
30const tone = (percent: number): string => (percent > 80 ? 'red' : percent > 70 ? 'yellow' : 'green')
31
32export const register: Register = on => {
33  on('session.start', async ($, e, next) => {
34    const usage = await $.session.usage()
35    await update($, limits, () => usage.rateLimits)
36
37    return next(e)
38  })
39
40  on('session.measure', async ($, e, next) => {
41    if (e.changed.includes('rateLimits')) {
42      await update($, limits, () => e.rateLimits)
43    }
44
45    return next(e)
46  })
47
48  on('turn.complete', async ($, e, next) => {
49    const u = e.usage
50
51    // Main loop only: a subagent's turns have an agentId.
52    if (u && !e.agentId) {
53      await update($, tokens, t => ({
54        input: t.input + u.input_tokens + u.cache_creation_input_tokens,
55        output: t.output + u.output_tokens,
56        cached: t.cached + u.cache_read_input_tokens
57      }))
58    }
59
60    return next(e)
61  })
62
63  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
64    if (e.props.hasSurvey) {
65      return next(e)
66    }
67
68    const { Box, Text } = $.ui.resolve(e)
69    const windows = await read($, limits)
70    const spent = await read($, tokens)
71    const now = await $.clock.now()
72    const five = windows.find(w => w.kind === 'five_hour')
73    const week = windows.find(w => w.kind === 'seven_day')
74
75    const label = (text: string) => (
76      <Text weight="bold" color="cyan" wrap="truncate-end">
77        {text}
78      </Text>
79    )
80
81    const meter = (w: Limit | undefined, name: string) => {
82      const percent = w ? w.percentUsed : 0
83      const color = tone(percent)
84      const filled = Math.min(BAR, Math.round((percent / 100) * BAR))
85
86      return (
87        <Box key={name} gap={1} flexShrink={0}>
88          {label(name)}
89          <Box flexShrink={0}>
90            <Text backgroundColor={w ? color : 'gray'} wrap="truncate-end">
91              {' '.repeat(filled)}
92            </Text>
93            <Text backgroundColor="gray" dimColor wrap="truncate-end">
94              {' '.repeat(BAR - filled)}
95            </Text>
96          </Box>
97          <Text weight="bold" color={color} tabular wrap="truncate-end">
98            {w ? `${Math.round(percent)}%` : '—'}
99          </Text>
100          <Text dimColor wrap="truncate-end">
101            {w ? untilReset(w.resetsAt, now) : ''}
102          </Text>
103        </Box>
104      )
105    }
106
107    return (
108      <Box borderStyle="round" borderColor="gray" borderDimColor paddingX={1} gap={3} flexWrap="wrap">
109        {meter(five, '5H')}
110        {meter(week, 'WEEK')}
111        <Box gap={1} flexShrink={0}>
112          {label('TOKENS')}
113          <Text weight="bold" color="green" wrap="truncate-end">
114            ↑ {compact(spent.input)}
115          </Text>
116          <Text weight="bold" color="red" wrap="truncate-end">
117            ↓ {compact(spent.output)}
118          </Text>
119          <Text dimColor wrap="truncate-end">
120            CACHE {compact(spent.cached)}
121          </Text>
122        </Box>
123      </Box>
124    )
125  })
126}
127

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3export type Tokens = { input: number; output: number; cached: number }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'usage-band': { limits: Limit[]; tokens: Tokens }
8  }
9}
10