quota-meter

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionRateLimit } from 'claude-code'
3
4import type { Limit, Reading } from '../types'
5import { renderBar } from './bar'
6import {
7  formatCountdown,
8  formatTokens,
9  formatUsd,
10  idealPercent,
11  isExpired,
12  paceStatus,
13  windowLabel,
14} from './pace'
15
16const COMMAND = 'quota'
17const TICK_MS = 60_000
18const READING_KEY = 'lastReading'
19const HIDDEN_KEY = 'isHidden'
20const LABEL_COLUMNS = 4
21const MIN_BAR = 10
22const MAX_BAR = 40
23
24const limits = atom({ plugin: 'quota-meter', key: 'limits' } as const, [])
25const isStale = atom({ plugin: 'quota-meter', key: 'isStale' } as const, false)
26const usd = atom({ plugin: 'quota-meter', key: 'usd' } as const, null)
27const tokens = atom({ plugin: 'quota-meter', key: 'tokens' } as const, { total: 0, output: 0 })
28const now = atom({ plugin: 'quota-meter', key: 'now' } as const, 0)
29const isHidden = atom({ plugin: 'quota-meter', key: 'isHidden' } as const, false)
30
31const copyLimits = (rateLimits: readonly SessionRateLimit[]): Limit[] =>
32  rateLimits.map(limit => ({ ...limit }))
33
34const isLimit = (value: unknown): value is Limit => {
35  if (typeof value !== 'object' || value === null) return false
36  const { kind, percentUsed, resetsAt } = value as Record<string, unknown>
37
38  return (
39    typeof kind === 'string' &&
40    typeof percentUsed === 'number' &&
41    Number.isFinite(percentUsed) &&
42    (resetsAt === undefined || typeof resetsAt === 'string')
43  )
44}
45
46// A reading an older version stored may differ: keep the entries that still read as limits.
47const storedLimits = (value: unknown): Limit[] => {
48  if (typeof value !== 'object' || value === null) return []
49  const stored = (value as { limits?: unknown }).limits
50
51  return Array.isArray(stored) ? stored.filter(isLimit) : []
52}
53
54const tick = async ($: EngineInterface): Promise<void> => {
55  const time = await $.clock.now()
56  await update($, now, () => time)
57}
58
59export const register: Register = on => {
60  on('session.start', async ($, e, next) => {
61    await $.command.register({ name: COMMAND, description: 'Show or hide the quota meter band' })
62
63    const usage = await $.session.usage()
64    const stored = storedLimits(await $.store.get(READING_KEY))
65    const hidden = (await $.store.get(HIDDEN_KEY)) === true
66
67    if (usage.rateLimits.length > 0) {
68      const fresh = copyLimits(usage.rateLimits)
69      await update($, limits, () => fresh)
70      await update($, isStale, () => false)
71    } else if (stored.length > 0) {
72      await update($, limits, () => stored)
73      await update($, isStale, () => true)
74    }
75
76    await update($, usd, () => usage.cost?.usd ?? null)
77    await update($, isHidden, () => hidden)
78    await tick($)
79    $.clock.every(TICK_MS, () => void tick($))
80
81    return next(e)
82  })
83
84  on('session.measure', async ($, e, next) => {
85    if (e.rateLimits.length > 0) {
86      const fresh = copyLimits(e.rateLimits)
87      const reading: Reading = { limits: fresh, at: await $.clock.now() }
88      await update($, limits, () => fresh)
89      await update($, isStale, () => false)
90      await $.store.set(READING_KEY, reading)
91    }
92
93    if (e.cost !== undefined) {
94      const total = e.cost.usd
95      await update($, usd, () => total)
96    }
97
98    return next(e)
99  })
100
101  on('turn.complete', async ($, e, next) => {
102    const usage = e.usage
103    if (usage !== undefined) {
104      await update($, tokens, current => ({
105        total:
106          current.total +
107          usage.input_tokens +
108          usage.output_tokens +
109          usage.cache_read_input_tokens +
110          usage.cache_creation_input_tokens,
111        output: current.output + usage.output_tokens,
112      }))
113    }
114
115    return next(e)
116  })
117
118  on('session.end', async ($, e, next) => {
119    // Neither /clear nor an in-process /resume is followed by session.start.
120    if (e.reason === 'clear' || e.reason === 'resume') {
121      await update($, tokens, () => ({ total: 0, output: 0 }))
122      await update($, usd, () => null)
123    }
124
125    return next(e)
126  })
127
128  on('command.run', { command: COMMAND }, async $ => {
129    // This session's state, not the store: another session may have toggled it.
130    const hidden = !(await read($, isHidden))
131    await $.store.set(HIDDEN_KEY, hidden)
132    await update($, isHidden, () => hidden)
133
134    return { text: hidden ? 'Quota meter hidden. Run /quota to show it again.' : 'Quota meter shown.' }
135  })
136
137  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
138    if (e.props.hasSurvey || (await read($, isHidden))) return next(e)
139
140    const list = await read($, limits)
141    const stale = await read($, isStale)
142    const cost = await read($, usd)
143    const used = await read($, tokens)
144    // Reading `now` subscribes the band to the minute tick.
145    const tickedAt = await read($, now)
146    const time = Math.max(tickedAt, await $.clock.now())
147
148    const session = [
149      used.total > 0 ? `${formatTokens(used.total)} tok (out ${formatTokens(used.output)})` : null,
150      cost !== null ? formatUsd(cost) : null,
151    ].filter(part => part !== null)
152
153    if (list.length === 0 && session.length === 0) return next(e)
154
155    const { Box, Text } = $.ui.resolve(e)
156
157    const rows = list.map(limit => {
158      const isOver = isExpired(limit.resetsAt, time)
159      const real = isOver ? 0 : limit.percentUsed
160      const ideal = isOver ? undefined : idealPercent(limit.kind, limit.resetsAt, time)
161      const details = isOver
162        ? ['new window on next msg']
163        : [
164            ideal === undefined ? null : `ideal ${Math.round(ideal)}%`,
165            limit.resetsAt === undefined
166              ? null
167              : `resets ${formatCountdown(Date.parse(limit.resetsAt) - time)}`,
168          ].filter(part => part !== null)
169      const tail = `  ${[`${Math.round(real)}%`.padStart(4), ...details].join('  ')}${stale ? '  (last seen)' : ''}`
170
171      return { limit, real, ideal, tail }
172    })
173
174    // The bar takes what the longest row's text leaves; below MIN_BAR the text truncates instead.
175    const tailColumns = Math.max(0, ...rows.map(row => row.tail.length))
176    const width = Math.min(MAX_BAR, Math.max(MIN_BAR, e.props.bodyColumns - LABEL_COLUMNS - tailColumns))
177
178    return (
179      <Box flexDirection="column">
180        {rows.map(({ limit, real, ideal, tail }) => (
181          <Box key={limit.kind}>
182            <Text dimColor={stale}>{windowLabel(limit.kind).padEnd(LABEL_COLUMNS)}</Text>
183            <Box key={`${limit.kind}-bar`} flexShrink={0}>
184              {renderBar(width, real, ideal, paceStatus(real, ideal)).map((segment, index) => (
185                <Text key={`${limit.kind}-${index}`} color={segment.color} dimColor={segment.dimColor === true || stale}>
186                  {segment.text}
187                </Text>
188              ))}
189            </Box>
190            <Text dimColor={stale} wrap="truncate-end">
191              {tail}
192            </Text>
193          </Box>
194        ))}
195        {session.length > 0 && <Text dimColor>{`session  ${session.join(' · ')}`}</Text>}
196      </Box>
197    )
198  })
199}
200

