token-meter

1import type { Register, EngineInterface } from 'claude-code'
2import {
3  emptyTotals,
4  addUsage,
5  claudeTotal,
6  parseJob,
7  localTotal,
8  statusLine,
9  spawnDecision,
10  paneLines,
11  spinnerSuffix,
12  freshTotal,
13  dotsOwed,
14  stepEaten,
15  contextSize,
16  ctxTarget,
17  stepToward,
18  ctxMaze,
19  ctxLabel,
20  scoreLine,
21} from './tally'
22import type { Totals, LocalJob, Run, Tone } from './tally'
23
24let totals: Totals = emptyTotals()
25let jobs: LocalJob[] = []
26let confirmed = new Set<string>()
27let sid = ''
28let jobsDir = ''
29const PANE = 'token-meter'
30let mainTurnId = ''
31let turnBase = 0
32let eaten = 0
33let ctx = 0
34let shown = 0
35
36async function save($: EngineInterface) {
37  await $.store.set('totals:' + sid, totals)
38}
39
40async function scanLocal($: EngineInterface) {
41  if (jobsDir === '' || !(await $.fs.exists(jobsDir))) return
42  try {
43    const entries = await $.fs.list(jobsDir)
44    const next: LocalJob[] = []
45    for (const e of entries) {
46      if (e.kind === 'file' && e.name.endsWith('.json')) {
47        const text = await $.fs.read(jobsDir + '/' + e.name)
48        const j = parseJob(text)
49        if (j !== null) next.push(j)
50      }
51    }
52    jobs = next
53  } catch {
54    // leave jobs unchanged
55  }
56}
57
58function show($: EngineInterface) {
59  $.ui.status(statusLine(claudeTotal(totals), localTotal(jobs)))
60  $.ui.invalidate('ui.render')
61}
62
63function toneProps(tone: Tone) {
64  if (tone === 'head') return { bold: true }
65  if (tone === 'dim') return { dimColor: true }
66  if (tone === 'amber') return { color: '#e8a33d' }
67  if (tone === 'green') return { color: '#6bcb77' }
68  if (tone === 'pellet') return { color: '#ffb8ae' }
69  if (tone === 'pac') return { color: '#ffff00', bold: true }
70  return {}
71}
72
73export const register: Register = on => {
74  on('session.start', async ($, e, next) => {
75    sid = await $.session.id()
76    const home = (await $.env.get('USERPROFILE')) ?? (await $.env.get('HOME')) ?? ''
77    jobsDir = home ? home.replace(/\\/g, '/') + '/.lanllm/jobs/' + sid : ''
78    const saved = await $.store.get('totals:' + sid)
79    if (
80      typeof saved === 'object' &&
81      saved !== null &&
82      typeof (saved as { byModel?: unknown }).byModel === 'object' &&
83      (saved as { byModel?: unknown }).byModel !== null
84    ) {
85      totals = saved as Totals
86    }
87    eaten = dotsOwed(freshTotal(totals))
88    const savedCtx = await $.store.get('ctx:' + sid)
89    if (typeof savedCtx === 'number') ctx = savedCtx
90    shown = ctxTarget(ctx)
91    await $.command.register({ name: 'tokens', description: 'Token use this session: Claude per model, local per box' })
92    await scanLocal($)
93    show($)
94    $.clock.every(20000, async () => {
95      await scanLocal($)
96      show($)
97    })
98    $.clock.every(120, () => {
99      const owed = dotsOwed(freshTotal(totals))
100      const target = ctxTarget(ctx)
101      let changed = false
102      if (eaten < owed) {
103        eaten = stepEaten(eaten, owed)
104        changed = true
105      }
106      if (shown !== target) {
107        shown = stepToward(shown, target)
108        changed = true
109      }
110      if (changed) $.ui.invalidate('ui.render')
111    })
112    return next(e)
113  })
114
115  on('turn.step', async function* ($, e, next) {
116    if (e.agentId === undefined && e.turnId !== mainTurnId) {
117      mainTurnId = e.turnId
118      turnBase = claudeTotal(totals).all
119    }
120    const r = yield* next(e)
121    if (r && r.usage) {
122      totals = addUsage(totals, r.usage)
123      if (e.agentId === undefined) {
124        ctx = contextSize(r.usage)
125        await $.store.set('ctx:' + sid, ctx)
126      }
127      show($)
128      await save($)
129    }
130    return r
131  })
132
133  on('agent.spawn', async ($, e, next) => {
134    const info = {
135      fork: e.fork,
136      model: e.model,
137      subagentType: e.subagentType,
138      description: e.description,
139      isTeammate: e.isTeammate,
140    }
141    const d = spawnDecision(info, confirmed)
142    if (d.action === 'pass') return next(e)
143    if (d.action === 'rewrite') return next({ ...e, model: 'haiku' })
144    confirmed.add(d.key)
145    return { deny: d.reason }
146  })
147
148  on('ui.render', { component: 'Spinner' }, async ($, e, next) => {
149    const n = claudeTotal(totals).all - turnBase
150    return next({ ...e, props: { ...e.props, suffix: spinnerSuffix(n) } })
151  })
152
153  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
154    const { Box, Text } = $.ui.resolve(e)
155    const fresh = freshTotal(totals)
156    const cols = e.props.bodyColumns
157    const lines: Run[][] = [
158      scoreLine(fresh),
159      ...ctxMaze(shown, eaten, dotsOwed(fresh)),
160      [{ text: ctxLabel(ctx), tone: 'dim' }],
161      [],
162      ...paneLines(totals, jobs, cols < 54),
163    ]
164    return (
165      <Box flexDirection="column">
166        {lines.map(line =>
167          line.length === 0 ? (
168            <Text> </Text>
169          ) : (
170            <Text wrap="wrap">
171              {line.map((r: Run) => (
172                <Text {...toneProps(r.tone)}>{r.text}</Text>
173              ))}
174            </Text>
175          ),
176        )}
177      </Box>
178    )
179  })
180
181  on('command.run', { command: 'tokens' }, async $ => {
182    await scanLocal($)
183    show($)
184    await $.ui.open({ id: PANE, title: 'Tokens · this session', columns: 56 })
185    return { text: 'Opened the tokens pane.' }
186  })
187}
188

