Merge pull request #60 from modelcontextprotocol/justin/remove-abort-after-timeout

Replace `abortAfterTimeout` with `RequestOptions.timeout`

README.md

@@ -1,4 +1,4 @@
-# MCP TypeScript SDK
+# MCP TypeScript SDK ![NPM Version](https://img.shields.io/npm/v/%40modelcontextprotocol%2Fsdk)
 
 TypeScript implementation of the Model Context Protocol (MCP), providing both client and server capabilities for integrating with LLM surfaces.
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@modelcontextprotocol/sdk",
-  "version": "0.5.0",
+  "version": "0.6.0",
   "description": "Model Context Protocol implementation for TypeScript",
   "license": "MIT",
   "author": "Anthropic, PBC (https://anthropic.com)",

src/client/index.test.ts

@@ -14,6 +14,7 @@ import {
   ListToolsRequestSchema,
   CreateMessageRequestSchema,
   ListRootsRequestSchema,
+  ErrorCode,
 } from "../types.js";
 import { Transport } from "../shared/transport.js";
 import { Server } from "../server/index.js";
@@ -491,3 +492,58 @@ test("should handle client cancelling a request", async () => {
   // Request should be rejected
   await expect(listResourcesPromise).rejects.toBe("Cancelled by test");
 });
+
+test("should handle request timeout", async () => {
+  const server = new Server(
+    {
+      name: "test server",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        resources: {},
+      },
+    },
+  );
+
+  // Set up server with a delayed response
+  server.setRequestHandler(
+    ListResourcesRequestSchema,
+    async (_request, extra) => {
+      const timer = new Promise((resolve) => {
+        const timeout = setTimeout(resolve, 100);
+        extra.signal.addEventListener("abort", () => clearTimeout(timeout));
+      });
+
+      await timer;
+      return {
+        resources: [],
+      };
+    },
+  );
+
+  const [clientTransport, serverTransport] =
+    InMemoryTransport.createLinkedPair();
+
+  const client = new Client(
+    {
+      name: "test client",
+      version: "1.0",
+    },
+    {
+      capabilities: {},
+    },
+  );
+
+  await Promise.all([
+    client.connect(clientTransport),
+    server.connect(serverTransport),
+  ]);
+
+  // Request with 0 msec timeout should fail immediately
+  await expect(
+    client.listResources(undefined, { timeout: 0 }),
+  ).rejects.toMatchObject({
+    code: ErrorCode.RequestTimeout,
+  });
+});

src/server/index.test.ts

@@ -14,6 +14,7 @@ import {
   ListResourcesRequestSchema,
   ListToolsRequestSchema,
   SetLevelRequestSchema,
+  ErrorCode,
 } from "../types.js";
 import { Transport } from "../shared/transport.js";
 import { InMemoryTransport } from "../inMemory.js";
@@ -475,3 +476,72 @@ test("should handle server cancelling a request", async () => {
   // Request should be rejected
   await expect(createMessagePromise).rejects.toBe("Cancelled by test");
 });
+test("should handle request timeout", async () => {
+  const server = new Server(
+    {
+      name: "test server",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        sampling: {},
+      },
+    },
+  );
+
+  // Set up client that delays responses
+  const client = new Client(
+    {
+      name: "test client",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        sampling: {},
+      },
+    },
+  );
+
+  client.setRequestHandler(
+    CreateMessageRequestSchema,
+    async (_request, extra) => {
+      await new Promise((resolve, reject) => {
+        const timeout = setTimeout(resolve, 100);
+        extra.signal.addEventListener("abort", () => {
+          clearTimeout(timeout);
+          reject(extra.signal.reason);
+        });
+      });
+
+      return {
+        model: "test",
+        role: "assistant",
+        content: {
+          type: "text",
+          text: "Test response",
+        },
+      };
+    },
+  );
+
+  const [clientTransport, serverTransport] =
+    InMemoryTransport.createLinkedPair();
+
+  await Promise.all([
+    client.connect(clientTransport),
+    server.connect(serverTransport),
+  ]);
+
+  // Request with 0 msec timeout should fail immediately
+  await expect(
+    server.createMessage(
+      {
+        messages: [],
+        maxTokens: 10,
+      },
+      { timeout: 0 },
+    ),
+  ).rejects.toMatchObject({
+    code: ErrorCode.RequestTimeout,
+  });
+});

