cc-model-router

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PromptSubmitInput, Register } from 'claude-code'
3
4import type { Decision, DecisionSource, RouterState, Tier } from '../types'
5import { logDecision } from './log'
6import { decideAgainstSticky, isFollowUp, parsePrefix, rankByRules } from './rank'
7import { routeFor } from './route'
8
9const router = atom({ plugin: 'cc-model-router', key: 'router' } as const, { isOn: true, stickyTier: null, decision: null } as RouterState)
10
11const TIERS: readonly Tier[] = ['simple', 'medium', 'hard']
12
13// Claude Code's own commands: they keep the sticky model and are never ranked.
14const BUILT_INS = new Set([
15  'add-dir', 'agents', 'bug', 'clear', 'compact', 'config', 'context', 'cost', 'doctor', 'effort', 'exit', 'export', 'fast',
16  'help', 'hooks', 'ide', 'init', 'login', 'logout', 'mcp', 'memory', 'model', 'permissions', 'plugin', 'release-notes',
17  'reload-plugins', 'resume', 'rewind', 'router', 'status', 'terminal-setup', 'theme', 'usage', 'vim',
18])
19
20const EDIT_TOOLS = new Set(['Edit', 'Write', 'NotebookEdit'])
21
22// Bash that changes files or history; `2>&1` and `>/dev/null` are not writes.
23const WRITING_BASH = /(^|[;&|]\s*)(rm|mv|cp|mkdir|touch|tee|sed -i|git (commit|push|checkout|reset|rebase|merge))\b|>>?\s*(?!&|\/dev\/null)[^\s&|;]/
24
25type Verdict = { tier: Tier; source: DecisionSource; reason: string }
26
27// The model the main loop's latest step was sent on: every tool call of a Haiku step is judged by it,
28// not by the decision, which the first denied edit already moved to Sonnet.
29let stepModel: string | null = null
30
31// Only a Haiku pick needs the size; the local estimate counts tool results the last response has not seen.
32async function contextTokens($: EngineInterface): Promise<number | undefined> {
33  const usage = await $.session.usage({ breakdown: 'summary' }).catch(() => undefined)
34
35  return usage?.context.breakdown?.totalTokens ?? usage?.context.tokens
36}
37
38// Only what the owner typed, at an idle prompt, is ranked; everything else rides the sticky tier.
39const isRankable = (e: PromptSubmitInput): boolean =>
40  e.isMeta !== true &&
41  e.turnId === undefined &&
42  (e.origin.kind === 'composer' || e.origin.kind === 'bridge') &&
43  (e.command === undefined || !BUILT_INS.has(e.command.name))
44
45// What Haiku reads before each prompt: the owner's hard categories and what a follow-up looks like.
46const CLASSIFIER_SYSTEM = [
47  'You route a coding assistant prompt to a model. Answer with exactly one word.',
48  'hard: planning, brainstorming, design or architecture; debugging, failing tests, crashes, root causes; changes across many files; long pasted logs, docs or code.',
49  'medium: ordinary code changes in one or two files, moderate how-to questions.',
50  'simple: short factual questions, lookups, explanations, chat.',
51  'keep: the prompt continues the current task (feedback, a reply to a question, a small tweak) and should stay on the current level.',
52  'Words: simple, medium, hard, keep.',
53].join('\n')
54
55const CLASSIFIER_TIMEOUT_MS = 3000
56
57const TIER_ORDER: Record<Tier, number> = { simple: 0, medium: 1, hard: 2 }
58
59const higher = (a: Tier, b: Tier): Tier => (TIER_ORDER[a] >= TIER_ORDER[b] ? a : b)
60
61// Haiku's own read of the prompt; undefined when it fails, times out or names no level.
62async function classify($: EngineInterface, text: string, sticky: Tier | null): Promise<Tier | 'keep' | undefined> {
63  const current = sticky === null ? 'Current level: none (first prompt, never answer keep).' : `Current level: ${sticky}.`
64  const result = await $.model
65    .complete({ model: 'haiku', system: CLASSIFIER_SYSTEM, prompt: `${current}\n\nPrompt:\n${text.slice(0, 4000)}`, maxTokens: 5, effort: 'low', timeoutMs: CLASSIFIER_TIMEOUT_MS })
66    .catch(() => undefined)
67  const word = result?.isAnswered ? result.text.trim().toLowerCase().match(/simple|medium|hard|keep/)?.[0] : undefined
68
69  return word === 'keep' ? (sticky === null ? undefined : 'keep') : TIERS.find(tier => tier === word)
70}
71
72// Instant answers first (a hard rule, a follow-up); otherwise Haiku and the rules both rank it and the higher wins.
73async function judge($: EngineInterface, text: string, sticky: Tier | null, command?: { command: string; args: string }): Promise<Verdict> {
74  const verdict = rankByRules(text, command)
75  const rules = verdict === undefined ? 'unsure' : `${verdict.tier}/${verdict.rule}`
76
77  if (verdict?.tier === 'hard') {
78    return { tier: 'hard', source: 'rule', reason: verdict.rule }
79  }
80
81  if (sticky !== null && isFollowUp(text)) {
82    return { tier: sticky, source: 'sticky', reason: 'follow-up' }
83  }
84
85  const haiku = await classify($, text, sticky)
86
87  if (haiku === 'keep' && sticky !== null) {
88    return { tier: sticky, source: 'classifier', reason: `rules: ${rules} · haiku: keep` }
89  }
90
91  if (haiku !== undefined && haiku !== 'keep') {
92    const tier = verdict === undefined ? haiku : higher(haiku, verdict.tier)
93
94    return { tier, source: 'classifier', reason: `rules: ${rules} · haiku: ${haiku} → ${tier}` }
95  }
96
97  // Haiku gave nothing: the rules alone, as before.
98  const decision = decideAgainstSticky(text, sticky, verdict)
99
100  if (decision.kind === 'switch') {
101    return { tier: decision.verdict.tier, source: 'rule', reason: `${decision.verdict.rule} (haiku failed)` }
102  }
103
104  if (sticky !== null) {
105    return { tier: sticky, source: 'sticky', reason: `rules: ${rules} (haiku failed)` }
106  }
107
108  throw new Error('no rule and no classifier answer')
109}
110
111async function decide($: EngineInterface, verdict: Verdict, text: string): Promise<void> {
112  const tokens = verdict.tier === 'simple' ? await contextTokens($) : undefined
113  const route = routeFor({ tier: verdict.tier, sessionModel: await $.session.model(), contextTokens: tokens })
114  const decision: Decision = { ...verdict, tier: route.tier, model: route.model, effort: route.effort, at: await $.clock.now(), contextTokens: tokens }
115
116  await update($, router, state => ({ ...state, stickyTier: route.tier, decision }))
117  await logDecision($, { ...decision, text })
118}
119
120function describe(state: RouterState): string {
121  const { decision } = state
122  const last = decision === null ? 'no decision yet' : `last: ${decision.model} (${decision.source}: ${decision.reason})`
123
124  return `Router ${state.isOn ? 'on' : 'off'} · sticky tier: ${state.stickyTier ?? 'none'} · ${last}`
125}
126
127export const register: Register = on => {
128  on('session.start', async ($, e, next) => {
129    await $.command.register({ name: 'router', description: 'Turn the model router on or off, or show its state' })
130    // Earlier versions set a status entry that outlives a reload; session-info draws the routed model now.
131    $.ui.status(undefined)
132
133    return next(e)
134  })
135
136  on('command.run', { command: 'router' }, async ($, e) => {
137    const arg = e.args.trim().toLowerCase()
138
139    if (arg === 'on' || arg === 'off') {
140      await update($, router, state => ({ ...state, isOn: arg === 'on' }))
141    } else if (arg !== '' && arg !== 'status') {
142      return { text: 'Usage: /router on | off | status' }
143    }
144
145    return { text: describe(await read($, router)) }
146  })
147
148  // Haiku answers; it does not edit. An edit on a Haiku turn is refused and the turn moves up to Sonnet, which retries it.
149  on('tool.call', async ($, e, next) => {
150    const state = await read($, router)
151    const tool = String(e.tool)
152    const isEdit = EDIT_TOOLS.has(tool) || (e.tool === 'Bash' && WRITING_BASH.test(e.command))
153
154    if (e.agentId !== undefined || !state.isOn || !isEdit || stepModel === null || !/haiku/i.test(stepModel)) {
155      return next(e)
156    }
157
158    const verdict: Verdict = { tier: 'medium', source: 'escalated', reason: `haiku tried ${tool}` }
159    await decide($, verdict, tool === 'Bash' && e.tool === 'Bash' ? e.command : tool)
160
161    return { deny: 'Routed up to Sonnet for edits; retry this call.' }
162  }).catch(($, e, next) => next(e))
163
164  on('prompt.submit', async ($, e, next) => {
165    const state = await read($, router)
166
167    if (!state.isOn || !isRankable(e)) {
168      return next(e)
169    }
170
171    // A command's text is not rewritable, so its prefix is not read.
172    const prefix = e.command === undefined ? parsePrefix(e.text) : undefined
173
174    try {
175      if (prefix !== undefined) {
176        await decide($, { tier: prefix.tier, source: 'prefix', reason: 'prefix' }, e.text)
177
178        return next({ ...e, text: prefix.text })
179      }
180
181      const command = e.command === undefined ? undefined : { command: e.command.name, args: e.command.args }
182      await decide($, await judge($, e.text, state.stickyTier, command), e.text)
183    } catch {
184      // Rules or classifier failed: the sticky tier, or the session's model, carries the prompt.
185    }
186
187    return next(e)
188  }).catch(($, e, next) => next(e))
189
190  on('turn.step', async function* ($, e, next) {
191    const state = await read($, router)
192
193    if (e.agentId !== undefined) {
194      return yield* next(e)
195    }
196
197    if (!state.isOn || state.decision === null) {
198      stepModel = e.model
199
200      return yield* next(e)
201    }
202
203    const tokens = state.decision.tier === 'simple' ? await contextTokens($) : undefined
204    const route = routeFor({ tier: state.decision.tier, sessionModel: e.model, sessionEffort: e.effort, contextTokens: tokens })
205
206    if (route.model !== state.decision.model || route.effort !== state.decision.effort || route.tier !== state.decision.tier) {
207      await update($, router, current =>
208        current.decision === null ? current : { ...current, stickyTier: route.tier, decision: { ...current.decision, ...route } },
209      )
210    }
211
212    // One log line whenever the model a step runs on differs from the step before it.
213    if (stepModel !== route.model) {
214      await logDecision($, { ...state.decision, ...route, source: 'step', reason: `step ${e.index} of ${e.turnId}: ${stepModel ?? 'none'} -> ${route.model}`, text: '' })
215    }
216
217    stepModel = route.model
218
219    return yield* next({ ...e, model: route.model, effort: route.effort ?? e.effort })
220  })
221
222  on('session.end', async ($, e, next) => {
223    if (e.reason === 'clear') {
224      await update($, router, state => ({ ...state, stickyTier: null, decision: null }))
225    }
226
227    return next(e)
228  })
229}
230

