diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index 19b0e9d..f0379eb 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -19,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"; @@ -250,20 +254,24 @@ export function useConnection({ }); if (onNotification) { - client.setNotificationHandler( + [ + CancelledNotificationSchema, ProgressNotificationSchema, - onNotification, - ); - - client.setNotificationHandler( - ResourceUpdatedNotificationSchema, - onNotification, - ); - - client.setNotificationHandler( LoggingMessageNotificationSchema, - onNotification, - ); + ResourceUpdatedNotificationSchema, + ResourceListChangedNotificationSchema, + ToolListChangedNotificationSchema, + PromptListChangedNotificationSchema, + ].forEach((notificationSchema) => { + client.setNotificationHandler(notificationSchema, onNotification); + }); + + client.fallbackNotificationHandler = ( + notification: Notification, + ): Promise => { + 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; export type Notification = z.infer;