model-tier

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Effort, Tier, TierState } from '../types'
5import { segment } from './segment'
6
7const PLUGIN = 'model-tier'
8const SET_TIER = `mcp__${PLUGIN}__set_tier`
9const TIERS: readonly Tier[] = ['haiku', 'sonnet', 'opus']
10const RANK: Record<Tier, number> = { haiku: 0, sonnet: 1, opus: 2 }
11const LABELS = ['trivial', 'standard', 'complex'] as const
12const LABEL_TIER: Record<(typeof LABELS)[number], Tier> = {
13  trivial: 'haiku',
14  standard: 'sonnet',
15  complex: 'opus',
16}
17// How much thinking the prompt needs, as the classifier names it. `max` is never
18// proposed: only the model (set_tier) or the person (/tier effort) asks for it.
19const EFFORTS: readonly Effort[] = ['low', 'medium', 'high', 'xhigh', 'max']
20const EFFORT_RANK: Record<Effort, number> = { low: 0, medium: 1, high: 2, xhigh: 3, max: 4 }
21// Named for how hard the problem is to think through, not how long the answer
22// is: measured live, 'quick-answer' caught a hard design question that asked
23// for a three-line reply.
24const EFFORT_LABELS = ['recall-or-lookup', 'routine-reasoning', 'tricky-problem', 'hard-open-problem'] as const
25const LABEL_EFFORT: Record<(typeof EFFORT_LABELS)[number], Effort> = {
26  'recall-or-lookup': 'low',
27  'routine-reasoning': 'medium',
28  'tricky-problem': 'high',
29  'hard-open-problem': 'xhigh',
30}
31// Subagent types that carry no model of their own; others keep their definition's.
32const ROUTABLE_AGENTS = new Set(['general-purpose', 'claude'])
33const CLASSIFY_CHARS = 4000
34
35// Full model ids per tier. A request built here goes to the API as written:
36// unlike /model or --model, an alias ('sonnet') is not resolved, and a router
37// in ANTHROPIC_BASE_URL answers it 404, which the engine then papers over with
38// fallbackModel ("Switched to Opus 5 because sonnet is not available").
39export const DEFAULT_TIER_MODELS: Record<Tier, string> = {
40  opus: 'claude-opus-5-5',
41  sonnet: 'claude-sonnet-5-5',
42  haiku: 'claude-haiku-5-5',
43}
44
45/** `opus=claude-opus-5-5,haiku=…` over the defaults; bad entries are ignored. */
46export function tierModels(spec: unknown): Record<Tier, string> {
47  const models = { ...DEFAULT_TIER_MODELS }
48  if (typeof spec !== 'string') return models
49  for (const pair of spec.split(',')) {
50    const [tier, id] = pair.split('=').map(part => part?.trim() ?? '')
51    if ((tier === 'opus' || tier === 'sonnet' || tier === 'haiku') && id) models[tier] = id
52  }
53  return models
54}
55
56export const INITIAL: TierState = {
57  tier: 'opus',
58  source: 'auto',
59  pendingCount: 0,
60  lastInputTokens: 0,
61  effortSource: 'auto',
62  pendingEffortCount: 0,
63}
64
65// Change the tag when TierState's fields change: it declines an older value.
66const state = atom({ plugin: 'model-tier', key: 'tier' } as const, INITIAL, {
67  shape: 'tier-2',
68})
69
70export type DecideOptions = {
71  downgradeAfterTurns: number
72  maxTokensForHaiku: number
73  maxTokensForDowngrade: number
74}
75
76/**
77 * Folds one classifier verdict into the state. Upgrades apply at once;
78 * downgrades wait for `downgradeAfterTurns` consecutive lower verdicts, since
79 * every model switch re-bills the conversation prefix outside the prompt cache.
80 */
81export function decide(s: TierState, proposed: Tier, o: DecideOptions): TierState {
82  if (s.source === 'pinned') return s
83  let target = proposed
84  if (target === 'haiku' && s.lastInputTokens > o.maxTokensForHaiku) target = 'sonnet'
85
86  if (RANK[target] >= RANK[s.tier]) {
87    const changed = target !== s.tier
88    return {
89      ...s,
90      tier: target,
91      source: changed ? 'auto' : s.source,
92      pendingDown: undefined,
93      pendingCount: 0,
94    }
95  }
96
97  // Nothing sent yet means no cached prefix to lose: step down at once.
98  // A large context re-billed outside the cache costs more than a cheaper model saves.
99  if (s.lastInputTokens > o.maxTokensForDowngrade) return { ...s, pendingDown: undefined, pendingCount: 0 }
100
101  const count = s.pendingDown === target ? s.pendingCount + 1 : 1
102  if (s.lastInputTokens === 0 || count >= o.downgradeAfterTurns) {
103    return { ...s, tier: target, source: 'auto', pendingDown: undefined, pendingCount: 0 }
104  }
105  return { ...s, pendingDown: target, pendingCount: count }
106}
107
108/**
109 * Folds one effort verdict into the state, the way `decide` folds a tier:
110 * more effort at once, less only after `downgradeAfterTurns` lower verdicts in
111 * a row (a request's thinking setting is part of what the prompt cache keys on).
112 */
113export function decideEffort(s: TierState, proposed: Effort, o: DecideOptions): TierState {
114  if (s.effortSource === 'pinned') return s
115  if (s.effort === undefined || EFFORT_RANK[proposed] >= EFFORT_RANK[s.effort]) {
116    const changed = proposed !== s.effort
117    return { ...s, effort: proposed, effortSource: changed ? 'auto' : s.effortSource, pendingEffortDown: undefined, pendingEffortCount: 0 }
118  }
119  const count = s.pendingEffortDown === proposed ? s.pendingEffortCount + 1 : 1
120  if (s.lastInputTokens === 0 || count >= o.downgradeAfterTurns) {
121    return { ...s, effort: proposed, effortSource: 'auto', pendingEffortDown: undefined, pendingEffortCount: 0 }
122  }
123  return { ...s, pendingEffortDown: proposed, pendingEffortCount: count }
124}
125
126const isEffort = (v: unknown): v is Effort => typeof v === 'string' && (EFFORTS as readonly string[]).includes(v)
127
128/** Whether a resolved model id (`claude-opus-5-5[1m]`) is the given tier. */
129export function isTier(model: string, tier: Tier): boolean {
130  return model === tier || model.toLowerCase().includes(tier)
131}
132
133export function tierOf(model: string): Tier | undefined {
134  return TIERS.find(t => isTier(model, t))
135}
136
137const isTierName = (v: unknown): v is Tier => typeof v === 'string' && (TIERS as readonly string[]).includes(v)
138
139const label = (s: TierState) => `${s.tier}${s.effort ? `·${s.effort}` : ''} (${s.source === s.effortSource || s.effort === undefined ? s.source : `${s.source}/${s.effortSource}`})`
140
141// The tier rides the person's status line; nothing about it needs a notice row.
142function show($: EngineInterface, s: TierState, enabled: boolean) {
143  const pinned = s.source === 'pinned' || s.effortSource === 'pinned'
144  // A row an earlier version pinned outlives a reload: take it down.
145  $.ui.status(undefined)
146  void segment($, enabled ? `${s.tier}${s.effort ? `·${s.effort}` : ''}${pinned ? ' (pin)' : ''}` : '')
147}
148
149/** Labels `text`, or gives up after `timeoutMs` so a slow classifier never holds the prompt. */
150async function classify($: EngineInterface, text: string, model: string, timeoutMs: number): Promise<Tier | undefined> {
151  const stop = new AbortController()
152  // A sleep that fails or is aborted never wins the race: only a real timeout does.
153  const timeout = $.clock.sleep(timeoutMs, { signal: stop.signal }).then(
154    () => undefined,
155    () => new Promise<undefined>(() => {}),
156  )
157  const labelled = $.model.classify(text.slice(0, CLASSIFY_CHARS), LABELS, { model }).then(label =>
158    label === undefined ? undefined : LABEL_TIER[label as (typeof LABELS)[number]],
159  )
160  try {
161    return await Promise.race([labelled, timeout])
162  } finally {
163    stop.abort()
164  }
165}
166
167/** The effort a prompt needs, or undefined on a timeout or a label outside the set. */
168async function classifyEffort($: EngineInterface, text: string, model: string, timeoutMs: number): Promise<Effort | undefined> {
169  const stop = new AbortController()
170  const timeout = $.clock.sleep(timeoutMs, { signal: stop.signal }).then(
171    () => undefined,
172    () => new Promise<undefined>(() => {}),
173  )
174  const labelled = $.model.classify(text.slice(0, CLASSIFY_CHARS), EFFORT_LABELS, { model }).then(found =>
175    found !== undefined && found in LABEL_EFFORT ? LABEL_EFFORT[found as (typeof EFFORT_LABELS)[number]] : undefined,
176  )
177  try {
178    return await Promise.race([labelled, timeout])
179  } finally {
180    stop.abort()
181  }
182}
183
184export const register: Register = (on, options) => {
185  const enabled = options.enabled !== false
186  const decideOptions: DecideOptions = {
187    downgradeAfterTurns: Number(options.downgradeAfterTurns ?? 2),
188    maxTokensForHaiku: Number(options.maxTokensForHaiku ?? 150000),
189    maxTokensForDowngrade: Number(options.maxTokensForDowngrade ?? 200000),
190  }
191  const classifierModel = String(options.classifierModel ?? 'haiku')
192  const models = tierModels(options.tierModels)
193  const classifierTimeoutMs = Number(options.classifierTimeoutMs ?? 3000)
194  const routeEffort = options.routeEffort !== false
195
196  on('session.start', async ($, e, next) => {
197    const started = await next(e)
198
199    await $.tool.register({
200      name: 'set_tier',
201      description:
202        'Switch the model and/or the reasoning effort that run the rest of this task: ' +
203        'opus for design, debugging and multi-file reasoning; sonnet for routine implementation; ' +
204        'haiku for lookups, formatting and mechanical edits. Effort: low for quick answers, ' +
205        'medium for ordinary work, high for careful reasoning, xhigh for hard investigations, ' +
206        'max only when nothing less will do. Call it when the classifier\'s pick is clearly too ' +
207        'weak or wasteful for the work ahead. Takes effect from the next model request.',
208      inputSchema: {
209        type: 'object',
210        properties: {
211          tier: { type: 'string', enum: TIERS },
212          effort: { type: 'string', enum: EFFORTS },
213          reason: { type: 'string', description: 'One line on why' },
214        },
215      },
216      isDeferred: false,
217    })
218    await $.command.register({
219      name: 'tier',
220      description: 'Pin the model tier (opus|sonnet|haiku) or return to auto routing',
221    })
222
223    // A fresh session starts on whatever /model is set to.
224    const s = await read($, state)
225    if (s === INITIAL || (s.lastInputTokens === 0 && s.source === 'auto')) {
226      const tier = tierOf(await $.session.model())
227      if (tier) await update($, state, v => ({ ...v, tier }))
228    }
229    // Loaded into a running session: its context is already large and cached.
230    const { context } = await $.session.usage()
231    const tokens = context.tokens ?? 0
232    if (tokens > 0) await update($, state, v => ({ ...v, lastInputTokens: Math.max(v.lastInputTokens, tokens) }))
233    show($, await read($, state), enabled)
234    return started
235  })
236
237  on('command.run', { command: 'tier' }, async ($, e) => {
238    const arg = e.args.trim().toLowerCase()
239    if (arg === '') {
240      const s = await read($, state)
241      return { text: `tier: ${label(s)}` }
242    }
243    if (arg.startsWith('effort')) {
244      const level = arg.slice('effort'.length).trim()
245      if (level === 'auto') {
246        const s = await update($, state, (v): TierState => ({ ...v, effortSource: 'auto', pendingEffortDown: undefined, pendingEffortCount: 0 }))
247        show($, s, enabled)
248        return { text: 'Effort follows the classifier again from the next prompt.' }
249      }
250      if (!isEffort(level)) return { text: 'Usage: /tier effort [auto|low|medium|high|xhigh|max]' }
251      const s = await update($, state, (v): TierState => ({ ...v, effort: level, effortSource: 'pinned', pendingEffortDown: undefined, pendingEffortCount: 0 }))
252      show($, s, enabled)
253      return { text: `Effort pinned to ${level}. /tier effort auto to release.` }
254    }
255    if (arg === 'auto') {
256      const s = await update($, state, (v): TierState => ({ ...v, source: 'auto', pendingDown: undefined, pendingCount: 0 }))
257      show($, s, enabled)
258      return { text: `Auto routing on; tier stays ${s.tier} until the next prompt.` }
259    }
260    if (!isTierName(arg)) return { text: 'Usage: /tier [auto|opus|sonnet|haiku] | /tier effort [auto|low|medium|high|xhigh|max]' }
261    const s = await update($, state, (v): TierState => ({ ...v, tier: arg, source: 'pinned', pendingDown: undefined, pendingCount: 0 }))
262    show($, s, enabled)
263    return { text: `Pinned to ${arg}. /tier auto to release.` }
264  })
265
266  on('prompt.submit', async ($, e, next) => {
267    if (!enabled || e.isMeta || e.command !== undefined || e.text.trim() === '') return next(e)
268    const current = await read($, state)
269    let s = current
270    const wantTier = current.source !== 'pinned'
271    const wantEffort = routeEffort && current.effortSource !== 'pinned'
272    const [proposed, effort] = await Promise.all([
273      wantTier ? classify($, e.text, classifierModel, classifierTimeoutMs).catch(() => undefined) : undefined,
274      wantEffort ? classifyEffort($, e.text, classifierModel, classifierTimeoutMs).catch(() => undefined) : undefined,
275    ])
276    if (proposed) s = await update($, state, v => decide(v, proposed, decideOptions))
277    if (effort) s = await update($, state, v => decideEffort(v, effort, decideOptions))
278    show($, s, enabled)
279    const effortPart = s.effort ? ` with ${s.effort} effort (${s.effortSource})` : ''
280    const note =
281      `[model-tier] This turn runs on ${s.tier} (${s.source})${effortPart}. Review that pick against the work ahead: ` +
282      `if it clearly needs a stronger model or more effort, or allows a cheaper model or less, call ${SET_TIER} with tier and/or effort.`
283    return next({ ...e, context: [...(e.context ?? []), note] })
284  }).catch(($, e, next) => next(e))
285
286  on('tool.call', { tool: 'mcp__model-tier__set_tier' }, async ($, e) => {
287    const input = e as unknown as { tier?: unknown; effort?: unknown; reason?: unknown }
288    if (input.tier !== undefined && !isTierName(input.tier)) return { deny: 'tier must be one of opus, sonnet, haiku' }
289    if (input.effort !== undefined && !isEffort(input.effort)) return { deny: 'effort must be one of low, medium, high, xhigh, max' }
290    if (input.tier === undefined && input.effort === undefined) return { deny: 'give tier, effort or both' }
291    const tier = input.tier as Tier | undefined
292    const effort = input.effort as Effort | undefined
293    const s = await update($, state, (v): TierState => {
294      let next = v
295      if (tier !== undefined && v.source !== 'pinned') next = { ...next, tier, source: 'model', pendingDown: undefined, pendingCount: 0 }
296      if (effort !== undefined && v.effortSource !== 'pinned') next = { ...next, effort, effortSource: 'model', pendingEffortDown: undefined, pendingEffortCount: 0 }
297      return next
298    })
299    show($, s, enabled)
300    const parts: string[] = []
301    if (tier !== undefined) parts.push(s.source === 'pinned' ? `The person pinned ${s.tier} with /tier; it stays.` : `Switched to ${tier} from the next request.`)
302    if (effort !== undefined) parts.push(s.effortSource === 'pinned' ? `The person pinned ${s.effort} effort with /tier effort; it stays.` : `Effort ${effort} from the next request.`)
303    return { result: parts.join(' ') }
304  }).catch(() => ({ deny: 'model-tier: set_tier failed; the tier is unchanged.' }))
305
306  on('turn.step', async function* ($, e, next) {
307    if (!enabled || e.agentId !== undefined) return yield* next(e)
308    const s = await read($, state)
309    const switching = !isTier(e.model, s.tier)
310    // A model that takes no effort reads `effort` absent: nothing to rewrite there.
311    const reEffort = s.effort !== undefined && e.effort !== undefined && e.effort !== s.effort
312    const stream = next({ ...e, ...(switching ? { model: models[s.tier] } : {}), ...(reEffort ? { effort: s.effort } : {}) })
313
314    let yielded = false
315    let result
316    try {
317      while (true) {
318        const step = await stream.next()
319        if (step.done) {
320          result = step.value
321          break
322        }
323        yielded = true
324        yield step.value
325      }
326    } catch (err) {
327      // Refused before anything streamed: retry on the session's own model.
328      if (!switching || yielded) throw err
329      $.ui.toast(`model-tier: ${models[s.tier]} was refused, staying on ${e.model}`)
330      result = yield* next(e)
331    }
332
333    const u = result.usage
334    if (u) {
335      const tokens = u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
336      await update($, state, v => ({ ...v, lastInputTokens: tokens }))
337    }
338    return result
339  })
340
341  on('agent.spawn', async ($, e, next) => {
342    if (!enabled || e.fork || e.model !== undefined || !ROUTABLE_AGENTS.has(e.subagentType)) return next(e)
343    const tier = await classify($, e.prompt, classifierModel, classifierTimeoutMs).catch(() => undefined)
344    return next(tier ? { ...e, model: models[tier] } : e)
345  }).catch(($, e, next) => next(e))
346}
347

