model-router

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, TurnUsage } from 'claude-code'
3
4import type { Mode } from '../types'
5
6// Routes sub-agents only. The main thread stays on whatever /model is set to:
7// switching it mid-conversation costs its prompt cache and risks handing a
8// hard task to a weaker model halfway through.
9
10const TIERS = ['opus', 'sonnet', 'haiku'] as const
11type Tier = (typeof TIERS)[number]
12
13// The model that makes the routing decision.
14const ROUTER_MODEL = 'haiku'
15const ROUTER_TIMEOUT_MS = 8000
16const MAX_TASK_CHARS = 6000
17const RECENT_KEPT = 50
18const STATS_KEY = 'stats'
19
20const RUBRIC = `You choose which model runs a coding sub-agent: the cheapest one that will do the job well.
21Reply with exactly one word: opus, sonnet or haiku.
22
23opus: hard or open-ended work. Architecture and design decisions, multi-file refactors,
24difficult debugging with unclear causes, security-sensitive changes, subtle concurrency or
25performance problems, long multi-step plans, or anything where a wrong answer is costly.
26
27sonnet: normal software work. Writing or changing a feature, fixing a bug with a clear cause,
28writing tests, code review, explaining a moderately complex piece of code, routine refactors.
29
30haiku: small, mechanical or lookup work. Finding files or symbols, listing usages, reading a few
31files and reporting what they say, running a command and summarising its output, renames,
32formatting, one-line edits.
33
34Examples:
35"Find every call site of parseConfig and list the files" -> haiku
36"Read the README and package.json and summarise the build steps" -> haiku
37"Add a --dry-run flag to the deploy script and update its tests" -> sonnet
38"Review this diff for correctness bugs" -> sonnet
39"Work out why the cache returns stale data under concurrent writes" -> opus
40"Design a migration plan from REST to event-driven messaging across the services" -> opus
41
42If unsure between two tiers, pick the higher one.`
43
44type Source = 'router' | 'pinned' | 'requested' | 'inherited'
45
46type Decision = {
47  at: number
48  agentId: string
49  subagentType: string
50  description: string
51  model: string
52  parentModel: string
53  source: Source
54}
55
56type TierTotals = {
57  spawns: number
58  input: number
59  output: number
60  cacheRead: number
61  cacheWrite: number
62}
63
64type Stats = { totals: Partial<Record<Tier, TierTotals>>; recent: Decision[] }
65
66const modeAtom = atom({ plugin: 'model-router', key: 'mode' } as const, 'auto' as Mode)
67
68const isTier = (s: string): s is Tier => (TIERS as readonly string[]).includes(s)
69const isMode = (s: string): s is Mode => s === 'auto' || s === 'off' || isTier(s)
70const tierOf = (model: string): Tier | undefined =>
71  TIERS.find(tier => model.toLowerCase().includes(tier))
72
73const emptyTotals = (): TierTotals => ({ spawns: 0, input: 0, output: 0, cacheRead: 0, cacheWrite: 0 })
74
75const loadStats = async ($: EngineInterface): Promise<Stats> => {
76  const raw = (await $.store.get(STATS_KEY)) as Stats | undefined
77
78  return raw ?? { totals: {}, recent: [] }
79}
80
81// Stats are best effort: a failed write never gets in the way of a spawn.
82const changeStats = async ($: EngineInterface, fn: (stats: Stats) => void) => {
83  try {
84    const stats = await loadStats($)
85    fn(stats)
86    await $.store.set(STATS_KEY, stats)
87  } catch {
88    // ignored
89  }
90}
91
92// Which tier each running sub-agent was started on, for its usage at the end.
93const running = new Map<string, Tier>()
94
95const showStatus = ($: EngineInterface, mode: Mode) =>
96  $.ui.status(mode === 'auto' ? undefined : mode === 'off' ? 'route: off' : `route: sub-agents pinned ${mode}`)
97
98// Asks Haiku which tier fits `task`; undefined when it cannot say.
99const decide = async ($: EngineInterface, task: string): Promise<Tier | undefined> => {
100  const reply = await $.model
101    .complete({
102      model: ROUTER_MODEL,
103      system: RUBRIC,
104      prompt: `Sub-agent task to route:\n<task>\n${task.slice(0, MAX_TASK_CHARS)}\n</task>`,
105      maxTokens: 5,
106      timeoutMs: ROUTER_TIMEOUT_MS,
107    })
108    .catch(() => undefined)
109  if (reply === undefined || !reply.isAnswered) return undefined
110  const word = reply.text.toLowerCase().match(/\b(opus|sonnet|haiku)\b/)?.[1]
111
112  return word !== undefined && isTier(word) ? word : undefined
113}
114
115const fmt = (n: number) => (n >= 1_000_000 ? `${(n / 1e6).toFixed(1)}M` : n >= 1000 ? `${(n / 1e3).toFixed(1)}k` : `${n}`)
116
117const statsText = (stats: Stats) => {
118  const rows = TIERS.map(tier => {
119    const t = stats.totals[tier] ?? emptyTotals()
120
121    return `  ${tier.padEnd(7)} ${String(t.spawns).padStart(5)} spawns   in ${fmt(t.input).padStart(7)}   out ${fmt(t.output).padStart(7)}   cache read ${fmt(t.cacheRead).padStart(7)}   cache write ${fmt(t.cacheWrite).padStart(7)}`
122  })
123  const recent = stats.recent
124    .slice(-10)
125    .reverse()
126    .map(d => `  ${(tierOf(d.model) ?? d.model).padEnd(7)} ${d.source.padEnd(9)} ${d.subagentType}: ${d.description}`)
127
128  return [
129    'Sub-agent routing stats (all sessions)',
130    ...rows,
131    '',
132    recent.length === 0 ? 'No sub-agents routed yet.' : 'Latest decisions (model, why, agent):',
133    ...recent,
134  ].join('\n')
135}
136
137const HELP = [
138  'Usage: /route [auto | off | opus | sonnet | haiku | status | stats | reset-stats]',
139  '  auto    let Haiku pick a model for each sub-agent (default)',
140  '  off     stop routing; sub-agents use their usual model',
141  '  opus|sonnet|haiku   run every sub-agent on that model',
142  '  stats   spawns and tokens per model, and the latest decisions',
143].join('\n')
144
145export const register: Register = on => {
146  on('session.start', async ($, e, next) => {
147    await $.command.register({
148      name: 'route',
149      description: 'Sub-agent model router: auto, off, pin a model, or stats',
150      argumentHint: '[auto|off|opus|sonnet|haiku|status|stats|reset-stats]',
151    })
152    showStatus($, await read($, modeAtom))
153
154    return next(e)
155  })
156
157  on('command.run', { command: 'route' }, async ($, e) => {
158    const arg = e.args.trim().toLowerCase()
159    if (arg === 'stats') return { text: statsText(await loadStats($)) }
160    if (arg === 'reset-stats') {
161      await $.store.delete(STATS_KEY)
162
163      return { text: 'Sub-agent routing stats cleared.' }
164    }
165    if (arg === '' || arg === 'status') return { text: `Sub-agent router: ${await read($, modeAtom)}\n${HELP}` }
166    if (!isMode(arg)) return { text: HELP }
167    await update($, modeAtom, () => arg)
168    showStatus($, arg)
169
170    return { text: `Sub-agent router set to ${arg}.` }
171  })
172
173  on('agent.spawn', async ($, e, next) => {
174    const mode = await read($, modeAtom)
175    // Forks always inherit, and a model the caller named explicitly wins.
176    const canRoute = mode !== 'off' && !e.fork && e.model === undefined
177    const tier = !canRoute ? undefined : mode === 'auto' ? await decide($, `${e.description}\n\n${e.prompt}`) : mode
178    const started = await next(tier === undefined ? e : { ...e, model: tier })
179    if (started.deny !== undefined || started.model === undefined) return started
180
181    const source: Source =
182      tier !== undefined ? (mode === 'auto' ? 'router' : 'pinned') : e.model !== undefined ? 'requested' : 'inherited'
183    const ranOn = tierOf(started.model)
184    const parentTier = tierOf(e.parentModel)
185    // A teammate carries a teammateId rather than an agentId.
186    const agentId = started.agentId ?? started.teammateId ?? ''
187    if (ranOn !== undefined && started.agentId !== undefined) running.set(started.agentId, ranOn)
188
189    if (ranOn !== undefined && ranOn !== parentTier) {
190      const why = source === 'router' ? 'router picked' : source === 'pinned' ? 'pinned' : 'requested'
191      $.ui.toast(`Sub-agent "${e.description}" runs on ${ranOn} (${why}; main thread: ${parentTier ?? e.parentModel})`, {
192        timeoutMs: 6000,
193      })
194    }
195
196    const decision: Decision = {
197      at: await $.clock.now().catch(() => 0),
198      agentId,
199      subagentType: e.subagentType,
200      description: e.description,
201      model: started.model,
202      parentModel: e.parentModel,
203      source,
204    }
205    await changeStats($, stats => {
206      stats.recent = [...stats.recent, decision].slice(-RECENT_KEPT)
207      if (ranOn === undefined) return
208      const t = (stats.totals[ranOn] ??= emptyTotals())
209      t.spawns += 1
210    })
211
212    return started
213    // Only work before `next` can throw (the code after it is guarded), so a
214    // failure here never starts the sub-agent twice: it just runs unrouted.
215  }).catch(($, e, next) => next(e))
216
217  // Adds each routed sub-agent's token usage to its tier's totals.
218  on('turn.complete', async ($, e, next) => {
219    const tier = e.agentId === undefined ? undefined : running.get(e.agentId)
220    if (tier !== undefined && e.agentId !== undefined) {
221      running.delete(e.agentId)
222      const usage: TurnUsage | undefined = e.usage
223      if (usage !== undefined) {
224        await changeStats($, stats => {
225          const t = (stats.totals[tier] ??= emptyTotals())
226          t.input += usage.input_tokens
227          t.output += usage.output_tokens
228          t.cacheRead += usage.cache_read_input_tokens
229          t.cacheWrite += usage.cache_creation_input_tokens
230        })
231      }
232    }
233
234    return next(e)
235  })
236}
237

1export type Mode = 'auto' | 'off' | 'opus' | 'sonnet' | 'haiku'
2
3declare module 'claude-code' {
4  interface PluginState {
5    'model-router': { mode: Mode }
6  }
7}
8