agent-sankey

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { SessionRecord } from '../types'
5import { buildGraph, renderSvg, textTree } from './sankey'
6import {
7  READ_SESSIONS,
8  familyGraphInput,
9  orchestrators,
10  outcomeOf,
11  ownRoot,
12  rootOf,
13  sessionLink,
14  sessionsFolder,
15} from './sessions'
16
17// The desktop measures the row in character cells; these turn cells into the
18// CSS pixels an Svg is sized in, so the diagram fills the row's width and most
19// of the visible height instead of a small default frame.
20const CELL_PX = { width: 8.4, height: 18 }
21const MIN_SIZE = { width: 640, height: 420 }
22const PICKER_SIZE = 8
23
24// Each state a session can end up in has a colour of its own.
25const OUTCOME_COLORS: Record<string, string> = {
26  Coordinating: '#6C8CFF',
27  Active: '#3DA5FF',
28  'PR open': '#A98BFA',
29  'PR merged': '#34C38F',
30  'PR closed': '#FF7A8A',
31  Idle: '#F7C548',
32  Archived: '#B9BEC9',
33}
34
35const sessions = atom({ plugin: 'agent-sankey', key: 'sessions' } as const, [])
36const root = atom({ plugin: 'agent-sankey', key: 'root' } as const, null)
37const note = atom({ plugin: 'agent-sankey', key: 'note' } as const, null)
38
39// Reads the desktop app's session records: titles, who started whom, turns
40// and state, never what was said.
41const load = async ($: EngineInterface): Promise<SessionRecord[]> => {
42  await update($, note, () => 'Reading sessions…')
43  try {
44    const home = (await $.env.get('HOME')) ?? ''
45    const { exitCode, stdout, stderr } = await $.process.run(
46      ['python3', '-c', READ_SESSIONS, sessionsFolder(home)],
47      { timeoutMs: 60_000 },
48    )
49    if (exitCode !== 0) throw new Error(stderr.trim() || `exit ${exitCode}`)
50    const records = JSON.parse(stdout) as SessionRecord[]
51    await update($, sessions, () => records)
52    await update($, note, () => null)
53    return records
54  } catch (error) {
55    await update($, note, () => `Could not read the sessions: ${String(error)}`)
56    return []
57  }
58}
59
60// The orchestrator to draw: one the command named, else this session's own
61// family, else the most recently active one.
62const pickRoot = (records: readonly SessionRecord[], args: string, ownCliId: string): string | null => {
63  const wanted = args.trim().toLowerCase()
64  const all = orchestrators(records)
65  if (wanted !== '') {
66    const named = all.find(r => r.title.toLowerCase().includes(wanted))
67    if (named) return named.id
68  }
69  return ownRoot(records, ownCliId) ?? all[0]?.id ?? null
70}
71
72const openSession = async ($: EngineInterface, id: string): Promise<void> => {
73  const { exitCode } = await $.process.run(['open', sessionLink(id)]).catch(() => ({ exitCode: 1 }))
74  if (exitCode !== 0) $.ui.toast('Could not open that session.')
75}
76
77export const register: Register = on => {
78  on('session.start', async ($, e, next) => {
79    await $.command.register({
80      name: 'sankey',
81      description: 'Show an orchestrator session and the sessions it started as a Sankey diagram',
82      argumentHint: '[orchestrator name]',
83    })
84
85    return next(e)
86  })
87
88  // The diagram draws inline in the conversation, in the command's own output row.
89  on('command.run', { command: 'sankey' }, async ($, e) => {
90    const records = await load($)
91    const picked = pickRoot(records, e.args, await $.session.id())
92    await update($, root, () => picked)
93    if (picked === null) return { text: 'No orchestrator sessions found: none of your sessions has started another.' }
94    const { agents, tally, rootTitle } = familyGraphInput(records, picked, await $.clock.now())
95    const graph = buildGraph(agents, tally, { rootLabel: rootTitle, unit: ['turn', 'turns'] })
96
97    return { text: `Session family of ${rootTitle}:\n${textTree(graph, 80).join('\n')}` }
98  })
99
100  on('ui.render', { component: 'CommandOutput', props: { command: 'sankey' } }, async ($, e) => {
101    const records = await read($, sessions)
102    const picked = await read($, root)
103    const message = await read($, note)
104    const now = await $.clock.now()
105    const { Box, Text, Button } = $.ui.resolve(e)
106    const columns = e.viewport?.columns ?? 100
107
108    if (picked === null || !records.some(r => r.id === picked)) {
109      return <Text dimColor>{message ?? 'No orchestrator sessions found.'}</Text>
110    }
111
112    const { agents, tally, rootTitle } = familyGraphInput(records, picked, now)
113    const graph = buildGraph(agents, tally, {
114      rootLabel: rootTitle,
115      unit: ['turn', 'turns'],
116      outcomeColors: OUTCOME_COLORS,
117    })
118    const family = records
119      .filter(r => rootOf(records, r.id) === picked)
120      .sort((a, b) => (a.id === picked ? -1 : b.id === picked ? 1 : b.lastActivityAt - a.lastActivityAt))
121    const others = orchestrators(records).slice(0, PICKER_SIZE)
122
123    const picker = (
124      <Box gap={1} flexWrap="wrap">
125        <Text dimColor>Orchestrator:</Text>
126        {others.map(one => (
127          <Button
128            key={`root-${one.id}`}
129            label={one.title}
130            variant={one.id === picked ? 'primary' : 'secondary'}
131            onPress={() => update($, root, () => one.id)}
132          />
133        ))}
134        <Button key="refresh" label="Refresh" onPress={() => load($)} />
135      </Box>
136    )
137    const opener = (
138      <Box flexDirection="column">
139        <Text dimColor>Click a session to open it:</Text>
140        <Box gap={1} flexWrap="wrap">
141          {family.map(one => (
142            <Button
143              key={`open-${one.id}`}
144              label={`${one.id === picked ? '★ ' : '↗ '}${one.title} · ${outcomeOf(one, now)}`}
145              onPress={() => openSession($, one.id)}
146            />
147          ))}
148        </Box>
149      </Box>
150    )
151
152    // The terminal draws no pictures: it gets the flow as text.
153    if (e.surface === 'terminal') {
154      return (
155        <Box flexDirection="column" gap={1}>
156          {picker}
157          {message !== null && <Text dimColor>{message}</Text>}
158          <Box flexDirection="column">
159            {textTree(graph, columns).map(line => (
160              <Text wrap="truncate-end">{line}</Text>
161            ))}
162          </Box>
163          {opener}
164        </Box>
165      )
166    }
167
168    // The Sankey as a picture, sized to the row: wide as the transcript and
169    // most of the visible height. Hover traces a lane; the buttons open sessions.
170    const { Svg } = $.ui.resolve(e)
171    const width = Math.max(MIN_SIZE.width, Math.round(columns * CELL_PX.width))
172    const rows = e.viewport?.rows ?? 40
173    const height = Math.max(MIN_SIZE.height, Math.round((rows - 10) * CELL_PX.height))
174    return (
175      <Box flexDirection="column" gap={1}>
176        {picker}
177        {message !== null && <Text dimColor>{message}</Text>}
178        <Svg
179          source={renderSvg(graph, width, null, height)}
180          alt={textTree(graph, 80).join('\n')}
181          width={width}
182          height={height}
183          isInteractive
184        />
185        {opener}
186      </Box>
187    )
188  })
189}
190

