diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 5c7f5d61ff..33b53b9c8b 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -208,6 +208,7 @@ from tools.mcp_tool_agent import ( # noqa: F401 restore_agent_tool_prefix, ) from tools.mcp_tool_transport import MCPServerTransportMixin +from tools.mcp_tool_server_run import MCPServerRunMixin from tools.mcp_tool_health import MCPServerHealthMixin @@ -551,7 +552,7 @@ def _mcp_types(): # Server task -- each MCP server lives in one long-lived asyncio Task # --------------------------------------------------------------------------- -class MCPServerTask(MCPServerTransportMixin, MCPServerHealthMixin): +class MCPServerTask(MCPServerRunMixin, MCPServerTransportMixin, MCPServerHealthMixin): """One MCP server connection living in one long-lived asyncio Task. Connect, discover, serve and disconnect all run in that Task so the @@ -658,554 +659,15 @@ class MCPServerTask(MCPServerTransportMixin, MCPServerHealthMixin): # POST/GET; anything else on a 2xx means the URL is not an MCP endpoint. _MCP_CONTENT_TYPES = ("application/json", "text/event-stream") - @staticmethod - async def _cancel_waiters(*tasks: asyncio.Task) -> None: - for t in tasks: - if not t.done(): - t.cancel() - try: - await t - except (asyncio.CancelledError, Exception): - pass - async def _wait_for_lifecycle_event(self) -> str: - """Serve the connection until a lifecycle event; return its kind. - ``"shutdown"`` exits the run loop; ``"reconnect"`` tears the session - down and re-enters the transport (event cleared before return); - ``"recycle"`` means a stdio idle/lifetime limit elapsed and the - transport restarts lazily on the next call. Shutdown wins a tie. - Between events a keepalive (``ping``, list_tools fallback) runs every - ``keepalive_interval`` — which must stay below the server's session - TTL — and a failure triggers a reconnect. ``ping`` is a few bytes - regardless of tool count; list_changed notifications still arrive - out-of-band. - """ - keepalive_interval = max( - _MIN_KEEPALIVE_INTERVAL, - float(self._config.get("keepalive_interval", _DEFAULT_KEEPALIVE_INTERVAL)), - ) - shutdown_task = asyncio.create_task(self._shutdown_event.wait()) - 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=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: probe for a stale session — but NEVER while an RPC - # is in flight (a concurrent ping can wedge the single stdio - # stream, and a busy server is provably alive anyway). - if self.session: - if self._rpc_lock.locked() or any( - not t.done() for t in self._inflight_tasks - ): - continue - try: - async def _probe_under_lock(): - async with self._rpc_lock: - await self._keepalive_probe() - await _probe_under_lock() - except Exception as exc: - root = _unwrap_exception_group(exc) - logger.warning( - "MCP server '%s' keepalive failed, triggering " - "reconnect (state: connected → degraded): %s: %s", - self.name, type(root).__name__, root, - ) - self.mark_suspect( - f"keepalive failed: {type(root).__name__}: {root}" - ) - self._reconnect_event.set() - break - # Survived a full keepalive interval: real proof of health. - self._mark_session_proven() - finally: - await self._cancel_waiters(shutdown_task, reconnect_task) - - if self._shutdown_event.is_set(): - self._fail_inflight_calls("shutdown") - return "shutdown" - # Deliberate teardown: fail in-flight RPCs NOW rather than letting - # them ride the dying transport to the full tool timeout. - self._fail_inflight_calls("reconnect") - self._reconnect_event.clear() - return "reconnect" - - async def _wait_for_reconnect_or_shutdown( - self, timeout: Optional[float] = None - ) -> str: - """Wait, while parked, for a reconnect request or shutdown. - - Returns ``"shutdown"`` or ``"reconnect"`` (explicit request or, with - ``timeout``, the periodic self-probe); the reconnect event is cleared - first. Shutdown wins a tie. - """ - shutdown_task = asyncio.ensure_future(self._shutdown_event.wait()) - reconnect_task = asyncio.ensure_future(self._reconnect_event.wait()) - try: - await asyncio.wait( - {shutdown_task, reconnect_task}, - return_when=asyncio.FIRST_COMPLETED, - timeout=timeout, - ) - finally: - await self._cancel_waiters(shutdown_task, reconnect_task) - if self._shutdown_event.is_set(): - return "shutdown" - self._reconnect_event.clear() - return "reconnect" - - async def _park(self, revival_reason: str) -> bool: - """Drop this server's tools and wait for a reconnect request. - - The run task must NOT exit: it is the only listener on - ``_reconnect_event``, so returning leaves the server unrevivable for - the life of the process. Parking deregisters the tools, so no call - can reach the breaker probe or ``_signal_reconnect``; the wait is - therefore TIMED (one self-probe per ``_PARKED_RETRY_INTERVAL``), and - an explicit ``_reconnect_event.set()`` wakes it immediately. Returns - True when shutdown was requested instead. - """ - self._was_parked = True - self._deregister_tools() - self._reconnect_event.clear() - parked = await self._wait_for_reconnect_or_shutdown( - timeout=_PARKED_RETRY_INTERVAL - ) - if parked == "shutdown": - return True - logger.debug( - "MCP server '%s': attempting revival %s (self-probe or explicit " - "reconnect request); rebuilding transport.", - self.name, revival_reason, - ) - return False - - async def _prepare_run(self, config: dict) -> bool: - """Bind config, build sampling/elicitation handlers, validate HTTP. - - Returns False when the server must not start: a bad remote URL or a - non-MCP endpoint (both fail fast, non-retryably, with ``_error`` set - and ``_ready`` fired) instead of burning the reconnect ladder inside - the SDK's httpx layer on every retry. - """ - self._config = config - self.tool_timeout = _resolve_tool_timeout(config) - 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") - - # The _MCP_*_TYPES flags are False until the lazy SDK import runs. - _ensure_mcp_sdk() - - sampling_config = config.get("sampling", {}) - if sampling_config.get("enabled", True) and _MCP_SAMPLING_TYPES: - self._sampling = SamplingHandler(self.name, sampling_config) - else: - self._sampling = None - - # elicitation/create lets a server ask for structured input mid-call; - # the handler routes it through Hermes' approval system. - elicitation_config = config.get("elicitation", {}) - if elicitation_config.get("enabled", True) and _MCP_ELICITATION_TYPES: - self._elicitation = ElicitationHandler(self.name, elicitation_config, owner=self) - else: - self._elicitation = None - - if "url" in config and "command" in config: - logger.warning( - "MCP server '%s' has both 'url' and 'command' in config. " - "Using HTTP transport ('url'). Remove 'command' to silence " - "this warning.", - self.name, - ) - - if not self._is_http(): - return True - try: - _validate_remote_mcp_url(self.name, config.get("url")) - except InvalidMcpUrlError as exc: - logger.warning("%s", exc) - self._error = exc - self._ready.set() - return False - - # Content-type preflight (Streamable HTTP only; SSE legitimately serves - # text/event-stream): a URL at a web-app root returns HTML and would - # make the SDK hang for the full connect_timeout. Skipped once _ready - # was ever set (endpoint already validated) and for OAuth servers, - # where a token-less probe sees HTML/401 and would block the flow. - if config.get("transport") != "sse" and not config.get("skip_preflight") and not self._ready.is_set() and self._auth_type != "oauth": - try: - _probe_headers = dict(config.get("headers") or {}) - await self._preflight_content_type( - config["url"], - headers=_probe_headers, - ssl_verify=config.get("ssl_verify", True), - client_cert=_resolve_client_cert(self.name, config), - ) - except NonMcpEndpointError as exc: - logger.warning("%s", exc) - self._error = exc - self._ready.set() - return False - return True - - async def run(self, config: dict): - """Long-lived coroutine: connect, discover, serve, reconnect. - - State machine: connecting -> connected -> (degraded -> parked -> - revived)*. Unproven drops and transport errors charge a rapid-drop - budget with jittered exponential backoff; exhausting it (or a - permanent error) parks the server via :meth:`_park` rather than - exiting, so it stays revivable. - """ - if not await self._prepare_run(config): - return - - self._reconnect_retries = 0 - initial_retries = 0 - backoff = 1.0 - - while True: - try: - if self._is_http(): - lifecycle_reason = await self._run_http(config) - else: - lifecycle_reason = await self._run_stdio(config) - # Clean transport return: shutdown, stdio recycle, or a - # requested rebuild (auth recovery / manual refresh / keepalive - # failure). A rebuild is not a failure for the retry counters. - 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 - # Per-cycle chatter stays DEBUG; WARNINGs mark state transitions. - logger.debug( - "MCP server '%s': reconnecting (OAuth recovery or " - "manual refresh)", - self.name, - ) - # A clean return is NOT proof of health (a flapping transport - # handshakes fine and drops moments later). Only a PROVEN - # session clears the budget; a teardown race is recovery, not - # a failure, and must never reach the park on its own. - if self._teardown_race and not self._session_proven: - logger.info( - "MCP server '%s': reconnect after teardown race " - "(in-flight calls were failed); not charging the " - "rapid-drop budget", - self.name, - ) - self._teardown_race = False - backoff = 1.0 - elif self._session_proven: - self._reconnect_retries = 0 - backoff = 1.0 - else: - self._reconnect_retries += 1 - if self._reconnect_retries > _MAX_RECONNECT_RETRIES: - logger.warning( - "MCP server '%s': %d consecutive reconnects " - "without a healthy session (rapid-drop budget " - "exhausted), parking; will self-probe every %ds " - "until it recovers (state: degraded → parked)", - self.name, _MAX_RECONNECT_RETRIES, - _PARKED_RETRY_INTERVAL, - ) - if await self._park("from parked state"): - break - # Budget of one probe per wake, so a still-dead server - # parks again instead of burning 5 rapid retries. - self._reconnect_retries = _MAX_RECONNECT_RETRIES - backoff = 1.0 - # Clear readiness too: a stale _ready lets handler-side - # recovery mistake the old session for a fresh one. - self._ready.clear() - self.session = None - continue - except asyncio.CancelledError: - # Not a connection failure: re-raise so cancellation reaches - # asyncio and shutdown()'s ``await self._task`` completes. - self.session = None - raise - except Exception as exc: - self.session = None - # Unwrap anyio TaskGroup wrappers: the group's str() is useless - # and hides the root cause from the classification below. - root = _unwrap_exception_group(exc) - failure_class = _classify_mcp_failure(root) - if self._is_recycled_stdio(): - logger.warning( - "MCP server '%s': lazy reconnect after stdio recycle " - "failed, marking unavailable while retrying: %s: %s", - self.name, type(root).__name__, root, - ) - self._recycled_reason = None - - # Initial-connect ladder: a transient blip at startup must not - # kill the server. Gated on _ever_connected (never cleared), - # not _ready (cleared every reconnect cycle). - if not self._ever_connected: - if failure_class == "permanent": - # Deterministic failure (bad command, non-MCP URL, - # 401/403): park at once instead of burning the ladder. - # Auth failures park rather than return so the task - # stays alive to pick up fresh tokens later. - if _is_auth_error(root): - logger.warning( - "MCP server '%s' failed initial authentication, " - "parking until credentials change; re-authenticate " - "with `hermes mcp login %s` " - "(state: connecting → parked): %s: %s", - self.name, self.name, - type(root).__name__, root, - ) - else: - logger.warning( - "MCP server '%s' failed initial connection with a " - "permanent error, parking without retries " - "(state: connecting → parked): %s: %s", - self.name, type(root).__name__, root, - ) - self._error = exc - self._ready.set() - if await self._park("after permanent initial failure"): - return - initial_retries = 0 - self._reconnect_retries = 0 - backoff = 1.0 - self._error = None - self._ready.clear() - continue - - initial_retries += 1 - if initial_retries > _MAX_INITIAL_CONNECT_RETRIES: - logger.warning( - "MCP server '%s' failed initial connection after " - "%d attempts, parking until a reconnect is " - "requested (state: connecting → parked): %s: %s", - self.name, _MAX_INITIAL_CONNECT_RETRIES, - type(root).__name__, root, - ) - self._error = exc - self._ready.set() - if await self._park("after initial connection failures"): - return - initial_retries = 0 - self._reconnect_retries = 0 - backoff = 1.0 - self._error = None - self._ready.clear() - continue - - logger.debug( - "MCP server '%s' initial connection failed " - "(attempt %d/%d), retrying in %.0fs: %s: %s", - self.name, initial_retries, - _MAX_INITIAL_CONNECT_RETRIES, backoff, - type(root).__name__, root, - ) - await asyncio.sleep(_jittered(backoff)) - backoff = min(backoff * 2, _MAX_BACKOFF_SECONDS) - - # Check if shutdown was requested during the sleep - if self._shutdown_event.is_set(): - self._error = exc - self._ready.set() - return - continue - - # If shutdown was requested, don't reconnect - if self._shutdown_event.is_set(): - logger.debug( - "MCP server '%s' disconnected during shutdown: %s: %s", - self.name, type(root).__name__, root, - ) - return - - if failure_class == "permanent": - # An auth failure on a PROVEN session is often a corrupt - # OAuth lock from a raced teardown, not revoked - # credentials: grant ONE suspect+reconnect cycle first. - if ( - _is_auth_error(root) - and self._session_proven - and not self._permanent_grace_used - ): - self._permanent_grace_used = True - self.mark_suspect( - f"auth error on proven session: {root}" - ) - logger.warning( - "MCP server '%s': auth error on a previously " - "healthy session — marking suspect and forcing " - "one reconnect instead of parking (state: " - "connected → suspect): %s: %s", - self.name, type(root).__name__, root, - ) - self._reconnect_retries = 0 - backoff = 1.0 - await asyncio.sleep(_jittered(1.0)) - if self._shutdown_event.is_set(): - return - continue - # Deterministic failure on a working server: park now. - logger.warning( - "MCP server '%s' hit a permanent error, parking " - "without retries; will self-probe every %ds " - "(state: connected → parked): %s: %s", - self.name, _PARKED_RETRY_INTERVAL, - type(root).__name__, root, - ) - if await self._park("from parked state (permanent error)"): - return - self._reconnect_retries = _MAX_RECONNECT_RETRIES - backoff = 1.0 - continue - - self._reconnect_retries += 1 - if self._reconnect_retries > _MAX_RECONNECT_RETRIES: - logger.warning( - "MCP server '%s' failed after %d reconnection attempts, " - "parking; will self-probe every %ds until it recovers " - "(state: degraded → parked): %s: %s", - self.name, _MAX_RECONNECT_RETRIES, - _PARKED_RETRY_INTERVAL, - type(root).__name__, root, - ) - if await self._park("from parked state"): - return - self._reconnect_retries = _MAX_RECONNECT_RETRIES - backoff = 1.0 - continue - - logger.debug( - "MCP server '%s' connection lost (attempt %d/%d), " - "reconnecting in %.0fs: %s: %s", - self.name, self._reconnect_retries, _MAX_RECONNECT_RETRIES, - backoff, type(root).__name__, root, - ) - await asyncio.sleep(_jittered(backoff)) - backoff = min(backoff * 2, _MAX_BACKOFF_SECONDS) - - # Check again after sleeping - if self._shutdown_event.is_set(): - return - finally: - self.session = None - # Stale PIDs must never fast-fail the NEXT transport's calls. - self._stdio_child_pids = set() - - async def start(self, config: dict): - """Create the background Task and wait until ready (or failed).""" - self._task = asyncio.ensure_future(self.run(config)) - try: - await self._ready.wait() - except asyncio.CancelledError: - # The caller's connect timeout (discover_mcp_tools wraps start() - # in asyncio.wait_for) cancels *this* coroutine, but the - # ensure_future'd run() task is independent and would otherwise - # keep running detached — parked on a hung transport with no - # owner to reap it (#59349). Propagate the cancellation so the - # transport context managers unwind and their finally blocks - # release the child process / FDs. - if self._task and not self._task.done(): - self._task.cancel() - raise - if self._error: - raise self._error - - async def shutdown(self): - """Signal the Task to exit and wait for clean resource teardown.""" - self._shutdown_event.set() - # Defensive: if _wait_for_lifecycle_event is blocking, we need ANY - # event to unblock it. _shutdown_event alone is sufficient (the - # helper checks shutdown first), but setting reconnect too ensures - # there's no race where the helper misses the shutdown flag after - # returning "reconnect". - self._reconnect_event.set() - if self._task and not self._task.done(): - try: - await asyncio.wait_for(self._task, timeout=10) - except asyncio.TimeoutError: - logger.warning( - "MCP server '%s' shutdown timed out, cancelling task", - self.name, - ) - self._task.cancel() - try: - await self._task - except asyncio.CancelledError: - pass - if self._pending_refresh_tasks: - for task in list(self._pending_refresh_tasks): - task.cancel() - await asyncio.gather(*self._pending_refresh_tasks, return_exceptions=True) - self._pending_refresh_tasks.clear() - self._deregister_tools() - self.session = None - - def _deregister_tools(self) -> None: - """Drop this server's tools from the global registry (idempotent). - - Pulls the server's tool schemas out of the registry so the agent - stops advertising them to the model. Called on shutdown AND when the - reconnect budget is exhausted, so a dead server never leaves phantom - tool definitions bloating the prompt cache and producing "not - connected" errors on every turn. - """ - from tools.registry import registry - - for tool_name in list(getattr(self, "_registered_tool_names", [])): - registry.deregister(tool_name, scope=_server_registry_scope(self.name)) - _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: - await self._cancel_waiters(shutdown_task, reconnect_task) # --------------------------------------------------------------------------- diff --git a/tools/mcp_tool_server_run.py b/tools/mcp_tool_server_run.py new file mode 100644 index 0000000000..77bbf099cf --- /dev/null +++ b/tools/mcp_tool_server_run.py @@ -0,0 +1,583 @@ +"""Lifecycle of :class:`tools.mcp_tool.MCPServerTask`: the long-lived ``run`` state machine +(connect -> serve -> reconnect/park/recycle), keepalive-driven lifecycle waits, start/shutdown +and tool deregistration. Split from tools/mcp_tool.py; origin state and patchable helpers are +read through ``_core`` so ``mock.patch("tools.mcp_tool.X")`` keeps working.""" + +import asyncio +import logging +import time +from dataclasses import dataclass +from typing import Optional +from tools.mcp_tool_common import _core + +logger = logging.getLogger("tools.mcp_tool") + + +@dataclass +class _RetryBudget: + """Per-run() retry counters shared by the branch helpers (``_reconnect_retries`` + lives on the task because handlers and tests read it).""" + + initial_retries: int = 0 + backoff: float = 1.0 + + +class MCPServerRunMixin: + """Methods of :class:`tools.mcp_tool.MCPServerTask` (mixed in; relies on its attributes).""" + + @staticmethod + async def _cancel_waiters(*tasks: asyncio.Task) -> None: + for t in tasks: + if not t.done(): + t.cancel() + try: + await t + except (asyncio.CancelledError, Exception): + pass + + async def _wait_for_lifecycle_event(self) -> str: + """Serve the connection until a lifecycle event; return its kind. + + ``"shutdown"`` exits the run loop; ``"reconnect"`` tears the session + down and re-enters the transport (event cleared before return); + ``"recycle"`` means a stdio idle/lifetime limit elapsed and the + transport restarts lazily on the next call. Shutdown wins a tie. + + Between events a keepalive (``ping``, list_tools fallback) runs every + ``keepalive_interval`` — which must stay below the server's session + TTL — and a failure triggers a reconnect. ``ping`` is a few bytes + regardless of tool count; list_changed notifications still arrive + out-of-band. + """ + keepalive_interval = max( + _core._MIN_KEEPALIVE_INTERVAL, + float(self._config.get("keepalive_interval", _core._DEFAULT_KEEPALIVE_INTERVAL)), + ) + + shutdown_task = asyncio.create_task(self._shutdown_event.wait()) + 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=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: probe for a stale session — but NEVER while an RPC + # is in flight (a concurrent ping can wedge the single stdio + # stream, and a busy server is provably alive anyway). + if self.session: + if self._rpc_lock.locked() or any( + not t.done() for t in self._inflight_tasks + ): + continue + try: + async def _probe_under_lock(): + async with self._rpc_lock: + await self._keepalive_probe() + + await _probe_under_lock() + except Exception as exc: + root = _core._unwrap_exception_group(exc) + logger.warning( + "MCP server '%s' keepalive failed, triggering " + "reconnect (state: connected → degraded): %s: %s", + self.name, type(root).__name__, root, + ) + self.mark_suspect( + f"keepalive failed: {type(root).__name__}: {root}" + ) + self._reconnect_event.set() + break + # Survived a full keepalive interval: real proof of health. + self._mark_session_proven() + finally: + await self._cancel_waiters(shutdown_task, reconnect_task) + + if self._shutdown_event.is_set(): + self._fail_inflight_calls("shutdown") + return "shutdown" + # Deliberate teardown: fail in-flight RPCs NOW rather than letting + # them ride the dying transport to the full tool timeout. + self._fail_inflight_calls("reconnect") + self._reconnect_event.clear() + return "reconnect" + + async def _wait_for_reconnect_or_shutdown( + self, timeout: Optional[float] = None + ) -> str: + """Wait, while parked, for a reconnect request or shutdown. + + Returns ``"shutdown"`` or ``"reconnect"`` (explicit request or, with + ``timeout``, the periodic self-probe); the reconnect event is cleared + first. Shutdown wins a tie. + """ + shutdown_task = asyncio.ensure_future(self._shutdown_event.wait()) + reconnect_task = asyncio.ensure_future(self._reconnect_event.wait()) + try: + await asyncio.wait( + {shutdown_task, reconnect_task}, + return_when=asyncio.FIRST_COMPLETED, + timeout=timeout, + ) + finally: + await self._cancel_waiters(shutdown_task, reconnect_task) + if self._shutdown_event.is_set(): + return "shutdown" + self._reconnect_event.clear() + return "reconnect" + + async def _park(self, revival_reason: str) -> bool: + """Drop this server's tools and wait for a reconnect request. + + The run task must NOT exit: it is the only listener on + ``_reconnect_event``, so returning leaves the server unrevivable for + the life of the process. Parking deregisters the tools, so no call + can reach the breaker probe or ``_signal_reconnect``; the wait is + therefore TIMED (one self-probe per ``_PARKED_RETRY_INTERVAL``), and + an explicit ``_reconnect_event.set()`` wakes it immediately. Returns + True when shutdown was requested instead. + """ + self._was_parked = True + self._deregister_tools() + self._reconnect_event.clear() + parked = await self._wait_for_reconnect_or_shutdown( + timeout=_core._PARKED_RETRY_INTERVAL + ) + if parked == "shutdown": + return True + logger.debug( + "MCP server '%s': attempting revival %s (self-probe or explicit " + "reconnect request); rebuilding transport.", + self.name, revival_reason, + ) + return False + + async def _prepare_run(self, config: dict) -> bool: + """Bind config, build sampling/elicitation handlers, validate HTTP. + + Returns False when the server must not start: a bad remote URL or a + non-MCP endpoint (both fail fast, non-retryably, with ``_error`` set + and ``_ready`` fired) instead of burning the reconnect ladder inside + the SDK's httpx layer on every retry. + """ + self._config = config + self.tool_timeout = _core._resolve_tool_timeout(config) + self._auth_type = (config.get("auth") or "").lower().strip() + self._idle_timeout_seconds = _core._get_lifecycle_seconds(config, "idle_timeout_seconds") + self._max_lifetime_seconds = _core._get_lifecycle_seconds(config, "max_lifetime_seconds") + + # The _MCP_*_TYPES flags are False until the lazy SDK import runs. + _core._ensure_mcp_sdk() + + sampling_config = config.get("sampling", {}) + if sampling_config.get("enabled", True) and _core._MCP_SAMPLING_TYPES: + self._sampling = _core.SamplingHandler(self.name, sampling_config) + else: + self._sampling = None + + # elicitation/create lets a server ask for structured input mid-call; + # the handler routes it through Hermes' approval system. + elicitation_config = config.get("elicitation", {}) + if elicitation_config.get("enabled", True) and _core._MCP_ELICITATION_TYPES: + self._elicitation = _core.ElicitationHandler(self.name, elicitation_config, owner=self) + else: + self._elicitation = None + + if "url" in config and "command" in config: + logger.warning( + "MCP server '%s' has both 'url' and 'command' in config. " + "Using HTTP transport ('url'). Remove 'command' to silence " + "this warning.", + self.name, + ) + + if not self._is_http(): + return True + try: + _core._validate_remote_mcp_url(self.name, config.get("url")) + except _core.InvalidMcpUrlError as exc: + logger.warning("%s", exc) + self._error = exc + self._ready.set() + return False + + # Content-type preflight (Streamable HTTP only; SSE legitimately serves + # text/event-stream): a URL at a web-app root returns HTML and would + # make the SDK hang for the full connect_timeout. Skipped once _ready + # was ever set (endpoint already validated) and for OAuth servers, + # where a token-less probe sees HTML/401 and would block the flow. + if config.get("transport") != "sse" and not config.get("skip_preflight") and not self._ready.is_set() and self._auth_type != "oauth": + try: + _probe_headers = dict(config.get("headers") or {}) + await self._preflight_content_type( + config["url"], + headers=_probe_headers, + ssl_verify=config.get("ssl_verify", True), + client_cert=_core._resolve_client_cert(self.name, config), + ) + except _core.NonMcpEndpointError as exc: + logger.warning("%s", exc) + self._error = exc + self._ready.set() + return False + return True + + async def run(self, config: dict): + """Long-lived coroutine: connect, discover, serve, reconnect. + + State machine: connecting -> connected -> (degraded -> parked -> + revived)*. Unproven drops and transport errors charge a rapid-drop + budget with jittered exponential backoff; exhausting it (or a + permanent error) parks the server via :meth:`_park` rather than + exiting, so it stays revivable. The branch helpers return True to + keep looping and False to exit the loop. + """ + if not await self._prepare_run(config): + return + + self._reconnect_retries = 0 + budget = _RetryBudget() + + while True: + try: + if self._is_http(): + lifecycle_reason = await self._run_http(config) + else: + lifecycle_reason = await self._run_stdio(config) + if not await self._on_clean_return(lifecycle_reason, budget): + break + except asyncio.CancelledError: + # Not a connection failure: re-raise so cancellation reaches + # asyncio and shutdown()'s ``await self._task`` completes. + self.session = None + raise + except Exception as exc: + self.session = None + if not await self._on_transport_error(exc, budget): + break + finally: + self.session = None + # Stale PIDs must never fast-fail the NEXT transport's calls. + self._stdio_child_pids = set() + + async def _on_clean_return(self, lifecycle_reason: str, budget: "_RetryBudget") -> bool: + """Transport returned cleanly: shutdown, stdio recycle, or a requested + rebuild (auth recovery / manual refresh / keepalive failure). A rebuild + is not a failure for the retry counters.""" + if self._shutdown_event.is_set(): + return False + 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(): + return False + self._reconnect_event.clear() + return True + # Per-cycle chatter stays DEBUG; WARNINGs mark state transitions. + logger.debug( + "MCP server '%s': reconnecting (OAuth recovery or " + "manual refresh)", + self.name, + ) + # A clean return is NOT proof of health (a flapping transport + # handshakes fine and drops moments later). Only a PROVEN + # session clears the budget; a teardown race is recovery, not + # a failure, and must never reach the park on its own. + if self._teardown_race and not self._session_proven: + logger.info( + "MCP server '%s': reconnect after teardown race " + "(in-flight calls were failed); not charging the " + "rapid-drop budget", + self.name, + ) + self._teardown_race = False + budget.backoff = 1.0 + elif self._session_proven: + self._reconnect_retries = 0 + budget.backoff = 1.0 + else: + self._reconnect_retries += 1 + if self._reconnect_retries > _core._MAX_RECONNECT_RETRIES: + logger.warning( + "MCP server '%s': %d consecutive reconnects " + "without a healthy session (rapid-drop budget " + "exhausted), parking; will self-probe every %ds " + "until it recovers (state: degraded → parked)", + self.name, _core._MAX_RECONNECT_RETRIES, + _core._PARKED_RETRY_INTERVAL, + ) + if not await self._park_and_rearm("from parked state", budget): + return False + # Clear readiness too: a stale _ready lets handler-side + # recovery mistake the old session for a fresh one. + self._ready.clear() + self.session = None + return True + + async def _park_and_rearm(self, revival_reason: str, budget: "_RetryBudget") -> bool: + """Park; on revival leave a budget of ONE probe per wake so a still-dead + server parks again instead of burning 5 rapid retries. False on shutdown.""" + if await self._park(revival_reason): + return False + self._reconnect_retries = _core._MAX_RECONNECT_RETRIES + budget.backoff = 1.0 + return True + + async def _park_initial_failure(self, exc: Exception, revival_reason: str, + budget: "_RetryBudget") -> bool: + """Publish ``exc`` to the waiting ``start()``, park, and on revival reset + every counter so the ladder starts fresh. False on shutdown.""" + self._error = exc + self._ready.set() + if await self._park(revival_reason): + return False + budget.initial_retries = 0 + self._reconnect_retries = 0 + budget.backoff = 1.0 + self._error = None + self._ready.clear() + return True + + async def _backoff_sleep(self, budget: "_RetryBudget") -> None: + await asyncio.sleep(_core._jittered(budget.backoff)) + budget.backoff = min(budget.backoff * 2, _core._MAX_BACKOFF_SECONDS) + + async def _on_transport_error(self, exc: Exception, budget: "_RetryBudget") -> bool: + """Transport raised: classify, then run the initial-connect or the + reconnect ladder. Returns False when the run loop must exit.""" + # Unwrap anyio TaskGroup wrappers: the group's str() is useless + # and hides the root cause from the classification below. + root = _core._unwrap_exception_group(exc) + failure_class = _core._classify_mcp_failure(root) + if self._is_recycled_stdio(): + logger.warning( + "MCP server '%s': lazy reconnect after stdio recycle " + "failed, marking unavailable while retrying: %s: %s", + self.name, type(root).__name__, root, + ) + self._recycled_reason = None + + # Initial-connect ladder: a transient blip at startup must not + # kill the server. Gated on _ever_connected (never cleared), + # not _ready (cleared every reconnect cycle). + if not self._ever_connected: + return await self._on_initial_connect_error(exc, root, failure_class, budget) + + # If shutdown was requested, don't reconnect + if self._shutdown_event.is_set(): + logger.debug( + "MCP server '%s' disconnected during shutdown: %s: %s", + self.name, type(root).__name__, root, + ) + return False + + if failure_class == "permanent": + return await self._on_permanent_error(root, budget) + + self._reconnect_retries += 1 + if self._reconnect_retries > _core._MAX_RECONNECT_RETRIES: + logger.warning( + "MCP server '%s' failed after %d reconnection attempts, " + "parking; will self-probe every %ds until it recovers " + "(state: degraded → parked): %s: %s", + self.name, _core._MAX_RECONNECT_RETRIES, + _core._PARKED_RETRY_INTERVAL, + type(root).__name__, root, + ) + return await self._park_and_rearm("from parked state", budget) + + logger.debug( + "MCP server '%s' connection lost (attempt %d/%d), " + "reconnecting in %.0fs: %s: %s", + self.name, self._reconnect_retries, _core._MAX_RECONNECT_RETRIES, + budget.backoff, type(root).__name__, root, + ) + await self._backoff_sleep(budget) + # Check again after sleeping + return not self._shutdown_event.is_set() + + async def _on_initial_connect_error(self, exc: Exception, root: BaseException, + failure_class: str, budget: "_RetryBudget") -> bool: + if failure_class == "permanent": + # Deterministic failure (bad command, non-MCP URL, + # 401/403): park at once instead of burning the ladder. + # Auth failures park rather than return so the task + # stays alive to pick up fresh tokens later. + if _core._is_auth_error(root): + logger.warning( + "MCP server '%s' failed initial authentication, " + "parking until credentials change; re-authenticate " + "with `hermes mcp login %s` " + "(state: connecting → parked): %s: %s", + self.name, self.name, + type(root).__name__, root, + ) + else: + logger.warning( + "MCP server '%s' failed initial connection with a " + "permanent error, parking without retries " + "(state: connecting → parked): %s: %s", + self.name, type(root).__name__, root, + ) + return await self._park_initial_failure(exc, "after permanent initial failure", budget) + + budget.initial_retries += 1 + if budget.initial_retries > _core._MAX_INITIAL_CONNECT_RETRIES: + logger.warning( + "MCP server '%s' failed initial connection after " + "%d attempts, parking until a reconnect is " + "requested (state: connecting → parked): %s: %s", + self.name, _core._MAX_INITIAL_CONNECT_RETRIES, + type(root).__name__, root, + ) + return await self._park_initial_failure(exc, "after initial connection failures", budget) + + logger.debug( + "MCP server '%s' initial connection failed " + "(attempt %d/%d), retrying in %.0fs: %s: %s", + self.name, budget.initial_retries, + _core._MAX_INITIAL_CONNECT_RETRIES, budget.backoff, + type(root).__name__, root, + ) + await self._backoff_sleep(budget) + # Check if shutdown was requested during the sleep + if self._shutdown_event.is_set(): + self._error = exc + self._ready.set() + return False + return True + + async def _on_permanent_error(self, root: BaseException, budget: "_RetryBudget") -> bool: + # An auth failure on a PROVEN session is often a corrupt + # OAuth lock from a raced teardown, not revoked + # credentials: grant ONE suspect+reconnect cycle first. + if ( + _core._is_auth_error(root) + and self._session_proven + and not self._permanent_grace_used + ): + self._permanent_grace_used = True + self.mark_suspect( + f"auth error on proven session: {root}" + ) + logger.warning( + "MCP server '%s': auth error on a previously " + "healthy session — marking suspect and forcing " + "one reconnect instead of parking (state: " + "connected → suspect): %s: %s", + self.name, type(root).__name__, root, + ) + self._reconnect_retries = 0 + budget.backoff = 1.0 + await asyncio.sleep(_core._jittered(1.0)) + return not self._shutdown_event.is_set() + # Deterministic failure on a working server: park now. + logger.warning( + "MCP server '%s' hit a permanent error, parking " + "without retries; will self-probe every %ds " + "(state: connected → parked): %s: %s", + self.name, _core._PARKED_RETRY_INTERVAL, + type(root).__name__, root, + ) + return await self._park_and_rearm("from parked state (permanent error)", budget) + + async def start(self, config: dict): + """Create the background Task and wait until ready (or failed).""" + self._task = asyncio.ensure_future(self.run(config)) + try: + await self._ready.wait() + except asyncio.CancelledError: + # The caller's connect timeout (discover_mcp_tools wraps start() + # in asyncio.wait_for) cancels *this* coroutine, but the + # ensure_future'd run() task is independent and would otherwise + # keep running detached — parked on a hung transport with no + # owner to reap it (#59349). Propagate the cancellation so the + # transport context managers unwind and their finally blocks + # release the child process / FDs. + if self._task and not self._task.done(): + self._task.cancel() + raise + if self._error: + raise self._error + + async def shutdown(self): + """Signal the Task to exit and wait for clean resource teardown.""" + self._shutdown_event.set() + # Defensive: if _wait_for_lifecycle_event is blocking, we need ANY + # event to unblock it. _shutdown_event alone is sufficient (the + # helper checks shutdown first), but setting reconnect too ensures + # there's no race where the helper misses the shutdown flag after + # returning "reconnect". + self._reconnect_event.set() + if self._task and not self._task.done(): + try: + await asyncio.wait_for(self._task, timeout=10) + except asyncio.TimeoutError: + logger.warning( + "MCP server '%s' shutdown timed out, cancelling task", + self.name, + ) + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + if self._pending_refresh_tasks: + for task in list(self._pending_refresh_tasks): + task.cancel() + await asyncio.gather(*self._pending_refresh_tasks, return_exceptions=True) + self._pending_refresh_tasks.clear() + self._deregister_tools() + self.session = None + + def _deregister_tools(self) -> None: + """Drop this server's tools from the global registry (idempotent). + + Pulls the server's tool schemas out of the registry so the agent + stops advertising them to the model. Called on shutdown AND when the + reconnect budget is exhausted, so a dead server never leaves phantom + tool definitions bloating the prompt cache and producing "not + connected" errors on every turn. + """ + from tools.registry import registry + + for tool_name in list(getattr(self, "_registered_tool_names", [])): + registry.deregister(tool_name, scope=_core._server_registry_scope(self.name)) + _core._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: + await self._cancel_waiters(shutdown_task, reconnect_task)