From 1b778cc70aacd537ada016fc40f3d879713497a8 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Sun, 13 Sep 2026 17:27:00 +0400 Subject: [PATCH] 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 5ffe2ee2..3298147d 100644 --- a/src/AirExtension.ts +++ b/src/AirExtension.ts @@ -17,6 +17,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); }