token-meter

1import { atom, read, update } from 'claude-code'
2import type { Register, EngineInterface } from 'claude-code'
3
4const prompt = atom({ plugin: 'token-meter', key: 'prompt' } as const, 0)
5const total = atom({ plugin: 'token-meter', key: 'total' } as const, 0)
6const ctxTokens = atom({ plugin: 'token-meter', key: 'ctxTokens' } as const, 0)
7const ctxWindow = atom({ plugin: 'token-meter', key: 'ctxWindow' } as const, 0)
8const ctxPercent = atom({ plugin: 'token-meter', key: 'ctxPercent' } as const, 0)
9
10const model = atom({ plugin: 'token-meter', key: 'model' } as const, '')
11const effort = atom({ plugin: 'token-meter', key: 'effort' } as const, '')
12
13// the host's cost ledger in USD (-1 where it keeps none, which hides the
14// segment), and where it stood at the last /clear and at the current turn's start
15const ledger = atom({ plugin: 'token-meter', key: 'ledger' } as const, -1)
16const clearBase = atom({ plugin: 'token-meter', key: 'clearBase' } as const, 0)
17const turnBase = atom({ plugin: 'token-meter', key: 'turnBase' } as const, 0)
18
19const ICON = '◆'
20const PROMPT_COLOR = '#0e7490'
21const TOTAL_COLOR = '#92400e'
22const MILLION_COLOR = '#b91c1c'
23const COST_COLOR = '#6d28d9'
24
25// context-fill colors by percentage: calm → warming → full
26const CTX_LOW = '#15803d'
27const CTX_MID = '#b45309'
28const CTX_HIGH = '#b91c1c'
29const ctxColor = (pct: number) => (pct >= 80 ? CTX_HIGH : pct >= 50 ? CTX_MID : CTX_LOW)
30
31// 1 decimal, right-aligned in 7 columns: "   2.8k", "   1.1M"
32const fmt = (n: number) =>
33  (n >= 999_950 ? `${(n / 1_000_000).toFixed(1)}M` : `${(n / 1000).toFixed(1)}k`).padStart(7)
34
35// compact, no padding, for the [ctx] segment: "92k", "1.1M"
36const fmtTight = (n: number) =>
37  n >= 999_950 ? `${(n / 1_000_000).toFixed(1)}M` : `${Math.round(n / 1000)}k`
38
39// collapse a model id to a compact label for narrow terminals:
40// "eu.anthropic.claude-opus-4-8" → "opus-4.8"; any region prefix and trailing
41// date stamp are dropped. Unrecognised shapes fall back to the bare id.
42const shortModel = (id: string) => {
43  const bare = id
44    .replace(/^[a-z]+\.anthropic\./, '') // strip "<region>.anthropic." prefix
45    .replace(/^claude-/, '') //            strip the "claude-" family prefix
46  // family + major.minor, e.g. "opus-4-8" → "opus-4.8" (ignore any -YYYYMMDD)
47  const m = bare.match(/^([a-z]+)-(\d+)-(\d+)/)
48  return m ? `${m[1]}-${m[2]}.${m[3]}` : bare
49}
50
51// "$0.42", "$12.35"; a nonzero amount under half a cent shows as "<$0.01"
52const fmtUsd = (usd: number) => (usd > 0 && usd < 0.005 ? '<$0.01' : `$${usd.toFixed(2)}`)
53
54// take the ledger's new reading; one below a baseline means the host reset it
55// (e.g. on /clear), so count that baseline from 0
56const syncCost = async ($: EngineInterface, usd: number | undefined) => {
57  if (usd === undefined) return
58  await update($, clearBase, b => (usd < b ? 0 : b))
59  await update($, turnBase, b => (usd < b ? 0 : b))
60  await update($, ledger, () => usd)
61}
62
63// pull the live context-window fill from $.session.usage() into the atoms;
64// `tokens`/`percent` are absent until the first response of a fresh/compacted window
65const syncContext = async ($: EngineInterface) => {
66  const { context, cost } = await $.session.usage()
67  await update($, ctxWindow, () => context.window ?? 0)
68  await update($, ctxTokens, () => context.tokens ?? 0)
69  await update($, ctxPercent, () => context.percent ?? 0)
70  await syncCost($, cost?.usd)
71}
72
73// one spelling for every form a model is named in: "eu.anthropic.claude-opus-5-5[1m]",
74// "Opus 5.5" and "claude-opus-5-5-20260101" all become "opus-5-5"
75const normModel = (id: string) =>
76  id
77    .toLowerCase()
78    .replace(/\[.*?\]/g, '')
79    .replace(/^[a-z]+\.anthropic\./, '')
80    .replace(/^claude[-\s]/, '')
81    .trim()
82    .replace(/[\s.]+/g, '-')
83    .replace(/-\d{8}$/, '')
84
85// show model + effort before the first prompt; the first turn.step then replaces
86// the effort with the one the request actually used
87const seedEffort = async ($: EngineInterface, current: string) => {
88  const s = (await $.settings.read()) as {
89    effortLevel?: string
90    modelSettings?: Record<string, { effortLevel?: string }>
91  }
92  // $.session.model() answers as /model shows it ("Opus 5.5", "opus", an id
93  // with "[1m]"), while modelSettings is keyed by id ("claude-opus-5-5")
94  const cur = normModel(current)
95  const key = Object.keys(s.modelSettings ?? {}).find(k => {
96    const nk = normModel(k)
97    return nk === cur || cur.startsWith(`${nk}-`) || nk.startsWith(`${cur}-`)
98  })
99  return (key && s.modelSettings?.[key]?.effortLevel) || s.effortLevel || ''
100}
101
102const syncModel = async ($: EngineInterface) => {
103  const current = await $.session.model()
104  if (!current) return
105  await update($, model, () => current)
106  const seeded = await seedEffort($, current)
107  if (seeded) await update($, effort, () => seeded)
108}
109
110// temporary: toasts which effort-change paths the mod can actually see
111const DEBUG = true
112const debug = ($: EngineInterface, text: string) => {
113  if (DEBUG) $.ui.toast(`token-meter debug: ${text}`)
114}
115
116const LEVELS = ['low', 'medium', 'high', 'xhigh', 'max']
117
118// no event fires when effort changes from a picker or the desktop dropdown, so
119// poll the saved effort and take it only when it moves — between polls the
120// value turn.step reported stays on screen
121const SETTINGS_POLL_MS = 1500
122let lastSaved = ''
123let poll: { cancel: () => void } | undefined
124const startSettingsPoll = ($: EngineInterface) => {
125  poll?.cancel()
126  poll = $.clock.every(SETTINGS_POLL_MS, async () => {
127    const current = await $.session.model()
128    if (!current) return
129    const saved = await seedEffort($, current)
130    if (saved === lastSaved) return
131    lastSaved = saved
132    debug($, `saved effort changed → ${saved || '(none)'}`)
133    if (saved) await update($, effort, () => saved)
134  })
135}
136
137export const register: Register = on => {
138  // clears the pinned status line the 0.1.0 version left under the prompt
139  on('session.start', async ($, e, next) => {
140    $.ui.status(undefined)
141    await syncModel($)
142    const current = await $.session.model()
143    lastSaved = current ? await seedEffort($, current) : ''
144    debug($, `session.start model="${current}" seeded effort="${lastSaved}"`)
145    startSettingsPoll($)
146    // a resumed session starts with the ledger already running
147    await syncCost($, (await $.session.usage()).cost?.usd)
148    return next(e)
149  })
150
151  // /model and /effort before any turn: no turn.step has run yet, so refresh here
152  on('command.run', async ($, e, next) => {
153    const r = await next(e)
154    debug($, `command.run /${e.command}${e.args ? ` ${e.args}` : ''}`)
155    if (e.command === 'effort') {
156      // "/effort high" names the level outright: show it without waiting for
157      // settings (a session-only level may never be saved there)
158      const level = e.args.trim().split(/\s+/)[0]?.toLowerCase() ?? ''
159      if (LEVELS.includes(level)) {
160        await update($, effort, () => level)
161        return r
162      }
163    }
164    if (e.command === 'model' || e.command === 'effort') await syncModel($)
165    return r
166  })
167
168  on('turn.start', async ($, e, next) => {
169    $.ui.status(undefined)
170    await update($, prompt, () => 0)
171    const l = await read($, ledger)
172    await update($, turnBase, () => Math.max(l, 0))
173    return next(e)
174  })
175
176  // input + cache writes + output; cache reads are left out (the whole context is re-read every step)
177  on('turn.step', async function* ($, e, next) {
178    // main loop only: a subagent's step may run on another model
179    if (!e.agentId) {
180      await update($, model, () => e.model)
181      await update($, effort, () => (e.effort === undefined ? '' : String(e.effort)))
182    }
183    const r = yield* next(e)
184    const u = r.usage
185    if (u) {
186      const n = u.input_tokens + u.output_tokens + u.cache_creation_input_tokens
187      await update($, prompt, v => v + n)
188      await update($, total, v => v + n)
189    }
190    // real-time context-window fill: refreshed after each agent step
191    await syncContext($)
192    return r
193  })
194
195  // the engine's own push after each turn and on compaction — catches the
196  // window shrinking (compaction / fresh window) outside a step
197  on('session.measure', async ($, e, next) => {
198    if (e.changed.includes('context')) {
199      await update($, ctxWindow, () => e.context.window ?? 0)
200      await update($, ctxTokens, () => e.context.tokens ?? 0)
201      await update($, ctxPercent, () => e.context.percent ?? 0)
202    }
203    if (e.changed.includes('cost')) await syncCost($, e.cost?.usd)
204    return next(e)
205  })
206
207  on('session.end', async ($, e, next) => {
208    await update($, prompt, () => 0)
209    await update($, total, () => 0)
210    await update($, ctxTokens, () => 0)
211    await update($, ctxPercent, () => 0)
212    // the ledger may run on across /clear: count cost from where it stands now
213    const l = Math.max(await read($, ledger), 0)
214    await update($, clearBase, () => l)
215    await update($, turnBase, () => l)
216    // /clear keeps the process alive and fires no session.start after it, so keep model + effort
217    if (e.reason === 'clear') await syncModel($)
218    else {
219      poll?.cancel()
220      poll = undefined
221      await update($, model, () => '')
222      await update($, effort, () => '')
223    }
224    return next(e)
225  })
226
227  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
228    const p = await read($, prompt)
229    const t = await read($, total)
230    const ct = await read($, ctxTokens)
231    const cw = await read($, ctxWindow)
232    const cp = await read($, ctxPercent)
233    const usd = await read($, ledger)
234    const clearUsd = usd - (await read($, clearBase))
235    const promptUsd = usd - (await read($, turnBase))
236    // /clear can wipe the atoms after session.end ran, so re-seed whenever the model is missing
237    // read-only: render may not write state
238    let cur = await read($, model)
239    let ef = await read($, effort)
240    if (!cur) {
241      cur = (await $.session.model()) ?? ''
242      if (cur && !ef) ef = await seedEffort($, cur)
243    }
244    const m = shortModel(cur)
245    const { Box, Text } = $.ui.resolve(e)
246    const isMillions = t >= 999_950
247    const below = await next(e)
248
249    return (
250      <Box flexDirection="column">
251        <Box>
252          <Text dimColor>{ICON} </Text>
253          <Text color={PROMPT_COLOR}>[prompt:{fmt(p)}]</Text>
254          <Text color={isMillions ? MILLION_COLOR : TOTAL_COLOR} bold={isMillions}>
255            [total from last clear:{fmt(t)}]
256          </Text>
257          <Text color={ctxColor(cp)} bold={cp >= 80}>
258            {' '}
259            [ctx:{cw > 0 ? ` ${cp}% ${fmtTight(ct)}/${fmtTight(cw)}` : ' —'}]
260          </Text>
261          {usd >= 0 && (
262            <Text color={COST_COLOR}>
263              {' '}
264              [cost: {fmtUsd(promptUsd)} / {fmtUsd(clearUsd)}]
265            </Text>
266          )}
267          {m && <Text dimColor> [{ef ? `${m} · ${ef}` : m}]</Text>}
268        </Box>
269        {below}
270      </Box>
271    )
272  })
273}
274

1export type Tokens = number
2
3declare module 'claude-code' {
4  interface PluginState {
5    'token-meter': {
6      prompt: Tokens
7      total: Tokens
8      ctxTokens: Tokens
9      ctxWindow: Tokens
10      ctxPercent: number
11      model: string
12      effort: string
13      ledger: number
14      clearBase: number
15      turnBase: number
16    }
17  }
18}
19