first-prompt-model

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { ModelPick } from '../types'
5
6const MODEL_CHOICES = [
7  { model: 'haiku', use: 'a short question, a lookup, a one-line edit, a rename, a format change' },
8  { model: 'sonnet', use: 'ordinary coding: a feature, a bug fix, tests, a refactor in a few files' },
9  { model: 'opus', use: 'hard work: architecture, design, a deep debug, a large refactor, research across many files' },
10] as const
11
12const RECOMMENDED_SUFFIX = ' (Recommended)'
13
14const pick = atom({ plugin: 'first-prompt-model', key: 'pick' } as const, null)
15const isDecided = atom({ plugin: 'first-prompt-model', key: 'isDecided' } as const, false)
16
17const ROUTER_SYSTEM = [
18  'You choose the model that will execute a coding-assistant prompt.',
19  'Choose the cheapest model that can do the task well.',
20  ...MODEL_CHOICES.map(c => `- ${c.model}: ${c.use}`),
21  `Answer with one word: ${MODEL_CHOICES.map(c => c.model).join(', ')}.`,
22].join('\n')
23
24export function parseModelAnswer(text: string): string | undefined {
25  const word = text.trim().toLowerCase().match(/[a-z]+/)?.[0]
26  return MODEL_CHOICES.find(c => c.model === word)?.model
27}
28
29export function modelOptions(suggested: string, sessionModel: string): string[] {
30  const others = MODEL_CHOICES.map(c => c.model).filter(m => m !== suggested)
31  return [suggested + RECOMMENDED_SUFFIX, ...others, `Keep ${sessionModel}`]
32}
33
34export function modelFromAnswer(answer: string): string | undefined {
35  const label = answer.endsWith(RECOMMENDED_SUFFIX) ? answer.slice(0, -RECOMMENDED_SUFFIX.length) : answer
36  return MODEL_CHOICES.find(c => c.model === label)?.model
37}
38
39export const register: Register = on => {
40  on('prompt.submit', async ($, e, next) => {
41    if (e.origin.kind !== 'composer' || (await read($, isDecided))) {
42      return next(e)
43    }
44    await update($, isDecided, () => true)
45    if ((await $.session.turns()) > 0) {
46      return next(e)
47    }
48
49    $.ui.status('first-prompt-model: Haiku is choosing a model')
50    const reply = await $.model.complete({
51      model: 'haiku',
52      system: ROUTER_SYSTEM,
53      prompt: `<prompt>\n${e.text}\n</prompt>`,
54      effort: 'low',
55      maxTokens: 16,
56      timeoutMs: 15000,
57    })
58    $.ui.status(undefined)
59    const suggested = reply.isAnswered ? parseModelAnswer(reply.text) : undefined
60    if (suggested === undefined) {
61      $.ui.toast('first-prompt-model: Haiku gave no model, the session model stays')
62      return next(e)
63    }
64
65    const sessionModel = await $.session.model()
66    const answer = await $.ui
67      .ask('Which model must run this session?', {
68        header: 'Model',
69        options: modelOptions(suggested, sessionModel),
70      })
71      .catch(() => '')
72    const model = modelFromAnswer(answer)
73    if (model === undefined) {
74      return next(e)
75    }
76
77    const chosen: ModelPick = { model, sessionModel }
78    await update($, pick, () => chosen)
79    $.ui.status(`model: ${model}`)
80    return next(e)
81  }).catch(($, e, next) => next(e))
82
83  on('turn.step', async function* ($, e, next) {
84    const chosen = await read($, pick)
85    const isMainLoop = e.agentId === undefined
86    const isModelUnchanged = chosen !== null && (await $.session.model()) === chosen.sessionModel
87    if (!isMainLoop || !isModelUnchanged) {
88      return yield* next(e)
89    }
90    return yield* next({ ...e, model: chosen.model })
91  })
92}
93

1export type ModelPick = {
2  model: string
3  sessionModel: string
4}
5
6declare module 'claude-code' {
7  interface PluginState {
8    'first-prompt-model': { pick: ModelPick | null; isDecided: boolean }
9  }
10}
11