agent-usage

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { AgentUsage } from '../types'
5import { SEQUENCE, VESUVIUS } from './palette'
6
7const PANE = 'agent-usage'
8export const MAIN = 'main'
9
10const agents = atom({ plugin: 'agent-usage', key: 'agents' } as const, [])
11const expanded = atom({ plugin: 'agent-usage', key: 'expanded' } as const, null)
12const costUsd = atom({ plugin: 'agent-usage', key: 'costUsd' } as const, null)
13
14export function emptyUsage(id: string, label: string, type: string): AgentUsage {
15  return { id, label, type, model: '', status: 'running', steps: 0, tools: 0, input: 0, output: 0, cacheRead: 0, cacheWrite: 0, durationMs: null }
16}
17
18export function totalOf(a: AgentUsage): number {
19  return a.input + a.output + a.cacheRead + a.cacheWrite
20}
21
22export function formatTokens(n: number): string {
23  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
24  if (n >= 10_000) return `${Math.round(n / 1000)}k`
25  if (n >= 1000) return `${(n / 1000).toFixed(1)}k`
26
27  return String(n)
28}
29
30export function formatDuration(ms: number | null): string {
31  if (ms === null) return '—'
32  const s = Math.round(ms / 1000)
33
34  return s >= 60 ? `${Math.floor(s / 60)}m ${s % 60}s` : `${s}s`
35}
36
37export function shortModel(model: string): string {
38  return model.replace(/^claude-/, '').replace(/-\d{8}$/, '')
39}
40
41// Adds or changes one agent's row, creating it when it is new.
42export function upsert(
43  list: readonly AgentUsage[],
44  id: string,
45  change: (a: AgentUsage) => AgentUsage,
46  fallback: () => AgentUsage,
47): AgentUsage[] {
48  const found = list.find(a => a.id === id)
49
50  return found === undefined ? [...list, change(fallback())] : list.map(a => (a.id === id ? change(a) : a))
51}
52
53export function addUsage(a: AgentUsage, usage: ModelUsage & { model?: string }): AgentUsage {
54  return {
55    ...a,
56    model: usage.model ?? a.model,
57    steps: a.steps + 1,
58    input: a.input + usage.input_tokens,
59    output: a.output + usage.output_tokens,
60    cacheRead: a.cacheRead + usage.cache_read_input_tokens,
61    cacheWrite: a.cacheWrite + usage.cache_creation_input_tokens,
62  }
63}
64
65// The main conversation first, then the subagents, biggest first.
66export function sorted(list: readonly AgentUsage[]): AgentUsage[] {
67  return [...list].sort((a, b) => (a.id === MAIN ? -1 : b.id === MAIN ? 1 : totalOf(b) - totalOf(a)))
68}
69
70const STATUS_GLYPH: Record<AgentUsage['status'], string> = {
71  pending: '◌',
72  running: '●',
73  waiting: '◐',
74  idle: '○',
75  completed: '✓',
76  failed: '✗',
77  killed: '✗',
78}
79
80function statusColor(status: AgentUsage['status']): string {
81  if (status === 'running' || status === 'waiting') return VESUVIUS.magma
82  if (status === 'failed' || status === 'killed') return VESUVIUS.lava
83  if (status === 'completed') return VESUVIUS.vineyard
84
85  return VESUVIUS.ash
86}
87
88const mainRow = () => ({ ...emptyUsage(MAIN, 'Main conversation', 'main'), status: 'idle' as const })
89
90// Reads the agents' status and cost from the engine.
91async function refresh($: EngineInterface): Promise<void> {
92  const known = await $.agent.list().catch(() => [])
93  const usage = await $.session.usage().catch(() => undefined)
94  await update($, agents, list =>
95    list.map(a => {
96      const info = known.find(k => k.id === a.id)
97
98      return info === undefined ? a : { ...a, status: info.status, label: a.label || info.description, type: a.type || info.type }
99    }),
100  )
101  await update($, costUsd, () => usage?.cost?.usd ?? null)
102}
103
104// The status line while subagents run: how many, and what they consumed so far.
105async function showStatus($: EngineInterface): Promise<void> {
106  const list = await read($, agents)
107  const running = list.filter(a => a.id !== MAIN && (a.status === 'running' || a.status === 'waiting'))
108  if (running.length === 0) {
109    $.ui.status(undefined)
110
111    return
112  }
113  const tokens = running.reduce((sum, a) => sum + totalOf(a), 0)
114  $.ui.status(`▲ ${running.length} ${running.length === 1 ? 'agent' : 'agents'} running · ${formatTokens(tokens)} tokens`)
115}
116
117export const register: Register = on => {
118  on('session.start', async ($, e, next) => {
119    await $.command.register({
120      name: 'agent-usage',
121      description: 'Open the pane with the tokens each agent consumed',
122      immediate: true,
123    })
124
125    return next(e)
126  })
127
128  on('command.run', { command: 'agent-usage' }, async $ => {
129    await refresh($)
130    await $.ui.open({ id: PANE, title: 'Agent usage', focus: true, closeOnEscape: true })
131
132    return {}
133  })
134
135  // A subagent starts: its row gets the task's description and type.
136  on('agent.spawn', async ($, e, next) => {
137    const result = await next(e)
138    if (result.agentId !== undefined) {
139      const id = result.agentId
140      await update($, agents, list =>
141        upsert(
142          list,
143          id,
144          a => ({ ...a, label: e.description, type: e.subagentType, model: result.model, status: 'running' }),
145          () => emptyUsage(id, e.description, e.subagentType),
146        ),
147      )
148      await showStatus($)
149    }
150
151    return result
152  })
153
154  // Every model request, in the main loop or in a subagent's: its tokens go to that loop.
155  on('turn.step', async function* ($, e, next) {
156    const result = yield* next(e)
157    const usage = result.usage
158    if (usage !== null) {
159      const id = e.agentId ?? MAIN
160      await update($, agents, list =>
161        upsert(list, id, a => addUsage(a, usage), () => (id === MAIN ? mainRow() : emptyUsage(id, '', ''))),
162      )
163      await showStatus($)
164    }
165
166    return result
167  })
168
169  // Every tool call counts for the loop it runs in.
170  on('tool.call', async ($, e, next) => {
171    const id = e.agentId ?? MAIN
172    await update($, agents, list =>
173      upsert(list, id, a => ({ ...a, tools: a.tools + 1 }), () => (id === MAIN ? mainRow() : emptyUsage(id, '', ''))),
174    )
175
176    return next(e)
177  })
178
179  // A foreground subagent ends: its Agent call reports how long it ran.
180  on('tool.call', { tool: 'Agent' }, async ($, e, next) => {
181    const ran = await next(e)
182    const result = ran.result as { status?: string; agentId?: string; totalDurationMs?: number } | undefined
183    if (result?.status === 'completed' && result.agentId !== undefined) {
184      const id = result.agentId
185      const finished = (await read($, agents)).find(a => a.id === id)
186      await update($, agents, list =>
187        list.map(a => (a.id === id ? { ...a, status: 'completed', durationMs: result.totalDurationMs ?? a.durationMs } : a)),
188      )
189      if (finished !== undefined) {
190        $.ui.toast(`✓ ${finished.label || finished.type} · ${formatTokens(totalOf(finished))} tokens`)
191      }
192      await showStatus($)
193    }
194
195    return ran
196  })
197
198  on('turn.complete', async ($, e, next) => {
199    const result = await next(e)
200    await refresh($)
201    await showStatus($)
202
203    return result
204  })
205
206  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
207    const { Box, Text, Button } = $.ui.resolve(e)
208    const list = sorted(await read($, agents))
209    const open = await read($, expanded)
210    const cost = await read($, costUsd)
211
212    if (list.length === 0) {
213      return <Text color={VESUVIUS.ash}>No model requests yet.</Text>
214    }
215
216    const grand = list.reduce((sum, a) => sum + totalOf(a), 0)
217    const subagents = list.filter(a => a.id !== MAIN)
218    const barWidth = Math.max(6, Math.min(24, e.props.bodyColumns - 60))
219    const labelWidth = Math.max(12, Math.min(28, e.props.bodyColumns - barWidth - 34))
220
221    return (
222      <Box flexDirection="column" rowGap={1}>
223        <Text>
224          <Text color={VESUVIUS.lava}>▲</Text> <Text bold>{formatTokens(grand)}</Text> tokens ·{' '}
225          <Text color={VESUVIUS.tuff}>
226            {subagents.length} {subagents.length === 1 ? 'subagent' : 'subagents'}
227          </Text>
228          {cost !== null ? <Text color={VESUVIUS.ash}> · session cost ${cost.toFixed(2)}</Text> : ''}
229        </Text>
230        <Box flexDirection="column">
231          {list.map((a, i) => {
232            const total = totalOf(a)
233            const share = grand > 0 ? total / grand : 0
234            const cells = total > 0 ? Math.max(1, Math.round(share * barWidth)) : 0
235            const color = SEQUENCE[i % SEQUENCE.length] ?? VESUVIUS.lava
236            const name = (a.label || a.type || a.id).slice(0, labelWidth).padEnd(labelWidth)
237
238            return (
239              <Box flexDirection="column">
240                <Button key={`agent-${a.id}`} plain onPress={() => update($, expanded, x => (x === a.id ? null : a.id))}>
241                  <Text color={statusColor(a.status)}>{STATUS_GLYPH[a.status]}</Text> <Text bold={a.id === open}>{name}</Text>{' '}
242                  <Text color={VESUVIUS.magma}>{formatTokens(total).padStart(6)}</Text>{' '}
243                  <Text color={VESUVIUS.tuff}>{`${Math.round(share * 100)}%`.padStart(4)}</Text> <Text color={color}>{'█'.repeat(cells)}</Text>
244                  <Text color={VESUVIUS.ash}>{'░'.repeat(Math.max(0, barWidth - cells))}</Text>
245                </Button>
246                {a.id === open && (
247                  <Box flexDirection="column" paddingLeft={2}>
248                    <Text color={VESUVIUS.ash}>
249                      {a.type !== 'main' ? `${a.type} · ` : ''}
250                      {a.model ? shortModel(a.model) : 'model unknown'} · {a.status} · ran {formatDuration(a.durationMs)}
251                    </Text>
252                    <Text>
253                      <Text color={VESUVIUS.gulf}>in</Text> {formatTokens(a.input)} · <Text color={VESUVIUS.vineyard}>out</Text> {formatTokens(a.output)} ·{' '}
254                      <Text color={VESUVIUS.sulfur}>cache read</Text> {formatTokens(a.cacheRead)} · <Text color={VESUVIUS.tuff}>cache write</Text>{' '}
255                      {formatTokens(a.cacheWrite)}
256                    </Text>
257                    <Text color={VESUVIUS.ash}>
258                      {a.steps} {a.steps === 1 ? 'request' : 'requests'} · {a.tools} tool {a.tools === 1 ? 'call' : 'calls'}
259                    </Text>
260                  </Box>
261                )}
262              </Box>
263            )
264          })}
265        </Box>
266        <Box flexDirection="row" columnGap={2}>
267          <Text color={VESUVIUS.ash}>● running ◐ waiting ✓ done ✗ failed · press a row for details</Text>
268          <Button key="reset" plain label="reset" onPress={() => update($, agents, () => [])} />
269        </Box>
270      </Box>
271    )
272  })
273}
274

