fix(server): one Codex task running out of memory no longer kills every Codex thread

apps/server/scripts/record-codex-app-server-replay-fixture.ts

@@ -1017,7 +1017,14 @@ const makeRecorder = Effect.fn("record-codex-app-server-replay-fixture.makeRecor
   return { path: outPath, setVersion, writeRecord, flush };
 });
 
-function layerCodex({ recorder }: { readonly recorder: Recorder }) {
+/** `connection` tags records from a second, concurrent app-server; replay gives it to the next session. */
+function layerCodex({
+  recorder,
+  connection = 0,
+}: {
+  readonly recorder: Recorder;
+  readonly connection?: number;
+}) {
   const clientRequestMethodById = new Map<string, string>();
   const serverRequestMethodById = new Map<string, string>();
   const clientOptions: CodexClient.CodexAppServerClientOptions = {
@@ -1063,11 +1070,13 @@ function layerCodex({ recorder }: { readonly recorder: Recorder }) {
               type: "expect_outbound",
               ...(label ? { label } : {}),
               frame: event.payload,
+              ...(connection === 0 ? {} : { connection }),
             }
           : {
               type: "emit_inbound",
               ...(label ? { label } : {}),
               frame: event.payload,
+              ...(connection === 0 ? {} : { connection }),
             };
 
       return Effect.gen(function* () {
@@ -1407,17 +1416,31 @@ const runReplaySession = Effect.fn("record-codex-app-server-replay-fixture.runRe
       return;
     }
 
+    // T3 runs one app-server per thread, so each fork records its own connection.
+    let connectionCount = 0;
+    const openConnection = Effect.gen(function* () {
+      const context = yield* Layer.build(layerCodex({ recorder, connection: connectionCount++ }));
+      return yield* initializeClient.pipe(Effect.provide(context));
+    });
+
     yield* Effect.gen(function* () {
-      const client = yield* initializeClient;
-      const thread = yield* client.request("thread/start", {
+      const sourceClient = yield* openConnection;
+      const thread = yield* sourceClient.request("thread/start", {
         ...threadRuntimeParams,
         config: { ...threadRuntimeParams.config, ...run.threadConfig },
       });
       let activeThreadId = thread.thread.id;
       const threadIds = new Map<string, string>([["source", thread.thread.id]]);
+      const clientByThread = new Map([[thread.thread.id, sourceClient]]);
+      const clientFor = (threadId: string) => {
+        const client = clientByThread.get(threadId);
+        if (client === undefined) throw new Error(`No app-server for replay thread ${threadId}.`);
+        return client;
+      };
 
       for (const [stepIndex, step] of run.steps.entries()) {
         if (step.type === "rollback") {
+          const client = clientFor(activeThreadId);
           yield* client.request("thread/read", { threadId: activeThreadId, includeTurns: false });
           yield* revertCodexThread(client, activeThreadId, step.numTurns);
           continue;
@@ -1429,11 +1452,13 @@ const runReplaySession = Effect.fn("record-codex-app-server-replay-fixture.runRe
             throw new Error(`Unknown replay thread alias '${step.from}'.`);
           }
           const lastTurnId = lastTurnIdByThread.get(sourceThreadId);
-          const forked = yield* client.request("thread/fork", {
+          const forkClient = yield* openConnection;
+          const forked = yield* forkClient.request("thread/fork", {
             threadId: sourceThreadId,
             ...(lastTurnId === undefined ? {} : { lastTurnId }),
             ...threadRuntimeParams,
           });
+          clientByThread.set(forked.thread.id, forkClient);
           activeThreadId = forked.thread.id;
           if (step.as !== undefined) {
             threadIds.set(step.as, forked.thread.id);
@@ -1447,16 +1472,10 @@ const runReplaySession = Effect.fn("record-codex-app-server-replay-fixture.runRe
         if (turnThreadId === undefined) {
           throw new Error(`Unknown replay thread alias '${threadAlias}'.`);
         }
-        yield* runTurnStep(client, turnThreadId, step);
+        yield* runTurnStep(clientFor(turnThreadId), turnThreadId, step);
         void stepIndex;
       }
-    }).pipe(
-      Effect.provide(
-        layerCodex({
-          recorder,
-        }),
-      ),
-    );
+    }).pipe(Effect.scoped);
   },
 );
 

apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts

@@ -862,6 +862,93 @@ describe("CodexAdapterV2 process exit", () => {
   );
 });
 
+describe("CodexAdapterV2 replay testkit", () => {
+  it.effect("gives each session its own recorded connection and continues it on reopen", () =>
+    Effect.scoped(
+      Effect.gen(function* () {
+        const initialize = (userAgent: string, connection?: number) =>
+          [
+            {
+              type: "expect_outbound",
+              label: "initialize",
+              frame: { id: 1, method: "initialize", params: {} },
+              ...(connection === undefined ? {} : { connection }),
+            },
+            {
+              type: "emit_inbound",
+              label: "initialize",
+              frame: { id: 1, result: { userAgent } },
+              ...(connection === undefined ? {} : { connection }),
+            },
+          ] satisfies Array<CodexReplay.CodexAppServerReplayEntry>;
+        const transcript = makeCodexReplayTranscript({
+          scenario: "testkit-session-connections",
+          entries: [
+            ...initialize("session A, first process"),
+            { type: "runtime_exit", status: "success" },
+            ...initialize("session A, after restart"),
+            ...initialize("session B", 1),
+          ],
+        });
+        const factory = yield* CodexAdapterV2Testkit.makeReplayClientFactory({ transcript });
+        const userAgent = (providerSessionId: string) =>
+          factory
+            .open({ providerSessionId: ProviderSessionId.make(providerSessionId) } as never)
+            .pipe(
+              Effect.flatMap((client) => client.raw.request("initialize", {})),
+              Effect.map((result) => (result as { readonly userAgent: string }).userAgent),
+            );
+
+        assert.equal(yield* userAgent("session-a"), "session A, first process");
+        assert.equal(yield* userAgent("session-b"), "session B");
+        assert.equal(yield* userAgent("session-a"), "session A, after restart");
+      }),
+    ),
+  );
+});
+
+describe("CodexAdapterV2 thread writer", () => {
+  const writerError = (message: string) =>
+    CodexError.CodexAppServerRequestError.fromProtocolError(
+      { code: -32600, message },
+      "thread/resume",
+      "3",
+    );
+
+  it.effect.each([
+    "thread saved-thread already has an active writer",
+    "thread history projection for saved-thread is behind durable rollout",
+  ])("retries while another app-server is busy with the thread: %s", (busyMessage) =>
+    Effect.gen(function* () {
+      let attempts = 0;
+      const request = Effect.suspend(() => {
+        attempts += 1;
+        return attempts === 1 ? Effect.fail(writerError(busyMessage)) : Effect.succeed("resumed");
+      });
+      const fiber = yield* CodexAdapterV2.retryWhileThreadBusyElsewhere(request).pipe(
+        Effect.forkChild({ startImmediately: true }),
+      );
+      assert.equal(attempts, 1);
+      yield* TestClock.adjust("500 millis");
+      assert.equal(yield* Fiber.join(fiber), "resumed");
+      assert.equal(attempts, 2);
+    }),
+  );
+
+  it.effect("does not retry other request errors", () =>
+    Effect.gen(function* () {
+      let attempts = 0;
+      const request = Effect.suspend(() => {
+        attempts += 1;
+        return Effect.fail(writerError("thread not found: saved-thread"));
+      });
+      const error = yield* Effect.flip(CodexAdapterV2.retryWhileThreadBusyElsewhere(request));
+      assert.nestedPropertyVal(error, "errorMessage", "thread not found: saved-thread");
+      assert.equal(attempts, 1);
+    }),
+  );
+});
+
 describe("CodexAdapterV2 session initialize", () => {
   const openReplaySession = (
     transcript: CodexReplay.CodexAppServerReplayTranscript,

apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.testkit.ts

@@ -4,7 +4,9 @@ import * as CodexClient from "effect-codex-app-server/client";
 import type * as CodexError from "effect-codex-app-server/errors";
 import * as CodexReplay from "effect-codex-app-server/replay";
 import * as Effect from "effect/Effect";
+import * as FileSystem from "effect/FileSystem";
 import * as Layer from "effect/Layer";
+import * as Ref from "effect/Ref";
 import * as Predicate from "effect/Predicate";
 import * as Schema from "effect/Schema";
 
@@ -94,45 +96,101 @@ function metadataFromTranscript(transcript: ProviderReplayTranscript): {
   };
 }
 
+/** Serves recorded app-server connections to the sessions a replay test opens. */
+export const makeReplayClientFactory = Effect.fn("CodexAdapterV2Testkit.makeReplayClientFactory")(
+  function* (input: {
+    readonly transcript: CodexReplay.CodexAppServerReplayTranscript;
+    readonly driver?: CodexReplay.CodexAppServerReplayDriver;
+  }) {
+    const connections = CodexReplay.splitReplayConnections(input.transcript);
+    // Each thread runs its own app-server, so each recorded connection belongs to one
+    // session: a session's first open takes the next connection, and a reopen continues
+    // its own (a recorded restart shows up there as a runtime_exit). A single-connection
+    // recording serves every session, as before. Only the caller's emit hook carries over.
+    const drivers =
+      connections.length <= 1
+        ? undefined
+        : yield* Effect.forEach(connections, (connection) =>
+            CodexReplay.makeReplayDriver(
+              connection,
+              input.driver?.beforeEmitInbound === undefined
+                ? {}
+                : { beforeEmitInbound: input.driver.beforeEmitInbound },
+            ),
+          );
+    const driverBySession = yield* Ref.make(
+      new Map<string, CodexReplay.CodexAppServerReplayDriver>(),
+    );
+    return CodexAdapterV2.CodexAppServerClientFactory.of({
+      open: (openInput) =>
+        Effect.gen(function* () {
+          const driver =
+            drivers === undefined
+              ? input.driver
+              : yield* Ref.modify(driverBySession, (assigned) => {
+                  const existing = assigned.get(openInput.providerSessionId);
+                  const next = existing ?? drivers[assigned.size];
+                  return next === undefined || existing !== undefined
+                    ? [next, assigned]
+                    : [next, new Map(assigned).set(openInput.providerSessionId, next)];
+                });
+          if (drivers !== undefined && driver === undefined) {
+            return yield* new ProviderAdapter.ProviderAdapterOpenSessionError({
+              driver: CodexAdapterV2.CODEX_DRIVER_KIND,
+              providerSessionId: openInput.providerSessionId,
+              cause: `The replay recorded ${drivers.length} app-server connections, and every one already belongs to another session.`,
+            });
+          }
+          const layerReplay =
+            driver === undefined
+              ? CodexReplay.layerReplay(input.transcript)
+              : CodexReplay.layerReplayWithDriver(driver);
+          const context = yield* Layer.build(layerReplay).pipe(
+            Effect.mapError(
+              (cause) =>
+                new ProviderAdapter.ProviderAdapterOpenSessionError({
+                  driver: CodexAdapterV2.CODEX_DRIVER_KIND,
+                  providerSessionId: openInput.providerSessionId,
+                  cause,
+                }),
+            ),
+          );
+          return yield* Effect.service(CodexClient.CodexAppServerClient).pipe(
+            Effect.map((client) => withCodexReplayChildMetadata(client, input.transcript)),
+            Effect.provide(context),
+          );
+        }),
+    });
+  },
+);
+
 export function layer(input: {
   readonly transcript: CodexReplay.CodexAppServerReplayTranscript;
   readonly driver?: CodexReplay.CodexAppServerReplayDriver;
 }) {
-  const layerReplay =
-    input.driver === undefined
-      ? CodexReplay.layerReplay(input.transcript)
-      : CodexReplay.layerReplayWithDriver(input.driver);
-  const layerReplayClientFactory = Layer.succeed(CodexAdapterV2.CodexAppServerClientFactory, {
-    open: (openInput) =>
-      Effect.gen(function* () {
-        const context = yield* Layer.build(layerReplay).pipe(
-          Effect.mapError(
-            (cause) =>
-              new ProviderAdapter.ProviderAdapterOpenSessionError({
-                driver: CodexAdapterV2.CODEX_DRIVER_KIND,
-                providerSessionId: openInput.providerSessionId,
-                cause,
-              }),
-          ),
-        );
-        return yield* Effect.service(CodexClient.CodexAppServerClient).pipe(
-          Effect.map((client) => withCodexReplayChildMetadata(client, input.transcript)),
-          Effect.provide(context),
-        );
-      }),
-  });
+  const layerReplayClientFactory = Layer.effect(
+    CodexAdapterV2.CodexAppServerClientFactory,
+    makeReplayClientFactory(input),
+  );
   const layerServerConfig = Layer.effect(
     ServerConfig.ServerConfig,
     makeReplayServerConfig(input.transcript.scenario).pipe(Effect.orDie),
   ).pipe(Layer.provide(NodeServices.layer));
   const layerRegistry = ProviderAdapterRegistry.layerFromAdaptersEffect(
-    CodexAdapterV2.createCodexAdapterV2({
-      instanceId: CodexAdapterV2.CODEX_DEFAULT_INSTANCE_ID,
-      displayName: undefined,
-      environment: [],
-      enabled: true,
-      config: DEFAULT_CODEX_SETTINGS,
-    }).pipe(Effect.map((adapter) => [adapter])),
+    Effect.gen(function* () {
+      // Keeps replay sessions away from the developer's real ~/.codex/auth.json.
+      const codexHome = yield* (yield* FileSystem.FileSystem)
+        .makeTempDirectoryScoped()
+        .pipe(Effect.orDie);
+      const adapter = yield* CodexAdapterV2.createCodexAdapterV2({
+        instanceId: CodexAdapterV2.CODEX_DEFAULT_INSTANCE_ID,
+        displayName: undefined,
+        environment: [{ name: "CODEX_HOME", value: codexHome, sensitive: false }],
+        enabled: true,
+        config: DEFAULT_CODEX_SETTINGS,
+      });
+      return [adapter];
+    }),
   ).pipe(
     Layer.provide(
       Layer.mergeAll(