Teknium
317 lines
17 KiB
Python
317 lines
17 KiB
Python
"""Session health for MCPServerTask: dynamic tool refresh on list_changed notifications, server log forwarding, keepalive probes, suspect-mark / lazy-verify, in-flight call fail-fast, stdio child liveness and stdio idle/lifetime recycling. Split from tools/mcp_tool.py."""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from typing import Iterable, Optional
|
|
from tools.mcp_tool_errors import _is_method_not_found_error, _unwrap_exception_group
|
|
from tools.mcp_tool_schema import mcp_prefixed_tool_name
|
|
from tools.mcp_tool_registration import _forget_mcp_tool_server
|
|
from tools.mcp_tool_common import _core
|
|
|
|
logger = logging.getLogger("tools.mcp_tool")
|
|
|
|
_KEEPALIVE_RPC_TIMEOUT = 30.0
|
|
|
|
|
|
def _stdio_children_dead_impl(pids, is_http: bool) -> bool:
|
|
"""True when every pid has exited. Best-effort: False (unknown → don't fail fast) for HTTP,
|
|
no captured PIDs, missing psutil, or a failed probe."""
|
|
if not pids or is_http:
|
|
return False
|
|
try:
|
|
import psutil
|
|
except ImportError:
|
|
return False
|
|
for pid in pids:
|
|
try:
|
|
if psutil.pid_exists(pid): # handles Windows without signal-permission noise
|
|
return False
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
|
|
class MCPServerHealthMixin:
|
|
"""Methods of :class:`tools.mcp_tool.MCPServerTask` (mixed in; relies on its attributes)."""
|
|
|
|
__slots__ = ()
|
|
|
|
def _is_http(self) -> bool:
|
|
return "url" in self._config
|
|
|
|
def _is_recycled_stdio(self) -> bool:
|
|
"""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:
|
|
self._lifecycle_started_at = self._last_tool_call_at = time.monotonic()
|
|
self._recycled_reason = None
|
|
|
|
# ------------------------------------------------------- stdio recycling
|
|
|
|
def _stdio_recycle_deadlines(self):
|
|
"""``[(deadline, reason), ...]`` for the configured lifetime/idle limits; empty for HTTP
|
|
servers or while an RPC holds the lock."""
|
|
if self._is_http() or self._rpc_lock.locked():
|
|
return []
|
|
deadlines = []
|
|
if self._max_lifetime_seconds is not None:
|
|
deadlines.append((self._lifecycle_started_at + self._max_lifetime_seconds, "max_lifetime_seconds"))
|
|
if self._idle_timeout_seconds is not None:
|
|
deadlines.append((self._last_tool_call_at + self._idle_timeout_seconds, "idle_timeout_seconds"))
|
|
return deadlines
|
|
|
|
def _stdio_recycle_reason(self, now: Optional[float] = None) -> Optional[str]:
|
|
"""The stdio recycle reason if idle/age limits have elapsed (lifetime wins), else None."""
|
|
now = time.monotonic() if now is None else now
|
|
return next((reason for deadline, reason in self._stdio_recycle_deadlines() if now >= deadline), None)
|
|
|
|
def _next_stdio_recycle_deadline(self) -> Optional[float]:
|
|
"""The next monotonic recycle deadline for stdio, if any."""
|
|
deadlines = self._stdio_recycle_deadlines()
|
|
return min(d for d, _ in 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
|
|
|
|
# -------------------------------------------------- notifications / logs
|
|
|
|
async def _refresh_tools_task(self):
|
|
"""Run a dynamic tool refresh and log failures from background tasks."""
|
|
try:
|
|
await self._refresh_tools()
|
|
except Exception:
|
|
logger.exception("MCP server '%s': dynamic tool refresh failed", self.name)
|
|
|
|
def _schedule_tools_refresh(self) -> asyncio.Task:
|
|
"""Schedule a background tool refresh and keep it strongly referenced."""
|
|
task = asyncio.create_task(self._refresh_tools_task())
|
|
self._pending_refresh_tasks.add(task)
|
|
task.add_done_callback(self._pending_refresh_tasks.discard)
|
|
return task
|
|
|
|
def _make_logging_callback(self):
|
|
"""``logging_callback`` forwarding server ``notifications/message`` into Hermes logging
|
|
tagged with the server name (the SDK default drops them)."""
|
|
async def _on_log(params):
|
|
try:
|
|
level = _core._MCP_LOG_LEVEL_MAP.get(str(getattr(params, "level", "info")).lower(), logging.INFO)
|
|
data = getattr(params, "data", None)
|
|
if not isinstance(data, str):
|
|
try:
|
|
data = json.dumps(data, ensure_ascii=False, default=str)
|
|
except (TypeError, ValueError):
|
|
data = str(data)
|
|
# Cap payloads so a chatty server can't flood agent.log.
|
|
if len(data) > 2000:
|
|
data = data[:2000] + "... [truncated]"
|
|
logger_name = getattr(params, "logger", None)
|
|
origin = f"{self.name}/{logger_name}" if logger_name else self.name
|
|
logger.log(level, "MCP server log [%s]: %s", origin, data)
|
|
except Exception:
|
|
logger.debug("Failed to handle MCP log notification from '%s'", self.name, exc_info=True)
|
|
return _on_log
|
|
|
|
def _make_message_handler(self):
|
|
"""``message_handler`` for ``ClientSession``: only ``ToolListChangedNotification`` triggers
|
|
a refresh; prompt/resource changes are logged."""
|
|
async def _handler(message):
|
|
try:
|
|
if isinstance(message, Exception):
|
|
logger.debug("MCP message handler (%s): exception: %s", self.name, message)
|
|
return
|
|
if _core._MCP_NOTIFICATION_TYPES and isinstance(message, _core.ServerNotification):
|
|
# mcp 2.0 made ServerNotification a plain union (payload IS the message)
|
|
# instead of a RootModel (payload under ``.root``). ``isinstance`` accepts
|
|
# both; only the unwrap differs — without it ``.root`` raises into the
|
|
# catch-all and refreshes stop.
|
|
match getattr(message, "root", message):
|
|
case _core.ToolListChangedNotification():
|
|
logger.info("MCP server '%s': received tools/list_changed notification", self.name)
|
|
# Refresh in a separate task: some servers emit list_changed right
|
|
# after initialize while another request is in flight, and refreshing
|
|
# synchronously inside the handler can wedge the stdio JSON-RPC stream.
|
|
self._schedule_tools_refresh()
|
|
# Yield one tick so short-lived notification contexts (and tests)
|
|
# can observe the scheduled refresh.
|
|
await asyncio.sleep(0)
|
|
case _core.PromptListChangedNotification():
|
|
logger.debug("MCP server '%s': prompts/list_changed (ignored)", self.name)
|
|
case _core.ResourceListChangedNotification():
|
|
logger.debug("MCP server '%s': resources/list_changed (ignored)", self.name)
|
|
case _:
|
|
pass
|
|
except Exception:
|
|
logger.exception("Error in MCP message handler for '%s'", self.name)
|
|
return _handler
|
|
|
|
def _deregister_owned(self, tool_names: Iterable[str]) -> None:
|
|
"""Deregister *tool_names* that this server's toolset still owns. Never removes a
|
|
colliding name currently owned by another server."""
|
|
from tools.registry import registry
|
|
|
|
toolset_name = f"mcp-{self.name}"
|
|
for tool_name in tool_names:
|
|
if registry.get_toolset_for_tool(tool_name) != toolset_name:
|
|
continue
|
|
registry.deregister(tool_name, scope=_core._server_registry_scope(self.name))
|
|
_forget_mcp_tool_server(tool_name)
|
|
|
|
async def _refresh_tools(self):
|
|
"""Re-fetch tools on ``tools/list_changed`` and update the registry. The lock serializes
|
|
rapid-fire notifications; after the list_tools ``await`` all mutations are synchronous —
|
|
atomic on the event loop."""
|
|
if not self._advertises_tools():
|
|
return # tools/list would raise MCPError(-32601)
|
|
async with self._refresh_lock:
|
|
old_tool_names = set(self._registered_tool_names)
|
|
async with self._rpc_lock:
|
|
new_mcp_tools = await _core._paginate_full_list(self.session.list_tools, "tools", self.name)
|
|
# Remove only stale names first — no nuke-and-repave: live agent turns may hold
|
|
# tool-call IDs pointing at existing handlers, and in-place replacement avoids
|
|
# transient "tool not connected" races.
|
|
self._deregister_owned(old_tool_names - {mcp_prefixed_tool_name(self.name, tool.name) for tool in new_mcp_tools})
|
|
# Re-register; the helper may skip names ambiguous after normalization. A raw name
|
|
# can become ambiguous without changing its normalized name, so the pre-pass misses
|
|
# it: drop any old entry the final collision-checked registration no longer owns.
|
|
self._tools = new_mcp_tools
|
|
registered_names = _core._register_server_tools(self.name, self, self._config)
|
|
self._deregister_owned(old_tool_names - set(registered_names))
|
|
self._registered_tool_names = registered_names
|
|
# Log what changed (user-visible).
|
|
new_tool_names = set(registered_names)
|
|
changes = [f"{label}: {', '.join(sorted(names))}" for label, names in
|
|
(("added", new_tool_names - old_tool_names), ("removed", old_tool_names - new_tool_names)) if names]
|
|
if changes:
|
|
logger.warning("MCP server '%s': tools changed dynamically — %s. "
|
|
"Verify these changes are expected.", self.name, "; ".join(changes))
|
|
else:
|
|
logger.info("MCP server '%s': dynamically refreshed %d tool(s) (no changes)",
|
|
self.name, len(self._registered_tool_names))
|
|
|
|
# ------------------------------------------------------ keepalive / health
|
|
|
|
async def _keepalive_probe(self) -> None:
|
|
"""Exercise the session; raise on a genuine connection failure. ``ping`` first (cheap,
|
|
OPTIONAL utility). On -32601 latch ``_ping_unsupported`` and fall back to ``list_tools``
|
|
when the server advertises tools; otherwise the -32601 propagates (no liveness primitive
|
|
left). The latch resets on each fresh transport connection."""
|
|
if not self._ping_unsupported:
|
|
try:
|
|
await asyncio.wait_for(self.session.send_ping(), timeout=_KEEPALIVE_RPC_TIMEOUT)
|
|
return
|
|
except Exception as exc:
|
|
if _is_method_not_found_error(exc):
|
|
# Ping is definitively unsupported.
|
|
if not self._advertises_tools():
|
|
raise
|
|
self._ping_unsupported = True
|
|
logger.info("MCP server '%s': does not implement the optional 'ping' utility (-32601); "
|
|
"using 'list_tools' for keepalive on this connection.", self.name)
|
|
elif isinstance(exc, (TimeoutError, asyncio.TimeoutError)) and self._advertises_tools():
|
|
# A server that silently drops ping looks like a dead transport. Confirm with
|
|
# list_tools before declaring it dead; if that also fails, propagate the
|
|
# original failure.
|
|
try:
|
|
await asyncio.wait_for(self.session.list_tools(), timeout=_KEEPALIVE_RPC_TIMEOUT)
|
|
except Exception:
|
|
raise exc from None
|
|
# Transport alive; latch so later keepalives skip the 30s wait.
|
|
self._ping_unsupported = True
|
|
logger.info("MCP server '%s': ping timed out but list_tools succeeded — server "
|
|
"silently drops ping; using 'list_tools' for keepalive on this connection.", self.name)
|
|
return
|
|
else:
|
|
raise # closed transport, expired session, etc. — real failure
|
|
# Fallback probe for servers without ping support.
|
|
await asyncio.wait_for(self.session.list_tools(), timeout=_KEEPALIVE_RPC_TIMEOUT)
|
|
|
|
def _mark_session_proven(self) -> None:
|
|
"""Record that the session demonstrated real health (keepalive or tool-call success).
|
|
Only then is the reconnect budget cleared: a handshake that drops moments later must
|
|
keep consuming ``_reconnect_retries`` so a flapping transport still reaches the park
|
|
instead of respawning forever."""
|
|
if self._session_proven:
|
|
return
|
|
self._session_proven = True
|
|
self._reconnect_retries = 0
|
|
if self._was_parked:
|
|
self._was_parked = False
|
|
logger.warning("MCP server '%s': revived — session healthy again after "
|
|
"parking (state: parked → connected)", self.name)
|
|
# A proven fresh transport clears the one-time permanent-failure grace and any race bookkeeping.
|
|
self._permanent_grace_used = False
|
|
self._teardown_race = False
|
|
|
|
def mark_suspect(self, reason: str) -> None:
|
|
"""Latch a suspicion (no I/O). The NEXT call verifies via :meth:`ensure_healthy` and
|
|
recycles the transport if the probe fails."""
|
|
if self._suspect_reason is None and reason:
|
|
logger.warning("MCP server '%s': connection marked suspect (%s); next call will health-check it",
|
|
self.name, reason)
|
|
self._suspect_reason = reason or None
|
|
|
|
async def ensure_healthy(self, timeout: float = 5.0) -> bool:
|
|
"""Verify a suspect connection before reuse; recycle if dead. True when healthy (suspicion
|
|
cleared). On failure requests a reconnect, drops the stale session so the caller's
|
|
no-session path takes over, and returns False. Never raises."""
|
|
reason = self._suspect_reason
|
|
if not reason:
|
|
return True
|
|
if self.session is None:
|
|
# Nothing to verify — the reconnect path owns recovery now.
|
|
self._suspect_reason = None
|
|
self._reconnect_event.set()
|
|
return False
|
|
try:
|
|
await asyncio.wait_for(self._keepalive_probe(), timeout=timeout)
|
|
except Exception as exc:
|
|
root = _unwrap_exception_group(exc)
|
|
logger.warning("MCP server '%s': suspect connection (%s) failed health check (%s: %s) — "
|
|
"requesting reconnect (state: suspect → degraded)",
|
|
self.name, reason, type(root).__name__, root)
|
|
self._suspect_reason = None
|
|
self.mark_suspect(f"health check failed after {reason}")
|
|
self.session = None
|
|
self._ready.clear()
|
|
self._reconnect_event.set()
|
|
return False
|
|
logger.info("MCP server '%s': suspect connection passed health check (%s) — clearing suspicion",
|
|
self.name, reason)
|
|
self._suspect_reason = None
|
|
self._mark_session_proven()
|
|
return True
|
|
|
|
def _fail_inflight_calls(self, reason: str) -> None:
|
|
"""Cancel every in-flight RPC on this connection. Called from lifecycle exits BEFORE the
|
|
transport unwinds: the SDK does not always fail pending requests when streams close, so
|
|
a call would otherwise wait out the full tool timeout. Cancelling anything flags
|
|
``_teardown_race`` so run() treats the next reconnect as recovery rather than charging
|
|
the rapid-drop budget."""
|
|
victims = [t for t in self._inflight_tasks if not t.done()]
|
|
if not victims:
|
|
return
|
|
self._reconnecting = True
|
|
self._teardown_race = True
|
|
self.mark_suspect(f"{reason} tore down {len(victims)} in-flight call(s)")
|
|
for task in victims:
|
|
task.cancel()
|
|
|
|
def _stdio_children_dead(self) -> bool:
|
|
"""True when every stdio child we spawned has exited (see :func:`_stdio_children_dead_impl`)."""
|
|
return _stdio_children_dead_impl(getattr(self, "_stdio_child_pids", None), self._is_http())
|
|
|
|
async def _watch_stdio_children(self) -> None:
|
|
"""Poll child liveness while a stdio RPC is in flight; resolves when a tracked child dies
|
|
so the caller cancels the RPC instead of waiting out the timeout."""
|
|
while not self._stdio_children_dead():
|
|
await asyncio.sleep(0.25)
|