model-effort

1import { atom, read, update } from 'claude-code'
2import type { Color, Engine, Register } from 'claude-code'
3
4const model = atom({ plugin: 'model-effort', key: 'model' } as const, '')
5const effort = atom({ plugin: 'model-effort', key: 'effort' } as const, '')
6
7// 'claude-sonnet-5-5' -> 'sonnet 5.5'; 'claude-3-5-sonnet-20241022' -> 'sonnet 3.5';
8// 'opus[1m]' -> 'opus 1m'. An id with no family word is shown as it came.
9export function modelName(raw: string): string {
10  const context = /\[(\w+)\]$/.exec(raw)?.[1]
11  const parts = raw
12    .replace(/\[\w+\]$/, '')
13    .replace(/^claude-/, '')
14    .replace(/-\d{8}$/, '')
15    .split('-')
16    .filter(part => part !== '')
17  const family = parts.find(part => /^[a-z]/i.test(part))
18  if (family === undefined) return raw
19  const version = parts.filter(part => /^\d+$/.test(part)).join('.')
20  return [family, version, context].filter(part => part !== undefined && part !== '').join(' ')
21}
22
23// Green for the lighter models and efforts, red for the heavier ones; anything else
24// (an alias with no family, a new effort name) keeps the footer's dim look.
25const light = 'success'
26const heavy = 'error'
27
28export function modelColor(rawModel: string): Color | undefined {
29  const family = modelName(rawModel).split(' ')[0]
30  if (family === 'haiku' || family === 'sonnet') return light
31  if (family === 'opus' || family === 'fable') return heavy
32  return undefined
33}
34
35export function effortColor(rawEffort: string): Color | undefined {
36  if (rawEffort === 'low' || rawEffort === 'medium' || rawEffort === 'high') return light
37  if (rawEffort === 'xhigh' || rawEffort === 'max') return heavy
38  return undefined
39}
40
41// The model is the session's (what /model shows); a failed read keeps the last one.
42async function sync($: Engine): Promise<void> {
43  const current = await $.session.model().catch(() => '')
44  if (current !== '') await update($, model, () => current).catch(() => {})
45}
46
47export const register: Register = on => {
48  // Read the model at the moments it can change:
49  // the session start, a /model run, a /config row, and each prompt as a net.
50  on('session.start', async ($, e, next) => {
51    await sync($)
52    return next(e)
53  })
54
55  // /model sets the model without a prompt or a request: read it once the command is done.
56  // A hook here gates the command, so a failed read must not break it.
57  on('command.run', { command: 'model' }, async ($, e, next) => {
58    const result = await next(e)
59    await sync($)
60    return result
61  })
62
63  // The /config Model row, when it is the one that changed.
64  on('config.set', { key: 'model' }, async ($, e, next) => {
65    const result = await next(e)
66    await sync($)
67    return result
68  })
69
70  // A hook on the prompt gates it: a failed read must not hold the prompt back.
71  on('prompt.submit', async ($, e, next) => {
72    await sync($)
73    return next(e)
74  })
75
76  // The effort is only known on a request: note what the main loop's step asks for.
77  on('turn.step', async function* ($, e, next) {
78    if (e.agentId === undefined) {
79      const used = e.effort === undefined ? '' : String(e.effort)
80      await update($, model, () => e.model)
81      await update($, effort, () => used)
82    }
83    return yield* next(e)
84  })
85
86  // The labels at the right of the prompt footer: one more, so it takes no extra row.
87  // The engine draws them as one dim string, so the colored label needs a tree of its own.
88  on('ui.render', { component: 'SessionMode' }, async ($, e, next) => {
89    const rawModel = await read($, model)
90    if (rawModel === '') return next(e)
91    const rawEffort = await read($, effort)
92    const { Text } = $.ui.resolve(e)
93    const tint = (color: Color | undefined) => (color === undefined ? { dimColor: true } : { color })
94    const others = e.props.modes.join(' & ')
95    return (
96      <Text>
97        {others === '' ? null : <Text dimColor>{`${others} & `}</Text>}
98        <Text {...tint(modelColor(rawModel))}>{modelName(rawModel)}</Text>
99        {rawEffort === '' ? null : <Text dimColor>{' · '}</Text>}
100        {rawEffort === '' ? null : <Text {...tint(effortColor(rawEffort))}>{rawEffort}</Text>}
101      </Text>
102    )
103  })
104}
105

1declare module 'claude-code' {
2  interface PluginState {
3    'model-effort': { model: string; effort: string }
4  }
5}
6