Check and add tests for notification capabilities too

Author: jspahrsummers

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

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);
   }

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

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<InitializeResult> {
@@ -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);
   }

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",