trace-waterfall

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, TurnStepResult } from 'claude-code'
3
4import type { AgentLink, Theme, Trace } from '../types'
5import { chartOf, INK_COLOR, NOW_COLOR, PALETTES } from './draw'
6import type { Palette, Run } from './draw'
7import {
8  AGENT_PREFIX,
9  addSpan,
10  changeSpan,
11  closeLoop,
12  closeSpan,
13  completeMain,
14  continueRequest,
15  costOf,
16  describeTool,
17  figuresOf,
18  findSpan,
19  formatCost,
20  formatDuration,
21  formatTokens,
22  isLive,
23  latestTrace,
24  modelLabel,
25  placeOf,
26  REQUEST_PREFIX,
27  startRequest,
28  subtreeOf,
29} from './trace'
30import type { Figures, NewSpan } from './trace'
31
32const PLUGIN = 'trace-waterfall'
33const PANE = 'trace'
34const COMMAND = 'trace-waterfall'
35const TITLE = 'Trace'
36const OPEN_TOOL = 'open_pane'
37const AGENT_TOOLS = /^(Agent|Task|Workflow)$/
38const OTHER_TOOLS = /^(?!(Agent|Task|Workflow)$)/
39
40const traces = atom({ plugin: 'trace-waterfall', key: 'traces' } as const, [])
41const agents = atom({ plugin: 'trace-waterfall', key: 'agents' } as const, {})
42const main = atom({ plugin: 'trace-waterfall', key: 'main' } as const, null)
43const view = atom({ plugin: 'trace-waterfall', key: 'view' } as const, null)
44const selected = atom({ plugin: 'trace-waterfall', key: 'selected' } as const, null)
45const now = atom({ plugin: 'trace-waterfall', key: 'now' } as const, 0)
46const theme = atom({ plugin: 'trace-waterfall', key: 'theme' } as const, 'dark')
47
48const REQUEST_ORIGINS = new Set(['composer', 'bridge', 'sdk', 'scheduled-trigger'])
49
50type Arrival = { isRequest: boolean; text: string; detail: string; agentId: string | null }
51
52let queue: Promise<unknown> = Promise.resolve()
53let upcoming: Arrival | null = null
54const reported = new Set<string>()
55
56function arrivalOf(kind: string | null, text: string): Arrival {
57  const agentId = /<agent-message from="([^"]+)"/.exec(text)?.[1] ?? null
58
59  if (kind === null ? !text.trimStart().startsWith('<') : REQUEST_ORIGINS.has(kind)) {
60    return { isRequest: true, text, detail: '', agentId: null }
61  }
62
63  if (agentId !== null || kind === 'peer' || kind === 'peer-send-message' || kind === 'coordinator') {
64    return { isRequest: false, text, detail: 'after an agent message', agentId }
65  }
66
67  if (kind === 'task-notification') {
68    return { isRequest: false, text, detail: 'after a background task', agentId }
69  }
70
71  return { isRequest: false, text, detail: 'without a typed prompt', agentId }
72}
73
74function reasonOf(error: unknown): string {
75  return error instanceof Error ? error.message : String(error)
76}
77
78function inOrder($: EngineInterface, label: string, job: () => Promise<unknown>): Promise<void> {
79  const run = queue.then(job).then(
80    () => undefined,
81    error => {
82      if (!reported.has(label)) {
83        reported.add(label)
84        $.ui.log(`${PLUGIN}: ${label} failed: ${reasonOf(error).replace(/\s+/g, ' ').slice(0, 300)}`)
85      }
86    },
87  )
88  queue = run
89
90  return run
91}
92
93function stamp($: EngineInterface): Promise<number> {
94  return $.clock.now().catch(() => Date.now())
95}
96
97function toTheme(value: unknown): Theme {
98  return String(value).startsWith('light') ? 'light' : 'dark'
99}
100
101async function refreshTheme($: EngineInterface): Promise<void> {
102  const row = (await $.config.list()).find(one => one.key === 'theme')
103  const wanted = toTheme(row?.value)
104
105  if ((await read($, theme)) !== wanted) {
106    await update($, theme, () => wanted)
107  }
108}
109
110async function startMainTurn($: EngineInterface, turnId: string, text: string, arrival: Arrival, at: number): Promise<void> {
111  const held = await read($, traces)
112  const links = await read($, agents)
113  const latest = latestTrace(held)
114
115  if (arrival.isRequest || latest === undefined) {
116    await update($, traces, current => startRequest(current, turnId, arrival.text || text, at))
117    await update($, main, () => ({ trace: REQUEST_PREFIX + turnId, span: turnId }))
118
119    return
120  }
121
122  const owner = arrival.agentId === null ? undefined : links[arrival.agentId]?.trace
123  const traceId = owner !== undefined && held.some(trace => trace.id === owner) ? owner : latest.id
124
125  await update($, traces, current => continueRequest(current, traceId, turnId, arrival.detail, at))
126  await update($, main, () => ({ trace: traceId, span: turnId }))
127}
128
129async function openSpan($: EngineInterface, loop: string | null, fields: NewSpan): Promise<void> {
130  const links = await read($, agents)
131  const current = await read($, main)
132
133  await update($, traces, held => {
134    const place = placeOf(held, links, current, loop)
135
136    return place === undefined ? held : addSpan(held, place, fields, loop)
137  })
138}
139
140async function endSpan($: EngineInterface, id: string, at: number, isError: boolean): Promise<void> {
141  await update($, traces, held => changeSpan(held, id, span => closeSpan(span, at, isError)))
142}
143
144async function endStep($: EngineInterface, id: string, at: number, result: TurnStepResult | undefined): Promise<void> {
145  await update($, traces, held =>
146    changeSpan(held, id, span => {
147      const closed = closeSpan(span, at, false)
148      const usage = result?.usage
149
150      if (usage === null || usage === undefined) {
151        return closed
152      }
153
154      const counts = {
155        input: usage.input_tokens,
156        output: usage.output_tokens,
157        cacheRead: usage.cache_read_input_tokens,
158        cacheWrite: usage.cache_creation_input_tokens,
159      }
160
161      return {
162        ...closed,
163        model: usage.model,
164        details: [modelLabel(usage.model)],
165        usage: counts,
166        cost: costOf(usage.model, counts),
167      }
168    }),
169  )
170}
171
172function linkedSpans(links: Record<string, AgentLink>): Set<string> {
173  return new Set(Object.values(links).map(link => link.span))
174}
175
176async function startAgentTurn($: EngineInterface, agentId: string, prompt: string, at: number): Promise<void> {
177  const links = await read($, agents)
178  let link = links[agentId]
179
180  if (link === undefined) {
181    const held = await read($, traces)
182    const taken = linkedSpans(links)
183
184    for (let index = held.length - 1; index >= 0 && link === undefined; index -= 1) {
185      const trace = held[index]
186      const match = trace?.spans.find(
187        span => span.kind === 'agent' && span.prompt !== null && span.prompt === prompt && !taken.has(span.id),
188      )
189
190      if (trace !== undefined && match !== undefined) {
191        link = { trace: trace.id, span: match.id }
192      }
193    }
194  }
195
196  if (link === undefined) {
197    const place = placeOf(await read($, traces), links, await read($, main), null)
198
199    if (place === undefined) {
200      return
201    }
202
203    const id = AGENT_PREFIX + agentId
204    link = { trace: place.trace, span: id }
205    await update($, traces, current =>
206      addSpan(current, place, { id, kind: 'agent', name: 'Agent', details: [], start: at, agentId, pending: 0 }, null),
207    )
208  }
209
210  const found = link
211
212  await update($, agents, current => ({ ...current, [agentId]: found }))
213  await update($, traces, held =>
214    changeSpan(held, found.span, span => ({ ...span, agentId, pending: span.pending + 1, end: Math.max(span.end, at) })),
215  )
216}
217
218async function linkAgent(
219  $: EngineInterface,
220  agentId: string,
221  toolUseId: string,
222  description: string,
223  subagentType: string,
224  model: string,
225  at: number,
226): Promise<void> {
227  if ((await read($, agents))[agentId] !== undefined) {
228    return
229  }
230
231  const held = await read($, traces)
232  const found = findSpan(held, toolUseId)
233
234  if (found === undefined) {
235    return
236  }
237
238  if (found.span.kind === 'agent' && found.span.agentId === null) {
239    await update($, agents, current => ({ ...current, [agentId]: { trace: found.trace.id, span: found.span.id } }))
240    await update($, traces, current => changeSpan(current, found.span.id, span => ({ ...span, agentId, model })))
241
242    return
243  }
244
245  const id = AGENT_PREFIX + agentId
246  const details = [description, subagentType].filter(text => text !== '')
247
248  await update($, traces, current =>
249    addSpan(current, { trace: found.trace.id, span: found.span.id }, { id, kind: 'agent', name: 'Agent', details, start: at, model, agentId, pending: 0 }, null),
250  )
251  await update($, agents, current => ({ ...current, [agentId]: { trace: found.trace.id, span: id } }))
252}
253
254async function completeAgent($: EngineInterface, agentId: string, at: number): Promise<void> {
255  const link = (await read($, agents))[agentId]
256
257  if (link === undefined) {
258    return
259  }
260
261  await update($, traces, held =>
262    changeSpan(
263      held.map(trace => (trace.id === link.trace ? closeLoop(trace, agentId, at) : trace)),
264      link.span,
265      span => closeSpan(span, at, false),
266    ),
267  )
268}
269
270async function tick($: EngineInterface): Promise<void> {
271  if (!(await read($, traces)).some(isLive)) {
272    return
273  }
274
275  const at = await stamp($)
276  await update($, now, () => at)
277}
278
279function followedIndex(list: Trace[], current: AgentLink | null): number {
280  const found = current === null ? -1 : list.findIndex(trace => trace.id === current.trace)
281
282  return found < 0 ? list.length - 1 : found
283}
284
285async function moveView($: EngineInterface, step: number): Promise<void> {
286  const list = await read($, traces)
287  const followed = followedIndex(list, await read($, main))
288
289  await update($, view, current => {
290    const found = current === null ? followed : list.findIndex(trace => trace.id === current)
291    const target = Math.min(list.length - 1, Math.max(0, (found < 0 ? followed : found) + step))
292
293    return target === followed ? null : (list[target]?.id ?? null)
294  })
295  await update($, selected, () => null)
296}
297
298function promptLine(trace: Trace): string {
299  const line = trace.prompt
300    .split(/\r?\n/)
301    .map(one => one.trim())
302    .find(one => one !== '')
303
304  return line === undefined || line.startsWith('<') ? 'A request without a typed prompt' : line
305}
306
307type Item = { label: string; value: string; width: number; swatch?: string }
308
309function itemsOf(figures: Figures, palette: Palette): Item[] {
310  return [
311    { label: 'Duration', value: formatDuration(figures.duration), width: 7 },
312    { label: 'Cost', value: formatCost(figures.cost) || 'none', width: 6 },
313    { label: 'Tokens in', value: formatTokens(figures.tokensIn), width: 6 },
314    { label: 'out', value: formatTokens(figures.tokensOut), width: 6 },
315    { label: 'Model calls', value: String(figures.modelCalls), width: 4, swatch: palette.model },
316    { label: 'Tool calls', value: String(figures.toolCalls), width: 4, swatch: palette.tool },
317    { label: 'Errors', value: String(figures.errors), width: 3, swatch: palette.error },
318  ]
319}
320
321export const register: Register = (on, options) => {
322  const opensAtStart = options?.start !== 'off'
323
324  on('session.start', async ($, e, next) => {
325    await $.command.register({
326      name: COMMAND,
327      description: 'Show where the time of each request goes, as a waterfall in a side pane',
328    })
329    await $.tool.register({
330      name: OPEN_TOOL,
331      description:
332        'Opens the Trace pane beside the chat: a live waterfall of where the time of each request goes. Call it when the person asks to see, show or reopen the trace, the waterfall or the Gantt pane.',
333    })
334    await refreshTheme($).catch(() => undefined)
335    $.clock.every(1000, () => {
336      void tick($).catch(() => undefined)
337    })
338    if (opensAtStart) void $.ui.open({ id: PANE, title: TITLE })
339
340    return next(e)
341  })
342
343  on('command.run', { command: COMMAND }, async $ => {
344    const opened = await $.ui.open({ id: PANE, title: TITLE })
345
346    return { text: opened.isPlaced ? 'Trace pane opened.' : `The trace pane waits: ${opened.reason}` }
347  })
348
349  on('config.set', { key: 'theme' }, async ($, e, next) => {
350    const changed = await next(e)
351
352    if (changed.deny === undefined) {
353      void inOrder($, 'storing the theme', () => update($, theme, () => toTheme(changed.value)))
354    }
355
356    return changed
357  })
358
359  on('session.end', async ($, e, next) => {
360    if (e.reason === 'clear') {
361      await inOrder($, 'clearing the traces', async () => {
362        await update($, traces, () => [])
363        await update($, agents, () => ({}))
364        await update($, main, () => null)
365        await update($, view, () => null)
366        await update($, selected, () => null)
367      })
368    }
369
370    return next(e)
371  })
372
373  on('prompt.submit', ($, e, next) => {
374    if (e.turnId === undefined) {
375      const typed = e.command === undefined ? e.text : `/${e.command.name} ${e.command.args}`.trim()
376      upcoming = arrivalOf(e.origin.kind, typed)
377    }
378
379    return next(e)
380  })
381
382  on('turn.start', async ($, e, next) => {
383    const at = await stamp($)
384    const agentId = e.agentId
385
386    if (agentId === undefined) {
387      const arrival = upcoming ?? arrivalOf(null, e.text)
388      upcoming = null
389      void inOrder($, 'starting a turn', () => startMainTurn($, e.turnId, e.text, arrival, at))
390      void inOrder($, 'reading the theme', () => refreshTheme($))
391    } else {
392      void inOrder($, 'starting an agent turn', () => startAgentTurn($, agentId, e.text, at))
393    }
394
395    return next(e)
396  })
397
398  on('turn.step', async function* ($, e, next) {
399    const at = await stamp($)
400    const id = `step:${e.turnId}:${e.index}`
401    const loop = e.agentId ?? null
402    let result: TurnStepResult | undefined
403
404    void inOrder($, 'starting a model call', () =>
405      openSpan($, loop, { id, kind: 'model', name: 'Model call', details: [modelLabel(e.model)], start: at, model: e.model }),
406    )
407
408    try {
409      result = yield* next(e)
410
411      return result
412    } finally {
413      const end = await stamp($)
414      void inOrder($, 'ending a model call', () => endStep($, id, end, result))
415    }
416  })
417
418  on('turn.complete', async ($, e, next) => {
419    const at = await stamp($)
420    const agentId = e.agentId
421
422    if (agentId === undefined) {
423      const isError = e.reason === 'error' || e.reason === 'refusal'
424      void inOrder($, 'ending a turn', () => update($, traces, held => completeMain(held, e.turnId, at, isError)))
425    } else {
426      void inOrder($, 'ending an agent turn', () => completeAgent($, agentId, at))
427    }
428
429    return next(e)
430  })
431
432  on('tool.call', { tool: /^mcp__trace-waterfall__open_pane$/ }, async $ => {
433    const opened = await $.ui.open({ id: PANE, title: TITLE })
434
435    return { result: opened.isPlaced ? 'The Trace pane is open.' : `The Trace pane waits: ${opened.reason}` }
436  })
437
438  on('tool.call', { tool: OTHER_TOOLS }, async ($, e, next) => {
439    const at = await stamp($)
440    const id = e.tool_use_id ?? `tool:${String(e.tool)}:${at}`
441    const { name, details } = describeTool(String(e.tool), e as unknown as Record<string, unknown>)
442    let isError = true
443
444    void inOrder($, 'starting a tool call', () => openSpan($, e.agentId ?? null, { id, kind: 'tool', name, details, start: at }))
445
446    try {
447      const ran = await next(e)
448      isError = ran.deny !== undefined || ran.isError === true
449
450      return ran
451    } finally {
452      const end = await stamp($)
453      void inOrder($, 'ending a tool call', () => endSpan($, id, end, isError))
454    }
455  })
456
457  on('tool.call', { tool: AGENT_TOOLS }, async ($, e, next) => {
458    const at = await stamp($)
459    const input = e as unknown as Record<string, unknown>
460    const id = e.tool_use_id ?? `tool:${String(e.tool)}:${at}`
461    const { name, details } = describeTool(String(e.tool), input)
462    const prompt = typeof input.prompt === 'string' ? input.prompt : null
463    const kind = e.tool === 'Workflow' ? 'tool' : 'agent'
464    let isError = true
465
466    void inOrder($, 'starting an agent call', () => openSpan($, e.agentId ?? null, { id, kind, name, details, start: at, prompt }))
467
468    try {
469      const ran = await next(e)
470      isError = ran.deny !== undefined || ran.isError === true
471
472      return ran
473    } finally {
474      const end = await stamp($)
475      void inOrder($, 'ending an agent call', () => endSpan($, id, end, isError))
476    }
477  })
478
479  on('agent.spawn', async ($, e, next) => {
480    const spawned = await next(e)
481
482    if (spawned.deny === undefined && spawned.agentId !== undefined) {
483      const agentId = spawned.agentId
484      const at = await stamp($)
485
486      void inOrder($, 'linking an agent', () =>
487        linkAgent($, agentId, e.tool_use_id, e.description, e.subagentType, spawned.model, at),
488      )
489    }
490
491    return spawned
492  })
493
494  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
495    const { Box, Button, Text } = $.ui.resolve(e)
496
497    const textOf = (run: Run) => (
498      <Text color={run.color} inverse={run.inverse} dimColor={run.dim} bold={run.bold}>
499        {run.text}
500      </Text>
501    )
502
503    try {
504      const held = await read($, traces)
505
506      if (held.length === 0) {
507        return <Text dimColor>No request has run yet in this session.</Text>
508      }
509
510      const viewed = await read($, view)
511      const found = viewed === null ? -1 : held.findIndex(trace => trace.id === viewed)
512      const index = found < 0 ? followedIndex(held, await read($, main)) : found
513      const trace = held[index]
514
515      if (trace === undefined) {
516        return <Text dimColor>No request has run yet in this session.</Text>
517      }
518
519      const live = isLive(trace)
520      const at = live ? await read($, now) : 0
521      const palette = PALETTES[await read($, theme)]
522      const chosen = await read($, selected)
523      const selectedId = chosen !== null && trace.spans.some(span => span.id === chosen) ? chosen : null
524      const scope = selectedId === null ? trace.spans : subtreeOf(trace, selectedId)
525      const figures = figuresOf(scope, at)
526      const chart = chartOf(trace, at, Math.max(24, e.props.bodyColumns), palette, selectedId)
527      const isFirst = index === 0
528      const isLast = index === held.length - 1
529
530      return (
531        <Box flexDirection="column">
532          <Box flexDirection="row" justifyContent="space-between">
533            <Text bold>
534              {`Request ${index + 1} of ${held.length}`}
535              {live ? <Text color={NOW_COLOR}> running</Text> : ''}
536            </Text>
537            <Box flexDirection="row" columnGap={1}>
538              <Button
539                key="earlier"
540                label="◀ Earlier"
541                dimColor={isFirst || undefined}
542                onPress={() => (isFirst ? undefined : moveView($, -1))}
543              />
544              <Button
545                key="later"
546                label="Later ▶"
547                dimColor={isLast || undefined}
548                onPress={() => (isLast ? undefined : moveView($, 1))}
549              />
550            </Box>
551          </Box>
552          <Text dimColor>{promptLine(trace)}</Text>
553          <Box flexDirection="row" flexWrap="wrap" columnGap={3}>
554            <Text bold>{(selectedId === null ? 'Whole request' : 'Selected').padEnd(13)}</Text>
555            {itemsOf(figures, palette).map(item => (
556              <Text>
557                {item.swatch === undefined ? '' : <Text color={item.swatch}>■ </Text>}
558                <Text dimColor>{`${item.label} `}</Text>
559                {item.value.padStart(item.width)}
560              </Text>
561            ))}
562          </Box>
563          <Text> </Text>
564          <Text>{chart.axis.map(textOf)}</Text>
565          {chart.rows.map(row => (
566            <Button
567              key={row.id}
568              plain
569              onPress={() => update($, selected, current => (current === row.id ? null : row.id))}
570            >
571              {row.runs.map(textOf)}
572            </Button>
573          ))}
574        </Box>
575      )
576    } catch (error) {
577      return <Text color={INK_COLOR}>{`The trace could not be drawn: ${reasonOf(error).replace(/\s+/g, ' ').slice(0, 300)}`}</Text>
578    }
579  })
580}
581

