Switch message timeline virtualization to Virtuoso

apps/web/package.json

@@ -39,6 +39,7 @@
     "react": "^19.0.0",
     "react-dom": "^19.0.0",
     "react-markdown": "^10.1.0",
+    "react-virtuoso": "^4.18.4",
     "remark-gfm": "^4.0.1",
     "tailwind-merge": "^3.4.0",
     "zustand": "^5.0.11"

apps/web/src/components/chat/MessagesTimeline.tsx

@@ -1,21 +1,17 @@
 import { type EnvironmentId, type MessageId, type TurnId } from "@t3tools/contracts";
 import {
+  forwardRef,
   memo,
   useCallback,
-  useEffect,
   useLayoutEffect,
   useMemo,
   useRef,
   useState,
   type ReactNode,
 } from "react";
-import {
-  measureElement as measureVirtualElement,
-  type VirtualItem,
-  useVirtualizer,
-} from "@tanstack/react-virtual";
+import { Virtuoso, type ItemProps, type ListProps, type VirtuosoHandle } from "react-virtuoso";
 import { deriveTimelineEntries, formatElapsed } from "../../session-logic";
-import { AUTO_SCROLL_BOTTOM_THRESHOLD_PX } from "../../chat-scroll";
+import { isScrollContainerNearBottom } from "../../chat-scroll";
 import { type TurnDiffSummary } from "../../types";
 import { summarizeTurnDiffStats } from "../../lib/turnDiffTree";
 import ChatMarkdown from "../ChatMarkdown";
@@ -34,7 +30,6 @@
   ZapIcon,
 } from "lucide-react";
 import { Button } from "../ui/button";
-import { clamp } from "effect/Number";
 import { buildExpandedImagePreview, ExpandedImagePreview } from "./ExpandedImagePreview";
 import { ProposedPlanCard } from "./ProposedPlanCard";
 import { ChangedFilesTree } from "./ChangedFilesTree";
@@ -64,7 +59,7 @@
   textContainsInlineTerminalContextLabels,
 } from "./userMessageTerminalContexts";
 
