subagent-tree

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, Timer } from 'claude-code'
3
4import type { AgentStats } from '../types'
5import { buildRows, formatTokens, liveCount, prune } from './tree'
6
7const PANE = 'subagent-tree'
8const stats = atom(
9  { plugin: 'subagent-tree', key: 'stats' } as const,
10  {},
11  { shape: 'stats-v2' },
12)
13
14const blank = (now: number): AgentStats => ({
15  tools: 0,
16  steps: 0,
17  tokensOut: 0,
18  activeMs: 0,
19  lastEventAt: now,
20})
21
22async function touch(
23  $: EngineInterface,
24  id: string,
25  change: (s: AgentStats) => AgentStats,
26) {
27  const now = await $.clock.now()
28  await update($, stats, all => {
29    const current = all[id] ?? blank(now)
30    // A run opens at the agent's first event: subagents raise no turn.start.
31    const open = { ...current, runStartedAt: current.runStartedAt ?? now, lastEventAt: now }
32    return { ...all, [id]: change(open) }
33  })
34}
35
36// Module state: a reload starts it over, and the next draw of the pane (a
37// reload redraws it) or /agents-tree restores it.
38let timer: Timer | undefined
39let lastRenderAt = 0
40let wasLive = false
41let idleTicks = 0
42// A pane that is drawn is redrawn on every invalidate, and at least every
43// IDLE_EVERY ticks even when nothing runs, so a pane not drawn for this long
44// is closed or was dropped without ui.close reaching this plugin; the timer
45// stops and the next draw restarts it.
46const IDLE_EVERY = 5
47const STALE_MS = 12_000
48
49function stopTimer() {
50  timer?.cancel()
51  timer = undefined
52  wasLive = false
53  idleTicks = 0
54}
55
56async function pruneUnlisted($: EngineInterface, agents: { id: string }[], now: number) {
57  const listed = new Set(agents.map(agent => agent.id))
58  const all = await read($, stats)
59  if (Object.keys(prune(all, listed, now)).length !== Object.keys(all).length) {
60    await update($, stats, current => prune(current, listed, now))
61  }
62}
63
64// Whether the engine still lists the pane; undefined when it cannot say.
65async function isPaneOpen($: EngineInterface): Promise<boolean | undefined> {
66  try {
67    return (await $.ui.panes()).some(pane => pane.id === PANE)
68  } catch {
69    return undefined
70  }
71}
72
73async function redrawTick($: EngineInterface) {
74  try {
75    const now = await $.clock.now()
76    if (now - lastRenderAt > STALE_MS) return stopTimer()
77    // The engine's own record: a pane it no longer lists is gone, whatever the clock says.
78    if ((await isPaneOpen($)) === false) return stopTimer()
79
80    const agents = await $.agent.list()
81    const isLive = liveCount(agents) > 0
82    idleTicks = isLive || wasLive ? 0 : idleTicks + 1
83    // Every second while something runs, and once more after the last agent
84    // finishes so the final frame is not left showing it running. When idle, a
85    // slow heartbeat: the draw reads the agent list, so a status that changed
86    // with no state write (a pending agent starting) still shows.
87    if (isLive || wasLive || idleTicks >= IDLE_EVERY) {
88      $.ui.invalidate('ui.render')
89      idleTicks = 0
90    }
91    wasLive = isLive
92    await pruneUnlisted($, agents, now)
93  } catch (error) {
94    // A failed period ends the interval; forget it so the next draw restarts it.
95    $.ui.log(`redraw tick failed: ${String(error)}`, { to: 'debug' })
96    stopTimer()
97  }
98}
99
100// A hook's refusal arrives as 'HooksError: <plugin>: $.ui.close: <reason>'.
101function reasonOf(error: unknown): string {
102  const text = error instanceof Error ? error.message : String(error)
103
104  return text.replace(/^.*?\$\.ui\.close: /, '')
105}
106
107async function watch($: EngineInterface) {
108  lastRenderAt = await $.clock.now()
109  // The timer keeps the `$` of the hook that started it, as the types' own
110  // example does; a tick that fails stops it and the next draw starts it again.
111  timer ??= $.clock.every(1000, () => redrawTick($))
112}
113
114export const register: Register = on => {
115  on('session.start', async ($, e, next) => {
116    await $.command.register({
117      name: 'agents-tree',
118      description: 'Show the session\'s subagents as a live tree in a pane',
119      argumentHint: '[close]',
120    })
121
122    return next(e)
123  })
124
125  on('command.run', { command: 'agents-tree' }, async ($, e) => {
126    const argument = e.args.trim()
127
128    if (argument === 'close') {
129      try {
130        await $.ui.close({ id: PANE })
131      } catch (error) {
132        // A hook beneath refused the close.
133        return { text: `Could not close the subagent tree: ${reasonOf(error)}` }
134      }
135      // A hook beneath may also keep the pane open without refusing.
136      if ((await isPaneOpen($)) === true) return { text: 'Subagent tree stays open.' }
137      stopTimer()
138
139      return { text: 'Subagent tree closed.' }
140    }
141    if (argument !== '' && argument !== 'open') {
142      return { text: 'Usage: /agents-tree [close]' }
143    }
144    await $.ui.open({ id: PANE, title: 'Subagents' })
145    await watch($)
146
147    return { text: 'Subagent tree opened.' }
148  })
149
150  on('tool.call', async ($, e, next) => {
151    if (e.agentId) {
152      // Written before the call runs, so the pane shows the tool while it runs;
153      // not awaited, so the call is not held up by the write. update() retries
154      // on a version miss, so the write is not lost to the turn.step and
155      // turn.complete writes; on an interrupt it may be aborted and that call
156      // goes uncounted.
157      touch($, e.agentId, s => ({ ...s, tools: s.tools + 1, lastTool: e.tool })).catch(
158        () => undefined,
159      )
160    }
161
162    return next(e)
163  })
164
165  on('turn.step', async function* ($, e, next) {
166    if (e.agentId) {
167      // A new request means the last tool call has finished.
168      await touch($, e.agentId, s => ({ ...s, steps: s.steps + 1, lastTool: undefined }))
169    }
170
171    return yield* next(e)
172  })
173
174  on('turn.complete', async ($, e, next) => {
175    const result = await next(e)
176    if (e.agentId) {
177      const id = e.agentId
178      const now = await $.clock.now()
179      const listed = new Set((await $.agent.list()).map(agent => agent.id))
180      await update($, stats, all => {
181        const current = all[id] ?? blank(now)
182        // Loaded mid-run: the run began when the turn's length says it did.
183        const began = current.runStartedAt ?? now - e.durationMs
184        const updated = {
185          ...all,
186          [id]: {
187            ...current,
188            runStartedAt: undefined,
189            activeMs: current.activeMs + Math.max(0, now - began),
190            lastEventAt: now,
191            // No usage on an interrupt or an API error.
192            tokensOut: current.tokensOut + (e.usage?.output_tokens ?? 0),
193          },
194        }
195
196        return prune(updated, listed, now)
197      })
198    }
199
200    return result
201  })
202
203  on('ui.close', { id: PANE }, async ($, e, next) => {
204    const result = await next(e)
205    // A hook beneath may refuse the close ({ deny }) or keep the pane open by
206    // answering without next; the engine's list of panes says which happened.
207    // When it cannot say, only a refusal keeps the timer running.
208    const isOpen = await isPaneOpen($)
209    const isRefused = Boolean((result as { deny?: string } | undefined)?.deny)
210    if (isOpen === false || (isOpen === undefined && !isRefused)) stopTimer()
211
212    return result
213  })
214
215  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
216    const { Box, Text } = $.ui.resolve(e)
217    await watch($)
218    const all = await read($, stats)
219    const agents = await $.agent.list()
220    const now = await $.clock.now()
221    const rows = buildRows(agents, all, now)
222    const live = liveCount(agents)
223    // Only agents the tree shows: forks and workflow agents carry ids no list names.
224    const tokens = agents.reduce((sum, agent) => sum + (all[agent.id]?.tokensOut ?? 0), 0)
225
226    return (
227      <Box flexDirection="column">
228        {rows.length === 0 && <Text dimColor>No subagents yet.</Text>}
229        {rows.map(row => (
230          <Box key={row.id} flexDirection="column" paddingLeft={row.depth * 2}>
231            <Text bold={row.isLive} dimColor={!row.isLive} wrap="truncate">
232              {row.glyph} {row.label}
233            </Text>
234            <Box paddingLeft={2}>
235              <Text dimColor wrap="truncate">
236                {row.detail}
237              </Text>
238            </Box>
239          </Box>
240        ))}
241        {rows.length > 0 && (
242          <Text dimColor>
243            {rows.length} agents · {live} active · {formatTokens(tokens)} tokens out
244          </Text>
245        )}
246      </Box>
247    )
248  })
249}
250

