turn-telemetry

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { StepStat } from '../types'
5import { HISTORY, SEPARATOR, fit, segments, stallCounter, summarize } from './stats'
6
7const COMMAND = 'tps'
8/** How often the band redraws while a response streams. */
9const LIVE_EVERY_MS = 1000
10
11const last = atom({ plugin: 'turn-telemetry', key: 'last' } as const, null)
12const history = atom({ plugin: 'turn-telemetry', key: 'history' } as const, [])
13const live = atom({ plugin: 'turn-telemetry', key: 'live' } as const, null)
14const isHidden = atom({ plugin: 'turn-telemetry', key: 'isHidden' } as const, false)
15
16export const register: Register = on => {
17  // The main-loop turn in flight; a hot reload mid-turn loses it, which
18  // turn.complete tolerates.
19  let turn: { id: string; steps: StepStat[] } | null = null
20
21  on('session.start', async ($, e, next) => {
22    await $.command.register({
23      name: COMMAND,
24      description: 'Toggle the turn telemetry band (`/tps reset` clears its history)',
25    })
26
27    return next(e)
28  })
29
30  on('command.run', { command: COMMAND }, async ($, e) => {
31    if (e.args.trim() === 'reset') {
32      await update($, last, () => null)
33      await update($, history, () => [])
34
35      return { text: 'Turn telemetry cleared.' }
36    }
37    const hidden = await update($, isHidden, value => !value)
38
39    return { text: hidden ? 'Turn telemetry hidden.' : 'Turn telemetry shown.' }
40  })
41
42  on('turn.start', async ($, e, next) => {
43    turn = { id: e.turnId, steps: [] }
44    await update($, live, () => null)
45
46    return next(e)
47  })
48
49  on('turn.step', async function* ($, e, next) {
50    if (e.agentId) return yield* next(e)
51
52    if (turn?.id !== e.turnId) turn = { id: e.turnId, steps: [] }
53    const current = turn
54    const sentAt = await $.clock.now()
55    let firstAt: number | null = null
56    // Chunk-level timing stays on the synchronous clock so forwarding is
57    // never held up by a host round trip.
58    let perfFirst = 0
59    let lastPush = 0
60    let chars = 0
61    const gaps = stallCounter()
62
63    const pushLive = (perfNow: number) => {
64      lastPush = perfNow
65      const partial: StepStat = {
66        ttftMs: firstAt === null ? null : firstAt - sentAt,
67        genMs: perfNow - perfFirst,
68        outTokens: Math.round(chars / 4),
69        stalls: gaps.stalls,
70        longestStallMs: gaps.longestStallMs,
71      }
72
73      return update($, live, () => summarize([...current.steps, partial], null))
74    }
75
76    const stream = next(e)
77    for await (const chunk of stream) {
78      let isFirst = false
79      if (chunk.kind === 'text' || chunk.kind === 'thinking' || chunk.kind === 'input' || chunk.kind === 'tool') {
80        const perfNow = performance.now()
81        gaps.tick(perfNow)
82        if (chunk.kind === 'text' || chunk.kind === 'thinking') chars += chunk.text.length
83        if (chunk.kind === 'input') chars += chunk.json.length
84        if (firstAt === null) {
85          firstAt = await $.clock.now()
86          perfFirst = perfNow
87          isFirst = true
88        }
89      }
90      yield chunk
91      const perfNow = performance.now()
92      if (firstAt !== null && (isFirst || perfNow - lastPush >= LIVE_EVERY_MS)) await pushLive(perfNow)
93    }
94    const result = await stream.result
95    const endedAt = await $.clock.now()
96
97    current.steps.push({
98      ttftMs: firstAt === null ? null : firstAt - sentAt,
99      genMs: firstAt === null ? 0 : endedAt - firstAt,
100      outTokens: result.usage?.output_tokens ?? Math.round(chars / 4),
101      stalls: gaps.stalls,
102      longestStallMs: gaps.longestStallMs,
103    })
104    await update($, live, () => summarize(current.steps, null))
105
106    return result
107  })
108
109  on('turn.complete', async ($, e, next) => {
110    const result = await next(e)
111    if (e.agentId) return result
112
113    const steps = turn?.id === e.turnId ? turn.steps : []
114    turn = null
115    await update($, live, () => null)
116    if (e.reason !== 'answer') return result
117
118    const stat = summarize(steps, e.durationMs, e.usage)
119    await update($, last, () => stat)
120    const tps = stat.tps
121    if (tps !== null) await update($, history, list => [...list, Math.round(tps)].slice(-HISTORY))
122
123    return result
124  })
125
126  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
127    if (e.props.hasSurvey || e.props.view.agentId || (await read($, isHidden))) return next(e)
128
129    const [done, now, spark] = await Promise.all([read($, last), read($, live), read($, history)])
130    const isLive = e.props.isWorking && now !== null && now.ttftMs !== null
131    const stat = isLive ? now : done
132    if (stat === null) return next(e)
133
134    const { Box, Text } = $.ui.resolve(e)
135    const parts = fit(segments(stat, spark, isLive), e.props.bodyColumns)
136
137    return (
138      <Box flexDirection="row">
139        {parts.flatMap((part, i) => [
140          ...(i === 0 ? [] : [<Text key={`sep-${i}`} dimColor>{SEPARATOR}</Text>]),
141          <Text key={`seg-${i}`} color={part.color} dimColor={part.dimColor} wrap="truncate-end">
142            {part.text}
143          </Text>,
144        ])}
145      </Box>
146    )
147  })
148}
149

