sub-agents

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { AgentRow } from '../types'
5import { CANCEL, elapsed, label, options, parse, shortModel, tokens } from './choices'
6import type { Choice } from './choices'
7
8const PANE = 'sub-agents'
9const TITLE = 'Sub-agents'
10
11const agents = atom({ plugin: 'sub-agents', key: 'agents' } as const, [] as AgentRow[])
12// On by default in every session; /sub-agents choose turns it off.
13const choose = atom({ plugin: 'sub-agents', key: 'choose' } as const, true)
14const now = atom({ plugin: 'sub-agents', key: 'now' } as const, 0)
15
16// Sent once to the model, beside the next prompt, while choose mode is on.
17// Goes in the prompt's context rather than the system prompt, so the prompt cache survives.
18const HINT =
19  'Subagent choice mode is on: on every Agent tool call, set `model` and `effort` to the pair you ' +
20  'recommend for that task (the user asked for this and confirms each choice before the subagent runs). ' +
21  'Prefer the cheapest pair that will get the task right.'
22
23// Module state: starts over on a reload, which only loses a pending hint and in-flight efforts.
24let hintPending = true
25// Agent tool call id -> effort chosen; agent.spawn does not carry the effort.
26const efforts = new Map<string, string | undefined>()
27
28async function setRow($: EngineInterface, id: string, fn: (row: AgentRow) => AgentRow) {
29  await update($, agents, list => list.map(row => (row.id === id ? fn(row) : row)))
30}
31
32async function refreshStatus($: EngineInterface) {
33  const running = (await read($, agents)).filter(a => a.status === 'running').length
34  $.ui.status(running === 0 ? undefined : `${running} sub-agent${running === 1 ? '' : 's'} running`)
35}
36
37async function toggleChoose($: EngineInterface): Promise<boolean> {
38  const isOn = !(await read($, choose))
39  await update($, choose, () => isOn)
40  hintPending = isOn
41  return isOn
42}
43
44export const register: Register = on => {
45  on('session.start', async ($, e, next) => {
46    const result = await next(e)
47    await $.command.register({
48      name: 'sub-agents',
49      description: 'Opens the sub-agents pane; "/sub-agents choose" turns model and effort choice on or off',
50    })
51    // Ticks the elapsed time and reconciles statuses, only while something runs.
52    $.clock.every(1000, () => {
53      void (async () => {
54        const list = await read($, agents)
55        if (!list.some(a => a.status === 'running')) return
56        const at = await $.clock.now()
57        const live = new Map((await $.agent.list()).map(a => [a.id, a.status]))
58        await update($, agents, rows =>
59          rows.map((row): AgentRow => {
60            const status = live.get(row.id)
61            if (row.status !== 'running' || status === undefined) return row
62            if (status === 'completed') return { ...row, status: 'done', endedAt: row.endedAt ?? at }
63            if (status === 'failed') return { ...row, status: 'failed', endedAt: row.endedAt ?? at }
64            if (status === 'killed') return { ...row, status: 'stopped', endedAt: row.endedAt ?? at }
65            return row
66          }),
67        )
68        await update($, now, () => at)
69        await refreshStatus($)
70      })()
71    })
72    return result
73  })
74
75  on('command.run', { command: 'sub-agents' }, async ($, e) => {
76    if (/^choose\b/i.test(e.args.trim())) {
77      const isOn = await toggleChoose($)
78      return {
79        text: isOn
80          ? 'Model and effort choice on: each sub-agent asks for your confirmation before it runs.'
81          : 'Model and effort choice off: Claude decides on its own.',
82      }
83    }
84    await $.ui.open({ id: PANE, title: TITLE })
85    return { text: 'Sub-agents pane opened.' }
86  })
87
88  on('prompt.submit', async ($, e, next) => {
89    if (!hintPending || !(await read($, choose))) return next(e)
90    hintPending = false
91    return next({ ...e, context: [...(e.context ?? []), HINT] })
92  })
93
94  on('tool.call', { tool: 'Agent' }, async ($, e, next) => {
95    if (!(await read($, choose))) {
96      efforts.set(e.tool_use_id, e.effort)
97      return next(e)
98    }
99    const suggested: Choice = { model: e.model, effort: e.effort }
100    const { labels, byLabel } = options(suggested)
101    const type = e.subagent_type ?? 'general-purpose'
102    let answer: string
103    try {
104      answer = await $.ui.ask(`Run "${e.description}" (${type}) with which model and effort?`, {
105        header: 'Sub-agent',
106        options: labels,
107      })
108    } catch {
109      answer = CANCEL
110    }
111    if (answer === CANCEL) {
112      return { deny: 'The user cancelled this subagent in the choice dialog. Do not spawn it again unless they ask.' }
113    }
114    let choice = byLabel.get(answer) ?? parse(answer)
115    if (choice === null) {
116      $.ui.toast(`Didn't understand "${answer}"; using ${label(suggested)}`)
117      choice = suggested
118    }
119    efforts.set(e.tool_use_id, choice.effort)
120    return next({ ...e, model: choice.model, effort: choice.effort })
121  })
122
123  on('agent.spawn', async ($, e, next) => {
124    const result = await next(e)
125    if (result.agentId === undefined) return result
126    const at = await $.clock.now()
127    const row: AgentRow = {
128      id: result.agentId,
129      description: e.description,
130      type: e.subagentType,
131      model: shortModel(result.model),
132      effort: efforts.get(e.tool_use_id) ?? null,
133      status: 'running',
134      startedAt: at,
135      endedAt: null,
136      tools: 0,
137      lastTool: null,
138      tokens: null,
139    }
140    efforts.delete(e.tool_use_id)
141    await update($, agents, list => [...list, row].slice(-50))
142    await update($, now, () => at)
143    await refreshStatus($)
144    void $.ui.open({ id: PANE, title: TITLE }).catch(() => undefined)
145    return result
146  })
147
148  on('tool.call', ($, e, next) => {
149    const id = e.agentId
150    if (id !== undefined) {
151      void setRow($, id, row => ({ ...row, tools: row.tools + 1, lastTool: e.tool }))
152    }
153    return next(e)
154  })
155
156  on('turn.complete', async ($, e, next) => {
157    const id = e.agentId
158    if (id !== undefined) {
159      const at = await $.clock.now()
160      const u = e.usage
161      const used = u
162        ? u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens + u.output_tokens
163        : null
164      const status = e.isAborted || e.reason === 'aborted' ? 'stopped' : e.reason === 'answer' ? 'done' : 'failed'
165      await setRow($, id, row => ({ ...row, status, endedAt: at, tokens: used ?? row.tokens }))
166      await refreshStatus($)
167    }
168    return next(e)
169  })
170
171  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
172    const { Box, Text, Button } = $.ui.resolve(e)
173    const list = await read($, agents)
174    const isChoosing = await read($, choose)
175    const at = Math.max(await read($, now), ...list.map(a => a.endedAt ?? a.startedAt))
176    const room = Math.max(1, (e.viewport?.rows ?? 24) - 3)
177
178    const mark = {
179      running: { glyph: '●', color: '#4a8fe7' },
180      done: { glyph: '✓', color: 'success' },
181      failed: { glyph: '✗', color: 'error' },
182      stopped: { glyph: '■', color: 'subtle' },
183    } as const
184
185    return (
186      <Box flexDirection="column">
187        <Box flexDirection="row" columnGap={2}>
188          <Button
189            key="choose"
190            label={isChoosing ? 'Choose model: on' : 'Choose model: off'}
191            onPress={() => toggleChoose($)}
192          />
193          <Button
194            key="clear"
195            label="Clear finished"
196            onPress={() => update($, agents, rows => rows.filter(r => r.status === 'running'))}
197          />
198        </Box>
199        {list.length === 0 && <Text dimColor>No sub-agents in this session.</Text>}
200        {list.slice(-room).map(a => (
201          <Box key={a.id} flexDirection="row" columnGap={1}>
202            <Text color={mark[a.status].color}>{mark[a.status].glyph}</Text>
203            <Box flexGrow={1} flexShrink={1}>
204              <Text wrap="truncate-end" dimColor={a.status !== 'running'}>
205                {a.description}
206              </Text>
207            </Box>
208            <Text dimColor>{a.type}</Text>
209            <Text color="#4a8fe7">{a.effort ? `${a.model} · ${a.effort}` : a.model}</Text>
210            <Text dimColor>{elapsed((a.endedAt ?? at) - a.startedAt)}</Text>
211            <Text dimColor>
212              {a.tools} tools{a.lastTool && a.status === 'running' ? ` · ${a.lastTool}` : ''}
213            </Text>
214            {a.tokens !== null && <Text dimColor>{tokens(a.tokens)} tok</Text>}
215          </Box>
216        ))}
217      </Box>
218    )
219  })
220}
221

