fix(mcp): keep transport cleanup in owner tasks

tests/agent/test_mcp_connection.py

@@ -74,15 +74,6 @@ class _FakeMcpTool(Tool):
         return "ok"
 
 
-class _CancelScopeStack:
-    def __init__(self) -> None:
-        self.closed = False
-
-    async def aclose(self) -> None:
-        self.closed = True
-        raise asyncio.CancelledError("Cancelled via cancel scope test")
-
-
 def _make_loop(tmp_path, *, mcp_servers: dict | None = None) -> AgentLoop:
     bus = MessageBus()
     provider = MagicMock()
@@ -125,15 +116,26 @@ async def test_mcp_read_filter_drops_progress_notifications_without_progress_tok
 
 
 @pytest.mark.asyncio
-async def test_close_mcp_swallows_sdk_cancel_scope_cleanup(tmp_path):
-    loop = _make_loop(tmp_path)
-    stack = _CancelScopeStack()
-    loop._mcp_stacks["remote"] = stack  # type: ignore[assignment]
+async def test_owned_mcp_connection_closes_from_its_owner_task():
+    close_requested = asyncio.Event()
+    ready = asyncio.Event()
+    tasks: dict[str, asyncio.Task] = {}
 
-    await loop.close_mcp()
+    async def own_connection() -> None:
+        tasks["open"] = asyncio.current_task()  # type: ignore[assignment]
+        ready.set()
+        await close_requested.wait()
+        tasks["close"] = asyncio.current_task()  # type: ignore[assignment]
 
-    assert stack.closed is True
-    assert loop._mcp_stacks == {}
+    owner = asyncio.create_task(own_connection())
+    connection = mcp_runtime._OwnedMCPConnection(owner, close_requested)
+    await ready.wait()
+
+    await connection.aclose()
+
+    assert tasks["open"] is owner
+    assert tasks["close"] is owner
+    assert tasks["close"] is not asyncio.current_task()
 
 
 @pytest.mark.asyncio
