diff --git a/packages/core/src/webmcp/declare_tool.ts b/packages/core/src/webmcp/declare_tool.ts index 7a4fe5ebc6f..08e84fcd008 100644 --- a/packages/core/src/webmcp/declare_tool.ts +++ b/packages/core/src/webmcp/declare_tool.ts @@ -59,13 +59,20 @@ export async function declareExperimentalWebMcpTool< const abortCtrl = new AbortController(); const wrappedTool: ToolDescriptor = { ...tool, - execute: (args, client) => - runInInjectionContext(currentInjector, () => + execute: (args, client) => { + // TODO: `@mcp-b/webmcp-polyfill` currently lacks `AbortSignal` in its mock client. + // Remove the optional chaining when it is updated to match Chrome 153 spec. + const signal = client?.signal + ? AbortSignal.any([abortCtrl.signal, client.signal]) + : abortCtrl.signal; + + return runInInjectionContext(currentInjector, () => tool.execute(args, { ...client, - signal: abortCtrl.signal, + signal, }), - ), + ); + }, }; // Unregister when the associated `Injector` is destroyed. diff --git a/packages/core/test/webmcp/declare_tool_spec.ts b/packages/core/test/webmcp/declare_tool_spec.ts index 3f3e2935d7c..e57cec921e4 100644 --- a/packages/core/test/webmcp/declare_tool_spec.ts +++ b/packages/core/test/webmcp/declare_tool_spec.ts @@ -10,7 +10,7 @@ import {initializeWebMCPPolyfill, cleanupWebMCPPolyfill} from '@mcp-b/webmcp-pol import type {JsonSchemaForInference} from '../../third_party/@mcp-b/webmcp-types'; import {inject, Injectable, Injector, runInInjectionContext} from '../../src/di'; import {declareExperimentalWebMcpTool} from '../../src/webmcp/declare_tool'; -import {Execute} from '../../src/webmcp/types'; +import {Execute, ModelContext, ToolDescriptor} from '../../src/webmcp/types'; import {RuntimeErrorCode} from '../../src/errors'; // Whether or not the input type is `any`. @@ -168,6 +168,38 @@ describe('declareExperimentalWebMcpTool', () => { expect(signal.aborted).toBeTrue(); }); + it('should pass an `AbortSignal` to the tool and abort it when the client signal aborts', async () => { + const injector = Injector.create({providers: []}); + const execute = jasmine + .createSpy>('execute') + .and.returnValue({content: []}); + + const modelContext = (globalThis.document as any).modelContext; + const registerToolSpy = spyOn(modelContext, 'registerTool').and.callThrough(); + + await declareExperimentalWebMcpTool( + { + name: 'testTool', + description: 'A test tool', + inputSchema: {type: 'object', properties: {}}, + execute, + }, + injector, + ); + + const wrappedTool = registerToolSpy.calls.first() + .args[0] as ToolDescriptor; + + const clientAbortCtrl = new AbortController(); + await wrappedTool.execute({}, {signal: clientAbortCtrl.signal}); + + const [, {signal}] = execute.calls.first().args; + expect(signal.aborted).toBeFalse(); + + clientAbortCtrl.abort(); + expect(signal.aborted).toBeTrue(); + }); + it('should run `execute` in an injection context', async () => { @Injectable() class TestService {}