session-hud

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register, SessionContextUsage } from 'claude-code'
3
4import type { HudContext, HudEdits, HudSpeedSample, HudTokens } from '../types'
5import {
6  addSpeedSample,
7  countLines,
8  countPatch,
9  enterLoop,
10  HUD_GAPS,
11  leaveLoop,
12  NO_ACTIVITY,
13  recentTps,
14  renderCells,
15} from './format'
16
17// The statusLine command (~/.claude/session-hud/statusline.js) prints
18// the cells in <HUD_DIR>/<session id>.json under the prompt, the model added.
19const HUD_DIR = '.claude/session-hud'
20const NO_EDITS: HudEdits = { added: 0, removed: 0, files: [] }
21const NO_TOKENS: HudTokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, requests: 0 }
22const NO_SPEED: HudSpeedSample[] = []
23const MAX_FILES = 500
24// Shorter streams than this give no meaningful tokens-per-second figure.
25const MIN_TPS_MS = 300
26
27const compactAt = atom({ plugin: 'session-hud', key: 'compactAt' } as const, null)
28const edits = atom({ plugin: 'session-hud', key: 'edits' } as const, NO_EDITS)
29const tokens = atom({ plugin: 'session-hud', key: 'tokens' } as const, NO_TOKENS)
30const speed = atom({ plugin: 'session-hud', key: 'speed' } as const, NO_SPEED)
31const active = atom({ plugin: 'session-hud', key: 'active' } as const, NO_ACTIVITY)
32
33const toHudContext = (c: SessionContextUsage): HudContext => ({
34  tokens: c.tokens ?? null,
35  window: c.window,
36})
37
38const record = ($: EngineInterface, filePath: string, delta: { added: number; removed: number }) =>
39  update($, edits, cur => ({
40    added: cur.added + delta.added,
41    removed: cur.removed + delta.removed,
42    files: cur.files.includes(filePath) ? cur.files : [...cur.files, filePath].slice(-MAX_FILES),
43  }))
44
45const addUsage = ($: EngineInterface, usage: ModelUsage) =>
46  update($, tokens, cur => ({
47    input: cur.input + usage.input_tokens,
48    output: cur.output + usage.output_tokens,
49    cacheRead: cur.cacheRead + usage.cache_read_input_tokens,
50    cacheWrite: cur.cacheWrite + usage.cache_creation_input_tokens,
51    requests: cur.requests + 1,
52  }))
53
54const refreshCompactAt = async ($: EngineInterface) => {
55  try {
56    const { breakdown } = (await $.session.usage({ breakdown: 'summary' })).context
57    if (breakdown === undefined) return
58    const at = breakdown.isAutoCompactEnabled ? (breakdown.autoCompactThreshold ?? null) : null
59    await update($, compactAt, () => at)
60  } catch {
61    // Keep the last threshold; the line draws without it.
62  }
63}
64
65let published = ''
66let windowSeen: number | null = null
67let compactCheckedAt = 0
68let estimate: { at: number; tokens: number | null } = { at: 0, tokens: null }
69const ESTIMATE_EVERY_MS = 5000
70// A compaction can read the threshold while the engine has none to give (it
71// came back off once and stayed so), so it is read again now and then too.
72const COMPACT_EVERY_MS = 10_000
73
74// Before the first response there is no reading of the window: the engine's
75// local estimate (the one /context draws, no API call) stands in, at most
76// every few seconds.
77const estimateTokens = async ($: EngineInterface, now: number) => {
78  if (now - estimate.at < ESTIMATE_EVERY_MS) return estimate.tokens
79  try {
80    const { breakdown } = (await $.session.usage({ breakdown: 'summary' })).context
81    estimate = { at: now, tokens: breakdown?.totalTokens ?? null }
82  } catch {
83    estimate = { at: now, tokens: null }
84  }
85  return estimate.tokens
86}
87
88// Renders the cells and writes them for the statusLine command when they
89// changed. The context and cost are read from the engine on every pass (what
90// the status line itself reads), so they are there as soon as the engine has
91// them; the compaction threshold is looked up again when the window changes
92// and every few seconds.
93const publish = async ($: EngineInterface) => {
94  const home = await $.env.get('HOME')
95  if (home === undefined) return
96  const usage = await $.session.usage()
97  const now = await $.clock.now()
98  if (usage.context.window !== windowSeen || now - compactCheckedAt >= COMPACT_EVERY_MS) {
99    windowSeen = usage.context.window
100    compactCheckedAt = now
101    await refreshCompactAt($)
102  }
103  let context = toHudContext(usage.context)
104  if (context.tokens === null) {
105    const tokens = await estimateTokens($, now)
106    if (tokens !== null) context = { ...context, tokens, isEstimate: true }
107  }
108  const cells = renderCells({
109    context,
110    compactAt: await read($, compactAt),
111    cost: usage.cost?.usd ?? null,
112    edits: await read($, edits),
113    tokens: await read($, tokens),
114    tps: recentTps(await read($, speed)),
115    active: await read($, active),
116    now,
117  })
118  const path = `${home}/${HUD_DIR}/${await $.session.id()}.json`
119  const text = JSON.stringify({ rows: cells, gaps: HUD_GAPS })
120  if (`${path}\n${text}` === published) return
121  published = `${path}\n${text}`
122  await $.fs.write(path, text)
123}
124
125export const register: Register = on => {
126  on('session.start', async ($, e, next) => {
127    const ran = await next(e)
128    // Right away, so a new session never shows the stdin-only fallback for long.
129    $.clock.after(1, () => void publish($))
130    $.clock.every(1000, () => void publish($))
131
132    return ran
133  })
134
135  on('session.end', async ($, e, next) => {
136    const ran = await next(e)
137    if (e.reason === 'clear') {
138      await update($, edits, () => NO_EDITS)
139      await update($, tokens, () => NO_TOKENS)
140      await update($, speed, () => NO_SPEED)
141      await update($, active, () => NO_ACTIVITY)
142    }
143
144    return ran
145  })
146
147  // The active clock: the main loop works from turn.start to its
148  // turn.complete; a subagent, which raises no turn.start, from its first
149  // model request to its turn.complete.
150  on('turn.start', async ($, e, next) => {
151    const t = await $.clock.now()
152    await update($, active, cur => enterLoop(cur, 'main', t))
153
154    return next(e)
155  })
156
157  on('turn.complete', async ($, e, next) => {
158    const t = await $.clock.now()
159    await update($, active, cur => leaveLoop(cur, e.agentId ?? 'main', t))
160
161    return next(e)
162  })
163
164  // Every model request, subagents' included, adds to the token totals; the
165  // TPS figure is over the main loop's last few requests, each timed from its
166  // first streamed chunk (the response envelope) to its stop.
167  on('turn.step', async function* ($, e, next) {
168    if (e.agentId !== undefined) {
169      const loop = e.agentId
170      const t = await $.clock.now()
171      await update($, active, cur => enterLoop(cur, loop, t))
172    }
173    const stream = next(e)
174    let firstAt: number | null = null
175    let stopAt: number | null = null
176    for await (const chunk of stream) {
177      if (firstAt === null) firstAt = performance.now()
178      if (chunk.kind === 'stop') stopAt = performance.now()
179      yield chunk
180    }
181    const result = await stream.result
182    if (result.usage !== null) {
183      await addUsage($, result.usage)
184      const ms = firstAt !== null && stopAt !== null ? stopAt - firstAt : 0
185      if (e.agentId === undefined && ms >= MIN_TPS_MS && result.usage.output_tokens > 0) {
186        const sample = { tokens: result.usage.output_tokens, ms }
187        await update($, speed, cur => addSpeedSample(cur, sample))
188      }
189    }
190
191    return result
192  })
193
194  // Edits made inside subagents count too: they change the same files.
195  on('tool.call', { tool: 'Edit' }, async ($, e, next) => {
196    const ran = await next(e)
197    if (ran.deny === undefined && ran.isError !== true && ran.result.staged !== true) {
198      await record($, ran.result.filePath, countPatch(ran.result.structuredPatch))
199    }
200
201    return ran
202  })
203
204  on('tool.call', { tool: 'Write' }, async ($, e, next) => {
205    const ran = await next(e)
206    if (ran.deny === undefined && ran.isError !== true && ran.result.staged !== true) {
207      const r = ran.result
208      const delta =
209        r.structuredPatch.length > 0
210          ? countPatch(r.structuredPatch)
211          : r.type === 'create'
212            ? { added: countLines(r.content), removed: 0 }
213            : { added: 0, removed: 0 }
214      await record($, r.filePath, delta)
215    }
216
217    return ran
218  })
219}
220

