tps-meter

1import type { EngineInterface, Register } from 'claude-code'
2
3// Times every main-loop model request and writes the figures to
4// <config dir>/tps/<session id>.json, which statusline-command.sh reads.
5//
6//   ttft_ms  request sent -> first stream event (network + queue + prefill)
7//   tps      output_tokens / (stop - first event); thinking tokens count too
8//
9// A step that generates only a few tokens arrives in one burst, so its rate
10// is noise: `tps` is null for it and `last_tps` keeps the latest usable one.
11
12const MIN_TOKENS = 50
13const MIN_DECODE_MS = 250
14
15type Totals = { steps: number; outputTokens: number; decodeMs: number; ttftMs: number }
16
17type Snapshot = {
18  model: string
19  updated_at: number
20  ttft_ms: number
21  output_tokens: number
22  decode_ms: number
23  tps: number | null
24  last_tps: number | null
25  avg_tps: number | null
26  avg_ttft_ms: number
27  steps: number
28  session_output_tokens: number
29  session_decode_ms: number
30  session_ttft_ms: number
31}
32
33const emptyTotals = (): Totals => ({ steps: 0, outputTokens: 0, decodeMs: 0, ttftMs: 0 })
34
35// Module state starts over on a hot reload; restore() picks the session's
36// totals back up from the file the previous load wrote.
37let sessionId: string | undefined
38let totals = emptyTotals()
39let lastTps: number | null = null
40let dir: string | undefined
41
42const round1 = (n: number) => Math.round(n * 10) / 10
43
44async function metricsDir($: EngineInterface): Promise<string> {
45  if (dir === undefined) {
46    const config = (await $.env.get('CLAUDE_CONFIG_DIR')) ?? `${await $.env.get('HOME')}/.claude`
47    dir = `${config}/tps`
48  }
49  return dir
50}
51
52async function restore($: EngineInterface, path: string): Promise<void> {
53  totals = emptyTotals()
54  lastTps = null
55  try {
56    const prev = JSON.parse(String(await $.fs.read(path))) as Partial<Snapshot>
57    totals = {
58      steps: prev.steps ?? 0,
59      outputTokens: prev.session_output_tokens ?? 0,
60      decodeMs: prev.session_decode_ms ?? 0,
61      ttftMs: prev.session_ttft_ms ?? 0,
62    }
63    lastTps = prev.last_tps ?? null
64  } catch {
65    // no file yet: a new session
66  }
67}
68
69async function record(
70  $: EngineInterface,
71  model: string,
72  ttftMs: number,
73  decodeMs: number,
74  outputTokens: number,
75): Promise<void> {
76  const id = await $.session.id()
77  const path = `${await metricsDir($)}/${id}.json`
78  if (id !== sessionId) {
79    sessionId = id
80    await restore($, path)
81  }
82
83  totals.steps += 1
84  totals.outputTokens += outputTokens
85  totals.decodeMs += decodeMs
86  totals.ttftMs += ttftMs
87
88  const tps = outputTokens >= MIN_TOKENS && decodeMs >= MIN_DECODE_MS ? round1((outputTokens / decodeMs) * 1000) : null
89  if (tps !== null) lastTps = tps
90
91  const snapshot: Snapshot = {
92    model,
93    updated_at: Date.now(),
94    ttft_ms: Math.round(ttftMs),
95    output_tokens: outputTokens,
96    decode_ms: Math.round(decodeMs),
97    tps,
98    last_tps: lastTps,
99    avg_tps: totals.decodeMs > 0 ? round1((totals.outputTokens / totals.decodeMs) * 1000) : null,
100    avg_ttft_ms: Math.round(totals.ttftMs / totals.steps),
101    steps: totals.steps,
102    session_output_tokens: totals.outputTokens,
103    session_decode_ms: Math.round(totals.decodeMs),
104    session_ttft_ms: Math.round(totals.ttftMs),
105  }
106  await $.fs.write(path, JSON.stringify(snapshot))
107}
108
109export const register: Register = (on) => {
110  on('turn.step', async function* ($, e, next) {
111    // Subagents run in parallel with the main loop; only the main loop's
112    // requests are what the status line describes.
113    if (e.agentId !== undefined) return yield* next(e)
114
115    const sentAt = performance.now()
116    const stream = next(e)
117    let firstAt: number | undefined
118    for await (const chunk of stream) {
119      const now = performance.now()
120      firstAt ??= now
121      if (chunk.kind === 'stop' && chunk.usage) {
122        // Written before the stop chunk is passed on, so the status line
123        // refresh that the engine's usage update triggers already reads it.
124        await record($, e.model, firstAt - sentAt, now - firstAt, chunk.usage.output_tokens).catch(() => {})
125      }
126      yield chunk
127    }
128    return await stream.result
129  })
130}
131