From 4ba0c2d339209f819e280ce07a0600af1c2d0286 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Fri, 17 Jul 2026 15:52:33 +0200 Subject: [PATCH] fix(app): close directory bootstrap races --- .../directory-sync/agent-replica.test.ts | 7 ++ .../runtime/directory-sync/agent-replica.ts | 7 +- .../src/runtime/directory-sync/index.test.ts | 113 ++++++++++++++++++ .../app/src/runtime/directory-sync/index.ts | 66 ++++++++-- packages/app/src/runtime/host-runtime.test.ts | 57 ++++++++- packages/app/src/runtime/host-runtime.ts | 19 +-- 6 files changed, 245 insertions(+), 24 deletions(-) create mode 100644 packages/app/src/runtime/directory-sync/index.test.ts 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 e0b22119d..4456ecc0a 100644 --- a/packages/app/src/runtime/directory-sync/agent-replica.test.ts +++ b/packages/app/src/runtime/directory-sync/agent-replica.test.ts @@ -57,6 +57,10 @@ describe("AgentDirectoryReplica", () => { store.initializeSession(serverId, null as unknown as DaemonClient); const replica = new AgentDirectoryReplica(serverId, () => undefined); replica.commitSnapshot([entry(payload("directory"))], []); + const directoryPlacement = useSessionStore + .getState() + .sessions[serverId]?.agents.get("agent")?.projectPlacement; + expect(directoryPlacement).toBeDefined(); const staleToken = replica.captureTimeline("agent"); replica.remove("agent"); @@ -74,6 +78,9 @@ describe("AgentDirectoryReplica", () => { expect(useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.title).toBe( "current", ); + expect( + useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.projectPlacement, + ).toEqual(directoryPlacement); store.clearSession(serverId); }); }); diff --git a/packages/app/src/runtime/directory-sync/agent-replica.ts b/packages/app/src/runtime/directory-sync/agent-replica.ts index e1bcab4c1..62de843cd 100644 --- a/packages/app/src/runtime/directory-sync/agent-replica.ts +++ b/packages/app/src/runtime/directory-sync/agent-replica.ts @@ -40,10 +40,15 @@ export class AgentDirectoryReplica { ) { return false; } - const normalized = applyLegacyDaemonWorkspaceOwnership({ + const existing = useSessionStore.getState().sessions[this.serverId]?.agents.get(token.agentId); + const timelineAgent = applyLegacyDaemonWorkspaceOwnership({ serverId: this.serverId, agent: normalizeAgentSnapshot(payload, this.serverId), }); + const normalized: Agent = { + ...timelineAgent, + projectPlacement: timelineAgent.projectPlacement ?? existing?.projectPlacement, + }; const accepted = upsertAgentReplica(this.serverId, normalized); replaceAgentPendingPermissions(this.serverId, accepted); useSessionStore.getState().setAgentLastActivity(accepted.id, accepted.lastActivityAt); diff --git a/packages/app/src/runtime/directory-sync/index.test.ts b/packages/app/src/runtime/directory-sync/index.test.ts new file mode 100644 index 000000000..d610f04be --- /dev/null +++ b/packages/app/src/runtime/directory-sync/index.test.ts @@ -0,0 +1,113 @@ +import { afterEach, describe, expect, it } from "vitest"; +import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; +import { useSessionStore } from "@/stores/session-store"; +import { DirectoryRefreshSupersededError, DirectorySync } from "./index"; + +class FakeDirectoryClient { + fetchAgentsCalls = 0; + fetchWorkspacesCalls = 0; + + on(): () => void { + return () => undefined; + } + + async fetchAgents(): Promise>> { + this.fetchAgentsCalls += 1; + return { + requestId: "agents", + entries: [], + pageInfo: { hasMore: false, nextCursor: null, prevCursor: null }, + }; + } + + async fetchWorkspaces(): Promise>> { + this.fetchWorkspacesCalls += 1; + return { + requestId: "workspaces", + entries: [], + emptyProjects: [], + pageInfo: { hasMore: false, nextCursor: null, prevCursor: null }, + }; + } +} + +const serverIds = new Set(); + +function createDirectory(serverId: string): { + client: FakeDirectoryClient; + directory: DirectorySync; +} { + serverIds.add(serverId); + const client = new FakeDirectoryClient(); + const directory = new DirectorySync(serverId, { + drainQueuedAgentMessage: () => undefined, + markAgentLoading: () => undefined, + markAgentReady: () => undefined, + markAgentError: () => undefined, + }); + directory.connectionChanged({ + client: client as unknown as DaemonClient, + status: "online", + source: { clientGeneration: 1, connectionEpoch: 1 }, + }); + return { client, directory }; +} + +afterEach(() => { + for (const serverId of serverIds) useSessionStore.getState().clearSession(serverId); + serverIds.clear(); +}); + +describe("DirectorySync session readiness", () => { + it("waits for workspace capability metadata before choosing the workspace protocol", async () => { + const serverId = "workspace-metadata"; + const { client, directory } = createDirectory(serverId); + + const refresh = directory.refreshWorkspaces({ subscribe: true }); + await Promise.resolve(); + expect(client.fetchWorkspacesCalls).toBe(0); + + const store = useSessionStore.getState(); + store.initializeSession(serverId, client as unknown as DaemonClient, 1); + await Promise.resolve(); + expect(client.fetchWorkspacesCalls).toBe(0); + + store.updateSessionServerInfo(serverId, { + serverId, + hostname: null, + version: "test", + features: { workspaceMultiplicity: true }, + }); + await refresh; + + expect(client.fetchWorkspacesCalls).toBe(1); + expect(useSessionStore.getState().sessions[serverId]?.hasHydratedWorkspaces).toBe(true); + directory.dispose(); + }); + + it("rejects a session wait on disconnect so the reconnect can refresh", async () => { + const serverId = "session-wait-reconnect"; + const { client, directory } = createDirectory(serverId); + const staleRefresh = directory.refreshAgents(); + await Promise.resolve(); + + directory.connectionChanged({ + client: null, + status: "offline", + source: { clientGeneration: 1, connectionEpoch: 1 }, + }); + await expect(staleRefresh).rejects.toBeInstanceOf(DirectoryRefreshSupersededError); + + directory.connectionChanged({ + client: client as unknown as DaemonClient, + status: "online", + source: { clientGeneration: 1, connectionEpoch: 2 }, + }); + const currentRefresh = directory.refreshAgents(); + useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient, 1); + await currentRefresh; + + expect(client.fetchAgentsCalls).toBe(1); + directory.dispose(); + }); +}); diff --git a/packages/app/src/runtime/directory-sync/index.ts b/packages/app/src/runtime/directory-sync/index.ts index 76658af2f..a181bcd0c 100644 --- a/packages/app/src/runtime/directory-sync/index.ts +++ b/packages/app/src/runtime/directory-sync/index.ts @@ -74,6 +74,7 @@ export class DirectorySync { source: { clientGeneration: 0, connectionEpoch: 0 }, }; private unsubscribe: (() => void) | null = null; + private readonly abortSessionWaits = new Set<() => void>(); constructor( private readonly serverId: string, @@ -88,7 +89,7 @@ export class DirectorySync { this.workspaces = new WorkspaceDirectoryReplica(serverId); } - connectionChanged(connection: DirectoryConnection): void { + connectionChanged(connection: DirectoryConnection): boolean { const changed = this.connection.client !== connection.client || this.connection.source.clientGeneration !== connection.source.clientGeneration || @@ -96,13 +97,14 @@ export class DirectorySync { const wentOffline = this.connection.status === "online" && connection.status === "offline"; if (!changed && !wentOffline) { this.connection = connection; - return; + return false; } this.flushAbortedTransactions(); this.unsubscribe?.(); this.unsubscribe = null; this.connection = connection; - if (!connection.client || connection.status !== "online") return; + this.abortPendingSessionWaits(); + if (!connection.client || connection.status !== "online") return true; const client = connection.client; const source = connection.source; const subscriptions = [ @@ -131,10 +133,12 @@ export class DirectorySync { this.unsubscribe = () => { for (const unsubscribe of subscriptions) unsubscribe(); }; + return true; } dispose(): void { this.flushAbortedTransactions(); + this.abortPendingSessionWaits(); this.unsubscribe?.(); this.unsubscribe = null; } @@ -216,14 +220,19 @@ export class DirectorySync { } async refreshWorkspaces(input?: { subscribe?: boolean }): Promise { - const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo; - if (serverInfo?.features?.workspaceMultiplicity !== true) return; const { client, source } = this.requireOnline(); const transaction = this.workspaceTransactions.begin(source, () => ({ workspaces: new Map(), emptyProjects: new Map(), })); try { + await this.waitForSessionMetadata(client, source); + const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo; + if (serverInfo?.features?.workspaceMultiplicity !== true) { + const deltas = this.workspaceTransactions.fail(transaction); + if (deltas) for (const delta of deltas) this.workspaces.applyDelta(delta); + return; + } await this.fetchWorkspaceSnapshot(client, source, transaction, input?.subscribe === true); if (!this.isCurrent(client, source) || !this.hasMatchingSession(client, source)) { throw new DirectoryRefreshSupersededError("workspace completion no longer current"); @@ -343,18 +352,47 @@ export class DirectorySync { } private async waitForSession(client: DaemonClient, source: DirectorySourceToken): Promise { - const matches = () => this.hasMatchingSession(client, source); + await this.waitForSessionState(client, source, () => this.hasMatchingSession(client, source)); + } + + private async waitForSessionMetadata( + client: DaemonClient, + source: DirectorySourceToken, + ): Promise { + await this.waitForSessionState(client, source, () => { + const session = useSessionStore.getState().sessions[this.serverId]; + return this.hasMatchingSession(client, source) && session?.serverInfo !== null; + }); + } + + private async waitForSessionState( + client: DaemonClient, + source: DirectorySourceToken, + matches: () => boolean, + ): Promise { if (matches()) return; await new Promise((resolve, reject) => { - const unsubscribe = useSessionStore.subscribe(() => { + let settled = false; + let unsubscribe: () => void = () => undefined; + const finish = (result: "ready" | "aborted") => { + if (settled) return; + settled = true; + unsubscribe(); + this.abortSessionWaits.delete(abort); + if (result === "ready") resolve(); + else reject(new DirectoryRefreshSupersededError("session wait no longer current")); + }; + const abort = () => finish("aborted"); + const check = () => { if (matches()) { - unsubscribe(); - resolve(); + finish("ready"); } else if (!this.isCurrent(client, source)) { - unsubscribe(); - reject(new DirectoryRefreshSupersededError("session wait no longer current")); + finish("aborted"); } - }); + }; + this.abortSessionWaits.add(abort); + unsubscribe = useSessionStore.subscribe(check); + check(); }); } @@ -367,6 +405,10 @@ export class DirectorySync { for (const delta of this.agentTransactions.abort()) this.agents.applyDelta(delta); for (const delta of this.workspaceTransactions.abort()) this.workspaces.applyDelta(delta); } + + private abortPendingSessionWaits(): void { + for (const abort of this.abortSessionWaits) abort(); + } } export class DirectoryRefreshSupersededError extends Error {} diff --git a/packages/app/src/runtime/host-runtime.test.ts b/packages/app/src/runtime/host-runtime.test.ts index e3141fb0a..4de541e34 100644 --- a/packages/app/src/runtime/host-runtime.test.ts +++ b/packages/app/src/runtime/host-runtime.test.ts @@ -1875,7 +1875,10 @@ describe("HostRuntimeStore", () => { }); it("buffers updates until the matching session generation exists", async () => { - const host = makeHost({ serverId: "srv_pre_session" }); + const host = makeHost({ + serverId: "srv_pre_session", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); const snapshotEntry = makeFetchAgentsEntry({ @@ -1904,9 +1907,15 @@ describe("HostRuntimeStore", () => { agent: { ...snapshotEntry.agent, title: "before-session" }, project: snapshotEntry.project, }); + const generation = store.getSnapshot(host.serverId)?.clientGeneration ?? 0; useSessionStore .getState() - .initializeSession(host.serverId, fakeClient as unknown as DaemonClient, 1); + .initializeSession(host.serverId, fakeClient as unknown as DaemonClient, generation); + useSessionStore.getState().updateSessionServerInfo(host.serverId, { + serverId: host.serverId, + hostname: null, + version: "test", + }); await fakeClient.waitForFetches(1); await waitForDirectoryReady(store, host.serverId); @@ -1918,6 +1927,50 @@ describe("HostRuntimeStore", () => { useSessionStore.getState().clearSession(host.serverId); }); + it("restarts directory bootstrap when reconnect supersedes a pending session wait", async () => { + const host = makeHost({ + serverId: "srv_session_wait_reconnect", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); + const fakeClient = new FakeDaemonClient(); + fakeClient.setConnectionState({ status: "connected" }); + const store = new HostRuntimeStore({ + deps: { + createClient: () => fakeClient as unknown as DaemonClient, + connectToDaemon: async () => ({ + client: fakeClient as unknown as DaemonClient, + serverId: host.serverId, + hostname: null, + }), + getClientId: async () => "cid_session_wait_reconnect", + }, + }); + + store.syncHosts([host]); + await fakeClient.waitForAgentUpdates(); + fakeClient.setConnectionState({ status: "disconnected", reason: "network" }); + fakeClient.setConnectionState({ status: "connected" }); + + const generation = store.getSnapshot(host.serverId)?.clientGeneration ?? 0; + const sessionStore = useSessionStore.getState(); + sessionStore.initializeSession( + host.serverId, + fakeClient as unknown as DaemonClient, + generation, + ); + sessionStore.updateSessionServerInfo(host.serverId, { + serverId: host.serverId, + hostname: null, + version: "test", + }); + await fakeClient.waitForFetches(1); + await waitForDirectoryReady(store, host.serverId); + + expect(fakeClient.fetchAgentsCalls).toHaveLength(1); + store.syncHosts([]); + sessionStore.clearSession(host.serverId); + }); + it("rejects a superseded refresh without overwriting the newer replica", async () => { const host = makeHost({ serverId: "srv_overlap", diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index d300dc968..85f4634ec 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -1962,14 +1962,15 @@ export class HostRuntimeStore { return; } const snapshot = controller.getSnapshot(); - this.directorySyncByServer.get(serverId)?.connectionChanged({ - client: snapshot.client, - status: snapshot.connectionStatus === "online" ? "online" : "offline", - source: { - clientGeneration: snapshot.clientGeneration, - connectionEpoch: snapshot.connectionEpoch, - }, - }); + const directorySourceChanged = + this.directorySyncByServer.get(serverId)?.connectionChanged({ + client: snapshot.client, + status: snapshot.connectionStatus === "online" ? "online" : "offline", + source: { + clientGeneration: snapshot.clientGeneration, + connectionEpoch: snapshot.connectionEpoch, + }, + }) ?? false; const previousStatus = this.lastConnectionStatusByServer.get(serverId); this.lastConnectionStatusByServer.set(serverId, snapshot.connectionStatus); const didTransitionOnline = @@ -1992,7 +1993,7 @@ export class HostRuntimeStore { if (!didTransitionOnline && snapshot.hasEverLoadedAgentDirectory) { return; } - if (this.agentDirectoryBootstrapInFlight.has(serverId)) { + if (this.agentDirectoryBootstrapInFlight.has(serverId) && !directorySourceChanged) { return; }