Merge pull request #13 from modelcontextprotocol/ochafik/easier-notification-handlers

sdk: add notification handlers (App.ontoolinput, .ontoolresult, AppBridge.onsizechange...)

Author: ochafik

examples/simple-server/src/ui-raw.ts

@@ -8,7 +8,7 @@
  * <code>
  * npx esbuild src/ui-raw.ts --bundle --outfile=dist/ui-raw.js --minify --sourcemap --platform=browser
  * </code>
- * 
+ *
  * We implement a barebones JSON-RPC message sender/receiver (see `app` object below),
  * but without timeouts or runtime type validation of any kind
  * (for that, use the Apps SDK / see ui-vanilla.ts or ui-react.ts).

examples/simple-server/src/ui-react.tsx

@@ -26,19 +26,9 @@ export function McpClientApp() {
     appInfo: APP_INFO,
     capabilities: {},
     onAppCreated: (app) => {
-      app.setNotificationHandler(
-        McpUiToolResultNotificationSchema,
-        async (notification) => {
-          setToolResults((prev) => [...prev, notification.params]);
-        },
-      );
-      app.setNotificationHandler(
-        McpUiSizeChangeNotificationSchema,
-        async (notification) => {
-          document.body.style.width = `${notification.params.width}px`;
-          document.body.style.height = `${notification.params.height}px`;
-        },
-      );
+      app.ontoolresult = async (params) => {
+        setToolResults((prev) => [...prev, params]);
+      };
     },
   });
 

examples/simple-server/src/ui-vanilla.ts

@@ -36,28 +36,17 @@ window.addEventListener("load", async () => {
     version: "1.0.0",
   });
 
