Model Router

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Route, Tier } from '../types'
5import { AGENT_DEFAULTS, classify } from './classify'
6
7const enabled = atom({ plugin: 'model-router', key: 'enabled' } as const, true)
8const last = atom({ plugin: 'model-router', key: 'last' } as const, null as Route | null)
9const overrides = atom({ plugin: 'model-router', key: 'overrides' } as const, {} as Record<string, Tier>)
10
11const TIERS: readonly Tier[] = ['haiku', 'sonnet', 'opus']
12const NAME: Record<Tier, string> = { haiku: 'Haiku', sonnet: 'Sonnet', opus: 'Opus' }
13const isTier = (v: unknown): v is Tier => typeof v === 'string' && (TIERS as readonly string[]).includes(v)
14
15const save = async ($: EngineInterface) => {
16  await $.store.set('enabled', await read($, enabled))
17  await $.store.set('overrides', await read($, overrides))
18}
19
20const statusText = async ($: EngineInterface) => {
21  const on = await read($, enabled)
22  const mine = await read($, overrides)
23  const table = { ...AGENT_DEFAULTS, ...mine }
24  const rows = Object.entries(table).map(([agent, m]) => `  ${agent}: ${NAME[m]}${agent in mine ? ' (yours)' : ''}`)
25  return [
26    `Model router: ${on ? 'Auto' : 'Off'}`,
27    'Subagent models (others are picked from their task):',
28    ...rows,
29    'Change: /router on | off | set <agent> <haiku|sonnet|opus|default>. Start a prompt with ~ to skip routing once.',
30  ].join('\n')
31}
32
33export const register: Register = on => {
34  // The model picked for the current main reply; cleared when the reply ends.
35  let turnModel: Tier | undefined
36
37  on('session.start', async ($, e, next) => {
38    const savedOn = await $.store.get('enabled').catch(() => undefined)
39    const savedOverrides = await $.store.get('overrides').catch(() => undefined)
40    if (typeof savedOn === 'boolean') await update($, enabled, () => savedOn)
41    if (savedOverrides && typeof savedOverrides === 'object') {
42      const clean = Object.fromEntries(
43        Object.entries(savedOverrides as Record<string, unknown>).filter(([, v]) => isTier(v)),
44      ) as Record<string, Tier>
45      await update($, overrides, () => clean)
46    }
47    await $.command.register({ name: 'router', description: 'Show or change the model router: on, off, or set a subagent model' })
48    return next(e)
49  })
50
51  on('command.run', { command: 'router' }, async ($, e) => {
52    const [verb, agent, model] = e.args.trim().split(/\s+/)
53    if (!verb) return { text: await statusText($) }
54    if (verb === 'on' || verb === 'off') {
55      await update($, enabled, () => verb === 'on')
56      await save($)
57      return { text: `Model router: ${verb === 'on' ? 'Auto' : 'Off'}` }
58    }
59    if (verb === 'set' && agent && (isTier(model) || model === 'default')) {
60      await update($, overrides, o => {
61        const copy = { ...o }
62        if (model === 'default') delete copy[agent]
63        else copy[agent] = model
64        return copy
65      })
66      await save($)
67      return { text: model === 'default' ? `${agent} is back to its default.` : `${agent} will run on ${NAME[model as Tier]}.` }
68    }
69    return { text: 'Use /router, /router on, /router off, or /router set <agent> <haiku|sonnet|opus|default>' }
70  })
71
72  // Pick the model for the person's own prompts.
73  on('prompt.submit', async ($, e, next) => {
74    if (e.origin.kind !== 'composer' || !(await read($, enabled))) return next(e)
75    const skip = e.text.startsWith('~')
76    const verdict = skip ? null : classify(e.text)
77    turnModel = verdict?.model
78    // Always overwrite, so a prompt with no pick never shows the previous prompt's model.
79    await update($, last, () => (verdict ? { where: 'prompt', model: verdict.model, reason: verdict.reason } : null))
80    return next(skip ? { ...e, text: e.text.slice(1).trimStart() } : e)
81  }).catch(($, e, next) => next(e))
82
83  // Run each step of the main reply on the picked model. Subagents are routed below.
84  on('turn.step', async function* ($, e, next) {
85    if (turnModel && !e.agentId && e.model !== turnModel) return yield* next({ ...e, model: turnModel })
86    return yield* next(e)
87  })
88
89  on('turn.complete', (_$, e, next) => {
90    if (!e.agentId) turnModel = undefined
91    return next(e)
92  })
93
94  // Give every subagent a model unless Claude chose one on purpose.
95  on('tool.call', { tool: 'Agent' }, async ($, e, next) => {
96    if (e.model || e.subagent_type === 'fork' || !(await read($, enabled))) return next(e)
97    const type = e.subagent_type ?? 'general-purpose'
98    const mine = await read($, overrides)
99    const fixed = mine[type] ?? AGENT_DEFAULTS[type]
100    // The short description says what kind of task it is; the full prompt is the fallback.
101    const verdict = fixed ? { model: fixed, reason: `${type} default` } : (classify(e.description) ?? classify(e.prompt))
102    if (!verdict) return next(e)
103    await update($, last, () => ({ where: type, model: verdict.model, reason: verdict.reason }))
104    return next({ ...e, model: verdict.model })
105  }).catch(($, e, next) => next(e))
106
107  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
108    const rest = await next(e)
109    if (e.props.hasSurvey) return rest
110
111    const { Box, Select } = $.ui.resolve(e)
112    const isOn = await read($, enabled)
113    const route = await read($, last)
114
115    return (
116      <Box flexDirection="row" flexWrap="wrap" alignItems="center" columnGap={2}>
117        <Select
118          key="router"
119          label="Router"
120          options={[
121            // The picked model rides in the option label so the strip stays one row.
122            { value: 'auto', label: route?.where === 'prompt' ? `Auto · ${NAME[route.model]}` : 'Auto' },
123            { value: 'off', label: 'Off' },
124          ]}
125          value={isOn ? 'auto' : 'off'}
126          onSelect={async v => {
127            await update($, enabled, () => v === 'auto')
128            await save($)
129          }}
130        />
131        {rest}
132      </Box>
133    )
134  })
135}
136

