feat: MCPTool and MCPToolset async resource management improvements (#1758)

* Add MCPClientSessionManager to connect/close mcp clients

* Update and refactor mcp tests

* More descriptive connection error raising

* Proper test cleanup

* Testing CI windows

* linting

* Improve connection error raise

* PR feedback

* Proper naming, and more precise cleanup sequence

---------

Co-authored-by: Michele Pangrazzi <xmikex83@gmail.com>

Author: vblagoje

integrations/mcp/src/haystack_integrations/tools/mcp/mcp_tool.py

@@ -7,7 +7,8 @@
 import threading
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Coroutine
+from collections.abc import Callable, Coroutine
+from concurrent.futures import Future
 from contextlib import AsyncExitStack
 from dataclasses import dataclass, fields
 from typing import Any, cast
@@ -73,6 +74,47 @@ def run(self, coro: Coroutine[Any, Any, Any], timeout: float | None = None) -> A
             message = f"Operation timed out after {timeout} seconds"
             raise TimeoutError(message) from e
 
+    def get_loop(self):
+        """
+        Get the event loop.
+
+        :returns: The event loop
+        """
+        return self._loop
+
+    def run_background(
+        self, coro_factory: Callable[[asyncio.Event], Coroutine[Any, Any, Any]], timeout: float | None = None
+    ) -> tuple[concurrent.futures.Future[Any], asyncio.Event]:
+        """
+        Schedule `coro_factory` to run in the executor's event loop **without** blocking the
+        caller thread.
+
+        The factory receives an :class:`asyncio.Event` that can be used to cooperatively shut
+        the coroutine down. The method returns **both** the concurrent future (to observe
+        completion or failure) and the created *stop_event* so that callers can signal termination.
+
+        :param coro_factory: A callable receiving the stop_event and returning the coroutine to execute.
+        :param timeout: Optional timeout while waiting for the stop_event to be created.
+        :returns: Tuple ``(future, stop_event)``.
+        """
+        # A promise that will be fulfilled from inside the coroutine_with_stop_event coroutine once the
+        # stop_event is created *inside* the target event loop to ensure it is bound to the
+        # correct loop and can safely be set from other threads via *call_soon_threadsafe*.
+        stop_event_promise: Future[asyncio.Event] = Future()
+
+        async def _coroutine_with_stop_event():
+            stop_event = asyncio.Event()
+            stop_event_promise.set_result(stop_event)
+            await coro_factory(stop_event)
+
+        # Schedule the coroutine
+        future = asyncio.run_coroutine_threadsafe(_coroutine_with_stop_event(), self._loop)
+
+        # This ensures that the stop_event is fully initialized and ready for use before
+        # the run_background method returns, allowing the caller to immediately
+        # use it to control the coroutine.
+        return future, stop_event_promise.result(timeout)
+
     def shutdown(self, timeout: float = 2):
         """
         Shut down the background event loop and thread.
@@ -213,7 +255,7 @@ async def call_tool(self, tool_name: str, tool_args: dict[str, Any]) -> Any:
             raise
         except Exception as e:
             # Wrap other exceptions with context about which tool failed
-            message = f"Failed to invoke tool '{tool_name}'"
+            message = f"Failed to invoke tool '{tool_name}' due to: {e}"
             raise MCPInvocationError(message, tool_name, tool_args) from e
 
     def _validate_response(self, tool_name: str, result: types.CallToolResult) -> types.CallToolResult:
@@ -254,7 +296,7 @@ def _validate_response(self, tool_name: str, result: types.CallToolResult) -> ty
         # Return the original result object
         return result
 
-    async def close(self) -> None:
+    async def aclose(self) -> None:
         """
         Close the connection and clean up resources.
 
@@ -273,15 +315,6 @@ async def close(self) -> None:
             self.stdio = None
             self.write = None
 
-    def close_sync(self) -> None:
-        """Synchronous version of close for use in __del__ - ensures resources are cleaned up."""
-        logger.debug("PROCESS: Closing StdioClient (sync)")
-
-        try:
-            AsyncExecutor.get_instance().run(self.close(), timeout=2)
-        except Exception as e:
-            logger.debug(f"PROCESS: Error during async cleanup in sync close: {e!s}")
-
     async def _initialize_session_with_transport(
         self,
         transport_tuple: tuple[
@@ -310,7 +343,7 @@ async def _initialize_session_with_transport(
             return response.tools
 
         except Exception as e:
-            await self.close()
+            # We'll clean up the session in the calling code, so we don't need to do it here.
             message = f"Failed to connect to {connection_type}: {e}"
             raise MCPConnectionError(message=message, operation="connect") from e
 
@@ -572,22 +605,17 @@ def __init__(
 
         logger.debug(f"TOOL: Initializing MCPTool '{name}'")
 
-        # Create client
-        self._client = server_info.create_client()
-        logger.debug(f"TOOL: Created client for MCPTool '{name}'")
-
         try:
+            # Create client and spin up a long-lived worker that keeps the
+            # connect/close lifecycle inside one coroutine.
+            self._client = server_info.create_client()
+            logger.debug(f"TOOL: Created client for MCPTool '{name}'")
 
-            async def connect():
-                logger.debug(f"TOOL: Inside connect coroutine for '{name}'")
-                result = await asyncio.wait_for(self._client.connect(), timeout=connection_timeout)
-                logger.debug(f"TOOL: Connect successful for '{name}', found {len(result)} tools")
-                return result
-
-            logger.debug(f"TOOL: About to run connect for '{name}'")
-            tools = AsyncExecutor.get_instance().run(connect(), timeout=connection_timeout)
-            logger.debug(f"TOOL: Connection complete for '{name}'")
+            # The worker starts immediately and blocks here until the connection
+            # is established (or fails), returning the tool list.
+            self._worker = _MCPClientSessionManager(self._client, timeout=connection_timeout)
 
+            tools = self._worker.tools()
             # Handle no tools case
             if not tools:
                 logger.debug(f"TOOL: No tools found for '{name}'")
@@ -617,17 +645,28 @@ async def connect():
             logger.debug(f"TOOL: Initialization complete for '{name}'")
 
         except Exception as e:
-            # Clean up resources on error
-            logger.debug(f"TOOL: Error during initialization of '{name}': {e!s}")
-            if self._client:
-                try:
-                    logger.debug(f"TOOL: Attempting cleanup after initialization failure for '{name}'")
-                    AsyncExecutor.get_instance().run(self._client.close(), timeout=5)
-                    logger.debug(f"TOOL: Cleanup successful for '{name}'")
-                except Exception as cleanup_error:
-                    logger.debug(f"TOOL: Error during cleanup after initialization failure: {cleanup_error!s}")
-
-            message = f"Failed to initialize MCPTool '{name}': {e}"
+            # We need to close because we could connect properly, retrieve tools yet
+            # fail because of an MCPToolNotFoundError
+            self.close()
+
+            # Extract more detailed error information from TaskGroup/ExceptionGroup exceptions
+            from exceptiongroup import ExceptionGroup
+
+            error_message = str(e)
+            # Handle ExceptionGroup to extract more useful error messages
+            if isinstance(e, ExceptionGroup):
+                if e.exceptions:
+                    first_exception = e.exceptions[0]
+                    error_message = (
+                        first_exception.message if hasattr(first_exception, "message") else str(first_exception)
+                    )
+
+            # Ensure we always have a meaningful error message
+            if not error_message or error_message.strip() == "":
+                # Provide platform-independent fallback message for connection errors
+                error_message = f"Connection failed to MCP server (using {type(server_info).__name__})"
+
+            message = f"Failed to initialize MCPTool '{name}': {error_message}"
             raise MCPConnectionError(message=message, server_info=server_info, operation="initialize") from e
 
     def _invoke_tool(self, **kwargs: Any) -> Any:
@@ -743,13 +782,106 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
             invocation_timeout=invocation_timeout,
         )
 
+    def close(self):
+        """Close the tool synchronously."""
+        if hasattr(self, "_client") and self._client:
+            try:
+                # Tell the background worker to shut down gracefully.
+                if hasattr(self, "_worker") and self._worker:
+                    self._worker.stop()
+            except Exception as e:
+                logger.debug(f"TOOL: Error during synchronous worker stop: {e!s}")
+
     def __del__(self):
         """Cleanup resources when the tool is garbage collected."""
         logger.debug(f"TOOL: __del__ called for MCPTool '{self.name if hasattr(self, 'name') else 'unknown'}'")
 
-        # Call synchronous close on the client
-        if hasattr(self, "_client") and self._client:
+        self.close()
+
+
+class _MCPClientSessionManager:
+    """Runs an MCPClient connect/close inside the AsyncExecutor's event loop.
+
+    Life-cycle:
+      1.  Create the worker to schedule a long-running coroutine in the
+          dedicated background loop.
+      2.  The coroutine calls *connect* on mcp client; when it has the tool list it fulfils
+          a concurrent future so the synchronous thread can continue.
+      3.  It then waits on an `asyncio.Event`.
+      4.  `stop()` sets the event from any thread. The same coroutine then calls
+          *close()* on mcp client and finishes without the dreaded
+          `Attempted to exit cancel scope in a different task than it was entered in` error
+          thus properly closing the client.
+    """
+
+    # Maximum time to wait for worker shutdown in seconds
+    WORKER_SHUTDOWN_TIMEOUT = 2.0
+
+    def __init__(self, client: "MCPClient", *, timeout: float | None = None):
+        self._client = client
+        self.executor = AsyncExecutor.get_instance()
+
+        # Where the tool list (or an exception) will be delivered.
+        self._tools_promise: Future[list[Tool]] = Future()
+
+        # Kick off the worker coroutine in the background loop
+        self._worker_future, self._stop_event = self.executor.run_background(self._run, timeout=None)
+
+        # Wait (in the caller thread) until connect() finishes or raises.
+        try:
+            self._tools_promise.result(timeout)
+        except BaseException:
+            # If connect failed we should cancel the worker so it doesn't hang.
+            self.stop()
+            raise
+
+    def tools(self) -> list[Tool]:
+        """Return the tool list already collected during startup."""
+
+        return self._tools_promise.result()
+
+    def stop(self) -> None:
+        """Request the worker to shut down and block until done."""
+
+        def _set(ev: asyncio.Event):
+            if not ev.is_set():
+                ev.set()
+
+        if self.executor.get_loop().is_closed():
+            return
+
+        # The stop event is created inside the worker *before* the connect
+        # promise is fulfilled, so at this point it must exist.
+        self.executor.get_loop().call_soon_threadsafe(_set, self._stop_event)  # type: ignore[attr-defined]
+
+        # Wait for the worker coroutine to finish so resources are fully
+        # released before returning. Swallow any errors during shutdown.
+        try:
+            self._worker_future.result(timeout=self.WORKER_SHUTDOWN_TIMEOUT)
+        except Exception as e:
+            logger.debug(f"Error during worker future result: {e}")
+            pass
+
+    async def _run(self, stop_event: asyncio.Event):
+        """Background coroutine living in AsyncExecutor's loop."""
+
+        try:
+            # logger.debug(f"TOOL: _run current task: {asyncio.current_task()}")
+            tools = await self._client.connect()
+            # Deliver the tool list to the waiting synchronous code.
+            if not self._tools_promise.done():
+                self._tools_promise.set_result(tools)
+            # Park until told to stop.
+            await stop_event.wait()
+        except Exception as exc:
+            logger.debug(f"Error during _run: {exc}")
+            if not self._tools_promise.done():
+                self._tools_promise.set_exception(exc)
+            raise
+        finally:
+            # logger.debug(f"TOOL: _run current task: {asyncio.current_task()}")
+            # Close the client in the same couroutine that connected it
             try:
-                self._client.close_sync()
+                await self._client.aclose()
             except Exception as e:
-                logger.debug(f"TOOL: Error during synchronous client close: {e!s}")
+                logger.debug(f"Error during MCP client cleanup: {e!s}")

integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py

@@ -18,6 +18,7 @@
     MCPToolNotFoundError,
     SSEServerInfo,
     StdioServerInfo,
+    _MCPClientSessionManager,
 )
 
 logger = logging.getLogger(__name__)
@@ -122,11 +123,12 @@ def __init__(
 
         # Connect and load tools
         try:
-            # Create the appropriate client using the factory method
+            # Create the client and spin up a worker so open/close happen in the
+            # same coroutine, avoiding AnyIO cancel-scope issues.
             client = self.server_info.create_client()
+            self._worker = _MCPClientSessionManager(client, timeout=self.connection_timeout)
 
-            # Connect and get available tools using AsyncExecutor
-            tools = AsyncExecutor.get_instance().run(client.connect(), timeout=self.connection_timeout)
+            tools = self._worker.tools()
 
             # If tool_names is provided, validate that all requested tools exist
             if self.tool_names:
@@ -175,6 +177,11 @@ def invoke_tool(**kwargs) -> Any:
             super().__init__(tools=haystack_tools)
 
         except Exception as e:
+            # We need to close because we could connect properly, retrieve tools yet
+            # fail because of an MCPToolNotFoundError
+            self.close()
+
+            # Create informative error message for SSE connection errors
             if isinstance(self.server_info, SSEServerInfo):
                 base_message = f"Failed to connect to SSE server at {self.server_info.url}"
                 checks = ["1. The server is running"]
@@ -205,6 +212,7 @@ def invoke_tool(**kwargs) -> Any:
                     message = f"{base_message}. Please check if:\n" + "\\n".join(checks)
                 else:
                     message = f"{base_message}: {e}"
+            # and for stdio connection errors
             elif isinstance(self.server_info, StdioServerInfo):  # stdio connection
                 base_message = "Failed to start MCP server process"
                 stdio_info = self.server_info
@@ -255,3 +263,14 @@ def from_dict(cls, data: dict[str, Any]) -> "MCPToolset":
             connection_timeout=inner_data.get("connection_timeout", 30.0),
             invocation_timeout=inner_data.get("invocation_timeout", 30.0),
         )
+
+    def close(self):
+        """Close the underlying MCP client safely."""
+        if hasattr(self, "_worker") and self._worker:
+            try:
+                self._worker.stop()
+            except Exception as e:
+                logger.debug(f"TOOLSET: error during worker stop: {e!s}")
+
+    def __del__(self):
+        self.close()

integrations/mcp/tests/conftest.py

@@ -0,0 +1,27 @@
+import pytest
+
+from haystack_integrations.tools.mcp import MCPTool, MCPToolset
+
+
+@pytest.fixture
+def mcp_tool_cleanup():
+    """Fixture to ensure all MCPTool and MCPToolset instances are properly closed after tests."""
+    tools = []
+    toolsets = []
+
+    def _register(item):
+        """Register an MCP component for cleanup."""
+        if isinstance(item, MCPTool):
+            tools.append(item)
+        elif isinstance(item, MCPToolset):
+            toolsets.append(item)
+        return item
+
+    yield _register
+
+    # Finalizer to close all tools and toolsets
+    for tool in tools:
+        tool.close()
+
+    for toolset in toolsets:
+        toolset.close()