usage-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { ContextSnapshot, LimitWindow } from '../types'
5
6const context = atom({ plugin: 'usage-band', key: 'context' } as const, null)
7const limits = atom({ plugin: 'usage-band', key: 'limits' } as const, [])
8
9const TOKYO_NIGHT = {
10  blue: '#7aa2f7',
11  magenta: '#bb9af7',
12  red: '#f7768e',
13  fg: '#a9b1d6',
14  track: '#414868',
15}
16
17const LIMIT_STYLE: Record<string, { label: string; color: string }> = {
18  five_hour: { label: '5h', color: TOKYO_NIGHT.blue },
19  seven_day: { label: 'week', color: TOKYO_NIGHT.magenta },
20}
21
22const LABEL_WIDTH = 8
23const META_WIDTH = 24
24
25export const formatTokens = (tokens: number): string => {
26  if (tokens >= 1_000_000) return `${+(tokens / 1_000_000).toFixed(1)}M`
27  if (tokens >= 10_000) return `${Math.round(tokens / 1000)}k`
28  if (tokens >= 1000) return `${+(tokens / 1000).toFixed(1)}k`
29  return `${tokens}`
30}
31
32export const formatResetIn = (resetsAt: string | undefined, now: number): string => {
33  if (!resetsAt) return ''
34  const minutes = Math.max(0, Math.round((Date.parse(resetsAt) - now) / 60_000))
35  const days = Math.floor(minutes / 1440)
36  const hours = Math.floor((minutes % 1440) / 60)
37  if (days > 0) return `resets in ${days}d ${hours}h`
38  if (hours > 0) return `resets in ${hours}h ${minutes % 60}m`
39  return `resets in ${minutes}m`
40}
41
42export const filledCells = (percent: number, width: number): number =>
43  Math.min(width, Math.max(0, Math.round((percent / 100) * width)))
44
45const refresh = async ($: EngineInterface): Promise<void> => {
46  const usage = await $.session.usage()
47  const windows: LimitWindow[] = usage.rateLimits.map(({ kind, percentUsed, resetsAt }) => ({ kind, percentUsed, resetsAt }))
48  await update($, limits, () => windows)
49
50  const { tokens, window, percent } = usage.context
51  if (tokens === undefined || percent === undefined) return
52  const snapshot: ContextSnapshot = { tokens, window, percent }
53  await update($, context, () => snapshot)
54}
55
56export const register: Register = on => {
57  on('session.start', async ($, e, next) => {
58    const result = await next(e)
59    await refresh($)
60    return result
61  })
62
63  on('turn.complete', async ($, e, next) => {
64    const result = await next(e)
65    await refresh($)
66    return result
67  })
68
69  on('session.compact', async ($, e, next) => {
70    const result = await next(e)
71    await refresh($)
72    return result
73  })
74
75  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
76    if (e.props.hasSurvey || e.props.view.agentId) return next(e)
77
78    const snapshot = await read($, context)
79    const windows = await read($, limits)
80    if (snapshot === null && windows.length === 0) return next(e)
81
82    const { Box, Text } = $.ui.resolve(e)
83    const now = await $.clock.now()
84    const barWidth = Math.max(10, e.props.bodyColumns - LABEL_WIDTH - META_WIDTH - 6)
85
86    const row = (key: string, label: string, color: string, percent: number, meta: string) => {
87      const filled = filledCells(percent, barWidth)
88      const isHigh = percent >= 80
89      return (
90        <Text key={key}>
91          <Text color={TOKYO_NIGHT.fg}>{label.padEnd(LABEL_WIDTH)}</Text>
92          <Text color={isHigh ? TOKYO_NIGHT.red : color}>{'▆'.repeat(filled)}</Text>
93          <Text color={TOKYO_NIGHT.track}>{'▆'.repeat(barWidth - filled)}</Text>
94          <Text bold color={isHigh ? TOKYO_NIGHT.red : color}>{` ${Math.round(percent)}%`.padStart(5)}</Text>
95          <Text dimColor>{`  ${meta}`}</Text>
96        </Text>
97      )
98    }
99
100    return (
101      <Box flexDirection="column" borderStyle="round" borderColor={TOKYO_NIGHT.track} paddingX={1}>
102        {snapshot
103          ? row('context', 'context', 'claude', snapshot.percent, `${formatTokens(snapshot.tokens)} / ${formatTokens(snapshot.window)}`)
104          : null}
105        {windows.map(w => {
106          const style = LIMIT_STYLE[w.kind] ?? { label: w.kind, color: TOKYO_NIGHT.blue }
107          return row(w.kind, style.label, style.color, w.percentUsed, formatResetIn(w.resetsAt, now))
108        })}
109      </Box>
110    )
111  })
112}
113

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