usage-meter

1import { atom, read, update } from 'claude-code'
2import type { Color, EngineInterface, Register, SessionUsage } from 'claude-code'
3
4import type { Limit, Reading, Tokens } from '../types'
5
6const ZERO: Tokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }
7const EMPTY: Reading = { limits: [] }
8const tokens = atom({ plugin: 'usage-meter', key: 'tokens' } as const, ZERO)
9const reading = atom({ plugin: 'usage-meter', key: 'reading' } as const, EMPTY)
10
11const LABELS: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
12
13export const label = (kind: string): string => LABELS[kind] ?? kind
14
15export const compact = (n: number): string =>
16  n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${(n / 1e3).toFixed(1)}k` : `${n}`
17
18export const severity = (percent: number): Color | undefined =>
19  percent >= 90 ? 'error' : percent >= 70 ? 'warning' : undefined
20
21const toReading = (u: Omit<SessionUsage, 'startedAt'>): Reading => ({
22  contextPercent: u.context.percent,
23  limits: u.rateLimits.map(({ kind, percentUsed }): Limit => ({ kind, percentUsed })),
24  usd: u.cost?.usd,
25})
26
27async function refresh($: EngineInterface): Promise<void> {
28  const usage = await $.session.usage()
29  await update($, reading, () => toReading(usage))
30}
31
32export const register: Register = on => {
33  on('session.start', async ($, e, next) => {
34    const r = await next(e)
35    await refresh($)
36    return r
37  })
38
39  on('turn.complete', async ($, e, next) => {
40    const u = e.usage
41    if (u) {
42      await update($, tokens, (t = ZERO) => ({
43        input: t.input + u.input_tokens,
44        output: t.output + u.output_tokens,
45        cacheRead: t.cacheRead + u.cache_read_input_tokens,
46        cacheWrite: t.cacheWrite + u.cache_creation_input_tokens,
47      }))
48    }
49    return next(e)
50  })
51
52  on('session.measure', async ($, e, next) => {
53    await update($, reading, () => toReading(e))
54    if (e.changed.includes('rateLimits')) {
55      for (const r of e.rateLimits.filter(r => r.percentUsed >= 90)) {
56        $.ui.toast(`Usage limit: ${label(r.kind)} window at ${r.percentUsed}%`)
57      }
58    }
59    return next(e)
60  })
61
62  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
63    if (e.props.hasSurvey) {
64      return next(e)
65    }
66
67    const t = await read($, tokens)
68    const r = await read($, reading)
69    const { Box, Text } = $.ui.resolve(e)
70    const sep = <Text dimColor>{'  ·  '}</Text>
71
72    return (
73      <Box flexDirection="row">
74        <Text color="claude">{'📊 '}</Text>
75        <Text bold>usage-meter</Text>
76        <Text>{'  '}</Text>
77        <Text dimColor>tokens </Text>
78        <Text bold>
79          {compact(t.input + t.output + t.cacheRead + t.cacheWrite)}
80        </Text>
81        <Text dimColor>
82          {` (in ${compact(t.input + t.cacheRead + t.cacheWrite)} / out ${compact(t.output)})`}
83        </Text>
84        {r.contextPercent !== undefined && sep}
85        {r.contextPercent !== undefined && <Text dimColor>ctx </Text>}
86        {r.contextPercent !== undefined && (
87          <Text bold color={severity(r.contextPercent)}>
88            {`${r.contextPercent}%`}
89          </Text>
90        )}
91        {r.limits.map(l => (
92          <Box key={`limit-${l.kind}`} flexDirection="row">
93            {sep}
94            <Text dimColor>{`${label(l.kind)} `}</Text>
95            <Text bold color={severity(l.percentUsed)}>
96              {`${l.percentUsed}%`}
97            </Text>
98          </Box>
99        ))}
100        {r.usd !== undefined && sep}
101        {r.usd !== undefined && (
102          <Text bold>
103            {`$${r.usd.toFixed(2)}`}
104          </Text>
105        )}
106      </Box>
107    )
108  })
109}
110

1export type Tokens = { input: number; output: number; cacheRead: number; cacheWrite: number }
2
3export type Limit = { kind: string; percentUsed: number }
4
5export type Reading = { contextPercent?: number; limits: Limit[]; usd?: number }
6
7declare module 'claude-code' {
8  interface PluginState {
9    'usage-meter': { tokens: Tokens; reading: Reading }
10  }
11}
12