token-buddy

1import { atom, read, update } from 'claude-code'
2import type { Register, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { Snapshot } from '../types'
5
6const PANE = 'token-buddy'
7const TITLE = 'Token Buddy'
8const snapshot = atom({ plugin: 'token-buddy', key: 'snapshot' } as const, null)
9
10const LIMIT_NAMES: Record<string, string> = {
11  five_hour: '5-hour limit',
12  seven_day: 'Weekly limit',
13  spend_limit: 'Spend limit',
14}
15
16function toSnapshot(context: SessionContextUsage, limits: SessionRateLimit[]): Snapshot {
17  return {
18    tokens: context.tokens,
19    window: context.window,
20    percent: context.percent,
21    limits: limits.map(l => ({ kind: l.kind, percentUsed: l.percentUsed, resetsAt: l.resetsAt })),
22  }
23}
24
25export function mood(left: number): { face: string; line: string; color: string } {
26  if (left > 60) return { face: '(｡◕‿◕｡)', line: 'Plenty of room to play!', color: 'success' }
27  if (left > 30) return { face: '(・ω・)', line: 'Doing just fine~', color: 'suggestion' }
28  if (left > 10) return { face: '(；・∀・)', line: 'Getting a little full...', color: 'warning' }
29  return { face: '(×_×)', line: 'So full! Try /compact', color: 'error' }
30}
31
32export function bar(left: number, width = 16): string {
33  const filled = Math.round((Math.max(0, Math.min(100, left)) / 100) * width)
34  return '♥'.repeat(filled) + '·'.repeat(width - filled)
35}
36
37export function statusLine(snap: Snapshot): string {
38  if (snap.percent === undefined) return '(｡･ω･｡) counting tokens after the first reply'
39  const left = 100 - snap.percent
40  const limits = snap.limits
41    .map(l => `${LIMIT_NAMES[l.kind] ?? l.kind} ${Math.max(0, Math.round(100 - l.percentUsed))}% left`)
42    .join(' · ')
43  return `${mood(left).face} context ${bar(left, 10)} ${left}% left${limits ? ' · ' + limits : ''}`
44}
45
46function compact(n: number): string {
47  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(n % 1_000_000 === 0 ? 0 : 1)}M`
48  if (n >= 1000) return `${Math.round(n / 1000)}k`
49  return String(n)
50}
51
52function untilReset(resetsAt: string | undefined, now: number): string {
53  if (!resetsAt) return ''
54  const mins = Math.max(0, Math.round((Date.parse(resetsAt) - now) / 60000))
55  if (Number.isNaN(mins)) return ''
56  if (mins < 60) return ` · resets in ${mins}m`
57  const hours = Math.floor(mins / 60)
58  if (hours < 48) return ` · resets in ${hours}h ${mins % 60}m`
59  return ` · resets in ${Math.round(hours / 24)}d`
60}
61
62export const register: Register = on => {
63  on('session.start', async ($, e, next) => {
64    await $.command.register({ name: 'tokens', description: 'Open the Token Buddy panel' })
65    const usage = await $.session.usage()
66    const snap = toSnapshot(usage.context, usage.rateLimits)
67    await update($, snapshot, () => snap)
68    $.ui.status(statusLine(snap))
69    void $.ui.open({ id: PANE, title: TITLE })
70
71    return next(e)
72  })
73
74  on('command.run', { command: 'tokens' }, async $ => {
75    const opened = await $.ui.open({ id: PANE, title: TITLE })
76    const surfaces = (await $.session.surfaces()).join(', ') || 'none'
77    const snap = await read($, snapshot)
78    const line = snap ? statusLine(snap) : 'no reading yet'
79
80    if (opened.isPlaced) return { text: `Token Buddy is here! (｡◕‿◕｡)  ${line}` }
81
82    return { text: `${line}\n(panel not shown: ${opened.reason}; surfaces: ${surfaces})` }
83  })
84
85  on('session.measure', async ($, e, next) => {
86    const snap = toSnapshot(e.context, e.rateLimits)
87    await update($, snapshot, () => snap)
88    $.ui.status(statusLine(snap))
89
90    return next(e)
91  })
92
93  on('session.attach', async ($, e, next) => {
94    void $.ui.open({ id: PANE, title: TITLE })
95
96    return next(e)
97  })
98
99  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
100    const { Box, Text } = $.ui.resolve(e)
101    const snap = await read($, snapshot)
102    const now = await $.clock.now()
103
104    if (!snap || snap.percent === undefined) {
105      return (
106        <Box flexDirection="column">
107          <Text color="claude">(｡･ω･｡)ﾉ</Text>
108          <Text dimColor>Hi! I'll count your tokens after the first reply.</Text>
109        </Box>
110      )
111    }
112
113    const left = 100 - snap.percent
114    const m = mood(left)
115    const tokensLeft = snap.tokens === undefined ? undefined : Math.max(0, snap.window - snap.tokens)
116
117    return (
118      <Box flexDirection="column">
119        <Box key="face"><Text color={m.color} bold>{m.face}</Text></Box>
120        <Box key="mood"><Text italic>{m.line}</Text></Box>
121        <Text> </Text>
122        <Text bold>Context left</Text>
123        <Text color={m.color}>{bar(left)}</Text>
124        <Box key="left">
125          <Text>
126            {left}% left
127            {tokensLeft !== undefined && ` · ~${compact(tokensLeft)} of ${compact(snap.window)} tokens`}
128          </Text>
129        </Box>
130        {snap.limits.length > 0 && <Text> </Text>}
131        {snap.limits.map(l => {
132          const limitLeft = Math.max(0, Math.round(100 - l.percentUsed))
133          const lm = mood(limitLeft)
134          return (
135            <Box key={`limit-${l.kind}`} flexDirection="column">
136              <Text bold>{LIMIT_NAMES[l.kind] ?? l.kind}</Text>
137              <Text color={lm.color}>{bar(limitLeft)}</Text>
138              <Text dimColor>
139                {limitLeft}% left{untilReset(l.resetsAt, now)}
140              </Text>
141            </Box>
142          )
143        })}
144      </Box>
145    )
146  })
147}
148

1export type Snapshot = {
2  tokens?: number
3  window: number
4  percent?: number
5  limits: { kind: string; percentUsed: number; resetsAt?: string }[]
6}
7
8declare module 'claude-code' {
9  interface PluginState {
10    'token-buddy': { snapshot: Snapshot | null }
11  }
12}
13