diff --git a/packages/app/src/screens/workspace/workspace-screen.tsx b/packages/app/src/screens/workspace/workspace-screen.tsx index c724b561c..33668dc47 100644 --- a/packages/app/src/screens/workspace/workspace-screen.tsx +++ b/packages/app/src/screens/workspace/workspace-screen.tsx @@ -1964,7 +1964,7 @@ function WorkspaceScreenContent({ const workspaceSetupSnapshot = useWorkspaceSetupStore((state) => persistenceKey ? (state.snapshots[persistenceKey] ?? null) : null, ); - const upsertWorkspaceSetupProgress = useWorkspaceSetupStore((state) => state.upsertProgress); + const ensureWorkspaceSetupStatus = useWorkspaceSetupStore((state) => state.ensureSetupStatus); const showWorkspaceSetup = shouldShowWorkspaceSetup(workspaceSetupSnapshot); const uiTabs = useMemo( () => (workspaceLayout ? collectAllTabs(workspaceLayout.root) : EMPTY_UI_TABS), @@ -2176,54 +2176,22 @@ function WorkspaceScreenContent({ const emptyWorkspaceSeedRef = useRef(null); const autoOpenedSetupTabWorkspaceRef = useRef(null); - const requestedWorkspaceSetupStatusKeyRef = useRef(null); useEffect(() => { - if (!isRouteFocused) { + if (!isRouteFocused || !client || !normalizedServerId || !normalizedWorkspaceId) { return; } - if (!client || !normalizedServerId || !normalizedWorkspaceId || !persistenceKey) { - return; - } - if (workspaceSetupSnapshot) { - return; - } - if (requestedWorkspaceSetupStatusKeyRef.current === persistenceKey) { - return; - } - - requestedWorkspaceSetupStatusKeyRef.current = persistenceKey; - let isCancelled = false; - - client - .fetchWorkspaceSetupStatus(normalizedWorkspaceId) - .then((response) => { - if (isCancelled || response.workspaceId !== normalizedWorkspaceId || !response.snapshot) { - return; - } - upsertWorkspaceSetupProgress({ - serverId: normalizedServerId, - payload: { workspaceId: response.workspaceId, ...response.snapshot }, - }); - return; - }) - .catch(() => { - if (requestedWorkspaceSetupStatusKeyRef.current === persistenceKey) { - requestedWorkspaceSetupStatusKeyRef.current = null; - } - }); - - return () => { - isCancelled = true; - }; + ensureWorkspaceSetupStatus({ + serverId: normalizedServerId, + workspaceId: normalizedWorkspaceId, + client, + }); }, [ client, + ensureWorkspaceSetupStatus, isRouteFocused, normalizedServerId, normalizedWorkspaceId, - persistenceKey, - upsertWorkspaceSetupProgress, - workspaceSetupSnapshot, ]); useEffect(() => { diff --git a/packages/app/src/stores/workspace-setup-store.test.ts b/packages/app/src/stores/workspace-setup-store.test.ts index 9c9a4718d..44c710b26 100644 --- a/packages/app/src/stores/workspace-setup-store.test.ts +++ b/packages/app/src/stores/workspace-setup-store.test.ts @@ -1,9 +1,111 @@ import { beforeEach, describe, expect, it } from "vitest"; -import { shouldShowWorkspaceSetup, useWorkspaceSetupStore } from "./workspace-setup-store"; +import { + shouldShowWorkspaceSetup, + useWorkspaceSetupStore, + type WorkspaceSetupStatusClient, + type WorkspaceSetupStatusResult, +} from "./workspace-setup-store"; + +const DEFAULT_SNAPSHOT: WorkspaceSetupStatusResult["snapshot"] = { + status: "running", + detail: { + type: "worktree_setup", + worktreePath: "/Users/test/project", + branchName: "main", + log: "", + commands: [], + }, + error: null, +}; + +function setupResult( + workspaceId: string, + snapshot: WorkspaceSetupStatusResult["snapshot"] = DEFAULT_SNAPSHOT, +): WorkspaceSetupStatusResult { + return { requestId: "req-1", workspaceId, snapshot }; +} + +function makeClient(handler: (workspaceId: string) => Promise) { + const calls: string[] = []; + const client: WorkspaceSetupStatusClient = { + fetchWorkspaceSetupStatus: (workspaceId) => { + calls.push(workspaceId); + return handler(workspaceId); + }, + }; + return { client, calls }; +} + +function createDeferred() { + let resolve!: (value: T) => void; + let reject!: (reason: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +function flush() { + return new Promise((resolve) => setTimeout(resolve, 0)); +} + +function storedSnapshots() { + return Object.values(useWorkspaceSetupStore.getState().snapshots); +} + +function resolveDefault(workspaceId: string): Promise { + return Promise.resolve(setupResult(workspaceId)); +} + +function rejectThenResolve() { + let attempt = 0; + return (workspaceId: string): Promise => { + attempt += 1; + if (attempt === 1) { + return Promise.reject(new Error("boom")); + } + return resolveDefault(workspaceId); + }; +} + +function nullThenResolve() { + let attempt = 0; + return (workspaceId: string): Promise => { + attempt += 1; + if (attempt === 1) { + return Promise.resolve(setupResult(workspaceId, null)); + } + return resolveDefault(workspaceId); + }; +} + +function mismatchThenResolve() { + let attempt = 0; + return (workspaceId: string): Promise => { + attempt += 1; + if (attempt === 1) { + return Promise.resolve(setupResult("999")); + } + return resolveDefault(workspaceId); + }; +} + +function ensureSetupStatus(client: WorkspaceSetupStatusClient) { + useWorkspaceSetupStore.getState().ensureSetupStatus({ + serverId: "server-1", + workspaceId: "42", + client, + }); +} describe("workspace-setup-store", () => { beforeEach(() => { - useWorkspaceSetupStore.setState({ pendingWorkspaceSetup: null }); + useWorkspaceSetupStore.setState({ + pendingWorkspaceSetup: null, + snapshots: {}, + requestedKeys: new Set(), + }); }); it("tracks deferred workspace setup by source directory and optional workspace id", () => { @@ -96,4 +198,118 @@ describe("workspace-setup-store", () => { }), ).toBe(true); }); + + it("ensureSetupStatus fetches setup status once and stores the snapshot", async () => { + const { client, calls } = makeClient(resolveDefault); + + ensureSetupStatus(client); + await flush(); + + expect(calls).toEqual(["42"]); + expect(storedSnapshots()).toEqual([ + expect.objectContaining({ workspaceId: "42", status: "running" }), + ]); + }); + + it("ensureSetupStatus does not refetch while a request is in flight", async () => { + const deferred = createDeferred(); + const { client, calls } = makeClient(() => deferred.promise); + + ensureSetupStatus(client); + ensureSetupStatus(client); + + expect(calls).toEqual(["42"]); + + deferred.resolve(setupResult("42")); + await flush(); + + ensureSetupStatus(client); + expect(calls).toEqual(["42"]); + }); + + it("ensureSetupStatus skips fetching when a snapshot already exists", () => { + useWorkspaceSetupStore.getState().upsertProgress({ + serverId: "server-1", + payload: { workspaceId: "42", ...DEFAULT_SNAPSHOT }, + }); + const { client, calls } = makeClient(resolveDefault); + + ensureSetupStatus(client); + + expect(calls).toEqual([]); + }); + + it("ensureSetupStatus ignores a response for a different workspace", async () => { + const { client } = makeClient(() => Promise.resolve(setupResult("999"))); + + ensureSetupStatus(client); + await flush(); + + expect(storedSnapshots()).toHaveLength(0); + }); + + it("ensureSetupStatus does not store a snapshot when the response snapshot is null", async () => { + const { client } = makeClient((workspaceId) => Promise.resolve(setupResult(workspaceId, null))); + + ensureSetupStatus(client); + await flush(); + + expect(storedSnapshots()).toHaveLength(0); + }); + + it("ensureSetupStatus retries after a null-snapshot response", async () => { + const { client, calls } = makeClient(nullThenResolve()); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42"]); + expect(storedSnapshots()).toHaveLength(0); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42", "42"]); + expect(storedSnapshots()).toHaveLength(1); + }); + + it("ensureSetupStatus retries after a mismatched-workspace response", async () => { + const { client, calls } = makeClient(mismatchThenResolve()); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42"]); + expect(storedSnapshots()).toHaveLength(0); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42", "42"]); + expect(storedSnapshots()).toHaveLength(1); + }); + + it("ensureSetupStatus clears the in-flight marker on error so a later call retries", async () => { + const { client, calls } = makeClient(rejectThenResolve()); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42"]); + expect(storedSnapshots()).toHaveLength(0); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42", "42"]); + expect(storedSnapshots()).toHaveLength(1); + }); + + it("ensureSetupStatus retries after the workspace is removed", async () => { + const { client, calls } = makeClient(resolveDefault); + + ensureSetupStatus(client); + await flush(); + expect(calls).toEqual(["42"]); + + useWorkspaceSetupStore.getState().removeWorkspace({ serverId: "server-1", workspaceId: "42" }); + ensureSetupStatus(client); + await flush(); + + expect(calls).toEqual(["42", "42"]); + }); }); diff --git a/packages/app/src/stores/workspace-setup-store.ts b/packages/app/src/stores/workspace-setup-store.ts index da5e815bb..23e69eedc 100644 --- a/packages/app/src/stores/workspace-setup-store.ts +++ b/packages/app/src/stores/workspace-setup-store.ts @@ -17,6 +17,15 @@ export type WorkspaceSetupProgressPayload = Extract< { type: "workspace_setup_progress" } >["payload"]; +export type WorkspaceSetupStatusResult = Extract< + SessionOutboundMessage, + { type: "workspace_setup_status_response" } +>["payload"]; + +export interface WorkspaceSetupStatusClient { + fetchWorkspaceSetupStatus: (workspaceId: string) => Promise; +} + export interface WorkspaceSetupSnapshot extends WorkspaceSetupProgressPayload { updatedAt: number; } @@ -31,9 +40,15 @@ export function shouldShowWorkspaceSetup(snapshot: WorkspaceSetupSnapshot | null interface WorkspaceSetupStoreState { pendingWorkspaceSetup: PendingWorkspaceSetup | null; snapshots: Record; + requestedKeys: Set; beginWorkspaceSetup: (value: PendingWorkspaceSetup) => void; clearWorkspaceSetup: () => void; upsertProgress: (input: { serverId: string; payload: WorkspaceSetupProgressPayload }) => void; + ensureSetupStatus: (input: { + serverId: string; + workspaceId: string; + client: WorkspaceSetupStatusClient; + }) => void; removeWorkspace: (input: { serverId: string; workspaceId: string }) => void; clearServer: (serverId: string) => void; } @@ -42,9 +57,10 @@ function buildWorkspaceSetupKey(input: { serverId: string; workspaceId: string } return buildWorkspaceTabPersistenceKey(input); } -export const useWorkspaceSetupStore = create()((set) => ({ +export const useWorkspaceSetupStore = create()((set, get) => ({ pendingWorkspaceSetup: null, snapshots: {}, + requestedKeys: new Set(), beginWorkspaceSetup: (value) => { set({ pendingWorkspaceSetup: value }); }, @@ -67,6 +83,40 @@ export const useWorkspaceSetupStore = create()((set) = }, })); }, + ensureSetupStatus: async ({ serverId, workspaceId, client }) => { + const key = buildWorkspaceSetupKey({ serverId, workspaceId }); + if (!key) { + return; + } + const state = get(); + if (state.snapshots[key] || state.requestedKeys.has(key)) { + return; + } + + // requestedKeys is a pure in-flight marker: it dedupes concurrent fetches and is + // released once the request settles. A settle that stored no snapshot (null snapshot, + // mismatched workspace, or error) leaves no marker, so a later call can retry; once a + // snapshot lands, the snapshots[key] guard above prevents redundant refetches. + set((current) => ({ requestedKeys: new Set(current.requestedKeys).add(key) })); + + try { + const response = await client.fetchWorkspaceSetupStatus(workspaceId); + if (response.workspaceId === workspaceId && response.snapshot) { + get().upsertProgress({ + serverId, + payload: { workspaceId: response.workspaceId, ...response.snapshot }, + }); + } + } catch { + // Swallowed: the finally clears the in-flight marker so a later call retries. + } finally { + set((current) => { + const next = new Set(current.requestedKeys); + next.delete(key); + return { requestedKeys: next }; + }); + } + }, removeWorkspace: ({ serverId, workspaceId }) => { const key = buildWorkspaceSetupKey({ serverId, workspaceId }); if (!key) {