model-router

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Decision, Effort, Mode, ModelKey } from '../types'
5import {
6  CLASSIFIER_SYSTEM,
7  EFFORTS,
8  MODELS,
9  bump,
10  clamp,
11  classifierPrompt,
12  classifyHeuristic,
13  colorAt,
14  commandFloor,
15  costUsd,
16  effortAt,
17  labelAt,
18  modelAt,
19  parseClassifier,
20  rungFromSession,
21  rungOf,
22} from './ladder'
23
24const PANE = 'model-router'
25const CACHE_WARM_MS = 5 * 60_000
26const EDIT_TOOLS = new Set(['Edit', 'Write', 'MultiEdit', 'NotebookEdit'])
27
28const mode = atom({ plugin: 'model-router', key: 'mode' } as const, 'auto' as Mode)
29const decision = atom({ plugin: 'model-router', key: 'decision' } as const, null as Decision | null)
30const active = atom({ plugin: 'model-router', key: 'active' } as const, null as number | null)
31const lastTurn = atom({ plugin: 'model-router', key: 'lastTurn' } as const, null as { rung: number; costUsd: number; endedAt: number } | null)
32const sessionUsd = atom({ plugin: 'model-router', key: 'sessionUsd' } as const, 0)
33const pin = atom({ plugin: 'model-router', key: 'pin' } as const, rungOf('opus', 'medium'))
34
35type Options = { classifier?: string; ceiling?: string; stickiness?: number }
36
37async function setMode($: EngineInterface, next: Mode, pinned?: number) {
38  await update($, mode, () => next)
39  await $.store.set('mode', next)
40  if (pinned === undefined) return
41  await update($, pin, () => pinned)
42  await $.store.set('pin', pinned)
43  if (next === 'pinned') await update($, decision, (): Decision => ({ rung: pinned, reason: 'pinned by you', source: 'pinned' }))
44}
45
46/** Raises this turn's rung to at least `floor` (auto mode only; never lowers it). */
47async function raise($: EngineInterface, floor: number, reason: string, ceiling: ModelKey) {
48  if ((await read($, mode)) !== 'auto') return
49  const target = clamp(floor, ceiling)
50  await update($, decision, (d): Decision => (d && d.rung >= target ? d : { rung: target, reason, source: 'escalated' }))
51}
52
53export const register: Register = (on, options) => {
54  const opts = (options ?? {}) as Options
55  const ceiling = (opts.ceiling ?? 'fable') as ModelKey
56  const stickiness = Math.max(0, Number(opts.stickiness ?? 3))
57
58  // Per-turn escalation signals; a reload resetting them is harmless.
59  let failures = 0
60  let edited = new Set<string>()
61
62  on('session.start', async ($, e, next) => {
63    await $.command.register({
64      name: 'route',
65      description: 'Model router: show the ladder, or set auto | off | pin <haiku|sonnet|opus|fable> [effort]',
66      argumentHint: '[auto | off | pin <haiku|sonnet|opus|fable> [effort]]',
67      immediate: true,
68    })
69    const savedMode = (await $.store.get('mode')) as Mode | undefined
70    const savedPin = (await $.store.get('pin')) as number | undefined
71    if (savedMode) await setMode($, savedMode, savedPin)
72    void $.ui.open({ id: PANE, title: 'Model router' })
73
74    return next(e)
75  })
76
77  on('prompt.submit', async ($, e, next) => {
78    const m = await read($, mode)
79    if (m !== 'auto' || e.text.trimStart().startsWith('/')) return next(e)
80
81    const current = await read($, active)
82    let pick = classifyHeuristic(e.text)
83    let source: Decision['source'] = 'heuristic'
84
85    if (opts.classifier !== 'heuristic') {
86      // A refused request (model not allowed, bad args) falls back to the heuristic.
87      const r = await $.model
88        .complete({
89          model: 'claude-haiku-5-5',
90          system: CLASSIFIER_SYSTEM,
91          prompt: classifierPrompt(e.text, current),
92          effort: 'low',
93          maxTokens: 120,
94          timeoutMs: 6000,
95        })
96        .catch(() => null)
97      const graded = r?.isAnswered ? parseClassifier(r.text) : null
98      if (graded) (pick = graded), (source = 'haiku')
99    }
100
101    let rung = clamp(pick.rung, ceiling)
102    let reason = pick.reason
103
104    // A prompt folded into a running turn may only raise the rung.
105    if (e.turnId !== undefined && current !== null) rung = Math.max(rung, current)
106
107    // Switching model or effort re-bills the cached prefix, so small downgrades
108    // aren't worth it while the cache is still warm.
109    const last = await read($, lastTurn)
110    const warm = last !== null && (await $.clock.now()) - last.endedAt < CACHE_WARM_MS
111    if (warm && current !== null && rung < current && current - rung < stickiness) {
112      reason = `cache warm; wanted ${labelAt(rung)} (${reason})`
113      rung = current
114      source = 'held'
115    }
116
117    failures = 0
118    edited = new Set()
119    await update($, decision, () => ({ rung, reason, source }))
120
121    return next(e)
122  }).catch(($, e, next) => next(e))
123
124  on('turn.step', async function* ($, e, next) {
125    // Subagents keep whatever model their definition asked for.
126    if (e.agentId !== undefined) return yield* next(e)
127
128    const m = await read($, mode)
129    const d = await read($, decision)
130    if (m === 'off' || d === null) {
131      const seen = rungFromSession(e.model, e.effort)
132      await update($, active, () => seen)
133      return yield* next(e)
134    }
135
136    const rung = m === 'pinned' ? d.rung : clamp(d.rung, ceiling)
137    if ((await read($, active)) !== rung) {
138      await update($, active, () => rung)
139      $.ui.status(`⇅ ${labelAt(rung)}`)
140    }
141
142    return yield* next({ ...e, model: modelAt(rung).id, effort: effortAt(rung) })
143  })
144
145  on('tool.call', async ($, e, next) => {
146    if (e.agentId !== undefined) return next(e)
147
148    const tool = String(e.tool)
149    if (e.tool === 'Bash') {
150      const floor = commandFloor(e.command)
151      if (floor) await raise($, floor.rung, floor.reason, ceiling)
152    }
153
154    const ran = await next(e)
155
156    if (ran.deny === undefined && ran.isError === true) {
157      failures += 1
158      // Every second failure in a turn, step up to the next sensible rung.
159      if (failures % 2 === 0) {
160        const d = await read($, decision)
161        await raise($, bump(d?.rung ?? 0), `${failures} failed tool calls this turn`, ceiling)
162      }
163    }
164
165    if (EDIT_TOOLS.has(tool) && 'file_path' in e && typeof e.file_path === 'string') {
166      edited.add(e.file_path)
167      if (edited.size === 6) await raise($, rungOf('opus', 'high'), 'editing 6+ files', ceiling)
168    }
169
170    return ran
171  }).catch(($, e, next) => next(e))
172
173  on('turn.complete', async ($, e, next) => {
174    if (e.agentId === undefined && e.usage) {
175      const cost = costUsd(e.usage)
176      const rung = (await read($, active)) ?? 0
177      const endedAt = await $.clock.now()
178      await update($, lastTurn, () => ({ rung, costUsd: cost, endedAt }))
179      await update($, sessionUsd, total => total + cost)
180    }
181
182    return next(e)
183  })
184
185  on('command.run', { command: 'route' }, async ($, e) => {
186    const [verb = '', model, effort] = e.args.trim().toLowerCase().split(/\s+/)
187
188    if (verb === 'auto' || verb === 'off') {
189      await setMode($, verb)
190      return { text: `Model router: ${verb === 'auto' ? 'routing automatically' : 'off, the session model runs as set'}.` }
191    }
192
193    if (verb === 'pin') {
194      if (!MODELS.some(x => x.key === model)) return { text: 'Usage: /route pin <haiku|sonnet|opus|fable> [low|medium|high|xhigh|max]' }
195      const ef = (EFFORTS.includes(effort as Effort) ? effort : 'medium') as Effort
196      const rung = rungOf(model as ModelKey, ef)
197      await setMode($, 'pinned', rung)
198      return { text: `Model router: pinned to ${labelAt(rung)}.` }
199    }
200
201    const d = await read($, decision)
202    const a = await read($, active)
203    await $.ui.open({ id: PANE, title: 'Model router' })
204    const now = a === null ? 'nothing yet' : labelAt(a)
205    return { text: `Model router (${await read($, mode)}): ${now}${d ? ` — ${d.source}: ${d.reason}` : ''}` }
206  })
207
208  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
209    const { Box, Text, Button } = $.ui.resolve(e)
210    const m = await read($, mode)
211    const a = await read($, active)
212    const d = await read($, decision)
213    const last = await read($, lastTurn)
214    const total = await read($, sessionUsd)
215
216    return (
217      <Box flexDirection="column">
218        <Text dimColor>effort  low → max</Text>
219        {MODELS.map((model, mi) => (
220          <Box flexDirection="column" marginTop={1}>
221            <Box>
222              <Text bold={a !== null && modelAt(a).key === model.key}>{model.label}</Text>
223              <Text dimColor>{`  $${model.input}/$${model.output}`}</Text>
224            </Box>
225            <Box>
226              {EFFORTS.map((_, ei) => {
227                const rung = mi * EFFORTS.length + ei
228                const isLit = a !== null && rung <= a
229                const isCurrent = rung === a
230                return (
231                  <Text color={isLit ? colorAt(rung) : undefined} dimColor={!isLit} bold={isCurrent}>
232                    {isCurrent ? '◉ ' : isLit ? '● ' : '○ '}
233                  </Text>
234                )
235              })}
236            </Box>
237          </Box>
238        ))}
239        <Box marginTop={1} flexDirection="column">
240          {a === null ? (
241            <Text dimColor>Waiting for the first turn…</Text>
242          ) : (
243            <Text bold color={colorAt(a)}>
244              ▶ {labelAt(a)}
245            </Text>
246          )}
247          <Text dimColor wrap="wrap">
248            {m}
249            {d ? ` · ${d.source}: ${d.reason}` : ''}
250          </Text>
251          {last && (
252            <Text dimColor>
253              last ${last.costUsd.toFixed(3)} · session ${total.toFixed(2)}
254            </Text>
255          )}
256        </Box>
257        <Box marginTop={1} gap={1}>
258          <Button key="auto" label="auto" variant={m === 'auto' ? 'primary' : 'secondary'} onPress={() => setMode($, 'auto')} />
259          <Button key="off" label="off" variant={m === 'off' ? 'primary' : 'secondary'} onPress={() => setMode($, 'off')} />
260          <Button
261            key="pin"
262            label={m === 'pinned' ? 'pinned' : 'pin here'}
263            variant={m === 'pinned' ? 'primary' : 'secondary'}
264            onPress={() => setMode($, 'pinned', a ?? rungOf('opus', 'medium'))}
265          />
266        </Box>
267      </Box>
268    )
269  })
270}
271

