auto-model

1import { atom, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import { decide, modelFor, tierFor, tierOf, type Complexity, type Decision, type Overrides, type Tier } from './lib'
5
6type Mode = 'auto' | 'off' | Tier
7const CFG = 'auto-model:mode'
8const LOG = 'auto-model:log'
9const IDS = 'auto-model:ids'
10// The usage meter reads this: a store is private to each plugin, an atom can be read by its name.
11const effective = atom({ plugin: 'auto-model', key: 'effective' } as const, '')
12// The model the session itself named when auto-model last routed: the meter compares it with the live one to see a /model.
13const base = atom({ plugin: 'auto-model', key: 'base' } as const, '')
14const KEEP = 200
15const TIERS: readonly string[] = ['haiku', 'sonnet', 'opus']
16
17type Store = { store: { get: (k: string) => Promise<unknown>; set: (k: string, v: unknown) => Promise<unknown> } }
18async function idsOf($: Store): Promise<Overrides> {
19  return ((await $.store.get(IDS)) as Overrides | undefined) ?? {}
20}
21async function modeOf($: Store): Promise<Mode> {
22  return ((await $.store.get(CFG)) as Mode | undefined) ?? 'auto'
23}
24
25export const register: Register = on => {
26  let lastSent: Tier | undefined // the model whose cache is warm
27  let sessionModel: string | undefined // what the session itself names, to notice /model
28  let lastDoneAt = 0
29  let sinceSwitch = 99
30  let plan: (Decision & { turnId: string }) | undefined
31
32  on('session.start', async ($, e, next) => {
33    await update($, effective, () => '')
34    await update($, base, () => '')
35    await $.command.register({
36      name: 'auto-model',
37      description: 'Auto model routing: on, off, haiku, sonnet, opus (pin), status, set <tier> <id>, reset',
38      argumentHint: '[on|off|haiku|sonnet|opus|status|set <tier> <model-id>|reset]',
39    })
40    return next(e)
41  })
42
43  on('command.run', { command: 'auto-model' }, async ($, e) => {
44    const raw = e.args.trim()
45    const arg = raw.toLowerCase()
46    const setCmd = /^set\s+(haiku|sonnet|opus)\s+(\S+)$/i.exec(raw)
47    if (setCmd) await $.store.set(IDS, { ...(await idsOf($)), [setCmd[1].toLowerCase()]: setCmd[2] })
48    else if (arg === 'reset') await $.store.set(IDS, {})
49    else if (arg === 'on') await $.store.set(CFG, 'auto')
50    else if (arg === 'off' || TIERS.includes(arg)) await $.store.set(CFG, arg)
51    else if (arg !== '' && arg !== 'status') return { text: 'Usage: /auto-model [on|off|haiku|sonnet|opus|status|set <tier> <model-id>|reset]' }
52
53    const mode = await modeOf($)
54    const log = ((await $.store.get(LOG)) as { to: Tier; reason: string; ctx: number }[] | undefined) ?? []
55    const counts = { haiku: 0, sonnet: 0, opus: 0 }
56    for (const l of log) counts[l.to] += 1
57    const recent = log.slice(-5).map(l => `  ${l.to}  ${l.reason}  (ctx ${l.ctx})`)
58    return {
59      text: [
60        `auto-model: ${mode === 'auto' ? 'on' : mode === 'off' ? 'off' : `pinned to ${mode}`}. Warm cache: ${lastSent ?? 'none yet'}.`,
61        `Last ${log.length} turns: haiku ${counts.haiku}, sonnet ${counts.sonnet}, opus ${counts.opus}.`,
62        ...(recent.length ? ['Recent:', ...recent] : []),
63      ].join('\n'),
64    }
65  })
66
67  on('turn.start', async ($, e, next) => {
68    if (e.agentId !== undefined) return next(e)
69    const result = await next(e)
70
71    const mode = await modeOf($)
72    if (mode === 'off') {
73      await update($, effective, () => '')
74      return result
75    }
76    const current = await $.session.model()
77    if (!tierOf(current)) {
78      await update($, effective, () => '')
79      return result // a model off the ladder is left alone
80    }
81
82    // The person ran /model: stand down rather than fight the choice.
83    if (sessionModel !== undefined && sessionModel !== current) {
84      sessionModel = current
85      await $.store.set(CFG, 'off')
86      await update($, effective, () => '')
87      $.ui.toast('auto-model paused: you picked a model. /auto-model on to resume.')
88      return result
89    }
90    sessionModel = current
91    await update($, base, () => current)
92
93    let target: Tier
94    let why: string
95    if (mode !== 'auto') {
96      target = mode
97      why = 'pinned'
98    } else if (e.text.trim() === '') {
99      if (!plan) return result
100      target = plan.tier // a continuation keeps the last route
101      why = 'continuation'
102    } else {
103      let c: Complexity
104      try {
105        const label = await $.model.classify(e.text.slice(0, 600), ['simple', 'moderate', 'complex'], { model: modelFor('haiku', sessionModel, await idsOf($)) })
106        c = (label as Complexity | undefined) ?? 'moderate'
107        why = 'classifier'
108      } catch {
109        c = 'moderate'
110        why = 'classifier failed'
111      }
112      target = tierFor(c)
113    }
114
115    const { context } = await $.session.usage()
116    const now = await $.clock.now()
117    const d = mode === 'auto'
118      ? decide({ target, cur: lastSent, ctx: context.tokens ?? 0, idleMs: lastDoneAt ? now - lastDoneAt : Infinity, sinceSwitch })
119      : { tier: target, reason: why }
120    plan = { ...d, turnId: e.turnId }
121
122    if (lastSent !== undefined && d.tier !== lastSent) {
123      sinceSwitch = 0
124      $.ui.toast(`auto-model: ${d.tier} (${d.reason})`)
125    } else sinceSwitch += 1
126
127    const log = ((await $.store.get(LOG)) as object[] | undefined) ?? []
128    await $.store.set(LOG, [...log, { to: d.tier, reason: `${why}: ${d.reason}`, ctx: context.tokens ?? 0 }].slice(-KEEP))
129    return result
130  })
131
132  on('turn.step', async function* ($, e, next) {
133    if (e.agentId !== undefined || !plan || plan.turnId !== e.turnId) return yield* next(e)
134    lastSent = plan.tier
135    const model = modelFor(plan.tier, sessionModel ?? e.model, await idsOf($))
136    await update($, effective, () => model) // the usage meter shows what is really sent
137    const send = (m: string) => (m === e.model ? e : { ...e, model: m })
138    const fallback = modelFor(plan.tier, undefined)
139    let started = false
140    try {
141      for await (const chunk of next(send(model))) {
142        started = true
143        yield chunk
144      }
145    } catch (err) {
146      // A derived id the API refuses before any output: retry once on the built-in id.
147      if (started || model === fallback) throw err
148      await update($, effective, () => fallback)
149      yield* next(send(fallback))
150    }
151    return
152  })
153
154  on('turn.complete', async ($, e, next) => {
155    lastDoneAt = await $.clock.now()
156    return next(e)
157  })
158}
159

1export type Tier = 'haiku' | 'sonnet' | 'opus'
2export type Complexity = 'simple' | 'moderate' | 'complex'
3export type Decision = { tier: Tier; reason: string }
4
5// Fallback ids, used only when neither an override nor the session's own generation gives one.
6export const MODEL_ID: Record<Tier, string> = {
7  haiku: 'claude-haiku-5-5',
8  sonnet: 'claude-sonnet-5-5',
9  opus: 'claude-opus-5-5',
10}
11const RANK: Record<Tier, number> = { haiku: 0, sonnet: 1, opus: 2 }
12
13// Tunable assumptions, not billing data. Prices are relative input prices per tier;
14// a cache write costs 1.25x and a cache read 0.1x the base input price.
15export const RELATIVE_PRICE: Record<Tier, number> = { haiku: 1, sonnet: 3, opus: 5 }
16const WRITE = 1.25
17const READ = 0.1
18const TURN_TOKENS = 3000 // new input plus output per turn, in full-price token equivalents
19const EXPECTED_TURNS = 5 // how long a switch is expected to hold
20const MIN_HOLD = 3 // turns before a downgrade may follow another switch
21const FREE_CTX = 20_000 // re-writing a prefix this small costs next to nothing
22const COLD_MS = 60 * 60_000 // idle this long: the cache has lapsed whatever its lifetime
23const HAIKU_MAX_CTX = 150_000 // keeps headroom under a 200k window
24
25export const tierOf = (model: string): Tier | undefined => /haiku|sonnet|opus/i.exec(model)?.[0].toLowerCase() as Tier | undefined
26
27export const tierFor = (c: Complexity): Tier => (c === 'simple' ? 'haiku' : c === 'moderate' ? 'sonnet' : 'opus')
28
29// Turns a downgrade must stay to repay the cache re-write on the cheaper model.
30export function breakEvenTurns(cur: Tier, target: Tier, ctx: number): number {
31  const saved = (ctx * READ + TURN_TOKENS) * (RELATIVE_PRICE[cur] - RELATIVE_PRICE[target])
32  return saved <= 0 ? Infinity : (ctx * WRITE * RELATIVE_PRICE[target]) / saved
33}
34
35export type DecideInput = { target: Tier; cur?: Tier; ctx: number; idleMs: number; sinceSwitch: number }
36
37export function decide({ target, cur, ctx, idleMs, sinceSwitch }: DecideInput): Decision {
38  if (cur === undefined) return { tier: target, reason: 'first turn' }
39  if (target === cur) return { tier: cur, reason: 'same tier' }
40  const isFree = ctx < FREE_CTX || idleMs > COLD_MS
41  if (RANK[target] > RANK[cur]) return { tier: target, reason: isFree ? 'upgrade, cache cheap' : 'upgrade for hard task' }
42  if (target === 'haiku' && ctx > HAIKU_MAX_CTX) return { tier: cur, reason: 'context too big for haiku' }
43  if (sinceSwitch < MIN_HOLD) return { tier: cur, reason: 'holding recent switch' }
44  if (isFree) return { tier: target, reason: 'downgrade, cache cheap' }
45  const n = breakEvenTurns(cur, target, ctx)
46  return n <= EXPECTED_TURNS
47    ? { tier: target, reason: `downgrade, repays in ~${Math.ceil(n)} turns` }
48    : { tier: cur, reason: 'switch costs more than it saves' }
49}
50
51export type Overrides = Partial<Record<Tier, string>>
52
53// The id to send for a tier: a person's override, else the session's own id with the
54// family swapped (so a new generation is followed), else the built-in default.
55export function modelFor(tier: Tier, sessionModel: string | undefined, overrides: Overrides = {}): string {
56  if (overrides[tier]) return overrides[tier]!
57  const cur = sessionModel === undefined ? undefined : tierOf(sessionModel)
58  if (sessionModel !== undefined && cur === tier) return sessionModel
59  if (sessionModel !== undefined && cur) return sessionModel.replace(/\[.*\]$/, '').replace(/haiku|sonnet|opus/i, tier)
60  return MODEL_ID[tier]
61}
62

1declare module 'claude-code' {
2  interface PluginState {
3    'auto-model': { effective: string; base: string }
4  }
5}
6