From dd460bd877eb9df356ce02a35b9e8cc2ef5222ef Mon Sep 17 00:00:00 2001
From: Nathan Arseneau <arseneau.nathan@gmail.com>
Date: Mon, 17 Mar 2025 02:44:59 -0400
Subject: [PATCH 1/2] Refactor notification handling to include all
 notifications

---
 client/src/lib/hooks/useConnection.ts | 23 ++++++-----------------
 client/src/lib/notificationTypes.ts   |  4 +++-
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts
index 19b0e9d..d432c28 100644
--- a/client/src/lib/hooks/useConnection.ts
+++ b/client/src/lib/hooks/useConnection.ts
@@ -8,9 +8,6 @@ import {
   ClientRequest,
   CreateMessageRequestSchema,
   ListRootsRequestSchema,
-  ProgressNotificationSchema,
-  ResourceUpdatedNotificationSchema,
-  LoggingMessageNotificationSchema,
   Request,
   Result,
   ServerCapabilities,
@@ -250,20 +247,12 @@ export function useConnection({
       });
 
       if (onNotification) {
-        client.setNotificationHandler(
-          ProgressNotificationSchema,
-          onNotification,
-        );
-
-        client.setNotificationHandler(
-          ResourceUpdatedNotificationSchema,
-          onNotification,
-        );
-
-        client.setNotificationHandler(
-          LoggingMessageNotificationSchema,
-          onNotification,
-        );
+        client.fallbackNotificationHandler = (
+          notification: Notification,
+        ): Promise<void> => {
+          onNotification(notification);
+          return Promise.resolve();
+        };
       }
 
       if (onStdErrNotification) {
diff --git a/client/src/lib/notificationTypes.ts b/client/src/lib/notificationTypes.ts
index 82c1fd8..8627ccc 100644
--- a/client/src/lib/notificationTypes.ts
+++ b/client/src/lib/notificationTypes.ts
@@ -14,7 +14,9 @@ export const StdErrNotificationSchema = BaseNotificationSchema.extend({
 
 export const NotificationSchema = ClientNotificationSchema.or(
   StdErrNotificationSchema,
-).or(ServerNotificationSchema);
+)
+  .or(ServerNotificationSchema)
+  .or(BaseNotificationSchema);
 
 export type StdErrNotification = z.infer<typeof StdErrNotificationSchema>;
 export type Notification = z.infer<typeof NotificationSchema>;

From 27b54104c1523f06e6fa3a9ad54a249a74f9f6b2 Mon Sep 17 00:00:00 2001
From: Nathan Arseneau <arseneau.nathan@gmail.com>
Date: Wed, 19 Mar 2025 19:51:01 -0400
Subject: [PATCH 2/2] Add support for additional notification schemas in
 useConnection hook

---
 client/src/lib/hooks/useConnection.ts | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts
index d432c28..f0379eb 100644
--- a/client/src/lib/hooks/useConnection.ts
+++ b/client/src/lib/hooks/useConnection.ts
@@ -8,6 +8,9 @@ import {
   ClientRequest,
   CreateMessageRequestSchema,
   ListRootsRequestSchema,
+  ProgressNotificationSchema,
+  ResourceUpdatedNotificationSchema,
+  LoggingMessageNotificationSchema,
   Request,
   Result,
   ServerCapabilities,
@@ -16,6 +19,10 @@ import {
   McpError,
   CompleteResultSchema,
   ErrorCode,
+  CancelledNotificationSchema,
+  ResourceListChangedNotificationSchema,
+  ToolListChangedNotificationSchema,
+  PromptListChangedNotificationSchema,
 } from "@modelcontextprotocol/sdk/types.js";
 import { useState } from "react";
 import { toast } from "react-toastify";
@@ -247,6 +254,18 @@ export function useConnection({
       });
 
       if (onNotification) {
+        [
+          CancelledNotificationSchema,
+          ProgressNotificationSchema,
+          LoggingMessageNotificationSchema,
+          ResourceUpdatedNotificationSchema,
+          ResourceListChangedNotificationSchema,
+          ToolListChangedNotificationSchema,
+          PromptListChangedNotificationSchema,
+        ].forEach((notificationSchema) => {
+          client.setNotificationHandler(notificationSchema, onNotification);
+        });
+
         client.fallbackNotificationHandler = (
           notification: Notification,
         ): Promise<void> => {