token-weather

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Forecast } from '../types'
5import { bar, chart, compact, growth, percentOf, push, weather } from './forecast'
6
7const forecast = atom({ plugin: 'token-weather', key: 'forecast' } as const, null)
8
9// Record where the context stands after a finished turn. A compaction
10// shrinks the context, so the delta can go negative; that is shown as ▼.
11async function record($: EngineInterface) {
12  const { context } = await $.session.usage()
13  await update($, forecast, (prev): Forecast => {
14    const history = prev?.history ?? []
15    if (context.tokens === undefined) {
16      return { history, tokens: null, window: context.window, delta: null }
17    }
18    const before = history[history.length - 1] ?? 0
19    return {
20      history: push(history, context.tokens),
21      tokens: context.tokens,
22      window: context.window,
23      delta: context.tokens - before,
24    }
25  })
26}
27
28export const register: Register = on => {
29  // Seed from the session's last response so a resumed session isn't blank.
30  on('session.start', async ($, e, next) => {
31    const result = await next(e)
32    await record($)
33    return result
34  })
35
36  on('turn.complete', async ($, e, next) => {
37    const result = await next(e)
38    // A subagent's turn is its own loop; the forecast follows the main one.
39    if (e.agentId === undefined) await record($)
40    return result
41  })
42
43  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
44    // Keep what a plugin beneath drew and stack ours on top.
45    const below = await next(e)
46    const f = await read($, forecast)
47    if (e.props.hasSurvey || f === null || f.tokens === null) return below
48
49    const { Box, Text, Button } = $.ui.resolve(e)
50    const percent = percentOf(f.tokens, f.window)
51    const w = weather(percent)
52    const [full, empty] = bar(percent, 24)
53    const g = f.delta === null ? null : growth(f.delta)
54
55    const band = (
56      <Box flexDirection="row">
57        <Text color={w.color} bold>
58          {w.icon} {w.word}
59        </Text>
60        <Text dimColor> · </Text>
61        <Text bold>{percent}%</Text>
62        <Text> </Text>
63        <Text color={w.color} bold>
64          {full}
65        </Text>
66        <Text color="gray">{empty}</Text>
67        <Text dimColor>
68          {' '}
69          {compact(f.tokens)} / {compact(f.window)}
70        </Text>
71        <Text dimColor> · </Text>
72        <Text color={w.color}>{chart(f.history, f.window)}</Text>
73        {g === null ? null : (
74          <Text dimColor>
75            {' · '}
76            {g.arrow} {g.text}
77          </Text>
78        )}
79        {/* One click compacts, no /compact typed. Hidden mid-turn, when a
80            compaction can't run. Claude terracotta, like Gaurav's Tools. */}
81        {e.props.isWorking ? null : (
82          <Box key="tw-compact-box" marginLeft={2} paddingX={1} backgroundColor="#d97757">
83            <Button
84              key="tw-compact"
85              label="⇣ Compact"
86              plain
87              hover={{ color: '#ffffff', backgroundColor: '#e8906f', bold: true }}
88              onPress={() => $.session.compact()}
89            />
90          </Box>
91        )}
92      </Box>
93    )
94
95    return below ? (
96      <Box flexDirection="column">
97        {band}
98        {below}
99      </Box>
100    ) : (
101      band
102    )
103  })
104}
105

1export const TURNS = 12
2
3const BARS = '▁▂▃▄▅▆▇█'
4
5export type Weather = { icon: string; word: string; color: string }
6
7export function weather(percent: number): Weather {
8  if (percent >= 90) return { icon: '↯', word: 'Compact soon', color: 'red' }
9  if (percent >= 75) return { icon: '☇', word: 'Storm', color: 'magenta' }
10  if (percent >= 50) return { icon: '☂', word: 'Showers', color: 'blue' }
11  if (percent >= 25) return { icon: '☁', word: 'Cloudy', color: 'cyan' }
12  return { icon: '☀', word: 'Clear', color: 'yellow' }
13}
14
15export function percentOf(tokens: number, window: number): number {
16  return window <= 0 ? 0 : Math.max(0, Math.min(100, Math.round((tokens / window) * 100)))
17}
18
19// 134400 -> "134.4k", 200000 -> "200k", 1000000 -> "1M", 850 -> "850".
20export function compact(n: number): string {
21  const abs = Math.abs(n)
22  if (abs >= 1_000_000) return trim(abs / 1_000_000) + 'M'
23  if (abs >= 1_000) return trim(abs / 1_000) + 'k'
24  return String(Math.round(abs))
25}
26
27function trim(x: number): string {
28  return x.toFixed(1).replace(/\.0$/, '')
29}
30
31// One block per turn, scaled to the window so the chart reads the same as the percentage.
32export function chart(history: number[], window: number): string {
33  return history
34    .map(t => {
35      const level = window <= 0 ? 0 : Math.min(1, Math.max(0, t / window))
36      return BARS[Math.min(BARS.length - 1, Math.floor(level * BARS.length))]
37    })
38    .join('')
39}
40
41export function growth(delta: number): { arrow: string; text: string } {
42  const sign = delta < 0 ? '-' : '+'
43  return { arrow: delta < 0 ? '▼' : '▲', text: `${sign}${compact(delta)} last turn` }
44}
45
46export function push(history: number[], tokens: number): number[] {
47  return [...history, tokens].slice(-TURNS)
48}
49
50// Eighths of a cell, so the fill's edge moves smoothly instead of a whole cell at a time.
51const EIGHTHS = ['', '▏', '▎', '▍', '▌', '▋', '▊', '▉']
52
53// A fill gauge for the percentage: solid full-height blocks on a shaded track.
54// Any use at all shows at least an eighth of a cell.
55export function bar(percent: number, width: number): [string, string] {
56  const p = Math.max(0, Math.min(100, percent))
57  const eighths = p === 0 ? 0 : Math.max(1, Math.round((p / 100) * width * 8))
58  const whole = Math.floor(eighths / 8)
59  const part = EIGHTHS[eighths % 8]
60  const used = whole + (part ? 1 : 0)
61  return ['█'.repeat(whole) + part, '░'.repeat(width - used)]
62}
63

1export type Forecast = {
2  /** Input tokens in the context after each of the last turns, oldest first (at most 12). */
3  history: number[]
4  /** Tokens in the context now; null until a response of the live window reports one. */
5  tokens: number | null
6  /** The model's context window, in tokens. */
7  window: number
8  /** What the last turn added (negative after a compaction); null before two readings. */
9  delta: number | null
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    'token-weather': { forecast: Forecast | null }
15  }
16}
17