diff --git a/docs/air-extensions.md b/docs/air-extensions.md index c59f7877..0aaa1f1d 100644 --- a/docs/air-extensions.md +++ b/docs/air-extensions.md @@ -29,6 +29,7 @@ It describes only this adapter. - [Session index](#session-index) - [Context compaction](#context-compaction) - [Question custom answers](#question-custom-answers) +- [Load replay end](#load-replay-end) - [Session fork point](#session-fork-point) - [Presentation hints](#presentation-hints) - [Removed keys](#removed-keys) @@ -158,6 +159,7 @@ The response to a client that is not AIR has no `_meta.jetbrains` key. | `sessionRename` | with `sessionIndex` | Nothing on its own: the agent lists it to tell that it supports `_session/rename`. | [Rename](#rename) | | `sessionListSubscribe` | with `sessionIndex` | Nothing on its own: the agent lists it to tell that it supports `_session/list/subscribe`. | [List subscription](#list-subscription) | | `sessionWorktrees` | with `sessionIndex` | Nothing on its own: the agent lists it to tell that rows carry `agentWorktree` for a worktree Claude Code made. | [List response](#list-response) | +| `loadReplayEnd` | no | Sends one `session_info_update` after the whole history replay of a loaded session. | [Load replay end](#load-replay-end) | | `rawInputRendering` | no | Sends no display copy of readable input in `content`. The client renders `rawInput`. | [Tool call contract](#tool-call-contract) | The adapter ignores `planContentDelta`. @@ -193,6 +195,7 @@ The adapter sends these keys only to an AIR client. "AIR" in the Gate column mea | `kind` | session mode `_meta.jetbrains.air` and mode config option value `_meta.jetbrains.air` | `standard`, `plan`, `auto_review`, or `full_access` | AIR | | `kind` | available command `_meta.jetbrains.air` of a skill | `skill` | AIR, when the skill's `SKILL.md` is found; not for a built-in command or an MCP prompt | | `skillPath` | available command `_meta.jetbrains.air` of a skill | absolute path of the skill's `SKILL.md` | AIR, together with `kind: "skill"` | +| `loadReplayEnd` | `session_info_update._meta.jetbrains.air` | `true` | `loadReplayEnd` | | `fork` | `session/fork` request `_meta.jetbrains.air` (client to agent) | `{version: 1, messageId, messageFingerprint?, messageOccurrence?}` | none | ## JetBrains shared keys @@ -1576,6 +1579,32 @@ For AIR, the companion field also carries the same value under the root key `_me Released AIR versions read only that key. A client that is not AIR gets neither key. +## Load replay end + +ACP has no signal that the history replay of `session/load` is over, other than the response. +The adapter replays the transcript while the Claude Code CLI starts, so the history can arrive long before the response. + +The client declares `loadReplayEnd` in its capabilities. The agent does not advertise it. +After the history replay of the loaded session, the agent then sends one more update for that session: + +```json +{ + "sessionId": "", + "update": { + "sessionUpdate": "session_info_update", + "_meta": { "jetbrains": { "air": { "version": 1, "loadReplayEnd": true } } } + } +} +``` + +- The update has no `title` and no `updatedAt`. A `null` there would clear the value. +- It comes after every replayed update, the subagent history that the replay sends to the child sessions included, and before the `session/load` response. +- It comes also for an empty history. +- It comes for `session/load` of a new or an already live session, and for the ACP v2 `session/resume` with `replayFrom: {type: "start"}`. +- It does not come when the replay stops or fails. The replay stops when the CLI start fails during the replay. +- It does not come for a live update, and it does not come to a client without the capability. +- The `session/load` can still fail after it, for example when the CLI start fails after the replay. The event means that the recorded history arrived whole. The client decides what to do with the failed load. + ## Session fork point AIR can fork a session at one agent message. diff --git a/src/acp-agent.ts b/src/acp-agent.ts index a8472173..213074b9 100644 --- a/src/acp-agent.ts +++ b/src/acp-agent.ts @@ -132,6 +132,8 @@ import { AIR_PLAN_FILE_CAPABILITY, AIR_GOAL_KEY, AIR_KIND_KEY, + AIR_LOAD_REPLAY_END_CAPABILITY, + AIR_LOAD_REPLAY_END_KEY, AIR_RECOMMENDED_CONFIG_VALUE_CAPABILITY, AIR_SKILL_PATH_KEY, clientSupportsAirCapability, @@ -9524,6 +9526,22 @@ export class ClaudeAcpAgent { this.logger.log( `[session/replay] sessionId=${sessionId} phase=publish durationMs=${Math.round(replayFinishedAt - historyLoadedAt)} totalMs=${Math.round(replayFinishedAt - replayStartedAt)} messages=${messages.length}`, ); + // A client with `loadReplayEnd` gets one more update after the whole + // replay, the subagent history included. A stopped replay (the CLI start + // failed) does not get it, and a failed replay throws before this line. + if ( + !pending?.stopped && + clientSupportsAirCapability(this.clientCapabilities, AIR_LOAD_REPLAY_END_CAPABILITY) + ) { + await this.client.sessionUpdate({ + sessionId, + // No `title` and no `updatedAt`: `null` would clear them. + update: { + sessionUpdate: "session_info_update", + _meta: withAirMeta(undefined, AIR_LOAD_REPLAY_END_KEY, true), + }, + }); + } } async readTextFile(params: ReadTextFileRequest): Promise { diff --git a/src/air-extension.ts b/src/air-extension.ts index 5d9731fc..d8c86aa4 100644 --- a/src/air-extension.ts +++ b/src/air-extension.ts @@ -40,6 +40,14 @@ export const AIR_SESSION_LIST_SUBSCRIBE_CAPABILITY = "sessionListSubscribe"; * {@link AIR_SESSION_INDEX_CAPABILITY} is; the client does not declare it. */ export const AIR_SESSION_WORKTREES_CAPABILITY = "sessionWorktrees"; +/** + * The client wants to know when the history replay of a loaded session is + * whole. After the last replayed update, the agent sends one + * `session_info_update` with `_meta.jetbrains.air.loadReplayEnd: true` (see + * {@link AIR_LOAD_REPLAY_END_KEY}). The client declares it; the agent does not + * advertise it. + */ +export const AIR_LOAD_REPLAY_END_CAPABILITY = "loadReplayEnd"; /** The `_meta.jetbrains.air` keys that the ACP tool call contract defines. */ export const AIR_COMMAND_TITLE_KEY = "commandTitle"; @@ -53,6 +61,8 @@ export const AIR_KIND_KEY = "kind"; export const AIR_SKILL_PATH_KEY = "skillPath"; export const AIR_PERMISSION_KEY = "permission"; export const AIR_CUSTOM_ANSWER_KEY = "customAnswer"; +/** The `_meta.jetbrains.air` key of the update that ends a whole history replay. */ +export const AIR_LOAD_REPLAY_END_KEY = "loadReplayEnd"; /** * The root `_meta` spelling of {@link AIR_CUSTOM_ANSWER_KEY} that released * AIR versions read. Remove it when the ACP registry no longer serves AIR diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts index ed0b2522..634da036 100644 --- a/src/tests/acp-agent.test.ts +++ b/src/tests/acp-agent.test.ts @@ -3656,7 +3656,14 @@ describe("subagent transcript replay", () => { } as unknown as AcpClient, { log: () => {}, error: () => {} }, ); - await initializeClient(agent, AIR_NATIVE_SUBAGENTS); + // With `loadReplayEnd`, the end event comes after the child history too. + await initializeClient(agent, { + _meta: { + jetbrains: { + air: { version: 1, capabilities: ["nativeSubagentSessions", "loadReplayEnd"] }, + }, + }, + } as ClientCapabilities); vi.mocked(getSessionInfo).mockResolvedValueOnce({ cwd: "/tmp/proj" } as any); vi.mocked(getSubagentMessages).mockClear(); await ( @@ -3688,7 +3695,11 @@ describe("subagent transcript replay", () => { [sessionId, "agent_message_chunk", "Done."], [sessionId, "subagent_state_update", child2], [sessionId, "subagent_state_update", child1], + [sessionId, "session_info_update", undefined], ]); + expect(updates.at(-1)!.update._meta).toEqual({ + jetbrains: { air: { version: 1, loadReplayEnd: true } }, + }); expect( updates.flatMap(({ update }) => update.sessionUpdate === "subagent_state_update" ? [(update as any).state] : [], diff --git a/src/tests/acp-v2.test.ts b/src/tests/acp-v2.test.ts index 301b386d..6ca63216 100644 --- a/src/tests/acp-v2.test.ts +++ b/src/tests/acp-v2.test.ts @@ -1091,6 +1091,55 @@ describe("ACP v2 prompts", () => { ]); }); + it("ends a resume with replay with one end event for a client with loadReplayEnd", async () => { + const client = v2Client(); + const isEnd = (notification: v2.UpdateSessionNotification) => + (notification.update._meta as any)?.jetbrains?.air?.loadReplayEnd === true; + let atResponse: v2.UpdateSessionNotification[] = []; + let resumedSessionId = ""; + await client.app.connectWith(connectRouter(), async (agent) => { + await initializeV2(agent, { + _meta: { jetbrains: { air: { version: 1, capabilities: ["loadReplayEnd"] } } }, + } as v2.ClientCapabilities); + const { sessionId } = await agent.request(v2.methods.agent.session.new, { cwd }); + resumedSessionId = sessionId; + vi.mocked(getSessionMessages).mockResolvedValueOnce([ + transcriptEntry("user", randomUUID(), { role: "user", content: "Say hello" }), + transcriptEntry("assistant", randomUUID(), { + id: "msg_hello", + role: "assistant", + content: [{ type: "text", text: "Hello!" }], + }), + ] as unknown as Awaited>); + client.sessionUpdates.length = 0; + await agent.request(v2.methods.agent.session.resume, { + sessionId, + cwd, + replayFrom: { type: "start" }, + }); + atResponse = [...client.sessionUpdates]; + // A resume without a replay sends no end event. + await agent.request(v2.methods.agent.session.resume, { sessionId, cwd }); + await client.authUpdate(1); + }); + + expect(atResponse.map(({ update }) => update.sessionUpdate)).toEqual([ + "user_message", + "user_message_chunk", + "agent_message", + "agent_message_chunk", + "session_info_update", + ]); + expect(atResponse.at(-1)).toEqual({ + sessionId: resumedSessionId, + update: { + sessionUpdate: "session_info_update", + _meta: { jetbrains: { air: { version: 1, loadReplayEnd: true } } }, + }, + }); + expect(client.sessionUpdates.filter(isEnd)).toHaveLength(1); + }); + it("replays the plan of an approved ExitPlanMode as a markdown plan", async () => { const client = v2Client(); await client.app.connectWith(connectRouter(), async (agent) => { diff --git a/src/tests/create-session-options.test.ts b/src/tests/create-session-options.test.ts index c2fffbb4..4f56bf8c 100644 --- a/src/tests/create-session-options.test.ts +++ b/src/tests/create-session-options.test.ts @@ -1882,5 +1882,167 @@ describe("createSession options merging", () => { expect(updates).toEqual([]); expect(getSessionMessages).not.toHaveBeenCalled(); }); + + describe("load replay end", () => { + const LOAD_REPLAY_END_CLIENT: ClientCapabilities = { + _meta: { jetbrains: { air: { version: 1, capabilities: ["loadReplayEnd"] } } }, + }; + + async function endAgent( + updates: SessionNotification[], + sessionUpdate?: (notification: SessionNotification) => Promise, + clientCapabilities: ClientCapabilities = LOAD_REPLAY_END_CLIENT, + ) { + const endAgent = new ClaudeAcpAgent({ + ...createMockClient(), + sessionUpdate: + sessionUpdate ?? + (async (notification: SessionNotification) => { + updates.push(notification); + }), + } as AcpClient); + await endAgent.initialize({ protocolVersion: 1, clientCapabilities }); + return endAgent; + } + + const isEnd = (notification: SessionNotification) => + (notification.update._meta as any)?.jetbrains?.air?.loadReplayEnd === true; + + /** The update kinds, with the end event as "end". */ + const kinds = (updates: SessionNotification[]) => + updates.map((update) => (isEnd(update) ? "end" : update.update.sessionUpdate)); + + it("sends one end event after the replay and before the response", async () => { + sessionMessages = [userMessage("u1", "first"), userMessage("u2", "second")]; + const updates: SessionNotification[] = []; + const loadAgent = await endAgent(updates); + let atResponse: string[] | undefined; + await loadAgent + .loadSession({ sessionId: "replay-probe", cwd: process.cwd(), mcpServers: [] }) + .then(() => (atResponse = kinds(updates))); + // `afterSetupResponse` sends the commands after the response. + await vi.waitFor(() => expect(kinds(updates)).toContain("available_commands_update")); + + expect(atResponse).toEqual(["user_message_chunk", "user_message_chunk", "end"]); + expect(updates.filter(isEnd)).toHaveLength(1); + const end = updates.find(isEnd)!; + expect(end.sessionId).toBe("replay-probe"); + // No `title` and no `updatedAt` key: a `null` would clear them. + expect(end.update).toEqual({ + sessionUpdate: "session_info_update", + _meta: { jetbrains: { air: { version: 1, loadReplayEnd: true } } }, + }); + }); + + it("sends the end event before the CLI start finishes", async () => { + let releaseInit!: () => void; + initGate = new Promise((resolve) => (releaseInit = resolve)); + sessionMessages = [userMessage("u1", "first")]; + const updates: SessionNotification[] = []; + const loading = (await endAgent(updates)).loadSession({ + sessionId: "replay-probe", + cwd: process.cwd(), + mcpServers: [], + }); + + await vi.waitFor(() => expect(kinds(updates)).toEqual(["user_message_chunk", "end"])); + + releaseInit(); + await loading; + }); + + it("sends the end event for an empty history", async () => { + sessionMessages = []; + const updates: SessionNotification[] = []; + let atResponse: string[] | undefined; + await ( + await endAgent(updates) + ) + .loadSession({ sessionId: "replay-probe", cwd: process.cwd(), mcpServers: [] }) + .then(() => (atResponse = kinds(updates))); + + expect(atResponse).toEqual(["end"]); + }); + + it("sends the end event after the replay of a session that is already live", async () => { + const updates: SessionNotification[] = []; + const loadAgent = await endAgent(updates); + const { sessionId } = await loadAgent.newSession({ cwd: process.cwd(), mcpServers: [] }); + sessionMessages = [userMessage("u1", "first"), userMessage("u2", "second")]; + await new Promise((resolve) => setTimeout(resolve, 0)); + updates.length = 0; + + let atResponse: string[] | undefined; + await loadAgent + .loadSession({ sessionId, cwd: process.cwd(), mcpServers: [] }) + .then(() => (atResponse = kinds(updates))); + + expect(atResponse).toEqual(["user_message_chunk", "user_message_chunk", "end"]); + }); + + it("sends no end event when the CLI start fails part way through the replay", async () => { + let failInit!: () => void; + initGate = new Promise((resolve) => (failInit = resolve)); + initError = new Error("Claude Code process exited with code 1"); + sessionMessages = [ + userMessage("u1", "first"), + userMessage("u2", "second"), + userMessage("u3", "third"), + ]; + const updates: SessionNotification[] = []; + let releaseClient!: () => void; + const clientGate = new Promise((resolve) => (releaseClient = resolve)); + const loading = ( + await endAgent(updates, async (notification) => { + updates.push(notification); + await clientGate; + }) + ).loadSession({ sessionId: "replay-probe", cwd: process.cwd(), mcpServers: [] }); + const failed = expect(loading).rejects.toThrow("exited with code 1"); + + // The client still handles the first replayed update when the CLI start fails. + await vi.waitFor(() => expect(updates).toHaveLength(1)); + failInit(); + await vi.waitFor(() => expect(closeSpy).toHaveBeenCalled()); + releaseClient(); + await failed; + for (let i = 0; i < 5; i++) await new Promise((resolve) => setImmediate(resolve)); + + expect(updates.length).toBeLessThan(3); + expect(updates.some(isEnd)).toBe(false); + }); + + it("sends no end event when the replay fails", async () => { + const updates: SessionNotification[] = []; + const loadAgent = await endAgent(updates); + const { sessionId } = await loadAgent.newSession({ cwd: process.cwd(), mcpServers: [] }); + await new Promise((resolve) => setTimeout(resolve, 0)); + updates.length = 0; + sessionMessagesResult = async () => { + throw new Error("history unreadable"); + }; + + await expect( + loadAgent.loadSession({ sessionId, cwd: process.cwd(), mcpServers: [] }), + ).rejects.toThrow("history unreadable"); + expect(updates.some(isEnd)).toBe(false); + }); + + it("sends no end event to a client without the capability", async () => { + sessionMessages = [userMessage("u1", "first")]; + for (const clientCapabilities of [ + {}, + { _meta: { jetbrains: { air: { version: 1, capabilities: [] } } } }, + ] as ClientCapabilities[]) { + const updates: SessionNotification[] = []; + await ( + await endAgent(updates, undefined, clientCapabilities) + ).loadSession({ sessionId: "replay-probe", cwd: process.cwd(), mcpServers: [] }); + await vi.waitFor(() => expect(kinds(updates)).toContain("available_commands_update")); + + expect(kinds(updates)).toEqual(["user_message_chunk", "available_commands_update"]); + } + }); + }); }); });