-  app.setNotificationHandler(
-    McpUiToolInputNotificationSchema,
-    async ({ params }) => {
-      appendText(
-        `Tool call input received: ${JSON.stringify(params.arguments)}`,
-      );
-    },
-  );
-  app.setNotificationHandler(
-    McpUiToolResultNotificationSchema,
-    async ({ params: { content, structuredContent, isError } }) => {
-      appendText(
-        `Tool call result received: isError=${isError}, content=${content}, structuredContent=${JSON.stringify(structuredContent)}`,
-      );
-    },
-  );
-  app.setNotificationHandler(
-    McpUiHostContextChangedNotificationSchema,
-    async (params) => {
-      appendText(`Host context changed: ${JSON.stringify(params)}`);
-    },
-  );
+  app.ontoolinput = (params) => {
+    appendText(`Tool call input received: ${JSON.stringify(params.arguments)}`);
+  };
+  app.ontoolresult = ({ content, structuredContent, isError }) => {
+    appendText(
+      `Tool call result received: isError=${isError}, content=${content}, structuredContent=${JSON.stringify(structuredContent)}`,
+    );
+  };
+  app.onhostcontextchanged = (params) => {
+    appendText(`Host context changed: ${JSON.stringify(params)}`);
+  };
 
   document.body.addEventListener("resize", () => {
     app.sendSizeChange({

src/app-bridge.ts

@@ -13,6 +13,7 @@ import {
   ListResourceTemplatesRequestSchema,
   ListResourceTemplatesResultSchema,
   Notification,
+  PingRequest,
   PingRequestSchema,
   PromptListChangedNotificationSchema,
   ResourceListChangedNotificationSchema,
@@ -34,12 +35,16 @@ import {
   LATEST_PROTOCOL_VERSION,
   McpUiAppCapabilities,
   McpUiHostCapabilities,
+  McpUiInitializedNotification,
   McpUiInitializedNotificationSchema,
   McpUiInitializeRequest,
   McpUiInitializeRequestSchema,
   McpUiInitializeResult,
   McpUiResourceTeardownRequest,
   McpUiResourceTeardownResultSchema,
+  McpUiSandboxProxyReadyNotification,
+  McpUiSandboxProxyReadyNotificationSchema,
+  McpUiSizeChangeNotificationSchema,
 } from "./types";
 export * from "./types";
 export { PostMessageTransport } from "./message-transport";
@@ -48,9 +53,11 @@ type HostOptions = ProtocolOptions;
 
 export const SUPPORTED_PROTOCOL_VERSIONS = [LATEST_PROTOCOL_VERSION];
 
-export class AppBridge extends Protocol<Request, Notification, Result> {
-  oninitialized?: () => void;
+type RequestExtra = Parameters<
+  Parameters<AppBridge["setRequestHandler"]>[1]
+>[1];
 
+export class AppBridge extends Protocol<Request, Notification, Result> {
   private _appCapabilities?: McpUiAppCapabilities;
 
   constructor(
@@ -64,16 +71,37 @@ export class AppBridge extends Protocol<Request, Notification, Result> {
     this.setRequestHandler(McpUiInitializeRequestSchema, (request) =>
       this._oninitialize(request),
     );
-    this.setNotificationHandler(McpUiInitializedNotificationSchema, () =>
-      this.oninitialized?.(),
-    );
 
-    this.setRequestHandler(PingRequestSchema, (request) => {
-      console.log("Received ping:", request.params);
+    this.setRequestHandler(PingRequestSchema, (request, extra) => {
+      this.onping?.(request.params, extra);
       return {};
     });
   }
 
+  onping?: (params: PingRequest["params"], extra: RequestExtra) => void;
+
+  set onsizechange(
+    callback: (params: McpUiSizeChangeNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(McpUiSizeChangeNotificationSchema, (n) =>
+      callback(n.params),
+    );
+  }
+  set onsandboxready(
+    callback: (params: McpUiSandboxProxyReadyNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(McpUiSandboxProxyReadyNotificationSchema, (n) =>
+      callback(n.params),
+    );
+  }
+  set oninitialized(
+    callback: (params: McpUiInitializedNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(McpUiInitializedNotificationSchema, (n) =>
+      callback(n.params),
+    );
+  }
+
   assertCapabilityForMethod(method: Request["method"]): void {
     // TODO
   }

src/app.ts

@@ -19,6 +19,8 @@ import {
   LATEST_PROTOCOL_VERSION,
   McpUiAppCapabilities,
   McpUiHostCapabilities,
+  McpUiHostContextChangedNotification,
+  McpUiHostContextChangedNotificationSchema,
   McpUiInitializedNotification,
   McpUiInitializeRequest,
   McpUiInitializeResultSchema,
@@ -27,6 +29,12 @@ import {
   McpUiOpenLinkRequest,
   McpUiOpenLinkResultSchema,
   McpUiSizeChangeNotification,
+  McpUiToolInputNotification,
+  McpUiToolInputNotificationSchema,
+  McpUiToolInputPartialNotification,
+  McpUiToolInputPartialNotificationSchema,
+  McpUiToolResultNotification,
+  McpUiToolResultNotificationSchema,
 } from "./types";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 
@@ -54,6 +62,36 @@ export class App extends Protocol<Request, Notification, Result> {
     });
   }
 
+  set ontoolinput(
+    callback: (params: McpUiToolInputNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(McpUiToolInputNotificationSchema, (n) =>
+      callback(n.params),
+    );
+  }
+  set ontoolinputpartial(
+    callback: (params: McpUiToolInputPartialNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(McpUiToolInputPartialNotificationSchema, (n) =>
+      callback(n.params),
+    );
+  }
+  set ontoolresult(
+    callback: (params: McpUiToolResultNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(McpUiToolResultNotificationSchema, (n) =>
+      callback(n.params),
+    );
+  }
+  set onhostcontextchanged(
+    callback: (params: McpUiHostContextChangedNotification["params"]) => void,
+  ) {
+    this.setNotificationHandler(
+      McpUiHostContextChangedNotificationSchema,
+      (n) => callback(n.params),
+    );
+  }
+
   assertCapabilityForMethod(method: Request["method"]): void {
     // TODO
   }

src/react/useApp.tsx

@@ -9,8 +9,7 @@ export interface UseAppOptions {
   capabilities: McpUiAppCapabilities;
   /**
    * Called after client is created but before connection.
-   * Use this to register request/notification handlers via
-   * client.setRequestHandler() and client.setNotificationHandler().
+   * Use this to register handlers via app.ontoolinput, app.toolresult, etc.
    */
   onAppCreated?: (app: App) => void;
 }