session-stats

1import type { EngineInterface, Register, SessionUsage, TurnUsage } from 'claude-code'
2import { atom, read, update } from 'claude-code'
3
4import type { Category, Snapshot, Totals } from '../types'
5
6const emptyTotals: Totals = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, turnMs: 0, byModel: {} }
7
8const snapshot = atom({ plugin: 'session-stats', key: 'snapshot' } as const, null)
9const breakdown = atom({ plugin: 'session-stats', key: 'breakdown' } as const, [] as Category[])
10const isOpen = atom({ plugin: 'session-stats', key: 'isOpen' } as const, false)
11const now = atom({ plugin: 'session-stats', key: 'now' } as const, 0)
12const totals = atom({ plugin: 'session-stats', key: 'totals' } as const, emptyTotals)
13const effort = atom({ plugin: 'session-stats', key: 'effort' } as const, null)
14const location = atom({ plugin: 'session-stats', key: 'location' } as const, null)
15
16const EFFORT_LEVELS = ['low', 'medium', 'high', 'xhigh', 'max'] as const
17const EFFORT_ICONS: Record<(typeof EFFORT_LEVELS)[number], string> = { low: '○', medium: '◐', high: '●', xhigh: '◉', max: '◈' }
18
19export const effortIcon = (level: string | number | null): string =>
20  level === null ? '' : (EFFORT_ICONS[String(level) as keyof typeof EFFORT_ICONS] ?? String(level))
21
22export const parseEffort = (args: string): string | null => {
23  const level = args.trim().toLowerCase()
24
25  return (EFFORT_LEVELS as readonly string[]).includes(level) ? level : null
26}
27
28const BAR_CELLS = 5
29const PANE_BAR_CELLS = 20
30const PANE = 'usage'
31const MAX_PANE_ROWS = 30
32const TICK_MS = 60_000
33const LOCATION_OVERHEAD = 4
34const BAND_RESERVE = 2
35
36type Level = 'green' | 'yellow' | 'red'
37
38let ticker: { cancel: () => void } | undefined
39export const prettyModel = (id: string): string => {
40  const match = /^claude-([a-z]+)-(\d+)(?:-(\d{1,2}))?(?!\d)/.exec(id)
41
42  if (!match) {
43    return id
44  }
45
46  const [, family = '', major, minor] = match
47
48  return `${family.charAt(0).toUpperCase()}${family.slice(1)} ${minor ? `${major}.${minor}` : major}`
49}
50
51export const levelColor = (percent: number): Level => (percent >= 85 ? 'red' : percent >= 60 ? 'yellow' : 'green')
52
53const capitalise = (word: string): string => `${word.charAt(0).toUpperCase()}${word.slice(1)}`
54
55const windowModel = (kind: string): string | undefined => /^seven_day_(.+)$/.exec(kind)?.[1]
56
57export const limitLabel = (kind: string): string => {
58  const model = windowModel(kind)
59
60  return model ? `W ${capitalise(model)}` : ({ five_hour: '5h', seven_day: 'W', spend_limit: 'Spend' })[kind] ?? kind
61}
62
63export const limitTitle = (kind: string): string => {
64  const model = windowModel(kind)
65
66  return model
67    ? `Weekly · ${capitalise(model)}`
68    : ({ five_hour: 'Session limit', seven_day: 'Weekly · all models', spend_limit: 'Spend limit' })[kind] ?? kind
69}
70
71export const formatDuration = (ms: number): string => {
72  const seconds = Math.round(ms / 1000)
73  const hours = Math.floor(seconds / 3600)
74  const minutes = Math.floor((seconds % 3600) / 60)
75
76  if (hours > 0) {
77    return `${hours}h ${minutes}m`
78  }
79
80  return minutes > 0 ? `${minutes}m ${seconds % 60}s` : `${seconds}s`
81}
82
83export const addTurn = (current: Totals, usage: TurnUsage, durationMs: number): Totals => ({
84  input: current.input + usage.input_tokens,
85  output: current.output + usage.output_tokens,
86  cacheRead: current.cacheRead + usage.cache_read_input_tokens,
87  cacheWrite: current.cacheWrite + usage.cache_creation_input_tokens,
88  turnMs: current.turnMs + durationMs,
89  byModel: {
90    ...current.byModel,
91    [usage.model]:
92      (current.byModel[usage.model] ?? 0) +
93      usage.input_tokens +
94      usage.output_tokens +
95      usage.cache_read_input_tokens +
96      usage.cache_creation_input_tokens,
97  },
98})
99
100export const cacheHit = (current: Totals): number | null => {
101  const prompt = current.input + current.cacheRead + current.cacheWrite
102
103  return prompt === 0 ? null : Math.round((current.cacheRead / prompt) * 100)
104}
105
106export const modelShares = (current: Totals): { model: string; percent: number }[] => {
107  const sum = Object.values(current.byModel).reduce((acc, tokens) => acc + tokens, 0)
108
109  return Object.entries(current.byModel)
110    .map(([model, tokens]) => ({ model, percent: sum === 0 ? 0 : Math.round((tokens / sum) * 100) }))
111    .filter(share => share.percent > 0)
112    .sort((a, b) => b.percent - a.percent)
113}
114
115export const fillCells = (percent: number, cells = BAR_CELLS): number =>
116  Math.max(0, Math.min(cells, Math.round((percent / 100) * cells)))
117
118export const countdown = (resetsAt: string | undefined, at: number): string => {
119  const minutes = Math.ceil((Date.parse(resetsAt ?? '') - at) / 60_000)
120
121  if (!Number.isFinite(minutes) || minutes <= 0) {
122    return ''
123  }
124
125  const days = Math.floor(minutes / 1440)
126  const hours = Math.floor((minutes % 1440) / 60)
127
128  if (days > 0) {
129    return `${days}d ${hours}h`
130  }
131
132  return hours > 0 ? `${hours}h ${minutes % 60}m` : `${minutes}m`
133}
134
135export const tokensLabel = (tokens: number): string =>
136  tokens >= 1_000_000 ? `${(tokens / 1_000_000).toFixed(1)}M` : tokens >= 1000 ? `${Math.round(tokens / 1000)}k` : `${tokens}`
137
138export const toSnapshot = (model: string, usage: SessionUsage): Snapshot => ({
139  model,
140  percent: usage.context.percent ?? null,
141  tokens: usage.context.tokens ?? null,
142  window: usage.context.window,
143  limits: usage.rateLimits.map(({ kind, percentUsed, resetsAt }) => ({ kind, percentUsed, resetsAt })),
144  usd: usage.cost?.usd ?? null,
145})
146
147export const parseLocation = (revParse: string, branch: string, head: string): string | null => {
148  const [toplevel, gitDir, commonDir] = revParse.split('\n').map(line => line.trim())
149
150  if (!toplevel || !gitDir || !commonDir) {
151    return null
152  }
153
154  const worktree = gitDir === commonDir ? null : toplevel.split('/').filter(Boolean).pop()
155
156  return worktree || branch.trim() || head.trim() || null
157}
158
159export const fitLocation = (name: string | null, room: number): string => {
160  const left = room - LOCATION_OVERHEAD
161
162  if (!name || left < 1) {
163    return ''
164  }
165
166  return name.length > left ? `${name.slice(0, left - 1)}…` : name
167}
168
169export const bandColumns = (current: Snapshot, icon: string, at: number): number => {
170  const percent = current.percent === null ? '–' : `${current.percent}%`
171  const limits = current.limits.reduce((sum, limit) => {
172    const left = countdown(limit.resetsAt, at)
173
174    return sum + 2 + limitLabel(limit.kind).length + 1 + `${Math.round(limit.percentUsed)}%`.length + (left ? left.length + 3 : 0)
175  }, 0)
176
177  return prettyModel(current.model).length + (icon ? icon.length + 1 : 0) + 2 + BAR_CELLS + 1 + percent.length + limits + 3
178}
179
180async function readLocation($: EngineInterface): Promise<string | null> {
181  const top = await $.process.run(['git', 'rev-parse', '--path-format=absolute', '--show-toplevel', '--git-dir', '--git-common-dir'])
182
183  if (top.exitCode !== 0) {
184    return null
185  }
186
187  const named = await $.process.run(['git', 'branch', '--show-current'])
188  const detached = named.exitCode === 0 && !named.stdout.trim() ? await $.process.run(['git', 'rev-parse', '--short', 'HEAD']) : null
189
190  return parseLocation(top.stdout, named.exitCode === 0 ? named.stdout : '', detached?.stdout ?? '')
191}
192
193async function refreshLocation($: EngineInterface) {
194  const next = await readLocation($).catch(() => null)
195
196  if ((await read($, location)) !== next) {
197    await update($, location, () => next)
198  }
199}
200
201async function refresh($: EngineInterface) {
202  const usage = await $.session.usage()
203  const model = await $.session.model()
204  await update($, snapshot, () => toSnapshot(model, usage))
205  await update($, now, () => Date.now())
206}
207
208async function tick($: EngineInterface) {
209  await update($, now, () => Date.now())
210}
211
212async function loadBreakdown($: EngineInterface) {
213  await refresh($)
214  const { context } = await $.session.usage({ breakdown: 'summary' })
215  const categories = (context.breakdown?.categories ?? []).map(({ name, tokens, kind }) => ({ name, tokens, kind }))
216  await update($, breakdown, () => categories)
217}
218
219async function openPane($: EngineInterface) {
220  await loadBreakdown($)
221  const current = await read($, snapshot)
222  const lines = (await read($, breakdown)).length + (current?.limits.length ?? 0) + 20
223  await $.ui.open({ id: PANE, title: 'Usage', focus: true, closeOnEscape: true, rows: Math.min(lines, MAX_PANE_ROWS) })
224  await update($, isOpen, () => true)
225}
226
227async function togglePane($: EngineInterface) {
228  const panes = await $.ui.panes()
229
230  if (panes.some(pane => pane.id === PANE)) {
231    await $.ui.close({ id: PANE })
232    await update($, isOpen, () => false)
233
234    return
235  }
236
237  await openPane($)
238}
239
240type TextElement = ReturnType<EngineInterface['ui']['resolve']>['Text']
241
242function squares(Text: TextElement, percent: number | null, cells: number) {
243  const filled = percent === null ? 0 : fillCells(percent, cells)
244  const color = percent === null ? undefined : levelColor(percent)
245
246  return [...Array(cells).keys()].map(index => (
247    <Text color={index < filled ? color : undefined} dimColor={index >= filled}>◼</Text>
248  ))
249}
250
251const RESEED_MS = 500
252
253async function seed($: EngineInterface) {
254  await update($, isOpen, () => false)
255  await refresh($)
256  await refreshLocation($)
257  ticker?.cancel()
258  ticker = await $.clock.every(TICK_MS, () => tick($))
259}
260
261export const register: Register = on => {
262  on('session.start', async ($, e, next) => {
263    await seed($)
264    await $.command
265      .register({ name: 'session-stats', description: 'Show model, effort and usage in a pane' })
266      .catch(() => undefined)
267
268    return next(e)
269  })
270
271  // /clear resets this plugin's state after session.end and fires no session.start, so setup reruns from a timer.
272  on('session.end', async ($, e, next) => {
273    if (e.reason === 'clear') {
274      $.clock.after(RESEED_MS, () => {
275        void seed($)
276      })
277    }
278
279    return next(e)
280  })
281
282  on('command.run', { command: 'session-stats' }, async $ => {
283    await openPane($)
284
285    return { text: 'Session stats pane opened.' }
286  })
287
288  on('command.run', { command: 'effort' }, async ($, e, next) => {
289    const result = await next(e)
290    const level = parseEffort(e.args)
291
292    if (level) {
293      await update($, effort, () => level)
294    }
295
296    return result
297  })
298
299  on('command.run', { command: 'model' }, async ($, e, next) => {
300    const before = (await read($, snapshot))?.model
301    const result = await next(e)
302    await refresh($).catch(() => undefined)
303
304    if ((await read($, snapshot))?.model !== before) {
305      await update($, effort, () => null)
306    }
307
308    return result
309  })
310
311  on('ui.close', async ($, e, next) => {
312    const closed = await next(e)
313
314    if (e.id === PANE) {
315      await update($, isOpen, () => false)
316    }
317
318    return closed
319  })
320
321  on('session.measure', async ($, e, next) => {
322    await refresh($)
323
324    return next(e)
325  })
326
327  on('turn.step', async function* ($, e, next) {
328    if (!e.agentId) {
329      await update($, effort, () => e.effort ?? null)
330    }
331
332    yield* next(e)
333  })
334
335  on('turn.complete', async ($, e, next) => {
336    if (e.usage) {
337      const usage = e.usage
338      await update($, totals, current => addTurn(current, usage, e.durationMs))
339    }
340
341    await refresh($)
342    await refreshLocation($)
343
344    return next(e)
345  })
346
347  on('tool.call', { tool: ['EnterWorktree', 'ExitWorktree'] }, async ($, e, next) => {
348    const result = await next(e)
349
350    $.clock.after(0, () => {
351      void refreshLocation($).catch(() => undefined)
352    })
353
354    return result
355  })
356
357  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
358    const other = await next(e)
359    const current = await read($, snapshot)
360
361    if (e.props.hasSurvey || current === null) {
362      return other
363    }
364
365    const { Box, Button, Text } = $.ui.resolve(e)
366    const at = await read($, now)
367    const paneIsOpen = await read($, isOpen)
368    const icon = effortIcon(await read($, effort))
369    const shown = fitLocation(await read($, location), e.props.bodyColumns - bandColumns(current, icon, at) - BAND_RESERVE)
370
371    const band = (
372      <Box>
373        <Text bold>{prettyModel(current.model)}</Text>
374        {icon && <Text>{` ${icon}`}</Text>}
375        <Text>  </Text>
376        {squares(Text, current.percent, BAR_CELLS)}
377        <Text color={current.percent === null ? undefined : levelColor(current.percent)} dimColor={current.percent === null}>
378          {` ${current.percent === null ? '–' : `${current.percent}%`}`}
379        </Text>
380        {current.limits.map(limit => {
381          const left = countdown(limit.resetsAt, at)
382
383          return (
384            <Box>
385              <Text dimColor>  {limitLabel(limit.kind)} </Text>
386              <Text color={levelColor(limit.percentUsed)}>{Math.round(limit.percentUsed)}%</Text>
387              {left && <Text dimColor> ({left})</Text>}
388            </Box>
389          )
390        })}
391        {shown ? <Text dimColor>{`  ⎇ ${shown}`}</Text> : null}
392        <Text>  </Text>
393        <Button key="toggle" label={paneIsOpen ? '−' : '+'} plain dimColor onPress={() => togglePane($)} />
394      </Box>
395    )
396
397    return other ? (
398      <Box flexDirection="column">
399        {band}
400        {other}
401      </Box>
402    ) : (
403      band
404    )
405  })
406
407  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
408    const { Box, Button, Text } = $.ui.resolve(e)
409    const current = await read($, snapshot)
410    const categories = await read($, breakdown)
411    const spent = await read($, totals)
412    const at = await read($, now)
413    const level = await read($, effort)
414
415    if (current === null) {
416      return <Text dimColor>No usage reading yet.</Text>
417    }
418
419    const shown = categories
420      .filter(category => category.kind !== 'deferred')
421      .sort((a, b) => Number(a.kind !== 'used') - Number(b.kind !== 'used') || b.tokens - a.tokens)
422    const nameWidth = Math.max(0, ...shown.map(category => category.name.length))
423    const tokenWidth = Math.max(0, ...shown.map(category => tokensLabel(category.tokens).length))
424    const hit = cacheHit(spent)
425    const shares = modelShares(spent)
426    const titleWidth = Math.max(8, ...current.limits.map(limit => limitTitle(limit.kind).length)) + 2
427    const tokenRows: [string, number][] = [
428      ['Input', spent.input],
429      ['Output', spent.output],
430      ['Cache read', spent.cacheRead],
431      ['Cache write', spent.cacheWrite],
432    ]
433
434    return (
435      <Box flexDirection="column">
436        <Box>
437          <Text bold>Usage  </Text>
438          <Text dimColor>{prettyModel(current.model)}  </Text>
439          {level !== null && <Text>{`${effortIcon(level)} ${level}  `}</Text>}
440          <Button key="refresh" label="Refresh (r)" hotkey="r" onPress={() => loadBreakdown($)} />
441        </Box>
442        <Text> </Text>
443        <Box>
444          <Text dimColor>{'Context'.padEnd(titleWidth)}</Text>
445          {squares(Text, current.percent, PANE_BAR_CELLS)}
446          <Text color={current.percent === null ? undefined : levelColor(current.percent)}>
447            {` ${current.percent === null ? '–' : `${current.percent}%`}`}
448          </Text>
449          <Text dimColor>
450            {current.tokens === null ? '' : `  ${tokensLabel(current.tokens)} / ${tokensLabel(current.window)}`}
451          </Text>
452        </Box>
453        {current.limits.map(limit => {
454          const left = countdown(limit.resetsAt, at)
455
456          return (
457            <Box>
458              <Text dimColor>{limitTitle(limit.kind).padEnd(titleWidth)}</Text>
459              {squares(Text, limit.percentUsed, PANE_BAR_CELLS)}
460              <Text color={levelColor(limit.percentUsed)}>{` ${Math.round(limit.percentUsed)}%`}</Text>
461              {left && <Text dimColor>  resets in {left}</Text>}
462            </Box>
463          )
464        })}
465        <Text> </Text>
466        <Text bold>This session</Text>
467        <Box>
468          {current.usd !== null && <Text dimColor>Cost </Text>}
469          {current.usd !== null && <Text>{`$${current.usd.toFixed(2)}`}  </Text>}
470          <Text dimColor>Turns </Text>
471          <Text>{formatDuration(spent.turnMs)}  </Text>
472          {hit !== null && <Text dimColor>Cache hit </Text>}
473          {hit !== null && <Text>{hit}%</Text>}
474        </Box>
475        {shares.length > 0 && (
476          <Box>
477            {shares.map(share => (
478              <Box>
479                <Text dimColor>{prettyModel(share.model)} </Text>
480                <Text>{share.percent}%  </Text>
481              </Box>
482            ))}
483          </Box>
484        )}
485        <Text> </Text>
486        <Text bold>Breakdown</Text>
487        {tokenRows.map(([label, tokens]) => (
488          <Box>
489            <Text dimColor>{label.padEnd(12)}</Text>
490            <Text>{tokensLabel(tokens)}</Text>
491          </Box>
492        ))}
493        <Text> </Text>
494        <Text bold>Context window</Text>
495        {shown.length === 0 && <Text dimColor>No context breakdown yet.</Text>}
496        {shown.map(category => (
497          <Box>
498            <Text dimColor={category.kind !== 'used'}>{category.name.padEnd(nameWidth)}  </Text>
499            <Text dimColor>{tokensLabel(category.tokens).padStart(tokenWidth)}  </Text>
500            <Text dimColor>{`${((category.tokens / current.window) * 100).toFixed(1)}%`}</Text>
501          </Box>
502        ))}
503      </Box>
504    )
505  })
506}
507

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2export type Snapshot = {
3  model: string
4  percent: number | null
5  tokens: number | null
6  window: number
7  limits: Limit[]
8  usd: number | null
9}
10export type Totals = {
11  input: number
12  output: number
13  cacheRead: number
14  cacheWrite: number
15  turnMs: number
16  byModel: Record<string, number>
17}
18export type Category ={ name: string; tokens: number; kind: 'used' | 'free' | 'buffer' | 'deferred' }
19
20declare module 'claude-code' {
21  interface PluginState {
22    'session-stats': {
23      snapshot: Snapshot | null
24      breakdown: Category[]
25      isOpen: boolean
26      now: number
27      totals: Totals
28      effort: string | number | null
29      location: string | null
30    }
31  }
32}
33