1import type { PaceStatus } from './pace'
2
3export type BarColor = 'green' | 'yellow' | 'red' | 'cyan'
4
5export type Segment = { text: string; color?: BarColor; dimColor?: boolean }
6
7const EIGHTHS = ['', '▏', '▎', '▍', '▌', '▋', '▊', '▉'] as const
8const FULL = '█'
9const EMPTY = '░'
10const MARKER = '│'
11const MARKER_COLOR: BarColor = 'cyan'
12
13const FILL_COLORS: Readonly<Record<PaceStatus, BarColor>> = {
14  under: 'green',
15  near: 'yellow',
16  over: 'red',
17}
18
19const clampPercent = (percent: number): number => Math.min(100, Math.max(0, percent))
20
21export const barCells = (width: number, real: number): string[] => {
22  const eighths = Math.round((clampPercent(real) / 100) * width * 8)
23  const cells: string[] = Array.from({ length: Math.floor(eighths / 8) }, () => FULL)
24  const partial = EIGHTHS[eighths % 8]
25  if (partial) cells.push(partial)
26  while (cells.length < width) cells.push(EMPTY)
27
28  return cells
29}
30
31export const markerIndex = (width: number, ideal: number): number =>
32  Math.min(width - 1, Math.round((clampPercent(ideal) / 100) * width))
33
34export const renderBar = (
35  width: number,
36  real: number,
37  ideal: number | undefined,
38  status: PaceStatus,
39): Segment[] => {
40  const marker = ideal === undefined ? -1 : markerIndex(width, ideal)
41  const segments: Segment[] = []
42
43  barCells(width, real).forEach((cell, index) => {
44    const next: Segment =
45      index === marker
46        ? { text: MARKER, color: MARKER_COLOR }
47        : cell === EMPTY
48          ? { text: cell, dimColor: true }
49          : { text: cell, color: FILL_COLORS[status] }
50    const last = segments[segments.length - 1]
51    const isSameStyle =
52      last !== undefined &&
53      last.color === next.color &&
54      last.dimColor === next.dimColor &&
55      next.text !== MARKER &&
56      last.text !== MARKER
57
58    if (isSameStyle) last.text += next.text
59    else segments.push(next)
60  })
61
62  return segments
63}
64