1import type { AgentNode, Tally } from '../types'
2
3export const MAIN = 'main'
4export const SELF = 'main:self'
5const MAX_TOOLS = 12
6
7export type Node = {
8  id: string
9  label: string
10  column: number
11  value: number
12  color: string
13  status?: string
14}
15
16export type Link = { source: string; target: string; value: number }
17
18export type Graph = {
19  nodes: Node[]
20  links: Link[]
21  columns: number
22  total: number
23  agentCount: number
24  // What a flow counts, singular and plural ("tool call", "turn").
25  unit: [string, string]
26  rootLabel: string
27}
28
29export type GraphOptions = {
30  rootLabel?: string
31  unit?: [string, string]
32  // Fixed colours for nodes in the last column, by label (a session's outcome).
33  outcomeColors?: Record<string, string>
34}
35
36export const count = (graph: Graph, n: number): string => `${n} ${n === 1 ? graph.unit[0] : graph.unit[1]}`
37
38// Light but saturated enough to tell apart on the white card.
39const PALETTE = ['#FF9F6B', '#5CCB8A', '#FF7A8A', '#A98BFA', '#3CC7C0', '#F7C548', '#F78FC4', '#6FB2FF', '#B5D86A', '#FFB86B']
40const ORCHESTRATOR = '#6C8CFF'
41
42// The label a call is counted under, or null for a call drawn as a spawn link.
43export const toolLabel = (tool: string, input: Record<string, unknown> = {}): string | null => {
44  if (tool === 'Agent' || tool === 'Task') return null
45  if (tool === 'SendMessage') {
46    const to = input.to ?? input.recipient
47    return typeof to === 'string' && to !== '' ? `Message → ${to}` : 'Message'
48  }
49  const mcp = /^mcp__([^_]+(?:_[^_]+)*)__(.+)$/.exec(tool)
50  if (mcp) return `MCP ${mcp[1]}`
51  return tool
52}
53
54const sum = (counts: Record<string, number> | undefined): number =>
55  Object.values(counts ?? {}).reduce((a, b) => a + b, 0)
56
57// Columns: the orchestrator, its children by depth, then the tools they called.
58// A link's width is tool calls: parent to child carries the child's whole subtree,
59// agent to tool what that agent called itself.
60export const buildGraph = (agents: readonly AgentNode[], given: Tally, options: GraphOptions = {}): Graph => {
61  const rootLabel = options.rootLabel ?? 'Orchestrator'
62  // The orchestrator's own calls flow through a node of their own beside its
63  // children, so every flow into the tools comes from the middle columns.
64  const { [MAIN]: ownCalls, ...rest } = given
65  const tally: Tally = sum(ownCalls) > 0 ? { ...rest, [SELF]: ownCalls ?? {} } : rest
66  const known = new Map<string, AgentNode>()
67  if (sum(ownCalls) > 0) known.set(SELF, { id: SELF, label: `${rootLabel} itself`, parentId: MAIN, status: 'running' })
68  known.set(MAIN, { id: MAIN, label: rootLabel, parentId: null, status: 'running' })
69  for (const agent of agents) if (agent.id !== MAIN) known.set(agent.id, agent)
70  for (const id of Object.keys(tally)) {
71    if (!known.has(id)) known.set(id, { id, label: `Agent ${id.slice(0, 8)}`, parentId: MAIN, status: 'completed' })
72  }
73
74  const parentOf = (id: string): string | null => {
75    if (id === MAIN) return null
76    const parent = known.get(id)?.parentId ?? MAIN
77    return known.has(parent) && parent !== id ? parent : MAIN
78  }
79  const children = new Map<string, string[]>()
80  for (const id of known.keys()) {
81    const parent = parentOf(id)
82    if (parent !== null) children.set(parent, [...(children.get(parent) ?? []), id])
83  }
84
85  const depth = new Map<string, number>([[MAIN, 0]])
86  const subtree = new Map<string, number>()
87  const walk = (id: string, d: number, seen: Set<string>): number => {
88    depth.set(id, d)
89    seen.add(id)
90    let total = sum(tally[id])
91    for (const child of children.get(id) ?? []) if (!seen.has(child)) total += walk(child, d + 1, seen)
92    subtree.set(id, total)
93    return total
94  }
95  const total = walk(MAIN, 0, new Set())
96  const agentDepth = Math.max(...depth.values())
97  const toolColumn = agentDepth + 1
98
99  const colorOf = new Map<string, string>([[MAIN, ORCHESTRATOR]])
100  const nodes: Node[] = []
101  const links: Link[] = []
102  const ordered = [...depth.keys()].sort((a, b) => (subtree.get(b) ?? 0) - (subtree.get(a) ?? 0))
103  let paint = 0
104  for (const id of ordered) {
105    if (id === SELF) colorOf.set(id, ORCHESTRATOR)
106    else if (id !== MAIN) colorOf.set(id, PALETTE[paint++ % PALETTE.length] ?? ORCHESTRATOR)
107    const agent = known.get(id)
108    nodes.push({
109      id,
110      label: agent?.label ?? id,
111      column: depth.get(id) ?? 0,
112      value: subtree.get(id) ?? 0,
113      color: colorOf.get(id) ?? ORCHESTRATOR,
114      status: agent?.status,
115    })
116    const parent = parentOf(id)
117    if (parent !== null) links.push({ source: parent, target: id, value: subtree.get(id) ?? 0 })
118  }
119
120  // The busiest tools by name; the rest fold into Other.
121  const toolTotals = new Map<string, number>()
122  for (const counts of Object.values(tally)) {
123    for (const [tool, n] of Object.entries(counts)) toolTotals.set(tool, (toolTotals.get(tool) ?? 0) + n)
124  }
125  const top = new Set([...toolTotals.entries()].sort((a, b) => b[1] - a[1]).slice(0, MAX_TOOLS).map(([t]) => t))
126  const toolOf = (tool: string): string => (top.has(tool) ? tool : 'Other')
127  const toolLinks = new Map<string, number>()
128  const biggest = new Map<string, { agent: string; n: number }>()
129  for (const [agent, counts] of Object.entries(tally)) {
130    if (!depth.has(agent)) continue
131    for (const [tool, n] of Object.entries(counts)) {
132      const key = `${agent}\u0000tool:${toolOf(tool)}`
133      toolLinks.set(key, (toolLinks.get(key) ?? 0) + n)
134    }
135  }
136  const toolNodes = new Map<string, number>()
137  for (const [key, n] of toolLinks) {
138    const [agent = MAIN, target = ''] = key.split('\u0000')
139    links.push({ source: agent, target, value: n })
140    toolNodes.set(target, (toolNodes.get(target) ?? 0) + n)
141    if ((biggest.get(target)?.n ?? -1) < n) biggest.set(target, { agent, n })
142  }
143  for (const [id, value] of [...toolNodes.entries()].sort((a, b) => b[1] - a[1])) {
144    nodes.push({
145      id,
146      label: id.slice('tool:'.length),
147      column: toolColumn,
148      value,
149      color: options.outcomeColors?.[id.slice('tool:'.length)] ?? colorOf.get(biggest.get(id)?.agent ?? MAIN) ?? ORCHESTRATOR,
150    })
151  }
152
153  return {
154    nodes,
155    links,
156    columns: toolColumn + 1,
157    total,
158    agentCount: [...depth.keys()].filter(id => id !== MAIN && id !== SELF).length,
159    unit: options.unit ?? ['tool call', 'tool calls'],
160    rootLabel,
161  }
162}
163
164const escape = (text: string): string =>
165  text.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;').replace(/"/g, '&quot;')
166
167const short = (text: string, max: number): string => (text.length > max ? `${text.slice(0, max - 1)}…` : text)
168
169type Box = { x: number; y: number; h: number }
170
171// The agents and tools a focused agent touches: its line up to the
172// orchestrator, everything beneath it, and the tools those agents called.
173export const relatedTo = (graph: Graph, focus: string | null): Set<string> | null => {
174  if (focus === null || !graph.nodes.some(n => n.id === focus)) return null
175  const related = new Set<string>([focus])
176  const parentOf = new Map(graph.links.filter(l => !l.target.startsWith('tool:')).map(l => [l.target, l.source]))
177  for (let up = parentOf.get(focus); up !== undefined; up = parentOf.get(up)) related.add(up)
178  const below = [focus]
179  while (below.length > 0) {
180    const id = below.pop() ?? ''
181    for (const link of graph.links) {
182      if (link.source !== id || related.has(link.target)) continue
183      related.add(link.target)
184      if (!link.target.startsWith('tool:')) below.push(link.target)
185    }
186  }
187  return related
188}
189
190export const renderSvg = (
191  graph: Graph,
192  width: number,
193  focus: string | null = null,
194  height?: number,
195): string => {
196  const W = Math.max(480, Math.round(width))
197  const perColumn = new Map<number, Node[]>()
198  for (const node of graph.nodes) perColumn.set(node.column, [...(perColumn.get(node.column) ?? []), node])
199  const busiest = Math.max(1, ...[...perColumn.values()].map(list => list.length))
200  const H = Math.max(340, Math.round(height ?? Math.min(900, busiest * 40 + 110)), busiest * 30 + 110)
201  const top = 54
202  const bottom = 34
203  const gap = 12
204  const nodeW = 16
205  const side = 28
206  const usable = H - top - bottom
207  const scale = Math.min(
208    ...[...perColumn.values()].map(list => {
209      const values = list.reduce((a, n) => a + Math.max(n.value, 0), 0)
210      return values === 0 ? Infinity : (usable - gap * (list.length - 1)) / values
211    }),
212  )
213  const k = Number.isFinite(scale) ? scale : 1
214  const heightOf = (value: number): number => Math.max(3, value * k)
215
216  const boxes = new Map<string, Box>()
217  const columnX = (column: number): number =>
218    graph.columns <= 1 ? side : side + (column * (W - 2 * side - nodeW)) / (graph.columns - 1)
219  // Columns left to right; the last column is ordered by where its flows come
220  // from, so bands cross as little as they can.
221  const middleOf = (id: string): number => {
222    const box = boxes.get(id)
223    return box === undefined ? 0 : box.y + box.h / 2
224  }
225  const lastColumn = graph.columns - 1
226  for (const [column, unordered] of [...perColumn.entries()].sort((a, b) => a[0] - b[0])) {
227    const list =
228      column !== lastColumn
229        ? unordered
230        : [...unordered].sort((a, b) => {
231            const at = (id: string): number => {
232              const into = graph.links.filter(l => l.target === id)
233              const weight = into.reduce((sum, l) => sum + l.value, 0) || 1
234              return into.reduce((sum, l) => sum + middleOf(l.source) * l.value, 0) / weight
235            }
236            return at(a.id) - at(b.id)
237          })
238    const span = list.reduce((a, n) => a + heightOf(n.value), 0) + gap * (list.length - 1)
239    let y = top + (usable - span) / 2
240    for (const node of list) {
241      boxes.set(node.id, { x: columnX(column), y, h: heightOf(node.value) })
242      y += heightOf(node.value) + gap
243    }
244  }
245
246  const related = relatedTo(graph, focus)
247  const ancestors = new Set<string>()
248  if (focus !== null) {
249    const parentOf = new Map(graph.links.filter(l => !l.target.startsWith('tool:')).map(l => [l.target, l.source]))
250    for (let up = parentOf.get(focus); up !== undefined; up = parentOf.get(up)) ancestors.add(up)
251  }
252  // Above the focus only the path down to it lights, not the ancestors' own tools.
253  const isLit = (link: Link): boolean =>
254    related !== null &&
255    related.has(link.source) &&
256    related.has(link.target) &&
257    !(ancestors.has(link.source) && link.target.startsWith('tool:'))
258  const index = new Map(graph.nodes.map((n, i) => [n.id, i]))
259  const colorOf = new Map(graph.nodes.map(n => [n.id, n.color]))
260  const outAt = new Map<string, number>()
261  const inAt = new Map<string, number>()
262  const sorted = [...graph.links].sort(
263    (a, b) => (boxes.get(a.target)?.y ?? 0) - (boxes.get(b.target)?.y ?? 0),
264  )
265  const name = (id: string) => graph.nodes.find(n => n.id === id)?.label ?? id
266  const ribbons: string[] = []
267  const gradients: string[] = []
268  // Bands leave a bar in the order of where they go, and arrive in the order of
269  // where they came from, so they do not cross at either end.
270  const arrival = new Map<Link, number>()
271  for (const link of [...graph.links].sort((a, b) => middleOf(a.source) - middleOf(b.source))) {
272    if (link.value <= 0 || !boxes.has(link.source)) continue
273    arrival.set(link, inAt.get(link.target) ?? 0)
274    inAt.set(link.target, (inAt.get(link.target) ?? 0) + link.value * k)
275  }
276  for (const link of sorted) {
277    const from = boxes.get(link.source)
278    const to = boxes.get(link.target)
279    if (!from || !to || link.value <= 0) continue
280    const h = link.value * k
281    const y0 = from.y + (outAt.get(link.source) ?? 0)
282    const y1 = to.y + (arrival.get(link) ?? 0)
283    outAt.set(link.source, (outAt.get(link.source) ?? 0) + h)
284    const x0 = from.x + nodeW
285    const x1 = to.x
286    const mid = (x0 + x1) / 2
287    const d = `M${x0},${y0} C${mid},${y0} ${mid},${y1} ${x1},${y1} L${x1},${y1 + h} C${mid},${y1 + h} ${mid},${y0 + h} ${x0},${y0 + h} Z`
288    const lit = related === null ? 'flow' : isLit(link) ? 'flow lit' : 'flow dim'
289    const gradient = `g${ribbons.length}`
290    gradients.push(
291      `<linearGradient id="${gradient}" gradientUnits="userSpaceOnUse" x1="${x0}" y1="0" x2="${x1}" y2="0"><stop offset="0" stop-color="${colorOf.get(link.source)}"/><stop offset="1" stop-color="${colorOf.get(link.target)}"/></linearGradient>`,
292    )
293    ribbons.push(
294      `<path class="${lit} s${index.get(link.source)} t${index.get(link.target)}" d="${d}" fill="url(#${gradient})"><title>${escape(
295        `${name(link.source)} → ${name(link.target)}: ${count(graph, link.value)}`,
296      )}</title></path>`,
297    )
298  }
299
300  // Labels may run over the bands to their right, up to most of the gap to the next column.
301  const labelRoom = Math.max(14, Math.floor((W / Math.max(2, graph.columns - 1)) * 0.85 / 7.6))
302  const rects: string[] = []
303  const labels: string[] = []
304  const hoverRules: string[] = []
305  for (const node of graph.nodes) {
306    const box = boxes.get(node.id)
307    if (!box) continue
308    const i = index.get(node.id)
309    hoverRules.push(`.chart:has(.n${i}:hover) .s${i},.chart:has(.n${i}:hover) .t${i}{fill-opacity:.8}`)
310    const isRunning = ['running', 'pending', 'waiting', 'Active'].includes(node.status ?? '')
311    const isFocus = node.id === focus
312    const isDim = related !== null && !related.has(node.id)
313    const tip = `${node.label}: ${count(graph, node.value)}${node.status && node.id.startsWith('tool:') === false ? ` · ${node.status}` : ''}`
314    const outline = isFocus
315      ? ' stroke="#2B2D33" stroke-width="2"'
316      : isRunning && node.id !== MAIN && node.id !== SELF
317        ? ' stroke="#2B2D33" stroke-width="1.2" stroke-dasharray="3 2"'
318        : ''
319    rects.push(
320      `<rect class="node n${i}${isDim ? ' dim' : ''}" x="${box.x}" y="${box.y}" width="${nodeW}" height="${box.h}" rx="3" fill="${node.color}"${outline}><title>${escape(tip)}</title></rect>`,
321    )
322    const isTool = node.column === lastColumn && node.id.startsWith('tool:')
323    const x = isTool ? box.x - 6 : box.x + nodeW + 6
324    const anchor = isTool ? 'end' : 'start'
325    const text = `${short(node.label, labelRoom)} ${node.value}`
326    labels.push(
327      `<text class="${isDim ? 'dim' : ''}" x="${x}" y="${box.y + box.h / 2}" text-anchor="${anchor}" dominant-baseline="middle"${isFocus ? ' font-weight="800"' : ''}>${escape(text)}</text>`,
328    )
329  }
330
331  const style = [
332    '.flow{fill-opacity:.42;transition:fill-opacity .15s}',
333    '.flow.lit{fill-opacity:.7}.flow.dim{fill-opacity:.08}',
334    '.node{transition:opacity .15s}.node.dim,text.dim{opacity:.35}',
335    '.chart:has(.node:hover) .flow,.chart:has(.flow:hover) .flow{fill-opacity:.1}',
336    '.chart .flow:hover{fill-opacity:.8}',
337    '.node:hover{filter:brightness(.9)}',
338    'text{pointer-events:none}',
339    ...hoverRules,
340  ].join('')
341  const heading = `${graph.rootLabel} · ${graph.agentCount} session${graph.agentCount === 1 ? '' : 's'} started · ${count(graph, graph.total)}`
342  return [
343    `<svg class="chart" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 ${W} ${H}" width="${W}" height="${H}" font-family="system-ui, -apple-system, sans-serif">`,
344    `<style>${style}</style>`,
345    `<defs>${gradients.join('')}</defs>`,
346    `<rect width="${W}" height="${H}" rx="12" fill="#FBFBFD" stroke="#E4E6EB"/>`,
347    `<text x="${W / 2}" y="32" text-anchor="middle" fill="#2B2D33" font-size="18" font-weight="700">${escape(heading)}</text>`,
348    `<g>${ribbons.join('')}</g>`,
349    `<g>${rects.join('')}</g>`,
350    `<g fill="#2B2D33" font-size="14" font-weight="600" stroke="#FBFBFD" stroke-width="3" paint-order="stroke">${labels.join('')}</g>`,
351    `<text x="${W / 2}" y="${H - 12}" text-anchor="middle" fill="#8A8D96" font-size="12">Width = ${graph.unit[1]} · hover a bar or band to trace it · click a session above to open it</text>`,
352    '</svg>',
353  ].join('')
354}
355
356// The same flow as text, for a surface that cannot draw SVG.
357export const textTree = (graph: Graph, columns: number): string[] => {
358  const lines: string[] = []
359  const barRoom = Math.max(6, Math.min(24, columns - 40))
360  const bar = (value: number): string => '█'.repeat(Math.max(value > 0 ? 1 : 0, Math.round((value / Math.max(1, graph.total)) * barRoom)))
361  const byId = new Map(graph.nodes.map(n => [n.id, n]))
362  const kids = (id: string) =>
363    graph.links.filter(l => l.source === id && !l.target.startsWith('tool:')).map(l => l.target)
364  const tools = (id: string) =>
365    graph.links
366      .filter(l => l.source === id && l.target.startsWith('tool:'))
367      .sort((a, b) => b.value - a.value)
368      .map(l => `${l.target.slice(5)} ${l.value}`)
369  const walk = (id: string, indent: string): void => {
370    const node = byId.get(id)
371    if (!node) return
372    lines.push(`${indent}${node.label} ${bar(node.value)} ${node.value}${node.status && id !== MAIN && id !== SELF ? ` (${node.status})` : ''}`)
373    const own = tools(id)
374    if (own.length > 0) lines.push(`${indent}  ${graph.unit[0] === 'turn' ? 'state' : 'tools'}: ${own.join(', ')}`)
375    for (const child of kids(id)) walk(child, `${indent}  `)
376  }
377  walk(MAIN, '')
378  return lines
379}
380
381// What the details line under the diagram says of one agent.
382export const detailsOf = (graph: Graph, id: string): string[] => {
383  const node = graph.nodes.find(n => n.id === id)
384  if (node === undefined) return []
385  const parent = graph.links.find(l => l.target === id && !l.source.startsWith('tool:'))?.source
386  const parentName = graph.nodes.find(n => n.id === parent)?.label
387  const children = graph.links.filter(l => l.source === id && !l.target.startsWith('tool:'))
388  const tools = graph.links
389    .filter(l => l.source === id && l.target.startsWith('tool:'))
390    .sort((a, b) => b.value - a.value)
391    .map(l => `${l.target.slice(5)} ${l.value}`)
392  const head = [node.label, node.status && node.id !== SELF ? node.status : undefined, count(graph, node.value)]
393  return [
394    head.filter(Boolean).join(' · '),
395    ...(parentName !== undefined && parent !== MAIN ? [`Spawned by ${parentName}`] : []),
396    ...(tools.length > 0 ? [`${graph.unit[0] === 'turn' ? 'State' : 'Called'}: ${tools.join(', ')}`] : []),
397    ...(children.length > 0
398      ? [`${graph.unit[0] === 'turn' ? 'Started' : 'Subagents'}: ${children.map(l => `${graph.nodes.find(n => n.id === l.target)?.label ?? l.target} ${l.value}`).join(', ')}`]
399      : []),
400  ]
401}
402

