fix(web): restore persisted file tree expansions (#1892)

apps/web/components/task/file-browser-hooks.ts

@@ -9,25 +9,20 @@ import { getFilesPanelExpandedPaths, setFilesPanelExpandedPaths } from "@/lib/lo
 import { useTree, type VisibleRow } from "@/hooks/use-tree";
 import { mergeTreeNodes } from "./file-browser-parts";
 import { compareTreeNodes, sortRootChildren } from "./file-tree-utils";
+import { restoredExpandedPaths } from "./file-browser-restore";
+import { useTreeLoader } from "./file-browser-tree-loader";
 import { createDebugLogger, isDebug } from "@/lib/debug/log";
 
 const debugLoad = createDebugLogger("file-browser:load");
 const debugChanges = createDebugLogger("file-browser:changes");
 
 const FB_GET_PATH = (n: FileTreeNode) => n.path;
-// Children are sorted (dirs first, then files, alphabetically) on every
-// access so the flat visibleRows list comes out in the same order the
-// recursive render produced. The backend doesn't guarantee order and tree
-// mutations (rename / create) can insert arbitrarily.
 const FB_GET_CHILDREN = (n: FileTreeNode) =>
   n.children ? [...n.children].sort(compareTreeNodes) : undefined;
 const FB_IS_DIR = (n: FileTreeNode) => n.is_dir;
 
 export type FileBrowserRow = VisibleRow<FileTreeNode>;
 
-const MAX_RETRY_ATTEMPTS = 4;
-const RETRY_DELAYS_MS = [1000, 2000, 5000, 10000];
-
 export type LoadState = "loading" | "waiting" | "loaded" | "manual" | "error";
 
 /** Hook encapsulating file search state and handlers. */
@@ -39,14 +34,12 @@ export function useFileBrowserSearch(sessionId: string) {
   const searchTimeoutRef = useRef<NodeJS.Timeout | null>(null);
   const searchInputRef = useRef<HTMLInputElement>(null);
 
-  // Focus search input when search opens
   useEffect(() => {
     if (isSearchActive && searchInputRef.current) {
       searchInputRef.current.focus();
     }
   }, [isSearchActive]);
 
-  // Clear search when closing
   useEffect(() => {
     if (!isSearchActive) {
       setLocalSearchQuery("");
@@ -56,7 +49,6 @@ export function useFileBrowserSearch(sessionId: string) {
     }
   }, [isSearchActive]);
 
-  // Cleanup on unmount
   useEffect(() => {
     return () => {
       if (searchTimeoutRef.current) clearTimeout(searchTimeoutRef.current);
@@ -123,7 +115,6 @@ export function applyFileChanges(ctx: {
   const { client, sessionId, expandedPaths, changes, setTree, setLoadState } = ctx;
   const foldersToRefresh = new Set<string>();
   for (const change of changes) {
-    // `refresh` events have empty path; refresh root + every expanded folder under the affected repo so new files show up.
     if (change.operation === "refresh") {
       foldersToRefresh.add("");
       const repo = change.repository_name;
@@ -191,7 +182,6 @@ export function applyFileChanges(ctx: {
           }
           return node.children ? { ...node, children: node.children.map(patchNode) } : node;
         };
-        // Skip root: already merged above; re-matching folderUpdates.has("") here would overwrite preserved subtrees.
         return { ...updated, children: updated.children?.map(patchNode) };
       });
       setLoadState("loaded");
@@ -234,109 +224,10 @@ function useLoadingTimers() {
   return { visibleLoadingPaths, showLoading, hideLoading, isLoading };
 }
 
-type TreeLoaderContext = {
-  sessionId: string;
-  clearRetryTimer: () => void;
-  retryAttemptRef: React.MutableRefObject<number>;
-  retryTimerRef: React.MutableRefObject<NodeJS.Timeout | null>;
-  setTree: React.Dispatch<React.SetStateAction<FileTreeNode | null>>;
-  setIsLoadingTree: React.Dispatch<React.SetStateAction<boolean>>;
-  setLoadState: React.Dispatch<React.SetStateAction<LoadState>>;
-  setLoadError: React.Dispatch<React.SetStateAction<string | null>>;
-};
-
-// Thin wrapper so loadTree callers don't each pay a complexity point for isDebug().
 function logLoad(event: string, data: Record<string, unknown>) {
   if (isDebug()) debugLoad(event, data);
 }
 
-function useTreeLoader(ctx: TreeLoaderContext) {
-  const {
-    clearRetryTimer,
-    retryAttemptRef,
-    retryTimerRef,
-    setTree,
-    setIsLoadingTree,
-    setLoadState,
-    setLoadError,
-  } = ctx;
-  // Use ref for sessionId so loadTree stays stable across session switches.
-  // This prevents the tree-clearing effect from re-firing when only sessionId changes.
-  const sessionIdRef = useRef(ctx.sessionId);
-  sessionIdRef.current = ctx.sessionId;
-  const loadInFlightRef = useRef(false);
-  const loadTree = useCallback(
-    async (options?: { resetRetry?: boolean }) => {
-      if (loadInFlightRef.current) {
-        logLoad("skip-in-flight", { sessionId: sessionIdRef.current });
-        return;
-      }
-      loadInFlightRef.current = true;
-      setIsLoadingTree(true);
-      setLoadState("loading");
-      setLoadError(null);
-      if (options?.resetRetry) {
-        retryAttemptRef.current = 0;
-        clearRetryTimer();
-      }
-      logLoad("start", {
-        sessionId: sessionIdRef.current,
-        resetRetry: options?.resetRetry === true,
-        retryAttempt: retryAttemptRef.current,
-      });
-      try {
-        const client = getWebSocketClient();
-        if (!client) throw new Error("WebSocket client not available");
-        const response = await requestFileTree(client, sessionIdRef.current, "", 1);
-        setTree(response.root ?? null);
-        setLoadState("loaded");
-        retryAttemptRef.current = 0;
-        clearRetryTimer();
-        logLoad("loaded", {
-          sessionId: sessionIdRef.current,
-          rootPath: response.root?.path ?? null,
-          children: response.root?.children?.length ?? 0,
-        });
-      } catch (error) {
-        const message = error instanceof Error ? error.message : "Failed to load file tree";
-        setLoadError(message);
-        if (retryAttemptRef.current < MAX_RETRY_ATTEMPTS) {
-          const delay =
-            RETRY_DELAYS_MS[Math.min(retryAttemptRef.current, RETRY_DELAYS_MS.length - 1)];
-          retryAttemptRef.current += 1;
-          setLoadState("waiting");
-          clearRetryTimer();
-          logLoad("retry", {
-            sessionId: sessionIdRef.current,
-            attempt: retryAttemptRef.current,
-            delayMs: delay,
-            error: message,
-          });
-          retryTimerRef.current = setTimeout(() => {
-            void loadTree();
-          }, delay);
-        } else {
-          setLoadState("manual");
-          logLoad("gave-up", { sessionId: sessionIdRef.current, error: message });
-        }
-      } finally {
-        setIsLoadingTree(false);
-        loadInFlightRef.current = false;
-      }
-    },
-    [
-      clearRetryTimer,
-      retryAttemptRef,
-      retryTimerRef,
-      setIsLoadingTree,
-      setLoadError,
-      setLoadState,
-      setTree,
-    ],
-  );
-  return loadTree;
-}
-
 type TreeLoadEffectsContext = {
   sessionId: string;
   effectiveResetKey: string;
@@ -346,16 +237,21 @@ type TreeLoadEffectsContext = {
   treeRef: React.MutableRefObject<FileTreeNode | null>;
   retryAttemptRef: React.MutableRefObject<number>;
   hasInitializedExpandedRef: React.MutableRefObject<string | null>;
+  restoreExpandedPathsRef: React.MutableRefObject<string[]>;
+  expandedPathsRef: React.MutableRefObject<ReadonlySet<string>>;
   clearRetryTimer: () => void;
-  loadTree: (options?: { resetRetry?: boolean }) => Promise<void> | void;
+  loadTree: (options?: {
+    resetRetry?: boolean;
+    restoreExpandedPaths?: string[];
+  }) => Promise<void> | void;
   setTree: React.Dispatch<React.SetStateAction<FileTreeNode | null>>;
   setIsLoadingTree: React.Dispatch<React.SetStateAction<boolean>>;
   setLoadState: React.Dispatch<React.SetStateAction<LoadState>>;
   setLoadError: React.Dispatch<React.SetStateAction<string | null>>;
-  setExpandedPaths: (paths: Set<string>) => void;
+  setExpandedPaths: React.Dispatch<React.SetStateAction<Set<string>>>;
+  lastResetKeyRef: React.MutableRefObject<string | null>;
 };
 
-/** Owns the two effects that drive initial load and the waiting→ready flip. */
 function useTreeLoadEffects(ctx: TreeLoadEffectsContext) {
   const {
     sessionId,
@@ -366,6 +262,8 @@ function useTreeLoadEffects(ctx: TreeLoadEffectsContext) {
     treeRef,
     retryAttemptRef,
     hasInitializedExpandedRef,
+    restoreExpandedPathsRef,
+    expandedPathsRef,
     clearRetryTimer,
     loadTree,
     setTree,
@@ -373,19 +271,28 @@ function useTreeLoadEffects(ctx: TreeLoadEffectsContext) {
     setLoadState,
     setLoadError,
     setExpandedPaths,
+    lastResetKeyRef,
   } = ctx;
 
   useEffect(() => {
-    setTree(null);
-    setIsLoadingTree(true);
-    setLoadState(agentctlIsReadyRef.current ? "loading" : "waiting");
-    setLoadError(null);
-    retryAttemptRef.current = 0;
+    const resetKeyChanged = lastResetKeyRef.current !== effectiveResetKey;
+    lastResetKeyRef.current = effectiveResetKey;
     clearRetryTimer();
-    hasInitializedExpandedRef.current = null;
-    const savedPaths = getFilesPanelExpandedPaths(effectiveResetKey);
-    setExpandedPaths(savedPaths.length > 0 ? new Set(savedPaths) : new Set());
-    if (savedPaths.length > 0) hasInitializedExpandedRef.current = effectiveResetKey;
+    retryAttemptRef.current = 0;
+    if (resetKeyChanged) {
+      setTree(null);
+      setIsLoadingTree(true);
+      setLoadState(agentctlIsReadyRef.current ? "loading" : "waiting");
+      setLoadError(null);
+      hasInitializedExpandedRef.current = null;
+      const savedPaths = restoredExpandedPaths(getFilesPanelExpandedPaths(effectiveResetKey));
+      restoreExpandedPathsRef.current = savedPaths;
+      setExpandedPaths(savedPaths.length > 0 ? new Set(savedPaths) : new Set());
+    }
+    const savedPaths = resetKeyChanged
+      ? restoreExpandedPathsRef.current
+      : Array.from(expandedPathsRef.current);
+    restoreExpandedPathsRef.current = savedPaths;
     logLoad("init-effect", {
       sessionId,
       effectiveResetKey,
@@ -393,20 +300,15 @@ function useTreeLoadEffects(ctx: TreeLoadEffectsContext) {
       savedPaths: savedPaths.length,
       willLoad: agentctlIsReadyRef.current,
     });
-    if (agentctlIsReadyRef.current) void loadTree({ resetRetry: true });
-    else setIsLoadingTree(false);
+    if (agentctlIsReadyRef.current) {
+      void loadTree({ resetRetry: true, restoreExpandedPaths: savedPaths });
+    } else setIsLoadingTree(false);
     return () => {
       clearRetryTimer();
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps -- refs intentionally omitted
-  }, [clearRetryTimer, loadTree, effectiveResetKey, setExpandedPaths]);
+  }, [clearRetryTimer, loadTree, effectiveResetKey, sessionId, setExpandedPaths]);
 
-  // Fire the initial load on the waiting → ready transition. `loadState` is
-  // read via a ref so a failed load (which sets state back to "waiting") does
-  // not re-fire this effect and cancel the retry timer via `resetRetry: true`.
-  // The `treeRef.current` check alongside "loaded" is load-bearing:
-  // `applyFileChanges` can flip state to "loaded" while the tree is still
-  // null, and without the tree check this effect would skip the real load.
   useEffect(() => {
     let reason: string | null = null;
     if (!agentctlIsReady) reason = "agentctl-not-ready";
@@ -417,7 +319,7 @@ function useTreeLoadEffects(ctx: TreeLoadEffectsContext) {
       return;
     }
     logLoad("ready-flip", { sessionId, loadState: loadStateRef.current });
-    void loadTree({ resetRetry: true });
+    void loadTree({ resetRetry: true, restoreExpandedPaths: restoreExpandedPathsRef.current });
     // eslint-disable-next-line react-hooks/exhaustive-deps -- refs intentionally omitted
   }, [agentctlIsReady, loadTree, sessionId]);
 }
@@ -433,9 +335,6 @@ function useFileChangeSubscription({
   setTree: React.Dispatch<React.SetStateAction<FileTreeNode | null>>;
   setLoadState: React.Dispatch<React.SetStateAction<LoadState>>;
 }) {
-  // expandedPaths is read via ref so that toggling a directory (which mints a
-  // new Set) does not tear down and re-attach this WS listener — any event
-  // arriving during the swap would otherwise be silently dropped.
   useEffect(() => {
     const client = getWebSocketClient();
     if (!client) return;
@@ -464,22 +363,11 @@ function useFileChangeSubscription({
   }, [sessionIdRef, expandedPathsRef, setTree, setLoadState]);
 }
 
-/**
- * Hook for tree loading with retry logic, file-change subscription, and expanded state.
- * @param sessionId - Session ID for API calls (agentctl routing).
- * @param resetKey - Optional stable key (e.g. environmentId) that controls when the tree
- *   does a full reset. When sessions share the same environment, this prevents the tree
- *   from flashing on tab switch.
- */
 export function useFileBrowserTree(sessionId: string, resetKey?: string) {
   const effectiveResetKey = resetKey ?? sessionId;
   const [tree, setTree] = useState<FileTreeNode | null>(null);
-  // Expansion + flat visibleRows come from the shared useTree hook. tree is
-  // a single root node with children; useTree wants an array, so feed it the
-  // root's children (pre-sorted — see sortRootChildren).
-  const treeNodes = useMemo(() => sortRootChildren(tree), [tree]);
   const treeApi = useTree<FileTreeNode>({
-    nodes: treeNodes,
+    nodes: useMemo(() => sortRootChildren(tree), [tree]),
     getPath: FB_GET_PATH,
     getChildren: FB_GET_CHILDREN,
     isDir: FB_IS_DIR,
@@ -487,41 +375,40 @@ export function useFileBrowserTree(sessionId: string, resetKey?: string) {
   const expandedPaths = treeApi.expanded;
   const setExpandedPaths = treeApi.setExpanded;
   const visibleRows = treeApi.visibleRows;
-  // Stable ref over `expandedPaths` so the WS file-change subscription does
-  // not re-attach on every toggle (each toggle mints a new Set reference).
   const expandedPathsRef = useRef<ReadonlySet<string>>(expandedPaths);
   expandedPathsRef.current = expandedPaths;
   const [isLoadingTree, setIsLoadingTree] = useState(true);
   const [loadState, setLoadState] = useState<LoadState>("loading");
   const [loadError, setLoadError] = useState<string | null>(null);
   const hasInitializedExpandedRef = useRef<string | null>(null);
+  const restoreExpandedPathsRef = useRef<string[]>([]);
+  const lastResetKeyRef = useRef<string | null>(null);
   const retryAttemptRef = useRef(0);
   const retryTimerRef = useRef<NodeJS.Timeout | null>(null);
   const sessionIdRef = useRef(sessionId);
   sessionIdRef.current = sessionId;
   const agentctlStatus = useSessionAgentctl(sessionId);
   const { visibleLoadingPaths, showLoading, hideLoading, isLoading } = useLoadingTimers();
-
   const clearRetryTimer = useCallback(() => {
     if (retryTimerRef.current) {
       clearTimeout(retryTimerRef.current);
       retryTimerRef.current = null;
     }
   }, []);
-
   const loadTree = useTreeLoader({
     sessionId,
+    effectiveResetKey,
     clearRetryTimer,
     retryAttemptRef,
     retryTimerRef,
+    restoreExpandedPathsRef,
+    hasInitializedExpandedRef,
     setTree,
+    setExpandedPaths,
     setIsLoadingTree,
     setLoadState,
     setLoadError,
   });
-
-  // Refs for effects below to read without adding deps — avoids infinite
-  // retry loops and wipe-on-reconnect when `isReady`/`loadState`/`tree` tick.
   const agentctlIsReadyRef = useRef(agentctlStatus.isReady);
   const loadStateRef = useRef(loadState);
   const treeRef = useRef(tree);
@@ -537,6 +424,8 @@ export function useFileBrowserTree(sessionId: string, resetKey?: string) {
     treeRef,
     retryAttemptRef,
     hasInitializedExpandedRef,
+    restoreExpandedPathsRef,
+    expandedPathsRef,
     clearRetryTimer,
     loadTree,
     setTree,
@@ -544,25 +433,13 @@ export function useFileBrowserTree(sessionId: string, resetKey?: string) {
     setLoadState,
     setLoadError,
     setExpandedPaths,
+    lastResetKeyRef,
   });
-
   useEffect(() => {
-    if (!tree || isLoadingTree || hasInitializedExpandedRef.current === effectiveResetKey) return;
-    hasInitializedExpandedRef.current = effectiveResetKey;
-    setExpandedPaths(new Set());
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [tree?.children?.length, isLoadingTree, effectiveResetKey]);
-
-  useEffect(() => {
-    if (expandedPaths.size > 0 || hasInitializedExpandedRef.current === effectiveResetKey) {
-      setFilesPanelExpandedPaths(effectiveResetKey, Array.from(expandedPaths));
-    }
-  }, [expandedPaths, effectiveResetKey]);
-
+    if (isLoadingTree || hasInitializedExpandedRef.current !== effectiveResetKey) return;
+    setFilesPanelExpandedPaths(effectiveResetKey, Array.from(expandedPaths));
+  }, [expandedPaths, effectiveResetKey, isLoadingTree]);
   useFileChangeSubscription({ sessionIdRef, expandedPathsRef, setTree, setLoadState });
-
-  const collapseAll = treeApi.collapseAll;
-
   return {
     tree,
     setTree,
@@ -577,13 +454,10 @@ export function useFileBrowserTree(sessionId: string, resetKey?: string) {
     showLoading,
     hideLoading,
     isLoading,
-    collapseAll,
+    collapseAll: treeApi.collapseAll,
   };
 }
 
-// useScrollPersistence, loadNodeChildren, toggleFolderExpand, fetchAndOpenFile,
-// and the ToggleFolderExpandDeps type live in ./file-browser-actions.ts to keep
-// this file within the 600-line lint limit.
 export {
   useScrollPersistence,
   loadNodeChildren,

apps/web/components/task/file-browser-restore-expanded.test.tsx

@@ -0,0 +1,129 @@
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const requestFileTreeMock = vi.fn();
+const getWebSocketClientMock = vi.fn(() => ({ on: vi.fn(() => () => {}) }));
+
+vi.mock("@/lib/ws/connection", () => ({
+  getWebSocketClient: () => getWebSocketClientMock(),
+}));
+vi.mock("@/lib/ws/workspace-files", () => ({
+  requestFileTree: (...args: unknown[]) => requestFileTreeMock(...args),
+  requestFileContent: vi.fn(),
+  searchWorkspaceFiles: vi.fn(),
+}));
+vi.mock("@/hooks/domains/session/use-session-agentctl", () => ({
+  useSessionAgentctl: () => ({ isReady: true }),
+}));
+
+import { useFileBrowserTree } from "./file-browser-hooks";
+
+const ROOT_PATH = "";
+const CODEX_PATH = ".codex";
+const AGENTS_PATH = `${CODEX_PATH}/agents`;
+const CONFIG_PATH = `${AGENTS_PATH}/config.toml`;
+const EXPANDED_PATHS = [CODEX_PATH, AGENTS_PATH];
+const ROOT = { name: "", path: ROOT_PATH, is_dir: true, size: 0 };
+const CODEX = { name: CODEX_PATH, path: CODEX_PATH, is_dir: true, size: 0 };
+const AGENTS = { name: "agents", path: AGENTS_PATH, is_dir: true, size: 0 };
+const CONFIG = { name: "config.toml", path: CONFIG_PATH, is_dir: false, size: 1 };
+const SESSION = "session-1";
+const ENVIRONMENT = "environment-1";
+const STORAGE_KEY = `kandev.filesPanel.expanded.${ENVIRONMENT}`;
+
+describe("useFileBrowserTree persisted expansion", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    sessionStorage.clear();
+    sessionStorage.setItem(STORAGE_KEY, JSON.stringify(EXPANDED_PATHS));
+    requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+      if (path === ROOT_PATH) return Promise.resolve({ root: { ...ROOT, children: [CODEX] } });
+      if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+      if (path === AGENTS_PATH) {
+        return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+      }
+      throw new Error(`Unexpected path: ${path}`);
+    });
+  });
+
+  afterEach(() => vi.useRealTimers());
+
+  it("hydrates every persisted expanded ancestor before marking the tree loaded", async () => {
+    const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+    await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+
+    expect(requestFileTreeMock.mock.calls.map((call) => call[2])).toEqual([
+      ROOT_PATH,
+      ...EXPANDED_PATHS,
+    ]);
+    expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+    expect(result.current.visibleRows.map((row) => row.path)).toEqual([
+      CODEX_PATH,
+      AGENTS_PATH,
+      CONFIG_PATH,
+    ]);
+  });
+
+  it.each([
+    ["deep-only", [AGENTS_PATH]],
+    ["deep-first", [AGENTS_PATH, CODEX_PATH]],
+  ])("synthesizes ancestor expansion for %s saved paths", async (_label, savedPaths) => {
+    sessionStorage.setItem(STORAGE_KEY, JSON.stringify(savedPaths));
+    const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+    await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+
+    expect(requestFileTreeMock.mock.calls.map((call) => call[2])).toEqual([
+      ROOT_PATH,
+      ...EXPANDED_PATHS,
+    ]);
+    expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+  });
+
+  it("retries a transient restored-folder failure without pruning persisted expansion", async () => {
+    vi.useFakeTimers();
+    let codexAttempts = 0;
+    requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+      if (path === ROOT_PATH) return Promise.resolve({ root: { ...ROOT, children: [CODEX] } });
+      if (path === CODEX_PATH && codexAttempts++ === 0) {
+        return Promise.reject(new Error("folder unavailable"));
+      }
+      if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+      if (path === AGENTS_PATH) return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+      throw new Error(`Unexpected path: ${path}`);
+    });
+
+    const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+    await act(async () => Promise.resolve());
+
+    expect(result.current.loadState).toBe("waiting");
+    expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+    expect(sessionStorage.getItem(STORAGE_KEY)).toBe(JSON.stringify(EXPANDED_PATHS));
+
+    await act(async () => vi.advanceTimersByTimeAsync(1000));
+
+    expect(result.current.loadState).toBe("loaded");
+    expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+    expect(result.current.visibleRows.map((row) => row.path)).toEqual([
+      CODEX_PATH,
+      AGENTS_PATH,
+      CONFIG_PATH,
+    ]);
+  });
+
+  it("prunes an expanded folder and descendants when its restore response is null", async () => {
+    requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+      if (path === ROOT_PATH) return Promise.resolve({ root: { ...ROOT, children: [CODEX] } });
+      if (path === CODEX_PATH) return Promise.resolve({ root: null });
+      throw new Error(`Unexpected path: ${path}`);
+    });
+
+    const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+    await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+
+    expect(result.current.expandedPaths).toEqual(new Set());
+    expect(result.current.visibleRows.map((row) => row.path)).toEqual([CODEX_PATH]);
+    expect(sessionStorage.getItem(STORAGE_KEY)).toBe("[]");
+  });
+});

apps/web/components/task/file-browser-restore-loader.test.tsx

@@ -0,0 +1,342 @@
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { afterEach, beforeEach, expect, it, vi } from "vitest";
+
+const requestFileTreeMock = vi.fn();
+const getWebSocketClientMock = vi.fn(() => ({ on: vi.fn(() => () => {}) }));
+
+vi.mock("@/lib/ws/connection", () => ({
+  getWebSocketClient: () => getWebSocketClientMock(),
+}));
+vi.mock("@/lib/ws/workspace-files", () => ({
+  requestFileTree: (...args: unknown[]) => requestFileTreeMock(...args),
+  requestFileContent: vi.fn(),
+  searchWorkspaceFiles: vi.fn(),
+}));
+vi.mock("@/hooks/domains/session/use-session-agentctl", () => ({
+  useSessionAgentctl: () => ({ isReady: true }),
+}));
+
+import { useFileBrowserTree } from "./file-browser-hooks";
+
+const SESSION = "session-1";
+const SUCCESSOR_SESSION = "session-2";
+const ENVIRONMENT = "environment-1";
+const SUCCESSOR_ENVIRONMENT = "environment-2";
+const ROOT_PATH = "";
+const CODEX_PATH = ".codex";
+const AGENTS_PATH = ".codex/agents";
+const CONFIG_PATH = ".codex/agents/config.toml";
+const EXPANDED_PATHS = [CODEX_PATH, AGENTS_PATH];
+const STORAGE_KEY = `kandev.filesPanel.expanded.${ENVIRONMENT}`;
+const SUCCESSOR_STORAGE_KEY = `kandev.filesPanel.expanded.${SUCCESSOR_ENVIRONMENT}`;
+const RETRY_DELAYS_MS = [1000, 2000, 5000, 10000];
+const ROOT = { name: "", path: ROOT_PATH, is_dir: true, size: 0 };
+const CODEX = { name: ".codex", path: CODEX_PATH, is_dir: true, size: 0 };
+const AGENTS = { name: "agents", path: AGENTS_PATH, is_dir: true, size: 0 };
+const CONFIG = { name: "config.toml", path: CONFIG_PATH, is_dir: false, size: 1 };
+const NEW_CONFIG = { name: "config.toml", path: ".new/config.toml", is_dir: false, size: 1 };
+
+function deferred<T>() {
+  let resolve!: (value: T) => void;
+  return { promise: new Promise<T>((done) => (resolve = done)), resolve };
+}
+
+function mockHydratedTree() {
+  requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+    if (path === ROOT_PATH) return Promise.resolve({ root: { ...ROOT, children: [CODEX] } });
+    if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+    if (path === AGENTS_PATH) return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+    throw new Error(`Unexpected path: ${path}`);
+  });
+}
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  sessionStorage.clear();
+  sessionStorage.setItem(STORAGE_KEY, JSON.stringify(EXPANDED_PATHS));
+  mockHydratedTree();
+});
+
+afterEach(() => vi.useRealTimers());
+
+it("starts a successor load after a session switch and ignores the stale predecessor", async () => {
+  const firstRoot = deferred<{ root: typeof ROOT & { children: (typeof CODEX)[] } }>();
+  requestFileTreeMock.mockImplementation((_client: unknown, sessionId: string, path: string) => {
+    if (sessionId === SESSION && path === ROOT_PATH) return firstRoot.promise;
+    if (sessionId === SUCCESSOR_SESSION && path === ROOT_PATH) {
+      return Promise.resolve({
+        root: { ...ROOT, children: [{ ...CODEX, name: ".new", path: ".new" }] },
+      });
+    }
+    return Promise.resolve({ root: { ...CODEX, children: [] } });
+  });
+  const { result, rerender } = renderHook(
+    ({ sessionId }) => useFileBrowserTree(sessionId, ENVIRONMENT),
+    { initialProps: { sessionId: SESSION } },
+  );
+  rerender({ sessionId: SUCCESSOR_SESSION });
+  firstRoot.resolve({ root: { ...ROOT, children: [CODEX] } });
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+
+  expect(requestFileTreeMock.mock.calls.map((call) => [call[1], call[2]])).toContainEqual([
+    SUCCESSOR_SESSION,
+    ROOT_PATH,
+  ]);
+  expect(result.current.visibleRows.map((row) => row.path)).toEqual([".new"]);
+});
+
+it("replaces a completed session tree with the successor's root and restored expansion", async () => {
+  requestFileTreeMock.mockImplementation((_client: unknown, sessionId: string, path: string) => {
+    if (sessionId === SESSION && path === ROOT_PATH) {
+      return Promise.resolve({
+        root: { ...ROOT, children: [{ ...CODEX, name: ".old", path: ".old" }, CODEX] },
+      });
+    }
+    if (path === ROOT_PATH) return Promise.resolve({ root: { ...ROOT, children: [CODEX] } });
+    if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+    if (path === AGENTS_PATH) return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+    throw new Error(`Unexpected successor path: ${path}`);
+  });
+  const { result, rerender } = renderHook(
+    ({ sessionId }) => useFileBrowserTree(sessionId, ENVIRONMENT),
+    { initialProps: { sessionId: SESSION } },
+  );
+
+  await waitFor(() => expect(result.current.visibleRows.map((row) => row.path)).toContain(".old"));
+  rerender({ sessionId: SUCCESSOR_SESSION });
+
+  await waitFor(() =>
+    expect(result.current.visibleRows.map((row) => row.path)).toEqual([
+      CODEX_PATH,
+      AGENTS_PATH,
+      CONFIG_PATH,
+    ]),
+  );
+
+  expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+  expect(sessionStorage.getItem(STORAGE_KEY)).toBe(JSON.stringify(EXPANDED_PATHS));
+  expect(result.current.visibleRows.map((row) => row.path)).toEqual([
+    CODEX_PATH,
+    AGENTS_PATH,
+    CONFIG_PATH,
+  ]);
+  expect(requestFileTreeMock.mock.calls.map((call) => [call[1], call[2]])).toContainEqual([
+    SUCCESSOR_SESSION,
+    ROOT_PATH,
+  ]);
+  expect(requestFileTreeMock.mock.calls.map((call) => [call[1], call[2]])).toContainEqual([
+    SUCCESSOR_SESSION,
+    AGENTS_PATH,
+  ]);
+});
+
+it("hydrates folders expanded after restore when a same-key session is replaced", async () => {
+  requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+    if (path === ROOT_PATH)
+      return Promise.resolve({
+        root: { ...ROOT, children: [CODEX, { ...CODEX, name: ".new", path: ".new" }] },
+      });
+    if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+    if (path === AGENTS_PATH) return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+    if (path === ".new")
+      return Promise.resolve({
+        root: { ...CODEX, name: ".new", path: ".new", children: [NEW_CONFIG] },
+      });
+    throw new Error(`Unexpected path: ${path}`);
+  });
+  const { result, rerender } = renderHook(
+    ({ sessionId }) => useFileBrowserTree(sessionId, ENVIRONMENT),
+    { initialProps: { sessionId: SESSION } },
+  );
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+  act(() => result.current.setExpandedPaths((paths) => new Set(paths).add(".new")));
+  rerender({ sessionId: SUCCESSOR_SESSION });
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+
+  expect(requestFileTreeMock.mock.calls.map((call) => [call[1], call[2]])).toContainEqual([
+    SUCCESSOR_SESSION,
+    ".new",
+  ]);
+  expect(result.current.visibleRows.map((row) => row.path)).toContain(NEW_CONFIG.path);
+});
+
+it("invalidates an in-flight request when the reset key changes", async () => {
+  const firstRoot = deferred<{ root: typeof ROOT & { children: (typeof CODEX)[] } }>();
+  requestFileTreeMock.mockImplementationOnce(() => firstRoot.promise);
+  mockHydratedTree();
+  sessionStorage.setItem(`kandev.filesPanel.expanded.${SUCCESSOR_ENVIRONMENT}`, "[]");
+  const { result, rerender } = renderHook(({ resetKey }) => useFileBrowserTree(SESSION, resetKey), {
+    initialProps: { resetKey: ENVIRONMENT },
+  });
+  rerender({ resetKey: SUCCESSOR_ENVIRONMENT });
+  firstRoot.resolve({ root: { ...ROOT, children: [CODEX] } });
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+  expect(result.current.visibleRows.map((row) => row.path)).toEqual([CODEX_PATH]);
+});
+
+it("replaces a completed reset-key tree with the new key's root and saved expansion", async () => {
+  sessionStorage.setItem(SUCCESSOR_STORAGE_KEY, JSON.stringify([".new"]));
+  let rootRequests = 0;
+  requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+    if (path === ROOT_PATH) {
+      rootRequests += 1;
+      return Promise.resolve(
+        rootRequests === 1
+          ? { root: { ...ROOT, children: [CODEX] } }
+          : { root: { ...ROOT, children: [{ ...CODEX, name: ".new", path: ".new" }] } },
+      );
+    }
+    if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+    if (path === AGENTS_PATH) return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+    if (path === ".new")
+      return Promise.resolve({
+        root: { ...CODEX, name: ".new", path: ".new", children: [NEW_CONFIG] },
+      });
+    throw new Error(`Unexpected path: ${path}`);
+  });
+  const { result, rerender } = renderHook(({ resetKey }) => useFileBrowserTree(SESSION, resetKey), {
+    initialProps: { resetKey: ENVIRONMENT },
+  });
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+  rerender({ resetKey: SUCCESSOR_ENVIRONMENT });
+
+  await waitFor(() => expect(rootRequests).toBe(2));
+  await waitFor(() =>
+    expect(requestFileTreeMock.mock.calls.map((call) => call[2])).toEqual([
+      ROOT_PATH,
+      CODEX_PATH,
+      AGENTS_PATH,
+      ROOT_PATH,
+      ".new",
+    ]),
+  );
+  await waitFor(() => expect(result.current.expandedPaths).toEqual(new Set([".new"])));
+
+  expect(sessionStorage.getItem(SUCCESSOR_STORAGE_KEY)).toBe(JSON.stringify([".new"]));
+  expect(result.current.visibleRows.map((row) => row.path)).toEqual([".new", NEW_CONFIG.path]);
+});
+
+it("does not persist the previous reset key's expansion under a new key", async () => {
+  const successorStorageKey = `kandev.filesPanel.expanded.${SUCCESSOR_ENVIRONMENT}`;
+  const setItemSpy = vi.spyOn(Storage.prototype, "setItem");
+  sessionStorage.setItem(successorStorageKey, JSON.stringify([".new"]));
+  const { result, rerender } = renderHook(({ resetKey }) => useFileBrowserTree(SESSION, resetKey), {
+    initialProps: { resetKey: ENVIRONMENT },
+  });
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+  setItemSpy.mockClear();
+
+  rerender({ resetKey: SUCCESSOR_ENVIRONMENT });
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+
+  expect(setItemSpy).not.toHaveBeenCalledWith(successorStorageKey, JSON.stringify(EXPANDED_PATHS));
+  setItemSpy.mockRestore();
+});
+
+it("recovers from manual retry after all automatic root retries preserve saved expansion", async () => {
+  vi.useFakeTimers();
+  let rootAttempts = 0;
+  requestFileTreeMock.mockImplementation((_client: unknown, _sessionId: string, path: string) => {
+    if (path === ROOT_PATH && rootAttempts++ <= RETRY_DELAYS_MS.length) {
+      return Promise.reject(new Error("root unavailable"));
+    }
+    if (path === ROOT_PATH) return Promise.resolve({ root: { ...ROOT, children: [CODEX] } });
+    if (path === CODEX_PATH) return Promise.resolve({ root: { ...CODEX, children: [AGENTS] } });
+    if (path === AGENTS_PATH) return Promise.resolve({ root: { ...AGENTS, children: [CONFIG] } });
+    throw new Error(`Unexpected path: ${path}`);
+  });
+  const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+  for (const delay of RETRY_DELAYS_MS) {
+    await act(async () => vi.advanceTimersByTimeAsync(delay));
+  }
+  expect(result.current.loadState).toBe("manual");
+  expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+
+  await act(async () => result.current.loadTree({ resetRetry: true }));
+
+  expect(result.current.loadState).toBe("loaded");
+  expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+  expect(result.current.visibleRows.map((row) => row.path)).toEqual([
+    CODEX_PATH,
+    AGENTS_PATH,
+    CONFIG_PATH,
+  ]);
+  expect(rootAttempts).toBe(RETRY_DELAYS_MS.length + 2);
+});
+
+it("deduplicates concurrent load requests for one generation", async () => {
+  const root = deferred<{ root: typeof ROOT & { children: (typeof CODEX)[] } }>();
+  mockHydratedTree();
+  requestFileTreeMock.mockImplementationOnce(() => root.promise);
+  const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+  const duplicate = result.current.loadTree();
+  expect(requestFileTreeMock).toHaveBeenCalledTimes(1);
+  root.resolve({ root: { ...ROOT, children: [CODEX] } });
+  await duplicate;
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+});
+
+it("retains restored paths while an automatic root retry recovers", async () => {
+  vi.useFakeTimers();
+  sessionStorage.setItem(STORAGE_KEY, JSON.stringify(EXPANDED_PATHS));
+  requestFileTreeMock.mockRejectedValueOnce(new Error("root unavailable"));
+  mockHydratedTree();
+  const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+  await act(async () => Promise.resolve());
+  expect(result.current.loadState).toBe("waiting");
+  expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+
+  await act(async () => vi.advanceTimersByTimeAsync(1000));
+  expect(result.current.expandedPaths).toEqual(new Set(EXPANDED_PATHS));
+  expect(result.current.loadState).toBe("loaded");
+});
+
+it("does not let a scheduled predecessor retry request for a replacement owner", async () => {
+  vi.useFakeTimers();
+  const clearTimeoutSpy = vi.spyOn(globalThis, "clearTimeout").mockImplementation(() => undefined);
+  requestFileTreeMock.mockRejectedValueOnce(new Error("root unavailable"));
+  mockHydratedTree();
+  const { result, rerender } = renderHook(
+    ({ sessionId }) => useFileBrowserTree(sessionId, ENVIRONMENT),
+    { initialProps: { sessionId: SESSION } },
+  );
+
+  await act(async () => Promise.resolve());
+  rerender({ sessionId: SUCCESSOR_SESSION });
+  await act(async () => vi.advanceTimersByTimeAsync(1000));
+
+  expect(
+    requestFileTreeMock.mock.calls.filter(
+      (call) => call[1] === SUCCESSOR_SESSION && call[2] === ROOT_PATH,
+    ),
+  ).toHaveLength(1);
+  expect(result.current.loadState).toBe("loaded");
+  clearTimeoutSpy.mockRestore();
+});
+
+it("prunes persisted folders when the restored root is null", async () => {
+  requestFileTreeMock.mockResolvedValue({ root: null });
+  const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+  expect(result.current.expandedPaths).toEqual(new Set());
+  await waitFor(() => expect(sessionStorage.getItem(STORAGE_KEY)).toBe("[]"));
+});
+
+it("does not request children for empty persisted paths", async () => {
+  sessionStorage.setItem(STORAGE_KEY, "[]");
+  requestFileTreeMock.mockResolvedValue({ root: { ...ROOT, children: [CODEX] } });
+  const { result } = renderHook(() => useFileBrowserTree(SESSION, ENVIRONMENT));
+
+  await waitFor(() => expect(result.current.loadState).toBe("loaded"));
+  expect(requestFileTreeMock.mock.calls.map((call) => call[2])).toEqual([ROOT_PATH]);
+  expect(result.current.expandedPaths).toEqual(new Set());
+});