1export type PaceStatus = 'under' | 'near' | 'over'
2
3const HOUR_MS = 3_600_000
4
5const DURATIONS_MS: Readonly<Record<string, number>> = {
6  five_hour: 5 * HOUR_MS,
7  seven_day: 168 * HOUR_MS,
8}
9
10const LABELS: Readonly<Record<string, string>> = {
11  five_hour: '5h',
12  seven_day: '7d',
13}
14
15// Own keys only: `kind` comes from outside, and `toString` is not a window.
16export const windowDuration = (kind: string): number | undefined =>
17  Object.hasOwn(DURATIONS_MS, kind) ? DURATIONS_MS[kind] : undefined
18
19export const windowLabel = (kind: string): string =>
20  (Object.hasOwn(LABELS, kind) ? LABELS[kind] : undefined) ?? 'lim'
21
22export const isExpired = (resetsAt: string | undefined, now: number): boolean =>
23  resetsAt !== undefined && Date.parse(resetsAt) <= now
24
25// Linear 24/7 pace: the share of the window that has elapsed.
26export const idealPercent = (
27  kind: string,
28  resetsAt: string | undefined,
29  now: number,
30): number | undefined => {
31  const duration = windowDuration(kind)
32  if (duration === undefined || resetsAt === undefined) return undefined
33
34  const end = Date.parse(resetsAt)
35  if (Number.isNaN(end) || end <= now) return undefined
36
37  const elapsed = now - (end - duration)
38
39  return Math.min(100, Math.max(0, (elapsed / duration) * 100))
40}
41
42export const formatCountdown = (ms: number): string => {
43  const minutes = Math.floor(ms / 60_000)
44  if (minutes < 1) return '<1m'
45
46  const days = Math.floor(minutes / 1_440)
47  const hours = Math.floor((minutes % 1_440) / 60)
48  const rest = minutes % 60
49
50  if (days > 0) return `${days}d ${hours}h`
51  if (hours > 0) return `${hours}h${String(rest).padStart(2, '0')}m`
52
53  return `${rest}m`
54}
55
56export const formatTokens = (count: number): string => {
57  if (count < 1_000) return String(count)
58  if (count < 10_000) return `${(Math.floor(count / 100) / 10).toFixed(1)}k`
59  if (count < 1_000_000) return `${Math.floor(count / 1_000)}k`
60
61  return `${(Math.floor(count / 10_000) / 100).toFixed(2)}M`
62}
63
64export const formatUsd = (usd: number): string => `$${usd.toFixed(2)}`
65
66// Decides the bar's color: how far ahead of the ideal pace is still fine.
67// `real` and `ideal` are 0–100; `ideal` is undefined when the window has no
68// known duration or reset time (e.g. `spend_limit`), so only `real` can decide.
69// Invariants the tests pin: 0% used is 'under', 100% used is 'over', and the
70// status never improves as `real` grows for a fixed `ideal`.
71export const paceStatus = (real: number, ideal: number | undefined): PaceStatus => {
72  if (real >= 90) return 'over'
73  if (ideal === undefined) return real >= 75 ? 'near' : 'under'
74  if (real <= ideal) return 'under'
75
76  return real - ideal <= 10 ? 'near' : 'over'
77}
78

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3export type Reading = { limits: Limit[]; at: number }
4
5export type Tokens = { total: number; output: number }
6
7declare module 'claude-code' {
8  interface PluginState {
9    'quota-meter': {
10      limits: Limit[]
11      isStale: boolean
12      usd: number | null
13      tokens: Tokens
14      now: number
15      isHidden: boolean
16    }
17  }
18}
19