1import type { ModelUsage } from 'claude-code'
2
3import type { StepStat, TurnStat } from '../types'
4
5/** A gap this long between two streamed chunks counts as a stall. */
6export const STALL_MS = 5000
7/** Steps shorter than this stream in one burst; their tok/s is noise. */
8export const MIN_GEN_MS = 250
9export const MIN_GEN_TOKENS = 16
10/** How many turns the sparkline remembers. */
11export const HISTORY = 16
12
13export type Segment = {
14  text: string
15  color?: string
16  dimColor?: boolean
17  /** Lower survives longer when the band is too narrow. */
18  priority: number
19}
20
21/** Counts stalls from monotonic chunk stamps. */
22export function stallCounter() {
23  let lastAt: number | null = null
24  const counter = {
25    stalls: 0,
26    longestStallMs: 0,
27    tick(now: number) {
28      if (lastAt !== null) {
29        const gap = now - lastAt
30        if (gap >= STALL_MS) {
31          counter.stalls += 1
32          counter.longestStallMs = Math.max(counter.longestStallMs, gap)
33        }
34      }
35      lastAt = now
36    },
37  }
38
39  return counter
40}
41
42export function summarize(
43  steps: readonly StepStat[],
44  durationMs: number | null,
45  usage?: ModelUsage | null,
46): TurnStat {
47  let genMs = 0
48  let genTokens = 0
49  let stepTokens = 0
50  let stalls = 0
51  let longestStallMs = 0
52  for (const step of steps) {
53    stepTokens += step.outTokens
54    stalls += step.stalls
55    longestStallMs = Math.max(longestStallMs, step.longestStallMs)
56    if (step.genMs >= MIN_GEN_MS && step.outTokens >= MIN_GEN_TOKENS) {
57      genMs += step.genMs
58      genTokens += step.outTokens
59    }
60  }
61
62  return {
63    durationMs,
64    ttftMs: steps.find(step => step.ttftMs !== null)?.ttftMs ?? null,
65    tps: genMs > 0 ? genTokens / (genMs / 1000) : null,
66    inTokens: usage
67      ? usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
68      : 0,
69    cachedTokens: usage?.cache_read_input_tokens ?? 0,
70    outTokens: usage?.output_tokens ?? stepTokens,
71    requests: steps.length,
72    stalls,
73    longestStallMs,
74  }
75}
76
77const BARS = '▁▂▃▄▅▆▇█'
78
79export function sparkline(values: readonly number[]): string {
80  if (values.length === 0) return ''
81  const min = Math.min(...values)
82  const max = Math.max(...values)
83  if (max === min) return BARS.charAt(3).repeat(values.length)
84
85  return values.map(v => BARS.charAt(Math.round(((v - min) / (max - min)) * (BARS.length - 1)))).join('')
86}
87
88export function fmtMs(ms: number): string {
89  if (ms < 1000) return `${Math.round(ms)}ms`
90  if (ms < 60_000) return `${(ms / 1000).toFixed(1)}s`
91  const seconds = Math.round(ms / 1000)
92
93  return `${Math.floor(seconds / 60)}m${String(seconds % 60).padStart(2, '0')}s`
94}
95
96export function fmtTokens(n: number): string {
97  if (n < 1000) return String(n)
98  if (n < 10_000) return `${(n / 1000).toFixed(1)}k`
99  if (n < 1_000_000) return `${Math.round(n / 1000)}k`
100
101  return `${(n / 1_000_000).toFixed(1)}M`
102}
103
104function ttftColor(ms: number): string {
105  if (ms < 2000) return 'green'
106  if (ms < 5000) return 'yellow'
107
108  return 'red'
109}
110
111/** The band's pieces in display order; `fit` drops the least important. */
112export function segments(stat: TurnStat, history: readonly number[], isLive: boolean): Segment[] {
113  const out: Segment[] = []
114  if (isLive) {
115    out.push({ text: '● live', color: 'cyan', priority: 0 })
116  } else if (stat.durationMs !== null) {
117    out.push({ text: `turn ${fmtMs(stat.durationMs)}`, priority: 0 })
118  }
119  if (stat.ttftMs !== null) {
120    out.push({ text: `ttft ${fmtMs(stat.ttftMs)}`, color: ttftColor(stat.ttftMs), priority: 1 })
121  }
122  if (stat.tps !== null) {
123    const spark = history.length >= 2 ? ` ${sparkline(history)}` : ''
124    out.push({ text: `${isLive ? '~' : ''}${Math.round(stat.tps)} tok/s${spark}`, priority: 1 })
125  }
126  if (stat.outTokens > 0) {
127    out.push({ text: `out ${fmtTokens(stat.outTokens)}`, dimColor: true, priority: 3 })
128  }
129  if (stat.inTokens > 0) {
130    const cached = Math.round((stat.cachedTokens / stat.inTokens) * 100)
131    const note = cached >= 1 ? ` (${cached}% cached)` : ''
132    out.push({ text: `in ${fmtTokens(stat.inTokens)}${note}`, dimColor: true, priority: 4 })
133  }
134  if (stat.requests > 1) {
135    out.push({ text: `${stat.requests} req`, dimColor: true, priority: 5 })
136  }
137  if (stat.stalls > 0) {
138    const plural = stat.stalls === 1 ? '' : 's'
139    out.push({
140      text: `${stat.stalls} stall${plural} (max ${fmtMs(stat.longestStallMs)})`,
141      color: 'yellow',
142      priority: 2,
143    })
144  }
145
146  return out
147}
148
149export const SEPARATOR = ' · '
150
151/** Drops the least important segments until the row fits `width` cells. */
152export function fit(parts: readonly Segment[], width: number): Segment[] {
153  const kept = [...parts]
154  const length = () =>
155    kept.reduce((sum, part) => sum + part.text.length, 0) + SEPARATOR.length * (kept.length - 1)
156  while (kept.length > 1 && length() > width) {
157    let drop = 0
158    kept.forEach((part, i) => {
159      if (part.priority >= (kept[drop]?.priority ?? 0)) drop = i
160    })
161    kept.splice(drop, 1)
162  }
163
164  return kept
165}
166

1/** One model request inside a turn, as the band measures it. */
2export type StepStat = {
3  /** Request sent → first streamed content chunk; null when nothing streamed. */
4  ttftMs: number | null
5  /** First content chunk → end of the response. */
6  genMs: number
7  outTokens: number
8  /** Gaps of at least STALL_MS between content chunks. */
9  stalls: number
10  longestStallMs: number
11}
12
13/** A whole turn, summed over its requests. */
14export type TurnStat = {
15  durationMs: number | null
16  ttftMs: number | null
17  tps: number | null
18  /** Every input token the turn sent, cached ones included. */
19  inTokens: number
20  cachedTokens: number
21  outTokens: number
22  requests: number
23  stalls: number
24  longestStallMs: number
25}
26
27declare module 'claude-code' {
28  interface PluginState {
29    'turn-telemetry': {
30      last: TurnStat | null
31      /** tok/s of the last turns, oldest first, for the sparkline. */
32      history: number[]
33      /** The turn in flight, refreshed while it streams; null when idle. */
34      live: TurnStat | null
35      isHidden: boolean
36    }
37  }
38}
39