-const ALWAYS_UNVIRTUALIZED_TAIL_ROWS = 8;
+const VIRTUALIZED_VIEWPORT_PADDING_PX = 640;
 
 interface MessagesTimelineProps {
   hasMessages: boolean;
@@ -92,17 +87,6 @@
   resolvedTheme: "light" | "dark";
   timestampFormat: TimestampFormat;
   workspaceRoot: string | undefined;
-  onVirtualizerSnapshot?: (snapshot: {
-    totalSize: number;
-    measurements: ReadonlyArray<{
-      id: string;
-      kind: MessagesTimelineRow["kind"];
-      index: number;
-      size: number;
-      start: number;
-      end: number;
-    }>;
-  }) => void;
 }
 
 export const MessagesTimeline = memo(function MessagesTimeline({
@@ -131,9 +115,9 @@
   resolvedTheme,
   timestampFormat,
   workspaceRoot,
-  onVirtualizerSnapshot,
 }: MessagesTimelineProps) {
   const timelineRootRef = useRef<HTMLDivElement | null>(null);
+  const virtuosoRef = useRef<VirtuosoHandle | null>(null);
   const [timelineWidthPx, setTimelineWidthPx] = useState<number | null>(null);
 
   useLayoutEffect(() => {
@@ -172,138 +156,28 @@
     [timelineEntries, completionDividerBeforeEntryId, isWorking, activeTurnStartedAt],
   );
 
-  const firstUnvirtualizedRowIndex = useMemo(() => {
-    const firstTailRowIndex = Math.max(rows.length - ALWAYS_UNVIRTUALIZED_TAIL_ROWS, 0);
-    if (!activeTurnInProgress) return firstTailRowIndex;
-
-    const turnStartedAtMs =
-      typeof activeTurnStartedAt === "string" ? Date.parse(activeTurnStartedAt) : Number.NaN;
-    let firstCurrentTurnRowIndex = -1;
-    if (!Number.isNaN(turnStartedAtMs)) {
-      firstCurrentTurnRowIndex = rows.findIndex((row) => {
-        if (row.kind === "working") return true;
-        if (!row.createdAt) return false;
-        const rowCreatedAtMs = Date.parse(row.createdAt);
-        return !Number.isNaN(rowCreatedAtMs) && rowCreatedAtMs >= turnStartedAtMs;
-      });
-    }
-
-    if (firstCurrentTurnRowIndex < 0) {
-      firstCurrentTurnRowIndex = rows.findIndex(
-        (row) => row.kind === "message" && row.message.streaming,
-      );
-    }
-
-    if (firstCurrentTurnRowIndex < 0) return firstTailRowIndex;
-
-    for (let index = firstCurrentTurnRowIndex - 1; index >= 0; index -= 1) {
-      const previousRow = rows[index];
-      if (!previousRow || previousRow.kind !== "message") continue;
-      if (previousRow.message.role === "user") {
-        return Math.min(index, firstTailRowIndex);
-      }
-      if (previousRow.message.role === "assistant" && !previousRow.message.streaming) {
-        break;
-      }
-    }
-
-    return Math.min(firstCurrentTurnRowIndex, firstTailRowIndex);
-  }, [activeTurnInProgress, activeTurnStartedAt, rows]);
-
-  const virtualizedRowCount = clamp(firstUnvirtualizedRowIndex, {
-    minimum: 0,
-    maximum: rows.length,
-  });
-  const virtualMeasurementScopeKey =
-    timelineWidthPx === null ? "width:unknown" : `width:${Math.round(timelineWidthPx)}`;
-
-  const rowVirtualizer = useVirtualizer({
-    count: virtualizedRowCount,
-    getScrollElement: () => scrollContainer,
-    // Scope cached row measurements to the current timeline width so offscreen
-    // rows do not keep stale heights after wrapping changes.
-    getItemKey: (index: number) => {
-      const rowId = rows[index]?.id ?? String(index);
-      return `${virtualMeasurementScopeKey}:${rowId}`;
-    },
-    estimateSize: (index: number) => {
-      const row = rows[index];
-      if (!row) return 96;
-      return estimateMessagesTimelineRowHeight(row, {
-        expandedWorkGroups,
-        timelineWidthPx,
-        turnDiffSummaryByAssistantMessageId,
-      });
-    },
-    measureElement: measureVirtualElement,
-    useAnimationFrameWithResizeObserver: true,
-    overscan: 8,
-  });
-  useEffect(() => {
-    if (timelineWidthPx === null) return;
-    rowVirtualizer.measure();
-  }, [rowVirtualizer, timelineWidthPx]);
-  useEffect(() => {
-    rowVirtualizer.shouldAdjustScrollPositionOnItemSizeChange = (item, _delta, instance) => {
-      const viewportHeight = instance.scrollRect?.height ?? 0;
-      const scrollOffset = instance.scrollOffset ?? 0;
-      const itemIntersectsViewport =
-        item.end > scrollOffset && item.start < scrollOffset + viewportHeight;
-      if (itemIntersectsViewport) {
-        return false;
-      }
-      const remainingDistance = instance.getTotalSize() - (scrollOffset + viewportHeight);
-      return remainingDistance > AUTO_SCROLL_BOTTOM_THRESHOLD_PX;
-    };
-    return () => {
-      rowVirtualizer.shouldAdjustScrollPositionOnItemSizeChange = undefined;
-    };
-  }, [rowVirtualizer]);
-  const pendingMeasureFrameRef = useRef<number | null>(null);
+  const canVirtualize = scrollContainer !== null && rows.length > 0;
+  const rowHeightEstimates = useMemo(
+    () =>
+      rows.map((row) =>
+        estimateMessagesTimelineRowHeight(row, {
+          expandedWorkGroups,
+          timelineWidthPx,
+          turnDiffSummaryByAssistantMessageId,
+        }),
+      ),
+    [expandedWorkGroups, rows, timelineWidthPx, turnDiffSummaryByAssistantMessageId],
+  );
+  const defaultRowHeight = rowHeightEstimates[0];
   const onTimelineImageLoad = useCallback(() => {
-    if (pendingMeasureFrameRef.current !== null) return;
-    pendingMeasureFrameRef.current = window.requestAnimationFrame(() => {
-      pendingMeasureFrameRef.current = null;
-      rowVirtualizer.measure();
-    });
-  }, [rowVirtualizer]);
-  useEffect(() => {
-    return () => {
-      const frame = pendingMeasureFrameRef.current;
-      if (frame !== null) {
-        window.cancelAnimationFrame(frame);
-      }
-    };
-  }, []);
-  useLayoutEffect(() => {
-    if (!onVirtualizerSnapshot) {
+    const activeScrollContainer = scrollContainer;
+    if (!activeScrollContainer || !isScrollContainerNearBottom(activeScrollContainer)) {
       return;
     }
-    onVirtualizerSnapshot({
-      totalSize: rowVirtualizer.getTotalSize(),
-      measurements: rowVirtualizer.measurementsCache
-        .slice(0, virtualizedRowCount)
-        .flatMap((measurement) => {
-          const row = rows[measurement.index];
-          if (!row) {
-            return [];
-          }
-          return [
-            {
-              id: row.id,
-              kind: row.kind,
-              index: measurement.index,
-              size: measurement.size,
-              start: measurement.start,
-              end: measurement.end,
-            },
-          ];
-        }),
+    window.requestAnimationFrame(() => {
+      virtuosoRef.current?.autoscrollToBottom();
     });
-  }, [onVirtualizerSnapshot, rowVirtualizer, rows, virtualizedRowCount]);
-
-  const virtualRows = rowVirtualizer.getVirtualItems();
-  const nonVirtualizedRows = rows.slice(virtualizedRowCount);
+  }, [scrollContainer]);
 
   const renderRowContent = (row: TimelineRow) => (
     <div
@@ -601,34 +475,28 @@
       data-timeline-root="true"
       className="mx-auto w-full min-w-0 max-w-3xl overflow-x-hidden"
     >
-      {virtualizedRowCount > 0 && (
-        <div className="relative" style={{ height: `${rowVirtualizer.getTotalSize()}px` }}>
-          {virtualRows.map((virtualRow: VirtualItem) => {
-            const row = rows[virtualRow.index];
-            if (!row) return null;
-
-            return (
-              <div
-                key={`virtual-row:${row.id}`}
-                data-index={virtualRow.index}
-                data-virtual-row-id={row.id}
-                data-virtual-row-kind={row.kind}
-                data-virtual-row-size={virtualRow.size}
-                data-virtual-row-start={virtualRow.start}
-                ref={rowVirtualizer.measureElement}
-                className="absolute left-0 top-0 w-full"
-                style={{ transform: `translateY(${virtualRow.start}px)` }}
-              >
-                {renderRowContent(row)}
-              </div>
-            );
-          })}
-        </div>
+      {canVirtualize && (
+        <Virtuoso
+          ref={virtuosoRef}
+          customScrollParent={scrollContainer ?? undefined}
+          data={rows}
+          computeItemKey={(_index, row) => row.id}
+          followOutput={(isAtBottom) => (isAtBottom ? "auto" : false)}
+          heightEstimates={rowHeightEstimates}
+          increaseViewportBy={{
+            top: VIRTUALIZED_VIEWPORT_PADDING_PX,
+            bottom: VIRTUALIZED_VIEWPORT_PADDING_PX,
+          }}
+          components={{
+            List: VirtualizedTimelineList,
+            Item: VirtualizedTimelineRowItem,
+          }}
+          itemContent={(_index, row) => renderRowContent(row)}
+          {...(typeof defaultRowHeight === "number" ? { defaultItemHeight: defaultRowHeight } : {})}
+        />
       )}
-
-      {nonVirtualizedRows.map((row) => (
-        <div key={`non-virtual-row:${row.id}`}>{renderRowContent(row)}</div>
-      ))}
+      {!canVirtualize &&
+        rows.map((row) => <div key={`row:${row.id}`}>{renderRowContent(row)}</div>)}
     </div>
   );
 });
@@ -638,6 +506,44 @@
 type TimelineWorkEntry = Extract<MessagesTimelineRow, { kind: "work" }>["groupedEntries"][number];
 type TimelineRow = MessagesTimelineRow;
 
+const VirtualizedTimelineList = forwardRef<HTMLDivElement, ListProps>(
+  function VirtualizedTimelineList(props, ref) {
+    const { children, style, ...domProps } = props;
+    return (
+      <div
+        {...domProps}
+        ref={ref}
+        style={{
+          ...style,
+          width: "100%",
+        }}
+      >
+        {children}
+      </div>
+    );
+  },
+);
+
+const VirtualizedTimelineRowItem = memo(function VirtualizedTimelineRowItem(
+  props: ItemProps<TimelineRow>,
+) {
+  const { children, item, style, ...domProps } = props;
+  return (
+    <div
+      {...domProps}
+      data-virtual-row-id={item.id}
+      data-virtual-row-kind={item.kind}
+      data-virtual-row-size={domProps["data-known-size"]}
+      style={{
+        ...style,
+        width: "100%",
+      }}
+    >
+      {children}
+    </div>
+  );
+});
+
 function formatWorkingTimer(startIso: string, endIso: string): string | null {
   const startedAtMs = Date.parse(startIso);
   const endedAtMs = Date.parse(endIso);