apps/web/components/task/file-browser-restore.ts

@@ -0,0 +1,77 @@
+import { getWebSocketClient } from "@/lib/ws/connection";
+import { requestFileTree } from "@/lib/ws/workspace-files";
+import type { FileTreeNode } from "@/lib/types/backend";
+import { findNodeByPath, mergeTreeNodes } from "./file-tree-utils";
+import type { TreeLoadOwner } from "./file-browser-tree-loader";
+export type RestoredTree = {
+  root: FileTreeNode | null;
+  tree: FileTreeNode | null;
+  failedPaths: string[];
+};
+
+export function restoredExpandedPaths(paths: string[]): string[] {
+  const restored = new Set<string>();
+  for (const path of paths) {
+    if (!path) continue;
+    const parts = path.split("/");
+    for (let i = 1; i <= parts.length; i++) restored.add(parts.slice(0, i).join("/"));
+  }
+  return Array.from(restored).sort((a, b) => {
+    const depth = a.split("/").length - b.split("/").length;
+    return depth || a.localeCompare(b);
+  });
+}
+
+export function removeFailedExpansions(paths: string[], failedPaths: string[]): Set<string> {
+  return new Set(
+    paths.filter(
+      (expanded) =>
+        !failedPaths.some((failed) => expanded === failed || expanded.startsWith(`${failed}/`)),
+    ),
+  );
+}
+
+function mergeLoadedFolder(tree: FileTreeNode, incoming: FileTreeNode): FileTreeNode {
+  if (tree.path === incoming.path) return mergeTreeNodes(tree, incoming);
+  if (!tree.children) return tree;
+  return { ...tree, children: tree.children.map((child) => mergeLoadedFolder(child, incoming)) };
+}
+
+export async function fetchRestoredTree({
+  client,
+  owner,
+  paths,
+  isCurrentLoad,
+}: {
+  client: NonNullable<ReturnType<typeof getWebSocketClient>>;
+  owner: TreeLoadOwner;
+  paths: string[];
+  isCurrentLoad: () => boolean;
+}): Promise<RestoredTree | null> {
+  const rootResponse = await requestFileTree(client, owner.sessionId, "", 1);
+  if (!isCurrentLoad()) return null;
+  let tree = rootResponse.root ?? null;
+  const failedPaths: string[] = [];
+  for (const path of paths) {
+    if (failedPaths.some((failed) => path === failed || path.startsWith(`${failed}/`))) continue;
+    if (!tree || !findNodeByPath(tree, path)?.is_dir) {
+      failedPaths.push(path);
+      continue;
+    }
+    const response = await requestFileTree(client, owner.sessionId, path, 1);
+    if (!isCurrentLoad()) return null;
+    if (response.root) tree = mergeLoadedFolder(tree, response.root);
+    else failedPaths.push(path);
+  }
+  return { root: rootResponse.root ?? null, tree, failedPaths };
+}
+
+export function completeRestoredTree(
+  restored: RestoredTree | null,
+  isCurrentLoad: () => boolean,
+  onFailedPaths: (paths: string[]) => void,
+): RestoredTree | null {
+  if (!restored || !isCurrentLoad()) return null;
+  if (restored.failedPaths.length > 0) onFailedPaths(restored.failedPaths);
+  return restored;
+}

