diff --git a/extensions/codex/src/app-server/run-attempt-resources.ts b/extensions/codex/src/app-server/run-attempt-resources.ts index 3e862b74fc66..ea0271253097 100644 --- a/extensions/codex/src/app-server/run-attempt-resources.ts +++ b/extensions/codex/src/app-server/run-attempt-resources.ts @@ -21,8 +21,8 @@ import type { CodexAttemptPrompt } from "./run-attempt-prompt.js"; import { releaseCodexSandboxExecServerEnvironment } from "./sandbox-exec-server.js"; import type { CodexAppServerThreadBinding } from "./session-binding.js"; import { + clearSharedCodexAppServerClientIfCurrentAndUnclaimed, retainSharedCodexAppServerClientIfCurrent, - retireSharedCodexAppServerClientIfCurrent, } from "./shared-client.js"; import type { CodexAppServerThreadLifecycleBinding } from "./thread-lifecycle.js"; import { createCodexTrajectoryRecorder, type CodexHostTrajectoryRecorder } from "./trajectory.js"; @@ -132,16 +132,17 @@ export function prepareCodexAttemptResources(prompt: CodexAttemptPrompt) { return; } state.sharedCodexClientRetiredForOneShotCleanup = true; - const retired = retireSharedCodexAppServerClientIfCurrent(state.client); - embeddedAgentLog.info("codex app-server one-shot cleanup retired shared client", { + const retired = clearSharedCodexAppServerClientIfCurrentAndUnclaimed(state.client); + embeddedAgentLog.info("codex app-server one-shot cleanup checked shared client retirement", { runId: params.runId, sessionId: params.sessionId, sessionKey: params.sessionKey, - activeLeases: retired?.activeLeases ?? null, - closed: retired?.closed ?? false, - matchedSharedClient: Boolean(retired), + activeLeases: retired.activeLeases, + pendingAcquires: retired.pendingAcquires, + closed: retired.closed, + matchedSharedClient: retired.found, }); - if (retired?.closed) { + if (retired.closed) { await state.client.closeAndWait({ exitTimeoutMs: 2_000, forceKillDelayMs: 250 }); } }; diff --git a/extensions/codex/src/app-server/run-attempt.test.ts b/extensions/codex/src/app-server/run-attempt.test.ts index 7b24be2a06a3..8e50a4483b33 100644 --- a/extensions/codex/src/app-server/run-attempt.test.ts +++ b/extensions/codex/src/app-server/run-attempt.test.ts @@ -2022,8 +2022,11 @@ describe("runCodexAppServerAttempt", () => { }); it("retires the shared Codex app-server client after one-shot cleanup turns", async () => { - const retireSpy = vi.spyOn(sharedClientModule, "retireSharedCodexAppServerClientIfCurrent"); - retireSpy.mockReturnValue({ activeLeases: 0, closed: true }); + const retireSpy = vi.spyOn( + sharedClientModule, + "clearSharedCodexAppServerClientIfCurrentAndUnclaimed", + ); + retireSpy.mockReturnValue({ found: true, activeLeases: 0, pendingAcquires: 0, closed: true }); const events: string[] = []; const closeAndWait = vi.fn(async () => { events.push("closeAndWait"); @@ -2086,8 +2089,11 @@ describe("runCodexAppServerAttempt", () => { }); it("retires the shared Codex app-server client after one-shot turn start failures", async () => { - const retireSpy = vi.spyOn(sharedClientModule, "retireSharedCodexAppServerClientIfCurrent"); - retireSpy.mockReturnValue({ activeLeases: 0, closed: true }); + const retireSpy = vi.spyOn( + sharedClientModule, + "clearSharedCodexAppServerClientIfCurrentAndUnclaimed", + ); + retireSpy.mockReturnValue({ found: true, activeLeases: 0, pendingAcquires: 0, closed: true }); const events: string[] = []; const closeAndWait = vi.fn(async () => { events.push("closeAndWait"); @@ -2132,7 +2138,10 @@ describe("runCodexAppServerAttempt", () => { }); it("keeps the shared Codex app-server client warm without one-shot cleanup", async () => { - const retireSpy = vi.spyOn(sharedClientModule, "retireSharedCodexAppServerClientIfCurrent"); + const retireSpy = vi.spyOn( + sharedClientModule, + "clearSharedCodexAppServerClientIfCurrentAndUnclaimed", + ); const closeAndWait = vi.fn(async () => true); let notify: ((notification: CodexServerNotification) => Promise) | undefined; const request = vi.fn(async (method: string) => { diff --git a/extensions/codex/src/app-server/shared-client.test.ts b/extensions/codex/src/app-server/shared-client.test.ts index 66e2500466e7..225fdf2c8715 100644 --- a/extensions/codex/src/app-server/shared-client.test.ts +++ b/extensions/codex/src/app-server/shared-client.test.ts @@ -88,6 +88,7 @@ import { let listCodexAppServerModels: typeof import("./models.js").listCodexAppServerModels; let clearSharedCodexAppServerClient: typeof import("./shared-client.js").clearSharedCodexAppServerClient; let clearSharedCodexAppServerClientIfCurrent: typeof import("./shared-client.js").clearSharedCodexAppServerClientIfCurrent; +let clearSharedCodexAppServerClientIfCurrentAndUnclaimed: typeof import("./shared-client.js").clearSharedCodexAppServerClientIfCurrentAndUnclaimed; let clearSharedCodexAppServerClientIfCurrentAndWait: typeof import("./shared-client.js").clearSharedCodexAppServerClientIfCurrentAndWait; let createIsolatedCodexAppServerClient: typeof import("./shared-client.js").createIsolatedCodexAppServerClient; let getLeasedSharedCodexAppServerClient: typeof import("./shared-client.js").getLeasedSharedCodexAppServerClient; @@ -188,6 +189,7 @@ describe("shared Codex app-server client", () => { ({ clearSharedCodexAppServerClient, clearSharedCodexAppServerClientIfCurrent, + clearSharedCodexAppServerClientIfCurrentAndUnclaimed, clearSharedCodexAppServerClientIfCurrentAndWait, createIsolatedCodexAppServerClient, getLeasedSharedCodexAppServerClient, @@ -1731,6 +1733,44 @@ describe("shared Codex app-server client", () => { expect(releaseLeasedSharedCodexAppServerClient(first.client)).toBe(false); }); + it("keeps the current client registered while a staggered sibling lease is active", async () => { + const first = createClientHarness(); + const replacement = createClientHarness(); + const startSpy = vi + .spyOn(CodexAppServerClient, "start") + .mockReturnValueOnce(first.client) + .mockReturnValueOnce(replacement.client); + + const completedRunLease = getLeasedSharedCodexAppServerClient({ timeoutMs: 1000 }); + const siblingRunLease = getLeasedSharedCodexAppServerClient({ timeoutMs: 1000 }); + await sendInitializeResult(first, "openclaw/0.143.0 (macOS; test)"); + await expect(completedRunLease).resolves.toBe(first.client); + await expect(siblingRunLease).resolves.toBe(first.client); + + expect(releaseLeasedSharedCodexAppServerClient(first.client)).toBe(true); + expect(clearSharedCodexAppServerClientIfCurrentAndUnclaimed(first.client)).toEqual({ + found: true, + closed: false, + activeLeases: 1, + pendingAcquires: 0, + }); + expect(first.process.stdin.destroyed).toBe(false); + + const staggeredLease = await getLeasedSharedCodexAppServerClient({ timeoutMs: 1000 }); + expect(staggeredLease).toBe(first.client); + expect(startSpy).toHaveBeenCalledTimes(1); + + expect(releaseLeasedSharedCodexAppServerClient(first.client)).toBe(true); + expect(releaseLeasedSharedCodexAppServerClient(first.client)).toBe(true); + expect(clearSharedCodexAppServerClientIfCurrentAndUnclaimed(first.client)).toEqual({ + found: true, + closed: true, + activeLeases: 0, + pendingAcquires: 0, + }); + expect(first.process.stdin.destroyed).toBe(true); + }); + it("rejects pending acquires during shared-client retirement", async () => { const first = createClientHarness(); const second = createClientHarness();