hermes-admin/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 29 Packages Projects Releases Wiki Activity

refactor(mcp): move connect/lazy-start/discovery/public API into mcp_tool_discovery.py

Browse Source
This commit is contained in:
Teknium
2026-09-02 15:50:54 -07:00
parent cbecd7d0c6
commit 46697da751
3 changed files with 678 additions and 651 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/tools/test_mcp_tool.py
View File

@@ -2358,7 +2358,7 @@ class TestDiscoveryFailedCount:
_ensure_mcp_loop()
# Capture the logger to verify failed_count in summary
with patch("tools.mcp_tool.logger") as mock_logger:
with patch("tools.mcp_tool_discovery.logger") as mock_logger:
discover_mcp_tools()
# Find the summary info call
@@ -2401,7 +2401,7 @@ class TestDiscoveryFailedCount:
patch("tools.mcp_tool._existing_tool_names", return_value=["mcp__ok1__t", "mcp__ok2__t"]):
_ensure_mcp_loop()
with patch("tools.mcp_tool.logger") as mock_logger:
with patch("tools.mcp_tool_discovery.logger") as mock_logger:
discover_mcp_tools()
info_calls = [str(call) for call in mock_logger.info.call_args_list]

667
tools/mcp_tool.py
View File

@@ -210,6 +210,24 @@ from tools.mcp_tool_agent import ( # noqa: F401
from tools.mcp_tool_transport import MCPServerTransportMixin
from tools.mcp_tool_server_run import MCPServerRunMixin
from tools.mcp_tool_health import MCPServerHealthMixin
from tools.mcp_tool_discovery import ( # noqa: F401 -- re-exported for callers and test patches
_record_connect_failure,
_clear_connect_failure,
_connect_cooldown_active,
_connect_server,
_request_lazy_reconnect,
_resolve_server_lazy,
_ensure_lazy_server_connected,
_get_connected_server_for_call,
_discover_and_register_server,
register_mcp_servers,
discover_mcp_tools,
is_mcp_tool_parallel_safe,
get_mcp_status,
probe_mcp_server_tools,
has_registered_mcp_tools,
get_registered_mcp_server_names,
)
# Wall-clock bound on the (fail-open) OSV malware preflight run off the loop
@@ -660,16 +678,6 @@ class MCPServerTask(MCPServerRunMixin, MCPServerTransportMixin, MCPServerHealthM
_MCP_CONTENT_TYPES = ("application/json", "text/event-stream")
# ---------------------------------------------------------------------------
# Module-level state (every mutation under ``_lock``)
# ---------------------------------------------------------------------------
@@ -704,28 +712,6 @@ _CONNECT_RETRY_BASE_BACKOFF_SEC = 30.0
_CONNECT_RETRY_MAX_BACKOFF_SEC = 600.0
def _record_connect_failure(server_name: str) -> None:
"""Stamp a geometric, capped retry cooldown after a failed connect (under ``_lock``)."""
n = _server_connect_failures.get(server_name, 0) + 1
_server_connect_failures[server_name] = n
backoff = min(
_CONNECT_RETRY_BASE_BACKOFF_SEC * (2 ** (n - 1)),
_CONNECT_RETRY_MAX_BACKOFF_SEC,
)
_server_connect_retry_after[server_name] = time.monotonic() + backoff
def _clear_connect_failure(server_name: str) -> None:
"""Clear the connect-cooldown state after a successful connection."""
_server_connect_failures.pop(server_name, None)
_server_connect_retry_after.pop(server_name, None)
def _connect_cooldown_active(server_name: str) -> bool:
"""Return True if ``server_name`` is still within its retry cooldown."""
deadline = _server_connect_retry_after.get(server_name)
return deadline is not None and time.monotonic() < deadline
# Circuit breaker per server: closed (count < threshold) -> open (calls
# short-circuit with a "stop retrying" message until the cooldown elapses) ->
# half-open (next call is a probe; success closes, failure re-arms). Mutate
@@ -1148,628 +1134,11 @@ def _run_on_mcp_loop(coro_or_factory, timeout: float = 30):
# Connecting, lazy start, discovery
# ---------------------------------------------------------------------------
async def _connect_server(name: str, config: dict) -> MCPServerTask:
"""Create an MCPServerTask, start it and return once ready.
Tear it down with ``server.shutdown()`` on the same loop. Raises on bad
config, missing HTTP support, or connect/initialize failure.
"""
server = MCPServerTask(name)
claim = _connect_server_claim.get()
claim_token = None
if claim is not None:
claim(server)
# The run task copies this context; the claim is for this attempt
# only, so don't retain the discovery closure for the server's life.
claim_token = _connect_server_claim.set(None)
try:
await server.start(config)
except asyncio.CancelledError:
# start() already reaps server._task; a shutdown() here could
# swallow the cancellation.
raise
except BaseException:
# Discovery owns claimed tasks (recoverable park vs terminal failure);
# standalone probes have no revival owner and must reap locally.
if claim is None:
try:
await server.shutdown()
except Exception as shutdown_exc: # noqa: BLE001 -- best-effort reap, don't mask the real error
logger.debug(
"MCP server '%s' shutdown during orphan-reap failed: %s",
name, shutdown_exc,
)
raise
finally:
if claim_token is not None:
_connect_server_claim.reset(claim_token)
return server
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 _resolve_server_lazy(name: str, config: dict) -> bool:
"""True when ``mcp_servers.<name>.lazy`` defers connect to first tool use (default off)."""
return _parse_boolish(config.get("lazy", False), default=False)
def _ensure_lazy_server_connected(server_name: str) -> bool:
"""Connect a lazily-registered server on demand (sync; blocks the caller).
Honours the connect cooldown and the ``_server_connecting`` dedup set and
routes through ``_discover_and_register_server`` so park/recycle/cooldown
bookkeeping stays in one place. True when a live session exists after.
"""
with _lock:
server = _servers.get(server_name)
if server is not None and server.session is not None:
return True
config = _lazy_server_configs.get(server_name)
if not config:
return False
if _connect_cooldown_active(server_name):
return False
if server_name in _server_connecting:
return False
_server_connecting.add(server_name)
_server_connect_errors.pop(server_name, None)
logger.info("MCP server '%s': lazy start on first use", server_name)
_ensure_mcp_loop()
connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT)
async def _connect():
return await _discover_and_register_server(server_name, config)
try:
_run_on_mcp_loop(_connect, timeout=float(connect_timeout) + 30.0)
except BaseException as exc:
message = _format_connect_error(exc)
with _lock:
_server_connecting.discard(server_name)
_server_connect_errors[server_name] = message
_record_connect_failure(server_name)
logger.warning(
"Lazy MCP connect failed for '%s': %s", server_name, message,
)
return False
with _lock:
_server_connecting.discard(server_name)
_clear_connect_failure(server_name)
_lazy_server_configs.pop(server_name, None)
stale_fingerprint = _lazy_server_fingerprints.pop(server_name, None)
cached_names = _lazy_server_tool_names.pop(server_name, None) or []
server = _servers.get(server_name)
live_names = set(
getattr(server, "_registered_tool_names", []) or []
)
# The cached manifest may advertise tools the live server no longer
# serves; deregister those phantoms.
phantom_names = [n for n in cached_names if n not in live_names]
if phantom_names:
from tools.registry import registry
for tool_name in phantom_names:
registry.deregister(tool_name, scope=_server_registry_scope(server_name))
_forget_mcp_tool_server(tool_name)
logger.info(
"MCP server '%s': deregistered %d phantom cached tool(s) not "
"served live (stale schema-cache fingerprint %s): %s",
server_name, len(phantom_names), stale_fingerprint,
", ".join(phantom_names),
)
return server is not None and server.session is not None
def _get_connected_server_for_call(server_name: str) -> Optional[MCPServerTask]:
"""Return a connected server; the single first-use connect point for lazy
servers and the wake-up point for recycled stdio ones."""
with _lock:
server = _servers.get(server_name)
is_lazy = server_name in _lazy_server_configs
if is_lazy and (server is None or server.session is None):
_ensure_lazy_server_connected(server_name)
with _lock:
server = _servers.get(server_name)
return server
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
async def _discover_and_register_server(name: str, config: dict) -> List[str]:
"""Connect one server, register its tools; return the registered names."""
connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT)
# The claim callback runs inside _connect_server while this frame is
# suspended; a list append avoids a nonlocal rebind.
claimed: List[MCPServerTask] = []
def _claim_server(created: MCPServerTask) -> None:
claimed.append(created)
claim_token = _connect_server_claim.set(_claim_server)
try:
server = await asyncio.wait_for(
_connect_server(name, config),
timeout=connect_timeout,
)
except BaseException:
server = claimed[0] if claimed else None
task = server._task if server is not None else None
task_cancelling = (
task.cancelling()
if task is not None and hasattr(task, "cancelling")
else 0
)
if (
server is not None
and server._error is not None
and task is not None
and not task.done()
and not task_cancelling
):
# Recoverable park: the run task stays alive to self-probe, so
# adopt it for shutdown/revival.
with _lock:
_servers[name] = server
_server_scope_keys[name] = _mcp_registry_scope()
elif server is not None:
await server.shutdown()
raise
finally:
_connect_server_claim.reset(claim_token)
with _lock:
_server_connecting.discard(name)
_server_connect_errors.pop(name, None)
_servers[name] = server
_server_scope_keys[name] = _mcp_registry_scope()
registered_names = _register_server_tools(name, server, config)
server._registered_tool_names = list(registered_names)
transport_type = "HTTP" if "url" in config else "stdio"
logger.info(
"MCP server '%s' (%s): registered %d tool(s): %s",
name, transport_type, len(registered_names),
", ".join(registered_names),
)
return registered_names
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
"""Connect the given ``{name: config}`` servers and register their tools.
Idempotent for connected names; ``enabled: false`` servers are skipped
without disconnecting existing sessions. Returns every registered MCP
tool name.
"""
if not _ensure_mcp_sdk():
logger.debug("MCP SDK not available -- skipping explicit MCP registration")
return []
servers = _filter_suspicious_mcp_servers(servers)
if not servers:
logger.debug("No explicit MCP servers provided")
return []
# Candidates: enabled, not connected, not connecting (dedups concurrent
# discovery entry points), not lazily registered, not in backoff.
with _lock:
connecting = set(_server_connecting)
new_servers = {
k: v
for k, v in servers.items()
if k not in _servers
and k not in connecting
and k not in _lazy_server_configs
and _parse_boolish(v.get("enabled", True), default=True)
and not _connect_cooldown_active(k)
}
# Known servers without a live session are parked or mid-reconnect;
# their tools are deregistered so nothing else can nudge them.
stale_cached = [
_servers[k]
for k in servers
if k in _servers and getattr(_servers[k], "session", None) is None
]
_server_connecting.update(new_servers)
for srv_name in new_servers:
_server_connect_errors.pop(srv_name, None)
# Track which servers opt-in to parallel tool calls (idempotent).
for srv_name, srv_cfg in servers.items():
if _parse_boolish(srv_cfg.get("supports_parallel_tool_calls", False), default=False):
_parallel_safe_servers.add(srv_name)
else:
_parallel_safe_servers.discard(srv_name)
for srv in stale_cached:
_signal_reconnect(srv)
if not new_servers:
return _existing_tool_names()
# ``lazy: true`` servers with a valid schema-cache entry register from
# cache without connecting; a missing/stale entry falls back to eager.
eager_servers: Dict[str, dict] = dict(new_servers)
lazy_registered = 0
lazy_server_count = 0
try:
from tools.mcp_schema_cache import config_fingerprint, get_cached_entry
except Exception: # pragma: no cover - cache module missing
config_fingerprint = None # type: ignore[assignment]
get_cached_entry = None # type: ignore[assignment]
if config_fingerprint is not None and get_cached_entry is not None:
for name, cfg in new_servers.items():
if not _resolve_server_lazy(name, cfg):
continue
entry = get_cached_entry(name, config_fingerprint(cfg))
if not entry:
continue
with _lock:
_server_connecting.discard(name)
try:
names = _register_from_cache_sync(name, cfg, entry)
except Exception as exc:
logger.warning(
"Failed lazy MCP registration for '%s': %s", name, exc,
)
with _lock:
_server_connecting.add(name)
continue
eager_servers.pop(name, None)
lazy_registered += len(names)
lazy_server_count += 1
new_servers = eager_servers
if not new_servers:
if lazy_registered:
logger.info(
"MCP: registered %d lazy tool(s) from schema cache "
"(no processes spawned)",
lazy_registered,
)
return _existing_tool_names()
_ensure_mcp_loop()
async def _discover_all():
server_names = list(new_servers.keys())
results = await asyncio.gather(
*(_discover_and_register_server(name, cfg) for name, cfg in new_servers.items()),
return_exceptions=True,
)
for name, result in zip(server_names, results):
if isinstance(result, BaseException):
command = new_servers.get(name, {}).get("command")
message = _format_connect_error(result)
with _lock:
_server_connecting.discard(name)
_server_connect_errors[name] = message
_record_connect_failure(name)
logger.warning(
"Failed to connect to MCP server '%s'%s: %s",
name,
f" (command={command})" if command else "",
message,
)
else:
with _lock:
_server_connecting.discard(name)
_server_connect_errors.pop(name, None)
_clear_connect_failure(name)
# Clear a stale interrupt flag (executor threads are reused) so a prior
# session's interrupt cannot cancel this discovery pass.
from tools.interrupt import is_interrupted as _is_interrupted, set_interrupt as _set_interrupt
_was_interrupted = _is_interrupted()
if _was_interrupted:
_set_interrupt(False)
try:
_run_on_mcp_loop(_discover_all, timeout=120)
except (TimeoutError, InterruptedError) as _e:
# Entries stranded in _server_connecting would block future
# reconnect attempts.
with _lock:
stale = [n for n in new_servers if n in _server_connecting]
if stale:
logger.warning(
"MCP discovery %s while %d server(s) were still "
"connecting; clearing stale connecting set: %s",
"timed out" if isinstance(_e, TimeoutError) else "interrupted",
len(stale),
", ".join(stale),
)
_server_connecting.difference_update(stale)
for _sn in stale:
_server_connect_errors.setdefault(
_sn,
f"Connection attempt {'timed out' if isinstance(_e, TimeoutError) else 'interrupted'} during discovery",
)
raise
finally:
if _was_interrupted:
_set_interrupt(True)
with _lock:
connected = [
n
for n in new_servers
if n in _servers and n not in _server_connect_errors
]
new_tool_count = sum(
len(getattr(_servers[n], "_registered_tool_names", []))
for n in connected
)
failed = len(new_servers) - len(connected)
new_tool_count += lazy_registered
connected_count = len(connected) + lazy_server_count
if new_tool_count or failed:
summary = f"MCP: registered {new_tool_count} tool(s) from {connected_count} server(s)"
if failed:
summary += f" ({failed} failed)"
logger.info(summary)
return _existing_tool_names()
def discover_mcp_tools() -> List[str]:
"""Entry point: load config, connect servers, register tools.
Safe without the ``mcp`` package (returns []). Idempotent: only servers
missing from a previous call are retried. Returns all MCP tool names.
"""
servers = _load_mcp_config()
if not servers:
logger.debug("No MCP servers configured")
return []
# SDK import deferred to here so a config without servers never pays it.
if not _ensure_mcp_sdk():
logger.debug("MCP SDK not available -- skipping MCP tool discovery")
return []
# Cross-process guard: a lock loser waits for the holder, then runs its
# own discovery; if locking is unavailable or the wait expires, run
# unguarded (fail-soft).
cookie = _try_acquire_mcp_discovery_lock()
if cookie is None:
logger.debug(
"Another process holds MCP discovery lock -- retrying with backoff"
)
for _ in range(_MCP_DISCOVERY_LOCK_MAX_RETRIES):
time.sleep(_MCP_DISCOVERY_LOCK_RETRY_DELAY_S)
cookie = _try_acquire_mcp_discovery_lock()
if cookie is not None:
break
if cookie is None:
logger.warning(
"MCP discovery lock still held after %d retries -- "
"running discovery unguarded",
_MCP_DISCOVERY_LOCK_MAX_RETRIES,
)
elif cookie is not _LOCK_UNAVAILABLE:
logger.debug("Retry succeeded -- acquired MCP discovery lock")
try:
with _lock:
connecting = set(_server_connecting)
new_server_names = [
name
for name, cfg in servers.items()
if name not in _servers
and name not in connecting
and _parse_boolish(cfg.get("enabled", True), default=True)
]
tool_names = register_mcp_servers(servers)
if not new_server_names:
return tool_names
with _lock:
connected_server_names = [
name
for name in new_server_names
if name in _servers and name not in _server_connect_errors
]
new_tool_count = sum(
len(getattr(_servers[name], "_registered_tool_names", []))
for name in connected_server_names
)
failed_count = len(new_server_names) - len(connected_server_names)
if new_tool_count or failed_count:
summary = f" MCP: {new_tool_count} tool(s) from {len(connected_server_names)} server(s)"
if failed_count:
summary += f" ({failed_count} failed)"
logger.info(summary)
return tool_names
finally:
if cookie not in (None, _LOCK_UNAVAILABLE):
cookie.release()
def is_mcp_tool_parallel_safe(tool_name: str) -> bool:
"""True when the tool's server opted into ``supports_parallel_tool_calls``.
Uses the provenance captured at registration, never the (ambiguous)
``mcp__{server}__{tool}`` string shape.
"""
if not tool_name.startswith(MCP_TOOL_NAME_PREFIX):
return False
with _lock:
server_name = _mcp_tool_server_names.get(tool_name)
return bool(server_name and server_name in _parallel_safe_servers)
def get_mcp_status() -> List[dict]:
"""Status of every configured server for banner/TUI display.
Each dict has name, transport, tools, connected, disabled, status (one of
connected / disabled / connecting / failed / configured) and, for failed,
error. ``enabled: false`` is reported as disabled, not failed.
"""
configured = _load_mcp_config()
if not configured:
return []
with _lock:
active_servers = dict(_servers)
connecting = set(_server_connecting)
connect_errors = dict(_server_connect_errors)
def _entry(name: str, transport: str, status: str, **extra) -> dict:
return {
"name": name,
"transport": transport,
"tools": 0,
"connected": False,
"disabled": status == "disabled",
"status": status,
**extra,
}
result: List[dict] = []
for name, cfg in configured.items():
transport = cfg.get("transport", "http") if "url" in cfg else "stdio"
enabled = _parse_boolish(cfg.get("enabled", True), default=True)
server = active_servers.get(name)
if server and server.session is not None:
entry = _entry(name, transport, "connected", connected=True)
entry["tools"] = (
len(server._registered_tool_names)
if hasattr(server, "_registered_tool_names")
else len(server._tools)
)
if server._sampling:
entry["sampling"] = dict(server._sampling.metrics)
elif not enabled:
entry = _entry(name, transport, "disabled")
elif name in connecting:
entry = _entry(name, transport, "connecting")
elif name in connect_errors:
entry = _entry(name, transport, "failed", error=connect_errors[name])
else:
entry = _entry(name, transport, "configured")
result.append(entry)
return result
def probe_mcp_server_tools() -> Dict[str, List[tuple]]:
"""Connect to each enabled server, list ``(tool_name, description)`` and
disconnect, without registering anything. Failed servers are omitted."""
if not _ensure_mcp_sdk():
return {}
servers_config = _load_mcp_config()
if not servers_config:
return {}
enabled = {
k: v for k, v in servers_config.items()
if _parse_boolish(v.get("enabled", True), default=True)
}
if not enabled:
return {}
_ensure_mcp_loop()
result: Dict[str, List[tuple]] = {}
probed_servers: List[MCPServerTask] = []
async def _probe_all():
names = list(enabled.keys())
coros = []
for name, cfg in enabled.items():
ct = cfg.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT)
coros.append(asyncio.wait_for(_connect_server(name, cfg), timeout=ct))
outcomes = await asyncio.gather(*coros, return_exceptions=True)
for name, outcome in zip(names, outcomes):
if isinstance(outcome, Exception):
logger.debug("Probe: failed to connect to '%s': %s", name, outcome)
continue
probed_servers.append(outcome)
tools = []
for t in outcome._tools:
desc = getattr(t, "description", "") or ""
tools.append((t.name, desc))
result[name] = tools
await asyncio.gather(
*(s.shutdown() for s in probed_servers),
return_exceptions=True,
)
try:
_run_on_mcp_loop(_probe_all, timeout=120)
except Exception as exc:
logger.debug("MCP probe failed: %s", exc)
finally:
_stop_mcp_loop_if_idle()
return result
def has_registered_mcp_tools() -> bool:
"""True if any MCP server has registered tools (cheap; no registry walk).
Checks registered TOOLS, not connected servers, so the per-turn refresh
hook stays idle for zero-tool servers.
"""
with _lock:
return bool(_mcp_tool_server_names)
def get_registered_mcp_server_names() -> set:
"""Server names that registered at least one tool (the live, filtered
signal — not merely what config.yaml lists)."""
with _lock:
return set(_mcp_tool_server_names.values())
def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool:
"""Stop the background event loop and join its thread."""

