Fix Codex rollout writer handoff races

plugins/provider-codex/src/bridge/app-server-connection.ts

@@ -44,7 +44,7 @@ interface CodexAppServerRequestArgs<TResult> {
 export interface CodexAppServerConnection {
   request<TResult>(args: CodexAppServerRequestArgs<TResult>): Promise<TResult>;
   notify(method: string, params?: unknown): void;
-  kill(): void;
+  kill(): Promise<void>;
   readonly exited: boolean;
 }
 
@@ -112,6 +112,10 @@ export function createCodexAppServerConnection(
   } | null = null;
   let closeGraceTimer: NodeJS.Timeout | null = null;
   let stdoutLines: Interface | null = null;
+  let resolveExit!: () => void;
+  const exitPromise = new Promise<void>((resolve) => {
+    resolveExit = resolve;
+  });
 
   function writeLine(message: object): void {
     const stdin = child.stdin;
@@ -155,7 +159,11 @@ export function createCodexAppServerConnection(
         { spawnFailed },
       ),
     );
-    options.onExit({ ...status, stderrTail, spawnFailed });
+    try {
+      options.onExit({ ...status, stderrTail, spawnFailed });
+    } finally {
+      resolveExit();
+    }
   }
 
   if (child.stdout) {
@@ -313,7 +321,7 @@ export function createCodexAppServerConnection(
 
     kill() {
       if (finalized) {
-        return;
+        return exitPromise;
       }
       const escalation = setTimeout(() => {
         if (!finalized) {
@@ -322,6 +330,7 @@ export function createCodexAppServerConnection(
       }, KILL_ESCALATION_MS);
       escalation.unref?.();
       child.kill("SIGTERM");
+      return exitPromise;
     },
   };
 }

plugins/provider-codex/src/bridge/bridge-process.test-support.ts

@@ -0,0 +1,130 @@
+import { existsSync, readFileSync, rmSync } from "node:fs";
+import type { experimental_createBridgeJsonRpcTestHarness as createBridgeJsonRpcTestHarness } from "@get-bb/plugin-sdk/provider-bridge/testing";
+
+type BridgeJsonRpcTestHarness = ReturnType<
+  typeof createBridgeJsonRpcTestHarness
+>;
+
+const PROCESS_POLL_INTERVAL_MS = 20;
+
+function readProcessLog(processLogPath: string): string {
+  if (!existsSync(processLogPath)) {
+    return "";
+  }
+  try {
+    return readFileSync(processLogPath, "utf8");
+  } catch (error) {
+    if (error instanceof Error && Reflect.get(error, "code") === "ENOENT") {
+      return "";
+    }
+    throw error;
+  }
+}
+
+export function spawnedAppServerPids(processLogPath: string): number[] {
+  return readProcessLog(processLogPath)
+    .split("\n")
+    .filter((line) => line.startsWith("spawn:"))
+    .map((line) => Number(line.split(":")[1]))
+    .filter((pid) => Number.isSafeInteger(pid) && pid > 0);
+}
+
+function processIsAlive(pid: number): boolean {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch (error) {
+    if (error instanceof Error && Reflect.get(error, "code") === "ESRCH") {
+      return false;
+    }
+    throw error;
+  }
+}
+
+export async function waitForAppServerChildrenToExit(
+  processLogPath: string,
+  timeoutMs = 8_000,
+): Promise<void> {
+  const childPids = spawnedAppServerPids(processLogPath);
+  const deadline = Date.now() + timeoutMs;
+  while (childPids.some(processIsAlive)) {
+    if (Date.now() > deadline) {
+      throw new Error(
+        `Timed out waiting for app-server children to exit: ${JSON.stringify(childPids.filter(processIsAlive))}`,
+      );
+    }
+    await new Promise((resolveTick) =>
+      setTimeout(resolveTick, PROCESS_POLL_INTERVAL_MS),
+    );
+  }
+}
+
+export async function waitForAppServerProcessStep(
+  processLogPath: string,
+  step: string,
+  timeoutMs = 5_000,
+): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    if (readProcessLog(processLogPath).includes(`${step}:`)) {
+      return;
+    }
+    await new Promise((resolveTick) =>
+      setTimeout(resolveTick, PROCESS_POLL_INTERVAL_MS),
+    );
+  }
+  throw new Error(`Timed out waiting for app-server process step: ${step}`);
+}
+
+export async function cleanupBridgeProcessTest(args: {
+  harness: BridgeJsonRpcTestHarness | undefined;
+  cleanupId: number;
+  threadId: string;
+  providerThreadId: string;
+  processLogPath: string;
+  workspaceDir: string;
+  unstubEnvs: () => void;
+}): Promise<void> {
+  const errors: unknown[] = [];
+  if (args.harness !== undefined) {
+    try {
+      args.harness.sendRequest(args.cleanupId, "thread/stop", {
+        threadId: args.threadId,
+        providerThreadId: args.providerThreadId,
+        intent: "release",
+        activeTurnId: null,
+      });
+      await args.harness.waitForResponse(args.cleanupId).catch(() => undefined);
+    } catch (error) {
+      errors.push(error);
+    }
+  }
+  try {
+    await waitForAppServerChildrenToExit(args.processLogPath);
+  } catch (error) {
+    errors.push(error);
+  }
+  try {
+    args.harness?.restore();
+  } catch (error) {
+    errors.push(error);
+  }
+  try {
+    args.unstubEnvs();
+  } catch (error) {
+    errors.push(error);
+  }
+  try {
+    if (args.workspaceDir !== "") {
+      rmSync(args.workspaceDir, { recursive: true, force: true });
+    }
+  } catch (error) {
+    errors.push(error);
+  }
+  if (errors.length === 1) {
+    throw errors[0];
+  }
+  if (errors.length > 1) {
+    throw new AggregateError(errors, "Bridge process test cleanup failed");
+  }
+}

