diff --git a/packages/app/src/runtime/directory-sync/index.test.ts b/packages/app/src/runtime/directory-sync/index.test.ts index 6fae0a7e92..e4db5dfb71 100644 --- a/packages/app/src/runtime/directory-sync/index.test.ts +++ b/packages/app/src/runtime/directory-sync/index.test.ts @@ -16,7 +16,8 @@ import { normalizeWorkspaceDescriptor, useSessionStore, } from "@/stores/session-store"; -import { normalizeAgentSnapshot } from "@/utils/agent-snapshots"; +import { normalizeAgentSnapshot, projectAgentSnapshot } from "@/utils/agent-snapshots"; +import { deriveProjectPlacementFromCwd } from "@/utils/project-placement"; import { selectWorkspaceDirectoryServerIds } from "@/stores/session-store-hooks/selectors"; import { useWorkspaceLabels } from "@/workspace-labels"; import type { DirectoryReplicaMutation } from "@/runtime/replica-cache"; @@ -41,12 +42,15 @@ class FakeDirectoryClient { supportsWorkspaceLabels = false; listWorkspaceLabelsCalls = 0; fetchAgentsCalls = 0; + observeAgentsCalls = 0; lastAgentOptions: unknown; fetchWorkspacesCalls = 0; + observeWorkspacesCalls = 0; lastWorkspaceOptions: unknown; listProjectsCalls = 0; lastProjectOptions: unknown; projectResult: ProjectListResult | null = null; + readonly listSnapshots: Promise[] = []; private pendingAgentFetch: Promise | null = null; private pendingWorkspaceFetch: Promise | null = null; private pendingProjectFetch: Promise | null = null; @@ -151,15 +155,22 @@ class FakeDirectoryClient { } observeAgents(options: Parameters[0]) { - return subscriptionFixture(this.fetchAgents({ ...options, subscribe: {} }), (receive) => - this.on("agent_update", receive), + this.observeAgentsCalls += 1; + const subscription = subscriptionFixture( + this.fetchAgents({ ...options, subscribe: {} }), + (receive) => this.on("agent_update", receive), ); + this.listSnapshots.push(subscription.ready); + return subscription; } observeWorkspaces(options: Parameters[0]) { - return subscriptionFixture(this.fetchWorkspaces(options), (receive) => + this.observeWorkspacesCalls += 1; + const subscription = subscriptionFixture(this.fetchWorkspaces(options), (receive) => this.on("workspace_update", receive), ); + this.listSnapshots.push(subscription.ready); + return subscription; } observeEvents(events: readonly SessionOutboundMessage["type"][]) { @@ -592,6 +603,83 @@ describe("DirectorySync session readiness", () => { directory.dispose(); }); + it.each([ + ["full", "before"], + ["full", "after"], + ["route", "before"], + ["route", "after"], + ] as const)( + "restores live directory updates when %s demand returns %s refresh completion", + async (demand, timing) => { + const serverId = `released-directory-${demand}-${timing}`; + const { client, directory } = createDirectory(serverId); + const store = useSessionStore.getState(); + store.initializeSession(serverId, client as unknown as DaemonClient, 1); + store.updateSessionServerInfo(serverId, { + serverId, + hostname: null, + version: "test", + features: { workspaceMultiplicity: true, projectList: true }, + }); + const source = {}; + function setDemand(demanded: boolean): void { + if (demand === "full") directory.setDemand(source, demanded); + else directory.setAgentRouteDemand(demanded ? ["agent-1"] : []); + } + const releaseProjects = client.holdProjectFetch(); + const projects: ProjectListResult = { requestId: "projects", projects: [] }; + + try { + setDemand(true); + const refresh = directory.refreshDemand(); + await expect + .poll(() => [client.observeAgentsCalls, client.observeWorkspacesCalls]) + .toEqual([1, 1]); + + await Promise.all(client.listSnapshots); + setDemand(false); + if (timing === "before") setDemand(true); + releaseProjects(projects); + await refresh; + if (timing === "after") setDemand(true); + + await expect + .poll(() => [client.observeAgentsCalls, client.observeWorkspacesCalls]) + .toEqual([2, 2]); + client.emit({ + type: "agent_update", + payload: { + kind: "upsert", + agent: { + ...projectAgentSnapshot(createAgent(serverId, "agent-1")), + title: "Live title", + }, + project: deriveProjectPlacementFromCwd("/repo"), + }, + }); + client.emit({ + type: "workspace_update", + payload: { + kind: "upsert", + workspace: { ...createWorkspaceEntry("workspace-1"), name: "Live workspace" }, + }, + }); + + await expect + .poll(() => { + const session = useSessionStore.getState().sessions[serverId]; + return [ + session.agents.get("agent-1")?.title, + session.workspaces.get("workspace-1")?.name, + ]; + }) + .toEqual(["Live title", "Live workspace"]); + } finally { + directory.dispose(); + } + }, + ); + it("publishes cached directory data before the authoritative request completes", async () => { const serverId = "cached-directory"; serverIds.add(serverId); diff --git a/packages/app/src/runtime/directory-sync/index.ts b/packages/app/src/runtime/directory-sync/index.ts index b59440d782..1ffb4fae5f 100644 --- a/packages/app/src/runtime/directory-sync/index.ts +++ b/packages/app/src/runtime/directory-sync/index.ts @@ -160,6 +160,7 @@ export class DirectorySync { private readonly fullDemandSources = new Set(); private readonly scriptStatusDemand = new Map(); private demandRefresh: Promise | null = null; + private subscriptionGeneration = 0; private satisfiedDemandSource: DirectorySourceToken | null = null; private cursors: DirectoryCheckpoint = {}; @@ -262,6 +263,7 @@ export class DirectorySync { } private releaseSubscriptions(): void { + this.subscriptionGeneration += 1; for (const subscription of [ this.agentSubscription, this.workspaceSubscription, @@ -362,6 +364,7 @@ export class DirectorySync { this.observeDirectoryEvents(); this.observeScriptStatusEvents(); const source = this.connection.source; + const subscriptionGeneration = this.subscriptionGeneration; if ( !force && this.satisfiedDemandSource?.clientGeneration === source.clientGeneration && @@ -378,17 +381,19 @@ export class DirectorySync { ]).then(() => undefined); this.demandRefresh = refresh .then(() => { - this.satisfiedDemandSource = source; + if (this.subscriptionGeneration === subscriptionGeneration) { + this.satisfiedDemandSource = source; + } return undefined; }) .finally(() => { this.demandRefresh = null; const current = this.connection.source; - if ( - this.hasDemand() && - (current.clientGeneration !== source.clientGeneration || - current.connectionEpoch !== source.connectionEpoch) - ) { + const sourceChanged = + current.clientGeneration !== source.clientGeneration || + current.connectionEpoch !== source.connectionEpoch; + const subscriptionsReleased = this.subscriptionGeneration !== subscriptionGeneration; + if (this.hasDemand() && (sourceChanged || subscriptionsReleased)) { void this.requestDemandRefresh().catch(() => undefined); } });