diff --git a/src/v2/acp.test.ts b/src/v2/acp.test.ts index 4ec569da..4a772470 100644 --- a/src/v2/acp.test.ts +++ b/src/v2/acp.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, expectTypeOf, it } from "vitest"; import { PROTOCOL_VERSION, @@ -19,10 +19,12 @@ import { import type { AgentContext, Annotations, + ClientContext, DiffPatch, InitializeResponse, McpServer, NewSessionRequest, + NewSessionResponse, SessionInfo, SessionInfoUpdate, SessionUpdate, @@ -31,6 +33,148 @@ import type { const clientInfo = { name: "test-client", version: "1.0.0" }; const agentInfo = { name: "test-agent", version: "1.0.0" }; +function assertV2MethodTypes( + agentContext: ClientContext, + clientContext: AgentContext, +): void { + // @ts-expect-error Built-in methods must not fall through the extension overload. + agentContext.request(methods.agent.session.new, { sessionId: "wrong" }); + // @ts-expect-error Built-in notifications must not fall through the extension overload. + agentContext.notify(methods.agent.session.cancel, {}); + agent().onRequest( + // @ts-expect-error Built-in handlers cannot replace their generated params parser. + methods.agent.session.new, + (params: unknown) => params, + () => ({ sessionId: "wrong-parser" }), + ); + // @ts-expect-error Request methods cannot be sent as notifications. + agentContext.notify(methods.agent.session.new, {}); + // @ts-expect-error Notification methods cannot be sent as requests. + agentContext.request(methods.agent.session.cancel, {}); + // @ts-expect-error Client-directed methods cannot be sent to an agent. + agentContext.request(methods.client.mcp.disconnect, { + connectionId: "connection-1", + }); + + const parseValue = (params: unknown): { value: string } => + params as { value: string }; + void agentContext.request("session/load", { sessionId: "session-1" }); + void agentContext.request< + { value: string }, + { sessionId: string }, + "session/load" + >("session/load", { sessionId: "session-1" }); + void agentContext.request<{ value: string }, { value: string }>( + "_vendor/acme/echo", + { value: "request" }, + ); + void agentContext.notify("session/set_model", { modelId: "model-1" }); + agent().onRequest("session/load", parseValue, ({ params }) => params); + agent().onNotification("session/set_model", parseValue, () => {}); + + const dynamicMethod: string = "method/from-another-draft"; + void agentContext.request(dynamicMethod, {}); + void agentContext.notify(dynamicMethod, {}); + agent().onRequest(dynamicMethod, parseValue, ({ params }) => params); + agent().onNotification(dynamicMethod, parseValue, () => {}); + + const outputs = agentContext.batch([ + batchRequest(methods.agent.session.new, { + cwd: "/workspace", + mcpServers: [], + }), + batchNotification(methods.agent.session.cancel, { + sessionId: "session-1", + }), + ] as const); + expectTypeOf(outputs).toEqualTypeOf>(); + void agentContext.notify(methods.protocol.cancelRequest, { requestId: 1 }); + + const clientRequest = batchRequest(methods.client.mcp.disconnect, { + connectionId: "connection-1", + }); + // @ts-expect-error Notification methods cannot be used as batch requests. + batchRequest(methods.agent.session.cancel, { sessionId: "session-1" }); + // @ts-expect-error Request methods cannot be used as batch notifications. + batchNotification(methods.agent.session.new, { + cwd: "/workspace", + mcpServers: [], + }); + // @ts-expect-error Client-directed methods cannot be sent in an agent-directed batch. + agentContext.batch([clientRequest] as const); + agentContext.batch([ + // @ts-expect-error Raw built-in batch entries must use method-specific params. + { + kind: "request", + method: methods.agent.session.new, + params: { sessionId: "wrong" }, + }, + ] as const); + agentContext.batch([ + batchRequest("session/load", { sessionId: "session-1" }), + batchNotification("session/set_model", { modelId: "model-1" }), + ] as const); + const legacyAgentEntry: sdk.AgentBatchEntry = batchRequest("session/load", { + sessionId: "session-1", + }); + const legacyClientEntry: sdk.ClientBatchEntry = batchNotification( + "authentication/status", + {}, + ); + void legacyAgentEntry; + void legacyClientEntry; + agentContext.batch([ + { kind: "request", method: dynamicMethod, params: {} }, + { kind: "notification", method: dynamicMethod, params: {} }, + ] as const); + + clientContext.batch([clientRequest] as const); +} + +void assertV2MethodTypes; + +function memoryWireStreamPair(): [sdk.Stream, sdk.Stream] { + const leftToRight = new TransformStream(); + const rightToLeft = new TransformStream(); + return [ + { + readable: rightToLeft.readable, + writable: leftToRight.writable, + }, + { + readable: leftToRight.readable, + writable: rightToLeft.writable, + }, + ]; +} + +async function respondToNextRequest( + stream: sdk.Stream, + result: unknown, +): Promise { + const reader = stream.readable.getReader(); + const request = await reader.read(); + reader.releaseLock(); + if ( + request.done || + Array.isArray(request.value) || + !("id" in request.value) + ) { + throw new Error("Expected one JSON-RPC request"); + } + + const writer = stream.writable.getWriter(); + try { + await writer.write({ + jsonrpc: "2.0", + id: request.value.id, + result, + }); + } finally { + writer.releaseLock(); + } +} + describe("experimental v2 date-time schemas", () => { it("preserves RFC 3339 timestamps with timezone offsets as strings", () => { const timestamp = "2026-07-20T01:00:00+01:00"; @@ -583,67 +727,161 @@ describe("experimental v2 app API", () => { ).rejects.toMatchObject({ code: -32600 }); }); - it("requires and preserves underscore-prefixed extension methods", async () => { + 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: [], + }), + ); + + await respondToNextRequest(peerStream, { sessionId: 42 }); + await expect(response).rejects.toThrow(); + }); + + 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 reader = peerStream.readable.getReader(); + const request = await reader.read(); + reader.releaseLock(); + if ( + request.done || + !Array.isArray(request.value) || + request.value.length !== 1 || + !("id" in request.value[0]) + ) { + throw new Error("Expected one JSON-RPC batch request"); + } + const writer = peerStream.writable.getWriter(); + try { + await writer.write([ + { + jsonrpc: "2.0", + id: request.value[0].id, + result: { sessionId: 42 }, + }, + ]); + } finally { + writer.releaseLock(); + } + + await expect(response).rejects.toThrow(); + expect(mapped).toBe(false); + }); + + 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", + }), + ); + + 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, { + sessionId: "session-1", + }), + ), + ).resolves.toEqual({}); + }); + + it("supports unrecognized protocol methods and underscore extensions", async () => { const parseValue = (params: unknown): { value: string } => params as { value: string }; const returnValue = ({ params }: { params: { value: string } }) => params; - expect(() => - agent().onRequest("vendor/echo", parseValue, returnValue), - ).toThrow("must start with '_'"); - expect(() => - agent().onNotification("vendor/event", parseValue, () => {}), - ).toThrow("must start with '_'"); - expect(() => - client().onRequest("vendor/echo", parseValue, returnValue), - ).toThrow("must start with '_'"); - expect(() => - client().onNotification("vendor/event", parseValue, () => {}), - ).toThrow("must start with '_'"); - - let notificationValue: string | undefined; - const clientApp = client().onNotification( - "_vendor/acme/event", - parseValue, - ({ params }) => { - notificationValue = params.value; - }, - ); + let agentNotificationValue: string | undefined; + let clientNotificationValue: string | undefined; + const clientApp = client() + .onRequest("authentication/logout", parseValue, returnValue) + .onNotification("authentication/status", parseValue, ({ params }) => { + clientNotificationValue = params.value; + }); const agentApp = agent() + .onRequest("session/load", parseValue, async ({ params, client }) => { + const response = await client.request< + { value: string }, + { value: string }, + "authentication/logout" + >("authentication/logout", params); + await client.notify("authentication/status", params); + return response; + }) .onRequest("_vendor/acme/echo", parseValue, returnValue) - .onRequest( - methods.agent.initialize, - async ({ client: clientContext }) => { - expect(() => - clientContext.request("vendor/client-request", {}), - ).toThrow("must start with '_'"); - expect(() => - clientContext.notify("vendor/client-notification", {}), - ).toThrow("must start with '_'"); - await clientContext.notify("_vendor/acme/event", { - value: "notification", - }); - return { protocolVersion: PROTOCOL_VERSION, info: agentInfo }; - }, - ); + .onNotification("session/set_model", parseValue, ({ params }) => { + agentNotificationValue = params.value; + }) + .onRequest(methods.agent.initialize, () => ({ + protocolVersion: PROTOCOL_VERSION, + info: agentInfo, + })); await clientApp.connectWith(agentApp, async (agentContext) => { - expect(() => agentContext.request("vendor/request", {})).toThrow( - "must start with '_'", - ); - expect(() => agentContext.notify("vendor/notification", {})).toThrow( - "must start with '_'", - ); + const wrongDirection: string = methods.client.mcp.disconnect; + expect(() => + agentContext.request(wrongDirection, { connectionId: "connection-1" }), + ).toThrow("not valid in this direction"); expect(() => agentContext.batch([ - batchNotification("vendor/batch-notification", {}), + { + kind: "request", + method: wrongDirection, + params: { connectionId: "connection-1" }, + }, ] as const), - ).toThrow("must start with '_'"); + ).toThrow("not valid in this direction"); await agentContext.request(methods.agent.initialize, { protocolVersion: PROTOCOL_VERSION, info: clientInfo, }); + await expect( + agentContext.request< + { value: string }, + { value: string }, + "session/load" + >("session/load", { value: "legacy response" }), + ).resolves.toEqual({ value: "legacy response" }); + expect(clientNotificationValue).toBe("legacy response"); + + await agentContext.notify("session/set_model", { + value: "direct notification", + }); + expect(agentNotificationValue).toBe("direct notification"); + + const [batchResponse] = await agentContext.batch([ + batchRequest<{ value: string }, { value: string }, "session/load">( + "session/load", + { value: "batch response" }, + ), + batchNotification("session/set_model", { + value: "batch notification", + }), + ] as const); + expect(batchResponse).toEqual({ value: "batch response" }); + expect(agentNotificationValue).toBe("batch notification"); + await expect( agentContext.request<{ value: string }, { value: string }>( "_vendor/acme/echo", @@ -651,6 +889,10 @@ describe("experimental v2 app API", () => { ), ).resolves.toEqual({ value: "response" }); }); - expect(notificationValue).toBe("notification"); + + const dynamicBuiltIn: string = methods.agent.session.new; + expect(() => + agent().onRequest(dynamicBuiltIn, parseValue, returnValue), + ).toThrow("Cannot replace the built-in"); }); }); diff --git a/src/v2/acp.ts b/src/v2/acp.ts index 1e653368..8259fff0 100644 --- a/src/v2/acp.ts +++ b/src/v2/acp.ts @@ -85,7 +85,7 @@ export function ndJsonStream( return createJsonStream(output, input); } -export { RequestError, batchNotification, batchRequest } from "../jsonrpc.js"; +export { RequestError } from "../jsonrpc.js"; export { AgentProtocolRouter, agentProtocolRouter, @@ -117,6 +117,8 @@ export type { } from "../jsonrpc.js"; import { + batchNotification as jsonRpcBatchNotification, + batchRequest as jsonRpcBatchRequest, Connection, Handled, HandlerRegistration, @@ -125,7 +127,9 @@ import { import type { AnyWireMessage, BatchEntry, + BatchNotification, BatchOutputs, + BatchRequest, ConnectionBuilder, ConnectionContext, HandleResult, @@ -136,11 +140,201 @@ import type { SendRequestOptions, } from "../jsonrpc.js"; +/** + * ACP v2 extension method name. + * + * New custom methods should begin with `_` so they cannot collide with present + * or future protocol methods. + * + * @experimental + */ +export type ExtensionMethod = `_${string}`; + +/** + * A method name that is not part of the current ACP v2 draft. + * + * This compatibility type permits methods from older or newer unstable ACP + * revisions while preventing current built-in method literals from falling + * through the untyped overloads. Prefer {@link ExtensionMethod} for new custom + * methods. + * + * @experimental + */ +export type UnrecognizedMethod = string extends Method + ? Method + : Method extends + | AgentRequestMethod + | AgentNotificationMethod + | ClientRequestMethod + | ClientNotificationMethod + | typeof schema.PROTOCOL_METHODS.cancel_request + ? never + : Method; + +/** + * Creates a typed request descriptor for an ACP v2 batch. + * + * Built-in method literals infer their params and response types. Extension + * methods retain the low-level helper's explicit params/response generics. + * + * @experimental + */ +export function batchRequest( + method: Method, + params: AgentRequestParamsByMethod[Method], + options?: SendRequestOptions, +): BatchRequest< + AgentRequestParamsByMethod[Method], + AgentRequestResponsesByMethod[Method] +> & { + readonly method: Method; +}; +export function batchRequest( + method: Method, + params: AgentRequestParamsByMethod[Method], + mapResponse: (response: AgentRequestResponsesByMethod[Method]) => Output, + options?: SendRequestOptions, +): BatchRequest< + AgentRequestParamsByMethod[Method], + AgentRequestResponsesByMethod[Method], + Output +> & { + readonly method: Method; +}; +export function batchRequest( + method: Method, + params: ClientRequestParamsByMethod[Method], + options?: SendRequestOptions, +): BatchRequest< + ClientRequestParamsByMethod[Method], + ClientRequestResponsesByMethod[Method] +> & { + readonly method: Method; +}; +export function batchRequest( + method: Method, + params: ClientRequestParamsByMethod[Method], + mapResponse: (response: ClientRequestResponsesByMethod[Method]) => Output, + options?: SendRequestOptions, +): BatchRequest< + ClientRequestParamsByMethod[Method], + ClientRequestResponsesByMethod[Method], + Output +> & { + readonly method: Method; +}; +export function batchRequest( + method: ExtensionMethod, + params?: Params, + options?: SendRequestOptions, +): BatchRequest & { + readonly method: ExtensionMethod; +}; +export function batchRequest( + method: ExtensionMethod, + params: Params | undefined, + mapResponse: (response: Response) => Output, + options?: SendRequestOptions, +): BatchRequest & { + readonly method: ExtensionMethod; +}; +export function batchRequest< + Params = unknown, + Response = unknown, + const Method extends string = never, +>( + method: UnrecognizedMethod, + params?: Params, + options?: SendRequestOptions, +): BatchRequest & { + readonly method: Method; +}; +export function batchRequest< + Params = unknown, + Response = unknown, + Output = Response, + const Method extends string = never, +>( + method: UnrecognizedMethod, + params: Params | undefined, + mapResponse: (response: Response) => Output, + options?: SendRequestOptions, +): BatchRequest & { + readonly method: Method; +}; +export function batchRequest( + method: string, + params?: Params, + mapResponseOrOptions?: ((response: Response) => unknown) | SendRequestOptions, + options?: SendRequestOptions, +): BatchRequest & { + readonly method: string; +} { + const request = + typeof mapResponseOrOptions === "function" + ? jsonRpcBatchRequest(method, params, mapResponseOrOptions, options) + : jsonRpcBatchRequest(method, params, mapResponseOrOptions); + return request; +} + +/** + * Creates a typed notification descriptor for an ACP v2 batch. + * + * @experimental + */ +export function batchNotification( + method: Method, + params: AgentNotificationParamsByMethod[Method], +): BatchNotification & { + readonly method: Method; +}; +export function batchNotification( + method: Method, + params: ClientNotificationParamsByMethod[Method], +): BatchNotification & { + readonly method: Method; +}; +export function batchNotification( + method: typeof schema.PROTOCOL_METHODS.cancel_request, + params: schema.CancelRequestNotification, +): BatchNotification & { + readonly method: typeof schema.PROTOCOL_METHODS.cancel_request; +}; +export function batchNotification( + method: ExtensionMethod, + params?: Params, +): BatchNotification & { + readonly method: ExtensionMethod; +}; +export function batchNotification< + Params = unknown, + const Method extends string = never, +>( + method: UnrecognizedMethod, + params?: Params, +): BatchNotification & { + readonly method: Method; +}; +export function batchNotification( + method: string, + params?: Params, +): BatchNotification & { + readonly method: string; +} { + return jsonRpcBatchNotification(method, params); +} + function emptyObjectResponse(response: T | null | undefined | void): T { return response ?? ({} as T); } -function assertV2Method( +const knownProtocolMethods = new Set([ + ...Object.values(schema.AGENT_METHODS), + ...Object.values(schema.CLIENT_METHODS), + ...Object.values(schema.PROTOCOL_METHODS), +]); + +function assertV2MethodDirection( method: string, builtIns: Record, kind: "request" | "notification", @@ -153,9 +347,20 @@ function assertV2Method( ) { return; } - if (!method.startsWith("_")) { + if (knownProtocolMethods.has(method)) { + throw new TypeError( + `ACP v2 ${kind} method '${method}' is not valid in this direction`, + ); + } +} + +function assertUnrecognizedV2Method( + method: string, + kind: "request" | "notification", +): void { + if (knownProtocolMethods.has(method)) { throw new TypeError( - `Custom ACP v2 ${kind} method '${method}' must start with '_'`, + `Cannot replace the built-in ACP v2 ${kind} parser for '${method}'`, ); } } @@ -166,7 +371,7 @@ function assertV2BatchMethods( notificationMethods: Record, ): void { for (const entry of entries) { - assertV2Method( + assertV2MethodDirection( entry.method, entry.kind === "request" ? requestMethods : notificationMethods, entry.kind, @@ -202,9 +407,7 @@ function normalizeOutgoingV2InitializeRequest( }); } -function mapV2InitializeResponse( - response: schema.InitializeResponse, -): schema.InitializeResponse { +function mapV2InitializeResponse(response: unknown): schema.InitializeResponse { const parsed = validate.zInitializeResponse.parse(response); if (parsed.protocolVersion !== schema.PROTOCOL_VERSION) { throw RequestError.invalidRequest( @@ -218,26 +421,41 @@ function mapV2InitializeResponse( return parsed; } -function normalizeClientBatch( +function parseRequestResponse( + spec: { response?: ParamsParser }, + response: unknown, +): unknown { + return parseParams(spec.response, response); +} + +function normalizeV2Batch( entries: Entries & { readonly 0: BatchEntry }, + requestSpecs: Record< + string, + { response?: ParamsParser } | undefined + >, + normalizeInitialize = false, ): Entries & { readonly 0: BatchEntry } { return entries.map((entry) => { - if ( - entry.kind !== "request" || - entry.method !== schema.AGENT_METHODS.initialize - ) { + if (entry.kind !== "request") { return entry; } + const spec = requestSpecs[entry.method]; const mapResponse = entry.mapResponse as - ((response: schema.InitializeResponse) => unknown) | undefined; + ((response: unknown) => unknown) | undefined; return { ...entry, - params: normalizeOutgoingV2InitializeRequest(entry.params), - mapResponse: (response: schema.InitializeResponse) => { - const parsed = mapV2InitializeResponse(response); - return mapResponse ? mapResponse(parsed) : parsed; - }, + params: + normalizeInitialize && entry.method === schema.AGENT_METHODS.initialize + ? normalizeOutgoingV2InitializeRequest(entry.params) + : entry.params, + mapResponse: spec + ? (response: unknown) => { + const parsed = parseRequestResponse(spec, response); + return mapResponse ? mapResponse(parsed) : parsed; + } + : mapResponse, }; }) as unknown as Entries & { readonly 0: BatchEntry }; } @@ -396,6 +614,70 @@ export interface ClientConnection extends AcpConnection { readonly agent: ClientContext; } +/** + * One batch entry sent to an ACP v2 agent. + * + * @experimental + */ +export type AgentBatchEntry = + | { + [Method in AgentRequestMethod]: BatchRequest< + AgentRequestParamsByMethod[Method], + AgentRequestResponsesByMethod[Method], + unknown + > & { + readonly method: Method; + }; + }[AgentRequestMethod] + | { + [Method in AgentNotificationMethod]: BatchNotification< + AgentNotificationParamsByMethod[Method] + > & { + readonly method: Method; + }; + }[AgentNotificationMethod] + | (BatchNotification & { + readonly method: typeof schema.PROTOCOL_METHODS.cancel_request; + }) + | (BatchRequest & { + readonly method: ExtensionMethod | UnrecognizedMethod; + }) + | (BatchNotification & { + readonly method: ExtensionMethod | UnrecognizedMethod; + }); + +/** + * One batch entry sent to an ACP v2 client. + * + * @experimental + */ +export type ClientBatchEntry = + | { + [Method in ClientRequestMethod]: BatchRequest< + ClientRequestParamsByMethod[Method], + ClientRequestResponsesByMethod[Method], + unknown + > & { + readonly method: Method; + }; + }[ClientRequestMethod] + | { + [Method in ClientNotificationMethod]: BatchNotification< + ClientNotificationParamsByMethod[Method] + > & { + readonly method: Method; + }; + }[ClientNotificationMethod] + | (BatchNotification & { + readonly method: typeof schema.PROTOCOL_METHODS.cancel_request; + }) + | (BatchRequest & { + readonly method: ExtensionMethod | UnrecognizedMethod; + }) + | (BatchNotification & { + readonly method: ExtensionMethod | UnrecognizedMethod; + }); + class AcpContext { /** @internal */ constructor( @@ -476,7 +758,16 @@ export class AgentContext extends AcpContext { options?: SendRequestOptions, ): Promise; request( - method: string, + method: ExtensionMethod, + params?: Params, + options?: SendRequestOptions, + ): Promise; + request< + Response = unknown, + Params = unknown, + const Method extends string = never, + >( + method: UnrecognizedMethod, params?: Params, options?: SendRequestOptions, ): Promise; @@ -485,10 +776,15 @@ export class AgentContext extends AcpContext { params?: unknown, options?: SendRequestOptions, ): Promise { - assertV2Method(method, clientRequestSpecsByMethod, "request"); + assertV2MethodDirection(method, clientRequestSpecsByMethod, "request"); const spec = clientRequestSpecsByMethod[method] as AcpRequestSpec | undefined; - return this.sendRequest(method, params, spec?.mapResponse, options); + return this.sendRequest( + method, + params, + spec ? (response) => parseRequestResponse(spec, response) : undefined, + options, + ); } /** @@ -501,9 +797,20 @@ export class AgentContext extends AcpContext { method: Method, params: ClientNotificationParamsByMethod[Method], ): Promise; - notify(method: string, params?: Params): Promise; + notify( + method: typeof schema.PROTOCOL_METHODS.cancel_request, + params: schema.CancelRequestNotification, + ): Promise; + notify( + method: ExtensionMethod, + params?: Params, + ): Promise; + notify( + method: UnrecognizedMethod, + params?: Params, + ): Promise; notify(method: string, params?: unknown): Promise { - assertV2Method( + assertV2MethodDirection( method, clientNotificationSpecsByMethod, "notification", @@ -516,14 +823,31 @@ export class AgentContext extends AcpContext { * Sends requests and notifications to the client as one JSON-RPC batch. */ batch( - entries: Entries & { readonly 0: BatchEntry }, + entries: Entries & { readonly 0: BatchEntry } & { + [Index in keyof Entries]: Entries[Index] extends BatchEntry + ? string extends Entries[Index]["method"] + ? Entries[Index] + : Entries[Index]["method"] extends + | AgentRequestMethod + | AgentNotificationMethod + | ClientRequestMethod + | ClientNotificationMethod + | typeof schema.PROTOCOL_METHODS.cancel_request + ? Entries[Index] extends ClientBatchEntry + ? Entries[Index] + : never + : Entries[Index] + : never; + }, ): Promise> { assertV2BatchMethods( entries, clientRequestSpecsByMethod, clientNotificationSpecsByMethod, ); - return this.sendBatch(entries); + return this.sendBatch( + normalizeV2Batch(entries, clientRequestSpecsByMethod), + ); } } @@ -551,15 +875,8 @@ export class ClientContext extends AcpContext { params: schema.NewSessionRequest, options?: SendRequestOptions, ): Promise { - return this.sendRequest< - schema.NewSessionRequest, - schema.NewSessionResponse, - ActiveSession - >( - schema.AGENT_METHODS.session_new, - params, + return this.request(schema.AGENT_METHODS.session_new, params, options).then( (response) => this.attachSession(response), - options, ); } @@ -655,7 +972,16 @@ export class ClientContext extends AcpContext { options?: SendRequestOptions, ): Promise; request( - method: string, + method: ExtensionMethod, + params?: Params, + options?: SendRequestOptions, + ): Promise; + request< + Response = unknown, + Params = unknown, + const Method extends string = never, + >( + method: UnrecognizedMethod, params?: Params, options?: SendRequestOptions, ): Promise; @@ -664,14 +990,19 @@ export class ClientContext extends AcpContext { params?: unknown, options?: SendRequestOptions, ): Promise { - assertV2Method(method, agentRequestSpecsByMethod, "request"); + 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?.mapResponse, options); + return this.sendRequest( + method, + wireParams, + spec ? (response) => parseRequestResponse(spec, response) : undefined, + options, + ); } /** @@ -684,9 +1015,20 @@ export class ClientContext extends AcpContext { method: Method, params: AgentNotificationParamsByMethod[Method], ): Promise; - notify(method: string, params?: Params): Promise; + notify( + method: typeof schema.PROTOCOL_METHODS.cancel_request, + params: schema.CancelRequestNotification, + ): Promise; + notify( + method: ExtensionMethod, + params?: Params, + ): Promise; + notify( + method: UnrecognizedMethod, + params?: Params, + ): Promise; notify(method: string, params?: unknown): Promise { - assertV2Method( + assertV2MethodDirection( method, agentNotificationSpecsByMethod, "notification", @@ -699,14 +1041,31 @@ export class ClientContext extends AcpContext { * Sends requests and notifications to the agent as one JSON-RPC batch. */ batch( - entries: Entries & { readonly 0: BatchEntry }, + entries: Entries & { readonly 0: BatchEntry } & { + [Index in keyof Entries]: Entries[Index] extends BatchEntry + ? string extends Entries[Index]["method"] + ? Entries[Index] + : Entries[Index]["method"] extends + | AgentRequestMethod + | AgentNotificationMethod + | ClientRequestMethod + | ClientNotificationMethod + | typeof schema.PROTOCOL_METHODS.cancel_request + ? Entries[Index] extends AgentBatchEntry + ? Entries[Index] + : never + : Entries[Index] + : never; + }, ): Promise> { assertV2BatchMethods( entries, agentRequestSpecsByMethod, agentNotificationSpecsByMethod, ); - return this.sendBatch(normalizeClientBatch(entries)); + return this.sendBatch( + normalizeV2Batch(entries, agentRequestSpecsByMethod, true), + ); } } @@ -1434,10 +1793,11 @@ function parseParams( return parser.parse(params); } -type AcpRequestSpec = { +type AcpRequestSpec = { method: string; params?: ParamsParser; - mapResponse?: (response: Response) => WireResponse; + response?: ParamsParser; + serializeResponse?: (response: HandlerResponse) => Response; }; type AcpNotificationSpec = { @@ -1445,12 +1805,13 @@ type AcpNotificationSpec = { params?: ParamsParser; }; -function requestSpec( +function requestSpec( method: string, params: ParamsParser, - mapResponse?: (response: Response) => WireResponse, -): AcpRequestSpec { - return { method, params, mapResponse }; + response: ParamsParser, + serializeResponse?: (response: HandlerResponse) => Response, +): AcpRequestSpec { + return { method, params, response, serializeResponse }; } function notificationSpec( @@ -1460,18 +1821,18 @@ function notificationSpec( return { method, params }; } -function registerAppRequest( +function registerAppRequest( builder: ConnectionBuilder, - spec: AcpRequestSpec, + spec: AcpRequestSpec, context: ( params: Params, cx: ConnectionContext, signal: AbortSignal, requestId: JsonRpcId, ) => Context, - handler: (context: Context) => MaybePromise, + handler: (context: Context) => MaybePromise, ): void { - builder.onReceiveRequest( + builder.onReceiveRequest( spec.method, (params) => parseParams(spec.params, params), async (params, responder, cx) => { @@ -1479,9 +1840,9 @@ function registerAppRequest( context(params, cx, responder.signal, responder.id), ); await responder.respond( - (spec.mapResponse - ? spec.mapResponse(response) - : response) as WireResponse, + spec.serializeResponse + ? spec.serializeResponse(response) + : (response as unknown as Response), ); }, ); @@ -1519,6 +1880,7 @@ const agentRequestSpecs = { schema.AGENT_METHODS.initialize, parseV2InitializeRequest, mapV2InitializeResponse, + mapV2InitializeResponse, ), loginAuth: requestSpec< schema.LoginAuthRequest, @@ -1527,12 +1889,17 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.auth_login, validate.zLoginAuthRequest, + validate.zLoginAuthResponse, emptyObjectResponse, ), unstable_listProviders: requestSpec< schema.ListProvidersRequest, schema.ListProvidersResponse - >(schema.AGENT_METHODS.providers_list, validate.zListProvidersRequest), + >( + schema.AGENT_METHODS.providers_list, + validate.zListProvidersRequest, + validate.zListProvidersResponse, + ), unstable_setProvider: requestSpec< schema.SetProviderRequest, schema.SetProviderResponse | void, @@ -1540,6 +1907,7 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.providers_set, validate.zSetProviderRequest, + validate.zSetProviderResponse, emptyObjectResponse, ), unstable_disableProvider: requestSpec< @@ -1549,11 +1917,13 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.providers_disable, validate.zDisableProviderRequest, + validate.zDisableProviderResponse, emptyObjectResponse, ), newSession: requestSpec( schema.AGENT_METHODS.session_new, validate.zNewSessionRequest, + validate.zNewSessionResponse, ), setSessionConfigOption: requestSpec< schema.SetSessionConfigOptionRequest, @@ -1561,6 +1931,7 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.session_set_config_option, validate.zSetSessionConfigOptionRequest, + validate.zSetSessionConfigOptionResponse, ), prompt: requestSpec< schema.PromptRequest, @@ -1569,16 +1940,25 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.session_prompt, validate.zPromptRequest, + validate.zPromptResponse, emptyObjectResponse, ), unstable_messageMcp: requestSpec< schema.MessageMcpRequest, schema.MessageMcpResponse - >(schema.AGENT_METHODS.mcp_message, validate.zMessageMcpRequest), + >( + schema.AGENT_METHODS.mcp_message, + validate.zMessageMcpRequest, + validate.zMessageMcpResponse, + ), listSessions: requestSpec< schema.ListSessionsRequest, schema.ListSessionsResponse - >(schema.AGENT_METHODS.session_list, validate.zListSessionsRequest), + >( + schema.AGENT_METHODS.session_list, + validate.zListSessionsRequest, + validate.zListSessionsResponse, + ), deleteSession: requestSpec< schema.DeleteSessionRequest, schema.DeleteSessionResponse | void, @@ -1586,16 +1966,25 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.session_delete, validate.zDeleteSessionRequest, + validate.zDeleteSessionResponse, emptyObjectResponse, ), unstable_forkSession: requestSpec< schema.ForkSessionRequest, schema.ForkSessionResponse - >(schema.AGENT_METHODS.session_fork, validate.zForkSessionRequest), + >( + schema.AGENT_METHODS.session_fork, + validate.zForkSessionRequest, + validate.zForkSessionResponse, + ), resumeSession: requestSpec< schema.ResumeSessionRequest, schema.ResumeSessionResponse - >(schema.AGENT_METHODS.session_resume, validate.zResumeSessionRequest), + >( + schema.AGENT_METHODS.session_resume, + validate.zResumeSessionRequest, + validate.zResumeSessionResponse, + ), closeSession: requestSpec< schema.CloseSessionRequest, schema.CloseSessionResponse | void, @@ -1603,6 +1992,7 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.session_close, validate.zCloseSessionRequest, + validate.zCloseSessionResponse, emptyObjectResponse, ), logoutAuth: requestSpec< @@ -1612,16 +2002,25 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.auth_logout, validate.zLogoutAuthRequest, + validate.zLogoutAuthResponse, emptyObjectResponse, ), unstable_startNes: requestSpec< schema.StartNesRequest, schema.StartNesResponse - >(schema.AGENT_METHODS.nes_start, validate.zStartNesRequest), + >( + schema.AGENT_METHODS.nes_start, + validate.zStartNesRequest, + validate.zStartNesResponse, + ), unstable_suggestNes: requestSpec< schema.SuggestNesRequest, schema.SuggestNesResponse - >(schema.AGENT_METHODS.nes_suggest, validate.zSuggestNesRequest), + >( + schema.AGENT_METHODS.nes_suggest, + validate.zSuggestNesRequest, + validate.zSuggestNesResponse, + ), unstable_closeNes: requestSpec< schema.CloseNesRequest, schema.CloseNesResponse | void, @@ -1629,6 +2028,7 @@ const agentRequestSpecs = { >( schema.AGENT_METHODS.nes_close, validate.zCloseNesRequest, + validate.zCloseNesResponse, emptyObjectResponse, ), }; @@ -1684,15 +2084,24 @@ const clientRequestSpecs = { >( schema.CLIENT_METHODS.session_request_permission, validate.zRequestPermissionRequest, + validate.zRequestPermissionResponse, ), unstable_connectMcp: requestSpec< schema.ConnectMcpRequest, schema.ConnectMcpResponse - >(schema.CLIENT_METHODS.mcp_connect, validate.zConnectMcpRequest), + >( + schema.CLIENT_METHODS.mcp_connect, + validate.zConnectMcpRequest, + validate.zConnectMcpResponse, + ), unstable_messageMcp: requestSpec< schema.MessageMcpRequest, schema.MessageMcpResponse - >(schema.CLIENT_METHODS.mcp_message, validate.zMessageMcpRequest), + >( + schema.CLIENT_METHODS.mcp_message, + validate.zMessageMcpRequest, + validate.zMessageMcpResponse, + ), unstable_disconnectMcp: requestSpec< schema.DisconnectMcpRequest, schema.DisconnectMcpResponse | void, @@ -1700,6 +2109,7 @@ const clientRequestSpecs = { >( schema.CLIENT_METHODS.mcp_disconnect, validate.zDisconnectMcpRequest, + validate.zDisconnectMcpResponse, emptyObjectResponse, ), unstable_createElicitation: requestSpec< @@ -1708,6 +2118,7 @@ const clientRequestSpecs = { >( schema.CLIENT_METHODS.elicitation_create, validate.zCreateElicitationRequest, + validate.zCreateElicitationResponse, ), }; @@ -2258,7 +2669,12 @@ export class AgentApp { handler: AgentRequestHandlersByMethod[Method], ): this; onRequest( - method: string, + method: ExtensionMethod, + params: ParamsParser, + handler: AgentRequestHandler, + ): this; + onRequest( + method: UnrecognizedMethod, params: ParamsParser, handler: AgentRequestHandler, ): this; @@ -2269,7 +2685,7 @@ export class AgentApp { handler?: AgentRequestHandler, ): this { if (handler) { - assertV2Method(method, agentRequestSpecsByMethod, "request"); + assertUnrecognizedV2Method(method, "request"); return this.request( { method, params: handlerOrParams as ParamsParser }, handler, @@ -2300,7 +2716,12 @@ export class AgentApp { handler: AgentNotificationHandlersByMethod[Method], ): this; onNotification( - method: string, + method: ExtensionMethod, + params: ParamsParser, + handler: AgentNotificationHandler, + ): this; + onNotification( + method: UnrecognizedMethod, params: ParamsParser, handler: AgentNotificationHandler, ): this; @@ -2312,12 +2733,7 @@ export class AgentApp { handler?: AgentNotificationHandler, ): this { if (handler) { - assertV2Method( - method, - agentNotificationSpecsByMethod, - "notification", - true, - ); + assertUnrecognizedV2Method(method, "notification"); return this.notification( { method, params: handlerOrParams as ParamsParser }, handler, @@ -2518,7 +2934,12 @@ export class ClientApp { handler: ClientRequestHandlersByMethod[Method], ): this; onRequest( - method: string, + method: ExtensionMethod, + params: ParamsParser, + handler: ClientRequestHandler, + ): this; + onRequest( + method: UnrecognizedMethod, params: ParamsParser, handler: ClientRequestHandler, ): this; @@ -2529,7 +2950,7 @@ export class ClientApp { handler?: ClientRequestHandler, ): this { if (handler) { - assertV2Method(method, clientRequestSpecsByMethod, "request"); + assertUnrecognizedV2Method(method, "request"); return this.request( { method, params: handlerOrParams as ParamsParser }, handler, @@ -2560,7 +2981,12 @@ export class ClientApp { handler: ClientNotificationHandlersByMethod[Method], ): this; onNotification( - method: string, + method: ExtensionMethod, + params: ParamsParser, + handler: ClientNotificationHandler, + ): this; + onNotification( + method: UnrecognizedMethod, params: ParamsParser, handler: ClientNotificationHandler, ): this; @@ -2572,12 +2998,7 @@ export class ClientApp { handler?: ClientNotificationHandler, ): this { if (handler) { - assertV2Method( - method, - clientNotificationSpecsByMethod, - "notification", - true, - ); + assertUnrecognizedV2Method(method, "notification"); return this.notification( { method, params: handlerOrParams as ParamsParser }, handler,