From 748d0277a8bc0c304d35946903c2f4cf6f0444ab Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Mon, 30 Jun 2025 09:30:03 +0200 Subject: [PATCH 1/2] Prevent MCPToolset GC during toolset add --- .../haystack_integrations/tools/mcp/mcp_toolset.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py b/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py index 00b131d6fd..146fba2f00 100644 --- a/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py +++ b/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py @@ -145,13 +145,14 @@ def __init__( ) # This is a factory that creates the invocation function for the Tool - def create_invoke_tool(mcp_client, tool_name, tool_timeout): + def create_invoke_tool(owner_toolset, mcp_client, tool_name, tool_timeout): + """Return a closure that keeps a strong reference to *owner_toolset* alive.""" + def invoke_tool(**kwargs) -> Any: - """Invoke a tool using the existing client and AsyncExecutor.""" - result = AsyncExecutor.get_instance().run( + _ = owner_toolset # strong reference so GC can't collect the toolset too early + return AsyncExecutor.get_instance().run( mcp_client.call_tool(tool_name, kwargs), timeout=tool_timeout ) - return result return invoke_tool @@ -170,7 +171,7 @@ def invoke_tool(**kwargs) -> Any: name=tool_info.name, description=tool_info.description, parameters=tool_info.inputSchema, - function=create_invoke_tool(client, tool_info.name, self.invocation_timeout), + function=create_invoke_tool(self, client, tool_info.name, self.invocation_timeout), ) haystack_tools.append(tool) From 3f2744a05795f8e5a4c03e8504ff51a6c0d73949 Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Tue, 1 Jul 2025 14:21:38 +0200 Subject: [PATCH 2/2] Improve typing --- .../src/haystack_integrations/tools/mcp/mcp_toolset.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py b/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py index 146fba2f00..39fcf3d982 100644 --- a/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py +++ b/integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py @@ -2,6 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable from typing import Any from urllib.parse import urlparse @@ -13,6 +14,7 @@ from .mcp_tool import ( AsyncExecutor, + MCPClient, MCPConnectionError, MCPServerInfo, MCPToolNotFoundError, @@ -145,7 +147,12 @@ def __init__( ) # This is a factory that creates the invocation function for the Tool - def create_invoke_tool(owner_toolset, mcp_client, tool_name, tool_timeout): + def create_invoke_tool( + owner_toolset: "MCPToolset", + mcp_client: MCPClient, + tool_name: str, + tool_timeout: float, + ) -> Callable[..., Any]: """Return a closure that keeps a strong reference to *owner_toolset* alive.""" def invoke_tool(**kwargs) -> Any: