model-router

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { RouteDecision, RouterStats } from '../types'
5
6/**
7 * model-router: quick questions go to Claude Haiku 5.5, everything else stays put.
8 *
9 * Each turn's prompt is classified once, at `turn.start`, by the engine's small
10 * fast model. A turn judged a plain question that can be answered in prose has
11 * its model requests (`turn.step`) renamed to Haiku, at low effort. Every doubt
12 * resolves toward the person's own model:
13 *
14 *  - the classifier failing, or answering neither label, keeps the turn on it;
15 *  - a long prompt, or a slash command, is never classified at all;
16 *  - a subagent keeps the model it was spawned with;
17 *  - a conversation past maxContextTokens is never switched (below);
18 *  - and the moment Haiku reaches for a tool, the turn was not as simple as it
19 *    looked, so its remaining steps go back to the person's model.
20 *
21 * WHY THE GATE IS TOKENS, AND WHY HAIKU ALONE. The prompt cache is per model,
22 * so a routed turn sends the whole conversation to Haiku uncached. At list
23 * prices that is still cheaper than the person's model reading its own cache:
24 * Haiku 5.5 input is $0.10/M, against $0.20/M for an Opus 5.5 or Sonnet 5.5
25 * cache read. But Haiku's input price steps to $0.50/M past 100K prompt tokens,
26 * at which point the cached read wins, so 100K is the default gate. Sonnet 5.5
27 * was dropped as a target: uncached at $2.00/M it is ten times an Opus cache
28 * read, so it saved money only in the first few turns of a conversation.
29 */
30
31export const SIMPLE = 'quick question or chat, answered in prose'
32export const COMPLEX = 'task needing tools, files, code or careful judgement'
33export const CHEAP_MODEL = 'claude-haiku-5-5'
34
35/** Past this, a prompt is real work by its length alone: not worth a classify call. */
36export const LONG_PROMPT = 600
37
38const decision = atom({ plugin: 'model-router', key: 'decision' } as const, {
39  turnId: '',
40  route: 'main',
41  reason: 'no turn yet',
42  escalated: false,
43} as RouteDecision)
44const stats = atom({ plugin: 'model-router', key: 'stats' } as const, {
45  turns: 0,
46  routed: 0,
47  escalated: 0,
48} as RouterStats)
49const isEnabled = atom({ plugin: 'model-router', key: 'isEnabled' } as const, true)
50const lastContextTokens = atom({ plugin: 'model-router', key: 'lastContextTokens' } as const, 0)
51
52const short = (model: string) => model.replace(/^claude-/, '')
53const thousands = (n: number) => `${Math.round(n / 1000)}K`
54
55/** Which model a turn should run on, and the reason, judged once at its start. */
56export async function decide(
57  $: EngineInterface,
58  text: string,
59  maxContextTokens: number,
60): Promise<Pick<RouteDecision, 'route' | 'reason'>> {
61  if (!(await read($, isEnabled))) return { route: 'main', reason: 'the router is off' }
62
63  const prompt = text.trim()
64  if (prompt === '') return { route: 'main', reason: 'no prompt to judge (a continuation)' }
65  // A slash command runs a skill or a workflow: real work whatever it says.
66  if (prompt.startsWith('/')) return { route: 'main', reason: 'a slash command is real work' }
67  if (prompt.length > LONG_PROMPT) return { route: 'main', reason: 'a long prompt is real work' }
68
69  const context = await read($, lastContextTokens)
70  // No reading yet. A brand-new session's first turn is small by definition,
71  // but one resumed from a long transcript is not, and its size is unknown
72  // until a response reports it. "I could not measure it" is not "it is small".
73  if (context === 0 && (await $.session.turns()) > 1) {
74    return { route: 'main', reason: 'the conversation size is not known yet (a resumed session)' }
75  }
76  if (context > maxContextTokens) {
77    return {
78      route: 'main',
79      reason: `the conversation is ${thousands(context)} tokens, past the ${thousands(maxContextTokens)} where Haiku stops being cheaper`,
80    }
81  }
82  if ((await $.session.model()) === CHEAP_MODEL) {
83    return { route: 'main', reason: `the session already runs ${short(CHEAP_MODEL)}` }
84  }
85
86  let label: string | undefined
87  try {
88    label = await $.model.classify(prompt, [SIMPLE, COMPLEX])
89  } catch {
90    return { route: 'main', reason: 'the classifier did not answer' }
91  }
92  if (label === SIMPLE) return { route: 'cheap', reason: 'judged a quick question' }
93  if (label === COMPLEX) return { route: 'main', reason: 'judged to need tools or care' }
94
95  return { route: 'main', reason: 'the classifier was unsure' }
96}
97
98export const register: Register = (on, options) => {
99  const maxContextTokens = Number(options.maxContextTokens ?? 100_000)
100
101  on('session.start', async ($, e, next) => {
102    await $.command.register({
103      name: 'route',
104      description: 'Model router: show what it did, or turn it on or off (/route on, /route off)',
105    })
106    // The on/off choice outlives the session; the rest is per session.
107    const saved = await $.store.get('isEnabled')
108    if (typeof saved === 'boolean') await update($, isEnabled, () => saved)
109
110    return next(e)
111  })
112
113  on('turn.start', async ($, e, next) => {
114    const choice = await decide($, e.text, maxContextTokens)
115    await update($, decision, () => ({ turnId: e.turnId, ...choice, escalated: false }))
116    await update($, stats, s => ({ ...s, turns: s.turns + 1 }))
117
118    return next(e)
119  })
120
121  on('turn.step', async function* ($, e, next) {
122    // A subagent was spawned with the model it should have; not ours to move.
123    if (e.agentId !== undefined) return yield* next(e)
124
125    const current = await read($, decision)
126    const isRouted =
127      current.turnId === e.turnId &&
128      current.route === 'cheap' &&
129      !current.escalated &&
130      e.model !== CHEAP_MODEL &&
131      (await read($, lastContextTokens)) <= maxContextTokens
132
133    if (isRouted) $.ui.status(`router: ${short(CHEAP_MODEL)}`)
134    // Low effort on a routed turn: it was judged simple, and lower effort means
135    // fewer tokens and a shorter answer. A step that is not routed keeps the
136    // person's own model and effort, untouched.
137    const result = isRouted
138      ? yield* next({ ...e, model: CHEAP_MODEL, effort: 'low' })
139      : yield* next(e)
140
141    // How big the conversation is, as the API itself reported it: everything the
142    // request sent, cached or not. The next turn's gate reads this.
143    if (result.usage !== null) {
144      const u = result.usage
145      const size = u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
146      await update($, lastContextTokens, () => size)
147    }
148    if (!isRouted) return result
149
150    if (e.index === 0) await update($, stats, s => ({ ...s, routed: s.routed + 1 }))
151
152    // A simple turn answers in prose. One that reaches for a tool was misjudged:
153    // let that call run as any call does (the permission checks are the same)
154    // and give every later step of the turn back to the person's model.
155    if (result.toolUses.length > 0) {
156      await update($, decision, d => ({ ...d, escalated: true }))
157      await update($, stats, s => ({ ...s, escalated: s.escalated + 1 }))
158      $.ui.status(`router: back to ${short(e.model)}`)
159    }
160
161    return result
162  })
163
164  on('turn.complete', async ($, e, next) => {
165    $.ui.status(undefined)
166
167    return next(e)
168  })
169
170  on('command.run', { command: 'route' }, async ($, e) => {
171    const arg = e.args.trim().toLowerCase()
172    if (arg === 'on' || arg === 'off') {
173      const value = arg === 'on'
174      await update($, isEnabled, () => value)
175      await $.store.set('isEnabled', value)
176
177      return { text: `Model router ${arg}.` }
178    }
179    if (arg !== '') return { text: 'Usage: /route, /route on, /route off' }
180
181    const s = await read($, stats)
182    const d = await read($, decision)
183    const isOn = await read($, isEnabled)
184    const context = await read($, lastContextTokens)
185    const lines = [
186      isOn
187        ? `Model router is on: quick questions go to ${short(CHEAP_MODEL)} while the conversation is under ${thousands(maxContextTokens)} tokens (now ${thousands(context)}).`
188        : 'Model router is off: every turn runs on your own model.',
189      `This session: ${s.routed} of ${s.turns} turn(s) routed` +
190        (s.escalated > 0 ? `, ${s.escalated} handed back after reaching for a tool.` : '.'),
191    ]
192    if (d.turnId !== '') {
193      lines.push(
194        `Last turn: ${d.route === 'cheap' ? short(CHEAP_MODEL) : 'your model'}, ${d.reason}` +
195          (d.escalated ? ', then handed back.' : '.'),
196      )
197    }
198
199    return { text: lines.join('\n') }
200  })
201}
202

1/** What the router decided for one turn, and why, so `/route` can say it. */
2export type RouteDecision = {
3  turnId: string
4  route: 'cheap' | 'main'
5  reason: string
6  /** The cheaper model reached for a tool, so the rest of the turn went back. */
7  escalated: boolean
8}
9
10export type RouterStats = {
11  turns: number
12  routed: number
13  escalated: number
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    'model-router': {
19      decision: RouteDecision
20      stats: RouterStats
21      isEnabled: boolean
22      /** The prompt size of the last main-loop request, in tokens, as the API reported it. */
23      lastContextTokens: number
24    }
25  }
26}
27