agent-model-badge

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { AgentBadge } from '../types'
5import { modelColor, prettyModel, resolveAlias } from './model'
6
7// Keyed by both the Agent call's tool_use_id (for its transcript row) and the
8// spawned agentId (for the viewer opened from the tasks list).
9const agents = atom({ plugin: 'agent-model-badge', key: 'agents' } as const, {})
10
11function stringField(value: unknown, field: string): string | undefined {
12  if (value && typeof value === 'object' && field in value) {
13    const v = (value as Record<string, unknown>)[field]
14    if (typeof v === 'string' && v) return v
15  }
16  return undefined
17}
18
19async function lookupAgent($: EngineInterface, agentId: string): Promise<AgentBadge | undefined> {
20  const known = (await read($, agents))[agentId]
21  if (known) return known
22  const row = (await $.agent.list()).find(a => a.id === agentId)
23  const model = row && (stringField(row, 'resolvedModel') ?? stringField(row, 'model'))
24  return row && model ? { model, type: row.type } : undefined
25}
26
27export const register: Register = on => {
28  // The tasks list under the prompt draws each agent's description, so the
29  // model goes there. It must be known before the spawn: the call's model, a
30  // fork's parent, the model this type resolved to last time, else the parent.
31  on('agent.spawn', async ($, e, next) => {
32    const learnedKey = `typeModel:${e.subagentType}`
33    const learned = e.model || e.fork ? undefined : await $.store.get(learnedKey)
34    const expected = e.fork
35      ? e.parentModel
36      : e.model
37        ? resolveAlias(e.model, e.parentModel)
38        : typeof learned === 'string'
39          ? learned
40          : e.parentModel
41    const label = ` · ${prettyModel(expected)}`
42    const description = e.description.endsWith(label) ? e.description : e.description + label
43
44    const result = await next({ ...e, description })
45    if ('model' in result && result.model) {
46      const badge: AgentBadge = { model: result.model, type: e.subagentType }
47      await update($, agents, all => ({
48        ...all,
49        ...(e.tool_use_id ? { [e.tool_use_id]: badge } : {}),
50        ...(result.agentId ? { [result.agentId]: badge } : {}),
51      }))
52      if (!e.model && !e.fork) await $.store.set(learnedKey, result.model)
53    }
54    return result
55  }).catch(($, e, next) => next(e))
56
57  on('ui.render', { component: 'ToolUse' }, async ($, e, next) => {
58    if (e.props.tool !== 'Agent') return next(e)
59
60    const known = (await read($, agents))[e.props.tool_use_id]
61    const model =
62      known?.model ??
63      stringField(e.props.output, 'resolvedModel') ??
64      stringField(e.props.input, 'model')
65    if (!model) return next(e)
66
67    const type = known?.type ?? stringField(e.props.input, 'subagent_type') ?? 'general-purpose'
68    const row = await next(e)
69    const { Box, Text } = $.ui.resolve(e)
70
71    return (
72      <Box flexDirection="column">
73        {row}
74        <Box>
75          <Text dimColor>{'  ⎿  '}</Text>
76          <Text color={modelColor(model)} bold>
77            {prettyModel(model)}
78          </Text>
79          <Text dimColor>{` · ${type}`}</Text>
80        </Box>
81      </Box>
82    )
83  })
84
85  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
86    const agentId = e.props.view.agentId
87    if (!agentId || e.props.hasSurvey) return next(e)
88
89    const badge = await lookupAgent($, agentId)
90    if (!badge) return next(e)
91
92    const { Box, Text } = $.ui.resolve(e)
93    return (
94      <Box>
95        <Text dimColor>{'◆ Viewing '}</Text>
96        <Text>{badge.type}</Text>
97        <Text dimColor>{' on '}</Text>
98        <Text color={modelColor(badge.model)} bold>
99          {prettyModel(badge.model)}
100        </Text>
101      </Box>
102    )
103  })
104}
105

1const FAMILY_COLORS: Record<string, string> = {
2  opus: 'magenta',
3  sonnet: 'cyan',
4  haiku: 'green',
5  fable: 'yellow',
6}
7
8/** `claude-sonnet-5-5` → `Sonnet 5.5`, `claude-haiku-4-5-20251001[1m]` → `Haiku 4.5 1M`, `haiku` → `Haiku`. */
9export function prettyModel(id: string): string {
10  let rest = id.trim().replace(/^(us\.|eu\.|global\.)?anthropic\./, '')
11  let suffix = ''
12  const ctx = rest.match(/\[(\d+)([km])\]$/i)
13  if (ctx) {
14    suffix = ` ${ctx[1]}${ctx[2]!.toUpperCase()}`
15    rest = rest.slice(0, ctx.index)
16  }
17  rest = rest.replace(/^claude-/, '').replace(/-v\d+(:\d+)?$/, '').replace(/-\d{8}$/, '')
18  const m = rest.match(/^([a-z]+)((?:-\d{1,2})*)$/i)
19  if (!m) return id
20  const family = m[1]!.charAt(0).toUpperCase() + m[1]!.slice(1)
21  const version = m[2] ? ' ' + m[2].slice(1).split('-').join('.') : ''
22  return family + version + suffix
23}
24
25export function modelColor(id: string): string | undefined {
26  const family = Object.keys(FAMILY_COLORS).find(f => id.toLowerCase().includes(f))
27  return family ? FAMILY_COLORS[family] : undefined
28}
29
30// What a bare alias in an Agent call resolves to on this account today.
31const ALIAS_LATEST: Record<string, string> = {
32  haiku: 'claude-haiku-5-5',
33  sonnet: 'claude-sonnet-5-5',
34  opus: 'claude-opus-5-5',
35  fable: 'claude-fable-5-1',
36}
37
38/** An alias as the parent's own model when it is that family, else the latest of the family. */
39export function resolveAlias(model: string, parentModel: string): string {
40  const alias = model.toLowerCase()
41  if (!(alias in ALIAS_LATEST)) return model
42  return parentModel.toLowerCase().includes(alias) ? parentModel : ALIAS_LATEST[alias]!
43}
44

1export type AgentBadge = { model: string; type: string }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'agent-model-badge': { agents: Record<string, AgentBadge> }
6  }
7}
8