io-meter

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Reading, Timing } from '../types'
5
6const EMPTY: Timing = { ms: 0, samples: [] }
7
8const tools = atom({ plugin: 'io-meter', key: 'toolTimes' } as const, {})
9const api = atom({ plugin: 'io-meter', key: 'apiTimes' } as const, { main: EMPTY, sub: EMPTY })
10const readings = atom({ plugin: 'io-meter', key: 'readings' } as const, [])
11const usd = atom({ plugin: 'io-meter', key: 'usd' } as const, null)
12
13const SKIP = new Set(['Agent'])
14const TOP = 5
15const HISTORY = 12
16const SAMPLES = 200
17const BARS = '▁▂▃▄▅▆▇█'
18
19const FORECAST = [
20  { upTo: 25, icon: '☀', word: 'Clear', color: 'yellow' },
21  { upTo: 50, icon: '☁', word: 'Cloudy', color: 'cyan' },
22  { upTo: 75, icon: '☂', word: 'Showers', color: 'blue' },
23  { upTo: 90, icon: '☇', word: 'Storm', color: 'magenta' },
24  { upTo: Infinity, icon: '↯', word: 'Compact soon', color: 'red' },
25]
26
27function fmt(ms: number): string {
28  if (ms < 10_000) return (ms / 1000).toFixed(1) + 's'
29  if (ms < 60_000) return Math.round(ms / 1000) + 's'
30  const m = Math.floor(ms / 60_000)
31  return m + 'm' + String(Math.round((ms % 60_000) / 1000)).padStart(2, '0')
32}
33
34function short(n: number): string {
35  if (n >= 1_000_000) return (n / 1_000_000).toFixed(1) + 'M'
36  if (n >= 1_000) return (n / 1_000).toFixed(1) + 'k'
37  return String(n)
38}
39
40function add(t: Timing, ms: number): Timing {
41  return { ms: t.ms + ms, samples: [...t.samples, ms].slice(-SAMPLES) }
42}
43
44function p90(t: Timing): number {
45  const sorted = [...t.samples].sort((x, y) => x - y)
46  return sorted[Math.ceil(sorted.length * 0.9) - 1] ?? 0
47}
48
49function forecastFor(percent: number) {
50  return FORECAST.find(f => percent < f.upTo) ?? FORECAST[FORECAST.length - 1]
51}
52
53function chart(rs: Reading[]): string {
54  const top = Math.max(...rs.map(r => r.tokens), 1)
55  return rs.map(r => BARS[Math.min(BARS.length - 1, Math.floor((r.tokens / top) * (BARS.length - 1)))]).join('')
56}
57
58function trend(rs: Reading[]): string {
59  const [prev, last] = rs.slice(-2)
60  if (!prev || !last) return ''
61  const delta = last.tokens - prev.tokens
62  if (delta > 0) return '▲ +' + short(delta)
63  if (delta < 0) return '▼ ' + short(-delta)
64  return 'steady'
65}
66
67async function measure($: EngineInterface) {
68  const { context: c, cost } = await $.session.usage({ breakdown: 'summary' })
69  await update($, usd, () => cost?.usd ?? null)
70  const window = Math.min(c.window, c.breakdown?.rawMaxTokens ?? c.window)
71  if (!window) return
72  const tokens = c.tokens ?? 0
73  const reading: Reading = { tokens, window, percent: Math.round((tokens / window) * 100) }
74  await update($, readings, rs => {
75    const last = rs[rs.length - 1]
76    if (last && last.tokens === reading.tokens && last.window === reading.window) return rs
77    return [...rs.filter(r => r.tokens > 0), reading].slice(-HISTORY)
78  })
79}
80
81export const register: Register = on => {
82  on('session.start', async ($, e, next) => {
83    const result = await next(e)
84    await measure($)
85    return result
86  })
87
88  on('session.measure', async ($, e, next) => {
89    await measure($)
90    return next(e)
91  })
92
93  on('tool.call', async ($, e, next) => {
94    const t0 = await $.clock.now()
95    const result = await next(e)
96    if (!SKIP.has(e.tool)) {
97      const ms = (await $.clock.now()) - t0
98      await update($, tools, all => ({ ...all, [e.tool]: add(all[e.tool] ?? EMPTY, ms) }))
99    }
100    return result
101  })
102
103  on('turn.step', async function* ($, e, next) {
104    const t0 = await $.clock.now()
105    const result = yield* next(e)
106    const ms = (await $.clock.now()) - t0
107    await update($, api, a => (e.agentId ? { ...a, sub: add(a.sub, ms) } : { ...a, main: add(a.main, ms) }))
108    return result
109  })
110
111  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
112    if (e.props.hasSurvey) return next(e)
113
114    const rs = await read($, readings)
115    const cost = await read($, usd)
116    const a = await read($, api)
117    const t = await read($, tools)
118    const wide = (e.viewport?.columns ?? 80) >= 60
119    const { Box, Text } = $.ui.resolve(e)
120
121    const top = Object.entries(t)
122      .sort((x, y) => y[1].ms - x[1].ms)
123      .slice(0, TOP)
124    const toolMs = Object.values(t).reduce((s, v) => s + v.ms, 0)
125
126    const now = rs[rs.length - 1]
127    const f = now && forecastFor(now.percent)
128    const tr = trend(rs)
129
130    const ctxRow = now && f && (
131      <Box flexDirection="row">
132        <Text color={f.color} bold>
133          {f.icon} {f.word}
134        </Text>
135        <Text>  {now.percent}% of context</Text>
136        <Text dimColor>  {short(now.tokens)} / {short(now.window)}</Text>
137        {wide && <Text dimColor>   last turns </Text>}
138        {wide && <Text color={f.color}>{chart(rs)}</Text>}
139        {wide && tr !== '' && <Text dimColor>  {tr}</Text>}
140        {cost !== null && <Text dimColor>  · ${cost.toFixed(2)}</Text>}
141      </Box>
142    )
143
144    const ioRow = (
145      <Box flexDirection="row">
146        <Text dimColor>total (p90) │ api </Text>
147        <Text>{fmt(a.main.ms)}</Text>
148        <Text dimColor> ({fmt(p90(a.main))})</Text>
149        {a.sub.samples.length > 0 && <Text dimColor> + sub </Text>}
150        {a.sub.samples.length > 0 && <Text>{fmt(a.sub.ms)}</Text>}
151        {a.sub.samples.length > 0 && <Text dimColor> ({fmt(p90(a.sub))})</Text>}
152        <Text dimColor> │ tools </Text>
153        <Text>{fmt(toolMs)}</Text>
154        {wide &&
155          top.map(([name, s]) => (
156            <Text key={name} dimColor>
157              {' │ '}{name} {fmt(s.ms)} ({fmt(p90(s))})
158            </Text>
159          ))}
160      </Box>
161    )
162
163    return (
164      <Box flexDirection="column" paddingX={1}>
165        {ctxRow}
166        {ioRow}
167      </Box>
168    )
169  })
170}
171

1export type Timing = { ms: number; samples: number[] }
2export type ApiTiming = { main: Timing; sub: Timing }
3export type Reading = { tokens: number; window: number; percent: number }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'io-meter': {
8      toolTimes: Record<string, Timing>
9      apiTimes: ApiTiming
10      readings: Reading[]
11      usd: number | null
12    }
13  }
14}
15