session-flow

1import { atom, read, update } from 'claude-code'
2import type { Register, TurnUsage } from 'claude-code'
3
4import type { FlowNode, FlowStatus } from '../types'
5
6import { childrenOf, phasesOf, secondsBetween } from './phases'
7
8const PANE = 'session-flow'
9const nodes = atom({ plugin: 'session-flow', key: 'nodes' } as const, [])
10const alias = atom({ plugin: 'session-flow', key: 'alias' } as const, {})
11// ponytail: keeps the last 5000 calls; each write copies the whole list, so move to a StateFamily per turn if long sessions get slow.
12const MAX_NODES = 5000
13const turn = atom({ plugin: 'session-flow', key: 'turn' } as const, '')
14const promptRow = atom({ plugin: 'session-flow', key: 'promptRow' } as const, '')
15
16// The first of these arguments that holds a string is the detail of a tool call.
17const DETAIL_KEYS = ['command', 'description', 'file_path', 'path', 'pattern', 'url', 'query', 'skill', 'prompt']
18const COLOR: Record<FlowStatus, string> = { running: 'warning', done: 'success', error: 'error' }
19// Columns of a phase row: '├─ ' + mark + label + time + gap, so details line up.
20const LABEL_W = 11
21const TIME_W = 7
22const DETAIL_COL = 3 + 2 + LABEL_W + TIME_W + 2
23// A gray at 12% opacity: subtle on a light and on a dark theme. The theme has no neutral background key.
24const HEADER_TINT = '#8080801f'
25
26const oneLine = (text: string, max = 60) => {
27  const flat = text.replace(/\s+/g, ' ').trim()
28  return flat.length > max ? `${flat.slice(0, max - 1)}…` : flat
29}
30
31// The first line of an answer that holds words, without Markdown marks.
32export const firstLine = (text: string) =>
33  oneLine(text.split('\n').map(line => line.replace(/[*_`#>|]/g, '').trim()).find(Boolean) ?? '', 100)
34
35// Kept whole up to 300 characters; the pane cuts it to its width, so a file name survives a long path.
36export const detailOf = (input: object) => {
37  const args = Object.fromEntries(Object.entries(input))
38  const questions = args.questions
39  if (Array.isArray(questions) && questions.length > 0) {
40    const first = Object.fromEntries(Object.entries(Object(questions[0])))
41    if (typeof first.question === 'string') return oneLine(first.question, 300)
42  }
43  const key = DETAIL_KEYS.find(k => typeof args[k] === 'string' && args[k] !== '')
44  return key === undefined ? '' : oneLine(String(args[key]), 300)
45}
46
47// A prompt as the person typed it: a slash command as `/name args`, a background task's notice as its summary.
48export const cleanPrompt = (text: string): { label: string; text: string } => {
49  const tag = (name: string) => new RegExp(`<${name}>([\\s\\S]*?)</${name}>`).exec(text)?.[1]?.trim()
50  if (text.includes('<task-notification>')) {
51    return { label: 'Task', text: tag('summary') ?? tag('status') ?? 'background task update' }
52  }
53  const command = tag('command-name') ?? tag('command-message')
54  if (command !== undefined) {
55    return { label: 'You', text: [`/${command.replace(/^\//, '')}`, tag('command-args')].filter(Boolean).join(' ') }
56  }
57  return { label: 'You', text: text.replace(/<[^>]+>/g, ' ') }
58}
59
60const cut = (text: string, max: number) =>
61  text.length <= max ? text : max <= 1 ? '' : `${text.slice(0, max - 1)}…`
62
63// The widths of the label and time columns. The terminal's fixed-width font lines them up; the desktop
64// draws a proportional font, where padding cannot line up, so it passes none and keeps one space.
65export type Cols = { label: number; time: number }
66
67// Each column as wide as its widest phase row, so the gap after the label is one space.
68export const columnsOf = (rows: readonly Row[]): Cols => {
69  const phases = rows.filter(r => r.mark !== undefined)
70  return {
71    label: Math.max(0, ...phases.map(r => r.label.slice(0, LABEL_W - 1).length)) + 1,
72    time: Math.max(0, ...phases.map(r => (r.seconds ?? '').length)),
73  }
74}
75
76// The label and time a row takes before its detail.
77export const labelText = (row: Row, cols?: Cols) =>
78  row.label === '' ? '' : row.mark === undefined ? `${row.label}  ` : cols ? row.label.slice(0, LABEL_W - 1).padEnd(cols.label) : `${row.label} `
79
80export const timeText = (row: Row, cols?: Cols) =>
81  row.mark === undefined ? '' : `${cols ? (row.seconds ?? '').padStart(cols.time) : row.seconds ?? ''}  `
82
83// Cuts a row's detail, and drops its cost, so the whole row fits in `width` cells and never wraps.
84// With columns, a subagent's note starts under the detail of the phase rows.
85export const fitRow = (input: Row, width: number, cols?: Cols): Row => {
86  const row = cols && input.isDim === true ? { ...input, prefix: `│${' '.repeat(3 + 2 + cols.label + cols.time + 2 - 1)}` } : input
87  const used = row.prefix.length + (row.mark === undefined ? 0 : row.mark.length + 1) + labelText(row, cols).length + timeText(row, cols).length
88  const room = Math.max(0, width - used)
89  const cost = row.cost !== undefined && row.detail.length + row.cost.length + 2 <= room ? row.cost : undefined
90  return { ...row, detail: cut(row.detail, room - (cost === undefined ? 0 : cost.length + 2)), cost }
91}
92
93export type Row = {
94  prefix: string
95  label: string
96  detail: string
97  mark?: string
98  status?: FlowStatus
99  seconds?: string
100  color?: string
101  target?: string
102  cost?: string
103  // A dim line above each prompt: its number, cost and time.
104  isHeader?: boolean
105  isDim?: boolean
106}
107
108// One row per line: each prompt, its phases, a subagent's answer under its row, and Claude's answer last.
109export const rowsOf = (list: readonly FlowNode[], aliases: Record<string, string>): Row[] => {
110  const children = childrenOf(list, aliases)
111  const rows: Row[] = []
112  let count = 0
113  for (const prompt of children.get('') ?? []) {
114    if (rows.length > 0) rows.push({ prefix: '', label: '', detail: '' })
115    const steps = children.get(prompt.id) ?? []
116    const isTask = prompt.label === 'Task'
117    if (!isTask) count += 1
118    const header = [
119      isTask ? 'Background task' : `Prompt ${count}`,
120      prompt.usd === undefined ? '' : `$${prompt.usd.toFixed(2)}`,
121      secondsBetween(prompt.startedAt, prompt.endedAt),
122    ]
123    rows.push({ prefix: '', label: '', detail: header.filter(Boolean).join(' · '), isHeader: true })
124    rows.push({
125      prefix: '', label: isTask ? 'Task' : 'You', detail: prompt.detail,
126      color: isTask ? 'subtle' : 'suggestion', target: prompt.row ?? steps[0]?.id,
127    })
128    for (const phase of phasesOf(steps, children)) {
129      rows.push({ prefix: '├─ ', ...phase })
130      if (phase.note) {
131        rows.push({ prefix: `│${' '.repeat(DETAIL_COL - 1)}`, label: '', detail: `“${phase.note}”`, target: phase.target, isDim: true })
132      }
133    }
134    const answer =
135      prompt.status === 'running' ? 'working…' : prompt.result ?? (prompt.status === 'error' ? 'stopped' : '')
136    rows.push({ prefix: '└─ ', label: 'Claude', detail: answer, color: 'claude', target: prompt.answerRow })
137  }
138  // ponytail: calls of a subagent whose spawn never reported its id stay hidden; add an "orphans" group if that shows up.
139  return rows
140}
141
142const addUsage = (node: FlowNode, usage: TurnUsage | undefined): FlowNode =>
143  usage === undefined
144    ? node
145    : {
146        ...node,
147        tokensIn: (node.tokensIn ?? 0) + usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens,
148        tokensOut: (node.tokensOut ?? 0) + usage.output_tokens,
149      }
150
151const finish = (id: string, status: FlowStatus) => (list: FlowNode[]) =>
152  list.map(n => (n.id === id ? { ...n, status, endedAt: Date.now() } : n))
153
154export const register: Register = on => {
155  on('session.start', async ($, e, next) => {
156    await $.command.register({ name: 'flow', description: 'Show this session as a live flow of prompts, phases and subagents' })
157    void $.ui.open({ id: PANE, title: 'Session flow' }).then(() => $.ui.scroll({ in: PANE, to: 'end' }))
158
159    return next(e)
160  })
161
162  on('command.run', { command: 'flow' }, async $ => {
163    await $.ui.open({ id: PANE, title: 'Session flow' })
164    await $.ui.scroll({ in: PANE, to: 'end' })
165
166    return { text: 'Session flow pane opened.' }
167  })
168
169  on('turn.start', async ($, e, next) => {
170    // Subagent turns may also start here: a main turn is one that starts while no main turn runs.
171    if ((await read($, turn)) === '') {
172      await update($, turn, () => e.turnId)
173      const row = (await read($, promptRow)) || undefined
174      await update($, promptRow, () => '')
175      const prompt = cleanPrompt(e.text)
176      const node: FlowNode = {
177        id: e.turnId, parent: '', kind: 'turn', label: prompt.label,
178        detail: oneLine(prompt.text, 300) || '(no text)', status: 'running', startedAt: Date.now(), row,
179        usdAtStart: (await $.session.usage()).cost?.usd,
180      }
181      await update($, nodes, list => [...list, node].slice(-MAX_NODES))
182    }
183
184    return next(e)
185  })
186
187  on('turn.complete', async ($, e, next) => {
188    const result = firstLine(e.answer)
189    if (e.agentId !== undefined) {
190      const id = (await read($, alias))[`agent:${e.agentId}`]
191      if (id !== undefined) {
192        await update($, nodes, list => list.map(n => (n.id === id ? { ...addUsage(n, e.usage), result } : n)))
193      }
194    } else if ((await read($, turn)) === e.turnId) {
195      const usdNow = (await $.session.usage()).cost?.usd
196      await update($, nodes, list =>
197        finish(e.turnId, e.reason === 'answer' ? 'done' : 'error')(list).map(n =>
198          n.id === e.turnId
199            ? { ...n, result, usd: usdNow === undefined || n.usdAtStart === undefined ? undefined : usdNow - n.usdAtStart }
200            : n,
201        ),
202      )
203      await update($, turn, () => '')
204    }
205
206    return next(e)
207  })
208
209  on('tool.call', async ($, e, next) => {
210    const id = e.tool_use_id ?? crypto.randomUUID()
211    const parent = e.agentId === undefined ? await read($, turn) : `agent:${e.agentId}`
212    const node: FlowNode = {
213      id, parent, kind: e.tool === 'Agent' ? 'agent' : 'tool', label: e.tool,
214      detail: detailOf(e), status: 'running', startedAt: Date.now(),
215    }
216    await update($, nodes, list => [...list, node].slice(-MAX_NODES))
217    const ran = await next(e)
218    await update($, nodes, finish(id, ran.deny !== undefined || ran.isError === true ? 'error' : 'done'))
219
220    return ran
221  })
222
223  // Keeps the transcript row ids a click jumps to: the person's prompt and Claude's text replies on the main loop.
224  on('session.append', async ($, e, next) => {
225    const isPrompt = e.door === 'prompt'
226    const isReply = e.door === 'response' && e.message.content.some(block => block.type === 'text')
227    if (e.agentId === undefined && (isPrompt || isReply)) {
228      const current = await read($, turn)
229      const node = (await read($, nodes)).find(n => n.id === current)
230      if (isReply && node !== undefined) {
231        await update($, nodes, list => list.map(n => (n.id === current ? { ...n, answerRow: e.uuid } : n)))
232      } else if (isPrompt && node !== undefined && node.row === undefined) {
233        await update($, nodes, list => list.map(n => (n.id === current ? { ...n, row: e.uuid } : n)))
234      } else if (isPrompt) {
235        await update($, promptRow, () => e.uuid)
236      }
237    }
238
239    return next(e)
240  })
241
242  on('agent.spawn', async ($, e, next) => {
243    const ran = await next(e)
244    if (ran.agentId !== undefined) {
245      const agentId = ran.agentId
246      await update($, alias, map => ({ ...map, [`agent:${agentId}`]: e.tool_use_id }))
247      await update($, nodes, list =>
248        list.map(n => (n.id === e.tool_use_id ? { ...n, label: `Agent · ${e.subagentType}` } : n)),
249      )
250    }
251
252    return ran
253  })
254
255  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
256    const { Box, Button, Text } = $.ui.resolve(e)
257    // ponytail: the desktop app refuses to scroll the conversation ("transcript not scrollable here"), so only the terminal gets clickable rows; drop this check once a build allows it.
258    const canJump = e.surface === 'terminal'
259    const jump = async (target: string) => {
260      const moved = await $.ui.scroll({ to: { requestId: target }, block: 'start' })
261      if (moved.deny !== undefined) $.ui.toast(`Cannot jump there: ${moved.deny}`)
262    }
263    // ponytail: the desktop wraps a long Text inside a one-line Box, so rows are cut here to the pane's width in cells; the terminal reports the conversation's width, not the pane's, and cuts by itself.
264    const width = e.surface === 'terminal' ? Infinity : (e.viewport?.columns ?? 80) - 4 - (canJump ? 2 : 0)
265    const all = rowsOf(await read($, nodes), await read($, alias))
266    const cols = e.surface === 'terminal' ? columnsOf(all) : undefined
267    const rows = all.map(row => fitRow(row, width, cols))
268
269    return (
270      // In the terminal: one cell clear before the ↗ column, and two at the right edge, where a cut row would touch the pane border.
271      <Box flexDirection="column" paddingLeft={cols ? 1 : 0} paddingRight={cols ? 2 : 0}>
272        {rows.length === 0 && (
273          // Starts where the rows' text starts: after the ↗ column in the terminal.
274          <Box paddingLeft={canJump ? 2 : 0}><Text dimColor>Nothing yet. Send a prompt.</Text></Box>
275        )}
276        {rows.map((row, i) => (
277          // fitRow cuts each row to the pane's width; a row that still overflows wraps to a second line and never covers the next.
278          <Box
279            backgroundColor={row.isHeader === true && e.surface !== 'terminal' ? HEADER_TINT : undefined}
280            // The terminal already draws the empty row above a header as a blank line, so only the desktop gets this margin.
281            marginTop={row.isHeader === true && i > 0 && cols === undefined ? 1 : 0}
282          >
283            {canJump && (row.target !== undefined
284              ? <Button plain dimColor key={`${row.target}:${row.label}`} onPress={() => jump(row.target ?? '')}>↗</Button>
285              : <Text>{' '}</Text>)}
286            {canJump && <Text>{' '}</Text>}
287            <Text wrap="truncate-end">
288              <Text dimColor>{row.prefix}</Text>
289              {row.mark !== undefined && row.status !== undefined && (
290                <Text color={COLOR[row.status]}>{`${row.mark} `}</Text>
291              )}
292              {row.label !== '' && (
293                <Text bold color={row.color}>{labelText(row, cols)}</Text>
294              )}
295              {row.mark !== undefined && <Text dimColor>{timeText(row, cols)}</Text>}
296              <Text dimColor={row.isHeader === true || row.isDim === true} italic={row.isDim}>{row.detail}</Text>
297              {row.cost ? <Text dimColor>{`  ${row.cost}`}</Text> : null}
298            </Text>
299          </Box>
300        ))}
301      </Box>
302    )
303  })
304}
305

1import type { FlowNode, FlowStatus } from '../types'
2
3export type Phase = {
4  label: string
5  // A theme color for the label, by what the phase does.
6  color: string
7  detail: string
8  // '✓', '✗', '◐', or '↻' when a later step of the phase passed after one failed.
9  mark: string
10  status: FlowStatus
11  seconds: string
12  // A subagent's first line of its answer.
13  note?: string
14  // The transcript row a click jumps to: the phase's first tool call.
15  target?: string
16  // A subagent's tokens.
17  cost?: string
18}
19
20const PHASE: [RegExp, string][] = [
21  [/^(Read|Glob|Grep|LS|LSP|NotebookRead)$/, 'Explored'],
22  [/^(WebFetch|WebSearch)$/, 'Researched'],
23  [/^(Edit|Write|MultiEdit|NotebookEdit)$/, 'Edited'],
24  [/^AskUserQuestion$/, 'Asked you'],
25  [/^TodoWrite$/, 'Planned'],
26  [/^(EnterPlanMode|ExitPlanMode)$/, 'Plan'],
27]
28const COLOR: Record<string, string> = {
29  Explored: 'suggestion',
30  Researched: 'planMode',
31  Edited: 'autoAccept',
32  Tested: 'bashBorder',
33  Ran: 'bashBorder',
34  Subagent: 'merged',
35  'Asked you': 'permission',
36}
37// Steps the engine makes for itself, which say nothing about the work.
38const HIDDEN = /^ToolSearch$/
39const SEARCH = /^(Grep|Glob|LSP)$/
40const CHECK = /\b(test|tests|jest|vitest|pytest|tsc|lint|eslint|validate|typecheck)\b/
41// mcp__<server>__<tool>; a server name such as claude_ai_Atlassian shows as its last part.
42const MCP = /^mcp__(.+?)__(.+)$/
43
44const toolOf = (node: FlowNode) => node.label.replace(/ · .*$/, '')
45
46const phaseOf = (node: FlowNode) => {
47  const tool = toolOf(node)
48  if (node.kind === 'agent') return 'Subagent'
49  if (tool === 'Bash') return CHECK.test(node.detail) ? 'Tested' : 'Ran'
50  const mcp = MCP.exec(tool)
51  if (mcp) return mcp[1]?.split('_').pop() ?? tool
52  return PHASE.find(([re]) => re.test(tool))?.[1] ?? tool
53}
54
55// What a step did, in a few words: an MCP step leads with its tool name.
56const stepText = (node: FlowNode) => {
57  const mcp = MCP.exec(toolOf(node))
58  return mcp ? [mcp[2], node.detail].filter(Boolean).join(' · ') : node.detail
59}
60
61const baseName = (path: string) => path.split('/').filter(Boolean).pop() ?? path
62
63const plural = (n: number, word: string) => `${n} ${word}${n === 1 ? '' : 's'}`
64
65const names = (paths: readonly string[]) => {
66  const unique = [...new Set(paths.map(baseName))]
67  return unique.length <= 3 ? unique.join(', ') : plural(unique.length, 'file')
68}
69
70export const tokens = (n: number) =>
71  n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${Math.round(n / 1e3)}k` : `${n}`
72
73export const duration = (ms: number) => {
74  const s = ms / 1000
75  return s < 60 ? `${s.toFixed(1)}s` : `${Math.floor(s / 60)}m ${Math.round(s % 60)}s`
76}
77
78export const secondsBetween = (start: number, end: number | undefined) =>
79  end === undefined ? '' : duration(end - start)
80
81// Groups the nodes by parent; a subagent's calls go under the Agent call that started it.
82export const childrenOf = (list: readonly FlowNode[], aliases: Record<string, string>) => {
83  const children = new Map<string, FlowNode[]>()
84  for (const node of list) {
85    const parent = aliases[node.parent] ?? node.parent
86    children.set(parent, [...(children.get(parent) ?? []), node])
87  }
88  return children
89}
90
91// The count leads, so a cut at the pane's edge never hides it.
92const describe = (label: string, steps: readonly FlowNode[]) => {
93  const times = steps.length > 1 ? `${steps.length}× ` : ''
94  if (label === 'Explored') {
95    const files = steps.filter(s => !SEARCH.test(toolOf(s))).map(s => s.detail)
96    const searches = steps.length - files.length
97    return [files.length > 0 ? names(files) : '', searches > 0 ? plural(searches, 'search') : '']
98      .filter(Boolean)
99      .join(', ')
100  }
101  if (label === 'Edited') {
102    const files = steps.map(s => s.detail)
103    return `${steps.length > new Set(files).size ? times : ''}${names(files)}`
104  }
105  const last = steps[steps.length - 1]
106  return `${times}${last === undefined ? '' : stepText(last)}`
107}
108
109// The built-in tools a phase used, as `Read, Grep`. Empty for an MCP phase, whose detail already leads with its tool, and where the label is the tool.
110const toolsOf = (label: string, steps: readonly FlowNode[]) => {
111  const tools = [...new Set(steps.map(toolOf))].join(', ')
112  return tools === label || steps.some(s => MCP.test(toolOf(s))) ? '' : tools
113}
114
115const markOf = (steps: readonly FlowNode[]): [string, FlowStatus] => {
116  if (steps.some(s => s.status === 'running')) return ['◐', 'running']
117  if (steps[steps.length - 1]?.status === 'error') return ['✗', 'error']
118  return steps.some(s => s.status === 'error') ? ['↻', 'done'] : ['✓', 'done']
119}
120
121// Turns the steps of one prompt into phases: runs of steps of one kind merge into one row; each subagent keeps a row.
122export const phasesOf = (steps: readonly FlowNode[], children: Map<string, FlowNode[]>): Phase[] => {
123  const runs: { label: string; steps: FlowNode[] }[] = []
124  for (const step of steps) {
125    if (HIDDEN.test(toolOf(step))) continue
126    const label = phaseOf(step)
127    const run = runs[runs.length - 1]
128    if (run !== undefined && run.label === label && label !== 'Subagent') run.steps.push(step)
129    else runs.push({ label, steps: [step] })
130  }
131
132  return runs.map(({ label, steps: run }) => {
133    const [mark, status] = markOf(run)
134    const start = Math.min(...run.map(s => s.startedAt))
135    const end = status === 'running' ? undefined : Math.max(...run.map(s => s.endedAt ?? s.startedAt))
136    const first = run[0]
137    const color = COLOR[label] ?? (MCP.test(first === undefined ? '' : toolOf(first)) ? 'ide' : 'text')
138    const base = { label, color, mark, status, seconds: secondsBetween(start, end), target: first?.id }
139    if (label === 'Subagent' && first !== undefined) {
140      const count = children.get(first.id)?.length ?? 0
141      return {
142        ...base,
143        detail: `${first.detail}${count > 0 ? ` · ${plural(count, 'step')}` : ''}`,
144        note: first.result,
145        cost: first.tokensIn === undefined ? undefined : `${tokens(first.tokensIn)} in · ${tokens(first.tokensOut ?? 0)} out`,
146      }
147    }
148    return { ...base, detail: [toolsOf(label, run), describe(label, run)].filter(Boolean).join(' · ') }
149  })
150}
151

1export type FlowStatus = 'running' | 'done' | 'error'
2
3export type FlowNode = {
4  id: string
5  // A node id, or `agent:<agentId>` for a call made inside a subagent.
6  parent: string
7  kind: 'turn' | 'tool' | 'agent'
8  label: string
9  detail: string
10  status: FlowStatus
11  startedAt: number
12  endedAt?: number
13  // The first line of the answer: Claude's for a turn, the subagent's for an Agent call.
14  result?: string
15  // Transcript row ids of a turn: the prompt, and Claude's last text reply.
16  row?: string
17  answerRow?: string
18  // A turn: the session's dollar total when it started, and what the turn cost, subagents included.
19  usdAtStart?: number
20  usd?: number
21  // A subagent: its tokens, summed over its turns. Input counts the cache reads and writes too.
22  tokensIn?: number
23  tokensOut?: number
24}
25
26declare module 'claude-code' {
27  interface PluginState {
28    'session-flow': {
29      nodes: FlowNode[]
30      // `agent:<agentId>` -> the id of the Agent tool call that started it.
31      alias: Record<string, string>
32      // The id of the main turn that runs now, '' between turns.
33      turn: string
34      // The id of a prompt row stored before its turn started, '' when none waits.
35      promptRow: string
36    }
37  }
38}
39