orchestrator

1import { atom, read, update } from 'claude-code'
2import type { Register, BoxProps, ElementConstructor, EngineInterface, RenderChildren, RenderElement, RenderNode, SelectOption, Timer } from 'claude-code'
3
4import type { Config, Field, Model } from '../types'
5
6const MODELS: Model[] = ['opus', 'sonnet', 'haiku']
7const FALLBACK: Config = { isOn: true, models: ['sonnet', 'haiku'], maxAgents: 2, contextK: 200 }
8const AGENT_PRESETS = [1, 2]
9const CONTEXT_PRESETS = [200, 400]
10const CUSTOM = 'custom'
11const SETTINGS = 'settings'
12const USER_ORIGINS = new Set(['composer', 'bridge', 'sdk'])
13
14const CONFIG_REF = { plugin: 'orchestrator', key: 'config' } as const
15const config = atom(CONFIG_REF, FALLBACK)
16const spawned = atom({ plugin: 'orchestrator', key: 'spawned' } as const, 0)
17const running = atom({ plugin: 'orchestrator', key: 'running' } as const, [])
18const spin = atom({ plugin: 'orchestrator', key: 'spin' } as const, 0)
19const openMenu = atom({ plugin: 'orchestrator', key: 'openMenu' } as const, null)
20
21// The spinner turns only while a subagent this session spawned is still running.
22const SPINNER = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
23const SPIN_MS = 100
24let spinner: Timer | undefined
25const track = async ($: EngineInterface, agentId: string, isRunning: boolean) => {
26  await update($, running, ids => (isRunning ? [...ids, agentId] : ids.filter(id => id !== agentId)))
27  const count = (await read($, running)).length
28  if (count > 0 && !spinner) spinner = $.clock.every(SPIN_MS, () => void update($, spin, n => n + 1))
29  if (count === 0) {
30    spinner?.cancel()
31    spinner = undefined
32  }
33}
34const editing = atom({ plugin: 'orchestrator', key: 'editing' } as const, null)
35const isOpen = atom({ plugin: 'orchestrator', key: 'isOpen' } as const, false)
36
37const contextLabel = (k: number) => (k >= 1000 ? `${k / 1000}M` : `${k}k`)
38const title = (m: Model) => m[0]!.toUpperCase() + m.slice(1)
39
40const summary = (c: Config) =>
41  `${c.isOn ? 'ON' : 'OFF'} | models: ${c.models.join(', ')} | max agents: ${c.maxAgents} | context: ${contextLabel(c.contextK)}`
42
43const brief = (c: Config) =>
44  [
45    'Orchestrator mode (chosen by the user for this request):',
46    '- You are the orchestrator. Break the work down, delegate it to subagents with the Agent tool, then verify and merge their results.',
47    `- Allowed subagent models: ${c.models.join(', ')}. Always set \`model\` on every Agent call to one of these, matched to the difficulty of the subtask.`,
48    `- Spawn at most ${c.maxAgents} subagents for this request. Calls past that are refused.`,
49    `- Keep each subagent within about ${contextLabel(c.contextK)} tokens of context: give it a self-contained brief, ask for a compact report, and split any subtask that needs more.`,
50    '- Subagents must not spawn subagents of their own. Tell them so in each brief.',
51  ].join('\n')
52
53// "300", "300k" and "1m" all read as thousands of tokens for ctx.
54const toNumber = (v: string) => Number(v.trim().replace(/k$/i, '').replace(/m$/i, '000'))
55
56// `/orch max=4 ctx=200 models=sonnet,haiku on` style edits; unknown words are reported back.
57const parse = (args: string, c: Config): { next: Config; bad: string[] } => {
58  const bad: string[] = []
59  let next = c
60  for (const word of args.split(/\s+/).filter(Boolean)) {
61    const [k, v = ''] = word.split('=')
62    const n = toNumber(v)
63    const picked = v.split(',').filter((m): m is Model => MODELS.includes(m as Model))
64    if (k === 'on' || k === 'off') next = { ...next, isOn: k === 'on' }
65    else if (k === 'max' && n >= 1) next = { ...next, maxAgents: Math.floor(n) }
66    else if (k === 'ctx' && n > 0) next = { ...next, contextK: n }
67    else if (k === 'models' && picked.length > 0) next = { ...next, models: picked }
68    else bad.push(word)
69  }
70  return { next, bad }
71}
72
73const withValue = (presets: number[], value: number) =>
74  (presets.includes(value) ? presets : [...presets, value].sort((a, b) => a - b))
75
76// Every modz mod places its tile in one shared band: a lead column (the bot), then a column of tiles, then an optional
77// right column. A column with no tiles is not drawn, and columns wrap under each other when the width runs out.
78const BAND = 'modz-band'
79const SLOTS: Record<string, 'lead' | 'left' | 'right'> = {
80  buddy: 'lead',
81  'usage-limits': 'left',
82  'cache-timer': 'left',
83  'model-info': 'left',
84  orchestrator: 'left',
85}
86const ORDER = Object.keys(SLOTS)
87const keyOf = (n: RenderNode) => (typeof n !== 'string' && n.type === 'Box' ? String(n.props?.key ?? '') : '')
88const rank = (n: RenderNode) => {
89  const i = ORDER.indexOf(keyOf(n))
90  return i < 0 ? ORDER.length : i
91}
92const pack = (Box: ElementConstructor<BoxProps>, below: RenderElement, tile: RenderElement) => {
93  const isEmpty = below.type === 'Box' && !below.children?.length
94  const tiles = keyOf(below) === BAND
95    ? (below.type === 'Box' ? below.children ?? [] : []).flatMap(col => (typeof col !== 'string' && col.type === 'Box' ? col.children ?? [] : [col]))
96    : isEmpty ? [] : [below]
97  const all = [...tiles, tile].sort((a, b) => rank(a) - rank(b))
98  const [lead, left, right] = (['lead', 'left', 'right'] as const).map(slot => all.filter(n => (SLOTS[keyOf(n)] ?? 'left') === slot))
99  return (
100    <Box key={BAND} flexDirection="row" flexWrap="wrap" columnGap={3}>
101      {lead!.length > 0 && <Box key="modz-lead" flexDirection="column" justifyContent="center">{lead}</Box>}
102      {left!.length > 0 && <Box key="modz-left" flexDirection="column" flexGrow={right!.length > 0 ? 1 : 0}>{left}</Box>}
103      {right!.length > 0 && <Box key="modz-right" flexDirection="column" alignItems="flex-end">{right}</Box>}
104    </Box>
105  )
106}
107
108
109export const register: Register = on => {
110  on('session.start', async ($, e, next) => {
111    await $.command.register({
112      name: 'orch',
113      description: 'Orchestrator settings: on, off, max=N, ctx=200k, models=sonnet,haiku, save, reset',
114      argumentHint: '[on|off] [max=N] [ctx=N] [models=a,b] [save|reset]',
115    })
116    const held = await $.state.get(CONFIG_REF)
117    if (held.version === 0) {
118      const saved = (await $.store.get('defaults')) as Config | undefined
119      if (saved) await update($, config, () => saved)
120    }
121    return next(e)
122  })
123
124  on('command.run', { command: 'orch' }, async ($, e) => {
125    const words = e.args.trim()
126    if (words === 'save') {
127      await $.store.set('defaults', await read($, config))
128      return { text: `Saved as default: ${summary(await read($, config))}` }
129    }
130    if (words === 'reset') {
131      const saved = ((await $.store.get('defaults')) as Config | undefined) ?? FALLBACK
132      await update($, config, () => saved)
133      return { text: `Reset: ${summary(saved)}` }
134    }
135    const { next, bad } = parse(words, await read($, config))
136    await update($, config, () => next)
137    return { text: `Orchestrator ${summary(next)}${bad.length ? `\nIgnored: ${bad.join(' ')}` : ''}` }
138  })
139
140  on('prompt.submit', async ($, e, next) => {
141    const c = await read($, config)
142    if (!c.isOn || !USER_ORIGINS.has(e.origin.kind)) return next(e)
143    await update($, spawned, () => 0)
144    return next({ ...e, context: [...(e.context ?? []), brief(c)] })
145  }).catch(($, e, next) => next(e))
146
147  on('agent.spawn', async ($, e, next) => {
148    const c = await read($, config)
149    if (!c.isOn || e.workflow) return next(e)
150    if (e.parentAgentId !== undefined) {
151      return { deny: 'Orchestrator mode: only the main agent may spawn subagents. Do this work yourself.' }
152    }
153    if (!e.fork && !c.models.some(m => e.model?.includes(m))) {
154      return { deny: `Orchestrator mode: set model to one of ${c.models.join(', ')}.` }
155    }
156    if ((await read($, spawned)) >= c.maxAgents) {
157      return { deny: `Orchestrator mode: the limit of ${c.maxAgents} subagents for this request is reached.` }
158    }
159    const ran = await next(e)
160    if (ran.deny === undefined) await update($, spawned, n => n + 1)
161    if (ran.agentId) await track($, ran.agentId, true)
162    return ran
163  }).catch(($, e, next) => next(e))
164
165  on('turn.complete', async ($, e, next) => {
166    if (e.agentId && (await read($, running)).includes(e.agentId)) await track($, e.agentId, false)
167    return next(e)
168  })
169
170  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
171    const below = await next(e)
172    if (e.props.hasSurvey) return below
173    const c = await read($, config)
174    const elements = $.ui.resolve(e)
175    if (!('Select' in elements)) return below
176    const { Box, Button, Input, Select, Text } = elements
177    const edit = await read($, editing)
178    const open = await read($, isOpen)
179    const [busy, turn, menu] = await Promise.all([read($, running), read($, spin), read($, openMenu)])
180    const set = (fn: (c: Config) => Config) => update($, config, fn)
181    const toggle = (m: Model) =>
182      set(c => {
183        const models = MODELS.filter(x => (x === m) !== c.models.includes(x))
184        return models.length ? { ...c, models } : c
185      })
186
187    const LABEL_WIDTH = Math.max(...['Models', 'Max agents', 'Context limit'].map(l => l.length))
188    const row = (label: string, ...controls: RenderChildren[]) => (
189      <Box flexDirection="row" flexWrap="wrap" columnGap={1} paddingLeft={2}>
190        <Text dimColor>{label.padEnd(LABEL_WIDTH)}</Text>
191        {controls}
192      </Box>
193    )
194
195    // In the terminal a Select answers only the keyboard, so there a dropdown is a button that opens a row of
196    // option buttons a click can pick from; every other surface keeps its native Select.
197    const dropdown = (key: string, value: string | undefined, options: readonly SelectOption[], onSelect: (v: string) => unknown) => {
198      if (e.surface !== 'terminal') return <Select key={key} value={value} options={options} onSelect={onSelect} />
199      const isShown = menu === key
200      const labelOf = (v: string | undefined) => options.find(o => o.value === v)?.label ?? v ?? '…'
201      return (
202        <Box flexDirection="row" columnGap={1}>
203          <Button key={key} plain label={`${labelOf(value)} ${isShown ? '▴' : '▾'}`} onPress={() => update($, openMenu, m => (m === key ? null : key))} />
204          {isShown && options.filter(o => o.value !== value).map(o => (
205            <Button key={`${key}:${o.value}`} plain dimColor label={o.label ?? o.value} onPress={async () => {
206              await update($, openMenu, () => null)
207              await onSelect(o.value)
208            }} />
209          ))}
210        </Box>
211      )
212    }
213
214    const numberPicker = (key: Field, value: number, presets: number[], label: (n: number) => string, apply: (n: number) => Config) =>
215      edit === key ? (
216        <Input
217          key={`${key}-custom`}
218          autoFocus
219          placeholder={key === 'max' ? 'number of agents' : 'tokens, e.g. 300k or 1m'}
220          submitLabel="Set"
221          onSubmit={async v => {
222            const n = toNumber(v)
223            if (!(n >= 1)) return $.ui.toast(`Not a number: ${v}`)
224            await set(() => apply(key === 'max' ? Math.floor(n) : n))
225            await update($, editing, () => null)
226          }}
227        />
228      ) : (
229        dropdown(
230          key,
231          String(value),
232          [...withValue(presets, value).map(n => ({ value: String(n), label: label(n) })), { value: CUSTOM, label: 'Custom…' }],
233          v => (v === CUSTOM ? update($, editing, () => key) : set(() => apply(Number(v)))),
234        )
235      )
236
237    return pack(Box, below, (
238      <Box key="orchestrator" flexDirection="column">
239        <Box flexDirection="row" flexWrap="wrap" columnGap={3}>
240          {dropdown(
241            'orch',
242            c.isOn ? 'on' : 'off',
243            [
244              { value: 'on', label: 'Orchestrator on' },
245              { value: 'off', label: 'Orchestrator off' },
246              { value: SETTINGS, label: open ? 'Hide settings' : 'Settings…' },
247            ],
248            v => (v === SETTINGS ? update($, isOpen, o => !o) : set(c => ({ ...c, isOn: v === 'on' }))),
249          )}
250          {busy.length > 0 && <Text color="planMode">{SPINNER[turn % SPINNER.length]} {busy.length} running</Text>}
251          {open && (
252            <Button key="save" plain dimColor label="Save as default" onPress={async () => {
253              await $.store.set('defaults', c)
254              $.ui.toast(`Saved: ${summary(c)}`)
255            }} />
256          )}
257        </Box>
258        {open && c.isOn && row('Models', ...MODELS.map(m => (
259          <Button key={`model-${m}`} plain label={`${c.models.includes(m) ? '☑' : '☐'} ${title(m)}`} dimColor={!c.models.includes(m)} onPress={() => toggle(m)} />
260        )))}
261        {open && c.isOn && row('Max agents', numberPicker('max', c.maxAgents, AGENT_PRESETS, n => `${n} ${n === 1 ? 'agent' : 'agents'}`, n => ({ ...c, maxAgents: n })))}
262        {open && c.isOn && row('Context limit', numberPicker('ctx', c.contextK, CONTEXT_PRESETS, k => `${contextLabel(k)} tokens`, n => ({ ...c, contextK: n })))}
263      </Box>
264    ))
265  })
266}
267

1export type Model = 'opus' | 'sonnet' | 'haiku'
2
3export type Field = 'max' | 'ctx'
4
5export type Config = {
6  isOn: boolean
7  models: Model[]
8  maxAgents: number
9  contextK: number
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    orchestrator: { config: Config; spawned: number; editing: Field | null; isOpen: boolean; running: string[]; spin: number; openMenu: string | null }
15  }
16}
17