Trace waterfall

npx shadcn@latest add "@assistant-ui/elements-trace-waterfall"

npx assistant-ui@latest init

import { AssistantRuntimeProvider } from "@assistant-ui/react";
import { useChatRuntime, AssistantChatTransport } from "@assistant-ui/ai-sdk";

export default function App() {
  const runtime = useChatRuntime({
    transport: new AssistantChatTransport({ api: "/api/chat" }),
  });

  return (
    <AssistantRuntimeProvider runtime={runtime}>
      {/* your components */}
    </AssistantRuntimeProvider>
  );
}

npx shadcn@latest add "@assistant-ui/elements-trace-waterfall"

TraceWaterfall lays a run's spans on one shared time axis: a header with the total duration, then one row per span with its name, a positioned bar, and its own duration in milliseconds. With a runtime you derive spans from a message's own timing and its tool calls; standalone you hold the span list yourself.

Getting started

"use client";

import { useAuiState, useMessageTiming } from "@assistant-ui/react";
import {
  TraceWaterfall,
  type TraceSpan,
} from "@/components/assistant-ui/elements/trace-waterfall";

function useMessageTrace(): { spans: TraceSpan[]; totalMs: number } {
  const timing = useMessageTiming();
  const parts = useAuiState((s) => s.message.parts);
  const totalMs = timing?.totalStreamTime ?? 0;

  const spans: TraceSpan[] = [
    {
      id: "message",
      name: "message",
      depth: 0,
      startMs: 0,
      durationMs: totalMs,
      status: totalMs > 0 ? "completed" : "running",
    },
  ];

  for (const part of parts) {
    if (part.type !== "tool-call" || !part.timing) continue;
    const streamStart = timing?.streamStartTime ?? part.timing.startedAt;
    spans.push({
      id: part.toolCallId,
      name: part.toolName,
      depth: 1,
      startMs: Math.max(0, part.timing.startedAt - streamStart),
      durationMs: part.timing.completedAt
        ? part.timing.completedAt - part.timing.startedAt
        : 0,
      status:
        part.status.type === "running"
          ? "running"
          : part.status.type === "complete"
            ? "completed"
            : "failed",
    });
  }

  return { spans, totalMs };
}

export function MessageTrace() {
  const { spans, totalMs } = useMessageTrace();
  if (spans.length <= 1) return null;
  return (
    <TraceWaterfall spans={spans} totalMs={totalMs} visibleCount={spans.length} />
  );
}

"use client";

import { useState } from "react";
import {
  TraceWaterfall,
  type TraceSpan,
} from "@/components/assistant-ui/elements/trace-waterfall";

const SPANS: TraceSpan[] = [
  { id: "run", name: "run", depth: 0, startMs: 0, durationMs: 1840, status: "completed" },
  { id: "model", name: "chat.completions", depth: 1, startMs: 40, durationMs: 720, status: "completed" },
  { id: "search", name: "web_search", depth: 1, startMs: 790, durationMs: 460, status: "completed" },
];

export function Run() {
  const [visibleCount, setVisibleCount] = useState(0);
  return <TraceWaterfall spans={SPANS} totalMs={1840} visibleCount={visibleCount} />;
}

useEffect(() => {
  if (visibleCount >= SPANS.length) return;
  const id = setTimeout(() => setVisibleCount((n) => n + 1), 700);
  return () => clearTimeout(id);
}, [visibleCount]);

Anatomy

<div data-slot="trace-waterfall">
  <div>
    <span>Trace</span>
    <span>{/* totalMs, suffixed "ms" */}</span>
  </div>
  <div>
    {/* one row per visible span: name, labelled positioned bar, duration */}
  </div>
</div>

Only the first visibleCount spans render, floored and clamped to the array's length, so a span can exist in spans before it's shown. A newly revealed row fades and slides in from the left; a row that was already visible does not replay the animation on a later re-render, since it's the same DOM node. Each bar is positioned at startMs / totalMs and sized to durationMs / totalMs, both as percentages, with a totalMs of 0 treated as 1 so the math never divides by zero. A bar's width never drops below 1.5% of the row, so a very short span still shows a visible sliver. The positioned bar is a labelled image carrying the status and timing that colour and position alone would otherwise hold; the row's own name and duration stay readable text, and neither is a completion progressbar. Row indentation follows depth, at 0.75rem per level; nothing computes depth for you, it's a field you set. Color follows status: blue and pulsing while "running", a flat neutral tone once "completed", red for "failed". The header's total carries an "ms" suffix; each row's own duration does not.

Examples

Restyle the waterfall

Both lanes take className on the root. Rows and the header's total use the shared mono surface from surfaces.tsx; the root itself uses paper.

<TraceWaterfall className="max-w-none" /* ... */ />

Marking a span as failed

{ id: "search", name: "web_search", depth: 1, startMs: 790, durationMs: 210, status: "failed" }

Showing only the most recent spans

visibleCount reveals from the start of the array, so cap the array itself, not just visibleCount, to show only the most recent spans of a long run. totalMs stays the run's real total, so the remaining bars keep their true position along the axis rather than rescaling to fill it.

<TraceWaterfall spans={spans.slice(-6)} totalMs={totalMs} visibleCount={6} />

API reference