usage-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, TurnUsage } from 'claude-code'
3
4import type { Figures } from '../types'
5import { cacheMiss, cacheTtl, money, row, SEPARATOR, thresholdsFrom, warnings } from './format'
6import type { Thresholds } from './format'
7
8const figures = atom({ plugin: 'usage-bar', key: 'figures' } as const, null)
9const now = atom({ plugin: 'usage-bar', key: 'now' } as const, 0)
10const warned = atom({ plugin: 'usage-bar', key: 'warned' } as const, [])
11const turnStartCost = atom({ plugin: 'usage-bar', key: 'turnStartCost' } as const, null)
12const cache = atom({ plugin: 'usage-bar', key: 'cache' } as const, null)
13
14const TICK_MS = 60_000
15
16async function measure($: EngineInterface): Promise<Figures> {
17  const usage = await $.session.usage()
18  const measured: Figures = {
19    startedAt: usage.startedAt,
20    limits: [...usage.rateLimits],
21    costUsd: usage.cost?.usd ?? null,
22    contextTokens: usage.context.tokens ?? null,
23  }
24
25  await update($, figures, () => measured)
26
27  return measured
28}
29
30async function tick($: EngineInterface): Promise<void> {
31  const time = await $.clock.now()
32  await update($, now, () => time)
33}
34
35async function warn($: EngineInterface, measured: Figures, thresholds: Thresholds): Promise<void> {
36  const given = new Set(await read($, warned))
37  const fresh = warnings({ ...measured, thresholds, now: await $.clock.now() }).filter(w => !given.has(w.key))
38
39  if (fresh.length === 0) {
40    return
41  }
42
43  await update($, warned, keys => [...keys, ...fresh.map(w => w.key)].slice(-200))
44
45  for (const warning of fresh) {
46    $.ui.toast(`⚠ ${warning.text}`, { timeoutMs: 10_000 })
47  }
48}
49
50async function currentCost($: EngineInterface): Promise<number | null> {
51  return (await $.session.usage()).cost?.usd ?? null
52}
53
54/** Records a main-thread request as the cache's last refresh, warning when it re-wrote the cache. */
55async function noteRequest($: EngineInterface, requestAt: number, usage: TurnUsage, thresholds: Thresholds) {
56  const { startedAt, rateLimits } = await $.session.usage()
57  const stored = await read($, cache)
58  const miss = cacheMiss({
59    previous: stored !== null && stored.sessionStartedAt === startedAt ? stored : null,
60    requestAt,
61    usage,
62    ttl: cacheTtl(rateLimits),
63    minTokens: thresholds.cacheMissWarnTokens,
64  })
65
66  await update($, cache, () => ({ sessionStartedAt: startedAt, requestAt, model: usage.model }))
67
68  if (miss !== null) {
69    $.ui.toast(`⚠ ${miss}`, { timeoutMs: 10_000 })
70  }
71}
72
73export const register: Register = (on, options) => {
74  const thresholds = thresholdsFrom(options)
75
76  on('session.start', async ($, e, next) => {
77    // The line the first version pinned under the prompt, if a reload left it.
78    $.ui.status(undefined)
79    // Countdowns move with the clock, not with the session's measurements.
80    $.clock.every(TICK_MS, () => tick($))
81    await measure($)
82    await tick($)
83
84    return next(e)
85  })
86
87  on('session.measure', async ($, e, next) => {
88    const measured = await measure($)
89    await tick($)
90    await warn($, measured, thresholds)
91
92    return next(e)
93  })
94
95  // Subagents keep caches of their own; the board follows the main thread's.
96  on('turn.step', async function* ($, e, next) {
97    if (e.agentId !== undefined) {
98      return yield* next(e)
99    }
100
101    const requestAt = await $.clock.now()
102    const result = yield* next(e)
103
104    if (result.usage !== null) {
105      try {
106        await noteRequest($, requestAt, result.usage, thresholds)
107      } catch {
108        // The board is a bystander: a failed note must not touch the turn.
109      }
110    }
111
112    return result
113  })
114
115  on('prompt.submit', async ($, e, next) => {
116    const cost = await currentCost($)
117    await update($, turnStartCost, () => cost)
118
119    return next(e)
120  })
121
122  on('turn.complete', async ($, e, next) => {
123    const result = await next(e)
124    const startedAt = await read($, turnStartCost)
125    const cost = await currentCost($)
126
127    if (startedAt !== null && cost !== null && cost - startedAt >= thresholds.turnCostWarnUsd) {
128      $.ui.toast(`⚠ That turn cost ${money(cost - startedAt)} (session ${money(cost)}).`, { timeoutMs: 10_000 })
129    }
130
131    await update($, turnStartCost, () => null)
132
133    return result
134  })
135
136  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
137    if (e.props.hasSurvey) {
138      return next(e)
139    }
140
141    const { Box, Text } = $.ui.resolve(e)
142    const measured = await read($, figures)
143    const mark = await read($, cache)
144    const segments = row({
145      limits: measured?.limits ?? [],
146      costUsd: measured?.costUsd ?? null,
147      contextTokens: measured?.contextTokens ?? null,
148      // A mark left by the conversation before a /clear says nothing about this one.
149      mark: mark !== null && mark.sessionStartedAt === measured?.startedAt ? mark : null,
150      now: await read($, now),
151      thresholds,
152    })
153
154    return (
155      <Box paddingLeft={1}>
156        <Text>
157          {segments.flatMap((segment, index) => [
158            ...(index === 0 ? [] : [<Text dimColor>{SEPARATOR}</Text>]),
159            ...segment.parts.map(part => (
160              <Text color={part.color} dimColor={part.dim}>
161                {part.text}
162              </Text>
163            )),
164          ])}
165        </Text>
166      </Box>
167    )
168  })
169}
170

