fix(codex): clean up early bootstrap returns

vincentkoc · vincentkoc · commit c99a46c76230 · 2026-05-27T13:05:35.000+02:00
diff --git a/extensions/codex/src/app-server/run-attempt.context-engine.test.ts b/extensions/codex/src/app-server/run-attempt.context-engine.test.ts
@@ -6,6 +6,7 @@ import { SessionManager } from "@earendil-works/pi-coding-agent";
 import type { EmbeddedRunAttemptParams } from "openclaw/plugin-sdk/agent-harness";
 import {
   embeddedAgentLog,
+  nativeHookRelayTesting,
   PREEMPTIVE_OVERFLOW_ERROR_TEXT,
   queueAgentHarnessMessage,
   type HarnessContextEngine as ContextEngine,
@@ -311,6 +312,22 @@ function getRequestInputTextAt(
     .join("\n");
 }
 
+function extractRelayIdFromThreadRequest(params: unknown): string {
+  const config = (params as { config?: Record<string, unknown> }).config;
+  const entries = config?.["hooks.PreToolUse"];
+  if (!Array.isArray(entries)) {
+    throw new Error("pre-tool hook config missing");
+  }
+  const command = (entries as Array<{ hooks?: Array<{ command?: string }> }>)
+    .flatMap((entry) => entry.hooks ?? [])
+    .find((hook) => typeof hook.command === "string")?.command;
+  const match = command?.match(/--relay-id ([^ ]+)/);
+  if (!match?.[1]) {
+    throw new Error(`relay id missing from command: ${command}`);
+  }
+  return match[1];
+}
+
 type CodexNativeThreadLifecycleEvent = Extract<
   DiagnosticEventPayload,
   { type: "codex.native_thread.lifecycle" }
@@ -337,6 +354,7 @@ describe("runCodexAppServerAttempt context-engine lifecycle", () => {
 
   afterEach(async () => {
     resetCodexAppServerClientFactoryForTest();
+    nativeHookRelayTesting.clearNativeHookRelaysForTests();
     resetDiagnosticEventsForTest();
     vi.restoreAllMocks();
     await fs.rm(tempDir, { recursive: true, force: true });
@@ -1568,14 +1586,19 @@ describe("runCodexAppServerAttempt context-engine lifecycle", () => {
     params.contextEngine = contextEngine;
     params.contextTokenBudget = 16_000;
 
-    const result = await runCodexAppServerAttempt(params);
+    const result = await runCodexAppServerAttempt(params, {
+      nativeHookRelay: { enabled: true, events: ["pre_tool_use"] },
+    });
 
     expect(result.promptError).toBe(PREEMPTIVE_OVERFLOW_ERROR_TEXT);
     expect(result.promptErrorSource).toBe("precheck");
     expect(result.preflightRecovery?.route).not.toBe("fits");
     expect(compact).not.toHaveBeenCalled();
     expect(assemble).toHaveBeenCalledTimes(2);
     expect(await readCodexAppServerBinding(sessionFile)).toBeUndefined();
+    const startRequest = harness.requests.find((request) => request.method === "thread/start");
+    const relayId = extractRelayIdFromThreadRequest(startRequest?.params);
+    expect(nativeHookRelayTesting.getNativeHookRelayRegistrationForTests(relayId)).toBeUndefined();
     expect(harness.requests.map((request) => request.method)).toEqual([
       "thread/resume",
       "thread/start",
diff --git a/extensions/codex/src/app-server/run-attempt.ts b/extensions/codex/src/app-server/run-attempt.ts
@@ -2036,6 +2036,15 @@ export async function runCodexAppServerAttempt(
       await releaseCodexSandboxExecServerEnvironment(sandbox);
     }
   };
+  const cleanupStartedThreadBeforeTurnReturn = async () => {
+    nativeHookRelay?.unregister();
+    nativeHookRelay = undefined;
+    await releaseSandboxExecEnvironment();
+    params.abortSignal?.removeEventListener("abort", abortFromUpstream);
+    releaseSharedClientLease?.();
+    releaseSharedClientLease = undefined;
+    startupClientForCleanup = undefined;
+  };
   let codexEnvironmentSelection: CodexTurnEnvironmentParams[] | undefined;
   let codexExecutionCwd = effectiveWorkspace;
   let codexSandboxPolicy: CodexSandboxPolicy | undefined;
@@ -2357,20 +2366,20 @@ export async function runCodexAppServerAttempt(
             threadId: thread.threadId,
             timeoutMs: CODEX_APP_SERVER_UNSUBSCRIBE_TIMEOUT_MS,
           });
-          await releaseSandboxExecEnvironment();
-          params.abortSignal?.removeEventListener("abort", abortFromUpstream);
-          releaseSharedClientLease?.();
-          releaseSharedClientLease = undefined;
+          await cleanupStartedThreadBeforeTurnReturn();
           return providerBoundaryPrecheckFailure;
         }
       } catch (error) {
-        params.abortSignal?.removeEventListener("abort", abortFromUpstream);
         if (runAbortController.signal.aborted || isAbortLikeError(error)) {
-          await releaseSandboxExecEnvironment();
-          releaseSharedClientLease?.();
-          releaseSharedClientLease = undefined;
+          await clearCodexAppServerBinding(activeSessionFile);
+          await unsubscribeCodexThreadBestEffort(client, {
+            threadId: thread.threadId,
+            timeoutMs: CODEX_APP_SERVER_UNSUBSCRIBE_TIMEOUT_MS,
+          });
+          await cleanupStartedThreadBeforeTurnReturn();
           return buildPreStartAbortResult(error);
         }
+        params.abortSignal?.removeEventListener("abort", abortFromUpstream);
         throw error;
       }
       systemPromptReport = buildCurrentSystemPromptReport();