1import type { HudActive, HudContext, HudEdits, HudSpeedSample, HudTokens } from '../types'
2
3export type HudSnapshot = {
4  context: HudContext | null
5  compactAt: number | null
6  cost: number | null
7  edits: HudEdits
8  tokens: HudTokens
9  tps: number | null
10  active: HudActive
11  now: number
12}
13
14const BAR_CELLS = 20
15
16// Claude Code draws the whole status line dimmed (`<Text dimColor>` around
17// it), so every piece here states its own intensity: 22 (normal) first, then a
18// bright color. Labels are bright white, figures bold in their own color
19// (green/yellow/red where it grades good to bad), separators and the bar's
20// empty cells light gray.
21const ansi = (code: string) => (text: string) => (text === '' ? '' : `\x1b[22;${code}m${text}\x1b[0m`)
22const label = ansi('97')
23const light = ansi('37')
24const green = ansi('1;92')
25const yellow = ansi('1;93')
26const red = ansi('1;91')
27const blue = ansi('1;94')
28const magenta = ansi('1;95')
29const cyan = ansi('1;96')
30const orange = ansi('1;38;5;208')
31const softBlue = ansi('1;38;5;111')
32
33export const countPatch = (hunks: readonly { lines: readonly string[] }[]) => {
34  let added = 0
35  let removed = 0
36  for (const hunk of hunks) {
37    for (const line of hunk.lines) {
38      if (line.startsWith('+')) added += 1
39      else if (line.startsWith('-')) removed += 1
40    }
41  }
42
43  return { added, removed }
44}
45
46export const countLines = (text: string) =>
47  text === '' ? 0 : text.split('\n').length - (text.endsWith('\n') ? 1 : 0)
48
49export const formatTokens = (n: number) => {
50  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1)}M`
51  if (n >= 10_000) return `${Math.round(n / 1000)}k`
52  if (n >= 1000) return `${+(n / 1000).toFixed(1)}k`
53
54  return String(n)
55}
56
57// The input and output totals to a tenth of the unit, the zero kept (615.0k,
58// 1.2M); a count under a thousand stays whole.
59export const formatTokensTenths = (n: number) => {
60  if (n >= 999_950) return `${(n / 1_000_000).toFixed(1)}M`
61  if (n >= 1000) return `${(n / 1000).toFixed(1)}k`
62
63  return String(n)
64}
65
66export const formatDuration = (ms: number) => {
67  const seconds = Math.floor(Math.max(0, ms) / 1000)
68  if (seconds < 60) return `${seconds}s`
69  const minutes = Math.floor(seconds / 60)
70  if (minutes < 60) return `${minutes}m${String(seconds % 60).padStart(2, '0')}s`
71
72  return `${Math.floor(minutes / 60)}h${String(minutes % 60).padStart(2, '0')}m`
73}
74
75// Share of all input tokens (uncached + cache write + cache read) the cache served.
76export const cacheHitPercent = (t: HudTokens) => {
77  const input = t.input + t.cacheWrite + t.cacheRead
78
79  return input === 0 ? null : (t.cacheRead / input) * 100
80}
81
82// TPS is over the main loop's last few requests: their output tokens over
83// their streaming time together, so a long answer weighs more than a request
84// that only called a tool.
85const SPEED_SAMPLES = 10
86
87export const addSpeedSample = (samples: readonly HudSpeedSample[], sample: HudSpeedSample) =>
88  [...samples, sample].slice(-SPEED_SAMPLES)
89
90export const recentTps = (samples: readonly HudSpeedSample[]) => {
91  const tokens = samples.reduce((n, s) => n + s.tokens, 0)
92  const ms = samples.reduce((n, s) => n + s.ms, 0)
93
94  return ms === 0 ? null : (tokens / ms) * 1000
95}
96
97// Active time is the union of the loops' busy spans: the clock runs while at
98// least one loop (the main one, `main`, or a subagent's) is working, once
99// however many are.
100export const NO_ACTIVITY: HudActive = { loops: [], since: null, totalMs: 0 }
101
102export const enterLoop = (a: HudActive, loop: string, now: number): HudActive =>
103  a.loops.includes(loop) ? a : { ...a, loops: [...a.loops, loop], since: a.since ?? now }
104
105export const leaveLoop = (a: HudActive, loop: string, now: number): HudActive => {
106  if (!a.loops.includes(loop)) return a
107  const loops = a.loops.filter(one => one !== loop)
108  if (loops.length > 0 || a.since === null) return { ...a, loops }
109
110  return { loops, since: null, totalMs: a.totalMs + (now - a.since) }
111}
112
113export const activeMs = (a: HudActive, now: number) => a.totalMs + (a.since === null ? 0 : now - a.since)
114
115// Cells the terminal draws for the text, as the statusLine command measures
116// it: CJK two, the rest one.
117const WIDE = /[ᄀ-ᅟ⺀-〾ぁ-㏿㐀-䶿一-鿿가-힣豈-﫿︰-﹏＀-｠￠-￦⚡]/
118const displayWidth = (text: string) =>
119  [...text.replace(/\x1b\[[0-9;]*m/g, '')].reduce((n, ch) => n + (WIDE.test(ch) ? 2 : 1), 0)
120
121// `上下文 207k ████░░░░░░░┃567k░░░░░ 1M`: used on the left in the fill's
122// color, the window at the bar's end, and the auto-compaction threshold as a
123// mark with its figure written into the bar beside it (before it when no room
124// is left after).
125const contextPart = (ctx: HudContext | null, threshold: number | null, cells = BAR_CELLS) => {
126  if (ctx === null || ctx.tokens === null) return `${label('上下文 ')}${label('等待首次响应')}`
127  const used = ctx.tokens
128  // The fill's color: how close `used` is to the threshold (or the window).
129  const pressure = used / (threshold ?? (ctx.window || 1))
130  const fill = pressure < 0.6 ? green : pressure < 0.85 ? yellow : red
131  const filled = ctx.window > 0 ? Math.min(cells, Math.round((used / ctx.window) * cells)) : 0
132  // Cells from..to of the bar: filled ones colored by pressure, the rest light gray.
133  const run = (from: number, to: number) =>
134    fill('█'.repeat(Math.max(0, Math.min(to, filled) - from))) +
135    light('░'.repeat(Math.max(0, to - Math.max(from, filled))))
136  let bar = run(0, cells)
137  if (threshold !== null && ctx.window > 0) {
138    const mark = Math.min(cells, Math.max(0, Math.round((threshold / ctx.window) * cells)))
139    const figure = formatTokens(threshold)
140    if (mark + figure.length <= cells) {
141      bar = `${run(0, mark)}${orange(`┃${figure}`)}${run(mark + figure.length, cells)}`
142    } else {
143      const start = Math.max(0, mark - figure.length)
144      bar = `${run(0, start)}${orange(`${figure}┃`)}${run(mark, cells)}`
145    }
146  }
147
148  return `${label('上下文 ')}${fill(`${ctx.isEstimate === true ? '~' : ''}${formatTokens(used)}`)} ${bar} ${softBlue(formatTokens(ctx.window))}`
149}
150
151// Two rows of cells, each cell a list of parts, ANSI-colored: the session's
152// work (model, context, changes), then what it has consumed (cost, tokens and
153// speed, active time and cache). The statusLine command writes the model
154// into the empty first cell, sizes each column to its widest cell, spreads a
155// narrower cell's parts across the width (the spare room goes between them, not
156// after them) and joins the columns with 丨.
157export type HudCells = string[][][]
158
159// Where a part gives way when its column lines parts up and it is narrower
160// than the widest there: the statusLine command puts the room at this mark
161// (a space when there is none), so the figure after it ends with the widest.
162export const FLEX = '\t'
163
164// The least room between the parts of a cell, per column: the cost's label
165// and figure need only a space (the figure ends under the model's effort), the
166// last column's parts stand further apart.
167export const HUD_GAPS = [1, 4, 6] as const
168
169export const renderCells = (s: HudSnapshot): HudCells => {
170  const hit = cacheHitPercent(s.tokens)
171  const hitText = hit === null ? label('—') : (hit >= 80 ? green : hit >= 50 ? yellow : red)(`${hit.toFixed(2)}%`)
172  const input = s.tokens.input + s.tokens.cacheWrite + s.tokens.cacheRead
173  const isActive = s.active.since !== null
174  const speed = [
175    `${label('输入 ')}${blue(formatTokensTenths(input))}`,
176    `${label('输出 ')}${magenta(formatTokensTenths(s.tokens.output))}`,
177    `${label('TPS ')}${s.tps === null ? label('—') : cyan(s.tps.toFixed(1))}`,
178  ]
179  // The bar grows by however much the speed cell below is wider, so the
180  // window's figure ends where the TPS does; a wider bar spreads that cell.
181  const speedWidth = speed.reduce((n, part) => n + displayWidth(part), 0) + HUD_GAPS[1] * (speed.length - 1)
182  const short = speedWidth - displayWidth(contextPart(s.context, s.compactAt))
183  const context = contextPart(s.context, s.compactAt, BAR_CELLS + Math.max(0, short))
184
185  return [
186    [
187      [''],
188      [context],
189      [
190        `${label('改动')}${FLEX}${blue(String(s.edits.files.length))}${label(' 文件')}`,
191        `${green(`+${s.edits.added}`)}${FLEX}${red(`−${s.edits.removed}`)}${label(' 行')}`,
192      ],
193    ],
194    [
195      [label('费用'), s.cost === null ? label('—') : yellow(`$${s.cost.toFixed(2)}`)],
196      speed,
197      [
198        `${label('活跃')}${FLEX}${(isActive ? cyan : softBlue)(formatDuration(activeMs(s.active, s.now)))}`,
199        `${label('缓存')}${FLEX}${hitText}`,
200      ],
201    ],
202  ]
203}
204

1// `isEstimate`: no response has reported the window yet, so `tokens` is the
2// engine's local estimate of what the first request will carry.
3export type HudContext = { tokens: number | null; window: number; isEstimate?: boolean }
4
5export type HudEdits = {
6  added: number
7  removed: number
8  files: string[]
9}
10
11// The loops working now (`main` or a subagent's id), since when at least one
12// has been, and the time banked from the busy spans that ended.
13export type HudActive = { loops: string[]; since: number | null; totalMs: number }
14
15// One main-loop request's output and how long it streamed, for the TPS figure.
16export type HudSpeedSample = { tokens: number; ms: number }
17
18export type HudTokens = {
19  input: number
20  output: number
21  cacheRead: number
22  cacheWrite: number
23  requests: number
24}
25
26declare module 'claude-code' {
27  interface PluginState {
28    'session-hud': {
29      compactAt: number | null
30      edits: HudEdits
31      tokens: HudTokens
32      speed: HudSpeedSample[]
33      active: HudActive
34    }
35  }
36}
37