fix(cli): restore codex review resume flow

saltbo · saltbo · commit 4760ad301039 · 2026-04-14T17:39:58.000-04:00
Persist the Codex thread id as a provider resume token so reject/resume reconnects to the same Codex conversation instead of passing the daemon session id to the SDK.

Also update the daemon smoke test to wait for async review-session cleanup and guard empty task cleanup arrays under set -u.
diff --git a/packages/cli/src/daemon/resumer.ts b/packages/cli/src/daemon/resumer.ts
@@ -7,7 +7,6 @@
  */
 
 import { existsSync } from "node:fs";
-import type { AgentInfo } from "../agent/systemPrompt.js";
 import { AgentClient, type ApiClient } from "../client/index.js";
 import { getCredentials } from "../config.js";
 import { createLogger } from "../logger.js";
@@ -75,6 +74,7 @@ export async function resumeSession(session: SessionFile, message: string, clien
     provider,
     taskId,
     sessionId: session.sessionId,
+    resumeToken: session.providerResumeToken,
     cwd: workspace.cwd,
     taskContext: message,
     agentClient,
diff --git a/packages/cli/src/daemon/runtimePool.ts b/packages/cli/src/daemon/runtimePool.ts
@@ -44,6 +44,7 @@ export interface SpawnRequest {
   provider: AgentProvider;
   taskId: string;
   sessionId: string;
+  resumeToken?: string;
   cwd: string;
   taskContext: string;
   agentClient: AgentClient;
@@ -118,6 +119,7 @@ export class RuntimePool {
     const handle: AgentHandle = await providerExecute(provider.name, () =>
       provider.execute({
         sessionId,
+        resumeToken: req.resumeToken,
         cwd: req.cwd,
         env: { ...(process.env as Record<string, string>), ...req.agentEnv },
         taskContext: req.taskContext,
@@ -431,6 +433,7 @@ async function finalize(agent: AgentProcess, opts: { crashed: boolean; error?: u
   //   - in_progress + no reject → release (agent forgot to submit review)
   //   - done/cancelled → cleanup
   const taskInReview = agent.resultReceived && !opts.crashed;
+  const providerResumeToken = agent.handle.getResumeToken?.();
 
   const event: SessionEvent = classifyIteratorEnd({
     resultReceived: agent.resultReceived,
@@ -440,7 +443,7 @@ async function finalize(agent: AgentProcess, opts: { crashed: boolean; error?: u
     transient,
   });
 
-  const next = await sessions.applyEvent(sessionId, event).catch((e) => {
+  const next = await sessions.applyEvent(sessionId, event, providerResumeToken ? { providerResumeToken } : undefined).catch((e) => {
     logger.error(`State transition failed for ${sessionId}: ${errMessage(e)}`);
     return null;
   });
diff --git a/packages/cli/src/providers/codex.ts b/packages/cli/src/providers/codex.ts
@@ -1,3 +1,4 @@
+import { execSync } from "node:child_process";
 import { readFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { join } from "node:path";
@@ -44,6 +45,27 @@ function calcCost(model: string, inputTokens: number, cachedInputTokens: number,
   return (inputTokens * price.input + cachedInputTokens * price.cached_input + outputTokens * price.output) / 1_000_000;
 }
 
+function resolveCodexPath(): string | undefined {
+  try {
+    const path = execSync("which codex", { encoding: "utf-8", stdio: ["ignore", "pipe", "ignore"] }).trim();
+    return path || undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function resolveCodexModel(opts: ExecuteOpts): string | undefined {
+  if (!opts.model) {
+    return readAccessToken() ? undefined : "o3";
+  }
+  if (readAccessToken() && !opts.env.OPENAI_API_KEY) {
+    // ChatGPT-backed Codex accounts reject explicit model overrides like "o3".
+    // Let the CLI choose the account-compatible default model.
+    return undefined;
+  }
+  return opts.model;
+}
+
 /** Map a single Codex thread event to an AgentEvent (or null to skip). */
 export function mapThreadEvent(event: ThreadEvent, model = "o3"): AgentEvent | null {
   switch (event.type) {
@@ -174,23 +196,25 @@ export const codexProvider: AgentProvider = {
   label: "Codex CLI",
 
   async execute(opts: ExecuteOpts): Promise<AgentHandle> {
-    const model = opts.model ?? "o3";
+    const model = resolveCodexModel(opts) ?? "o3";
+    let resumeToken: string | undefined = opts.resumeToken;
 
-    const codex = new Codex({ env: opts.env });
+    const codex = new Codex({ env: opts.env, codexPathOverride: resolveCodexPath() });
     const threadOpts = {
-      model: opts.model,
+      model: resolveCodexModel(opts),
       workingDirectory: opts.cwd,
       sandboxMode: "danger-full-access" as const,
       approvalPolicy: "never" as const,
     };
-    const thread = opts.resume ? codex.resumeThread(opts.sessionId, threadOpts) : codex.startThread(threadOpts);
+    const thread = opts.resume ? codex.resumeThread(opts.resumeToken ?? opts.sessionId, threadOpts) : codex.startThread(threadOpts);
 
     const abortController = new AbortController();
     const streamed = await thread.runStreamed(opts.taskContext, { signal: abortController.signal });
 
     const events = (async function* () {
       const turnOpen = { value: false };
       for await (const event of streamed.events) {
+        if (event.type === "thread.started") resumeToken = event.thread_id;
         yield* mapThreadEventStream(event, model, turnOpen);
       }
     })();
@@ -206,6 +230,9 @@ export const codexProvider: AgentProvider = {
       async send() {
         throw new Error("Codex multi-turn send not implemented");
       },
+      getResumeToken() {
+        return resumeToken;
+      },
     };
   },
 
diff --git a/packages/cli/src/providers/types.ts b/packages/cli/src/providers/types.ts
@@ -4,6 +4,7 @@ export type { AgentEvent, AgentRuntime, ContentBlock, UsageInfo, UsageWindow };
 
 export interface ExecuteOpts {
   sessionId: string;
+  resumeToken?: string;
   cwd: string;
   env: Record<string, string>;
   taskContext: string;
@@ -28,6 +29,7 @@ export interface AgentHandle {
   events: AsyncIterable<AgentEvent>;
   abort(): Promise<void>;
   send(message: string): Promise<void>;
+  getResumeToken?(): string | undefined;
 }
 
 export interface AgentProvider {
diff --git a/packages/cli/src/session/types.ts b/packages/cli/src/session/types.ts
@@ -34,6 +34,7 @@ export interface SessionFile {
   workspace?: WorkspaceInfo;
   status?: SessionStatus;
   model?: string;
+  providerResumeToken?: string;
   gpgSubkeyId?: string | null;
   agentUsername?: string;
   agentName?: string;
diff --git a/packages/cli/tests/providers.test.ts b/packages/cli/tests/providers.test.ts
@@ -1,5 +1,5 @@
 // @vitest-environment node
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 
 // Mock node:fs so readFileSync never touches disk (gemini reads system prompt file, codex reads auth)
 vi.mock("node:fs", () => ({
@@ -8,6 +8,12 @@ vi.mock("node:fs", () => ({
   }),
 }));
 
+vi.mock("node:child_process", () => ({
+  execSync: vi.fn().mockImplementation(() => {
+    throw new Error("missing");
+  }),
+}));
+
 // Mock logger to suppress output
 vi.mock("../src/logger.js", () => ({
   createLogger: () => ({
@@ -62,6 +68,18 @@ import {
 import { getAvailableProviders, getProvider, registerProvider } from "../src/providers/registry.js";
 import type { AgentProvider } from "../src/providers/types.js";
 
+beforeEach(async () => {
+  const fsModule = await import("node:fs");
+  vi.mocked(fsModule.readFileSync).mockImplementation(() => {
+    throw new Error("ENOENT");
+  });
+
+  const childProcess = await import("node:child_process");
+  vi.mocked(childProcess.execSync).mockImplementation(() => {
+    throw new Error("missing");
+  });
+});
+
 // ---------------------------------------------------------------------------
 // mapSDKMessage — rate_limit_event
 // ---------------------------------------------------------------------------
@@ -591,6 +609,8 @@ describe("codexProvider.execute — handle shape", () => {
 describe("codexProvider.execute — thread selection", () => {
   it("calls startThread when resume is false or absent", async () => {
     const { Codex } = await import("@openai/codex-sdk");
+    const childProcess = await import("node:child_process");
+    vi.mocked(childProcess.execSync).mockReturnValueOnce("/opt/homebrew/bin/codex\n" as any);
     const startThreadSpy = vi.fn().mockReturnValue({
       runStreamed: vi.fn().mockResolvedValue({ events: (async function* () {})() }),
     });
@@ -603,10 +623,15 @@ describe("codexProvider.execute — thread selection", () => {
     );
     await codexProvider.execute({ sessionId: "s1", cwd: "/tmp", env: {}, taskContext: "ctx" });
     expect(startThreadSpy).toHaveBeenCalledOnce();
+    expect(vi.mocked(Codex)).toHaveBeenCalledWith(expect.objectContaining({ codexPathOverride: "/opt/homebrew/bin/codex" }));
   });
 
   it("calls resumeThread when resume is true", async () => {
     const { Codex } = await import("@openai/codex-sdk");
+    const childProcess = await import("node:child_process");
+    vi.mocked(childProcess.execSync).mockImplementationOnce(() => {
+      throw new Error("missing");
+    });
     const resumeThreadSpy = vi.fn().mockReturnValue({
       runStreamed: vi.fn().mockResolvedValue({ events: (async function* () {})() }),
     });
@@ -617,8 +642,63 @@ describe("codexProvider.execute — thread selection", () => {
           resumeThread: resumeThreadSpy,
         }) as any,
     );
-    await codexProvider.execute({ sessionId: "sess-77", cwd: "/tmp", env: {}, taskContext: "ctx", resume: true });
-    expect(resumeThreadSpy).toHaveBeenCalledWith("sess-77", expect.any(Object));
+    await codexProvider.execute({
+      sessionId: "sess-77",
+      resumeToken: "codex-thread-1",
+      cwd: "/tmp",
+      env: {},
+      taskContext: "ctx",
+      resume: true,
+    });
+    expect(resumeThreadSpy).toHaveBeenCalledWith("codex-thread-1", expect.any(Object));
+  });
+
+  it("omits explicit model for ChatGPT-backed Codex sessions", async () => {
+    const { Codex } = await import("@openai/codex-sdk");
+    const fsModule = await import("node:fs");
+    const childProcess = await import("node:child_process");
+    vi.mocked(childProcess.execSync).mockReturnValueOnce("/opt/homebrew/bin/codex\n" as any);
+    vi.mocked(fsModule.readFileSync).mockReturnValue(JSON.stringify({ tokens: { access_token: "chatgpt-token" } }) as any);
+    const startThreadSpy = vi.fn().mockReturnValue({
+      runStreamed: vi.fn().mockResolvedValue({ events: (async function* () {})() }),
+    });
+    vi.mocked(Codex).mockImplementationOnce(
+      () =>
+        ({
+          startThread: startThreadSpy,
+          resumeThread: vi.fn(),
+        }) as any,
+    );
+
+    await codexProvider.execute({ sessionId: "s1", cwd: "/tmp", env: {}, taskContext: "ctx", model: "o3" });
+
+    expect(startThreadSpy).toHaveBeenCalledWith(expect.objectContaining({ model: undefined }));
+  });
+
+  it("keeps explicit model when using API-key-backed Codex sessions", async () => {
+    const { Codex } = await import("@openai/codex-sdk");
+    const childProcess = await import("node:child_process");
+    vi.mocked(childProcess.execSync).mockReturnValueOnce("/opt/homebrew/bin/codex\n" as any);
+    const startThreadSpy = vi.fn().mockReturnValue({
+      runStreamed: vi.fn().mockResolvedValue({ events: (async function* () {})() }),
+    });
+    vi.mocked(Codex).mockImplementationOnce(
+      () =>
+        ({
+          startThread: startThreadSpy,
+          resumeThread: vi.fn(),
+        }) as any,
+    );
+
+    await codexProvider.execute({
+      sessionId: "s1",
+      cwd: "/tmp",
+      env: { OPENAI_API_KEY: "sk-test" },
+      taskContext: "ctx",
+      model: "o3",
+    });
+
+    expect(startThreadSpy).toHaveBeenCalledWith(expect.objectContaining({ model: "o3" }));
   });
 
   it("events yields mapped AgentEvents from SDK thread events", async () => {
@@ -646,6 +726,29 @@ describe("codexProvider.execute — thread selection", () => {
     expect(events[1]).toEqual({ type: "block.done", block: { type: "text", text: "codex message" } });
   });
 
+  it("captures the Codex thread id as a resume token from thread.started", async () => {
+    const { Codex } = await import("@openai/codex-sdk");
+    vi.mocked(Codex).mockImplementationOnce(
+      () =>
+        ({
+          startThread: vi.fn().mockReturnValue({
+            runStreamed: vi.fn().mockResolvedValue({
+              events: (async function* () {
+                yield { type: "thread.started", thread_id: "thread-123" } as any;
+                yield { type: "turn.completed", usage: {} } as any;
+              })(),
+            }),
+          }),
+          resumeThread: vi.fn(),
+        }) as any,
+    );
+    const handle = await codexProvider.execute({ sessionId: "s1", cwd: "/tmp", env: {}, taskContext: "ctx" });
+    for await (const _ev of handle.events) {
+      // Drain the stream so thread.started is observed.
+    }
+    expect(handle.getResumeToken?.()).toBe("thread-123");
+  });
+
   it("send() throws not-implemented error", async () => {
     const { Codex } = await import("@openai/codex-sdk");
     const runStreamedSpy = vi.fn().mockResolvedValue({ events: (async function* () {})() });
diff --git a/scripts/daemon-smoke-test.sh b/scripts/daemon-smoke-test.sh
@@ -57,6 +57,19 @@ task_session_exists() {
     | xargs grep -l "\"taskId\": *\"$task_id\"" 2>/dev/null | head -1
 }
 
+wait_session_cleanup() {
+  local task_id="$1" timeout_secs="${2:-120}"
+  local elapsed=0
+  while [ "$elapsed" -lt "$timeout_secs" ]; do
+    if [ -z "$(task_session_exists "$task_id")" ]; then
+      return 0
+    fi
+    sleep 2
+    elapsed=$((elapsed + 2))
+  done
+  return 1
+}
+
 pass() { echo "  PASS: $1"; PASS=$((PASS + 1)); }
 fail() { echo "  FAIL: $1"; FAIL=$((FAIL + 1)); }
 
@@ -122,20 +135,17 @@ echo ""
 echo "[Test 3/4] Complete — mark task done, verify cleanup"
 ak task complete "$T1" >/dev/null 2>&1
 
-# Give daemon a poll cycle to clean up
-sleep 15
-
 STATUS_AFTER_COMPLETE=$(task_status "$T1")
 if [ "$STATUS_AFTER_COMPLETE" = "done" ]; then
   pass "task is done"
 else
   fail "expected done, got: $STATUS_AFTER_COMPLETE"
 fi
 
-if [ -z "$(task_session_exists "$T1")" ]; then
+if wait_session_cleanup "$T1" 120; then
   pass "session cleaned up after completion"
 else
-  fail "session still exists after completion"
+  fail "session still exists after completion timeout"
 fi
 echo ""
 
@@ -156,9 +166,6 @@ fi
 sleep 3
 ak task cancel "$T4" >/dev/null 2>&1
 
-# Give daemon a poll cycle to detect and kill
-sleep 12
-
 STATUS_AFTER_CANCEL=$(task_status "$T4")
 if [ "$STATUS_AFTER_CANCEL" = "cancelled" ]; then
   pass "task is cancelled"
@@ -184,9 +191,11 @@ echo "  Failed: $FAIL"
 echo "==============================="
 
 # Cleanup test tasks
-for tid in "${TASKS[@]}"; do
-  ak task cancel "$tid" >/dev/null 2>&1 || true
-done
+if [ "${#TASKS[@]}" -gt 0 ]; then
+  for tid in "${TASKS[@]}"; do
+    ak task cancel "$tid" >/dev/null 2>&1 || true
+  done
+fi
 
 if [ "$FAIL" -gt 0 ]; then
   exit 1
