context-meter

1import { atom, read, update } from 'claude-code';
2import type { Register } from 'claude-code';
3
4import type { MeterAction, Totals } from '../types';
5import {
6  EMPTY_TOTALS,
7  IDLE_HANDOFF,
8  LEVEL_COLOR,
9  addUsage,
10  bar,
11  costWeight,
12  formatTokens,
13  handoffStep,
14  levelFor,
15  percentOfWindow,
16  resolveThresholds,
17} from './meter.ts';
18
19const totalsAtom = atom({ plugin: 'context-meter', key: 'totals' } as const, EMPTY_TOTALS);
20const confirmAtom = atom({ plugin: 'context-meter', key: 'confirm' } as const, null);
21const handoffAtom = atom({ plugin: 'context-meter', key: 'handoff' } as const, IDLE_HANDOFF);
22const noteAtom = atom({ plugin: 'context-meter', key: 'note' } as const, null);
23
24const HANDOFF_PROMPT =
25  'Write your dated handoff memory file now (state, open work, next steps; index it in MEMORY.md), ' +
26  'then reply "handoff written". Do no other work; the session will be cleared after this reply.';
27
28const LABELS: Record<MeterAction, string> = {
29  compact: '/compact',
30  clear: '/clear',
31  handoff: 'hand off, then /clear',
32};
33
34export const register: Register = (on, options) => {
35  const thresholds = resolveThresholds(options as Record<string, unknown>);
36
37  on('turn.complete', async ($, e, next) => {
38    if (e.usage) {
39      await update($, totalsAtom, (t: Totals) => addUsage(t, e.usage!));
40    }
41
42    // A pending confirmation does not outlive the turn it was asked in.
43    await update($, confirmAtom, () => null);
44
45    const step = handoffStep(
46      await read($, handoffAtom),
47      { kind: 'turn.complete', turnId: e.turnId, reason: e.reason, agentId: e.agentId },
48      HANDOFF_PROMPT,
49    );
50    await update($, handoffAtom, () => step.state);
51    if (step.clear) {
52      void $.command.run({ command: 'clear' });
53    }
54
55    return next(e);
56  });
57
58  // Only the handoff prompt's own turn arms the clear; match it by its text.
59  on('turn.start', async ($, e, next) => {
60    const step = handoffStep(
61      await read($, handoffAtom),
62      { kind: 'turn.start', turnId: e.turnId, text: e.text },
63      HANDOFF_PROMPT,
64    );
65    await update($, handoffAtom, () => step.state);
66    return next(e);
67  });
68
69  // /clear starts a fresh session: totals count from there.
70  on('session.end', async ($, e, next) => {
71    await update($, confirmAtom, () => null);
72    await update($, handoffAtom, () => IDLE_HANDOFF);
73    await update($, noteAtom, () => null);
74    if (e.reason === 'clear') {
75      await update($, totalsAtom, () => EMPTY_TOTALS);
76    }
77
78    return next(e);
79  });
80
81  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
82    if (e.props.hasSurvey) {
83      return next(e);
84    }
85
86    // Degrade to no output, never an error, when the host exposes no metrics.
87    let tokens: number | undefined;
88    let window: number | undefined;
89    let usd: number | undefined;
90    try {
91      const usage = await $.session.usage();
92      tokens = usage.context.tokens;
93      window = usage.context.window;
94      usd = usage.cost?.usd;
95    } catch {
96      return next(e);
97    }
98
99    const percent = percentOfWindow(tokens, window);
100    const totals = await read($, totalsAtom);
101    if (percent === null || tokens === undefined) {
102      return next(e);
103    }
104
105    const confirm = await read($, confirmAtom);
106    const note = await read($, noteAtom);
107    const level = levelFor(percent, thresholds);
108    const color = LEVEL_COLOR[level];
109    const { Box, Button, Text } = $.ui.resolve(e);
110    const weight = costWeight(totals);
111
112    return (
113      <Box flexDirection="column">
114        <Box>
115          <Text color={color}>
116            ctx {formatTokens(tokens)} / {formatTokens(window ?? 0)} ({percent.toFixed(1)}%){' '}
117            {bar(percent, thresholds)}
118            {level === 'hot' ? ' hand off soon' : ''}
119            {level === 'red' ? ' clear now' : ''}
120          </Text>
121        </Box>
122        <Box>
123          <Text dimColor>
124            session in {formatTokens(totals.input)} / cache-write {formatTokens(totals.cacheWrite)}{' '}
125            / cache-read {formatTokens(totals.cacheRead)} / out {formatTokens(totals.output)} ·
126            weight ~{formatTokens(weight)} input-eq
127            {usd === undefined ? '' : ` · $${usd.toFixed(2)}`}{' '}
128          </Text>
129        </Box>
130        {note === null ? null : (
131          <Box>
132            <Text color="red">{note}</Text>
133          </Box>
134        )}
135        {confirm === null ? (
136          <Box>
137            <Button
138              key="compact"
139              label="Compact"
140              onPress={() => update($, confirmAtom, () => 'compact')}
141            />
142            <Button
143              key="clear"
144              label="Clear"
145              onPress={() => update($, confirmAtom, () => 'clear')}
146            />
147            <Button
148              key="handoff"
149              label="Hand off + clear"
150              onPress={() => update($, confirmAtom, () => 'handoff')}
151            />
152          </Box>
153        ) : (
154          <Box>
155            <Text color="yellow">Run {LABELS[confirm]}? </Text>
156            <Button
157              key="yes"
158              label="Yes"
159              onPress={async () => {
160                await update($, confirmAtom, () => null);
161                await update($, noteAtom, () => null);
162                try {
163                  if (confirm === 'compact') {
164                    await $.session.compact();
165                  } else if (confirm === 'clear') {
166                    await $.command.run({ command: 'clear' });
167                  } else {
168                    const armed = handoffStep(IDLE_HANDOFF, { kind: 'submit' }, HANDOFF_PROMPT);
169                    await update($, handoffAtom, () => armed.state);
170                    try {
171                      await $.prompt.submit({ text: HANDOFF_PROMPT });
172                    } catch (err) {
173                      await update($, handoffAtom, () => IDLE_HANDOFF);
174                      throw err;
175                    }
176                  }
177                } catch (err) {
178                  const msg = err instanceof Error ? err.message : String(err);
179                  await update($, noteAtom, () => `${LABELS[confirm]} failed: ${msg}`);
180                }
181              }}
182            />
183            <Button key="no" label="Cancel" onPress={() => update($, confirmAtom, () => null)} />
184          </Box>
185        )}
186      </Box>
187    );
188  });
189};
190

