Trace Inspector

"use client";

import { type ComponentProps, useMemo } from "react";
import { GraphCanvas, type GraphCanvasEdge, type GraphCanvasNode, useGraphLayout } from "./graph-canvas";
import { adjacency, GRAPH_STATUS_LABELS, type GraphDirection, type GraphStatus, layoutEdges } from "./graph-layout";
import {
  cx,
  DetailHeader,
  DetailNote,
  DetailProperties,
  DetailSection,
  DetailStack,
  divider,
  EdgeLegend,
  EmptyDetails,
  ErrorNote,
  formatDuration,
  GraphHeader,
  GraphListbox,
  GraphListRow,
  graphSurfaceClass,
  RelatedButton,
  StatusCounts,
  useControllableState,
} from "./graph-ui";

export type TraceStepKind = "agent" | "llm" | "tool" | "approval" | "handoff" | "step";
/** calls: the source started the target as a child; follows: the target ran after the source; handoff: work passed to another agent. */
export type TraceEdgeType = "calls" | "follows" | "handoff";

export interface TraceStep {
  /** Stable id, for example the span id. */
  id: string;
  name: string;
  kind: TraceStepKind;
  status: GraphStatus;
  /** Start time in milliseconds from the start of the trace. */
  startMs: number;
  /** Omit while the step is still running. */
  durationMs?: number;
  summary?: string;
  input?: unknown;
  /** Rendered as text, never as HTML. */
  output?: string;
  error?: string;
  /** Extra key-value details such as model, tokens, or cache hits. */
  attributes?: Record<string, string | number | boolean>;
}

export interface TraceEdge {
  id: string;
  source: string;
  target: string;
  type: TraceEdgeType;
  label?: string;
}

export interface TraceInspectorProps extends Omit<ComponentProps<"section">, "children" | "onSelect" | "title"> {
  nodes: readonly TraceStep[];
  edges: readonly TraceEdge[];
  title?: string;
  /** Controlled selection. Pass `null` for none; leave undefined to let the component manage it. */
  selectedId?: string | null;
  defaultSelectedId?: string | null;
  onSelect?: (id: string | null) => void;
  /** Height of the timeline and graph in pixels. */
  height?: number;
  direction?: GraphDirection;
}

export const TRACE_KIND_LABELS: Record<TraceStepKind, string> = {
  agent: "Agent",
  llm: "Model call",
  tool: "Tool call",
  approval: "Approval",
  handoff: "Hand-off",
  step: "Step",
};

export const TRACE_EDGE_LABELS: Record<TraceEdgeType, string> = { calls: "Calls", follows: "Follows", handoff: "Hand-off" };

// Bars stay neutral for finished work; color marks only what needs a look.
const barColor: Record<GraphStatus, string> = {
  queued: "bg-zinc-950/15 dark:bg-white/15",
  running: "bg-sky-500",
  "waiting-for-approval": "bg-amber-500",
  succeeded: "bg-zinc-950/40 dark:bg-white/45",
  failed: "bg-red-500",
  skipped: "bg-zinc-950/15 dark:bg-white/15",
  cancelled: "bg-zinc-950/15 dark:bg-white/15",
};

/** Steps in start order; ties keep their input order. This is the timeline the list shows. */
export function timelineOrder<T extends Pick<TraceStep, "startMs">>(steps: readonly T[]): T[] {
  return steps
    .map((step, index) => ({ step, index }))
    .sort((a, b) => a.step.startMs - b.step.startMs || a.index - b.index)
    .map(({ step }) => step);
}

