fix(codex): retire shared app-server client only when unclaimed (#110227)

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

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<void>) | undefined;
     const request = vi.fn(async (method: string) => {

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