fix(app): preserve chat rows across virtual layout changes (#6172)

* fix(app): preserve timeline rows across virtual layout changes Keep recent and virtual history rows under one keyed host so entering virtual history preserves their DOM and loaded image state. Initialize the virtual range from the mounted viewport to avoid an empty handoff commit. Measure committed row dimensions before accepting reading geometry, including child layout changes flushed during resize observation. Separate row content from position updates while retaining live revision signals. Add a regression covering row identity through virtualization and the position of content following a growing row. Keep the existing browser assertions unchanged. * test(app): verify child state survives timeline virtualization
2026-10-14 13:22:37 +00:00 · 2026-10-06 12:31:49 +01:00
parent b7c8be4081
commit 543e3712e1
2 changed files with 186 additions and 39 deletions
@@ -32,7 +32,7 @@ vi.mock("react-native-unistyles", () => ({
  withUnistyles: (Component: React.ComponentType) => Component,
 }));

-function userMessage(index: number): StreamItem {
+function userMessage(index: number): Extract<StreamItem, { kind: "user_message" }> {
  return {
    kind: "user_message",
    id: `message-${index}`,
@@ -55,6 +55,30 @@ function createRenderers(onRowRender: () => void): StreamSegmentRenderers {
  };
 }

+function MeasuredTestRow({ item }: { item: StreamItem }) {
+  const [loaded, setLoaded] = React.useState(false);
+  const load = React.useCallback(() => setLoaded(true), []);
+  const height = item.kind === "user_message" && item.text === "grown" ? 120 : 24;
+  const measure = React.useCallback(
+    (node: HTMLDivElement | null) => {
+      if (!node?.parentElement) return;
+      Object.defineProperty(node.parentElement, "offsetHeight", {
+        configurable: true,
+        value: height,
+      });
+    },
+    [height],
+  );
+  return (
+    <div ref={measure}>
+      <button type="button" onClick={load}>
+        {loaded ? "Loaded" : "Load"}
+      </button>
+      {item.id}
+    </div>
+  );
+}
+
 describe("createWebStreamStrategy", () => {
  let root: Root | null = null;
  let container: HTMLDivElement | null = null;
@@ -155,6 +179,75 @@ describe("createWebStreamStrategy", () => {
    expect(rowRenderCount.mock.calls.length).toBeLessThanOrEqual(historyVirtualized.length);
  });

+  it("retains a row's DOM and local state when recent history becomes virtualized", () => {
+    const strategy = createWebStreamStrategy({ isMobileBreakpoint: true });
+    container = document.createElement("div");
+    document.body.appendChild(container);
+    root = createRoot(container);
+    const item = userMessage(0);
+    const renderRow = (rowItem: StreamItem) => <MeasuredTestRow item={rowItem} />;
+    const input: StreamRenderInput = {
+      agentId: "agent",
+      segments: { historyVirtualized: [], historyMounted: [item], liveHead: [] },
+      boundary: { hasVirtualizedHistory: false, hasMountedHistory: true, hasLiveHead: false },
+      renderers: {
+        renderHistoryVirtualizedRow: renderRow,
+        renderHistoryMountedRow: renderRow,
+        renderLiveHeadRow: renderRow,
+        renderLiveAuxiliary: () => null,
+      },
+      listEmptyComponent: null,
+      viewportRef: React.createRef<StreamViewportHandle>(),
+      routeBottomAnchorRequest: null,
+      isAuthoritativeHistoryReady: true,
+      onNearBottomChange: () => {},
+      onNearHistoryStart: () => true,
+      isLoadingOlderHistory: false,
+      hasOlderHistory: false,
+      olderHistoryProgressKey: null,
+      scrollEnabled: true,
+      listStyle: null,
+      baseListContentContainerStyle: null,
+      forwardListContentContainerStyle: null,
+      contentMaxWidth: DEFAULT_CONTENT_MAX_WIDTH,
+    };
+    act(() => root?.render(strategy.render(input)));
+    const row = container.querySelector('[data-history-row-id="message-0"]')!;
+    act(() => row.querySelector("button")!.click());
+    expect(row.querySelector("button")!.textContent).toBe("Loaded");
+    const viewport = container.querySelector('[data-testid="agent-chat-scroll"]')!;
+    Object.defineProperty(viewport, "clientHeight", { configurable: true, value: 500 });
+    act(() =>
+      root?.render(
+        strategy.render({
+          ...input,
+          segments: { historyVirtualized: [item], historyMounted: [userMessage(1)], liveHead: [] },
+          boundary: { ...input.boundary, hasVirtualizedHistory: true },
+        }),
+      ),
+    );
+    expect(container.querySelector('[data-history-row-id="message-0"]')).toBe(row);
+    expect(row.isConnected).toBe(true);
+    expect(row.querySelector("button")!.textContent).toBe("Loaded");
+    act(() =>
+      root?.render(
+        strategy.render({
+          ...input,
+          segments: {
+            historyVirtualized: [{ ...item, text: "grown" }, userMessage(1)],
+            historyMounted: [userMessage(2)],
+            liveHead: [],
+          },
+          boundary: { ...input.boundary, hasVirtualizedHistory: true },
+        }),
+      ),
+    );
+    expect(
+      (container.querySelector('[data-history-row-id="message-1"]') as HTMLElement).style.transform,
+    ).toBe("translateY(120px)");
+    expect(row.querySelector("button")!.textContent).toBe("Loaded");
+  });
+
  it("keeps the timeline width stable with an overlay scrollbar", () => {
    const strategy = createWebStreamStrategy({ isMobileBreakpoint: false });
    container = document.createElement("div");
@@ -20,6 +20,7 @@ import { withUnistyles } from "react-native-unistyles";
 import { useRetainedPanelActive } from "@/components/retained-panel";
 import { LoadingSpinner } from "@/components/ui/loading-spinner";
 import { useStableEvent } from "@/hooks/use-stable-event";
+import type { StreamItem } from "@/types/stream";
 import type { Theme } from "@/styles/theme";
 import { WEB_SCROLLBAR_SIZE_PX } from "@/styles/web-scrollbar";
 import { DomOverlayScrollbar } from "@/components/ui/overlay-scrollbar/dom-overlay-scrollbar";
@@ -76,6 +77,8 @@ const historyStartSlotStyle: CSSProperties = {
  flexShrink: 0,
 };

+const streamRowsHostStyle: CSSProperties = { position: "relative", width: "100%" };
+
 const streamRowStyle: CSSProperties = {
  display: "flex",
  flexDirection: "column",
@@ -268,6 +271,23 @@ function isScrollContainerOverscrolledPastBottom(
  return getScrollContainerDistanceFromBottom(scrollContainer) < -BOTTOM_OVERSCROLL_TOLERANCE_PX;
 }

+// Position updates do not change a row's content. Keep renderer work separate
+// from virtual measurements while retaining the same component across lanes.
+const StreamRowContent = React.memo(function StreamRowContent({
+  item,
+  index,
+  items,
+  render,
+}: {
+  revision?: StreamRenderInput["liveHeadRowRevision"];
+  item: StreamItem;
+  index: number;
+  items: StreamItem[];
+  render: StreamRenderInput["renderers"]["renderHistoryVirtualizedRow"];
+}) {
+  return render(item, index, items);
+});
+
 function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: boolean }) {
  const {
    segments: inputSegments,
@@ -425,6 +445,10 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
  const rowVirtualizer = useVirtualizer({
    count: segments.historyVirtualized.length,
    enabled: shouldUseVirtualizer,
+    initialRect: {
+      width: scrollContainerRef.current?.clientWidth ?? 0,
+      height: scrollContainerRef.current?.clientHeight ?? 0,
+    },
    getScrollElement: () => scrollContainerRef.current,
    getItemKey: (index: number) => segments.historyVirtualized[index]?.id ?? index,
    estimateSize: (index: number) => {
@@ -536,25 +560,31 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
    scheduler.schedule();
  });

-  // ResizeObserver runs before paint. Batch measurements and commit React's new
-  // positions here so a resized image cannot overlap rows for an extra frame.
-  // Use resizeItem exclusively; measureElement would install a second observer.
+  const measureVirtualizedRows = useStableEvent(() => {
+    let resized = false;
+    for (const element of contentRef.current?.querySelectorAll<HTMLElement>(
+      "[data-history-row-id][data-index]",
+    ) ?? []) {
+      const index = Number(element.dataset.index);
+      const height = measureVirtualElement(element, undefined, rowVirtualizer);
+      if (rowVirtualizer.measurementsCache[index]?.size === height) continue;
+      rowVirtualizer.resizeItem(index, height);
+      resized = true;
+    }
+    return resized;
+  });
+
+  // Measure the committed DOM rather than queued observer entries: flushSync can
+  // also commit an image's pending intrinsic dimensions, invalidating those entries.
  const virtualRowObserver = useMemo(
    () =>
-      new ResizeObserver((entries) => {
+      new ResizeObserver(() => {
        if (!isActiveRef.current) return;
        flushSync(() => {
-          for (const entry of entries) {
-            const element = entry.target as HTMLElement;
-            if (!element.isConnected) continue;
-            rowVirtualizer.resizeItem(
-              Number(element.dataset.index),
-              measureVirtualElement(element, entry, rowVirtualizer),
-            );
-          }
+          measureVirtualizedRows();
        });
      }),
-    [rowVirtualizer],
+    [measureVirtualizedRows],
  );
  const measureVirtualizedRowElement = useCallback(
    (node: HTMLDivElement | null) => {
@@ -731,8 +761,18 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
  // Apply row positions and their scroll correction in the same paint, including
  // the ordinary-to-virtualized handoff and measurements of newly mounted rows.
  useLayoutEffect(() => {
+    // Commit positions for every resized row before accepting reading geometry.
+    // Child layout can change during this commit, ahead of ResizeObserver delivery.
+    if (!isActive || measureVirtualizedRows()) return;
    reconcileReadingPosition();
-  }, [contentWidth, isActive, segments, virtualRows, reconcileReadingPosition]);
+  }, [
+    contentWidth,
+    isActive,
+    segments,
+    virtualRows,
+    measureVirtualizedRows,
+    reconcileReadingPosition,
+  ]);

  const stopFollowingOutputFromUserIntent = useStableEvent(() => {
    cancelPendingStickToBottom();
@@ -915,6 +955,9 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
      evaluateHistoryStart();
    }
    const observer = new ResizeObserver(() => {
+      flushSync(() => {
+        measureVirtualizedRows();
+      });
      const nextGeometry = getObservedViewportGeometry(scrollContainer);
      if (pendingResumeGeometryCheckRef.current) {
        pendingResumeGeometryCheckRef.current = false;
@@ -953,6 +996,7 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
    scheduleHistoryStartPrependSettle,
    scheduleStickToBottom,
    updateScrollMetrics,
+    measureVirtualizedRows,
  ]);

  useLayoutEffect(() => {
@@ -1171,13 +1215,13 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
    }),
    [],
  );
-  const virtualRowsContainerStyle = useMemo((): CSSProperties => {
-    return {
-      position: "relative",
-      width: "100%",
+  const virtualRowsContainerStyle = useMemo(
+    (): CSSProperties => ({
      height: virtualTotalSize,
-    };
-  }, [virtualTotalSize]);
+      flexShrink: 0,
+    }),
+    [virtualTotalSize],
+  );
  const renderVirtualRowStyle = useCallback(
    (start: number): CSSProperties => ({
      position: "absolute",
@@ -1198,12 +1242,16 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
        data-message-id={getStreamItemMessageId(item)}
        style={streamRowStyle}
      >
-        {renderHistoryMountedRow(item, index, segments.historyMounted)}
+        <StreamRowContent
+          item={item}
+          index={index}
+          items={segments.historyMounted}
+          render={renderHistoryMountedRow}
+        />
      </div>
    ));
  }, [renderHistoryMountedRow, segments.historyMounted]);
  const liveHeadRows = useMemo(() => {
-    void liveHeadRowRevision;
    return segments.liveHead.map((item, index) => (
      <div
        key={item.id}
@@ -1211,7 +1259,13 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
        data-message-id={getStreamItemMessageId(item)}
        style={streamRowStyle}
      >
-        {renderLiveHeadRow(item, index, segments.liveHead)}
+        <StreamRowContent
+          item={item}
+          index={index}
+          items={segments.liveHead}
+          revision={liveHeadRowRevision}
+          render={renderLiveHeadRow}
+        />
      </div>
    ));
  }, [liveHeadRowRevision, renderLiveHeadRow, segments.liveHead]);
@@ -1251,13 +1305,12 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
      >
        <div ref={handleContentRef} style={contentContainerStyle}>
          {historyStartSlot}
-          {shouldUseVirtualizer ? (
-            <div style={virtualRowsContainerStyle}>
-              {virtualRows.map((virtualRow) => {
+          <div style={streamRowsHostStyle}>
+            <div style={virtualRowsContainerStyle} />
+            {[
+              ...virtualRows.map((virtualRow) => {
                const item = segments.historyVirtualized[virtualRow.index];
-                if (!item) {
-                  return null;
-                }
+                if (!item) return null;
                return (
                  <div
                    key={virtualRow.key}
@@ -1267,17 +1320,18 @@ function WebStreamViewport(props: StreamRenderInput & { isMobileBreakpoint: bool
                    ref={measureVirtualizedRowElement}
                    style={renderVirtualRowStyle(virtualRow.start)}
                  >
-                    {renderHistoryVirtualizedRow(
-                      item,
-                      virtualRow.index,
-                      segments.historyVirtualized,
-                    )}
+                    <StreamRowContent
+                      item={item}
+                      index={virtualRow.index}
+                      items={segments.historyVirtualized}
+                      render={renderHistoryVirtualizedRow}
+                    />
                  </div>
                );
-              })}
-            </div>
-          ) : null}
-          {mountedRows}
+              }),
+              ...mountedRows,
+            ]}
+          </div>
          {liveAuxiliary}
          {shouldRenderEmpty ? listEmptyComponent : null}
        </div>