usage-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PluginOptions, Register, SessionRateLimit } from 'claude-code'
3
4import type { Limit } from '../types'
5import { DEFAULT_STYLE, barSvg, textLine } from './bar'
6import type { Style } from './bar'
7
8const snapshot = atom({ plugin: 'usage-bar', key: 'snapshot' } as const, {
9  limits: [],
10  limitsAt: null,
11  costUsd: null,
12})
13const counts = atom({ plugin: 'usage-bar', key: 'counts' } as const, {
14  input: 0,
15  output: 0,
16  cache: 0,
17})
18
19const isHidden = atom({ plugin: 'usage-bar', key: 'isHidden' } as const, false)
20
21// The store key every session writes its latest reading under, so an idle
22// session can show what a busy one last saw.
23const SHARED = 'limits'
24
25type Shared = { limits: Limit[]; at: number }
26
27function limitsOf(rateLimits: SessionRateLimit[]): Limit[] {
28  return rateLimits.map(one => ({
29    kind: one.kind,
30    percentUsed: one.percentUsed,
31    resetsAt: one.resetsAt ?? null,
32  }))
33}
34
35function sharedOf(stored: unknown): Shared | null {
36  if (typeof stored !== 'object' || stored === null) {
37    return null
38  }
39
40  const { limits, at } = stored as { limits?: unknown; at?: unknown }
41  const isValid =
42    typeof at === 'number' &&
43    Number.isFinite(at) &&
44    Array.isArray(limits) &&
45    limits.every(
46      (one: unknown) =>
47        typeof one === 'object' &&
48        one !== null &&
49        typeof (one as Limit).kind === 'string' &&
50        Number.isFinite((one as Limit).percentUsed) &&
51        (typeof (one as Limit).resetsAt === 'string' || (one as Limit).resetsAt === null),
52    )
53
54  return isValid ? { limits: limits as Limit[], at: at as number } : null
55}
56
57// The bar only observes the session: nothing it does may fail an event of
58// the session's, so each piece of its own work is tried on its own.
59async function attempt(work: () => Promise<unknown>): Promise<void> {
60  try {
61    await work()
62  } catch {
63    // The bar keeps what it last drew.
64  }
65}
66
67// Takes the reading another session stored when it is newer than this one's.
68async function adopt($: EngineInterface): Promise<void> {
69  const shared = sharedOf(await $.store.get(SHARED))
70  const now = await $.clock.now()
71
72  // A reading stamped ahead of the clock would outrank every real one.
73  if (shared === null || shared.limits.length === 0 || shared.at > now + 60_000) {
74    return
75  }
76
77  await update($, snapshot, held =>
78    held.limitsAt !== null && held.limitsAt >= shared.at
79      ? held
80      : { ...held, limits: shared.limits, limitsAt: shared.at },
81  )
82}
83
84function numberOf(value: unknown, fallback: number, min: number, max: number): number {
85  return typeof value === 'number' && Number.isFinite(value) ? Math.min(max, Math.max(min, value)) : fallback
86}
87
88function styleOf(options: PluginOptions): Style {
89  return {
90    theme: options.theme === 'dark' ? 'dark' : 'light',
91    staleAfterMs: numberOf(options.staleMinutes, DEFAULT_STYLE.staleAfterMs / 60_000, 1, 1440) * 60_000,
92    warnBelow: numberOf(options.warnBelow, DEFAULT_STYLE.warnBelow, 0, 100),
93    dangerBelow: numberOf(options.dangerBelow, DEFAULT_STYLE.dangerBelow, 0, 100),
94    showUpdated: options.showUpdated !== false,
95  }
96}
97
98export const register: Register = (on, options) => {
99  const style = styleOf(options)
100  // Off, the up arrow counts uncached input alone and cache writes join the reads.
101  const countsCacheWrites = options.countCacheWrites !== false
102
103  on('session.start', async ($, e, next) => {
104    // The countdowns move with the clock, and another session's reading
105    // arrives through the store: neither raises an event here.
106    $.clock.every(60_000, () => {
107      void attempt(() => adopt($)).then(() => $.ui.invalidate('ui.render'))
108    })
109    await attempt(() =>
110      $.command.register({
111        name: 'usage-bar',
112        description: 'Show or hide the usage bar above the prompt',
113      }),
114    )
115    // How old this reading is cannot be known here, so it carries no time and
116    // any stored one replaces it.
117    await attempt(async () => {
118      const usage = await $.session.usage()
119      await update($, snapshot, held => ({
120        limits: held.limitsAt === null ? limitsOf(usage.rateLimits) : held.limits,
121        limitsAt: held.limitsAt,
122        costUsd: usage.cost?.usd ?? held.costUsd,
123      }))
124    })
125    await attempt(() => adopt($))
126
127    return next(e)
128  })
129
130  on('command.run', { command: 'usage-bar' }, async $ => {
131    const hidden = !(await read($, isHidden))
132    await update($, isHidden, () => hidden)
133
134    return { text: hidden ? 'Usage bar hidden.' : 'Usage bar shown.' }
135  })
136
137  on('session.measure', async ($, e, next) => {
138    const limits = limitsOf(e.rateLimits)
139    await attempt(async () => {
140      const at = await $.clock.now()
141      await update($, snapshot, held => ({
142        limits: limits.length > 0 ? limits : held.limits,
143        limitsAt: limits.length > 0 ? at : held.limitsAt,
144        costUsd: e.cost?.usd ?? held.costUsd,
145      }))
146
147      if (limits.length > 0) {
148        const shared: Shared = { limits, at }
149        await $.store.set(SHARED, shared)
150      }
151    })
152
153    return next(e)
154  })
155
156  on('turn.complete', async ($, e, next) => {
157    const { usage } = e
158
159    if (usage !== undefined) {
160      await attempt(() =>
161        update($, counts, sum => ({
162          input: sum.input + usage.input_tokens + (countsCacheWrites ? usage.cache_creation_input_tokens : 0),
163          output: sum.output + usage.output_tokens,
164          cache:
165            sum.cache + usage.cache_read_input_tokens + (countsCacheWrites ? 0 : usage.cache_creation_input_tokens),
166        })),
167      )
168    }
169
170    return next(e)
171  })
172
173  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
174    // Hidden, the band draws nothing at all; /usage-bar brings it back.
175    if (e.props.hasSurvey || (await read($, isHidden))) {
176      return next(e)
177    }
178
179    const { Box, Button, Text } = $.ui.resolve(e)
180
181    const now = await $.clock.now()
182    const shown = await read($, snapshot)
183    const sum = await read($, counts)
184    const line = textLine(shown, sum, now, style)
185    const hide = (
186      <Button
187        key="hide"
188        plain
189        dimColor
190        role="dismiss"
191        label="×"
192        onPress={async () => {
193          await update($, isHidden, () => true)
194          $.ui.toast('Usage bar hidden. Type /usage-bar to show it again.')
195        }}
196      />
197    )
198
199    if (e.surface === 'desktop') {
200      const { Svg } = $.ui.resolve(e)
201
202      return (
203        <Box alignItems="center" gap={1}>
204          <Svg source={barSvg(shown, sum, now, style)} alt={line} />
205          {hide}
206        </Box>
207      )
208    }
209
210    return (
211      <Box gap={1}>
212        <Text dimColor wrap="truncate-end">
213          {line}
214        </Text>
215        {hide}
216      </Box>
217    )
218  })
219}
220

