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>

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<T>(): Deferred<T> {
   return { promise, resolve, reject };
 }
 
+async function drainAsyncGenerator<T>(generator: AsyncGenerator<T>): Promise<void> {
+  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<void> {}
 }
 
+class ResumeTrackingTestAgentClient extends TestAgentClient {
+  private readonly retryStarted = deferred<void>();
+
+  override async resumeSession(
+    _handle: AgentPersistenceHandle,
+    config?: Partial<AgentSessionConfig>,
+  ): Promise<AgentSession> {
+    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<void> {
+    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<AgentSession> {
+      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<AgentSession> {
+      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<void> {
+      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<AgentSession> {
+      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");

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 };

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<unknown>;
+};
 
 export interface StartAgentRunOptions {
   replaceRunning?: boolean;
@@ -78,6 +81,14 @@ async function startOrReplaceRun(
   return { iterator, replaced };
 }
 
+async function drainAgentRunIterator(
+  iterator: AsyncGenerator<import("./agent-sdk-types.js").AgentStreamEvent>,
+): Promise<void> {
+  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(
         {

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<void>((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());

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<Value> {
   promise: Promise<Value>;
@@ -130,6 +131,10 @@ class ProviderRuntime {
     return this.connection?.capabilities ?? [];
   }
 
+  get isClosed(): boolean {
+    return this.closed;
+  }
+
   async isAvailable(): Promise<boolean> {
     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<Extract<ProviderEvent, { type: "session.prompt_result" }>>();
     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);
     }

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();

packages/server/src/server/agent/provider-snapshot-manager.ts

@@ -191,6 +191,7 @@ export interface AgentManagerProviderState {
     >
   >;
   clients: Partial<Record<AgentProvider, AgentClient>>;
+  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<AgentProvider>();
     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(

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;
+}

packages/server/src/server/schedule/service.ts

@@ -202,6 +202,7 @@ function buildRunOutput(params: {
 type ScheduleAgentManager = Pick<
   AgentRunController,
   | "getAgent"
+  | "reloadAgentSession"
   | "tryRunOutOfBand"
   | "hasInFlightRun"
   | "replaceAgentRun"