fix(codex): scope app-server migration cleanup

steipete · steipete · commit 191bd7dc9ab6 · 2026-05-17T01:46:39.000+01:00
diff --git a/extensions/codex/src/app-server/shared-client.test.ts b/extensions/codex/src/app-server/shared-client.test.ts
@@ -39,6 +39,7 @@ vi.mock("openclaw/plugin-sdk/agent-runtime", () => ({
 let listCodexAppServerModels: typeof import("./models.js").listCodexAppServerModels;
 let clearSharedCodexAppServerClient: typeof import("./shared-client.js").clearSharedCodexAppServerClient;
 let clearSharedCodexAppServerClientIfCurrent: typeof import("./shared-client.js").clearSharedCodexAppServerClientIfCurrent;
+let clearSharedCodexAppServerClientIfCurrentAndWait: typeof import("./shared-client.js").clearSharedCodexAppServerClientIfCurrentAndWait;
 let createIsolatedCodexAppServerClient: typeof import("./shared-client.js").createIsolatedCodexAppServerClient;
 let getSharedCodexAppServerClient: typeof import("./shared-client.js").getSharedCodexAppServerClient;
 let resetSharedCodexAppServerClientForTests: typeof import("./shared-client.js").resetSharedCodexAppServerClientForTests;
@@ -111,6 +112,7 @@ describe("shared Codex app-server client", () => {
     ({
       clearSharedCodexAppServerClient,
       clearSharedCodexAppServerClientIfCurrent,
+      clearSharedCodexAppServerClientIfCurrentAndWait,
       createIsolatedCodexAppServerClient,
       getSharedCodexAppServerClient,
       resetSharedCodexAppServerClientForTests,
@@ -476,6 +478,44 @@ describe("shared Codex app-server client", () => {
     expect(second.process.stdin.destroyed).toBe(true);
   });
 
+  it("waits only for the shared client that is still current", async () => {
+    const first = createClientHarness();
+    const second = createClientHarness();
+    vi.spyOn(CodexAppServerClient, "start")
+      .mockReturnValueOnce(first.client)
+      .mockReturnValueOnce(second.client);
+    const firstCloseAndWait = vi.spyOn(first.client, "closeAndWait");
+    const secondCloseAndWait = vi.spyOn(second.client, "closeAndWait");
+
+    const firstList = listCodexAppServerModels({
+      timeoutMs: 1000,
+      agentDir: "/tmp/openclaw-agent-one",
+    });
+    await sendInitializeResult(first, "openclaw/0.125.0 (macOS; test)");
+    await sendEmptyModelList(first);
+    await expect(firstList).resolves.toEqual({ models: [] });
+
+    const secondList = listCodexAppServerModels({
+      timeoutMs: 1000,
+      agentDir: "/tmp/openclaw-agent-two",
+    });
+    await sendInitializeResult(second, "openclaw/0.125.0 (macOS; test)");
+    await sendEmptyModelList(second);
+    await expect(secondList).resolves.toEqual({ models: [] });
+
+    await expect(
+      clearSharedCodexAppServerClientIfCurrentAndWait(first.client, {
+        exitTimeoutMs: 25,
+        forceKillDelayMs: 5,
+      }),
+    ).resolves.toBe(true);
+
+    expect(firstCloseAndWait).toHaveBeenCalledTimes(1);
+    expect(secondCloseAndWait).not.toHaveBeenCalled();
+    expect(first.process.stdin.destroyed).toBe(true);
+    expect(second.process.stdin.destroyed).toBe(false);
+  });
+
   it("uses a fresh websocket Authorization header after shared-client token rotation", async () => {
     const server = new WebSocketServer({ host: "127.0.0.1", port: 0 });
     const authHeaders: Array<string | undefined> = [];
diff --git a/extensions/codex/src/app-server/shared-client.ts b/extensions/codex/src/app-server/shared-client.ts
@@ -198,6 +198,27 @@ export function clearSharedCodexAppServerClientIfCurrent(
   return false;
 }
 
+export async function clearSharedCodexAppServerClientIfCurrentAndWait(
+  client: CodexAppServerClient | undefined,
+  options?: {
+    exitTimeoutMs?: number;
+    forceKillDelayMs?: number;
+  },
+): Promise<boolean> {
+  if (!client) {
+    return false;
+  }
+  const state = getSharedCodexAppServerClientState();
+  for (const [key, entry] of state.clients) {
+    if (entry.client === client) {
+      state.clients.delete(key);
+      await client.closeAndWait(options);
+      return true;
+    }
+  }
+  return false;
+}
+
 export async function clearSharedCodexAppServerClientAndWait(options?: {
   exitTimeoutMs?: number;
   forceKillDelayMs?: number;
diff --git a/extensions/codex/src/migration/apply.ts b/extensions/codex/src/migration/apply.ts
@@ -40,7 +40,7 @@ import { buildCodexPluginAppCacheKey } from "../app-server/plugin-app-cache-key.
 import type { v2 } from "../app-server/protocol.js";
 import { requestCodexAppServerJson } from "../app-server/request.js";
 import {
-  clearSharedCodexAppServerClientAndWait,
+  clearSharedCodexAppServerClientIfCurrentAndWait,
   getSharedCodexAppServerClient,
 } from "../app-server/shared-client.js";
 import { buildCodexMigrationPlan } from "./plan.js";
@@ -84,19 +84,22 @@ export function prepareTargetCodexAppServer(
 ): CodexMigrationTargetAppServerPreparation {
   const appServer = resolveTargetCodexAppServer(ctx);
   const targets = resolveCodexMigrationTargets(ctx);
+  let warmedClient: Awaited<ReturnType<typeof getSharedCodexAppServerClient>> | undefined;
   const ready = getSharedCodexAppServerClient({
     startOptions: appServer.start,
     timeoutMs: 60_000,
     agentDir: targets.agentDir,
     config: ctx.config,
   }).then(
-    () => undefined,
+    (client) => {
+      warmedClient = client;
+    },
     () => undefined,
   );
   return {
     async dispose() {
       await ready;
-      await clearSharedCodexAppServerClientAndWait({
+      await clearSharedCodexAppServerClientIfCurrentAndWait(warmedClient, {
         exitTimeoutMs: 2_000,
         forceKillDelayMs: 250,
       });
