diff --git a/packages/server/src/server/daemon-e2e/concurrent-agent-refresh.e2e.test.ts b/packages/server/src/server/daemon-e2e/concurrent-agent-refresh.e2e.test.ts new file mode 100644 index 0000000000..3958add65b --- /dev/null +++ b/packages/server/src/server/daemon-e2e/concurrent-agent-refresh.e2e.test.ts @@ -0,0 +1,150 @@ +import { mkdirSync, mkdtempSync, rmSync } from "node:fs"; +import path from "node:path"; +import { afterEach, expect, test } from "vitest"; +import type { AgentClient, AgentSession } from "../agent/agent-sdk-types.js"; +import { createTestAgentClient } from "../test-utils/fake-agent-client.js"; +import { DaemonClient } from "../test-utils/daemon-client.js"; +import { createTestPaseoDaemon, type TestPaseoDaemon } from "../test-utils/paseo-daemon.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +const artifactRoot = path.resolve(".artifacts"); +mkdirSync(artifactRoot, { recursive: true }); +const resources: { daemon: TestPaseoDaemon; clients: DaemonClient[]; cwd: string }[] = []; +const releases: (() => void)[] = []; +afterEach(async () => { + releases.splice(0).forEach((release) => release()); + for (const { daemon, clients, cwd } of resources.splice(0)) { + await Promise.all(clients.map((client) => client.close())); + await daemon.close(); + rmSync(cwd, { recursive: true, force: true }); + } +}); + +function heldHistoryProvider() { + const historyStarted = deferred(); + const releaseHistory = deferred(); + releases.push(releaseHistory.resolve); + const base = createTestAgentClient("claude"); + let resumes = 0; + const provider: AgentClient = { + provider: base.provider, + capabilities: base.capabilities, + createSession: (...args) => base.createSession(...args), + fetchCatalog: (...args) => base.fetchCatalog(...args), + isAvailable: (...args) => base.isAvailable(...args), + async resumeSession(...args) { + resumes += 1; + const session = await base.resumeSession(...args); + // Delay provider history through its public interface, keeping the manager real. + return new Proxy(session, { + get(target, key: keyof AgentSession) { + if (key === "streamHistory") + return async function* () { + historyStarted.resolve(); + await releaseHistory.promise; + yield* target.streamHistory(); + }; + const value = Reflect.get(target, key); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + }, + }; + return { provider, historyStarted, releaseHistory, resumes: () => resumes }; +} + +async function setup(provider: AgentClient, agentId?: string) { + const cwd = mkdtempSync(path.join(artifactRoot, "concurrent-refresh-")); + const daemon = await createTestPaseoDaemon({ agentClients: { claude: provider } }); + const clients = [0, 1].map(() => new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` })); + resources.push({ daemon, clients, cwd }); + await Promise.all(clients.map((client) => client.connect())); + const agent = await daemon.daemon.agentManager.createAgent({ provider: "claude", cwd }, agentId, { + workspaceId: undefined, + }); + return { clients, agent }; +} + +test("refreshes from separate connections share the reload until history is ready", async () => { + const provider = heldHistoryProvider(); + const { clients, agent } = await setup(provider.provider); + const first = clients[0].refreshAgent(agent.id); + await provider.historyStarted.promise; + const second = clients[1].refreshAgent(agent.id); + // The following RPC traverses the same connection after the second refresh request. + await clients[1].fetchAgents(); + provider.releaseHistory.resolve(); + const results = await Promise.all([first, second]); + expect(results.map((result) => result.agentId)).toEqual([agent.id, agent.id]); + expect(provider.resumes()).toBe(1); + await clients[0].refreshAgent(agent.id); + expect(provider.resumes()).toBe(2); +}); + +test("refreshing the same id in separate daemons does not share a provider reload", async () => { + const provider = heldHistoryProvider(); + const firstDaemon = await setup(provider.provider); + const secondDaemon = await setup(createTestAgentClient("claude"), firstDaemon.agent.id); + const first = firstDaemon.clients[0].refreshAgent(firstDaemon.agent.id); + await provider.historyStarted.promise; + const second = secondDaemon.clients[0].refreshAgent(secondDaemon.agent.id); + await secondDaemon.clients[0].fetchAgents(); + // The other daemon must finish without releasing the first one's provider history. + let deadline: ReturnType | undefined; + let independent: boolean; + try { + independent = await Promise.race([ + second.then(() => true), + new Promise((resolve) => { + deadline = setTimeout(() => resolve(false), 5000); + }), + ]); + } finally { + clearTimeout(deadline); + } + provider.releaseHistory.resolve(); + await Promise.all([first, second]); + expect(independent).toBe(true); +}); + +test("concurrent refreshes share a failure and a later retry starts a new reload", async () => { + const resumeStarted = deferred(); + const releaseResume = deferred(); + releases.push(releaseResume.resolve); + let resumes = 0; + const provider = createTestAgentClient("claude", { + beforeResumeSession: async () => { + resumes += 1; + if (resumes === 1) { + resumeStarted.resolve(); + await releaseResume.promise; + throw new Error("provider resume failed"); + } + }, + }); + const { clients, agent } = await setup(provider); + const first = clients[0].refreshAgent(agent.id); + // Attach handlers before releasing the rejected provider operation. + const outcome = Promise.allSettled([first]); + await resumeStarted.promise; + const second = clients[1].refreshAgent(agent.id); + const retryOutcome = Promise.allSettled([second]); + await clients[1].fetchAgents(); + releaseResume.resolve(); + const results = [...(await outcome), ...(await retryOutcome)]; + expect(results.map((result) => result.status)).toEqual(["rejected", "rejected"]); + for (const result of results) { + if (result.status === "rejected") + expect(String(result.reason)).toContain("provider resume failed"); + } + expect(resumes).toBe(1); + await expect(clients[0].refreshAgent(agent.id)).resolves.toMatchObject({ agentId: agent.id }); + expect(resumes).toBe(2); +}); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 55b46aa9f8..02b94a3643 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -682,8 +682,8 @@ interface ClientActivity { appVisibilityChangedAt: Date; } -// The agent refreshes in flight on this daemon, by agent id; see handleRefreshAgentRequest. -const pendingAgentRefreshes = new Map>(); +// Connections on one daemon share refreshes; independent managers own independent agents. +const pendingAgentRefreshes = new WeakMap>>(); export class Session { private readonly connectedHostSession: ConnectedHostSession | null; @@ -4761,20 +4761,21 @@ export class Session { this.sessionLogger.info({ agentId }, `Refreshing agent ${agentId} from persistence`); try { - // A refresh takes the agent out of the manager while it resumes the provider session, so a - // second refresh arriving meanwhile, from any connection, found no agent and failed with - // "Unknown agent". A client that gives up waiting and retries makes exactly that second - // request. It waits for the refresh in flight and answers with its outcome instead. - let refresh = pendingAgentRefreshes.get(agentId); + // Runtime replacement is serialized, but the reply also needs provider history. + // Another replacement during hydration can remove the agent before that reply. + // Share the whole reload and hydration across connections, including failures. + const refreshes = + pendingAgentRefreshes.get(this.agentManager) ?? new Map>(); + pendingAgentRefreshes.set(this.agentManager, refreshes); + let refresh = refreshes.get(agentId); if (!refresh) { const started = this.refreshAgentFromPersistence(agentId); refresh = started; - pendingAgentRefreshes.set(agentId, started); + refreshes.set(agentId, started); void started .catch(() => undefined) .finally(() => { - if (pendingAgentRefreshes.get(agentId) === started) - pendingAgentRefreshes.delete(agentId); + if (refreshes.get(agentId) === started) refreshes.delete(agentId); }); } const snapshot = await refresh; diff --git a/packages/server/src/server/session.workspaces.test.ts b/packages/server/src/server/session.workspaces.test.ts index f875dcba66..9e73c9cad4 100644 --- a/packages/server/src/server/session.workspaces.test.ts +++ b/packages/server/src/server/session.workspaces.test.ts @@ -5426,77 +5426,6 @@ test("refresh_agent_request leaves workspace archival independent when its direc expect(findByType(emitted, "rpc_error")).toBeUndefined(); }); -test("refresh_agent_request arriving while a refresh of the same agent resumes waits for it", async () => { - // A client gives up on a slow refresh after its request timeout and asks again. The refresh in - // flight has taken the agent out of the manager, so a second one used to fail with "Unknown - // agent", and every retry after it did the same while the first was still resuming. - const emitted: SessionOutboundMessage[] = []; - const session = createSessionForWorkspaceTests({ - onMessage: (message) => { - if (isSessionOutboundMessage(message)) emitted.push(message); - }, - }); - const agentId = "agent-refreshed-twice"; - const cwd = path.resolve("/tmp/paseo-refresh-twice"); - session.filesystem.isDirectory = async () => true; - session.agentStorage.get = async (id: string) => - id === agentId - ? makeStoredAgent({ id: agentId, cwd, updatedAt: "2026-03-10T00:00:00.000Z" }) - : null; - session.agentStorage.upsert = async () => {}; - const managed = makeManagedAgent({ - id: agentId, - cwd, - lifecycle: "idle", - updatedAt: "2026-03-10T00:00:00.000Z", - }); - let loaded: ManagedAgent | null = managed; - let releaseResume!: () => void; - const resumed = new Promise((resolve) => { - releaseResume = resolve; - }); - let reloads = 0; - session.agentManager.getAgent = () => loaded; - session.interruptAgentIfRunning = async () => undefined; - session.agentManager.reloadAgentSession = async () => { - reloads += 1; - loaded = null; - await resumed; - loaded = managed; - return managed; - }; - session.agentManager.hydrateTimelineFromProvider = async () => { - if (!loaded) throw new Error(`Unknown agent '${agentId}'`); - }; - session.agentManager.getTimeline = () => []; - session.agentUpdates.forwardLiveAgent = async () => undefined; - - const first = session.handleMessage({ - type: "refresh_agent_request", - agentId, - requestId: "req-first", - }); - await waitForImmediate(); - const retry = session.handleMessage({ - type: "refresh_agent_request", - agentId, - requestId: "req-retry", - }); - await waitForImmediate(); - releaseResume(); - await Promise.all([first, retry]); - - expect(reloads).toBe(1); - expect(findByType(emitted, "rpc_error")).toBeUndefined(); - const refreshed = emitted.filter( - (m) => m.type === "status" && (m.payload as { status?: string }).status === "agent_refreshed", - ); - expect(refreshed.map((m) => (m.payload as { requestId?: string }).requestId).sort()).toEqual([ - "req-first", - "req-retry", - ]); -}); - test("refresh_agent_request leaves workspace archival independent when its directory is missing", async () => { const emitted: SessionOutboundMessage[] = []; const session = createSessionForWorkspaceTests({