fix(mcp): close the HTTP transport when a remote connect fails

_connect_http entered streamablehttp_client and ClientSession on an
AsyncExitStack but never closed it when initialize or list_tools
failed. The half-open transport was left to the garbage collector,
which exited its AnyIO cancel scope from the wrong task, and the event
loop kept spinning: about 18% of a core per failed connect, growing
with each one, until restart.

Close the stack in the connecting task on any failure or cancellation,
the same way _connect_sse already does.

Fixes #5518
This commit is contained in:
harshit-ojha0324
2026-10-06 19:01:10 -04:00
parent ed0c8d877d
commit 50b05bf54f
2 changed files with 86 additions and 5 deletions
+15 -5
View File
@@ -10,6 +10,7 @@ import logging
import os import os
import re import re
import asyncio import asyncio
import contextlib
from typing import Any, Dict, List, Optional, Set, Tuple from typing import Any, Dict, List, Optional, Set, Tuple
from src.database import McpServer, SessionLocal from src.database import McpServer, SessionLocal
@@ -383,12 +384,21 @@ class McpManager:
provider = build_provider(server_id, url, on_redirect=_on_redirect) provider = build_provider(server_id, url, on_redirect=_on_redirect)
stack = AsyncExitStack() stack = AsyncExitStack()
transport = await stack.enter_async_context(streamablehttp_client(url, auth=provider)) try:
read_stream, write_stream, _get_session_id = transport transport = await stack.enter_async_context(streamablehttp_client(url, auth=provider))
session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) read_stream, write_stream, _get_session_id = transport
await session.initialize() 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 = [] tools = []
for tool in tools_result.tools: for tool in tools_result.tools:
tools.append({ tools.append({
+71
View File
@@ -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")) result = asyncio.run(mgr.connect_server("id1", "n", "http", url="https://x/mcp"))
assert result == "ROUTED" assert result == "ROUTED"
m.assert_called_once() 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]