feat: support message-specific ACP session forks (#1046)

## 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

@@ -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<ForkSessionResponse> {
     if (this.providerUpdate) await this.providerUpdate;
-    const response = await this.createSession(
+    return forkSession(params, {
-      {
+      liveMessageIdToUuid: this.sessions[params.sessionId]?.messageIdToUuid,
-        cwd: params.cwd,
+      messageIdForGrouping,
-        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;
   }
 
   async resumeSession(params: ResumeSessionRequest): Promise<ResumeSessionResponse> {

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<string, string>;
+  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<ForkSessionResponse> {
+  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 };
+}

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 = {