agent-governor

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Gov, Settings } from '../types'
5import { bandRow, crossings, empty, finished, parseCommand, reconcile, running, settingsFrom, spawned, statusText, stepped, tokensOf } from './govern'
6
7// Watches the session's subagents: which run, how many at once, and how many
8// tokens each and all of them used. A [GOV] row above the prompt while any
9// runs, and a toast once each time a limit is crossed. Warn-only by default:
10// `/governor caps on` asks before a spawn past the concurrency limit and
11// denies the tool calls of an agent past its budget; `/governor route` picks
12// a model per agent type when the spawn named none. No network, no
13// processes, no model calls.
14//
15// Running vs finished: an agent is running from its spawn (agent.spawn's
16// agentId) and on every model step of its loop (turn.step's agentId); its own
17// turn.complete ends it. `$.agent.list()` is the authority behind that: asked
18// a few seconds after tool activity while anything runs, its statuses correct
19// what the events missed (a killed agent, one spawned before this loaded).
20
21const STORE_SETTINGS = 'settings'
22const RECONCILE_MS = 3000
23const OVER_BUDGET = 'agent-governor: this agent used its token budget; stop and hand back your findings now.'
24const gov = atom({ plugin: 'agent-governor', key: 'gov' } as const, empty())
25
26// Mirrors for the paths that must not read `$` (a .catch handler) or should
27// not read it on every call (tool.call). A reload clears them; `ensure` and
28// the next change fill them again.
29let isLoaded = false
30let isCapsOn = false
31let isLive = false
32let isQueued = false
33// Spawns past the check whose agent is not recorded yet (recording waits for
34// `next`): spawns made at once each read the state before any records its
35// agent, so these count toward the limit too.
36let inFlight = 0
37
38// An agent over its budget may still wind down: hand its report back, message,
39// stop a task. Denying these loses the report and costs a request per retry.
40const WIND_DOWN = new Set(['SubagentHandback', 'SendMessage', 'TaskStop'])
41
42const tooMany = (n: number, max: number) =>
43  `agent-governor: ${n} subagents are already running (limit ${max}). Wait for the running agents to finish and hand back their results, then start this one if it is still needed.`
44
45/** The governor's state with the stored settings in it, read from the store once per load. */
46async function ensure($: EngineInterface): Promise<Gov> {
47  if (!isLoaded) {
48    const settings = settingsFrom(await $.store.get(STORE_SETTINGS).catch(() => undefined))
49    isLoaded = true
50    isCapsOn = settings.caps
51    return update($, gov, g => ({ ...g, settings }))
52  }
53  return read($, gov)
54}
55
56/** Applies a change and toasts each limit it crossed. */
57async function change($: EngineInterface, fn: (g: Gov, now: number) => Gov): Promise<Gov> {
58  const now = await $.clock.now()
59  let before: Gov | undefined
60  const after = await update($, gov, g => {
61    before = g
62    return fn(g, now)
63  })
64  if (before && before !== after) for (const text of crossings(before, after)) $.ui.toast(text)
65  isLive = running(after).length > 0
66  return after
67}
68
69async function setSettings($: EngineInterface, fn: (s: Settings) => Settings): Promise<Gov> {
70  await ensure($)
71  const g = await update($, gov, g => ({ ...g, settings: fn(g.settings) }))
72  isCapsOn = g.settings.caps
73  await $.store.set(STORE_SETTINGS, g.settings)
74  return g
75}
76
77// Ask the engine's list once activity settles; never two asks queued.
78function soon($: EngineInterface) {
79  if (isQueued) return
80  isQueued = true
81  $.clock.after(RECONCILE_MS, () => {
82    isQueued = false
83    void reconcileNow($).catch(() => undefined)
84  })
85}
86
87async function reconcileNow($: EngineInterface) {
88  const fetchedAt = await $.clock.now()
89  const list = await $.agent.list()
90  await change($, (g, now) => reconcile(g, list, fetchedAt, now))
91}
92
93export const register: Register = on => {
94  on('session.start', async ($, e, next) => {
95    await $.command.register({
96      name: 'governor',
97      description: 'Subagent governor: status, limits (max/budget/session), opt-in caps and model routing per agent type',
98      argumentHint: 'max N | budget N | session N | caps on|off | route <type>=<model>|off',
99    })
100    await ensure($).catch(() => undefined)
101    return next(e)
102  })
103
104  on('command.run', { command: 'governor' }, async ($, e) => {
105    const cmd = parseCommand(e.args)
106    const now = await $.clock.now()
107    switch (cmd.kind) {
108      case 'status':
109        return { text: statusText(await ensure($), now) }
110      case 'error':
111        return { text: cmd.text }
112      case 'set': {
113        const g = await setSettings($, s => ({ ...s, [cmd.key]: cmd.value }))
114        return { text: `agent-governor: ${cmd.key === 'maxConcurrent' ? `max ${g.settings.maxConcurrent} concurrent` : cmd.key === 'agentBudget' ? `budget ${g.settings.agentBudget} tokens per agent` : `session budget ${g.settings.sessionBudget} tokens`}.` }
115      }
116      case 'caps':
117        await setSettings($, s => ({ ...s, caps: cmd.isOn }))
118        return {
119          text: cmd.isOn
120            ? 'agent-governor: caps ON. Spawns past the concurrency limit ask first; an agent past its token budget has its tool calls denied.'
121            : 'agent-governor: caps off. Watching and warning only.',
122        }
123      case 'routes-off':
124        await setSettings($, s => ({ ...s, routes: {} }))
125        return { text: 'agent-governor: routing off; agents run on the model their call or definition picks.' }
126      case 'route': {
127        const g = await setSettings($, s => {
128          const routes = { ...s.routes }
129          if (cmd.model) routes[cmd.type] = cmd.model
130          else delete routes[cmd.type]
131          return { ...s, routes }
132        })
133        const all = Object.entries(g.settings.routes)
134        return {
135          text: `agent-governor: ${cmd.model ? `${cmd.type} spawns with no model given run on ${cmd.model}` : `no route for ${cmd.type}`}. Routes: ${all.length ? all.map(([t, m]) => `${t}=${m}`).join(', ') : 'none'}.`,
136        }
137      }
138    }
139  })
140
141  on('agent.spawn', async ($, e, next) => {
142    const g = await ensure($)
143    const s = g.settings
144    const route = s.routes[e.subagentType]
145    const input = route && !e.model && !e.fork && !e.workflow ? { ...e, model: route } : e
146    if (s.caps) {
147      const n = running(g).length + inFlight
148      if (n >= s.maxConcurrent) {
149        const answer = await $.ui.ask(`${n} agents already running. Start another?`, { options: ['Start it', 'Deny'], header: 'Governor' }).catch(() => 'Deny')
150        if (answer !== 'Start it') return { deny: tooMany(n, s.maxConcurrent) }
151      }
152    }
153    inFlight++
154    try {
155      const started = await next(input)
156      if (started.agentId) {
157        const id = started.agentId
158        await change($, (g, now) => spawned(g, { id, type: e.subagentType || 'agent', description: e.description }, now)).catch(() => undefined)
159      }
160      return started
161    } finally {
162      inFlight--
163    }
164  }).catch(($, e, next) => (next.called ? next(e) : isCapsOn ? { deny: 'agent-governor: its check failed while caps are on; try the spawn again.' } : next(e)))
165
166  on('turn.step', async function* ($, e, next) {
167    const r = yield* next(e)
168    if (e.agentId) {
169      const id = e.agentId
170      const g = await change($, (g, now) => stepped(g, id, tokensOf(r.usage), now)).catch(() => null)
171      if (g && !g.agents[id]) soon($) // a loop we never saw spawn: the list may name it
172    }
173    return r
174  })
175
176  on('turn.complete', async ($, e, next) => {
177    const r = await next(e)
178    if (e.agentId) {
179      const id = e.agentId
180      await change($, (g, now) => finished(g, id, now)).catch(() => undefined)
181    }
182    return r
183  }).catch(($, e, next) => next(e))
184
185  on('tool.call', async ($, e, next) => {
186    if (e.agentId && (isCapsOn || !isLoaded) && !WIND_DOWN.has(String(e.tool))) {
187      const g = await ensure($)
188      const a = g.agents[e.agentId]
189      if (g.settings.caps && a && a.tokens > g.settings.agentBudget) return { deny: OVER_BUDGET }
190    }
191    const r = await next(e)
192    if (isLive) soon($)
193    return r
194  }).catch(($, e, next) => next(e))
195
196  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
197    const below = await next(e)
198    if (e.props.hasSurvey) return below
199    const row = bandRow(await read($, gov))
200    if (!row) return below
201    const { Box, Text } = $.ui.resolve(e)
202    return (
203      <Box flexDirection="column">
204        <Box width={e.props.bodyColumns}>
205          <Text wrap="truncate-end">
206            <Text bold color={row.warnings.length ? '#f87171' : '#a78bfa'}>
207              [GOV]{' '}
208            </Text>
209            {row.text}
210            {row.warnings.length > 0 && <Text color="#f87171"> · ⚠ {row.warnings.join(' · ')}</Text>}
211          </Text>
212        </Box>
213        {below}
214      </Box>
215    )
216  })
217}
218

