token-meter

1import { atom, read, update } from 'claude-code'
2import type { Register, TurnUsage } from 'claude-code'
3
4import type { Usage } from '../types'
5
6// API list prices in USD per million tokens. Cache writes are priced at the
7// 1-hour TTL (2x input), which is what Claude Code sessions use.
8type Price = { input: number; output: number; cacheRead: number; cacheWrite: number }
9
10const PRICES: ReadonlyArray<[string, Price]> = [
11  ['fable', { input: 10, output: 50, cacheRead: 0.25, cacheWrite: 20 }],
12  ['mythos', { input: 10, output: 50, cacheRead: 0.25, cacheWrite: 20 }],
13  ['opus-5-5', { input: 4, output: 20, cacheRead: 0.2, cacheWrite: 8 }],
14  ['opus', { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 10 }],
15  ['sonnet-5-5', { input: 2, output: 10, cacheRead: 0.2, cacheWrite: 4 }],
16  ['sonnet-5', { input: 2, output: 10, cacheRead: 0.2, cacheWrite: 4 }],
17  ['sonnet', { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 6 }],
18  ['haiku', { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 2 }],
19]
20
21const DEFAULT_PRICE: Price = { input: 4, output: 20, cacheRead: 0.2, cacheWrite: 8 }
22
23const EMPTY: Usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, usd: 0, turns: 0 }
24
25const last = atom({ plugin: 'token-meter', key: 'last' } as const, null)
26const session = atom({ plugin: 'token-meter', key: 'session' } as const, EMPTY)
27const pending = atom({ plugin: 'token-meter', key: 'pending' } as const, EMPTY)
28
29export const priceFor = (model: string): Price =>
30  PRICES.find(([name]) => model.includes(name))?.[1] ?? DEFAULT_PRICE
31
32export const toUsage = (u: TurnUsage): Usage => {
33  const p = priceFor(u.model)
34  const usd =
35    (u.input_tokens * p.input +
36      u.output_tokens * p.output +
37      u.cache_read_input_tokens * p.cacheRead +
38      u.cache_creation_input_tokens * p.cacheWrite) /
39    1e6
40
41  return {
42    input: u.input_tokens,
43    output: u.output_tokens,
44    cacheRead: u.cache_read_input_tokens,
45    cacheWrite: u.cache_creation_input_tokens,
46    usd,
47    turns: 1,
48  }
49}
50
51export const add = (a: Usage, b: Usage): Usage => ({
52  input: a.input + b.input,
53  output: a.output + b.output,
54  cacheRead: a.cacheRead + b.cacheRead,
55  cacheWrite: a.cacheWrite + b.cacheWrite,
56  usd: a.usd + b.usd,
57  turns: a.turns + b.turns,
58})
59
60const tokens = (u: Usage) => u.input + u.output + u.cacheRead + u.cacheWrite
61
62const short = (n: number) =>
63  n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${Math.round(n / 1e3)}k` : `${n}`
64
65const usd = (n: number) => `$${n.toFixed(2)}`
66
67// Share of the prompt that came from the cache: high means cheap.
68const cacheShare = (u: Usage) => {
69  const prompt = u.input + u.cacheRead + u.cacheWrite
70  return prompt === 0 ? 0 : Math.round((u.cacheRead / prompt) * 100)
71}
72
73export const statusText = (lastPrompt: Usage | null, total: Usage) =>
74  lastPrompt === null
75    ? `Session: ${short(tokens(total))} tok · ~${usd(total.usd)}`
76    : `Last: ${short(tokens(lastPrompt))} tok · ~${usd(lastPrompt.usd)} (cache ${cacheShare(lastPrompt)}%) | Session: ~${usd(total.usd)}`
77
78const lines = (label: string, u: Usage) => [
79  `${label}: ~${usd(u.usd)}`,
80  `  cache read: ${short(u.cacheRead)} · cache write: ${short(u.cacheWrite)}`,
81  `  uncached input: ${short(u.input)} · output: ${short(u.output)}`,
82  `  from cache: ${cacheShare(u)}% of the prompt`,
83]
84
85export const register: Register = on => {
86  on('session.start', async ($, e, next) => {
87    await $.command.register({
88      name: 'token-cost',
89      description: 'Tokens and estimated cost of the last prompt and the whole session',
90    })
91    $.ui.status(statusText(await read($, last), await read($, session)))
92
93    return next(e)
94  })
95
96  on('turn.complete', async ($, e, next) => {
97    if (e.usage === undefined) {
98      return next(e)
99    }
100
101    const spent = toUsage(e.usage)
102    const total = await update($, session, s => add(s, spent))
103
104    if (e.agentId !== undefined) {
105      // A subagent's turn: count it toward the prompt that started it.
106      await update($, pending, p => add(p, spent))
107      $.ui.status(statusText(await read($, last), total))
108
109      return next(e)
110    }
111
112    const sub = await read($, pending)
113    const prompt = add(spent, sub)
114    await update($, pending, () => EMPTY)
115    await update($, last, () => ({ ...prompt, turns: 1 }))
116    $.ui.status(statusText(prompt, total))
117
118    return next(e)
119  })
120
121  on('command.run', { command: 'token-cost' }, async $ => {
122    const lastPrompt = await read($, last)
123    const total = await read($, session)
124    const engine = (await $.session.usage()).cost?.usd
125    const text = [
126      ...(lastPrompt === null ? ['Last prompt: no data yet'] : lines('Last prompt', lastPrompt)),
127      ...lines(`Session (${total.turns} responses since the mod loaded)`, total),
128      ...(engine === undefined ? [] : [`Claude Code's own total for the session: ${usd(engine)}`]),
129      'Estimate at API list prices; on the Max plan this uses your limits, not dollars.',
130    ].join('\n')
131
132    return { text }
133  })
134}
135

1export type Usage = {
2  input: number
3  output: number
4  cacheRead: number
5  cacheWrite: number
6  usd: number
7  turns: number
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'token-meter': { last: Usage | null; session: Usage; pending: Usage }
13  }
14}
15