model-pick

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Combo, Current } from '../types'
5import { DEFAULT_MODELS, allCombos, comboOfKey, keyOf, labelOf, rank } from './fuzzy'
6
7const PANE = 'model-pick'
8const MAX_RECENTS = 5
9const MAX_MATCHES = 12
10const HOTKEYS = ['1', '2', '3', '4', '5', '6', '7', '8', '9']
11const ALIAS_RE = /^[a-z0-9][a-z0-9_.-]{0,15}$/
12
13const query = atom({ plugin: 'model-pick', key: 'query' } as const, '')
14const assign = atom({ plugin: 'model-pick', key: 'assign' } as const, null)
15const notice = atom({ plugin: 'model-pick', key: 'notice' } as const, null)
16const favorites = atom({ plugin: 'model-pick', key: 'favorites' } as const, [])
17const recents = atom({ plugin: 'model-pick', key: 'recents' } as const, [])
18const aliases = atom({ plugin: 'model-pick', key: 'aliases' } as const, {})
19const current = atom({ plugin: 'model-pick', key: 'current' } as const, null)
20
21// The combos offered: the built-in aliases, plus any model ids the person's
22// settings name (`availableModels`, `model`). Rebuilt at session start.
23let combos: Combo[] = allCombos(DEFAULT_MODELS)
24let byKey = new Map<string, Combo>(combos.map(c => [keyOf(c), c]))
25
26const setModels = (models: readonly string[]) => {
27  combos = allCombos([...DEFAULT_MODELS, ...models])
28  byKey = new Map(combos.map(c => [keyOf(c), c]))
29}
30
31/** Model names the settings name: the `availableModels` allowlist and the default `model`. */
32const modelsInSettings = (settings: Readonly<Record<string, unknown>>): string[] => {
33  const out: string[] = []
34  const listed = settings['availableModels']
35  if (Array.isArray(listed)) out.push(...listed.filter(isString))
36  const chosen = settings['model']
37  if (isString(chosen)) out.push(chosen)
38  return out.map(m => m.trim()).filter(m => m !== '' && !m.includes(' '))
39}
40
41const loadModels = async ($: EngineInterface) => {
42  try {
43    setModels(modelsInSettings(await $.settings.read()))
44  } catch {
45    setModels([])
46  }
47}
48
49const isString = (v: unknown): v is string => typeof v === 'string'
50
51const isCombo = (v: unknown): v is Combo =>
52  typeof v === 'object' &&
53  v !== null &&
54  isString((v as Combo).model) &&
55  isString((v as Combo).effort)
56
57const asAliases = (v: unknown): Record<string, string> => {
58  if (typeof v !== 'object' || v === null) return {}
59  const out: Record<string, string> = {}
60  for (const [alias, key] of Object.entries(v as Record<string, unknown>)) {
61    if (isString(key)) out[alias] = key
62  }
63  return out
64}
65
66const resolveKey = (key: string): Combo | null => byKey.get(key) ?? comboOfKey(key)
67
68const aliasFor = (alis: Record<string, string>, key: string): string | undefined =>
69  Object.keys(alis).find(alias => alis[alias] === key)
70
71const boostOf =
72  (favs: string[], recs: Combo[]) =>
73  (key: string): number =>
74    (favs.includes(key) ? 0.5 : 0) + (recs.some(c => keyOf(c) === key) ? 0.25 : 0)
75
76/**
77 * A trailing row number: "3" picks row 3 of the unfiltered list, "op hi 2"
78 * row 2 of the matches for "op hi". `index` is 0-based; null when none.
79 */
80const splitRowNumber = (text: string): { base: string; index: number | null } => {
81  const m = /^(?:(.*\S)\s+)?([1-9])$/.exec(text.trim())
82  if (m === null) return { base: text, index: null }
83  return { base: m[1] ?? '', index: Number(m[2]) - 1 }
84}
85
86const loadStore = async ($: EngineInterface) => {
87  const fav = await $.store.get('favorites')
88  const rec = await $.store.get('recents')
89  const ali = await $.store.get('aliases')
90  await update($, favorites, () => (Array.isArray(fav) ? fav.filter(isString) : []))
91  await update($, recents, () => (Array.isArray(rec) ? rec.filter(isCombo).slice(0, MAX_RECENTS) : []))
92  await update($, aliases, () => asAliases(ali))
93}
94
95const say = ($: EngineInterface, text: string | null) => update($, notice, () => text)
96
97const pickBy = async ($: EngineInterface, text: string): Promise<Combo | null> => {
98  const [favs, recs, alis] = await Promise.all([read($, favorites), read($, recents), read($, aliases)])
99  const { base, index } = splitRowNumber(text)
100  if (index !== null) {
101    const rows = sectionsFor(base, favs, recs, alis, 999).flatMap(s => s.rows)
102    return rows[index] ?? null
103  }
104  const q = text.trim().toLowerCase()
105  if (q === '') {
106    const first = favs[0] !== undefined ? resolveKey(favs[0]) : null
107    return first ?? recs[0] ?? combos[0] ?? null
108  }
109  const aliased = alis[q]
110  if (aliased !== undefined) {
111    const found = resolveKey(aliased)
112    if (found !== null) return found
113  }
114  const ranked = rank(q, combos, key => aliasFor(alis, key), boostOf(favs, recs))
115  return ranked[0]?.combo ?? null
116}
117
118const pushRecent = async ($: EngineInterface, combo: Combo) => {
119  const list = await update($, recents, old =>
120    [combo, ...old.filter(c => keyOf(c) !== keyOf(combo))].slice(0, MAX_RECENTS),
121  )
122  await $.store.set('recents', list)
123}
124
125/** Whether `picked` (an alias or id) names the model the session reports as `current`. */
126const sameModel = (picked: string, current: string): boolean => {
127  const p = picked.toLowerCase().replace(/\[1m\]$/, '')
128  const c = current.toLowerCase()
129  return p === c || c.includes(p) || p.includes(c)
130}
131
132/**
133 * Whether the session's model became `picked`: at once when it already is, or
134 * when it changes from `before` within 30 s (the confirm dialog's answer).
135 */
136const modelSettled = async ($: EngineInterface, before: string, picked: string): Promise<boolean> => {
137  const deadline = (await $.clock.now()) + 30_000
138  for (;;) {
139    const now = await $.session.model()
140    if (now !== before || sameModel(picked, now)) return true
141    if ((await $.clock.now()) >= deadline) return false
142    await $.clock.sleep(500)
143  }
144}
145
146const applyCombo = ($: EngineInterface, combo: Combo) => {
147  // Detached: a command may not run inside the frame of the hook that asks.
148  $.clock.after(0, async () => {
149    try {
150      const before = await $.session.model()
151      const model = await $.command.run({ command: 'model', args: combo.model })
152      const text = model.text ?? ''
153      // "/model" prints "Set model to" or "Kept model as" when it answers at once;
154      // through its confirm dialog it prints nothing, so watch the session's model.
155      if (text.startsWith('Kept model as')) {
156        $.ui.toast(`${labelOf(combo)} not applied: ${text}`)
157        return
158      }
159      if (!text.startsWith('Set model to') && !(await modelSettled($, before, combo.model))) {
160        $.ui.toast(`${labelOf(combo)} not applied: the model switch was not confirmed`)
161        return
162      }
163      // "/effort" prints its result to the transcript, not to `text`: resolved is set.
164      await $.command.run({ command: 'effort', args: combo.effort })
165      await pushRecent($, combo)
166      $.ui.toast(`now ${labelOf(combo)}`)
167    } catch (err) {
168      $.ui.toast(`${err instanceof Error ? err.message : String(err)}`)
169    }
170  })
171}
172
173const choose = async ($: EngineInterface, combo: Combo, close = true) => {
174  await update($, query, () => '')
175  await update($, assign, () => null)
176  await say($, null)
177  if (close) await $.ui.close({ id: PANE })
178  applyCombo($, combo)
179}
180
181const toggleFav = async ($: EngineInterface, combo: Combo) => {
182  const key = keyOf(combo)
183  let isFav = false
184  const list = await update($, favorites, old => {
185    isFav = !old.includes(key)
186    return isFav ? [...old, key] : old.filter(k => k !== key)
187  })
188  await $.store.set('favorites', list)
189  await say($, `${isFav ? '★ favorited' : '☆ unfavorited'} ${labelOf(combo)}`)
190}
191
192const setAlias = async ($: EngineInterface, combo: Combo, text: string) => {
193  const key = keyOf(combo)
194  const alias = text.trim().toLowerCase()
195  if (alias !== '' && !ALIAS_RE.test(alias)) {
196    await say($, `alias "${alias}" is refused: 1-16 of a-z 0-9 _ . -`)
197    return
198  }
199  const map = await update($, aliases, old => {
200    const next: Record<string, string> = {}
201    for (const [a, k] of Object.entries(old)) if (k !== key && a !== alias) next[a] = k
202    if (alias !== '') next[alias] = key
203    return next
204  })
205  await $.store.set('aliases', map)
206  await update($, assign, () => null)
207  await update($, query, () => '')
208  await say($, alias === '' ? `alias cleared for ${labelOf(combo)}` : `/pick ${alias} → ${labelOf(combo)}`)
209}
210
211const startAssign = async ($: EngineInterface, combo: Combo) => {
212  await update($, assign, () => keyOf(combo))
213  await update($, query, () => '')
214  await say($, `type an alias for ${labelOf(combo)}, Enter saves, empty clears`)
215}
216
217const submit = async ($: EngineInterface, text: string) => {
218  const pending = await read($, assign)
219  if (pending !== null) {
220    const combo = resolveKey(pending)
221    if (combo !== null) await setAlias($, combo, text)
222    else await update($, assign, () => null)
223    return
224  }
225  const t = text.trim()
226  if (t.startsWith('*')) {
227    const combo = await pickBy($, t.slice(1))
228    if (combo === null) await say($, `no match for "${t.slice(1).trim()}"`)
229    else await toggleFav($, combo)
230    await update($, query, () => '')
231    return
232  }
233  if (t.startsWith('=')) {
234    const [alias = '', ...rest] = t.slice(1).trim().split(/\s+/)
235    const combo = await pickBy($, rest.join(' '))
236    if (alias === '') await say($, 'usage: =alias model effort')
237    else if (combo === null) await say($, `no match for "${rest.join(' ')}"`)
238    else await setAlias($, combo, alias)
239    return
240  }
241  const combo = await pickBy($, t)
242  if (combo === null) {
243    await say($, `no match for "${t}"`)
244    return
245  }
246  await choose($, combo)
247}
248
249type Section = { title: string; rows: Combo[] }
250
251const sectionsFor = (
252  q: string,
253  favs: string[],
254  recs: Combo[],
255  alis: Record<string, string>,
256  room: number,
257): Section[] => {
258  if (q.trim() !== '') {
259    const ranked = rank(q, combos, key => aliasFor(alis, key), boostOf(favs, recs))
260    return [
261      {
262        title: ranked.length === 0 ? 'no matches' : 'matches',
263        rows: ranked.slice(0, MAX_MATCHES).map(r => r.combo),
264      },
265    ]
266  }
267  const favRows = favs.map(resolveKey).filter((c): c is Combo => c !== null)
268  const out: Section[] = []
269  if (favRows.length > 0) out.push({ title: 'favorites', rows: favRows })
270  if (recs.length > 0) out.push({ title: 'recent', rows: recs.slice(0, MAX_RECENTS) })
271  const used = favRows.length + Math.min(recs.length, MAX_RECENTS)
272  out.push({ title: 'all (type to filter)', rows: combos.slice(0, Math.max(3, room - used)) })
273  return out
274}
275
276const HINT = 'Enter applies ▸ · a row number picks that row · *q favorites · =alias q binds · Esc closes'
277
278export const register: Register = on => {
279  on('session.start', async ($, e, next) => {
280    await $.command.register({
281      name: 'pick',
282      description: 'Fuzzy-pick model + effort; favorites, aliases, 5 recents',
283      argumentHint: '[query or alias]',
284      immediate: true,
285    })
286    await loadModels($)
287    await loadStore($)
288    return next(e)
289  })
290
291  // /clear, /resume and /branch reset $.state without a session.start: load again.
292  on('classic.SessionStart', { source: ['clear', 'resume', 'fork'] }, async ($, e, next) => {
293    await loadModels($)
294    await loadStore($)
295    return next(e)
296  })
297
298  on('command.run', { command: 'pick' }, async ($, e) => {
299    const args = e.args.trim()
300    if (args === '') {
301      await update($, query, () => '')
302      await update($, assign, () => null)
303      await say($, null)
304      const opened = await $.ui.open({
305        id: PANE,
306        title: 'model · effort',
307        focus: true,
308        closeOnEscape: true,
309        rows: 18,
310      })
311      return { text: opened.isPlaced ? 'picker opened' : opened.reason }
312    }
313    const combo = await pickBy($, args)
314    if (combo === null) return { text: `no match for "${args}"` }
315    await choose($, combo, false)
316    return { text: labelOf(combo) }
317  })
318
319  on('turn.step', async function* ($, e, next) {
320    if (e.agentId === undefined) {
321      const now: Current = { model: e.model, effort: e.effort === undefined ? null : String(e.effort) }
322      await update($, current, () => now)
323    }
324    return yield* next(e)
325  })
326
327  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
328    const [q, pending, note, favs, recs, alis, cur] = await Promise.all([
329      read($, query),
330      read($, assign),
331      read($, notice),
332      read($, favorites),
333      read($, recents),
334      read($, aliases),
335      read($, current),
336    ])
337    const pendingCombo = pending === null ? null : resolveKey(pending)
338    const room = Math.max(6, (e.props.scroll?.bodyRows ?? 18) - 5)
339    const { base, index } = pendingCombo === null ? splitRowNumber(q) : { base: q, index: null }
340    const sections = sectionsFor(base, favs, recs, alis, room)
341    const topIndex = index ?? 0
342    const nowLine =
343      cur === null ? 'now: unknown until the next request' : `now: ${cur.model} · ${cur.effort ?? 'default'}`
344    let hot = 0
345
346    const { Box, Text, Button } = $.ui.resolve(e)
347    const rows = sections.map(section => (
348      <Box flexDirection="column">
349        <Text bold>{section.title}</Text>
350        {section.rows.map(combo => {
351          const key = keyOf(combo)
352          const isFav = favs.includes(key)
353          const alias = aliasFor(alis, key)
354          const isTop = hot === topIndex
355          const hotkey = hot < HOTKEYS.length ? HOTKEYS[hot++] : undefined
356          return (
357            <Box flexDirection="row" gap={1}>
358              <Text color="suggestion" bold>{isTop ? '▸' : ' '}</Text>
359              <Button
360                key={`apply:${key}`}
361                plain
362                {...(hotkey === undefined ? {} : { hotkey })}
363                label={labelOf(combo)}
364                onPress={() => void choose($, combo)}
365              />
366              <Button
367                key={`fav:${key}`}
368                plain
369                dimColor={!isFav}
370                label={isFav ? '★' : '☆'}
371                onPress={() => void toggleFav($, combo)}
372              />
373              <Button
374                key={`alias:${key}`}
375                plain
376                dimColor={alias === undefined}
377                label={alias === undefined ? '=…' : `=${alias}`}
378                onPress={() => void startAssign($, combo)}
379              />
380            </Box>
381          )
382        })}
383      </Box>
384    ))
385
386    if (e.surface === 'mobile') {
387      // No Input on this surface yet: the rows and their buttons still work.
388      return (
389        <Box flexDirection="column">
390          <Text dimColor>{nowLine}</Text>
391          {rows}
392          <Text dimColor>{note ?? 'tap a row to apply'}</Text>
393        </Box>
394      )
395    }
396
397    const { Input } = $.ui.resolve(e)
398    return (
399      <Box flexDirection="column">
400        <Input
401          key="q"
402          label={pendingCombo === null ? 'find' : `alias for ${labelOf(pendingCombo)}`}
403          placeholder={pendingCombo === null ? 'model effort, e.g. "op hi"' : 'a-z 0-9, empty clears'}
404          value={q}
405          submitLabel={pendingCombo === null ? 'apply' : 'save'}
406          autoFocus
407          onInput={value => void update($, query, () => value)}
408          onSubmit={value => void submit($, value)}
409        />
410        <Text dimColor>{nowLine}</Text>
411        {rows}
412        <Text dimColor>{note ?? HINT}</Text>
413      </Box>
414    )
415  })
416}
417