1import type { AgentStats } from '../types'
2
3/** The part of `$.agent.list()`'s entries the tree reads. */
4export type AgentLike = {
5  id: string
6  description: string
7  type: string
8  status: string
9  parentId?: string
10  name?: string
11}
12
13export type Row = {
14  id: string
15  depth: number
16  glyph: string
17  label: string
18  detail: string
19  isLive: boolean
20}
21
22const LIVE = new Set(['pending', 'running', 'waiting'])
23
24const GLYPHS: Record<string, string> = {
25  pending: '·',
26  running: '●',
27  waiting: '◐',
28  idle: '○',
29  completed: '✓',
30  failed: '✗',
31  killed: '✗',
32}
33
34export const isLive = (status: string): boolean => LIVE.has(status)
35
36export const formatElapsed = (ms: number): string => {
37  const seconds = Math.max(0, Math.floor(ms / 1000))
38  if (seconds < 60) return `${seconds}s`
39  const minutes = Math.floor(seconds / 60)
40  if (minutes < 60) return `${minutes}m${String(seconds % 60).padStart(2, '0')}s`
41
42  return `${Math.floor(minutes / 60)}h${String(minutes % 60).padStart(2, '0')}m`
43}
44
45export const formatTokens = (tokens: number): string =>
46  tokens >= 1_000_000
47    ? `${(tokens / 1_000_000).toFixed(1)}M`
48    : tokens >= 1_000
49      ? `${Math.round(tokens / 1_000)}k`
50      : String(tokens)
51
52/**
53 * How long the agent has been active: its finished runs, plus the run in
54 * progress. A run left open by an agent that is no longer live stops at its
55 * last event, so a killed or failed agent's clock does not keep moving.
56 */
57export const elapsedMs = (stats: AgentStats, status: string, now: number): number => {
58  if (stats.runStartedAt === undefined) return stats.activeMs
59  const end = isLive(status) ? now : stats.lastEventAt
60
61  return stats.activeMs + Math.max(0, end - stats.runStartedAt)
62}
63
64const PRUNE_AFTER_MS = 10 * 60 * 1000
65const PRUNE_CAP = 200
66
67/**
68 * Drops the stats of ids the agent list does not show once they have been
69 * idle for ten minutes, and past 200 entries drops unlisted ones first, oldest
70 * first. A listed agent's stats are never dropped.
71 */
72export const prune = (
73  stats: Record<string, AgentStats>,
74  listed: Set<string>,
75  now: number,
76): Record<string, AgentStats> => {
77  const kept = Object.entries(stats).filter(
78    ([id, s]) => listed.has(id) || now - s.lastEventAt < PRUNE_AFTER_MS,
79  )
80  if (kept.length <= PRUNE_CAP) return Object.fromEntries(kept)
81
82  const rank = ([id, s]: [string, AgentStats]) => (listed.has(id) ? Infinity : s.lastEventAt)
83  const newest = [...kept].sort((a, b) => rank(b) - rank(a)).slice(0, PRUNE_CAP)
84
85  return Object.fromEntries(newest)
86}
87
88const detailOf = (agent: AgentLike, stats: AgentStats | undefined, now: number): string => {
89  if (!stats) return agent.status
90  const parts = [
91    agent.status,
92    formatElapsed(elapsedMs(stats, agent.status, now)),
93    `${stats.tools} tool${stats.tools === 1 ? '' : 's'}`,
94    `${stats.steps} step${stats.steps === 1 ? '' : 's'}`,
95  ]
96  if (isLive(agent.status) && stats.lastTool) parts.push(`on ${stats.lastTool}`)
97  if (stats.tokensOut > 0) parts.push(`${formatTokens(stats.tokensOut)} out`)
98
99  return parts.join(' · ')
100}
101
102/**
103 * Flattens the agents into display rows, children under their parent in the
104 * order the list gives them. An agent whose parent is not in the list (or is
105 * the main loop) is a root; a cycle in the list cannot loop forever.
106 */
107export const buildRows = (
108  agents: AgentLike[],
109  stats: Record<string, AgentStats>,
110  now: number,
111): Row[] => {
112  const ids = new Set(agents.map(agent => agent.id))
113  const children = new Map<string, AgentLike[]>()
114  const roots: AgentLike[] = []
115
116  for (const agent of agents) {
117    if (agent.parentId && ids.has(agent.parentId) && agent.parentId !== agent.id) {
118      children.set(agent.parentId, [...(children.get(agent.parentId) ?? []), agent])
119    } else {
120      roots.push(agent)
121    }
122  }
123
124  const rows: Row[] = []
125  const seen = new Set<string>()
126  const visit = (agent: AgentLike, depth: number) => {
127    if (seen.has(agent.id)) return
128    seen.add(agent.id)
129    const title = agent.description || agent.name || agent.id
130    rows.push({
131      id: agent.id,
132      depth,
133      glyph: GLYPHS[agent.status] ?? '?',
134      label: `${agent.type}: ${title}`,
135      detail: detailOf(agent, stats[agent.id], now),
136      isLive: isLive(agent.status),
137    })
138    for (const child of children.get(agent.id) ?? []) visit(child, depth + 1)
139  }
140  for (const root of roots) visit(root, 0)
141  // Agents left over are in a parent cycle: show them rather than hide them.
142  for (const agent of agents) visit(agent, 0)
143
144  return rows
145}
146
147export const liveCount = (agents: AgentLike[]): number =>
148  agents.filter(agent => isLive(agent.status)).length
149

1export type AgentStats = {
2  /** Tool calls the agent's loop has made. */
3  tools: number
4  /** Model requests its loop has made. */
5  steps: number
6  /** The tool it called last. */
7  lastTool?: string
8  /** Output tokens summed over its finished turns. */
9  tokensOut: number
10  /** Milliseconds its finished runs were active, idle gaps between runs left out. */
11  activeMs: number
12  /** Milliseconds since the epoch when the run in progress began; unset between runs. */
13  runStartedAt?: number
14  /** Milliseconds since the epoch of the latest event seen from it. */
15  lastEventAt: number
16}
17
18declare module 'claude-code' {
19  interface PluginState {
20    'subagent-tree': {
21      stats: Shaped<Record<string, AgentStats>>
22    }
23  }
24}
25