Fix cross-thread Codex identity binding and rollout writer handoffs

packages/agent-runtime/src/runtime-thread-identity.test.ts

@@ -109,6 +109,33 @@ describe("RuntimeThreadIdentityRegistry", () => {
     ).toBeUndefined();
   });
 
+  it("consumes a thread's pending identity slot when its identity is recorded", () => {
+    const registry = new RuntimeThreadIdentityRegistry();
+    const providerState = registry.createProviderState({ providerId: "codex" });
+    for (const threadId of ["thread-1", "thread-2"]) {
+      registry.registerThreadProvider({
+        providerId: "codex",
+        providerState,
+        expectsIdentityNotification: true,
+        threadId,
+      });
+    }
+
+    registry.recordProviderThreadIdentity({
+      providerState,
+      threadId: "thread-1",
+      providerThreadId: "provider-thread-1",
+    });
+
+    expect(providerState.pendingIdentityThreadIds).toEqual(["thread-2"]);
+    expect(registry.resolvePendingProviderThreadIdentity(providerState)).toBe(
+      "thread-2",
+    );
+    expect(
+      registry.resolvePendingProviderThreadIdentity(providerState),
+    ).toBeUndefined();
+  });
+
   it("stamps projected events with the resolved bb thread id", () => {
     const event: ThreadEvent = {
       type: "turn/started",

packages/agent-runtime/src/runtime-thread-identity.ts

@@ -91,6 +91,10 @@ export class RuntimeThreadIdentityRegistry {
 
   recordProviderThreadIdentity(args: RecordProviderThreadIdentityArgs): void {
     this.threadToProviderThread.set(args.threadId, args.providerThreadId);
+    args.providerState.pendingIdentityThreadIds =
+      args.providerState.pendingIdentityThreadIds.filter(
+        (pendingThreadId) => pendingThreadId !== args.threadId,
+      );
   }
 
   resolveBbThreadIdForProviderThread(

packages/agent-runtime/src/runtime.codex-topology.test.ts

@@ -373,6 +373,93 @@ describe("codex process topology", () => {
     expect(runtime.hasThread("t1")).toBe(false);
     expect(runtime.hasThread("t2")).toBe(false);
   }, 30_000);
+  it("refuses to resume a provider thread that another hosted thread already owns", async () => {
+    const topology = createCodexTopologyRuntime();
+    const { runtime } = topology;
+
+    const providerThreadId1 = await startCodexThread(runtime, "t1");
+
+    await expect(
+      runtime.resumeThread({
+        environmentId: "env-1",
+        projectId: "p1",
+        providerId: "codex",
+        providerThreadId: providerThreadId1,
+        threadId: "t2",
+        options: fullRuntimeOptions,
+      }),
+    ).rejects.toThrow(
+      `provider thread "${providerThreadId1}" is already hosted by thread "t1"`,
+    );
+    expect(runtime.hasThread("t2")).toBe(false);
+    expect(runtime.getProviderSession("t1")?.providerThreadId).toBe(
+      providerThreadId1,
+    );
+    expect(topology.spawned()).toBe(1);
+  }, 30_000);
+
+  it("keeps a rewound thread's provider identity when another thread starts on the same bridge", async () => {
+    const topology = createCodexTopologyRuntime();
+    const { runtime, events } = topology;
+
+    const originalProviderThreadId1 = await startCodexThread(runtime, "t1");
+    const staged = await runtime.prepareThreadRewind({
+      environmentId: "env-1",
+      threadId: "t1",
+      leaseId: "lease-1",
+      projectId: "p1",
+      providerId: "codex",
+      sourceProviderThreadId: originalProviderThreadId1,
+      retainThroughProviderCheckpoint: "turn-1",
+      options: fullRuntimeOptions,
+      instructionMode: "append",
+    });
+    const { providerThreadId: providerThreadId1 } = await runtime.startThread({
+      environmentId: "env-1",
+      projectId: "p1",
+      providerId: "codex",
+      threadId: "t1",
+      options: fullRuntimeOptions,
+      fork: { sourceProviderThreadId: staged.providerThreadId },
+    });
+    await runtime.discardThreadRewind({ leaseId: "lease-1" });
+    expect(providerThreadId1).not.toBe(originalProviderThreadId1);
+    const eventsBeforeSecondStart = events.length;
+
+    const providerThreadId2 = await startCodexThread(runtime, "t2");
+    expect(providerThreadId2).not.toBe(providerThreadId1);
+    expect(runtime.getProviderSession("t1")?.providerThreadId).toBe(
+      providerThreadId1,
+    );
+    expect(runtime.getProviderSession("t2")?.providerThreadId).toBe(
+      providerThreadId2,
+    );
+
+    await runtime.runTurn({
+      clientRequestId: "creq_cdxidnty22",
+      threadId: "t1",
+      input: [promptTextInput({ text: "hello" })],
+      options: fullRuntimeOptions,
+    });
+    await waitForThreadAgentMessageText({
+      events,
+      providerId: "codex",
+      runtime,
+      text: "hello from codex turn",
+      threadId: "t1",
+    });
+    const stampedProviderThreadIds = new Set(
+      events
+        .slice(eventsBeforeSecondStart)
+        .filter((event) => event.threadId === "t1")
+        .flatMap((event) =>
+          "providerThreadId" in event && event.providerThreadId
+            ? [event.providerThreadId]
+            : [],
+        ),
+    );
+    expect([...stampedProviderThreadIds]).toEqual([providerThreadId1]);
+  }, 30_000);
 });
 
 function isAlive(pid: number): boolean {

packages/agent-runtime/src/runtime.ts

@@ -778,6 +778,25 @@ export function createAgentRuntime(options: AgentRuntimeOptions): AgentRuntime {
     });
   }
 
+  function resolveProviderIdentityTargetThreadId(args: {
+    eventThreadId: string;
+    proc: ProviderProcess;
+    sourceThreadId: string | undefined;
+  }): string | undefined {
+    if (
+      args.sourceThreadId !== undefined &&
+      args.proc.identity.threadIds.has(args.sourceThreadId)
+    ) {
+      return args.sourceThreadId;
+    }
+    if (args.proc.identity.threadIds.has(args.eventThreadId)) {
+      return args.eventThreadId;
+    }
+    return threadIdentityRegistry.resolvePendingProviderThreadIdentity(
+      args.proc.identity,
+    );
+  }
+
   function forgetThreadRuntimeStateForProviderState(
     providerState: RuntimeProviderProcess["identity"],
     threadId: string,
@@ -1184,27 +1203,33 @@ export function createAgentRuntime(options: AgentRuntimeOptions): AgentRuntime {
       if (event.type !== "thread/identity" || !event.providerThreadId) {
         continue;
       }
-
-      if (args.proc.identity.threadIds.has(event.threadId)) {
-        recordProviderThreadIdentity(
-          args.proc,
-          event.threadId,
-          event.providerThreadId,
+      const targetThreadId = resolveProviderIdentityTargetThreadId({
+        eventThreadId: event.threadId,
+        proc: args.proc,
+        sourceThreadId: args.sourceThreadId,
+      });
+      if (targetThreadId === undefined) {
+        options.onStderr?.(
+          `Dropping thread/identity for provider thread "${event.providerThreadId}" from "${args.proc.providerId}"; no bb thread could be resolved`,
         );
         continue;
       }
-
-      const bbThreadId =
-        threadIdentityRegistry.resolvePendingProviderThreadIdentity(
-          args.proc.identity,
-        );
-      if (bbThreadId) {
-        recordProviderThreadIdentity(
-          args.proc,
-          bbThreadId,
-          event.providerThreadId,
+      const ownerThreadId =
+        threadIdentityRegistry.resolveBbThreadIdForProviderThread({
+          providerState: args.proc.identity,
+          providerThreadId: event.providerThreadId,
+        });
+      if (ownerThreadId !== undefined && ownerThreadId !== targetThreadId) {
+        options.onStderr?.(
+          `Ignoring thread/identity that would bind provider thread "${event.providerThreadId}" to bb thread "${targetThreadId}"; bb thread "${ownerThreadId}" already owns it on "${args.proc.providerId}"`,
         );
+        continue;
       }
+      recordProviderThreadIdentity(
+        args.proc,
+        targetThreadId,
+        event.providerThreadId,
+      );
     }
 
     for (const event of args.events) {
@@ -1822,6 +1847,18 @@ export function createAgentRuntime(options: AgentRuntimeOptions): AgentRuntime {
             options: execOpts,
             providerId,
           });
+          if (providerThreadId !== undefined) {
+            const ownerThreadId =
+              threadIdentityRegistry.resolveBbThreadIdForProviderThread({
+                providerState: proc.identity,
+                providerThreadId,
+              });
+            if (ownerThreadId !== undefined && ownerThreadId !== threadId) {
+              throw new Error(
+                `Cannot resume thread "${threadId}" on "${providerId}": provider thread "${providerThreadId}" is already hosted by thread "${ownerThreadId}"`,
+              );
+            }
+          }
           const resolvedEnvironment = resolveRuntimeThreadEnvironment({
             contributedEnv,
             environmentId,

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