mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-08-10 08:21:54 +00:00
Preserve root causes during MCP connection cleanup
This commit is contained in:
@@ -52,6 +52,38 @@ _mcp_schema_cache: dict[str, tuple[list[dict[str, Any]], float]] = {}
|
||||
_cache_ttl = 300 # 5 minutes
|
||||
|
||||
|
||||
def _find_meaningful_exception(error: BaseException | None) -> Exception | None:
|
||||
"""Return the first actionable leaf from an async exception tree."""
|
||||
if error is None:
|
||||
return None
|
||||
|
||||
if isinstance(error, BaseExceptionGroup):
|
||||
for child in error.exceptions:
|
||||
if meaningful := _find_meaningful_exception(child):
|
||||
return meaningful
|
||||
return None
|
||||
|
||||
if isinstance(error, (asyncio.CancelledError, GeneratorExit)):
|
||||
return None
|
||||
|
||||
if isinstance(error, RuntimeError):
|
||||
message = str(error).lower()
|
||||
if "cancel scope" in message or "different task" in message:
|
||||
return None
|
||||
|
||||
return error if isinstance(error, Exception) else None
|
||||
|
||||
|
||||
def _connection_error_type(error: BaseException) -> str:
|
||||
"""Classify an MCP connection error for emitted events."""
|
||||
message = str(error).lower()
|
||||
if "401" in message or "unauthorized" in message:
|
||||
return "authentication"
|
||||
if "certificate verify failed" in message or "ssl" in message:
|
||||
return "tls"
|
||||
return "network"
|
||||
|
||||
|
||||
class MCPClient:
|
||||
"""MCP client with session management.
|
||||
|
||||
@@ -196,29 +228,9 @@ class MCPClient:
|
||||
|
||||
# MCP requires initialize() before any other request.
|
||||
await self._session.initialize()
|
||||
except asyncio.CancelledError:
|
||||
# If initialization was cancelled (e.g., event loop closing),
|
||||
# cleanup and re-raise - don't suppress cancellation
|
||||
await self._cleanup_on_error()
|
||||
raise
|
||||
except BaseExceptionGroup as eg:
|
||||
# Handle exception groups from anyio task groups
|
||||
# Extract the actual meaningful error (not GeneratorExit)
|
||||
actual_error = None
|
||||
for exc in eg.exceptions:
|
||||
if isinstance(exc, Exception) and not isinstance(
|
||||
exc, GeneratorExit
|
||||
):
|
||||
# Check if it's an HTTP error (like 401)
|
||||
error_msg = str(exc).lower()
|
||||
if "401" in error_msg or "unauthorized" in error_msg:
|
||||
actual_error = exc
|
||||
break
|
||||
if "cancel scope" not in error_msg and "task" not in error_msg:
|
||||
actual_error = exc
|
||||
break
|
||||
|
||||
await self._cleanup_on_error()
|
||||
actual_error = _find_meaningful_exception(eg)
|
||||
if actual_error:
|
||||
raise ConnectionError(
|
||||
f"Failed to connect to MCP server: {actual_error}"
|
||||
@@ -245,7 +257,7 @@ class MCPClient:
|
||||
|
||||
return self
|
||||
except ImportError as e:
|
||||
await self._cleanup_on_error()
|
||||
await self._cleanup_on_error(e)
|
||||
error_msg = (
|
||||
"MCP library not available. Please install with: pip install mcp"
|
||||
)
|
||||
@@ -259,7 +271,7 @@ class MCPClient:
|
||||
)
|
||||
raise ImportError(error_msg) from e
|
||||
except asyncio.TimeoutError as e:
|
||||
await self._cleanup_on_error()
|
||||
await self._cleanup_on_error(e)
|
||||
error_msg = f"MCP connection timed out after {self.connect_timeout} seconds. The server may be slow or unreachable."
|
||||
self._emit_connection_failed(
|
||||
server_name,
|
||||
@@ -270,9 +282,22 @@ class MCPClient:
|
||||
started_at,
|
||||
)
|
||||
raise ConnectionError(error_msg) from e
|
||||
except asyncio.CancelledError:
|
||||
# Re-raise cancellation - don't suppress it
|
||||
await self._cleanup_on_error()
|
||||
except asyncio.CancelledError as e:
|
||||
cleanup_error = await self._cleanup_on_error(e, log_error=False)
|
||||
if actual_error := _find_meaningful_exception(cleanup_error):
|
||||
error_msg = str(actual_error)
|
||||
self._emit_connection_failed(
|
||||
server_name,
|
||||
server_url,
|
||||
transport_type,
|
||||
error_msg,
|
||||
_connection_error_type(actual_error),
|
||||
started_at,
|
||||
)
|
||||
raise ConnectionError(
|
||||
f"Failed to connect to MCP server: {actual_error}"
|
||||
) from actual_error
|
||||
|
||||
self._emit_connection_failed(
|
||||
server_name,
|
||||
server_url,
|
||||
@@ -284,26 +309,10 @@ class MCPClient:
|
||||
raise
|
||||
except BaseExceptionGroup as eg:
|
||||
# Handle exception groups from anyio task groups at outer level
|
||||
actual_error = None
|
||||
for exc in eg.exceptions:
|
||||
if isinstance(exc, Exception) and not isinstance(exc, GeneratorExit):
|
||||
error_msg = str(exc).lower()
|
||||
if "401" in error_msg or "unauthorized" in error_msg:
|
||||
actual_error = exc
|
||||
break
|
||||
if "cancel scope" not in error_msg and "task" not in error_msg:
|
||||
actual_error = exc
|
||||
break
|
||||
|
||||
await self._cleanup_on_error()
|
||||
actual_error = _find_meaningful_exception(eg)
|
||||
await self._cleanup_on_error(eg)
|
||||
error_type = (
|
||||
"authentication"
|
||||
if actual_error
|
||||
and (
|
||||
"401" in str(actual_error).lower()
|
||||
or "unauthorized" in str(actual_error).lower()
|
||||
)
|
||||
else "network"
|
||||
_connection_error_type(actual_error) if actual_error else "network"
|
||||
)
|
||||
error_msg = str(actual_error) if actual_error else str(eg)
|
||||
self._emit_connection_failed(
|
||||
@@ -320,12 +329,8 @@ class MCPClient:
|
||||
) from actual_error
|
||||
raise ConnectionError(f"Failed to connect to MCP server: {eg}") from eg
|
||||
except Exception as e:
|
||||
await self._cleanup_on_error()
|
||||
error_type = (
|
||||
"authentication"
|
||||
if "401" in str(e).lower() or "unauthorized" in str(e).lower()
|
||||
else "network"
|
||||
)
|
||||
await self._cleanup_on_error(e)
|
||||
error_type = _connection_error_type(e)
|
||||
self._emit_connection_failed(
|
||||
server_name, server_url, transport_type, str(e), error_type, started_at
|
||||
)
|
||||
@@ -355,19 +360,37 @@ class MCPClient:
|
||||
),
|
||||
)
|
||||
|
||||
async def _cleanup_on_error(self) -> None:
|
||||
"""Cleanup resources when an error occurs during connection."""
|
||||
async def _cleanup_on_error(
|
||||
self,
|
||||
error: BaseException | None = None,
|
||||
*,
|
||||
log_error: bool = True,
|
||||
) -> BaseException | None:
|
||||
"""Cleanup resources while preserving the initiating exception context."""
|
||||
cleanup_error: BaseException | None = None
|
||||
try:
|
||||
await self._exit_stack.aclose()
|
||||
|
||||
except Exception as e:
|
||||
if error is None:
|
||||
await self._exit_stack.aclose()
|
||||
else:
|
||||
await self._exit_stack.__aexit__(
|
||||
type(error), error, error.__traceback__
|
||||
)
|
||||
except (KeyboardInterrupt, SystemExit):
|
||||
raise
|
||||
except BaseException as e:
|
||||
# Do not replace the connection failure with a secondary cleanup
|
||||
# failure. Preserve the root cause and leave cleanup observable.
|
||||
self._logger.warning("Error during MCP client cleanup: %s", e)
|
||||
cleanup_error = e
|
||||
if log_error:
|
||||
visible_error = _find_meaningful_exception(e) or e
|
||||
self._logger.warning(
|
||||
"Error during MCP client cleanup: %s", visible_error
|
||||
)
|
||||
finally:
|
||||
self._session = None
|
||||
self._initialized = False
|
||||
self._exit_stack = AsyncExitStack()
|
||||
return cleanup_error
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
"""Disconnect from MCP server and cleanup resources."""
|
||||
|
||||
@@ -1,17 +1,10 @@
|
||||
"""HTTP and Streamable HTTP transport for MCP servers."""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from builtins import BaseExceptionGroup
|
||||
else:
|
||||
from exceptiongroup import BaseExceptionGroup
|
||||
|
||||
from crewai.mcp._utils import async_timeout
|
||||
from crewai.mcp.transports.base import BaseTransport, TransportType
|
||||
|
||||
@@ -87,7 +80,9 @@ class HTTPTransport(BaseTransport):
|
||||
# later trigger "cancel scope in a different task" failures.
|
||||
async with async_timeout(self.connect_timeout):
|
||||
read, write, _ = await self._transport_context.__aenter__()
|
||||
except TimeoutError as e:
|
||||
# async-timeout raises asyncio.TimeoutError on Python 3.10, where
|
||||
# it is distinct from the built-in TimeoutError.
|
||||
except (TimeoutError, asyncio.TimeoutError) as e:
|
||||
self._transport_context = None
|
||||
raise ConnectionError(
|
||||
f"Transport context entry timed out after {self.connect_timeout} seconds. "
|
||||
@@ -109,57 +104,25 @@ class HTTPTransport(BaseTransport):
|
||||
self._transport_context = None
|
||||
raise ConnectionError(f"Failed to connect to MCP server: {e}") from e
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
"""Close HTTP connection."""
|
||||
async def _disconnect(
|
||||
self,
|
||||
exc_type: type[BaseException] | None = None,
|
||||
exc_val: BaseException | None = None,
|
||||
exc_tb: Any = None,
|
||||
) -> None:
|
||||
"""Close the SDK context with the exception that triggered unwinding."""
|
||||
if not self._connected:
|
||||
return
|
||||
|
||||
try:
|
||||
# Clear streams first
|
||||
self._clear_streams()
|
||||
# await self._exit_stack.aclose()
|
||||
self._clear_streams()
|
||||
transport_context = self._transport_context
|
||||
self._transport_context = None
|
||||
if transport_context is not None:
|
||||
await transport_context.__aexit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
# Exit transport context - this will clean up background tasks
|
||||
# Give a small delay to allow background tasks to complete
|
||||
if self._transport_context is not None:
|
||||
try:
|
||||
# Wait a tiny bit for any pending operations
|
||||
await asyncio.sleep(0.1)
|
||||
await self._transport_context.__aexit__(None, None, None)
|
||||
except (RuntimeError, asyncio.CancelledError) as e:
|
||||
# Ignore "exit cancel scope in different task" errors and cancellation
|
||||
# These happen when asyncio.run() closes the event loop
|
||||
# while background tasks are still running
|
||||
error_msg = str(e).lower()
|
||||
if "cancel scope" not in error_msg and "task" not in error_msg:
|
||||
# Only suppress cancel scope/task errors, re-raise others
|
||||
if isinstance(e, RuntimeError):
|
||||
raise
|
||||
# For CancelledError, just suppress it
|
||||
except BaseExceptionGroup as eg:
|
||||
# Handle exception groups from anyio task groups
|
||||
# Suppress if they contain cancel scope errors
|
||||
should_suppress = False
|
||||
for exc in eg.exceptions:
|
||||
error_msg = str(exc).lower()
|
||||
if "cancel scope" in error_msg or "task" in error_msg:
|
||||
should_suppress = True
|
||||
break
|
||||
if not should_suppress:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Error during HTTP transport disconnect: {e}"
|
||||
) from e
|
||||
|
||||
self._connected = False
|
||||
|
||||
except Exception as e:
|
||||
# Log but don't raise - cleanup should be best effort
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"Error during HTTP transport disconnect: {e}")
|
||||
async def disconnect(self) -> None:
|
||||
"""Close HTTP connection."""
|
||||
await self._disconnect()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Async context manager entry."""
|
||||
@@ -172,5 +135,4 @@ class HTTPTransport(BaseTransport):
|
||||
exc_tb: Any,
|
||||
) -> None:
|
||||
"""Async context manager exit."""
|
||||
|
||||
await self.disconnect()
|
||||
await self._disconnect(exc_type, exc_val, exc_tb)
|
||||
|
||||
@@ -5,17 +5,25 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
import logging
|
||||
import sys
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
from crewai.events.types.mcp_events import MCPConnectionFailedEvent
|
||||
from crewai.mcp.client import MCPClient
|
||||
from crewai.mcp.transports.base import BaseTransport, TransportType
|
||||
from crewai.mcp.transports.http import HTTPTransport
|
||||
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from builtins import BaseExceptionGroup
|
||||
else:
|
||||
from exceptiongroup import BaseExceptionGroup
|
||||
|
||||
|
||||
class ConnectedTransport(BaseTransport):
|
||||
"""Minimal connected transport for client operation tests."""
|
||||
|
||||
@@ -37,6 +45,14 @@ class ConnectedTransport(BaseTransport):
|
||||
await self.disconnect()
|
||||
|
||||
|
||||
class LifecycleTransport(ConnectedTransport):
|
||||
"""Connected transport with placeholder streams for session startup tests."""
|
||||
|
||||
async def connect(self) -> LifecycleTransport:
|
||||
self._set_streams(MagicMock(), MagicMock())
|
||||
return self
|
||||
|
||||
|
||||
class FailingSession:
|
||||
"""Session that models a lost response after a side effect committed."""
|
||||
|
||||
@@ -161,3 +177,92 @@ async def test_cleanup_failure_is_logged_without_masking_connection_error():
|
||||
log_template, log_error = logger.warning.call_args.args
|
||||
assert log_template == "Error during MCP client cleanup: %s"
|
||||
assert str(log_error) == "cleanup failed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_startup_reports_nested_tls_error():
|
||||
"""A transport task failure must beat the cancellation used to signal it."""
|
||||
|
||||
class TLSFailingTransportContext(AbstractAsyncContextManager):
|
||||
async def __aenter__(self):
|
||||
return MagicMock(), MagicMock(), None
|
||||
|
||||
async def __aexit__(self, exc_type, _exc, _tb):
|
||||
assert exc_type is asyncio.CancelledError
|
||||
raise BaseExceptionGroup(
|
||||
"transport failed",
|
||||
[
|
||||
ConnectionError(
|
||||
"[SSL: CERTIFICATE_VERIFY_FAILED] unable to get local issuer"
|
||||
),
|
||||
GeneratorExit(),
|
||||
],
|
||||
)
|
||||
|
||||
class TLSFailingSession:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def initialize(self) -> None:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
emitted_events = []
|
||||
client = MCPClient(HTTPTransport("https://mcp.example.com"))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"mcp.client.streamable_http.streamablehttp_client",
|
||||
return_value=TLSFailingTransportContext(),
|
||||
),
|
||||
patch("mcp.ClientSession", return_value=TLSFailingSession()),
|
||||
patch.object(
|
||||
crewai_event_bus,
|
||||
"emit",
|
||||
side_effect=lambda _source, event: emitted_events.append(event),
|
||||
),
|
||||
pytest.raises(ConnectionError, match="CERTIFICATE_VERIFY_FAILED"),
|
||||
):
|
||||
await client.connect()
|
||||
|
||||
failed_event = next(
|
||||
event for event in emitted_events if isinstance(event, MCPConnectionFailedEvent)
|
||||
)
|
||||
assert failed_event.error_type == "tls"
|
||||
assert "CERTIFICATE_VERIFY_FAILED" in failed_event.error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_external_startup_cancellation_remains_cancelled():
|
||||
"""Cancellation without an underlying transport error must propagate."""
|
||||
|
||||
class CancelledSession:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def initialize(self) -> None:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
emitted_events = []
|
||||
client = MCPClient(LifecycleTransport())
|
||||
|
||||
with (
|
||||
patch("mcp.ClientSession", return_value=CancelledSession()),
|
||||
patch.object(
|
||||
crewai_event_bus,
|
||||
"emit",
|
||||
side_effect=lambda _source, event: emitted_events.append(event),
|
||||
),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await client.connect()
|
||||
|
||||
failed_event = next(
|
||||
event for event in emitted_events if isinstance(event, MCPConnectionFailedEvent)
|
||||
)
|
||||
assert failed_event.error_type == "cancelled"
|
||||
|
||||
Reference in New Issue
Block a user