From 16da3300e3ceaeca950cd97336c5c2f52f0a1c21 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 12 Mar 2026 22:24:41 +0700 Subject: [PATCH] fix: prevent Codex replacement stream from being killed by stale turn_completed notification The turn_completed notification handler called eventQueue.end(), which could terminate a replacement stream's queue if the old interrupted turn's completion arrived after the new stream had set up its queue. The streamInternal loop already breaks on terminal events, making the end() call redundant and unsafe during replaceAgentRun. --- .../providers/codex-app-server-agent.test.ts | 147 ++++++++++++++++++ .../agent/providers/codex-app-server-agent.ts | 1 - 2 files changed, 147 insertions(+), 1 deletion(-) diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index b67b1b12c..19b52eaee 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts @@ -11,9 +11,12 @@ import { } from "./codex-app-server-agent.js"; import { createTestLogger } from "../../../test-utils/test-logger.js"; import { agentConfigs } from "../../daemon-e2e/agent-configs.js"; +import { AgentManager } from "../agent-manager.js"; +import { AgentStorage } from "../agent-storage.js"; import type { AgentPermissionRequest, AgentPromptContentBlock, + AgentStreamEvent, AgentTimelineItem, } from "../agent-sdk-types.js"; @@ -179,6 +182,22 @@ function readRolloutTurnContextEfforts(rolloutPath: string): string[] { return efforts; } +type Deferred = { + promise: Promise; + resolve: (value: T) => void; + reject: (reason?: unknown) => void; +}; + +function deferred(): Deferred { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + describe("Codex app-server provider (integration)", () => { const logger = createTestLogger(); @@ -1457,6 +1476,134 @@ describe("Codex app-server provider (integration)", () => { 120000 ); + test.runIf(isCodexInstalled())( + "replaceAgentRun keeps the replacement Codex stream alive when the previous interrupted turn completes late", + async () => { + const cleanup = useTempCodexSessionDir(); + const cwd = tmpCwd("codex-replace-run-"); + const storageDir = tmpCwd("codex-replace-run-storage-"); + + try { + const manager = new AgentManager({ + clients: { + codex: new CodexAppServerAgentClient(logger), + }, + registry: new AgentStorage(storageDir, logger), + logger, + }); + + const snapshot = await manager.createAgent({ + provider: "codex", + cwd, + modeId: "auto", + model: CODEX_TEST_MODEL, + thinkingOptionId: CODEX_TEST_THINKING_OPTION_ID, + }); + + const managedAgent = ((manager as any).agents.get(snapshot.id) ?? null) as + | { session?: any } + | null; + const session = managedAgent?.session as + | { + client?: { + request: (method: string, params?: unknown, timeoutMs?: number) => Promise; + } | null; + handleNotification?: (method: string, params: unknown) => void; + } + | undefined; + if (!session?.client || !session.handleNotification) { + throw new Error("Codex session internals unavailable for replaceAgentRun regression test"); + } + + let turnStartCount = 0; + const replacementTurnInjected = deferred(); + const originalRequest = session.client.request.bind(session.client); + session.client.request = async (method: string, params?: unknown, timeoutMs?: number) => { + if (method === "turn/start") { + turnStartCount += 1; + if (turnStartCount === 1) { + queueMicrotask(() => { + session.handleNotification?.("turn/started", { + turn: { id: "initial-turn" }, + }); + }); + } else if (turnStartCount === 2) { + queueMicrotask(() => { + session.handleNotification?.("turn/completed", { + turn: { status: "interrupted" }, + }); + session.handleNotification?.("turn/started", { + turn: { id: "replacement-turn" }, + }); + session.handleNotification?.("turn/completed", { + turn: { status: "completed" }, + }); + replacementTurnInjected.resolve(undefined); + }); + } + } + return originalRequest(method, params, timeoutMs); + }; + + const firstRun = manager.streamAgent( + snapshot.id, + "Keep working until you are interrupted." + ); + const firstRunReady = deferred(); + const firstRunDrain = (async () => { + for await (const event of firstRun) { + if (event.type === "turn_started") { + firstRunReady.resolve(undefined); + } + } + })(); + + await Promise.race([ + firstRunReady.promise, + new Promise((_, reject) => + setTimeout( + () => reject(new Error("Timed out waiting for initial Codex turn to start")), + 15_000 + ) + ), + ]); + + const replacementEvents: AgentStreamEvent[] = []; + await Promise.race([ + (async () => { + for await (const event of manager.replaceAgentRun( + snapshot.id, + "Reply exactly REPLACED and stop." + )) { + replacementEvents.push(event); + } + })(), + new Promise((_, reject) => + setTimeout( + () => + reject( + new Error("Timed out waiting for replacement Codex stream to finish") + ), + 60_000 + ) + ), + ]); + + await replacementTurnInjected.promise; + await firstRunDrain; + + expect(turnStartCount).toBeGreaterThanOrEqual(2); + expect(replacementEvents.some((event) => event.type === "turn_started")).toBe(true); + expect(replacementEvents.some((event) => event.type === "turn_completed")).toBe(true); + } finally { + cleanup(); + rmSync(storageDir, { recursive: true, force: true }); + rmSync(cwd, { recursive: true, force: true }); + } + }, + 120000 + ); + test.runIf(isCodexInstalled())( "persists session metadata and resumes with history", async () => { diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index 299bf99ae..ea20b1ad8 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.ts @@ -2626,7 +2626,6 @@ class CodexAppServerAgentSession implements AgentSession { this.pendingCommandOutputDeltas.clear(); this.pendingFileChangeOutputDeltas.clear(); this.warnedIncompleteEditToolCallIds.clear(); - this.eventQueue?.end(); return; }