diff --git a/src/jsonrpc.ts b/src/jsonrpc.ts index 5f709cc..bf03afe 100644 --- a/src/jsonrpc.ts +++ b/src/jsonrpc.ts @@ -724,6 +724,15 @@ export class RequestResponder { } } +const requestBatchSizes = new WeakMap, number>(); + +/** @internal */ +export function requestBatchSize( + responder: RequestResponder, +): number | undefined { + return requestBatchSizes.get(responder); +} + /** * Disposable handle returned when a handler is registered dynamically. */ @@ -1324,6 +1333,7 @@ export class Connection { const processing = this.receiveMessage( message, isRequestMessage(message) ? collectResponse : undefined, + batch.length, ); if (isNotificationMessage(message)) { void processing.finally(() => { @@ -1337,6 +1347,7 @@ export class Connection { private receiveMessage( message: AnyMessage, sendResponse?: (response: AnyResponse) => Promise, + batchSize?: number, ): Promise { if (this.abortController.signal.aborted) { return Promise.resolve(); @@ -1354,7 +1365,7 @@ export class Connection { this.handleProtocolNotification(message); } return this.processIncomingMessage( - this.toIncomingMessage(message, sendResponse), + this.toIncomingMessage(message, sendResponse, batchSize), ).catch((error) => this.close(error)); } else if ("id" in message) { this.handleResponse(message); @@ -1427,6 +1438,7 @@ export class Connection { private toIncomingMessage( message: AnyRequest | AnyNotification, sendResponse?: (response: AnyResponse) => Promise, + batchSize?: number, ): IncomingMessage { if ("id" in message) { const abortController = new AbortController(); @@ -1437,27 +1449,32 @@ export class Connection { } }; + const responder = new RequestResponder( + message.id, + (result) => { + const response: AnyResponse = { + jsonrpc: "2.0", + id: message.id, + ...result, + }; + return sendResponse + ? sendResponse(response) + : this.sendWireMessage(response); + }, + abortController.signal, + finishRequest, + ); + if (batchSize !== undefined) { + requestBatchSizes.set(responder, batchSize); + } + return { kind: "request", method: message.method, params: message.params, raw: message, signal: abortController.signal, - responder: new RequestResponder( - message.id, - (result) => { - const response: AnyResponse = { - jsonrpc: "2.0", - id: message.id, - ...result, - }; - return sendResponse - ? sendResponse(response) - : this.sendWireMessage(response); - }, - abortController.signal, - finishRequest, - ), + responder, }; } diff --git a/src/v2/acp.test.ts b/src/v2/acp.test.ts index 4a77247..3dee946 100644 --- a/src/v2/acp.test.ts +++ b/src/v2/acp.test.ts @@ -21,6 +21,7 @@ import type { Annotations, ClientContext, DiffPatch, + InitializeRequest, InitializeResponse, McpServer, NewSessionRequest, @@ -33,6 +34,13 @@ import type { const clientInfo = { name: "test-client", version: "1.0.0" }; const agentInfo = { name: "test-agent", version: "1.0.0" }; +function testAgent(): sdk.AgentApp { + return agent().onRequest(methods.agent.initialize, () => ({ + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + })); +} + function assertV2MethodTypes( agentContext: ClientContext, clientContext: AgentContext, @@ -303,7 +311,7 @@ describe("experimental v2 app API", () => { } satisfies NewSessionRequest; const expected = structuredClone(request); - await client().connectWith(agent(), (agentClient) => { + await client().connectWith(testAgent(), (agentClient) => { const builder = agentClient.buildSession(request); additionalDirectories[0] = "/mutated-input"; @@ -368,7 +376,7 @@ describe("experimental v2 app API", () => { } satisfies McpServer; const expected = structuredClone(mcpServer); - await client().connectWith(agent(), (agentClient) => { + await client().connectWith(testAgent(), (agentClient) => { const builder = agentClient .buildSession("/workspace") .withMcpServer(mcpServer); @@ -400,6 +408,577 @@ describe("experimental v2 app API", () => { } }); + it("initializes exactly once, queues later calls, and exposes the exchange", async () => { + const initializeGate = Promise.withResolvers(); + const agentReady = Promise.withResolvers(); + const clientReady = Promise.withResolvers(); + const events: string[] = []; + let newSessionCalls = 0; + let receivedInitialize: InitializeRequest | undefined; + const requestMeta = { nested: { source: "client" } }; + + const initializeRequest = { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + capabilities: { + _meta: requestMeta, + }, + } satisfies InitializeRequest; + const initializeResponse = { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + capabilities: { + session: {}, + _meta: { source: "agent" }, + }, + } satisfies InitializeResponse; + const expectedInitializeRequest = structuredClone(initializeRequest); + + const agentApp = agent() + .onConnect(async (connection) => { + events.push("agent-connect"); + const initialization = await connection.initialized; + events.push("agent-initialized"); + expect(initialization).toMatchObject({ + request: expectedInitializeRequest, + response: initializeResponse, + }); + agentReady.resolve(); + }) + .onRequest(methods.agent.initialize, async ({ params }) => { + receivedInitialize = params; + await initializeGate.promise; + return initializeResponse; + }) + .onRequest(methods.agent.session.new, () => { + newSessionCalls += 1; + return { sessionId: "session-1" }; + }); + const clientApp = client().onConnect(async (connection) => { + events.push("client-connect"); + const initialization = await connection.initialized; + events.push("client-initialized"); + expect(initialization).toMatchObject({ + request: expectedInitializeRequest, + response: initializeResponse, + }); + clientReady.resolve(); + }); + + await clientApp.connectWith(agentApp, async (agentContext) => { + await expect( + agentContext.request(methods.agent.session.new, { + cwd: "/workspace", + mcpServers: [], + }), + ).rejects.toMatchObject({ code: -32600 }); + await expect( + agentContext.notify(methods.agent.session.cancel, { + sessionId: "session-1", + }), + ).rejects.toMatchObject({ code: -32600 }); + await expect( + agentContext.request("_vendor/pre-initialize", {}), + ).rejects.toMatchObject({ code: -32600 }); + expect(() => + agentContext.request(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + _meta: { notCloneable: () => undefined }, + }), + ).toThrow(); + + const initialized = agentContext.request( + methods.agent.initialize, + initializeRequest, + ); + requestMeta.nested.source = "mutated"; + const queuedSession = agentContext.request(methods.agent.session.new, { + cwd: "/workspace", + mcpServers: [], + }); + await Promise.resolve(); + expect(newSessionCalls).toBe(0); + expect(() => + agentContext.request(methods.agent.initialize, initializeRequest), + ).toThrow("Invalid request"); + + initializeGate.resolve(); + await expect(initialized).resolves.toMatchObject(initializeResponse); + await expect(queuedSession).resolves.toEqual({ sessionId: "session-1" }); + + await expect(agentContext.initialized).resolves.toMatchObject({ + request: expectedInitializeRequest, + response: initializeResponse, + }); + expect(receivedInitialize).toMatchObject(expectedInitializeRequest); + expect(events.slice(0, 2)).toEqual(["agent-connect", "client-connect"]); + expect(events).toContain("client-initialized"); + expect(() => + agentContext.request(methods.agent.initialize, initializeRequest), + ).toThrow("Invalid request"); + }); + await Promise.all([agentReady.promise, clientReady.promise]); + expect(events).toContain("agent-initialized"); + }); + + it("supports standalone initialize batches and does not retry failures", async () => { + let initializedConnections = 0; + let receivedInitialize: InitializeRequest | undefined; + const agentInitialized = Promise.withResolvers(); + const validAgent = agent() + .onConnect(async (connection) => { + await connection.initialized; + initializedConnections += 1; + agentInitialized.resolve(); + }) + .onRequest(methods.agent.initialize, ({ params }) => { + receivedInitialize = params; + return { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }; + }); + + const validConnection = client().connect(validAgent); + const agentContext = validConnection.agent; + try { + await expect( + agentContext.batch([ + batchRequest(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }), + batchNotification(methods.agent.session.cancel, { + sessionId: "session-1", + }), + ] as const), + ).rejects.toMatchObject({ code: -32600 }); + await expect( + agentContext.batch([ + batchRequest(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }), + batchRequest(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }), + ] as const), + ).rejects.toMatchObject({ code: -32600 }); + + let metaReads = 0; + const initializeRequest = { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + get _meta() { + metaReads += 1; + return { seen: metaReads }; + }, + } satisfies InitializeRequest; + const mapperError = new Error("initialize mapper failed"); + await expect( + agentContext.batch([ + batchRequest(methods.agent.initialize, initializeRequest, () => { + throw mapperError; + }), + ] as const), + ).rejects.toBe(mapperError); + expect(validConnection.signal.aborted).toBe(false); + validConnection.close(); + await validConnection.closed; + await agentInitialized.promise; + expect(initializedConnections).toBe(1); + expect(metaReads).toBe(1); + expect(receivedInitialize?._meta).toEqual({ seen: 1 }); + await expect(agentContext.initialized).resolves.toMatchObject({ + request: { _meta: { seen: 1 } }, + }); + } finally { + validConnection.close(); + await validConnection.closed; + } + + const invalidAgent = agent().onRequest(methods.agent.initialize, () => ({ + protocolVersion: 1, + info: agentInfo, + })); + const invalidConnection = client().connect(invalidAgent); + try { + const initialized = invalidConnection.agent.request( + methods.agent.initialize, + { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }, + ); + const queued = invalidConnection.agent.request( + methods.agent.session.new, + { + cwd: "/workspace", + mcpServers: [], + }, + ); + const initializeError = await initialized.catch((error) => error); + expect(initializeError).toMatchObject({ code: -32600 }); + await expect(invalidConnection.initialized).rejects.toBe(initializeError); + await expect(queued).rejects.toBe(initializeError); + await expect(invalidConnection.agent.initialized).rejects.toBe( + initializeError, + ); + await invalidConnection.closed; + expect(invalidConnection.signal.reason).toBe(initializeError); + expect(() => + invalidConnection.agent.request(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }), + ).toThrow("Invalid request"); + } finally { + invalidConnection.close(); + await invalidConnection.closed; + } + }); + + it("rejects initialize after a pre-initialize request fails the connection", async () => { + let initializeCalls = 0; + let newSessionCalls = 0; + const agentApp = agent() + .onRequest(methods.agent.initialize, () => { + initializeCalls += 1; + return { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }; + }) + .onRequest(methods.agent.session.new, () => { + newSessionCalls += 1; + return { sessionId: "session-1" }; + }); + const [agentStream, peerStream] = memoryWireStreamPair(); + const connection = agentApp.connect(agentStream); + const writer = peerStream.writable.getWriter(); + const reader = peerStream.readable.getReader(); + + try { + const preInitializeWrite = writer.write({ + jsonrpc: "2.0", + id: 1, + method: methods.agent.session.new, + params: { cwd: "/workspace", mcpServers: [] }, + }); + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: { + id: 1, + error: { code: -32600 }, + }, + }); + await preInitializeWrite; + + const initializeWrite = writer.write({ + jsonrpc: "2.0", + id: 2, + method: methods.agent.initialize, + params: { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }, + }); + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: { + id: 2, + error: { code: -32600 }, + }, + }); + await initializeWrite; + + expect(initializeCalls).toBe(0); + expect(newSessionCalls).toBe(0); + await expect(connection.initialized).rejects.toMatchObject({ + code: -32600, + }); + } finally { + writer.releaseLock(); + reader.releaseLock(); + connection.close(); + await connection.closed; + } + }); + + it("rejects mixed raw initialize batches before dispatching handlers", async () => { + let initializeCalls = 0; + let newSessionCalls = 0; + const agentApp = agent() + .onRequest(methods.agent.initialize, () => { + initializeCalls += 1; + return { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }; + }) + .onRequest(methods.agent.session.new, () => { + newSessionCalls += 1; + return { sessionId: "session-1" }; + }); + const [agentStream, peerStream] = memoryWireStreamPair(); + const connection = agentApp.connect(agentStream); + const writer = peerStream.writable.getWriter(); + const reader = peerStream.readable.getReader(); + + try { + const batchWrite = writer.write([ + { + jsonrpc: "2.0", + id: 1, + method: methods.agent.initialize, + params: { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }, + }, + { + jsonrpc: "2.0", + id: 2, + method: methods.agent.session.new, + params: { cwd: "/workspace", mcpServers: [] }, + }, + ]); + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: expect.arrayContaining([ + expect.objectContaining({ + id: 1, + error: expect.objectContaining({ code: -32600 }), + }), + expect.objectContaining({ + id: 2, + error: expect.objectContaining({ code: -32600 }), + }), + ]), + }); + await batchWrite; + + expect(initializeCalls).toBe(0); + expect(newSessionCalls).toBe(0); + await expect(connection.initialized).rejects.toMatchObject({ + code: -32600, + }); + } finally { + writer.releaseLock(); + reader.releaseLock(); + connection.close(); + await connection.closed; + } + }); + + it("requires an initialize handler before opening a connection", async () => { + const agentApp = agent(); + const [firstStream] = memoryWireStreamPair(); + expect(() => agentApp.connect(firstStream)).toThrow( + "requires an initialize request handler", + ); + + agentApp.onRequest(methods.agent.initialize, () => ({ + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + })); + const [secondStream] = memoryWireStreamPair(); + const connection = agentApp.connect(secondStream); + connection.close(); + await connection.closed; + await expect(connection.initialized).rejects.toThrow( + "ACP connection closed", + ); + }); + + it("treats a cancellation notification before initialize as the first message", async () => { + let initializeCalls = 0; + const agentApp = agent().onRequest(methods.agent.initialize, () => { + initializeCalls += 1; + return { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }; + }); + const [agentStream, peerStream] = memoryWireStreamPair(); + const connection = agentApp.connect(agentStream); + const writer = peerStream.writable.getWriter(); + const reader = peerStream.readable.getReader(); + + try { + await writer.write({ + jsonrpc: "2.0", + method: methods.protocol.cancelRequest, + params: { requestId: 999 }, + }); + const initializeWrite = writer.write({ + jsonrpc: "2.0", + id: 1, + method: methods.agent.initialize, + params: { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }, + }); + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: { id: 1, error: { code: -32600 } }, + }); + await initializeWrite; + expect(initializeCalls).toBe(0); + } finally { + writer.releaseLock(); + reader.releaseLock(); + connection.close(); + await connection.closed; + } + }); + + it("queues raw peer requests behind an in-flight initialize", async () => { + const initializeGate = Promise.withResolvers(); + const initializeStarted = Promise.withResolvers(); + let newSessionCalls = 0; + const agentApp = agent() + .onRequest(methods.agent.initialize, async () => { + initializeStarted.resolve(); + await initializeGate.promise; + return { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }; + }) + .onRequest(methods.agent.session.new, () => { + newSessionCalls += 1; + return { sessionId: "session-1" }; + }); + const [agentStream, peerStream] = memoryWireStreamPair(); + const connection = agentApp.connect(agentStream); + const writer = peerStream.writable.getWriter(); + const reader = peerStream.readable.getReader(); + + try { + const initializeWrite = writer.write({ + jsonrpc: "2.0", + id: 1, + method: methods.agent.initialize, + params: { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }, + }); + await initializeStarted.promise; + + const malformedDuplicateWrite = writer.write({ + jsonrpc: "2.0", + id: 3, + method: methods.agent.initialize, + params: { protocolVersion: "invalid" }, + }); + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: { + id: 3, + error: { code: -32600 }, + }, + }); + await malformedDuplicateWrite; + + await writer.write({ + jsonrpc: "2.0", + method: methods.protocol.cancelRequest, + params: { requestId: 999 }, + }); + const queuedWrite = writer.write({ + jsonrpc: "2.0", + id: 2, + method: methods.agent.session.new, + params: { cwd: "/workspace", mcpServers: [] }, + }); + await Promise.resolve(); + expect(newSessionCalls).toBe(0); + + const initializeResponse = reader.read(); + initializeGate.resolve(); + await expect(initializeResponse).resolves.toMatchObject({ + done: false, + value: { + id: 1, + result: { protocolVersion: PROTOCOL_VERSION }, + }, + }); + await initializeWrite; + + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: { + id: 2, + result: { sessionId: "session-1" }, + }, + }); + await queuedWrite; + expect(newSessionCalls).toBe(1); + + const malformedCancelWrite = writer.write({ + jsonrpc: "2.0", + id: 4, + method: methods.protocol.cancelRequest, + params: { requestId: 999 }, + }); + await expect(reader.read()).resolves.toMatchObject({ + done: false, + value: { id: 4, error: { code: -32601 } }, + }); + await malformedCancelWrite; + } finally { + initializeGate.resolve(); + writer.releaseLock(); + reader.releaseLock(); + connection.close(); + await connection.closed; + } + }); + + it("rejects queued peer calls when an in-flight initialize connection closes", async () => { + const initializeGate = Promise.withResolvers(); + const initializeStarted = Promise.withResolvers(); + const agentApp = agent().onRequest(methods.agent.initialize, async () => { + initializeStarted.resolve(); + await initializeGate.promise; + return { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }; + }); + const [agentStream, peerStream] = memoryWireStreamPair(); + const connection = agentApp.connect(agentStream); + const writer = peerStream.writable.getWriter(); + + try { + const initializeWrite = writer.write({ + jsonrpc: "2.0", + id: 1, + method: methods.agent.initialize, + params: { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }, + }); + await initializeStarted.promise; + await initializeWrite; + + const queuedCall = connection.client.request("_vendor/queued", {}); + connection.close(); + await expect(queuedCall).rejects.toThrow("ACP connection closed"); + } finally { + initializeGate.resolve(); + writer.releaseLock(); + connection.close(); + await connection.closed; + } + }); + it("does not complete a prompt from an idle update received before it", async () => { let updateClient: AgentContext | undefined; @@ -729,13 +1308,24 @@ describe("experimental v2 app API", () => { it("validates every built-in direct response before returning it", async () => { const [clientStream, peerStream] = memoryWireStreamPair(); - const response = client().connectWith(clientStream, (agentContext) => - agentContext.request(methods.agent.session.new, { - cwd: "/workspace", - mcpServers: [], - }), + const response = client().connectWith( + clientStream, + async (agentContext) => { + await agentContext.request(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }); + return agentContext.request(methods.agent.session.new, { + cwd: "/workspace", + mcpServers: [], + }); + }, ); + await respondToNextRequest(peerStream, { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }); await respondToNextRequest(peerStream, { sessionId: 42 }); await expect(response).rejects.toThrow(); }); @@ -743,19 +1333,30 @@ describe("experimental v2 app API", () => { it("validates built-in batch responses before applying caller mappings", async () => { const [clientStream, peerStream] = memoryWireStreamPair(); let mapped = false; - const response = client().connectWith(clientStream, (agentContext) => - agentContext.batch([ - batchRequest( - methods.agent.session.new, - { cwd: "/workspace", mcpServers: [] }, - (session) => { - mapped = true; - return session.sessionId; - }, - ), - ] as const), + const response = client().connectWith( + clientStream, + async (agentContext) => { + await agentContext.request(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }); + return agentContext.batch([ + batchRequest( + methods.agent.session.new, + { cwd: "/workspace", mcpServers: [] }, + (session) => { + mapped = true; + return session.sessionId; + }, + ), + ] as const); + }, ); + await respondToNextRequest(peerStream, { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }); const reader = peerStream.readable.getReader(); const request = await reader.read(); reader.releaseLock(); @@ -786,22 +1387,43 @@ describe("experimental v2 app API", () => { it("rejects peer null for empty responses but preserves local void handlers", async () => { const [clientStream, peerStream] = memoryWireStreamPair(); - const invalidResponse = client().connectWith(clientStream, (agentContext) => - agentContext.request(methods.agent.session.delete, { - sessionId: "session-1", - }), + const invalidResponse = client().connectWith( + clientStream, + async (agentContext) => { + await agentContext.request(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }); + return agentContext.request(methods.agent.session.delete, { + sessionId: "session-1", + }); + }, ); + await respondToNextRequest(peerStream, { + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + }); await respondToNextRequest(peerStream, null); await expect(invalidResponse).rejects.toThrow(); await expect( client().connectWith( - agent().onRequest(methods.agent.session.delete, () => {}), - (agentContext) => - agentContext.request(methods.agent.session.delete, { + agent() + .onRequest(methods.agent.initialize, () => ({ + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + })) + .onRequest(methods.agent.session.delete, () => {}), + async (agentContext) => { + await agentContext.request(methods.agent.initialize, { + protocolVersion: PROTOCOL_VERSION, + info: clientInfo, + }); + return agentContext.request(methods.agent.session.delete, { sessionId: "session-1", - }), + }); + }, ), ).resolves.toEqual({}); }); @@ -813,10 +1435,15 @@ describe("experimental v2 app API", () => { let agentNotificationValue: string | undefined; let clientNotificationValue: string | undefined; + let initializedNotificationValue: string | undefined; + const notificationSent = Promise.withResolvers(); const clientApp = client() .onRequest("authentication/logout", parseValue, returnValue) .onNotification("authentication/status", parseValue, ({ params }) => { clientNotificationValue = params.value; + }) + .onNotification("_vendor/acme/event", parseValue, ({ params }) => { + initializedNotificationValue = params.value; }); const agentApp = agent() .onRequest("session/load", parseValue, async ({ params, client }) => { @@ -832,6 +1459,18 @@ describe("experimental v2 app API", () => { .onNotification("session/set_model", parseValue, ({ params }) => { agentNotificationValue = params.value; }) + .onConnect(async (connection) => { + await connection.initialized; + try { + await connection.client.notify("_vendor/acme/event", { + value: "initialized notification", + }); + notificationSent.resolve(); + } catch (error) { + notificationSent.reject(error); + throw error; + } + }) .onRequest(methods.agent.initialize, () => ({ protocolVersion: PROTOCOL_VERSION, info: agentInfo, @@ -856,6 +1495,8 @@ describe("experimental v2 app API", () => { protocolVersion: PROTOCOL_VERSION, info: clientInfo, }); + await notificationSent.promise; + expect(initializedNotificationValue).toBe("initialized notification"); await expect( agentContext.request< { value: string }, diff --git a/src/v2/acp.ts b/src/v2/acp.ts index 8259fff..8b75d55 100644 --- a/src/v2/acp.ts +++ b/src/v2/acp.ts @@ -122,6 +122,7 @@ import { Connection, Handled, HandlerRegistration, + requestBatchSize, RequestError, } from "../jsonrpc.js"; import type { @@ -391,7 +392,7 @@ function parseV2InitializeRequest(params: unknown): schema.InitializeRequest { `The v2 API only supports protocol version ${schema.PROTOCOL_VERSION}`, ); } - return request; + return structuredClone(request); } function normalizeOutgoingV2InitializeRequest( @@ -418,7 +419,240 @@ function mapV2InitializeResponse(response: unknown): schema.InitializeResponse { `The v2 API only supports protocol version ${schema.PROTOCOL_VERSION}`, ); } - return parsed; + return structuredClone(parsed); +} + +/** + * Validated initialize request and response for one ACP v2 connection. + * + * The snapshot becomes available only after the initialize response has been + * validated and sent or received successfully. + * + * @experimental + */ +export type InitializationSnapshot = Readonly<{ + request: schema.InitializeRequest; + response: schema.InitializeResponse; +}>; + +type InitializationPhase = + "uninitialized" | "initializing" | "initialized" | "failed"; + +function cloneInitialization( + initialization: InitializationSnapshot, +): InitializationSnapshot { + return structuredClone(initialization); +} + +function deferred(): { + promise: Promise; + resolve(value: T | PromiseLike): void; + reject(reason?: unknown): void; +} { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + +class InitializationState { + private phase: InitializationPhase = "uninitialized"; + private request?: schema.InitializeRequest; + private readonly barrier = deferred(); + + constructor() { + void this.barrier.promise.catch(() => {}); + } + + get status(): InitializationPhase { + return this.phase; + } + + get initialized(): Promise { + return this.barrier.promise.then(cloneInitialization); + } + + begin(request: schema.InitializeRequest): void { + if (this.phase !== "uninitialized") { + throw RequestError.invalidRequest( + "ACP v2 initialize may only be requested once per connection", + ); + } + + const requestSnapshot = structuredClone(request); + this.request = requestSnapshot; + this.phase = "initializing"; + } + + complete(response: schema.InitializeResponse): void { + if (this.phase !== "initializing" || !this.request) { + throw RequestError.invalidRequest( + "ACP v2 initialization is not in progress", + ); + } + + const initialization = { + request: structuredClone(this.request), + response: structuredClone(response), + }; + this.phase = "initialized"; + this.barrier.resolve(initialization); + } + + fail(error?: unknown): void { + if (this.phase === "uninitialized" || this.phase === "initializing") { + this.phase = "failed"; + this.request = undefined; + this.barrier.reject( + error ?? + RequestError.invalidRequest( + "ACP v2 connection initialization failed", + ), + ); + } + } + + waitUntilInitialized(method: string): Promise { + if (this.phase === "initialized") { + return Promise.resolve(); + } + if (this.phase === "initializing") { + return this.barrier.promise.then(() => {}); + } + return Promise.reject(initializationUnavailable(method)); + } +} + +const initializationStates = new WeakMap< + ConnectionContext, + InitializationState +>(); + +function initializationState(cx: ConnectionContext): InitializationState { + let state = initializationStates.get(cx); + if (!state) { + state = new InitializationState(); + initializationStates.set(cx, state); + if (cx.signal.aborted) { + state.fail(cx.signal.reason); + } else { + cx.signal.addEventListener("abort", () => state?.fail(cx.signal.reason), { + once: true, + }); + } + } + return state; +} + +function initializationUnavailable(method: string): RequestError { + return RequestError.invalidRequest( + `ACP v2 connection must be initialized before '${method}'`, + ); +} + +async function blockUninitializedIncoming( + message: IncomingMessage, + state: InitializationState, +): Promise { + const error = initializationUnavailable(message.method); + state.fail(error); + return rejectIncoming(message, error); +} + +async function rejectIncoming( + message: IncomingMessage, + error: RequestError, +): Promise { + if (message.kind === "request") { + await message.responder.respondWithError(error); + } + return Handled.yes(); +} + +function agentInitializationGuard(): JsonRpcHandler { + return { + async handleMessage(message, cx) { + const state = initializationState(cx); + if ( + message.kind === "notification" && + message.method === schema.PROTOCOL_METHODS.cancel_request + ) { + return state.status === "initializing" || state.status === "initialized" + ? Handled.yes() + : blockUninitializedIncoming(message, state); + } + + if ( + message.kind === "request" && + message.method === schema.AGENT_METHODS.initialize + ) { + if (state.status !== "uninitialized") { + return rejectIncoming( + message, + RequestError.invalidRequest( + "ACP v2 initialize may only be requested once per connection", + ), + ); + } + const batchSize = requestBatchSize(message.responder); + if (batchSize !== undefined && batchSize !== 1) { + const error = RequestError.invalidRequest( + "ACP v2 initialize must be the only entry in its batch", + ); + state.fail(error); + return rejectIncoming(message, error); + } + + let request: schema.InitializeRequest; + try { + request = parseV2InitializeRequest(message.params); + } catch (error) { + state.fail(error); + throw error; + } + state.begin(request); + return Handled.no(message); + } + + if (state.status === "initialized") { + return Handled.no(message); + } + if (state.status === "initializing") { + await state.waitUntilInitialized(message.method); + return Handled.no(message); + } + return blockUninitializedIncoming(message, state); + }, + describe: () => "agent-initialization", + }; +} + +function clientInitializationGuard(): JsonRpcHandler { + return { + async handleMessage(message, cx) { + const state = initializationState(cx); + if ( + message.kind === "notification" && + message.method === schema.PROTOCOL_METHODS.cancel_request + ) { + return state.status === "initializing" || state.status === "initialized" + ? Handled.yes() + : blockUninitializedIncoming(message, state); + } + if (state.status === "initialized") { + return Handled.no(message); + } + if (state.status === "initializing") { + await state.waitUntilInitialized(message.method); + return Handled.no(message); + } + return blockUninitializedIncoming(message, state); + }, + describe: () => "client-initialization", + }; } function parseRequestResponse( @@ -434,7 +668,7 @@ function normalizeV2Batch( string, { response?: ParamsParser } | undefined >, - normalizeInitialize = false, + onInitializeResponse?: (response: schema.InitializeResponse) => void, ): Entries & { readonly 0: BatchEntry } { return entries.map((entry) => { if (entry.kind !== "request") { @@ -446,13 +680,15 @@ function normalizeV2Batch( ((response: unknown) => unknown) | undefined; return { ...entry, - params: - normalizeInitialize && entry.method === schema.AGENT_METHODS.initialize - ? normalizeOutgoingV2InitializeRequest(entry.params) - : entry.params, mapResponse: spec ? (response: unknown) => { const parsed = parseRequestResponse(spec, response); + if ( + onInitializeResponse && + entry.method === schema.AGENT_METHODS.initialize + ) { + onInitializeResponse(parsed as schema.InitializeResponse); + } return mapResponse ? mapResponse(parsed) : parsed; } : mapResponse, @@ -566,6 +802,12 @@ const startActiveSession = Symbol("startActiveSession"); * @experimental */ export interface AcpConnection { + /** + * Resolves with the validated initialize exchange once initialization + * succeeds, or rejects if initialization or the connection fails. + */ + readonly initialized: Promise; + /** * AbortSignal that aborts when the connection closes. */ @@ -685,6 +927,18 @@ class AcpContext { private readonly currentRequestId?: JsonRpcId, ) {} + /** + * Validated initialize exchange, once this connection is initialized. + */ + get initialized(): Promise { + return initializationState(this.cx).initialized; + } + + /** @internal */ + protected get initializationLifecycle(): InitializationState { + return initializationState(this.cx); + } + /** * JSON-RPC id of the request currently being handled. * @@ -779,12 +1033,16 @@ export class AgentContext extends AcpContext { assertV2MethodDirection(method, clientRequestSpecsByMethod, "request"); const spec = clientRequestSpecsByMethod[method] as AcpRequestSpec | undefined; - return this.sendRequest( - method, - params, - spec ? (response) => parseRequestResponse(spec, response) : undefined, - options, - ); + return this.initializationLifecycle + .waitUntilInitialized(method) + .then(() => + this.sendRequest( + method, + params, + spec ? (response) => parseRequestResponse(spec, response) : undefined, + options, + ), + ); } /** @@ -816,7 +1074,9 @@ export class AgentContext extends AcpContext { "notification", true, ); - return this.sendNotification(method, params); + return this.initializationLifecycle + .waitUntilInitialized(method) + .then(() => this.sendNotification(method, params)); } /** @@ -845,9 +1105,11 @@ export class AgentContext extends AcpContext { clientRequestSpecsByMethod, clientNotificationSpecsByMethod, ); - return this.sendBatch( - normalizeV2Batch(entries, clientRequestSpecsByMethod), - ); + return this.initializationLifecycle + .waitUntilInitialized("batch") + .then(() => + this.sendBatch(normalizeV2Batch(entries, clientRequestSpecsByMethod)), + ); } } @@ -861,13 +1123,21 @@ export class AgentContext extends AcpContext { * @experimental */ export class ClientContext extends AcpContext { - private constructor(cx: ConnectionContext, requestId?: JsonRpcId) { + private constructor( + cx: ConnectionContext, + requestId?: JsonRpcId, + private readonly closeOnInitializationFailure?: (error: unknown) => void, + ) { super(cx, requestId); } /** @internal */ - static create(cx: ConnectionContext, requestId?: JsonRpcId): ClientContext { - return new ClientContext(cx, requestId); + static create( + cx: ConnectionContext, + requestId?: JsonRpcId, + closeOnInitializationFailure?: (error: unknown) => void, + ): ClientContext { + return new ClientContext(cx, requestId, closeOnInitializationFailure); } /** @internal */ @@ -993,16 +1263,47 @@ export class ClientContext extends AcpContext { assertV2MethodDirection(method, agentRequestSpecsByMethod, "request"); const spec = agentRequestSpecsByMethod[method] as AcpRequestSpec | undefined; - const wireParams = - method === schema.AGENT_METHODS.initialize - ? normalizeOutgoingV2InitializeRequest(params) - : params; - return this.sendRequest( - method, - wireParams, - spec ? (response) => parseRequestResponse(spec, response) : undefined, - options, - ); + const state = this.initializationLifecycle; + if (method === schema.AGENT_METHODS.initialize) { + const request = normalizeOutgoingV2InitializeRequest(params); + state.begin(request); + + let response: Promise; + try { + response = this.sendRequest( + method, + request, + (value) => { + const parsed = parseRequestResponse(spec!, value); + state.complete(parsed as schema.InitializeResponse); + return parsed; + }, + options, + ); + } catch (error) { + state.fail(error); + this.closeOnInitializationFailure?.(error); + throw error; + } + void response.catch((error) => { + if (state.status !== "initialized") { + state.fail(error); + this.closeOnInitializationFailure?.(error); + } + }); + return response; + } + + return state + .waitUntilInitialized(method) + .then(() => + this.sendRequest( + method, + params, + spec ? (response) => parseRequestResponse(spec, response) : undefined, + options, + ), + ); } /** @@ -1034,7 +1335,9 @@ export class ClientContext extends AcpContext { "notification", true, ); - return this.sendNotification(method, params); + return this.initializationLifecycle + .waitUntilInitialized(method) + .then(() => this.sendNotification(method, params)); } /** @@ -1063,15 +1366,67 @@ export class ClientContext extends AcpContext { agentRequestSpecsByMethod, agentNotificationSpecsByMethod, ); - return this.sendBatch( - normalizeV2Batch(entries, agentRequestSpecsByMethod, true), + const initializeEntries = entries.filter( + (entry) => + entry.kind === "request" && + entry.method === schema.AGENT_METHODS.initialize, ); + if (initializeEntries.length > 0) { + if (entries.length !== 1 || initializeEntries.length !== 1) { + return Promise.reject( + RequestError.invalidRequest( + "ACP v2 initialize must be the only entry in its batch", + ), + ); + } + + const state = this.initializationLifecycle; + const request = normalizeOutgoingV2InitializeRequest( + initializeEntries[0].params, + ); + state.begin(request); + const normalizedEntries = [ + { ...initializeEntries[0], params: request }, + ] as unknown as Entries & { readonly 0: BatchEntry }; + + let response: Promise>; + try { + response = this.sendBatch( + normalizeV2Batch( + normalizedEntries, + agentRequestSpecsByMethod, + (value) => state.complete(value), + ), + ); + } catch (error) { + state.fail(error); + this.closeOnInitializationFailure?.(error); + throw error; + } + void response.catch((error) => { + if (state.status !== "initialized") { + state.fail(error); + this.closeOnInitializationFailure?.(error); + } + }); + return response; + } + + return this.initializationLifecycle + .waitUntilInitialized("batch") + .then(() => + this.sendBatch(normalizeV2Batch(entries, agentRequestSpecsByMethod)), + ); } } class AcpConnectionHandle implements AcpConnection { constructor(private readonly connection: Connection) {} + get initialized(): Promise { + return initializationState(this.connection.getContext()).initialized; + } + get signal(): AbortSignal { return this.connection.signal; } @@ -1123,7 +1478,11 @@ class ClientConnectionHandle private readonly connectHandlers: readonly ClientConnectHandler[] = [], ) { super(connection); - this.agent = ClientContext.create(connection.getContext()); + this.agent = ClientContext.create( + connection.getContext(), + undefined, + (error) => connection.close(error), + ); } /** @internal */ @@ -1831,19 +2190,32 @@ function registerAppRequest( requestId: JsonRpcId, ) => Context, handler: (context: Context) => MaybePromise, + lifecycle?: { + afterResponse( + params: Params, + response: Response, + cx: ConnectionContext, + ): void; + onError(cx: ConnectionContext, error: unknown): void; + }, ): void { builder.onReceiveRequest( spec.method, (params) => parseParams(spec.params, params), async (params, responder, cx) => { - const response = await handler( - context(params, cx, responder.signal, responder.id), - ); - await responder.respond( - spec.serializeResponse + try { + const response = await handler( + context(params, cx, responder.signal, responder.id), + ); + const sentResponse = spec.serializeResponse ? spec.serializeResponse(response) - : (response as unknown as Response), - ); + : (response as unknown as Response); + await responder.respond(sentResponse); + lifecycle?.afterResponse(params, sentResponse, cx); + } catch (error) { + lifecycle?.onError(cx, error); + throw error; + } }, ); } @@ -2559,6 +2931,15 @@ type ClientConnectionState = { connection: ClientConnection; }; +function bridgeConnectionClose(source: Connection, peer: AcpConnection): void { + void source.closed.then(async () => { + if (initializationState(source.getContext()).status === "initialized") { + await peer.initialized.catch(() => {}); + } + peer.close(source.signal.reason); + }); +} + /** * Creates an agent-side app for the experimental draft ACP v2 API. * @@ -2584,8 +2965,10 @@ export function agent(options?: AppOptions): AgentApp { export class AgentApp { private readonly builder = Connection.builder(); private readonly connectHandlers: AgentConnectHandler[] = []; + private hasInitializeHandler = false; constructor(options: AppOptions = {}) { + this.builder.withHandler(agentInitializationGuard()); if (options.name) { this.builder.name(options.name); } @@ -2593,6 +2976,7 @@ export class AgentApp { /** @internal */ [appBuilder](): ConnectionBuilder { + this.assertInitializeHandler(); return this.builder; } @@ -2699,6 +3083,9 @@ export class AgentApp { ); } + if (method === schema.AGENT_METHODS.initialize) { + this.hasInitializeHandler = true; + } return this.request( spec as AcpRequestSpec, handlerOrParams as AgentRequestHandler, @@ -2768,6 +3155,16 @@ export class AgentApp { requestId, ), handler, + spec.method === schema.AGENT_METHODS.initialize + ? { + afterResponse: (_params, response, cx) => { + initializationState(cx).complete( + response as unknown as schema.InitializeResponse, + ); + }, + onError: (cx, error) => initializationState(cx).fail(error), + } + : undefined, ); return this; } @@ -2790,6 +3187,7 @@ export class AgentApp { target: Stream | ClientApp, options: AppConnectOptions = {}, ): AgentConnectionState { + this.assertInitializeHandler(); if (isStream(target)) { const state = this.openStreamConnection(target); if (!options.deferConnectHandlers) { @@ -2802,8 +3200,8 @@ export class AgentApp { const peerRawConnection = target[appBuilder]().connect(peerStream); const peerConnection = clientConnection(peerRawConnection); const state = this.openStreamConnection(thisStream); - void state.rawConnection.closed.then(() => peerConnection.close()); - void peerRawConnection.closed.then(() => state.connection.close()); + bridgeConnectionClose(state.rawConnection, peerConnection); + bridgeConnectionClose(peerRawConnection, state.connection); try { target[runClientConnectHandlers](peerConnection); this[runAgentConnectHandlers](state.connection); @@ -2822,6 +3220,14 @@ export class AgentApp { connection: agentConnection(rawConnection, this.connectHandlers), }; } + + private assertInitializeHandler(): void { + if (!this.hasInitializeHandler) { + throw new Error( + "AgentApp requires an initialize request handler before connecting", + ); + } + } } /** @@ -2854,6 +3260,7 @@ export class ClientApp { if (options.name) { this.builder.name(options.name); } + this.builder.withHandler(clientInitializationGuard()); this.builder.withHandler({ handleMessage: (message, cx) => sessionUpdateRouter(cx).handleMessage(message), @@ -3062,8 +3469,8 @@ export class ClientApp { const peerRawConnection = target[appBuilder]().connect(peerStream); const peerConnection = agentConnection(peerRawConnection); const state = this.openStreamConnection(thisStream); - void state.rawConnection.closed.then(() => peerConnection.close()); - void peerRawConnection.closed.then(() => state.connection.close()); + bridgeConnectionClose(state.rawConnection, peerConnection); + bridgeConnectionClose(peerRawConnection, state.connection); try { target[runAgentConnectHandlers](peerConnection); this[runClientConnectHandlers](state.connection);