plugins/provider-codex/src/bridge/bridge.archived-rebuild.test.ts

@@ -1,11 +1,15 @@
-import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { mkdtempSync, writeFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { fileURLToPath } from "node:url";
 import { afterEach, beforeEach, expect, it, vi } from "vitest";
 import { experimental_createBridgeJsonRpcTestHarness as createBridgeJsonRpcTestHarness } from "@get-bb/plugin-sdk/provider-bridge/testing";
 import type { BridgeJsonRpcOutputMessage } from "@get-bb/plugin-sdk/provider-bridge/testing";
 import { handleLine } from "./bridge.js";
+import {
+  cleanupBridgeProcessTest,
+  spawnedAppServerPids,
+} from "./bridge-process.test-support.js";
 
 const THREAD_ID = "thr_archived_rebuild_1";
 const PROVIDER_THREAD_ID = "rebuild-rollout-1";
@@ -31,9 +35,9 @@ const changedSessionOptions = {
 const turnInput = [{ type: "text", text: "hello", mentions: [] }];
 
 let harness: ReturnType<typeof createBridgeJsonRpcTestHarness>;
-let workspaceDir: string;
-let archiveStatePath: string;
-let processLogPath: string;
+let workspaceDir = "";
+let archiveStatePath = "";
+let processLogPath = "";
 
 beforeEach(() => {
   workspaceDir = mkdtempSync(join(tmpdir(), "bb-codex-archived-rebuild-"));
@@ -53,52 +57,17 @@ beforeEach(() => {
 });
 
 afterEach(async () => {
-  const cleanupId = 993_001;
-  harness.sendRequest(cleanupId, "thread/stop", {
+  await cleanupBridgeProcessTest({
+    harness,
+    cleanupId: 993_001,
     threadId: THREAD_ID,
     providerThreadId: PROVIDER_THREAD_ID,
-    intent: "release",
-    activeTurnId: null,
+    processLogPath,
+    workspaceDir,
+    unstubEnvs: vi.unstubAllEnvs,
   });
-  await harness.waitForResponse(cleanupId).catch(() => undefined);
-  await waitForAppServerChildrenToExit();
-  harness.restore();
-  vi.unstubAllEnvs();
-  rmSync(workspaceDir, { recursive: true, force: true });
 });
 
-function spawnedAppServerPids(): number[] {
-  return readFileSync(processLogPath, "utf8")
-    .split("\n")
-    .filter((line) => line.startsWith("spawn:"))
-    .map((line) => Number(line.split(":")[1]));
-}
-
-function processIsAlive(pid: number): boolean {
-  try {
-    process.kill(pid, 0);
-    return true;
-  } catch (error) {
-    if (error instanceof Error && Reflect.get(error, "code") === "ESRCH") {
-      return false;
-    }
-    throw error;
-  }
-}
-
-async function waitForAppServerChildrenToExit(): Promise<void> {
-  const childPids = spawnedAppServerPids();
-  const deadline = Date.now() + 15_000;
-  while (childPids.some(processIsAlive)) {
-    if (Date.now() > deadline) {
-      throw new Error(
-        `Timed out waiting for app-server children to exit: ${JSON.stringify(childPids.filter(processIsAlive))}`,
-      );
-    }
-    await new Promise((resolveTick) => setTimeout(resolveTick, 20));
-  }
-}
-
 async function resumeThread(): Promise<void> {
   harness.sendRequest(1, "thread/resume", {
     threadId: THREAD_ID,
@@ -197,11 +166,11 @@ it("keeps the thread resumable when a settings-change rebuild hits an externally
 
 it("keeps the thread resumable when the rebuild after the child died hits an externally archived rollout", async () => {
   await resumeThread();
-  const spawnLine = readFileSync(processLogPath, "utf8")
-    .split("\n")
-    .find((line) => line.startsWith("spawn:"));
-  const childPid = Number(spawnLine?.split(":")[1]);
-  expect(Number.isInteger(childPid)).toBe(true);
+  const [childPid] = spawnedAppServerPids(processLogPath);
+  expect(childPid).toBeDefined();
+  if (childPid === undefined) {
+    throw new Error("Expected the fake app-server child to have spawned");
+  }
   process.kill(childPid, "SIGKILL");
   const deadline = Date.now() + 15_000;
   while (