Support passing an AbortSignal to request()

Author: jspahrsummers

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,
     );
   }
 

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,
     );
   }
 

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<T extends ZodType<object>>(
     request: SendRequestT,
     resultSchema: T,
-    onprogress?: ProgressCallback,
+    options?: RequestOptions,
   ): Promise<z.infer<T>> {
     return new Promise((resolve, reject) => {
       if (!this._transport) {
@@ -304,22 +319,28 @@ export abstract class Protocol<
         this.assertCapabilityForMethod(request.method);
       }
 
+      options?.signal?.throwIfAborted();
+
       const messageId = this._requestMessageId++;
       const jsonrpcRequest: JSONRPCRequest = {
         ...request,
         jsonrpc: "2.0",
         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 },
         };
       }
 
       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);
     });
   }