turn-stats

1// Session statistics band: the desktop Code tab does not show a Stop hook's
2// systemMessage, so the session's totals are drawn above the prompt instead.
3import { atom, read, update } from 'claude-code'
4import type { EngineInterface, ModelUsage, Register, RenderElement, RenderNode, SessionRateLimit } from 'claude-code'
5
6import type { SessionTotals, TokenUsage } from '../types'
7
8const totals = atom({ plugin: 'turn-stats', key: 'totals' } as const, null)
9const runningSince = atom({ plugin: 'turn-stats', key: 'runningSince' } as const, null)
10
11// The band is shared, and the engine chains mods in no fixed order: each mod
12// files its section under a rank (time 10, PR 20, progress 30), the sections
13// beneath come up as one column, and the column is kept sorted by rank, the
14// rest (core's own drawing) last.
15const COLUMN = 'band-sections'
16const SECTION = 'band-section:'
17const RANK = 10
18
19function rankOf(node: RenderNode): number {
20  const key = typeof node === 'object' && node.type === 'Box' ? node.props?.key : undefined
21  const rank = typeof key === 'string' && key.startsWith(SECTION) ? Number(key.slice(SECTION.length)) : NaN
22  return Number.isFinite(rank) ? rank : Number.MAX_SAFE_INTEGER
23}
24
25function stacked(below: RenderElement, section: RenderElement): RenderElement {
26  const rest = below.type === 'Box' && below.props?.key === COLUMN ? (below.children ?? []) : [below]
27  const children = [section, ...rest].sort((a, b) => rankOf(a) - rankOf(b))
28  return { type: 'Box', props: { key: COLUMN, flexDirection: 'column' }, children }
29}
30
31const STORE_PREFIX = 'session:'
32const MAX_STORED_SESSIONS = 200
33const SEPARATOR = ' · '
34
35// One styled run of the stats row; the row is a line of runs.
36type Part = { text: string; color?: string; isDim?: boolean; isBold?: boolean }
37
38// A dim label and its value in the figure's color: "вх 213.8k".
39const labelled = (label: string, value: string, color: string): Part[] => [
40  { text: `${label} `, isDim: true },
41  { text: value, color },
42]
43
44const QUOTA_LABELS: [kind: string, label: string][] = [
45  ['five_hour', '5ч'],
46  ['seven_day', 'W'],
47]
48
49function trim(n: number, digits: number): string {
50  return String(Number(n.toFixed(digits)))
51}
52
53// Thresholds sit where rounding would otherwise print "1000k" or "1000".
54function formatTokens(n: number): string {
55  if (n >= 999_950) return trim(n / 1e6, 2) + 'M'
56  if (n >= 999.5) return trim(n / 1e3, 1) + 'k'
57  return String(Math.round(n))
58}
59
60function formatDuration(ms: number): string {
61  const total = Math.floor(ms / 1000)
62  const hours = Math.floor(total / 3600)
63  const minutes = Math.floor((total % 3600) / 60)
64  const seconds = total % 60
65  const pad = (v: number) => String(v).padStart(2, '0')
66  if (hours > 0) return `${hours}ч${pad(minutes)}м`
67  if (minutes > 0) return `${minutes}м${pad(seconds)}с`
68  return `${seconds}с`
69}
70
71const emptyTokens = (): TokenUsage => ({ input: 0, cacheWrite: 0, cacheRead: 0, output: 0 })
72
73const freshTotals = (sessionId: string): SessionTotals => ({
74  sessionId,
75  turns: 0,
76  workMs: 0,
77  lastTurnMs: 0,
78  main: emptyTokens(),
79  agents: emptyTokens(),
80  quotaSpent: {},
81})
82
83const addTokens = (sum: TokenUsage, u: ModelUsage): TokenUsage => ({
84  input: sum.input + u.input_tokens,
85  cacheWrite: sum.cacheWrite + u.cache_creation_input_tokens,
86  cacheRead: sum.cacheRead + u.cache_read_input_tokens,
87  output: sum.output + u.output_tokens,
88})
89
90const tokenTotal = (u: TokenUsage) => u.input + u.cacheWrite + u.cacheRead + u.output
91
92// Percentage points each window grew by over one turn. A window that reset
93// mid-turn counts its whole new reading.
94function quotaSpend(before: SessionRateLimit[], after: SessionRateLimit[]): Record<string, number> {
95  const spent: Record<string, number> = {}
96  for (const now of after) {
97    const was = before.find(l => l.kind === now.kind)
98    if (!was) continue
99    const delta = was.resetsAt === now.resetsAt ? now.percentUsed - was.percentUsed : now.percentUsed
100    if (delta > 0) spent[now.kind] = delta
101  }
102  return spent
103}
104
105function addSpent(sum: Record<string, number>, spent: Record<string, number>): Record<string, number> {
106  const result = { ...sum }
107  for (const [kind, delta] of Object.entries(spent)) result[kind] = (result[kind] ?? 0) + delta
108  return result
109}
110
111const storeKey = (sessionId: string) => STORE_PREFIX + sessionId
112
113function isTotals(value: unknown, sessionId: string): value is SessionTotals {
114  return typeof value === 'object' && value !== null && (value as SessionTotals).sessionId === sessionId
115}
116
117async function loadTotals($: EngineInterface, sessionId: string): Promise<SessionTotals | null> {
118  const stored = await $.store.get(storeKey(sessionId))
119  return isTotals(stored, sessionId) ? stored : null
120}
121
122// Applies `change` to this session's totals and persists the result, so a
123// resumed session or a relaunched app keeps counting from where it stopped.
124async function record($: EngineInterface, change: (t: SessionTotals) => SessionTotals): Promise<void> {
125  const sessionId = await $.session.id()
126  const current = await read($, totals)
127  const seed = current?.sessionId === sessionId ? null : await loadTotals($, sessionId)
128  const next = await update($, totals, t =>
129    change(t?.sessionId === sessionId ? t : (seed ?? freshTotals(sessionId))),
130  )
131  await $.store.set(storeKey(sessionId), next)
132}
133
134export const register: Register = on => {
135  // The main loop runs one turn at a time; a hot reload mid-turn only drops
136  // that turn's quota delta and live ticking, never its time or tokens.
137  let turnBefore: { turnId: string; limits: SessionRateLimit[] } | null = null
138  let ticker: { cancel: () => void } | null = null
139  const stopTicker = () => {
140    ticker?.cancel()
141    ticker = null
142  }
143
144  on('session.start', async ($, e, next) => {
145    const result = await next(e)
146    const sessionId = await $.session.id()
147    if ((await read($, totals))?.sessionId !== sessionId) {
148      const stored = await loadTotals($, sessionId)
149      await update($, totals, () => stored)
150    }
151    await update($, runningSince, () => null)
152
153    const keys = (await $.store.keys()).filter(k => k.startsWith(STORE_PREFIX))
154    for (const key of keys.slice(0, Math.max(0, keys.length - MAX_STORED_SESSIONS))) {
155      await $.store.delete(key)
156    }
157    return result
158  })
159
160  // /clear goes on under a new session id with no session.start: start over.
161  on('session.end', async ($, e, next) => {
162    if (e.reason === 'clear') {
163      stopTicker()
164      turnBefore = null
165      await update($, totals, () => null)
166      await update($, runningSince, () => null)
167    }
168    return next(e)
169  })
170
171  // turn.start is raised for the main loop only; subagent runs carry agentId on turn.complete.
172  on('turn.start', async ($, e, next) => {
173    turnBefore = { turnId: e.turnId, limits: (await $.session.usage()).rateLimits }
174    const startedAt = await $.clock.now()
175    await update($, runningSince, () => startedAt)
176    stopTicker()
177    ticker = $.clock.every(1000, () => $.ui.invalidate('ui.render'))
178    return next(e)
179  })
180
181  on('turn.complete', async ($, e, next) => {
182    const result = await next(e)
183    const usage = e.usage
184    if (e.agentId) {
185      if (usage) await record($, t => ({ ...t, agents: addTokens(t.agents, usage) }))
186      return result
187    }
188
189    stopTicker()
190    const before = turnBefore?.turnId === e.turnId ? turnBefore.limits : null
191    turnBefore = null
192    const spent = before ? quotaSpend(before, (await $.session.usage()).rateLimits) : {}
193    await record($, t => ({
194      ...t,
195      turns: t.turns + 1,
196      workMs: t.workMs + e.durationMs,
197      lastTurnMs: e.durationMs,
198      main: usage ? addTokens(t.main, usage) : t.main,
199      quotaSpent: addSpent(t.quotaSpent, spent),
200    }))
201    await update($, runningSince, () => null)
202    return result
203  })
204
205  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
206    const session = await read($, totals)
207    const since = await read($, runningSince)
208    const liveMs = e.props.isWorking && since !== null ? Math.max(0, (await $.clock.now()) - since) : null
209    if (e.props.hasSurvey || (session === null && liveMs === null)) {
210      return next(e)
211    }
212
213    const below = await next(e)
214    const { Box, Text } = $.ui.resolve(e)
215    const { rateLimits } = await $.session.usage()
216    const segments: Part[][] = []
217
218    const workMs = (session?.workMs ?? 0) + (liveMs ?? 0)
219    const turnMs = liveMs ?? (session !== null && session.turns > 0 ? session.lastTurnMs : null)
220    segments.push([
221      { text: '⏱ ', isDim: true },
222      { text: formatDuration(workMs), isBold: true },
223      ...(turnMs !== null ? [{ text: ` (ход ${formatDuration(turnMs)})`, isDim: true }] : []),
224    ])
225
226    if (session !== null) {
227      const { main, agents } = session
228      if (tokenTotal(main) + tokenTotal(agents) > 0) {
229        segments.push(labelled('вх', formatTokens(main.input + main.cacheWrite + agents.input + agents.cacheWrite), 'cyan'))
230        segments.push(labelled('вых', formatTokens(main.output + agents.output), 'green'))
231        segments.push(labelled('кэш', formatTokens(main.cacheRead + agents.cacheRead), 'blue'))
232        if (tokenTotal(agents) > 0) segments.push(labelled('агенты', formatTokens(tokenTotal(agents)), 'magenta'))
233      }
234    }
235
236    const quotas: Part[] = []
237    for (const [kind, label] of QUOTA_LABELS) {
238      const limit = rateLimits.find(l => l.kind === kind)
239      if (!limit) continue
240      const spent = session?.quotaSpent[kind] ?? 0
241      if (quotas.length > 0) quotas.push({ text: ' ' })
242      quotas.push(...labelled(label, `${trim(limit.percentUsed, 1)}%`, 'yellow'))
243      if (spent >= 0.05) quotas.push({ text: ` +${trim(spent, 1)}`, isDim: true })
244    }
245    if (quotas.length > 0) segments.push(quotas)
246
247    const parts = segments.flatMap((segment, i) => (i === 0 ? segment : [{ text: SEPARATOR, isDim: true }, ...segment]))
248
249    return stacked(
250      below,
251      <Box key={`${SECTION}${RANK}`} flexDirection="column">
252        <Text wrap="truncate-end">
253          {parts.map(part => (
254            <Text color={part.color} dimColor={part.isDim} bold={part.isBold}>
255              {part.text}
256            </Text>
257          ))}
258        </Text>
259      </Box>,
260    )
261  })
262}
263

1export type TokenUsage = {
2  input: number
3  cacheWrite: number
4  cacheRead: number
5  output: number
6}
7
8// Everything counted since the session began; kept in $.store under the
9// session id, so a resume or an app relaunch continues the same totals.
10export type SessionTotals = {
11  sessionId: string
12  turns: number
13  workMs: number
14  lastTurnMs: number
15  main: TokenUsage
16  agents: TokenUsage
17  // Percentage points of each rate-limit window spent during this session's turns.
18  quotaSpent: Record<string, number>
19}
20
21declare module 'claude-code' {
22  interface PluginState {
23    'turn-stats': { totals: SessionTotals | null; runningSince: number | null }
24  }
25}
26