Fix initial file tree expansion persistence

Author: pallyoung

packages/web/src/features/workspace/views/shared/file-tree-panel.test.tsx

@@ -267,6 +267,104 @@ describe("FileTreePanel", () => {
     expect(store.get(expandedDirsAtomFamily("ws-test"))).toBeNull();
   });
 
+  it("preserves default root expansion when the user first expands a non-default directory", async () => {
+    const sendCommand = vi
+      .fn()
+      .mockImplementation(async (op: string, args?: { subPath?: string }) => {
+        if (op === "workspace.uiState.set") {
+          return {
+            id: "ws-test",
+            path: "/workspace",
+            targetRuntime: "native",
+            openedAt: 1,
+            lastActiveAt: 1,
+            uiState: args?.uiState,
+          };
+        }
+
+        if (args?.subPath === "src") {
+          return {
+            path: "src",
+            children: [{ path: "src/index.ts", name: "index.ts", kind: "file" }],
+          };
+        }
+
+        if (args?.subPath === "lib") {
+          return {
+            path: "lib",
+            children: [{ path: "lib/foo.ts", name: "foo.ts", kind: "file" }],
+          };
+        }
+
+        return {
+          path: "/workspace",
+          children: [
+            { path: "lib", name: "lib", kind: "dir" },
+            { path: "src", name: "src", kind: "dir" },
+          ],
+        };
+      });
+    const store = createStore();
+    store.set(wsClientAtom, { sendCommand } as never);
+    store.set(workspacesAtom, {
+      "ws-test": {
+        id: "ws-test",
+        path: "/workspace",
+        targetRuntime: "native",
+        openedAt: 1,
+        lastActiveAt: 1,
+        uiState: {
+          leftPanelWidth: 280,
+          bottomPanelHeight: 200,
+          focusMode: false,
+        },
+      },
+    } as never);
+    store.set(
+      fileTreeAtomFamily("ws-test"),
+      new Map([
+        [
+          ".",
+          [
+            {
+              path: "lib",
+              name: "lib",
+              kind: "dir",
+            },
+            {
+              path: "src",
+              name: "src",
+              kind: "dir",
+            },
+          ],
+        ],
+      ])
+    );
+
+    render(
+      <Provider store={store}>
+        <FileTreePanel workspaceId="ws-test" />
+      </Provider>
+    );
+
+    fireEvent.click(screen.getByText("lib"));
+
+    await waitFor(() => {
+      expect(sendCommand).toHaveBeenCalledWith(
+        "workspace.uiState.set",
+        expect.objectContaining({
+          workspaceId: "ws-test",
+          uiState: expect.objectContaining({
+            fileTreeExpandedDirs: expect.arrayContaining(["lib", "src"]),
+          }),
+        }),
+        undefined
+      );
+    });
+
+    expect(Array.from(store.get(expandedDirsAtomFamily("ws-test")) ?? [])).toEqual(["lib", "src"]);
+  });
+
   it("reloads the file tree after creating a file from the toolbar", async () => {
     const sendCommand = vi
       .fn()

packages/web/src/features/workspace/views/shared/file-tree-panel.tsx

@@ -157,6 +157,15 @@ export const FileTreePanel: FC<FileTreePanelProps> = ({
     () => (hasSearch ? searchResults.length : countVisibleFiles(treeNodes)),
     [hasSearch, searchResults.length, treeNodes]
   );
+  const defaultExpandedRootPaths = useMemo(
+    () =>
+      treeNodes
+        .filter(
+          (node) => node.kind === "dir" && DEFAULT_EXPANDED_ROOT_DIRS.has(node.name.toLowerCase())
+        )
+        .map((node) => node.path),
+    [treeNodes]
+  );
 
   useEffect(() => {
     if (expandedDirs !== null || !workspace) {
@@ -276,6 +285,7 @@ export const FileTreePanel: FC<FileTreePanelProps> = ({
                 onSelectFile={handleSelectFile}
                 onLoadChildren={loadChildren}
                 onToggleDirs={applyExpandedDirs}
+                defaultExpandedRootPaths={defaultExpandedRootPaths}
                 isLoadingDir={isLoadingDir}
               />
             ))
@@ -356,6 +366,7 @@ interface FileTreeNodeProps {
   node: FileNode;
   depth: number;
   expandedDirs: Set<string> | null;
+  defaultExpandedRootPaths: string[];
   selectedPath: string | null;
   onRequestCreate: (mode: "file" | "folder", baseDir: string | null) => void;
   onRequestDelete: (path: string, name: string) => void;
@@ -369,6 +380,7 @@ const FileTreeNode: FC<FileTreeNodeProps> = ({
   node,
   depth,
   expandedDirs,
+  defaultExpandedRootPaths,
   selectedPath,
   onRequestCreate,
   onRequestDelete,
@@ -406,7 +418,7 @@ const FileTreeNode: FC<FileTreeNodeProps> = ({
       return;
     }
 
-    const nextExpanded = new Set(expandedDirs ?? []);
+    const nextExpanded = new Set(expandedDirs ?? defaultExpandedRootPaths);
     if (isExpanded) {
       nextExpanded.delete(node.path);
     } else {
@@ -509,6 +521,7 @@ const FileTreeNode: FC<FileTreeNodeProps> = ({
               onSelectFile={onSelectFile}
               onLoadChildren={onLoadChildren}
               onToggleDirs={onToggleDirs}
+              defaultExpandedRootPaths={defaultExpandedRootPaths}
               isLoadingDir={isLoadingDir}
             />
           ))}