From 019bf657908aac8120379386f6acb4063be3e004 Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Wed, 22 Apr 2026 16:45:57 +0700
Subject: [PATCH] Fix running-agent replacement lifecycle

---
 .../src/server/agent/agent-manager.test.ts    | 30 +++++++++++++++----
 .../server/src/server/agent/agent-manager.ts  | 30 +++++++++++++++++--
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts
index 455666184..dc8ca9f3f 100644
--- a/packages/server/src/server/agent/agent-manager.test.ts
+++ b/packages/server/src/server/agent/agent-manager.test.ts
@@ -2253,6 +2253,8 @@ describe("AgentManager", () => {
     const storagePath = join(workdir, "agents");
     const storage = new AgentStorage(storagePath, logger);
     const secondStartEntered = deferred<void>();
+    const interruptStarted = deferred<void>();
+    const allowInterruptToFinish = deferred<void>();
     const allowSecondStartToResolve = deferred<void>();
     let capturedSession: StaleReplacementSession | null = null;
 
@@ -2280,6 +2282,8 @@ describe("AgentManager", () => {
       }
 
       override async interrupt(): Promise<void> {
+        interruptStarted.resolve();
+        await allowInterruptToFinish.promise;
         this.pushEvent({
           type: "turn_canceled",
           provider: this.provider,
@@ -2310,13 +2314,16 @@ describe("AgentManager", () => {
       cwd: workdir,
     });
 
-    const lifecycleUpdates: string[] = [];
+    const stateUpdates: Array<{ lifecycle: string; updatedAt: number }> = [];
     const unsubscribe = manager.subscribe(
       (event) => {
         if (event.type !== "agent_state" || event.agent.id !== snapshot.id) {
           return;
         }
-        lifecycleUpdates.push(event.agent.lifecycle);
+        stateUpdates.push({
+          lifecycle: event.agent.lifecycle,
+          updatedAt: event.agent.updatedAt.getTime(),
+        });
       },
       { agentId: snapshot.id, replayState: false },
     );
@@ -2330,6 +2337,8 @@ describe("AgentManager", () => {
 
     await manager.waitForAgentRunStart(snapshot.id);
 
+    const replaceUpdatesStart = stateUpdates.length;
+    const beforeReplaceUpdatedAt = manager.getAgent(snapshot.id)?.updatedAt.getTime() ?? 0;
     const secondRun = manager.replaceAgentRun(snapshot.id, "replacement run");
     const secondRunDrain = (async () => {
       for await (const _event of secondRun) {
@@ -2337,6 +2346,16 @@ describe("AgentManager", () => {
       }
     })();
 
+    await interruptStarted.promise;
+    const replacementUpdates = stateUpdates.slice(replaceUpdatesStart);
+    expect(
+      replacementUpdates.some(
+        (update) => update.lifecycle === "running" && update.updatedAt > beforeReplaceUpdatedAt,
+      ),
+    ).toBe(true);
+    expect(replacementUpdates.map((update) => update.lifecycle)).not.toContain("idle");
+    allowInterruptToFinish.resolve();
+
     await secondStartEntered.promise;
 
     const replaceGapSnapshot = manager.getAgent(snapshot.id) as
@@ -2349,10 +2368,11 @@ describe("AgentManager", () => {
     capturedSession!.pushEvent({ type: "turn_completed", provider: "codex", turnId: "turn-1" });
     await new Promise((resolve) => setTimeout(resolve, 0));
 
-    // eslint-disable-next-line no-console
-    console.log("replace-gap-lifecycle", lifecycleUpdates, manager.getAgent(snapshot.id));
     expect(manager.getAgent(snapshot.id)?.lifecycle).toBe("running");
-    expect(lifecycleUpdates.at(-1)).toBe("running");
+    expect(stateUpdates.at(-1)?.lifecycle).toBe("running");
+    expect(stateUpdates.slice(replaceUpdatesStart).map((update) => update.lifecycle)).not.toContain(
+      "idle",
+    );
 
     allowSecondStartToResolve.resolve();
 
diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts
index 3764623cf..072c1309a 100644
--- a/packages/server/src/server/agent/agent-manager.ts
+++ b/packages/server/src/server/agent/agent-manager.ts
@@ -187,6 +187,7 @@ type ManagedAgentBase = {
   lastError?: string;
   attention: AttentionState;
   foregroundTurnWaiters: Set<ForegroundTurnWaiter>;
+  finalizedForegroundTurnIds: Set<string>;
   unsubscribeSession: (() => void) | null;
   /**
    * Internal agents are hidden from listings and don't trigger notifications.
@@ -1302,8 +1303,11 @@ export class AgentManager {
     return streamForwarder;
   }
 
-  private finalizeForegroundTurn(agent: ActiveManagedAgent): void {
+  private finalizeForegroundTurn(agent: ActiveManagedAgent, turnId?: string): void {
     const mutableAgent = agent as ActiveManagedAgent;
+    if (turnId) {
+      this.rememberFinalizedForegroundTurn(mutableAgent, turnId);
+    }
     mutableAgent.activeForegroundTurnId = null;
     const terminalError = mutableAgent.lastError;
     const shouldHoldBusyForReplacement = mutableAgent.pendingReplacement && !terminalError;
@@ -1351,6 +1355,9 @@ export class AgentManager {
 
     const agent = this.requireSessionAgent(agentId);
     agent.pendingReplacement = true;
+    agent.lifecycle = "running";
+    this.touchUpdatedAt(agent);
+    this.emitState(agent);
 
     const self = this;
     return (async function* replaceRunForwarder() {
@@ -1990,6 +1997,7 @@ export class AgentManager {
       pendingReplacement: false,
       activeForegroundTurnId: null,
       foregroundTurnWaiters: new Set<ForegroundTurnWaiter>(),
+      finalizedForegroundTurnIds: new Set<string>(),
       unsubscribeSession: null,
       persistence: attachPersistenceCwd(session.describePersistence(), config.cwd),
       historyPrimed: options?.historyPrimed ?? durableTimelineHasRows,
@@ -2150,6 +2158,17 @@ export class AgentManager {
     waiter.resolveSettled();
   }
 
+  private rememberFinalizedForegroundTurn(agent: ActiveManagedAgent, turnId: string): void {
+    agent.finalizedForegroundTurnIds.add(turnId);
+    if (agent.finalizedForegroundTurnIds.size <= 50) {
+      return;
+    }
+    const oldest = agent.finalizedForegroundTurnIds.values().next().value;
+    if (oldest) {
+      agent.finalizedForegroundTurnIds.delete(oldest);
+    }
+  }
+
   private createPendingForegroundRun(): PendingForegroundRun {
     let resolveSettled!: () => void;
     const settledPromise = new Promise<void>((resolve) => {
@@ -2334,6 +2353,13 @@ export class AgentManager {
   ): Promise<boolean> {
     const eventTurnId = (event as { turnId?: string }).turnId;
     const isForegroundEvent = Boolean(eventTurnId && agent.activeForegroundTurnId === eventTurnId);
+    if (
+      eventTurnId &&
+      isTurnTerminalEvent(event) &&
+      agent.finalizedForegroundTurnIds.has(eventTurnId)
+    ) {
+      return false;
+    }
 
     // Only update timestamp for live events, not history replay
     if (!options?.fromHistory) {
@@ -2540,7 +2566,7 @@ export class AgentManager {
     }
 
     if (!options?.fromHistory && isForegroundEvent && isTurnTerminalEvent(event)) {
-      this.finalizeForegroundTurn(agent);
+      this.finalizeForegroundTurn(agent, eventTurnId);
     }
 
     // Skip dispatching individual stream events during history replay.