Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 91 additions & 1 deletion src/app-bridge.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import {
import { z } from "zod/v4";

import { App } from "./app";
import { LATEST_PROTOCOL_VERSION } from "./types";
import {
AppBridge,
getToolUiResourceUri,
Expand Down Expand Up @@ -2021,6 +2022,45 @@ describe("App <-> AppBridge integration", () => {
expect(warnings()).toEqual([]);
});

it("does not warn on re-registration when a handler existed pre-connect", async () => {
const h = () => {};
app.addEventListener("toolresult", h);
await bridge.connect(bridgeTransport);
await app.connect(appTransport);

// React-style: useEffect cleanup removes, dep change re-adds.
app.removeEventListener("toolresult", h);
app.addEventListener("toolresult", () => {});

expect(warnings().filter((m) => lateMsg.test(m))).toEqual([]);
});

it("warns once for the first late registration, not on subsequent ones", async () => {
await bridge.connect(bridgeTransport);
await app.connect(appTransport);

app.addEventListener("toolinput", () => {});
app.addEventListener("toolinput", () => {});
app.addEventListener("toolinput", () => {});

const late = warnings().filter((m) => lateMsg.test(m));
expect(late).toHaveLength(1);
expect(late[0]).toContain('"toolinput"');
});

it("tracks first-handler per event independently", async () => {
app.addEventListener("toolinput", () => {});
await bridge.connect(bridgeTransport);
await app.connect(appTransport);

app.addEventListener("toolinput", () => {}); // had pre-connect handler → silent
app.addEventListener("toolresult", () => {}); // first reg, late → warns

const late = warnings().filter((m) => lateMsg.test(m));
expect(late).toHaveLength(1);
expect(late[0]).toContain('"toolresult"');
});

it("throws instead of warning when strict: true", async () => {
const [strictAppT, strictBridgeT] =
InMemoryTransport.createLinkedPair();
Expand All @@ -2034,6 +2074,9 @@ describe("App <-> AppBridge integration", () => {
{},
{ autoResize: false, strict: true },
);
// Pre-connect registration for toolcancelled — later swap must not throw.
strictApp.addEventListener("toolcancelled", () => {});

await strictBridge.connect(strictBridgeT);
await strictApp.connect(strictAppT);

Expand All @@ -2043,8 +2086,55 @@ describe("App <-> AppBridge integration", () => {
expect(() => {
strictApp.addEventListener("toolinput", () => {});
}).toThrow(lateMsg);
expect(warnings()).toEqual([]);
// Swapping a handler that existed pre-connect is allowed under strict.
expect(() => {
strictApp.addEventListener("toolcancelled", () => {});
}).not.toThrow();
expect(warnings().filter((m) => lateMsg.test(m))).toEqual([]);
});
});

it("AppBridge warns on a second ui/initialize (View double-mount)", async () => {
await bridge.connect(bridgeTransport);
await app.connect(appTransport);
expect(warnings()).toEqual([]);

// Simulate a second View instance re-running the handshake.
appTransport.send({
jsonrpc: "2.0",
id: 99,
method: "ui/initialize",
params: {
protocolVersion: LATEST_PROTOCOL_VERSION,
appInfo: testAppInfo,
appCapabilities: {},
},
});
await flush();

const doubleInit = warnings().filter((m) =>
/second ui\/initialize/.test(m),
);
expect(doubleInit).toHaveLength(1);
});

it("close() stops further notification delivery (StrictMode cleanup relies on this)", async () => {
const received: unknown[] = [];
app.addEventListener("toolresult", (r) => received.push(r));
await bridge.connect(bridgeTransport);
await app.connect(appTransport);

await bridge.sendToolResult({
content: [{ type: "text", text: "before" }],
});
expect(received).toHaveLength(1);

await app.close();

await bridge
.sendToolResult({ content: [{ type: "text", text: "after" }] })
.catch(() => {});
expect(received).toHaveLength(1);
});

it("AppBridge warns on tools/call from a View that skipped the handshake", async () => {
Expand Down
9 changes: 9 additions & 0 deletions src/app-bridge.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1463,6 +1463,15 @@ export class AppBridge extends ProtocolWithEvents<
): Promise<McpUiInitializeResult> {
const requestedVersion = request.params.protocolVersion;

if (this._appInfo !== undefined) {
console.warn(
"[ext-apps] AppBridge received a second ui/initialize. The View may " +
"be double-mounting (e.g. React StrictMode in dev) without closing " +
"the previous App instance. Responding normally; the latest " +
"appInfo/appCapabilities replace the previous values.",
);
}

this._appCapabilities = request.params.appCapabilities;
this._appInfo = request.params.appInfo;

Expand Down
22 changes: 18 additions & 4 deletions src/app.ts
Original file line number Diff line number Diff line change
Expand Up @@ -395,14 +395,28 @@ export class App extends ProtocolWithEvents<
new Set(["toolinput", "toolinputpartial", "toolresult", "toolcancelled"]);

/**
* Warn (or throw under `strict`) when a one-shot event handler is registered
* after the `ui/initialize` → `ui/notifications/initialized` handshake has
* completed. The host may have already fired the notification by then.
* One-shot events that have had at least one handler registered (via `on*`
* setter or `addEventListener`) at any point. Once an event is in this set,
* subsequent late registrations are not flagged — only the *first* handler
* matters for the missed-notification race, and re-registration (e.g. React
* `useEffect` cleanup → re-add on dep change) is a legitimate pattern.
*/
private readonly _everHadListener = new Set<keyof AppEventMap>();

/**
* Warn (or throw under `strict`) when the *first* handler for a one-shot
* event is registered after the `ui/initialize` → `ui/notifications/initialized`
* handshake has completed. The host may have already fired the notification
* by then. Subsequent registrations for the same event are not flagged.
*
* Mirrors {@link _assertInitialized `_assertInitialized`} (the outbound-side guard).
*/
private _assertHandlerTiming(event: keyof AppEventMap): void {
if (!this._initializedSent || !App.ONE_SHOT_EVENTS.has(event)) return;
if (!App.ONE_SHOT_EVENTS.has(event) || this._everHadListener.has(event)) {
return;
}
this._everHadListener.add(event);
if (!this._initializedSent) return;
const msg =
`[ext-apps] "${String(event)}" handler registered after connect() ` +
`completed the ui/initialize handshake. The host may have already sent ` +
Expand Down
22 changes: 15 additions & 7 deletions src/react/useApp.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -138,25 +138,32 @@ export function useApp({

useEffect(() => {
let mounted = true;
let appInstance: App | undefined;

async function connect() {
try {
const transport = new PostMessageTransport(
window.parent,
window.parent,
);
const app = new App(appInfo, capabilities, { autoResize, strict });
appInstance = new App(appInfo, capabilities, { autoResize, strict });

// Register handlers BEFORE connecting
onAppCreated?.(app);
onAppCreated?.(appInstance);

await app.connect(transport);
await appInstance.connect(transport);

if (mounted) {
setApp(app);
setIsConnected(true);
setError(null);
if (!mounted) {
// Unmounted during the handshake (e.g. React StrictMode dev
// double-invoke). Close so the abandoned transport's window
// `message` listener doesn't keep receiving host postMessages
// alongside the second mount's instance.
void appInstance.close();
return;
}
setApp(appInstance);
setIsConnected(true);
setError(null);
} catch (error) {
if (mounted) {
setApp(null);
Expand All @@ -172,6 +179,7 @@ export function useApp({

return () => {
mounted = false;
void appInstance?.close();
};
}, []); // Intentionally not including options to avoid reconnection

Expand Down
Loading