1import type { AgentNode, SessionRecord, Tally } from '../types'
2
3// Reads the desktop app's session records and prints only the fields the
4// diagram needs: titles, who started whom, turns and outcome. Nothing of what
5// was said in a session is read out.
6export const READ_SESSIONS = `
7import json, glob, os, sys
8out = []
9for path in glob.glob(os.path.join(sys.argv[1], '*', '*', 'local_*.json')):
10    try:
11        with open(path) as f:
12            d = json.load(f)
13    except Exception:
14        continue
15    spawned = d.get('spawnedFrom') or {}
16    prs = [p for p in (d.get('prs') or []) if isinstance(p, dict)]
17    pr_state = d.get('prState') or (prs[-1].get('state') if prs else None)
18    out.append({
19        'id': d.get('sessionId'),
20        'title': d.get('title') or 'Untitled session',
21        'parentId': spawned.get('sessionId'),
22        'turns': d.get('completedTurns') or 0,
23        'isArchived': bool(d.get('isArchived')),
24        'prState': pr_state,
25        'lastActivityAt': d.get('lastActivityAt') or 0,
26        'cliSessionIds': [i for i in [d.get('cliSessionId')] + list(d.get('priorCliSessionIds') or []) if i],
27    })
28print(json.dumps(out))
29`
30
31export const sessionsFolder = (home: string): string =>
32  `${home}/Library/Application Support/Claude/claude-code-sessions`
33
34// A link that opens a session in the desktop app.
35export const sessionLink = (id: string): string => `claude://claude.ai/epitaxy/${id}`
36
37const ACTIVE_MS = 10 * 60 * 1000
38
39export const outcomeOf = (session: SessionRecord, now: number): string => {
40  if (session.prState === 'MERGED') return 'PR merged'
41  if (session.prState === 'OPEN') return 'PR open'
42  if (session.prState === 'CLOSED') return 'PR closed'
43  if (session.isArchived) return 'Archived'
44  if (now - session.lastActivityAt < ACTIVE_MS) return 'Active'
45  return 'Idle'
46}
47
48const byId = (records: readonly SessionRecord[]) => new Map(records.map(r => [r.id, r]))
49
50export const rootOf = (records: readonly SessionRecord[], id: string): string => {
51  const all = byId(records)
52  let at = id
53  for (let hops = 0; hops < 50; hops++) {
54    const parent = all.get(at)?.parentId
55    if (parent === null || parent === undefined || !all.has(parent)) return at
56    at = parent
57  }
58  return at
59}
60
61// Sessions that started others, the most recently active family first.
62export const orchestrators = (records: readonly SessionRecord[]): SessionRecord[] => {
63  const all = byId(records)
64  const latest = new Map<string, number>()
65  for (const r of records) {
66    if (!r.parentId || !all.has(r.parentId)) continue
67    const root = rootOf(records, r.id)
68    latest.set(root, Math.max(latest.get(root) ?? 0, r.lastActivityAt, all.get(root)?.lastActivityAt ?? 0))
69  }
70  return [...latest.entries()]
71    .sort((a, b) => b[1] - a[1])
72    .map(([id]) => all.get(id))
73    .filter((r): r is SessionRecord => r !== undefined)
74}
75
76// The family the session running this mod belongs to, by its transcript id.
77export const ownRoot = (records: readonly SessionRecord[], cliSessionId: string): string | null => {
78  const own = records.find(r => r.cliSessionIds.includes(cliSessionId))
79  if (own === undefined) return null
80  const root = rootOf(records, own.id)
81  return records.some(r => r.parentId !== null && rootOf(records, r.id) === root && r.id !== root) ? root : null
82}
83
84// The family as the diagram's input: the root stands as the orchestrator
85// (`main`), each session's turns flow on to the column of outcomes.
86export const familyGraphInput = (
87  records: readonly SessionRecord[],
88  root: string,
89  now: number,
90): { agents: AgentNode[]; tally: Tally; rootTitle: string } => {
91  const family = records.filter(r => rootOf(records, r.id) === root)
92  const key = (id: string): string => (id === root ? 'main' : id)
93  const agents: AgentNode[] = family
94    .filter(r => r.id !== root)
95    .map(r => ({
96      id: r.id,
97      label: r.title,
98      parentId: r.parentId ? key(r.parentId) : 'main',
99      status: outcomeOf(r, now),
100    }))
101  const tally: Tally = {}
102  for (const r of family) {
103    const outcome = r.id === root ? 'Coordinating' : outcomeOf(r, now)
104    tally[key(r.id)] = { [outcome]: Math.max(1, r.turns) }
105  }
106  return { agents, tally, rootTitle: byId(records).get(root)?.title ?? 'Orchestrator' }
107}
108

1// One agent of the diagram; the orchestrator is `main`.
2export type AgentNode = { id: string; label: string; parentId: string | null; status: string }
3
4// Flow by agent id, then by what it flows into (a tool, or a session's state).
5export type Tally = Record<string, Record<string, number>>
6
7// One desktop session as the diagram reads it: no conversation content.
8export type SessionRecord = {
9  id: string
10  title: string
11  parentId: string | null
12  turns: number
13  isArchived: boolean
14  prState: string | null
15  lastActivityAt: number
16  cliSessionIds: string[]
17}
18
19declare module 'claude-code' {
20  interface PluginState {
21    'agent-sankey': {
22      sessions: SessionRecord[]
23      // The orchestrator session whose family is drawn.
24      root: string | null
25      note: string | null
26    }
27  }
28}
29