diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 27a54de33..a4d3aad5d 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -975,11 +975,8 @@ async def connect_mcp_servers( from mcp.client.streamable_http import streamable_http_client async def open_single_server( - name: str, cfg: "MCPServerConfig" - ) -> tuple[str, AsyncExitStack | None]: - server_stack = AsyncExitStack() - await server_stack.__aenter__() - + name: str, cfg: "MCPServerConfig", server_stack: AsyncExitStack + ) -> bool: try: transport_type = cfg.type if not transport_type: @@ -991,8 +988,7 @@ async def connect_mcp_servers( ) else: logger.warning("MCP server '{}': no command or url configured, skipping", name) - await server_stack.aclose() - return name, None + return False if transport_type in {"sse", "streamableHttp"}: ok, error = validate_url_target(cfg.url) @@ -1003,8 +999,7 @@ async def connect_mcp_servers( _redact_url(cfg.url), error, ) - await server_stack.aclose() - return name, None + return False if transport_type == "stdio": command, args, env = _normalize_windows_stdio_command( @@ -1022,8 +1017,7 @@ async def connect_mcp_servers( elif transport_type == "sse": if not await _probe_http_url(cfg.url): logger.warning("MCP server '{}': {} unreachable, skipping", name, _redact_url(cfg.url)) - await server_stack.aclose() - return name, None + return False def httpx_client_factory( headers: dict[str, str] | None = None, @@ -1050,8 +1044,7 @@ async def connect_mcp_servers( elif transport_type == "streamableHttp": if not await _probe_http_url(cfg.url): logger.warning("MCP server '{}': {} unreachable, skipping", name, _redact_url(cfg.url)) - await server_stack.aclose() - return name, None + return False http_client = await server_stack.enter_async_context( httpx.AsyncClient( @@ -1067,8 +1060,7 @@ async def connect_mcp_servers( ) else: logger.warning("MCP server '{}': unknown transport type '{}'", name, transport_type) - await server_stack.aclose() - return name, None + return False read = _filter_malformed_mcp_progress_notifications(read, name) session = await server_stack.enter_async_context(ClientSession(read, write)) @@ -1171,7 +1163,7 @@ async def connect_mcp_servers( logger.info( "MCP server '{}': connected, {} capabilities registered", name, registered_count ) - return name, server_stack + return True except Exception as e: hint = "" @@ -1191,9 +1183,7 @@ async def connect_mcp_servers( "only JSON-RPC to stdout and sends logs/debug output to stderr instead." ) logger.exception("MCP server '{}': failed to connect: {}", name, hint) - with suppress(Exception): - await server_stack.aclose() - return name, None + return False async def connect_single_server( name: str, cfg: "MCPServerConfig" @@ -1203,30 +1193,30 @@ async def connect_mcp_servers( close_requested = asyncio.Event() async def own_connection() -> None: - stack: AsyncExitStack | None = None try: - _, stack = await open_single_server(name, cfg) - if not ready.done(): - ready.set_result(stack is not None) - if stack is not None: - await close_requested.wait() + async with AsyncExitStack() as stack: + connected = await open_single_server(name, cfg, stack) + if not ready.done(): + ready.set_result(connected) + if connected: + await close_requested.wait() except BaseException as exc: if not ready.done(): ready.set_exception(exc) raise - finally: - if stack is not None: - await stack.aclose() owner = asyncio.create_task(own_connection(), name=f"mcp:{name}") connection = _OwnedMCPConnection(owner, close_requested) try: connected = await ready - except BaseException: + except BaseException as exc: close_requested.set() owner.cancel() with suppress(BaseException): await asyncio.shield(owner) + if isinstance(exc, asyncio.CancelledError) and not task_is_cancelling(): + logger.warning("MCP server '{}': connection cancelled by server/SDK", name) + return name, None raise if not connected: await connection.aclose() diff --git a/tests/tools/test_mcp_probe.py b/tests/tools/test_mcp_probe.py index c64a1172d..0cad030ac 100644 --- a/tests/tools/test_mcp_probe.py +++ b/tests/tools/test_mcp_probe.py @@ -5,10 +5,13 @@ import asyncio import socket from unittest.mock import MagicMock, patch +import httpx import pytest +from nanobot.agent.tools import mcp as mcp_mod from nanobot.agent.tools.mcp import _probe_http_url, connect_mcp_servers from nanobot.agent.tools.registry import ToolRegistry +from nanobot.config.schema import MCPServerConfig from nanobot.security.network import configure_ssrf_whitelist _PROXY_ENV_VARS = ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy") @@ -171,6 +174,58 @@ async def test_connect_skips_unreachable_sse(): assert len(registry._tools) == 0 +@pytest.mark.asyncio +async def test_connect_isolates_streamable_http_status_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A reachable endpoint returning HTTP 530 must not poison the event loop.""" + async def _reachable(_url: str) -> bool: + return True + + def _return_http_530(request: httpx.Request) -> httpx.Response: + return httpx.Response(530, text="cloudflare error 1033", request=request) + + monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, "")) + monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable) + monkeypatch.setattr( + mcp_mod, + "PinnedDNSAsyncTransport", + lambda: httpx.MockTransport(_return_http_530), + ) + + loop = asyncio.get_running_loop() + previous_exception_handler = loop.get_exception_handler() + unhandled: list[BaseException] = [] + + def _capture_unhandled(_loop: asyncio.AbstractEventLoop, context: dict) -> None: + if isinstance(context.get("exception"), BaseException): + unhandled.append(context["exception"]) + + loop.set_exception_handler(_capture_unhandled) + try: + registry = ToolRegistry() + stacks = await asyncio.wait_for( + connect_mcp_servers( + { + "cloudflare": MCPServerConfig( + type="streamableHttp", + url="https://mcp.example.com/mcp", + ) + }, + registry, + ), + timeout=5.0, + ) + await asyncio.sleep(0) + + assert stacks == {} + assert registry.tool_names == [] + assert unhandled == [] + assert not any(task.get_name() == "mcp:cloudflare" for task in asyncio.all_tasks()) + finally: + loop.set_exception_handler(previous_exception_handler) + + @pytest.mark.asyncio async def test_probe_not_called_for_stdio(): """stdio transport should not be probed — it spawns a local process.""" diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 8ece82788..e2c282ebd 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -1082,10 +1082,22 @@ async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets( @pytest.mark.asyncio +@pytest.mark.parametrize("failure_mode", ["exception", "cancellation"]) async def test_connect_mcp_servers_one_failure_does_not_block_others( monkeypatch: pytest.MonkeyPatch, + failure_mode: str, ) -> None: - sessions = {"good": _make_fake_session(["demo"])} + bad_session = _make_fake_session([]) + + async def _cancel_initialize() -> None: + raise asyncio.CancelledError("cancelled by SDK") + + if failure_mode == "cancellation": + bad_session.initialize = _cancel_initialize + sessions = { + "bad": bad_session, + "good": _make_fake_session(["demo"]), + } class _SelectiveClientSession: def __init__(self, read: object, _write: object) -> None: @@ -1099,7 +1111,7 @@ async def test_connect_mcp_servers_one_failure_does_not_block_others( @asynccontextmanager async def _selective_stdio_client(params: object): - if params.command == "bad": + if params.command == "bad" and failure_mode == "exception": raise RuntimeError("boom") yield params.command, object() @@ -1109,8 +1121,8 @@ async def test_connect_mcp_servers_one_failure_does_not_block_others( registry = ToolRegistry() stacks = await connect_mcp_servers( { - "good": MCPServerConfig(command="good"), "bad": MCPServerConfig(command="bad"), + "good": MCPServerConfig(command="good"), }, registry, ) @@ -1121,6 +1133,36 @@ async def test_connect_mcp_servers_one_failure_does_not_block_others( assert set(stacks) == {"good"} +@pytest.mark.asyncio +async def test_connect_mcp_servers_propagates_external_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + closed = asyncio.Event() + + @asynccontextmanager + async def _blocking_stdio_client(_params: object): + try: + started.set() + await asyncio.Event().wait() + yield object(), object() + finally: + closed.set() + + monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _blocking_stdio_client) + + task = asyncio.create_task( + connect_mcp_servers({"slow": MCPServerConfig(command="slow")}, ToolRegistry()) + ) + await asyncio.wait_for(started.wait(), timeout=1.0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.wait_for(closed.wait(), timeout=1.0) + + @pytest.mark.asyncio async def test_connect_mcp_servers_streamable_http_uses_finite_timeout( fake_mcp_runtime: dict[str, object | None],