Make worktree archive optimistic

packages/app/src/stores/checkout-git-actions-store.test.ts

@@ -1,12 +1,23 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { QueryClient } from "@tanstack/react-query";
+import type { DaemonClient } from "@server/client/daemon-client";
+import { queryClient as appQueryClient } from "@/query/query-client";
 import { useSessionStore } from "@/stores/session-store";
+import type { WorkspaceDescriptor } from "@/stores/session-store";
 import {
   __resetCheckoutGitActionsStoreForTests,
   invalidateCheckoutGitQueriesForClient,
   useCheckoutGitActionsStore,
 } from "@/stores/checkout-git-actions-store";
 
+vi.mock("@react-native-async-storage/async-storage", () => ({
+  default: {
+    getItem: vi.fn(async () => null),
+    setItem: vi.fn(async () => undefined),
+    removeItem: vi.fn(async () => undefined),
+  },
+}));
+
 function createDeferred<T>() {
   let resolve!: (value: T) => void;
   let reject!: (reason?: unknown) => void;
@@ -17,6 +28,22 @@ function createDeferred<T>() {
   return { promise, resolve, reject };
 }
 
+function workspace(input: Partial<WorkspaceDescriptor> & Pick<WorkspaceDescriptor, "id">) {
+  return {
+    id: input.id,
+    projectId: input.projectId ?? "project-1",
+    projectDisplayName: input.projectDisplayName ?? "Project",
+    projectRootPath: input.projectRootPath ?? "/tmp/repo",
+    workspaceDirectory: input.workspaceDirectory ?? input.id,
+    projectKind: input.projectKind ?? "git",
+    workspaceKind: input.workspaceKind ?? "worktree",
+    name: input.name ?? input.id,
+    status: input.status ?? "done",
+    diffStat: input.diffStat ?? null,
+    scripts: input.scripts ?? [],
+  } satisfies WorkspaceDescriptor;
+}
+
 describe("checkout-git-actions-store", () => {
   const serverId = "server-1";
   const cwd = "/tmp/repo";
@@ -24,12 +51,14 @@ describe("checkout-git-actions-store", () => {
   beforeEach(() => {
     vi.useFakeTimers();
     __resetCheckoutGitActionsStoreForTests();
+    appQueryClient.clear();
     useSessionStore.setState((state) => ({ ...state, sessions: {} as any }));
   });
 
   afterEach(() => {
     vi.useRealTimers();
     __resetCheckoutGitActionsStoreForTests();
+    appQueryClient.clear();
     useSessionStore.setState((state) => ({ ...state, sessions: {} as any }));
   });
 
@@ -89,4 +118,53 @@ describe("checkout-git-actions-store", () => {
 
     queryClient.clear();
   });
+
+  it("hides an archived worktree optimistically while the archive RPC is in flight", async () => {
+    const deferred = createDeferred<Record<string, never>>();
+    const client = {
+      archivePaseoWorktree: vi.fn(() => deferred.promise),
+    };
+    const featureWorkspace = workspace({ id: cwd, name: "feature" });
+    useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient);
+    useSessionStore.getState().setWorkspaces(serverId, new Map([[cwd, featureWorkspace]]));
+    appQueryClient.setQueryData(
+      ["sidebarPaseoWorktreeList", serverId, "/tmp"],
+      [{ worktreePath: cwd }, { worktreePath: "/tmp/other" }],
+    );
+
+    const archive = useCheckoutGitActionsStore
+      .getState()
+      .archiveWorktree({ serverId, cwd, worktreePath: cwd });
+
+    expect(client.archivePaseoWorktree).toHaveBeenCalledWith({ worktreePath: cwd });
+    expect(useSessionStore.getState().sessions[serverId]?.workspaces.has(cwd)).toBe(false);
+    expect(appQueryClient.getQueryData(["sidebarPaseoWorktreeList", serverId, "/tmp"])).toEqual([
+      { worktreePath: "/tmp/other" },
+    ]);
+
+    deferred.resolve({});
+    await archive;
+  });
+
+  it("restores an optimistically hidden worktree when archive fails", async () => {
+    const client = {
+      archivePaseoWorktree: vi.fn(async () => ({ error: { message: "archive failed" } })),
+    };
+    const featureWorkspace = workspace({ id: cwd, name: "feature" });
+    const listSnapshot = [{ worktreePath: cwd }, { worktreePath: "/tmp/other" }];
+    useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient);
+    useSessionStore.getState().setWorkspaces(serverId, new Map([[cwd, featureWorkspace]]));
+    appQueryClient.setQueryData(["sidebarPaseoWorktreeList", serverId, "/tmp"], listSnapshot);
+
+    await expect(
+      useCheckoutGitActionsStore.getState().archiveWorktree({ serverId, cwd, worktreePath: cwd }),
+    ).rejects.toThrow("archive failed");
+
+    expect(useSessionStore.getState().sessions[serverId]?.workspaces.get(cwd)).toEqual(
+      featureWorkspace,
+    );
+    expect(appQueryClient.getQueryData(["sidebarPaseoWorktreeList", serverId, "/tmp"])).toEqual(
+      listSnapshot,
+    );
+  });
 });

packages/app/src/stores/checkout-git-actions-store.ts

@@ -1,4 +1,4 @@
-import type { QueryClient } from "@tanstack/react-query";
+import type { QueryClient, QueryKey } from "@tanstack/react-query";
 import { create } from "zustand";
 import { queryClient as appQueryClient } from "@/query/query-client";
 import {
@@ -6,6 +6,7 @@ import {
   useWorkspaceLayoutStore,
 } from "@/stores/workspace-layout-store";
 import { useSessionStore } from "@/stores/session-store";
+import type { WorkspaceDescriptor } from "@/stores/session-store";
 import { useWorkspaceTabsStore } from "@/stores/workspace-tabs-store";
 
 const SUCCESS_DISPLAY_MS = 1000;
@@ -149,6 +150,58 @@ function removeWorktreeFromCachedLists(input: { serverId: string; worktreePath:
   );
 }
 
+interface WorktreeArchiveSnapshot {
+  workspace: WorkspaceDescriptor | null;
+  worktreeLists: Array<[QueryKey, unknown]>;
+}
+
+function isWorktreeListQuery(input: { queryKey: QueryKey; serverId: string }): boolean {
+  return (
+    Array.isArray(input.queryKey) &&
+    (input.queryKey[0] === "paseoWorktreeList" ||
+      input.queryKey[0] === "sidebarPaseoWorktreeList") &&
+    input.queryKey[1] === input.serverId
+  );
+}
+
+function snapshotWorktreeArchiveState(input: {
+  serverId: string;
+  worktreePath: string;
+}): WorktreeArchiveSnapshot {
+  return {
+    workspace:
+      useSessionStore.getState().sessions[input.serverId]?.workspaces.get(input.worktreePath) ??
+      null,
+    worktreeLists: appQueryClient.getQueriesData({
+      predicate: (query) =>
+        isWorktreeListQuery({ queryKey: query.queryKey, serverId: input.serverId }),
+    }),
+  };
+}
+
+function removeWorktreeFromSessionStore(input: { serverId: string; worktreePath: string }): void {
+  const serverId = input.serverId.trim();
+  const worktreePath = input.worktreePath.trim();
+  if (!serverId || !worktreePath) {
+    return;
+  }
+  useSessionStore.getState().removeWorkspace(serverId, worktreePath);
+}
+
+function restoreWorktreeArchiveState(input: {
+  serverId: string;
+  worktreePath: string;
+  snapshot: WorktreeArchiveSnapshot;
+}): void {
+  if (input.snapshot.workspace) {
+    useSessionStore.getState().mergeWorkspaces(input.serverId, [input.snapshot.workspace]);
+  }
+
+  for (const [queryKey, data] of input.snapshot.worktreeLists) {
+    appQueryClient.setQueryData(queryKey, data);
+  }
+}
+
 function purgeArchivedWorkspaceState(input: { serverId: string; worktreePath: string }): void {
   const serverId = input.serverId.trim();
   const workspaceId = input.worktreePath.trim();
@@ -353,11 +406,18 @@ export const useCheckoutGitActionsStore = create<CheckoutGitActionsStoreState>()
       actionId: "archive-worktree",
       run: async () => {
         const client = resolveClient(serverId);
-        const payload = await client.archivePaseoWorktree({ worktreePath });
-        if (payload.error) {
-          throw new Error(payload.error.message);
-        }
+        const snapshot = snapshotWorktreeArchiveState({ serverId, worktreePath });
         removeWorktreeFromCachedLists({ serverId, worktreePath });
+        removeWorktreeFromSessionStore({ serverId, worktreePath });
+        try {
+          const payload = await client.archivePaseoWorktree({ worktreePath });
+          if (payload.error) {
+            throw new Error(payload.error.message);
+          }
+        } catch (error) {
+          restoreWorktreeArchiveState({ serverId, worktreePath, snapshot });
+          throw error;
+        }
         invalidateWorktreeList();
         purgeArchivedWorkspaceState({ serverId, worktreePath });
       },