model-router

1import type { EngineInterface, Register } from 'claude-code'
2
3// ターンごとに TypeSafe Jev でプロンプトを判定し、メインループのモデルを切り替える。
4// 既定は Sonnet。Jev の確信度が低い・失敗・タイムアウトのときも Sonnet に倒す。
5
6type Tier = 'haiku' | 'sonnet' | 'opus'
7type Mode = 'auto' | 'off' | Tier
8
9const MODELS: Record<Tier, string> = {
10  haiku: 'claude-haiku-4-5-20251001',
11  sonnet: 'claude-sonnet-5-5',
12  opus: 'claude-opus-5-5',
13}
14const DEFAULT_TIER: Tier = 'sonnet'
15// 選ばれた選択肢の確率がこれ未満なら既定 (Sonnet) にする
16const MIN_PROBABILITY = 0.55
17const JEV_TIMEOUT_MS = 4000
18const JEV_URL = 'https://api.typesafe.ai/v1/systemone'
19const KEYCHAIN_SERVICE = 'TYPESAFE_API_KEY'
20
21type Route = { tier: Tier; reason: string }
22
23const ROUTE_QUESTION = {
24  type: 'choice',
25  instructions: {
26    question:
27      'A developer typed `prompt` to a coding agent working in their repository. ' +
28      'Which model tier should handle this request? Pick the cheapest tier that can do the job well.',
29    note: '`previous_tier` is the tier that handled the previous request; a short follow-up usually needs the same tier.',
30  },
31  criteria: {
32    haiku:
33      'Small, mechanical or lookup work: answer a quick factual question, rename or reformat, run a command and report, ' +
34      'fix a typo, list files, simple one-file edits with an obvious change, translate or summarize short text.',
35    sonnet:
36      'Ordinary software work: implement a feature, fix a bug with a findable cause, write tests, explain code, ' +
37      'refactor a module, review a diff, or anything that does not clearly fit the other two.',
38    opus:
39      'Hard work that needs deep reasoning: system or architecture design, debugging an unclear cause across many components, ' +
40      'large multi-file refactors or migrations, subtle concurrency, performance or security analysis, or planning an ambiguous open-ended project.',
41  },
42} as const
43
44let apiKey: string | undefined
45let mode: Mode = 'auto'
46let lastTier: Tier = DEFAULT_TIER
47// turnId -> 判定（turn.start で投げ、turn.step で待つ）
48const routes = new Map<string, Promise<Route>>()
49
50async function loadKey($: EngineInterface): Promise<string | undefined> {
51  if (apiKey) return apiKey
52  const r = await $.process.run(['security', 'find-generic-password', '-s', KEYCHAIN_SERVICE, '-w'])
53  if (r.exitCode !== 0) return undefined
54  apiKey = r.stdout.trim() || undefined
55  return apiKey
56}
57
58async function askJev($: EngineInterface, prompt: string): Promise<Route> {
59  const key = await loadKey($)
60  if (!key) return { tier: DEFAULT_TIER, reason: 'no TYPESAFE_API_KEY' }
61
62  const res = await $.http.fetch(JEV_URL, {
63    method: 'POST',
64    headers: { Authorization: `Bearer ${key}`, 'Content-Type': 'application/json' },
65    body: JSON.stringify({
66      model: 'jev-latest',
67      state: { prompt: prompt.slice(0, 8000), previous_tier: lastTier },
68      questions: { tier: ROUTE_QUESTION },
69    }),
70  })
71  if (!res.ok) return { tier: DEFAULT_TIER, reason: `jev HTTP ${res.status}` }
72
73  const answer = JSON.parse(res.text)?.answers?.tier as
74    | { choice: Tier; probabilities: Record<Tier, number> }
75    | undefined
76  if (!answer || !(answer.choice in MODELS)) return { tier: DEFAULT_TIER, reason: 'jev: no answer' }
77
78  const p = answer.probabilities[answer.choice] ?? 0
79  if (p < MIN_PROBABILITY) {
80    return { tier: DEFAULT_TIER, reason: `jev unsure (${answer.choice} ${p.toFixed(2)})` }
81  }
82  return { tier: answer.choice, reason: `jev ${p.toFixed(2)}` }
83}
84
85function route($: EngineInterface, text: string): Promise<Route> {
86  if (mode === 'off') return Promise.resolve({ tier: DEFAULT_TIER, reason: 'router off' })
87  if (mode !== 'auto') return Promise.resolve({ tier: mode, reason: 'pinned' })
88  // 続きのターン（テキストなし）は直前の判定を引き継ぐ
89  if (!text.trim()) return Promise.resolve({ tier: lastTier, reason: 'continuation' })
90
91  const fallback = new Promise<Route>(resolve => {
92    $.clock.after(JEV_TIMEOUT_MS, () => resolve({ tier: DEFAULT_TIER, reason: 'jev timeout' }))
93  })
94  const judged = askJev($, text).catch(() => ({ tier: DEFAULT_TIER, reason: 'jev error' }) as Route)
95  return Promise.race([judged, fallback])
96}
97
98export const register: Register = on => {
99  on('session.start', async ($, e, next) => {
100    await $.command.register({
101      name: 'router',
102      description: 'Model router: /router [auto|off|haiku|sonnet|opus]',
103    })
104    return next(e)
105  })
106
107  on('command.run', { command: 'router' }, async ($, e) => {
108    const arg = e.args.trim().toLowerCase()
109    if (arg === '') return { text: `model-router: mode=${mode}, last=${lastTier}` }
110    if (!['auto', 'off', 'haiku', 'sonnet', 'opus'].includes(arg)) {
111      return { text: 'usage: /router [auto|off|haiku|sonnet|opus]' }
112    }
113    mode = arg as Mode
114    $.ui.status(`router: ${mode}`)
115    return { text: `model-router: mode=${mode}` }
116  })
117
118  // 判定はターン開始と同時に投げ、最初のリクエストで待つ
119  on('turn.start', async ($, e, next) => {
120    const started = await next(e)
121    routes.set(started.turnId, route($, e.text))
122    return started
123  })
124
125  on('turn.step', async function* ($, e, next) {
126    // サブエージェントは自分のモデル設定に任せる
127    if (e.agentId !== undefined) return yield* next(e)
128
129    const pending = routes.get(e.turnId)
130    const { tier, reason } = pending ? await pending : { tier: lastTier, reason: 'no route' }
131    lastTier = tier
132
133    if (e.index === 0) $.ui.status(`router: ${tier} (${reason})`)
134
135    // Haiku 4.5 は effort 非対応
136    const effort = tier === 'haiku' ? undefined : e.effort
137    return yield* next({ ...e, model: MODELS[tier], effort })
138  })
139
140  on('turn.complete', async ($, e, next) => {
141    const done = await next(e)
142    if (e.turnId) routes.delete(e.turnId)
143    return done
144  })
145}
146