Fix lazy file tree pagination

Co-Authored-By: Paperclip <noreply@paperclip.ing>

ui/src/components/WorkspaceFileBrowser.test.tsx

@@ -17,6 +17,7 @@ function act(callback: () => void | Promise<void>) {
 }
 
 const useQueryMock = vi.fn();
+const LIST_LIMIT = 100;
 
 vi.mock("@tanstack/react-query", async () => {
   const actual = await vi.importActual<typeof import("@tanstack/react-query")>("@tanstack/react-query");
@@ -85,6 +86,17 @@ function availableResponse(items: WorkspaceFileListItem[], truncated = false): W
   };
 }
 
+function availableAllResponse(
+  items: WorkspaceFileListItem[],
+  path: string | null,
+  truncated = false,
+  offset = 0,
+): WorkspaceFileListResponse {
+  const response = availableResponse(items, truncated);
+  response.query = { ...response.query, mode: "all", path, offset };
+  return response;
+}
+
 function createWorkspace(overrides: Partial<ProjectWorkspace> = {}): ProjectWorkspace {
   return {
     id: "workspace-content",
@@ -405,6 +417,46 @@ describe("WorkspaceFileBrowser", () => {
     expect(container.textContent).toContain("refine the search to narrow");
   });
 
+  it("renders newly loaded current-folder rows after Load more", () => {
+    const folderPath = "ui/src/components";
+    const pageZero = availableAllResponse([
+      createItem({
+        title: "IssueLinkQuicklook.tsx",
+        relativePath: `${folderPath}/IssueLinkQuicklook.tsx`,
+        displayPath: `${folderPath}/IssueLinkQuicklook.tsx`,
+      }),
+    ], folderPath, true, 0);
+    const pageOne = availableAllResponse([
+      createItem({
+        title: "SourceTrustBadge.tsx",
+        relativePath: `${folderPath}/SourceTrustBadge.tsx`,
+        displayPath: `${folderPath}/SourceTrustBadge.tsx`,
+      }),
+    ], folderPath, true, LIST_LIMIT);
+    useQueryMock.mockImplementation((options: { queryKey: readonly unknown[] }) => {
+      const query = options.queryKey[4] as { path?: string | null; offset?: number } | undefined;
+      if (query?.path === folderPath && query.offset === LIST_LIMIT) return ok(pageOne);
+      return ok(pageZero);
+    });
+
+    renderBrowser(vi.fn(), {
+      compact: true,
+      autoFocusSearch: false,
+      initialFolderPath: folderPath,
+    });
+
+    expect(container.textContent).toContain("IssueLinkQuicklook.tsx");
+    expect(container.textContent).not.toContain("SourceTrustBadge.tsx");
+    const loadMore = Array.from(container.querySelectorAll("button")).find(
+      (el) => el.textContent === "Load more from this folder",
+    );
+    act(() => {
+      loadMore!.dispatchEvent(new MouseEvent("click", { bubbles: true, cancelable: true }));
+    });
+
+    expect(container.textContent).toContain("SourceTrustBadge.tsx");
+  });
+
   it("opens the highlighted row when Enter is pressed in the search field", () => {
     useQueryMock.mockReturnValue(ok(availableResponse([createItem()])));
     const { onOpen } = renderBrowser();
@@ -668,6 +720,58 @@ describe("WorkspaceFileBrowser", () => {
     expect(container.textContent).toContain("setup-commands.md");
   });
 
+  it("auto-pages the selected file's current folder until the selected row is loaded", async () => {
+    const folderPath = "ui/src/components";
+    const selectedPath = `${folderPath}/WorkspaceFileBrowser.tsx`;
+    const pageZero = availableAllResponse([
+      createItem({
+        title: "ActivityCharts.tsx",
+        relativePath: `${folderPath}/ActivityCharts.tsx`,
+        displayPath: `${folderPath}/ActivityCharts.tsx`,
+      }),
+    ], folderPath, true, 0);
+    const pageOne = availableAllResponse([
+      createItem({
+        title: "SourceTrustBadge.tsx",
+        relativePath: `${folderPath}/SourceTrustBadge.tsx`,
+        displayPath: `${folderPath}/SourceTrustBadge.tsx`,
+      }),
+    ], folderPath, true, LIST_LIMIT);
+    const pageTwo = availableAllResponse([
+      createItem({
+        title: "WorkspaceFileBrowser.tsx",
+        relativePath: selectedPath,
+        displayPath: selectedPath,
+      }),
+    ], folderPath, false, LIST_LIMIT * 2);
+    useQueryMock.mockImplementation((options: { queryKey: readonly unknown[] }) => {
+      const query = options.queryKey[4] as { path?: string | null; offset?: number } | undefined;
+      if (query?.path === folderPath && query.offset === LIST_LIMIT) return ok(pageOne);
+      if (query?.path === folderPath && query.offset === LIST_LIMIT * 2) return ok(pageTwo);
+      return ok(pageZero);
+    });
+
+    renderBrowser(vi.fn(), {
+      compact: true,
+      autoFocusSearch: false,
+      initialFolderPath: folderPath,
+      selectedPath,
+    });
+
+    await new Promise((resolve) => window.setTimeout(resolve, 0));
+    await new Promise((resolve) => window.setTimeout(resolve, 0));
+
+    const pageTwoCall = useQueryMock.mock.calls.find(
+      ([options]) => options.queryKey?.[4]?.path === folderPath && options.queryKey[4].offset === LIST_LIMIT * 2,
+    );
+    expect(pageTwoCall?.[0].queryKey[4]).toMatchObject({ path: folderPath, offset: LIST_LIMIT * 2 });
+    const selected = Array.from(container.querySelectorAll('[role="treeitem"]')).find(
+      (el) => el.getAttribute("title") === selectedPath,
+    );
+    expect(selected?.getAttribute("aria-selected")).toBe("true");
+    expect(container.textContent).toContain("WorkspaceFileBrowser.tsx");
+  });
+
   it("lets breadcrumb folders navigate to parent directories", () => {
     const folderPath = "content-os/cases/active/";
     useQueryMock.mockReturnValue(ok(availableResponse([createItem({

ui/src/components/WorkspaceFileBrowser.tsx

@@ -711,6 +711,13 @@ export function WorkspaceFileBrowser({
     });
     return set;
   }, [folderPageQueries, folderPageSpecs]);
+  const lazyErroredFolders = useMemo(() => {
+    const set = new Set<string>();
+    folderPageSpecs.forEach((spec, index) => {
+      if (folderPageQueries[index]?.isError) set.add(spec.path);
+    });
+    return set;
+  }, [folderPageQueries, folderPageSpecs]);
   const allLoadedItems = useMemo(() => {
     if (!isLazyBrowse) return items;
     return Array.from(lazyItemsByFolder.values()).flat();
@@ -719,10 +726,12 @@ export function WorkspaceFileBrowser({
     () => allLoadedItems.filter((item): item is WorkspaceFileListFileItem => item.kind === "file"),
     [allLoadedItems],
   );
-  const treeNodes = useMemo(
-    () => isLazyBrowse ? buildWorkspaceDirectoryTree(items) : buildWorkspaceFileTree(items, folderPath),
-    [folderPath, isLazyBrowse, items],
-  );
+  const treeNodes = useMemo(() => {
+    if (isLazyBrowse) {
+      return buildWorkspaceDirectoryTree(lazyItemsByFolder.get(currentFolderKey) ?? items);
+    }
+    return buildWorkspaceFileTree(items, folderPath);
+  }, [currentFolderKey, folderPath, isLazyBrowse, items, lazyItemsByFolder]);
   const selectedItemIndex = selectedPath
     ? allLoadedFileItems.findIndex((item) =>
       item.relativePath === selectedPath &&
@@ -747,6 +756,37 @@ export function WorkspaceFileBrowser({
     setHighlightedIndex(allLoadedFileItems.length > 0 ? 0 : -1);
   }, [allLoadedFileItems.length, q, workspace, source, selectedProjectId, selectedWorkspaceId, folderPath, selectedPath, selectedItemIndex]);
 
+  useEffect(() => {
+    if (!selectedPath || !isLazyBrowse) return;
+    const selectedFolderKey = folderKey(parentFolderPath(selectedPath));
+    if (selectedFolderKey !== currentFolderKey && !loadedLazyFolders.includes(selectedFolderKey)) return;
+    if (!lazyTruncatedFolders.has(selectedFolderKey)) return;
+    if (lazyFetchingFolders.has(selectedFolderKey) || lazyErroredFolders.has(selectedFolderKey)) return;
+    const loadedItems = lazyItemsByFolder.get(selectedFolderKey) ?? [];
+    const selectedLoaded = loadedItems.some((item) =>
+      item.kind === "file" &&
+      item.relativePath === selectedPath &&
+      (activeProjectId ? item.projectId === activeProjectId : true) &&
+      (activeWorkspaceId ? item.workspaceId === activeWorkspaceId : true)
+    );
+    if (selectedLoaded) return;
+    setFolderPageCounts((current) => ({
+      ...current,
+      [selectedFolderKey]: (current[selectedFolderKey] ?? 1) + 1,
+    }));
+  }, [
+    activeProjectId,
+    activeWorkspaceId,
+    currentFolderKey,
+    isLazyBrowse,
+    lazyErroredFolders,
+    lazyFetchingFolders,
+    lazyItemsByFolder,
+    lazyTruncatedFolders,
+    loadedLazyFolders,
+    selectedPath,
+  ]);
+
   const announcement = useMemo(() => {
     if (listQuery.isFetching) return "Loading workspace files…";
     if (listQuery.isError) return "Unable to load workspace files.";