diff --git a/src/message-transport.examples.ts b/src/message-transport.examples.ts index b42184bf..afa85b1b 100644 --- a/src/message-transport.examples.ts +++ b/src/message-transport.examples.ts @@ -56,3 +56,51 @@ function PostMessageTransport_constructor_host() { ); //#endregion PostMessageTransport_constructor_host } + +/** + * Example: Host with deferred target for srcdoc iframes. + * + * When loading View HTML via `srcdoc`, `contentWindow` is not available until + * the iframe loads. Create the transport with `null` target, connect the bridge + * (which starts listening for messages), set `srcdoc`, then call `setTarget()` + * after the iframe loads to flush queued outgoing messages. + */ +async function PostMessageTransport_deferred( + bridge: AppBridge, + htmlContent: string, +) { + //#region PostMessageTransport_deferred + const iframe = document.createElement("iframe"); + iframe.sandbox.add("allow-scripts"); + document.body.appendChild(iframe); + + const transport = new PostMessageTransport(null, null); + await bridge.connect(transport); + + iframe.srcdoc = htmlContent; + iframe.onload = () => { + transport.setTarget(iframe.contentWindow!); + }; + //#endregion PostMessageTransport_deferred +} + +/** + * Example: Creating deferred transport (constructor only). + */ +function PostMessageTransport_constructor_deferred() { + //#region PostMessageTransport_constructor_deferred + const transport = new PostMessageTransport(null, null); + //#endregion PostMessageTransport_constructor_deferred +} + +/** + * Example: Setting the target after iframe loads. + */ +function PostMessageTransport_setTarget(transport: PostMessageTransport) { + //#region PostMessageTransport_setTarget + const iframe = document.getElementById("app-iframe") as HTMLIFrameElement; + iframe.onload = () => { + transport.setTarget(iframe.contentWindow!); + }; + //#endregion PostMessageTransport_setTarget +} diff --git a/src/message-transport.test.ts b/src/message-transport.test.ts new file mode 100644 index 00000000..3aa486bf --- /dev/null +++ b/src/message-transport.test.ts @@ -0,0 +1,108 @@ +import { describe, it, expect, beforeEach, afterEach } from "bun:test"; +import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js"; + +import { PostMessageTransport } from "./message-transport"; + +const makeMessage = (id: number): JSONRPCMessage => ({ + jsonrpc: "2.0", + id, + method: "test", + params: {}, +}); + +type MessageHandler = (ev: MessageEvent) => void; + +let messageListeners: Set; + +const mockWindow = { + addEventListener: (type: string, listener: MessageHandler) => { + if (type === "message") messageListeners.add(listener); + }, + removeEventListener: (type: string, listener: MessageHandler) => { + if (type === "message") messageListeners.delete(listener); + }, +} as unknown as Window & typeof globalThis; + +function dispatchMessage(data: unknown, source?: MessageEventSource | null) { + const event = { data, source, origin: "null" } as MessageEvent; + for (const listener of messageListeners) { + listener(event); + } +} + +describe("PostMessageTransport", () => { + let savedWindow: typeof globalThis.window; + + beforeEach(() => { + messageListeners = new Set(); + savedWindow = globalThis.window; + (globalThis as any).window = mockWindow; + }); + + afterEach(() => { + (globalThis as any).window = savedWindow; + }); + + it("queues messages when target is null and flushes on setTarget", async () => { + // Arrange + const transport = new PostMessageTransport(null, null); + await transport.start(); + const posted: unknown[] = []; + const fakeTarget = { + postMessage: (data: unknown, _origin: string) => posted.push(data), + } as unknown as Window; + + // Act + await transport.send(makeMessage(1)); + await transport.send(makeMessage(2)); + transport.setTarget(fakeTarget); + + // Assert + expect(posted).toEqual([makeMessage(1), makeMessage(2)]); + expect((transport as any)._sendQueue).toHaveLength(0); + }); + + it("receives messages with null eventSource before setTarget", async () => { + // Arrange + const transport = new PostMessageTransport(null, null); + const received: JSONRPCMessage[] = []; + transport.onmessage = (msg) => received.push(msg); + await transport.start(); + + // Act + dispatchMessage(makeMessage(1)); + + // Assert + expect(received).toHaveLength(1); + }); + + it("sends directly when target is provided (backward compat)", async () => { + // Arrange + const posted: unknown[] = []; + const fakeTarget = { + postMessage: (data: unknown, _origin: string) => posted.push(data), + } as unknown as Window; + const transport = new PostMessageTransport(fakeTarget, fakeTarget); + await transport.start(); + + // Act + await transport.send(makeMessage(1)); + + // Assert + expect(posted).toEqual([makeMessage(1)]); + expect((transport as any)._sendQueue).toHaveLength(0); + }); + + it("clears the send queue on close", async () => { + // Arrange + const transport = new PostMessageTransport(null, null); + await transport.start(); + await transport.send(makeMessage(1)); + + // Act + await transport.close(); + + // Assert + expect((transport as any)._sendQueue).toHaveLength(0); + }); +}); diff --git a/src/message-transport.ts b/src/message-transport.ts index dded7046..b252d424 100644 --- a/src/message-transport.ts +++ b/src/message-transport.ts @@ -17,9 +17,31 @@ import { * * ## Security * - * The `eventSource` parameter is required and validates the message source window - * by checking `event.source`. For views, pass `window.parent`. + * The `eventSource` parameter validates the message source window by checking + * `event.source`. For views, pass `window.parent`. * For hosts, pass `iframe.contentWindow` to validate the iframe source. + * When `null`, all sources are accepted (useful when the iframe hasn't loaded yet). + * + * ## Deferred Target (srcdoc iframes) + * + * When the host creates an iframe dynamically (e.g., via `srcdoc`), `contentWindow` + * is not available until the iframe loads. Pass `null` as `eventTarget` and call + * {@link setTarget `setTarget`} after the iframe loads. Outgoing messages are queued + * until the target is set. + * + * ```ts source="./message-transport.examples.ts#PostMessageTransport_deferred" + * const iframe = document.createElement("iframe"); + * iframe.sandbox.add("allow-scripts"); + * document.body.appendChild(iframe); + * + * const transport = new PostMessageTransport(null, null); + * await bridge.connect(transport); + * + * iframe.srcdoc = htmlContent; + * iframe.onload = () => { + * transport.setTarget(iframe.contentWindow!); + * }; + * ``` * * ## Usage * @@ -43,6 +65,9 @@ import { * @see {@link app-bridge!AppBridge.connect `AppBridge.connect`} for Host usage */ export class PostMessageTransport implements Transport { + private _eventTarget: Window | null; + private _eventSource: MessageEventSource | null; + private _sendQueue: JSONRPCMessage[] = []; private messageListener: ( this: Window, ev: WindowEventMap["message"], @@ -51,9 +76,12 @@ export class PostMessageTransport implements Transport { /** * Create a new PostMessageTransport. * - * @param eventTarget - Target window to send messages to (default: `window.parent`) + * @param eventTarget - Target window to send messages to. Pass `null` to defer + * — outgoing messages will be queued until {@link setTarget `setTarget`} is called. + * Defaults to `window.parent` for View usage. * @param eventSource - Source window for message validation. For views, pass - * `window.parent`. For hosts, pass `iframe.contentWindow`. + * `window.parent`. For hosts, pass `iframe.contentWindow`. Pass `null` to + * accept messages from any source (useful for deferred/srcdoc iframes). * * @example View connecting to parent * ```ts source="./message-transport.examples.ts#PostMessageTransport_constructor_view" @@ -68,13 +96,20 @@ export class PostMessageTransport implements Transport { * iframe.contentWindow!, * ); * ``` + * + * @example Host with deferred target (srcdoc) + * ```ts source="./message-transport.examples.ts#PostMessageTransport_constructor_deferred" + * const transport = new PostMessageTransport(null, null); + * ``` */ constructor( - private eventTarget: Window = window.parent, - private eventSource: MessageEventSource, + eventTarget: Window | null = window.parent, + eventSource: MessageEventSource | null, ) { + this._eventTarget = eventTarget; + this._eventSource = eventSource; this.messageListener = (event) => { - if (eventSource && event.source !== this.eventSource) { + if (this._eventSource && event.source !== this._eventSource) { console.debug("Ignoring message from unknown source", event); return; } @@ -103,6 +138,35 @@ export class PostMessageTransport implements Transport { }; } + /** + * Set or update the target window for outgoing messages. + * + * When the transport was created with a `null` target (deferred mode), call + * this method after the iframe loads to provide `contentWindow` and flush + * any queued messages. Also updates the event source for incoming message + * validation. + * + * @param target - The iframe's `contentWindow` to send messages to + * @param eventSource - Optional new event source for message validation. + * Defaults to `target`, which is correct for most iframe setups. + * + * @example + * ```ts source="./message-transport.examples.ts#PostMessageTransport_setTarget" + * const iframe = document.getElementById("app-iframe") as HTMLIFrameElement; + * iframe.onload = () => { + * transport.setTarget(iframe.contentWindow!); + * }; + * ``` + */ + setTarget(target: Window, eventSource?: MessageEventSource): void { + this._eventTarget = target; + this._eventSource = eventSource ?? target; + for (const message of this._sendQueue) { + this._eventTarget.postMessage(message, "*"); + } + this._sendQueue = []; + } + /** * Begin listening for messages from the event source. * @@ -116,24 +180,32 @@ export class PostMessageTransport implements Transport { /** * Send a JSON-RPC message to the target window. * - * Messages are sent using `postMessage` with `"*"` origin, meaning they are visible - * to all frames. The receiver should validate the message source for security. + * When the target is set, messages are sent immediately using `postMessage` + * with `"*"` origin. When the target is `null` (deferred mode), messages are + * queued and flushed when {@link setTarget `setTarget`} is called. * * @param message - JSON-RPC message to send * @param options - Optional send options (currently unused) */ async send(message: JSONRPCMessage, options?: TransportSendOptions) { - console.debug("Sending message", message); - this.eventTarget.postMessage(message, "*"); + if (this._eventTarget) { + console.debug("Sending message", message); + this._eventTarget.postMessage(message, "*"); + } else { + console.debug("Queuing message (target not set)", message); + this._sendQueue.push(message); + } } /** * Stop listening for messages and cleanup. * - * Removes the message event listener and calls the {@link onclose `onclose`} callback if set. + * Removes the message event listener, clears any queued messages, and calls + * the {@link onclose `onclose`} callback if set. */ async close() { window.removeEventListener("message", this.messageListener); + this._sendQueue = []; this.onclose?.(); }