1import type { EngineInterface } from 'claude-code'
2
3// This mod's piece of the person's own status line: ~/.claude/statusline-command.sh
4// joins ~/.claude/statusline/<sessionId>/<mod>.txt after its own fields, so the
5// engine's one-row-per-plugin notices (`$.ui.status`) stay for what needs the
6// person's eye. A leading '!' asks the status line to draw the piece in red;
7// '' clears it. Copied per mod: a mod imports nothing outside its folder.
8
9const MOD = 'model-tier'
10let last: string | undefined
11
12export async function segment($: EngineInterface, text: string): Promise<void> {
13  if (text === last) return
14  last = text
15  try {
16    const home = await $.env.get('HOME')
17    if (home === undefined) return
18    await $.fs.write(`${home}/.claude/statusline/${await $.session.id()}/${MOD}.txt`, text)
19  } catch {
20    // No status line piece is worth failing a hook over.
21    last = undefined
22  }
23}
24

1export type Tier = 'opus' | 'sonnet' | 'haiku'
2
3export type Effort = 'low' | 'medium' | 'high' | 'xhigh' | 'max'
4
5export type TierState = {
6  /** The tier the main loop runs on now. */
7  tier: Tier
8  /** Who set `tier`: the classifier, the model's set_tier call, or /tier. */
9  source: 'auto' | 'model' | 'pinned'
10  /** The lower tier the classifier keeps proposing, and how many turns in a row. */
11  pendingDown?: Tier
12  pendingCount: number
13  /** Prompt tokens of the last main-loop request (input + cache read + cache write). */
14  lastInputTokens: number
15  /** The effort requests are sent with; unset keeps the session's own (/effort). */
16  effort?: Effort
17  /** Who set `effort`: the classifier, the model's set_tier call, or /tier effort. */
18  effortSource: 'auto' | 'model' | 'pinned'
19  /** The lower effort the classifier keeps proposing, and how many turns in a row. */
20  pendingEffortDown?: Effort
21  pendingEffortCount: number
22}
23
24declare module 'claude-code' {
25  interface PluginState {
26    'model-tier': { tier: Shaped<TierState> }
27  }
28}
29