1import type { EngineInterface } from 'claude-code'
2
3import type { Decision } from '../types'
4
5const KEEP_LINES = 1000
6
7// Outside the plugin folder: a write inside it would hot-reload the mod.
8async function logPath($: EngineInterface): Promise<string | undefined> {
9  const home = await $.env.get('HOME')
10
11  return home === undefined ? undefined : `${home}/.claude/cc-model-router/decisions.jsonl`
12}
13
14export async function logDecision($: EngineInterface, decision: Decision & { text: string }): Promise<void> {
15  try {
16    const path = await logPath($)
17
18    if (path === undefined) {
19      return
20    }
21
22    const line = JSON.stringify({ ...decision, at: new Date(decision.at).toISOString(), text: decision.text.slice(0, 200) })
23    const previous = (await $.fs.exists(path)) ? await $.fs.read(path) : ''
24    const lines = [...previous.split('\n').filter(Boolean), line].slice(-KEEP_LINES)
25
26    await $.fs.write(path, `${lines.join('\n')}\n`)
27  } catch {
28    // The log is a convenience; a failed write never touches routing.
29  }
30}
31

1import type { Tier } from '../types'
2
3export type RuleVerdict = { tier: Tier; rule: string }
4
5export type StickyDecision = { kind: 'keep' } | { kind: 'switch'; verdict: RuleVerdict } | { kind: 'ask' }
6
7const PREFIX = /^\s*(haiku|sonnet|opus):\s*/i
8
9const PREFIX_TIERS: Record<string, Tier> = { haiku: 'simple', sonnet: 'medium', opus: 'hard' }
10
11const FOLLOW_UP = /^(y|yes|yep|ok|okay|sure|go|go ahead|do it|continue|proceed|next|lgtm|looks good|thanks|thank you|nice|great|done)\b/i
12
13// The owner's hard cases, first match wins.
14const HARD_RULES: readonly { rule: string; test: (text: string) => boolean }[] = [
15  { rule: 'long-paste', test: text => text.length > 4000 || text.split('\n').length > 80 },
16  { rule: 'planning', test: text => /\b(plan|brainstorm|architect(ure)?|design|refactor|migrat(e|ion)|rfc|spec)\b/i.test(text) },
17  {
18    rule: 'debugging',
19    test: text =>
20      /\b(debug|bug|broken|crash(es|ing)?|stack ?trace|traceback|exception|failing|fails?|regression|root cause|why (is|does|doesn'?t))\b/i.test(text),
21  },
22  {
23    rule: 'multi-file',
24    test: text => /\b(across|all|every|multiple|several) (files?|modules?|components?|packages?)\b/i.test(text) || /\bcodebase\b/i.test(text),
25  },
26]
27
28const CODE_CHANGE = /\b(add|implement|write|create|update|change|rename|edit|modify|fix|remove|delete)\b/i
29
30const QUESTION_START = /^(what|how|where|which|who|when|is|are|can|does|do|explain|show|list)\b/i
31
32const HARD_COMMAND = /brainstorm|plan|architect|debug|review/i
33
34// The minimum length a prompt needs to count as a clearly new task rather than an aside.
35const NEW_TASK_LENGTH = 60
36
37export function parsePrefix(text: string): { tier: Tier; text: string } | undefined {
38  const match = PREFIX.exec(text)
39  const tier = match === null ? undefined : PREFIX_TIERS[match[1]!.toLowerCase()]
40
41  return tier === undefined ? undefined : { tier, text: text.slice(match![0].length) }
42}
43
44export function isFollowUp(text: string): boolean {
45  const trimmed = text.trim()
46
47  return trimmed.length <= 40 && FOLLOW_UP.test(trimmed)
48}
49
50export function rankByRules(text: string, command?: { command: string; args: string }): RuleVerdict | undefined {
51  if (command !== undefined && HARD_COMMAND.test(command.command)) {
52    return { tier: 'hard', rule: 'command' }
53  }
54
55  const subject = command === undefined ? text : `${command.command} ${command.args}`
56  const hard = HARD_RULES.find(({ test }) => test(subject))
57
58  if (hard !== undefined) {
59    return { tier: 'hard', rule: hard.rule }
60  }
61
62  if (CODE_CHANGE.test(subject)) {
63    return { tier: 'medium', rule: 'code-change' }
64  }
65
66  const trimmed = subject.trim()
67  const isShortQuestion = trimmed.length <= 200 && !trimmed.includes('\n') && (trimmed.endsWith('?') || QUESTION_START.test(trimmed))
68
69  return isShortQuestion ? { tier: 'simple', rule: 'short-question' } : undefined
70}
71
72export function decideAgainstSticky(text: string, sticky: Tier | null, verdict: RuleVerdict | undefined): StickyDecision {
73  if (sticky === null) {
74    return verdict === undefined ? { kind: 'ask' } : { kind: 'switch', verdict }
75  }
76
77  if (verdict === undefined || verdict.tier === sticky) {
78    return { kind: 'keep' }
79  }
80
81  // A hard signal always moves up, even behind "ok"; a follow-up only blocks moving down or sideways.
82  if (verdict.tier === 'hard') {
83    return { kind: 'switch', verdict }
84  }
85
86  if (!isFollowUp(text) && text.trim().length >= NEW_TASK_LENGTH) {
87    return { kind: 'switch', verdict }
88  }
89
90  return { kind: 'keep' }
91}
92

1import type { Effort, Tier } from '../types'
2
3export type RouteInput = { tier: Tier; sessionModel: string; sessionEffort?: Effort | number; contextTokens?: number }
4
5export type Route = { tier: Tier; model: string; effort?: Effort }
6
7// Haiku's window is 200k; past this the context no longer fits it with room to answer.
8export const HAIKU_CONTEXT_LIMIT = 160_000
9
10const EFFORT_ORDER: readonly Effort[] = ['low', 'medium', 'high', 'xhigh', 'max']
11
12const MODELS: Record<Tier, string> = { simple: 'haiku', medium: 'sonnet', hard: 'opus' }
13
14const EFFORTS: Record<Tier, Effort> = { simple: 'low', medium: 'medium', hard: 'high' }
15
16// Opus never runs weaker than the session's own effort; a numeric budget counts as none.
17const opusEffort = (session: Effort | number | undefined): Effort =>
18  typeof session === 'string' && EFFORT_ORDER.indexOf(session) > EFFORT_ORDER.indexOf('high') ? session : 'high'
19
20export function routeFor({ tier, sessionModel, sessionEffort, contextTokens }: RouteInput): Route {
21  // An unknown size (a fresh session has none yet) counts as small; a known one past the limit never goes to Haiku.
22  const fits = tier !== 'simple' || (contextTokens ?? 0) <= HAIKU_CONTEXT_LIMIT
23  const routed: Tier = fits ? tier : 'medium'
24  const isLong = /\[1m\]/i.test(sessionModel) && routed !== 'simple'
25
26  return {
27    tier: routed,
28    model: isLong ? `${MODELS[routed]}[1m]` : MODELS[routed],
29    effort: routed === 'hard' ? opusEffort(sessionEffort) : EFFORTS[routed],
30  }
31}
32

1export type Tier = 'simple' | 'medium' | 'hard'
2
3export type Effort = 'low' | 'medium' | 'high' | 'xhigh' | 'max'
4
5export type DecisionSource = 'prefix' | 'rule' | 'classifier' | 'sticky' | 'escalated' | 'step'
6
7export type Decision = { tier: Tier; model: string; effort?: Effort; source: DecisionSource; reason: string; at: number; contextTokens?: number }
8
9export type RouterState = { isOn: boolean; stickyTier: Tier | null; decision: Decision | null }
10
11declare module 'claude-code' {
12  interface PluginState {
13    'cc-model-router': { router: RouterState }
14  }
15}
16