agent-flow

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { AgentNode, Flow, FlowKind } from '../types'
5
6const PANE = 'agent-flow'
7const MAIN = 'main'
8const ORCHESTRATOR: AgentNode = { id: MAIN, label: 'orchestrator', type: 'main', status: 'idle', tools: 0 }
9
10const agents = atom({ plugin: 'agent-flow', key: 'agents' } as const, [ORCHESTRATOR])
11const flows = atom({ plugin: 'agent-flow', key: 'flows' } as const, [])
12
13const clip = (text: string, max = 80) => {
14  const line = text.replace(/\s+/g, ' ').trim()
15  return line.length > max ? line.slice(0, max - 1) + '…' : line
16}
17
18const PROMPT_MAX = 2000
19const ANSWER_MAX = 1200
20
21const keep = (text: string, max: number) => (text.length > max ? text.slice(0, max - 1) + '…' : text)
22
23let flowSeq = 0
24
25function patch($: EngineInterface, id: string, fn: (a: AgentNode) => AgentNode) {
26  return update($, agents, list => list.map(a => (a.id === id ? fn(a) : a)))
27}
28
29async function flow($: EngineInterface, from: string, to: string, kind: FlowKind, text: string) {
30  const now = await $.clock.now()
31  const one: Flow = { id: `${now}-${flowSeq++}`, from, to, kind, text: clip(text) }
32  await update($, flows, list => [...list, one].slice(-60))
33}
34
35async function known($: EngineInterface, id: string) {
36  return (await read($, agents)).some(a => a.id === id)
37}
38
39function openPane($: EngineInterface) {
40  return $.ui.open({ id: PANE, title: 'Agent flow' })
41}
42
43type AutoOpen = 'never' | 'on-subagent' | 'always'
44const AUTO_OPEN: readonly AutoOpen[] = ['never', 'on-subagent', 'always']
45
46export const register: Register = (on, options) => {
47  const autoOpen = AUTO_OPEN.find(v => v === options.autoOpen) ?? 'on-subagent'
48
49  on('session.start', async ($, e, next) => {
50    await $.command.register({ name: 'agent-flow', description: 'Show the orchestrator and its subagents in an animated pane' })
51    if (autoOpen === 'always') void openPane($).catch(() => undefined)
52    return next(e)
53  })
54
55  on('command.run', { command: 'agent-flow' }, async $ => {
56    await openPane($)
57    return { text: 'Agent flow pane opened.' }
58  })
59
60  on('turn.start', async ($, e, next) => {
61    await patch($, MAIN, a => ({
62      ...a,
63      status: 'working',
64      tool: undefined,
65      prompt: e.text ? keep(e.text, PROMPT_MAX) : a.prompt,
66    }))
67    return next(e)
68  })
69
70  on('agent.spawn', async ($, e, next) => {
71    const started = await next(e)
72    if (started.deny !== undefined || started.agentId === undefined) return started
73
74    const parent = e.parentAgentId ?? MAIN
75    const now = await $.clock.now()
76    const node: AgentNode = {
77      id: started.agentId,
78      parentId: parent,
79      label: e.description || e.subagentType,
80      type: e.subagentType,
81      name: e.name,
82      status: 'working',
83      tools: 0,
84      model: started.model,
85      prompt: keep(e.prompt, PROMPT_MAX),
86      background: e.background,
87      startedAt: now,
88    }
89    await update($, agents, list => [...list.filter(a => a.id !== node.id), node])
90    await flow($, parent, node.id, 'task', e.prompt)
91
92    if (autoOpen !== 'never' && !(await $.ui.panes()).some(p => p.id === PANE)) void openPane($).catch(() => undefined)
93    return started
94  })
95
96  on('tool.call', async ($, e, next) => {
97    const id = e.agentId ?? MAIN
98    if (!(await known($, id))) return next(e)
99
100    const skill = e.tool === 'Skill' ? e.skill : undefined
101    await patch($, id, a => ({
102      ...a,
103      status: 'working',
104      tool: e.tool,
105      tools: a.tools + 1,
106      toolCounts: { ...a.toolCounts, [e.tool]: (a.toolCounts?.[e.tool] ?? 0) + 1 },
107      skills: skill && !a.skills?.includes(skill) ? [...(a.skills ?? []), skill] : a.skills,
108    }))
109    const ran = await next(e)
110    await patch($, id, a => (a.tool === e.tool ? { ...a, tool: undefined } : a))
111    return ran
112  })
113
114  on('turn.step', async function* ($, e, next) {
115    const id = e.agentId ?? MAIN
116    if (!(await known($, id))) return yield* next(e)
117
118    await patch($, id, a => ({ ...a, model: e.model, effort: e.effort }))
119    const answered = yield* next(e)
120    const usage = answered.usage
121    if (usage) {
122      const input = usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
123      await patch($, id, a => ({
124        ...a,
125        model: usage.model,
126        tokensIn: (a.tokensIn ?? 0) + input,
127        tokensOut: (a.tokensOut ?? 0) + usage.output_tokens,
128        cacheRead: (a.cacheRead ?? 0) + usage.cache_read_input_tokens,
129      }))
130    }
131    return answered
132  })
133
134  on('session.send', async ($, e, next) => {
135    const sent = await next(e)
136    if (!sent.isDelivered) return sent
137
138    const list = await read($, agents)
139    const to = list.find(a => a.id === e.to || a.name === e.to)
140    if (to) await flow($, e.agentId ?? MAIN, to.id, 'message', e.text)
141    return sent
142  })
143
144  on('turn.complete', async ($, e, next) => {
145    const id = e.agentId ?? MAIN
146    const list = await read($, agents)
147    const node = list.find(a => a.id === id)
148    if (!node) return next(e)
149
150    if (id === MAIN) {
151      await patch($, MAIN, a => ({ ...a, status: 'idle', tool: undefined, answer: keep(e.answer, ANSWER_MAX) }))
152    } else {
153      const now = await $.clock.now()
154      const isOk = e.reason === 'answer'
155      await patch($, id, a => ({
156        ...a,
157        status: isOk ? 'done' : 'failed',
158        tool: undefined,
159        endedAt: now,
160        answer: keep(e.answer || e.reason, ANSWER_MAX),
161      }))
162      await flow($, id, node.parentId ?? MAIN, 'result', e.answer || e.reason)
163    }
164    return next(e)
165  })
166
167  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
168    const list = await read($, agents)
169    const recent = await read($, flows)
170
171    if (e.surface === 'terminal' || e.surface === 'desktop') {
172      const { Box, Client, Button } = $.ui.resolve(e)
173      return (
174        <Box flexDirection="column">
175          <Client key="graph" module="./graph.tsx" props={{ agents: list, flows: recent }} />
176          <Box marginTop={1}>
177            <Button key="clear" onPress={async () => {
178              await update($, agents, all => all.filter(a => a.id === MAIN || a.status === 'working'))
179              await update($, flows, () => [])
180            }}>clear finished</Button>
181          </Box>
182        </Box>
183      )
184    }
185
186    const { Box, Text } = $.ui.resolve(e)
187    return (
188      <Box flexDirection="column">
189        {list.map(a => (
190          <Text dimColor={a.status !== 'working'}>
191            {a.id === MAIN ? '' : '  └ '}
192            {a.label} · {a.status}
193            {a.tool ? ` · ${a.tool}` : ''}
194          </Text>
195        ))}
196      </Box>
197    )
198  })
199}
200

