context-tokens

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4const TOKENS_WARN = 100_000
5const TOKENS_CRIT = 150_000
6const GREEN = '#4eba65'
7const YELLOW = '#ffc107'
8const RED = '#e5484d'
9
10const usedTokens = atom({ plugin: 'context-tokens', key: 'tokens' } as const, null)
11
12export const register: Register = on => {
13  on('session.start', async ($, e, next) => {
14    await refreshTokens($)
15    return next(e)
16  })
17
18  on('tool.call', async ($, e, next) => {
19    const result = await next(e)
20    await refreshTokens($)
21    return result
22  })
23
24  on('turn.complete', async ($, e, next) => {
25    const result = await next(e)
26    await refreshTokens($)
27    return result
28  })
29
30  on('ui.render', { component: 'SessionMode', surface: 'desktop' }, async ($, e, next) => {
31    const tokens = await read($, usedTokens)
32    if (tokens === null) {
33      return next(e)
34    }
35
36    const { Box, Text } = $.ui.resolve(e)
37    const modes = e.props.modes.join(' & ')
38
39    return (
40      <Box>
41        <Text color={tokensColor(tokens)}>{formatTokens(tokens)}</Text>
42        {modes && <Text dimColor> • {modes}</Text>}
43      </Box>
44    )
45  })
46}
47
48async function refreshTokens($: EngineInterface): Promise<void> {
49  const { context } = await $.session.usage()
50  await update($, usedTokens, () => context.tokens ?? null)
51}
52
53function tokensColor(tokens: number): string {
54  if (tokens >= TOKENS_CRIT) return RED
55  if (tokens >= TOKENS_WARN) return YELLOW
56  return GREEN
57}
58
59function formatTokens(tokens: number): string {
60  if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`
61  if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1)}k`
62  return String(tokens)
63}
64

1export type UsedTokens = number | null
2
3declare module 'claude-code' {
4  interface PluginState {
5    'context-tokens': { tokens: UsedTokens }
6  }
7}
8