1import type { Combo } from '../types'
2
3export const DEFAULT_MODELS = [
4  'default',
5  'opus',
6  'sonnet',
7  'haiku',
8  'fable',
9  'opusplan',
10  'opus[1m]',
11  'sonnet[1m]',
12  'fable[1m]',
13  'opusplan[1m]',
14]
15
16export const EFFORTS = ['low', 'medium', 'high', 'xhigh', 'max', 'auto']
17
18export const keyOf = (c: Combo): string => `${c.model}/${c.effort}`
19
20export const labelOf = (c: Combo): string => `${c.model} · ${c.effort}`
21
22export const comboOfKey = (key: string): Combo | null => {
23  const at = key.lastIndexOf('/')
24  if (at <= 0) return null
25  return { model: key.slice(0, at), effort: key.slice(at + 1) }
26}
27
28export const allCombos = (models: readonly string[]): Combo[] => {
29  const seen = new Set<string>()
30  const out: Combo[] = []
31  for (const model of models) {
32    const m = model.trim()
33    if (m === '' || seen.has(m)) continue
34    seen.add(m)
35    for (const effort of EFFORTS) out.push({ model: m, effort })
36  }
37  return out
38}
39
40const isWordChar = (ch: string | undefined): boolean =>
41  ch !== undefined && /[a-z0-9]/i.test(ch)
42
43/**
44 * Subsequence score of `needle` in `hay`, higher is better; null when
45 * `needle` is not a subsequence. Adjacent hits and word starts score more.
46 */
47export const fuzzyScore = (needle: string, hay: string): number | null => {
48  if (needle === '') return 0
49  let score = 0
50  let from = 0
51  let prev = -2
52  for (const ch of needle) {
53    const at = hay.indexOf(ch, from)
54    if (at < 0) return null
55    if (at === prev + 1) score += 3
56    else if (at === 0 || !isWordChar(hay[at - 1])) score += 2
57    else score += 1
58    prev = at
59    from = at + 1
60  }
61  if (hay.startsWith(needle)) score += 2
62  return score - (prev + 1 - needle.length) * 0.01
63}
64
65export type Ranked = { combo: Combo; score: number }
66
67/**
68 * Every combo each whitespace-separated token of `query` fuzzy-matches, best
69 * first. A token matches the model, the effort, or the combo's alias.
70 */
71export const rank = (
72  query: string,
73  combos: readonly Combo[],
74  aliasOf: (key: string) => string | undefined,
75  boost: (key: string) => number,
76): Ranked[] => {
77  const tokens = query.toLowerCase().split(/\s+/).filter(t => t !== '')
78  const out: Ranked[] = []
79  for (const combo of combos) {
80    const key = keyOf(combo)
81    const alias = aliasOf(key)
82    const hay = `${combo.model} ${combo.effort}`.toLowerCase()
83    let total = 0
84    let ok = true
85    for (const token of tokens) {
86      const best = Math.max(
87        fuzzyScore(token, hay) ?? -Infinity,
88        alias !== undefined ? (fuzzyScore(token, alias) ?? -Infinity) + 1 : -Infinity,
89      )
90      if (best === -Infinity) {
91        ok = false
92        break
93      }
94      total += best
95    }
96    if (ok) out.push({ combo, score: total + boost(key) })
97  }
98  out.sort((a, b) => b.score - a.score)
99  return out
100}
101

1export type Combo = { model: string; effort: string }
2
3export type Current = { model: string; effort: string | null } | null
4
5declare module 'claude-code' {
6  interface PluginState {
7    'model-pick': {
8      query: string
9      assign: string | null
10      notice: string | null
11      favorites: string[]
12      recents: Combo[]
13      aliases: Record<string, string>
14      current: Current
15    }
16  }
17}
18