fix: stabilize messages timeline virtualization

apps/web/src/components/chat/ChangedFilesTree.test.tsx

@@ -0,0 +1,116 @@
+import { TurnId } from "@t3tools/contracts";
+import { renderToStaticMarkup } from "react-dom/server";
+import { describe, expect, it } from "vitest";
+
+import { ChangedFilesTree } from "./ChangedFilesTree";
+
+describe("ChangedFilesTree", () => {
+  it.each([
+    {
+      name: "a compacted single-chain directory",
+      files: [
+        { path: "apps/web/src/index.ts", additions: 2, deletions: 1 },
+        { path: "apps/web/src/main.ts", additions: 3, deletions: 0 },
+      ],
+      visibleLabels: ["apps/web/src"],
+      hiddenLabels: ["index.ts", "main.ts"],
+    },
+    {
+      name: "a branch point after a compacted prefix",
+      files: [
+        { path: "apps/server/src/git/Layers/GitCore.ts", additions: 4, deletions: 3 },
+        { path: "apps/server/src/provider/Layers/CodexAdapter.ts", additions: 7, deletions: 2 },
+      ],
+      visibleLabels: ["apps/server/src"],
+      hiddenLabels: ["git", "provider", "GitCore.ts", "CodexAdapter.ts"],
+    },
+    {
+      name: "mixed root files and nested compacted directories",
+      files: [
+        { path: "README.md", additions: 1, deletions: 0 },
+        { path: "packages/shared/src/git.ts", additions: 8, deletions: 2 },
+        { path: "packages/contracts/src/orchestration.ts", additions: 13, deletions: 3 },
+      ],
+      visibleLabels: ["README.md", "packages"],
+      hiddenLabels: ["shared/src", "contracts/src", "git.ts", "orchestration.ts"],
+    },
+  ])(
+    "renders $name collapsed on the first render when collapse-all is active",
+    ({ files, visibleLabels, hiddenLabels }) => {
+      const markup = renderToStaticMarkup(
+        <ChangedFilesTree
+          turnId={TurnId.makeUnsafe("turn-1")}
+          files={files}
+          allDirectoriesExpanded={false}
+          resolvedTheme="light"
+          onOpenTurnDiff={() => {}}
+        />,
+      );
+
+      for (const label of visibleLabels) {
+        expect(markup).toContain(label);
+      }
+      for (const label of hiddenLabels) {
+        expect(markup).not.toContain(label);
+      }
+    },
+  );
+
+  it.each([
+    {
+      name: "a compacted single-chain directory",
+      files: [
+        { path: "apps/web/src/index.ts", additions: 2, deletions: 1 },
+        { path: "apps/web/src/main.ts", additions: 3, deletions: 0 },
+      ],
+      visibleLabels: ["apps/web/src", "index.ts", "main.ts"],
+    },
+    {
+      name: "a branch point after a compacted prefix",
+      files: [
+        { path: "apps/server/src/git/Layers/GitCore.ts", additions: 4, deletions: 3 },
+        { path: "apps/server/src/provider/Layers/CodexAdapter.ts", additions: 7, deletions: 2 },
+      ],
+      visibleLabels: [
+        "apps/server/src",
+        "git/Layers",
+        "provider/Layers",
+        "GitCore.ts",
+        "CodexAdapter.ts",
+      ],
+    },
+    {
+      name: "mixed root files and nested compacted directories",
+      files: [
+        { path: "README.md", additions: 1, deletions: 0 },
+        { path: "packages/shared/src/git.ts", additions: 8, deletions: 2 },
+        { path: "packages/contracts/src/orchestration.ts", additions: 13, deletions: 3 },
+      ],
+      visibleLabels: [
+        "README.md",
+        "packages",
+        "shared/src",
+        "contracts/src",
+        "git.ts",
+        "orchestration.ts",
+      ],
+    },
+  ])(
+    "renders $name expanded on the first render when expand-all is active",
+    ({ files, visibleLabels }) => {
+      const markup = renderToStaticMarkup(
+        <ChangedFilesTree
+          turnId={TurnId.makeUnsafe("turn-1")}
+          files={files}
+          allDirectoriesExpanded
+          resolvedTheme="light"
+          onOpenTurnDiff={() => {}}
+        />,
+      );
+
+      for (const label of visibleLabels) {
+        expect(markup).toContain(label);
+      }
+    },
+  );
+});

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

@@ -1,12 +1,14 @@
 import { type TurnId } from "@t3tools/contracts";
