cost-info

1// Cost Info: what this session has cost so far, and the tokens it took, live.
2//
3// The figure is the one /cost shows ($.session.usage().cost.usd): every priced
4// API response this session, subagents included. On a Pro or Max plan it is
5// what the same usage would have cost on the API. Tokens are every model
6// request's four counts summed (input, cache writes, cache reads, output),
7// subagents included, counted from when the mod loaded into the session.
8
9import { atom, read, update } from 'claude-code'
10import type { BoxProps, ElementConstructor, EngineInterface, ModelUsage, Register, TextProps } from 'claude-code'
11
12import type { Limit, Totals } from '../types'
13
14const EMPTY: Totals = {
15  session: null,
16  total: 0,
17  turnBase: 0,
18  last: null,
19  turns: 0,
20  turnsCost: 0,
21  priciest: 0,
22  warned: false,
23  tokens: 0,
24  turnTokens: 0,
25  isWorking: false,
26  limits: [],
27}
28// Held by the host, so the totals survive a hot reload of this file.
29const meter = atom({ plugin: 'cost-info', key: 'meter' } as const, EMPTY)
30// VS Code draws no prompt footer, so there the meter opens in a pane of its own.
31const PANE = 'cost-info'
32const TITLE = 'Cost'
33const COMMAND = 'spend'
34const NOTHING = 'Nothing spent yet this session.'
35// Below this many columns the footer leaves the turn out.
36const WIDE = 120
37
38type Kit = { Box: ElementConstructor<BoxProps>; Text: ElementConstructor<TextProps> }
39
40// One run of the meter's text: a figure, drawn green, or the words around the figures.
41type Piece = { text: string; isFigure?: boolean; color?: string }
42
43const money = (usd: number): string => `$${usd < 0.01 ? usd.toFixed(4) : usd.toFixed(2)}`
44
45const count = (n: number): string =>
46  n < 1000 ? `${n}` : n < 999_500 ? `${(n / 1000).toFixed(n < 9_950 ? 1 : 0)}k` : `${(n / 1_000_000).toFixed(2)}M`
47
48const tok = (n: number): string => `${count(n)} tok`
49
50// A plan's rate-limit window, read from whatever shape the host reports it in. Plans that have
51// none (API billing) report an empty list, which is how the meter knows to show cost instead.
52const WINDOWS: Record<string, string> = {
53  five_hour: '5h',
54  seven_day: 'wk',
55  seven_day_opus: 'wk Opus',
56  seven_day_sonnet: 'wk Sonnet',
57}
58
59const limitOf = (raw: unknown): Limit | null => {
60  const r = raw as Record<string, unknown> | null
61  if (r === null || typeof r !== 'object') {
62    return null
63  }
64  const pick = (...keys: string[]): unknown => keys.map(k => r[k]).find(v => v !== undefined && v !== null)
65  const used = pick('percentUsed', 'usedPercentage', 'used_percentage', 'usedPercent', 'percent', 'utilization')
66  if (typeof used !== 'number') {
67    return null
68  }
69  const name = String(pick('kind', 'label', 'name', 'type', 'rateLimitType', 'window', 'id') ?? 'limit')
70  const at = pick('resetsAt', 'resets_at', 'resetAt', 'reset')
71  const ms = typeof at === 'number' ? (at < 1e11 ? at * 1000 : at) : typeof at === 'string' ? Date.parse(at) : NaN
72
73  return {
74    label: WINDOWS[name] ?? name,
75    percent: Math.round(used <= 1 && !Number.isInteger(used) ? used * 100 : used),
76    resetsAt: Number.isNaN(ms) ? null : ms,
77  }
78}
79
80const limitsOf = (raw: unknown): Limit[] => (Array.isArray(raw) ? raw.map(limitOf).filter((l): l is Limit => l !== null) : [])
81
82// Time until a window resets, in its largest units: 6d, 3h 20m, 45m. Empty when unknown or past.
83const left = (resetsAt: number | null, now = Date.now()): string => {
84  if (resetsAt === null || resetsAt <= now) {
85    return ''
86  }
87  const min = Math.ceil((resetsAt - now) / 60_000)
88  if (min >= 1440) {
89    const days = Math.floor(min / 1440)
90    const hours = Math.floor((min % 1440) / 60)
91
92    return days >= 2 || hours === 0 ? `${days}d` : `${days}d ${hours}h`
93  }
94
95  return min >= 60 ? `${Math.floor(min / 60)}h ${min % 60}m` : `${min}m`
96}
97
98const tokensOf = (u: ModelUsage): number =>
99  u.input_tokens + u.output_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
100
101const BAR = 4
102
103// A thin gauge: the used share filled (at least a dot once anything is used), the rest a track.
104// It turns yellow from 70% and red from 90%.
105const gauge = (percent: number): Piece[] => {
106  const filled = percent <= 0 ? 0 : Math.min(BAR, Math.max(1, Math.round((percent / 100) * BAR)))
107
108  return [
109    { text: '━'.repeat(filled), color: percent >= 90 ? 'red' : percent >= 70 ? 'yellow' : 'green' },
110    { text: '─'.repeat(BAR - filled), color: 'gray' },
111  ]
112}
113
114// The plan's windows as gauges, and when the soonest one resets.
115const limitsPieces = (limits: Limit[]): Piece[] => {
116  const next = limits.filter(l => left(l.resetsAt) !== '').sort((a, b) => a.resetsAt! - b.resetsAt!)[0]
117  const reset = next === undefined ? '' : left(next.resetsAt)
118
119  return [
120    ...limits.flatMap((l, i): Piece[] => [
121      { text: i > 0 ? '  ' : '' },
122      { text: `${l.label} `, color: 'gray' },
123      ...gauge(l.percent),
124      { text: ` ${l.percent}%`, isFigure: false },
125    ]),
126    ...(reset === '' ? [] : [{ text: `  Resets in ${reset}`, color: 'gray' }]),
127  ]
128}
129
130// Tokens used: the session's total, then the running or last turn's, as `Token Usage 104k / 70k`.
131const usageOf = (m: Totals, withLast: boolean): string =>
132  `Token Usage ${count(m.tokens)}${withLast && m.turnTokens > 0 ? ` / ${count(m.turnTokens)}` : ''}`
133
134// The session's cost and tokens behind its mark, and the budget where one is shown.
135// On a plan with rate limits the dollars mean nothing, so the limits' used share stands in for them.
136const sessionOf = (m: Totals, budget = 0, withLast = false): Piece[] => [
137  { text: '◉ ' },
138  ...(m.limits.length > 0 ? limitsPieces(m.limits) : [{ text: money(m.total), isFigure: true }]),
139  ...(m.tokens > 0 ? [{ text: ' | ' }, { text: usageOf(m, withLast), isFigure: true }] : []),
140  ...(budget > 0 && m.limits.length === 0 ? [{ text: ` of ${money(budget)} budget` }] : []),
141]
142
143// The running turn's cost while it works, the last turn's after; null before the first and on a
144// plan, where the dollars mean nothing and the turn's tokens already sit beside the session's.
145const turnOf = (m: Totals, withCount = false): Piece[] | null => {
146  const cost = m.isWorking ? m.total - m.turnBase : m.last
147  if (cost === null || m.limits.length > 0) {
148    return null
149  }
150
151  return [
152    { text: m.isWorking ? 'now ' : 'last ' },
153    { text: money(cost), isFigure: true },
154    ...(withCount && !m.isWorking ? [{ text: ` · ${m.turns} ${m.turns === 1 ? 'turn' : 'turns'}` }] : []),
155  ]
156}
157
158const draw = ({ Box, Text }: Kit, key: string, pieces: Piece[]) => (
159  <Box key={key} flexDirection="row">
160    {pieces.map(piece =>
161      piece.text === '' ? null : piece.isFigure || piece.color !== undefined ? (
162        <Text color={piece.isFigure ? 'green' : piece.color}>{piece.text}</Text>
163      ) : (
164        <Text>{piece.text}</Text>
165      ),
166    )}
167  </Box>
168)
169
170export const register: Register = (on, options) => {
171  const budget = typeof options.budget === 'number' ? options.budget : 0
172
173  on('session.start', async ($, e, next) => {
174    const result = await next(e)
175    const { startedAt, cost, rateLimits } = await $.session.usage()
176    await spend($, startedAt, cost?.usd, budget, undefined, limitsOf(rateLimits)) // a new session starts from zero, a reload keeps its totals
177    await $.command.register({ name: COMMAND, description: 'Show what this session has cost, turn by turn' })
178
179    return result
180  })
181
182  // Fires whenever the status line's figures move, mid-turn too, so the meter keeps up live.
183  on('session.measure', async ($, e, next) => {
184    const result = await next(e)
185    if (e.changed.includes('cost') && e.cost !== undefined) {
186      await spend($, (await $.session.usage()).startedAt, e.cost.usd, budget, undefined, limitsOf(e.rateLimits))
187    } else if (e.changed.some(c => (c as string) === 'rateLimits')) {
188      await update($, meter, m => ({ ...m, limits: limitsOf(e.rateLimits) }))
189    }
190
191    return result
192  })
193
194  on('turn.start', async ($, e, next) => {
195    await update($, meter, m => ({ ...m, turnTokens: 0, isWorking: true }))
196
197    return next(e)
198  })
199
200  // One model request, the main loop's or a subagent's: its tokens count once its response is in.
201  on('turn.step', async function* ($, e, next) {
202    const result = yield* next(e)
203    if (result.usage !== null) {
204      const n = tokensOf(result.usage)
205      await update($, meter, m => ({
206        ...m,
207        tokens: m.tokens + n,
208        turnTokens: m.isWorking ? m.turnTokens + n : m.turnTokens,
209      }))
210    }
211
212    return result
213  })
214
215  on('turn.complete', async ($, e, next) => {
216    const result = await next(e)
217    if (e.agentId === undefined) {
218      const { startedAt, cost, rateLimits } = await $.session.usage()
219      await spend($, startedAt, cost?.usd, budget, m => {
220        const last = m.total - m.turnBase
221
222        return {
223          ...m,
224          turnBase: m.total,
225          last,
226          turns: m.turns + 1,
227          turnsCost: m.turnsCost + last,
228          priciest: Math.max(m.priciest, last),
229          isWorking: false,
230        }
231      }, limitsOf(rateLimits))
232    }
233
234    return result
235  })
236
237  on('session.attach', { surface: 'vscode' }, async ($, e, next) => {
238    const result = await next(e)
239    void $.ui.open({ id: PANE, title: TITLE })
240
241    return result
242  })
243
244  on('command.run', { command: COMMAND }, async $ => {
245    const m = await read($, meter)
246    if ((await $.session.surfaces()).includes('vscode')) {
247      await $.ui.open({ id: PANE, title: TITLE }) // asked for, so it is placed at any width
248    }
249
250    const text = report(m, budget)
251    if (m.limits.length > 0) {
252      return { text }
253    }
254    // No limits parsed: show what the host reports, unparsed, to find out why.
255    const u = await $.session.usage()
256
257    return { text: `${text}\n\n(no plan limits) usage keys: ${Object.keys(u).join(', ')}\nrateLimits: ${JSON.stringify(u.rateLimits)}` }
258  })
259
260  // Under the prompt, left of the engine's own mode labels, which stay as the engine drew them.
261  on('ui.render', { component: 'SessionMode' }, async ($, e, next) => {
262    const modes = await next(e)
263    const m = await read($, meter)
264    if (m.total === 0 && m.tokens === 0 && m.limits.length === 0) {
265      return modes
266    }
267    const isWide = (e.viewport?.columns ?? WIDE) >= WIDE
268    const turn = isWide ? turnOf(m) : null
269    const pieces = turn === null ? sessionOf(m, 0, isWide) : [...sessionOf(m, 0, true), { text: ' | ' }, ...turn]
270    const kit = $.ui.resolve(e)
271
272    return (
273      <kit.Box flexDirection="row">
274        {draw(kit, 'meter', pieces)}
275        {e.props.modes.length > 0 && <kit.Text>{'  '}</kit.Text>}
276        {modes}
277      </kit.Box>
278    )
279  })
280
281  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
282    const m = await read($, meter)
283    const kit = $.ui.resolve(e)
284    if (m.total === 0 && m.limits.length === 0) {
285      return <kit.Text>{NOTHING}</kit.Text>
286    }
287    const turn = turnOf(m, true)
288
289    return (
290      <kit.Box flexDirection="column" paddingX={1}>
291        {draw(kit, 'session', sessionOf(m, budget, true))}
292        {turn !== null && draw(kit, 'turn', turn)}
293      </kit.Box>
294    )
295  })
296}
297
298// The stored totals, or empty ones if they belong to an earlier session.
299const fresh = (m: Totals, startedAt: number): Totals =>
300  m.session === startedAt ? { ...EMPTY, ...m } : { ...EMPTY, session: startedAt }
301
302// Moves the session's cost to `usd` (kept where the host has no ledger) and applies `then`,
303// warning once when the cost passes the budget.
304const spend = async (
305  $: EngineInterface,
306  startedAt: number,
307  usd: number | undefined,
308  budget: number,
309  then: (m: Totals) => Totals = m => m,
310  limits?: Limit[],
311): Promise<Totals> => {
312  let isCrossed = false
313  const m = await update($, meter, value => {
314    const m = fresh(value, startedAt)
315    const total = usd ?? m.total
316    const isOver = budget > 0 && total >= budget && (limits ?? m.limits).length === 0
317    isCrossed = isOver && !m.warned
318    // A meter that starts mid-session (installed into it, a resumed session) counts turns from here.
319    const turnBase = value.session === startedAt ? Math.min(m.turnBase, total) : total
320
321    return then({ ...m, total, turnBase, warned: m.warned || isOver, limits: limits ?? m.limits })
322  })
323  if (isCrossed) {
324    await $.ui.toast(`Cost Info: this session passed your ${money(budget)} budget`)
325  }
326
327  return m
328}
329
330const report = (m: Totals, budget: number): string => {
331  if (m.total === 0 && m.limits.length === 0) {
332    return NOTHING
333  }
334  if (m.limits.length > 0) {
335    return [
336      'Plan usage:',
337      ...m.limits.map(l => `  ${l.label.padEnd(10)} ${l.percent}% used${left(l.resetsAt) === '' ? '' : ` (resets in ${left(l.resetsAt)})`}`),
338      ...(m.tokens > 0 ? [`  Tokens     ${tok(m.tokens)} this session`] : []),
339      ...(m.turns > 0 ? [`  Turns      ${m.turns}`] : []),
340    ].join('\n')
341  }
342  const lines = [`This session: ${money(m.total)}`]
343  if (m.tokens > 0) {
344    lines.push(`  Tokens        ${tok(m.tokens)}`)
345  }
346  if (m.turns > 0) {
347    lines.push(`  Turns         ${m.turns}`)
348    lines.push(`  Per turn      ${money(m.turnsCost / m.turns)} on average`)
349    lines.push(`  Priciest turn ${money(m.priciest)}`)
350  }
351  if (budget > 0) {
352    lines.push(`  Budget        ${money(budget)} (${Math.round((m.total / budget) * 100)}% used)`)
353  }
354
355  return lines.join('\n')
356}
357

1export type Limit = { label: string; percent: number; resetsAt: number | null }
2
3export type Totals = {
4  session: number | null
5  total: number
6  turnBase: number
7  last: number | null
8  turns: number
9  turnsCost: number
10  priciest: number
11  warned: boolean
12  tokens: number
13  turnTokens: number
14  isWorking: boolean
15  limits: Limit[]
16}
17
18declare module 'claude-code' {
19  interface PluginState {
20    'cost-info': { meter: Totals }
21  }
22}
23