auto-model

1import type { EngineInterface, Register } from 'claude-code'
2
3type Tier = 'haiku' | 'sonnet' | 'opus'
4
5// Edit these ids when new models ship.
6export const MODELS: Record<Tier, string> = {
7  haiku: 'claude-haiku-5-5',
8  sonnet: 'claude-sonnet-5-5',
9  opus: 'claude-opus-5-5',
10}
11const CLASSIFIER = MODELS.haiku
12const RANK: Record<Tier, number> = { haiku: 1, sonnet: 2, opus: 3 }
13const LABEL: Record<Tier, string> = { haiku: 'Haiku', sonnet: 'Sonnet', opus: 'Opus' }
14
15const SYSTEM = `You route coding-assistant prompts to the cheapest model that can do the job well.
16Answer with exactly one word: haiku, sonnet or opus.
17- haiku: simple questions, lookups, explanations, renaming, small single-file edits, formatting, running a command, short text.
18- sonnet: ordinary multi-step coding, bug fixes, writing tests, features touching a few files, moderate analysis or writing.
19- opus: hard reasoning, architecture or design decisions, large refactors across many files, tricky debugging, security review, long research.
20When unsure between two, pick the stronger one.`
21
22export function tierOf(model: string): Tier | undefined {
23  const m = model.toLowerCase()
24  if (m.includes('haiku')) return 'haiku'
25  if (m.includes('sonnet')) return 'sonnet'
26  if (m.includes('opus')) return 'opus'
27  return undefined
28}
29
30export function parseTier(text: string): Tier | undefined {
31  const word = text.trim().toLowerCase().match(/haiku|sonnet|opus/)
32  return word ? (word[0] as Tier) : undefined
33}
34
35async function isEnabled($: EngineInterface): Promise<boolean> {
36  return (await $.store.get('enabled')) !== false
37}
38
39export const register: Register = on => {
40  // Per-turn choice; module state is fine here, a reload only drops the turn in flight.
41  const chosen = new Map<string, Tier>()
42  let last: Tier | undefined
43
44  on('session.start', async ($, e, next) => {
45    await $.command.register({
46      name: 'auto-model',
47      description: 'Automatic model choice per prompt: /auto-model [on|off|status]',
48      argumentHint: '[on|off|status]',
49    })
50    $.ui.status((await isEnabled($)) ? 'Auto-model: on' : undefined)
51    return next(e)
52  })
53
54  on('command.run', { command: 'auto-model' }, async ($, e) => {
55    const arg = e.args.trim().toLowerCase()
56    const was = await isEnabled($)
57
58    if (arg === 'status') {
59      const lastText = last ? ` Last choice: ${LABEL[last]}.` : ''
60      return { text: `Auto-model is ${was ? 'on' : 'off'}.${lastText}` }
61    }
62
63    const now = arg === 'on' ? true : arg === 'off' ? false : !was
64    await $.store.set('enabled', now)
65    $.ui.status(now ? 'Auto-model: on' : undefined)
66    return {
67      text: now
68        ? 'Auto-model is on: each prompt runs on the cheapest model that fits (never above your /model choice).'
69        : 'Auto-model is off: every prompt runs on your /model choice.',
70    }
71  })
72
73  on('turn.start', async ($, e, next) => {
74    if (!(await isEnabled($))) return next(e)
75
76    if (e.text.trim() === '') {
77      if (last) chosen.set(e.turnId, last)
78      return next(e)
79    }
80
81    const reply = await $.model.complete({
82      model: CLASSIFIER,
83      system: SYSTEM,
84      prompt: `Prompt:\n"""\n${e.text.slice(0, 4000)}\n"""\nModel:`,
85      maxTokens: 5,
86      timeoutMs: 8000,
87    })
88    const tier = reply.isAnswered ? parseTier(reply.text) : undefined
89    if (tier) {
90      chosen.set(e.turnId, tier)
91      last = tier
92    }
93    return next(e)
94  })
95
96  on('turn.step', async function* ($, e, next) {
97    const want = e.agentId === undefined ? chosen.get(e.turnId) : undefined
98    const current = tierOf(e.model)
99    if (!want || !current) return yield* next(e)
100
101    // Only ever step down from the model the person picked with /model.
102    const tier = RANK[want] < RANK[current] ? want : current
103    $.ui.status(`Auto-model: ${LABEL[tier]}${tier === current ? '' : ` (instead of ${LABEL[current]})`}`)
104    return yield* next(tier === current ? e : { ...e, model: MODELS[tier] })
105  })
106
107  on('turn.complete', ($, e, next) => {
108    chosen.delete(e.turnId)
109    return next(e)
110  })
111}
112