Check capabilities when request handlers are set

Author: jspahrsummers

src/client/index.test.ts

@@ -12,6 +12,8 @@ import {
   InitializeRequestSchema,
   ListResourcesRequestSchema,
   ListToolsRequestSchema,
+  CreateMessageRequestSchema,
+  ListRootsRequestSchema,
 } from "../types.js";
 import { Transport } from "../shared/transport.js";
 import { Server } from "../server/index.js";
@@ -325,6 +327,37 @@ test("should respect server notification capabilities", async () => {
   );
 });
 
+test("should only allow setRequestHandler for declared capabilities", () => {
+  const client = new Client(
+    {
+      name: "test client",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        sampling: {},
+      },
+    },
+  );
+
+  // This should work because sampling is a declared capability
+  expect(() => {
+    client.setRequestHandler(CreateMessageRequestSchema, () => ({
+      model: "test-model",
+      role: "assistant",
+      content: {
+        type: "text",
+        text: "Test response",
+      },
+    }));
+  }).not.toThrow();
+
+  // This should throw because roots listing is not a declared capability
+  expect(() => {
+    client.setRequestHandler(ListRootsRequestSchema, () => ({}));
+  }).toThrow("Client does not support roots capability");
+});
+
 /*
   Test that custom request/notification/result schemas can be used with the Client class.
   */

src/client/index.ts

@@ -248,6 +248,26 @@ export class Client<
     }
   }
 
+  protected assertRequestHandlerCapability(method: string): void {
+    switch (method) {
+      case "sampling/createMessage":
+        if (!this._capabilities.sampling) {
+          throw new Error("Client does not support sampling capability");
+        }
+        break;
+
+      case "roots/list":
+        if (!this._capabilities.roots) {
+          throw new Error("Client does not support roots capability");
+        }
+        break;
+
+      case "ping":
+        // No specific capability required for ping
+        break;
+    }
+  }
+
   async ping() {
     return this.request({ method: "ping" }, EmptyResultSchema);
   }

src/server/index.test.ts

@@ -10,6 +10,10 @@ import {
   LATEST_PROTOCOL_VERSION,
   SUPPORTED_PROTOCOL_VERSIONS,
   CreateMessageRequestSchema,
+  ListPromptsRequestSchema,
+  ListResourcesRequestSchema,
+  ListToolsRequestSchema,
+  SetLevelRequestSchema,
 } from "../types.js";
 import { Transport } from "../shared/transport.js";
 import { InMemoryTransport } from "../inMemory.js";
@@ -293,6 +297,41 @@ test("should respect server notification capabilities", async () => {
   ).rejects.toThrow(/^Server does not support/);
 });
 
+test("should only allow setRequestHandler for declared capabilities", () => {
+  const server = new Server(
+    {
+      name: "test server",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        prompts: {},
+        resources: {},
+      },
+    },
+  );
+
+  // These should work because the capabilities are declared
+  expect(() => {
+    server.setRequestHandler(ListPromptsRequestSchema, () => ({ prompts: [] }));
+  }).not.toThrow();
+
+  expect(() => {
+    server.setRequestHandler(ListResourcesRequestSchema, () => ({
+      resources: [],
+    }));
+  }).not.toThrow();
+
+  // These should throw because the capabilities are not declared
+  expect(() => {
+    server.setRequestHandler(ListToolsRequestSchema, () => ({ tools: [] }));
+  }).toThrow(/^Server does not support tools/);
+
+  expect(() => {
+    server.setRequestHandler(SetLevelRequestSchema, () => ({}));
+  }).toThrow(/^Server does not support logging/);
+});
+
 /*
   Test that custom request/notification/result schemas can be used with the Server class.
   */

src/server/index.ts

@@ -163,6 +163,59 @@ export class Server<
     }
   }
 
+  protected assertRequestHandlerCapability(method: string): void {
+    switch (method) {
+      case "sampling/createMessage":
+        if (!this._capabilities.sampling) {
+          throw new Error(
+            `Server does not support sampling (required for ${method})`,
+          );
+        }
+        break;
+
+      case "logging/setLevel":
+        if (!this._capabilities.logging) {
+          throw new Error(
+            `Server does not support logging (required for ${method})`,
+          );
+        }
+        break;
+
+      case "prompts/get":
+      case "prompts/list":
+        if (!this._capabilities.prompts) {
+          throw new Error(
+            `Server does not support prompts (required for ${method})`,
+          );
+        }
+        break;
+
+      case "resources/list":
+      case "resources/templates/list":
+      case "resources/read":
+        if (!this._capabilities.resources) {
+          throw new Error(
+            `Server does not support resources (required for ${method})`,
+          );
+        }
+        break;
+
+      case "tools/call":
+      case "tools/list":
+        if (!this._capabilities.tools) {
+          throw new Error(
+            `Server does not support tools (required for ${method})`,
+          );
+        }
+        break;
+
+      case "ping":
+      case "initialize":
+        // No specific capability required for these methods
+        break;
+    }
+  }
+
   private async _oninitialize(
     request: InitializeRequest,
   ): Promise<InitializeResult> {

src/shared/protocol.ts

@@ -277,6 +277,13 @@ export abstract class Protocol<
     method: SendNotificationT["method"],
   ): void;
 
+  /**
+   * A method to check if a request handler is supported by the local side, for the given method to be handled.
+   *
+   * This should be implemented by subclasses.
+   */
+  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).
    *
@@ -360,7 +367,9 @@ export abstract class Protocol<
     requestSchema: T,
     handler: (request: z.infer<T>) => SendResultT | Promise<SendResultT>,
   ): void {
-    this._requestHandlers.set(requestSchema.shape.method.value, (request) =>
+    const method = requestSchema.shape.method.value;
+    this.assertRequestHandlerCapability(method);
+    this._requestHandlers.set(method, (request) =>
       Promise.resolve(handler(requestSchema.parse(request))),
     );
   }