-import { memo, useCallback, useEffect, useMemo, useState } from "react";
+import { memo, useCallback, useMemo, useState } from "react";
 import { type TurnDiffFileChange } from "../../types";
 import { buildTurnDiffTree, type TurnDiffTreeNode } from "../../lib/turnDiffTree";
 import { ChevronRightIcon, FolderIcon, FolderClosedIcon } from "lucide-react";
 import { cn } from "~/lib/utils";
 import { DiffStatLabel, hasNonZeroStat } from "./DiffStatLabel";
 import { VscodeEntryIcon } from "./VscodeEntryIcon";
 
+const EMPTY_DIRECTORY_OVERRIDES: Record<string, boolean> = {};
+
 export const ChangedFilesTree = memo(function ChangedFilesTree(props: {
   turnId: TurnId;
   files: ReadonlyArray<TurnDiffFileChange>;
@@ -20,40 +22,47 @@ export const ChangedFilesTree = memo(function ChangedFilesTree(props: {
     () => collectDirectoryPaths(treeNodes).join("\u0000"),
     [treeNodes],
   );
-  const allDirectoryExpansionState = useMemo(
-    () =>
-      buildDirectoryExpansionState(
-        directoryPathsKey ? directoryPathsKey.split("\u0000") : [],
-        allDirectoriesExpanded,
-      ),
-    [allDirectoriesExpanded, directoryPathsKey],
-  );
-  const [expandedDirectories, setExpandedDirectories] = useState<Record<string, boolean>>(() =>
-    buildDirectoryExpansionState(directoryPathsKey ? directoryPathsKey.split("\u0000") : [], true),
-  );
-  useEffect(() => {
-    setExpandedDirectories(allDirectoryExpansionState);
-  }, [allDirectoryExpansionState]);
+  const expansionStateKey = `${allDirectoriesExpanded ? "expanded" : "collapsed"}\u0000${directoryPathsKey}`;
+  const [directoryExpansionState, setDirectoryExpansionState] = useState<{
+    key: string;
+    overrides: Record<string, boolean>;
+  }>(() => ({
+    key: expansionStateKey,
+    overrides: {},
+  }));
+  const expandedDirectories =
+    directoryExpansionState.key === expansionStateKey
+      ? directoryExpansionState.overrides
+      : EMPTY_DIRECTORY_OVERRIDES;
 
-  const toggleDirectory = useCallback((pathValue: string, fallbackExpanded: boolean) => {
-    setExpandedDirectories((current) => ({
-      ...current,
-      [pathValue]: !(current[pathValue] ?? fallbackExpanded),
-    }));
-  }, []);
+  const toggleDirectory = useCallback(
+    (pathValue: string) => {
+      setDirectoryExpansionState((current) => {
+        const nextOverrides = current.key === expansionStateKey ? current.overrides : {};
+        return {
+          key: expansionStateKey,
+          overrides: {
+            ...nextOverrides,
+            [pathValue]: !(nextOverrides[pathValue] ?? allDirectoriesExpanded),
+          },
+        };
+      });
+    },
+    [allDirectoriesExpanded, expansionStateKey],
+  );
 
   const renderTreeNode = (node: TurnDiffTreeNode, depth: number) => {
     const leftPadding = 8 + depth * 14;
     if (node.kind === "directory") {
-      const isExpanded = expandedDirectories[node.path] ?? depth === 0;
+      const isExpanded = expandedDirectories[node.path] ?? allDirectoriesExpanded;
       return (
         <div key={`dir:${node.path}`}>
           <button
             type="button"
             data-scroll-anchor-ignore
             className="group flex w-full items-center gap-1.5 rounded-md py-1 pr-2 text-left hover:bg-background/80"
             style={{ paddingLeft: `${leftPadding}px` }}
-            onClick={() => toggleDirectory(node.path, depth === 0)}
+            onClick={() => toggleDirectory(node.path)}
           >
             <ChevronRightIcon
               aria-hidden="true"
@@ -124,14 +133,3 @@ function collectDirectoryPaths(nodes: ReadonlyArray<TurnDiffTreeNode>): string[]
   }
   return paths;
 }
-
-function buildDirectoryExpansionState(
-  directoryPaths: ReadonlyArray<string>,
-  expanded: boolean,
-): Record<string, boolean> {
-  const expandedState: Record<string, boolean> = {};
-  for (const directoryPath of directoryPaths) {
-    expandedState[directoryPath] = expanded;
-  }
-  return expandedState;
-}