subagent-budget

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { TurnCount } from '../types'
5import { DEFAULT_MAX, FRESH, bump, dayKey, peak, summary, toastText } from './budget'
6
7const turn = atom({ plugin: 'subagent-budget', key: 'turn' } as const, FRESH)
8
9async function isEnabled($: EngineInterface): Promise<boolean> {
10  return (await $.store.get('enabled')) !== false
11}
12
13// Today's peak: what the store holds for today, or the live count when higher.
14async function todayMax($: EngineInterface, live: number): Promise<number> {
15  return peak(await $.store.get('today'), dayKey(await $.clock.now()), live).max
16}
17
18// Records the turn's count into today's peak. Called from sequential points only
19// (prompt.submit, main turn.complete), never from concurrent spawns.
20async function savePeak($: EngineInterface): Promise<void> {
21  const current: TurnCount = await read($, turn)
22  await $.store.set('today', peak(await $.store.get('today'), dayKey(await $.clock.now()), current.count))
23}
24
25async function counted($: EngineInterface, max: number): Promise<void> {
26  const enabled = await isEnabled($)
27  let toast = false
28  let count = 0
29  await update($, turn, (t: TurnCount) => {
30    const out = bump(t, max, enabled)
31    toast = out.toast
32    count = out.turn.count
33    return out.turn
34  })
35  if (toast) $.ui.toast(toastText(count, max))
36}
37
38export const register: Register = (on, options) => {
39  const max = typeof options?.max === 'number' ? options.max : DEFAULT_MAX
40
41  on('session.start', async ($, e, next) => {
42    await $.command.register({
43      name: 'subagent-budget',
44      description: 'Subagents spawned this turn, today\'s peak and the budget (on | off toggles the alert)',
45    })
46
47    return next(e)
48  })
49
50  // A prompt that starts a turn resets the count; one delivered into a running turn (turnId set) does not.
51  on('prompt.submit', async ($, e, next) => {
52    if (e.turnId === undefined) {
53      await savePeak($)
54      await update($, turn, () => FRESH)
55    }
56
57    return next(e)
58  })
59
60  // Every started spawn counts: main loop, nested (spawned by a subagent), teammates, plugins' own.
61  on('agent.spawn', async ($, e, next) => {
62    const started = await next(e)
63    if (started.deny === undefined && started.agentId !== undefined) await counted($, max)
64
65    return started
66  })
67
68  on('turn.complete', async ($, e, next) => {
69    if (e.agentId === undefined) await savePeak($)
70
71    return next(e)
72  })
73
74  on('command.run', { command: 'subagent-budget' }, async ($, e) => {
75    const arg = e.args.trim()
76    if (arg === 'on' || arg === 'off') {
77      await $.store.set('enabled', arg === 'on')
78      return { text: `subagent-budget alerts ${arg}.` }
79    }
80    const current: TurnCount = await read($, turn)
81
82    return { text: summary(current.count, await todayMax($, current.count), max, await isEnabled($)) }
83  })
84}
85

1// Pure logic: per-turn counting, the toast decision, today's peak and the summary.
2import type { TurnCount } from '../types'
3
4export const DEFAULT_MAX = 10
5export const FRESH: TurnCount = { count: 0, toasted: false }
6
7export type Today = { day: string; max: number }
8
9// One more spawn this turn; toasts once, the first time the count goes over the budget.
10export function bump(turn: TurnCount, max: number, enabled: boolean): { turn: TurnCount; toast: boolean } {
11  const count = turn.count + 1
12  const toast = enabled && count > max && !turn.toasted
13
14  return { turn: { count, toasted: turn.toasted || toast }, toast }
15}
16
17export const toastText = (count: number, max: number): string =>
18  `subagent-budget: ${count} subagents this turn (budget ${max})`
19
20// Local calendar day of a clock reading, YYYY-MM-DD.
21export function dayKey(ms: number): string {
22  const d = new Date(ms)
23  const pad = (n: number) => String(n).padStart(2, '0')
24
25  return `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())}`
26}
27
28const isToday = (v: unknown): v is Today =>
29  typeof v === 'object' && v !== null && typeof (v as Today).day === 'string' && typeof (v as Today).max === 'number'
30
31// Today's peak per-turn count, given what the store held (another day's record starts over).
32export function peak(stored: unknown, day: string, count: number): Today {
33  const prev = isToday(stored) && stored.day === day ? stored.max : 0
34
35  return { day, max: Math.max(prev, count) }
36}
37
38export function summary(count: number, today: number, max: number, enabled: boolean): string {
39  return [
40    `Subagents this turn: ${count}`,
41    `Max in one turn today: ${today}`,
42    `Budget: ${max} per turn${enabled ? '' : ' (alerts off)'}`,
43  ].join('\n')
44}
45

1export type TurnCount = {
2  /** Subagents started since the last prompt that began a turn. */
3  count: number
4  /** Whether this turn already toasted. */
5  toasted: boolean
6}
7
8declare module 'claude-code' {
9  interface PluginState {
10    'subagent-budget': { turn: TurnCount }
11  }
12}
13