fleet

1import { atom, read, update } from 'claude-code'
2import type { AgentStatus, EngineInterface, Register, Timer } from 'claude-code'
3import type { Budget, BudgetArgs, FleetAgent } from '../types'
4
5const PANE = 'fleet'
6const TOOL = 'set_agent_budget'
7const ACTIVE = new Set<AgentStatus>(['pending', 'running', 'waiting', 'idle'])
8const FAILED = new Set<AgentStatus>(['failed', 'killed'])
9
10const budget = atom({ plugin: 'fleet', key: 'budget' } as const, {})
11const agents = atom({ plugin: 'fleet', key: 'agents' } as const, [])
12const doneOffset = atom({ plugin: 'fleet', key: 'doneOffset' } as const, 0)
13
14const MODEL = /(opus|sonnet|haiku|fable)(?:-(\d{1,2})(?!\d)(?:-(\d{1,2})(?!\d))?)?/i
15
16// `claude-opus-5-5` and the alias `opus` are both family `opus`; any other model is its own.
17const familyOf = (model: string) => MODEL.exec(model)?.[1]?.toLowerCase() ?? model
18
19// `claude-opus-5-5` (with or without a date suffix) reads `Opus 5.5`, the alias `opus` reads `Opus`.
20const labelOf = (model: string) => {
21  const [, family, major, minor] = MODEL.exec(model) ?? []
22  if (!family) return model
23
24  const name = family.charAt(0).toUpperCase() + family.slice(1).toLowerCase()
25  if (!major) return name
26
27  return `${name} ${major}${minor ? `.${minor}` : ''}`
28}
29
30// A row keeps the full model id it already knows when the budget names only the alias.
31const modelFor = (known: string | undefined, given: string) =>
32  known && given.toLowerCase() === familyOf(given) ? known : given
33
34// `new` replaces the caps and restarts the count; `add` raises the caps and keeps it.
35const applyBudget = (current: Budget, { mode, caps }: BudgetArgs): Budget => {
36  const next: Budget = mode === 'new' ? {} : { ...current }
37  for (const [given, n] of Object.entries(caps)) {
38    const family = familyOf(given)
39    const line = next[family]
40    next[family] = { model: modelFor(current[family]?.model, given), cap: (line?.cap ?? 0) + n, spent: line?.spent ?? 0 }
41  }
42
43  return next
44}
45
46const elapsed = (ms: number) => {
47  const s = Math.max(0, Math.floor(ms / 1000))
48  if (s < 60) return `${s}s`
49
50  const m = Math.floor(s / 60)
51  return m < 60 ? `${m}m` : `${Math.floor(m / 60)}h${String(m % 60).padStart(2, '0')}m`
52}
53
54// Agent statuses and elapsed times are not state, so the open pane asks again every second.
55let ticker: Timer | undefined
56const tick = ($: EngineInterface) => {
57  ticker ??= $.clock.every(1000, () => $.ui.invalidate('ui.render'))
58}
59
60const DESCRIPTION = [
61  'Call this every time the user sets, changes, adds to or clears a subagent budget ("use up to 3 opus and 2 sonnet agents", "add 2 more opus").',
62  'It shows the budget in their /fleet pane; it does not enforce it, so you still keep to the budget yourself.',
63  'mode "new" replaces the caps and restarts the count; "add" adds to the caps and keeps the count; "new" with empty caps clears the budget.',
64].join(' ')
65
66const SCHEMA = {
67  type: 'object',
68  properties: {
69    mode: { type: 'string', enum: ['new', 'add'], description: '"new" for a new budget, "add" to add to the current one' },
70    caps: {
71      type: 'object',
72      description: 'Agents allowed per model, by exact model id: {"claude-opus-5-5": 3, "claude-sonnet-5-5": 2}',
73      additionalProperties: { type: 'integer', minimum: 0 },
74    },
75  },
76  required: ['mode', 'caps'],
77}
78
79export const register: Register = on => {
80  on('session.start', async ($, e, next) => {
81    await $.command.register({
82      name: 'fleet',
83      description: 'Toggle the pane with the subagent budget and the agents running now',
84      immediate: true,
85    })
86    await $.tool.register({ name: TOOL, description: DESCRIPTION, inputSchema: SCHEMA, isDeferred: false })
87    // A reload keeps the pane open but stops the module's timers.
88    if ((await $.ui.panes()).some(pane => pane.id === PANE)) tick($)
89
90    return next(e)
91  })
92
93  on('command.run', { command: 'fleet' }, async $ => {
94    if ((await $.ui.panes()).some(pane => pane.id === PANE)) {
95      await $.ui.close({ id: PANE })
96      return {}
97    }
98
99    const opened = await $.ui.open({ id: PANE, title: 'Fleet' })
100    tick($)
101    return opened.isPlaced ? {} : { text: `The fleet pane waits: ${opened.reason}` }
102  })
103
104  // Every close, the person's included, stops the redraws.
105  on('ui.close', { id: PANE }, ($, e, next) => {
106    ticker?.cancel()
107    ticker = undefined
108    return next(e)
109  })
110
111  on('tool.call', { tool: 'mcp__fleet__set_agent_budget' }, async ($, e) => {
112    // The schema only guides the model, so the caps are checked again here.
113    const mode = e.mode === 'add' ? 'add' : 'new'
114    const caps = Object.fromEntries(Object.entries(e.caps ?? {}).filter(([, n]) => Number.isInteger(n) && n >= 0))
115    const lines = await update($, budget, current => applyBudget(current, { mode, caps }))
116
117    // A new budget starts the done list again; agents still running stay listed.
118    if (mode === 'new') {
119      const live = new Set((await $.agent.list()).filter(one => ACTIVE.has(one.status)).map(one => one.id))
120      await Promise.all([update($, agents, list => list.filter(one => live.has(one.id))), update($, doneOffset, () => 0)])
121    }
122
123    const summary = Object.values(lines)
124      .map(line => `${labelOf(line.model)} ${line.spent}/${line.cap ?? '-'}`)
125      .join(', ')
126
127    return { result: summary ? `The /fleet pane shows: ${summary}.` : 'Budget cleared.' }
128  }).catch(() => ({ deny: 'The fleet pane could not record the budget.' }))
129
130  // Counts every spawn the engine made, nested ones too; never changes or blocks one.
131  // Bookkeeping that fails leaves the spawn as the engine answered it.
132  on('agent.spawn', async ($, e, next) => {
133    const spawned = await next(e)
134    if (spawned.deny !== undefined) return spawned
135
136    const { model, agentId } = spawned
137    const family = familyOf(model)
138    const [startedAt] = await Promise.all([
139      $.clock.now(),
140      update($, budget, current => {
141        const line = current[family]
142        return { ...current, [family]: { model: modelFor(line?.model, model), cap: line?.cap, spent: (line?.spent ?? 0) + 1 } }
143      }),
144    ])
145    if (agentId) {
146      await update($, agents, list => [...list, { id: agentId, model, description: e.description, startedAt }].slice(-200))
147    }
148
149    return spawned
150  }).catch(($, e, next) => next(e))
151
152  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
153    const { Box, Text, Button } = $.ui.resolve(e)
154    const [lines, spawned, scrolled, listed, now] = await Promise.all([
155      read($, budget),
156      read($, agents),
157      read($, doneOffset),
158      $.agent.list(),
159      $.clock.now(),
160    ])
161    // Ended agents leave `$.agent.list()` after a while, so a missing one is done too.
162    const status = new Map(listed.map(one => [one.id, one.status]))
163    const statusOf = (id: string) => status.get(id) ?? 'completed'
164    const active: FleetAgent[] = []
165    const finished: FleetAgent[] = []
166    for (const one of spawned) (ACTIVE.has(statusOf(one.id)) ? active : finished).push(one)
167    finished.reverse()
168    const rows = Object.values(lines).filter(line => line.cap !== undefined || line.spent > 0)
169    if (rows.length === 0 && spawned.length === 0) return <Text dimColor>No subagents yet.</Text>
170
171    const models = new Set([...rows.map(line => line.model), ...spawned.map(one => one.model)])
172    const labelWidth = Math.max(0, ...[...models].map(model => labelOf(model).length))
173
174    // The done list takes the rows left under the budget and the running agents, and scrolls past that.
175    const room = Math.max(3, e.props.scroll.bodyRows - rows.length - active.length - 1)
176    const isScrolling = finished.length > room
177    const size = isScrolling ? room - 2 : finished.length
178    const maxOffset = finished.length - size
179    const clamp = (n: number) => Math.max(0, Math.min(maxOffset, n))
180    const offset = clamp(scrolled)
181    const scroll = (by: number) => update($, doneOffset, n => clamp(clamp(n) + by))
182
183    return (
184      <Box flexDirection="column">
185        {rows.map(line => {
186          const tone = line.cap === undefined || line.spent < line.cap ? undefined : line.spent > line.cap ? 'error' : 'warning'
187          return (
188            <Text key={`budget-${line.model}`} color={tone} wrap="truncate-end">
189              {`${labelOf(line.model).padEnd(labelWidth)}  ${line.spent}${line.cap === undefined ? '' : `/${line.cap}`}`}
190            </Text>
191          )
192        })}
193        {active.map(one => {
194          const isRunning = statusOf(one.id) === 'running'
195          return (
196            <Box key={`agent-${one.id}`} flexDirection="row">
197              <Text color={isRunning ? 'success' : undefined} dimColor={!isRunning}>●</Text>
198              <Text>{` ${labelOf(one.model).padEnd(labelWidth)}  `}</Text>
199              <Box flexGrow={1}>
200                <Text wrap="truncate-end">{one.description}</Text>
201              </Box>
202              <Box flexShrink={0}>
203                <Text>{`  ${elapsed(now - one.startedAt)}`}</Text>
204              </Box>
205            </Box>
206          )
207        })}
208        {finished.length > 0 && <Text dimColor>done: {finished.length}</Text>}
209        {isScrolling && (
210          <Button key="done-up" label={`▲ ${offset} more`} hotkey="k" plain dimColor={offset === 0} onPress={() => scroll(-size)} />
211        )}
212        {finished.slice(offset, offset + size).map(one => {
213          const isFailed = FAILED.has(statusOf(one.id))
214          return (
215            <Text key={`done-${one.id}`} color={isFailed ? 'error' : undefined} dimColor={!isFailed} wrap="truncate-end">
216              {`${isFailed ? '✗' : ' '} ${labelOf(one.model).padEnd(labelWidth)}  ${one.description}`}
217            </Text>
218          )
219        })}
220        {isScrolling && (
221          <Button key="done-down" label={`▼ ${maxOffset - offset} more`} hotkey="j" plain dimColor={offset === maxOffset} onPress={() => scroll(size)} />
222        )}
223      </Box>
224    )
225  })
226}
227

1// One budget row: a model family's cap and how many agents it spawned since the budget was set.
2export type BudgetLine = { model: string; cap?: number; spent: number }
3
4// Keyed by model family (`opus`, `sonnet`, ...).
5export type Budget = Record<string, BudgetLine>
6
7// What Claude passes to `set_agent_budget`: caps keyed by model id or alias.
8export type BudgetArgs = { mode: 'new' | 'add'; caps: Record<string, number> }
9
10// A spawned subagent; its live status comes from `$.agent.list()` while drawing.
11export type FleetAgent = { id: string; model: string; description: string; startedAt: number }
12
13declare module 'claude-code' {
14  interface PluginState {
15    // `doneOffset`: how many finished agents the done list has scrolled past.
16    fleet: { budget: Budget; agents: FleetAgent[]; doneOffset: number }
17  }
18
19  interface McpToolInputs {
20    mcp__fleet__set_agent_budget: BudgetArgs
21  }
22}
23