auto-model

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { AutoModelPick } from '../types'
5
6// Auto Model — picks the right model ONCE, from the first prompt of a session,
7// and keeps it for the whole session. It never switches mid-session: the prompt
8// cache is tied to the model, so switching back and forth would make the new
9// model re-read the whole conversation and cost more than it saves.
10
11const MODELS = {
12  haiku: 'claude-haiku-5-5',
13  sonnet: 'claude-sonnet-5-5',
14  opus: 'claude-opus-5-5',
15} as const
16type Tier = keyof typeof MODELS
17const NAMES: Record<Tier, string> = { haiku: 'Haiku', sonnet: 'Sonnet', opus: 'Opus' }
18
19const pickAtom = atom({ plugin: 'auto-model', key: 'pick' } as const, null as AutoModelPick | null)
20const decidedAtom = atom({ plugin: 'auto-model', key: 'decided' } as const, false)
21
22const RULES = `You route a user's first message in a coding/assistant session to a model tier.
23Reply with exactly one word: haiku, sonnet, or opus.
24- haiku: a quick question, a definition, a one-line fix, a small lookup, a short rewrite, chit-chat.
25- sonnet: normal work — writing or editing a file or two, debugging, a script, a summary, a page of writing, research with a few steps.
26- opus: a big build or hard problem — a new app/site/feature across many files, architecture, a large refactor, a long multi-step plan, a deep analysis, anything ambiguous and high-stakes.
27When unsure between two, pick the larger one.`
28
29function isTier(x: unknown): x is Tier {
30  return x === 'haiku' || x === 'sonnet' || x === 'opus'
31}
32
33async function isEnabled($: any): Promise<boolean> {
34  const v = await $.store.get('enabled')
35  return v !== false
36}
37
38async function chooseTier($: any, text: string): Promise<Tier> {
39  const r = await $.model.complete({
40    model: 'haiku',
41    system: RULES,
42    prompt: 'Message:\n"""\n' + text.slice(0, 4000) + '\n"""\nTier:',
43    maxTokens: 5,
44    effort: 'low',
45    timeoutMs: 8000,
46  })
47  const word = r.isAnswered ? r.text.trim().toLowerCase().replace(/[^a-z]/g, '') : ''
48  return isTier(word) ? word : 'sonnet'
49}
50
51export const register: Register = on => {
52  on('session.start', async ($, e, next) => {
53    const out = await next(e)
54    try {
55      await $.command.register({
56        name: 'auto-model',
57        description: 'Auto Model: show the pick, or on | off | haiku | sonnet | opus',
58        argumentHint: '[on|off|haiku|sonnet|opus]',
59        immediate: true,
60      })
61    } catch (err) {
62      $.ui.log('auto-model: could not add /auto-model (' + String(err) + ')')
63    }
64    return out
65  })
66
67  // Decide once, at the first main-loop turn of the session.
68  on('turn.start', async ($, e, next) => {
69    if (e.agentId !== undefined) return next(e)
70    if (await read($, decidedAtom)) return next(e)
71    await update($, decidedAtom, () => true)
72
73    if (!(await isEnabled($))) return next(e)
74
75    // A resumed session already has a cache on its model: keep it.
76    const turns = await $.session.turns()
77    // The person picked a model with /model: respect it.
78    const rows = await $.config.list()
79    const modelRow = rows.find(row => row.key === 'model')
80    const manual = modelRow !== undefined && !/default/i.test(String(modelRow.value))
81    if (turns > 1 || manual) return next(e)
82
83    const text = (e.text || '').trim()
84    if (text === '') return next(e)
85
86    let tier: Tier = 'sonnet'
87    try {
88      tier = await chooseTier($, text)
89    } catch {
90      tier = 'sonnet'
91    }
92    const pick: AutoModelPick = { tier, model: MODELS[tier], how: 'auto' }
93    await update($, pickAtom, () => pick)
94    $.ui.toast('Auto Model picked ' + NAMES[tier] + ' for this session')
95    $.ui.status('Auto Model: ' + NAMES[tier])
96    return next(e)
97  })
98
99  // Every main-loop request of the session goes to the picked model.
100  on('turn.step', async function* ($, e, next) {
101    if (e.agentId !== undefined) return yield* next(e)
102    const pick = await read($, pickAtom)
103    if (pick === null || pick.how === 'kept') return yield* next(e)
104    return yield* next({ ...e, model: pick.model })
105  })
106
107  on('command.run', { command: 'auto-model' }, async ($, e) => {
108    const arg = (e.args || '').trim().toLowerCase()
109    const pick = await read($, pickAtom)
110    const turns = await $.session.turns()
111
112    if (arg === 'off') {
113      await $.store.set('enabled', false)
114      return { text: 'Auto Model is off for new sessions. This session keeps ' + (pick ? NAMES[pick.tier] : 'its current model') + '.' }
115    }
116    if (arg === 'on') {
117      await $.store.set('enabled', true)
118      return { text: 'Auto Model is on. It picks a model from the first prompt of each new session.' }
119    }
120    if (isTier(arg)) {
121      if (turns > 0 && pick !== null) {
122        return {
123          text: 'This session already runs on ' + NAMES[pick.tier] + '. Switching mid-session breaks the prompt cache and usually costs more — start a new session and run /auto-model ' + arg + ' before your first prompt, or use /model if you really want to switch.',
124        }
125      }
126      await update($, decidedAtom, () => true)
127      await update($, pickAtom, () => ({ tier: arg, model: MODELS[arg], how: 'forced' as const }))
128      $.ui.status('Auto Model: ' + NAMES[arg])
129      return { text: 'This session will use ' + NAMES[arg] + '.' }
130    }
131
132    const enabled = await isEnabled($)
133    const now = pick ? NAMES[pick.tier] + (pick.how === 'forced' ? ' (you chose it)' : ' (picked from your first prompt)') : 'not picked yet — it picks from your first prompt'
134    return {
135      text: 'Auto Model is ' + (enabled ? 'on' : 'off') + '. This session: ' + now + '.\nUse /auto-model off, on, or haiku | sonnet | opus (before your first prompt).',
136    }
137  })
138}
139

1export type AutoModelPick = {
2  /** The tier picked for this session */
3  tier: 'haiku' | 'sonnet' | 'opus'
4  /** The model id every main-loop request of this session is sent to */
5  model: string
6  /** Why: 'auto' (classified), 'forced' (/auto-model <tier>), 'kept' (manual /model or resumed session) */
7  how: 'auto' | 'forced' | 'kept'
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'auto-model': { pick: AutoModelPick | null; decided: boolean }
13  }
14}
15