token-garden

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { Forecast, Limit } from '../types'
5
6const forecast = atom({ plugin: 'token-garden', key: 'forecast' } as const, null)
7
8const BAR_CELLS = 20
9const BAR_COLOR = '#FF8C00'
10// '↯ Compact soon' is the widest label, '999.9k / 999.9k' the widest value.
11const LABEL_CELLS = 15
12const VALUE_CELLS = 16
13
14const SKIES = [
15  { from: 90, icon: '↯', word: 'Compact soon', color: 'red' },
16  { from: 75, icon: '☇', word: 'Storm', color: 'magenta' },
17  { from: 50, icon: '☂', word: 'Showers', color: 'blue' },
18  { from: 25, icon: '☁', word: 'Cloudy', color: 'cyan' },
19  { from: 0, icon: '☀', word: 'Clear', color: 'yellow' },
20]
21
22const skyFor = (percent: number) => SKIES.find(sky => percent >= sky.from) ?? SKIES[SKIES.length - 1]
23
24const bar = (percent: number) => {
25  const filled = Math.max(0, Math.min(BAR_CELLS, Math.round((percent / 100) * BAR_CELLS)))
26
27  return '█'.repeat(filled) + '░'.repeat(BAR_CELLS - filled)
28}
29
30const compact = (tokens: number) => {
31  const abs = Math.abs(tokens)
32
33  if (abs >= 1_000_000) {
34    return `${(tokens / 1_000_000).toFixed(1).replace(/\.0$/, '')}M`
35  }
36
37  if (abs >= 1000) {
38    return `${(tokens / 1000).toFixed(1).replace(/\.0$/, '')}k`
39  }
40
41  return String(tokens)
42}
43
44const clock = (iso: string) => {
45  const at = new Date(iso)
46  const hours = at.getHours()
47  const hh = String(hours % 12 || 12).padStart(2, '0')
48  const mm = String(at.getMinutes()).padStart(2, '0')
49
50  return `${hh}:${mm} ${hours < 12 ? 'AM' : 'PM'}`
51}
52
53const DAYS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat']
54
55const limitOf = (limits: readonly SessionRateLimit[], kind: string): Limit | null => {
56  const limit = limits.find(each => each.kind === kind)
57
58  return limit ? { percent: limit.percentUsed, resetsAt: limit.resetsAt ?? null } : null
59}
60
61const measure = async (
62  $: EngineInterface,
63  context: SessionContextUsage,
64  limits: readonly SessionRateLimit[],
65  isTurn: boolean,
66) => {
67  await update($, forecast, (prev: Forecast | null): Forecast | null => {
68    const tokens = context.tokens ?? prev?.tokens
69
70    if (tokens === undefined) {
71      return prev
72    }
73
74    const moved = prev && isTurn ? tokens - prev.tokens : null
75
76    return {
77      tokens,
78      window: context.window,
79      percent: context.percent ?? Math.round((tokens / context.window) * 100),
80      delta: moved === null || moved === 0 ? (prev?.delta ?? null) : moved,
81      session: limitOf(limits, 'five_hour') ?? prev?.session ?? null,
82      week: limitOf(limits, 'seven_day') ?? prev?.week ?? null,
83    }
84  })
85}
86
87export const register: Register = on => {
88  on('session.start', async ($, e, next) => {
89    const usage = await $.session.usage()
90    await measure($, usage.context, usage.rateLimits, false)
91
92    return next(e)
93  })
94
95  on('session.measure', async ($, e, next) => {
96    await measure($, e.context, e.rateLimits, e.changed.includes('context'))
97
98    return next(e)
99  })
100
101  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
102    const now = await read($, forecast)
103
104    if (e.props.hasSurvey || now === null) {
105      return next(e)
106    }
107
108    const { Box, Text } = $.ui.resolve(e)
109    const sky = skyFor(now.percent)
110
111    const row = (label: unknown, value: string, note: string) => (
112      <Text>
113        {label}
114        <Text color={BAR_COLOR}>{value === '' ? bar(0) : bar(parseFloat(value))}</Text>
115        <Text> {value.padEnd(VALUE_CELLS)}</Text>
116        <Text dimColor>{note}</Text>
117      </Text>
118    )
119
120    const limitRow = (label: string, limit: Limit | null, withDay: boolean) => {
121      const name = <Text dimColor>{label.padEnd(LABEL_CELLS)}</Text>
122
123      if (limit === null) {
124        return row(name, '', 'no reading')
125      }
126
127      const resets = limit.resetsAt
128        ? `resets ${withDay ? `${DAYS[new Date(limit.resetsAt).getDay()]} ` : ''}${clock(limit.resetsAt)}`
129        : ''
130
131      return row(name, `${Math.round(limit.percent)}%`, resets)
132    }
133
134    const used = `${compact(now.tokens)} / ${compact(now.window)}`
135    const turn = now.delta === null ? '' : `${now.delta < 0 ? '▼ ' : '▲ +'}${compact(now.delta)} last turn`
136
137    return (
138      <Box flexDirection="column">
139        <Text>
140          <Text color={sky.color} bold>
141            {`${sky.icon} ${sky.word}`.padEnd(LABEL_CELLS)}
142          </Text>
143          <Text color={BAR_COLOR}>{bar(now.percent)}</Text>
144          <Text> {used.padEnd(VALUE_CELLS)}</Text>
145          <Text dimColor>{turn}</Text>
146        </Text>
147        {limitRow('Session', now.session, false)}
148        {limitRow('Week', now.week, true)}
149      </Box>
150    )
151  })
152}
153

1export type Limit = { percent: number; resetsAt: string | null }
2
3export type Forecast = {
4  tokens: number
5  window: number
6  percent: number
7  delta: number | null
8  session: Limit | null
9  week: Limit | null
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    'token-garden': { forecast: Forecast | null }
15  }
16}
17