From 591f41cb0b4b68ae039dda8eb893f96b48295f7a Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 1 Sep 2026 16:14:46 +0200 Subject: [PATCH] fix(app): preserve turn ownership through session updates Stream turn closure now wakes queued-message delivery, while same-turn snapshots retain local cancellation state. Offline liveness remains visible until an authoritative close arrives. --- .../browser/agent-message-submission.spec.ts | 2 +- .../chat-outline/use-chat-outline.test.tsx | 2 + .../sidebar-workspace-list.test.tsx | 4 +- ...space-shortcut-targets-subscriber.test.tsx | 10 +-- packages/app/src/git/actions-store.test.ts | 67 +++++------------ .../sidebar-workspaces-view-model.test.ts | 17 ++++- .../app/src/hooks/use-archive-agent.test.ts | 7 +- .../directory-sync/agent-replica.test.ts | 75 +++++++++++++++++++ .../runtime/directory-sync/agent-replica.ts | 5 +- .../app/src/runtime/directory-sync/index.ts | 18 ++--- .../directory-sync/internal/agent-store.ts | 15 +++- .../workspace-scripts-button.test.tsx | 4 + packages/app/src/test/seed-session.ts | 60 +++++++++++++-- .../src/workspace/workspace-archive.test.ts | 16 ++-- 14 files changed, 213 insertions(+), 89 deletions(-) diff --git a/packages/app/e2e/browser/agent-message-submission.spec.ts b/packages/app/e2e/browser/agent-message-submission.spec.ts index 295e07ba88..41397e9d5b 100644 --- a/packages/app/e2e/browser/agent-message-submission.spec.ts +++ b/packages/app/e2e/browser/agent-message-submission.spec.ts @@ -551,8 +551,8 @@ async function expectProviderAcknowledgementBeforeRpcAcceptanceSettlesSubmission await gate.waitForAgentStreamItem("user_message"); gate.releaseHeldServerMessage(); await gate.drop(); - await expect(page.getByTestId("turn-working-indicator")).toHaveCount(0); await expect(userMessage).toHaveAttribute("aria-busy", "false"); + await expect(page.getByTestId("turn-working-indicator")).toBeVisible(); } finally { gate.restore(); await agent.cleanup(); diff --git a/packages/app/src/agent-stream/chat-outline/use-chat-outline.test.tsx b/packages/app/src/agent-stream/chat-outline/use-chat-outline.test.tsx index 6200b8b0f9..71035fbce1 100644 --- a/packages/app/src/agent-stream/chat-outline/use-chat-outline.test.tsx +++ b/packages/app/src/agent-stream/chat-outline/use-chat-outline.test.tsx @@ -24,6 +24,8 @@ const runtime = vi.hoisted(() => ({ vi.mock("@/constants/platform", () => ({ isWeb: true })); vi.mock("@/runtime/host-runtime", () => ({ + fetchAgentTimeline: runtime.fetchAgentTimeline, + getHostClient: () => runtime, getHostRuntimeStore: () => ({ getClient: () => runtime, fetchAgentTimeline: runtime.fetchAgentTimeline, diff --git a/packages/app/src/components/sidebar-workspace-list.test.tsx b/packages/app/src/components/sidebar-workspace-list.test.tsx index 94e6f36466..df83474db5 100644 --- a/packages/app/src/components/sidebar-workspace-list.test.tsx +++ b/packages/app/src/components/sidebar-workspace-list.test.tsx @@ -38,7 +38,7 @@ import { } from "@/runtime/host-runtime"; import type { HostProfile } from "@/types/host-connection"; import { useSessionStore, type WorkspaceDescriptor } from "@/stores/session-store"; -import { seedSessionWorkspaces } from "@/test/seed-session"; +import { publishSessionWorkspaces } from "@/test/seed-session"; import { useSidebarOrderStore } from "@/stores/sidebar-order-store"; import { useWorkspaceFields } from "@/stores/session-store-hooks"; import { useActiveWorkspaceSelection } from "@/stores/navigation-active-workspace-store"; @@ -167,7 +167,7 @@ function initializeSidebarState(workspaces: WorkspaceDescriptor[]): void { act(() => { setHostProfiles([makeHost()]); useSessionStore.getState().initializeSession(SERVER_ID, null as unknown as DaemonClient); - seedSessionWorkspaces(SERVER_ID, new Map(workspaces.map((entry) => [entry.id, entry]))); + publishSessionWorkspaces(SERVER_ID, new Map(workspaces.map((entry) => [entry.id, entry]))); useSessionStore.getState().setHasHydratedWorkspaces(SERVER_ID, true); useSidebarOrderStore.setState({ projectOrder: ["project-a", "project-b"], diff --git a/packages/app/src/components/workspace-shortcut-targets-subscriber.test.tsx b/packages/app/src/components/workspace-shortcut-targets-subscriber.test.tsx index 2cceeae961..ba428e94ce 100644 --- a/packages/app/src/components/workspace-shortcut-targets-subscriber.test.tsx +++ b/packages/app/src/components/workspace-shortcut-targets-subscriber.test.tsx @@ -9,7 +9,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { getHostRuntimeStore } from "@/runtime/host-runtime"; import { useKeyboardShortcutsStore } from "@/stores/keyboard-shortcuts-store"; import { useSessionStore, type WorkspaceDescriptor } from "@/stores/session-store"; -import { seedSessionWorkspaces } from "@/test/seed-session"; +import { publishSessionWorkspaces } from "@/test/seed-session"; import { useSidebarCollapsedSectionsStore } from "@/stores/sidebar-collapsed-sections-store"; import { useSidebarOrderStore } from "@/stores/sidebar-order-store"; import { useSidebarViewStore } from "@/stores/sidebar-view-store"; @@ -96,7 +96,7 @@ describe("WorkspaceShortcutTargetsSubscriber", () => { act(() => { setHostProfiles([hostProfile()]); useSessionStore.getState().initializeSession("srv", null as unknown as DaemonClient); - seedSessionWorkspaces( + publishSessionWorkspaces( "srv", new Map([ ["ws-1", workspaceDescriptor({ id: "ws-1", name: "Workspace 1" })], @@ -142,7 +142,7 @@ describe("WorkspaceShortcutTargetsSubscriber", () => { it("publishes status-mode shortcut targets in visual status order", async () => { act(() => { useSidebarViewStore.getState().setGroupMode("status"); - seedSessionWorkspaces( + publishSessionWorkspaces( "srv", new Map([ [ @@ -214,11 +214,11 @@ describe("WorkspaceShortcutTargetsSubscriber", () => { setHostProfiles([hostProfile("host-a"), hostProfile("host-b")]); useSessionStore.getState().initializeSession("host-a", null as unknown as DaemonClient); useSessionStore.getState().initializeSession("host-b", null as unknown as DaemonClient); - seedSessionWorkspaces( + publishSessionWorkspaces( "host-a", new Map([["a-1", workspaceDescriptor({ id: "a-1", name: "Host A" })]]), ); - seedSessionWorkspaces( + publishSessionWorkspaces( "host-b", new Map([["b-1", workspaceDescriptor({ id: "b-1", name: "Host B" })]]), ); diff --git a/packages/app/src/git/actions-store.test.ts b/packages/app/src/git/actions-store.test.ts index 9ab27ba40f..2b1b3af20e 100644 --- a/packages/app/src/git/actions-store.test.ts +++ b/packages/app/src/git/actions-store.test.ts @@ -7,6 +7,12 @@ import { useCheckoutGitActionsStore, } from "@/git/actions-store"; +const hostClients = vi.hoisted(() => new Map()); + +vi.mock("@/runtime/host-runtime", () => ({ + getHostClient: (serverId: string) => hostClients.get(serverId) ?? null, +})); + vi.mock("@react-native-async-storage/async-storage", () => ({ default: { getItem: vi.fn(async () => null), @@ -33,6 +39,7 @@ describe("checkout-git-actions-store", () => { vi.useFakeTimers(); __resetCheckoutGitActionsStoreForTests(); appQueryClient.clear(); + hostClients.clear(); useSessionStore.setState((state) => ({ ...state, sessions: {} })); }); @@ -40,6 +47,7 @@ describe("checkout-git-actions-store", () => { vi.useRealTimers(); __resetCheckoutGitActionsStoreForTests(); appQueryClient.clear(); + hostClients.clear(); useSessionStore.setState((state) => ({ ...state, sessions: {} })); }); @@ -49,13 +57,7 @@ describe("checkout-git-actions-store", () => { checkoutCommit: vi.fn(() => deferred.promise), }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); const store = useCheckoutGitActionsStore.getState(); @@ -85,13 +87,7 @@ describe("checkout-git-actions-store", () => { return {}; }), }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); await useCheckoutGitActionsStore.getState().pullAndPush({ serverId, cwd }); @@ -106,13 +102,7 @@ describe("checkout-git-actions-store", () => { checkoutPull: vi.fn(async () => ({ error: { message: "pull conflict" } })), checkoutPush: vi.fn(async () => ({})), }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); await expect( useCheckoutGitActionsStore.getState().pullAndPush({ serverId, cwd }), @@ -127,13 +117,7 @@ describe("checkout-git-actions-store", () => { checkoutPull: vi.fn(async () => ({})), checkoutPush: vi.fn(async () => ({ error: { message: "push rejected" } })), }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); await expect( useCheckoutGitActionsStore.getState().pullAndPush({ serverId, cwd }), @@ -147,13 +131,7 @@ describe("checkout-git-actions-store", () => { const client = { checkoutRefresh: vi.fn(async () => ({ success: true, error: null })), }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); await useCheckoutGitActionsStore.getState().refresh({ serverId, cwd }); @@ -167,13 +145,7 @@ describe("checkout-git-actions-store", () => { const client = { checkoutRefresh: vi.fn(async () => ({ error: { message: "not a git repository" } })), }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); await expect(useCheckoutGitActionsStore.getState().refresh({ serverId, cwd })).rejects.toThrow( "not a git repository", @@ -186,13 +158,7 @@ describe("checkout-git-actions-store", () => { it("discards selected paths through the shared checkout action workflow", async () => { const checkoutDiscardChanges = vi.fn(async () => ({ success: true, error: null })); const client = { checkoutDiscardChanges }; - useSessionStore.setState((state) => ({ - ...state, - sessions: { - ...state.sessions, - [serverId]: { client } as unknown as (typeof state.sessions)[string], - }, - })); + hostClients.set(serverId, client as unknown as DaemonClient); await useCheckoutGitActionsStore .getState() @@ -227,6 +193,7 @@ describe("checkout-git-actions-store", () => { error: null, })); const client = { [rpc.method]: setAutoMerge }; + hostClients.set(serverId, client as unknown as DaemonClient); useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient); useSessionStore.getState().updateSessionServerInfo(serverId, { serverId, @@ -257,6 +224,7 @@ describe("checkout-git-actions-store", () => { error: null, })); const client = { [rpc.method]: setAutoMerge }; + hostClients.set(serverId, client as unknown as DaemonClient); useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient); useSessionStore.getState().updateSessionServerInfo(serverId, { serverId, @@ -284,6 +252,7 @@ describe("checkout-git-actions-store", () => { error: null, })), }; + hostClients.set(serverId, client as unknown as DaemonClient); useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient); useSessionStore.getState().updateSessionServerInfo(serverId, { serverId, diff --git a/packages/app/src/hooks/sidebar-workspaces-view-model.test.ts b/packages/app/src/hooks/sidebar-workspaces-view-model.test.ts index 01d5c6d6cd..7381dbbc3a 100644 --- a/packages/app/src/hooks/sidebar-workspaces-view-model.test.ts +++ b/packages/app/src/hooks/sidebar-workspaces-view-model.test.ts @@ -692,6 +692,7 @@ function agent(input: { id: string; workspaceId: string; status: Agent["status"]; + turn?: Agent["turn"]; updatedAt?: Date; parentAgentId?: string | null; archivedAt?: Date | null; @@ -701,7 +702,7 @@ function agent(input: { id: input.id, provider: "claude" as Agent["provider"], status: input.status, - turn: { phase: "idle", cancellationRequestId: null }, + turn: input.turn ?? { phase: "idle", cancellationRequestId: null }, createdAt: new Date(0), updatedAt: input.updatedAt ?? new Date(1_000), lastUserMessageAt: null, @@ -891,7 +892,19 @@ describe("deriveProjectStatusBucket", () => { sessions: { srv: sessionWith({ workspaces: [projectWorkspace("ws-1", "done")], - agents: [agent({ id: "a1", workspaceId: "ws-1", status: "running" })], + agents: [ + agent({ + id: "a1", + workspaceId: "ws-1", + status: "running", + turn: { + phase: "open", + turnId: "turn-1", + startedAt: new Date(1_000), + cancellationRequestId: null, + }, + }), + ], }), }, }), diff --git a/packages/app/src/hooks/use-archive-agent.test.ts b/packages/app/src/hooks/use-archive-agent.test.ts index d6112644ff..9bda09b9fd 100644 --- a/packages/app/src/hooks/use-archive-agent.test.ts +++ b/packages/app/src/hooks/use-archive-agent.test.ts @@ -3,6 +3,8 @@ import { QueryClient } from "@tanstack/react-query"; import { beforeEach, describe, expect, it } from "vitest"; import type { Agent } from "@/stores/session-store"; import { useSessionStore } from "@/stores/session-store"; +import { acceptAgentSnapshot } from "@/runtime/session-data"; +import { installSessionDataTestOwner } from "@/test/seed-session"; import { agentHistoryQueryKey, allAgentHistoryQueryKey } from "./agent-history-query-key"; import { applyArchivedAgentCloseResults, @@ -48,6 +50,7 @@ function makeAgent(overrides: Partial = {}): Agent { describe("useArchiveAgent", () => { beforeEach(() => { useSessionStore.setState((state) => ({ ...state, sessions: {} })); + installSessionDataTestOwner(["server-a"]); }); it("tracks pending archive state in shared react-query cache", () => { @@ -127,7 +130,7 @@ describe("useArchiveAgent", () => { it("applies archived agent close results to session state and cached lists", async () => { const queryClient = new QueryClient(); useSessionStore.getState().initializeSession("server-a", {} as DaemonClient); - useSessionStore.getState().setAgents("server-a", new Map([["agent-1", makeAgent()]])); + acceptAgentSnapshot("server-a", makeAgent()); queryClient.setQueryData(["sidebarAgentsList", "server-a"], { entries: [{ agent: { id: "agent-1" } }, { agent: { id: "agent-2" } }], }); @@ -209,7 +212,7 @@ describe("useArchiveAgent", () => { it("can apply archived agent close results without invalidating cached lists", () => { const queryClient = new QueryClient(); useSessionStore.getState().initializeSession("server-a", {} as DaemonClient); - useSessionStore.getState().setAgents("server-a", new Map([["agent-1", makeAgent()]])); + acceptAgentSnapshot("server-a", makeAgent()); queryClient.setQueryData(["sidebarAgentsList", "server-a"], { entries: [{ agent: { id: "agent-1" } }, { agent: { id: "agent-2" } }], }); diff --git a/packages/app/src/runtime/directory-sync/agent-replica.test.ts b/packages/app/src/runtime/directory-sync/agent-replica.test.ts index 3a30f6b6c5..489de68d3f 100644 --- a/packages/app/src/runtime/directory-sync/agent-replica.test.ts +++ b/packages/app/src/runtime/directory-sync/agent-replica.test.ts @@ -135,6 +135,81 @@ describe("AgentDirectoryReplica", () => { store.clearSession(serverId); }); + it("notifies the queue owner exactly once when a stream turn closes", () => { + const serverId = "agent-replica-stream-stop"; + const store = useSessionStore.getState(); + store.initializeSession(serverId, null as unknown as DaemonClient); + const stopped: string[] = []; + const replica = new AgentDirectoryReplica( + serverId, + (agentId) => stopped.push(agentId), + () => undefined, + ); + replica.commitSnapshot([entry(payload("agent"))], []); + replica.applyTurnLiveness("agent", { + type: "stream_open", + turn: { turnId: "turn-1", startedAt: null }, + }); + + replica.applyTurnLiveness("agent", { type: "stream_close", turnId: "turn-1" }); + replica.applyTurnLiveness("agent", { type: "stream_close", turnId: "turn-1" }); + + expect(stopped).toEqual(["agent"]); + store.clearSession(serverId); + }); + + it("preserves cancellation state when the same active turn is snapshotted", () => { + const serverId = "agent-replica-cancellation-snapshot"; + const store = useSessionStore.getState(); + store.initializeSession(serverId, null as unknown as DaemonClient); + const replica = new AgentDirectoryReplica( + serverId, + () => undefined, + () => undefined, + ); + replica.commitSnapshot( + [ + entry({ + ...payload("running"), + activeTurn: { turnId: "turn-1", startedAt: null }, + }), + ], + [], + ); + replica.applyTurnLiveness("agent", { type: "cancellation_started", requestId: 7 }); + + replica.applyDelta({ + kind: "upsert", + agent: { + ...payload("same running turn"), + activeTurn: { turnId: "turn-1", startedAt: null }, + }, + project: entry(payload("project")).project, + }); + + expect(useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.turn).toMatchObject({ + phase: "open", + turnId: "turn-1", + cancellationRequestId: 7, + }); + + replica.commitSnapshot( + [ + entry({ + ...payload("same turn from refresh"), + activeTurn: { turnId: "turn-1", startedAt: null }, + }), + ], + [], + ); + expect(useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.turn).toMatchObject({ + phase: "open", + turnId: "turn-1", + cancellationRequestId: 7, + }); + store.clearSession(serverId); + }); + it("uses turn liveness for stopped transitions even when protocol status disagrees", () => { const serverId = "agent-replica-mismatched-status"; const store = useSessionStore.getState(); diff --git a/packages/app/src/runtime/directory-sync/agent-replica.ts b/packages/app/src/runtime/directory-sync/agent-replica.ts index 38d1a0f6f6..5ddb4dce3e 100644 --- a/packages/app/src/runtime/directory-sync/agent-replica.ts +++ b/packages/app/src/runtime/directory-sync/agent-replica.ts @@ -209,8 +209,11 @@ export class AgentDirectoryReplica { agentId: string, transition: TurnLivenessTransition | readonly TurnLivenessTransition[], ): void { + const wasRunning = this.storeProjection.get(agentId)?.turn.phase === "open"; const accepted = this.storeProjection.applyTurn(agentId, transition); - if (accepted) this.persist([this.agentUpsert(accepted)]); + if (!accepted) return; + this.persist([this.agentUpsert(accepted)]); + if (wasRunning && accepted.turn.phase === "idle") this.onStoppedRunning(agentId); } private agentUpsert(agent: Agent): DirectoryReplicaMutation { diff --git a/packages/app/src/runtime/directory-sync/index.ts b/packages/app/src/runtime/directory-sync/index.ts index c39f8b716e..4c663cf025 100644 --- a/packages/app/src/runtime/directory-sync/index.ts +++ b/packages/app/src/runtime/directory-sync/index.ts @@ -636,10 +636,8 @@ export class DirectorySync { } acceptWorkspaces(workspaces: readonly WorkspaceDescriptor[]): void { - this.checkpoints?.commitDirectoryMutations( - this.serverId, - this.workspaces.acceptWorkspaces(workspaces), - ); + const mutations = this.workspaces.acceptWorkspaces(workspaces); + this.checkpoints?.commitDirectoryMutations(this.serverId, mutations); } archiveAgent(agentId: string, archivedAt: string): void { @@ -651,17 +649,13 @@ export class DirectorySync { } acceptProject(project: ProjectDescriptor): void { - this.checkpoints?.commitDirectoryMutations( - this.serverId, - this.workspaces.acceptProject(project), - ); + const mutations = this.workspaces.acceptProject(project); + this.checkpoints?.commitDirectoryMutations(this.serverId, mutations); } removeWorkspace(workspaceId: string): void { - this.checkpoints?.commitDirectoryMutations( - this.serverId, - this.workspaces.removeWorkspaceSnapshot(workspaceId), - ); + const mutations = this.workspaces.removeWorkspaceSnapshot(workspaceId); + this.checkpoints?.commitDirectoryMutations(this.serverId, mutations); } markWorkspacesHydrated(hydrated: boolean): void { diff --git a/packages/app/src/runtime/directory-sync/internal/agent-store.ts b/packages/app/src/runtime/directory-sync/internal/agent-store.ts index 9be2192824..0566da5aaf 100644 --- a/packages/app/src/runtime/directory-sync/internal/agent-store.ts +++ b/packages/app/src/runtime/directory-sync/internal/agent-store.ts @@ -12,6 +12,16 @@ import { useDraftStore } from "@/stores/draft-store"; import { getInitDeferred, getInitKey, rejectInitDeferred } from "@/utils/agent-initialization"; import { reduceTurnLiveness, type TurnLivenessTransition } from "@/timeline/turn-liveness"; +function mergeSnapshotTurn(previous: Agent | undefined, incoming: Agent): Agent { + if (!previous) return incoming; + const activeTurn = + incoming.turn.phase === "open" + ? { turnId: incoming.turn.turnId, startedAt: incoming.turn.startedAt } + : null; + const turn = reduceTurnLiveness(previous.turn, { type: "snapshot", activeTurn }); + return turn === incoming.turn ? incoming : { ...incoming, turn }; +} + export class AgentStoreProjection { constructor(private readonly serverId: string) {} @@ -40,7 +50,7 @@ export class AgentStoreProjection { let accepted = agent; useSessionStore.getState().setAgents(this.serverId, (current) => { const previous = current.get(agent.id); - accepted = acceptAgentDirectoryUpdate(previous, agent); + accepted = acceptAgentDirectoryUpdate(previous, mergeSnapshotTurn(previous, agent)); if (accepted === previous) return current; const next = new Map(current); next.set(agent.id, accepted); @@ -97,7 +107,8 @@ export class AgentStoreProjection { const agents = new Map(); for (const [agentId, fetchedAgent] of fetched) { const existing = current.get(agentId); - agents.set(agentId, existing && equal(existing, fetchedAgent) ? existing : fetchedAgent); + const merged = mergeSnapshotTurn(existing, fetchedAgent); + agents.set(agentId, existing && equal(existing, merged) ? existing : merged); if (fetchedAgent.archivedAt) { clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId }); } diff --git a/packages/app/src/screens/workspace/workspace-scripts-button.test.tsx b/packages/app/src/screens/workspace/workspace-scripts-button.test.tsx index 82e0132cd3..1cba95bd3a 100644 --- a/packages/app/src/screens/workspace/workspace-scripts-button.test.tsx +++ b/packages/app/src/screens/workspace/workspace-scripts-button.test.tsx @@ -91,6 +91,10 @@ vi.mock("@/constants/platform", () => ({ })); vi.mock("@/runtime/host-runtime", () => ({ + useHostRuntimeClient: () => ({ + startWorkspaceScript: startWorkspaceScriptMock, + killTerminal: killTerminalMock, + }), useHostRuntimeSnapshot: () => ({ activeConnection: null }), })); diff --git a/packages/app/src/test/seed-session.ts b/packages/app/src/test/seed-session.ts index 23848acd24..61a75b068f 100644 --- a/packages/app/src/test/seed-session.ts +++ b/packages/app/src/test/seed-session.ts @@ -3,6 +3,14 @@ import { type ProjectDescriptor, type WorkspaceDescriptor, } from "@/stores/session-store"; +import { DirectorySync } from "@/runtime/directory-sync"; +import { + acceptProjectSnapshot, + acceptWorkspaceSnapshots, + removeWorkspaceSnapshot, + registerDefaultSessionDataOwner, + SessionDataOwner, +} from "@/runtime/session-data"; function projectFromWorkspace(workspace: WorkspaceDescriptor): ProjectDescriptor { return { @@ -17,6 +25,19 @@ function projectFromWorkspace(workspace: WorkspaceDescriptor): ProjectDescriptor }; } +function collectProjects( + workspaces: Iterable, + projects: Iterable = [], +): Map { + const byProjectId = new Map(Array.from(projects, (project) => [project.projectId, project])); + for (const workspace of workspaces) { + if (!byProjectId.has(workspace.projectId)) { + byProjectId.set(workspace.projectId, projectFromWorkspace(workspace)); + } + } + return byProjectId; +} + /** * Seeds a session the way a daemon populates it: every workspace's project is * published on the project channel, so the store never holds a workspace whose @@ -28,14 +49,39 @@ export function seedSessionWorkspaces( workspaces: Map, projects?: Iterable, ): void { - const byProjectId = new Map(); - for (const project of projects ?? []) byProjectId.set(project.projectId, project); - for (const workspace of workspaces.values()) { - if (!byProjectId.has(workspace.projectId)) { - byProjectId.set(workspace.projectId, projectFromWorkspace(workspace)); - } - } + const byProjectId = collectProjects(workspaces.values(), projects); const store = useSessionStore.getState(); store.setProjects(serverId, byProjectId.values()); store.setWorkspaces(serverId, workspaces); } + +export function publishSessionWorkspaces( + serverId: string, + workspaces: Map, + projects?: Iterable, +): void { + const byProjectId = collectProjects(workspaces.values(), projects); + for (const workspaceId of useSessionStore.getState().sessions[serverId]?.workspaces.keys() ?? + []) { + if (!workspaces.has(workspaceId)) removeWorkspaceSnapshot(serverId, workspaceId); + } + for (const project of byProjectId.values()) acceptProjectSnapshot(serverId, project); + acceptWorkspaceSnapshots(serverId, Array.from(workspaces.values())); +} + +export function installSessionDataTestOwner(serverIds: readonly string[]): SessionDataOwner { + const owner = new SessionDataOwner(); + for (const serverId of serverIds) { + owner.registerDirectory( + serverId, + new DirectorySync(serverId, { + onAgentStoppedRunning: () => undefined, + markAgentLoading: () => undefined, + markAgentReady: () => undefined, + markAgentError: () => undefined, + }), + ); + } + registerDefaultSessionDataOwner(owner); + return owner; +} diff --git a/packages/app/src/workspace/workspace-archive.test.ts b/packages/app/src/workspace/workspace-archive.test.ts index b4db600e60..5288efad32 100644 --- a/packages/app/src/workspace/workspace-archive.test.ts +++ b/packages/app/src/workspace/workspace-archive.test.ts @@ -5,6 +5,8 @@ import { isWorkspaceArchivePending, } from "@/contexts/session-workspace-upserts"; import { useSessionStore, type WorkspaceDescriptor } from "@/stores/session-store"; +import { installSessionDataTestOwner } from "@/test/seed-session"; +import type { SessionDataOwner } from "@/runtime/session-data"; import { archiveWorkspaceOptimistically, archiveWorkspacesOptimistically, @@ -13,6 +15,7 @@ import { const SERVER_ID = "workspace-archive-test"; const SECOND_SERVER_ID = "workspace-archive-test-2"; +let sessionDataOwner: SessionDataOwner; type ArchiveWorkspacePayload = Awaited>; @@ -86,6 +89,8 @@ function storedWorkspace(id: string): WorkspaceDescriptor | undefined { beforeEach(() => { useSessionStore.getState().initializeSession(SERVER_ID, {} as DaemonClient); + useSessionStore.getState().initializeSession(SECOND_SERVER_ID, {} as DaemonClient); + sessionDataOwner = installSessionDataTestOwner([SERVER_ID, SECOND_SERVER_ID]); }); afterEach(() => { @@ -99,7 +104,7 @@ afterEach(() => { describe("archiveWorkspaceOptimistically", () => { it("hides the workspace and marks the archive pending while the daemon call runs", async () => { const archived = workspace(); - useSessionStore.getState().mergeWorkspaces(SERVER_ID, [archived]); + sessionDataOwner.acceptWorkspaces(SERVER_ID, [archived]); const releaseArchive = deferred(); const client = createClient(vi.fn(async () => releaseArchive.promise)); @@ -124,7 +129,7 @@ describe("archiveWorkspaceOptimistically", () => { it("restores the workspace and clears pending state when the daemon rejects the archive", async () => { const archived = workspace(); - useSessionStore.getState().mergeWorkspaces(SERVER_ID, [archived]); + sessionDataOwner.acceptWorkspaces(SERVER_ID, [archived]); const client = createClient( vi.fn(async () => archivePayload({ workspaceId: archived.id, error: "nope" })), ); @@ -154,7 +159,7 @@ describe("archiveWorkspacesOptimistically", () => { workspaceDirectory: "/repo/project/workspace-2", name: "workspace-2", }); - useSessionStore.getState().mergeWorkspaces(SERVER_ID, [first, second]); + sessionDataOwner.acceptWorkspaces(SERVER_ID, [first, second]); const client = createClient( vi.fn(async (workspaceId) => archivePayload({ @@ -182,9 +187,8 @@ describe("archiveWorkspacesOptimistically", () => { workspaceDirectory: "/repo/project/workspace-2", name: "workspace-2", }); - useSessionStore.getState().initializeSession(SECOND_SERVER_ID, {} as DaemonClient); - useSessionStore.getState().mergeWorkspaces(SERVER_ID, [first]); - useSessionStore.getState().mergeWorkspaces(SECOND_SERVER_ID, [second]); + sessionDataOwner.acceptWorkspaces(SERVER_ID, [first]); + sessionDataOwner.acceptWorkspaces(SECOND_SERVER_ID, [second]); const archivedByServer = new Map(); const clientFor = (serverId: string) =>