apps/web/components/task/file-browser-tree-loader.ts

@@ -0,0 +1,211 @@
+import { useCallback, useRef } from "react";
+import type React from "react";
+import { getWebSocketClient } from "@/lib/ws/connection";
+import type { FileTreeNode } from "@/lib/types/backend";
+import { createDebugLogger, isDebug } from "@/lib/debug/log";
+import {
+  completeRestoredTree,
+  fetchRestoredTree,
+  removeFailedExpansions,
+} from "./file-browser-restore";
+import type { LoadState } from "./file-browser-hooks";
+
+const debugLoad = createDebugLogger("file-browser:load");
+const MAX_RETRY_ATTEMPTS = 4;
+const RETRY_DELAYS_MS = [1000, 2000, 5000, 10000];
+
+export type TreeLoadOwner = { sessionId: string; resetKey: string; generation: number };
+export type TreeLoaderContext = {
+  sessionId: string;
+  effectiveResetKey: string;
+  clearRetryTimer: () => void;
+  retryAttemptRef: React.MutableRefObject<number>;
+  retryTimerRef: React.MutableRefObject<NodeJS.Timeout | null>;
+  restoreExpandedPathsRef: React.MutableRefObject<string[]>;
+  hasInitializedExpandedRef: React.MutableRefObject<string | null>;
+  setTree: React.Dispatch<React.SetStateAction<FileTreeNode | null>>;
+  setExpandedPaths: React.Dispatch<React.SetStateAction<Set<string>>>;
+  setIsLoadingTree: React.Dispatch<React.SetStateAction<boolean>>;
+  setLoadState: React.Dispatch<React.SetStateAction<LoadState>>;
+  setLoadError: React.Dispatch<React.SetStateAction<string | null>>;
+};
+
+function logLoad(event: string, data: Record<string, unknown>) {
+  if (isDebug()) debugLoad(event, data);
+}
+
+function useTreeLoadOwner(sessionId: string, resetKey: string) {
+  const ownerRef = useRef<TreeLoadOwner>({ sessionId, resetKey, generation: 0 });
+  if (ownerRef.current.sessionId !== sessionId || ownerRef.current.resetKey !== resetKey) {
+    ownerRef.current = { sessionId, resetKey, generation: ownerRef.current.generation + 1 };
+  }
+  return ownerRef;
+}
+
+function requireWebSocketClient() {
+  const client = getWebSocketClient();
+  if (!client) throw new Error("WebSocket client not available");
+  return client;
+}
+
+function scheduleRetry({
+  error,
+  isCurrentLoad,
+  owner,
+  retryAttemptRef,
+  retryTimerRef,
+  clearRetryTimer,
+  setLoadError,
+  setLoadState,
+  retry,
+}: {
+  error: unknown;
+  isCurrentLoad: () => boolean;
+  owner: TreeLoadOwner;
+  retryAttemptRef: React.MutableRefObject<number>;
+  retryTimerRef: React.MutableRefObject<NodeJS.Timeout | null>;
+  clearRetryTimer: () => void;
+  setLoadError: React.Dispatch<React.SetStateAction<string | null>>;
+  setLoadState: React.Dispatch<React.SetStateAction<LoadState>>;
+  retry: () => void;
+}) {
+  if (!isCurrentLoad()) return;
+  const message = error instanceof Error ? error.message : "Failed to load file tree";
+  setLoadError(message);
+  if (retryAttemptRef.current >= MAX_RETRY_ATTEMPTS) {
+    setLoadState("manual");
+    logLoad("gave-up", { sessionId: owner.sessionId, error: message });
+    return;
+  }
+  const delay = RETRY_DELAYS_MS[Math.min(retryAttemptRef.current, RETRY_DELAYS_MS.length - 1)];
+  retryAttemptRef.current += 1;
+  setLoadState("waiting");
+  clearRetryTimer();
+  logLoad("retry", {
+    sessionId: owner.sessionId,
+    attempt: retryAttemptRef.current,
+    delayMs: delay,
+    error: message,
+  });
+  retryTimerRef.current = setTimeout(retry, delay);
+}
+
+async function restoreTree({
+  owner,
+  paths,
+  isCurrentLoad,
+  setExpandedPaths,
+}: {
+  owner: TreeLoadOwner;
+  paths: string[];
+  isCurrentLoad: () => boolean;
+  setExpandedPaths: React.Dispatch<React.SetStateAction<Set<string>>>;
+}) {
+  const hydrated = await fetchRestoredTree({
+    client: requireWebSocketClient(),
+    owner,
+    paths,
+    isCurrentLoad,
+  });
+  return completeRestoredTree(hydrated, isCurrentLoad, (failedPaths) =>
+    setExpandedPaths((previous) => removeFailedExpansions(Array.from(previous), failedPaths)),
+  );
+}
+
+export function useTreeLoader(ctx: TreeLoaderContext) {
+  const {
+    clearRetryTimer,
+    retryAttemptRef,
+    retryTimerRef,
+    restoreExpandedPathsRef,
+    hasInitializedExpandedRef,
+    setTree,
+    setExpandedPaths,
+    setIsLoadingTree,
+    setLoadState,
+    setLoadError,
+  } = ctx;
+  const ownerRef = useTreeLoadOwner(ctx.sessionId, ctx.effectiveResetKey);
+  const loadInFlightGenerationRef = useRef<number | null>(null);
+  const loadRequestRef = useRef(0);
+  const loadTree = useCallback(
+    async (options?: { resetRetry?: boolean; restoreExpandedPaths?: string[] }) => {
+      const owner = ownerRef.current;
+      if (loadInFlightGenerationRef.current === owner.generation) {
+        logLoad("skip-in-flight", { sessionId: owner.sessionId, resetKey: owner.resetKey });
+        return;
+      }
+      const requestId = ++loadRequestRef.current;
+      const isCurrentLoad = () =>
+        loadRequestRef.current === requestId && ownerRef.current.generation === owner.generation;
+      const restorePaths = options?.restoreExpandedPaths ?? restoreExpandedPathsRef.current;
+      loadInFlightGenerationRef.current = owner.generation;
+      setIsLoadingTree(true);
+      setLoadState("loading");
+      setLoadError(null);
+      if (options?.resetRetry) {
+        retryAttemptRef.current = 0;
+        clearRetryTimer();
+      }
+      logLoad("start", {
+        sessionId: owner.sessionId,
+        resetKey: owner.resetKey,
+        resetRetry: options?.resetRetry === true,
+        retryAttempt: retryAttemptRef.current,
+      });
+      try {
+        const completed = await restoreTree({
+          owner,
+          paths: restorePaths,
+          isCurrentLoad,
+          setExpandedPaths,
+        });
+        if (!completed) return;
+        hasInitializedExpandedRef.current = owner.resetKey;
+        setTree(completed.tree);
+        setLoadState("loaded");
+        retryAttemptRef.current = 0;
+        clearRetryTimer();
+        logLoad("loaded", {
+          sessionId: owner.sessionId,
+          rootPath: completed.root?.path ?? null,
+          children: completed.root?.children?.length ?? 0,
+        });
+      } catch (error) {
+        scheduleRetry({
+          error,
+          isCurrentLoad,
+          owner,
+          retryAttemptRef,
+          retryTimerRef,
+          clearRetryTimer,
+          setLoadError,
+          setLoadState,
+          retry: () => {
+            retryTimerRef.current = null;
+            if (!isCurrentLoad()) return;
+            void loadTree({ restoreExpandedPaths: restorePaths });
+          },
+        });
+      } finally {
+        if (isCurrentLoad()) {
+          setIsLoadingTree(false);
+          loadInFlightGenerationRef.current = null;
+        }
+      }
+    },
+    [
+      clearRetryTimer,
+      retryAttemptRef,
+      retryTimerRef,
+      restoreExpandedPathsRef,
+      hasInitializedExpandedRef,
+      setTree,
+      setExpandedPaths,
+      setIsLoadingTree,
+      setLoadError,
+      setLoadState,
+    ],
+  );
+  return loadTree;
+}