Fix subprocess cleanup for multiple MCP clients (#231)

lorenss-m · Copilot · web-flow · commit 9c6dc85c5de2 · 2025-08-26T22:26:36.000-07:00
* init

* add logging callback

* cleaner shutdown of multiple clients

* linting

* do not clear manager and fix stop for http error

* http fallback change

* Update mcp_use/connectors/http.py

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
Signed-off-by: lorenss &lt;31991968+lorenss-m@users.noreply.github.com&gt;

---------

Signed-off-by: lorenss &lt;31991968+lorenss-m@users.noreply.github.com&gt;
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/mcp_use/connectors/base.py b/mcp_use/connectors/base.py
@@ -111,29 +111,34 @@ async def _cleanup_resources(self) -> None:
         """Clean up all resources associated with this connector."""
         errors = []
 
-        # First close the client session
-        if self.client_session:
+        # First stop the connection manager, this closes the ClientSession inside
+        # the same task where it was opened, avoiding cancel-scope mismatches.
+        if self._connection_manager:
             try:
-                logger.debug("Closing client session")
-                await self.client_session.__aexit__(None, None, None)
+                logger.debug("Stopping connection manager")
+                await self._connection_manager.stop()
             except Exception as e:
-                error_msg = f"Error closing client session: {e}"
+                error_msg = f"Error stopping connection manager: {e}"
                 logger.warning(error_msg)
                 errors.append(error_msg)
             finally:
-                self.client_session = None
+                self._connection_manager = None
 
-        # Then stop the connection manager
-        if self._connection_manager:
+        # Ensure the client_session reference is cleared (it should already be
+        # closed by the connection manager). Only attempt a direct __aexit__ if
+        # the connection manager did *not* exist, this covers edge-cases like
+        # failed connections where no manager was started.
+        if self.client_session:
             try:
-                logger.debug("Stopping connection manager")
-                await self._connection_manager.stop()
+                if not self._connection_manager:
+                    logger.debug("Closing client session (no connection manager)")
+                    await self.client_session.__aexit__(None, None, None)
             except Exception as e:
-                error_msg = f"Error stopping connection manager: {e}"
+                error_msg = f"Error closing client session: {e}"
                 logger.warning(error_msg)
                 errors.append(error_msg)
             finally:
-                self._connection_manager = None
+                self.client_session = None
 
         # Reset tools
         self._tools = None
diff --git a/mcp_use/connectors/http.py b/mcp_use/connectors/http.py
@@ -8,6 +8,7 @@
 import httpx
 from mcp import ClientSession
 from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+from mcp.shared.exceptions import McpError
 
 from ..logging import logger
 from ..task_managers import SseConnectionManager, StreamableHttpConnectionManager
@@ -125,6 +126,21 @@ async def connect(self) -> None:
                 else:
                     self._prompts = []
 
+            except McpError as mcp_error:
+                # This is a protocol error, not a transport error
+                # The server is reachable and speaking MCP, but rejecting our request
+                logger.error("MCP protocol error during initialization: %s", mcp_error)
+
+                # Clean up the test client
+                try:
+                    await test_client.__aexit__(None, None, None)
+                except Exception:
+                    pass
+
+                # Don't try SSE fallback for protocol errors - the server is working,
+                # it just doesn't like our request
+                raise mcp_error
+
             except Exception as init_error:
                 # Clean up the test client
                 try:
diff --git a/mcp_use/task_managers/base.py b/mcp_use/task_managers/base.py
@@ -22,13 +22,14 @@ class ConnectionManager(Generic[T], ABC):
     used with MCP connectors.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize a new connection manager."""
         self._ready_event = asyncio.Event()
         self._done_event = asyncio.Event()
+        self._stop_event = asyncio.Event()
         self._exception: Exception | None = None
         self._connection: T | None = None
-        self._task: asyncio.Task | None = None
+        self._task: asyncio.Task[None] | None = None
 
     @abstractmethod
     async def _establish_connection(self) -> T:
@@ -86,20 +87,15 @@ async def start(self) -> T:
 
     async def stop(self) -> None:
         """Stop the connection manager and close the connection."""
+        # Signal stop to the connection task instead of cancelling it, avoids
+        # propagating CancelledError to unrelated tasks.
         if self._task and not self._task.done():
-            # Cancel the task
-            logger.debug(f"Cancelling {self.__class__.__name__} task")
-            self._task.cancel()
-
-            # Wait for it to complete
-            try:
-                await self._task
-            except asyncio.CancelledError:
-                logger.debug(f"{self.__class__.__name__} task cancelled successfully")
-            except Exception as e:
-                logger.warning(f"Error stopping {self.__class__.__name__} task: {e}")
-
-        # Wait for the connection to be done
+            logger.debug(f"Signaling stop to {self.__class__.__name__} task")
+            self._stop_event.set()
+            # Wait for it to finish gracefully
+            await self._task
+
+        # Ensure cleanup completed
         await self._done_event.wait()
         logger.debug(f"{self.__class__.__name__} task completed")
 
@@ -125,14 +121,8 @@ async def _connection_task(self) -> None:
             # Signal that the connection is ready
             self._ready_event.set()
 
-            # Wait indefinitely until cancelled
-            try:
-                # This keeps the connection open until cancelled
-                await asyncio.Event().wait()
-            except asyncio.CancelledError:
-                # Expected when stopping
-                logger.debug(f"{self.__class__.__name__} task received cancellation")
-                pass
+            # Wait until stop is requested
+            await self._stop_event.wait()
 
         except Exception as e:
             # Store the exception
