feat: implement subscribe/unsubscribe methods for multiple event handlers

Author: Avcharov

src/app-bridge.test.ts

@@ -13,7 +13,7 @@ import {
   ToolListChangedNotificationSchema,
 } from "@modelcontextprotocol/sdk/types.js";
 
-import { App } from "./app";
+import { App, type NotificationEventMap } from "./app";
 import {
   AppBridge,
   getToolUiResourceUri,
@@ -942,6 +942,144 @@ describe("App <-> AppBridge integration", () => {
       expect(receivedNotifications).toHaveLength(1);
     });
   });
+
+  describe("App.subscribe and App.unsubscribe", () => {
+    it("multiple subscribers all receive the same notification", async () => {
+      const received1: unknown[] = [];
+      const received2: unknown[] = [];
+
+      app.subscribe("hostcontextchanged", (params) => received1.push(params));
+      app.subscribe("hostcontextchanged", (params) => received2.push(params));
+
+      await bridge.connect(bridgeTransport);
+      await app.connect(appTransport);
+
+      bridge.setHostContext({ theme: "dark" });
+      await flush();
+
+      expect(received1).toEqual([{ theme: "dark" }]);
+      expect(received2).toEqual([{ theme: "dark" }]);
+    });
+
+    it("unsubscribe removes only the targeted subscriber", async () => {
+      const received1: unknown[] = [];
+      const received2: unknown[] = [];
+
+      const handler1 = (params: NotificationEventMap["hostcontextchanged"]) =>
+        received1.push(params);
+      const handler2 = (params: NotificationEventMap["hostcontextchanged"]) =>
+        received2.push(params);
+
+      app.subscribe("hostcontextchanged", handler1);
+      app.subscribe("hostcontextchanged", handler2);
+
+      await bridge.connect(bridgeTransport);
+      await app.connect(appTransport);
+
+      bridge.setHostContext({ theme: "dark" });
+      await flush();
+
+      // Both fired
+      expect(received1).toHaveLength(1);
+      expect(received2).toHaveLength(1);
+
+      app.unsubscribe("hostcontextchanged", handler1);
+
+      bridge.setHostContext({ theme: "light" });
+      await flush();
+
+      // Only handler2 fired for the second update
+      expect(received1).toHaveLength(1);
+      expect(received2).toHaveLength(2);
+    });
+
+    it("subscribe returns an unsubscribe function that works correctly", async () => {
+      const received: unknown[] = [];
+
+      const unsubscribe = app.subscribe("toolinput", (params) =>
+        received.push(params),
+      );
+
+      await bridge.connect(bridgeTransport);
+      await app.connect(appTransport);
+
+      bridge.sendToolInput({ arguments: { x: 1 } });
+      await flush();
+      expect(received).toHaveLength(1);
+
+      unsubscribe();
+
+      bridge.sendToolInput({ arguments: { x: 2 } });
+      await flush();
+      // No new events after unsubscribing
+      expect(received).toHaveLength(1);
+    });
+
+    it("setter and subscribers fire independently", async () => {
+      const setterReceived: unknown[] = [];
+      const subscriberReceived: unknown[] = [];
+
+      app.onhostcontextchanged = (params) => setterReceived.push(params);
+      app.subscribe("hostcontextchanged", (params) =>
+        subscriberReceived.push(params),
+      );
+
+      await bridge.connect(bridgeTransport);
+      await app.connect(appTransport);
+
+      bridge.setHostContext({ theme: "dark" });
+      await flush();
+
+      expect(setterReceived).toEqual([{ theme: "dark" }]);
+      expect(subscriberReceived).toEqual([{ theme: "dark" }]);
+    });
+
+    it("context merge runs before setter and subscriber callbacks", async () => {
+      let contextInSetter: unknown;
+      let contextInSubscriber: unknown;
+
+      app.onhostcontextchanged = () => {
+        contextInSetter = app.getHostContext();
+      };
+      app.subscribe("hostcontextchanged", () => {
+        contextInSubscriber = app.getHostContext();
+      });
+
+      await bridge.connect(bridgeTransport);
+      await app.connect(appTransport);
+
+      bridge.setHostContext({ theme: "dark" });
+      await flush();
+
+      expect((contextInSetter as { theme: string })?.theme).toBe("dark");
+      expect((contextInSubscriber as { theme: string })?.theme).toBe("dark");
+    });
+
+    it("unsubscribing all subscribers does not affect the setter callback", async () => {
+      const setterReceived: unknown[] = [];
+      const subscriberReceived: unknown[] = [];
+
+      app.onhostcontextchanged = (params) => setterReceived.push(params);
+      const unsub = app.subscribe("hostcontextchanged", (params) =>
+        subscriberReceived.push(params),
+      );
+
+      await bridge.connect(bridgeTransport);
+      await app.connect(appTransport);
+
+      bridge.setHostContext({ theme: "dark" });
+      await flush();
+
+      unsub();
+
+      bridge.setHostContext({ theme: "light" });
+      await flush();
+
+      // Setter still fires after subscriber is removed
+      expect(setterReceived).toHaveLength(2);
+      expect(subscriberReceived).toHaveLength(1);
+    });
+  });
 });
 
 describe("getToolUiResourceUri", () => {

src/app.examples.ts

@@ -18,6 +18,7 @@ import {
   RESOURCE_URI_META_KEY,
   McpUiToolMeta,
 } from "./app.js";
+import type { McpUiHostContextChangedNotification } from "./types.js";
 import { registerAppTool } from "./server/index.js";
 
 /**
@@ -263,6 +264,55 @@ function App_onhostcontextchanged_respondToDisplayMode(app: App) {
   //#endregion App_onhostcontextchanged_respondToDisplayMode
 }
 
+/**
+ * Example: Subscribe to the same event from multiple places without conflict.
+ */
+async function App_subscribe_multipleHandlers(app: App) {
+  //#region App_subscribe_multipleHandlers
+  // Both handlers receive every notification — neither overrides the other
+  app.subscribe("hostcontextchanged", (ctx) => {
+    if (ctx.theme) applyTheme(ctx.theme);
+  });
+  app.subscribe("hostcontextchanged", (ctx) => {
+    if (ctx.styles?.css?.fonts) applyFonts(ctx.styles.css.fonts);
+  });
+  await app.connect();
+  //#endregion App_subscribe_multipleHandlers
+}
+
+/**
+ * Example: Unsubscribe using the function returned by subscribe.
+ */
+function App_subscribe_cleanup(app: App) {
+  //#region App_subscribe_cleanup
+  const unsubscribe = app.subscribe("toolinput", (params) => {
+    console.log("Tool input received:", params.arguments);
+  });
+
+  // Later, when cleanup is needed:
+  unsubscribe();
+  //#endregion App_subscribe_cleanup
+}
+
+/**
+ * Example: Unsubscribe by passing the original handler reference.
+ */
+function App_unsubscribe_explicit(app: App) {
+  //#region App_unsubscribe_explicit
+  const handler = (ctx: McpUiHostContextChangedNotification["params"]) => {
+    applyTheme(ctx.theme);
+  };
+  app.subscribe("hostcontextchanged", handler);
+
+  // Later, remove only this handler:
+  app.unsubscribe("hostcontextchanged", handler);
+  //#endregion App_unsubscribe_explicit
+}
+
+// Stubs for subscribe examples
+declare function applyTheme(theme: string | undefined): void;
+declare function applyFonts(fonts: string): void;
+
 /**
  * Example: Perform cleanup before teardown.
  */