1import type { ClientModule, RenderElement } from 'claude-code'
2
3import type { AgentNode, Flow } from '../types'
4
5type Props = { agents: AgentNode[]; flows: Flow[] }
6type State = { frame: number; expanded: string[]; cursor?: string }
7
8const TICK_MS = 90
9const TRAVEL = 10 // frames a packet takes to cross a connector
10const FLASH = 40 // frames a new interaction stays highlighted in the log
11const WIRE = 4 // connector width in cells
12const SPINNER = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
13const PULSE = ['cyan', 'cyanBright', 'whiteBright', 'cyanBright']
14const KIND_COLOR = { task: 'yellow', result: 'green', message: 'magenta' } as const
15const KIND_GLYPH = { task: '●', result: '◆', message: '✉' } as const
16
17// First frame each flow was seen on, so a packet animates once from when it
18// arrives. Flows present at mount count as already delivered.
19const seenAt = new Map<string, number>()
20let mounted = false
21
22// What the latest drawing laid out, for the pointer and key listeners: the
23// agent each drawn line belongs to, and the agents in tree order.
24const layout: { owners: (string | undefined)[]; order: string[] } = { owners: [], order: [] }
25
26const PROMPT_LINES = 12
27const ANSWER_LINES = 8
28
29// Breaks text into lines of at most `width` cells, on spaces where it can.
30const wrapText = (text: string, width: number) => {
31  const out: string[] = []
32  for (const para of text.split('\n')) {
33    let line = ''
34    for (const word of para.split(/\s+/).filter(Boolean)) {
35      let rest = word
36      while (rest.length > width) {
37        if (line) out.push(line)
38        out.push(rest.slice(0, width))
39        rest = rest.slice(width)
40        line = ''
41      }
42      if (!line) line = rest
43      else if (line.length + 1 + rest.length <= width) line += ' ' + rest
44      else {
45        out.push(line)
46        line = rest
47      }
48    }
49    out.push(line)
50  }
51  while (out.length > 0 && out[out.length - 1] === '') out.pop()
52  return out
53}
54
55const duration = (a: AgentNode) => {
56  if (a.startedAt === undefined || a.endedAt === undefined) return ''
57  const s = Math.round((a.endedAt - a.startedAt) / 1000)
58  return s < 60 ? ` ${s}s` : ` ${Math.floor(s / 60)}m${s % 60}s`
59}
60
61const shortModel = (model: string) => model.replace(/^claude-/, '').replace(/-\d{8}$/, '').replace(/\[1m\]$/, '')
62
63const count = (n: number) => (n < 1000 ? `${n}` : n < 1e6 ? `${(n / 1000).toFixed(1)}k` : `${(n / 1e6).toFixed(2)}M`)
64
65// A level as is; a number is a thinking budget in tokens.
66const effort = (a: AgentNode) =>
67  a.effort === undefined ? 'n/a' : typeof a.effort === 'number' ? `budget ${count(a.effort)}` : a.effort
68
69const usage = (a: AgentNode) => {
70  const model = a.model ? ` · ${shortModel(a.model)}` : ''
71  const tokens = a.tokensIn || a.tokensOut ? ` · ↑${count(a.tokensIn ?? 0)} ↓${count(a.tokensOut ?? 0)}` : ''
72  return model + tokens
73}
74
75const GraphView: ClientModule<Props, State> = (props, surface) => {
76  const { Box, Text } = surface.elements
77  const { agents, flows } = props
78
79  const state: State = surface.state ?? { frame: 0, expanded: [] }
80  const set = (change: Partial<State>) => surface.setState({ ...(surface.state ?? state), ...change })
81  const toggle = (id: string) => {
82    const open = surface.state?.expanded ?? []
83    set({ cursor: id, expanded: open.includes(id) ? open.filter(x => x !== id) : [...open, id] })
84  }
85
86  if (surface.state === undefined) {
87    surface.every(TICK_MS, () => set({ frame: (surface.state?.frame ?? 0) + 1 }))
88    surface.onPointer(e => {
89      if (e.type !== 'down' || e.button !== 'left') return
90      const id = layout.owners[e.y]
91      if (id) toggle(id)
92    })
93    surface.onKey(e => {
94      const order = layout.order
95      if (order.length === 0) return
96      const at = order.indexOf(surface.state?.cursor ?? '')
97      if (e.key === 'down' || e.key === 'j') set({ cursor: order[Math.min(order.length - 1, at + 1)] })
98      else if (e.key === 'up' || e.key === 'k') set({ cursor: order[Math.max(0, at - 1)] })
99      else if ((e.key === 'return' || e.key === ' ') && at >= 0) toggle(order[at] as string)
100    })
101    surface.setState(state)
102  }
103  const { frame, expanded, cursor } = state
104  const width = surface.columns > 0 ? surface.columns : 80
105
106  if (!mounted) {
107    for (const f of flows) seenAt.set(f.id, -1e6)
108    mounted = true
109  }
110  for (const f of flows) if (!seenAt.has(f.id)) seenAt.set(f.id, frame)
111
112  const byId = new Map(agents.map(a => [a.id, a]))
113  const children = (id: string) => agents.filter(a => a.parentId === id)
114  const name = (id: string) => byId.get(id)?.label ?? id
115
116  // The packet in flight on the wire between `child` and its parent, if any.
117  const inFlight = (child: AgentNode) => {
118    for (let i = flows.length - 1; i >= 0; i--) {
119      const f = flows[i]
120      if (!f) continue
121      const age = frame - (seenAt.get(f.id) ?? frame)
122      if (age >= TRAVEL) continue
123      const down = f.from === child.parentId && f.to === child.id
124      const up = f.from === child.id && f.to === child.parentId
125      if (down || up) return { flow: f, pos: down ? age / TRAVEL : 1 - age / TRAVEL }
126    }
127    return undefined
128  }
129
130  const wire = (child: AgentNode) => {
131    const packet = inFlight(child)
132    if (packet) {
133      const at = Math.min(WIRE - 1, Math.floor(packet.pos * WIRE))
134      const color = KIND_COLOR[packet.flow.kind]
135      return (
136        <Text>
137          <Text color={color} dimColor>{'─'.repeat(at)}</Text>
138          <Text color={color} bold>{KIND_GLYPH[packet.flow.kind]}</Text>
139          <Text color={color} dimColor>{'─'.repeat(WIRE - 1 - at)}</Text>
140        </Text>
141      )
142    }
143    if (child.status === 'working') {
144      // Marching dots toward the child while it works.
145      const cells = Array.from({ length: WIRE }, (_, i) => ((i - frame) % 4 === 0 ? '•' : '─')).join('')
146      return <Text color="cyan">{cells}</Text>
147    }
148    return <Text dimColor>{'─'.repeat(WIRE)}</Text>
149  }
150
151  const badge = (a: AgentNode) => {
152    switch (a.status) {
153      case 'working':
154        return <Text color={PULSE[frame % PULSE.length]} bold>{SPINNER[frame % SPINNER.length]}</Text>
155      case 'done':
156        return <Text color="green">✓</Text>
157      case 'failed':
158        return <Text color="red">✗</Text>
159      default:
160        return <Text dimColor>○</Text>
161    }
162  }
163
164  const detail = (a: AgentNode) => {
165    if (a.status === 'working') return a.tool ? ` ${a.tool}` : ' thinking'
166    if (a.status === 'idle') return ' idle'
167    return `${duration(a)} · ${a.tools} tools`
168  }
169
170  const lines: RenderElement[] = []
171  const owners: (string | undefined)[] = []
172  const order: string[] = []
173  const line = (el: RenderElement, owner?: string) => {
174    lines.push(el)
175    owners.push(owner)
176  }
177
178  const details = (a: AgentNode, indent: string) => {
179    const room = Math.max(20, width - indent.length - 2)
180    const field = (label: string, value: string) =>
181      line(
182        <Text wrap="truncate">
183          <Text dimColor>{indent}</Text>
184          <Text color="gray">{label.padEnd(7)}</Text>
185          {value}
186        </Text>,
187        a.id,
188      )
189    const block = (label: string, text: string, max: number) => {
190      const wrapped = wrapText(text, room)
191      line(
192        <Text>
193          <Text dimColor>{indent}</Text>
194          <Text color="gray">{label}</Text>
195        </Text>,
196        a.id,
197      )
198      for (const one of wrapped.slice(0, max)) line(<Text wrap="truncate"><Text dimColor>{indent}</Text>  {one}</Text>, a.id)
199      if (wrapped.length > max)
200        line(<Text dimColor>{indent}  … {wrapped.length - max} more lines</Text>, a.id)
201    }
202
203    const tokens = a.tokensIn || a.tokensOut
204      ? `in ${count(a.tokensIn ?? 0)} (cache read ${count(a.cacheRead ?? 0)}) · out ${count(a.tokensOut ?? 0)}`
205      : 'no requests yet'
206    const tools = Object.entries(a.toolCounts ?? {})
207      .sort((x, y) => y[1] - x[1])
208      .map(([tool, n]) => `${tool}×${n}`)
209      .join(' · ')
210
211    field('type', `${a.type}${a.background ? ' · background' : ''}${a.name ? ` · @${a.name}` : ''}`)
212    field('model', a.model ?? 'unknown')
213    field('effort', effort(a))
214    field('tokens', tokens)
215    if (a.endedAt !== undefined) field('time', duration(a).trim())
216    field('tools', tools || 'none')
217    field('skills', a.skills?.join(' · ') || 'none')
218    block(a.id === 'main' ? 'last prompt' : 'task', a.prompt || '(not recorded)', PROMPT_LINES)
219    if (a.answer) block(a.id === 'main' ? 'last answer' : 'result', a.answer, ANSWER_LINES)
220  }
221
222  const walk = (a: AgentNode, guides: string, isLast: boolean, depth: number) => {
223    const isWorking = a.status === 'working'
224    const isOpen = expanded.includes(a.id)
225    const isCursor = cursor === a.id
226    order.push(a.id)
227    line(
228      <Text key={a.id} wrap="truncate">
229        {depth > 0 && <Text dimColor>{guides + (isLast ? '└' : '├')}</Text>}
230        {depth > 0 && wire(a)}
231        {depth > 0 ? ' ' : ''}
232        {badge(a)}
233        <Text bold={isWorking || depth === 0} dimColor={!isWorking && depth > 0 && !isCursor} inverse={isCursor}>
234          {' '}
235          {a.label}
236          {' '}
237        </Text>
238        <Text dimColor>{isOpen ? '▾' : '▸'}</Text>
239        <Text dimColor>{depth > 0 && a.type !== a.label ? ` (${a.type})` : ''}</Text>
240        <Text color={isWorking ? 'cyan' : undefined} dimColor={!isWorking}>{detail(a)}</Text>
241        <Text color="gray">{usage(a)}</Text>
242      </Text>,
243      a.id,
244    )
245    const kids = children(a.id)
246    const nextGuides = depth === 0 ? '' : guides + (isLast ? ' ' : '│') + ' '.repeat(WIRE + 1)
247    if (isOpen) details(a, nextGuides + (kids.length > 0 ? '│ ' : '  ') + ' ')
248    kids.forEach((k, i) => walk(k, nextGuides, i === kids.length - 1, depth + 1))
249  }
250
251  const root = byId.get('main')
252  if (root) walk(root, '', true, 0)
253
254  const totalIn = agents.reduce((n, a) => n + (a.tokensIn ?? 0), 0)
255  const totalOut = agents.reduce((n, a) => n + (a.tokensOut ?? 0), 0)
256  const working = agents.filter(a => a.status === 'working').length
257  const finished = agents.filter(a => a.status === 'done' || a.status === 'failed').length
258
259  // Every line is one row, so a pointer's `y` indexes `owners` directly.
260  const all: RenderElement[] = []
261  const allOwners: (string | undefined)[] = []
262  const push = (el: RenderElement, owner?: string) => {
263    all.push(el)
264    allOwners.push(owner)
265  }
266  push(
267    <Text wrap="truncate" dimColor>
268      {agents.length - 1} subagents · <Text color="cyan">{working} working</Text> · {finished} finished · ↑{count(totalIn)} in ↓{count(totalOut)} out
269    </Text>,
270  )
271  push(<Text dimColor wrap="truncate">click a row (or ↑/↓ and enter) for details</Text>)
272  lines.forEach((el, i) => push(el, owners[i]))
273  push(<Text> </Text>)
274  push(<Text bold>Interactions</Text>)
275
276  const logRoom = Math.max(3, surface.rows > 0 ? surface.rows - all.length : 8)
277  const log = flows.slice(-logRoom)
278  if (log.length === 0) push(<Text dimColor>None yet.</Text>)
279  for (const f of log) {
280    const isNew = frame - (seenAt.get(f.id) ?? frame) < FLASH
281    const arrow = f.kind === 'result' ? '◀── result ──' : f.kind === 'task' ? '── task ──▶' : '── msg ──▶'
282    const [left, right] = f.kind === 'result' ? [f.to, f.from] : [f.from, f.to]
283    push(
284      <Text key={f.id} wrap="truncate" dimColor={!isNew}>
285        <Text bold={isNew}>{name(left)}</Text>
286        <Text color={KIND_COLOR[f.kind]}> {arrow} </Text>
287        <Text bold={isNew}>{name(right)}</Text>
288        <Text dimColor> {f.text}</Text>
289      </Text>,
290    )
291  }
292
293  layout.owners = allOwners
294  layout.order = order
295
296  return <Box flexDirection="column">{all}</Box>
297}
298
299export default GraphView
300

1export type AgentStatus = 'working' | 'idle' | 'done' | 'failed'
2
3export type AgentNode = {
4  id: string
5  parentId?: string
6  label: string
7  type: string
8  name?: string
9  status: AgentStatus
10  tool?: string
11  tools: number
12  model?: string
13  effort?: string | number
14  tokensIn?: number
15  tokensOut?: number
16  cacheRead?: number
17  prompt?: string
18  answer?: string
19  toolCounts?: Record<string, number>
20  skills?: string[]
21  background?: boolean
22  startedAt?: number
23  endedAt?: number
24}
25
26export type FlowKind = 'task' | 'result' | 'message'
27
28export type Flow = {
29  id: string
30  from: string
31  to: string
32  kind: FlowKind
33  text: string
34}
35
36declare module 'claude-code' {
37  interface PluginState {
38    'agent-flow': { agents: AgentNode[]; flows: Flow[] }
39  }
40}
41