feat: support in-place session rewind

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.

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.

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<string, unknown> }): request is ExtMethodRequest {
     return request.method === "authentication/status"
@@ -87,7 +92,8 @@ export function isExtMethodRequest(request: { method: string, params: Record<str
         || request.method === GOAL_CONTROL_METHOD
         || request.method === LEGACY_GOAL_CONTROL_METHOD
         || request.method === SESSION_STEERING_METHOD
-        || request.method === ASYNC_TASK_STOP_METHOD;
+        || request.method === ASYNC_TASK_STOP_METHOD
+        || request.method === SESSION_REWIND_METHOD;
 }
 
 /**
@@ -133,6 +139,13 @@ export type SessionSteeringExtRequest = {
     params: SessionSteerRequest;
 }
 
+export type SessionRewindExtRequest = {
+    method: typeof SESSION_REWIND_METHOD;
+    params: SessionRewindRequest;
+}
+
+export {SESSION_REWIND_METHOD, type SessionRewindRequest, type SessionRewindResponse} from "./SessionRewind";
+
 export async function steerSessionWithFallback(
     connection: Pick<ClientContext, "request">,
     params: SessionSteerRequest,

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;

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<SessionMetadataWithThread> {
         const additionalDirectories = readAdditionalDirectories(request.cwd, request.additionalDirectories, request._meta);
         await this.refreshSkills(request.cwd, additionalDirectories);

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);

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<ThreadRevertResponse> {
+        return await this.sendRequest({method: "thread/revert", params});
+    }
+
     getThreadSettings(threadId: string): ThreadSettings | undefined {
         return this.threadSettings.get(threadId);
     }

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<SessionRewindResponse> {
+    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];
+}

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"],
                     },
                 },
             },

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();
+    });
+});

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);
 }