Thread Virtualization

Render very long threads with @tanstack/react-virtual, one row per user message, assistant part and turn end, with createThreadRowsSelector and ThreadPrimitive.Row.

Virtualization mounts only the messages near the viewport and represents the rest as empty space, so a thread with thousands of messages scrolls like one with twenty. assistant-ui does not ship a virtualized thread component; this guide shows the supported composition, extracted from a production consumer and available as a runnable example: examples/with-virtualized-thread.

Do you need this?

Probably not. The default kit already renders message bodies with content-visibility: auto and contain-intrinsic-size, which skips paint work for off-screen messages, and ThreadPrimitive.Viewport handles auto-scroll. That covers typical threads. Reach for virtualization when React mount and update cost itself becomes the bottleneck: threads with hundreds to thousands of messages, or very heavy per-message content, where typing latency degrades because every message stays mounted.

Rows, not messages

Virtualizing per message leaves one very large cell whenever an assistant message is long: an agent turn is often a single message with a reasoning run and hundreds of tool calls. Virtualize rows instead. createThreadRowsSelector flattens the thread into rows, and useAuiState reads them:

  • a "message" row per user or system message (and per assistant message being edited),
  • a "part" row per top-level node of each assistant message's part grouping, so every part is its own row unless groupBy coalesces adjacent ones (a reasoning run, for example),
  • a "turn-end" row after each turn's last message once the turn has an assistant reply, where a turn footer or the action bar goes. A turn still waiting for its first reply has no turn-end row yet, so render a pending indicator after the last row while s.thread.isRunning is true.
import {
  createThreadRowsSelector,
  groupPartByType,
  useAuiState,
} from "@assistant-ui/react";

const selectRows = createThreadRowsSelector({
  groupBy: groupPartByType({ reasoning: ["group-reasoning"] }),
});

function VirtualThread() {
  const rows = useAuiState(selectRows);
  const virtualizer = useVirtualizer({
    count: rows.length,
    getItemKey: (index) => rows[index]!.key,
    estimateSize: () => 80,
    getScrollElement: () => scrollerRef.current,
  });
  // ...
}

row.key is stable while tokens stream, when parts arrive and when earlier messages are prepended (as long as messages keep their ids), so use it as the virtualizer's item key. The selector returns the same array while a streamed token leaves the rows unchanged, so the list re-renders only when rows appear, disappear, regroup or a turn completes. Create one selector per list, at module scope or with useState; groupBy takes the same functions as MessagePrimitive.GroupedParts.

Rendering a row

ThreadPrimitive.Row renders one row and mounts only the scopes it needs: the row's message, and for a part row only that row's parts. Its render function switches on type:

const renderRow: ComponentProps<typeof ThreadPrimitive.Row>["children"] = (
  info,
) => {
  switch (info.type) {
    case "message":
      if (info.message.composer.isEditing) return <EditComposer />;
      return info.message.role === "system" ? <SystemMessage /> : <UserMessage />;
    case "part":
      switch (info.part.type) {
        case "group-reasoning":
          return <Reasoning>{info.children}</Reasoning>;
        case "text":
          return <MarkdownText />;
        case "tool-call":
          return info.part.toolUI ?? <ToolFallback {...info.part} />;
        default:
          return null;
      }
    case "turn-end":
      return info.row.completedAt === undefined ? null : (
        <p>Worked for {formatDuration(info.row.completedAt - info.row.startedAt)}</p>
      );
  }
};

{virtualizer.getVirtualItems().map((item) => (
  <div key={item.key} data-index={item.index} ref={virtualizer.measureElement}>
    <ThreadPrimitive.Row row={rows[item.index]!}>{renderRow}</ThreadPrimitive.Row>
  </div>
))}

A message row is a user or system message, or any message being edited. A part row's part and children follow MessagePrimitive.GroupedParts: a group part wraps children, a leaf part renders itself, and returning null for a tool or data leaf renders its registered UI. A turn-end row renders in the scope of the turn's last message and carries the turn's span: row.startedAt is the turn's first message createdAt, and row.completedAt is the latest end the turn records, from an assistant message's stream timing (metadata.timing) or a reasoning or tool call part's timing.completedAt. completedAt is absent while the turn runs or waits on the user, and when nothing in the turn records an end, so a host that follows an agent supplies those times on its messages. A row whose message has left the thread renders nothing, so a row the virtualizer keeps for a frame after a deletion is safe.

ThreadPrimitive.MessageByIndex and ThreadPrimitive.Unstable_MessageById still render one whole message, for a list that virtualizes per message.

Padding spacers, not absolute positioning

Render the virtual items in normal document flow inside a spacer div whose paddingTop/paddingBottom represent the unmounted regions. Items remain regular flow children, so message CSS (including position: sticky patterns and the kit styling) keeps working, and virtualizer.measureElement records real heights as items mount:

const items = virtualizer.getVirtualItems();
const paddingTop = items[0]?.start ?? 0;
const paddingBottom = Math.max(
  0,
  virtualizer.getTotalSize() - (items.at(-1)?.end ?? 0),
);

Owning the scroll element

The composition owns its scroll container instead of using ThreadPrimitive.Viewport: the built-in auto-scroll assumes every message is mounted, and its resize-driven re-pin can fight the virtualizer's measurement adjustments. Three pieces replace it:

  1. Auto-follow. A ResizeObserver on the content wrapper re-pins the scroller to the bottom while a sticky flag is armed. The flag disarms when the user scrolls up (detected as scrollTop decreasing while scrollHeight and clientHeight are stable, the same heuristic the built-in viewport uses, plus wheel-up and touchmove) and re-arms when the user returns to the bottom.
  2. Measurement guard. While pinned at the bottom, the virtualizer's own scroll adjustments on item re-measurement are suppressed via a custom scrollToFn; without this the two scroll writers fight and the view rubber-bands during streaming.
  3. Run-start jump. A useLayoutEffect observes s.thread.isRunning flipping to true and jumps to the bottom before paint, so a just-sent message never flashes below the fold.

The production consumer this is extracted from disables streaming auto-follow entirely as a product choice; the example keeps following because it matches ThreadPrimitive.Viewport's default behavior. Both are valid, and the disarm guard makes either safe.

Try it

npx assistant-ui@latest create my-app

Then copy app/VirtualizedThread.tsx, app/MyRuntimeProvider.tsx, and app/seed-messages.ts from the example, or clone the repo and run pnpm --filter with-virtualized-thread dev.