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 08e883fad5..c697d6fee1 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 @@ -31,6 +31,7 @@ import { normalizeCodexOutputSchema, toAgentUsage, } from "./codex-app-server-agent.js"; +import { AgentTimeline, type TimelineEntry } from "../timeline/index.js"; describe("mapCodexPlanUpdateToTodo", () => { test("preserves checklist progress without creating a plan card", () => { @@ -2604,6 +2605,89 @@ describe("Codex app-server provider", () => { await session.close(); }); + test("a steered prompt's rewind reports the native turn the fork removed, and the timeline keeps what Codex kept", async () => { + const appServer = createFakeCodexAppServer({ + initialize: () => ({ userAgent: "paseo/0.159.0 (Ubuntu 26.4.0; x86_64) (paseo; 0)" }), + "thread/read": () => ({ + thread: { id: "thread-1", historyMode: "legacy", turns: [] }, + }), + "thread/rollback": () => ({ + __jsonRpcError: { + code: -32600, + message: "Invalid request: unknown variant `thread/rollback`", + }, + }), + }); + const session = new CodexAppServerAgentSession( + createConfig({ cwd: "/workspace/project" }), + null, + createTestLogger(), + async () => appServer.child, + ); + + await session.startTurn("remember zero"); + emitCodexUserMessage(appServer, { + id: "codex-zero", + text: "remember zero", + turnId: "turn-zero", + }); + appServer.completeTurn(); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + await session.startTurn("remember first"); + emitCodexUserMessage(appServer, { + id: "codex-first", + text: "remember first", + turnId: "turn-first", + }); + // A steer joins the running native turn. + emitCodexUserMessage(appServer, { + id: "codex-second", + text: "remember second", + turnId: "turn-first", + }); + appServer.completeTurn(); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + const cut = await session.revertConversation({ messageId: "codex-second" }); + + expect( + appServer + .requests() + .filter((request) => request.method === "thread/fork") + .map((request) => (request.params as { beforeTurnId?: string }).beforeTurnId), + ).toEqual(["turn-first"]); + expect(cut).toEqual({ kind: "after", retainedMessageId: "codex-zero" }); + + const timeline = new AgentTimeline(); + const scope = { agentId: "agent" }; + await timeline.create(scope); + await timeline.append(scope, [ + prompt("zero"), + { item: { type: "assistant_message", text: "reply zero" } }, + prompt("first"), + prompt("second"), + { item: { type: "assistant_message", text: "reply second" } }, + ]); + const intent = await timeline.beginRewind(scope, "second"); + await timeline.recordRewindCut(scope, intent.id, cut); + await timeline.completeRewind(scope, intent.id); + expect((await timeline.readRows(scope)).map((row) => row.item)).toEqual([ + { type: "user_message", text: "zero", messageId: "zero" }, + { type: "assistant_message", text: "reply zero" }, + ]); + appServer.assertNoErrors(); + await session.close(); + + function prompt(name: string): TimelineEntry { + return { + item: { type: "user_message", text: name, messageId: name }, + providerMessageId: `codex-${name}`, + }; + } + }); + test("keeps a background terminal controllable on its original thread after rewind", async () => { let live = true; const appServer = createFakeCodexAppServer({ 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 045590c4ff..c0c587a65d 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 @@ -42,6 +42,7 @@ import { type ProviderRefreshContext, type ResolveAgentDefaultModeInput, type ConversationRewindCut, + type ConversationRewindResult, RewindNotAppliedError, } from "../agent-sdk-types.js"; import { importSessionFromPersistence } from "../provider-session-import.js"; @@ -4691,6 +4692,11 @@ export class CodexAppServerAgentSession implements AgentSession { : { index, turnId: this.userMessageProviderTurnIds.get(messageId) ?? null }; }, count: () => this.userMessageTurnIds.length, + list: () => + this.userMessageTurnIds.map((messageId) => ({ + messageId, + turnId: this.userMessageProviderTurnIds.get(messageId) ?? null, + })), }; } @@ -5128,7 +5134,7 @@ export class CodexAppServerAgentSession implements AgentSession { : null; } - async revertConversation(input: { messageId: string }): Promise { + async revertConversation(input: { messageId: string }): Promise { let threadBefore: string | null; let client: CodexAppServerClient; try { @@ -5151,8 +5157,9 @@ export class CodexAppServerAgentSession implements AgentSession { // The rewind forks; the thread it started from stays as it was. It happened once this session // continues the fork. let failure: unknown = null; + let cut: ConversationRewindCut | null = null; try { - await revertCodexConversation({ + cut = await revertCodexConversation({ client, threadId: threadBefore, messageId: input.messageId, @@ -5182,7 +5189,8 @@ export class CodexAppServerAgentSession implements AgentSession { if (failure) { this.logger.warn({ err: failure }, "Codex rewound the conversation but failed afterwards"); } - return { kind: "at_target" }; + // A failure after the switch leaves what the fork kept unreported; the resumed history shows it. + return cut ?? { kind: "unreported" }; } async interrupt(): Promise { 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 7a0b2dc0f9..445785be87 100644 --- a/packages/server/src/server/agent/providers/codex/rewind.test.ts +++ b/packages/server/src/server/agent/providers/codex/rewind.test.ts @@ -7,6 +7,7 @@ import type { CodexThreadRollbackResponse, } from "./app-server-transport.js"; import { + type CodexUserMessageTurn, type CodexUserMessageTurnIndex, type CodexRewindClient, revertCodexConversation, @@ -75,6 +76,15 @@ class CodexMessageTurns implements CodexUserMessageTurnIndex { count(): number { return this.indexesByMessageId.size; } + + list(): CodexUserMessageTurn[] { + return [...this.indexesByMessageId] + .sort(([, left], [, right]) => left - right) + .map(([messageId]) => ({ + messageId, + turnId: this.turnIdsByMessageId.get(messageId) ?? null, + })); + } } describe("Codex Rewind", () => { diff --git a/packages/server/src/server/agent/providers/codex/rewind.ts b/packages/server/src/server/agent/providers/codex/rewind.ts index a22931b92b..0cd606641e 100644 --- a/packages/server/src/server/agent/providers/codex/rewind.ts +++ b/packages/server/src/server/agent/providers/codex/rewind.ts @@ -9,6 +9,7 @@ import { parseCodexThreadRollbackResponse, } from "./app-server-transport.js"; import { readCodexThreadHistoryMode } from "./thread-history.js"; +import type { ConversationRewindCut } from "../../agent-sdk-types.js"; export interface CodexRewindClient { forkThread?(params: CodexThreadForkParams): Promise; @@ -16,9 +17,48 @@ export interface CodexRewindClient { request(method: string, params?: unknown, timeoutMs?: number): Promise; } +export interface CodexUserMessageTurn { + messageId: string; + /** The native turn the message started or joined; null when Codex did not say. */ + turnId: string | null; +} + export interface CodexUserMessageTurnIndex { resolve(messageId: string): { index: number; turnId: string | null } | null; count(): number; + /** The thread's user messages, in order. A steered message shares its native turn's id. */ + list(): readonly CodexUserMessageTurn[]; +} + +/** + * What the provider kept, given the first user message the rewind removed. Codex removes whole + * native turns, so a steered target takes the earlier messages of its turn with it. + */ +function rewindCut( + messages: readonly CodexUserMessageTurn[], + targetIndex: number, + firstRemoved: number, +): ConversationRewindCut { + if (firstRemoved === targetIndex) return { kind: "at_target" }; + return { kind: "after", retainedMessageId: messages[firstRemoved - 1]?.messageId ?? null }; +} + +/** The index of the first message in the native turn of `messages[index]`. */ +function turnStart(messages: readonly CodexUserMessageTurn[], index: number): number { + const turnId = messages[index]?.turnId ?? null; + if (turnId === null) return index; + let start = index; + while (start > 0 && messages[start - 1]!.turnId === turnId) start -= 1; + return start; +} + +/** The index of the first message a rollback of the last `numTurns` native turns removes. */ +function rollbackStart(messages: readonly CodexUserMessageTurn[], numTurns: number): number { + let start = messages.length; + for (let turns = 0; turns < numTurns && start > 0; turns += 1) { + start = turnStart(messages, start - 1); + } + return start; } async function forkCodexThread( @@ -52,7 +92,7 @@ export async function revertCodexConversation(input: { userMessageTurns: CodexUserMessageTurnIndex; threadRollbackAvailable: boolean; setThreadId: (threadId: string) => void | Promise; -}): Promise { +}): Promise { if (!input.threadId) { throw new Error("Codex thread is not ready for rewind"); } @@ -62,6 +102,7 @@ export async function revertCodexConversation(input: { throw new Error(`Codex could not find user message ${input.messageId} in the current thread`); } + const messages = input.userMessageTurns.list(); const currentUserTurnCount = input.userMessageTurns.count(); const numTurns = currentUserTurnCount - targetTurn.index; if (numTurns < 0) { @@ -93,7 +134,7 @@ export async function revertCodexConversation(input: { beforeTurnId: targetTurn.turnId, }); await input.setThreadId(forked.thread.id); - return; + return rewindCut(messages, targetTurn.index, turnStart(messages, targetTurn.index)); } // Legacy threads on Codex before 0.156 fork and then roll back. Fork is @@ -109,4 +150,5 @@ export async function revertCodexConversation(input: { numTurns, }); await input.setThreadId(rolledBack.thread.id); + return rewindCut(messages, targetTurn.index, rollbackStart(messages, numTurns)); }