diff --git a/.changeset/clever-mice-wait.md b/.changeset/clever-mice-wait.md new file mode 100644 index 000000000000..29eba9bb0164 --- /dev/null +++ b/.changeset/clever-mice-wait.md @@ -0,0 +1,5 @@ +--- +'@ai-sdk/mcp': patch +--- + +Reject in-flight MCP requests when their abort signal fires and remove the pending response handler. diff --git a/packages/mcp/src/tool/mcp-client.test.ts b/packages/mcp/src/tool/mcp-client.test.ts index 8268f740aaf6..ffa102cd5737 100644 --- a/packages/mcp/src/tool/mcp-client.test.ts +++ b/packages/mcp/src/tool/mcp-client.test.ts @@ -201,6 +201,47 @@ class FailsFirstToolCallTransport implements MCPTransport { } } +class HangingToolCallTransport implements MCPTransport { + sentMessages: JSONRPCMessage[] = []; + + onmessage?: (message: JSONRPCMessage) => void; + onclose?: () => void; + onerror?: (error: Error) => void; + + async start(): Promise {} + + async close(): Promise { + this.onclose?.(); + } + + async send(message: JSONRPCMessage): Promise { + this.sentMessages.push(message); + + if (!('method' in message) || !('id' in message)) { + return; + } + + if (message.method === 'initialize') { + this.onmessage?.({ + jsonrpc: '2.0', + id: message.id, + result: { + protocolVersion: LATEST_PROTOCOL_VERSION, + serverInfo: { name: 'hanging-tool-call-server', version: '1.0.0' }, + capabilities: { tools: {} }, + }, + }); + return; + } + + if (message.method === 'tools/call') { + // Intentionally never respond. This exercises aborting an in-flight + // request after it has been sent to a slow or hung MCP server. + return; + } + } +} + vi.mock('./mcp-transport.ts', async importOriginal => { const actual = await importOriginal(); return { @@ -1523,6 +1564,49 @@ describe('MCPClient', () => { ); }); + it('should reject in-flight tool call request when aborted', async () => { + const transport = new HangingToolCallTransport(); + client = await createMCPClient({ transport }); + + const abortController = new AbortController(); + const abortReason = new Error('abort after send'); + const toolCallPromise = ( + client as unknown as { + callTool: (args: { + name: string; + args: Record; + options?: { abortSignal?: AbortSignal }; + }) => Promise; + } + ).callTool({ + name: 'hanging-tool', + args: {}, + options: { abortSignal: abortController.signal }, + }); + + expect( + transport.sentMessages.some( + message => 'method' in message && message.method === 'tools/call', + ), + ).toBe(true); + + abortController.abort(abortReason); + + await expect(toolCallPromise).rejects.toSatisfy( + error => + MCPClientError.isInstance(error) && + error.message === 'Request was aborted' && + error.cause === abortReason, + ); + expect( + ( + client as unknown as { + responseHandlers: Map; + } + ).responseHandlers.size, + ).toBe(0); + }); + describe('elicitation support', () => { it('should handle elicitation requests from the server', async () => { client = await createMCPClient({ diff --git a/packages/mcp/src/tool/mcp-client.ts b/packages/mcp/src/tool/mcp-client.ts index b87e3c4a6b7d..1b3c3bdd3e72 100644 --- a/packages/mcp/src/tool/mcp-client.ts +++ b/packages/mcp/src/tool/mcp-client.ts @@ -508,39 +508,57 @@ class DefaultMCPClient implements MCPClient { id: messageId, }; + const rejectWithAbortError = () => { + reject( + new MCPClientError({ + message: 'Request was aborted', + cause: signal?.reason, + }), + ); + }; + const cleanup = () => { this.responseHandlers.delete(messageId); + signal?.removeEventListener('abort', onAbort); + }; + + const rejectAndCleanup = (error: unknown) => { + cleanup(); + reject(error); + }; + + const onAbort = () => { + cleanup(); + rejectWithAbortError(); }; this.responseHandlers.set(messageId, response => { if (signal?.aborted) { - return reject( - new MCPClientError({ - message: 'Request was aborted', - cause: signal.reason, - }), - ); + cleanup(); + return rejectWithAbortError(); } if (response instanceof Error) { - return reject(response); + return rejectAndCleanup(response); } try { const result = resultSchema.parse(response.result); + cleanup(); resolve(result); } catch (error) { const parseError = new MCPClientError({ message: 'Failed to parse server response', cause: error, }); - reject(parseError); + rejectAndCleanup(parseError); } }); + signal?.addEventListener('abort', onAbort, { once: true }); + this.transport.send(jsonrpcRequest).catch(error => { - cleanup(); - reject(error); + rejectAndCleanup(error); }); }); }