From 8b307a5d62e9b0a27950c5e6d81130fca1ecf0cd Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Tue, 12 Nov 2024 12:16:43 +0000 Subject: [PATCH] Check capabilities when request handlers are set --- src/client/index.test.ts | 33 +++++++++++++++++++++++++ src/client/index.ts | 20 +++++++++++++++ src/server/index.test.ts | 39 +++++++++++++++++++++++++++++ src/server/index.ts | 53 ++++++++++++++++++++++++++++++++++++++++ src/shared/protocol.ts | 11 ++++++++- 5 files changed, 155 insertions(+), 1 deletion(-) diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 87174308..5610a629 100644 --- a/src/client/index.test.ts +++ b/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. */ diff --git a/src/client/index.ts b/src/client/index.ts index 313c1218..dd5d2ced 100644 --- a/src/client/index.ts +++ b/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); } diff --git a/src/server/index.test.ts b/src/server/index.test.ts index 1d9e0190..d30c670b 100644 --- a/src/server/index.test.ts +++ b/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. */ diff --git a/src/server/index.ts b/src/server/index.ts index 5d4e5141..ecb525b5 100644 --- a/src/server/index.ts +++ b/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 { diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 8b1705be..85610a9d 100644 --- a/src/shared/protocol.ts +++ b/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) => SendResultT | Promise, ): 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))), ); }