mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-10-31 02:02:22 +01:00
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:
+15
-5
@@ -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({
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
Reference in New Issue
Block a user