agent-model-badge

1import type { EngineInterface, Register, RenderElement } from 'claude-code'
2
3type Badge = {
4  toolUseId: string
5  label: string
6  model: string
7  effort?: string
8}
9
10const SHORT: Record<string, string> = {
11  'claude-fable-5-1': 'fable',
12  'claude-opus-5-5': 'opus',
13  'claude-sonnet-5-5': 'sonnet',
14  'claude-haiku-4-5-20251001': 'haiku',
15}
16
17// agentId → badge; module state resets on reload, which is fine for a live view.
18const byAgent = new Map<string, Badge>()
19
20// The family name, so a new version of a model (claude-haiku-5-5) still reads as the haiku that was asked for.
21const shortModel = (id: string): string =>
22  SHORT[id] ?? /^claude-([a-z]+)/.exec(id)?.[1] ?? id
23
24// One colour per effort level, from the engine's theme keys.
25const EFFORT_COLOR: Record<string, 'success' | 'warning' | 'claude' | 'error' | 'permission'> = {
26  low: 'success',
27  medium: 'warning',
28  high: 'claude',
29  xhigh: 'error',
30  max: 'permission',
31}
32
33// A wide middle dot keeps the label clear of the description: "Count lines  ·  sonnet  ·  low". The engine folds
34// any run of whitespace in a task's description, so each side pairs a space with a braille blank (U+2800),
35// which draws as an empty cell but is not whitespace.
36const SEP = ' \u2800·\u2800 '
37const badgeText = (b: Badge): string =>
38  `${shortModel(b.model)}${b.effort ? `${SEP}${b.effort}` : ''}`
39
40// The effort each model family last ran at here, for a spawn whose call names none.
41const effortByFamily = new Map<string, string>()
42
43function refreshNotice($: EngineInterface, toolUseId: string): void {
44  if (!openCalls.has(toolUseId)) return
45  const rows = [...byAgent.values()].filter(b => b.toolUseId === toolUseId)
46  if (rows.length === 0) return
47  const [first] = rows
48  let text: string | undefined
49  if (first && rows.length === 1) text = `model: ${badgeText(first)}`
50  else if (rows.length > 1) text = rows.map(b => `${b.label}: ${badgeText(b)}`).join('  |  ')
51  try {
52    $.ui.notice(toolUseId, text)
53  } catch {
54    // the call may already be closed (a resumed workflow run); nothing to draw under
55  }
56}
57
58// What each Agent call asked for, by tool_use_id, and which call each labelled description came from: the
59// engine lists a running agent by its description, which is how a subagent's first request finds its call.
60const requested = new Map<string, { model?: string }>()
61const callByDescription = new Map<string, string>()
62// Agent ids already looked up without a match (a workflow's agents), so each is listed once.
63const unmatched = new Set<string>()
64// Agent calls still running. A notice can only sit under an open call: a background agent's call returns at
65// once, and a notice sent after that is dropped with a line in the transcript.
66const openCalls = new Set<string>()
67
68// A subagent's first request: find the Agent call that started it and give its row a badge.
69async function adopt($: EngineInterface, agentId: string, model: string, effort: string | undefined): Promise<void> {
70  if (unmatched.has(agentId)) return
71  const info = (await $.agent.list()).find(a => a.id === agentId)
72  const toolUseId = info ? callByDescription.get(info.description) : undefined
73  const ask = toolUseId ? requested.get(toolUseId) : undefined
74  if (!info || !toolUseId || !ask) {
75    unmatched.add(agentId)
76    return
77  }
78  const label = info.name ?? info.description.split(SEP)[0] ?? info.type
79  byAgent.set(agentId, {
80    toolUseId,
81    label: label.length > 28 ? `${label.slice(0, 27)}…` : label,
82    model,
83    effort,
84  })
85  if (ask.model && shortModel(model) !== ask.model) {
86    $.ui.toast(`${label}: asked for ${ask.model}, running on ${shortModel(model)}`)
87  }
88  refreshNotice($, toolUseId)
89  $.ui.invalidate('ui.render')
90}
91
92export const register: Register = on => {
93  // The tasks list above the prompt draws the call's description, so the label rides in it: the model the
94  // call names, else the session's own, and the effort the call names, else the one that model last ran at.
95  on('tool.call', { tool: 'Agent' }, async ($, e, next) => {
96    if (!e.tool_use_id || !e.description) return next(e)
97    const asked = e.model ?? shortModel(await $.session.model())
98    const effort = e.effort ?? effortByFamily.get(shortModel(asked))
99    const input = { ...e, description: `${e.description}${SEP}${asked}${effort ? `${SEP}${effort}` : ''}` }
100    requested.set(e.tool_use_id, { model: e.model })
101    callByDescription.set(input.description, e.tool_use_id)
102    openCalls.add(e.tool_use_id)
103    try {
104      return await next(input)
105    } finally {
106      openCalls.delete(e.tool_use_id)
107    }
108  })
109
110  on('turn.step', async function* ($, e, next) {
111    if (e.effort !== undefined) effortByFamily.set(shortModel(e.model), String(e.effort))
112    const effort = e.effort === undefined ? undefined : String(e.effort)
113    if (e.agentId) {
114      const badge = byAgent.get(e.agentId)
115      if (!badge) {
116        try {
117          await adopt($, e.agentId, e.model, effort)
118        } catch {
119          // the agent list is a lookup; the request goes on regardless
120        }
121      } else if (badge.model !== e.model || badge.effort !== effort) {
122        badge.model = e.model
123        badge.effort = effort
124        refreshNotice($, badge.toolUseId)
125        $.ui.invalidate('ui.render')
126      }
127    }
128    return yield* next(e)
129  })
130
131  // The Agent tool's transcript row: the engine's own drawing, then the badge beside it,
132  // the model in white and the effort in its level's colour.
133  on('ui.render', { component: 'ToolUse' }, async ($, e, next) => {
134    if (e.component !== 'ToolUse' || e.props.tool !== 'Agent') return next(e)
135    const badge = [...byAgent.values()].find(b => b.toolUseId === e.props.tool_use_id)
136    if (!badge) return next(e)
137    const { Box, Text } = $.ui.resolve(e)
138    const theirs = await next(e)
139    return h(
140      Box,
141      { flexDirection: 'row', columnGap: 1 },
142      theirs ?? '',
143      h(Text, { bold: true, color: 'text' }, shortModel(badge.model)),
144      badge.effort ? h(Text, { bold: true, color: EFFORT_COLOR[badge.effort] ?? 'text' }, badge.effort) : null,
145    ) as RenderElement
146  })
147}
148