1/**
2 * Pure logic for the context-meter mod (AISDLC-743): thresholds, colour
3 * mapping, token totals and cost weight. No host imports, so it runs under
4 * plain `node --test`.
5 */
6
7export type Thresholds = {
8  /** Percent of the window where the bar turns amber. */
9  amber: number;
10  /** Percent where the band says to hand off soon. */
11  hot: number;
12  /** Percent where the bar turns red (the clearing level). */
13  red: number;
14};
15
16/** The one place the defaults live: 10 / 13 / 15 percent of the window. */
17export const DEFAULT_THRESHOLDS: Thresholds = { amber: 10, hot: 13, red: 15 };
18
19export type Level = 'ok' | 'amber' | 'hot' | 'red';
20
21export const LEVEL_COLOR: Record<Level, string> = {
22  ok: 'green',
23  amber: 'yellow',
24  hot: 'yellow',
25  red: 'red',
26};
27
28export type Totals = { input: number; cacheWrite: number; cacheRead: number; output: number };
29
30export type UsageLike = {
31  input_tokens?: number;
32  cache_creation_input_tokens?: number;
33  cache_read_input_tokens?: number;
34  output_tokens?: number;
35};
36
37export const EMPTY_TOTALS: Totals = { input: 0, cacheWrite: 0, cacheRead: 0, output: 0 };
38
39/**
40 * Read the three thresholds from userConfig options. All-or-nothing: an
41 * invalid field falls back to its own default, but if the resulting trio is
42 * not ordered amber <= hot <= red, every override is discarded.
43 */
44export function resolveThresholds(options: Record<string, unknown> | undefined): Thresholds {
45  const num = (v: unknown, d: number) =>
46    typeof v === 'number' && Number.isFinite(v) && v > 0 ? v : d;
47  const t: Thresholds = {
48    amber: num(options?.amberAtPercent, DEFAULT_THRESHOLDS.amber),
49    hot: num(options?.hotAtPercent, DEFAULT_THRESHOLDS.hot),
50    red: num(options?.redAtPercent, DEFAULT_THRESHOLDS.red),
51  };
52  return t.amber <= t.hot && t.hot <= t.red ? t : { ...DEFAULT_THRESHOLDS };
53}
54
55/**
56 * "Hand off, then clear" state. `armed`: the handoff prompt was submitted and
57 * its own turn has not started yet. `running`: that turn started with `turnId`.
58 */
59export type HandoffState =
60  | { phase: 'idle' }
61  | { phase: 'armed' }
62  | { phase: 'running'; turnId: string };
63
64export type HandoffEvent =
65  | { kind: 'submit' }
66  | { kind: 'turn.start'; turnId: string; text: string }
67  | { kind: 'turn.complete'; turnId: string; reason: string; agentId?: string }
68  | { kind: 'reset' };
69
70export const IDLE_HANDOFF: HandoffState = { phase: 'idle' };
71
72/**
73 * Pure transition for the hand-off flow. The clear fires only when the
74 * handoff prompt's own turn ends with an answer: an earlier in-flight turn
75 * completing while armed is ignored, and a handoff turn that ends any other
76 * way (aborted, refusal, error) disarms without clearing.
77 */
78export function handoffStep(
79  state: HandoffState,
80  event: HandoffEvent,
81  handoffText: string,
82): { state: HandoffState; clear: boolean } {
83  switch (event.kind) {
84    case 'submit':
85      return { state: { phase: 'armed' }, clear: false };
86    case 'reset':
87      return { state: IDLE_HANDOFF, clear: false };
88    case 'turn.start':
89      if (state.phase === 'armed' && event.text === handoffText) {
90        return { state: { phase: 'running', turnId: event.turnId }, clear: false };
91      }
92      return { state, clear: false };
93    case 'turn.complete':
94      if (
95        state.phase !== 'running' ||
96        event.agentId !== undefined ||
97        event.turnId !== state.turnId
98      ) {
99        return { state, clear: false };
100      }
101      return { state: IDLE_HANDOFF, clear: event.reason === 'answer' };
102  }
103}
104
105/** Share of the window in use, as a percentage; null when it cannot be computed. */
106export function percentOfWindow(tokens: number | undefined, window: number | undefined) {
107  if (tokens === undefined || !window || window <= 0) return null;
108  return (tokens / window) * 100;
109}
110
111export function levelFor(percent: number, t: Thresholds = DEFAULT_THRESHOLDS): Level {
112  if (percent >= t.red) return 'red';
113  if (percent >= t.hot) return 'hot';
114  if (percent >= t.amber) return 'amber';
115  return 'ok';
116}
117
118/** Add one response's usage into running totals (missing counts read as zero). */
119export function addUsage(totals: Totals, usage: UsageLike): Totals {
120  return {
121    input: totals.input + (usage.input_tokens ?? 0),
122    cacheWrite: totals.cacheWrite + (usage.cache_creation_input_tokens ?? 0),
123    cacheRead: totals.cacheRead + (usage.cache_read_input_tokens ?? 0),
124    output: totals.output + (usage.output_tokens ?? 0),
125  };
126}
127
128/**
129 * Approximate cost weight in "input-token equivalents", using list-price
130 * ratios against uncached input: cache write 1.25x, cache read 0.1x, output 5x.
131 * A relative weight, not dollars.
132 */
133export function costWeight(totals: Totals): number {
134  return totals.input + totals.cacheWrite * 1.25 + totals.cacheRead * 0.1 + totals.output * 5;
135}
136
137export function formatTokens(n: number): string {
138  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(2)}M`;
139  if (n >= 1_000) return `${(n / 1_000).toFixed(n >= 100_000 ? 0 : 1)}k`;
140  return String(Math.round(n));
141}
142
143/** Text progress bar, `width` cells, filled up to the share of the red threshold. */
144export function bar(percent: number, t: Thresholds = DEFAULT_THRESHOLDS, width = 20): string {
145  const filled = Math.min(width, Math.max(0, Math.round((percent / t.red) * width)));
146  return '█'.repeat(filled) + '░'.repeat(width - filled);
147}
148

1export type Totals = {
2  input: number;
3  cacheWrite: number;
4  cacheRead: number;
5  output: number;
6};
7
8export type HandoffState =
9  | { phase: 'idle' }
10  | { phase: 'armed' }
11  | { phase: 'running'; turnId: string };
12
13export type MeterAction = 'compact' | 'clear' | 'handoff';
14
15declare module 'claude-code' {
16  interface PluginState {
17    'context-meter': {
18      totals: Totals;
19      confirm: MeterAction | null;
20      handoff: HandoffState;
21      note: string | null;
22    };
23  }
24}
25