/** A trace's timeline and graph side by side, with selection kept in sync, and details for the selected step. */
export function TraceInspector({
  nodes,
  edges,
  title = "Trace",
  selectedId,
  defaultSelectedId = null,
  onSelect,
  height = 440,
  direction = "TB",
  className,
  ...props
}: TraceInspectorProps) {
  const [selected, setSelected] = useControllableState<string | null>(selectedId, defaultSelectedId, onSelect);
  const shownEdges = useMemo(() => layoutEdges(nodes, edges), [nodes, edges]);
  const layout = useGraphLayout(nodes, shownEdges, { direction });
  const byId = useMemo(() => new Map(nodes.map((node) => [node.id, node])), [nodes]);
  const links = useMemo(() => adjacency(shownEdges), [shownEdges]);
  const timeline = useMemo(() => timelineOrder(nodes), [nodes]);
  const total = useMemo(() => Math.max(1, ...nodes.map((node) => node.startMs + (node.durationMs ?? 0))), [nodes]);
  const current = selected ? byId.get(selected) : undefined;
  const name = (id: string) => byId.get(id)?.name ?? id;

  const canvasNodes = useMemo<GraphCanvasNode[]>(
    () =>
      nodes.map((node) => ({
        id: node.id,
        title: node.name,
        subtitle: node.durationMs != null ? `${TRACE_KIND_LABELS[node.kind]} · ${formatDuration(node.durationMs)}` : TRACE_KIND_LABELS[node.kind],
        mono: node.kind === "tool",
        status: node.status,
      })),
    [nodes],
  );
  const canvasEdges = useMemo<GraphCanvasEdge[]>(
    () =>
      shownEdges.map((edge) => ({
        id: edge.id,
        source: edge.source,
        target: edge.target,
        type: edge.type,
        label: edge.label,
        dashed: edge.type !== "calls",
        animated: byId.get(edge.target)?.status === "running",
        tone: selected && (edge.source === selected || edge.target === selected) ? "highlight" : "default",
      })),
    [shownEdges, byId, selected],
  );

  const select = (id: string) => setSelected(id);

  // Rows do not depend on the selection, so reuse them: selecting re-renders only the option wrappers.
  const rows = useMemo(() => {
    const label = (id: string) => byId.get(id)?.name ?? id;
    return new Map(
      timeline.map((node) => {
        const parents = links.incoming.get(node.id) ?? [];
        return [
          node.id,
          <GraphListRow
            status={node.status}
            title={node.name}
            mono={node.kind === "tool"}
            muted={node.status === "queued" || node.status === "skipped" || node.status === "cancelled"}
            struck={node.status === "skipped"}
            trailing={node.durationMs != null ? formatDuration(node.durationMs) : node.status === "running" ? "running" : node.status === "waiting-for-approval" ? "waiting" : undefined}
            meta={
              <>
                {GRAPH_STATUS_LABELS[node.status]} · {TRACE_KIND_LABELS[node.kind]} · +{formatDuration(node.startMs)}
                {parents.length ? ` · ${parents.map((edge) => `${TRACE_EDGE_LABELS[edge.type].toLowerCase()} from ${label(edge.source)}`).join(", ")}` : ""}
              </>
            }
          >
            <div aria-hidden="true" className="relative mt-2 h-1 rounded-full bg-zinc-950/[0.06] dark:bg-white/[0.07]">
              <div
                className={cx("absolute inset-y-0 min-w-[3px] rounded-full", barColor[node.status])}
                style={{
                  left: `${(node.startMs / total) * 100}%`,
                  width: `${Math.min(100 - (node.startMs / total) * 100, ((node.durationMs ?? total - node.startMs) / total) * 100)}%`,
                }}
              />
            </div>
          </GraphListRow>,
        ] as const;
      }),
    );
  }, [timeline, links, byId, total]);

  const related = (label: string, list: { id: string; other: string; type: TraceEdgeType }[]) =>
    list.length ? (
      <DetailSection label={label}>
        {list.map((item) => (
          <RelatedButton
            key={item.id}
            label={name(item.other)}
            mono={byId.get(item.other)?.kind === "tool"}
            status={byId.get(item.other)?.status}
            note={TRACE_EDGE_LABELS[item.type]}
            onClick={() => select(item.other)}
          />
        ))}
      </DetailSection>
    ) : null;

  const timing = current
    ? [
        { label: "Kind", value: TRACE_KIND_LABELS[current.kind] },
        { label: "Started", value: `+${formatDuration(current.startMs)}`, mono: true },
        { label: "Duration", value: current.durationMs != null ? formatDuration(current.durationMs) : "Still running", mono: current.durationMs != null },
      ]
    : [];

  return (
    <section aria-label={title} className={cx(graphSurfaceClass, className)} {...props}>
      <GraphHeader
        title={title}
        meta={
          <StatusCounts
            lead={
              <>
                {nodes.length} {nodes.length === 1 ? "step" : "steps"} · <span className="font-mono tabular-nums">{formatDuration(total)}</span>
              </>
            }
            statuses={nodes.map((node) => node.status)}
          />
        }
      />
      <div className="grid @2xl:grid-cols-[19rem_minmax(0,1fr)]">
        <div style={{ maxHeight: height }} className={cx("min-w-0 overflow-y-auto border-b @2xl:border-r @2xl:border-b-0", divider)}>
          <div
            aria-hidden="true"
            className={cx("sticky top-0 z-10 flex justify-between border-b bg-white py-2 pr-4 pl-[2.625rem] font-mono text-[10px] text-zinc-500 tabular-nums dark:bg-black", divider)}
          >
            <span>0</span>
            <span>{formatDuration(total / 2)}</span>
            <span>{formatDuration(total)}</span>
          </div>
          <GraphListbox
            label={`${title} timeline`}
            items={timeline}
            getId={(node) => node.id}
            selectedId={selected}
            onSelect={select}
            empty="No steps yet."
            renderItem={(node) => rows.get(node.id)}
          />
        </div>
        <GraphCanvas
          label={`${title} graph: ${nodes.length} steps. The timeline list beside it has the same steps and is keyboard accessible.`}
          nodes={canvasNodes}
          edges={canvasEdges}
          layout={layout}
          selectedNodeId={selected}
          onNodeClick={select}
          onPaneClick={() => setSelected(null)}
          height={height}
        />
      </div>
      <div role="region" aria-label="Step details" className={cx("border-t", divider)}>
        {current ? (
          <div className="grid @2xl:grid-cols-2 @2xl:divide-x @2xl:divide-zinc-950/[0.07] @2xl:dark:divide-white/[0.08]">
            <DetailStack>
              <DetailHeader title={current.name} mono={current.kind === "tool"} status={current.status} meta={[TRACE_KIND_LABELS[current.kind]]} />
              {current.summary || (current.status === "failed" && current.error) ? (
                <div className="space-y-3">
                  {current.summary ? <p className="text-sm text-pretty text-zinc-600 dark:text-zinc-400">{current.summary}</p> : null}
                  {current.status === "failed" && current.error ? <ErrorNote>{current.error}</ErrorNote> : null}
                </div>
              ) : null}
              <div className="space-y-3.5">
                <DetailProperties items={timing} />
                {current.attributes && Object.keys(current.attributes).length ? (
                  <DetailSection label="Attributes">
                    <DetailProperties items={Object.entries(current.attributes).map(([key, value]) => ({ label: key, value: String(value), mono: true }))} />
                  </DetailSection>
                ) : null}
              </div>
            </DetailStack>
            <DetailStack className={cx("border-t @2xl:border-t-0", divider)}>
              {current.input !== undefined || current.output ? (
                <div className="space-y-3.5">
                  {current.input !== undefined ? <DetailSection label="Input">{JSON.stringify(current.input, null, 2)}</DetailSection> : null}
                  {current.output ? <DetailSection label="Output">{current.output}</DetailSection> : null}
                </div>
              ) : null}
              <div className="space-y-3.5">
                {related(
                  "From",
                  (links.incoming.get(current.id) ?? []).map((edge) => ({ id: edge.id, other: edge.source, type: edge.type })),
                ) ?? (
                  <DetailSection label="From">
                    <DetailNote>Nothing. This step starts the trace.</DetailNote>
                  </DetailSection>
                )}
                {related(
                  "Leads to",
                  (links.outgoing.get(current.id) ?? []).map((edge) => ({ id: edge.id, other: edge.target, type: edge.type })),
                ) ?? (
                  <DetailSection label="Leads to">
                    <DetailNote>Nothing comes after this step.</DetailNote>
                  </DetailSection>
                )}
              </div>
            </DetailStack>
          </div>
        ) : (
          <EmptyDetails hint="In the timeline, arrow keys, Home, and End move the selection.">
            Select a step in the timeline or the graph to see its timing, attributes, input, and output.
          </EmptyDetails>
        )}
      </div>
      <footer className={cx("border-t px-4 py-2.5", divider)}>
        <EdgeLegend items={[{ label: "Calls" }, { label: "Follows or hand-off", dashed: true }]} />
      </footer>
    </section>
  );
}