fix(codex): rewind whole turns so a steer never drops earlier prompts

The turn-boundary rule now lives in the Codex rewind helper. It refuses a
steer message on both history paths and counts turns rather than user
messages for thread/rollback, so a steer in a later turn no longer rolls
back one turn too many on Codex versions that still roll back.

Co-authored-by: 3ae3ae <ogyshjsm@naver.com>

packages/server/src/server/agent/providers/codex-app-server-agent.local.e2e.test.ts

@@ -217,6 +217,104 @@ function waitForEvent<TEvent extends AgentStreamEvent>(params: {
 }
 
 describe("Codex app-server provider (local e2e)", () => {
+  test.runIf(isCodexInstalled())(
+    "rewinding a steer message keeps the earlier prompt of its turn in model context",
+    async () => {
+      const cwd = mkdtempSync(path.join(os.tmpdir(), "codex-steer-rewind-cwd-"));
+      const codexHome = mkdtempSync(path.join(os.tmpdir(), "codex-steer-rewind-home-"));
+      const requestBodies: string[] = [];
+      let releaseFirstResponse: (() => void) | null = null;
+      const server = createServer((req, res) => {
+        const chunks: Buffer[] = [];
+        req.on("data", (chunk) => chunks.push(Buffer.from(chunk)));
+        req.on("end", () => {
+          if (req.method !== "POST" || req.url !== "/v1/responses") {
+            res.statusCode = 404;
+            res.end("not found");
+            return;
+          }
+          requestBodies.push(Buffer.concat(chunks).toString("utf8"));
+          const reply = () => {
+            res.statusCode = 200;
+            res.setHeader("content-type", "text/event-stream");
+            res.end(assistantMessageSse("done"));
+          };
+          if (requestBodies.length === 1) {
+            releaseFirstResponse = reply;
+          } else {
+            reply();
+          }
+        });
+      });
+      server.listen(0, "127.0.0.1");
+      await once(server, "listening");
+      const address = server.address();
+      if (!address || typeof address === "string") {
+        throw new Error("Expected TCP address for mock responses server");
+      }
+
+      try {
+        writeMockCodexConfig(codexHome, `http://127.0.0.1:${address.port}`);
+        const client = new CodexAppServerAgentClient(createTestLogger());
+        const session = await client.createSession(
+          { provider: "codex", cwd, modeId: "auto", model: "mock-model" },
+          { env: { CODEX_HOME: codexHome } },
+        );
+
+        try {
+          const userMessageIds = new Map<string, string>();
+          session.subscribe((event) => {
+            if (
+              event.type === "timeline" &&
+              event.item.type === "user_message" &&
+              event.item.messageId
+            ) {
+              userMessageIds.set(event.item.text, event.item.messageId);
+            }
+          });
+          const turnFinished = waitForEvent({
+            session,
+            timeoutMs: 15_000,
+            label: "turn completion",
+            predicate: (event): event is Extract<AgentStreamEvent, { type: "turn_completed" }> =>
+              event.type === "turn_completed",
+          });
+
+          const started = await session.startTurn("ORIGINAL_PROMPT");
+          await vi.waitFor(() => expect(releaseFirstResponse).not.toBeNull(), {
+            timeout: 15_000,
+          });
+          await expect(
+            session.steerActiveTurn!("STEER_PROMPT", { expectedTurnId: started.turnId }),
+          ).resolves.toEqual({ status: "accepted" });
+          releaseFirstResponse!();
+          await turnFinished;
+
+          await vi.waitFor(() => expect(userMessageIds.get("STEER_PROMPT")).toBeDefined(), {
+            timeout: 15_000,
+          });
+          // Codex can only fork at a turn boundary, so a rewind that cannot keep
+          // ORIGINAL_PROMPT has to refuse rather than silently drop it.
+          await session.revertConversation!({
+            messageId: userMessageIds.get("STEER_PROMPT")!,
+          }).catch(() => undefined);
+
+          await session.run("AFTER_REWIND");
+          expect(requestBodies.at(-1)).toContain("AFTER_REWIND");
+          expect(requestBodies.at(-1)).toContain("ORIGINAL_PROMPT");
+        } finally {
+          await session.close();
+        }
+      } finally {
+        server.closeAllConnections();
+        server.close();
+        rmSync(cwd, { recursive: true, force: true });
+        rmSync(codexHome, { recursive: true, force: true });
+      }
+    },
+    45_000,
+  );
+
   test.runIf(isCodexInstalled())(
     "reloads a persisted idle thread repeatedly without overlapping writers",
     async () => {

packages/server/src/server/agent/providers/codex-app-server-agent.test.ts

@@ -105,10 +105,7 @@ interface CollaborationModeRecord {
 }
 
 interface CodexSessionTestAccess {
-  codexUserMessageTurns(): {
-    resolve(messageId: string): { index: number; turnId: string | null } | null;
-    count(): number;
-  };
+  codexUserMessages(): { messageId: string; turnId: string | null }[];
   ensureThreadLoaded(): Promise<void>;
   handleToolApprovalRequest(params: unknown): Promise<unknown>;
   handleNotification(method: string, params: unknown): void;
@@ -4692,10 +4689,9 @@ describe("Codex app-server provider", () => {
 
     await asInternals(session).loadPersistedHistory(session.client);
 
-    expect(asInternals(session).codexUserMessageTurns().resolve("message-history")).toEqual({
-      index: 0,
-      turnId: "native-turn-1",
-    });
+    expect(asInternals(session).codexUserMessages()).toEqual([
+      { messageId: "message-history", turnId: "native-turn-1" },
+    ]);
   });
 
   test("loads mixed legacy and MultiAgentV2 sub-agent history", async () => {

packages/server/src/server/agent/providers/codex-app-server-agent.ts

@@ -91,7 +91,7 @@ import {
   type CodexThreadRollbackResponse,
   type CodexAppServerTraceContext,
 } from "./codex/app-server-transport.js";
-import { type CodexUserMessageTurnIndex, revertCodexConversation } from "./codex/rewind.js";
+import { type CodexUserMessage, revertCodexConversation } from "./codex/rewind.js";
 import {
   materializeProviderImage,
   renderProviderImageOutputAsAssistantMarkdown,
@@ -4529,16 +4529,11 @@ export class CodexAppServerAgentSession implements AgentSession {
     });
   }
 
-  private codexUserMessageTurns(): CodexUserMessageTurnIndex {
-    return {
-      resolve: (messageId) => {
-        const index = this.userMessageTurnIndexes.get(messageId);
-        return index === undefined
-          ? null
-          : { index, turnId: this.userMessageProviderTurnIds.get(messageId) ?? null };
-      },
-      count: () => this.userMessageTurnIds.length,
-    };
+  private codexUserMessages(): CodexUserMessage[] {
+    return this.userMessageTurnIds.map((messageId) => ({
+      messageId,
+      turnId: this.userMessageProviderTurnIds.get(messageId) ?? null,
+    }));
   }
 
   subscribe(callback: (event: AgentStreamEvent) => void): () => void {
@@ -4947,19 +4942,6 @@ export class CodexAppServerAgentSession implements AgentSession {
       await this.ensureThread();
     }
 
-    const targetTurnId = this.userMessageProviderTurnIds.get(input.messageId);
-    if (targetTurnId) {
-      const firstMessageId = this.userMessageTurnIds.find(
-        (messageId) => this.userMessageProviderTurnIds.get(messageId) === targetTurnId,
-      );
-      // Codex forks at turn boundaries; a steer cannot be removed independently.
-      if (firstMessageId !== input.messageId) {
-        throw new Error(
-          "Codex cannot rewind a message sent during a turn without removing earlier messages. Select the first message in the turn instead.",
-        );
-      }
-    }
-
     await revertCodexConversation({
       client: this.client,
       threadId: this.currentThreadId,
@@ -4968,7 +4950,7 @@ export class CodexAppServerAgentSession implements AgentSession {
       model: this.config.model ?? null,
       serviceTier: this.serviceTier,
       config: this.buildCodexInnerConfig(),
-      userMessageTurns: this.codexUserMessageTurns(),
+      userMessages: this.codexUserMessages(),
       threadRollbackAvailable: this.threadRollbackAvailable,
       setThreadId: async (threadId) => {
         this.currentThreadId = threadId;

packages/server/src/server/agent/providers/codex/rewind.test.ts

@@ -6,11 +6,7 @@ import type {
   CodexThreadRollbackParams,
   CodexThreadRollbackResponse,
 } from "./app-server-transport.js";
-import {
-  type CodexUserMessageTurnIndex,
-  type CodexRewindClient,
-  revertCodexConversation,
-} from "./rewind.js";
+import { type CodexRewindClient, revertCodexConversation } from "./rewind.js";
 
 class FakeCodex implements CodexRewindClient {
   readonly recordedForks: CodexThreadForkParams[] = [];
@@ -59,33 +55,13 @@ class FakeCodex implements CodexRewindClient {
   }
 }
 
-class CodexMessageTurns implements CodexUserMessageTurnIndex {
-  constructor(
-    private readonly indexesByMessageId: Map<string, number>,
-    private readonly turnIdsByMessageId: Map<string, string> = new Map(),
-  ) {}
-
-  resolve(messageId: string): { index: number; turnId: string | null } | null {
-    const index = this.indexesByMessageId.get(messageId);
-    return index === undefined
-      ? null
-      : { index, turnId: this.turnIdsByMessageId.get(messageId) ?? null };
-  }
-
-  count(): number {
-    return this.indexesByMessageId.size;
-  }
-}
-
 describe("Codex Rewind", () => {
   test("rewinds the conversation by forking the thread and rolling back past the native user message", async () => {
     const codex = new FakeCodex();
-    const userMessageTurns = new CodexMessageTurns(
-      new Map([
-        ["codex-first", 0],
-        ["codex-second", 1],
-      ]),
-    );
+    const userMessages = [
+      { messageId: "codex-first", turnId: null },
+      { messageId: "codex-second", turnId: null },
+    ];
     let reboundThreadId: string | null = null;
 
     await revertCodexConversation({
@@ -95,7 +71,7 @@ describe("Codex Rewind", () => {
       cwd: "/workspace/project",
       model: "gpt-5.4-mini",
       serviceTier: null,
-      userMessageTurns,
+      userMessages,
       threadRollbackAvailable: true,
       setThreadId: (threadId) => {
         reboundThreadId = threadId;
@@ -118,20 +94,18 @@ describe("Codex Rewind", () => {
 
   test("rewinds the conversation using native user message ids hydrated from app-server history", async () => {
     const codex = new FakeCodex();
-    const userMessageTurns = new CodexMessageTurns(
-      new Map([
-        ["codex-first", 0],
-        ["codex-second", 1],
-        ["codex-third", 2],
-      ]),
-    );
+    const userMessages = [
+      { messageId: "codex-first", turnId: null },
+      { messageId: "codex-second", turnId: null },
+      { messageId: "codex-third", turnId: null },
+    ];
     let reboundThreadId: string | null = null;
 
     await revertCodexConversation({
       client: codex,
       threadId: "source-thread",
       messageId: "codex-second",
-      userMessageTurns,
+      userMessages,
       threadRollbackAvailable: true,
       setThreadId: (threadId) => {
         reboundThreadId = threadId;
@@ -157,16 +131,10 @@ describe("Codex Rewind", () => {
     }
 
     const codex = new PaginatedCodex();
-    const userMessageTurns = new CodexMessageTurns(
-      new Map([
-        ["codex-first", 0],
-        ["codex-second", 1],
-      ]),
-      new Map([
-        ["codex-first", "turn-first"],
-        ["codex-second", "turn-second"],
-      ]),
-    );
+    const userMessages = [
+      { messageId: "codex-first", turnId: "turn-first" },
+      { messageId: "codex-second", turnId: "turn-second" },
+    ];
     let reboundThreadId: string | null = null;
 
     await revertCodexConversation({
@@ -176,7 +144,7 @@ describe("Codex Rewind", () => {
       cwd: "/workspace/project",
       model: "gpt-5.4-mini",
       serviceTier: null,
-      userMessageTurns,
+      userMessages,
       threadRollbackAvailable: true,
       setThreadId: (threadId) => {
         reboundThreadId = threadId;
@@ -209,7 +177,7 @@ describe("Codex Rewind", () => {
     }
 
     const codex = new PaginatedCodex();
-    const userMessageTurns = new CodexMessageTurns(new Map([["codex-first", 0]]));
+    const userMessages = [{ messageId: "codex-first", turnId: null }];
     let reboundThreadId: string | null = null;
 
     await expect(
@@ -220,7 +188,7 @@ describe("Codex Rewind", () => {
         cwd: "/workspace/project",
         model: "gpt-5.4-mini",
         serviceTier: null,
-        userMessageTurns,
+        userMessages,
         threadRollbackAvailable: true,
         setThreadId: (threadId) => {
           reboundThreadId = threadId;
@@ -235,14 +203,14 @@ describe("Codex Rewind", () => {
 
   test("declines to rewind when the user message is not in the Codex thread", async () => {
     const codex = new FakeCodex();
-    const userMessageTurns = new CodexMessageTurns(new Map([["codex-first", 0]]));
+    const userMessages = [{ messageId: "codex-first", turnId: null }];
 
     await expect(
       revertCodexConversation({
         client: codex,
         threadId: "source-thread",
         messageId: "missing-message",
-        userMessageTurns,
+        userMessages,
         threadRollbackAvailable: true,
         setThreadId: () => undefined,
       }),
@@ -250,4 +218,47 @@ describe("Codex Rewind", () => {
     expect(codex.recordedForks).toEqual([]);
     expect(codex.recordedRollbacks).toEqual([]);
   });
+
+  test("rolls back whole turns when a later turn holds a steer message", async () => {
+    const codex = new FakeCodex();
+
+    await revertCodexConversation({
+      client: codex,
+      threadId: "source-thread",
+      messageId: "codex-second",
+      userMessages: [
+        { messageId: "codex-first", turnId: "turn-first" },
+        { messageId: "codex-second", turnId: "turn-second" },
+        { messageId: "codex-steer", turnId: "turn-second" },
+      ],
+      threadRollbackAvailable: true,
+      setThreadId: () => undefined,
+    });
+
+    expect(codex.recordedRollbacks).toEqual([{ threadId: "forked-thread", numTurns: 1 }]);
+  });
+
+  test("declines to rewind a steer message without forking", async () => {
+    const codex = new FakeCodex();
+    let reboundThreadId: string | null = null;
+
+    await expect(
+      revertCodexConversation({
+        client: codex,
+        threadId: "source-thread",
+        messageId: "codex-steer",
+        userMessages: [
+          { messageId: "codex-first", turnId: "turn-first" },
+          { messageId: "codex-steer", turnId: "turn-first" },
+        ],
+        threadRollbackAvailable: true,
+        setThreadId: (threadId) => {
+          reboundThreadId = threadId;
+        },
+      }),
+    ).rejects.toThrow("Select the first message in the turn instead");
+    expect(codex.recordedForks).toEqual([]);
+    expect(codex.recordedRollbacks).toEqual([]);
+    expect(reboundThreadId).toBeNull();
+  });
 });

packages/server/src/server/agent/providers/codex/rewind.ts

@@ -15,9 +15,9 @@ export interface CodexRewindClient {
   request(method: string, params?: unknown, timeoutMs?: number): Promise<unknown>;
 }
 
-export interface CodexUserMessageTurnIndex {
-  resolve(messageId: string): { index: number; turnId: string | null } | null;
-  count(): number;
+export interface CodexUserMessage {
+  messageId: string;
+  turnId: string | null;
 }
 
 type CodexThreadHistoryMode = "legacy" | "paginated";
@@ -60,6 +60,12 @@ async function rollbackCodexThread(
   return parseCodexThreadRollbackResponse(await client.request("thread/rollback", params));
 }
 
+function countTurns(userMessages: readonly CodexUserMessage[]): number {
+  return userMessages.filter(
+    (message, index) => !message.turnId || message.turnId !== userMessages[index - 1]?.turnId,
+  ).length;
+}
+
 export async function revertCodexConversation(input: {
   client: CodexRewindClient;
   threadId: string | null;
@@ -68,7 +74,7 @@ export async function revertCodexConversation(input: {
   model?: string | null;
   serviceTier?: string | null;
   config?: Record<string, unknown> | null;
-  userMessageTurns: CodexUserMessageTurnIndex;
+  userMessages: readonly CodexUserMessage[];
   threadRollbackAvailable: boolean;
   setThreadId: (threadId: string) => void | Promise<void>;
 }): Promise<void> {
@@ -76,16 +82,21 @@ export async function revertCodexConversation(input: {
     throw new Error("Codex thread is not ready for rewind");
   }
 
-  const targetTurn = input.userMessageTurns.resolve(input.messageId);
-  if (targetTurn === null) {
+  const targetIndex = input.userMessages.findIndex(
+    (message) => message.messageId === input.messageId,
+  );
+  if (targetIndex === -1) {
     throw new Error(`Codex could not find user message ${input.messageId} in the current thread`);
   }
-
-  const currentUserTurnCount = input.userMessageTurns.count();
-  const numTurns = currentUserTurnCount - targetTurn.index;
-  if (numTurns < 0) {
-    throw new Error(`Codex user message ${input.messageId} is outside the current thread`);
+  const targetTurn = input.userMessages[targetIndex];
+  // Codex rewinds whole turns; a steer cannot be removed without the earlier
+  // messages of its turn.
+  if (targetTurn.turnId && input.userMessages[targetIndex - 1]?.turnId === targetTurn.turnId) {
+    throw new Error(
+      "Codex cannot rewind a message sent during a turn without removing earlier messages. Select the first message in the turn instead.",
+    );
   }
+  const numTurns = countTurns(input.userMessages.slice(targetIndex));
 
   // Codex does not carry the parent thread's config into a fork; without it the
   // forked thread falls back to the default model provider.