model-picker

1// Asks how hard a new session is before its first prompt runs, switches to the
2// matching model, and logs one JSON file per session in
3// ~/.local/state/model-picker/: the answer, how long it took, the model
4// chosen, each turn's model and duration, and the session's end and cost.
5
6const MODEL = {
7  'Trivial · Haiku': 'haiku',
8  'Easy · Sonnet': 'sonnet',
9  'Hard · Opus': 'opus',
10  'Intense · Fable': 'fable',
11}
12
13const QUESTION = [
14  'Trivial: mechanical; a wrong result is obvious at a glance.',
15  'Easy: goal and method known; tests or a quick look catch mistakes.',
16  'Hard: it can look right and be wrong, or the design is not settled.',
17  'Intense: hard, and every claim or word will be judged, or it runs long.',
18  '',
19  'How hard is this session?',
20].join('\n')
21
22async function recordPath($, sessionId) {
23  const home = await $.env.get('HOME')
24  if (!home) throw new Error('HOME is not set')
25  return home + '/.local/state/model-picker/' + sessionId + '.json'
26}
27
28// Models that answered each turn's requests, keyed by turn id, until the turn ends.
29const answeredBy = new Map()
30
31export function register(on) {
32  on('prompt.submit', async ($, e, next) => {
33    // Only prompts a person typed; -p runs, peers and notifications pass through.
34    if (e.origin.kind !== 'composer' && e.origin.kind !== 'bridge') return next(e)
35    const sessionId = await $.session.id()
36    const path = await recordPath($, sessionId)
37    if (await $.fs.exists(path)) return next(e)
38    // A resumed session already has its model and its cache.
39    if ((await $.session.messages()).length > 0) return next(e)
40
41    const record = {
42      sessionId,
43      cwd: await $.session.cwd(),
44      askedAt: new Date().toISOString(),
45      modelBefore: await $.session.model(),
46      prompt: e.text.slice(0, 300),
47      promptChars: e.text.length,
48      answer: null,
49      difficulty: null,
50      model: null,
51      turns: [],
52    }
53    try {
54      record.answer = await $.ui.ask(QUESTION, { options: Object.keys(MODEL), header: 'Difficulty' })
55    } catch {
56      // Dismissed: keep the current model, and log the skip so it counts.
57      record.skipped = true
58    }
59    record.answeredAt = new Date().toISOString()
60    record.decisionMs = Date.parse(record.answeredAt) - Date.parse(record.askedAt)
61    if (record.answer !== null) {
62      record.model = MODEL[record.answer] ?? null
63      // A label from the menu, or the words typed under Other.
64      record.difficulty = record.model ? record.answer.split(' · ')[0] : record.answer
65    }
66    await $.fs.write(path, JSON.stringify(record))
67    // Skipped, typed under Other, or already on the chosen model: send as is.
68    if (!record.model || record.modelBefore.includes(record.model)) return next(e)
69
70    // /model can't run while this hook holds the turn, so hand the prompt back
71    // and switch once the session is idle; nothing is cached yet.
72    $.clock.after(0, async () => {
73      try {
74        await $.command.run({ command: 'model', args: record.model })
75        record.switchedAt = new Date().toISOString()
76      } catch (err) {
77        $.ui.status('model-picker: switch to ' + record.model + ' failed: ' + String(err))
78        record.switchError = String(err)
79      }
80      await $.fs.write(path, JSON.stringify(record))
81    })
82    return { drop: 'Switching to ' + record.answer.split(' · ')[1] + '. Press Enter to send.' }
83  }).catch(async ($, e, next) => {
84    $.ui.status('model-picker failed: ' + next.error.kind + ' ' + (next.error.message ?? ''))
85    return next(e)
86  })
87
88  // The models that actually answered each main-conversation request, by turn;
89  // the session's reported model can differ from the one a request went to.
90  on('turn.step', async function* ($, e, next) {
91    const result = yield* next(e)
92    if (!e.agentId && result.usage?.model) {
93      const models = answeredBy.get(e.turnId) ?? new Set()
94      models.add(result.usage.model)
95      answeredBy.set(e.turnId, models)
96    }
97    return result
98  })
99
100  // Each main-conversation turn: when it ended, the session's model, the
101  // models that answered, how long it took.
102  on('turn.complete', async ($, e, next) => {
103    if (e.agentId) return next(e)
104    const answered = [...(answeredBy.get(e.turnId) ?? [])]
105    answeredBy.delete(e.turnId)
106    const path = await recordPath($, await $.session.id())
107    if (!(await $.fs.exists(path))) return next(e)
108    const record = JSON.parse(await $.fs.read(path))
109    record.turns.push({ at: new Date().toISOString(), model: await $.session.model(), answeredBy: answered, durationMs: e.durationMs, isAborted: e.isAborted ?? false })
110    await $.fs.write(path, JSON.stringify(record))
111    return next(e)
112  })
113
114  // The session ends, or /clear or /resume starts another: close the record.
115  on('session.end', async ($, e, next) => {
116    const path = await recordPath($, e.sessionId)
117    if (await $.fs.exists(path)) {
118      const record = JSON.parse(await $.fs.read(path))
119      const usage = await $.session.usage()
120      record.endedAt = new Date().toISOString()
121      record.endReason = e.reason
122      record.costUsd = usage.cost?.usd ?? null
123      record.contextTokens = usage.context?.tokens ?? null
124      await $.fs.write(path, JSON.stringify(record))
125    }
126    return next(e)
127  })
128}
129