diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index eeeeef7e48c..a27baf5e79b 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -172,6 +172,56 @@ class TestMCPStatus: assert statuses["disabled"]["disabled"] is True +class TestLifecycleConfig: + def test_get_lifecycle_seconds_accepts_top_level_and_nested_values(self): + from tools.mcp_tool import _get_lifecycle_seconds + + assert ( + _get_lifecycle_seconds( + {"idle_timeout_seconds": "3.5"}, + "idle_timeout_seconds", + ) + == 3.5 + ) + assert _get_lifecycle_seconds( + {"lifecycle": {"max_lifetime_seconds": 42}}, + "max_lifetime_seconds", + ) == 42.0 + + def test_get_lifecycle_seconds_treats_zero_as_disabled(self): + from tools.mcp_tool import _get_lifecycle_seconds + + assert ( + _get_lifecycle_seconds( + {"idle_timeout_seconds": 0}, + "idle_timeout_seconds", + ) + is None + ) + + def test_get_lifecycle_seconds_ignores_invalid_values(self, caplog): + from tools.mcp_tool import _get_lifecycle_seconds + + assert ( + _get_lifecycle_seconds( + {"idle_timeout_seconds": "soon"}, + "idle_timeout_seconds", + ) + is None + ) + assert ( + _get_lifecycle_seconds( + {"idle_timeout_seconds": -1}, + "idle_timeout_seconds", + ) + is None + ) + + messages = [record.getMessage() for record in caplog.records] + assert any("must be a number of seconds" in msg for msg in messages) + assert any("must be positive" in msg for msg in messages) + + # --------------------------------------------------------------------------- # Schema conversion # --------------------------------------------------------------------------- @@ -568,6 +618,19 @@ class TestCheckFunction: finally: _servers.pop("test_server", None) + def test_recycled_stdio_server_remains_available_for_lazy_reconnect(self): + from tools.mcp_tool import _make_check_fn, _servers + + server = _make_mock_server("test_server", session=None) + server._config = {"command": "npx"} + server._recycled_reason = "idle_timeout_seconds" + _servers["test_server"] = server + try: + check = _make_check_fn("test_server") + assert check() is True + finally: + _servers.pop("test_server", None) + # --------------------------------------------------------------------------- # MCP loop runner @@ -742,6 +805,36 @@ class TestToolHandler: finally: _servers.pop("test_srv", None) + def test_recycled_stdio_server_reconnects_lazily_on_tool_call(self): + from tools.mcp_tool import _make_tool_handler, _servers + + mock_session = MagicMock() + mock_session.call_tool = AsyncMock( + return_value=_make_call_result("reconnected", is_error=False) + ) + server = _make_mock_server("test_srv", session=None) + server._config = {"command": "npx"} + server._recycled_reason = "idle_timeout_seconds" + _servers["test_srv"] = server + + def fake_lazy_reconnect(server_name, srv): + assert server_name == "test_srv" + assert srv is server + srv.session = mock_session + srv._recycled_reason = None + return True + + try: + handler = _make_tool_handler("test_srv", "greet", 120) + with patch("tools.mcp_tool._request_lazy_reconnect", side_effect=fake_lazy_reconnect) as reconnect, \ + self._patch_mcp_loop(): + result = json.loads(handler({"name": "world"})) + assert result["result"] == "reconnected" + reconnect.assert_called_once() + mock_session.call_tool.assert_called_once_with("greet", arguments={"name": "world"}) + finally: + _servers.pop("test_srv", None) + class TestRunOnMCPLoopInterrupts: def test_interrupt_cancels_waiting_mcp_call(self): @@ -1148,6 +1241,52 @@ class TestMCPServerTask: asyncio.run(_test()) + def test_wait_for_lifecycle_event_recycles_idle_stdio_server(self): + from tools.mcp_tool import MCPServerTask + + async def _test(): + server = MCPServerTask("srv") + server._config = {"command": "npx"} + server.session = MagicMock() + server._idle_timeout_seconds = 0.01 + server._last_tool_call_at = time.monotonic() - 1.0 + + reason = await server._wait_for_lifecycle_event() + + assert reason == "recycle" + assert server._recycled_reason == "idle_timeout_seconds" + assert server.session is None + + asyncio.run(_test()) + + def test_stdio_recycle_reason_uses_max_lifetime(self): + from tools.mcp_tool import MCPServerTask + + async def _test(): + server = MCPServerTask("srv") + server._config = {"command": "npx"} + server._max_lifetime_seconds = 0.01 + server._lifecycle_started_at = time.monotonic() - 1.0 + + assert server._stdio_recycle_reason() == "max_lifetime_seconds" + + asyncio.run(_test()) + + def test_stdio_recycle_deadline_pauses_while_rpc_active(self): + from tools.mcp_tool import MCPServerTask + + async def _test(): + server = MCPServerTask("srv") + server._config = {"command": "npx"} + server._idle_timeout_seconds = 0.01 + server._last_tool_call_at = time.monotonic() - 1.0 + + async with server._rpc_lock: + assert server._stdio_recycle_reason() is None + assert server._next_stdio_recycle_deadline() is None + + asyncio.run(_test()) + # --------------------------------------------------------------------------- # discover_mcp_tools toolset injection @@ -2000,6 +2139,48 @@ class TestReconnection: asyncio.run(_test()) + def test_failed_recycle_reconnect_no_longer_checks_available(self): + """A failed lazy reconnect should not leave recycled availability true.""" + from tools.mcp_tool import ( + MCPServerTask, + _MAX_INITIAL_CONNECT_RETRIES, + _make_check_fn, + _servers, + ) + + run_count = 0 + target_server = None + + original_run_stdio = MCPServerTask._run_stdio + + async def patched_run_stdio(self_srv, config): + nonlocal run_count, target_server + run_count += 1 + if target_server is not self_srv: + return await original_run_stdio(self_srv, config) + raise ConnectionError("recycle reconnect failed") + + async def _test(): + nonlocal target_server + server = MCPServerTask("test_srv") + target_server = server + server._config = {"command": "test"} + server._ready.clear() + server._recycled_reason = "idle_timeout_seconds" + _servers["test_srv"] = server + try: + with patch.object(MCPServerTask, "_run_stdio", patched_run_stdio), \ + patch("asyncio.sleep", new_callable=AsyncMock): + await server.run({"command": "test"}) + + assert run_count == _MAX_INITIAL_CONNECT_RETRIES + 1 + assert server._recycled_reason is None + assert _make_check_fn("test_srv")() is False + finally: + _servers.pop("test_srv", None) + + asyncio.run(_test()) + # --------------------------------------------------------------------------- # Configurable timeouts diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 5f5726d769b..114fc172fc1 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -23,6 +23,10 @@ Example config:: # 180). Set below the server's session TTL for # servers that GC idle sessions quickly (e.g. # Unreal Engine editor MCP, ~15s). Floored at 5s. + idle_timeout_seconds: 3600 # optional stdio recycle after idle + max_lifetime_seconds: 86400 # optional stdio recycle after age + # The recycle settings may also live under lifecycle: {...}. + # Use 0 to disable either recycle limit. github: command: "npx" args: ["-y", "@modelcontextprotocol/server-github"] @@ -330,6 +334,7 @@ _MAX_BACKOFF_SECONDS = 60 # server is unrevivable: its tools are out of the registry, so no tool call # can ever reach the circuit-breaker half-open probe or _signal_reconnect. _PARKED_RETRY_INTERVAL = 300 # seconds between parked self-probes +_RECYCLED_RECONNECT_TIMEOUT = 15.0 # Keepalive cadence for HTTP/SSE sessions. The MCP spec lets a server expire # idle sessions on any TTL it chooses (Streamable HTTP "Session Management"), @@ -1518,6 +1523,8 @@ class MCPServerTask: "_registered_tool_names", "_auth_type", "_refresh_lock", "_rpc_lock", "_pending_refresh_tasks", "_pending_call_context", + "_lifecycle_started_at", "_last_tool_call_at", + "_idle_timeout_seconds", "_max_lifetime_seconds", "_recycled_reason", "initialize_result", "_ping_unsupported", "_reconnect_retries", ) @@ -1562,6 +1569,12 @@ class MCPServerTask: # gateway-platform attribution and routes the approval prompt # to the right surface (Telegram, Slack, etc.). self._pending_call_context: Optional[contextvars.Context] = None + now = time.monotonic() + self._lifecycle_started_at: float = now + self._last_tool_call_at: float = now + self._idle_timeout_seconds: Optional[float] = None + self._max_lifetime_seconds: Optional[float] = None + self._recycled_reason: Optional[str] = None # Captures the ``InitializeResult`` returned by # ``await session.initialize()`` so downstream code can inspect the # server's real advertised capabilities (``.capabilities.resources``, @@ -1599,6 +1612,53 @@ class MCPServerTask: return True return getattr(caps, "tools", None) is not None + def _is_recycled_stdio(self) -> bool: + """Return True when a stdio server was intentionally recycled.""" + return not self._is_http() and self._recycled_reason is not None + + def mark_tool_call(self) -> None: + """Record that a user-visible MCP operation is starting.""" + self._last_tool_call_at = time.monotonic() + + def _mark_lifecycle_started(self) -> None: + now = time.monotonic() + self._lifecycle_started_at = now + self._last_tool_call_at = now + self._recycled_reason = None + + def _stdio_recycle_reason(self, now: Optional[float] = None) -> Optional[str]: + """Return the stdio recycle reason if idle/age limits have elapsed.""" + if self._is_http() or self._rpc_lock.locked(): + return None + now = time.monotonic() if now is None else now + if ( + self._max_lifetime_seconds is not None + and now - self._lifecycle_started_at >= self._max_lifetime_seconds + ): + return "max_lifetime_seconds" + if ( + self._idle_timeout_seconds is not None + and now - self._last_tool_call_at >= self._idle_timeout_seconds + ): + return "idle_timeout_seconds" + return None + + def _next_stdio_recycle_deadline(self) -> Optional[float]: + """Return the next monotonic recycle deadline for stdio, if any.""" + if self._is_http() or self._rpc_lock.locked(): + return None + deadlines = [] + if self._max_lifetime_seconds is not None: + deadlines.append(self._lifecycle_started_at + self._max_lifetime_seconds) + if self._idle_timeout_seconds is not None: + deadlines.append(self._last_tool_call_at + self._idle_timeout_seconds) + return min(deadlines) if deadlines else None + + def _mark_stdio_recycled(self, reason: str) -> None: + """Mark a stdio session dormant before its transport finishes closing.""" + self._recycled_reason = reason + self.session = None + # ----- Dynamic tool discovery (notifications/tools/list_changed) ----- async def _refresh_tools_task(self): @@ -1811,6 +1871,9 @@ class MCPServerTask: tokens, new session ID, etc.). The reconnect event is cleared before return so the next cycle starts with a fresh signal. + "recycle" if a stdio idle/max-lifetime limit elapsed. The + current transport is torn down and restarted lazily + on the next tool call. Shutdown takes precedence if both events are set simultaneously. @@ -1842,14 +1905,29 @@ class MCPServerTask: reconnect_task = asyncio.create_task(self._reconnect_event.wait()) try: while True: + recycle_reason = self._stdio_recycle_reason() + if recycle_reason is not None: + self._mark_stdio_recycled(recycle_reason) + return "recycle" + + timeout = keepalive_interval + recycle_deadline = self._next_stdio_recycle_deadline() + if recycle_deadline is not None: + timeout = max(0.0, min(timeout, recycle_deadline - time.monotonic())) + done, _pending = await asyncio.wait( {shutdown_task, reconnect_task}, - timeout=keepalive_interval, + timeout=timeout, return_when=asyncio.FIRST_COMPLETED, ) if done: break + recycle_reason = self._stdio_recycle_reason() + if recycle_reason is not None: + self._mark_stdio_recycled(recycle_reason) + return "recycle" + # Timeout — no lifecycle event fired. Probe the connection # to detect stale/expired sessions. Prefer ``ping`` (MCP base # protocol liveness): it works uniformly and stays a few bytes @@ -2075,6 +2153,7 @@ class MCPServerTask: session.initialize(), timeout=connect_timeout ) self.session = session + self._mark_lifecycle_started() await self._discover_tools() self._ready.set() # Session is live again: clear any breaker state from a @@ -2088,7 +2167,7 @@ class MCPServerTask: # stdio transport does not use OAuth, but we still honor # _reconnect_event (e.g. future manual /mcp refresh) for # consistency with _run_http. - await self._wait_for_lifecycle_event() + return await self._wait_for_lifecycle_event() finally: # Runs on clean exit, exceptions, AND asyncio cancellation. # If any of the spawned PIDs are still alive, the SDK's @@ -2383,7 +2462,7 @@ class MCPServerTask: "MCP server '%s': reconnect requested — " "tearing down SSE session", self.name, ) - return + return reason if _MCP_NEW_HTTP: # New API (mcp >= 1.24.0): build an explicit httpx.AsyncClient @@ -2440,6 +2519,7 @@ class MCPServerTask: "MCP server '%s': reconnect requested — " "tearing down HTTP session", self.name, ) + return reason else: # Deprecated API (mcp < 1.24.0): manages httpx client internally. _http_kwargs: dict = { @@ -2471,6 +2551,7 @@ class MCPServerTask: "MCP server '%s': reconnect requested — " "tearing down legacy HTTP session", self.name, ) + return reason async def _discover_tools(self): """Discover tools from the connected session. @@ -2531,6 +2612,8 @@ class MCPServerTask: self._config = config self.tool_timeout = config.get("timeout", _DEFAULT_TOOL_TIMEOUT) self._auth_type = (config.get("auth") or "").lower().strip() + self._idle_timeout_seconds = _get_lifecycle_seconds(config, "idle_timeout_seconds") + self._max_lifetime_seconds = _get_lifecycle_seconds(config, "max_lifetime_seconds") # Set up sampling handler if enabled and SDK types are available sampling_config = config.get("sampling", {}) @@ -2606,9 +2689,9 @@ class MCPServerTask: while True: try: if self._is_http(): - await self._run_http(config) + lifecycle_reason = await self._run_http(config) else: - await self._run_stdio(config) + lifecycle_reason = await self._run_stdio(config) # Transport returned cleanly. Two cases: # - _shutdown_event was set: exit the run loop entirely. # - _reconnect_event was set (auth recovery): loop back and @@ -2616,6 +2699,18 @@ class MCPServerTask: # touch the retry counters — this is not a failure. if self._shutdown_event.is_set(): break + if lifecycle_reason == "recycle": + logger.info( + "MCP server '%s': stdio session recycled after %s; " + "waiting for lazy reconnect", + self.name, self._recycled_reason, + ) + self.session = None + await self._wait_for_lazy_reconnect() + if self._shutdown_event.is_set(): + break + self._reconnect_event.clear() + continue logger.info( "MCP server '%s': reconnecting (OAuth recovery or " "manual refresh)", @@ -2651,6 +2746,13 @@ class MCPServerTask: raise except Exception as exc: self.session = None + if self._is_recycled_stdio(): + logger.warning( + "MCP server '%s': lazy reconnect after stdio recycle " + "failed, marking unavailable while retrying: %s", + self.name, exc, + ) + self._recycled_reason = None # If this is the first connection attempt, retry with backoff # before giving up. A transient DNS/network blip at startup @@ -2839,6 +2941,24 @@ class MCPServerTask: _forget_mcp_tool_server(tool_name) self._registered_tool_names = [] + async def _wait_for_lazy_reconnect(self) -> None: + """Wait while an intentionally recycled stdio server is dormant.""" + shutdown_task = asyncio.create_task(self._shutdown_event.wait()) + reconnect_task = asyncio.create_task(self._reconnect_event.wait()) + try: + await asyncio.wait( + {shutdown_task, reconnect_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + for task in (shutdown_task, reconnect_task): + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + # --------------------------------------------------------------------------- # Module-level state @@ -3684,6 +3804,51 @@ async def _connect_server(name: str, config: dict) -> MCPServerTask: # Handler / check-fn factories # --------------------------------------------------------------------------- +def _request_lazy_reconnect(server_name: str, server: MCPServerTask) -> bool: + """Wake a recycled stdio server and wait briefly for a fresh session.""" + if not server._is_recycled_stdio(): + return False + + with _lock: + loop = _mcp_loop + if loop is None or not loop.is_running(): + return False + + def _signal_reconnect() -> None: + server._ready.clear() + server._reconnect_event.set() + + loop.call_soon_threadsafe(_signal_reconnect) + + async def _await_ready() -> bool: + deadline = time.monotonic() + _RECYCLED_RECONNECT_TIMEOUT + while time.monotonic() < deadline: + if server.session is not None and server._ready.is_set(): + return True + await asyncio.sleep(0.05) + return False + + try: + return bool(_run_on_mcp_loop(_await_ready, timeout=_RECYCLED_RECONNECT_TIMEOUT)) + except Exception as exc: + logger.warning( + "MCP server '%s': lazy reconnect after stdio recycle failed: %s", + server_name, exc, + ) + return False + + +def _get_connected_server_for_call(server_name: str) -> Optional[MCPServerTask]: + """Return a connected server, lazily reconnecting recycled stdio state.""" + with _lock: + server = _servers.get(server_name) + if server is not None and server.session is None and server._is_recycled_stdio(): + _request_lazy_reconnect(server_name, server) + with _lock: + server = _servers.get(server_name) + return server + + def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): """Return a sync handler that calls an MCP tool via the background loop. @@ -3718,8 +3883,7 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): }, ensure_ascii=False) # Cooldown elapsed → fall through as a half-open probe. - with _lock: - server = _servers.get(server_name) + server = _get_connected_server_for_call(server_name) if not server: _bump_server_error(server_name) return json.dumps({ @@ -3761,6 +3925,7 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): }, ensure_ascii=False) async def _call(): + server.mark_tool_call() async with server._rpc_lock: # Snapshot the agent's context so an elicitation callback # triggered during this call (fired on the MCP recv loop @@ -3874,14 +4039,14 @@ def _make_list_resources_handler(server_name: str, tool_timeout: float): """Return a sync handler that lists resources from an MCP server.""" def _handler(args: dict, **kwargs) -> str: - with _lock: - server = _servers.get(server_name) + server = _get_connected_server_for_call(server_name) if not server or not server.session: return json.dumps({ "error": f"MCP server '{server_name}' is not connected" }, ensure_ascii=False) async def _call(): + server.mark_tool_call() async with server._rpc_lock: result = await server.session.list_resources() resources = [] @@ -3934,8 +4099,7 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float): def _handler(args: dict, **kwargs) -> str: from tools.registry import tool_error - with _lock: - server = _servers.get(server_name) + server = _get_connected_server_for_call(server_name) if not server or not server.session: return json.dumps({ "error": f"MCP server '{server_name}' is not connected" @@ -3946,6 +4110,7 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float): return tool_error("Missing required parameter 'uri'") async def _call(): + server.mark_tool_call() async with server._rpc_lock: result = await server.session.read_resource(uri) # read_resource returns ReadResourceResult with .contents list @@ -3992,14 +4157,14 @@ def _make_list_prompts_handler(server_name: str, tool_timeout: float): """Return a sync handler that lists prompts from an MCP server.""" def _handler(args: dict, **kwargs) -> str: - with _lock: - server = _servers.get(server_name) + server = _get_connected_server_for_call(server_name) if not server or not server.session: return json.dumps({ "error": f"MCP server '{server_name}' is not connected" }, ensure_ascii=False) async def _call(): + server.mark_tool_call() async with server._rpc_lock: result = await server.session.list_prompts() prompts = [] @@ -4057,8 +4222,7 @@ def _make_get_prompt_handler(server_name: str, tool_timeout: float): def _handler(args: dict, **kwargs) -> str: from tools.registry import tool_error - with _lock: - server = _servers.get(server_name) + server = _get_connected_server_for_call(server_name) if not server or not server.session: return json.dumps({ "error": f"MCP server '{server_name}' is not connected" @@ -4070,6 +4234,7 @@ def _make_get_prompt_handler(server_name: str, tool_timeout: float): arguments = args.get("arguments", {}) async def _call(): + server.mark_tool_call() async with server._rpc_lock: result = await server.session.get_prompt(name, arguments=arguments) # GetPromptResult has .messages list @@ -4128,7 +4293,10 @@ def _make_check_fn(server_name: str): def _check() -> bool: with _lock: server = _servers.get(server_name) - return server is not None and server.session is not None + return ( + server is not None + and (server.session is not None or server._is_recycled_stdio()) + ) return _check @@ -4428,6 +4596,27 @@ def _parse_boolish(value: Any, default: bool = True) -> bool: return default +def _get_lifecycle_seconds(config: dict, key: str) -> Optional[float]: + """Return an optional positive lifecycle timeout from top-level/nested config.""" + raw = config.get(key) + lifecycle = config.get("lifecycle") + if raw is None and isinstance(lifecycle, dict): + raw = lifecycle.get(key) + if raw is None: + return None + try: + seconds = float(raw) + except (TypeError, ValueError): + logger.warning("MCP config %s must be a number of seconds; ignoring %r", key, raw) + return None + if seconds == 0: + return None + if seconds < 0: + logger.warning("MCP config %s must be positive; ignoring %r", key, raw) + return None + return seconds + + _UTILITY_CAPABILITY_METHODS = { "list_resources": "list_resources", "read_resource": "read_resource",