From f56d6048469d2f92e75b2fbd4eed4a8937efda4a Mon Sep 17 00:00:00 2001
From: hezihong <hezihong@bytedance.com>
Date: Mon, 10 Aug 2026 16:54:44 +0800
Subject: [PATCH] fix(agent): keep MCP connection failures from cancelling
 gateway

---
 raven/agent/tools/mcp.py   |  7 +++++
 tests/test_sandbox_unit.py | 56 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/raven/agent/tools/mcp.py b/raven/agent/tools/mcp.py
index 3c50d7cf..9e586c01 100644
--- a/raven/agent/tools/mcp.py
+++ b/raven/agent/tools/mcp.py
@@ -165,6 +165,13 @@ def httpx_client_factory(
                 logger.debug("MCP: registered tool '{}' from server '{}'", wrapper.name, name)
 
             logger.info("MCP server '{}': connected, {} tools registered", name, len(tools.tools))
+        except asyncio.CancelledError:
+            # MCP SDK cancel scopes can leak CancelledError when a transport fails.
+            # Preserve cancellation only when Raven's own task was cancelled.
+            task = asyncio.current_task()
+            if task is not None and task.cancelling() > 0:
+                raise
+            logger.error("MCP server '{}': connection cancelled by server/SDK", name)
         except (Exception, BaseExceptionGroup) as e:
             # BaseExceptionGroup is raised by anyio task groups (e.g. streamableHttp cancel
             # scope failures) and is not a subclass of Exception in Python 3.11+.
diff --git a/tests/test_sandbox_unit.py b/tests/test_sandbox_unit.py
index a66561b0..dd4af3e3 100644
--- a/tests/test_sandbox_unit.py
+++ b/tests/test_sandbox_unit.py
@@ -1052,6 +1052,62 @@ def fake_stdio_client(params):
 
         assert reached == ["mcp-server"]
 
+    async def test_stdio_sdk_cancellation_does_not_escape(self, monkeypatch):
+        """A bare transport cancellation is isolated to the failing MCP server."""
+        from contextlib import AsyncExitStack
+
+        import mcp.client.stdio
+
+        from raven.agent.tools.mcp import connect_mcp_servers
+        from raven.agent.tools.registry import ToolRegistry
+
+        def fake_stdio_client(params):
+            raise asyncio.CancelledError()
+
+        monkeypatch.setattr(mcp.client.stdio, "stdio_client", fake_stdio_client)
+
+        cfg = MagicMock()
+        cfg.type = "stdio"
+        cfg.command = "mcp-server"
+        cfg.args = []
+        cfg.env = None
+        cfg.tool_timeout = 30
+
+        await connect_mcp_servers({"svc": cfg}, ToolRegistry(), AsyncExitStack(), executor=None)
+
+    async def test_stdio_external_cancellation_propagates(self, monkeypatch):
+        """Cancellation of Raven's connection task is never treated as an MCP failure."""
+        from contextlib import AsyncExitStack, asynccontextmanager
+
+        import mcp.client.stdio
+
+        from raven.agent.tools.mcp import connect_mcp_servers
+        from raven.agent.tools.registry import ToolRegistry
+
+        entered = asyncio.Event()
+
+        @asynccontextmanager
+        async def fake_stdio_client(params):
+            entered.set()
+            await asyncio.Event().wait()
+            yield
+
+        monkeypatch.setattr(mcp.client.stdio, "stdio_client", fake_stdio_client)
+
+        cfg = MagicMock()
+        cfg.type = "stdio"
+        cfg.command = "mcp-server"
+        cfg.args = []
+        cfg.env = None
+        cfg.tool_timeout = 30
+
+        task = asyncio.create_task(connect_mcp_servers({"svc": cfg}, ToolRegistry(), AsyncExitStack(), executor=None))
+        await entered.wait()
+        task.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
     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