1export type Usage = {
2  model: string
3  input_tokens: number
4  output_tokens: number
5  cache_read_input_tokens: number
6  cache_creation_input_tokens: number
7}
8export type ModelRow = { requests: number; input: number; output: number; cacheRead: number; cacheWrite: number }
9export type Totals = { byModel: Record<string, ModelRow> }
10export type LocalJob = { id: string; box: string; model: string; prompt: number; completion: number; done: boolean }
11export type SpawnInfo = { fork: boolean; model?: string; subagentType: string; description: string; isTeammate?: true }
12export type SpawnDecision =
13  | { action: 'pass' }
14  | { action: 'rewrite'; model: 'haiku' }
15  | { action: 'deny'; key: string; reason: string }
16
17function num(v: unknown): number {
18  return typeof v === 'number' && Number.isFinite(v) ? v : 0
19}
20
21export const CTX_WINDOW = 1000000 // context window in tokens
22export const CTX_DOT = 5000       // context tokens per dot
23export const CTX_DOTS = 200       // CTX_WINDOW / CTX_DOT
24export const ROW = 20   // dots per maze row
25export const ROWS = 10  // CTX_DOTS / ROW
26
27// The context one step's request carried: input + cache read + cache write.
28export function contextSize(u: Usage): number {
29  return num(u.input_tokens) + num(u.cache_read_input_tokens) + num(u.cache_creation_input_tokens)
30}
31
32// Dots eaten for a context size: Math.min(CTX_DOTS, Math.floor(ctx / CTX_DOT)).
33export function ctxTarget(ctx: number): number {
34  return Math.min(CTX_DOTS, Math.floor(ctx / CTX_DOT))
35}
36
37// One animation step of the eater's position: shown < target ? shown + 1 : target
38// (moving up walks one dot at a time; a drop, e.g. after compaction, jumps straight there).
39export function stepToward(shown: number, target: number): number {
40  return shown < target ? shown + 1 : target
41}
42
43// The maze: ROWS lines of ROW cells, two characters per cell, walked as a snake
44// (even rows left to right, odd rows right to left).
45// s = Math.max(0, Math.min(CTX_DOTS, shown)); p = Math.min(s, CTX_DOTS - 1)  (the eater's cell)
46// moving = owed > eaten
47// For row r (0..ROWS-1) and column j (0..ROW-1): idx = r * ROW + (r % 2 === 0 ? j : ROW - 1 - j)
48//   idx === p -> { text: ' ' + mouth, tone: 'pac' }
49//                mouth = moving && eaten % 2 === 0 ? 'O' : (r % 2 === 0 ? 'C' : 'Ɔ')   ('Ɔ' is U+0186)
50//   idx < s   -> { text: '  ', tone: 'plain' }
51//   otherwise -> { text: ' •', tone: 'pellet' }   ('•' is U+2022)
52// Each line is its cells left to right, with ADJACENT RUNS OF THE SAME TONE MERGED into one run (texts concatenated).
53export function ctxMaze(shown: number, eaten: number, owed: number): Run[][] {
54  const s = Math.max(0, Math.min(CTX_DOTS, shown))
55  const p = Math.min(s, CTX_DOTS - 1)
56  const moving = owed > eaten
57  const rows: Run[][] = []
58  for (let r = 0; r < ROWS; r++) {
59    const cells: Run[] = []
60    for (let j = 0; j < ROW; j++) {
61      const idx = r * ROW + (r % 2 === 0 ? j : ROW - 1 - j)
62      let run: Run
63      if (idx === p) {
64        const mouth = moving && eaten % 2 === 0 ? 'O' : (r % 2 === 0 ? 'C' : '\u0186')
65        run = { text: ' ' + mouth, tone: 'pac' }
66      } else if (idx < s) {
67        run = { text: '  ', tone: 'plain' }
68      } else {
69        run = { text: ' \u2022', tone: 'pellet' }
70      }
71      const last = cells[cells.length - 1]
72      if (last !== undefined && last.tone === run.tone) last.text += run.text
73      else cells.push(run)
74    }
75    rows.push(cells)
76  }
77  return rows
78}
79
80// Label under the maze: '1 dot = 5k context · ' + formatCount(ctx) + ' of 1M'  (middle dot U+00B7)
81export function ctxLabel(ctx: number): string {
82  return '1 dot = 5k context \u00B7 ' + formatCount(ctx) + ' of 1M'
83}
84
85export function emptyTotals(): Totals {
86  return { byModel: {} }
87}
88
89export function addUsage(t: Totals, u: Usage): Totals {
90  const byModel: Record<string, ModelRow> = {}
91  for (const [m, r] of Object.entries(t.byModel)) byModel[m] = { ...r }
92  const model = typeof u.model === 'string' ? u.model : '?'
93  const row = byModel[model] ?? { requests: 0, input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }
94  row.requests += 1
95  row.input += num(u.input_tokens)
96  row.output += num(u.output_tokens)
97  row.cacheRead += num(u.cache_read_input_tokens)
98  row.cacheWrite += num(u.cache_creation_input_tokens)
99  byModel[model] = row
100  return { byModel }
101}
102
103export function claudeTotal(t: Totals): { all: number; cached: number } {
104  let all = 0
105  let cached = 0
106  for (const r of Object.values(t.byModel)) {
107    all += r.input + r.output + r.cacheRead + r.cacheWrite
108    cached += r.cacheRead
109  }
110  return { all, cached }
111}
112
113export function parseJob(text: string): LocalJob | null {
114  let obj: unknown
115  try {
116    obj = JSON.parse(text)
117  } catch {
118    return null
119  }
120  if (typeof obj !== 'object' || obj === null) return null
121  const o = obj as Record<string, any>
122  const result = o.result
123  if (typeof result !== 'object' || result === null || typeof result.usage !== 'object' || result.usage === null) return null
124  const usage = result.usage as Record<string, any>
125  const server = (typeof o.server === 'object' && o.server !== null ? o.server : {}) as Record<string, any>
126  let model = '?'
127  if (typeof server.model === 'string' && server.model.length > 0) {
128    const base = server.model.split(/[\\/]/).pop() ?? ''
129    model = base.endsWith('.gguf') ? base.slice(0, -'.gguf'.length) : base
130    if (model.length === 0) model = '?'
131  }
132  return {
133    id: typeof o.id === 'string' ? o.id : String(o.id ?? '?'),
134    box: typeof server.name === 'string' && server.name.length > 0 ? server.name : '?',
135    model,
136    prompt: num(usage.promptTokens),
137    completion: num(usage.completionTokens),
138    done: o.state === 'done',
139  }
140}
141
142export function localTotal(jobs: LocalJob[]): { all: number; byBox: Record<string, number> } {
143  let all = 0
144  const byBox: Record<string, number> = {}
145  for (const j of jobs) {
146    const s = j.prompt + j.completion
147    all += s
148    byBox[j.box] = (byBox[j.box] ?? 0) + s
149  }
150  return { all, byBox }
151}
152
153export function formatCount(n: number): string {
154  if (n < 1000) return String(n)
155  if (n < 1_000_000) return (n / 1000).toFixed(1) + 'k'
156  return (n / 1_000_000).toFixed(2) + 'M'
157}
158
159export function statusLine(claude: { all: number; cached: number }, local: { all: number }): string {
160  return `Claude ${formatCount(claude.all)} (${formatCount(claude.cached)} cached) · Local ${formatCount(local.all)}`
161}
162
163export function spawnDecision(s: SpawnInfo, confirmed: ReadonlySet<string>): SpawnDecision {
164  const key = `${s.subagentType}|${s.model ?? 'inherit'}|${s.fork ? 'fork' : 'agent'}|${s.description}`
165  if (s.isTeammate === true) return { action: 'pass' }
166  if (confirmed.has(key)) return { action: 'pass' }
167  if (s.fork === true) {
168    return {
169      action: 'deny',
170      key,
171      reason:
172        'token-meter: routing rule is Haiku first unless Haiku already failed at this task; forks bypass that rule. re-issue the identical call to confirm this fork.',
173    }
174  }
175  if (s.model === undefined) return { action: 'rewrite', model: 'haiku' }
176  if (s.model === 'haiku' || s.model.includes('haiku')) return { action: 'pass' }
177  return {
178    action: 'deny',
179    key,
180    reason:
181      `token-meter: routing rule is Haiku first unless Haiku already failed at this task; ` +
182      `${s.model} is not Haiku. re-issue the identical call to confirm it.`,
183  }
184}
185
186export type Tone = 'head' | 'name' | 'dim' | 'amber' | 'green' | 'plain' | 'pellet' | 'pac'
187export type Run = { text: string; tone: Tone }
188
189// Drop a leading "claude-" and a trailing "-YYYYMMDD" (8 digits).
190// "claude-opus-5-5" -> "opus-5-5"; "claude-haiku-4-5-20251001" -> "haiku-4-5"; "qwen" -> "qwen".
191export function shortModel(id: string): string {
192  let s = id
193  if (s.startsWith('claude-')) s = s.slice('claude-'.length)
194  s = s.replace(/-\d{8}$/, '')
195  return s
196}
197
198// Lines for the pane. Each line is a list of runs; an empty array is a blank line.
199// compact narrows the table: NAME 10 / COL 7 / names sliced to 10, instead of 14 / 8 / 13.
200export function paneLines(t: Totals, jobs: LocalJob[], compact = false): Run[][] {
201  const NAME = compact ? 10 : 14
202  const COL = compact ? 7 : 8
203  const SLICE = compact ? 10 : 13
204  const lines: Run[][] = []
205  lines.push([
206    { text: 'Claude'.padEnd(NAME) + ['req', 'in', 'write', 'read', 'out'].map(s => s.padStart(COL)).join(''), tone: 'head' },
207  ])
208  const models = Object.entries(t.byModel)
209  if (models.length === 0) {
210    lines.push([{ text: '  (none yet)', tone: 'dim' }])
211  } else {
212    const sorted = models
213      .map(([m, r]) => ({ m, r, total: r.input + r.output + r.cacheRead + r.cacheWrite }))
214      .sort((a, b) => b.total - a.total)
215    for (const { m, r } of sorted) {
216      lines.push([
217        { text: shortModel(m).slice(0, SLICE).padEnd(NAME), tone: 'name' },
218        {
219          text:
220            String(r.requests).padStart(COL) +
221            formatCount(r.input).padStart(COL) +
222            formatCount(r.cacheWrite).padStart(COL),
223          tone: 'plain',
224        },
225        { text: formatCount(r.cacheRead).padStart(COL), tone: 'dim' },
226        { text: formatCount(r.output).padStart(COL), tone: 'plain' },
227      ])
228    }
229  }
230  lines.push([])
231  lines.push([{ text: 'Local'.padEnd(NAME) + 'jobs'.padStart(COL) + 'total'.padStart(COL), tone: 'head' }])
232  const boxes: Record<string, { count: number; sum: number }> = {}
233  for (const j of jobs) {
234    const b = boxes[j.box] ?? { count: 0, sum: 0 }
235    b.count += 1
236    b.sum += j.prompt + j.completion
237    boxes[j.box] = b
238  }
239  const boxEntries = Object.entries(boxes)
240  if (boxEntries.length === 0) {
241    lines.push([{ text: '  (none yet)', tone: 'dim' }])
242  } else {
243    const sorted = boxEntries.sort((a, b) => b[1].sum - a[1].sum)
244    for (const [box, b] of sorted) {
245      lines.push([
246        { text: box.slice(0, SLICE).padEnd(NAME), tone: 'name' },
247        { text: String(b.count).padStart(COL) + formatCount(b.sum).padStart(COL), tone: 'plain' },
248      ])
249    }
250  }
251  lines.push([])
252  const claude = claudeTotal(t)
253  const local = localTotal(jobs)
254  lines.push([
255    { text: 'Total   ', tone: 'plain' },
256    { text: 'Claude ' + formatCount(claude.all), tone: 'amber' },
257    { text: ' \u00B7 ', tone: 'plain' },
258    { text: 'Local ' + formatCount(local.all), tone: 'green' },
259  ])
260  return lines
261}
262
263// Spinner suffix: '… · Claude +' + formatCount(n) + ' this turn'
264export function spinnerSuffix(n: number): string {
265  return '\u2026 \u00B7 Claude +' + formatCount(n) + ' this turn'
266}
267
268export const TRACK = 40 // track width in cells
269export const DOT = 1000 // fresh tokens per dot
270
271// Fresh tokens: input + output + cacheWrite over all models (cache reads NOT counted).
272// Equals claudeTotal(t).all - claudeTotal(t).cached.
273export function freshTotal(t: Totals): number {
274  const c = claudeTotal(t)
275  return c.all - c.cached
276}
277
278// Dots earned so far: Math.floor(fresh / DOT).
279export function dotsOwed(fresh: number): number {
280  return Math.floor(fresh / DOT)
281}
282
283// One tick of the eater. If eaten >= owed, return eaten unchanged.
284// Otherwise skip any backlog beyond TRACK, then eat one:
285//   return Math.max(eaten, owed - TRACK) + 1
286export function stepEaten(eaten: number, owed: number): number {
287  if (eaten >= owed) return eaten
288  return Math.max(eaten, owed - TRACK) + 1
289}
290
291// Arcade score line. The number is fresh with commas every three digits,
292// built WITHOUT toLocaleString/Intl (the runtime may lack them):
293//   String(fresh).replace(/\B(?=(\d{3})+(?!\d))/g, ',')
294// Returns exactly two runs:
295// [{ text: 'SCORE ', tone: 'head' }, { text: <grouped number>, tone: 'amber' }]
296export function scoreLine(fresh: number): Run[] {
297  const grouped = String(fresh).replace(/\B(?=(\d{3})+(?!\d))/g, ',')
298  return [
299    { text: 'SCORE ', tone: 'head' },
300    { text: grouped, tone: 'amber' },
301  ]
302}
303