cost-meter

1import type { EngineInterface, Register } from 'claude-code'
2
3import {
4  contextToast,
5  crossMarks,
6  formatUsd,
7  parseBudget,
8  report,
9  statusLine,
10  type Reading,
11  type Turn,
12} from './format'
13
14const KEEP = 'Keep going'
15const STOP = 'Stop here'
16
17// The session's figures. A reload starts them over; the next measurement fills them in.
18let reading: Reading = { limits: [] }
19let announced = new Set<number>()
20const turns: Turn[] = []
21let current: { turnId: string; prompt: string; startUsd: number | undefined } | null = null
22// Finished turns waiting for their `TurnDuration` line, and the lines already matched.
23const unclaimed: { durationMs: number; usd: number }[] = []
24const claimed = new Map<string, number>()
25// The budget: `bar` is where the next check fires; `stoppedTurn` refuses the rest of a stopped turn.
26const budget = { usd: 0, bar: 0, isOver: false }
27let stoppedTurn: string | null = null
28let pendingAsk: Promise<boolean> | null = null
29let warnContext = true
30
31function take(usage: { context: { percent?: number }; rateLimits: readonly { kind: string; percentUsed: number; resetsAt?: string }[]; cost?: { usd: number } }): void {
32  reading = {
33    usd: usage.cost?.usd,
34    percent: usage.context.percent,
35    limits: usage.rateLimits.map(l => ({ kind: l.kind, percentUsed: l.percentUsed, resetsAt: l.resetsAt })),
36  }
37}
38
39function react($: EngineInterface): void {
40  const line = statusLine(reading)
41  $.ui.status(line === '' ? undefined : line)
42
43  if (warnContext) {
44    const marks = crossMarks(reading.percent, announced)
45    announced = marks.announced
46    for (const mark of marks.crossed) $.ui.toast(contextToast(mark), { timeoutMs: 8000 })
47  }
48
49  const usd = reading.usd
50  if (budget.usd > 0 && usd !== undefined && usd >= budget.bar && !budget.isOver) {
51    budget.isOver = true
52    $.ui.toast(`Session cost ${formatUsd(usd)} passed your ${formatUsd(budget.bar)} budget`, { timeoutMs: 8000 })
53  }
54}
55
56function raiseBar(): void {
57  if (budget.usd <= 0) {
58    budget.isOver = false
59    return
60  }
61  const usd = reading.usd ?? 0
62  while (budget.bar <= usd) budget.bar += budget.usd
63  budget.isOver = false
64}
65
66async function askToContinue($: EngineInterface): Promise<boolean> {
67  let answer = STOP
68  try {
69    answer = await $.ui.ask(
70      `Session cost ${formatUsd(reading.usd ?? 0)} passed your ${formatUsd(budget.bar)} budget. Keep going?`,
71      { header: 'cost-meter', options: [KEEP, STOP] },
72    )
73  } catch {
74    // Nobody to ask (a -p run) or the dialog was dismissed: stop.
75  }
76  if (answer !== KEEP) return false
77  raiseBar()
78  return true
79}
80
81async function measure($: EngineInterface): Promise<void> {
82  try {
83    take(await $.session.usage())
84  } catch {
85    // A host without figures: keep what the last push said.
86  }
87}
88
89export const register: Register = (on, options) => {
90  budget.usd = Math.max(0, Number(options.budgetUsd ?? 0) || 0)
91  budget.bar = budget.usd
92  const showTurnCost = options.showTurnCost !== false
93  warnContext = options.contextWarnings !== false
94
95  on('session.start', async ($, e, next) => {
96    await $.command.register({
97      name: 'spend',
98      description: 'cost-meter: what this session has cost, turn by turn; /spend budget <usd> sets a budget',
99      argumentHint: '[budget <usd>|off]',
100    })
101    await measure($)
102    react($)
103    return next(e)
104  })
105
106  on('session.measure', ($, e, next) => {
107    take(e)
108    react($)
109    return next(e)
110  })
111
112  on('turn.start', async ($, e, next) => {
113    await measure($)
114    current = { turnId: e.turnId, prompt: e.text, startUsd: reading.usd }
115    return next(e)
116  })
117
118  on('turn.complete', async ($, e, next) => {
119    if (e.agentId !== undefined) return next(e)
120    await measure($)
121    react($)
122    const started = current
123    current = null
124    if (started !== null && started.turnId === e.turnId && started.startUsd !== undefined && reading.usd !== undefined) {
125      const usd = Math.max(0, reading.usd - started.startUsd)
126      turns.push({ prompt: started.prompt, usd, durationMs: e.durationMs })
127      if (turns.length > 200) turns.shift()
128      unclaimed.push({ durationMs: e.durationMs, usd })
129      if (unclaimed.length > 20) unclaimed.shift()
130      if (showTurnCost) $.ui.invalidate('ui.render')
131    }
132    return next(e)
133  })
134
135  // The line that closes a turn (`Baked for 1m 3s`) gets the turn's cost after it.
136  on('ui.render', { component: 'TurnDuration' }, async ($, e, next) => {
137    if (!showTurnCost) return next(e)
138    let usd = claimed.get(e.requestId)
139    if (usd === undefined) {
140      let best = -1
141      for (let i = 0; i < unclaimed.length; i++) {
142        const gap = Math.abs(unclaimed[i]!.durationMs - e.props.durationMs)
143        if (gap <= 1000 && (best < 0 || gap < Math.abs(unclaimed[best]!.durationMs - e.props.durationMs))) best = i
144      }
145      if (best < 0) return next(e)
146      usd = unclaimed.splice(best, 1)[0]!.usd
147      claimed.set(e.requestId, usd)
148    }
149    const { Box, Text } = $.ui.resolve(e)
150    // Claude Code's own line takes the full width, so the cost goes on a line of its own beneath it.
151    return (
152      <Box flexDirection="column">
153        {await next(e)}
154        <Text dimColor>{`  ${formatUsd(usd)} this turn`}</Text>
155      </Box>
156    )
157  })
158
159  on('prompt.submit', async ($, e, next) => {
160    if (!budget.isOver) return next(e)
161    // Asking here could hold the prompt box; the person's second Enter is the yes.
162    raiseBar()
163    stoppedTurn = null
164    if (e.origin.kind === 'composer') void $.prompt.fill({ text: e.text })
165    return {
166      drop: `cost-meter: this session has cost ${formatUsd(reading.usd ?? 0)}, past your budget. Press Enter to send it anyway (next check at ${formatUsd(budget.bar)}), or /spend budget <usd> to change the budget.`,
167    }
168  })
169
170  on('tool.call', async ($, e, next) => {
171    if (budget.usd <= 0 || e.tool === 'AskUserQuestion') return next(e)
172    // session.measure only fires between turns; a long turn is checked call by call (a free read).
173    if (!budget.isOver) {
174      await measure($)
175      react($)
176    }
177    if (!budget.isOver) return next(e)
178    const turnId = current?.turnId ?? 'idle'
179    if (stoppedTurn === turnId) {
180      return { deny: 'cost-meter: the user stopped this turn at their budget. End your turn now with a short summary of where things stand.' }
181    }
182    if (pendingAsk === null) pendingAsk = askToContinue($).finally(() => (pendingAsk = null))
183    if (await pendingAsk) return next(e)
184    stoppedTurn = turnId
185    return {
186      deny: `cost-meter: the session passed the user's ${formatUsd(budget.usd)} budget and they chose to stop. Do not call more tools; end your turn with a short summary of where things stand.`,
187    }
188  })
189
190  on('command.run', { command: 'spend' }, async ($, e) => {
191    const args = e.args.trim()
192    if (args.startsWith('budget')) {
193      const rest = args.slice('budget'.length)
194      if (rest.trim() === '') {
195        return { text: budget.usd > 0 ? `Budget: ${formatUsd(budget.usd)} per step, next check at ${formatUsd(budget.bar)}.` : 'No budget set. /spend budget 5 sets one.' }
196      }
197      const usd = parseBudget(rest)
198      if (usd === undefined) return { text: `cost-meter: "${rest.trim()}" is not an amount. Try /spend budget 5 or /spend budget off.` }
199      budget.usd = usd
200      budget.bar = usd
201      budget.isOver = false
202      if (usd > 0) raiseBar()
203      stoppedTurn = null
204      return { text: usd > 0 ? `Budget set: ${formatUsd(usd)} per step, next check at ${formatUsd(budget.bar)}.` : 'Budget off.' }
205    }
206    await measure($)
207    return { text: report(reading, turns, budget) }
208  })
209}
210

