From 8598b6aab9fd29696d595fb5ea612836f0299c54 Mon Sep 17 00:00:00 2001 From: Matt Cowger Date: Thu, 10 Sep 2026 00:20:25 -0700 Subject: [PATCH] fix(plugin-provider): recover stale sessions after reload (#4629) * fix(plugin-provider): recover stale sessions after reload * fix(plugin-provider): auto-recover stale sessions on next prompt * fix(plugin-provider): retry stale prompt streams * fix(plugin-provider): recover replacement prompts * test(plugin-provider): await runtime retirement --------- Co-authored-by: paseo-bot[bot] <266920839+paseo-bot[bot]@users.noreply.github.com> --- .../src/server/agent/agent-manager.test.ts | 228 ++++++++++++++++++ .../server/src/server/agent/agent-manager.ts | 21 ++ .../server/src/server/agent/agent-prompt.ts | 46 +++- .../src/server/agent/plugin-provider.test.ts | 94 +++++++- .../src/server/agent/plugin-provider.ts | 11 + .../agent/provider-snapshot-manager.test.ts | 1 + .../server/agent/provider-snapshot-manager.ts | 6 +- .../agent/stale-provider-session-error.ts | 11 + .../server/src/server/schedule/service.ts | 1 + 9 files changed, 414 insertions(+), 5 deletions(-) create mode 100644 packages/server/src/server/agent/stale-provider-session-error.ts diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 1e18ba9de4..f4d4069d26 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -19,6 +19,7 @@ import { toAgentPayload } from "./agent-projections.js"; import { projectTimelineRows } from "./timeline-projection.js"; import { getOpenAgentTabLabel, PARENT_AGENT_ID_LABEL } from "@getpaseo/protocol/agent-labels"; import { formatSystemNotificationPrompt, startAgentRun } from "./agent-prompt.js"; +import { StaleProviderSessionError } from "./stale-provider-session-error.js"; import { ensureAgentLoaded, ensureUnarchivedAgentLoaded } from "./agent-loading.js"; import type { StoredAgentRecord } from "./agent-storage.js"; import type { @@ -68,6 +69,12 @@ function deferred(): Deferred { return { promise, resolve, reject }; } +async function drainAsyncGenerator(generator: AsyncGenerator): Promise { + for await (const _ of generator) { + // Drain provider events while AgentManager subscribers observe them. + } +} + function waitForAgentLifecycle( manager: AgentManager, agentId: string, @@ -519,6 +526,32 @@ class TestAgentSession implements AgentSession { async close(): Promise {} } +class ResumeTrackingTestAgentClient extends TestAgentClient { + private readonly retryStarted = deferred(); + + override async resumeSession( + _handle: AgentPersistenceHandle, + config?: Partial, + ): Promise { + this.resumeOverrides.push(config); + const signalRetryStarted = () => this.retryStarted.resolve(); + return new (class extends TestAgentSession { + override async startTurn(): Promise<{ turnId: string }> { + signalRetryStarted(); + return await super.startTurn(); + } + })({ + provider: this.provider, + cwd: config?.cwd ?? process.cwd(), + daemonAppendSystemPrompt: config?.daemonAppendSystemPrompt, + }); + } + + waitForRetryStart(): Promise { + return this.retryStarted.promise; + } +} + class McpCapableTestAgentSession extends TestAgentSession { override readonly capabilities = { ...TEST_CAPABILITIES, @@ -3677,6 +3710,201 @@ test("updateProviderRegistry removes providers omitted from the next registry", expect(removedClient.createSessionCalls).toBe(0); }); +test("retires loaded agents when their plugin provider is replaced", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-plugin-provider-reload-")); + const storage = new AgentStorage(join(workdir, "agents"), logger); + const provider = "plugin-provider"; + + class OriginalPluginClient extends TestAgentClient { + session: CloseRecordingTestAgentSession | null = null; + + constructor() { + super(provider); + } + + override async createSession(config: AgentSessionConfig): Promise { + const session = new CloseRecordingTestAgentSession(config); + session.describePersistence = () => ({ + provider, + sessionId: "plugin-session", + }); + this.session = session; + return session; + } + } + + const original = new OriginalPluginClient(); + const replacement = new TestAgentClient(provider); + const manager = new AgentManager({ + clients: { [provider]: original }, + providerDefinitions: { [provider]: { enabled: true } }, + registry: storage, + logger, + }); + const created = await manager.createAgent({ provider, cwd: workdir }, undefined, { + workspaceId: undefined, + }); + + try { + manager.updateProviderRegistry({ + providerDefinitions: { [provider]: { enabled: true } }, + clients: { [provider]: replacement }, + retiredProviders: [provider], + }); + + await manager.waitForAgentClose(created.id); + expect(original.session?.closed).toBe(true); + expect(manager.getAgent(created.id)).toBeNull(); + + const resumed = await ensureAgentLoaded(created.id, { + agentManager: manager, + agentStorage: storage, + logger, + }); + expect(resumed.id).toBe(created.id); + expect(replacement.resumeOverrides).toEqual([ + expect.objectContaining({ cwd: workdir, provider }), + ]); + } finally { + await manager.closeAgent(created.id).catch(() => undefined); + await manager.flush().catch(() => undefined); + await storage.flush().catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + +test("a prompt after provider replacement reopens the stale session", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-stale-prompt-")); + const storage = new AgentStorage(join(workdir, "agents"), logger); + const provider = "plugin-provider"; + + class StalePluginSession extends CloseRecordingTestAgentSession { + override async startTurn(): Promise<{ turnId: string }> { + throw new StaleProviderSessionError("stale-bridge"); + } + } + + class StalePluginClient extends TestAgentClient { + constructor() { + super(provider); + } + + override async createSession(config: AgentSessionConfig): Promise { + const session = new StalePluginSession(config); + session.describePersistence = () => ({ + provider, + sessionId: "plugin-session", + }); + return session; + } + } + + const staleClient = new StalePluginClient(); + const replacement = new ResumeTrackingTestAgentClient(provider); + const manager = new AgentManager({ + clients: { [provider]: staleClient }, + providerDefinitions: { [provider]: { enabled: true } }, + registry: storage, + logger, + }); + const created = await manager.createAgent({ provider, cwd: workdir }, undefined, { + workspaceId: undefined, + }); + + try { + manager.updateProviderRegistry({ + providerDefinitions: { [provider]: { enabled: true } }, + clients: { [provider]: replacement }, + }); + + const dispatch = await startAgentRun(manager, created.id, "continue after reload", logger); + expect(dispatch.disposition).toBe("turn_started"); + await replacement.waitForRetryStart(); + const result = await manager.waitForAgentEvent(created.id); + expect(result.status).toBe("idle"); + expect(replacement.resumeOverrides).toHaveLength(1); + } finally { + await manager.closeAgent(created.id).catch(() => undefined); + await manager.flush().catch(() => undefined); + await storage.flush().catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + +test("a replacement prompt recovers when the retired session fails to start", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-stale-replacement-")); + const storage = new AgentStorage(join(workdir, "agents"), logger); + const provider = "plugin-provider"; + + class StaleReplacementSession extends TestAgentSession { + private starts = 0; + + override async startTurn(): Promise<{ turnId: string }> { + this.starts += 1; + if (this.starts > 1) throw new StaleProviderSessionError("stale-bridge"); + return { turnId: "initial-turn" }; + } + + override async interrupt(): Promise { + if (this.starts > 1) throw new StaleProviderSessionError("stale-bridge"); + this.pushEvent({ + type: "turn_canceled", + provider: this.provider, + turnId: "initial-turn", + }); + } + } + + class StaleReplacementClient extends TestAgentClient { + constructor() { + super(provider); + } + + override async createSession(config: AgentSessionConfig): Promise { + const session = new StaleReplacementSession(config); + session.describePersistence = () => ({ provider, sessionId: "plugin-session" }); + return session; + } + } + + const replacement = new ResumeTrackingTestAgentClient(provider); + const manager = new AgentManager({ + clients: { [provider]: new StaleReplacementClient() }, + providerDefinitions: { [provider]: { enabled: true } }, + registry: storage, + logger, + rescueTimeouts: { interruptSessionMs: 20 }, + }); + const created = await manager.createAgent({ provider, cwd: workdir }, undefined, { + workspaceId: undefined, + }); + + try { + const initial = manager.streamAgent(created.id, "initial"); + void drainAsyncGenerator(initial); + await manager.waitForAgentRunStart(created.id); + + manager.updateProviderRegistry({ + providerDefinitions: { [provider]: { enabled: true } }, + clients: { [provider]: replacement }, + }); + + const dispatch = await startAgentRun(manager, created.id, "replacement", logger, { + replaceRunning: true, + }); + expect(dispatch.disposition).toBe("turn_started"); + await replacement.waitForRetryStart(); + const result = await manager.waitForAgentEvent(created.id); + expect(result.status).toBe("idle"); + expect(replacement.resumeOverrides).toHaveLength(1); + } finally { + await manager.closeAgent(created.id).catch(() => undefined); + await manager.flush().catch(() => undefined); + await storage.flush().catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + test("createAgent passes explicit model strings through to the provider", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-")); const storagePath = join(workdir, "agents"); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 5fd1c85847..70b4e62dbc 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -77,6 +77,7 @@ import { } from "./agent-run-state.js"; import { invokeRewindCapability, type RewindMode } from "./rewind/rewind.js"; import { isSystemInjectedEnvelope } from "./agent-prompt.js"; +import { isStaleProviderSessionError } from "./stale-provider-session-error.js"; import { stripInternalPaseoMcpServer, withRuntimePaseoMcpServer } from "./runtime-mcp-config.js"; import { resolveCreateAgentTitles } from "./create-agent-title.js"; import type { PaseoToolCatalogFactory } from "./tools/types.js"; @@ -775,6 +776,7 @@ export class AgentManager { updateProviderRegistry(input: { providerDefinitions: ProviderEnabledMap; clients: ProviderClientMap; + retiredProviders?: readonly AgentProvider[]; }): void { this.providerEnabled.clear(); this.providerDefinitions.clear(); @@ -791,6 +793,18 @@ export class AgentManager { this.clients.set(provider, client); } } + + for (const provider of input.retiredProviders ?? []) { + for (const agent of this.agents.values()) { + if (agent.provider !== provider) continue; + void this.closeAgent(agent.id).catch((error) => { + this.logger.warn( + { err: error, agentId: agent.id, provider }, + "Failed to close agent after provider retirement", + ); + }); + } + } } getRegisteredProviderIds(): AgentProvider[] { @@ -2360,6 +2374,13 @@ export class AgentManager { if (pendingRun.settled) { throw error; } + if (isStaleProviderSessionError(error)) { + pendingRun.start = { status: "failed", error: error.message }; + agent.pendingReplacement = false; + if (!agent.activeForegroundTurnId) agent.lifecycle = "idle"; + this.runs.settleForegroundRun(agentId, pendingRun.token); + throw error; + } agent.pendingReplacement = false; const errorMsg = error instanceof Error ? error.message : "Failed to start turn"; pendingRun.start = { status: "failed", error: errorMsg }; diff --git a/packages/server/src/server/agent/agent-prompt.ts b/packages/server/src/server/agent/agent-prompt.ts index 96f5ce5ee3..f46822d924 100644 --- a/packages/server/src/server/agent/agent-prompt.ts +++ b/packages/server/src/server/agent/agent-prompt.ts @@ -8,6 +8,7 @@ import type { import type { AgentManager, ManagedAgent } from "./agent-manager.js"; import type { AgentStorage } from "./agent-storage.js"; import { ensureAgentLoaded } from "./agent-loading.js"; +import { isStaleProviderSessionError } from "./stale-provider-session-error.js"; import { getParentAgentIdFromLabels } from "@getpaseo/protocol/agent-labels"; import type { ActiveTurnBehavior } from "@getpaseo/protocol/messages"; @@ -21,7 +22,9 @@ export type AgentRunController = Pick< | "replaceAgentRun" | "steerOrReplaceActiveTurn" | "streamAgent" ->; +> & { + reloadAgentSession(agentId: string): Promise; +}; export interface StartAgentRunOptions { replaceRunning?: boolean; @@ -78,6 +81,14 @@ async function startOrReplaceRun( return { iterator, replaced }; } +async function drainAgentRunIterator( + iterator: AsyncGenerator, +): Promise { + for await (const _ of iterator) { + // Events are broadcast via AgentManager subscribers. + } +} + export async function startAgentRun( agentManager: AgentRunController, agentId: string, @@ -104,6 +115,26 @@ export async function startAgentRun( if (agentManager.tryRunOutOfBand(agentId, prompt, options?.runOptions)) { return { disposition: "out_of_band" }; } + try { + return await startAgentRunInner(agentManager, agentId, prompt, logger, options); + } catch (error) { + if (!isStaleProviderSessionError(error)) throw error; + logger.info({ agentId, err: error }, "Provider session went stale; reopening from persistence"); + // The live session belongs to a retired plugin runtime. Reload swaps in a + // fresh session on the current runtime while preserving history and labels. + await agentManager.reloadAgentSession(agentId); + return await startAgentRunInner(agentManager, agentId, prompt, logger, options); + } +} + +async function startAgentRunInner( + agentManager: AgentRunController, + agentId: string, + prompt: AgentPromptInput, + logger: Logger, + options?: StartAgentRunOptions, +): Promise<{ disposition: PromptDispatchDisposition }> { + const snapshot = agentManager.getAgent(agentId); const steered = await steerOrReplaceActiveRun(agentManager, agentId, prompt, options); if (steered?.disposition === "steered") { return steered; @@ -122,8 +153,17 @@ export async function startAgentRun( ); void (async () => { try { - for await (const _ of iterator) { - // Events are broadcast via AgentManager subscribers. + try { + await drainAgentRunIterator(iterator); + } catch (error) { + if (!isStaleProviderSessionError(error)) throw error; + logger.info( + { agentId, err: error }, + "Provider session went stale; reopening from persistence", + ); + await agentManager.reloadAgentSession(agentId); + const retry = await startOrReplaceRun(agentManager, agentId, prompt, options); + await drainAgentRunIterator(retry.iterator); } logger.trace( { diff --git a/packages/server/src/server/agent/plugin-provider.test.ts b/packages/server/src/server/agent/plugin-provider.test.ts index 75046396be..99d156473d 100644 --- a/packages/server/src/server/agent/plugin-provider.test.ts +++ b/packages/server/src/server/agent/plugin-provider.test.ts @@ -8,6 +8,10 @@ import { describe, expect, test } from "vitest"; import { createTestLogger } from "../../test-utils/test-logger.js"; import type { AgentStreamEvent } from "./agent-sdk-types.js"; import { PluginAgentClientRegistry } from "./plugin-provider.js"; +import { + isStaleProviderSessionError, + StaleProviderSessionError, +} from "./stale-provider-session-error.js"; const CAPABILITIES = [ "prompt.message", @@ -25,6 +29,10 @@ interface ProviderHarnessOptions { function createProviderHarness(options: ProviderHarnessOptions = {}) { let listener: ((event: ProviderEvent) => void) | null = null; let closeCount = 0; + let resolveClosed!: () => void; + const closed = new Promise((resolve) => { + resolveClosed = resolve; + }); const inputs: ProviderInput[] = []; const emit = (event: ProviderEvent) => listener?.(event); const capabilities = options.capabilities ?? CAPABILITIES; @@ -180,6 +188,7 @@ function createProviderHarness(options: ProviderHarnessOptions = {}) { }, async close() { closeCount += 1; + resolveClosed(); }, }; @@ -191,7 +200,12 @@ function createProviderHarness(options: ProviderHarnessOptions = {}) { }, }; - return { registration, inputs, closeCount: () => closeCount }; + return { + registration, + inputs, + closeCount: () => closeCount, + waitForClose: () => closed, + }; } function eventsOfType(events: AgentStreamEvent[], type: AgentStreamEvent["type"]) { @@ -258,6 +272,84 @@ describe("PluginAgentClientRegistry", () => { expect(eventsOfType(events, "turn_failed")).toHaveLength(1); }); + test("closes a stale session after its plugin provider is replaced", async () => { + const old = createProviderHarness(); + const next = createProviderHarness(); + const registry = new PluginAgentClientRegistry(createTestLogger()); + + try { + registry.replace([old.registration]); + + const stale = await registry.clients()[old.registration.id]!.createSession({ + provider: old.registration.id, + cwd: "/workspace", + }); + const persistence = stale.describePersistence(); + expect(persistence).not.toBeNull(); + + registry.replace([next.registration]); + await old.waitForClose(); + expect(old.closeCount()).toBe(1); + + await expect(stale.close()).resolves.toBeUndefined(); + + const replacement = registry.clients()[next.registration.id]; + expect(replacement).toBeDefined(); + const resumed = await replacement!.resumeSession(persistence!, { + cwd: "/workspace", + }); + + await expect( + resumed.startTurn("after reload", { clientMessageId: "after-reload" }), + ).resolves.toEqual({ turnId: "turn-1" }); + + expect(next.inputs).toContainEqual( + expect.objectContaining({ + type: "session.open", + history: "replay", + persistence: { + version: 1, + data: { token: "root" }, + }, + }), + ); + + await resumed.close(); + } finally { + await registry.shutdown(); + } + }); + + test("prompting a stale session raises StaleProviderSessionError", async () => { + const old = createProviderHarness(); + const registry = new PluginAgentClientRegistry(createTestLogger()); + + try { + registry.replace([old.registration]); + const stale = await registry.clients()[old.registration.id]!.createSession({ + provider: old.registration.id, + cwd: "/workspace", + }); + + registry.replace([]); + await old.waitForClose(); + expect(old.closeCount()).toBe(1); + + const failure = await stale + .startTurn("after reload", { clientMessageId: "after-reload" }) + .then( + () => null, + (error: unknown) => error, + ); + expect(failure).toBeInstanceOf(StaleProviderSessionError); + expect(isStaleProviderSessionError(failure)).toBe(true); + expect(isStaleProviderSessionError(new Error("Provider connection is closed"))).toBe(false); + expect(isStaleProviderSessionError(new Error("boom"))).toBe(false); + } finally { + await registry.shutdown(); + } + }); + test("adapts callback providers into the existing AgentClient and AgentSession path", async () => { const harness = createProviderHarness(); const registry = new PluginAgentClientRegistry(createTestLogger()); diff --git a/packages/server/src/server/agent/plugin-provider.ts b/packages/server/src/server/agent/plugin-provider.ts index b5f5ff7b92..0ab3ddecd7 100644 --- a/packages/server/src/server/agent/plugin-provider.ts +++ b/packages/server/src/server/agent/plugin-provider.ts @@ -58,6 +58,7 @@ import { } from "./create-agent-mode.js"; import type { ProviderDefinition } from "./provider-registry.js"; import { runProviderTurn } from "./providers/provider-runner.js"; +import { StaleProviderSessionError } from "./stale-provider-session-error.js"; interface Deferred { promise: Promise; @@ -130,6 +131,10 @@ class ProviderRuntime { return this.connection?.capabilities ?? []; } + get isClosed(): boolean { + return this.closed; + } + async isAvailable(): Promise { await this.getConnection(); return true; @@ -559,6 +564,7 @@ class ProviderRuntimeSession { sessionId: this.providerSessionId, prompt, }); + if (this.terminal) throw new StaleProviderSessionError(this.id); const pending = deferred>(); this.prompts.set(prompt.clientMessageId, pending); try { @@ -568,6 +574,9 @@ class ProviderRuntimeSession { prompt, }); return (await pending.promise).result; + } catch (error) { + if (this.terminal || this.runtime.isClosed) throw new StaleProviderSessionError(this.id); + throw error; } finally { this.prompts.delete(prompt.clientMessageId); } @@ -623,6 +632,8 @@ class ProviderRuntimeSession { requestId: randomUUID(), sessionId: this.providerSessionId, }); + } catch (error) { + if (!this.runtime.isClosed) throw error; } finally { this.runtime.removeSession(this.id, this.providerSessionId); } diff --git a/packages/server/src/server/agent/provider-snapshot-manager.test.ts b/packages/server/src/server/agent/provider-snapshot-manager.test.ts index 7ca0510f51..68837c8c5f 100644 --- a/packages/server/src/server/agent/provider-snapshot-manager.test.ts +++ b/packages/server/src/server/agent/provider-snapshot-manager.test.ts @@ -2281,6 +2281,7 @@ describe("ProviderSnapshotManager cwd routing", () => { const withoutPlugin = manager.replacePluginProviders([]); expect(manager.hasProvider(registration.id)).toBe(false); expect(withoutPlugin.clients[registration.id]).toBeUndefined(); + expect(withoutPlugin.retiredProviders).toEqual([registration.id]); } finally { await manager.shutdown(); manager.destroy(); diff --git a/packages/server/src/server/agent/provider-snapshot-manager.ts b/packages/server/src/server/agent/provider-snapshot-manager.ts index abf88d4488..610f7c8322 100644 --- a/packages/server/src/server/agent/provider-snapshot-manager.ts +++ b/packages/server/src/server/agent/provider-snapshot-manager.ts @@ -191,6 +191,7 @@ export interface AgentManagerProviderState { > >; clients: Partial>; + retiredProviders?: readonly AgentProvider[]; } interface ProviderLoadOptions { @@ -402,6 +403,9 @@ export class ProviderSnapshotManager { this.createAgentManagerState(this.generation.definitions, clients); this.pluginProviders.replace(registrations); const plugins = this.pluginProviders.definitions(); + const retiredProviders = Object.keys(previousPlugins).filter( + (provider) => previousPlugins[provider] !== plugins[provider], + ); const definitions = { ...this.generation.definitions }; const changed = new Set(); for (const provider of new Set([...Object.keys(previousPlugins), ...Object.keys(plugins)])) { @@ -415,7 +419,7 @@ export class ProviderSnapshotManager { const generation = this.createGeneration(definitions, this.providerOverrides); const state = this.createAgentManagerState(definitions, clients); this.installGeneration(generation, clients, changed); - return state; + return { ...state, retiredProviders }; } private ensureClient( diff --git a/packages/server/src/server/agent/stale-provider-session-error.ts b/packages/server/src/server/agent/stale-provider-session-error.ts new file mode 100644 index 0000000000..7039da3c61 --- /dev/null +++ b/packages/server/src/server/agent/stale-provider-session-error.ts @@ -0,0 +1,11 @@ +/** Thrown when a prompt reaches a plugin session whose connection died with its plugin. */ +export class StaleProviderSessionError extends Error { + constructor(sessionId: string) { + super(`Provider session ${sessionId} is stale after its plugin reloaded`); + this.name = "StaleProviderSessionError"; + } +} + +export function isStaleProviderSessionError(error: unknown): error is StaleProviderSessionError { + return error instanceof StaleProviderSessionError; +} diff --git a/packages/server/src/server/schedule/service.ts b/packages/server/src/server/schedule/service.ts index e3e9504e46..192b424337 100644 --- a/packages/server/src/server/schedule/service.ts +++ b/packages/server/src/server/schedule/service.ts @@ -202,6 +202,7 @@ function buildRunOutput(params: { type ScheduleAgentManager = Pick< AgentRunController, | "getAgent" + | "reloadAgentSession" | "tryRunOutOfBand" | "hasInFlightRun" | "replaceAgentRun"