1import type { Span, Theme, Trace } from '../types'
2import {
3  barCells,
4  displayEnd,
5  domainOf,
6  formatCost,
7  formatDuration,
8  formatTick,
9  isLive,
10  rowsOf,
11  ticksOf,
12} from './trace'
13
14export type Run = { text: string; color?: string; inverse?: boolean; dim?: boolean; bold?: boolean }
15
16export type ChartRow = { id: string; runs: Run[] }
17
18export type Chart = { axis: Run[]; rows: ChartRow[] }
19
20export type Palette = { model: string; tool: string; error: string; neutral: string }
21
22export const PALETTES: Record<Theme, Palette> = {
23  dark: { model: '#8B7BFF', tool: '#1FA2B8', error: '#E5504A', neutral: '#7A7A78' },
24  light: { model: '#6650E0', tool: '#0086A3', error: '#C8352F', neutral: '#858583' },
25}
26
27export const GRID_COLOR = 'subtle'
28export const NOW_COLOR = 'suggestion'
29export const INK_COLOR = 'text'
30
31const GUTTER = 2
32const DURATION_WIDTH = 7
33const COST_WIDTH = 6
34const MIN_TIMELINE = 8
35const MIN_LABEL = 12
36const MAX_LABEL = 34
37const INDENT = 2
38const MAX_INDENT_DEPTH = 4
39const COST_FROM_WIDTH = 56
40
41export type Layout = { label: number; timeline: number; hasCost: boolean }
42
43export function indentOf(depth: number): number {
44  return Math.min(depth, MAX_INDENT_DEPTH) * INDENT
45}
46
47export function layoutOf(width: number, nameWidth: number): Layout {
48  const hasCost = width >= COST_FROM_WIDTH
49  const numbers = DURATION_WIDTH + (hasCost ? 2 + COST_WIDTH : 0)
50  const fixed = GUTTER + 1 + 1 + numbers
51  const wanted = Math.max(MIN_LABEL, nameWidth, Math.min(Math.round(width * 0.36), MAX_LABEL))
52  const label = Math.min(wanted, width - fixed - MIN_TIMELINE)
53
54  if (label < nameWidth) {
55    return { label: Math.max(nameWidth, width - fixed), timeline: 0, hasCost }
56  }
57
58  return { label, timeline: width - fixed - label, hasCost }
59}
60
61export function colorOf(span: Span, palette: Palette): string {
62  if (span.isError) {
63    return palette.error
64  }
65
66  if (span.kind === 'model') {
67    return palette.model
68  }
69
70  return span.kind === 'tool' ? palette.tool : palette.neutral
71}
72
73function sameStyle(a: Run, b: Run): boolean {
74  return a.color === b.color && a.inverse === b.inverse && a.dim === b.dim && a.bold === b.bold
75}
76
77export function merged(runs: Run[]): Run[] {
78  const out: Run[] = []
79
80  for (const run of runs) {
81    const last = out[out.length - 1]
82
83    if (run.text === '') {
84      continue
85    }
86
87    if (last !== undefined && sameStyle(last, run)) {
88      out[out.length - 1] = { ...last, text: last.text + run.text }
89    } else {
90      out.push({ ...run })
91    }
92  }
93
94  return out
95}
96
97function labelRuns(span: Span, depth: number, width: number, isSelected: boolean): Run[] {
98  const indent = ' '.repeat(indentOf(depth))
99  const used = indent.length + span.name.length
100  const detail = span.details.find(text => used + 1 + text.length <= width)
101  const pad = width - used - (detail === undefined ? 0 : detail.length + 1)
102
103  return [
104    { text: indent },
105    { text: span.name, bold: isSelected || undefined },
106    { text: detail === undefined ? '' : ` ${detail}`, dim: true },
107    { text: ' '.repeat(Math.max(0, pad)) },
108  ]
109}
110
111function slowestLeaf(spans: Span[], now: number): string | null {
112  let best: Span | null = null
113
114  for (const span of spans) {
115    if (span.kind !== 'model' && span.kind !== 'tool') {
116      continue
117    }
118
119    if (best === null || displayEnd(span, now) - span.start > displayEnd(best, now) - best.start) {
120      best = span
121    }
122  }
123
124  return best === null ? null : best.id
125}
126
127export function rowCosts(trace: Trace): Map<string, number> {
128  const parents = new Map(trace.spans.map(span => [span.id, span.id === trace.id ? null : span.parent]))
129  const sums = new Map<string, number>()
130
131  for (const span of trace.spans) {
132    if (span.cost === null) {
133      continue
134    }
135
136    let id: string | null | undefined = span.id
137    const seen = new Set<string>()
138
139    while (id !== null && id !== undefined && !seen.has(id)) {
140      seen.add(id)
141      sums.set(id, (sums.get(id) ?? 0) + span.cost)
142      id = parents.get(id)
143    }
144  }
145
146  return sums
147}
148
149export function nameWidthOf(trace: Trace): number {
150  return rowsOf(trace).reduce((widest, row) => Math.max(widest, indentOf(row.depth) + row.span.name.length), 0)
151}
152
153export function chartOf(trace: Trace, now: number, width: number, palette: Palette, selected: string | null): Chart {
154  const rows = rowsOf(trace)
155  const layout = layoutOf(width, nameWidthOf(trace))
156  const live = isLive(trace)
157  const last = Math.max(trace.start, ...trace.spans.map(span => displayEnd(span, now)))
158  const extent = Math.max(1000, last - trace.start)
159  const domain = domainOf(extent)
160  const cells = layout.timeline
161  const ticks = ticksOf(domain, cells)
162  const column = (at: number): number => Math.floor(((at - trace.start) / domain) * cells)
163  const eighth = (at: number): number => Math.round(((at - trace.start) / domain) * cells * 8)
164  const gridColumns = new Set(ticks.filter(tick => tick > 0).map(tick => Math.floor(((tick * 1000) / domain) * cells)))
165  const nowColumn = live ? column(now) : -1
166  const slowest = slowestLeaf(trace.spans, now)
167  const costs = rowCosts(trace)
168
169  const axisLine = Array.from({ length: cells }, () => ' ')
170  let free = 0
171
172  for (const tick of ticks) {
173    const at = Math.floor(((tick * 1000) / domain) * cells)
174    const text = formatTick(tick)
175
176    if (at >= free && at + text.length <= cells) {
177      text.split('').forEach((char, index) => {
178        axisLine[at + index] = char
179      })
180      free = at + text.length + 1
181    }
182  }
183
184  const heads = 'time'.padStart(DURATION_WIDTH) + (layout.hasCost ? '  ' + 'cost'.padStart(COST_WIDTH) : '')
185  const axis = merged([
186    { text: ' '.repeat(GUTTER + layout.label + 1) },
187    { text: cells > 0 ? axisLine.join('') + ' ' : '', dim: true },
188    { text: heads, dim: true },
189  ])
190
191  const chartRows = rows.map(({ span, depth }) => {
192    const isSelected = span.id === selected
193    const color = colorOf(span, palette)
194    const end = span.id === trace.id ? last : displayEnd(span, now)
195    const runs: Run[] = isSelected ? [{ text: '▌', color: INK_COLOR }, { text: ' ' }] : [{ text: '  ' }]
196
197    runs.push(...labelRuns(span, depth, layout.label, isSelected), { text: ' ' })
198
199    if (cells > 0) {
200      barCells(eighth(span.start), eighth(end), cells).forEach((cell, index) => {
201        if (cell.mode === 'solid') {
202          runs.push({ text: cell.glyph, color })
203        } else if (cell.mode === 'inverse') {
204          runs.push({ text: cell.glyph, color, inverse: true })
205        } else if (index === nowColumn) {
206          runs.push({ text: '│', color: NOW_COLOR })
207        } else if (gridColumns.has(index)) {
208          runs.push({ text: '┊', color: GRID_COLOR })
209        } else {
210          runs.push({ text: ' ' })
211        }
212      })
213      runs.push({ text: ' ' })
214    }
215
216    runs.push({ text: formatDuration(end - span.start).padStart(DURATION_WIDTH), bold: span.id === slowest || undefined })
217
218    if (layout.hasCost) {
219      runs.push({ text: '  ' + formatCost(costs.get(span.id) ?? null).padStart(COST_WIDTH) })
220    }
221
222    return { id: span.id, runs: merged(runs) }
223  })
224
225  return { axis, rows: chartRows }
226}
227