1import type { SessionRateLimit, TurnUsage } from 'claude-code'
2
3import type { CacheMark } from '../types'
4
5export type Thresholds = {
6  costWarnUsd: number
7  turnCostWarnUsd: number
8  limitWarnPct: number
9  cacheMissWarnTokens: number
10}
11
12export const DEFAULTS: Thresholds = {
13  costWarnUsd: 5,
14  turnCostWarnUsd: 2,
15  limitWarnPct: 80,
16  cacheMissWarnTokens: 30_000,
17}
18
19const FIVE_HOUR = 'five_hour'
20const SEVEN_DAY = 'seven_day'
21
22const MINUTE = 60_000
23/** Claude Code's prompt-cache lifetime on a subscription. */
24const CACHE_TTL_MS = 60 * MINUTE
25/** Past a plan limit (usage overage) Claude Code drops to the API's default lifetime. */
26const OVERAGE_CACHE_TTL_MS = 5 * MINUTE
27
28export const SEPARATOR = ' | '
29
30export type Level = 'success' | 'warning' | 'error'
31
32/** A run of one style inside a segment. */
33export type Part = { text: string; color?: Level; dim?: boolean }
34
35/** One `|`-separated piece of the row. */
36export type Segment = { key: string; parts: Part[] }
37
38export const thresholdsFrom = (options: Readonly<Record<string, unknown>>): Thresholds => {
39  const pick = (key: keyof Thresholds): number => {
40    const value = options[key]
41
42    return typeof value === 'number' && Number.isFinite(value) && value > 0 ? value : DEFAULTS[key]
43  }
44
45  return {
46    costWarnUsd: pick('costWarnUsd'),
47    turnCostWarnUsd: pick('turnCostWarnUsd'),
48    limitWarnPct: pick('limitWarnPct'),
49    cacheMissWarnTokens: pick('cacheMissWarnTokens'),
50  }
51}
52
53const find = (limits: readonly SessionRateLimit[], kind: string): SessionRateLimit | undefined =>
54  limits.find(limit => limit.kind === kind)
55
56/** `42m`, `2h 14m`, `1d 8h`; `now` once the time has passed. */
57export const duration = (ms: number): string => {
58  const minutes = Math.ceil(ms / MINUTE)
59
60  if (minutes <= 0) {
61    return 'now'
62  }
63
64  const days = Math.floor(minutes / 1440)
65  const hours = Math.floor((minutes % 1440) / 60)
66  const mins = minutes % 60
67
68  if (days > 0) {
69    return `${days}d ${hours}h`
70  }
71
72  if (hours > 0) {
73    return `${hours}h ${mins}m`
74  }
75
76  return `${mins}m`
77}
78
79export const countdown = (resetsAt: string | undefined, now: number): string | null => {
80  if (resetsAt === undefined) {
81    return null
82  }
83
84  const at = Date.parse(resetsAt)
85
86  return Number.isNaN(at) ? null : duration(at - now)
87}
88
89/** Green below half, yellow from half, red from 80%: the person's own status line scheme. */
90export const levelOf = (percentUsed: number): Level =>
91  percentUsed >= 80 ? 'error' : percentUsed >= 50 ? 'warning' : 'success'
92
93const percent = (value: number): string => `${Math.round(value)}%`
94
95export const money = (usd: number): string => `$${usd.toFixed(2)}`
96
97/** `950`, `48k`, `1.2M`. */
98export const tokens = (count: number): string =>
99  count >= 1_000_000
100    ? `${(count / 1_000_000).toFixed(1)}M`
101    : count >= 1000
102      ? `${Math.round(count / 1000)}k`
103      : `${count}`
104
105const limitSegment = (key: string, title: string, limit: SessionRateLimit | undefined, now: number): Segment => {
106  if (limit === undefined) {
107    return { key, parts: [{ text: `${title}: waiting for first reply`, dim: true }] }
108  }
109
110  const level = levelOf(limit.percentUsed)
111  const resets = countdown(limit.resetsAt, now)
112  const tail = resets === null ? '' : resets === 'now' ? ', resets now' : `, resets in ${resets}`
113
114  return {
115    key,
116    parts: [{ text: `${title} ` }, { text: `${percent(limit.percentUsed)} used`, color: level }, { text: tail, dim: true }],
117  }
118}
119
120export const sessionText = (costUsd: number, thresholds: Thresholds): { text: string; isExpensive: boolean } => {
121  const isExpensive = costUsd >= thresholds.costWarnUsd
122
123  return { text: `${isExpensive ? '⚠ ' : ''}This session ${money(costUsd)}`, isExpensive }
124}
125
126/**
127 * How long a prompt-cache entry lives: an hour on a subscription, five minutes
128 * once a plan limit is spent (overage); unknown without a plan reading.
129 */
130export const cacheTtl = (limits: readonly SessionRateLimit[]): number | null => {
131  const plan = [find(limits, FIVE_HOUR), find(limits, SEVEN_DAY)].filter(limit => limit !== undefined)
132
133  if (plan.length === 0) {
134    return null
135  }
136
137  return plan.some(limit => limit.percentUsed >= 100) ? OVERAGE_CACHE_TTL_MS : CACHE_TTL_MS
138}
139
140export type CacheInput = {
141  mark: CacheMark | null
142  limits: readonly SessionRateLimit[]
143  /** Tokens the next request re-sends: the live context window's fill. */
144  contextTokens: number | null
145  now: number
146}
147
148/** The entry lives from the start of the last request that read or wrote it. */
149const cacheSegment = ({ mark, limits, contextTokens, now }: CacheInput): Segment | null => {
150  if (mark === null) {
151    return { key: 'cache', parts: [{ text: 'Cache: waiting for first reply', dim: true }] }
152  }
153
154  const ttl = cacheTtl(limits)
155
156  if (ttl === null) {
157    return null
158  }
159
160  const left = duration(mark.requestAt + ttl - now)
161
162  if (left !== 'now') {
163    return { key: 'cache', parts: [{ text: 'Cache ' }, { text: `expires in ${left}`, dim: true }] }
164  }
165
166  const rewrite = contextTokens === null ? '' : `: next reply re-caches ~${tokens(contextTokens)} tokens`
167
168  return { key: 'cache', parts: [{ text: 'Cache ' }, { text: `expired${rewrite}`, color: 'warning' }] }
169}
170
171export type RowInput = CacheInput & {
172  costUsd: number | null
173  thresholds: Thresholds
174}
175
176export const row = (input: RowInput): Segment[] => {
177  const { limits, now, costUsd, thresholds } = input
178  const segments = [
179    limitSegment('five-hour', '5h limit', find(limits, FIVE_HOUR), now),
180    limitSegment('seven-day', 'Weekly limit', find(limits, SEVEN_DAY), now),
181  ]
182
183  if (costUsd !== null) {
184    const session = sessionText(costUsd, thresholds)
185    segments.push({
186      key: 'session',
187      parts: [session.isExpensive ? { text: session.text, color: 'warning' } : { text: session.text }],
188    })
189  }
190
191  const cache = cacheSegment(input)
192
193  if (cache !== null) {
194    segments.push(cache)
195  }
196
197  return segments
198}
199
200export const plainText = (segments: readonly Segment[]): string =>
201  segments.map(segment => segment.parts.map(part => part.text).join('')).join(SEPARATOR)
202
203export type MissInput = {
204  /** The main thread's previous answered request this session; null for its first. */
205  previous: CacheMark | null
206  requestAt: number
207  usage: TurnUsage
208  ttl: number | null
209  minTokens: number
210}
211
212/**
213 * The warning for a request that re-wrote the cache instead of reading it, or
214 * null. The first request of a session writes by design and never warns. The
215 * cause is read from what the mod can see: a gap past the lifetime, another
216 * model; anything else changed the start of the prompt.
217 */
218export const cacheMiss = ({ previous, requestAt, usage, ttl, minTokens }: MissInput): string | null => {
219  const written = usage.cache_creation_input_tokens
220
221  if (previous === null || written < minTokens || usage.cache_read_input_tokens >= written) {
222    return null
223  }
224
225  const gap = requestAt - previous.requestAt
226  const cause =
227    ttl !== null && gap >= ttl
228      ? ` after ${duration(gap)} idle`
229      : previous.model !== usage.model
230        ? ' after the model changed'
231        : '; the start of the prompt changed (/compact, CLAUDE.md, MCP tools or effort)'
232
233  return `Cache miss: re-cached ${tokens(written)} tokens${cause}.`
234}
235
236export type Warning = { key: string; text: string }
237
238export type WarningInput = {
239  limits: readonly SessionRateLimit[]
240  costUsd: number | null
241  thresholds: Thresholds
242  now: number
243}
244
245const LABELS: Record<string, string> = { [FIVE_HOUR]: 'Session (5h)', [SEVEN_DAY]: 'Weekly' }
246
247/**
248 * Every warning the figures call for now, each under a key that says when it
249 * was already given: a limit's warnings re-arm with its next window, the cost
250 * ones last the session.
251 */
252export const warnings = ({ limits, costUsd, thresholds, now }: WarningInput): Warning[] => {
253  const found: Warning[] = []
254
255  if (costUsd !== null && costUsd >= thresholds.costWarnUsd) {
256    // $5, $10, $20, ...: the highest step passed, so a jump warns once.
257    let step = thresholds.costWarnUsd
258
259    while (step * 2 <= costUsd) {
260      step *= 2
261    }
262
263    found.push({
264      key: `cost:${step}`,
265      text: `This session has cost ${money(costUsd)} so far (past ${money(step)}). Consider /compact or a fresh session.`,
266    })
267  }
268
269  const levels = [...new Set([95, thresholds.limitWarnPct])].sort((a, b) => b - a)
270
271  for (const limit of limits) {
272    const label = LABELS[limit.kind]
273
274    if (label === undefined) {
275      continue
276    }
277
278    for (const level of levels) {
279      if (limit.percentUsed >= level) {
280        const resets = countdown(limit.resetsAt, now)
281
282        found.push({
283          key: `limit:${limit.kind}:${level}:${limit.resetsAt ?? ''}`,
284          text: `${label} limit at ${percent(limit.percentUsed)} used${resets ? `, resets in ${resets}` : ''}.`,
285        })
286
287        break
288      }
289    }
290  }
291
292  return found
293}
294

1/** One rate-limit window as `$.session.usage()` reports it. */
2export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
3
4/** What the board draws from: the session's last measurement. */
5export type Figures = {
6  /** `$.session.usage().startedAt`: which session the figures belong to. */
7  startedAt: number
8  limits: Limit[]
9  costUsd: number | null
10  /** Tokens the next request re-sends: the live context window's fill. */
11  contextTokens: number | null
12}
13
14/** The main thread's last answered model request, which last refreshed the prompt cache. */
15export type CacheMark = {
16  /** `$.session.usage().startedAt` of the session the request belongs to. */
17  sessionStartedAt: number
18  /** When the request began, in `$.clock.now()`'s milliseconds. */
19  requestAt: number
20  /** The model that answered it; the cache is per model. */
21  model: string
22}
23
24declare module 'claude-code' {
25  interface PluginState {
26    'usage-bar': {
27      figures: Figures | null
28      now: number
29      warned: string[]
30      turnStartCost: number | null
31      cache: CacheMark | null
32    }
33  }
34}
35