From 3acb318042363389949f8ee716966ec5ca5e82a6 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Sun, 13 Sep 2026 17:27:00 +0400 Subject: [PATCH 1/4] feat: support in-place session rewind Map the AIR session rewind extension to Codex thread/revert. Keep the same thread ID and avoid a provider fork. --- README.md | 1 + docs/session-rewind-extension.md | 56 ++++++++++++++ src/AcpExtensions.ts | 15 +++- src/AirExtension.ts | 1 + src/CodexAcpClient.ts | 5 ++ src/CodexAcpServer.ts | 8 ++ src/CodexAppServerClient.ts | 6 ++ src/SessionRewind.ts | 57 ++++++++++++++ .../CodexACPAgent/initialize.test.ts | 2 +- src/__tests__/SessionRewind.test.ts | 74 +++++++++++++++++++ src/index.ts | 14 ++++ 11 files changed, 237 insertions(+), 2 deletions(-) create mode 100644 docs/session-rewind-extension.md create mode 100644 src/SessionRewind.ts create mode 100644 src/__tests__/SessionRewind.test.ts diff --git a/README.md b/README.md index f43ad986..856c55a5 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Use [OpenAI Codex](https://github.com/openai/codex) from [Agent Client Protocol] - [Native ACP subagent sessions](docs/subagent-sessions.md) (after capability negotiation) with separate child histories and root-routed permissions; a legacy tool-call fallback otherwise. - [Background terminal tasks](docs/async-tasks.md) in AIR, with task status and targeted stop support after capability negotiation. - Session-scoped long-running goals through the provider-neutral [goal extension](docs/goal-extension.md). +- In-place message editing through the AIR [session rewind extension](docs/session-rewind-extension.md), without a provider fork. - A per-turn [agent file-change report](docs/agent-file-change-report.md) after capability negotiation. - Client-provided MCP servers over command-based stdio config and HTTP transport. - Slash commands: `/status`, `/mcp`, `/skills`, `/goal`, `/review`, `/review-branch`, `/review-commit`, `/compact`, and `/logout`, as well as configured skills. diff --git a/docs/session-rewind-extension.md b/docs/session-rewind-extension.md new file mode 100644 index 00000000..b26b2c61 --- /dev/null +++ b/docs/session-rewind-extension.md @@ -0,0 +1,56 @@ +# Session rewind extension + +Standard ACP can fork a session, but it cannot remove a transcript suffix from the same provider session. The experimental AIR session rewind extension adds that operation without creating another session. + +## Capability negotiation + +The adapter advertises `sessionRewind` in its `initialize` response: + +```json +{ + "_meta": { + "jetbrains": { + "air": { + "version": 1, + "capabilities": ["sessionRewind"] + } + } + } +} +``` + +A client must send `_session/rewind` only when the adapter advertises this capability. The leading underscore identifies a method outside standard ACP. + +## Request and response + +The request names the current ACP session and the first user message to remove: + +```json +{ + "sessionId": "thread-1", + "beforeMessage": { + "messageId": "user-message-2", + "messageFingerprint": "sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "messageOccurrence": 1 + }, + "resumeAtMessage": { + "messageId": "assistant-message-1", + "messageFingerprint": "sha256:abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789", + "messageOccurrence": 1 + } +} +``` + +`beforeMessage` is excluded from the retained history. `resumeAtMessage` identifies the last visible assistant message to retain. It is absent when the client rewinds the first user turn. + +Each history point contains the ACP message ID, the SHA-256 fingerprint of its complete text, and the one-based occurrence of that fingerprint for its role. The adapter uses the message ID first. It uses the fingerprint occurrence when restored provider history has different message IDs. + +The adapter returns `{ "rewound": true }` only after Codex accepts the rewind. A false response or an error leaves the client transcript unchanged. + +## Codex mapping + +The adapter reads the existing Codex thread history and resolves `beforeMessage` to its containing turn. It then calls `thread/revert` with that turn as the exclusive boundary. + +The Codex thread ID remains the ACP session ID. The adapter does not call `thread/fork`, create a thread, or add a session-list entry. `resumeAtMessage` is not needed for this mapping because Codex reverts at a turn boundary. + +After a successful response, the client can remove the same transcript suffix and place the selected user text in its editor. diff --git a/src/AcpExtensions.ts b/src/AcpExtensions.ts index b450c8bd..cca4297f 100644 --- a/src/AcpExtensions.ts +++ b/src/AcpExtensions.ts @@ -15,6 +15,10 @@ import { ASYNC_TASK_STOP_METHOD, type AsyncTaskStopExtRequest, } from "./async-tasks/AsyncTaskExtension"; +import { + SESSION_REWIND_METHOD, + type SessionRewindRequest, +} from "./SessionRewind"; export { AUTH_STATUS_META_KEY, @@ -79,6 +83,7 @@ export type ExtMethodRequest = | SessionSteeringExtRequest | GoalControlExtRequest | AsyncTaskStopExtRequest + | SessionRewindExtRequest export function isExtMethodRequest(request: { method: string, params: Record }): request is ExtMethodRequest { return request.method === "authentication/status" @@ -87,7 +92,8 @@ export function isExtMethodRequest(request: { method: string, params: Record, params: SessionSteerRequest, diff --git a/src/AirExtension.ts b/src/AirExtension.ts index 28af8784..1aaf9048 100644 --- a/src/AirExtension.ts +++ b/src/AirExtension.ts @@ -18,6 +18,7 @@ export const AIR_AGENT_FILE_CHANGE_REPORT_KEY = "agentFileChangeReport"; export const AIR_NATIVE_SUBAGENT_SESSIONS_KEY = "nativeSubagentSessions"; export const AIR_ASYNC_TASKS_KEY = "asyncTasks"; export const AIR_RECOMMENDED_CONFIG_VALUE_KEY = "recommendedValue"; +export const AIR_SESSION_REWIND_KEY = "sessionRewind"; export const AIR_ASYNC_TASKS_BACKGROUNDED_KEY = "backgrounded"; export const AIR_AGENT_FILE_CHANGE_REPORT_REQUEST_KEY = "agentFileChangeReportRequest"; export const AIR_EXTENSION_VERSION = 1; diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 7040057f..8166c5a4 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -69,6 +69,7 @@ import { } from "./AgentFileChangeReport"; import {CodexSubagentSubscriptions} from "./subagents/CodexSubagentSubscriptions"; import {forkSession as runForkSession} from "./SessionFork"; +import {rewindSession as runRewindSession, type SessionRewindRequest} from "./SessionRewind"; import type {SessionMetadata, SessionMetadataWithThread} from "./SessionMetadata"; export type {SessionMetadata, SessionMetadataWithThread} from "./SessionMetadata"; @@ -567,6 +568,10 @@ export class CodexAcpClient { }); } + async rewindSession(request: SessionRewindRequest): Promise<{rewound: boolean}> { + return await runRewindSession(request, this.codexClient); + } + async loadSession(request: acp.LoadSessionRequest, onSubscribed?: () => void): Promise { const additionalDirectories = readAdditionalDirectories(request.cwd, request.additionalDirectories, request._meta); await this.refreshSkills(request.cwd, additionalDirectories); diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 4ff0e1b4..120b77c7 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -78,6 +78,8 @@ import { type LegacySetSessionModelRequest, type LegacySetSessionModelResponse, SESSION_STEERING_METHOD, + SESSION_REWIND_METHOD, + type SessionRewindRequest, type SessionSteeringResponse, type SessionSteerRequest, } from "./AcpExtensions"; @@ -134,6 +136,7 @@ import { AIR_ASYNC_TASKS_KEY, AIR_NATIVE_SUBAGENT_SESSIONS_KEY, AIR_RECOMMENDED_CONFIG_VALUE_KEY, + AIR_SESSION_REWIND_KEY, AIR_EXTENSION_CAPABILITIES_KEY, AIR_EXTENSION_VERSION, AIR_EXTENSION_VERSION_KEY, @@ -394,6 +397,7 @@ export class CodexAcpServer { AIR_NATIVE_SUBAGENT_SESSIONS_KEY, AIR_ASYNC_TASKS_KEY, AIR_RECOMMENDED_CONFIG_VALUE_KEY, + AIR_SESSION_REWIND_KEY, ], }, }, @@ -429,6 +433,10 @@ export class CodexAcpServer { ), }; } + case SESSION_REWIND_METHOD: + return await this.runWithProcessCheck( + () => this.codexAcpClient.rewindSession(methodRequest.params as SessionRewindRequest), + ); case GOAL_CONTROL_METHOD: case LEGACY_GOAL_CONTROL_METHOD: { const sessionState = this.sessions.get(methodRequest.params.sessionId); diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts index daa7e875..71f91c34 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -60,6 +60,8 @@ import type { ThreadTurnsListResponse, ThreadResumeParams, ThreadResumeResponse, + ThreadRevertParams, + ThreadRevertResponse, ThreadSettings, ThreadStartParams, ThreadStartResponse, @@ -558,6 +560,10 @@ export class CodexAppServerClient { return await this.sendRequest({ method: "thread/fork", params: params }); } + async threadRevert(params: ThreadRevertParams): Promise { + return await this.sendRequest({method: "thread/revert", params}); + } + getThreadSettings(threadId: string): ThreadSettings | undefined { return this.threadSettings.get(threadId); } diff --git a/src/SessionRewind.ts b/src/SessionRewind.ts new file mode 100644 index 00000000..69f8d464 --- /dev/null +++ b/src/SessionRewind.ts @@ -0,0 +1,57 @@ +import {createHash} from "node:crypto"; +import {RequestError} from "@agentclientprotocol/sdk"; +import type {CodexAppServerClient} from "./CodexAppServerClient"; + +export const SESSION_REWIND_METHOD = "_session/rewind"; +export const SESSION_REWIND_CAPABILITY = "sessionRewind"; + +export type SessionHistoryPoint = { + messageId: string; + messageFingerprint: string; + messageOccurrence: number; +}; + +export type SessionRewindRequest = { + sessionId: string; + beforeMessage: SessionHistoryPoint; + resumeAtMessage?: SessionHistoryPoint; +}; + +export type SessionRewindResponse = {rewound: boolean}; + +export async function rewindSession( + request: SessionRewindRequest, + client: CodexAppServerClient, +): Promise { + const history = await client.threadReadWithHistory(request.sessionId); + const userTurns = history.thread.turns.flatMap(turn => turn.items + .filter(item => item.type === "userMessage") + .map(item => ({turn, item}))); + const candidates = messageIdCandidates(request.beforeMessage.messageId); + const exact = userTurns.find(({item}) => candidates.includes(item.id)); + const fingerprintMatches = userTurns.filter(({item}) => + fingerprint(userMessageText(item.content)) === request.beforeMessage.messageFingerprint, + ); + const turn = exact?.turn ?? fingerprintMatches[request.beforeMessage.messageOccurrence - 1]?.turn; + if (!turn) { + throw RequestError.invalidParams( + {messageId: request.beforeMessage.messageId}, + `Rewind message ${request.beforeMessage.messageId} was not found in session ${request.sessionId}`, + ); + } + await client.threadRevert({threadId: request.sessionId, beforeTurnId: turn.id}); + return {rewound: true}; +} + +function userMessageText(content: Array<{type: string; text?: string}>): string { + return content.filter(item => item.type === "text").map(item => item.text ?? "").join(""); +} + +function fingerprint(text: string): string { + return `sha256:${createHash("sha256").update(text, "utf8").digest("hex")}`; +} + +function messageIdCandidates(messageId: string): string[] { + const protocolMessageId = messageId.replace(/:segment:\d+$/, ""); + return protocolMessageId === messageId ? [messageId] : [messageId, protocolMessageId]; +} diff --git a/src/__tests__/CodexACPAgent/initialize.test.ts b/src/__tests__/CodexACPAgent/initialize.test.ts index e6bdb8bb..76c73690 100644 --- a/src/__tests__/CodexACPAgent/initialize.test.ts +++ b/src/__tests__/CodexACPAgent/initialize.test.ts @@ -78,7 +78,7 @@ describe('CodexACPAgent - initialize', () => { jetbrains: { air: { version: 1, - capabilities: ["sessionFailure", "agentFileChangeReport", "nativeSubagentSessions", "asyncTasks", "recommendedValue"], + capabilities: ["sessionFailure", "agentFileChangeReport", "nativeSubagentSessions", "asyncTasks", "recommendedValue", "sessionRewind"], }, }, }, diff --git a/src/__tests__/SessionRewind.test.ts b/src/__tests__/SessionRewind.test.ts new file mode 100644 index 00000000..37f8f535 --- /dev/null +++ b/src/__tests__/SessionRewind.test.ts @@ -0,0 +1,74 @@ +import {describe, expect, it, vi} from "vitest"; +import type {CodexAppServerClient} from "../CodexAppServerClient"; +import {rewindSession} from "../SessionRewind"; + +describe("session rewind", () => { + it("reverts the same Codex thread before the selected user turn", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + turns: [ + {id: "turn-1", items: [{type: "userMessage", id: "user-1", content: [{type: "text", text: "one"}]}]}, + {id: "turn-2", items: [{type: "userMessage", id: "user-2", content: [{type: "text", text: "two"}]}]}, + ], + }, + }), + threadRevert: vi.fn().mockResolvedValue({}), + } as unknown as CodexAppServerClient; + + const result = await rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "user-2", + messageFingerprint: "sha256:3fc4ccfe745870e2c0d99f71f30ff0656c8d1ed5d3f3b71b17a64d1c0d9a4f5f", + messageOccurrence: 1, + }, + }, client); + + expect(result).toEqual({rewound: true}); + expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); + }); + + it("resolves a restored message through its fingerprint occurrence", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + turns: [ + {id: "turn-1", items: [{type: "userMessage", id: "new-1", content: [{type: "text", text: "repeat"}]}]}, + {id: "turn-2", items: [{type: "userMessage", id: "new-2", content: [{type: "text", text: "repeat"}]}]}, + ], + }, + }), + threadRevert: vi.fn().mockResolvedValue({}), + } as unknown as CodexAppServerClient; + + const result = await rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "stale-id", + messageFingerprint: "sha256:25e2b6b106523880e27763084ffa6a0756335be0d7106022535365b9ad39b4b1", + messageOccurrence: 2, + }, + }, client); + + expect(result).toEqual({rewound: true}); + expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); + }); + + it("does not revert when the selected message is absent", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({thread: {turns: []}}), + threadRevert: vi.fn(), + } as unknown as CodexAppServerClient; + + await expect(rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "missing", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + }, client)).rejects.toThrow("Rewind message missing was not found"); + expect(client.threadRevert).not.toHaveBeenCalled(); + }); +}); diff --git a/src/index.ts b/src/index.ts index 19759300..db31da96 100644 --- a/src/index.ts +++ b/src/index.ts @@ -17,6 +17,7 @@ import { SESSION_STEERING_METHOD, } from "./AcpExtensions"; import {ASYNC_TASK_STOP_METHOD} from "./async-tasks/AsyncTaskExtension"; +import {SESSION_REWIND_METHOD} from "./SessionRewind"; const emptyExtensionParamsParser = z.preprocess( (params) => params ?? {}, @@ -50,6 +51,18 @@ const asyncTaskStopParamsParser = z.object({ asyncTaskId: z.string().trim().min(1), }).passthrough(); +const sessionHistoryPointParser = z.object({ + messageId: z.string().trim().min(1), + messageFingerprint: z.string().regex(/^sha256:[0-9a-f]{64}$/), + messageOccurrence: z.number().int().positive(), +}); + +const sessionRewindParamsParser = z.object({ + sessionId: z.string().trim().min(1), + beforeMessage: sessionHistoryPointParser, + resumeAtMessage: sessionHistoryPointParser.optional(), +}).passthrough(); + if (process.argv.includes("--version")) { console.log(`${packageJson.name} ${packageJson.version}`); process.exit(0); @@ -168,6 +181,7 @@ function startAcpServer() { .onRequest(LEGACY_SET_SESSION_MODEL_METHOD, legacySetSessionModelParamsParser, (ctx) => getAgent().extMethod(LEGACY_SET_SESSION_MODEL_METHOD, ctx.params)) .onRequest(SESSION_STEERING_METHOD, sessionSteerParamsParser, (ctx) => getAgent().extMethod(SESSION_STEERING_METHOD, ctx.params)) .onRequest(ASYNC_TASK_STOP_METHOD, asyncTaskStopParamsParser, (ctx) => getAgent().extMethod(ASYNC_TASK_STOP_METHOD, ctx.params)) + .onRequest(SESSION_REWIND_METHOD, sessionRewindParamsParser, (ctx) => getAgent().extMethod(SESSION_REWIND_METHOD, ctx.params)) .onRequest(GOAL_CONTROL_METHOD, goalControlParamsParser, (ctx) => getAgent().extMethod(GOAL_CONTROL_METHOD, ctx.params)) .connect(acpJsonStream); } From 28d0be7c0c0703c36a57e7391cb529eda4c59018 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Tue, 15 Sep 2026 13:23:45 +0400 Subject: [PATCH 2/4] fix: keep session rewind histories aligned --- docs/session-rewind-extension.md | 4 +- src/CodexAcpServer.ts | 31 +-------- src/CodexAppServerClient.ts | 6 ++ src/SessionRewind.ts | 33 ++++++---- src/UserInputContent.ts | 40 ++++++++++++ src/__tests__/SessionRewind.test.ts | 97 +++++++++++++++++++++++++++++ 6 files changed, 169 insertions(+), 42 deletions(-) create mode 100644 src/UserInputContent.ts diff --git a/docs/session-rewind-extension.md b/docs/session-rewind-extension.md index b26b2c61..54df9e92 100644 --- a/docs/session-rewind-extension.md +++ b/docs/session-rewind-extension.md @@ -49,8 +49,8 @@ The adapter returns `{ "rewound": true }` only after Codex accepts the rewind. A ## Codex mapping -The adapter reads the existing Codex thread history and resolves `beforeMessage` to its containing turn. It then calls `thread/revert` with that turn as the exclusive boundary. +The adapter reads the existing Codex thread history and resolves `beforeMessage` to its containing turn. Rewind is rejected when the selected message is a steer inside an existing turn because Codex cannot remove only that suffix. For paginated history, the adapter calls `thread/revert` with the containing turn as the exclusive boundary; legacy history uses the equivalent turn-count rollback operation. -The Codex thread ID remains the ACP session ID. The adapter does not call `thread/fork`, create a thread, or add a session-list entry. `resumeAtMessage` is not needed for this mapping because Codex reverts at a turn boundary. +The Codex thread ID remains the ACP session ID. The adapter does not call `thread/fork`, create a thread, or add a session-list entry. `resumeAtMessage` is not needed after the turn-boundary validation. After a successful response, the client can remove the same transcript suffix and place the selected user text in its editor. diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 120b77c7..39c6c8b2 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -60,6 +60,7 @@ import type {QuotaMeta} from "./QuotaMeta"; import {logger} from "./Logger"; import {sanitizeMcpServerName} from "./McpServerName"; import {createResponseItemHistoryFallbackUpdates} from "./ResponseItemHistoryFallback"; +import {userInputToContentBlocks} from "./UserInputContent"; import { AUTH_STATUS_META_KEY, AUTH_STATUS_UPDATE_METHOD, @@ -2294,7 +2295,7 @@ export class CodexAcpServer { const updates: UpdateSessionEvent[] = []; const messageId = item.id; for (const input of item.content) { - const blocks = this.userInputToContentBlocks(input); + const blocks = userInputToContentBlocks(input); for (const block of blocks) { updates.push(createUserMessageChunk(block, messageId)); } @@ -2357,34 +2358,6 @@ export class CodexAcpServer { ); } - private userInputToContentBlocks(input: UserInput): acp.ContentBlock[] { - switch (input.type) { - case "text": - return input.text.length > 0 ? [{ type: "text", text: input.text }] : []; - case "image": - return [{ type: "text", text: this.formatUriAsLink("image", input.url) }]; - case "localImage": { - const uri = input.path.startsWith("file://") ? input.path : `file://${input.path}`; - return [{ type: "text", text: this.formatUriAsLink(null, uri) }]; - } - case "skill": - return [{ type: "text", text: `skill:${input.name} (${input.path})` }]; - } - return []; - } - - private formatUriAsLink(name: string | null, uri: string): string { - if (name && name.length > 0) { - return `[@${name}](${uri})`; - } - if (uri.startsWith("file://")) { - const path = uri.replace("file://", ""); - const fileName = path.split("/").pop() ?? path; - return `[@${fileName}](${uri})`; - } - return uri; - } - getSessionState(sessionId: string): SessionState { const sessionState = this.sessions.get(sessionId); if (!sessionState) { diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts index 71f91c34..39364c96 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -62,6 +62,8 @@ import type { ThreadResumeResponse, ThreadRevertParams, ThreadRevertResponse, + ThreadRollbackParams, + ThreadRollbackResponse, ThreadSettings, ThreadStartParams, ThreadStartResponse, @@ -564,6 +566,10 @@ export class CodexAppServerClient { return await this.sendRequest({method: "thread/revert", params}); } + async threadRollback(params: ThreadRollbackParams): Promise { + return await this.sendRequest({method: "thread/rollback", params}); + } + getThreadSettings(threadId: string): ThreadSettings | undefined { return this.threadSettings.get(threadId); } diff --git a/src/SessionRewind.ts b/src/SessionRewind.ts index 69f8d464..2ccbed23 100644 --- a/src/SessionRewind.ts +++ b/src/SessionRewind.ts @@ -1,6 +1,7 @@ import {createHash} from "node:crypto"; import {RequestError} from "@agentclientprotocol/sdk"; import type {CodexAppServerClient} from "./CodexAppServerClient"; +import {userInputVisibleText} from "./UserInputContent"; export const SESSION_REWIND_METHOD = "_session/rewind"; export const SESSION_REWIND_CAPABILITY = "sessionRewind"; @@ -24,29 +25,39 @@ export async function rewindSession( client: CodexAppServerClient, ): Promise { const history = await client.threadReadWithHistory(request.sessionId); - const userTurns = history.thread.turns.flatMap(turn => turn.items - .filter(item => item.type === "userMessage") - .map(item => ({turn, item}))); + const userTurns = history.thread.turns.flatMap((turn, turnIndex) => turn.items + .flatMap((item, itemIndex) => item.type === "userMessage" + ? [{turn, turnIndex, item, itemIndex}] + : [])); const candidates = messageIdCandidates(request.beforeMessage.messageId); const exact = userTurns.find(({item}) => candidates.includes(item.id)); const fingerprintMatches = userTurns.filter(({item}) => - fingerprint(userMessageText(item.content)) === request.beforeMessage.messageFingerprint, + fingerprint(userInputVisibleText(item.content)) === request.beforeMessage.messageFingerprint, ); - const turn = exact?.turn ?? fingerprintMatches[request.beforeMessage.messageOccurrence - 1]?.turn; - if (!turn) { + const match = exact ?? fingerprintMatches[request.beforeMessage.messageOccurrence - 1]; + if (!match) { throw RequestError.invalidParams( {messageId: request.beforeMessage.messageId}, `Rewind message ${request.beforeMessage.messageId} was not found in session ${request.sessionId}`, ); } - await client.threadRevert({threadId: request.sessionId, beforeTurnId: turn.id}); + if (match.itemIndex !== 0) { + throw RequestError.invalidParams( + {messageId: request.beforeMessage.messageId}, + `Rewind message ${request.beforeMessage.messageId} does not start a turn`, + ); + } + if (history.thread.historyMode === "legacy") { + await client.threadRollback({ + threadId: request.sessionId, + numTurns: history.thread.turns.length - match.turnIndex, + }); + } else { + await client.threadRevert({threadId: request.sessionId, beforeTurnId: match.turn.id}); + } return {rewound: true}; } -function userMessageText(content: Array<{type: string; text?: string}>): string { - return content.filter(item => item.type === "text").map(item => item.text ?? "").join(""); -} - function fingerprint(text: string): string { return `sha256:${createHash("sha256").update(text, "utf8").digest("hex")}`; } diff --git a/src/UserInputContent.ts b/src/UserInputContent.ts new file mode 100644 index 00000000..e7502d65 --- /dev/null +++ b/src/UserInputContent.ts @@ -0,0 +1,40 @@ +import * as acp from "@agentclientprotocol/sdk"; +import type {UserInput} from "./app-server/v2"; + +export function userInputToContentBlocks(input: UserInput): acp.ContentBlock[] { + switch (input.type) { + case "text": + return input.text.length > 0 ? [{type: "text", text: input.text}] : []; + case "image": + return [{type: "text", text: formatUriAsLink("image", input.url)}]; + case "localImage": { + const uri = input.path.startsWith("file://") ? input.path : `file://${input.path}`; + return [{type: "text", text: formatUriAsLink(null, uri)}]; + } + case "skill": + return [{type: "text", text: `skill:${input.name} (${input.path})`}]; + case "audio": + case "localAudio": + case "mention": + return []; + } +} + +export function userInputVisibleText(content: UserInput[]): string { + return content.flatMap(userInputToContentBlocks) + .filter((block): block is Extract => block.type === "text") + .map(block => block.text) + .join(""); +} + +function formatUriAsLink(name: string | null, uri: string): string { + if (name && name.length > 0) { + return `[@${name}](${uri})`; + } + if (uri.startsWith("file://")) { + const path = uri.replace("file://", ""); + const fileName = path.split("/").pop() ?? path; + return `[@${fileName}](${uri})`; + } + return uri; +} diff --git a/src/__tests__/SessionRewind.test.ts b/src/__tests__/SessionRewind.test.ts index 37f8f535..fd91e891 100644 --- a/src/__tests__/SessionRewind.test.ts +++ b/src/__tests__/SessionRewind.test.ts @@ -55,6 +55,103 @@ describe("session rewind", () => { expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); }); + it("uses the visible replay text when fingerprinting multimodal and skill inputs", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + historyMode: "paginated", + turns: [{ + id: "turn-1", + items: [{ + type: "userMessage", + id: "new-id", + content: [ + {type: "text", text: "look"}, + {type: "image", url: "https://example.com/image.png"}, + {type: "skill", name: "review", path: "/tmp/SKILL.md"}, + ], + }], + }], + }, + }), + threadRevert: vi.fn().mockResolvedValue({}), + } as unknown as CodexAppServerClient; + + await rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "stale-id", + messageFingerprint: "sha256:d0425f232dd6a5d6a18eee0fb305ff976368b93fb5ad3da67a4b41d919f7e2de", + messageOccurrence: 1, + }, + }, client); + + expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-1"}); + }); + + it("uses turn-count rollback for legacy thread history", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + historyMode: "legacy", + turns: [ + {id: "turn-1", items: [{type: "userMessage", id: "user-1", content: [{type: "text", text: "one"}]}]}, + {id: "turn-2", items: [{type: "userMessage", id: "user-2", content: [{type: "text", text: "two"}]}]}, + {id: "turn-3", items: [{type: "userMessage", id: "user-3", content: [{type: "text", text: "three"}]}]}, + ], + }, + }), + threadRollback: vi.fn().mockResolvedValue({}), + threadRevert: vi.fn(), + } as unknown as CodexAppServerClient; + + await rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "user-2", + messageFingerprint: "sha256:3fc4ccfe745870e2c0d99f71f30ff0656c8d1ed5d3f3b71b17a64d1c0d9a4f5f", + messageOccurrence: 1, + }, + }, client); + + expect(client.threadRollback).toHaveBeenCalledWith({threadId: "thread-1", numTurns: 2}); + expect(client.threadRevert).not.toHaveBeenCalled(); + }); + + it("rejects rewinding a steer inside an existing turn", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + historyMode: "paginated", + turns: [{ + id: "turn-1", + items: [ + {type: "userMessage", id: "user-1", content: [{type: "text", text: "first"}]}, + {type: "agentMessage", id: "assistant-1", text: "working"}, + {type: "userMessage", id: "steer-1", content: [{type: "text", text: "steer"}]}, + ], + }], + }, + }), + threadRevert: vi.fn(), + } as unknown as CodexAppServerClient; + + await expect(rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "steer-1", + messageFingerprint: "sha256:57fce44d7c6df51ad8525da1580a246e9d1142d79d1d1f176b1d29643d61ed44", + messageOccurrence: 1, + }, + resumeAtMessage: { + messageId: "assistant-1", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + }, client)).rejects.toThrow("does not start a turn"); + expect(client.threadRevert).not.toHaveBeenCalled(); + }); + it("does not revert when the selected message is absent", async () => { const client = { threadReadWithHistory: vi.fn().mockResolvedValue({thread: {turns: []}}), From 10ef83cce6a72a2ee0df3dcad3d88cf5cfc29638 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Tue, 15 Sep 2026 13:33:34 +0400 Subject: [PATCH 3/4] fix: serialize session rewind lifecycle --- src/CodexAcpClient.ts | 4 +- src/CodexAcpServer.ts | 3 ++ src/SessionRewind.ts | 4 +- src/__tests__/CodexACPAgent/providers.test.ts | 49 +++++++++++++++++ src/__tests__/SessionRewind.test.ts | 52 +++++++++++++++++++ 5 files changed, 110 insertions(+), 2 deletions(-) diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 8166c5a4..0c3c10c4 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -569,7 +569,9 @@ export class CodexAcpClient { } async rewindSession(request: SessionRewindRequest): Promise<{rewound: boolean}> { - return await runRewindSession(request, this.codexClient); + const response = await runRewindSession(request, this.codexClient); + await this.waitForSessionNotifications(request.sessionId); + return response; } async loadSession(request: acp.LoadSessionRequest, onSubscribed?: () => void): Promise { diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 39c6c8b2..9d9e6586 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -435,6 +435,9 @@ export class CodexAcpServer { }; } case SESSION_REWIND_METHOD: + if (this.providerUpdate !== null) { + await this.providerUpdate; + } return await this.runWithProcessCheck( () => this.codexAcpClient.rewindSession(methodRequest.params as SessionRewindRequest), ); diff --git a/src/SessionRewind.ts b/src/SessionRewind.ts index 2ccbed23..84eba107 100644 --- a/src/SessionRewind.ts +++ b/src/SessionRewind.ts @@ -30,7 +30,9 @@ export async function rewindSession( ? [{turn, turnIndex, item, itemIndex}] : [])); const candidates = messageIdCandidates(request.beforeMessage.messageId); - const exact = userTurns.find(({item}) => candidates.includes(item.id)); + const exact = candidates + .map(candidate => userTurns.find(({item}) => item.id === candidate)) + .find(match => match !== undefined); const fingerprintMatches = userTurns.filter(({item}) => fingerprint(userInputVisibleText(item.content)) === request.beforeMessage.messageFingerprint, ); diff --git a/src/__tests__/CodexACPAgent/providers.test.ts b/src/__tests__/CodexACPAgent/providers.test.ts index 5a724039..db05eb52 100644 --- a/src/__tests__/CodexACPAgent/providers.test.ts +++ b/src/__tests__/CodexACPAgent/providers.test.ts @@ -2,6 +2,15 @@ import {describe, expect, it, vi} from "vitest"; import * as acp from "@agentclientprotocol/sdk"; import {createCodexMockTestFixture, createTestSessionState} from "../acp-test-utils"; import {CodexAcpClient, CUSTOM_GATEWAY_PROVIDER_ID, OPENAI_PROVIDER_ID} from "../../CodexAcpClient"; +import {SESSION_REWIND_METHOD} from "../../SessionRewind"; + +function deferred(): {promise: Promise, resolve: (value: T) => void} { + let resolve!: (value: T) => void; + const promise = new Promise(resolvePromise => { + resolve = resolvePromise; + }); + return {promise, resolve}; +} async function expectInvalidParams(fn: () => unknown): Promise { const caught = await Promise.resolve().then(fn).catch((err: unknown) => err); @@ -304,6 +313,46 @@ describe("Configurable LLM providers (providers/*)", () => { }); }); + it("waits for an in-flight provider update before rewinding", async () => { + const replacement = createCodexMockTestFixture().getCodexAcpClient(); + const resumeStarted = deferred(); + const resume = deferred(); + vi.spyOn(replacement, "initialize").mockResolvedValue(); + vi.spyOn(replacement, "resumeSession").mockImplementation(async () => { + resumeStarted.resolve(); + return await resume.promise; + }); + const replacementRewind = vi.spyOn(replacement, "rewindSession").mockResolvedValue({rewound: true}); + const fixture = createCodexMockTestFixture(vi.fn().mockResolvedValue(replacement)); + const agent = fixture.getCodexAcpAgent(); + await agent.initialize({protocolVersion: acp.PROTOCOL_VERSION}); + const sessions = (agent as unknown as {sessions: Map>}).sessions; + sessions.set("thread-1", createTestSessionState({sessionId: "thread-1", cwd: "/workspace"})); + + const providerUpdate = agent.setProvider({ + providerId: OPENAI_PROVIDER_ID, + apiType: "openai", + baseUrl: "https://gateway.example/v1", + }); + await resumeStarted.promise; + const rewind = agent.extMethod(SESSION_REWIND_METHOD, { + sessionId: "thread-1", + beforeMessage: { + messageId: "user-1", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + }); + + await Promise.resolve(); + expect(replacementRewind).not.toHaveBeenCalled(); + + resume.resolve({} as never); + await providerUpdate; + await expect(rewind).resolves.toEqual({rewound: true}); + expect(replacementRewind).toHaveBeenCalledOnce(); + }); + it("shares state with the legacy gateway auth method", async () => { const fixture = createCodexMockTestFixture(); const codexAcpClient = fixture.getCodexAcpClient(); diff --git a/src/__tests__/SessionRewind.test.ts b/src/__tests__/SessionRewind.test.ts index fd91e891..389c75d1 100644 --- a/src/__tests__/SessionRewind.test.ts +++ b/src/__tests__/SessionRewind.test.ts @@ -1,5 +1,6 @@ import {describe, expect, it, vi} from "vitest"; import type {CodexAppServerClient} from "../CodexAppServerClient"; +import {CodexAcpClient} from "../CodexAcpClient"; import {rewindSession} from "../SessionRewind"; describe("session rewind", () => { @@ -55,6 +56,31 @@ describe("session rewind", () => { expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); }); + it("prefers the exact segmented message id over its protocol id fallback", async () => { + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + turns: [ + {id: "turn-1", items: [{type: "userMessage", id: "user-1", content: [{type: "text", text: "fallback"}]}]}, + {id: "turn-2", items: [{type: "userMessage", id: "user-1:segment:0", content: [{type: "text", text: "exact"}]}]}, + ], + }, + }), + threadRevert: vi.fn().mockResolvedValue({}), + } as unknown as CodexAppServerClient; + + await rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "user-1:segment:0", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + }, client); + + expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); + }); + it("uses the visible replay text when fingerprinting multimodal and skill inputs", async () => { const client = { threadReadWithHistory: vi.fn().mockResolvedValue({ @@ -168,4 +194,30 @@ describe("session rewind", () => { }, client)).rejects.toThrow("Rewind message missing was not found"); expect(client.threadRevert).not.toHaveBeenCalled(); }); + + it("drains queued session notifications before acknowledging rewind", async () => { + const appServerClient = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + historyMode: "paginated", + turns: [{id: "turn-1", items: [{type: "userMessage", id: "user-1", content: [{type: "text", text: "one"}]}]}], + }, + }), + threadRevert: vi.fn().mockResolvedValue({}), + } as unknown as CodexAppServerClient; + const client = new CodexAcpClient(appServerClient); + const waitForNotifications = vi.spyOn(client, "waitForSessionNotifications").mockResolvedValue(); + + await client.rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "user-1", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + }); + + expect(waitForNotifications).toHaveBeenCalledWith("thread-1"); + expect(appServerClient.threadRevert).toHaveBeenCalledBefore(waitForNotifications); + }); }); From c2c9026b315a83dc3cf8b659afbf2108ddedf73e Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Tue, 15 Sep 2026 13:41:42 +0400 Subject: [PATCH 4/4] fix: validate session rewind targets --- src/CodexAcpServer.ts | 6 ++++ src/SessionRewind.ts | 5 ++- src/__tests__/CodexACPAgent/providers.test.ts | 15 +++++++++ src/__tests__/SessionRewind.test.ts | 31 +++++++++++++++++++ 4 files changed, 54 insertions(+), 3 deletions(-) diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 9d9e6586..3172f4de 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -438,6 +438,12 @@ export class CodexAcpServer { if (this.providerUpdate !== null) { await this.providerUpdate; } + if (!this.sessions.has(methodRequest.params.sessionId)) { + throw RequestError.invalidParams( + undefined, + `Unknown session: ${methodRequest.params.sessionId}`, + ); + } return await this.runWithProcessCheck( () => this.codexAcpClient.rewindSession(methodRequest.params as SessionRewindRequest), ); diff --git a/src/SessionRewind.ts b/src/SessionRewind.ts index 84eba107..a905d2a8 100644 --- a/src/SessionRewind.ts +++ b/src/SessionRewind.ts @@ -33,10 +33,9 @@ export async function rewindSession( const exact = candidates .map(candidate => userTurns.find(({item}) => item.id === candidate)) .find(match => match !== undefined); - const fingerprintMatches = userTurns.filter(({item}) => + const match = exact ?? userTurns.filter(({item}) => fingerprint(userInputVisibleText(item.content)) === request.beforeMessage.messageFingerprint, - ); - const match = exact ?? fingerprintMatches[request.beforeMessage.messageOccurrence - 1]; + )[request.beforeMessage.messageOccurrence - 1]; if (!match) { throw RequestError.invalidParams( {messageId: request.beforeMessage.messageId}, diff --git a/src/__tests__/CodexACPAgent/providers.test.ts b/src/__tests__/CodexACPAgent/providers.test.ts index db05eb52..31428894 100644 --- a/src/__tests__/CodexACPAgent/providers.test.ts +++ b/src/__tests__/CodexACPAgent/providers.test.ts @@ -353,6 +353,21 @@ describe("Configurable LLM providers (providers/*)", () => { expect(replacementRewind).toHaveBeenCalledOnce(); }); + it("rejects rewind for a thread that is not a loaded ACP session", async () => { + const fixture = createCodexMockTestFixture(); + const rewind = vi.spyOn(fixture.getCodexAcpClient(), "rewindSession"); + + await expect(fixture.getCodexAcpAgent().extMethod(SESSION_REWIND_METHOD, { + sessionId: "persisted-but-not-loaded", + beforeMessage: { + messageId: "user-1", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + })).rejects.toThrow("Unknown session: persisted-but-not-loaded"); + expect(rewind).not.toHaveBeenCalled(); + }); + it("shares state with the legacy gateway auth method", async () => { const fixture = createCodexMockTestFixture(); const codexAcpClient = fixture.getCodexAcpClient(); diff --git a/src/__tests__/SessionRewind.test.ts b/src/__tests__/SessionRewind.test.ts index 389c75d1..bbc43e01 100644 --- a/src/__tests__/SessionRewind.test.ts +++ b/src/__tests__/SessionRewind.test.ts @@ -81,6 +81,37 @@ describe("session rewind", () => { expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); }); + it("does not fingerprint history when an exact message id is found", async () => { + const unreadableItem = Object.defineProperty({type: "userMessage", id: "other"}, "content", { + enumerable: true, + get: () => { + throw new Error("fingerprint fallback should not run"); + }, + }); + const client = { + threadReadWithHistory: vi.fn().mockResolvedValue({ + thread: { + turns: [ + {id: "turn-1", items: [unreadableItem]}, + {id: "turn-2", items: [{type: "userMessage", id: "exact", content: [{type: "text", text: "selected"}]}]}, + ], + }, + }), + threadRevert: vi.fn().mockResolvedValue({}), + } as unknown as CodexAppServerClient; + + await rewindSession({ + sessionId: "thread-1", + beforeMessage: { + messageId: "exact", + messageFingerprint: `sha256:${"0".repeat(64)}`, + messageOccurrence: 1, + }, + }, client); + + expect(client.threadRevert).toHaveBeenCalledWith({threadId: "thread-1", beforeTurnId: "turn-2"}); + }); + it("uses the visible replay text when fingerprinting multimodal and skill inputs", async () => { const client = { threadReadWithHistory: vi.fn().mockResolvedValue({