1// Pure formatting and bookkeeping: no `$`, so the tests can drive it directly.
2
3export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
4
5export type Reading = {
6  usd?: number
7  percent?: number
8  limits: readonly Limit[]
9}
10
11export type Turn = {
12  prompt: string
13  usd: number
14  durationMs: number
15}
16
17const LIMIT_LABELS: Record<string, string> = {
18  five_hour: '5h',
19  seven_day: '7d',
20  spend_limit: 'spend',
21}
22
23const isNumber = (n: unknown): n is number => typeof n === 'number' && Number.isFinite(n)
24
25/** `$0.42`, `$12.3`, `$123`; never NaN. */
26export function formatUsd(usd: number): string {
27  if (!isNumber(usd) || usd <= 0) return '$0.00'
28  if (usd < 0.01) return '<$0.01'
29  if (usd >= 100) return `$${Math.round(usd)}`
30  if (usd >= 10) return `$${usd.toFixed(1)}`
31  return `$${usd.toFixed(2)}`
32}
33
34/** `3s`, `1m 4s`, `2h 5m`, as the engine's own turn line spells it. */
35export function formatDuration(ms: number): string {
36  if (!isNumber(ms) || ms < 0) return '0s'
37  const seconds = Math.round(ms / 1000)
38  if (seconds < 60) return `${seconds}s`
39  const minutes = Math.floor(seconds / 60)
40  if (minutes < 60) return `${minutes}m ${seconds % 60}s`
41  return `${Math.floor(minutes / 60)}h ${minutes % 60}m`
42}
43
44export function limitLabel(kind: string): string {
45  return LIMIT_LABELS[kind] ?? kind.replace(/_/g, ' ')
46}
47
48function round(n: number): string {
49  return Number.isInteger(n) ? String(n) : n.toFixed(0)
50}
51
52/** The status line: only the parts that have a reading, `''` when none does. */
53export function statusLine(reading: Reading): string {
54  const parts: string[] = []
55  if (isNumber(reading.usd)) parts.push(formatUsd(reading.usd))
56  if (isNumber(reading.percent)) parts.push(`ctx ${round(reading.percent)}%`)
57  for (const limit of reading.limits) {
58    if (isNumber(limit.percentUsed)) parts.push(`${limitLabel(limit.kind)} ${round(limit.percentUsed)}%`)
59  }
60  return parts.length === 0 ? '' : `💸 ${parts.join(' · ')}`
61}
62
63export const CONTEXT_MARKS = [70, 85, 95] as const
64
65/**
66 * Which marks `percent` newly crossed, given the ones already announced; marks
67 * the context has fallen well below again (a /compact) are re-armed.
68 */
69export function crossMarks(percent: number | undefined, announced: ReadonlySet<number>): { crossed: number[]; announced: Set<number> } {
70  const next = new Set(announced)
71  if (!isNumber(percent)) return { crossed: [], announced: next }
72  for (const mark of CONTEXT_MARKS) {
73    if (next.has(mark) && percent < mark - 10) next.delete(mark)
74  }
75  const crossed = CONTEXT_MARKS.filter(mark => percent >= mark && !next.has(mark))
76  for (const mark of crossed) next.add(mark)
77  return { crossed: crossed.length > 0 ? [crossed[crossed.length - 1]!] : [], announced: next }
78}
79
80export function contextToast(mark: number): string {
81  if (mark >= 95) return `context ${mark}% full: /compact now or the next turn may be cut off`
82  if (mark >= 85) return `context ${mark}% full: /compact soon`
83  return `context ${mark}% full`
84}
85
86export function preview(text: string, width = 48): string {
87  const one = text.replace(/\s+/g, ' ').trim()
88  if (one === '') return '(no prompt)'
89  return one.length > width ? `${one.slice(0, width - 1)}…` : one
90}
91
92/** The `/spend` report. */
93export function report(reading: Reading, turns: readonly Turn[], budget: { usd: number; bar: number }): string {
94  const lines: string[] = []
95  lines.push(isNumber(reading.usd) ? `Session cost: ${formatUsd(reading.usd)}` : 'Session cost: not reported by this host')
96  if (isNumber(reading.percent)) lines.push(`Context: ${round(reading.percent)}% full`)
97  for (const limit of reading.limits) {
98    const reset = limit.resetsAt === undefined ? '' : `, resets ${limit.resetsAt.replace('T', ' ').replace(/:\d\d(\.\d+)?Z$/, ' UTC')}`
99    lines.push(`Rate limit ${limitLabel(limit.kind)}: ${round(limit.percentUsed)}% used${reset}`)
100  }
101  if (budget.usd > 0) lines.push(`Budget: ${formatUsd(budget.usd)} (next check at ${formatUsd(budget.bar)})`)
102
103  if (turns.length > 0) {
104    const total = turns.reduce((sum, t) => sum + t.usd, 0)
105    const top = turns.reduce((a, b) => (b.usd > a.usd ? b : a))
106    lines.push('')
107    lines.push(`Turns measured: ${turns.length}, ${formatUsd(total)} in all, ${formatUsd(total / turns.length)} on average`)
108    lines.push(`Most expensive: ${formatUsd(top.usd)} for "${preview(top.prompt)}"`)
109    lines.push('')
110    lines.push('Last turns:')
111    for (const turn of turns.slice(-10)) {
112      lines.push(`  ${formatUsd(turn.usd).padStart(7)}  ${formatDuration(turn.durationMs).padStart(7)}  ${preview(turn.prompt)}`)
113    }
114  } else {
115    lines.push('', 'No turns measured yet.')
116  }
117  return lines.join('\n')
118}
119
120/** Parses `/spend budget 5`, `$5`, `off`; undefined when it is not a budget. */
121export function parseBudget(arg: string): number | undefined {
122  const text = arg.trim().toLowerCase()
123  if (text === 'off' || text === '0' || text === 'none') return 0
124  const match = /^\$?(\d+(?:\.\d+)?)$/.exec(text)
125  if (match === null) return undefined
126  const usd = Number(match[1])
127  return Number.isFinite(usd) ? usd : undefined
128}
129