token-usage

1import type { Register } from 'claude-code'
2
3import type { Meter } from '../types'
4
5const TOTALS = { plugin: 'token-usage', key: 'totals' } as const
6const METER = { plugin: 'token-usage', key: 'meter' } as const
7const ZERO = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, turns: 0 }
8const LIMIT_LABEL: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
9
10const fmt = (n: number) =>
11  n >= 1_000_000 ? `${(n / 1_000_000).toFixed(2)}M` : n >= 1000 ? `${(n / 1000).toFixed(1)}k` : `${n}`
12
13const level = (percent: number) => (percent >= 85 ? 'red' : percent >= 60 ? 'yellow' : 'green')
14
15const toMeter = (u: {
16  context: { tokens?: number; window: number; percent?: number }
17  rateLimits: { kind: string; percentUsed: number }[]
18  cost?: { usd: number }
19}): Meter => ({
20  contextTokens: u.context.tokens,
21  contextWindow: u.context.window,
22  contextPercent: u.context.percent,
23  costUsd: u.cost?.usd,
24  rateLimits: u.rateLimits.map(({ kind, percentUsed }) => ({ kind, percentUsed })),
25})
26
27export const register: Register = on => {
28  on('session.start', async ($, e, next) => {
29    await $.state.set(METER, toMeter(await $.session.usage()))
30
31    return next(e)
32  })
33
34  on('turn.complete', async ($, e, next) => {
35    const u = e.usage
36    if (u) {
37      const { value: t = ZERO } = await $.state.get(TOTALS)
38      await $.state.set(TOTALS, {
39        input: t.input + u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens,
40        output: t.output + u.output_tokens,
41        cacheRead: t.cacheRead + u.cache_read_input_tokens,
42        cacheWrite: t.cacheWrite + u.cache_creation_input_tokens,
43        turns: t.turns + 1,
44      })
45    }
46    await $.state.set(METER, toMeter(await $.session.usage()))
47
48    return next(e)
49  })
50
51  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
52    const { value: m } = await $.state.get(METER)
53    if (e.props.hasSurvey || !m) {
54      return next(e)
55    }
56
57    const { value: t = ZERO } = await $.state.get(TOTALS)
58    const { Box, Text } = $.ui.resolve(e)
59    const percent = m.contextPercent ?? 0
60    const width = e.props.bodyColumns >= 100 ? 12 : 6
61    const filled = Math.round((Math.min(percent, 100) / 100) * width)
62    const sep = <Text dimColor> · </Text>
63
64    return (
65      <Box>
66        <Text wrap="truncate-end">
67          <Text dimColor>Context </Text>
68          <Text color={level(percent)}>{'█'.repeat(filled)}</Text>
69          <Text dimColor>{'░'.repeat(width - filled)} </Text>
70          <Text color={level(percent)} bold>
71            {percent}%
72          </Text>
73          <Text dimColor>
74            {' '}
75            ({fmt(m.contextTokens ?? 0)}/{fmt(m.contextWindow)})
76          </Text>
77          {sep}
78          <Text dimColor>in </Text>
79          <Text color="cyan">{fmt(t.input)}</Text>
80          <Text dimColor> out </Text>
81          <Text color="magenta">{fmt(t.output)}</Text>
82          {t.input > 0 && (
83            <>
84              {sep}
85              <Text dimColor>cache </Text>
86              <Text color="blue">{Math.round((t.cacheRead / t.input) * 100)}%</Text>
87            </>
88          )}
89          {m.costUsd !== undefined && (
90            <>
91              {sep}
92              <Text color="green" bold>
93                ${m.costUsd.toFixed(2)}
94              </Text>
95            </>
96          )}
97          {m.rateLimits.map(r => (
98            <>
99              {sep}
100              <Text dimColor>{LIMIT_LABEL[r.kind] ?? r.kind} </Text>
101              <Text color={level(r.percentUsed)}>{r.percentUsed}%</Text>
102            </>
103          ))}
104        </Text>
105      </Box>
106    )
107  })
108}
109

1export type Totals = {
2  input: number
3  output: number
4  cacheRead: number
5  cacheWrite: number
6  turns: number
7}
8
9export type Meter = {
10  contextTokens?: number
11  contextWindow: number
12  contextPercent?: number
13  costUsd?: number
14  rateLimits: { kind: string; percentUsed: number }[]
15}
16
17declare module 'claude-code' {
18  interface PluginState {
19    'token-usage': {
20      totals: Totals
21      meter: Meter
22    }
23  }
24}
25