diff --git a/raven/agent/tools/mcp.py b/raven/agent/tools/mcp.py index 3c50d7cf..32c6585e 100644 --- a/raven/agent/tools/mcp.py +++ b/raven/agent/tools/mcp.py @@ -1,7 +1,7 @@ """MCP client: connects to MCP servers and wraps their tools as native Raven tools.""" import asyncio -from contextlib import AsyncExitStack +from contextlib import AsyncExitStack, asynccontextmanager from typing import TYPE_CHECKING, Any import httpx @@ -16,6 +16,54 @@ from raven.sandbox import SandboxExecutor +@asynccontextmanager +async def _mcp_server_connection(cfg, transport_type: str, executor: "SandboxExecutor | None"): + async with AsyncExitStack() as stack: + from mcp import ClientSession, StdioServerParameters + from mcp.client.sse import sse_client + from mcp.client.stdio import stdio_client + from mcp.client.streamable_http import streamable_http_client + + if transport_type == "stdio": + if executor is not None and executor.supports_process_spawning: + read, write = await executor.start_process(cfg.command, cfg.args, env=cfg.env or None) + else: + params = StdioServerParameters(command=cfg.command, args=cfg.args, env=cfg.env or None) + read, write = await stack.enter_async_context(stdio_client(params)) + elif transport_type == "sse": + + def httpx_client_factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + merged_headers = {**(cfg.headers or {}), **(headers or {})} + return httpx.AsyncClient( + headers=merged_headers or None, + follow_redirects=True, + timeout=timeout, + auth=auth, + ) + + read, write = await stack.enter_async_context( + sse_client(cfg.url, httpx_client_factory=httpx_client_factory) + ) + else: + http_client = await stack.enter_async_context( + httpx.AsyncClient( + headers=cfg.headers or None, + follow_redirects=True, + timeout=None, + ) + ) + read, write, _ = await stack.enter_async_context(streamable_http_client(cfg.url, http_client=http_client)) + + session = await stack.enter_async_context(ClientSession(read, write)) + await session.initialize() + tools = await session.list_tools() + yield session, tools + + class MCPToolWrapper(Tool): """Wraps a single MCP server tool as an Raven Tool.""" @@ -110,55 +158,12 @@ async def connect_mcp_servers( "Either switch to an HTTP/SSE MCP server or set sandbox.backend='none'." ) - try: - from mcp import ClientSession, StdioServerParameters - from mcp.client.sse import sse_client - from mcp.client.stdio import stdio_client - from mcp.client.streamable_http import streamable_http_client - - if transport_type == "stdio": - if executor is not None and executor.supports_process_spawning: - read, write = await executor.start_process(cfg.command, cfg.args, env=cfg.env or None) - else: - params = StdioServerParameters(command=cfg.command, args=cfg.args, env=cfg.env or None) - read, write = await stack.enter_async_context(stdio_client(params)) - elif transport_type == "sse": - - def httpx_client_factory( - headers: dict[str, str] | None = None, - timeout: httpx.Timeout | None = None, - auth: httpx.Auth | None = None, - ) -> httpx.AsyncClient: - merged_headers = {**(cfg.headers or {}), **(headers or {})} - return httpx.AsyncClient( - headers=merged_headers or None, - follow_redirects=True, - timeout=timeout, - auth=auth, - ) - - read, write = await stack.enter_async_context( - sse_client(cfg.url, httpx_client_factory=httpx_client_factory) - ) - elif transport_type == "streamableHttp": - http_client = await stack.enter_async_context( - httpx.AsyncClient( - headers=cfg.headers or None, - follow_redirects=True, - timeout=None, - ) - ) - read, write, _ = await stack.enter_async_context( - streamable_http_client(cfg.url, http_client=http_client) - ) - else: - logger.warning("MCP server '{}': unknown transport type '{}'", name, transport_type) - continue + if transport_type not in {"stdio", "sse", "streamableHttp"}: + logger.warning("MCP server '{}': unknown transport type '{}'", name, transport_type) + continue - session = await stack.enter_async_context(ClientSession(read, write)) - await session.initialize() - - tools = await session.list_tools() + try: + session, tools = await stack.enter_async_context(_mcp_server_connection(cfg, transport_type, executor)) for tool_def in tools.tools: wrapper = MCPToolWrapper(session, name, tool_def, tool_timeout=cfg.tool_timeout) registry.register(wrapper) diff --git a/tests/test_sandbox_unit.py b/tests/test_sandbox_unit.py index 93945e63..8f2a5f32 100644 --- a/tests/test_sandbox_unit.py +++ b/tests/test_sandbox_unit.py @@ -1052,6 +1052,221 @@ def fake_stdio_client(params): assert reached == ["mcp-server"] + async def test_streamable_failure_does_not_cancel_following_server(self, monkeypatch): + """A transport task failure is isolated to the server being initialized.""" + from contextlib import AsyncExitStack, asynccontextmanager + from types import SimpleNamespace + + import anyio + import mcp + import mcp.client.streamable_http + + from raven.agent.tools.mcp import connect_mcp_servers + from raven.agent.tools.registry import ToolRegistry + + attempted = [] + + @asynccontextmanager + async def fake_streamable_http_client(url, http_client): + attempted.append(url) + if url == "https://bad.example/mcp": + async with anyio.create_task_group() as group: + + async def fail_transport(): + await anyio.sleep(0) + raise RuntimeError("transport failed") + + group.start_soon(fail_transport) + yield url, object(), None + else: + yield url, object(), None + + class FakeSession: + def __init__(self, read, write): + self.read = read + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + async def initialize(self): + if self.read == "https://bad.example/mcp": + await asyncio.Event().wait() + + async def list_tools(self): + return SimpleNamespace(tools=[]) + + monkeypatch.setattr(mcp, "ClientSession", FakeSession) + monkeypatch.setattr( + mcp.client.streamable_http, + "streamable_http_client", + fake_streamable_http_client, + ) + + def config(url): + return SimpleNamespace(type="streamableHttp", url=url, headers=None, tool_timeout=30) + + async with AsyncExitStack() as stack: + await connect_mcp_servers( + { + "bad": config("https://bad.example/mcp"), + "good": config("https://good.example/mcp"), + }, + ToolRegistry(), + stack, + ) + + assert attempted == ["https://bad.example/mcp", "https://good.example/mcp"] + assert asyncio.current_task().cancelling() == 0 + + async def test_streamable_external_cancellation_propagates(self, monkeypatch): + """Cancellation of Raven's connection task is not treated as a server failure.""" + from contextlib import AsyncExitStack, asynccontextmanager + from types import SimpleNamespace + + import mcp + import mcp.client.streamable_http + + from raven.agent.tools.mcp import connect_mcp_servers + from raven.agent.tools.registry import ToolRegistry + + entered = asyncio.Event() + + @asynccontextmanager + async def fake_streamable_http_client(url, http_client): + yield object(), object(), None + + class FakeSession: + def __init__(self, read, write): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + async def initialize(self): + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(mcp, "ClientSession", FakeSession) + monkeypatch.setattr( + mcp.client.streamable_http, + "streamable_http_client", + fake_streamable_http_client, + ) + + cfg = SimpleNamespace( + type="streamableHttp", + url="https://wait.example/mcp", + headers=None, + tool_timeout=30, + ) + task = asyncio.create_task(connect_mcp_servers({"svc": cfg}, ToolRegistry(), AsyncExitStack())) + await entered.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + async def test_sse_connection_uses_merged_headers(self, monkeypatch): + """The SSE transport keeps config headers when the SDK adds its own.""" + from contextlib import AsyncExitStack, asynccontextmanager + from types import SimpleNamespace + + import httpx + import mcp + import mcp.client.sse + + from raven.agent.tools.mcp import connect_mcp_servers + from raven.agent.tools.registry import ToolRegistry + + clients = [] + + class FakeHttpClient: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + def fake_http_client(**kwargs): + clients.append(kwargs) + return FakeHttpClient() + + @asynccontextmanager + async def fake_sse_client(url, httpx_client_factory): + client = httpx_client_factory( + headers={"X-SDK": "sdk", "X-Shared": "sdk"}, + timeout="timeout", + auth="auth", + ) + async with client: + yield object(), object() + + class FakeSession: + def __init__(self, read, write): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + async def initialize(self): + pass + + async def list_tools(self): + return SimpleNamespace(tools=[]) + + monkeypatch.setattr(httpx, "AsyncClient", fake_http_client) + monkeypatch.setattr(mcp, "ClientSession", FakeSession) + monkeypatch.setattr(mcp.client.sse, "sse_client", fake_sse_client) + + cfg = SimpleNamespace( + type="sse", + url="https://example.test/sse", + headers={"X-Config": "config", "X-Shared": "config"}, + tool_timeout=30, + ) + async with AsyncExitStack() as stack: + await connect_mcp_servers({"svc": cfg}, ToolRegistry(), stack) + + assert clients == [ + { + "headers": {"X-Config": "config", "X-Shared": "sdk", "X-SDK": "sdk"}, + "follow_redirects": True, + "timeout": "timeout", + "auth": "auth", + } + ] + + async def test_unknown_transport_is_skipped(self, monkeypatch): + """An unknown transport does not attempt to open an MCP connection.""" + from contextlib import AsyncExitStack + from types import SimpleNamespace + + from raven.agent.tools import mcp as mcp_tools + from raven.agent.tools.registry import ToolRegistry + + attempted = False + + def fake_connection(cfg, transport_type, executor): + nonlocal attempted + attempted = True + raise AssertionError("unknown transport attempted a connection") + + monkeypatch.setattr(mcp_tools, "_mcp_server_connection", fake_connection) + cfg = SimpleNamespace(type="websocket", command=None, url="wss://example.test/mcp") + + await mcp_tools.connect_mcp_servers({"svc": cfg}, ToolRegistry(), AsyncExitStack()) + + assert attempted is False + async def test_stdio_sandboxed_with_spawning_does_not_raise(self): """Sandboxed executor that supports spawning does not trigger the guard.""" from contextlib import AsyncExitStack