1import type { AgentInfo, AgentStatus, ModelUsage } from 'claude-code'
2
3import type { Gov, Settings, Tracked } from '../types'
4
5// Pure helpers: every change to the governor's state is a function from the
6// old state to the new, so register.tsx can hand them to `update` (which may
7// run one twice) and compare before with after to find what crossed a limit.
8
9export const DEFAULTS: Settings = { maxConcurrent: 6, agentBudget: 300_000, sessionBudget: 3_000_000, caps: false, routes: {} }
10export const KEEP_FINISHED_MS = 60 * 60 * 1000
11const MAX_TRACKED = 200
12const LIVE: readonly AgentStatus[] = ['pending', 'running', 'waiting']
13
14export const empty = (settings: Settings = DEFAULTS): Gov => ({ agents: {}, peak: 0, spawned: 0, sessionTokens: 0, perType: {}, settings })
15
16export const running = (g: Gov): Tracked[] => Object.values(g.agents).filter(a => a.isRunning)
17
18/**
19 * New tokens a request cost: uncached input, cache writes and output. Cache reads are left out: a
20 * long agent loop re-reads its whole transcript every step at a tenth of the price, and counting
21 * those would trip any budget on cheap re-reads rather than on work.
22 */
23export const tokensOf = (u: ModelUsage | null | undefined): number =>
24  u ? (u.input_tokens || 0) + (u.output_tokens || 0) + (u.cache_creation_input_tokens || 0) : 0
25
26export function fmtTok(n: number): string {
27  if (n < 1000) return String(n)
28  if (n < 1_000_000) return `${Math.round(n / 1000)}k`
29  return `${(n / 1_000_000).toFixed(1).replace(/\.0$/, '')}M`
30}
31
32/** `300000`, `300k`, `3m`, `1.5M`, `3_000_000` -> a positive whole number, or null. */
33export function parseCount(text: string): number | null {
34  const m = /^(\d+(?:\.\d+)?)([km])?$/i.exec(text.trim().replace(/[_,]/g, ''))
35  if (!m) return null
36  const n = Math.round(Number(m[1]) * (m[2] ? (m[2].toLowerCase() === 'k' ? 1000 : 1_000_000) : 1))
37  return n > 0 && Number.isFinite(n) ? n : null
38}
39
40const ALIASES = ['sonnet', 'opus', 'haiku', 'fable']
41
42/**
43 * A route's model, lowercased: an alias the Agent tool takes or a full `claude-` id; null for
44 * anything else, so a typo (`sonet`) never reaches the store and breaks every spawn of the type.
45 */
46export function modelOf(text: string): string | null {
47  const m = text.trim().toLowerCase()
48  return ALIASES.includes(m) || /^claude-[a-z0-9.-]+(\[1m\])?$/.test(m) ? m : null
49}
50
51/** Settings as stored, each field checked, the defaults where one is missing or bad. */
52export function settingsFrom(stored: unknown): Settings {
53  const s = (stored && typeof stored === 'object' ? stored : {}) as Partial<Record<keyof Settings, unknown>>
54  const num = (v: unknown, d: number) => (typeof v === 'number' && Number.isFinite(v) && v > 0 ? Math.floor(v) : d)
55  const routes: Record<string, string> = {}
56  if (s.routes && typeof s.routes === 'object') {
57    for (const [type, model] of Object.entries(s.routes as Record<string, unknown>)) {
58      const valid = typeof model === 'string' ? modelOf(model) : null
59      if (valid) routes[type] = valid
60    }
61  }
62  return {
63    maxConcurrent: num(s.maxConcurrent, DEFAULTS.maxConcurrent),
64    agentBudget: num(s.agentBudget, DEFAULTS.agentBudget),
65    sessionBudget: num(s.sessionBudget, DEFAULTS.sessionBudget),
66    caps: s.caps === true,
67    routes,
68  }
69}
70
71function withPeak(g: Gov): Gov {
72  const n = running(g).length
73  return n > g.peak ? { ...g, peak: n } : g
74}
75
76/** Drops finished agents older than an hour (and the oldest finished past 200 tracked); totals stay. */
77export function prune(g: Gov, now: number): Gov {
78  const kept = Object.values(g.agents).filter(a => a.isRunning || (a.endedAt ?? now) > now - KEEP_FINISHED_MS)
79  const finished = kept.filter(a => !a.isRunning).sort((a, b) => (a.endedAt ?? 0) - (b.endedAt ?? 0))
80  const drop = new Set(finished.slice(0, Math.max(0, kept.length - MAX_TRACKED)).map(a => a.id))
81  const agents = Object.fromEntries(kept.filter(a => !drop.has(a.id)).map(a => [a.id, a]))
82  return Object.keys(agents).length === Object.keys(g.agents).length ? g : { ...g, agents }
83}
84
85export function spawned(g: Gov, a: { id: string; type: string; description: string }, now: number): Gov {
86  if (g.agents[a.id]) return withPeak({ ...g, agents: { ...g.agents, [a.id]: { ...g.agents[a.id], isRunning: true, endedAt: undefined, touchedAt: now } } })
87  const n = (g.perType[a.type] ?? 0) + 1
88  const t: Tracked = { id: a.id, type: a.type, label: `${a.type}#${n}`, description: a.description, startedAt: now, tokens: 0, isRunning: true, touchedAt: now }
89  return withPeak(prune({ ...g, spawned: g.spawned + 1, perType: { ...g.perType, [a.type]: n }, agents: { ...g.agents, [a.id]: t } }, now))
90}
91
92/** One model request of an agent's loop answered: its tokens count, and it is running. Unknown ids are left alone. */
93export function stepped(g: Gov, id: string, tokens: number, now: number): Gov {
94  const a = g.agents[id]
95  if (!a) return g
96  return withPeak({ ...g, sessionTokens: g.sessionTokens + tokens, agents: { ...g.agents, [id]: { ...a, tokens: a.tokens + tokens, isRunning: true, endedAt: undefined, touchedAt: now } } })
97}
98
99/** The agent's run ended (its own turn.complete); a later step marks it running again. */
100export function finished(g: Gov, id: string, now: number): Gov {
101  const a = g.agents[id]
102  if (!a || !a.isRunning) return g
103  return prune({ ...g, agents: { ...g.agents, [id]: { ...a, isRunning: false, endedAt: now, touchedAt: now } } }, now)
104}
105
106/**
107 * The engine's list is the authority on who is alive: a listed agent takes its
108 * status from it, unless an event of its own touched it after the list was
109 * asked for (that event is newer). A live agent nobody saw spawn is adopted;
110 * one missing from the list is left as it was.
111 */
112export function reconcile(g: Gov, list: readonly Pick<AgentInfo, 'id' | 'type' | 'description' | 'status'>[], fetchedAt: number, now: number): Gov {
113  let agents = g.agents
114  let perType = g.perType
115  let spawnedN = g.spawned
116  for (const info of list) {
117    const isLive = LIVE.includes(info.status)
118    const a = agents[info.id]
119    if (a) {
120      if (a.touchedAt > fetchedAt || a.isRunning === isLive) continue
121      agents = { ...agents, [info.id]: { ...a, isRunning: isLive, endedAt: isLive ? undefined : now } }
122    } else if (isLive) {
123      const n = (perType[info.type] ?? 0) + 1
124      perType = { ...perType, [info.type]: n }
125      spawnedN++
126      agents = { ...agents, [info.id]: { id: info.id, type: info.type, label: `${info.type}#${n}`, description: info.description, startedAt: now, tokens: 0, isRunning: true, touchedAt: now } }
127    }
128  }
129  return agents === g.agents ? g : withPeak(prune({ ...g, agents, perType, spawned: spawnedN }, now))
130}
131
132/** The toasts owed for what went over a limit between `before` and `after`: once per crossing. */
133export function crossings(before: Gov, after: Gov): string[] {
134  const s = after.settings
135  const out: string[] = []
136  const rb = running(before).length
137  const ra = running(after).length
138  if (rb <= s.maxConcurrent && ra > s.maxConcurrent) {
139    out.push(`agent-governor: ${ra} subagents running (limit ${s.maxConcurrent})${s.caps ? '; new spawns will ask first' : ''}`)
140  }
141  for (const a of Object.values(after.agents)) {
142    const tb = before.agents[a.id]?.tokens ?? 0
143    if (tb <= s.agentBudget && a.tokens > s.agentBudget) {
144      out.push(`agent-governor: ${a.label} used ${fmtTok(a.tokens)} tokens (budget ${fmtTok(s.agentBudget)})${s.caps ? '; its next tool call is denied' : ''}`)
145    }
146  }
147  if (before.sessionTokens <= s.sessionBudget && after.sessionTokens > s.sessionBudget) {
148    out.push(`agent-governor: subagents used ${fmtTok(after.sessionTokens)} tokens this session (budget ${fmtTok(s.sessionBudget)})`)
149  }
150  return out
151}
152
153/** What is over a limit now, among the running agents. */
154export function warnings(g: Gov): string[] {
155  const s = g.settings
156  const run = running(g)
157  const out: string[] = []
158  if (run.length > s.maxConcurrent) out.push(`${run.length} > ${s.maxConcurrent} running`)
159  const over = run.filter(a => a.tokens > s.agentBudget).sort((a, b) => b.tokens - a.tokens)
160  if (over.length > 0) out.push(`${over[0].label} ${fmtTok(over[0].tokens)} > ${fmtTok(s.agentBudget)}${over.length > 1 ? ` (+${over.length - 1})` : ''}`)
161  if (g.sessionTokens > s.sessionBudget) out.push(`session ${fmtTok(g.sessionTokens)} > ${fmtTok(s.sessionBudget)}`)
162  return out
163}
164
165/** The band's row while agents run; null when there is nothing to show. */
166export function bandRow(g: Gov): { text: string; warnings: string[] } | null {
167  const n = running(g).length
168  if (n === 0) return null
169  return { text: `agents ${n} running · peak ${g.peak} · ${fmtTok(g.sessionTokens)} tok`, warnings: warnings(g) }
170}
171
172export function statusText(g: Gov, now: number): string {
173  const s = g.settings
174  const run = running(g)
175  const routes = Object.entries(s.routes)
176  const top = Object.values(g.agents)
177    .sort((a, b) => b.tokens - a.tokens)
178    .slice(0, 5)
179  const lines = [
180    `agent-governor: ${s.caps ? 'caps ON (asks before over-limit spawns, stops over-budget agents)' : 'watch and warn only'}`,
181    `Running ${run.length} · peak ${g.peak} · spawned ${g.spawned} · subagent tokens ${fmtTok(g.sessionTokens)} of ${fmtTok(s.sessionBudget)}`,
182    `Limits: ${s.maxConcurrent} concurrent · ${fmtTok(s.agentBudget)} per agent · ${fmtTok(s.sessionBudget)} per session`,
183    `Routes: ${routes.length ? routes.map(([t, m]) => `${t}=${m}`).join(', ') : 'none'}`,
184  ]
185  const warn = warnings(g)
186  if (warn.length) lines.push(`Over: ${warn.join(' · ')}`)
187  if (top.length) {
188    lines.push('Top agents by tokens:')
189    for (const a of top) {
190      const when = a.isRunning ? `running ${Math.round((now - a.startedAt) / 60000)}m` : 'done'
191      lines.push(`  ${a.label}  ${fmtTok(a.tokens)}  ${when}  ${a.description}`)
192    }
193  }
194  lines.push('Commands: /governor max N | budget N | session N | caps on|off | route <type>=<model>|off')
195  return lines.join('\n')
196}
197
198export type Command =
199  | { kind: 'status' }
200  | { kind: 'set'; key: 'maxConcurrent' | 'agentBudget' | 'sessionBudget'; value: number }
201  | { kind: 'caps'; isOn: boolean }
202  | { kind: 'route'; type: string; model: string | null }
203  | { kind: 'routes-off' }
204  | { kind: 'error'; text: string }
205
206const KEYS = { max: 'maxConcurrent', budget: 'agentBudget', session: 'sessionBudget' } as const
207
208export function parseCommand(args: string): Command {
209  const [word = '', ...rest] = args.trim().split(/\s+/).filter(Boolean)
210  const arg = rest.join(' ')
211  const verb = word.toLowerCase()
212  if (verb === '' || verb === 'status') return { kind: 'status' }
213  if (verb === 'max' || verb === 'budget' || verb === 'session') {
214    const value = parseCount(arg)
215    return value === null ? { kind: 'error', text: `/governor ${verb} takes a positive number (e.g. ${verb === 'max' ? '6' : '300k'}).` } : { kind: 'set', key: KEYS[verb], value }
216  }
217  if (verb === 'caps') {
218    if (/^on$/i.test(arg)) return { kind: 'caps', isOn: true }
219    if (/^off$/i.test(arg)) return { kind: 'caps', isOn: false }
220    return { kind: 'error', text: '/governor caps takes on or off.' }
221  }
222  if (verb === 'route') {
223    if (/^off$/i.test(arg)) return { kind: 'routes-off' }
224    const usage = '/governor route takes <type>=<model> (e.g. Explore=sonnet), <type>=off, or off; the model is sonnet, opus, haiku, fable or a full claude- id.'
225    const m = /^([^=\s]+)=(\S*)$/.exec(arg)
226    if (!m) return { kind: 'error', text: usage }
227    const [, type = '', given = ''] = m
228    if (!given || given.toLowerCase() === 'off') return { kind: 'route', type, model: null }
229    const model = modelOf(given)
230    return model ? { kind: 'route', type, model } : { kind: 'error', text: `"${given}" is not a model. ${usage}` }
231  }
232  return { kind: 'error', text: `Unknown /governor command "${word}". Try: /governor, max N, budget N, session N, caps on|off, route <type>=<model>|off.` }
233}
234

