From 12ad72e8a4319557a97862cd5e99e2115519c9f3 Mon Sep 17 00:00:00 2001 From: Kashish Hora Date: Tue, 23 Jun 2026 18:22:34 +0200 Subject: [PATCH 1/4] feat: registry returns rejecting Promise; abort-only unregistration; listener fan-out --- src/__tests__/context.test.tsx | 22 +-- src/__tests__/smoke.test.tsx | 19 +-- src/hooks/__tests__/useMcpTool.test.tsx | 21 +-- src/hooks/useMcpTool.ts | 9 +- src/polyfill/__tests__/index.test.ts | 10 +- src/polyfill/__tests__/registry.test.ts | 165 +++++++++++----------- src/polyfill/__tests__/validation.test.ts | 29 +++- src/polyfill/index.ts | 1 - src/polyfill/registry.ts | 65 ++++++--- src/polyfill/testing-shim.ts | 4 +- src/polyfill/validation.ts | 22 +++ src/types.ts | 4 +- 12 files changed, 215 insertions(+), 156 deletions(-) diff --git a/src/__tests__/context.test.tsx b/src/__tests__/context.test.tsx index 854c4b4..f27b1c0 100644 --- a/src/__tests__/context.test.tsx +++ b/src/__tests__/context.test.tsx @@ -52,8 +52,9 @@ describe("WebMCPProvider availability", () => { it("available is true when native API exists", async () => { const native = { - registerTool() {}, - unregisterTool() {}, + registerTool() { + return Promise.resolve(undefined); + }, }; Object.defineProperty(navigator, "modelContext", { value: native, @@ -79,10 +80,11 @@ describe("WebMCPProvider availability", () => { } }); - it("available is true when native API lacks unregisterTool (Chrome 148+)", async () => { + it("available is true with a signal-only native API (Chrome 148+)", async () => { const native = { - registerTool() {}, - // no unregisterTool — simulates Chrome 148+ + registerTool() { + return Promise.resolve(undefined); + }, }; Object.defineProperty(navigator, "modelContext", { value: native, @@ -126,8 +128,9 @@ describe("polyfill lifecycle", () => { it("does not install polyfill when native API present", async () => { const native = { - registerTool() {}, - unregisterTool() {}, + registerTool() { + return Promise.resolve(undefined); + }, }; Object.defineProperty(navigator, "modelContext", { value: native, @@ -171,8 +174,9 @@ describe("polyfill lifecycle", () => { it("does not clean up native API on unmount", async () => { const native = { - registerTool() {}, - unregisterTool() {}, + registerTool() { + return Promise.resolve(undefined); + }, }; Object.defineProperty(navigator, "modelContext", { value: native, diff --git a/src/__tests__/smoke.test.tsx b/src/__tests__/smoke.test.tsx index 9ebbefa..2b2217d 100644 --- a/src/__tests__/smoke.test.tsx +++ b/src/__tests__/smoke.test.tsx @@ -86,7 +86,7 @@ describe("smoke: SSR + client mount availability", () => { // ─── Tool registration lifecycle ────────────────────────────────── describe("smoke: tool registration lifecycle", () => { - it("tool visible in listTools after mount, unregisterTool called on unmount", async () => { + it("tool visible in listTools after mount, removed on unmount", async () => { const { unmount } = render( { expect(tools[0].name).toBe("smoke_greet"); expect(tools[0].description).toBe("Smoke test greeter"); - // Spy before unmount — provider cleanup removes the polyfill - const mc = navigator.modelContext; - expect(mc).toBeDefined(); - const spy = vi.spyOn(mc as NonNullable, "unregisterTool"); - unmount(); - expect(spy).toHaveBeenCalledWith("smoke_greet"); + // Provider cleanup removes the polyfill, so the tool is gone. + await waitFor(() => expect(navigator.modelContextTesting?.listTools() ?? []).toHaveLength(0)); }); }); // ─── StrictMode safety ──────────────────────────────────────────── describe("smoke: StrictMode safety", () => { - it("one tool registered after double-mount, clean unregister on real unmount", async () => { + it("one tool registered after double-mount, removed on real unmount", async () => { const { unmount } = render( @@ -143,14 +139,9 @@ describe("smoke: StrictMode safety", () => { expect(tools).toHaveLength(1); expect(tools[0].name).toBe("strict_tool"); - const mc = navigator.modelContext; - expect(mc).toBeDefined(); - const spy = vi.spyOn(mc as NonNullable, "unregisterTool"); - unmount(); - expect(spy).toHaveBeenCalledTimes(1); - expect(spy).toHaveBeenCalledWith("strict_tool"); + await waitFor(() => expect(navigator.modelContextTesting?.listTools() ?? []).toHaveLength(0)); }); }); diff --git a/src/hooks/__tests__/useMcpTool.test.tsx b/src/hooks/__tests__/useMcpTool.test.tsx index 2a4510b..7f73530 100644 --- a/src/hooks/__tests__/useMcpTool.test.tsx +++ b/src/hooks/__tests__/useMcpTool.test.tsx @@ -185,7 +185,7 @@ describe("registration lifecycle", () => { expect(descriptor.outputSchema?.properties?.greeting).toBeDefined(); }); - it("unregisters tool on unmount", async () => { + it("removes tool on unmount", async () => { const { unmount } = renderWithProvider( { await waitForRegistration(); expect(navigator.modelContextTesting?.listTools()).toHaveLength(1); - // Spy before unmount — provider cleanup will remove the polyfill - const mc = navigator.modelContext; - expect(mc).toBeDefined(); - const spy = vi.spyOn(mc as NonNullable, "unregisterTool"); - unmount(); - expect(spy).toHaveBeenCalledWith("greet"); + await waitFor(() => expect(navigator.modelContextTesting?.listTools() ?? []).toHaveLength(0)); }); it("re-registers when description changes", async () => { @@ -546,7 +541,7 @@ describe("Strict Mode safety", () => { expect(result.content[0].text).toBe("hello world"); }); - it("tool unregistered on real unmount in Strict Mode", async () => { + it("tool removed on real unmount in Strict Mode", async () => { const { unmount } = render( @@ -565,13 +560,9 @@ describe("Strict Mode safety", () => { await waitForRegistration(); expect(navigator.modelContextTesting?.listTools()).toHaveLength(1); - const mc = navigator.modelContext; - expect(mc).toBeDefined(); - const spy = vi.spyOn(mc as NonNullable, "unregisterTool"); - unmount(); - expect(spy).toHaveBeenCalledWith("greet"); + await waitFor(() => expect(navigator.modelContextTesting?.listTools() ?? []).toHaveLength(0)); }); }); @@ -1093,7 +1084,7 @@ describe("provider warning", () => { // ─── Signal-only native API (Chrome 148+) ──────────────────────── -describe("signal-only native API (no unregisterTool)", () => { +describe("signal-only native API (Chrome 148+)", () => { type ToolEntry = { name: string; abortCleanup: () => void }; function installSignalOnlyNative() { @@ -1114,8 +1105,8 @@ describe("signal-only native API (no unregisterTool)", () => { opts.signal.addEventListener("abort", handler, { once: true }); entry.abortCleanup = () => opts.signal?.removeEventListener("abort", handler); } + return Promise.resolve(undefined); }, - // no unregisterTool — simulates Chrome 148+ }; Object.defineProperty(navigator, "modelContext", { diff --git a/src/hooks/useMcpTool.ts b/src/hooks/useMcpTool.ts index 3ea3f59..abe997e 100644 --- a/src/hooks/useMcpTool.ts +++ b/src/hooks/useMcpTool.ts @@ -225,16 +225,12 @@ export function useMcpTool( const controller = new AbortController(); - try { - mc.registerTool(descriptor, { signal: controller.signal }); - } catch (err) { + mc.registerTool(descriptor, { signal: controller.signal }).catch((err) => { warnOnce( `register-${cfg.name}`, `Failed to register tool "${cfg.name}": ${err instanceof Error ? err.message : String(err)}`, ); - return; - } - + }); TOOL_OWNER_BY_NAME.set(cfg.name, ownerToken); return () => { @@ -243,7 +239,6 @@ export function useMcpTool( return; } TOOL_OWNER_BY_NAME.delete(cfg.name); - mc.unregisterTool?.(cfg.name); controller.abort(); }; }, [ diff --git a/src/polyfill/__tests__/index.test.ts b/src/polyfill/__tests__/index.test.ts index 4871fe6..4c3450e 100644 --- a/src/polyfill/__tests__/index.test.ts +++ b/src/polyfill/__tests__/index.test.ts @@ -33,8 +33,9 @@ describe("installPolyfill / cleanupPolyfill", () => { it("installPolyfill skips when native API exists", () => { const native = { - registerTool() {}, - unregisterTool() {}, + registerTool() { + return Promise.resolve(undefined); + }, }; Object.defineProperty(navigator, "modelContext", { value: native, @@ -70,8 +71,9 @@ describe("installPolyfill / cleanupPolyfill", () => { it("cleanupPolyfill restores previous property descriptors", () => { const original = { - registerTool() {}, - unregisterTool() {}, + registerTool() { + return Promise.resolve(undefined); + }, }; Object.defineProperty(navigator, "modelContext", { value: original, diff --git a/src/polyfill/__tests__/registry.test.ts b/src/polyfill/__tests__/registry.test.ts index 892f1e9..9fc6917 100644 --- a/src/polyfill/__tests__/registry.test.ts +++ b/src/polyfill/__tests__/registry.test.ts @@ -12,92 +12,119 @@ function makeTool(overrides?: Partial): ToolDescriptor { } describe("createRegistry", () => { - it("registers a tool visible in getTools()", () => { + it("registers a tool visible in getTools()", async () => { const registry = createRegistry(); - registry.registerTool(makeTool()); + await registry.registerTool(makeTool()); expect(registry.getTools().has("test_tool")).toBe(true); }); - it("defaults inputSchema when absent", () => { + it("resolves and inserts synchronously on success", async () => { const registry = createRegistry(); - registry.registerTool(makeTool({ inputSchema: undefined })); + const p = registry.registerTool(makeTool()); + expect(registry.getTools().has("test_tool")).toBe(true); + await expect(p).resolves.toBeUndefined(); + }); + + it("defaults inputSchema when absent", async () => { + const registry = createRegistry(); + await registry.registerTool(makeTool({ inputSchema: undefined })); const stored = registry.getTools().get("test_tool"); expect(stored?.inputSchema).toEqual({ type: "object", properties: {} }); }); - it("preserves provided inputSchema", () => { + it("preserves provided inputSchema", async () => { const registry = createRegistry(); const schema = { type: "object", properties: { q: { type: "string" } }, required: ["q"] as const, }; - registry.registerTool(makeTool({ inputSchema: schema })); + await registry.registerTool(makeTool({ inputSchema: schema })); const stored = registry.getTools().get("test_tool"); expect(stored?.inputSchema).toEqual(schema); }); - it("throws on duplicate name", () => { + it("rejects with InvalidStateError on duplicate name", async () => { const registry = createRegistry(); - registry.registerTool(makeTool()); - expect(() => registry.registerTool(makeTool())).toThrow( - 'Tool "test_tool" is already registered', - ); + await registry.registerTool(makeTool()); + await expect(registry.registerTool(makeTool())).rejects.toMatchObject({ + name: "InvalidStateError", + }); }); - it("throws on empty name", () => { + it("rejects on empty / too-long / illegal name", async () => { const registry = createRegistry(); - expect(() => registry.registerTool(makeTool({ name: "" }))).toThrow( - "Tool name must be a non-empty string", - ); + await expect(registry.registerTool(makeTool({ name: "" }))).rejects.toMatchObject({ + name: "InvalidStateError", + }); + await expect(registry.registerTool(makeTool({ name: "a".repeat(129) }))).rejects.toMatchObject({ + name: "InvalidStateError", + }); + await expect(registry.registerTool(makeTool({ name: "bad name" }))).rejects.toMatchObject({ + name: "InvalidStateError", + }); }); - it("throws on empty description", () => { + it("rejects on empty description and non-function execute", async () => { const registry = createRegistry(); - expect(() => registry.registerTool(makeTool({ description: "" }))).toThrow( - "Tool description must be a non-empty string", - ); + await expect(registry.registerTool(makeTool({ description: "" }))).rejects.toMatchObject({ + name: "InvalidStateError", + }); + await expect( + registry.registerTool(makeTool({ execute: "x" as unknown as ToolDescriptor["execute"] })), + ).rejects.toMatchObject({ name: "InvalidStateError" }); }); - it("throws on non-function execute", () => { + it("rejects on non-serializable inputSchema with TypeError", async () => { const registry = createRegistry(); - expect(() => - registry.registerTool( - makeTool({ execute: "not a function" as unknown as ToolDescriptor["execute"] }), - ), - ).toThrow("Tool execute must be a function"); + const circular: Record = { type: "object" }; + circular.self = circular; + await expect( + registry.registerTool(makeTool({ inputSchema: circular as never })), + ).rejects.toBeInstanceOf(TypeError); }); - it("throws DOMException with InvalidStateError name", () => { + it("rejects on a non-trustworthy exposedTo origin with SecurityError", async () => { const registry = createRegistry(); - try { - registry.registerTool(makeTool({ name: "" })); - } catch (e) { - expect(e).toBeInstanceOf(DOMException); - expect((e as DOMException).name).toBe("InvalidStateError"); - } + await expect( + registry.registerTool(makeTool(), { exposedTo: ["http://evil.example"] }), + ).rejects.toMatchObject({ name: "SecurityError" }); }); - it("unregisters a tool", () => { + it("rejects when the signal is already aborted", async () => { const registry = createRegistry(); - registry.registerTool(makeTool()); - registry.unregisterTool("test_tool"); + await expect( + registry.registerTool(makeTool(), { signal: AbortSignal.abort() }), + ).rejects.toBeDefined(); expect(registry.getTools().has("test_tool")).toBe(false); }); - it("unregisterTool no-ops for unknown name", () => { + it("fans out to multiple listeners and supports unsubscribe", async () => { const registry = createRegistry(); - expect(() => registry.unregisterTool("nonexistent")).not.toThrow(); + const a = vi.fn(); + const b = vi.fn(); + const offA = registry.addChangeListener(a); + registry.addChangeListener(b); + await registry.registerTool(makeTool()); + await Promise.resolve(); + expect(a).toHaveBeenCalledTimes(1); + expect(b).toHaveBeenCalledTimes(1); + offA(); + await registry.registerTool(makeTool({ name: "second" })); + await Promise.resolve(); + expect(a).toHaveBeenCalledTimes(1); + expect(b).toHaveBeenCalledTimes(2); }); it("batches rapid registrations into one notification", async () => { const registry = createRegistry(); const cb = vi.fn(); - registry.onToolsChanged(cb); + registry.addChangeListener(cb); - registry.registerTool(makeTool({ name: "a" })); - registry.registerTool(makeTool({ name: "b" })); - registry.registerTool(makeTool({ name: "c" })); + // Synchronous insert; the notification is deferred to a microtask. + void registry.registerTool(makeTool({ name: "a" })); + void registry.registerTool(makeTool({ name: "b" })); + void registry.registerTool(makeTool({ name: "c" })); expect(cb).not.toHaveBeenCalled(); await Promise.resolve(); @@ -107,18 +134,18 @@ describe("createRegistry", () => { it("fires notification after microtask for single registration", async () => { const registry = createRegistry(); const cb = vi.fn(); - registry.onToolsChanged(cb); + registry.addChangeListener(cb); - registry.registerTool(makeTool()); + void registry.registerTool(makeTool()); expect(cb).not.toHaveBeenCalled(); await Promise.resolve(); expect(cb).toHaveBeenCalledTimes(1); }); - it("stores a shallow copy isolated from the original object", () => { + it("stores a shallow copy isolated from the original object", async () => { const registry = createRegistry(); const tool = makeTool(); - registry.registerTool(tool); + await registry.registerTool(tool); (tool as Record).description = "mutated"; const stored = registry.getTools().get("test_tool"); @@ -129,10 +156,10 @@ describe("createRegistry", () => { it("removes tool when abort signal fires", async () => { const registry = createRegistry(); const cb = vi.fn(); - registry.onToolsChanged(cb); + registry.addChangeListener(cb); const controller = new AbortController(); - registry.registerTool(makeTool(), { signal: controller.signal }); + await registry.registerTool(makeTool(), { signal: controller.signal }); expect(registry.getTools().has("test_tool")).toBe(true); await Promise.resolve(); @@ -145,40 +172,13 @@ describe("createRegistry", () => { expect(cb).toHaveBeenCalledTimes(2); }); - it("skips registration when signal is already aborted", () => { - const registry = createRegistry(); - registry.registerTool(makeTool(), { signal: AbortSignal.abort() }); - expect(registry.getTools().has("test_tool")).toBe(false); - }); - - it("abort after manual unregisterTool is a safe no-op", async () => { - const registry = createRegistry(); - const cb = vi.fn(); - registry.onToolsChanged(cb); - - const controller = new AbortController(); - registry.registerTool(makeTool(), { signal: controller.signal }); - - await Promise.resolve(); - expect(cb).toHaveBeenCalledTimes(1); - - registry.unregisterTool("test_tool"); - await Promise.resolve(); - expect(cb).toHaveBeenCalledTimes(2); - - controller.abort(); - await Promise.resolve(); - // no extra notification — tool was already removed - expect(cb).toHaveBeenCalledTimes(2); - }); - it("fires notification via microtask when abort removes a tool", async () => { const registry = createRegistry(); const cb = vi.fn(); - registry.onToolsChanged(cb); + registry.addChangeListener(cb); const controller = new AbortController(); - registry.registerTool(makeTool(), { signal: controller.signal }); + await registry.registerTool(makeTool(), { signal: controller.signal }); await Promise.resolve(); controller.abort(); @@ -191,24 +191,25 @@ describe("createRegistry", () => { it("stale abort does not remove a same-name re-registration", async () => { const registry = createRegistry(); const cb = vi.fn(); - registry.onToolsChanged(cb); + registry.addChangeListener(cb); const controller1 = new AbortController(); - registry.registerTool(makeTool(), { signal: controller1.signal }); + await registry.registerTool(makeTool(), { signal: controller1.signal }); await Promise.resolve(); expect(cb).toHaveBeenCalledTimes(1); - // Unregister, then re-register with a new signal - registry.unregisterTool("test_tool"); + // Abort the first registration, then re-register with a new signal + controller1.abort(); await Promise.resolve(); expect(cb).toHaveBeenCalledTimes(2); + expect(registry.getTools().has("test_tool")).toBe(false); const controller2 = new AbortController(); - registry.registerTool(makeTool(), { signal: controller2.signal }); + await registry.registerTool(makeTool(), { signal: controller2.signal }); await Promise.resolve(); expect(cb).toHaveBeenCalledTimes(3); - // Abort the OLD signal — should NOT remove the new registration + // Abort the OLD signal again — should NOT remove the new registration controller1.abort(); await Promise.resolve(); expect(registry.getTools().has("test_tool")).toBe(true); diff --git a/src/polyfill/__tests__/validation.test.ts b/src/polyfill/__tests__/validation.test.ts index 347281e..c69bf1d 100644 --- a/src/polyfill/__tests__/validation.test.ts +++ b/src/polyfill/__tests__/validation.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; import type { InputSchema } from "../../types"; -import { validateArgs } from "../validation"; +import { isPotentiallyTrustworthyOrigin, isValidToolName, validateArgs } from "../validation"; const schema: InputSchema = { type: "object", @@ -113,3 +113,30 @@ describe("validateArgs", () => { expect(() => validateArgs({ anything: true }, s)).not.toThrow(); }); }); + +describe("isValidToolName", () => { + it("accepts snake_case, dot, and dash within 1-128 chars", () => { + expect(isValidToolName("search_catalog")).toBe(true); + expect(isValidToolName("a.b-c_1")).toBe(true); + expect(isValidToolName("a".repeat(128))).toBe(true); + }); + it("rejects empty, over-128, and illegal characters", () => { + expect(isValidToolName("")).toBe(false); + expect(isValidToolName("a".repeat(129))).toBe(false); + expect(isValidToolName("has space")).toBe(false); + expect(isValidToolName("emoji😀")).toBe(false); + }); +}); + +describe("isPotentiallyTrustworthyOrigin", () => { + it("accepts https, wss, file, and localhost", () => { + expect(isPotentiallyTrustworthyOrigin("https://example.com")).toBe(true); + expect(isPotentiallyTrustworthyOrigin("wss://example.com")).toBe(true); + expect(isPotentiallyTrustworthyOrigin("http://localhost")).toBe(true); + expect(isPotentiallyTrustworthyOrigin("http://127.0.0.1")).toBe(true); + }); + it("rejects insecure http origins and unparseable strings", () => { + expect(isPotentiallyTrustworthyOrigin("http://example.com")).toBe(false); + expect(isPotentiallyTrustworthyOrigin("not a url")).toBe(false); + }); +}); diff --git a/src/polyfill/index.ts b/src/polyfill/index.ts index 86677b4..183e1c2 100644 --- a/src/polyfill/index.ts +++ b/src/polyfill/index.ts @@ -27,7 +27,6 @@ export function installPolyfill(): void { const modelContext: PolyfillModelContext = { registerTool: registry.registerTool, - unregisterTool: registry.unregisterTool, __isWebMCPPolyfill: true, }; diff --git a/src/polyfill/registry.ts b/src/polyfill/registry.ts index f6b3699..a9e6622 100644 --- a/src/polyfill/registry.ts +++ b/src/polyfill/registry.ts @@ -1,15 +1,15 @@ import type { RegisterToolOptions, ToolDescriptor } from "../types"; +import { isPotentiallyTrustworthyOrigin, isValidToolName } from "./validation"; export interface RegistryInternal { - registerTool(tool: ToolDescriptor, options?: RegisterToolOptions): void; - unregisterTool(name: string): void; + registerTool(tool: ToolDescriptor, options?: RegisterToolOptions): Promise; getTools(): ReadonlyMap; - onToolsChanged(callback: (() => void) | null): void; + addChangeListener(callback: () => void): () => void; } export function createRegistry(): RegistryInternal { const tools = new Map(); - let changeCallback: (() => void) | null = null; + const listeners = new Set<() => void>(); let notificationPending = false; function scheduleNotification(): void { @@ -17,27 +17,53 @@ export function createRegistry(): RegistryInternal { notificationPending = true; queueMicrotask(() => { notificationPending = false; - changeCallback?.(); + for (const listener of listeners) listener(); }); } return { - registerTool(tool: ToolDescriptor, options?: RegisterToolOptions): void { - if (typeof tool.name !== "string" || tool.name === "") { - throw new DOMException("Tool name must be a non-empty string", "InvalidStateError"); + registerTool(tool: ToolDescriptor, options?: RegisterToolOptions): Promise { + if (!isValidToolName(tool.name)) { + return Promise.reject( + new DOMException( + "Tool name must be 1-128 characters of [A-Za-z0-9_.-]", + "InvalidStateError", + ), + ); } if (typeof tool.description !== "string" || tool.description === "") { - throw new DOMException("Tool description must be a non-empty string", "InvalidStateError"); + return Promise.reject( + new DOMException("Tool description must be a non-empty string", "InvalidStateError"), + ); } if (typeof tool.execute !== "function") { - throw new DOMException("Tool execute must be a function", "InvalidStateError"); + return Promise.reject( + new DOMException("Tool execute must be a function", "InvalidStateError"), + ); } if (tools.has(tool.name)) { - throw new DOMException(`Tool "${tool.name}" is already registered`, "InvalidStateError"); + return Promise.reject( + new DOMException(`Tool "${tool.name}" is already registered`, "InvalidStateError"), + ); + } + if (tool.inputSchema !== undefined) { + try { + JSON.stringify(tool.inputSchema); + } catch { + return Promise.reject(new TypeError("Tool inputSchema is not JSON-serializable")); + } + } + if (options?.exposedTo) { + for (const origin of options.exposedTo) { + if (!isPotentiallyTrustworthyOrigin(origin)) { + return Promise.reject( + new DOMException(`exposedTo origin "${origin}" is not trustworthy`, "SecurityError"), + ); + } + } } - if (options?.signal?.aborted) { - return; + return Promise.reject(options.signal.reason); } tools.set(tool.name, { @@ -59,20 +85,19 @@ export function createRegistry(): RegistryInternal { { once: true }, ); } - }, - unregisterTool(name: string): void { - if (tools.delete(name)) { - scheduleNotification(); - } + return Promise.resolve(undefined); }, getTools(): ReadonlyMap { return tools; }, - onToolsChanged(callback: (() => void) | null): void { - changeCallback = callback; + addChangeListener(callback: () => void): () => void { + listeners.add(callback); + return () => { + listeners.delete(callback); + }; }, }; } diff --git a/src/polyfill/testing-shim.ts b/src/polyfill/testing-shim.ts index e5c2c62..23718d9 100644 --- a/src/polyfill/testing-shim.ts +++ b/src/polyfill/testing-shim.ts @@ -7,6 +7,7 @@ import type { RegistryInternal } from "./registry"; import { validateArgs } from "./validation"; export function createTestingShim(registry: RegistryInternal): ModelContextTesting { + let offChange: (() => void) | null = null; return { listTools() { return Array.from(registry.getTools().values()).map((tool) => ({ @@ -94,7 +95,8 @@ export function createTestingShim(registry: RegistryInternal): ModelContextTesti }, registerToolsChangedCallback(callback: () => void) { - registry.onToolsChanged(callback); + offChange?.(); + offChange = registry.addChangeListener(callback); }, getCrossDocumentScriptToolResult() { diff --git a/src/polyfill/validation.ts b/src/polyfill/validation.ts index e8155db..a7aa204 100644 --- a/src/polyfill/validation.ts +++ b/src/polyfill/validation.ts @@ -1,5 +1,27 @@ import type { InputSchema } from "../types"; +const TOOL_NAME_RE = /^[A-Za-z0-9_.-]{1,128}$/; + +export function isValidToolName(name: unknown): name is string { + return typeof name === "string" && TOOL_NAME_RE.test(name); +} + +export function isPotentiallyTrustworthyOrigin(origin: string): boolean { + let url: URL; + try { + url = new URL(origin); + } catch { + return false; + } + if (url.protocol === "https:" || url.protocol === "wss:" || url.protocol === "file:") { + return true; + } + const host = url.hostname; + return ( + host === "localhost" || host.endsWith(".localhost") || host === "127.0.0.1" || host === "::1" + ); +} + const TYPE_CHECKERS: Record boolean> = { string: (val) => typeof val === "string", number: (val) => typeof val === "number" && !Number.isNaN(val), diff --git a/src/types.ts b/src/types.ts index ba21e9a..966cbc1 100644 --- a/src/types.ts +++ b/src/types.ts @@ -138,11 +138,11 @@ export interface WebMCPStatus { export interface RegisterToolOptions { signal?: AbortSignal; + exposedTo?: string[]; } export interface ModelContext { - registerTool(tool: ToolDescriptor, options?: RegisterToolOptions): void; - unregisterTool?(name: string): void; + registerTool(tool: ToolDescriptor, options?: RegisterToolOptions): Promise; } export interface ModelContextTestingToolInfo { From eb9cd5efa606691dfd1a87d0f527e998f3f7615a Mon Sep 17 00:00:00 2001 From: Kashish Hora Date: Tue, 23 Jun 2026 18:29:53 +0200 Subject: [PATCH 2/4] refactor: tighten exposedTo origin validation to bare origins --- src/polyfill/__tests__/registry.test.ts | 2 +- src/polyfill/__tests__/validation.test.ts | 9 +++++++++ src/polyfill/registry.ts | 1 + src/polyfill/validation.ts | 10 +++++++++- 4 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/polyfill/__tests__/registry.test.ts b/src/polyfill/__tests__/registry.test.ts index 9fc6917..72e0564 100644 --- a/src/polyfill/__tests__/registry.test.ts +++ b/src/polyfill/__tests__/registry.test.ts @@ -95,7 +95,7 @@ describe("createRegistry", () => { const registry = createRegistry(); await expect( registry.registerTool(makeTool(), { signal: AbortSignal.abort() }), - ).rejects.toBeDefined(); + ).rejects.toMatchObject({ name: "AbortError" }); expect(registry.getTools().has("test_tool")).toBe(false); }); diff --git a/src/polyfill/__tests__/validation.test.ts b/src/polyfill/__tests__/validation.test.ts index c69bf1d..b7d08af 100644 --- a/src/polyfill/__tests__/validation.test.ts +++ b/src/polyfill/__tests__/validation.test.ts @@ -139,4 +139,13 @@ describe("isPotentiallyTrustworthyOrigin", () => { expect(isPotentiallyTrustworthyOrigin("http://example.com")).toBe(false); expect(isPotentiallyTrustworthyOrigin("not a url")).toBe(false); }); + it("rejects non-bare origins (path, query, or credentials)", () => { + expect(isPotentiallyTrustworthyOrigin("https://example.com/path")).toBe(false); + expect(isPotentiallyTrustworthyOrigin("https://example.com?x=1")).toBe(false); + expect(isPotentiallyTrustworthyOrigin("https://user:pass@example.com")).toBe(false); + }); + it("still accepts a bare https origin and file: urls", () => { + expect(isPotentiallyTrustworthyOrigin("https://example.com")).toBe(true); + expect(isPotentiallyTrustworthyOrigin("file:///Users/x/page.html")).toBe(true); + }); }); diff --git a/src/polyfill/registry.ts b/src/polyfill/registry.ts index a9e6622..df43d43 100644 --- a/src/polyfill/registry.ts +++ b/src/polyfill/registry.ts @@ -47,6 +47,7 @@ export function createRegistry(): RegistryInternal { ); } if (tool.inputSchema !== undefined) { + // Only rejects hard-fail serialization (cycles, BigInt); lossy cases (functions, undefined) are not caught. try { JSON.stringify(tool.inputSchema); } catch { diff --git a/src/polyfill/validation.ts b/src/polyfill/validation.ts index a7aa204..1be2bfb 100644 --- a/src/polyfill/validation.ts +++ b/src/polyfill/validation.ts @@ -13,7 +13,15 @@ export function isPotentiallyTrustworthyOrigin(origin: string): boolean { } catch { return false; } - if (url.protocol === "https:" || url.protocol === "wss:" || url.protocol === "file:") { + // file: is an intentional allowance; file: URLs have a null origin so check it first. + if (url.protocol === "file:") { + return true; + } + // Must be a bare origin — reject inputs carrying a path, query, or credentials. + if (url.origin !== origin) { + return false; + } + if (url.protocol === "https:" || url.protocol === "wss:") { return true; } const host = url.hostname; From b277c10dfc1fd0acaef9e214da134d6ccaf0946c Mon Sep 17 00:00:00 2001 From: Kashish Hora Date: Tue, 23 Jun 2026 18:33:50 +0200 Subject: [PATCH 3/4] feat: migrate modelContext surface from navigator to document --- src/__tests__/context.test.tsx | 46 +++++++++++---------- src/context.tsx | 4 +- src/hooks/__tests__/useMcpTool.test.tsx | 34 ++++++++-------- src/hooks/useMcpTool.ts | 4 +- src/polyfill/__tests__/index.test.ts | 53 +++++++++++++++---------- src/polyfill/index.ts | 10 ++--- src/types.ts | 4 +- 7 files changed, 87 insertions(+), 68 deletions(-) diff --git a/src/__tests__/context.test.tsx b/src/__tests__/context.test.tsx index f27b1c0..5aeb0d9 100644 --- a/src/__tests__/context.test.tsx +++ b/src/__tests__/context.test.tsx @@ -14,14 +14,14 @@ function StatusDisplay() { } function deleteNativeModelContext() { - const desc = Object.getOwnPropertyDescriptor(navigator, "modelContext"); + const desc = Object.getOwnPropertyDescriptor(document, "modelContext"); if (desc) { - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: undefined, configurable: true, writable: true, }); - delete navigator.modelContext; + delete document.modelContext; } } @@ -56,7 +56,7 @@ describe("WebMCPProvider availability", () => { return Promise.resolve(undefined); }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: native, configurable: true, enumerable: true, @@ -74,7 +74,7 @@ describe("WebMCPProvider availability", () => { expect(getByTestId("status")).toHaveTextContent("yes"); }); // Native API should still be the original, not replaced by polyfill - expect(navigator.modelContext).toBe(native); + expect(document.modelContext).toBe(native); } finally { deleteNativeModelContext(); } @@ -86,7 +86,7 @@ describe("WebMCPProvider availability", () => { return Promise.resolve(undefined); }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: native, configurable: true, enumerable: true, @@ -103,7 +103,7 @@ describe("WebMCPProvider availability", () => { await waitFor(() => { expect(getByTestId("status")).toHaveTextContent("yes"); }); - expect(navigator.modelContext).toBe(native); + expect(document.modelContext).toBe(native); } finally { deleteNativeModelContext(); } @@ -121,9 +121,11 @@ describe("polyfill lifecycle", () => { ); await waitFor(() => { - expect(navigator.modelContext).toBeDefined(); + expect(document.modelContext).toBeDefined(); }); - expect((navigator.modelContext as Record).__isWebMCPPolyfill).toBe(true); + expect((document.modelContext as unknown as Record).__isWebMCPPolyfill).toBe( + true, + ); }); it("does not install polyfill when native API present", async () => { @@ -132,7 +134,7 @@ describe("polyfill lifecycle", () => { return Promise.resolve(undefined); }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: native, configurable: true, enumerable: true, @@ -147,10 +149,10 @@ describe("polyfill lifecycle", () => { ); await waitFor(() => { - expect(navigator.modelContext).toBe(native); + expect(document.modelContext).toBe(native); }); expect( - (navigator.modelContext as Record).__isWebMCPPolyfill, + (document.modelContext as unknown as Record).__isWebMCPPolyfill, ).toBeUndefined(); } finally { deleteNativeModelContext(); @@ -165,11 +167,11 @@ describe("polyfill lifecycle", () => { ); await waitFor(() => { - expect(navigator.modelContext).toBeDefined(); + expect(document.modelContext).toBeDefined(); }); unmount(); - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); }); it("does not clean up native API on unmount", async () => { @@ -178,7 +180,7 @@ describe("polyfill lifecycle", () => { return Promise.resolve(undefined); }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: native, configurable: true, enumerable: true, @@ -193,11 +195,11 @@ describe("polyfill lifecycle", () => { ); await waitFor(() => { - expect(navigator.modelContext).toBe(native); + expect(document.modelContext).toBe(native); }); unmount(); - expect(navigator.modelContext).toBe(native); + expect(document.modelContext).toBe(native); } finally { deleteNativeModelContext(); } @@ -226,12 +228,14 @@ describe("multi-provider lifecycle", () => { // Unmount first provider — polyfill should still be alive result1.unmount(); - expect(navigator.modelContext).toBeDefined(); - expect((navigator.modelContext as Record).__isWebMCPPolyfill).toBe(true); + expect(document.modelContext).toBeDefined(); + expect((document.modelContext as unknown as Record).__isWebMCPPolyfill).toBe( + true, + ); // Unmount second provider — now polyfill should be cleaned up result2.unmount(); - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); }); }); @@ -312,6 +316,6 @@ describe("Strict Mode", () => { }); unmount(); - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); }); }); diff --git a/src/context.tsx b/src/context.tsx index e703944..4279aa9 100644 --- a/src/context.tsx +++ b/src/context.tsx @@ -30,7 +30,7 @@ function WebMCPProvider({ name, version, children }: WebMCPProviderProps) { useEffect(() => { const hasNativeApi = - !!navigator.modelContext && !("__isWebMCPPolyfill" in navigator.modelContext); + !!document.modelContext && !("__isWebMCPPolyfill" in document.modelContext); if (!hasNativeApi) { installPolyfill(); @@ -38,7 +38,7 @@ function WebMCPProvider({ name, version, children }: WebMCPProviderProps) { polyfillConsumerCount++; } - setAvailable(!!navigator.modelContext); + setAvailable(!!document.modelContext); return () => { if (usesPolyfillRef.current) { diff --git a/src/hooks/__tests__/useMcpTool.test.tsx b/src/hooks/__tests__/useMcpTool.test.tsx index 7f73530..7f3a76b 100644 --- a/src/hooks/__tests__/useMcpTool.test.tsx +++ b/src/hooks/__tests__/useMcpTool.test.tsx @@ -153,7 +153,7 @@ describe("registration lifecycle", () => { // listTools() only returns name/description/inputSchema, so spy on registerTool // and trigger re-registration by changing description. - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -230,7 +230,7 @@ describe("registration lifecycle", () => { await waitForRegistration(); - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -276,7 +276,7 @@ describe("registration lifecycle", () => { await waitForRegistration(); - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -322,7 +322,7 @@ describe("registration lifecycle", () => { await waitForRegistration(); - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -370,7 +370,7 @@ describe("registration lifecycle", () => { await waitForRegistration(); - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -414,7 +414,7 @@ describe("registration lifecycle", () => { await waitForRegistration(); - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -458,7 +458,7 @@ describe("registration lifecycle", () => { await waitForRegistration(); - const mc = navigator.modelContext; + const mc = document.modelContext; expect(mc).toBeDefined(); const spy = vi.spyOn(mc as NonNullable, "registerTool"); @@ -980,9 +980,9 @@ describe("SSR safety", () => { it("registration effect does not fire during SSR", () => { // During SSR, typeof navigator is "undefined" in Node, but jsdom provides it. // Verify no tool is registered synchronously during renderToString. - // After renderToString, navigator.modelContext should not exist (no provider effect ran). - const prevMc = navigator.modelContext; - delete navigator.modelContext; + // After renderToString, document.modelContext should not exist (no provider effect ran). + const prevMc = document.modelContext; + delete document.modelContext; try { renderToString( @@ -998,10 +998,10 @@ describe("SSR safety", () => { ); // No modelContext should exist — effects don't run during renderToString - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); } finally { if (prevMc) { - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: prevMc, configurable: true, enumerable: true, @@ -1109,7 +1109,7 @@ describe("signal-only native API (Chrome 148+)", () => { }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: native, configurable: true, enumerable: true, @@ -1120,14 +1120,14 @@ describe("signal-only native API (Chrome 148+)", () => { } function deleteModelContext() { - const desc = Object.getOwnPropertyDescriptor(navigator, "modelContext"); + const desc = Object.getOwnPropertyDescriptor(document, "modelContext"); if (desc) { - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: undefined, configurable: true, writable: true, }); - delete navigator.modelContext; + delete document.modelContext; } } @@ -1183,7 +1183,7 @@ describe("signal-only native API (Chrome 148+)", () => { it("re-registers with fresh signal on prop change", async () => { const registered = installSignalOnlyNative(); const registerSpy = vi.spyOn( - navigator.modelContext as NonNullable, + document.modelContext as NonNullable, "registerTool", ); diff --git a/src/hooks/useMcpTool.ts b/src/hooks/useMcpTool.ts index abe997e..9ab3c7d 100644 --- a/src/hooks/useMcpTool.ts +++ b/src/hooks/useMcpTool.ts @@ -137,11 +137,11 @@ export function useMcpTool( // biome-ignore lint/correctness/useExhaustiveDependencies: schema objects are tracked via fingerprints, handler/callbacks via refs useEffect(() => { - if (typeof navigator === "undefined" || !navigator.modelContext) { + if (typeof document === "undefined" || !document.modelContext) { return; } - const mc = navigator.modelContext; + const mc = document.modelContext; const cfg = configRef.current; const ownerToken = Symbol(cfg.name); const zodPath = "input" in cfg && cfg.input instanceof z.ZodObject; diff --git a/src/polyfill/__tests__/index.test.ts b/src/polyfill/__tests__/index.test.ts index 4c3450e..5788455 100644 --- a/src/polyfill/__tests__/index.test.ts +++ b/src/polyfill/__tests__/index.test.ts @@ -7,9 +7,12 @@ describe("installPolyfill / cleanupPolyfill", () => { cleanupPolyfill(); }); - it("installPolyfill creates navigator.modelContext", () => { + it("installPolyfill creates document.modelContext", () => { installPolyfill(); - expect(navigator.modelContext).toBeDefined(); + expect(document.modelContext).toBeDefined(); + expect((document.modelContext as unknown as Record).__isWebMCPPolyfill).toBe( + true, + ); }); it("installPolyfill creates navigator.modelContextTesting", () => { @@ -19,25 +22,35 @@ describe("installPolyfill / cleanupPolyfill", () => { it("modelContext has __isWebMCPPolyfill marker", () => { installPolyfill(); - expect((navigator.modelContext as Record).__isWebMCPPolyfill).toBe(true); + expect((document.modelContext as unknown as Record).__isWebMCPPolyfill).toBe( + true, + ); }); it("installPolyfill is idempotent", () => { installPolyfill(); - const first = navigator.modelContext; + const first = document.modelContext; const firstTesting = navigator.modelContextTesting; installPolyfill(); - expect(navigator.modelContext).toBe(first); + expect(document.modelContext).toBe(first); expect(navigator.modelContextTesting).toBe(firstTesting); }); + it("does not install when native document.modelContext exists", () => { + const native = { registerTool: () => Promise.resolve(undefined) }; + Object.defineProperty(document, "modelContext", { value: native, configurable: true }); + installPolyfill(); + expect(document.modelContext).toBe(native); + delete (document as { modelContext?: unknown }).modelContext; + }); + it("installPolyfill skips when native API exists", () => { const native = { registerTool() { return Promise.resolve(undefined); }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: native, configurable: true, enumerable: true, @@ -45,27 +58,27 @@ describe("installPolyfill / cleanupPolyfill", () => { }); installPolyfill(); - expect(navigator.modelContext).toBe(native); + expect(document.modelContext).toBe(native); // Manual cleanup since our polyfill wasn't installed - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: undefined, configurable: true, writable: true, }); - delete navigator.modelContext; + delete (document as { modelContext?: unknown }).modelContext; }); - it("navigator.modelContext is not writable", () => { + it("document.modelContext is not writable", () => { installPolyfill(); - const desc = Object.getOwnPropertyDescriptor(navigator, "modelContext"); + const desc = Object.getOwnPropertyDescriptor(document, "modelContext"); expect(desc?.writable).toBe(false); }); it("cleanupPolyfill removes polyfill properties", () => { installPolyfill(); cleanupPolyfill(); - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); expect(navigator.modelContextTesting).toBeUndefined(); }); @@ -75,7 +88,7 @@ describe("installPolyfill / cleanupPolyfill", () => { return Promise.resolve(undefined); }, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: original, configurable: true, enumerable: true, @@ -86,18 +99,18 @@ describe("installPolyfill / cleanupPolyfill", () => { (original as Record).__isWebMCPPolyfill = true; installPolyfill(); - expect(navigator.modelContext).not.toBe(original); + expect(document.modelContext).not.toBe(original); cleanupPolyfill(); - expect(navigator.modelContext).toBe(original); + expect(document.modelContext).toBe(original); // Final cleanup - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: undefined, configurable: true, writable: true, }); - delete navigator.modelContext; + delete (document as { modelContext?: unknown }).modelContext; }); it("cleanupPolyfill restores previous modelContextTesting descriptor", () => { @@ -131,7 +144,7 @@ describe("installPolyfill / cleanupPolyfill", () => { it("cleanupPolyfill is no-op when not installed", () => { expect(() => cleanupPolyfill()).not.toThrow(); - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); }); it("full lifecycle: install → register tool → execute via testing shim → cleanup", async () => { @@ -150,7 +163,7 @@ describe("installPolyfill / cleanupPolyfill", () => { }), }; - const mc = navigator.modelContext as NonNullable; + const mc = document.modelContext as NonNullable; mc.registerTool(tool); const testing = navigator.modelContextTesting as NonNullable< @@ -166,7 +179,7 @@ describe("installPolyfill / cleanupPolyfill", () => { }); cleanupPolyfill(); - expect(navigator.modelContext).toBeUndefined(); + expect(document.modelContext).toBeUndefined(); expect(navigator.modelContextTesting).toBeUndefined(); }); }); diff --git a/src/polyfill/index.ts b/src/polyfill/index.ts index 183e1c2..e576805 100644 --- a/src/polyfill/index.ts +++ b/src/polyfill/index.ts @@ -13,13 +13,13 @@ let previousModelContextTesting: PropertyDescriptor | undefined; export function installPolyfill(): void { if (typeof window === "undefined") return; - if (navigator.modelContext && !("__isWebMCPPolyfill" in navigator.modelContext)) { + if (document.modelContext && !("__isWebMCPPolyfill" in document.modelContext)) { return; } if (installed) return; - previousModelContext = Object.getOwnPropertyDescriptor(navigator, "modelContext"); + previousModelContext = Object.getOwnPropertyDescriptor(document, "modelContext"); previousModelContextTesting = Object.getOwnPropertyDescriptor(navigator, "modelContextTesting"); const registry = createRegistry(); @@ -30,7 +30,7 @@ export function installPolyfill(): void { __isWebMCPPolyfill: true, }; - Object.defineProperty(navigator, "modelContext", { + Object.defineProperty(document, "modelContext", { value: modelContext, configurable: true, enumerable: true, @@ -51,9 +51,9 @@ export function cleanupPolyfill(): void { if (!installed) return; if (previousModelContext) { - Object.defineProperty(navigator, "modelContext", previousModelContext); + Object.defineProperty(document, "modelContext", previousModelContext); } else { - delete navigator.modelContext; + delete document.modelContext; } if (previousModelContextTesting) { diff --git a/src/types.ts b/src/types.ts index 966cbc1..65823d0 100644 --- a/src/types.ts +++ b/src/types.ts @@ -167,8 +167,10 @@ export interface ModelContextTesting { } declare global { - interface Navigator { + interface Document { modelContext?: ModelContext; + } + interface Navigator { modelContextTesting?: ModelContextTesting; } } From dfb4671a01993838b6ae4d5d6e3eea6b0d32c362 Mon Sep 17 00:00:00 2001 From: Kashish Hora Date: Tue, 23 Jun 2026 18:39:06 +0200 Subject: [PATCH 4/4] feat!: single-arg execute; remove ModelContextClient --- src/hooks/useMcpTool.ts | 14 ++---- src/polyfill/__tests__/testing-shim.test.ts | 54 +++++++-------------- src/polyfill/testing-shim.ts | 24 +-------- src/types.ts | 16 ++---- 4 files changed, 25 insertions(+), 83 deletions(-) diff --git a/src/hooks/useMcpTool.ts b/src/hooks/useMcpTool.ts index 9ab3c7d..e584324 100644 --- a/src/hooks/useMcpTool.ts +++ b/src/hooks/useMcpTool.ts @@ -5,7 +5,6 @@ import type { CallToolResult, McpToolConfigJsonSchema, McpToolConfigZod, - ModelContextClient, ToolDescriptor, ToolExecutionState, UseMcpToolReturn, @@ -92,11 +91,7 @@ export function useMcpTool( ); } - const client: ModelContextClient = { - requestUserInteraction: (callback) => callback(), - }; - - const result = await handlerRef.current(validatedInput as Record, client); + const result = await handlerRef.current(validatedInput as Record); if (isMountedRef.current) { inFlightCountRef.current--; @@ -164,10 +159,7 @@ export function useMcpTool( ...(resolvedInputSchema && { inputSchema: resolvedInputSchema }), ...(resolvedOutputSchema && { outputSchema: resolvedOutputSchema }), ...(cfg.annotations && { annotations: cfg.annotations }), - execute: async ( - args: Record, - client: ModelContextClient, - ): Promise => { + execute: async (args: Record): Promise => { inFlightCountRef.current++; if (isMountedRef.current) { setState((prev) => ({ ...prev, isExecuting: true, error: null })); @@ -183,7 +175,7 @@ export function useMcpTool( validatedArgs = (currentConfig as McpToolConfigZod).input.parse(args); } - const result = await handlerRef.current(validatedArgs as Record, client); + const result = await handlerRef.current(validatedArgs as Record); if (isMountedRef.current) { inFlightCountRef.current--; diff --git a/src/polyfill/__tests__/testing-shim.test.ts b/src/polyfill/__tests__/testing-shim.test.ts index d588d93..daf0d24 100644 --- a/src/polyfill/__tests__/testing-shim.test.ts +++ b/src/polyfill/__tests__/testing-shim.test.ts @@ -53,7 +53,7 @@ describe("createTestingShim", () => { }); describe("executeTool", () => { - it("calls handler with parsed args and client", async () => { + it("calls handler with parsed args", async () => { const handler = vi.fn(async () => makeResult()); const { shim } = setup([makeTool({ execute: handler })]); @@ -61,7 +61,22 @@ describe("createTestingShim", () => { expect(handler).toHaveBeenCalledTimes(1); expect(handler.mock.calls[0][0]).toEqual({ query: "hello" }); - expect(handler.mock.calls[0][1]).toHaveProperty("requestUserInteraction"); + }); + + it("calls execute with a single input argument (no client)", async () => { + const registry = createRegistry(); + const shim = createTestingShim(registry); + let argCount = -1; + await registry.registerTool({ + name: "arity_tool", + description: "checks arity", + execute: (...args: unknown[]) => { + argCount = args.length; + return { content: [{ type: "text", text: "ok" }] }; + }, + }); + await shim.executeTool("arity_tool", "{}"); + expect(argCount).toBe(1); }); it("returns stringified CallToolResult", async () => { @@ -197,41 +212,6 @@ describe("createTestingShim", () => { }); }); - describe("requestUserInteraction", () => { - it("works during execution", async () => { - let interactionResult: unknown; - const { shim } = setup([ - makeTool({ - execute: async (_args, client) => { - interactionResult = await client.requestUserInteraction(async () => "user said yes"); - return makeResult(); - }, - }), - ]); - - await shim.executeTool("test_tool", '{"query":"x"}'); - expect(interactionResult).toBe("user said yes"); - }); - - it("throws InvalidStateError after execution completes", async () => { - let savedClient: { requestUserInteraction: (cb: () => Promise) => Promise }; - const { shim } = setup([ - makeTool({ - execute: async (_args, client) => { - savedClient = client; - return makeResult(); - }, - }), - ]); - - await shim.executeTool("test_tool", '{"query":"x"}'); - - expect(() => savedClient.requestUserInteraction(async () => "late")).toThrow( - expect.objectContaining({ name: "InvalidStateError" }), - ); - }); - }); - describe("registerToolsChangedCallback", () => { it("receives change notifications", async () => { const registry = createRegistry(); diff --git a/src/polyfill/testing-shim.ts b/src/polyfill/testing-shim.ts index 23718d9..c06e5e9 100644 --- a/src/polyfill/testing-shim.ts +++ b/src/polyfill/testing-shim.ts @@ -1,8 +1,4 @@ -import type { - ModelContextClient, - ModelContextTesting, - ModelContextTestingExecuteToolOptions, -} from "../types"; +import type { ModelContextTesting, ModelContextTestingExecuteToolOptions } from "../types"; import type { RegistryInternal } from "./registry"; import { validateArgs } from "./validation"; @@ -46,19 +42,6 @@ export function createTestingShim(registry: RegistryInternal): ModelContextTesti validateArgs(parsed as Record, tool.inputSchema); } - let contextActive = true; - const client: ModelContextClient = { - requestUserInteraction(callback) { - if (!contextActive) { - throw new DOMException( - "Tool execution context is no longer active", - "InvalidStateError", - ); - } - return callback(); - }, - }; - const signal = options?.signal; let onAbort: (() => void) | undefined; let abortPromise: Promise | undefined; @@ -77,9 +60,7 @@ export function createTestingShim(registry: RegistryInternal): ModelContextTesti } try { - const resultPromise = Promise.resolve( - tool.execute(parsed as Record, client), - ); + const resultPromise = Promise.resolve(tool.execute(parsed as Record)); const result = abortPromise ? await Promise.race([abortPromise, resultPromise]) @@ -87,7 +68,6 @@ export function createTestingShim(registry: RegistryInternal): ModelContextTesti return JSON.stringify(result); } finally { - contextActive = false; if (onAbort && signal) { signal.removeEventListener("abort", onAbort); } diff --git a/src/types.ts b/src/types.ts index 65823d0..f777310 100644 --- a/src/types.ts +++ b/src/types.ts @@ -70,17 +70,13 @@ export interface ToolAnnotations { openWorldHint?: boolean; } -export interface ModelContextClient { - requestUserInteraction(callback: () => Promise): Promise; -} - export interface ToolDescriptor> { name: string; description: string; inputSchema?: InputSchema; outputSchema?: InputSchema; annotations?: ToolAnnotations; - execute: (args: TArgs, client: ModelContextClient) => MaybePromise; + execute: (input: TArgs) => MaybePromise; } interface McpToolConfigBase { @@ -96,10 +92,7 @@ export interface McpToolConfigZod extends McpToolConfig inputSchema?: never; output?: z.ZodObject; outputSchema?: never; - handler: ( - args: z.infer>, - client: ModelContextClient, - ) => MaybePromise; + handler: (args: z.infer>) => MaybePromise; } export interface McpToolConfigJsonSchema extends McpToolConfigBase { @@ -107,10 +100,7 @@ export interface McpToolConfigJsonSchema extends McpToolConfigBase { inputSchema?: InputSchema; output?: never; outputSchema?: InputSchema; - handler: ( - args: Record, - client: ModelContextClient, - ) => MaybePromise; + handler: (args: Record) => MaybePromise; } export interface ToolExecutionState {