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..54df9e92 --- /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. 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 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/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..0c3c10c4 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,12 @@ export class CodexAcpClient { }); } + async rewindSession(request: SessionRewindRequest): Promise<{rewound: boolean}> { + const response = await runRewindSession(request, this.codexClient); + await this.waitForSessionNotifications(request.sessionId); + return response; + } + 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..3172f4de 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, @@ -78,6 +79,8 @@ import { type LegacySetSessionModelRequest, type LegacySetSessionModelResponse, SESSION_STEERING_METHOD, + SESSION_REWIND_METHOD, + type SessionRewindRequest, type SessionSteeringResponse, type SessionSteerRequest, } from "./AcpExtensions"; @@ -134,6 +137,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 +398,7 @@ export class CodexAcpServer { AIR_NATIVE_SUBAGENT_SESSIONS_KEY, AIR_ASYNC_TASKS_KEY, AIR_RECOMMENDED_CONFIG_VALUE_KEY, + AIR_SESSION_REWIND_KEY, ], }, }, @@ -429,6 +434,19 @@ export class CodexAcpServer { ), }; } + case SESSION_REWIND_METHOD: + 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), + ); case GOAL_CONTROL_METHOD: case LEGACY_GOAL_CONTROL_METHOD: { const sessionState = this.sessions.get(methodRequest.params.sessionId); @@ -2286,7 +2304,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)); } @@ -2349,34 +2367,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 daa7e875..39364c96 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -60,6 +60,10 @@ import type { ThreadTurnsListResponse, ThreadResumeParams, ThreadResumeResponse, + ThreadRevertParams, + ThreadRevertResponse, + ThreadRollbackParams, + ThreadRollbackResponse, ThreadSettings, ThreadStartParams, ThreadStartResponse, @@ -558,6 +562,14 @@ 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}); + } + + 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 new file mode 100644 index 00000000..a905d2a8 --- /dev/null +++ b/src/SessionRewind.ts @@ -0,0 +1,69 @@ +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"; + +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, turnIndex) => turn.items + .flatMap((item, itemIndex) => item.type === "userMessage" + ? [{turn, turnIndex, item, itemIndex}] + : [])); + const candidates = messageIdCandidates(request.beforeMessage.messageId); + const exact = candidates + .map(candidate => userTurns.find(({item}) => item.id === candidate)) + .find(match => match !== undefined); + const match = exact ?? userTurns.filter(({item}) => + fingerprint(userInputVisibleText(item.content)) === request.beforeMessage.messageFingerprint, + )[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}`, + ); + } + 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 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/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__/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__/CodexACPAgent/providers.test.ts b/src/__tests__/CodexACPAgent/providers.test.ts index 5a724039..31428894 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,61 @@ 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("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 new file mode 100644 index 00000000..bbc43e01 --- /dev/null +++ b/src/__tests__/SessionRewind.test.ts @@ -0,0 +1,254 @@ +import {describe, expect, it, vi} from "vitest"; +import type {CodexAppServerClient} from "../CodexAppServerClient"; +import {CodexAcpClient} from "../CodexAcpClient"; +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("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("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({ + 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: []}}), + 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(); + }); + + 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); + }); +}); 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); }