From c3ff3438844f5249d6a7f5c297906e2cd3d5fa7f Mon Sep 17 00:00:00 2001 From: nikita-ashihmin Date: Fri, 28 Aug 2026 12:33:14 +0400 Subject: [PATCH] feat: support message-specific ACP session forks (#1046) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - keep the standard latest-turn ACP fork behavior - accept an optional AIR fork point at `_meta.jetbrains.air.fork` - translate the ACP `messageId` to the Claude transcript UUID - use the current Agent SDK `forkSession(..., { upToMessageId })` API - resume the new SDK session with the same working directory and additional directories No Agent SDK dependency update is required. Version `0.3.238` already provides the required API. ## Tests - `npm run build` - `npm run lint` - `env -u ANTHROPIC_BASE_URL npm run test:run` — 981 passed, 27 skipped --- src/acp-agent.ts | 22 +++------- src/fork-session.ts | 80 +++++++++++++++++++++++++++++++++++++ src/tests/acp-agent.test.ts | 53 ++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 17 deletions(-) create mode 100644 src/fork-session.ts diff --git a/src/acp-agent.ts b/src/acp-agent.ts index 4cad440e..cd2fff2f 100644 --- a/src/acp-agent.ts +++ b/src/acp-agent.ts @@ -138,6 +138,7 @@ import { refusalFallbackResultFromResponse, refusalFallbackToCreateRequest, } from "./elicitation.js"; +import { forkSession } from "./fork-session.js"; import { ALLOW_BYPASS, resolvePermissionMode } from "./permissions/modes.js"; import { normalizeDurablePermissionChangeSet } from "./permissions/normalization.js"; import { buildClaudePermissionOptions } from "./permissions/options.js"; @@ -1771,23 +1772,10 @@ export class ClaudeAcpAgent { async unstable_forkSession(params: ForkSessionRequest): Promise { if (this.providerUpdate) await this.providerUpdate; - const response = await this.createSession( - { - cwd: params.cwd, - mcpServers: params.mcpServers ?? [], - additionalDirectories: params.additionalDirectories, - _meta: params._meta, - }, - { - resume: params.sessionId, - forkSession: true, - }, - ); - // Needs to happen after we return the session - setTimeout(() => { - this.sendAvailableCommandsUpdate(response.sessionId); - }, 0); - return response; + return forkSession(params, { + liveMessageIdToUuid: this.sessions[params.sessionId]?.messageIdToUuid, + messageIdForGrouping, + }); } async resumeSession(params: ResumeSessionRequest): Promise { diff --git a/src/fork-session.ts b/src/fork-session.ts new file mode 100644 index 00000000..a9686980 --- /dev/null +++ b/src/fork-session.ts @@ -0,0 +1,80 @@ +import { ForkSessionRequest, ForkSessionResponse, RequestError } from "@agentclientprotocol/sdk"; +import { + forkSession as forkClaudeSession, + getSessionMessages, +} from "@anthropic-ai/claude-agent-sdk"; + +type ForkSessionMeta = { + [key: string]: unknown; + jetbrains?: { + air?: { + fork?: { + version?: number; + messageId?: string; + }; + }; + }; +}; + +type ForkSessionDependencies = { + liveMessageIdToUuid?: ReadonlyMap; + messageIdForGrouping: (message: { + type?: string; + uuid?: string | null; + message?: unknown; + }) => string | undefined; +}; + +function forkPointMessageId(meta: unknown): string | undefined { + const fork = (meta as ForkSessionMeta | null | undefined)?.jetbrains?.air?.fork; + if (fork?.version !== 1) return undefined; + const messageId = fork.messageId?.trim(); + return messageId || undefined; +} + +function forkPointMessageIdCandidates(messageId: string): string[] { + // Older AIR builds sent their visible segment id. Prefer the exact id before its ACP source id. + const protocolMessageId = messageId.replace(/:segment:\d+$/, ""); + return protocolMessageId === messageId ? [messageId] : [messageId, protocolMessageId]; +} + +export async function forkSession( + params: ForkSessionRequest, + dependencies: ForkSessionDependencies, +): Promise { + const messageId = forkPointMessageId(params._meta); + if (!messageId) { + const forked = await forkClaudeSession(params.sessionId, { dir: params.cwd }); + return { sessionId: forked.sessionId }; + } + + const candidateIds = forkPointMessageIdCandidates(messageId); + const liveUuid = candidateIds + .map((candidateId) => dependencies.liveMessageIdToUuid?.get(candidateId)) + .find(Boolean); + const history = liveUuid + ? undefined + : await getSessionMessages(params.sessionId, { dir: params.cwd }); + const messageUuid = + liveUuid ?? + candidateIds + .map( + (candidateId) => + history?.find((message) => dependencies.messageIdForGrouping(message) === candidateId) + ?.uuid, + ) + .find(Boolean); + + if (!messageUuid) { + throw RequestError.invalidParams( + { messageId }, + `Fork point message ${messageId} was not found in session ${params.sessionId}`, + ); + } + + const forked = await forkClaudeSession(params.sessionId, { + dir: params.cwd, + upToMessageId: messageUuid, + }); + return { sessionId: forked.sessionId }; +} diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts index 98a21243..894b7f54 100644 --- a/src/tests/acp-agent.test.ts +++ b/src/tests/acp-agent.test.ts @@ -52,6 +52,7 @@ import { SessionTitles } from "../session-titles.js"; import { Pushable } from "../utils.js"; import { deleteSession, + forkSession, getSessionInfo, getSessionMessages, PermissionUpdate, @@ -78,6 +79,7 @@ vi.mock("@anthropic-ai/claude-agent-sdk", async (importOriginal) => { return { ...actual, deleteSession: vi.fn(), + forkSession: vi.fn(), getSessionInfo: vi.fn(), // Delegates to the real implementation so integration tests that read // actual transcripts keep working; unit tests override per-call with @@ -7636,6 +7638,57 @@ describe("logout", () => { }); }); +describe("session/fork", () => { + it("forks the latest turn in the requested workspace", async () => { + const client = { sessionUpdate: async () => {} } as unknown as AcpClient; + const agent = new ClaudeAcpAgent(client, { log: () => {}, error: () => {} }); + vi.mocked(forkSession).mockResolvedValueOnce({ sessionId: "fork-id" }); + + const response = await agent.unstable_forkSession({ + sessionId: "source-id", + cwd: "/workspace", + additionalDirectories: ["/workspace/extra"], + mcpServers: [], + }); + + expect(response.sessionId).toBe("fork-id"); + expect(forkSession).toHaveBeenCalledWith("source-id", { dir: "/workspace" }); + }); + + it("forks at an AIR message id through the current SDK", async () => { + const client = { sessionUpdate: async () => {} } as unknown as AcpClient; + const agent = new ClaudeAcpAgent(client, { log: () => {}, error: () => {} }); + vi.mocked(getSessionMessages).mockResolvedValueOnce([ + { + type: "assistant", + uuid: "assistant-uuid", + session_id: "source-id", + message: { id: "msg_123", role: "assistant", content: [] }, + parent_tool_use_id: null, + parent_agent_id: null, + }, + ]); + vi.mocked(forkSession).mockResolvedValueOnce({ sessionId: "fork-id" }); + const meta = { + jetbrains: { air: { fork: { version: 1, messageId: "msg_123:segment:0" } } }, + }; + const response = await agent.unstable_forkSession({ + sessionId: "source-id", + cwd: "/workspace", + additionalDirectories: ["/workspace/extra"], + mcpServers: [], + _meta: meta, + }); + + expect(response.sessionId).toBe("fork-id"); + expect(getSessionMessages).toHaveBeenCalledWith("source-id", { dir: "/workspace" }); + expect(forkSession).toHaveBeenCalledWith("source-id", { + dir: "/workspace", + upToMessageId: "assistant-uuid", + }); + }); +}); + describe("session/close", () => { function createMockAgent() { const mockClient = {