effort-router

1import type { EngineInterface, Register } from 'claude-code'
2
3// モデルは変えず、ターンごとにプロンプトを $.model.classify で分類して effort だけを切り替える。
4// effort を変えると会話部分のプロンプトキャッシュは作り直しになるため、
5// - 1ターンの中では effort を固定する（判定は最初のリクエストで一度だけ）
6// - 前のターンと同じ effort になるときは何も変えない
7// ことで切り替え回数を抑える。
8
9type Effort = 'low' | 'medium' | 'high' | 'xhigh' | 'max'
10type Mode = 'auto' | 'off' | Effort
11
12type Category = {
13  // classify に渡すラベル。小さなモデルがそのまま書き写せる短い英語にする
14  label: string
15  // undefined はセッション本来の effort のまま
16  effort: Effort | undefined
17}
18
19const CATEGORIES: readonly Category[] = [
20  { label: 'greeting or small talk', effort: 'low' },
21  { label: 'translation, summary or proofreading of text', effort: 'low' },
22  { label: 'quick question with a short answer', effort: 'low' },
23  { label: 'small edit or running a command', effort: 'medium' },
24  { label: 'coding task', effort: 'high' },
25  { label: 'hard problem needing deep thought', effort: 'xhigh' },
26  { label: 'other', effort: undefined },
27]
28const LABELS = CATEGORIES.map(c => c.label)
29const CLASSIFY_TIMEOUT_MS = 5000
30
31type Decision = { label: string; effort: Effort | undefined }
32
33let mode: Mode = 'auto'
34let last: Decision = { label: 'other', effort: undefined }
35// turnId -> 判定（turn.start で投げ、turn.step で待つ）
36const decisions = new Map<string, Promise<Decision>>()
37
38async function classify($: EngineInterface, text: string): Promise<Decision> {
39  const prompt =
40    'A developer sent this message to a coding agent. Classify what kind of request it is.\n\n' +
41    text.slice(0, 4000)
42  const label = await $.model.classify(prompt, LABELS)
43  const hit = CATEGORIES.find(c => c.label === label)
44  return hit ? { label: hit.label, effort: hit.effort } : { label: 'other', effort: undefined }
45}
46
47function decide($: EngineInterface, text: string): Promise<Decision> {
48  if (mode === 'off') return Promise.resolve({ label: 'off', effort: undefined })
49  if (mode !== 'auto') return Promise.resolve({ label: 'pinned', effort: mode })
50  // 続きのターン（テキストなし）は直前の判定を引き継ぐ
51  if (!text.trim()) return Promise.resolve(last)
52
53  const timeout = new Promise<Decision>(resolve => {
54    $.clock.after(CLASSIFY_TIMEOUT_MS, () => resolve({ label: 'timeout', effort: undefined }))
55  })
56  const judged = classify($, text).catch(() => ({ label: 'error', effort: undefined }) as Decision)
57  return Promise.race([judged, timeout])
58}
59
60export const register: Register = on => {
61  on('session.start', async ($, e, next) => {
62    await $.command.register({
63      name: 'effort-router',
64      description: 'Effort router: /effort-router [auto|off|low|medium|high|xhigh|max]',
65    })
66    return next(e)
67  })
68
69  on('command.run', { command: 'effort-router' }, async ($, e) => {
70    const arg = e.args.trim().toLowerCase()
71    if (arg === '') return { text: `effort-router: mode=${mode}, last=${last.label} (${last.effort ?? 'default'})` }
72    if (!['auto', 'off', 'low', 'medium', 'high', 'xhigh', 'max'].includes(arg)) {
73      return { text: 'usage: /effort-router [auto|off|low|medium|high|xhigh|max]' }
74    }
75    mode = arg as Mode
76    $.ui.status(`effort: ${mode}`)
77    return { text: `effort-router: mode=${mode}` }
78  })
79
80  // 分類はターン開始と同時に投げ、最初のリクエストで待つ
81  on('turn.start', async ($, e, next) => {
82    const started = await next(e)
83    decisions.set(started.turnId, decide($, e.text))
84    return started
85  })
86
87  on('turn.step', async function* ($, e, next) {
88    // サブエージェントと、effort を持たないモデルには触らない
89    if (e.agentId !== undefined || e.effort === undefined) return yield* next(e)
90
91    const pending = decisions.get(e.turnId)
92    const decision = pending ? await pending : last
93    last = decision
94
95    const effort = decision.effort ?? e.effort
96    if (e.index === 0) {
97      $.ui.status(`effort: ${effort} (${decision.label})`)
98      $.ui.log(`effort-router: ${decision.label} -> ${effort}`, { to: 'debug' })
99    }
100    return yield* next({ ...e, effort })
101  })
102
103  on('turn.complete', async ($, e, next) => {
104    const done = await next(e)
105    decisions.delete(e.turnId)
106    return done
107  })
108}
109