fleet-guard

1// fleet-guard — keeps subagents and agent-team teammates on allowed models and
2// caps how many teammates run at once.
3//
4// Hooks `agent.spawn`, which fires before a subagent starts and, since Claude
5// Code v2.1.289, before an agent-team teammate starts (e.isTeammate === true).
6// The decisions are pure and live in hooks/fleet.ts.
7//
8// Configuration (environment, read on each spawn so a change applies at once):
9//   FLEET_GUARD_MODELS         comma list of model ids or aliases, e.g. sonnet,haiku.
10//                              A spawn on any other model is moved to the first one.
11//   FLEET_GUARD_MAX_TEAMMATES  integer. A teammate spawn while that many teammates
12//                              are live (pending, running, waiting or idle) is denied.
13// Both unset: the mod changes nothing.
14//
15// If $.agent.list() fails the teammate count is unknown, and the spawn is
16// allowed with one log line: a guard that cannot count must not deny blind.
17//
18// Functions that take `$` are top-level declarations: the engine scans the
19// module before loading it and refuses `$` handed to anything else.
20
21import type { EngineInterface, Register } from 'claude-code'
22import { decide, liveTeammates, parseCap, parseModels, type Policy } from './fleet'
23
24async function policy($: EngineInterface): Promise<Policy> {
25  const models = parseModels(await $.env.get('FLEET_GUARD_MODELS'))
26  const cap = parseCap(await $.env.get('FLEET_GUARD_MAX_TEAMMATES'))
27  if (cap === 'invalid') {
28    $.ui.log('fleet-guard: FLEET_GUARD_MAX_TEAMMATES is not a whole number; teammates are not capped')
29    return { models, cap: null }
30  }
31  return { models, cap }
32}
33
34// The number of live teammates, or null when the list could not be read.
35async function countLive($: EngineInterface): Promise<number | null> {
36  try {
37    return liveTeammates(await $.agent.list())
38  } catch {
39    $.ui.log('fleet-guard: teammate count unknown ($.agent.list failed); allowing this teammate')
40    return null
41  }
42}
43
44export const register: Register = on => {
45  on('agent.spawn', async ($, e, next) => {
46    const p = await policy($)
47    if (p.models === null && p.cap === null) return next(e)
48    const isTeammate = e.isTeammate === true
49    const live = isTeammate && p.cap !== null ? await countLive($) : null
50    const d = decide({ isTeammate, fork: e.fork, model: e.model, parentModel: e.parentModel }, p, live)
51    if (d.kind === 'deny') return { deny: d.reason }
52    if (d.kind === 'rewrite') {
53      const who = isTeammate ? 'teammate' : 'subagent'
54      $.ui.log(`fleet-guard: ${who} model ${d.from} is not in FLEET_GUARD_MODELS; using ${d.model}`)
55      return next({ ...e, model: d.model })
56    }
57    return next(e)
58  })
59}
60

1// Pure decisions for fleet-guard: no `$`, so they unit-test without the kit.
2//
3// Two rules on every `agent.spawn` (a subagent, or an agent-team teammate when
4// e.isTeammate is true):
5//   models  FLEET_GUARD_MODELS, a comma list of model ids or aliases. A spawn
6//           whose model (the one it asked for, else its parent's) matches none
7//           is rewritten to the first entry. Unset or empty: every model passes.
8//   cap     FLEET_GUARD_MAX_TEAMMATES, an integer. A teammate spawn while that
9//           many teammates are already live is denied. Subagents are never
10//           capped. Unset: no cap.
11
12/** Agent states that hold a teammate's resources: started and not yet ended. */
13export const LIVE_STATES: ReadonlySet<string> = new Set(['pending', 'running', 'waiting', 'idle'])
14
15/** The allow list, lower-cased, or null when unset or empty (the rule is off). */
16export function parseModels(raw: string | undefined): string[] | null {
17  if (raw === undefined) return null
18  const list = raw.split(',').map(s => s.trim().toLowerCase()).filter(s => s.length > 0)
19  return list.length > 0 ? list : null
20}
21
22/** The cap, null when unset or empty, or 'invalid' for anything but a whole number ≥ 0. */
23export function parseCap(raw: string | undefined): number | null | 'invalid' {
24  if (raw === undefined || raw.trim() === '') return null
25  const t = raw.trim()
26  return /^\d+$/.test(t) ? Number(t) : 'invalid'
27}
28
29/**
30 * True when `model` is on the list: the same id or alias, or an id that has the
31 * alias as one of its dash-separated words (`claude-sonnet-5-5` matches
32 * `sonnet`). An alias asked for is never taken to match a full id on the list,
33 * because which id the alias resolves to is the host's choice, not ours.
34 */
35export function modelAllowed(model: string, allowed: readonly string[]): boolean {
36  const m = model.trim().toLowerCase()
37  const words = m.split(/[-_.:/\s[\]]+/)
38  return allowed.some(a => a === m || words.includes(a))
39}
40
41/** How many live teammates a `$.agent.list()` answer holds. Subagents are not counted. */
42export function liveTeammates(agents: readonly { teammateId?: string; status: string }[]): number {
43  return agents.filter(a => a.teammateId !== undefined && LIVE_STATES.has(a.status)).length
44}
45
46export type SpawnFacts = {
47  isTeammate: boolean
48  fork: boolean
49  /** The model the spawn asked for (alias or id), if any. */
50  model?: string
51  /** The parent's effective model: what a spawn that names none runs on. */
52  parentModel: string
53}
54
55export type Policy = {
56  models: string[] | null
57  cap: number | null
58}
59
60export type Decision =
61  | { kind: 'pass' }
62  | { kind: 'deny'; reason: string }
63  | { kind: 'rewrite'; model: string; from: string }
64
65/**
66 * What to do with one spawn. `live` is the number of live teammates, or null
67 * when it could not be counted (then the cap is not applied: never deny blind).
68 */
69export function decide(spawn: SpawnFacts, policy: Policy, live: number | null): Decision {
70  if (spawn.isTeammate && policy.cap !== null && live !== null && live >= policy.cap) {
71    return { kind: 'deny', reason: `fleet-guard: ${live} teammates already running (cap ${policy.cap})` }
72  }
73  // A fork always runs on its parent's model and ignores `model`, so there is nothing to rewrite.
74  if (policy.models === null || spawn.fork) return { kind: 'pass' }
75  const asked = spawn.model && spawn.model.toLowerCase() !== 'inherit' ? spawn.model : spawn.parentModel
76  if (modelAllowed(asked, policy.models)) return { kind: 'pass' }
77  return { kind: 'rewrite', model: policy.models[0] as string, from: asked }
78}
79