src/shared/protocol.ts

@@ -37,6 +37,11 @@ export type ProtocolOptions = {
   enforceStrictCapabilities?: boolean;
 };
 
+/**
+ * The default request timeout, in miliseconds.
+ */
+export const DEFAULT_REQUEST_TIMEOUT_MSEC = 60000;
+
 /**
  * Options that can be given per request.
  */
@@ -48,10 +53,15 @@ export type RequestOptions = {
 
   /**
    * Can be used to cancel an in-flight request. This will cause an AbortError to be raised from request().
-   *
-   * Use abortAfterTimeout() to easily implement timeouts using this signal.
    */
   signal?: AbortSignal;
+
+  /**
+   * A timeout (in milliseconds) for this request. If exceeded, an McpError with code `RequestTimeout` will be raised from request().
+   *
+   * If not specified, `DEFAULT_REQUEST_TIMEOUT_MSEC` will be used as the timeout.
+   */
+  timeout?: number;
 };
 
 /**
@@ -381,7 +391,13 @@ export abstract class Protocol<
         };
       }
 
+      let timeoutId: ReturnType<typeof setTimeout> | undefined = undefined;
+
       this._responseHandlers.set(messageId, (response) => {
+        if (timeoutId !== undefined) {
+          clearTimeout(timeoutId);
+        }
+
         if (options?.signal?.aborted) {
           return;
         }
@@ -398,24 +414,52 @@ export abstract class Protocol<
         }
       });
 
-      options?.signal?.addEventListener("abort", () => {
-        const reason = options?.signal?.reason;
+      const cancel = (reason: unknown) => {
         this._responseHandlers.delete(messageId);
         this._progressHandlers.delete(messageId);
 
-        this._transport?.send({
-          jsonrpc: "2.0",
-          method: "cancelled",
-          params: {
-            requestId: messageId,
-            reason: String(reason),
-          },
-        });
+        this._transport
+          ?.send({
+            jsonrpc: "2.0",
+            method: "cancelled",
+            params: {
+              requestId: messageId,
+              reason: String(reason),
+            },
+          })
+          .catch((error) =>
+            this._onerror(new Error(`Failed to send cancellation: ${error}`)),
+          );
 
         reject(reason);
+      };
+
+      options?.signal?.addEventListener("abort", () => {
+        if (timeoutId !== undefined) {
+          clearTimeout(timeoutId);
+        }
+
+        cancel(options?.signal?.reason);
       });
 
-      this._transport.send(jsonrpcRequest).catch(reject);
+      const timeout = options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC;
+      timeoutId = setTimeout(
+        () =>
+          cancel(
+            new McpError(ErrorCode.RequestTimeout, "Request timed out", {
+              timeout,
+            }),
+          ),
+        timeout,
+      );
+
+      this._transport.send(jsonrpcRequest).catch((error) => {
+        if (timeoutId !== undefined) {
+          clearTimeout(timeoutId);
+        }
+
+        reject(error);
+      });
     });
   }
 

src/types.ts

@@ -105,6 +105,7 @@ export const JSONRPCResponseSchema = z
 export enum ErrorCode {
   // SDK error codes
   ConnectionClosed = -1,
+  RequestTimeout = -2,
 
   // Standard JSON-RPC error codes
   ParseError = -32700,

src/utils.test.ts

@@ -1,15 +0,0 @@
-import { abortAfterTimeout } from "./utils.js";
-
-describe("abortAfterTimeout", () => {
-  it("should abort after timeout", () => {
-    const signal = abortAfterTimeout(0);
-    expect(signal.aborted).toBe(false);
-
-    return new Promise((resolve) => {
-      setTimeout(() => {
-        expect(signal.aborted).toBe(true);
-        resolve(true);
-      }, 0);
-    });
-  });
-});

src/utils.ts

@@ -1,11 +0,0 @@
-/**
- * Returns an AbortSignal that will enter aborted state after `timeoutMs` milliseconds.
- */
-export function abortAfterTimeout(timeoutMs: number): AbortSignal {
-  const controller = new AbortController();
-  setTimeout(() => {
-    controller.abort("Timeout exceeded");
-  }, timeoutMs);
-
-  return controller.signal;
-}