1import type { Effort, ModelKey } from '../types'
2
3export const EFFORTS: readonly Effort[] = ['low', 'medium', 'high', 'xhigh', 'max']
4
5// Cheapest first. Prices are $/MTok (input, output, cache read). Haiku 5.5 bills a higher
6// tier once a prompt passes 100k tokens.
7export const MODELS = [
8  {
9    key: 'haiku', id: 'claude-haiku-5-5', label: 'Haiku 5.5', input: 0.1, output: 0.5, cacheRead: 0.01,
10    long: { input: 0.5, output: 2.5, cacheRead: 0.05 },
11  },
12  { key: 'sonnet', id: 'claude-sonnet-5-5', label: 'Sonnet 5.5', input: 2, output: 10, cacheRead: 0.2 },
13  { key: 'opus', id: 'claude-opus-5-5', label: 'Opus 5.5', input: 4, output: 20, cacheRead: 0.2 },
14  { key: 'fable', id: 'claude-fable-5-1', label: 'Fable 5.1', input: 10, output: 50, cacheRead: 0.25 },
15] as const satisfies readonly {
16  key: ModelKey
17  id: string
18  label: string
19  input: number
20  output: number
21  cacheRead: number
22  long?: { input: number; output: number; cacheRead: number }
23}[]
24
25export type ModelInfo = (typeof MODELS)[number]
26
27export const RUNGS = MODELS.length * EFFORTS.length
28export const TOP = RUNGS - 1
29
30export const rungOf = (model: ModelKey, effort: Effort) =>
31  MODELS.findIndex(m => m.key === model) * EFFORTS.length + EFFORTS.indexOf(effort)
32
33export const modelAt = (rung: number): ModelInfo => MODELS[Math.floor(rung / EFFORTS.length)] ?? MODELS[0]
34export const effortAt = (rung: number): Effort => EFFORTS[rung % EFFORTS.length] ?? 'medium'
35export const labelAt = (rung: number) => `${modelAt(rung).label} · ${effortAt(rung)}`
36
37export const clamp = (rung: number, ceiling: ModelKey = 'fable') =>
38  Math.max(0, Math.min(rung, rungOf(ceiling, 'max')))
39
40/** Maps whatever model id the engine reports back onto the ladder; null when it's off-ladder (e.g. Haiku). */
41export function rungFromSession(modelId: string, effort: unknown): number | null {
42  const m = MODELS.findIndex(x => modelId.includes(x.id) || modelId.includes(x.key))
43  if (m < 0) return null
44  const ei = typeof effort === 'string' ? EFFORTS.indexOf(effort as Effort) : -1
45  return m * EFFORTS.length + (ei < 0 ? 1 : ei)
46}
47
48// Green → orange → red along the ladder.
49export function colorAt(rung: number): string {
50  const green = [0x22, 0xc5, 0x5e]
51  const orange = [0xf5, 0x9e, 0x0b]
52  const red = [0xef, 0x44, 0x44]
53  const t = rung / TOP
54  const [a, b, u] = t < 0.5 ? [green, orange, t * 2] : [orange, red, (t - 0.5) * 2]
55  const hex = a.map((c, i) => Math.round(c + ((b[i] ?? c) - c) * u).toString(16).padStart(2, '0'))
56  return `#${hex.join('')}`
57}
58
59// Rungs worth stopping on when escalating: max effort on a cheaper model is rarely
60// better value than the next model up at medium. Haiku low/medium/high, Sonnet medium/high,
61// Opus medium..xhigh, Fable high..max.
62const STOPS = [0, 1, 2, 6, 7, 11, 12, 13, 17, 18, 19]
63export const bump = (rung: number) => STOPS.find(s => s > rung) ?? TOP
64
65/* ------------------------------------------------------------------ prompts */
66
67type Rule = { re: RegExp; rung: number; why: string }
68
69// Highest matching rung wins. Ordered roughly by rung for readability only.
70const RULES: Rule[] = [
71  { re: /\b(typo|rename|reformat|format this|lint|spelling|bump (the )?version)\b/i, rung: 1, why: 'mechanical edit' },
72  { re: /^(what|who|when|where|which|is|are|does|do|can|how do i)\b.{0,160}\??$/is, rung: 1, why: 'short question' },
73  { re: /\b(explain|summari[sz]e|what does|tl;?dr|describe)\b/i, rung: 1, why: 'explanation' },
74  { re: /\b(git (status|log|diff|add|commit|checkout|branch)|ls|cat|grep|run (the )?(tests?|build))\b/i, rung: 1, why: 'routine command' },
75  { re: /\b(write|add|create|generate)\b.{0,40}\b(script|function|method|test|endpoint|component|query|regex|dockerfile|yaml|readme)\b/i, rung: 7, why: 'single-unit code' },
76  { re: /\b(fix|debug)\b.{0,40}\b(bug|error|exception|failing|crash|warning)\b/i, rung: 7, why: 'targeted fix' },
77  { re: /\b(implement|build|feature|integrate|wire up|hook up|add support for)\b/i, rung: 11, why: 'feature work' },
78  { re: /\b(agent(ic)?|llm|rag|embedding|tool[- ]?use|mcp server|semantic kernel|ai foundry|azure openai)\b/i, rung: 11, why: 'LLM/agent system' },
79  { re: /\b(refactor|restructure|across (the|all|every)|multi[- ]?file|whole (codebase|repo|project)|migrate|migration|upgrade to)\b/i, rung: 12, why: 'multi-file / migration' },
80  { re: /\b(architect(ure)?|design (the|a) system|trade-?offs?|scalab|security|auth(n|z|entication|orization)?|vulnerab|performance|optimi[sz]e|memory leak)\b/i, rung: 12, why: 'design / security / perf' },
81  { re: /\b(race condition|deadlock|concurren|thread[- ]safe|distributed|consensus|flaky|intermittent|heisenbug|non-?determin)/i, rung: 13, why: 'concurrency / nondeterminism' },
82  { re: /\b(from scratch|end[- ]to[- ]end|entire (app|system|platform)|production[- ]ready|long[- ]running|until (it|all) (works|pass))/i, rung: 13, why: 'long-horizon build' },
83  { re: /\b(prove|proof|formal(ly)?|invariant|novel algorithm|compiler|type system|cryptograph|research)\b/i, rung: 17, why: 'deep reasoning' },
84]
85
86/** Scores a prompt on the ladder from its wording and shape alone. */
87export function classifyHeuristic(text: string): { rung: number; reason: string } {
88  const t = text.trim()
89  let rung = t.length < 40 ? 1 : 6
90  const why: string[] = []
91
92  for (const r of RULES) {
93    if (r.re.test(t) && r.rung >= rung) {
94      if (r.rung > rung) why.length = 0
95      rung = r.rung
96      why.push(r.why)
97    }
98  }
99
100  // Shape: long specs, pasted code and stack traces all mean more to hold in mind.
101  const codeBlocks = (t.match(/```/g) ?? []).length / 2
102  const hasTrace = /(at [\w.$<>]+\(.*:\d+\)|Traceback \(most recent|Exception:|panicked at|\bE\d{4}\b|CS\d{4})/.test(t)
103  if (t.length > 2500 || codeBlocks >= 2) (rung = bump(rung)), why.push('long / code-heavy prompt')
104  if (hasTrace && rung < 11) (rung = Math.max(rung, 7)), why.push('stack trace')
105
106  // Explicit user dials.
107  if (/\bultrathink\b|\bthink (really|very) hard\b/i.test(t)) (rung = Math.max(rung, 18)), why.push('asked for max thinking')
108  else if (/\b(think hard|carefully|thorough(ly)?|deep dive|double[- ]check)\b/i.test(t)) (rung = bump(rung)), why.push('asked for care')
109  if (/\b(quick(ly)?|just|briefly|tiny|small|simple)\b/i.test(t) && rung > 1) (rung = Math.max(0, rung - (rung >= 11 ? 4 : 1))), why.push('asked for quick')
110
111  return { rung, reason: why.length ? why.join(', ') : 'general chat' }
112}
113
114/* ---------------------------------------------------------- Haiku grading */
115
116export const CLASSIFIER_SYSTEM = `You route a coding assistant's next turn to the cheapest model and reasoning effort that will do it well.
117Models, cheapest first: haiku (Haiku 5.5, $0.10/$0.50 per MTok — fastest; classification, lookups, short answers, mechanical edits), sonnet (Sonnet 5.5, $2/$10 per MTok — strong everyday coding), opus (Opus 5.5, $4/$20 — hard multi-file engineering, design, subtle debugging), fable (Fable 5.1, $10/$50 — only for the hardest reasoning, novel algorithms, long-horizon autonomous builds, or after cheaper models failed).
118Effort: low | medium | high | xhigh | max. Prefer a stronger model at lower effort over a weaker model at max.
119Guide: chat, lookups, short explanations, git/shell one-liners, mechanical edits -> haiku low/medium (haiku skips checks at low, so use medium for anything with tool use). Longer explanations, small scripts -> haiku high or sonnet medium. Single-file edits, scripts, small fixes, tests -> sonnet high. Features touching a few files, LLM/agent code, non-trivial debugging -> opus medium. Refactors, migrations, architecture, security, performance -> opus high. Concurrency, flaky failures, end-to-end builds -> opus xhigh. Proofs, novel algorithms, repeated failure -> fable high/xhigh.
120Reply with JSON only: {"model":"haiku|sonnet|opus|fable","effort":"low|medium|high|xhigh|max","reason":"<= 8 words"}`
121
122export function classifierPrompt(text: string, current: number | null): string {
123  const clipped = text.length > 6000 ? `${text.slice(0, 3000)}\n…\n${text.slice(-3000)}` : text
124  const now = current === null ? 'unknown' : labelAt(current)
125  return `Currently running: ${now}.\n\n<next_user_message>\n${clipped}\n</next_user_message>`
126}
127
128export function parseClassifier(reply: string): { rung: number; reason: string } | null {
129  const json = reply.match(/\{[\s\S]*\}/)?.[0]
130  if (!json) return null
131  try {
132    const v = JSON.parse(json) as { model?: string; effort?: string; reason?: string }
133    if (!MODELS.some(m => m.key === v.model) || !EFFORTS.includes(v.effort as Effort)) return null
134    return { rung: rungOf(v.model as ModelKey, v.effort as Effort), reason: String(v.reason ?? '').slice(0, 80) || 'graded' }
135  } catch {
136    return null
137  }
138}
139
140/* --------------------------------------------------------- shell commands */
141
142const RISKY = /\b(terraform (apply|destroy)|kubectl (apply|delete|rollout)|helm (upgrade|install|uninstall)|az (deployment|webapp|functionapp|aks|sql)\b|func azure functionapp publish|dotnet ef database update|prisma migrate|alembic upgrade|drop (table|database)|rm -rf|git push (-f|--force)|docker push)/i
143
144/** Rung floor a shell command implies, with why; null when the command says nothing. */
145export function commandFloor(cmd: string): { rung: number; reason: string } | null {
146  if (RISKY.test(cmd)) return { rung: rungOf('opus', 'high'), reason: `risky command: ${cmd.slice(0, 32)}` }
147  return null
148}
149
150/* -------------------------------------------------------------------- cost */
151
152type Usage = {
153  input_tokens?: number
154  output_tokens?: number
155  cache_creation_input_tokens?: number
156  cache_read_input_tokens?: number
157  model?: string
158}
159
160export function costUsd(u: Usage): number {
161  const base = MODELS.find(x => u.model?.includes(x.id)) ?? MODELS.find(x => x.key === 'opus') ?? MODELS[0]
162  const promptTokens = (u.input_tokens ?? 0) + (u.cache_creation_input_tokens ?? 0) + (u.cache_read_input_tokens ?? 0)
163  const m = 'long' in base && promptTokens > 100_000 ? { ...base, ...base.long } : base
164  const mtok =
165    (u.input_tokens ?? 0) * m.input +
166    (u.cache_creation_input_tokens ?? 0) * m.input * 1.25 +
167    (u.cache_read_input_tokens ?? 0) * m.cacheRead +
168    (u.output_tokens ?? 0) * m.output
169  return mtok / 1e6
170}
171

1export type ModelKey = 'haiku' | 'sonnet' | 'opus' | 'fable'
2export type Effort = 'low' | 'medium' | 'high' | 'xhigh' | 'max'
3export type Mode = 'auto' | 'pinned' | 'off'
4
5export type Decision = {
6  /** Rung on the 20-step ladder: model-major, effort-minor, cheapest first. */
7  rung: number
8  reason: string
9  source: 'heuristic' | 'haiku' | 'pinned' | 'escalated' | 'held' | 'session'
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    'model-router': {
15      mode: Mode
16      decision: Decision | null
17      /** The rung the last main-loop request actually ran on. */
18      active: number | null
19      lastTurn: { rung: number; costUsd: number; endedAt: number } | null
20      sessionUsd: number
21      /** Rung the person pinned with /route pin. */
22      pin: number
23    }
24  }
25}
26