model-router

1import type { EngineInterface, Register } from 'claude-code'
2
3import type { ModelTier } from '../types'
4
5const MODELS: Record<ModelTier, string> = {
6  haiku: 'claude-haiku-5-5',
7  sonnet: 'claude-sonnet-5-5',
8  opus: 'claude-opus-5-5',
9}
10
11const RUBRIC = `You route the opening prompt of a coding assistant's conversation to a model tier; the conversation stays on that tier. Answer with one word: haiku, sonnet or opus.
12haiku: trivial asks. Quick factual questions, a single shell command, renaming, small mechanical edits, formatting, "commit this".
13sonnet: ordinary engineering. Implementing a feature, fixing a normal bug, writing tests, explaining code, moderate refactors.
14opus: hard reasoning. Architecture or design decisions, subtle or cross-cutting bugs, concurrency, algorithms, proofs, large refactors, reviews of tricky code, ambiguous specs.
15When unsure between two tiers, pick the higher one.`
16
17const ENABLED_KEY = 'enabled'
18const TIER = { plugin: 'model-router', key: 'tier' } as const
19
20const isEnabled = async ($: EngineInterface) =>
21  ((await $.store.get(ENABLED_KEY)) ?? true) === true
22
23// The tier this conversation runs on, fixed at its first prompt; undefined for
24// the session's own model. The on/off switch only matters at the next first prompt.
25const activeTier = async ($: EngineInterface): Promise<ModelTier | undefined> => {
26  const { value } = await $.state.get(TIER)
27  return value === 'default' ? undefined : value
28}
29
30const showStatus = async ($: EngineInterface) => {
31  const [tier, isOn] = await Promise.all([activeTier($), isEnabled($)])
32  if (tier === undefined) {
33    $.ui.status(isOn ? 'session' : 'off')
34    return
35  }
36  $.ui.status(isOn ? tier : `${tier} (off next)`)
37}
38
39const pick = async ($: EngineInterface, text: string): Promise<ModelTier | undefined> => {
40  const r = await $.model.complete({
41    model: 'haiku',
42    system: RUBRIC,
43    prompt: `<prompt>\n${text.slice(0, 6000)}\n</prompt>\nTier:`,
44    maxTokens: 5,
45    effort: 'low',
46    timeoutMs: 5000,
47  })
48  if (!r.isAnswered) return undefined
49  const word = r.text.trim().toLowerCase().match(/haiku|sonnet|opus/)?.[0]
50  return word as ModelTier | undefined
51}
52
53export const register: Register = on => {
54  on('session.start', async ($, e, next) => {
55    await $.command.register({
56      name: 'router',
57      description: 'Automatic model routing: on, off, or toggle',
58      argumentHint: '[on|off]',
59      immediate: true,
60    })
61    await showStatus($)
62
63    return next(e)
64  })
65
66  on('command.run', { command: 'router' }, async ($, e) => {
67    const arg = e.args.trim().toLowerCase()
68    const isOn = arg === 'on' ? true : arg === 'off' ? false : !(await isEnabled($))
69    await $.store.set(ENABLED_KEY, isOn)
70    await showStatus($)
71
72    return { text: `Model router ${isOn ? 'on' : 'off'} from the next conversation (/clear).` }
73  })
74
75  // Pick only on a conversation's first prompt (new session or after /clear):
76  // switching models later forfeits the prompt cache. For the same reason
77  // /router takes effect at the next first prompt, not mid-conversation.
78  on('prompt.submit', async ($, e, next) => {
79    const isTyped = e.origin.kind === 'composer' || e.origin.kind === 'bridge'
80    if (!isTyped || e.turnId !== undefined || (await $.session.turns()) > 0) {
81      return next(e)
82    }
83    const tier = (await isEnabled($)) ? await pick($, e.text).catch(() => undefined) : undefined
84    await $.state.set(TIER, tier ?? 'default')
85    await showStatus($)
86
87    return next(e)
88  })
89
90  on('turn.step', async function* ($, e, next) {
91    const tier = e.agentId === undefined ? await activeTier($) : undefined
92    if (tier === undefined) {
93      return yield* next(e)
94    }
95
96    return yield* next({ ...e, model: MODELS[tier] })
97  })
98}
99

1export type ModelTier = 'haiku' | 'sonnet' | 'opus'
2
3declare module 'claude-code' {
4  interface PluginState {
5    'model-router': {
6      // The tier picked for this conversation, or 'default' for the session's own model.
7      tier: ModelTier | 'default'
8    }
9  }
10}
11