diff --git a/tests/tools/test_mcp_initial_connect_shutdown.py b/tests/tools/test_mcp_initial_connect_shutdown.py new file mode 100644 index 00000000000..049e319bef6 --- /dev/null +++ b/tests/tools/test_mcp_initial_connect_shutdown.py @@ -0,0 +1,273 @@ +"""Regression tests for initial MCP failure ownership and teardown.""" + +import asyncio +import json +import threading +from types import SimpleNamespace + +import pytest + + +def _reset_mcp_state(mcp_tool) -> None: + mcp_tool.shutdown_mcp_servers() + with mcp_tool._lock: + mcp_tool._servers.clear() + mcp_tool._server_connecting.clear() + mcp_tool._server_connect_errors.clear() + + +def _cleanup_mcp_state(mcp_tool, extra_servers=()) -> None: + with mcp_tool._lock: + loop = mcp_tool._mcp_loop + if loop is not None and loop.is_running(): + for server in extra_servers: + task = getattr(server, "_task", None) + if task is not None and not task.done(): + mcp_tool._run_on_mcp_loop(server.shutdown, timeout=5) + mcp_tool.shutdown_mcp_servers() + with mcp_tool._lock: + mcp_tool._servers.clear() + mcp_tool._server_connecting.clear() + mcp_tool._server_connect_errors.clear() + + +def test_initial_connect_failure_is_registry_owned_and_reaped(monkeypatch, tmp_path): + """Normal discovery must retain the parked task for clean shutdown.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + + from tools import mcp_tool + + _reset_mcp_state(mcp_tool) + created = [] + + class _FailingServerTask(mcp_tool.MCPServerTask): + def __init__(self, name): + super().__init__(name) + created.append(self) + + async def _run_stdio(self, config): + raise ConnectionError("deterministic initial failure") + + monkeypatch.setattr(mcp_tool, "MCPServerTask", _FailingServerTask) + monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True) + monkeypatch.setattr(mcp_tool, "_MAX_INITIAL_CONNECT_RETRIES", 0) + monkeypatch.setattr(mcp_tool, "_PARKED_RETRY_INTERVAL", 3600) + + real_stop = mcp_tool._stop_mcp_loop + pending_at_stop = [] + + async def _pending_tasks(): + current = asyncio.current_task() + return sorted( + task.get_coro().__qualname__ + for task in asyncio.all_tasks() + if task is not current and not task.done() + ) + + def _observed_stop(*, only_if_idle=False): + pending_at_stop.extend( + mcp_tool._run_on_mcp_loop(_pending_tasks, timeout=5) + ) + return real_stop(only_if_idle=only_if_idle) + + monkeypatch.setattr(mcp_tool, "_stop_mcp_loop", _observed_stop) + + try: + assert mcp_tool.register_mcp_servers({ + "initial-failure": {"command": "unused", "connect_timeout": 5} + }) == [] + + assert len(created) == 1 + server = created[0] + with mcp_tool._lock: + assert mcp_tool._servers["initial-failure"] is server + assert "deterministic initial failure" in ( + mcp_tool._server_connect_errors["initial-failure"] + ) + assert server._task is not None + assert not server._task.done(), "recoverable initial failure was not parked" + + mcp_tool.shutdown_mcp_servers() + + assert pending_at_stop == [], ( + "shutdown left MCP tasks pending at loop stop: " + f"{pending_at_stop!r}" + ) + assert server._task.done() + with mcp_tool._lock: + assert mcp_tool._mcp_loop is None + assert mcp_tool._mcp_thread is None + finally: + monkeypatch.setattr(mcp_tool, "_stop_mcp_loop", real_stop) + _cleanup_mcp_state(mcp_tool, created) + + +def test_initial_connect_failure_revives_same_registered_server(monkeypatch, tmp_path): + """A cached parked failure must revive through register_mcp_servers().""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + + from tools import mcp_tool + from tools.registry import ToolRegistry + import tools.registry as registry_module + + _reset_mcp_state(mcp_tool) + created = [] + backend_up = threading.Event() + revived = threading.Event() + state = {"transport_calls": 0, "tool_calls": 0} + mock_registry = ToolRegistry() + + class _Session: + async def call_tool(self, name, arguments): + state["tool_calls"] += 1 + return SimpleNamespace( + isError=False, + content=[SimpleNamespace(text=f"revived:{arguments['value']}")], + structuredContent=None, + ) + + class _RecoveringServerTask(mcp_tool.MCPServerTask): + def __init__(self, name): + super().__init__(name) + created.append(self) + + async def _run_stdio(self, config): + assert mcp_tool._connect_server_claim.get() is None + state["transport_calls"] += 1 + if not backend_up.is_set(): + raise ConnectionError("backend still booting") + + self.session = _Session() + self._tools = [SimpleNamespace( + name="ping", + description="Return a deterministic revival result", + inputSchema={ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + )] + # Match the real transports: discovery runs before _ready is set. + self._register_discovered_tools_if_needed() + self._ready.set() + revived.set() + return await self._wait_for_lifecycle_event() + + monkeypatch.setattr(mcp_tool, "MCPServerTask", _RecoveringServerTask) + monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True) + monkeypatch.setattr(mcp_tool, "_MAX_INITIAL_CONNECT_RETRIES", 0) + monkeypatch.setattr(mcp_tool, "_PARKED_RETRY_INTERVAL", 3600) + monkeypatch.setattr(registry_module, "registry", mock_registry) + + config = { + "recovering": {"command": "unused", "connect_timeout": 5} + } + + try: + assert mcp_tool.register_mcp_servers(config) == [] + assert len(created) == 1 + server = created[0] + with mcp_tool._lock: + assert mcp_tool._servers["recovering"] is server + assert "backend still booting" in ( + mcp_tool._server_connect_errors["recovering"] + ) + assert not server._task.done() + + backend_up.set() + mcp_tool.register_mcp_servers(config) + + assert revived.wait(timeout=5), "cached parked server did not revive" + assert len(created) == 1, "revival created a duplicate server task" + with mcp_tool._lock: + assert mcp_tool._servers["recovering"] is server + assert "recovering" not in mcp_tool._server_connect_errors + assert state["transport_calls"] == 2 + assert server.session is not None + assert server._error is None + + entry = mock_registry.get_entry("mcp__recovering__ping") + assert entry is not None + assert entry.check_fn() is True + assert json.loads(entry.handler({"value": "ok"})) == { + "result": "revived:ok" + } + assert state["tool_calls"] == 1 + finally: + _cleanup_mcp_state(mcp_tool, created) + + +def test_terminal_initial_failure_is_not_retained(monkeypatch, tmp_path): + """A non-recoverable startup error must not leave a dead cache entry.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + + from tools import mcp_tool + + _reset_mcp_state(mcp_tool) + created = [] + + class _AuthFailingServerTask(mcp_tool.MCPServerTask): + def __init__(self, name): + super().__init__(name) + created.append(self) + + async def _run_stdio(self, config): + raise PermissionError("terminal authentication failure") + + monkeypatch.setattr(mcp_tool, "MCPServerTask", _AuthFailingServerTask) + monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True) + monkeypatch.setattr(mcp_tool, "_is_auth_error", lambda exc: True) + + try: + assert mcp_tool.register_mcp_servers({ + "auth-failure": {"command": "unused", "connect_timeout": 5} + }) == [] + assert len(created) == 1 + assert created[0]._task.done() + with mcp_tool._lock: + assert "auth-failure" not in mcp_tool._servers + assert "terminal authentication failure" in ( + mcp_tool._server_connect_errors["auth-failure"] + ) + finally: + _cleanup_mcp_state(mcp_tool, created) + + +def test_standalone_failed_connect_is_reaped_without_global_owner(monkeypatch, tmp_path): + """Probe-only _connect_server failures must not publish parked servers.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + + from tools import mcp_tool + + _reset_mcp_state(mcp_tool) + created = [] + + class _ProbeServerTask(mcp_tool.MCPServerTask): + def __init__(self, name): + super().__init__(name) + created.append(self) + + async def _run_stdio(self, config): + raise ConnectionError("probe target unavailable") + + monkeypatch.setattr(mcp_tool, "MCPServerTask", _ProbeServerTask) + monkeypatch.setattr(mcp_tool, "_MAX_INITIAL_CONNECT_RETRIES", 0) + monkeypatch.setattr(mcp_tool, "_PARKED_RETRY_INTERVAL", 3600) + mcp_tool._ensure_mcp_loop() + + try: + with pytest.raises(ConnectionError, match="probe target unavailable"): + mcp_tool._run_on_mcp_loop( + lambda: mcp_tool._connect_server( + "probe-only", {"command": "unused"} + ), + timeout=5, + ) + + assert len(created) == 1 + assert created[0]._task.done() + with mcp_tool._lock: + assert "probe-only" not in mcp_tool._servers + assert "probe-only" not in mcp_tool._server_connect_errors + finally: + _cleanup_mcp_state(mcp_tool, created) diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index dd3f353a906..9895da2bc6f 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -3021,7 +3021,7 @@ class MCPServerTask: self._register_discovered_tools_if_needed() def _register_discovered_tools_if_needed(self) -> None: - """Re-register tools after a post-ready reconnect if needed. + """Re-register tools after an owned server reconnects if needed. Initial registration is performed by ``_discover_and_register_server`` after ``start()`` completes. During a later reconnect, outage handling @@ -3029,6 +3029,9 @@ class MCPServerTask: A managed server can still be identified by its entry in ``_servers``; publish its freshly discovered tools before transport readiness is restored so a successful revival cannot come back with zero tools. + A server retained after a recoverable initial failure is likewise + registry-owned before its first successful session, so ownership also + authorizes its first publication. """ if self._registered_tool_names: return @@ -3039,6 +3042,12 @@ class MCPServerTask: self._registered_tool_names = _register_server_tools( self.name, self, self._config ) + # A retained initial-failure server that just published tools has + # recovered: drop its stale connect error so status surfaces stop + # reporting it as failed. + with _lock: + if _servers.get(self.name) is self: + _server_connect_errors.pop(self.name, None) async def run(self, config: dict): """Long-lived coroutine: connect, discover tools, wait, disconnect. @@ -3523,6 +3532,12 @@ class MCPServerTask: _servers: Dict[str, MCPServerTask] = {} _server_connecting: set[str] = set() _server_connect_errors: Dict[str, str] = {} +# Discovery installs a task-local claim before calling ``_connect_server`` so +# it can retain a recoverable parked task without making standalone probe calls +# publish failed servers into module-global ownership. +_connect_server_claim: contextvars.ContextVar[ + Optional[Callable[[MCPServerTask], None]] +] = contextvars.ContextVar("mcp_connect_server_claim", default=None) # Connection-retry cooldown (per-server isolation against restart storms). # @@ -3632,9 +3647,8 @@ def _signal_reconnect(server: Any) -> bool: The tool handlers run on caller threads, while the server task and its ``_reconnect_event`` live on the background MCP loop. Setting an asyncio.Event from another thread must go through - ``loop.call_soon_threadsafe``; only fall back to a direct ``.set()`` - when the loop isn't running (e.g. unit tests that drive the handler - synchronously). + ``loop.call_soon_threadsafe``; non-async adapters and tests without a + running loop can use a direct ``.set()``. Returns True if a reconnect signal was delivered, False if the server has no reconnect machinery (nothing to revive). @@ -3643,7 +3657,11 @@ def _signal_reconnect(server: Any) -> bool: if event is None: return False loop = _mcp_loop - if loop is not None and loop.is_running(): + if ( + isinstance(event, asyncio.Event) + and loop is not None + and loop.is_running() + ): loop.call_soon_threadsafe(event.set) else: event.set() @@ -4616,7 +4634,26 @@ async def _connect_server(name: str, config: dict) -> MCPServerTask: Exception: on connection or initialization failure. """ server = MCPServerTask(name) - await server.start(config) + claim = _connect_server_claim.get() + claim_token = None + if claim is not None: + claim(server) + # ``start()`` creates the long-lived run task by copying this context. + # The ownership callback is only for this connection attempt; do not + # retain its discovery closure for the server's lifetime. + claim_token = _connect_server_claim.set(None) + try: + await server.start(config) + except BaseException: + # Discovery owns claimed tasks and decides whether a failed start is a + # live recoverable park or a terminal failure. Standalone probes have + # no revival owner, so they must reap their failed task locally. + if claim is None: + await server.shutdown() + raise + finally: + if claim_token is not None: + _connect_server_claim.reset(claim_token) return server @@ -5788,10 +5825,42 @@ async def _discover_and_register_server(name: str, config: dict) -> List[str]: Returns list of registered tool names. """ connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT) - server = await asyncio.wait_for( - _connect_server(name, config), - timeout=connect_timeout, - ) + server: Optional[MCPServerTask] = None + + def _claim_server(created: MCPServerTask) -> None: + nonlocal server + server = created + + claim_token = _connect_server_claim.set(_claim_server) + try: + server = await asyncio.wait_for( + _connect_server(name, config), + timeout=connect_timeout, + ) + except BaseException: + task = server._task if server is not None else None + task_cancelling = ( + task.cancelling() + if task is not None and hasattr(task, "cancelling") + else 0 + ) + recoverable_park = ( + server is not None + and server._error is not None + and task is not None + and not task.done() + and not task_cancelling + ) + if recoverable_park: + with _lock: + _servers[name] = server + elif server is not None: + await server.shutdown() + raise + finally: + _connect_server_claim.reset(claim_token) + + assert server is not None with _lock: _server_connecting.discard(name) _server_connect_errors.pop(name, None) @@ -5932,7 +6001,11 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]: # Log a summary so ACP callers get visibility into what was registered. with _lock: - connected = [n for n in new_servers if n in _servers] + connected = [ + n + for n in new_servers + if n in _servers and n not in _server_connect_errors + ] new_tool_count = sum( len(getattr(_servers[n], "_registered_tool_names", [])) for n in connected @@ -6005,7 +6078,11 @@ def discover_mcp_tools() -> List[str]: return tool_names with _lock: - connected_server_names = [name for name in new_server_names if name in _servers] + connected_server_names = [ + name + for name in new_server_names + if name in _servers and name not in _server_connect_errors + ] new_tool_count = sum( len(getattr(_servers[name], "_registered_tool_names", [])) for name in connected_server_names