1// Vesuvius palette (placeholder): replace the hex values with the design system's.
2// Every mod in affreschi keeps its own copy, because a plugin cannot import another plugin's files.
3export const VESUVIUS = {
4  lava: '#E4572E',
5  magma: '#F28C28',
6  sulfur: '#F2C14E',
7  tuff: '#C9B79C',
8  ash: '#8A8D91',
9  gulf: '#2E86AB',
10  vineyard: '#7A9E4E',
11} as const
12
13// The order categories take their colors in, warmest first.
14export const SEQUENCE = [
15  VESUVIUS.lava,
16  VESUVIUS.magma,
17  VESUVIUS.sulfur,
18  VESUVIUS.tuff,
19  VESUVIUS.gulf,
20  VESUVIUS.vineyard,
21] as const
22

1// What one loop consumed: the main conversation, or one subagent.
2export type AgentUsage = {
3  id: string
4  label: string
5  type: string
6  model: string
7  status: 'running' | 'waiting' | 'idle' | 'completed' | 'failed' | 'killed' | 'pending'
8  steps: number
9  tools: number
10  input: number
11  output: number
12  cacheRead: number
13  cacheWrite: number
14  durationMs: number | null
15}
16
17declare module 'claude-code' {
18  interface PluginState {
19    'agent-usage': { agents: AgentUsage[]; expanded: string | null; costUsd: number | null }
20  }
21}
22