diff --git a/src/mcp_manager.py b/src/mcp_manager.py index 53112009d..b47d0e1d7 100644 --- a/src/mcp_manager.py +++ b/src/mcp_manager.py @@ -10,6 +10,7 @@ import logging import os import re import asyncio +import contextlib from typing import Any, Dict, List, Optional, Set, Tuple from src.database import McpServer, SessionLocal @@ -383,12 +384,21 @@ class McpManager: provider = build_provider(server_id, url, on_redirect=_on_redirect) stack = AsyncExitStack() - transport = await stack.enter_async_context(streamablehttp_client(url, auth=provider)) - read_stream, write_stream, _get_session_id = transport - session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) - await session.initialize() + try: + transport = await stack.enter_async_context(streamablehttp_client(url, auth=provider)) + read_stream, write_stream, _get_session_id = transport + session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) + await session.initialize() - tools_result = await session.list_tools() + tools_result = await session.list_tools() + except BaseException: + # Close the half-open transport here, in the task that entered + # it. Left for the garbage collector, its AnyIO cancel scope is + # exited from the wrong task and the event loop keeps spinning + # (#5518). Its own teardown error would only mask the real one. + with contextlib.suppress(Exception): + await stack.aclose() + raise tools = [] for tool in tools_result.tools: tools.append({ diff --git a/tests/test_mcp_manager.py b/tests/test_mcp_manager.py index a879f95eb..11e655422 100644 --- a/tests/test_mcp_manager.py +++ b/tests/test_mcp_manager.py @@ -39,3 +39,74 @@ def test_http_transport_routes_to_start_http_connect(): result = asyncio.run(mgr.connect_server("id1", "n", "http", url="https://x/mcp")) assert result == "ROUTED" m.assert_called_once() + + +def _patch_http_transport(monkeypatch, initialize): + """Fake the Streamable HTTP transport and record, for each close, whether + it ran in the task that opened it (AnyIO requires that).""" + import contextlib + import mcp + import mcp.client.streamable_http as streamable_http + import src.mcp_oauth as mcp_oauth + + closes = [] + + @contextlib.asynccontextmanager + async def fake_client(url, auth=None): + opener = asyncio.current_task() + try: + yield (object(), object(), lambda: None) + finally: + closes.append(asyncio.current_task() is opener) + + class FakeSession: + def __init__(self, *a): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + FakeSession.initialize = initialize + monkeypatch.setattr(streamable_http, "streamablehttp_client", fake_client) + monkeypatch.setattr(mcp, "ClientSession", FakeSession) + monkeypatch.setattr(mcp_oauth, "build_provider", lambda *a, **k: None) + return closes + + +def test_failed_http_connect_closes_transport_in_its_own_task(monkeypatch): + """#5518: a failed connect left its transport for the GC, which exited the + AnyIO cancel scope from the wrong task and kept the event loop spinning.""" + async def initialize(self): + raise RuntimeError("Session terminated") + + closes = _patch_http_transport(monkeypatch, initialize) + mgr = McpManager() + + ok = asyncio.run(mgr._start_http_connect("s1", "remote", "http://x/mcp/", wait=2)) + + assert ok is False + assert mgr._connections["s1"]["status"] == "error" + assert "Session terminated" in mgr._connections["s1"]["error"] + assert closes == [True] + assert "s1" not in mgr._stacks + + +def test_cancelled_http_connect_closes_transport(monkeypatch): + """Deleting a server mid-connect cancels the connect task; the transport + must still be closed rather than leaked.""" + async def initialize(self): + await asyncio.sleep(3600) + + closes = _patch_http_transport(monkeypatch, initialize) + mgr = McpManager() + + async def run(): + await mgr._start_http_connect("s1", "remote", "http://x/mcp/", wait=0.05) + await mgr.disconnect_server("s1") + await asyncio.sleep(0) + + asyncio.run(run()) + assert closes == [True]