658
tools/mcp_tool_discovery.py Normal file
View File

@@ -0,0 +1,658 @@
"""Connecting and discovery for tools.mcp_tool: per-server connect cooldown, connect /
lazy-start / recycled-stdio wake-up, ``register_mcp_servers`` / ``discover_mcp_tools``
and the status / probe public API. Split from tools/mcp_tool.py; origin state
(``_servers``, ``_lock``, the loop, patchable helpers) is read through ``_core`` so
``mock.patch("tools.mcp_tool.X")`` keeps working."""
from __future__ import annotations
import asyncio
import logging
import time
from typing import Dict, List, Optional
from tools.mcp_tool_common import _core
logger = logging.getLogger("tools.mcp_tool")
def _record_connect_failure(server_name: str) -> None:
"""Stamp a geometric, capped retry cooldown after a failed connect (under ``_lock``)."""
n = _core._server_connect_failures.get(server_name, 0) + 1
_core._server_connect_failures[server_name] = n
backoff = min(
_core._CONNECT_RETRY_BASE_BACKOFF_SEC * (2 ** (n - 1)),
_core._CONNECT_RETRY_MAX_BACKOFF_SEC,
)
_core._server_connect_retry_after[server_name] = time.monotonic() + backoff
def _clear_connect_failure(server_name: str) -> None:
"""Clear the connect-cooldown state after a successful connection."""
_core._server_connect_failures.pop(server_name, None)
_core._server_connect_retry_after.pop(server_name, None)
def _connect_cooldown_active(server_name: str) -> bool:
"""Return True if ``server_name`` is still within its retry cooldown."""
deadline = _core._server_connect_retry_after.get(server_name)
return deadline is not None and time.monotonic() < deadline
async def _connect_server(name: str, config: dict) -> _core.MCPServerTask:
"""Create an MCPServerTask, start it and return once ready.
Tear it down with ``server.shutdown()`` on the same loop. Raises on bad
config, missing HTTP support, or connect/initialize failure.
"""
server = _core.MCPServerTask(name)
claim = _core._connect_server_claim.get()
claim_token = None
if claim is not None:
claim(server)
# The run task copies this context; the claim is for this attempt
# only, so don't retain the discovery closure for the server's life.
claim_token = _core._connect_server_claim.set(None)
try:
await server.start(config)
except asyncio.CancelledError:
# start() already reaps server._task; a shutdown() here could
# swallow the cancellation.
raise
except BaseException:
# Discovery owns claimed tasks (recoverable park vs terminal failure);
# standalone probes have no revival owner and must reap locally.
if claim is None:
try:
await server.shutdown()
except Exception as shutdown_exc: # noqa: BLE001 -- best-effort reap, don't mask the real error
logger.debug(
"MCP server '%s' shutdown during orphan-reap failed: %s",
name, shutdown_exc,
)
raise
finally:
if claim_token is not None:
_core._connect_server_claim.reset(claim_token)
return server
def _request_lazy_reconnect(server_name: str, server: _core.MCPServerTask) -> bool:
"""Wake a recycled stdio server and wait briefly for a fresh session."""
if not server._is_recycled_stdio():
return False
with _core._lock:
loop = _core._mcp_loop
if loop is None or not loop.is_running():
return False
def _wake() -> None:
server._ready.clear()
server._reconnect_event.set()
loop.call_soon_threadsafe(_wake)
async def _await_ready() -> bool:
deadline = time.monotonic() + _core._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(_core._run_on_mcp_loop(_await_ready, timeout=_core._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 _resolve_server_lazy(name: str, config: dict) -> bool:
"""True when ``mcp_servers.<name>.lazy`` defers connect to first tool use (default off)."""
return _core._parse_boolish(config.get("lazy", False), default=False)
def _ensure_lazy_server_connected(server_name: str) -> bool:
"""Connect a lazily-registered server on demand (sync; blocks the caller).
Honours the connect cooldown and the ``_server_connecting`` dedup set and
routes through ``_discover_and_register_server`` so park/recycle/cooldown
bookkeeping stays in one place. True when a live session exists after.
"""
with _core._lock:
server = _core._servers.get(server_name)
if server is not None and server.session is not None:
return True
config = _core._lazy_server_configs.get(server_name)
if not config:
return False
if _core._connect_cooldown_active(server_name):
return False
if server_name in _core._server_connecting:
return False
_core._server_connecting.add(server_name)
_core._server_connect_errors.pop(server_name, None)
logger.info("MCP server '%s': lazy start on first use", server_name)
_core._ensure_mcp_loop()
connect_timeout = config.get("connect_timeout", _core._DEFAULT_CONNECT_TIMEOUT)
async def _connect():
return await _core._discover_and_register_server(server_name, config)
try:
_core._run_on_mcp_loop(_connect, timeout=float(connect_timeout) + 30.0)
except BaseException as exc:
message = _core._format_connect_error(exc)
with _core._lock:
_core._server_connecting.discard(server_name)
_core._server_connect_errors[server_name] = message
_core._record_connect_failure(server_name)
logger.warning(
"Lazy MCP connect failed for '%s': %s", server_name, message,
)
return False
with _core._lock:
_core._server_connecting.discard(server_name)
_core._clear_connect_failure(server_name)
_core._lazy_server_configs.pop(server_name, None)
stale_fingerprint = _core._lazy_server_fingerprints.pop(server_name, None)
cached_names = _core._lazy_server_tool_names.pop(server_name, None) or []
server = _core._servers.get(server_name)
live_names = set(
getattr(server, "_registered_tool_names", []) or []
)
# The cached manifest may advertise tools the live server no longer
# serves; deregister those phantoms.
phantom_names = [n for n in cached_names if n not in live_names]
if phantom_names:
from tools.registry import registry
for tool_name in phantom_names:
registry.deregister(tool_name, scope=_core._server_registry_scope(server_name))
_core._forget_mcp_tool_server(tool_name)
logger.info(
"MCP server '%s': deregistered %d phantom cached tool(s) not "
"served live (stale schema-cache fingerprint %s): %s",
server_name, len(phantom_names), stale_fingerprint,
", ".join(phantom_names),
)
return server is not None and server.session is not None
def _get_connected_server_for_call(server_name: str) -> Optional[_core.MCPServerTask]:
"""Return a connected server; the single first-use connect point for lazy
servers and the wake-up point for recycled stdio ones."""
with _core._lock:
server = _core._servers.get(server_name)
is_lazy = server_name in _core._lazy_server_configs
if is_lazy and (server is None or server.session is None):
_core._ensure_lazy_server_connected(server_name)
with _core._lock:
server = _core._servers.get(server_name)
return server
if server is not None and server.session is None and server._is_recycled_stdio():
_core._request_lazy_reconnect(server_name, server)
with _core._lock:
server = _core._servers.get(server_name)
return server
async def _discover_and_register_server(name: str, config: dict) -> List[str]:
"""Connect one server, register its tools; return the registered names."""
connect_timeout = config.get("connect_timeout", _core._DEFAULT_CONNECT_TIMEOUT)
# The claim callback runs inside _connect_server while this frame is
# suspended; a list append avoids a nonlocal rebind.
claimed: List[_core.MCPServerTask] = []
def _claim_server(created: _core.MCPServerTask) -> None:
claimed.append(created)
claim_token = _core._connect_server_claim.set(_claim_server)
try:
server = await asyncio.wait_for(
_core._connect_server(name, config),
timeout=connect_timeout,
)
except BaseException:
server = claimed[0] if claimed else None
task = server._task if server is not None else None
task_cancelling = (
task.cancelling()
if task is not None and hasattr(task, "cancelling")
else 0
)
if (
server is not None
and server._error is not None
and task is not None
and not task.done()
and not task_cancelling
):
# Recoverable park: the run task stays alive to self-probe, so
# adopt it for shutdown/revival.
with _core._lock:
_core._servers[name] = server
_core._server_scope_keys[name] = _core._mcp_registry_scope()
elif server is not None:
await server.shutdown()
raise
finally:
_core._connect_server_claim.reset(claim_token)
with _core._lock:
_core._server_connecting.discard(name)
_core._server_connect_errors.pop(name, None)
_core._servers[name] = server
_core._server_scope_keys[name] = _core._mcp_registry_scope()
registered_names = _core._register_server_tools(name, server, config)
server._registered_tool_names = list(registered_names)
transport_type = "HTTP" if "url" in config else "stdio"
logger.info(
"MCP server '%s' (%s): registered %d tool(s): %s",
name, transport_type, len(registered_names),
", ".join(registered_names),
)
return registered_names
def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
"""Connect the given ``{name: config}`` servers and register their tools.
Idempotent for connected names; ``enabled: false`` servers are skipped
without disconnecting existing sessions. Returns every registered MCP
tool name.
"""
if not _core._ensure_mcp_sdk():
logger.debug("MCP SDK not available -- skipping explicit MCP registration")
return []
servers = _core._filter_suspicious_mcp_servers(servers)
if not servers:
logger.debug("No explicit MCP servers provided")
return []
# Candidates: enabled, not connected, not connecting (dedups concurrent
# discovery entry points), not lazily registered, not in backoff.
with _core._lock:
connecting = set(_core._server_connecting)
new_servers = {
k: v
for k, v in servers.items()
if k not in _core._servers
and k not in connecting
and k not in _core._lazy_server_configs
and _core._parse_boolish(v.get("enabled", True), default=True)
and not _core._connect_cooldown_active(k)
}
# Known servers without a live session are parked or mid-reconnect;
# their tools are deregistered so nothing else can nudge them.
stale_cached = [
_core._servers[k]
for k in servers
if k in _core._servers and getattr(_core._servers[k], "session", None) is None
]
_core._server_connecting.update(new_servers)
for srv_name in new_servers:
_core._server_connect_errors.pop(srv_name, None)
# Track which servers opt-in to parallel tool calls (idempotent).
for srv_name, srv_cfg in servers.items():
if _core._parse_boolish(srv_cfg.get("supports_parallel_tool_calls", False), default=False):
_core._parallel_safe_servers.add(srv_name)
else:
_core._parallel_safe_servers.discard(srv_name)
for srv in stale_cached:
_core._signal_reconnect(srv)
if not new_servers:
return _core._existing_tool_names()
# ``lazy: true`` servers with a valid schema-cache entry register from
# cache without connecting; a missing/stale entry falls back to eager.
eager_servers: Dict[str, dict] = dict(new_servers)
lazy_registered = 0
lazy_server_count = 0
try:
from tools.mcp_schema_cache import config_fingerprint, get_cached_entry
except Exception: # pragma: no cover - cache module missing
config_fingerprint = None # type: ignore[assignment]
get_cached_entry = None # type: ignore[assignment]
if config_fingerprint is not None and get_cached_entry is not None:
for name, cfg in new_servers.items():
if not _core._resolve_server_lazy(name, cfg):
continue
entry = get_cached_entry(name, config_fingerprint(cfg))
if not entry:
continue
with _core._lock:
_core._server_connecting.discard(name)
try:
names = _core._register_from_cache_sync(name, cfg, entry)
except Exception as exc:
logger.warning(
"Failed lazy MCP registration for '%s': %s", name, exc,
)
with _core._lock:
_core._server_connecting.add(name)
continue
eager_servers.pop(name, None)
lazy_registered += len(names)
lazy_server_count += 1
new_servers = eager_servers
if not new_servers:
if lazy_registered:
logger.info(
"MCP: registered %d lazy tool(s) from schema cache "
"(no processes spawned)",
lazy_registered,
)
return _core._existing_tool_names()
_core._ensure_mcp_loop()
async def _discover_all():
server_names = list(new_servers.keys())
results = await asyncio.gather(
*(_core._discover_and_register_server(name, cfg) for name, cfg in new_servers.items()),
return_exceptions=True,
)
for name, result in zip(server_names, results):
if isinstance(result, BaseException):
command = new_servers.get(name, {}).get("command")
message = _core._format_connect_error(result)
with _core._lock:
_core._server_connecting.discard(name)
_core._server_connect_errors[name] = message
_core._record_connect_failure(name)
logger.warning(
"Failed to connect to MCP server '%s'%s: %s",
name,
f" (command={command})" if command else "",
message,
)
else:
with _core._lock:
_core._server_connecting.discard(name)
_core._server_connect_errors.pop(name, None)
_core._clear_connect_failure(name)
# Clear a stale interrupt flag (executor threads are reused) so a prior
# session's interrupt cannot cancel this discovery pass.
from tools.interrupt import is_interrupted as _is_interrupted, set_interrupt as _set_interrupt
_was_interrupted = _is_interrupted()
if _was_interrupted:
_set_interrupt(False)
try:
_core._run_on_mcp_loop(_discover_all, timeout=120)
except (TimeoutError, InterruptedError) as _e:
# Entries stranded in _server_connecting would block future
# reconnect attempts.
with _core._lock:
stale = [n for n in new_servers if n in _core._server_connecting]
if stale:
logger.warning(
"MCP discovery %s while %d server(s) were still "
"connecting; clearing stale connecting set: %s",
"timed out" if isinstance(_e, TimeoutError) else "interrupted",
len(stale),
", ".join(stale),
)
_core._server_connecting.difference_update(stale)
for _sn in stale:
_core._server_connect_errors.setdefault(
_sn,
f"Connection attempt {'timed out' if isinstance(_e, TimeoutError) else 'interrupted'} during discovery",
)
raise
finally:
if _was_interrupted:
_set_interrupt(True)
with _core._lock:
connected = [
n
for n in new_servers
if n in _core._servers and n not in _core._server_connect_errors
]
new_tool_count = sum(
len(getattr(_core._servers[n], "_registered_tool_names", []))
for n in connected
)
failed = len(new_servers) - len(connected)
new_tool_count += lazy_registered
connected_count = len(connected) + lazy_server_count
if new_tool_count or failed:
summary = f"MCP: registered {new_tool_count} tool(s) from {connected_count} server(s)"
if failed:
summary += f" ({failed} failed)"
logger.info(summary)
return _core._existing_tool_names()
def discover_mcp_tools() -> List[str]:
"""Entry point: load config, connect servers, register tools.
Safe without the ``mcp`` package (returns []). Idempotent: only servers
missing from a previous call are retried. Returns all MCP tool names.
"""
servers = _core._load_mcp_config()
if not servers:
logger.debug("No MCP servers configured")
return []
# SDK import deferred to here so a config without servers never pays it.
if not _core._ensure_mcp_sdk():
logger.debug("MCP SDK not available -- skipping MCP tool discovery")
return []
# Cross-process guard: a lock loser waits for the holder, then runs its
# own discovery; if locking is unavailable or the wait expires, run
# unguarded (fail-soft).
cookie = _core._try_acquire_mcp_discovery_lock()
if cookie is None:
logger.debug(
"Another process holds MCP discovery lock -- retrying with backoff"
)
for _ in range(_core._MCP_DISCOVERY_LOCK_MAX_RETRIES):
time.sleep(_core._MCP_DISCOVERY_LOCK_RETRY_DELAY_S)
cookie = _core._try_acquire_mcp_discovery_lock()
if cookie is not None:
break
if cookie is None:
logger.warning(
"MCP discovery lock still held after %d retries -- "
"running discovery unguarded",
_core._MCP_DISCOVERY_LOCK_MAX_RETRIES,
)
elif cookie is not _core._LOCK_UNAVAILABLE:
logger.debug("Retry succeeded -- acquired MCP discovery lock")
try:
with _core._lock:
connecting = set(_core._server_connecting)
new_server_names = [
name
for name, cfg in servers.items()
if name not in _core._servers
and name not in connecting
and _core._parse_boolish(cfg.get("enabled", True), default=True)
]
tool_names = _core.register_mcp_servers(servers)
if not new_server_names:
return tool_names
with _core._lock:
connected_server_names = [
name
for name in new_server_names
if name in _core._servers and name not in _core._server_connect_errors
]
new_tool_count = sum(
len(getattr(_core._servers[name], "_registered_tool_names", []))
for name in connected_server_names
)
failed_count = len(new_server_names) - len(connected_server_names)
if new_tool_count or failed_count:
summary = f" MCP: {new_tool_count} tool(s) from {len(connected_server_names)} server(s)"
if failed_count:
summary += f" ({failed_count} failed)"
logger.info(summary)
return tool_names
finally:
if cookie not in (None, _core._LOCK_UNAVAILABLE):
cookie.release()
def is_mcp_tool_parallel_safe(tool_name: str) -> bool:
"""True when the tool's server opted into ``supports_parallel_tool_calls``.
Uses the provenance captured at registration, never the (ambiguous)
``mcp__{server}__{tool}`` string shape.
"""
if not tool_name.startswith(_core.MCP_TOOL_NAME_PREFIX):
return False
with _core._lock:
server_name = _core._mcp_tool_server_names.get(tool_name)
return bool(server_name and server_name in _core._parallel_safe_servers)
def get_mcp_status() -> List[dict]:
"""Status of every configured server for banner/TUI display.
Each dict has name, transport, tools, connected, disabled, status (one of
connected / disabled / connecting / failed / configured) and, for failed,
error. ``enabled: false`` is reported as disabled, not failed.
"""
configured = _core._load_mcp_config()
if not configured:
return []
with _core._lock:
active_servers = dict(_core._servers)
connecting = set(_core._server_connecting)
connect_errors = dict(_core._server_connect_errors)
def _entry(name: str, transport: str, status: str, **extra) -> dict:
return {
"name": name,
"transport": transport,
"tools": 0,
"connected": False,
"disabled": status == "disabled",
"status": status,
**extra,
}
result: List[dict] = []
for name, cfg in configured.items():
transport = cfg.get("transport", "http") if "url" in cfg else "stdio"
enabled = _core._parse_boolish(cfg.get("enabled", True), default=True)
server = active_servers.get(name)
if server and server.session is not None:
entry = _entry(name, transport, "connected", connected=True)
entry["tools"] = (
len(server._registered_tool_names)
if hasattr(server, "_registered_tool_names")
else len(server._tools)
)
if server._sampling:
entry["sampling"] = dict(server._sampling.metrics)
elif not enabled:
entry = _entry(name, transport, "disabled")
elif name in connecting:
entry = _entry(name, transport, "connecting")
elif name in connect_errors:
entry = _entry(name, transport, "failed", error=connect_errors[name])
else:
entry = _entry(name, transport, "configured")
result.append(entry)
return result
def probe_mcp_server_tools() -> Dict[str, List[tuple]]:
"""Connect to each enabled server, list ``(tool_name, description)`` and
disconnect, without registering anything. Failed servers are omitted."""
if not _core._ensure_mcp_sdk():
return {}
servers_config = _core._load_mcp_config()
if not servers_config:
return {}
enabled = {
k: v for k, v in servers_config.items()
if _core._parse_boolish(v.get("enabled", True), default=True)
}
if not enabled:
return {}
_core._ensure_mcp_loop()
result: Dict[str, List[tuple]] = {}
probed_servers: List[_core.MCPServerTask] = []
async def _probe_all():
names = list(enabled.keys())
coros = []
for name, cfg in enabled.items():
ct = cfg.get("connect_timeout", _core._DEFAULT_CONNECT_TIMEOUT)
coros.append(asyncio.wait_for(_core._connect_server(name, cfg), timeout=ct))
outcomes = await asyncio.gather(*coros, return_exceptions=True)
for name, outcome in zip(names, outcomes):
if isinstance(outcome, Exception):
logger.debug("Probe: failed to connect to '%s': %s", name, outcome)
continue
probed_servers.append(outcome)
tools = []
for t in outcome._tools:
desc = getattr(t, "description", "") or ""
tools.append((t.name, desc))
result[name] = tools
await asyncio.gather(
*(s.shutdown() for s in probed_servers),
return_exceptions=True,
)
try:
_core._run_on_mcp_loop(_probe_all, timeout=120)
except Exception as exc:
logger.debug("MCP probe failed: %s", exc)
finally:
_core._stop_mcp_loop_if_idle()
return result
def has_registered_mcp_tools() -> bool:
"""True if any MCP server has registered tools (cheap; no registry walk).
Checks registered TOOLS, not connected servers, so the per-turn refresh
hook stays idle for zero-tool servers.
"""
with _core._lock:
return bool(_core._mcp_tool_server_names)
def get_registered_mcp_server_names() -> set:
"""Server names that registered at least one tool (the live, filtered
signal — not merely what config.yaml lists)."""
with _core._lock:
return set(_core._mcp_tool_server_names.values())