budget-watch

1import type {
2  EngineInterface,
3  PluginOptions,
4  Register,
5  SessionRateLimit,
6} from 'claude-code';
7
8import type { BudgetWatchResume, BudgetWatchSample } from '../types';
9
10// Policy-free by design: the mod reports facts as `[budget]` lines and the
11// role's skill decides what to do with them. The only action it takes on its
12// own is the resume after a rate-limit stop, and that is a config switch.
13
14const MAIN = 'main';
15const SAMPLE_WINDOW_MS = 60 * 60 * 1000;
16const MIN_SAMPLE_SPAN_MS = 5 * 60 * 1000;
17const RESUME_MARGIN_MS = 60 * 1000;
18const LIMIT_HIT_PERCENT = 99;
19const HOUR_MS = 60 * 60 * 1000;
20
21type Config = {
22  contextStart: number
23  contextStep: number
24  fiveHourBands: number[]
25  sevenDayBands: number[]
26  autoResume: boolean
27  sevenDayResumeHours: number
28  resumePrompt: string
29};
30
31const resumeRef = { plugin: 'budget-watch', key: 'resume' } as const;
32
33// The armed resume timer. A hot reload drops timers with the old environment,
34// and `session.start` re-arms from `$.state`, so a module variable is enough.
35let pending: { cancel: () => void } | undefined;
36
37export const register: Register = (on, options) => {
38  const config = readConfig(options);
39
40  on('session.start', async ($, e, next) => {
41    await $.command.register({
42      name: 'budget',
43      description:
44        'Shows the context and session budget readings and any planned resume; `/budget resume off` cancels the resume.',
45    });
46    const { value: plan } = await $.state.get(resumeRef);
47    if (plan) {
48      arm($, config, plan, await $.clock.now());
49    }
50    return next(e);
51  });
52
53  on('turn.step', async function* ($, e, next) {
54    const result = yield* next(e);
55    const usage = result.usage;
56    if (usage) {
57      const used =
58        usage.input_tokens +
59        usage.cache_read_input_tokens +
60        usage.cache_creation_input_tokens +
61        usage.output_tokens;
62      const { context } = await $.session.usage();
63      if (context.window > 0) {
64        const percent = Math.floor((used / context.window) * 100);
65        await noteContext($, config, e.agentId, percent);
66      }
67    }
68    return result;
69  });
70
71  on('session.measure', async ($, e, next) => {
72    if (e.changed.includes('rateLimits')) {
73      const now = await $.clock.now();
74      for (const limit of e.rateLimits) {
75        await noteLimit($, config, limit, now);
76      }
77    }
78    return next(e);
79  });
80
81  on('turn.complete', async ($, e, next) => {
82    if (e.agentId === undefined && e.reason === 'error' && config.autoResume) {
83      await planResume($, config);
84    }
85    return next(e);
86  });
87
88  on('command.run', { command: 'budget' }, async ($, e) => {
89    if (e.args.trim() === 'resume off') {
90      pending?.cancel();
91      pending = undefined;
92      await $.state.set(resumeRef, null);
93      return { text: 'Planned resume cancelled.' };
94    }
95    return { text: await report($, config) };
96  });
97};
98
99async function noteContext(
100  $: EngineInterface,
101  config: Config,
102  agentId: string | undefined,
103  percent: number,
104): Promise<void> {
105  const id = agentId ?? MAIN;
106  const band = contextBand(percent, config.contextStart, config.contextStep);
107  if (!(await crossedContext($, id, band))) {
108    return;
109  }
110  const who = agentId === undefined ? MAIN : await describeAgent($, agentId);
111  const text = `[budget] context ${percent}% (${who})`;
112  await append($, text, agentId);
113  if (agentId === undefined) {
114    $.ui.toast(text);
115  }
116}
117
118async function noteLimit(
119  $: EngineInterface,
120  config: Config,
121  limit: SessionRateLimit,
122  now: number,
123): Promise<void> {
124  const bands = bandsFor(config, limit.kind);
125  if (bands === null) {
126    return;
127  }
128  const samples = await recordSample($, limit.kind, limit.percentUsed, now);
129  const band = listBand(limit.percentUsed, bands);
130  if (!(await crossedSession($, limit.kind, band))) {
131    return;
132  }
133  const text = describeLimit(limit, samples, now);
134  await appendEverywhere($, text);
135  $.ui.toast(text);
136}
137
138/**
139 * Records the band now standing for a loop's context and says whether it
140 * rose: the one case that earns a line. A drop (compaction) re-arms silently.
141 */
142async function crossedContext(
143  $: EngineInterface,
144  id: string,
145  band: number | null,
146): Promise<boolean> {
147  const { value, version } = await $.state.get({ plugin: 'budget-watch', key: 'context', id });
148  const last = value ?? null;
149  if (band === last) {
150    return false;
151  }
152  await $.state.set({ plugin: 'budget-watch', key: 'context', id }, band, { ifVersion: version });
153  return rose(last, band);
154}
155
156/**
157 * The same for a rate-limit window; a drop (the window reset) re-arms.
158 */
159async function crossedSession(
160  $: EngineInterface,
161  id: string,
162  band: number | null,
163): Promise<boolean> {
164  const { value, version } = await $.state.get({ plugin: 'budget-watch', key: 'session', id });
165  const last = value ?? null;
166  if (band === last) {
167    return false;
168  }
169  await $.state.set({ plugin: 'budget-watch', key: 'session', id }, band, { ifVersion: version });
170  return rose(last, band);
171}
172
173function rose(last: number | null, band: number | null): boolean {
174  return band !== null && (last === null || band > last);
175}
176
177async function recordSample(
178  $: EngineInterface,
179  kind: string,
180  percent: number,
181  now: number,
182): Promise<BudgetWatchSample[]> {
183  const { value } = await $.state.get({ plugin: 'budget-watch', key: 'samples', id: kind });
184  const previous = value ?? [];
185  const last = previous[previous.length - 1];
186  const hasReset = last !== undefined && percent < last.percent;
187  const kept = hasReset ? [] : previous.filter(sample => now - sample.at <= SAMPLE_WINDOW_MS);
188  const samples = [...kept, { at: now, percent }];
189  await $.state.set({ plugin: 'budget-watch', key: 'samples', id: kind }, samples);
190  return samples;
191}
192
193/**
194 * Appends the line to a loop's conversation as a user-role row the model
195 * reads, and notes the same line and its destination in the debug log.
196 */
197async function append(
198  $: EngineInterface,
199  text: string,
200  agentId?: string,
201): Promise<void> {
202  $.ui.log(`${text} -> ${agentId ?? MAIN}`, { to: 'debug' });
203  try {
204    await $.session.append({
205      message: { type: 'user', content: [{ type: 'text', text }] },
206      ...(agentId === undefined ? {} : { agentId }),
207    });
208  }
209  catch {
210    // The loop ended between the reading and the append; nothing to tell it.
211  }
212}
213
214async function appendEverywhere($: EngineInterface, text: string): Promise<void> {
215  await append($, text);
216  const agents = await listAgents($);
217  for (const agent of agents) {
218    if (agent.status === 'running') {
219      await append($, text, agent.id);
220    }
221  }
222}
223
224async function listAgents($: EngineInterface) {
225  try {
226    return await $.agent.list();
227  }
228  catch {
229    return [];
230  }
231}
232
233async function describeAgent($: EngineInterface, agentId: string): Promise<string> {
234  const agent = (await listAgents($)).find(candidate => candidate.id === agentId);
235  const short = agentId.slice(0, 8);
236  return agent === undefined ? `agent ${short}` : `agent ${short}, ${agent.type}`;
237}
238
239async function report($: EngineInterface, config: Config): Promise<string> {
240  const { context, rateLimits } = await $.session.usage();
241  const now = await $.clock.now();
242  const lines: string[] = [];
243  const percent = context.percent === undefined ? 'unknown' : `${context.percent}%`;
244  lines.push(
245    `context: ${percent} of ${context.window} tokens (bands from ${config.contextStart}% every ${config.contextStep}%)`,
246  );
247  if (rateLimits.length === 0) {
248    lines.push('session: no rate-limit reading yet');
249  }
250  for (const limit of rateLimits) {
251    const bands = bandsFor(config, limit.kind);
252    const bandText = bands === null ? 'not watched' : `bands ${bands.join(', ')}`;
253    lines.push(`session: ${describeLimit(limit, [], now).replace('[budget] session ', '')} (${bandText})`);
254  }
255  const { value: plan } = await $.state.get(resumeRef);
256  lines.push(
257    plan
258      ? `resume: planned at ${formatTime(plan.at, now)} after the ${plan.windows} limit`
259      : `resume: none planned (auto-resume ${config.autoResume ? 'on' : 'off'})`,
260  );
261  return lines.join('\n');
262}
263
264function describeLimit(limit: SessionRateLimit, samples: BudgetWatchSample[], now: number): string {
265  const parts = [`[budget] session ${labelFor(limit.kind)} ${limit.percentUsed}%`];
266  const resetsAt = limit.resetsAt === undefined ? NaN : Date.parse(limit.resetsAt);
267  if (!Number.isNaN(resetsAt)) {
268    parts.push(`resets ${formatTime(resetsAt, now)}`);
269  }
270  const remaining = minutesRemaining(samples, limit.percentUsed);
271  if (remaining !== null) {
272    parts.push(`about ${formatDuration(remaining)} at current rate`);
273  }
274  return parts.join(', ');
275}
276
277/**
278 * Minutes until the window fills at the rate the samples show, or null while
279 * the samples are too few or too close together to say.
280 */
281function minutesRemaining(samples: BudgetWatchSample[], percent: number): number | null {
282  const first = samples[0];
283  const last = samples[samples.length - 1];
284  if (first === undefined || last === undefined) {
285    return null;
286  }
287  const spanMs = last.at - first.at;
288  const rise = last.percent - first.percent;
289  if (spanMs < MIN_SAMPLE_SPAN_MS || rise <= 0) {
290    return null;
291  }
292  const percentPerMinute = rise / (spanMs / 60000);
293  return (100 - percent) / percentPerMinute;
294}
295
296function readConfig(options: PluginOptions): Config {
297  return {
298    contextStart: numberOption(options.contextStart, 40),
299    contextStep: Math.max(1, numberOption(options.contextStep, 5)),
300    fiveHourBands: bandList(options.fiveHourBands, [70, 80, 90, 95]),
301    sevenDayBands: bandList(options.sevenDayBands, [90, 95]),
302    autoResume: options.autoResume !== false,
303    sevenDayResumeHours: numberOption(options.sevenDayResumeHours, 3),
304    resumePrompt:
305      typeof options.resumePrompt === 'string' && options.resumePrompt.trim() !== ''
306        ? options.resumePrompt
307        : 'The session rate limit has reset. Resume the work in progress from the handoff on the active bead. If there is no handoff, report the current state and stop.',
308  };
309}
310
311function numberOption(value: unknown, fallback: number): number {
312  return typeof value === 'number' && Number.isFinite(value) ? value : fallback;
313}
314
315function bandList(value: unknown, fallback: number[]): number[] {
316  if (typeof value !== 'string') {
317    return fallback;
318  }
319  const bands = value
320    .split(',')
321    .map(part => Number(part.trim()))
322    .filter(n => Number.isFinite(n) && n > 0 && n <= 100)
323    .sort((a, b) => a - b);
324  return bands.length === 0 ? fallback : bands;
325}
326
327function bandsFor(config: Config, kind: string): number[] | null {
328  if (kind === 'five_hour') {
329    return config.fiveHourBands;
330  }
331  if (kind === 'seven_day') {
332    return config.sevenDayBands;
333  }
334  return null;
335}
336
337function contextBand(percent: number, start: number, step: number): number | null {
338  if (percent < start) {
339    return null;
340  }
341  return start + Math.floor((percent - start) / step) * step;
342}
343
344function listBand(percent: number, bands: number[]): number | null {
345  let band: number | null = null;
346  for (const candidate of bands) {
347    if (percent >= candidate) {
348      band = candidate;
349    }
350  }
351  return band;
352}
353
354function labelFor(kind: string): string {
355  if (kind === 'five_hour') {
356    return 'five-hour';
357  }
358  if (kind === 'seven_day') {
359    return 'seven-day';
360  }
361  return kind.replace(/_/g, '-');
362}
363
364function formatTime(ms: number, now: number): string {
365  const date = new Date(ms);
366  try {
367    const sameDay = date.toDateString() === new Date(now).toDateString();
368    return date.toLocaleTimeString(undefined, {
369      hour: '2-digit',
370      minute: '2-digit',
371      ...(sameDay ? {} : { weekday: 'short' }),
372    });
373  }
374  catch {
375    return date.toISOString().slice(11, 16);
376  }
377}
378
379function formatDuration(minutes: number): string {
380  if (minutes < 90) {
381    return `${Math.max(1, Math.round(minutes))} min`;
382  }
383  const hours = minutes / 60;
384  return `${hours < 10 ? hours.toFixed(1) : Math.round(hours)} h`;
385}
386
387function arm(
388  $: EngineInterface,
389  cfg: Config,
390  plan: BudgetWatchResume,
391  now: number,
392): void {
393  pending?.cancel();
394  pending = $.clock.after(Math.max(plan.at - now, 0), () => {
395    pending = undefined;
396    void fire($, cfg);
397  });
398  $.ui.toast(`[budget] session limit hit; resuming at ${formatTime(plan.at, now)}`);
399}
400
401async function fire($: EngineInterface, cfg: Config): Promise<void> {
402  const { value: plan } = await $.state.get(resumeRef);
403  if (!plan) {
404    return;
405  }
406  await $.state.set(resumeRef, null);
407  await $.prompt.submit({ text: cfg.resumePrompt });
408}
409
410async function planResume($: EngineInterface, cfg: Config): Promise<void> {
411  const { rateLimits } = await $.session.usage();
412  const now = await $.clock.now();
413  const hit = rateLimits.filter(limit => limit.percentUsed >= LIMIT_HIT_PERCENT);
414  if (hit.length === 0) {
415    return;
416  }
417  let at: number | null = null;
418  for (const limit of hit) {
419    const resetsAt = limit.resetsAt === undefined ? NaN : Date.parse(limit.resetsAt);
420    if (limit.kind !== 'five_hour' && limit.kind !== 'seven_day') {
421      $.ui.toast(`[budget] ${labelFor(limit.kind)} limit hit; not resuming automatically`);
422      return;
423    }
424    if (Number.isNaN(resetsAt)) {
425      continue;
426    }
427    if (limit.kind === 'seven_day' && resetsAt - now > cfg.sevenDayResumeHours * HOUR_MS) {
428      $.ui.toast(
429        `[budget] seven-day limit hit, resets ${formatTime(resetsAt, now)}; not resuming automatically`,
430      );
431      return;
432    }
433    at = Math.max(at ?? 0, resetsAt);
434  }
435  if (at === null) {
436    return;
437  }
438  const plan: BudgetWatchResume = {
439    at: at + RESUME_MARGIN_MS,
440    windows: hit.map(limit => labelFor(limit.kind)).join('+'),
441  };
442  await $.state.set(resumeRef, plan);
443  arm($, cfg, plan, now);
444}
445

1export type BudgetWatchSample = { at: number; percent: number };
2export type BudgetWatchResume = { at: number; windows: string };
3
4declare module 'claude-code' {
5  interface PluginState {
6    'budget-watch': {
7      /** The last context band appended per loop: 'main' or an agent id. */
8      context: StateFamily<number | null>
9      /** The last session band appended per rate-limit window kind. */
10      session: StateFamily<number | null>
11      /** Recent readings per window kind, for the burn-rate estimate. */
12      samples: StateFamily<BudgetWatchSample[]>
13      /** The resume the mod has planned, or null. */
14      resume: BudgetWatchResume | null
15    }
16  }
17}
18