usage-band

1import { atom, read, update } from 'claude-code'
2import type { Color, EngineInterface, Register, SessionContextBreakdown } from 'claude-code'
3
4import type { Breakdown, Limit, Snapshot } from '../types'
5
6const snapshot = atom({ plugin: 'usage-band', key: 'snapshot' } as const, null)
7const effort = atom({ plugin: 'usage-band', key: 'effort' } as const, null)
8
9// Mid-tone colours that read on both dark and light terminals.
10const PALETTE = {
11  model: '#d97757',
12  folder: '#0ea5e9',
13  branch: '#84cc16',
14  ok: '#22c55e',
15  warn: '#eab308',
16  high: '#ef4444',
17}
18
19const EFFORT_COLORS: Record<string, string> = {
20  low: '#94a3b8',
21  medium: '#3b82f6',
22  high: '#a855f7',
23  xhigh: '#d946ef',
24  max: '#ec4899',
25}
26
27// Each context category gets a colour of its own, so neighbours in the bar never blur.
28const CATEGORIES: Record<string, { label: string; color: string }> = {
29  Messages: { label: 'chat', color: '#3b82f6' },
30  'System tools': { label: 'tools', color: '#f59e0b' },
31  'System prompt': { label: 'sys', color: '#8b5cf6' },
32  Skills: { label: 'skills', color: '#10b981' },
33  'MCP server instructions': { label: 'mcp', color: '#ec4899' },
34  'MCP tools': { label: 'mcp tools', color: '#d946ef' },
35  'Memory files': { label: 'memory', color: '#14b8a6' },
36  'Custom agents': { label: 'agents', color: '#f97316' },
37}
38
39function category(name: string, fallback: string): { label: string; color: string } {
40  return CATEGORIES[name] ?? { label: name.toLowerCase(), color: fallback }
41}
42
43const LIMIT_NAMES: Record<string, string> = {
44  five_hour: 'session',
45  seven_day: 'weekly',
46  spend_limit: 'spend',
47}
48
49const BAR_WIDTH = 24
50
51// Terminal cells a plain string takes; every character the rows draw is one cell wide.
52function width(text: string): number {
53  return [...text].length
54}
55
56// Always in thousands: whole k from 1000 up, one decimal below it (631 → 0.6k).
57function tokens(n: number): string {
58  if (n <= 0) return '0k'
59  if (n >= 950) return `${Math.round(n / 1000)}k`
60
61  return `${Math.max(0.1, Math.round(n / 100) / 10).toFixed(1)}k`
62}
63
64function color(percent: number): string {
65  return percent >= 90 ? PALETTE.high : percent >= 70 ? PALETTE.warn : PALETTE.ok
66}
67
68function resetsIn(limit: Limit, now: number): string {
69  if (!limit.resetsAt) return ''
70  const minutes = Math.max(0, Math.round((Date.parse(limit.resetsAt) - now) / 60000))
71  const hours = Math.floor(minutes / 60)
72
73  if (hours >= 24) return ` (resets ${Math.floor(hours / 24)}d${hours % 24}h)`
74
75  return ` (resets ${hours}h${minutes % 60}m)`
76}
77
78function basename(path: string): string {
79  const parts = path.split('/').filter(Boolean)
80
81  return parts[parts.length - 1] ?? '/'
82}
83
84async function branchOf($: EngineInterface): Promise<string | null> {
85  try {
86    const { exitCode, stdout } = await $.process.run(['git', 'branch', '--show-current'], { timeoutMs: 3000 })
87    const name = stdout.trim()
88
89    return exitCode === 0 && name !== '' ? name : null
90  } catch {
91    return null
92  }
93}
94
95async function effortFromSettings($: EngineInterface): Promise<string | null> {
96  try {
97    const value = (await $.settings.read())['effortLevel']
98
99    return typeof value === 'string' ? value : null
100  } catch {
101    return null
102  }
103}
104
105function toBreakdown(b: SessionContextBreakdown | undefined): Breakdown | null {
106  if (!b) return null
107
108  return {
109    parts: b.categories
110      .filter(c => !c.isDeferred)
111      .map(c => ({ name: c.name, tokens: c.tokens, color: c.color, kind: c.kind })),
112    total: b.totalTokens,
113    max: b.rawMaxTokens,
114  }
115}
116
117// One bar for the whole window: a coloured run per used category, then free space.
118function segments(b: Breakdown, width: number): { color: string; text: string }[] {
119  const out: { color: string; text: string }[] = []
120  let left = width
121  for (const part of b.parts) {
122    if (part.kind !== 'used' || part.tokens <= 0 || left <= 0) continue
123    const cells = Math.min(left, Math.max(1, Math.round((part.tokens / b.max) * width)))
124    out.push({ color: category(part.name, part.color).color, text: '█'.repeat(cells) })
125    left -= cells
126  }
127  if (left > 0) out.push({ color: 'inactive', text: '░'.repeat(left) })
128
129  return out
130}
131
132async function refresh($: EngineInterface) {
133  const [model, usage, cwd, branch] = await Promise.all([
134    $.session.model(),
135    $.session.usage({ breakdown: 'summary' }),
136    $.session.cwd(),
137    branchOf($),
138  ])
139  const snap: Snapshot = {
140    model,
141    folder: basename(cwd),
142    branch,
143    context: { tokens: usage.context.tokens, window: usage.context.window, percent: usage.context.percent },
144    rateLimits: usage.rateLimits.map(l => ({ kind: l.kind, percentUsed: l.percentUsed, resetsAt: l.resetsAt })),
145    costUsd: usage.cost?.usd,
146    breakdown: toBreakdown(usage.context.breakdown),
147  }
148  await update($, snapshot, () => snap)
149}
150
151export const register: Register = on => {
152  on('session.start', async ($, e, next) => {
153    const result = await next(e)
154    const level = await effortFromSettings($)
155    await update($, effort, current => current ?? level)
156    await refresh($)
157
158    return result
159  })
160
161  on('session.measure', async ($, e, next) => {
162    await refresh($)
163
164    return next(e)
165  })
166
167  on('turn.step', async function* ($, e, next) {
168    if (e.agentId === undefined && e.effort !== undefined) {
169      const level = String(e.effort)
170      await update($, effort, () => level)
171    }
172
173    return yield* next(e)
174  })
175
176  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
177    const own = await next(e)
178    const snap = await read($, snapshot)
179
180    if (snap === null) {
181      return own
182    }
183
184    const level = await read($, effort)
185    const { Box, Text } = $.ui.resolve(e)
186    const ctx = snap.context
187    const b = snap.breakdown
188    // The exact count from the last reply when there is one, the estimate before it.
189    const used = ctx.tokens ?? b?.total
190    const window = ctx.tokens === undefined ? (b?.max ?? ctx.window) : ctx.window
191    const percent = ctx.percent ?? (used === undefined ? undefined : Math.round((used / window) * 100))
192    const now = await $.clock.now()
193
194    const sep = <Text dimColor>{'  │  '}</Text>
195    const ctxText =
196      percent === undefined
197        ? ` — / ${tokens(window)}`
198        : ` ${percent}% ${tokens(used ?? 0)}/${tokens(window)}`
199
200    // Pad the first part of each row to one width, so the next │ lines up across both rows.
201    const topLeft = width(
202      `${snap.model}${level === null ? '' : ` ${level}`}  │  ${snap.folder}${snap.branch === null ? '' : ` ⎇ ${snap.branch}`}`,
203    )
204    const bottomLeft = width(`ctx ${b === null ? '' : ' '.repeat(BAR_WIDTH)}${ctxText}`)
205    const column = Math.max(topLeft, bottomLeft)
206
207    return (
208      <Box flexDirection="column">
209        {own}
210        <Box>
211          <Text bold color={PALETTE.model}>{snap.model}</Text>
212          {level === null ? null : (
213            <Text color={EFFORT_COLORS[level] ?? EFFORT_COLORS.medium}>{` ${level}`}</Text>
214          )}
215          {sep}
216          <Text color={PALETTE.folder}>{snap.folder}</Text>
217          {snap.branch === null ? null : <Text color={PALETTE.branch}>{` ⎇ ${snap.branch}`}</Text>}
218          <Text>{' '.repeat(column - topLeft)}</Text>
219          {snap.rateLimits.length === 0 ? null : sep}
220          {snap.rateLimits.map((limit, i) => (
221            <Text key={limit.kind}>
222              <Text dimColor>{`${i === 0 ? '' : '  '}${LIMIT_NAMES[limit.kind] ?? limit.kind} `}</Text>
223              <Text bold color={color(limit.percentUsed)}>{`${limit.percentUsed}%`}</Text>
224              <Text dimColor>{resetsIn(limit, now)}</Text>
225            </Text>
226          ))}
227          {snap.costUsd === undefined ? null : sep}
228          {snap.costUsd === undefined ? null : <Text>{`$${snap.costUsd.toFixed(2)}`}</Text>}
229        </Box>
230        <Box>
231          <Text dimColor>{'ctx '}</Text>
232          {b === null
233            ? null
234            : segments(b, BAR_WIDTH).map((seg, i) => (
235                <Text key={`seg${i}`} color={seg.color as Color} dimColor={seg.color === 'inactive'}>{seg.text}</Text>
236              ))}
237          {percent === undefined ? (
238            <Text dimColor>{ctxText}</Text>
239          ) : (
240            <Text>
241              <Text bold color={color(percent)}>{` ${percent}%`}</Text>
242              <Text>{` ${tokens(used ?? 0)}`}</Text>
243              <Text dimColor>{`/${tokens(window)}`}</Text>
244            </Text>
245          )}
246          <Text>{' '.repeat(column - bottomLeft)}</Text>
247          {b === null ? null : sep}
248          {b === null
249            ? null
250            : b.parts
251                .filter(part => part.kind === 'used' && part.tokens > 0)
252                .sort((x, y) => y.tokens - x.tokens)
253                .map((part, i) => {
254                  const { label, color: hue } = category(part.name, part.color)
255
256                  return (
257                    <Text key={part.name}>
258                      <Text color={hue as Color}>{`${i === 0 ? '' : '  '}■ ${label}`}</Text>
259                      <Text dimColor>{` ${tokens(part.tokens)}`}</Text>
260                    </Text>
261                  )
262                })}
263        </Box>
264      </Box>
265    )
266  })
267}
268

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3export type Part = { name: string; tokens: number; color: string; kind: string }
4
5export type Breakdown = { parts: Part[]; total: number; max: number }
6
7export type Snapshot = {
8  model: string
9  folder: string
10  branch: string | null
11  context: { tokens?: number; window: number; percent?: number }
12  rateLimits: Limit[]
13  costUsd?: number
14  breakdown: Breakdown | null
15}
16
17declare module 'claude-code' {
18  interface PluginState {
19    'usage-band': { snapshot: Snapshot | null; effort: string | null }
20  }
21}
22