fix(codex): close one-shot app server clients

aliahnaf2013-max · web-flow · commit 870ec6dee291 · 2026-06-14T17:55:31.000-07:00
Fix gateway-routed one-shot Codex app-server teardown so owned shared clients are retired after run cleanup. Verified with focused tests, Showboat proof, and green PR CI.
diff --git a/extensions/codex/src/app-server/run-attempt.test.ts b/extensions/codex/src/app-server/run-attempt.test.ts
@@ -1810,6 +1810,162 @@ describe("runCodexAppServerAttempt", () => {
     expect(request.mock.calls.map(([method]) => method)).not.toContain("app/list");
   });
 
+  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 events: string[] = [];
+    const closeAndWait = vi.fn(async () => {
+      events.push("closeAndWait");
+      return true;
+    });
+    let startedClient: unknown;
+    let notify: ((notification: CodexServerNotification) => Promise<void>) | undefined;
+    setCodexAppServerClientFactoryForTest(async () => {
+      const client = {
+        request: vi.fn(async (method: string) => {
+          events.push(`request:${method}`);
+          if (method === "thread/start") {
+            return threadStartResult();
+          }
+          if (method === "turn/start") {
+            return turnStartResult();
+          }
+          return {};
+        }),
+        addNotificationHandler: vi.fn((handler) => {
+          notify = handler;
+          return () => undefined;
+        }),
+        addRequestHandler: vi.fn(() => () => undefined),
+        addCloseHandler: vi.fn(() => () => undefined),
+        closeAndWait,
+      };
+      startedClient = client;
+      return client as never;
+    });
+    const params = createParams(
+      path.join(tempDir, "session.jsonl"),
+      path.join(tempDir, "workspace"),
+    );
+    params.cleanupBundleMcpOnRunEnd = true;
+
+    const run = runCodexAppServerAttempt(params);
+    await vi.waitFor(() => expect(notify).toBeDefined(), fastWait);
+    if (!notify) {
+      throw new Error("expected turn notification handler");
+    }
+    await notify({
+      method: "turn/completed",
+      params: {
+        threadId: "thread-1",
+        turnId: "turn-1",
+        turn: { id: "turn-1", status: "completed" },
+      },
+    });
+    await run;
+
+    expect(retireSpy).toHaveBeenCalledWith(startedClient);
+    expect(closeAndWait).toHaveBeenCalledWith({ exitTimeoutMs: 2_000, forceKillDelayMs: 250 });
+    expect(events.indexOf("request:thread/unsubscribe")).toBeGreaterThan(-1);
+    expect(events.indexOf("closeAndWait")).toBeGreaterThan(
+      events.indexOf("request:thread/unsubscribe"),
+    );
+  });
+
+  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 events: string[] = [];
+    const closeAndWait = vi.fn(async () => {
+      events.push("closeAndWait");
+      return true;
+    });
+    let startedClient: unknown;
+    setCodexAppServerClientFactoryForTest(async () => {
+      const client = {
+        request: vi.fn(async (method: string) => {
+          events.push(`request:${method}`);
+          if (method === "thread/start") {
+            return threadStartResult();
+          }
+          if (method === "turn/start") {
+            throw new Error("turn start failed");
+          }
+          return {};
+        }),
+        addNotificationHandler: vi.fn(() => () => undefined),
+        addRequestHandler: vi.fn(() => () => undefined),
+        addCloseHandler: vi.fn(() => () => undefined),
+        closeAndWait,
+      };
+      startedClient = client;
+      return client as never;
+    });
+    const params = createParams(
+      path.join(tempDir, "session.jsonl"),
+      path.join(tempDir, "workspace"),
+    );
+    params.cleanupBundleMcpOnRunEnd = true;
+
+    await expect(runCodexAppServerAttempt(params)).rejects.toThrow("turn start failed");
+
+    expect(retireSpy).toHaveBeenCalledWith(startedClient);
+    expect(closeAndWait).toHaveBeenCalledWith({ exitTimeoutMs: 2_000, forceKillDelayMs: 250 });
+    expect(events.indexOf("request:thread/unsubscribe")).toBeGreaterThan(-1);
+    expect(events.indexOf("closeAndWait")).toBeGreaterThan(
+      events.indexOf("request:thread/unsubscribe"),
+    );
+  });
+
+  it("keeps the shared Codex app-server client warm without one-shot cleanup", async () => {
+    const retireSpy = vi.spyOn(sharedClientModule, "retireSharedCodexAppServerClientIfCurrent");
+    const closeAndWait = vi.fn(async () => true);
+    let notify: ((notification: CodexServerNotification) => Promise<void>) | undefined;
+    setCodexAppServerClientFactoryForTest(
+      async () =>
+        ({
+          request: vi.fn(async (method: string) => {
+            if (method === "thread/start") {
+              return threadStartResult();
+            }
+            if (method === "turn/start") {
+              return turnStartResult();
+            }
+            return {};
+          }),
+          addNotificationHandler: vi.fn((handler) => {
+            notify = handler;
+            return () => undefined;
+          }),
+          addRequestHandler: vi.fn(() => () => undefined),
+          addCloseHandler: vi.fn(() => () => undefined),
+          closeAndWait,
+        }) as never,
+    );
+    const params = createParams(
+      path.join(tempDir, "session.jsonl"),
+      path.join(tempDir, "workspace"),
+    );
+
+    const run = runCodexAppServerAttempt(params);
+    await vi.waitFor(() => expect(notify).toBeDefined(), fastWait);
+    if (!notify) {
+      throw new Error("expected turn notification handler");
+    }
+    await notify({
+      method: "turn/completed",
+      params: {
+        threadId: "thread-1",
+        turnId: "turn-1",
+        turn: { id: "turn-1", status: "completed" },
+      },
+    });
+    await run;
+
+    expect(retireSpy).not.toHaveBeenCalled();
+    expect(closeAndWait).not.toHaveBeenCalled();
+  });
+
   it("keeps searchable Codex dynamic tools canonical in mirrored transcript snapshots", async () => {
     const params = createParams(
       path.join(tempDir, "session.jsonl"),
diff --git a/extensions/codex/src/app-server/run-attempt.ts b/extensions/codex/src/app-server/run-attempt.ts
@@ -224,6 +224,7 @@ import {
   writeCodexAppServerBinding,
   type CodexAppServerThreadBinding,
 } from "./session-binding.js";
+import { retireSharedCodexAppServerClientIfCurrent } from "./shared-client.js";
 import { rotateOversizedCodexAppServerStartupBinding } from "./startup-binding.js";
 import {
   buildDeveloperInstructions,
@@ -1119,6 +1120,7 @@ export async function runCodexAppServerAttempt(
   };
   let nativeHookRelay: NativeHookRelayRegistrationHandle | undefined;
   let releaseSharedClientLease: (() => void) | undefined;
+  let sharedCodexClientRetiredForOneShotCleanup = false;
   const releaseSharedClientLeaseOnce = () => {
     const release = releaseSharedClientLease;
     if (!release) {
@@ -1127,6 +1129,31 @@ export async function runCodexAppServerAttempt(
     releaseSharedClientLease = undefined;
     release();
   };
+  const retireSharedCodexClientForOneShotCleanup = async () => {
+    if (params.cleanupBundleMcpOnRunEnd !== true) {
+      return;
+    }
+    if (sharedCodexClientRetiredForOneShotCleanup) {
+      return;
+    }
+    sharedCodexClientRetiredForOneShotCleanup = true;
+    const retired = retireSharedCodexAppServerClientIfCurrent(client);
+    embeddedAgentLog.info("codex app-server one-shot cleanup retired shared client", {
+      runId: params.runId,
+      sessionId: params.sessionId,
+      sessionKey: params.sessionKey,
+      activeLeases: retired?.activeLeases ?? null,
+      closed: retired?.closed ?? false,
+      matchedSharedClient: Boolean(retired),
+    });
+    if (retired?.closed) {
+      await client.closeAndWait({ exitTimeoutMs: 2_000, forceKillDelayMs: 250 });
+    }
+  };
+  const releaseSharedClientLeaseAndRetireOneShotClient = async () => {
+    releaseSharedClientLeaseOnce();
+    await retireSharedCodexClientForOneShotCleanup();
+  };
   let sandboxExecEnvironmentAcquired = false;
   const releaseSandboxExecEnvironment = async () => {
     if (sandboxExecEnvironmentAcquired) {
@@ -2302,7 +2329,7 @@ export async function runCodexAppServerAttempt(
         },
       });
       params.abortSignal?.removeEventListener("abort", abortFromUpstream);
-      releaseSharedClientLeaseOnce();
+      await releaseSharedClientLeaseAndRetireOneShotClient();
       if (usageLimitError) {
         await markCodexAuthProfileBlockedFromRateLimits({
           params,
@@ -2322,7 +2349,7 @@ export async function runCodexAppServerAttempt(
     }
   }
   if (!turn) {
-    releaseSharedClientLeaseOnce();
+    await releaseSharedClientLeaseAndRetireOneShotClient();
     throw new Error("codex app-server turn/start failed without an error");
   }
   turnIdRef.current = turn.turn.id;
@@ -2772,7 +2799,7 @@ export async function runCodexAppServerAttempt(
     notificationCleanup();
     requestCleanup();
     closeCleanup?.();
-    releaseSharedClientLeaseOnce();
+    await releaseSharedClientLeaseAndRetireOneShotClient();
     if (nativeHookRelay) {
       if (shouldDelayNativeHookRelayUnregister) {
         // Codex hook subprocesses can outlive a completed app-server turn by a
diff --git a/packages/gateway-protocol/src/schema/agent.ts b/packages/gateway-protocol/src/schema/agent.ts
@@ -205,8 +205,8 @@ export const AgentParamsSchema = Type.Object(
     timeout: Type.Optional(Type.Integer({ minimum: 0 })),
     bestEffortDeliver: Type.Optional(Type.Boolean()),
     lane: Type.Optional(Type.String()),
-    // Backward-compatible no-op. Older CLI clients sent this field on gateway
-    // agent requests; the gateway accepts but intentionally ignores it.
+    // One-shot CLI gateway requests can ask the gateway to close process-wide
+    // bundle MCP resources after the run instead of keeping them warm.
     cleanupBundleMcpOnRunEnd: Type.Optional(Type.Boolean()),
     modelRun: Type.Optional(Type.Boolean()),
     promptMode: Type.Optional(
diff --git a/src/agents/embedded-agent-runner/run.before-agent-reply-cron.test.ts b/src/agents/embedded-agent-runner/run.before-agent-reply-cron.test.ts
@@ -23,12 +23,20 @@ function firstBeforeAgentReplyCall() {
 }
 
 function firstAttemptParams(): {
+  cleanupBundleMcpOnRunEnd?: boolean;
   modelRun?: boolean;
   promptMode?: string;
   promptCacheKey?: string;
 } {
   const call = mockedRunEmbeddedAttempt.mock.calls[0] as
-    | [{ modelRun?: boolean; promptMode?: string; promptCacheKey?: string }]
+    | [
+        {
+          cleanupBundleMcpOnRunEnd?: boolean;
+          modelRun?: boolean;
+          promptMode?: string;
+          promptCacheKey?: string;
+        },
+      ]
     | undefined;
   if (!call) {
     throw new Error("expected embedded attempt call");
@@ -155,6 +163,17 @@ describe("runEmbeddedAgent cron before_agent_reply seam", () => {
     expect(attemptParams.promptMode).toBe("none");
   });
 
+  it("forwards one-shot bundle MCP cleanup into the embedded attempt", async () => {
+    mockedRunEmbeddedAttempt.mockResolvedValueOnce(makeAttemptResult({ promptError: null }));
+
+    await runEmbeddedAgent({
+      ...overflowBaseRunParams,
+      cleanupBundleMcpOnRunEnd: true,
+    });
+
+    expect(firstAttemptParams().cleanupBundleMcpOnRunEnd).toBe(true);
+  });
+
   it("forwards prompt cache identity into the embedded attempt", async () => {
     mockedRunEmbeddedAttempt.mockResolvedValueOnce(makeAttemptResult({ promptError: null }));
 
diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts
@@ -1869,6 +1869,7 @@ async function runEmbeddedAgentInternal(
             bootstrapContextRunKind: params.bootstrapContextRunKind,
             jobId: params.jobId,
             toolsAllow: params.toolsAllow,
+            cleanupBundleMcpOnRunEnd: params.cleanupBundleMcpOnRunEnd,
             disableMessageTool: params.disableMessageTool,
             forceMessageTool: params.forceMessageTool,
             enableHeartbeatTool: params.enableHeartbeatTool,
diff --git a/src/commands/agent-via-gateway.test.ts b/src/commands/agent-via-gateway.test.ts
@@ -301,7 +301,7 @@ describe("agentCliCommand", () => {
       expect(request.clientName).toBe("cli");
       expect(request.mode).toBe("cli");
       expect(request).not.toHaveProperty("scopes");
-      expect(request.params).not.toHaveProperty("cleanupBundleMcpOnRunEnd");
+      expect(request.params).toHaveProperty("cleanupBundleMcpOnRunEnd", true);
       expect(agentCommand).not.toHaveBeenCalled();
       expect(agentModuleLoadCount).not.toHaveBeenCalled();
       expect(runtime.log).toHaveBeenCalledWith("hello");
diff --git a/src/commands/agent-via-gateway.ts b/src/commands/agent-via-gateway.ts
@@ -704,6 +704,7 @@ async function agentViaGatewayCommand(
             timeout: timeoutSeconds,
             lane: opts.lane,
             extraSystemPrompt: opts.extraSystemPrompt,
+            cleanupBundleMcpOnRunEnd: true,
             idempotencyKey,
           },
           expectFinal: true,
