diff --git a/src/client/index.ts b/src/client/index.ts index e0df322b..97e6ee07 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -1,7 +1,7 @@ import { - ProgressCallback, Protocol, ProtocolOptions, + RequestOptions, } from "../shared/protocol.js"; import { Transport } from "../shared/transport.js"; import { @@ -278,14 +278,11 @@ export class Client< return this.request({ method: "ping" }, EmptyResultSchema); } - async complete( - params: CompleteRequest["params"], - onprogress?: ProgressCallback, - ) { + async complete(params: CompleteRequest["params"], options?: RequestOptions) { return this.request( { method: "completion/complete", params }, CompleteResultSchema, - onprogress, + options, ); } @@ -298,56 +295,56 @@ export class Client< async getPrompt( params: GetPromptRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "prompts/get", params }, GetPromptResultSchema, - onprogress, + options, ); } async listPrompts( params?: ListPromptsRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "prompts/list", params }, ListPromptsResultSchema, - onprogress, + options, ); } async listResources( params?: ListResourcesRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "resources/list", params }, ListResourcesResultSchema, - onprogress, + options, ); } async listResourceTemplates( params?: ListResourceTemplatesRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "resources/templates/list", params }, ListResourceTemplatesResultSchema, - onprogress, + options, ); } async readResource( params: ReadResourceRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "resources/read", params }, ReadResourceResultSchema, - onprogress, + options, ); } @@ -370,23 +367,23 @@ export class Client< resultSchema: | typeof CallToolResultSchema | typeof CompatibilityCallToolResultSchema = CallToolResultSchema, - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "tools/call", params }, resultSchema, - onprogress, + options, ); } async listTools( params?: ListToolsRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "tools/list", params }, ListToolsResultSchema, - onprogress, + options, ); } diff --git a/src/server/index.ts b/src/server/index.ts index ecb525b5..29a3207a 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1,7 +1,7 @@ import { - ProgressCallback, Protocol, ProtocolOptions, + RequestOptions, } from "../shared/protocol.js"; import { ClientCapabilities, @@ -257,23 +257,23 @@ export class Server< async createMessage( params: CreateMessageRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "sampling/createMessage", params }, CreateMessageResultSchema, - onprogress, + options, ); } async listRoots( params?: ListRootsRequest["params"], - onprogress?: ProgressCallback, + options?: RequestOptions, ) { return this.request( { method: "roots/list", params }, ListRootsResultSchema, - onprogress, + options, ); } diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 85610a9d..e83fe44e 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -35,6 +35,21 @@ export type ProtocolOptions = { enforceStrictCapabilities?: boolean; }; +/** + * Options that can be given per request. + */ +export type RequestOptions = { + /** + * If set, requests progress notifications from the remote end (if supported). When progress notifications are received, this callback will be invoked. + */ + onprogress?: ProgressCallback; + + /** + * Can be used to cancel an in-flight request. This will cause an AbortError to be raised from request(). + */ + signal?: AbortSignal; +}; + /** * Implements MCP protocol framing on top of a pluggable transport, including * features like request/response linking, notifications, and progress. @@ -285,14 +300,14 @@ export abstract class Protocol< protected abstract assertRequestHandlerCapability(method: string): void; /** - * Sends a request and wait for a response, with optional progress notifications in the meantime (if supported by the server). + * Sends a request and wait for a response. * * Do not use this method to emit notifications! Use notification() instead. */ request>( request: SendRequestT, resultSchema: T, - onprogress?: ProgressCallback, + options?: RequestOptions, ): Promise> { return new Promise((resolve, reject) => { if (!this._transport) { @@ -304,6 +319,8 @@ export abstract class Protocol< this.assertCapabilityForMethod(request.method); } + options?.signal?.throwIfAborted(); + const messageId = this._requestMessageId++; const jsonrpcRequest: JSONRPCRequest = { ...request, @@ -311,8 +328,8 @@ export abstract class Protocol< id: messageId, }; - if (onprogress) { - this._progressHandlers.set(messageId, onprogress); + if (options?.onprogress) { + this._progressHandlers.set(messageId, options.onprogress); jsonrpcRequest.params = { ...request.params, _meta: { progressToken: messageId }, @@ -320,6 +337,10 @@ export abstract class Protocol< } this._responseHandlers.set(messageId, (response) => { + if (options?.signal?.aborted) { + return; + } + if (response instanceof Error) { return reject(response); } @@ -332,6 +353,23 @@ export abstract class Protocol< } }); + options?.signal?.addEventListener("abort", () => { + const reason = options?.signal?.reason; + this._responseHandlers.delete(messageId); + this._progressHandlers.delete(messageId); + + this._transport?.send({ + jsonrpc: "2.0", + method: "cancelled", + params: { + requestId: messageId, + reason: String(reason), + }, + }); + + reject(reason); + }); + this._transport.send(jsonrpcRequest).catch(reject); }); }