setup notification handlers in App & AppBridge

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