Files
hermes-agent/tools/mcp_tool_server_run.py

568 lines
25 KiB
Python

"""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
def _recycle_if_due(self) -> bool:
"""Latch a stdio idle/lifetime recycle when its deadline has passed."""
recycle_reason = self._stdio_recycle_reason()
if recycle_reason is None:
return False
self._mark_stdio_recycled(recycle_reason)
return True
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:
if self._recycle_if_due():
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
if self._recycle_if_due():
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 with self._rpc_lock:
await self._keepalive_probe()
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"))
# 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"
):
await self._preflight_content_type(
config["url"],
headers=dict(config.get("headers") or {}),
ssl_verify=config.get("ssl_verify", True),
client_cert=_core._resolve_client_cert(self.name, config),
)
except (_core.InvalidMcpUrlError, _core.NonMcpEndpointError) as exc:
# Fail fast and non-retryably: publish the error to start().
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
# Dormant until a lazy call wakes it (untimed: nothing to self-probe).
return await self._wait_for_reconnect_or_shutdown() != "shutdown"
# 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 = []