From 64c28628839cc961d6940dfd7da03f16d6c34a6b Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Tue, 12 Nov 2024 12:02:10 +0000 Subject: [PATCH] Check and add tests for notification capabilities too --- src/client/index.test.ts | 102 +++++++++++++++++++++++++++++++++++++++ src/client/index.ts | 22 +++++++++ src/server/index.test.ts | 33 +++++++++++++ src/server/index.ts | 56 ++++++++++++++++----- src/shared/protocol.ts | 13 +++++ 5 files changed, 214 insertions(+), 12 deletions(-) diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 7cbafd17..87174308 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -223,6 +223,108 @@ test("should respect server capabilities", async () => { ); }); +test("should respect client notification capabilities", async () => { + const server = new Server( + { + name: "test server", + version: "1.0", + }, + { + capabilities: {}, + }, + ); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + roots: { + listChanged: true, + }, + }, + }, + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + // This should work because the client has the roots.listChanged capability + await expect(client.sendRootsListChanged()).resolves.not.toThrow(); + + // Create a new client without the roots.listChanged capability + const clientWithoutCapability = new Client( + { + name: "test client without capability", + version: "1.0", + }, + { + capabilities: {}, + enforceStrictCapabilities: true, + }, + ); + + await clientWithoutCapability.connect(clientTransport); + + // This should throw because the client doesn't have the roots.listChanged capability + await expect(clientWithoutCapability.sendRootsListChanged()).rejects.toThrow( + /^Client does not support/, + ); +}); + +test("should respect server notification capabilities", async () => { + const server = new Server( + { + name: "test server", + version: "1.0", + }, + { + capabilities: { + logging: {}, + resources: { + listChanged: true, + }, + }, + }, + ); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: {}, + }, + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + // These should work because the server has the corresponding capabilities + await expect( + server.sendLoggingMessage({ level: "info", data: "Test" }), + ).resolves.not.toThrow(); + await expect(server.sendResourceListChanged()).resolves.not.toThrow(); + + // This should throw because the server doesn't have the tools capability + await expect(server.sendToolListChanged()).rejects.toThrow( + "Server does not support notifying of tool list changes", + ); +}); + /* 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 402c3133..313c1218 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -226,6 +226,28 @@ export class Client< } } + protected assertNotificationCapability( + method: NotificationT["method"], + ): void { + switch (method as ClientNotification["method"]) { + case "notifications/roots/list_changed": + if (!this._capabilities.roots?.listChanged) { + throw new Error( + "Client does not support roots list changed notifications", + ); + } + break; + + case "notifications/initialized": + // No specific capability required for initialized + break; + + case "notifications/progress": + // Progress notifications are always allowed + break; + } + } + async ping() { return this.request({ method: "ping" }, EmptyResultSchema); } diff --git a/src/server/index.test.ts b/src/server/index.test.ts index 274bb07c..1d9e0190 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -260,6 +260,39 @@ test("should respect client capabilities", async () => { await expect(server.listRoots()).rejects.toThrow(/^Client does not support/); }); +test("should respect server notification capabilities", async () => { + const server = new Server( + { + name: "test server", + version: "1.0", + }, + { + capabilities: { + logging: {}, + }, + enforceStrictCapabilities: true, + }, + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await server.connect(serverTransport); + + // This should work because logging is supported by the server + await expect( + server.sendLoggingMessage({ + level: "info", + data: "Test log message", + }), + ).resolves.not.toThrow(); + + // This should throw because resource notificaitons are not supported by the server + await expect( + server.sendResourceUpdated({ uri: "test://resource" }), + ).rejects.toThrow(/^Server does not support/); +}); + /* 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 4558b348..5d4e5141 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -32,7 +32,7 @@ export type ServerOptions = ProtocolOptions & { /** * Capabilities to advertise as being supported by this server. */ - capabilities: ClientCapabilities; + capabilities: ServerCapabilities; }; /** @@ -120,6 +120,49 @@ export class Server< } } + protected assertNotificationCapability( + method: (ServerNotification | NotificationT)["method"], + ): void { + switch (method as ServerNotification["method"]) { + case "notifications/message": + if (!this._capabilities.logging) { + throw new Error( + `Server does not support logging (required for ${method})`, + ); + } + break; + + case "notifications/resources/updated": + case "notifications/resources/list_changed": + if (!this._capabilities.resources) { + throw new Error( + `Server does not support notifying about resources (required for ${method})`, + ); + } + break; + + case "notifications/tools/list_changed": + if (!this._capabilities.tools) { + throw new Error( + `Server does not support notifying of tool list changes (required for ${method})`, + ); + } + break; + + case "notifications/prompts/list_changed": + if (!this._capabilities.prompts) { + throw new Error( + `Server does not support notifying of prompt list changes (required for ${method})`, + ); + } + break; + + case "notifications/progress": + // Progress notifications are always allowed + break; + } + } + private async _oninitialize( request: InitializeRequest, ): Promise { @@ -155,17 +198,6 @@ export class Server< return this._capabilities; } - private assertCapability( - capability: keyof ClientCapabilities, - method: string, - ) { - if (!this._clientCapabilities?.[capability]) { - throw new Error( - `Client does not support ${capability} (required for ${method})`, - ); - } - } - async ping() { return this.request({ method: "ping" }, EmptyResultSchema); } diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index ae305413..8b1705be 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -28,6 +28,8 @@ export type ProtocolOptions = { /** * Whether to restrict emitted requests to only those that the remote side has indicated that they can handle, through their advertised capabilities. * + * Note that this DOES NOT affect checking of _local_ side capabilities, as it is considered a logic error to mis-specify those. + * * Currently this defaults to false, for backwards compatibility with SDK versions that did not advertise capabilities correctly. In future, this will default to true. */ enforceStrictCapabilities?: boolean; @@ -266,6 +268,15 @@ export abstract class Protocol< method: SendRequestT["method"], ): void; + /** + * A method to check if a notification is supported by the local side, for the given method to be sent. + * + * This should be implemented by subclasses. + */ + protected abstract assertNotificationCapability( + method: SendNotificationT["method"], + ): void; + /** * Sends a request and wait for a response, with optional progress notifications in the meantime (if supported by the server). * @@ -326,6 +337,8 @@ export abstract class Protocol< throw new Error("Not connected"); } + this.assertNotificationCapability(notification.method); + const jsonrpcNotification: JSONRPCNotification = { ...notification, jsonrpc: "2.0",