smart-model-router

1import type { AgentSpawnInput, EngineInterface, Register } from 'claude-code'
2
3import type { RoutedModel, RouteRecord, RouterMode } from '../types'
4
5const MODE = { plugin: 'smart-model-router', key: 'mode' } as const
6const ROUTES = { plugin: 'smart-model-router', key: 'routes' } as const
7
8async function getMode($: EngineInterface): Promise<RouterMode> {
9  return (await $.state.get(MODE)).value ?? 'off'
10}
11
12async function getRoutes($: EngineInterface): Promise<RouteRecord[]> {
13  return (await $.state.get(ROUTES)).value ?? []
14}
15
16async function editRoutes($: EngineInterface, fn: (list: RouteRecord[]) => RouteRecord[]) {
17  await $.state.set(ROUTES, fn(await getRoutes($)))
18}
19
20// The mode survives sessions through the plugin's store.
21const STORE_MODE = 'mode'
22const MAX_ROUTES = 200
23const LOG_LINES = 15
24
25const MODES: readonly RouterMode[] = ['off', 'sonnet', 'haiku', 'both']
26const ARGS = 'on|off|sonnet|haiku|both|status|log|clear'
27
28const MODE_LABEL: Record<RouterMode, string> = {
29  off: 'off — subagents use their own or the parent model',
30  sonnet: 'sonnet — every subagent runs on Sonnet',
31  haiku: 'haiku — every subagent runs on Haiku',
32  both: 'both — Haiku for lookups and inspection, Sonnet for building and reasoning',
33}
34
35// Agent types that only read and report: Haiku handles them in `both`.
36const LIGHT_AGENTS = new Set(['Explore', 'claude-code-guide', 'statusline-setup'])
37// Agent types that plan or judge: Sonnet in `both`.
38const HEAVY_AGENTS = new Set(['Plan'])
39
40// For general-purpose and plugin agents, the task's own words decide.
41const HEAVY_WORDS =
42  /\b(implement|write|build|create|refactor|fix|debug|design|architect|migrate|rewrite|review|plan|optimi[sz]e|edit|modify|test strategy|diagnos)/i
43const LIGHT_WORDS =
44  /\b(find|search|locate|list|inspect|look up|lookup|grep|read|scan|summari[sz]e|count|check|explore|map|survey|where|which|what)/i
45
46function classify(e: AgentSpawnInput): { routed: RoutedModel; reason: string } {
47  if (LIGHT_AGENTS.has(e.subagentType)) {
48    return { routed: 'haiku', reason: `${e.subagentType} is read-only` }
49  }
50  if (HEAVY_AGENTS.has(e.subagentType)) {
51    return { routed: 'sonnet', reason: `${e.subagentType} needs reasoning` }
52  }
53
54  const task = `${e.description}\n${e.prompt}`
55  const heavy = HEAVY_WORDS.exec(task)
56  if (heavy) {
57    return { routed: 'sonnet', reason: `task says "${heavy[1].toLowerCase()}"` }
58  }
59  const light = LIGHT_WORDS.exec(task)
60  if (light) {
61    return { routed: 'haiku', reason: `task says "${light[1].toLowerCase()}"` }
62  }
63
64  return { routed: 'sonnet', reason: 'no lookup signal; default to Sonnet' }
65}
66
67function route(m: RouterMode, e: AgentSpawnInput): { routed: RoutedModel; reason: string } {
68  if (m === 'both') {
69    return classify(e)
70  }
71  return { routed: m as RoutedModel, reason: `mode ${m}` }
72}
73
74function tally(list: readonly RouteRecord[]) {
75  const haiku = list.filter(r => r.routed === 'haiku').length
76  const sonnet = list.length - haiku
77  const running = list.filter(r => r.isRunning).length
78  return { haiku, sonnet, running }
79}
80
81async function showStatus($: EngineInterface) {
82  const m = await getMode($)
83
84  if (m === 'off') {
85    $.ui.status(undefined)
86    return
87  }
88
89  const { haiku, sonnet, running } = tally(await getRoutes($))
90  const live = running > 0 ? ` • ${running} running` : ''
91  $.ui.status(`router ${m} • ${sonnet} sonnet / ${haiku} haiku${live}`)
92}
93
94// What the model is told when the mode changes, so it plans around it.
95function modelNote(m: RouterMode) {
96  if (m === 'off') {
97    return 'The model router is off: subagents run on the model you pass or their default.'
98  }
99  return [
100    `The model router is on (mode: ${m}). Every subagent you start with the Agent tool is assigned a model by the router; any \`model\` you pass is overridden.`,
101    'So delegate freely: hand self-contained lookups, inspections and summaries to subagents, and start independent subagents in parallel (several Agent calls in one message) rather than one after another.',
102  ].join(' ')
103}
104
105function formatLog(list: readonly RouteRecord[]) {
106  if (list.length === 0) {
107    return 'No subagents routed yet.'
108  }
109  return list
110    .slice(-LOG_LINES)
111    .map(r => {
112      const time = new Date(r.at).toISOString().slice(11, 19)
113      const state = r.isRunning ? 'running' : 'done'
114      const asked = r.requested && r.requested !== r.routed ? ` (asked ${r.requested})` : ''
115      return `${time}  ${r.routed.padEnd(6)} ${r.subagentType}: ${r.description} — ${r.reason}${asked} [${state}]`
116    })
117    .join('\n')
118}
119
120async function setMode($: EngineInterface, m: RouterMode) {
121  await $.state.set(MODE, m)
122  await $.store.set(STORE_MODE, m)
123  await showStatus($)
124}
125
126export const register: Register = on => {
127  on('session.start', async ($, e, next) => {
128    const result = await next(e)
129
130    const saved = await $.store.get(STORE_MODE)
131    if (typeof saved === 'string' && (MODES as readonly string[]).includes(saved)) {
132      await $.state.set(MODE, saved as RouterMode)
133    }
134
135    await $.command.register({
136      name: 'router',
137      description: 'Smart Model Router: choose which model subagents run on',
138      argumentHint: `[${ARGS}]`,
139      immediate: true,
140    })
141    await showStatus($)
142
143    return result
144  })
145
146  on('command.run', { command: 'router' }, async ($, e) => {
147    const arg = e.args.trim().toLowerCase()
148    const current = await getMode($)
149
150    if (arg === '' || arg === 'status') {
151      const { haiku, sonnet, running } = tally(await getRoutes($))
152      return {
153        text: [
154          `Router: ${MODE_LABEL[current]}`,
155          `Routed this session: ${sonnet} to Sonnet, ${haiku} to Haiku, ${running} running.`,
156          `Usage: /router ${ARGS}`,
157        ].join('\n'),
158      }
159    }
160
161    if (arg === 'log') {
162      return { text: formatLog(await getRoutes($)) }
163    }
164
165    if (arg === 'clear') {
166      await $.state.set(ROUTES, [])
167      await showStatus($)
168      return { text: 'Router log cleared.' }
169    }
170
171    const chosen: RouterMode | undefined =
172      arg === 'on' ? (current === 'off' ? 'both' : current) : MODES.find(m => m === arg)
173
174    if (chosen === undefined) {
175      return { text: `Unknown option "${arg}". Usage: /router ${ARGS}` }
176    }
177
178    await setMode($, chosen)
179    return { text: `Router: ${MODE_LABEL[chosen]}`, context: [modelNote(chosen)] }
180  })
181
182  // The routing itself: the Agent tool's spawn, before its model resolves.
183  on('agent.spawn', async ($, e, next) => {
184    const m = await getMode($)
185
186    // Forks always inherit the parent's model; nothing to route.
187    if (m === 'off' || e.fork) {
188      return next(e)
189    }
190
191    const { routed, reason } = route(m, e)
192    const record: RouteRecord = {
193      toolUseId: e.tool_use_id,
194      at: await $.clock.now(),
195      subagentType: e.subagentType,
196      description: e.description,
197      requested: e.model,
198      routed,
199      reason,
200      isRunning: true,
201    }
202    await editRoutes($, list => [...list, record].slice(-MAX_ROUTES))
203
204    const result = await next({ ...e, model: routed })
205
206    if (result.deny !== undefined) {
207      await editRoutes($, list => list.filter(r => r.toolUseId !== record.toolUseId))
208    } else {
209      await editRoutes($, list =>
210        list.map(r => (r.toolUseId === record.toolUseId ? { ...r, agentId: result.agentId } : r)),
211      )
212      $.ui.toast(`Router → ${routed}: ${e.subagentType} "${e.description}"`)
213    }
214
215    await showStatus($)
216    return result
217  })
218
219  // A subagent's last turn marks its route done.
220  on('turn.complete', async ($, e, next) => {
221    const result = await next(e)
222
223    if (e.agentId === undefined) {
224      return result
225    }
226
227    const list = await getRoutes($)
228    if (list.some(r => r.agentId === e.agentId && r.isRunning)) {
229      await editRoutes($, all =>
230        all.map(r => (r.agentId === e.agentId ? { ...r, isRunning: false } : r)),
231      )
232      await showStatus($)
233    }
234
235    return result
236  })
237}
238

1/** How subagents are routed. `both` (what `on` turns on) picks per task. */
2export type RouterMode = 'off' | 'sonnet' | 'haiku' | 'both'
3
4export type RoutedModel = 'sonnet' | 'haiku'
5
6export type RouteRecord = {
7  /** The Agent tool call the spawn belongs to. */
8  toolUseId: string
9  /** When the spawn was routed, ms since epoch. */
10  at: number
11  subagentType: string
12  description: string
13  /** The model the caller asked for, if any, before routing. */
14  requested?: string
15  /** The model the router assigned. */
16  routed: RoutedModel
17  /** Why: the rule that picked it. */
18  reason: string
19  /** True while the subagent runs; false once its turn completed. */
20  isRunning: boolean
21  /** The subagent's id once started, to match its turn.complete. */
22  agentId?: string
23}
24
25declare module 'claude-code' {
26  interface PluginState {
27    'smart-model-router': {
28      mode: RouterMode
29      routes: RouteRecord[]
30    }
31  }
32}
33