1import type { Limit, Snapshot, Counts } from '../types'
2
3const HOUR = 3_600_000
4
5// What the person may set through the mod's options; `styleOf` fills these in.
6export type Style = {
7  theme: 'light' | 'dark'
8  // A reading older than this is drawn faded: the session has been idle since.
9  staleAfterMs: number
10  // At or below these percentages left, the bar is drawn orange, then red.
11  warnBelow: number
12  dangerBelow: number
13  // Whether the bar says how long ago the rate-limit reading was taken.
14  showUpdated: boolean
15}
16
17export const DEFAULT_STYLE: Style = {
18  theme: 'light',
19  staleAfterMs: 10 * 60_000,
20  warnBelow: 30,
21  dangerBelow: 10,
22  showUpdated: true,
23}
24const WINDOWS: Record<string, { label: string; ms: number }> = {
25  five_hour: { label: '5h', ms: 5 * HOUR },
26  seven_day: { label: '7d', ms: 7 * 24 * HOUR },
27}
28
29export type Window = {
30  label: string
31  // What is left of the window's allowance, 0 to 100.
32  percent: number | null
33  // How much of the window's time is left, 0 to 1; null without a reset time.
34  timeLeft: number | null
35  left: string
36  isStale: boolean
37}
38
39export function formatCount(n: number): string {
40  if (n < 1000) {
41    return String(n)
42  }
43
44  // 999,950 and up would round to "1000.0k".
45  return n < 999_950 ? `${(n / 1000).toFixed(1)}k` : `${(n / 1_000_000).toFixed(2)}M`
46}
47
48export function formatLeft(ms: number): string {
49  const minutes = Math.max(0, Math.floor(ms / 60_000))
50  const days = Math.floor(minutes / 1440)
51  const hours = Math.floor((minutes % 1440) / 60)
52
53  if (days > 0) {
54    return `${days}d ${hours}h`
55  }
56
57  return hours > 0 ? `${hours}h ${minutes % 60}m` : `${minutes}m`
58}
59
60export function windowOf(kind: string, snapshot: Snapshot, now: number, style: Style): Window {
61  const { limits, limitsAt } = snapshot
62  const { label, ms } = WINDOWS[kind] ?? { label: kind, ms: 0 }
63  const limit = limits.find(one => one.kind === kind)
64
65  if (limit === undefined) {
66    return { label, percent: null, timeLeft: null, left: '--', isStale: false }
67  }
68
69  const resetsIn = limit.resetsAt === null ? NaN : Date.parse(limit.resetsAt) - now
70  const hasReset = Number.isFinite(resetsIn) && ms > 0
71
72  return {
73    label,
74    percent: Math.max(0, 100 - Math.round(limit.percentUsed)),
75    timeLeft: hasReset ? Math.min(1, Math.max(0, resetsIn / ms)) : null,
76    left: hasReset ? formatLeft(resetsIn) : '--',
77    // Past its reset the percentage describes a window that is over.
78    isStale: limitsAt === null || now - limitsAt > style.staleAfterMs || (hasReset && resetsIn <= 0),
79  }
80}
81
82export function formatCost(usd: number | null): string {
83  return usd === null ? '$--' : `$${usd.toFixed(2)}`
84}
85
86// How long ago the rate-limit reading was taken: `now` under a minute, `--`
87// when there is no reading or its time is unknown.
88export function formatAge(snapshot: Snapshot, now: number): string {
89  if (snapshot.limits.length === 0 || snapshot.limitsAt === null) {
90    return '--'
91  }
92
93  const age = now - snapshot.limitsAt
94
95  return age < 60_000 ? 'now' : formatLeft(age)
96}
97
98export function textLine(snapshot: Snapshot, counts: Counts, now: number, style: Style): string {
99  const windows = ['five_hour', 'seven_day'].map(kind => {
100    const w = windowOf(kind, snapshot, now, style)
101
102    return `${w.label} ${w.isStale ? '~' : ''}${w.percent ?? '--'}% left (${w.left})`
103  })
104
105  const age = formatAge(snapshot, now)
106  const updated = style.showUpdated && snapshot.limits.length > 0
107    ? [`updated ${age === 'now' || age === '--' ? age : `${age} ago`}`]
108    : []
109
110  return [
111    ...windows,
112    ...updated,
113    `in ${formatCount(counts.input)}`,
114    `out ${formatCount(counts.output)}`,
115    `cache ${formatCount(counts.cache)}`,
116    formatCost(snapshot.costUsd),
117  ].join(' | ')
118}
119
120// The desktop drawing: one SVG, laid out left to right in CSS pixels.
121const FONT = 14
122const CHAR = FONT * 0.6
123const HEIGHT = 32
124const PILL = 30
125const TOP = (HEIGHT - PILL) / 2
126const MID = HEIGHT / 2
127
128type Palette = {
129  ink: string
130  // The refresh mark and its age, drawn on the band itself and not on a pill.
131  muted: string
132  warn: string
133  danger: string
134  // Per pill, in the order drawn: its background and its icon and bar color.
135  pills: readonly { background: string; accent: string }[]
136}
137
138const PALETTES: Record<Style['theme'], Palette> = {
139  light: {
140    ink: '#1f2328',
141    muted: '#6b7076',
142    warn: '#c77d1a',
143    danger: '#c0392b',
144    pills: [
145      { background: '#cfe3d8', accent: '#3f8f6b' },
146      { background: '#dad5f0', accent: '#6b55c9' },
147      { background: '#f1d6d0', accent: '#c0503c' },
148      { background: '#d6e7d6', accent: '#3f8f5a' },
149      { background: '#d4daf3', accent: '#4a5fd0' },
150      { background: '#efe4c4', accent: '#a8811c' },
151    ],
152  },
153  dark: {
154    ink: '#e6e8eb',
155    muted: '#9aa0a6',
156    warn: '#e0a040',
157    danger: '#f0705f',
158    pills: [
159      { background: '#1f3a2e', accent: '#6fcf9f' },
160      { background: '#2e2850', accent: '#a594f0' },
161      { background: '#4a2620', accent: '#ef8f7c' },
162      { background: '#22382a', accent: '#7fcf95' },
163      { background: '#232c52', accent: '#8a9cf5' },
164      { background: '#40361a', accent: '#e0b84a' },
165    ],
166  },
167}
168
169const ICONS = {
170  gauge: '<path d="M2.5 12a5.5 5.5 0 1 1 11 0"/><path d="M8 12l3-4"/>',
171  clock: '<circle cx="8" cy="8" r="5.5"/><path d="M8 5v3l2 1.5"/>',
172  calendar: '<rect x="2.5" y="3.5" width="11" height="10" rx="2"/><path d="M2.5 7h11M5.5 2v3M10.5 2v3"/>',
173  up: '<path d="M8 10V2.5M5 5.5l3-3 3 3M3 10v3.5h10V10"/>',
174  down: '<path d="M8 2.5V10M5 7l3 3 3-3M3 10v3.5h10V10"/>',
175  layers: '<path d="M8 2l6 3-6 3-6-3zM2 8l6 3 6-3M2 11l6 3 6-3"/>',
176  refresh: '<path d="M13 8a5 5 0 1 1-1.6-3.6M13 2.8v3.2H9.8"/>',
177  coin: '<circle cx="8" cy="8" r="5.5"/><path d="M9.8 6.2c-.3-.6-1-.9-1.8-.9-1 0-1.8.5-1.8 1.3 0 1.9 3.7.8 3.7 2.8 0 .8-.8 1.3-1.9 1.3-.9 0-1.6-.4-1.9-1M8 4.2v7.6"/>',
178}
179
180type Part =
181  | { icon: keyof typeof ICONS; color: string }
182  | { text: string; isBold?: boolean; isFaded?: boolean; color?: string }
183  | { bar: Window; color: string }
184  | { rule: true }
185
186function widthOf(part: Part): number {
187  if ('icon' in part) {
188    return 16
189  }
190
191  if ('text' in part) {
192    return Math.ceil(part.text.length * CHAR)
193  }
194
195  return 'bar' in part ? 84 : 1
196}
197
198function drawPart(part: Part, x: number, style: Style): string {
199  const { ink, warn, danger } = PALETTES[style.theme]
200
201  if ('icon' in part) {
202    return `<g transform="translate(${x} ${MID - 8})" fill="none" stroke="${part.color}" stroke-width="1.4" stroke-linecap="round" stroke-linejoin="round">${ICONS[part.icon]}</g>`
203  }
204
205  if ('text' in part) {
206    return `<text x="${x}" y="${MID + 5}" fill="${part.color ?? ink}" font-weight="${part.isBold ? 700 : 400}"${part.isFaded ? ' opacity="0.4"' : ''}>${part.text}</text>`
207  }
208
209  if ('rule' in part) {
210    return `<rect x="${x}" y="${MID - 9}" width="1" height="18" fill="${ink}" opacity="0.25"/>`
211  }
212
213  const { percent, timeLeft, isStale } = part.bar
214  const fill = Math.round((Math.min(100, percent ?? 0) / 100) * 84)
215  const left = percent ?? 100
216  const color = left <= style.dangerBelow ? danger : left <= style.warnBelow ? warn : part.color
217  const mark =
218    timeLeft === null
219      ? ''
220      : `<rect x="${x + Math.round(timeLeft * 82)}" y="${MID - 8}" width="2.5" height="16" rx="1" fill="${ink}"/>`
221
222  return (
223    `<rect x="${x}" y="${MID - 4}" width="84" height="8" rx="4" fill="${ink}" opacity="0.14"/>` +
224    `<rect x="${x}" y="${MID - 4}" width="${fill}" height="8" rx="4" fill="${color}"${isStale ? ' opacity="0.4"' : ''}/>` +
225    mark
226  )
227}
228
229function drawPill(parts: Part[], background: string, x: number, style: Style): { svg: string; width: number } {
230  let cursor = x + 12
231  const drawn = parts.map(part => {
232    const svg = drawPart(part, cursor, style)
233    cursor += widthOf(part) + 8
234
235    return svg
236  })
237  const width = cursor - 8 + 12 - x
238
239  return {
240    svg: `<rect x="${x}" y="${TOP}" width="${width}" height="${PILL}" rx="${PILL / 2}" fill="${background}"/>${drawn.join('')}`,
241    width,
242  }
243}
244
245function windowParts(w: Window, icon: 'gauge' | 'calendar', color: string): Part[] {
246  return [
247    { icon, color },
248    { text: w.label },
249    { bar: w, color },
250    { text: `${w.percent ?? '--'}%`, isBold: true, isFaded: w.isStale },
251    { rule: true },
252    { icon: 'clock', color },
253    { text: w.left },
254  ]
255}
256
257export function barSvg(snapshot: Snapshot, counts: Counts, now: number, style: Style): string {
258  const { pills: colors, muted } = PALETTES[style.theme]
259  // `pill` picks the palette entry; without one the parts sit on the band itself.
260  const contents: { pill?: number; parts: (accent: string) => Part[]; gapAfter: number }[] = [
261    { pill: 0, parts: accent => windowParts(windowOf('five_hour', snapshot, now, style), 'gauge', accent), gapAfter: 10 },
262    { pill: 1, parts: accent => windowParts(windowOf('seven_day', snapshot, now, style), 'calendar', accent), gapAfter: 30 },
263    { pill: 2, parts: accent => [{ icon: 'up', color: accent }, { text: formatCount(counts.input), isBold: true }], gapAfter: 10 },
264    { pill: 3, parts: accent => [{ icon: 'down', color: accent }, { text: formatCount(counts.output), isBold: true }], gapAfter: 10 },
265    { pill: 4, parts: accent => [{ icon: 'layers', color: accent }, { text: formatCount(counts.cache), isBold: true }], gapAfter: 30 },
266    { pill: 5, parts: accent => [{ icon: 'coin', color: accent }, { text: formatCost(snapshot.costUsd), isBold: true }], gapAfter: 0 },
267  ]
268
269  if (style.showUpdated && snapshot.limits.length > 0) {
270    const seven = contents[1]
271
272    if (seven !== undefined) {
273      seven.gapAfter = 0
274    }
275
276    contents.splice(2, 0, {
277      parts: () => [{ icon: 'refresh', color: muted }, { text: formatAge(snapshot, now), color: muted }],
278      gapAfter: 18,
279    })
280  }
281
282  let x = 0
283  const drawn = contents.map(content => {
284    const { background, accent } = (content.pill === undefined ? undefined : colors[content.pill]) ?? {
285      background: 'none',
286      accent: muted,
287    }
288    const { svg, width } = drawPill(content.parts(accent), background, x, style)
289    x += width + content.gapAfter
290
291    return svg
292  })
293  const width = x
294
295  // xmlns is the namespace identifier a standalone SVG must carry to be
296  // drawn as an image; it is a name, and nothing is requested from it.
297  return (
298    `<svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${HEIGHT}" viewBox="0 0 ${width} ${HEIGHT}" ` +
299    `font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="${FONT}">` +
300    `${drawn.join('')}</svg>`
301  )
302}
303

1export type Limit = { kind: string; percentUsed: number; resetsAt: string | null }
2
3// `limitsAt` is when the limits were read, in `$.clock.now()`'s milliseconds; null when unknown.
4export type Snapshot = { limits: Limit[]; limitsAt: number | null; costUsd: number | null }
5
6export type Counts = { input: number; output: number; cache: number }
7
8declare module 'claude-code' {
9  interface PluginState {
10    'usage-bar': { snapshot: Snapshot; counts: Counts; isHidden: boolean }
11  }
12}
13