tps

1// tps: tokens per second for every model response.
2//   status line:  ⚡ 62 t/s · sess 55 · opus
3//   /tps:         last, session, today, 7-day and all-time averages per model
4// Speed is output tokens over first chunk to response end, per `turn.step` (a model request).
5import type { Register } from 'claude-code'
6
7import { byModel, record, rate, table } from './stats'
8import type { Bucket, History } from './stats'
9
10const KEY = 'history'
11
12export const register: Register = (on, _options) => {
13  let session: Record<string, Bucket> = {}
14  const last: Record<string, number> = {}
15
16  on('session.start', async ($, e, next) => {
17    await $.command.register({ name: 'tps', description: 'Tokens per second per model: last, session, today, 7d, all-time.', immediate: true })
18    return next(e)
19  })
20
21  on('turn.step', async function* ($, e, next) {
22    // Time from the first streamed chunk to the response's end: the wait for the first token is not generation.
23    let first = 0
24    const stream = next(e)
25    let step = await stream.next()
26    while (!step.done) {
27      if (first === 0) first = await $.clock.now()
28      yield step.value
29      step = await stream.next()
30    }
31    const r = step.value
32    const u = r.usage
33    if (u && first > 0) {
34      const now = await $.clock.now()
35      const ms = now - first
36      const h = ((await $.store.get(KEY)) ?? {}) as History
37      const after = record(h, u.model, u.output_tokens, ms, now)
38      if (after !== h) {
39        await $.store.set(KEY, after)
40        const b = session[u.model] ?? { tokens: 0, ms: 0, steps: 0 }
41        session = { ...session, [u.model]: { tokens: b.tokens + u.output_tokens, ms: b.ms + ms, steps: b.steps + 1 } }
42        last[u.model] = (u.output_tokens / ms) * 1000
43        $.ui.status(`⚡ ${last[u.model]!.toFixed(0)} t/s · sess ${rate(session[u.model]!).toFixed(0)} · ${u.model.replace(/^claude-/, '')}`)
44      }
45    }
46    return r
47  })
48
49  on('command.run', { command: 'tps' }, async $ => {
50    const h = ((await $.store.get(KEY)) ?? {}) as History
51    return { text: table(h, session, last, await $.clock.now()) }
52  })
53}
54

1// Pure bookkeeping for the tps mod: buckets of output tokens and generation time, per day per model.
2export type Bucket = { tokens: number; ms: number; steps: number }
3export type History = Record<string, Record<string, Bucket>> // day -> model -> bucket
4
5// Responses shorter than this are mostly latency, not speed.
6export const MIN_TOKENS = 20
7export const MIN_MS = 300
8const KEEP_DAYS = 90
9
10export const rate = (b: Bucket): number => (b.ms > 0 ? (b.tokens / b.ms) * 1000 : 0)
11
12export function day(ms: number): string {
13  return new Date(ms).toISOString().slice(0, 10)
14}
15
16export function record(h: History, model: string, tokens: number, ms: number, now: number): History {
17  if (tokens < MIN_TOKENS || ms < MIN_MS) return h
18  const d = day(now)
19  const b = h[d]?.[model] ?? { tokens: 0, ms: 0, steps: 0 }
20  const next: History = { ...h, [d]: { ...h[d], [model]: { tokens: b.tokens + tokens, ms: b.ms + ms, steps: b.steps + 1 } } }
21  const days = Object.keys(next).sort()
22  for (const old of days.slice(0, Math.max(0, days.length - KEEP_DAYS))) delete next[old]
23  return next
24}
25
26// Sums every bucket of each model over the last `days` days (all of them when omitted).
27export function byModel(h: History, now: number, days?: number): Record<string, Bucket> {
28  const from = days === undefined ? '' : day(now - (days - 1) * 86_400_000)
29  const out: Record<string, Bucket> = {}
30  for (const [d, models] of Object.entries(h)) {
31    if (d < from) continue
32    for (const [m, b] of Object.entries(models)) {
33      const t = out[m] ?? { tokens: 0, ms: 0, steps: 0 }
34      out[m] = { tokens: t.tokens + b.tokens, ms: t.ms + b.ms, steps: t.steps + b.steps }
35    }
36  }
37  return out
38}
39
40export function table(h: History, session: Record<string, Bucket>, last: Record<string, number>, now: number): string {
41  const today = byModel(h, now, 1)
42  const week = byModel(h, now, 7)
43  const all = byModel(h, now)
44  const models = Object.keys(all).sort((a, b) => rate(all[b]!) - rate(all[a]!))
45  if (models.length === 0) return 'No samples yet. Speeds appear after a response of 20+ tokens.'
46  const f = (b?: Bucket) => (b && b.steps ? rate(b).toFixed(1) : '·').padStart(7)
47  const rows = models.map(m => `${m.padEnd(26)}${(last[m]?.toFixed(1) ?? '·').padStart(7)}${f(session[m])}${f(today[m])}${f(week[m])}${f(all[m])}${String(all[m]!.steps).padStart(7)}`)
48  return ['tokens/s (output tokens over streaming time, per model request)', `${'model'.padEnd(26)}${'last'.padStart(7)}${'sess'.padStart(7)}${'today'.padStart(7)}${'7d'.padStart(7)}${'all'.padStart(7)}${'steps'.padStart(7)}`, ...rows].join('\n')
49}
50