1export type Settings = {
2  /** More subagents running than this warns (and, with caps on, asks before another starts). */
3  maxConcurrent: number
4  /** Tokens one subagent may use before it warns (and, with caps on, is told to hand back). */
5  agentBudget: number
6  /** Tokens all subagents of the session may use before it warns. */
7  sessionBudget: number
8  /** Opt-in enforcement: ask before spawning past maxConcurrent, deny tool calls of an agent over budget. */
9  caps: boolean
10  /** Opt-in routing: subagent type -> model, applied only when the spawn named no model. */
11  routes: Record<string, string>
12}
13
14export type Tracked = {
15  id: string
16  type: string
17  /** `Explore#2`: the type and its ordinal this session. */
18  label: string
19  description: string
20  startedAt: number
21  endedAt?: number
22  tokens: number
23  isRunning: boolean
24  /** When this record last changed from an event of its own (spawn, step, completion). */
25  touchedAt: number
26}
27
28export type Gov = {
29  agents: Record<string, Tracked>
30  peak: number
31  spawned: number
32  /** Every subagent token counted this session, pruned agents included. */
33  sessionTokens: number
34  perType: Record<string, number>
35  settings: Settings
36}
37
38declare module 'claude-code' {
39  interface PluginState {
40    'agent-governor': { gov: Gov }
41  }
42}
43