beat

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { BeatTokens, BeatTurn } from '../types'
5
6const turns = atom({ plugin: 'beat', key: 'turns' } as const, [])
7
8const MAX_TURNS = 200
9const MAX_SESSIONS = 50
10const STORE_VERSION = 1
11const SESSIONS_KEY = 'sessions'
12const TURNS_PREFIX = 'turns:'
13const BAR = 12
14const MARGIN = 2
15const GAP = 1
16const CAP_LEFT = ''
17const CAP_RIGHT = ''
18const PILL_BG = '#313244'
19const PILL_FG = '#cdd6f4'
20const STAT_COLOR = '#89b4fa'
21
22type Part = 'think' | 'tools' | 'write' | 'other'
23const PARTS: readonly { part: Part; icon: string; color: string }[] = [
24  { part: 'think', icon: '', color: '#89b4fa' },
25  { part: 'tools', icon: '', color: '#f9e2af' },
26  { part: 'write', icon: '', color: '#a6e3a1' },
27  { part: 'other', icon: '\uf254', color: '#cba6f7' },
28]
29
30export const formatDuration = (ms: number) => {
31  const s = Math.max(0, Math.round(ms / 1000))
32  if (s < 60) return `${s}s`
33  const m = Math.floor(s / 60)
34  if (m < 60) return `${m}m ${s % 60}s`
35  return `${Math.floor(m / 60)}h ${m % 60}m`
36}
37
38const widthOf = (s: string) => [...s].length
39
40export const formatTokens = (n: number) => {
41  if (n < 1_000) return `${n}`
42  if (n < 999_950) return `${Number((n / 1_000).toFixed(1))}k`
43  return `${Number((n / 1_000_000).toFixed(1))}M`
44}
45
46export const tokensOf = (usage: ModelUsage): BeatTokens => ({
47  input: usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens,
48  output: usage.output_tokens,
49  cacheRead: usage.cache_read_input_tokens,
50})
51
52export type Stat = { key: string; parts: readonly (readonly [icon: string, value: string])[] }
53
54export const statsOf = (turn: BeatTurn): Stat[] => {
55  const stats: Stat[] = []
56  if (turn.tokens !== undefined) {
57    stats.push({ key: 'io', parts: [['↑', formatTokens(turn.tokens.input)], ['↓', formatTokens(turn.tokens.output)]] })
58    if (turn.tokens.input > 0) stats.push({ key: 'cache', parts: [['\uf1c0 ', `${Math.round((turn.tokens.cacheRead / turn.tokens.input) * 100)}%`]] })
59  }
60  if (turn.costUsd !== undefined) stats.push({ key: 'cost', parts: [['$', turn.costUsd.toFixed(2)]] })
61  if (turn.contextPercent !== undefined) stats.push({ key: 'context', parts: [['\uf2db ', `${Math.round(turn.contextPercent)}%`]] })
62  return stats
63}
64
65export const pillWidthOf = (stats: readonly Stat[]) =>
66  stats.length === 0
67    ? 0
68    : widthOf(CAP_LEFT + CAP_RIGHT) +
69      stats.reduce((a, s) => a + s.parts.reduce((b, [icon, value]) => b + widthOf(icon + value), 0) + s.parts.length - 1, 0) +
70      2 * (stats.length - 1)
71
72export const fitStats = (stats: readonly Stat[], room: number) => {
73  let kept = [...stats]
74  while (kept.length > 0 && pillWidthOf(kept) > room) kept = kept.slice(0, -1)
75  return kept
76}
77
78const settled = <T,>(p: Promise<T>) => p.then(v => v, () => undefined)
79
80export const storedTurnsOf = (raw: unknown): BeatTurn[] => {
81  if (typeof raw !== 'object' || raw === null) return []
82  const list: unknown = Reflect.get(raw, 'turns')
83  return Reflect.get(raw, 'version') === STORE_VERSION && Array.isArray(list) ? list : []
84}
85
86export const keptSessionsOf = (raw: unknown, id: string) =>
87  [...(Array.isArray(raw) ? raw.filter((s): s is string => typeof s === 'string' && s !== id) : []), id].slice(-MAX_SESSIONS)
88
89const load = async ($: EngineInterface, cache: Map<string, BeatTurn[]>, id: string) => {
90  const cached = cache.get(id)
91  if (cached !== undefined) return cached
92  const list = storedTurnsOf(await settled($.store.get(TURNS_PREFIX + id)))
93  cache.set(id, list)
94  return list
95}
96
97const persist = async ($: EngineInterface, id: string, list: readonly BeatTurn[]) => {
98  await $.store.set(TURNS_PREFIX + id, { version: STORE_VERSION, turns: list })
99  const kept = keptSessionsOf(await $.store.get(SESSIONS_KEY), id)
100  await $.store.set(SESSIONS_KEY, kept)
101  for (const key of await $.store.keys()) {
102    if (key.startsWith(TURNS_PREFIX) && !kept.includes(key.slice(TURNS_PREFIX.length))) await $.store.delete(key)
103  }
104}
105
106export const unionMs = (spans: readonly (readonly [number, number])[]) => {
107  const sorted = [...spans].sort((a, b) => a[0] - b[0])
108  let total = 0
109  let end = -Infinity
110  for (const [s, e] of sorted) {
111    if (e <= end) continue
112    total += e - Math.max(s, end)
113    end = e
114  }
115  return total
116}
117
118export const exclusiveMs = (spans: readonly (readonly [number, number])[], cut: readonly (readonly [number, number])[]) =>
119  unionMs([...spans, ...cut]) - unionMs(cut)
120
121export const partsOf = (
122  durationMs: number,
123  wallMs: number,
124  think: readonly (readonly [number, number])[],
125  write: readonly (readonly [number, number])[],
126  tools: readonly (readonly [number, number])[],
127) => {
128  const paused = Math.max(0, wallMs - durationMs)
129  const turn = {
130    think: exclusiveMs(think, tools),
131    tools: Math.max(0, unionMs(tools) - paused),
132    write: exclusiveMs(write, tools),
133  }
134  return { durationMs, ...turn, other: Math.max(0, durationMs - turn.think - turn.tools - turn.write) }
135}
136
137export const secondsOf = (values: readonly number[]) => {
138  const exact = values.map(v => v / 1000)
139  const secs = exact.map(Math.floor)
140  let left = Math.round(exact.reduce((a, b) => a + b, 0)) - secs.reduce((a, b) => a + b, 0)
141  for (const [, i] of exact.map((v, i) => [v - Math.floor(v), i] as const).sort((a, b) => b[0] - a[0])) {
142    if (left <= 0) break
143    secs[i] = (secs[i] ?? 0) + 1
144    left -= 1
145  }
146  return secs
147}
148
149export const cellsOf = (values: readonly number[], width: number) => {
150  const sum = values.reduce((a, b) => a + b, 0)
151  if (sum <= 0) return values.map(() => 0)
152  const cells = values.map(v => (v > 0 ? Math.max(1, Math.round((v / sum) * width)) : 0))
153  let over = cells.reduce((a, b) => a + b, 0) - width
154  while (over !== 0) {
155    const i = cells.indexOf(Math.max(...cells))
156    cells[i] = (cells[i] ?? 0) - Math.sign(over)
157    over -= Math.sign(over)
158  }
159  return cells
160}
161
162export const register: Register = on => {
163  let turnStart = 0
164  let costStart: number | undefined
165  let think: [number, number][] = []
166  let write: [number, number][] = []
167  let tools: [number, number][] = []
168  const stored = new Map<string, BeatTurn[]>()
169
170  on('session.end', async ($, e, next) => {
171    if (e.reason === 'clear' || e.reason === 'resume') await update($, turns, () => [])
172    return next(e)
173  })
174
175  on('turn.start', async ($, e, next) => {
176    turnStart = await $.clock.now()
177    costStart = (await settled($.session.usage()))?.cost?.usd
178    think = []
179    write = []
180    tools = []
181    return next(e)
182  })
183
184  on('turn.step', async function* ($, e, next) {
185    if (e.agentId !== undefined) return yield* next(e)
186
187    let phase: 'think' | 'write' = 'think'
188    let mark = await $.clock.now()
189    const close = (now: number) => {
190      ;(phase === 'think' ? think : write).push([mark, now])
191      mark = now
192    }
193
194    for await (const chunk of next(e)) {
195      const want: 'think' | 'write' = chunk.kind === 'thinking' ? 'think' : chunk.kind === 'engine' || chunk.kind === 'stop' ? phase : 'write'
196      if (want !== phase) {
197        close(await $.clock.now())
198        phase = want
199      }
200      yield chunk
201    }
202    close(await $.clock.now())
203  })
204
205  on('tool.call', async ($, e, next) => {
206    if (e.agentId !== undefined) return next(e)
207    const start = await $.clock.now()
208    try {
209      return await next(e)
210    } finally {
211      tools.push([start, await $.clock.now()])
212    }
213  })
214
215  on('turn.complete', async ($, e, next) => {
216    const done = await next(e)
217    if (e.agentId === undefined) {
218      const wallMs = (await $.clock.now()) - turnStart
219      const session = await settled($.session.usage())
220      const costEnd = session?.cost?.usd
221      const contextPercent = session?.context.percent
222      const turn: BeatTurn = {
223        ...partsOf(e.durationMs, wallMs, think, write, tools),
224        ...(e.usage === undefined ? {} : { tokens: tokensOf(e.usage) }),
225        ...(costStart === undefined || costEnd === undefined ? {} : { costUsd: Math.max(0, costEnd - costStart) }),
226        ...(contextPercent === undefined ? {} : { contextPercent }),
227      }
228      await update($, turns, list => [...list, turn].slice(-MAX_TURNS))
229      const id = await settled($.session.id())
230      if (id !== undefined) {
231        const list = [...(await load($, stored, id)), turn].slice(-MAX_TURNS)
232        stored.set(id, list)
233        await settled(persist($, id, list))
234      }
235    }
236    return done
237  })
238
239  on('ui.render', { component: 'TurnDuration' }, async ($, e) => {
240    const { Box, Text } = $.ui.resolve(e)
241    const match = (list: readonly BeatTurn[]) => [...list].reverse().find(t => t.durationMs === e.props.durationMs)
242    let hit = match(await read($, turns))
243    if (hit === undefined) {
244      const id = await settled($.session.id())
245      if (id !== undefined) hit = match(await load($, stored, id))
246    }
247    const values = secondsOf(hit === undefined ? [0, 0, 0, 0] : [hit.think, hit.tools, hit.write, hit.other])
248    const cells = cellsOf(values, BAR)
249    const shown = PARTS.map((p, i) => ({ ...p, seconds: values[i] ?? 0, cells: cells[i] ?? 0 })).filter(p => p.cells > 0)
250    const time = formatDuration(e.props.durationMs)
251    const legends = shown.map(p => ({ ...p, head: `  ${p.icon} `, tail: formatDuration(p.seconds * 1000) }))
252
253    const timeWidth =
254      widthOf(CAP_LEFT + time + CAP_RIGHT) +
255      (shown.length > 0 ? 1 + shown.reduce((a, p) => a + p.cells, 0) + legends.reduce((a, p) => a + widthOf(p.head + p.tail), 0) : 0)
256    const all = hit === undefined ? [] : statsOf(hit)
257    const columns = e.viewport?.columns
258    const stats = columns === undefined ? all : fitStats(all, columns - 2 * MARGIN - timeWidth - GAP)
259
260    return (
261      <Box flexDirection="row" marginTop={1} marginLeft={MARGIN} marginRight={MARGIN}>
262        <Box flexDirection="row" flexShrink={0}>
263          <Text color={PILL_BG}>{CAP_LEFT}</Text>
264          <Text backgroundColor={PILL_BG} color={PILL_FG} bold>
265            {time}
266          </Text>
267          {shown.length > 0 && <Text backgroundColor={PILL_BG}> </Text>}
268          {shown.map(p => (
269            <Text key={`bar-${p.part}`} backgroundColor={PILL_BG} color={p.color}>
270              {'━'.repeat(p.cells)}
271            </Text>
272          ))}
273          {legends.map(p => (
274            <Text key={`legend-${p.part}`} backgroundColor={PILL_BG}>
275              <Text color={p.color}>{p.head}</Text>
276              <Text color={PILL_FG}>{p.tail}</Text>
277            </Text>
278          ))}
279          <Text color={PILL_BG}>{CAP_RIGHT}</Text>
280        </Box>
281        {stats.length > 0 && <Box flexGrow={1} minWidth={GAP} />}
282        {stats.length > 0 && (
283          <Box flexDirection="row" flexShrink={0}>
284            <Text color={PILL_BG}>{CAP_LEFT}</Text>
285            {stats.map((stat, i) => (
286              <Text key={stat.key} backgroundColor={PILL_BG}>
287                {i > 0 && '  '}
288                {stat.parts.map(([icon, value], j) => (
289                  <Text key={`${stat.key}-${j}`}>
290                    {j > 0 && ' '}
291                    <Text color={STAT_COLOR}>{icon}</Text>
292                    <Text color={PILL_FG}>{value}</Text>
293                  </Text>
294                ))}
295              </Text>
296            ))}
297            <Text color={PILL_BG}>{CAP_RIGHT}</Text>
298          </Box>
299        )}
300      </Box>
301    )
302  })
303}
304

1export type BeatTokens = { input: number; output: number; cacheRead: number }
2
3export type BeatTurn = { durationMs: number; think: number; tools: number; write: number; other: number; tokens?: BeatTokens; costUsd?: number; contextPercent?: number }
4
5declare module 'claude-code' {
6  interface PluginState {
7    beat: { turns: BeatTurn[] }
8  }
9}
10