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] 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> => {