context-meter

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Reading } from '../types'
5import { budgetRow, budgetToast, overBudget, turnCost } from './cost'
6
7// Every turn re-sends the whole context, so absolute size drives cost more than percent of the window.
8const WARN_TOKENS = 100_000
9const ALERT_TOKENS = 200_000
10
11// Warn when a single turn costs at least this much (USD). Warning only, never blocks.
12const TURN_WARN_USD = 0.5
13// Off by default. Set a number to warn once when the session total reaches it.
14const SESSION_BUDGET_USD: number | undefined = undefined
15
16const PANE = 'context-meter'
17const TITLE = 'Context'
18
19const reading = atom({ plugin: 'context-meter', key: 'reading' } as const, null)
20const hasWarned = atom({ plugin: 'context-meter', key: 'hasWarned' } as const, false)
21const lastUsd = atom({ plugin: 'context-meter', key: 'lastUsd' } as const, undefined as number | undefined)
22const turnWarned = atom({ plugin: 'context-meter', key: 'turnWarned' } as const, false)
23const budgetWarned = atom({ plugin: 'context-meter', key: 'budgetWarned' } as const, false)
24
25const k = (n: number) => (n >= 1_000_000 ? `${+(n / 1_000_000).toFixed(1)}M` : `${Math.round(n / 1000)}k`)
26
27const level = (tokens: number) =>
28  tokens >= ALERT_TOKENS
29    ? ({ color: 'error', status: 'Compact now' } as const)
30    : tokens >= WARN_TOKENS
31      ? ({ color: 'warning', status: 'Growing' } as const)
32      : ({ color: 'success', status: 'Healthy' } as const)
33
34export const register: Register = on => {
35  on('session.start', async ($, e, next) => {
36    await $.command.register({ name: 'context-meter', description: 'Show the context meter pane' })
37    void $.ui.open({ id: PANE, title: TITLE })
38
39    return next(e)
40  })
41
42  on('command.run', { command: 'context-meter' }, async $ => {
43    await $.ui.open({ id: PANE, title: TITLE, focus: true })
44
45    return { text: 'Context meter opened.' }
46  })
47
48  on('session.measure', async ($, e, next) => {
49    const { tokens, window, percent } = e.context
50    const usd = e.cost?.usd
51    const turn = usd === undefined ? undefined : turnCost(await read($, lastUsd), usd)
52    const latest: Reading | null =
53      tokens === undefined ? null : { tokens, window, percent: percent ?? 0, usd, turnUsd: turn }
54    await update($, reading, () => latest)
55
56    // Toast once per crossing; compacting or clearing below the line re-arms it.
57    const isOver = latest !== null && latest.tokens >= ALERT_TOKENS
58    if (isOver && !(await read($, hasWarned))) {
59      $.ui.toast(`Context is ${k(latest.tokens)} and re-sent every turn: consider /compact or /clear`)
60    }
61    await update($, hasWarned, () => isOver)
62
63    if (usd !== undefined) {
64      // A reading without cost leaves the baseline alone, so the next turn is measured from the last known total.
65      await update($, lastUsd, () => usd)
66
67      // Dropping back under the line re-arms the turn warning.
68      const isTurnOver = turn !== undefined && turn >= TURN_WARN_USD
69      if (isTurnOver && !(await read($, turnWarned))) {
70        $.ui.toast(`That turn cost $${turn.toFixed(2)}`)
71      }
72      await update($, turnWarned, () => isTurnOver)
73
74      // The budget warning stays set for the rest of the session.
75      if (overBudget(usd, SESSION_BUDGET_USD) && !(await read($, budgetWarned))) {
76        $.ui.toast(budgetToast(usd, SESSION_BUDGET_USD!))
77        await update($, budgetWarned, () => true)
78      }
79    }
80
81    return next(e)
82  })
83
84  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
85    const { Box, Text } = $.ui.resolve(e)
86    const r = await read($, reading)
87
88    if (r === null) {
89      return (
90        <Box paddingX={1}>
91          <Text dimColor>Waiting for the first response…</Text>
92        </Box>
93      )
94    }
95
96    const { color, status } = level(r.tokens)
97    const cells = Math.max(10, e.props.bodyColumns - 2)
98    const filled = Math.min(cells, Math.ceil((r.tokens / r.window) * cells))
99    const row = (label: string, value: string, valueColor?: typeof color) => (
100      <Box flexDirection="row" justifyContent="space-between">
101        <Text dimColor>{label}</Text>
102        <Text color={valueColor} bold={valueColor !== undefined}>{value}</Text>
103      </Box>
104    )
105
106    return (
107      <Box flexDirection="column" paddingX={1}>
108        <Text bold color={color}>{`${r.percent}%`}</Text>
109        <Box flexDirection="row">
110          <Text color={color}>{'█'.repeat(filled)}</Text>
111          <Text dimColor>{'░'.repeat(cells - filled)}</Text>
112        </Box>
113        <Text dimColor>{`${k(r.tokens)} of ${k(r.window)} tokens`}</Text>
114        <Box marginTop={1} flexDirection="column">
115          {row('Status', status, color)}
116          {row('Re-sent each turn', k(r.tokens))}
117          {r.usd !== undefined && row('Session cost', `$${r.usd.toFixed(2)}`)}
118          {r.turnUsd !== undefined && (
119            <>
120              {row('Last turn', `$${r.turnUsd.toFixed(2)}`, r.turnUsd >= TURN_WARN_USD ? 'warning' : undefined)}
121              {row('Next turn ≈', `$${r.turnUsd.toFixed(2)}`)}
122            </>
123          )}
124          {r.usd !== undefined && budgetRow(r.usd, SESSION_BUDGET_USD) && row('Budget', budgetRow(r.usd, SESSION_BUDGET_USD)!)}
125        </Box>
126        {r.tokens >= ALERT_TOKENS && (
127          <Box marginTop={1}>
128            <Text color="error">Run /compact or /clear</Text>
129          </Box>
130        )}
131      </Box>
132    )
133  })
134}
135

1// Pure cost arithmetic for the cost guardian. No engine imports, so it runs under plain tests.
2
3// Increase since the last cumulative reading. Undefined when there is no baseline yet,
4// or when the total went down (a new session, or a clear that resets cost).
5export const turnCost = (lastUsd: number | undefined, usd: number): number | undefined => {
6  if (lastUsd === undefined || usd < lastUsd) return undefined
7  return usd - lastUsd
8}
9
10// No budget set means never over.
11export const overBudget = (usd: number, budget: number | undefined): boolean =>
12  budget !== undefined && usd >= budget
13
14// Pane row value, or undefined when no budget is set.
15export const budgetRow = (usd: number, budget: number | undefined): string | undefined =>
16  budget === undefined ? undefined : `$${usd.toFixed(2)} of $${budget.toFixed(2)}`
17
18export const budgetToast = (usd: number, budget: number): string =>
19  `Session cost is $${usd.toFixed(2)}, over your $${budget.toFixed(2)} budget`
20

1export type Reading = { tokens: number; window: number; percent: number; usd?: number; turnUsd?: number }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'context-meter': {
6      reading: Reading | null
7      hasWarned: boolean
8      lastUsd: number | undefined
9      turnWarned: boolean
10      budgetWarned: boolean
11    }
12  }
13}
14