@@ -241,9 +243,6 @@ async def test_reload_mcp_servers_adds_and_removes_tools_without_restart(
     assert removed["removed"] == ["browserbase"]
     assert not loop.tools.has("mcp_browserbase_navigate")
     assert "browserbase" not in loop._mcp_stacks
-
-    assert closed == []
-    await loop.close_mcp()
     assert closed == ["browserbase"]
 
 
@@ -305,9 +304,6 @@ async def test_request_mcp_reload_reaches_runtime_control_without_restart(
     assert result["removed"] == ["browserbase"]
     assert result["requires_restart"] is False
     assert not loop.tools.has("mcp_browserbase_navigate")
-
-    assert closed == []
-    await loop.close_mcp()
     assert closed == ["browserbase"]
 
 
@@ -403,15 +399,12 @@ async def test_mcp_tool_reconnects_after_session_terminated(
 
     assert output == "recovered"
     assert connect_count == 2
-    assert closed == []
+    assert closed == ["remote"]
     assert sessions[0].call_count == 1
     assert sessions[1].call_count == 1
     assert "remote" in loop._mcp_stacks
     assert loop.tools.get("mcp_remote_quote") is not old_tool
 
-    await loop.close_mcp()
-    assert closed == ["remote", "remote"]
-
 
 @pytest.mark.asyncio
 async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
@@ -521,7 +514,4 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session(
 
     assert outputs == ["fresh:alpha", "fresh:beta"]
     assert connect_count == 2
-    assert closed == []
-
-    await loop.close_mcp()
-    assert closed == ["remote", "remote"]
+    assert closed == ["remote"]

tests/agent/test_mcp_reconnect_crash.py

@@ -15,7 +15,6 @@ import asyncio
 import multiprocessing
 import socket
 import time
-from contextlib import suppress
 from unittest.mock import MagicMock
 
 import httpx
@@ -70,7 +69,7 @@ async def _wait_for_server(url: str, timeout: float = 10.0) -> bool:
     deadline = time.monotonic() + timeout
     while time.monotonic() < deadline:
         try:
-            async with httpx.AsyncClient(timeout=2.0) as client:
+            async with httpx.AsyncClient(timeout=2.0, trust_env=False) as client:
                 response = await client.get(
                     url,
                     headers={"Accept": "text/event-stream"},
@@ -170,26 +169,30 @@ async def test_mcp_reconnect_after_session_timeout(tmp_path, mcp_server_url):
     )
     loop = _make_loop(tmp_path, mcp_servers={"repro": cfg})
 
-    await loop._connect_mcp()
+    await asyncio.create_task(loop._connect_mcp())
     assert "repro" in loop._mcp_stacks
 
     tool = loop.tools.get("mcp_repro_greet")
     assert isinstance(tool, MCPToolWrapper)
 
-    output = await tool.execute(name="first")
+    output = await asyncio.create_task(tool.execute(name="first"))
     assert "Hello, first" in output
 
     # Wait for the server-side idle timeout to terminate the session.
     await asyncio.sleep(_IDLE_TIMEOUT_SECONDS + 1)
 
-    output = await tool.execute(name="second")
+    output = await asyncio.create_task(tool.execute(name="second"))
     assert "Hello, second" in output
 
-    await loop.close_mcp()
+    await asyncio.create_task(loop.close_mcp())
 
 
 @pytest.mark.asyncio
-async def test_mcp_reconnect_during_shutdown_does_not_crash(tmp_path, mcp_server_url):
+async def test_mcp_reconnect_during_shutdown_does_not_crash(
+    tmp_path,
+    mcp_server_url,
+    monkeypatch: pytest.MonkeyPatch,
+):
     """Simulate the production crash: shutdown while reconnect is in flight."""
     cfg = MCPServerConfig(
         type="streamableHttp",
@@ -199,15 +202,28 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash(tmp_path, mcp_server
     )
     loop = _make_loop(tmp_path, mcp_servers={"repro": cfg})
 
-    await loop._connect_mcp()
+    await asyncio.create_task(loop._connect_mcp())
     tool = loop.tools.get("mcp_repro_greet")
     assert isinstance(tool, MCPToolWrapper)
 
-    await tool.execute(name="first")
+    await asyncio.create_task(tool.execute(name="first"))
     await asyncio.sleep(_IDLE_TIMEOUT_SECONDS + 1)
 
+    reconnect_started = asyncio.Event()
+    finish_reconnect = asyncio.Event()
+    real_connect = mcp_module.connect_mcp_servers
+
+    async def gated_connect(*args, **kwargs):
+        reconnect_started.set()
+        await finish_reconnect.wait()
+        return await real_connect(*args, **kwargs)
+
+    monkeypatch.setattr(mcp_module, "connect_mcp_servers", gated_connect)
     call_task = asyncio.create_task(tool.execute(name="second"))
-    loop.stop()
+    await asyncio.wait_for(reconnect_started.wait(), timeout=5)
+    close_task = asyncio.create_task(loop.close_mcp())
+    await asyncio.sleep(0)
+    finish_reconnect.set()
 
     unhandled: list[BaseException] = []
 
@@ -219,13 +235,11 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash(tmp_path, mcp_server
     asyncio.get_running_loop().set_exception_handler(capture_unhandled)
 
     try:
-        await asyncio.wait_for(asyncio.shield(call_task), timeout=15)
+        await asyncio.wait_for(asyncio.gather(call_task, close_task), timeout=15)
     except asyncio.CancelledError:
         unhandled.append(asyncio.CancelledError("main task cancelled by leaked MCP cancel scope"))
     except Exception as exc:
         unhandled.append(exc)
 
-    with suppress(Exception):
-        await loop.close_mcp()
-
     assert not unhandled, f"Unhandled exception leaked during reconnect/shutdown: {unhandled[0]}"
+    assert loop._mcp_stacks == {}