token-usage

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import { addTurn, EMPTY, fmt, pct, SERIES, sum } from './lib'
5
6const PANE = 'token-usage'
7const COLUMNS = 28 // a narrow dock beside the transcript
8const BAR = 24 // cells per bar: the dock minus border and padding
9const tally = atom({ plugin: 'token-usage', key: 'tally' } as const, EMPTY)
10
11const ROWS = [
12  { key: 'input', name: 'New input' },
13  { key: 'cacheRead', name: 'Cache read' },
14  { key: 'cacheWrite', name: 'Cache write' },
15  { key: 'output', name: 'Output' },
16] as const
17
18export const register: Register = on => {
19  on('session.start', async ($, e, next) => {
20    await $.command.register({ name: 'tokens', description: 'Show the token usage pane' })
21    void $.ui.open({ id: PANE, title: 'Tokens', columns: COLUMNS })
22    return next(e)
23  })
24
25  on('command.run', { command: 'tokens' }, async $ => {
26    await $.ui.open({ id: PANE, title: 'Tokens', columns: COLUMNS })
27    return { text: 'Token pane opened.' }
28  })
29
30  // Every turn that made requests, the main loop's and each subagent's, counts toward the session.
31  on('turn.complete', async ($, e, next) => {
32    if (e.usage) {
33      const u = e.usage
34      await update($, tally, t => addTurn(t, u))
35    }
36    return next(e)
37  })
38
39  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
40    const { Box, Text } = $.ui.resolve(e)
41    const { total } = await read($, tally)
42    const all = sum(total)
43    const prompt = total.input + total.cacheRead + total.cacheWrite
44    const hit = total.cacheRead / Math.max(1, prompt)
45    const real = total.input + total.cacheWrite + total.output
46
47    const bar = (share: number, color: string) => {
48      const filled = Math.round(share * BAR)
49      return (
50        <Text>
51          <Text color={color}>{'█'.repeat(filled)}</Text>
52          <Text dimColor>{'─'.repeat(BAR - filled)}</Text>
53        </Text>
54      )
55    }
56    const rule = <Text dimColor>{'─'.repeat(BAR)}</Text>
57
58    return (
59      <Box flexDirection="column" borderStyle="round" borderDimColor paddingX={1}>
60        <Text bold>Where tokens went</Text>
61        <Text dimColor>{fmt(all)} tokens, by type</Text>
62
63        <Box flexDirection="column" marginTop={1}>
64          <Text dimColor>Input served from cache</Text>
65          <Text bold color={SERIES.cacheRead.color}>{pct(hit)}</Text>
66          {bar(hit, SERIES.cacheRead.color)}
67          <Text dimColor>{fmt(real)} tokens of real work</Text>
68        </Box>
69
70        {ROWS.map(r => {
71          const share = total[r.key] / Math.max(1, all)
72          return (
73            <Box key={r.key} flexDirection="column" marginTop={1}>
74              <Box justifyContent="space-between">
75                <Text>
76                  <Text color={SERIES[r.key].color}>● </Text>
77                  <Text>{r.name}</Text>
78                </Text>
79                <Text>
80                  <Text bold>{fmt(total[r.key])} </Text>
81                  <Text dimColor>{pct(share)}</Text>
82                </Text>
83              </Box>
84              {bar(share, SERIES[r.key].color)}
85            </Box>
86          )
87        })}
88
89        <Box flexDirection="column">
90          {rule}
91          <Box justifyContent="space-between">
92            <Text dimColor>Total</Text>
93            <Text>
94              <Text bold>{fmt(all)} </Text>
95              <Text dimColor>100%</Text>
96            </Text>
97          </Box>
98        </Box>
99      </Box>
100    )
101  })
102}
103

1import type { Counts, Tally } from '../types'
2
3export const KEYS = ['input', 'output', 'cacheRead', 'cacheWrite'] as const
4export type Key = (typeof KEYS)[number]
5
6// Colorblind-safe (Okabe-Ito) so the four series stay apart on dark and light themes.
7export const SERIES: Record<Key, { label: string; glyph: string; color: string }> = {
8  input: { label: 'in', glyph: '↓', color: '#56B4E9' },
9  output: { label: 'out', glyph: '↑', color: '#E69F00' },
10  cacheRead: { label: 'hit', glyph: '⟲', color: '#009E73' },
11  cacheWrite: { label: 'new', glyph: '✎', color: '#CC79A7' },
12}
13
14export const EMPTY: Tally = { total: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, turns: [] }
15const KEEP = 60
16
17type ApiUsage = {
18  input_tokens: number
19  output_tokens: number
20  cache_read_input_tokens: number
21  cache_creation_input_tokens: number
22}
23
24export function addTurn(t: Tally, u: ApiUsage): Tally {
25  const turn: Counts = {
26    input: u.input_tokens,
27    output: u.output_tokens,
28    cacheRead: u.cache_read_input_tokens,
29    cacheWrite: u.cache_creation_input_tokens,
30  }
31  const total = { ...t.total }
32  for (const k of KEYS) total[k] += turn[k]
33  return { total, turns: [...t.turns, turn].slice(-KEEP) }
34}
35
36export const sum = (c: Counts): number => KEYS.reduce((n, k) => n + c[k], 0)
37
38// 268123 -> "268k", 1234567 -> "1.2M": short enough for one line
39export function fmt(n: number): string {
40  if (n >= 1_000_000_000) return `${(n / 1_000_000_000).toFixed(1)}B`
41  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
42  if (n >= 10_000) return `${Math.round(n / 1000)}k`
43  if (n >= 1000) return `${(n / 1000).toFixed(1)}k`
44  return String(n)
45}
46
47// 0.956 -> "96%"; a sliver under 1% keeps a decimal so it never reads "0%": 0.005 -> "0.5%"
48export function pct(share: number): string {
49  const p = share * 100
50  return p > 0 && p < 1 ? `${p.toFixed(1)}%` : `${Math.round(p)}%`
51}
52

1export type Counts = { input: number; output: number; cacheRead: number; cacheWrite: number }
2export type Tally = { total: Counts; turns: Counts[] }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'token-usage': { tally: Tally }
7  }
8}
9