model-router

1import { atom, read, update } from 'claude-code'
2import type { Register, SessionRateLimit } from 'claude-code'
3import type { Effort, Pick, Tier, Usage } from '../types'
4
5const last = atom({ plugin: 'model-router', key: 'last' } as const, null)
6const auto = atom({ plugin: 'model-router', key: 'auto' } as const, true)
7const usage = atom({ plugin: 'model-router', key: 'usage' } as const, [])
8
9
10const NAME: Record<Tier, string> = { haiku: 'Haiku', sonnet: 'Sonnet', opus: 'Opus' }
11const COLOR: Record<Tier, string> = { haiku: '#7FB77E', sonnet: '#6FA8DC', opus: '#C9A0DC' }
12
13// Free, local routing: no model call, so picking a model never spends usage.
14const OPUS = /\b(build|design|architect\w*|refactor|strateg\w*|plan|research|analy[sz]e|develop|implement|migrate|debug|investigate|from scratch|end.to.end|entire|everything|optimi[sz]e|mod|plugin|system|pipeline|automat\w*)\b/i
15const SONNET = /\b(fix|edit|update|change|write|draft|rewrite|add|remove|create|make|script|code|function|caption|email|summari[sz]e|convert|post|schedule)\b/i
16const DEEP = /\b(debug|investigate|architect\w*|refactor|migrate|from scratch|end.to.end|entire|everything|why (is|does|did)|root cause)\b/i
17const QUICK = /^(hi|hey|yo|ok|okay|thanks|thx|ty|yes|no|nah|yep|sure|cool|nice|lol)\b|^(what|when|where|who|which|is|are|did|does|do|can|how much|how many)\b/i
18
19const DOWN: Record<Effort, Effort> = { xhigh: 'high', high: 'medium', medium: 'low', low: 'low' }
20
21export function route(text: string, weekUsed: number | null): { tier: Tier; effort: Effort; reason: string } {
22  const t = text.trim(), len = t.length
23  const bullets = (t.match(/^\s*([-*•]|\d+[.)])\s+/gm) || []).length
24  let tier: Tier, effort: Effort, reason: string
25  if (len > 900 || bullets >= 4 || (OPUS.test(t) && len > 60)) {
26    tier = 'opus'; reason = len > 900 ? 'long, multi-part ask' : bullets >= 4 ? `${bullets}-item task list` : 'build / plan / research'
27    effort = len > 900 || bullets >= 4 || DEEP.test(t) ? 'xhigh' : 'high'
28  } else if (len < 140 && !OPUS.test(t) && (QUICK.test(t) || !SONNET.test(t))) {
29    tier = 'haiku'; effort = 'low'; reason = 'quick question'
30  } else {
31    tier = 'sonnet'; effort = DEEP.test(t) ? 'high' : 'medium'; reason = OPUS.test(t) ? 'short build ask' : 'edit / write task'
32  }
33  // Protect the weekly limit: step down (model and effort) when it is nearly gone.
34  if (weekUsed != null && weekUsed >= 95 && tier !== 'haiku') return { tier: 'haiku', effort: 'low', reason: `${reason} · week almost gone` }
35  if (weekUsed != null && weekUsed >= 85 && tier === 'opus') return { tier: 'sonnet', effort: DOWN[effort], reason: `${reason} · saving the week` }
36  return { tier, effort, reason }
37}
38
39// "yes" after my offer to build something is a build, not a quick question. A short reply is
40// routed on what Claude just proposed (the end of its last message) plus what the person typed.
41const AFFIRM = /^(y|yes+|yeah|yep|yup|ya|sure|ok(ay)?|k|go|go ahead|go for it|do it|do that|lets? go|let's do it|continue|proceed|start|ship it|send it|approved?|perfect|all good|sounds good|👍|✅)\b/i
42const SMALLTALK = /^(thanks|thank you|thx|ty|lol|haha|cool|nice|good night|gn|bye)\b/i
43export function routeWithContext(text: string, prev: string, weekUsed: number | null): { tier: Tier; effort: Effort; reason: string } {
44  const t = text.trim(), tail = prev.trim().slice(-600)
45  const own = route(t, weekUsed)
46  if (!tail || t.length >= 160 || SMALLTALK.test(t) || own.tier === 'opus') return own
47  const ctx = route(tail, weekUsed)
48  if (AFFIRM.test(t)) return ctx.tier === 'haiku' ? own : { ...ctx, reason: `${ctx.reason} · you said “${t.slice(0, 24)}” to my plan` }
49  // short follow-up in the middle of a bigger task: at least Sonnet
50  if (own.tier === 'haiku' && ctx.tier !== 'haiku' && !/\?\s*$/.test(t)) return { tier: 'sonnet', effort: 'medium', reason: 'short reply inside a bigger task' }
51  return own
52}
53
54const toUsage = (r: readonly SessionRateLimit[]): Usage =>
55  r.map(x => ({ kind: x.kind, left: Math.max(0, Math.round((100 - x.percentUsed) * 10) / 10), resetsAt: x.resetsAt }))
56const tierOf = (id: string): Tier | null => (/haiku/i.test(id) ? 'haiku' : /sonnet/i.test(id) ? 'sonnet' : /opus/i.test(id) ? 'opus' : null)
57const label = (k: string) => (k === 'five_hour' ? '5h' : k === 'seven_day' ? 'Week' : k === 'spend_limit' ? 'Spend' : k)
58const bar = (left: number) => { const n = Math.round(left / 10); return '█'.repeat(n) + '░'.repeat(10 - n) }
59const tone = (left: number) => (left > 50 ? '#7FB77E' : left > 20 ? '#E5C07B' : '#E06C75')
60const resets = (iso?: string) => {
61  if (!iso) return ''
62  const ms = new Date(iso).getTime() - Date.now(); if (!(ms > 0)) return ''
63  const h = Math.floor(ms / 36e5); return h >= 24 ? ` · ${Math.floor(h / 24)}d ${h % 24}h` : ` · ${h}h ${Math.floor((ms % 36e5) / 6e4)}m`
64}
65
66export const register: Register = on => {
67  on('session.start', async ($, e, next) => {
68    const u = await $.session.usage().catch(() => null)
69    if (u) { const v = toUsage(u.rateLimits); await update($, usage, () => v) }
70    return next(e)
71  })
72
73  on('session.measure', async ($, e, next) => {
74    if (e.changed.includes('rateLimits')) { const v = toUsage(e.rateLimits); await update($, usage, () => v) }
75    return next(e)
76  })
77
78  on('prompt.submit', async ($, e, next) => {
79    const text = e.text.trim()
80    if (!text || text.startsWith('/') || e.origin.kind === 'plugin' || !(await read($, auto))) return next(e)
81    const week = (await read($, usage)).find(x => x.kind === 'seven_day')
82    const msgs = await $.session.messages().catch(() => [] as any[])
83    const prev = (msgs as any[]).filter(m => m.role === 'assistant' && typeof m.text === 'string' && m.text.trim()).at(-1)?.text || ''
84    const { tier, effort, reason } = routeWithContext(text, prev, week ? 100 - week.left : null)
85    const now = tierOf(await $.session.model().catch(() => ''))
86    let switched = false
87    if (now !== tier) {
88      const r = await $.command.run({ command: 'model', args: tier }).catch(() => null)
89      switched = !!r
90    }
91    const pick: Pick = { tier, effort, reason, switched }
92    await update($, last, () => pick)
93    return next(e)
94  })
95
96  // Effort rides on each model request of the turn (main thread only; subagents keep theirs).
97  // A model without effort support has no e.effort, so nothing is sent for it.
98  on('turn.step', async function* ($, e, next) {
99    const pick = await read($, last)
100    if (e.agentId !== undefined || e.effort === undefined || !pick?.effort || !(await read($, auto))) return yield* next(e)
101    return yield* next({ ...e, effort: pick.effort })
102  })
103
104  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
105    // Other mods' bands (autopilot-live) stay: draw them first, ours under.
106    const rest: any = await next(e)
107    if (e.props.hasSurvey) return rest
108    const { Box, Text, Button } = $.ui.resolve(e)
109    const pick = await read($, last), on_ = await read($, auto), u = await read($, usage)
110    const wins = u.filter(x => x.kind === 'five_hour' || x.kind === 'seven_day' || x.kind === 'spend_limit')
111    const mine = (
112      <Box flexDirection="column">
113        <Box>
114          <Text bold color={pick ? COLOR[pick.tier] : undefined}>◆ {on_ ? 'Auto model' : 'Auto model off'}</Text>
115          {on_ && pick ? <Text> → <Text bold color={COLOR[pick.tier]}>{NAME[pick.tier]}</Text>{pick.effort ? <Text color={COLOR[pick.tier]}> · {pick.effort}</Text> : null}<Text dimColor> · {pick.reason}{pick.switched ? '' : ' (already on it)'}</Text></Text> : null}
116          {on_ && !pick ? <Text dimColor> · picks Haiku / Sonnet / Opus when you send</Text> : null}
117          <Text> </Text>
118          <Button key="toggle" label={on_ ? 'Pause' : 'Resume'} onPress={() => update($, auto, v => !v)} />
119        </Box>
120        {wins.length
121          ? <Box>{wins.map(w => <Text key={w.kind}><Text dimColor>{label(w.kind)} </Text><Text color={tone(w.left)}>{bar(w.left)} {w.left}% left</Text><Text dimColor>{resets(w.resetsAt)}   </Text></Text>)}</Box>
122          : <Text dimColor>Usage: shows after your first reply</Text>}
123      </Box>
124    )
125    return rest && rest.type !== 'engine' ? <Box flexDirection="column">{rest}{mine}</Box> : mine
126  })
127}
128

1export type Tier = 'haiku' | 'sonnet' | 'opus'
2export type Effort = 'low' | 'medium' | 'high' | 'xhigh'
3export type Pick = { tier: Tier; effort: Effort; reason: string; switched: boolean }
4export type Usage = { kind: string; left: number; resetsAt?: string }[]
5
6declare module 'claude-code' {
7  interface PluginState {
8    'model-router': { last: Pick | null; auto: boolean; usage: Usage }
9  }
10}
11