1export type Model = 'haiku' | 'sonnet' | 'opus' | 'fable'
2export type Effort = 'low' | 'medium' | 'high' | 'xhigh' | 'max'
3export type Choice = { model?: Model; effort?: Effort }
4
5const MODELS: readonly Model[] = ['haiku', 'sonnet', 'opus', 'fable']
6const EFFORTS: readonly Effort[] = ['low', 'medium', 'high', 'xhigh', 'max']
7
8// Cheapest to strongest; the alternatives offered are the neighbours of the suggestion.
9const LADDER: readonly Required<Choice>[] = [
10  { model: 'haiku', effort: 'low' },
11  { model: 'haiku', effort: 'medium' },
12  { model: 'sonnet', effort: 'medium' },
13  { model: 'sonnet', effort: 'high' },
14  { model: 'opus', effort: 'medium' },
15  { model: 'opus', effort: 'high' },
16  { model: 'opus', effort: 'xhigh' },
17  { model: 'fable', effort: 'high' },
18]
19
20export const CANCEL = 'Cancel'
21
22export function label(c: Choice): string {
23  if (!c.model && !c.effort) return 'agent default'
24  return [c.model ?? 'default model', c.effort ?? 'default effort'].join(' · ')
25}
26
27function rung(c: Choice): number {
28  const exact = LADDER.findIndex(r => r.model === c.model && r.effort === c.effort)
29  if (exact >= 0) return exact
30  const sameModel = LADDER.findIndex(r => r.model === c.model)
31  return sameModel >= 0 ? sameModel : 4
32}
33
34/** Up to four options: the suggestion, one rung cheaper, one rung stronger, and Cancel. */
35export function options(suggested: Choice): { labels: string[]; byLabel: Map<string, Choice> } {
36  const byLabel = new Map<string, Choice>()
37  byLabel.set(`${label(suggested)} (suggested)`, suggested)
38  const i = rung(suggested)
39  for (const j of [i - 1, i + 1]) {
40    const alt = LADDER[j]
41    if (alt && label(alt) !== label(suggested)) byLabel.set(label(alt), alt)
42  }
43  return { labels: [...byLabel.keys(), CANCEL], byLabel }
44}
45
46/** Reads text typed under "Other", such as "sonnet high" or "opus". */
47export function parse(text: string): Choice | null {
48  const words = text.toLowerCase().split(/[^a-z]+/)
49  const model = MODELS.find(m => words.includes(m))
50  const effort = EFFORTS.find(x => words.includes(x))
51  return model || effort ? { model, effort } : null
52}
53
54export function shortModel(id: string): string {
55  return MODELS.find(m => id.toLowerCase().includes(m)) ?? id
56}
57
58export function elapsed(ms: number): string {
59  const s = Math.max(0, Math.floor(ms / 1000))
60  if (s < 60) return `${s}s`
61  const m = Math.floor(s / 60)
62  if (m < 60) return `${m}m${String(s % 60).padStart(2, '0')}s`
63  return `${Math.floor(m / 60)}h${String(m % 60).padStart(2, '0')}m`
64}
65
66export function tokens(n: number): string {
67  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
68  if (n >= 1_000) return `${Math.round(n / 1_000)}k`
69  return String(n)
70}
71

1export type AgentStatusRow = 'running' | 'done' | 'failed' | 'stopped'
2
3export type AgentRow = {
4  id: string
5  description: string
6  type: string
7  model: string
8  effort: string | null
9  status: AgentStatusRow
10  startedAt: number
11  endedAt: number | null
12  tools: number
13  lastTool: string | null
14  tokens: number | null
15}
16
17declare module 'claude-code' {
18  interface PluginState {
19    'sub-agents': { agents: AgentRow[]; choose: boolean; now: number }
20  }
21}
22