model-badge

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Effort } from '../types'
5
6const FAMILIES = ['opus', 'sonnet', 'haiku', 'fable'] as const
7const LEVELS = ['low', 'medium', 'high', 'xhigh', 'max']
8// Below these widths the badge drops the effort, then hides altogether.
9const FULL_COLUMNS = 60
10const NAME_COLUMNS = 40
11const POLL_MS = 2000
12
13const model = atom({ plugin: 'model-badge', key: 'model' } as const, null)
14const effort = atom({ plugin: 'model-badge', key: 'effort' } as const, null)
15
16const capitalize = (word: string) => word.charAt(0).toUpperCase() + word.slice(1)
17
18// `claude-opus-5-5`, `opus`, `Opus 5.5`, `claude-3-5-sonnet-20241022` and
19// `claude-haiku-4-5-20251001[1m]` all come out as a short display name.
20export function displayName(raw: string) {
21  const isLong = /\[1m\]/i.test(raw)
22  const id = raw.replace(/\[1m\]/gi, '').trim()
23  const lower = id.toLowerCase()
24  const family = FAMILIES.find(name => lower.includes(name))
25  if (family === undefined) return id
26
27  const after = lower.match(new RegExp(`${family}[-_ ]?(\\d+)(?:[-.](\\d{1,2}))?(?!\\d)`))
28  const before = lower.match(new RegExp(`(\\d+)[-.](\\d{1,2})[-_ ]${family}`))
29  const match = after?.[1] !== undefined && after[1].length < 4 ? after : before
30  const version = match ? [match[1], match[2]].filter(part => part !== undefined).join('.') : ''
31
32  return [capitalize(family), version, isLong ? '1M' : ''].filter(Boolean).join(' ')
33}
34
35export function effortLevel(value: unknown): Effort {
36  return typeof value === 'string' && LEVELS.includes(value) ? value : null
37}
38
39// The badge's text for the room the footer has; null when there is none.
40export function label(name: string | null, level: Effort, columns: number | undefined) {
41  if (name === null) return null
42  const width = columns ?? FULL_COLUMNS
43  if (width < NAME_COLUMNS) return null
44  if (level === null || width < FULL_COLUMNS) return name
45
46  return `${name} · ${level}`
47}
48
49// Re-reads the session's model; true when it moved since the last read.
50async function refresh($: EngineInterface) {
51  const name = displayName(await $.session.model())
52  const previous = await read($, model)
53  if (previous === name) return false
54  await update($, model, () => name)
55
56  return previous !== null
57}
58
59export const register: Register = on => {
60  let poll: { cancel: () => void } | undefined
61
62  on('session.start', async ($, e, next) => {
63    await refresh($)
64
65    // `/model`, `--model`, a settings change or a fallback: whatever moved it,
66    // the next tick redraws the badge and flashes the change top-right.
67    if (e.isInteractive) {
68      poll?.cancel()
69      poll = $.clock.every(POLL_MS, async () => {
70        if (!(await refresh($))) return
71        const text = label(await read($, model), await read($, effort), undefined)
72        if (text !== null) $.ui.toast(text, { timeoutMs: 2500 })
73      })
74    }
75
76    return next(e)
77  })
78
79  on('session.end', ($, e, next) => {
80    poll?.cancel()
81    poll = undefined
82
83    return next(e)
84  })
85
86  // `/model` and `/effort` redraw at once rather than on the next tick.
87  on('command.run', { command: 'model' }, async ($, e, next) => {
88    const result = await next(e)
89    await refresh($)
90
91    return result
92  })
93
94  on('command.run', { command: 'effort' }, async ($, e, next) => {
95    const result = await next(e)
96    const level = effortLevel(e.args.trim().toLowerCase())
97    if (level !== null) await update($, effort, () => level)
98
99    return result
100  })
101
102  // The main loop's requests carry the effort actually sent, after any
103  // downgrade for the model; subagents' steps are left alone.
104  on('turn.step', async function* ($, e, next) {
105    if (e.agentId === undefined) {
106      const level = effortLevel(e.effort)
107      if ((await read($, effort)) !== level) await update($, effort, () => level)
108    }
109
110    return yield* next(e)
111  })
112
113  // The footer's right-hand corner, ahead of whatever the plugins beneath
114  // (context-bar) or the engine draw there.
115  on('ui.render', { component: 'SessionMode' }, async ($, e, next) => {
116    const text = label(await read($, model), await read($, effort), e.viewport?.columns)
117    const rest = await next(e)
118    if (text === null) return rest
119
120    const { Box, Text } = $.ui.resolve(e)
121
122    return (
123      <Box flexDirection="row" gap={1}>
124        <Text dimColor inverse wrap="truncate">
125          {` ${text} `}
126        </Text>
127        {rest}
128      </Box>
129    )
130  })
131}
132
133

1export type Effort = string | null
2
3declare module 'claude-code' {
4  interface PluginState {
5    'model-badge': { model: string | null; effort: Effort }
6  }
7}
8