Token Usage Line: per-response & session tokens

1import type { Register, TurnUsage } from 'claude-code'
2
3import type { Tokens } from '../types'
4
5const ZERO: Tokens = { input: 0, cacheRead: 0, cacheWrite: 0, output: 0 }
6
7const LAST = { plugin: 'token-usage-line', key: 'last' } as const
8const TOTAL = { plugin: 'token-usage-line', key: 'total' } as const
9const PENDING = { plugin: 'token-usage-line', key: 'pending' } as const
10
11const ACCENT = 'claude'
12const LABEL = 'inactive'
13const VALUE = 'text'
14
15const fromUsage = (u: TurnUsage): Tokens => ({
16  input: u.input_tokens,
17  cacheRead: u.cache_read_input_tokens,
18  cacheWrite: u.cache_creation_input_tokens,
19  output: u.output_tokens,
20})
21
22const add = (a: Tokens, b: Tokens): Tokens => ({
23  input: a.input + b.input,
24  cacheRead: a.cacheRead + b.cacheRead,
25  cacheWrite: a.cacheWrite + b.cacheWrite,
26  output: a.output + b.output,
27})
28
29const k = (n: number) => (n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${(n / 1e3).toFixed(1)}k` : `${n}`)
30
31const inputOf = (t: Tokens) => t.input + t.cacheRead + t.cacheWrite
32
33const cachedPercent = (t: Tokens) => (inputOf(t) === 0 ? 0 : Math.round((t.cacheRead / inputOf(t)) * 100))
34
35export const register: Register = on => {
36  on('session.start', async ($, e, next) => {
37    $.ui.status(undefined)
38    return next(e)
39  })
40
41  on('session.end', async ($, e, next) => {
42    if (e.reason === 'clear') {
43      await $.state.set(LAST, null)
44      await $.state.set(PENDING, ZERO)
45      await $.state.set(TOTAL, ZERO)
46    }
47    return next(e)
48  })
49
50  on('turn.complete', async ($, e, next) => {
51    const result = await next(e)
52    const used = e.usage ? fromUsage(e.usage) : ZERO
53
54    const { value: total = ZERO } = await $.state.get(TOTAL)
55    await $.state.set(TOTAL, add(total, used))
56
57    const { value: subagents = ZERO } = await $.state.get(PENDING)
58    if (e.agentId !== undefined) {
59      await $.state.set(PENDING, add(subagents, used))
60      return result
61    }
62
63    await $.state.set(LAST, add(used, subagents))
64    await $.state.set(PENDING, ZERO)
65    return result
66  })
67
68  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
69    const below = await next(e)
70    if (e.props.hasSurvey) return below
71
72    const [{ value: last = null }, { value: total = ZERO }, usage] = await Promise.all([
73      $.state.get(LAST),
74      $.state.get(TOTAL),
75      $.session.usage(),
76    ])
77    const { Box, Text } = $.ui.resolve(e)
78
79    const pair = (label: string, value: string) => (
80      <Box flexDirection="row" gap={1}>
81        <Text color={LABEL}>{label}</Text>
82        <Text color={VALUE} bold>
83          {value}
84        </Text>
85      </Box>
86    )
87    const group = (title: string, t: Tokens) => (
88      <Box flexDirection="row" gap={1}>
89        <Text color={ACCENT} bold>
90          {title}
91        </Text>
92        {pair('in', k(inputOf(t)))}
93        {inputOf(t) > 0 ? <Text color={LABEL}>{`(${cachedPercent(t)}% cached)`}</Text> : null}
94        {pair('out', k(t.output))}
95      </Box>
96    )
97
98    const empty = inputOf(total) + total.output === 0
99    const usd = usage.cost?.usd ?? 0
100    const contextPercent = Math.round(usage.context.percent ?? 0)
101
102    const sections = empty
103      ? [<Text color={LABEL}>tokens: waiting for the first response</Text>]
104      : [
105          last ? group('Last', last) : null,
106          group('Session', total),
107          usd > 0 ? <Text color={VALUE} bold>{`~$${usd.toFixed(2)}`}</Text> : null,
108          contextPercent > 0 ? pair('context', `${contextPercent}%`) : null,
109        ].filter(section => section !== null)
110
111    return (
112      <Box flexDirection="column">
113        <Box flexDirection="row" flexWrap="nowrap" gap={1}>
114          <Text color={ACCENT}>✻</Text>
115          {sections.flatMap((section, i) => (i === 0 ? [section] : [<Text color={LABEL}>│</Text>, section]))}
116        </Box>
117        {below}
118      </Box>
119    )
120  })
121}
122

1export type Tokens = { input: number; cacheRead: number; cacheWrite: number; output: number }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'token-usage-line': { last: Tokens | null; pending: Tokens; total: Tokens }
6  }
7}
8