1import type { Tier } from '../types'
2
3export type Verdict = { model: Tier; reason: string } | null
4
5type Rule = readonly [RegExp, string]
6
7// Deep work: reasoning across code or systems, or anything high-stakes.
8const DEEP: readonly Rule[] = [
9  [/\b(debug|root cause|stack ?trace|flaky|regression|race condition|deadlock|memory leak)\b/, 'debugging'],
10  [/\bwhy (is|does|do|did|are)\b.*\b(fail|failing|broken|crash|slow|wrong|error)/, 'debugging'],
11  [/\b(architect|architecture|system design|trade-?offs?|redesign|rfc|design doc)\b/, 'architecture'],
12  [/\b(refactor|migrat(e|ion)|rewrite|overhaul)\b/, 'large change'],
13  [/\b(security|vulnerab|audit|threat model)\b/, 'security'],
14  [/\b(investigate|deep dive|research|analy[sz]e|compare)\b/, 'research and analysis'],
15  [/\b(plan|strategy|roadmap)\b/, 'planning'],
16  [/\b(optimi[sz]e|performance|bottleneck)\b/, 'performance'],
17  [/\breview (the |this |my )?(code|pr|pull request|diff|branch)\b/, 'code review'],
18]
19
20// Mechanical work: short, well-defined, little judgement.
21const QUICK: readonly Rule[] = [
22  [/\b(typo|rename|reformat|format|lint|prettier)\b/, 'mechanical edit'],
23  [/\b(commit message|git (add|commit|push|pull|status|log|checkout|branch)|bump (the )?version)\b/, 'git chore'],
24  [/^(list|show|print|count|sort|find|locate) /, 'quick lookup'],
25  [/\b(list|find|locate|count) (all )?(the )?(files|folders|directories|usages|references)\b/, 'quick lookup'],
26  [/^(what's|what is|what are|where is|which) /, 'quick question'],
27  [/\b(convert|translate) /, 'conversion'],
28]
29
30// Everyday coding.
31const STANDARD: readonly Rule[] = [
32  [/\b(implement|add|create|write|build|make|update|change|edit|fix|test|function|component|endpoint|script|feature|bug)\b/, 'coding'],
33]
34
35const match = (rules: readonly Rule[], text: string) => rules.find(([re]) => re.test(text))?.[1]
36
37/** Picks a model tier for a piece of work, or null when the text gives no clear signal. */
38export const classify = (raw: string): Verdict => {
39  const text = raw.trim().toLowerCase()
40  if (!text) return null
41
42  const deep = match(DEEP, text)
43  if (deep) return { model: 'opus', reason: deep }
44  if (text.length > 2000) return { model: 'opus', reason: 'long, detailed request' }
45
46  const quick = match(QUICK, text)
47  if (quick && text.length < 200) return { model: 'haiku', reason: quick }
48
49  const standard = match(STANDARD, text)
50  if (standard) return { model: 'sonnet', reason: standard }
51
52  return null
53}
54
55/** Default models for known subagent types. Anything else is classified from its task. */
56export const AGENT_DEFAULTS: Record<string, Tier> = {
57  'swarm:researcher': 'sonnet',
58  'swarm:data-analyst': 'sonnet',
59  'swarm:docs-writer': 'sonnet',
60  'swarm:slides-designer': 'sonnet',
61  'swarm:video-maker': 'sonnet',
62  'swarm:image-maker': 'haiku',
63  Explore: 'haiku',
64  'claude-code-guide': 'haiku',
65  Plan: 'opus',
66}
67

1export type Tier = 'haiku' | 'sonnet' | 'opus'
2export type Route = { where: string; model: Tier; reason: string }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'model-router': { enabled: boolean; last: Route | null; overrides: Record<string, Tier> }
7  }
8}
9