1import type { AgentLink, Span, SpanKind, SpanUsage, Trace } from '../types'
2
3export const MAX_TRACES = 20
4export const MAX_SPANS = 3000
5export const AGENT_PREFIX = 'agent:'
6export const REQUEST_PREFIX = 'request:'
7
8type Price = { input: number; output: number; cacheWrite: number; cacheRead: number }
9
10const PRICES: Record<string, Price> = {
11  'fable 5.1': { input: 10, output: 50, cacheWrite: 12.5, cacheRead: 0.25 },
12  'mythos 5.1': { input: 10, output: 50, cacheWrite: 12.5, cacheRead: 0.25 },
13  'fable 5': { input: 10, output: 50, cacheWrite: 12.5, cacheRead: 1 },
14  'mythos 5': { input: 10, output: 50, cacheWrite: 12.5, cacheRead: 1 },
15  'opus 5.5': { input: 4, output: 20, cacheWrite: 5, cacheRead: 0.2 },
16  'opus 5': { input: 5, output: 25, cacheWrite: 6.25, cacheRead: 0.5 },
17  'opus 4.8': { input: 5, output: 25, cacheWrite: 6.25, cacheRead: 0.5 },
18  'opus 4.7': { input: 5, output: 25, cacheWrite: 6.25, cacheRead: 0.5 },
19  'opus 4.6': { input: 5, output: 25, cacheWrite: 6.25, cacheRead: 0.5 },
20  'opus 4.5': { input: 5, output: 25, cacheWrite: 6.25, cacheRead: 0.5 },
21  'sonnet 5.5': { input: 2, output: 10, cacheWrite: 2.5, cacheRead: 0.1 },
22  'sonnet 5': { input: 2, output: 10, cacheWrite: 2.5, cacheRead: 0.2 },
23  'sonnet 4.6': { input: 3, output: 15, cacheWrite: 3.75, cacheRead: 0.3 },
24  'sonnet 4.5': { input: 3, output: 15, cacheWrite: 3.75, cacheRead: 0.3 },
25  'haiku 5.5': { input: 0.1, output: 0.5, cacheWrite: 0.125, cacheRead: 0.01 },
26  'haiku 4.5': { input: 1, output: 5, cacheWrite: 1.25, cacheRead: 0.1 },
27}
28
29const LONG_PROMPT_PRICES: Record<string, { over: number; price: Price }> = {
30  'haiku 5.5': { over: 100_000, price: { input: 0.5, output: 2.5, cacheWrite: 0.625, cacheRead: 0.05 } },
31}
32
33type ModelName = { family: string; version: string }
34
35export function parseModel(id: string): ModelName | null {
36  const found = /claude-([a-z]+)-(\d+)(?:-(\d{1,2}))?(?!\d)/.exec(id)
37
38  if (found === null || found[1] === undefined || found[2] === undefined) {
39    return null
40  }
41
42  return { family: found[1], version: found[3] === undefined ? found[2] : `${found[2]}.${found[3]}` }
43}
44
45export function modelLabel(id: string): string {
46  const parsed = parseModel(id)
47
48  if (parsed === null) {
49    return id
50  }
51
52  return `${parsed.family.charAt(0).toUpperCase()}${parsed.family.slice(1)} ${parsed.version}`
53}
54
55export function costOf(model: string, usage: SpanUsage): number | null {
56  const parsed = parseModel(model)
57  const name = parsed === null ? '' : `${parsed.family} ${parsed.version}`
58  const long = LONG_PROMPT_PRICES[name]
59  const prompt = usage.input + usage.cacheWrite + usage.cacheRead
60  const price = long !== undefined && prompt > long.over ? long.price : PRICES[name]
61
62  if (price === undefined) {
63    return null
64  }
65
66  return (
67    (usage.input * price.input +
68      usage.output * price.output +
69      usage.cacheWrite * price.cacheWrite +
70      usage.cacheRead * price.cacheRead) /
71    1e6
72  )
73}
74
75function oneLine(value: unknown): string {
76  if (typeof value !== 'string') {
77    return ''
78  }
79
80  const trimmed = value.trim()
81
82  return /[\r\n]/.test(trimmed) ? '' : trimmed
83}
84
85function pathForms(path: string): string[] {
86  const parts = path.replace(/\\/g, '/').split('/').filter(part => part !== '')
87  const last = parts[parts.length - 1]
88  const parent = parts[parts.length - 2]
89
90  if (last === undefined) {
91    return []
92  }
93
94  return parent === undefined ? [last] : [`${parent}/${last}`, last]
95}
96
97function hostOf(url: string): string {
98  return /^[a-z]+:\/\/([^/?#]+)/i.exec(url)?.[1] ?? ''
99}
100
101function firstWord(command: string): string {
102  return command.split(/\s+/)[0] ?? ''
103}
104
105function unique(texts: string[]): string[] {
106  return texts.filter((text, index) => text !== '' && texts.indexOf(text) === index)
107}
108
109export function describeTool(tool: string, input: Record<string, unknown>): { name: string; details: string[] } {
110  if (tool.startsWith('mcp__')) {
111    const parts = tool.split('__')
112
113    return { name: parts.slice(2).join('__') || tool, details: unique([parts[1] ?? '']) }
114  }
115
116  const text = (key: string): string => oneLine(input[key])
117
118  switch (tool) {
119    case 'Read':
120    case 'Edit':
121    case 'Write':
122    case 'NotebookEdit':
123      return { name: tool, details: unique(pathForms(text('file_path') || text('notebook_path'))) }
124    case 'Bash':
125    case 'PowerShell': {
126      const command = text('command')
127
128      return { name: tool, details: unique([text('description'), command, firstWord(command)]) }
129    }
130    case 'Grep':
131    case 'Glob':
132      return { name: tool, details: unique([text('pattern')]) }
133    case 'WebFetch':
134      return { name: tool, details: unique([text('url'), hostOf(text('url'))]) }
135    case 'WebSearch':
136    case 'ToolSearch':
137      return { name: tool, details: unique([text('query')]) }
138    case 'Skill':
139      return { name: tool, details: unique([text('skill')]) }
140    case 'Agent':
141    case 'Task':
142      return { name: 'Agent', details: unique([text('description'), text('subagent_type')]) }
143    case 'Workflow':
144      return { name: tool, details: unique([text('name'), text('description')]) }
145    default:
146      return { name: tool, details: [] }
147  }
148}
149
150export type NewSpan = {
151  id: string
152  kind: SpanKind
153  name: string
154  details: string[]
155  start: number
156  model?: string | null
157  prompt?: string | null
158  agentId?: string | null
159  pending?: number
160}
161
162function makeSpan(fields: NewSpan, parent: string | null, loop: string | null): Span {
163  return {
164    id: fields.id,
165    parent,
166    kind: fields.kind,
167    name: fields.name,
168    details: fields.details,
169    loop,
170    start: fields.start,
171    end: fields.start,
172    pending: fields.pending ?? 1,
173    isError: false,
174    model: fields.model ?? null,
175    usage: null,
176    cost: null,
177    prompt: fields.prompt ?? null,
178    agentId: fields.agentId ?? null,
179  }
180}
181
182export function startRequest(traces: Trace[], turnId: string, prompt: string, at: number): Trace[] {
183  const id = REQUEST_PREFIX + turnId
184  const root = makeSpan({ id, kind: 'turn', name: 'Request', details: [], start: at, pending: 0 }, null, null)
185  const turn = makeSpan({ id: turnId, kind: 'turn', name: 'Turn', details: [], start: at }, id, null)
186  const kept = traces.filter(trace => trace.id !== id)
187
188  return [...kept, { id, prompt, start: at, spans: [root, turn] }].slice(-MAX_TRACES)
189}
190
191export function continueRequest(traces: Trace[], traceId: string, turnId: string, detail: string, at: number): Trace[] {
192  const details = detail === '' ? [] : [detail]
193
194  return addSpan(traces, { trace: traceId, span: traceId }, { id: turnId, kind: 'turn', name: 'Turn', details, start: at }, null)
195}
196
197export function latestTrace(traces: Trace[]): Trace | undefined {
198  return traces[traces.length - 1]
199}
200
201export function findSpan(traces: Trace[], id: string): { trace: Trace; span: Span } | undefined {
202  for (let index = traces.length - 1; index >= 0; index -= 1) {
203    const trace = traces[index]
204    const span = trace?.spans.find(one => one.id === id)
205
206    if (trace !== undefined && span !== undefined) {
207      return { trace, span }
208    }
209  }
210
211  return undefined
212}
213
214export function placeOf(
215  traces: Trace[],
216  links: Record<string, AgentLink>,
217  main: AgentLink | null,
218  loop: string | null,
219): AgentLink | undefined {
220  const link = loop === null ? (main ?? undefined) : links[loop]
221
222  if (link !== undefined && traces.some(trace => trace.id === link.trace && trace.spans.some(span => span.id === link.span))) {
223    return link
224  }
225
226  const latest = latestTrace(traces)
227
228  return latest === undefined ? undefined : { trace: latest.id, span: latest.id }
229}
230
231export function addSpan(traces: Trace[], place: AgentLink, fields: NewSpan, loop: string | null): Trace[] {
232  return traces.map(trace => {
233    if (trace.id !== place.trace || trace.spans.length >= MAX_SPANS || trace.spans.some(one => one.id === fields.id)) {
234      return trace
235    }
236
237    return { ...trace, spans: [...trace.spans, makeSpan(fields, place.span, loop)] }
238  })
239}
240
241export function changeSpan(traces: Trace[], id: string, change: (span: Span) => Span): Trace[] {
242  const found = findSpan(traces, id)
243
244  if (found === undefined) {
245    return traces
246  }
247
248  return traces.map(trace =>
249    trace.id !== found.trace.id ? trace : { ...trace, spans: trace.spans.map(span => (span.id === id ? change(span) : span)) },
250  )
251}
252
253export function closeSpan(span: Span, at: number, isError: boolean): Span {
254  return {
255    ...span,
256    pending: Math.max(0, span.pending - 1),
257    end: Math.max(span.end, at),
258    isError: span.isError || isError,
259  }
260}
261
262export function closeLoop(trace: Trace, loop: string | null, at: number): Trace {
263  return {
264    ...trace,
265    spans: trace.spans.map(span =>
266      span.loop === loop && span.pending > 0 && (span.kind === 'model' || span.kind === 'tool')
267        ? { ...span, pending: 0, end: Math.max(span.end, at) }
268        : span,
269    ),
270  }
271}
272
273export function completeMain(traces: Trace[], id: string, at: number, isError: boolean): Trace[] {
274  return traces.map(trace => {
275    if (!trace.spans.some(span => span.id === id)) {
276      return trace
277    }
278
279    const closed = closeLoop(trace, null, at)
280
281    return {
282      ...closed,
283      spans: closed.spans.map(span => {
284        if (span.id === id) {
285          return { ...span, pending: 0, end: Math.max(span.end, at), isError: span.isError || isError }
286        }
287
288        if (span.kind === 'agent' && span.agentId === null && span.pending > 0) {
289          return { ...span, pending: 0, end: Math.max(span.end, at) }
290        }
291
292        return span
293      }),
294    }
295  })
296}
297
298export function displayEnd(span: Span, now: number): number {
299  return span.pending > 0 ? Math.max(now, span.end, span.start) : span.end
300}
301
302export function isLive(trace: Trace): boolean {
303  return trace.spans.some(span => span.pending > 0)
304}
305
306export type Row = { span: Span; depth: number }
307
308export function rowsOf(trace: Trace): Row[] {
309  const children = new Map<string | null, Span[]>()
310
311  for (const span of trace.spans) {
312    const parent = span.id === trace.id ? null : span.parent
313    const list = children.get(parent) ?? []
314    list.push(span)
315    children.set(parent, list)
316  }
317
318  const rows: Row[] = []
319  const seen = new Set<string>()
320
321  const walk = (parent: string | null, depth: number): void => {
322    const list = [...(children.get(parent) ?? [])].sort((a, b) => a.start - b.start)
323
324    for (const span of list) {
325      if (seen.has(span.id)) {
326        continue
327      }
328
329      seen.add(span.id)
330      rows.push({ span, depth })
331      walk(span.id, depth + 1)
332    }
333  }
334
335  walk(null, 0)
336
337  for (const span of trace.spans) {
338    if (!seen.has(span.id)) {
339      seen.add(span.id)
340      rows.push({ span, depth: 1 })
341    }
342  }
343
344  return rows
345}
346
347export function subtreeOf(trace: Trace, id: string): Span[] {
348  const rows = rowsOf(trace)
349  const at = rows.findIndex(row => row.span.id === id)
350  const head = rows[at]
351
352  if (head === undefined) {
353    return []
354  }
355
356  const out = [head.span]
357
358  for (const row of rows.slice(at + 1)) {
359    if (row.depth <= head.depth) {
360      break
361    }
362
363    out.push(row.span)
364  }
365
366  return out
367}
368
369export type Figures = {
370  duration: number
371  cost: number | null
372  tokensIn: number
373  tokensOut: number
374  modelCalls: number
375  toolCalls: number
376  errors: number
377}
378
379export function isToolCall(span: Span): boolean {
380  return span.kind === 'tool' || (span.kind === 'agent' && !span.id.startsWith(AGENT_PREFIX))
381}
382
383export function figuresOf(spans: Span[], now: number): Figures {
384  const head = spans[0]
385  const start = head?.start ?? 0
386  let end = start
387  let cost: number | null = null
388  let tokensIn = 0
389  let tokensOut = 0
390  let modelCalls = 0
391  let toolCalls = 0
392  let errors = 0
393
394  for (const span of spans) {
395    end = Math.max(end, displayEnd(span, now))
396
397    if (span.cost !== null) {
398      cost = (cost ?? 0) + span.cost
399    }
400
401    if (span.usage !== null) {
402      tokensIn += span.usage.input + span.usage.cacheRead + span.usage.cacheWrite
403      tokensOut += span.usage.output
404    }
405
406    modelCalls += span.kind === 'model' ? 1 : 0
407    toolCalls += isToolCall(span) ? 1 : 0
408    errors += span.isError ? 1 : 0
409  }
410
411  return { duration: end - start, cost, tokensIn, tokensOut, modelCalls, toolCalls, errors }
412}
413
414export function formatDuration(ms: number): string {
415  const safe = Math.max(0, ms)
416
417  if (safe < 1000) {
418    return `${Math.round(safe)}ms`
419  }
420
421  const seconds = safe / 1000
422
423  if (seconds < 10) {
424    return `${seconds.toFixed(1)}s`
425  }
426
427  const whole = Math.floor(seconds)
428
429  if (whole < 60) {
430    return `${whole}s`
431  }
432
433  if (whole < 3600) {
434    return `${Math.floor(whole / 60)}m ${String(whole % 60).padStart(2, '0')}s`
435  }
436
437  return `${Math.floor(whole / 3600)}h ${String(Math.floor((whole % 3600) / 60)).padStart(2, '0')}m`
438}
439
440export function formatTick(seconds: number): string {
441  if (seconds < 60) {
442    return `${seconds}s`
443  }
444
445  if (seconds < 3600) {
446    const rest = seconds % 60
447
448    return rest === 0 ? `${seconds / 60}m` : `${Math.floor(seconds / 60)}m ${rest}s`
449  }
450
451  const minutes = Math.floor((seconds % 3600) / 60)
452
453  return minutes === 0 ? `${Math.floor(seconds / 3600)}h` : `${Math.floor(seconds / 3600)}h ${minutes}m`
454}
455
456export function formatCost(usd: number | null): string {
457  if (usd === null) {
458    return ''
459  }
460
461  if (usd === 0) {
462    return '$0'
463  }
464
465  if (usd < 0.01) {
466    return '<$0.01'
467  }
468
469  if (usd < 100) {
470    return `$${usd.toFixed(2)}`
471  }
472
473  return `$${Math.round(usd)}`
474}
475
476export function formatTokens(count: number): string {
477  if (count < 1000) {
478    return String(count)
479  }
480
481  if (count < 1e6) {
482    return `${(count / 1e3).toFixed(count < 1e4 ? 1 : 0)}K`
483  }
484
485  return `${(count / 1e6).toFixed(1)}M`
486}
487
488const DOMAINS = [
489  5, 10, 15, 20, 30, 45, 60, 90, 120, 180, 240, 300, 450, 600, 900, 1200, 1800, 2700, 3600, 5400, 7200, 10800, 14400,
490  21600, 43200, 86400,
491]
492
493export function domainOf(extent: number): number {
494  const wanted = (extent * 1.1) / 1000
495
496  return (DOMAINS.find(seconds => seconds >= wanted) ?? Math.ceil(wanted / 86400) * 86400) * 1000
497}
498
499const STEPS = [1, 2, 5, 10, 15, 30, 60, 120, 300, 600, 900, 1800, 3600, 7200, 10800, 21600, 43200, 86400]
500
501export function ticksOf(domain: number, cells: number): number[] {
502  const seconds = domain / 1000
503  const room = Math.max(1, Math.floor(cells / 8))
504  const step = STEPS.find(one => seconds / one <= room) ?? seconds
505  const ticks: number[] = []
506
507  for (let at = 0; at < seconds; at += step) {
508    ticks.push(at)
509  }
510
511  return ticks
512}
513
514export type Cell = { glyph: string; mode: 'empty' | 'solid' | 'inverse' }
515
516const LEFT_EIGHTHS = ' ▏▎▍▌▋▊▉'
517
518export function barCells(from: number, to: number, cells: number): Cell[] {
519  const start = Math.max(0, Math.min(from, cells * 8))
520  const end = Math.max(start + 1, Math.min(to, cells * 8))
521  const out: Cell[] = []
522
523  for (let cell = 0; cell < cells; cell += 1) {
524    const left = Math.max(start, cell * 8)
525    const right = Math.min(end, cell * 8 + 8)
526
527    if (right <= left) {
528      out.push({ glyph: ' ', mode: 'empty' })
529    } else if (left === cell * 8 && right === cell * 8 + 8) {
530      out.push({ glyph: '█', mode: 'solid' })
531    } else if (left === cell * 8) {
532      out.push({ glyph: LEFT_EIGHTHS[right - cell * 8] ?? '█', mode: 'solid' })
533    } else if (right === cell * 8 + 8) {
534      out.push({ glyph: LEFT_EIGHTHS[left - cell * 8] ?? ' ', mode: 'inverse' })
535    } else {
536      out.push({ glyph: LEFT_EIGHTHS[Math.max(1, right - left)] ?? '▏', mode: 'solid' })
537    }
538  }
539
540  return out
541}
542

1export type SpanKind = 'turn' | 'agent' | 'model' | 'tool'
2
3export type SpanUsage = {
4  input: number
5  output: number
6  cacheRead: number
7  cacheWrite: number
8}
9
10export type Span = {
11  id: string
12  parent: string | null
13  kind: SpanKind
14  name: string
15  details: string[]
16  loop: string | null
17  start: number
18  end: number
19  pending: number
20  isError: boolean
21  model: string | null
22  usage: SpanUsage | null
23  cost: number | null
24  prompt: string | null
25  agentId: string | null
26}
27
28export type Trace = {
29  id: string
30  prompt: string
31  start: number
32  spans: Span[]
33}
34
35export type AgentLink = { trace: string; span: string }
36
37export type Theme = 'dark' | 'light'
38
39declare module 'claude-code' {
40  interface PluginState {
41    'trace-waterfall': {
42      traces: Trace[]
43      agents: Record<string, AgentLink>
44      main: AgentLink | null
45      view: string | null
46      selected: string | null
47      now: number
48      theme: Theme
49    }
50  }
51}
52