Check capabilities when request handlers are set

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.
   */