Model Router

1import type { On, PluginOptions } from 'claude-code'
2
3import {
4  parseModelMap,
5  parsePreset,
6  presetModel,
7  routedModel,
8  statusText,
9  type Preset,
10} from './routing.js'
11
12export function register(on: On, options: PluginOptions): void {
13  const map = parseModelMap(String(options.subagentModels ?? ''))
14  const models = {
15    fast: String(options.fastModel ?? 'sonnet'),
16    smart: String(options.smartModel ?? 'opus'),
17  }
18  const routed = new Map<string, number>()
19  let preset: Preset = 'off'
20
21  on('session.start', async ($, e, next) => {
22    await $.command.register({
23      name: 'router',
24      description: 'Show model routing, or set the main session to fast, smart or off',
25      argumentHint: '[fast|smart|off]',
26      immediate: true,
27    })
28    return next(e)
29  })
30
31  on('agent.spawn', ($, e, next) => {
32    const model = routedModel(e, map)
33    if (model === undefined) return next(e)
34    routed.set(e.subagentType, (routed.get(e.subagentType) ?? 0) + 1)
35    return next({ ...e, model })
36  })
37
38  on('turn.step', async function* ($, e, next) {
39    const model = e.agentId === undefined ? presetModel(preset, models) : undefined
40    return yield* next(model === undefined ? e : { ...e, model })
41  })
42
43  on('ui.render', { component: 'Spinner' }, ($, e, next) => {
44    const model = presetModel(preset, models)
45    if (model === undefined) return next(e)
46    return next({ ...e, props: { ...e.props, suffix: `${e.props.suffix ?? ''} · ${model}` } })
47  })
48
49  on('command.run', { command: 'router' }, ($, e) => {
50    const choice = parsePreset(e.args)
51    if (choice === 'invalid') return { text: 'Usage: /router [fast|smart|off]' }
52    if (choice !== 'status') {
53      preset = choice
54      $.ui.invalidate('ui.render')
55    }
56    return { text: statusText(preset, models, routed) }
57  })
58}
59

1export type Preset = 'off' | 'fast' | 'smart'
2
3export type PresetModels = { readonly fast: string; readonly smart: string }
4
5export type SpawnFacts = {
6  readonly subagentType: string
7  readonly model?: string
8  readonly fork: boolean
9}
10
11/** `Explore=haiku,Plan=sonnet` → Map { Explore → haiku, Plan → sonnet }. */
12export function parseModelMap(spec: string): ReadonlyMap<string, string> {
13  const map = new Map<string, string>()
14  for (const pair of spec.split(',')) {
15    const [agent, model] = pair.split('=').map(part => part.trim())
16    if (agent && model) map.set(agent, model)
17  }
18  return map
19}
20
21/** The model to give a subagent, or undefined to leave the spawn as it is. */
22export function routedModel(
23  spawn: SpawnFacts,
24  map: ReadonlyMap<string, string>,
25): string | undefined {
26  if (spawn.fork || spawn.model !== undefined) return undefined
27  return map.get(spawn.subagentType)
28}
29
30export function parsePreset(args: string): Preset | 'status' | 'invalid' {
31  const word = args.trim().toLowerCase()
32  if (word === '') return 'status'
33  if (word === 'off' || word === 'fast' || word === 'smart') return word
34  return 'invalid'
35}
36
37export function presetModel(preset: Preset, models: PresetModels): string | undefined {
38  return preset === 'off' ? undefined : models[preset]
39}
40
41export function statusText(
42  preset: Preset,
43  models: PresetModels,
44  routed: ReadonlyMap<string, number>,
45): string {
46  const main = preset === 'off' ? 'main session: unchanged' : `main session: ${preset} (${models[preset]})`
47  const counts = [...routed].map(([type, n]) => `${type} ×${n}`).join(', ')
48  return `${main} · subagents routed: ${counts || 'none yet'}`
49}
50