diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.local.e2e.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.local.e2e.test.ts index 0469d9fbe1..99a4956078 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.local.e2e.test.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.local.e2e.test.ts @@ -217,6 +217,104 @@ function waitForEvent(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(); + 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 => + 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 () => { diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index 645fbd8726..983800fdab 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/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; handleToolApprovalRequest(params: unknown): Promise; 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 () => { diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index d392bb1349..3131d64018 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/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; diff --git a/packages/server/src/server/agent/providers/codex/rewind.test.ts b/packages/server/src/server/agent/providers/codex/rewind.test.ts index da7fe6a8f5..11349c9a15 100644 --- a/packages/server/src/server/agent/providers/codex/rewind.test.ts +++ b/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, - private readonly turnIdsByMessageId: Map = 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(); + }); }); diff --git a/packages/server/src/server/agent/providers/codex/rewind.ts b/packages/server/src/server/agent/providers/codex/rewind.ts index 021d57645c..76a6a41ebe 100644 --- a/packages/server/src/server/agent/providers/codex/rewind.ts +++ b/packages/server/src/server/agent/providers/codex/rewind.ts @@ -15,9 +15,9 @@ export interface CodexRewindClient { request(method: string, params?: unknown, timeoutMs?: number): Promise; } -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 | null; - userMessageTurns: CodexUserMessageTurnIndex; + userMessages: readonly CodexUserMessage[]; threadRollbackAvailable: boolean; setThreadId: (threadId: string) => void | Promise; }): Promise { @@ -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.