ouyangbo/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 9 Packages Projects Releases Wiki Activity
Files
5e4f30bc9d47022b93e8f8eb7729a60dc17f0496
hermes-agent/tools/mcp_tool.py
T

577 lines
29 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""MCP (Model Context Protocol) client: connects to the ``mcp_servers`` configured in
~/.hermes/config.yaml over stdio, Streamable HTTP or SSE, discovers their tools and
registers them into the hermes tool registry. The ``mcp`` package is optional; without
it this module is a no-op.
Architecture: one background event loop (``_mcp_loop``) in a daemon thread; each server
is a long-lived Task on it (``MCPServerTask``) so the transport's anyio cancel scopes are
entered and exited in the same Task. Tool calls are scheduled onto the loop via
``run_coroutine_threadsafe``; every mutation of ``_servers`` / loop handles holds ``_lock``.
This module keeps the SDK loader, the ``MCPServerTask`` shell and all shared module state;
the ``mcp_tool_*`` siblings read it back through ``tools.mcp_tool`` at call time (never by
value) and every one of their names is re-exported here so ``from tools.mcp_tool import X``
and ``mock.patch("tools.mcp_tool.X")`` keep working.
"""
import asyncio
import contextvars
import importlib
import importlib.util
import inspect
import logging
import os # noqa: F401 — tests patch ``tools.mcp_tool.os.*``
import shutil # noqa: F401 — tests patch ``tools.mcp_tool.shutil.which``
import threading
import time
from typing import Any, Callable, Dict, List, Optional, Set
logger = logging.getLogger(__name__)
from tools.mcp_tool_common import ( # noqa: F401
_DEFAULT_TOOL_TIMEOUT, _env_ref_name, _exc_str, _get_lifecycle_seconds, _jittered,
_parse_boolish, _resolve_tool_timeout, _safe_numeric, _sanitize_error, mcp_field,
)
from tools.mcp_tool_schema import ( # noqa: F401
MCP_TOOL_NAME_PREFIX, _build_utility_schemas, _convert_mcp_schema, _normalize_mcp_input_schema,
_scan_mcp_description, matches_name_filter, mcp_prefixed_tool_name, sanitize_mcp_name_component,
)
from tools.mcp_tool_content import ( # noqa: F401
_MCP_HARD_RESULT_CAP_CHARS, _MCP_RESOURCE_MAX_B64_CHARS, _MCP_RESOURCE_MAX_BYTES,
_cache_mcp_audio_block, _cache_mcp_image_block, _is_reserved_mcp_meta_key,
_mcp_image_extension_for_mime_type, _mcp_resource_filename, _render_mcp_resource_block,
_truncate_mcp_text_result,
)
from tools.mcp_tool_errors import ( # noqa: F401
InvalidMcpUrlError, NonMcpEndpointError, _EXC_TRAVERSAL_MAX_NODES,
_JSONRPC_UNSUPPORTED_PROTOCOL_VERSION, _classify_mcp_failure, _format_connect_error,
_handshake_rejected_as_modern, _is_auth_error, _is_method_not_found_error,
_is_session_expired_error, _make_redirect_header_stripper, _resolve_client_cert,
_resolve_identity_header, _unwrap_exception_group, _validate_remote_mcp_url,
)
from tools.mcp_tool_config import ( # noqa: F401
_ENV_VAR_PATTERN, _build_safe_env, _filter_suspicious_mcp_servers, _get_mcp_stderr_log,
_interpolate_env_vars, _load_mcp_config, _resolve_stdio_command, _warn_hidden_whitespace,
_whitespace_warned, _workspace_folder, _wrap_command_with_watchdog, _write_stderr_log_header,
)
from tools.mcp_tool_sampling import ( # noqa: F401
ElicitationHandler, SamplingHandler, _format_elicitation_schema_summary,
)
from tools.mcp_tool_handlers import ( # noqa: F401
_handle_auth_error_and_retry, _handle_session_expired_and_retry, _make_check_fn,
_make_get_prompt_handler, _make_list_prompts_handler, _make_list_resources_handler,
_make_read_resource_handler, _make_tool_handler,
)
from tools.mcp_tool_registration import ( # noqa: F401
_annotation_read_only_hint, _existing_tool_names, _forget_mcp_tool_server,
_normalize_server_trust, _register_from_cache_sync, _register_server_tools,
_select_utility_schemas, _track_mcp_tool_server,
)
from tools.mcp_tool_lifecycle import ( # noqa: F401
_drain_and_stop_mcp_loop, _drain_mcp_loop_tasks, _filter_mcp_children,
_kill_orphaned_mcp_children, _orphan_stdio_pid_servers, _orphan_stdio_pids,
_snapshot_child_pids, _stdio_pgids, _stdio_pids, _stop_mcp_loop_if_idle, shutdown_mcp_servers,
)
from tools.mcp_tool_agent import ( # noqa: F401
_reinject_post_build_tools, persist_agent_tool_names, refresh_agent_mcp_tools,
reprobe_tool_availability, 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
from tools.mcp_tool_loop import ( # noqa: F401
_running_loop, _LockCookie, _acquire_lock_on_fh, _try_acquire_mcp_discovery_lock,
_mcp_loop_exception_handler, _wrap_with_home_override, _wrap_with_dashboard_oauth_flow,
_run_on_mcp_loop, _signal_reconnect, reconnect_mcp_server, _wait_for_server_session_ready,
_signal_reconnect_and_wait, _ensure_mcp_loop, _stop_mcp_loop,
)
from tools.mcp_tool_discovery import ( # noqa: F401
_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 before a
# stdio spawn. Kept just ABOVE osv_check._TIMEOUT (10s) so the inner socket timeout
# normally fires first; this only bites when a stalled SSL handshake defeats it.
_OSV_MALWARE_CHECK_TIMEOUT_S = 12.0
# ---------------------------------------------------------------------------
# Optional MCP SDK: availability probe now, symbol import on first use
# ---------------------------------------------------------------------------
_MCP_AVAILABLE = False
_MCP_HTTP_AVAILABLE = False
_MCP_NEW_HTTP = False
_MCP_LEGACY_HTTP = False
_MCP_SAMPLING_TYPES = False
_MCP_NOTIFICATION_TYPES = False
_MCP_ELICITATION_TYPES = False
_MCP_MESSAGE_HANDLER_SUPPORTED = False
_MCP_LOGGING_CALLBACK_SUPPORTED = False
sse_client = None
# Fallback for SDKs that don't export LATEST_PROTOCOL_VERSION (Streamable HTTP arrived
# with 2025-03-26, so this stays valid for the HTTP path).
LATEST_PROTOCOL_VERSION = "2025-03-26"
# Newest revision ``ClientSession.initialize()`` actually speaks. From 2026-07-28 the
# handshake is replaced by a per-request envelope, so this can be OLDER than
# LATEST_PROTOCOL_VERSION; the MCP-Protocol-Version header must be seeded from this one.
LATEST_HANDSHAKE_VERSION = LATEST_PROTOCOL_VERSION
# Importing ``mcp`` costs ~260ms, so it is deferred to first real use (_ensure_mcp_sdk).
# Availability is decided here with a metadata-only find_spec probe so every
# ``if not _MCP_AVAILABLE`` gate / test patch / skipif keeps its exact semantics.
try:
_MCP_AVAILABLE = importlib.util.find_spec("mcp") is not None
except Exception:
_MCP_AVAILABLE = False
if not _MCP_AVAILABLE:
logger.debug("mcp package not installed -- MCP tool support disabled")
ClientSession: Any = None
_MCP_SDK_IMPORT_ATTEMPTED = False
_MCP_SDK_IMPORT_LOCK = threading.Lock()
# SDK symbols bound by _ensure_mcp_sdk(). Module __getattr__ (PEP 562) imports the SDK on
# first external access, so mock.patch("tools.mcp_tool.stdio_client") sees a real original
# and the mock is never clobbered (_ensure is idempotent).
_MCP_SDK_LAZY_SYMBOLS = frozenset({
"StdioServerParameters", "stdio_client", "streamablehttp_client", "streamable_http_client",
"CreateMessageResult", "CreateMessageResultWithTools", "ErrorData", "SamplingCapability",
"SamplingToolsCapability", "TextContent", "ToolUseContent", "ElicitRequestParams",
"ElicitResult", "ServerNotification", "ToolListChangedNotification",
"PromptListChangedNotification", "ResourceListChangedNotification",
})
# Optional SDK type families: (module, names, debug message when absent). Each is gated
# separately so an older SDK only loses that feature, not MCP.
_SAMPLING_TYPE_NAMES = ("CreateMessageResult", "CreateMessageResultWithTools", "ErrorData",
"SamplingCapability", "SamplingToolsCapability", "TextContent", "ToolUseContent")
_NOTIFICATION_TYPE_NAMES = ("ServerNotification", "ToolListChangedNotification",
"PromptListChangedNotification", "ResourceListChangedNotification")
def __getattr__(name: str):
if name in _MCP_SDK_LAZY_SYMBOLS:
_ensure_mcp_sdk()
try:
return globals()[name]
except KeyError:
pass # SDK missing or symbol absent on this SDK build
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
def _import_sdk_names(module: str, names: tuple, missing_msg: Optional[str] = None) -> bool:
"""Bind ``names`` from SDK ``module`` into this module's globals; False (nothing bound,
optional debug line) when this SDK build lacks the module or any of the names."""
try:
mod = importlib.import_module(module)
values = {n: getattr(mod, n) for n in names}
except (ImportError, AttributeError):
if missing_msg:
logger.debug(missing_msg)
return False
globals().update(values)
return True
def _ensure_mcp_sdk() -> bool:
"""Import the optional ``mcp`` SDK on first use; return availability.
Idempotent and thread-safe. Honors a test-patched ``_MCP_AVAILABLE=False`` (no import)
and pre-installed mock symbols (``ClientSession`` already set means no re-import, so
mocks are never clobbered).
"""
global _MCP_SDK_IMPORT_ATTEMPTED, _MCP_AVAILABLE, _MCP_HTTP_AVAILABLE
global _MCP_SAMPLING_TYPES, _MCP_NOTIFICATION_TYPES, _MCP_ELICITATION_TYPES
global _MCP_MESSAGE_HANDLER_SUPPORTED, _MCP_LOGGING_CALLBACK_SUPPORTED
global _MCP_NEW_HTTP, _MCP_LEGACY_HTTP, LATEST_HANDSHAKE_VERSION, sse_client
global _JSONRPC_METHOD_NOT_FOUND
if not _MCP_AVAILABLE:
return False
if _MCP_SDK_IMPORT_ATTEMPTED or ClientSession is not None:
return _MCP_AVAILABLE
with _MCP_SDK_IMPORT_LOCK:
if _MCP_SDK_IMPORT_ATTEMPTED or ClientSession is not None:
return _MCP_AVAILABLE
if (_import_sdk_names("mcp", ("ClientSession", "StdioServerParameters"))
and _import_sdk_names("mcp.client.stdio", ("stdio_client",))):
_MCP_AVAILABLE = True
# mcp >= 1.24 ships streamable_http_client; 2.0 dropped the deprecated
# streamablehttp_client alias. Either one gives HTTP.
_MCP_NEW_HTTP = _import_sdk_names("mcp.client.streamable_http", ("streamable_http_client",))
_MCP_LEGACY_HTTP = _import_sdk_names("mcp.client.streamable_http", ("streamablehttp_client",))
_MCP_HTTP_AVAILABLE = _MCP_NEW_HTTP or _MCP_LEGACY_HTTP
_import_sdk_names("mcp.types", ("LATEST_PROTOCOL_VERSION",),
"mcp.types.LATEST_PROTOCOL_VERSION not available -- using fallback protocol version")
if not _import_sdk_names("mcp.client.session", ("LATEST_HANDSHAKE_VERSION",)):
# Pre-2.x SDKs: newest revision IS the handshake revision.
LATEST_HANDSHAKE_VERSION = LATEST_PROTOCOL_VERSION
if not _import_sdk_names("mcp.client.sse", ("sse_client",),
"mcp.client.sse.sse_client not available -- SSE transport disabled"):
sse_client = None
_MCP_SAMPLING_TYPES = _import_sdk_names(
"mcp.types", _SAMPLING_TYPE_NAMES, "MCP sampling types not available -- sampling disabled")
_MCP_ELICITATION_TYPES = _import_sdk_names(
"mcp.types", ("ElicitRequestParams", "ElicitResult"),
"MCP elicitation types not available -- elicitation disabled")
_MCP_NOTIFICATION_TYPES = _import_sdk_names(
"mcp.types", _NOTIFICATION_TYPE_NAMES,
"MCP notification types not available -- dynamic tool discovery disabled")
else:
logger.debug("mcp package not installed -- MCP tool support disabled")
if _MCP_AVAILABLE:
try:
_JSONRPC_METHOD_NOT_FOUND = importlib.import_module("mcp.types").METHOD_NOT_FOUND
except Exception: # pragma: no cover — SDK without the constant
pass
_MCP_MESSAGE_HANDLER_SUPPORTED = _client_session_accepts("message_handler")
if _MCP_AVAILABLE and not _MCP_MESSAGE_HANDLER_SUPPORTED:
logger.debug("MCP SDK does not support message_handler -- dynamic tool discovery disabled")
_MCP_LOGGING_CALLBACK_SUPPORTED = _client_session_accepts("logging_callback")
_MCP_SDK_IMPORT_ATTEMPTED = True
return _MCP_AVAILABLE
_SDK_HTTPX_MOD = None
def sdk_httpx():
"""Return the httpx module the *installed* MCP SDK is built against.
mcp 2.0 moved to ``httpx2`` (same API, separate distribution). Every object crossing the
SDK boundary (the ``AsyncClient`` passed to the transport, OAuth ``Request`` objects, the
exception classes) must come from the module the SDK itself imports, or it fails at the
transport layer. Resolved from the SDK's transport module, not a version number; falls
back to the newest module present. ``None`` only when neither module is importable.
"""
global _SDK_HTTPX_MOD
if _SDK_HTTPX_MOD is not None:
return _SDK_HTTPX_MOD
try:
from mcp.client import streamable_http as _transport
_SDK_HTTPX_MOD = getattr(_transport, "httpx2", None) or getattr(_transport, "httpx", None)
except ImportError:
_SDK_HTTPX_MOD = None
if _SDK_HTTPX_MOD is None:
try:
import httpx2 as _fallback
except ImportError:
try:
import httpx as _fallback # type: ignore[no-redef]
except ImportError:
return None
_SDK_HTTPX_MOD = _fallback
return _SDK_HTTPX_MOD
def _client_session_accepts(kwarg: str) -> bool:
"""Whether this SDK's ``ClientSession.__init__`` takes ``kwarg`` (older SDKs lack
``message_handler`` and ``logging_callback``)."""
if not _MCP_AVAILABLE:
return False
try:
return kwarg in inspect.signature(ClientSession).parameters
except (TypeError, ValueError):
return False
# MCP logging levels (RFC 5424 syslog severities) -> Python logging levels.
_MCP_LOG_LEVEL_MAP = {
"debug": logging.DEBUG, "info": logging.INFO, "notice": logging.INFO,
"warning": logging.WARNING, "error": logging.ERROR, "critical": logging.ERROR,
"alert": logging.ERROR, "emergency": logging.ERROR,
}
# ---------------------------------------------------------------------------
# Reconnect / keepalive tuning
# ---------------------------------------------------------------------------
_DEFAULT_CONNECT_TIMEOUT = 60 # seconds for initial connection per server
_MAX_RECONNECT_RETRIES = 5
_MAX_INITIAL_CONNECT_RETRIES = 3 # retries for the very first connection attempt
_MAX_BACKOFF_SECONDS = 60
# Parked servers (budget exhausted, tools deregistered) self-probe on this cadence: with
# no tools registered nothing else can ever revive them.
_PARKED_RETRY_INTERVAL = 300
_RECYCLED_RECONNECT_TIMEOUT = 15.0
# Bounded wait for a respawned stdio child when a call finds it dead (gateway restarts kill
# every MCP child); bounded so a broken server still parks via run()'s rapid-drop budget.
_STDIO_RESPAWN_WAIT_SEC = 15.0
# Servers may expire idle sessions on any TTL, so the client MUST ping faster than that TTL;
# short-TTL servers (~15s) need a smaller configured ``keepalive_interval``. The floor stops
# a tiny interval from busy-looping.
_DEFAULT_KEEPALIVE_INTERVAL = 180
_MIN_KEEPALIVE_INTERVAL = 5
# One bounded cancellation cycle for pending loop tasks at final shutdown, so
# cancellation-resistant tasks cannot hang process exit.
_MCP_LOOP_DRAIN_TIMEOUT = 3.0
# JSON-RPC 2.0 "method not found" (e.g. a server without the optional ``ping``).
# _ensure_mcp_sdk() overrides it from mcp.types once the SDK is loaded.
_JSONRPC_METHOD_NOT_FOUND = -32601
# Cap on nextCursor pagination so a server returning a cursor forever cannot spin
# discovery; 50 pages at 50-100 items/page covers thousands of entries.
_MCP_LIST_MAX_PAGES = 50
async def _paginate_full_list(list_method, items_attr: str, server_name: str,
cache_meta_out: Optional[dict] = None):
"""Drain a paginated ``list_*`` call by following ``nextCursor`` (the SDK fetches one
page per call). ``cache_meta_out`` receives the first page's SEP-2549 hints (``ttl_ms``,
``cache_scope``). Callers must hold the server's ``_rpc_lock`` for a consistent snapshot."""
items: list = []
cursor = None
for _ in range(_MCP_LIST_MAX_PAGES):
if not cursor:
result = await list_method()
else:
# mcp 2.0 takes params=PaginatedRequestParams, 1.x takes cursor=.
try:
_params_cls = getattr(_mcp_types(), "PaginatedRequestParams", None)
if _params_cls is not None:
result = await list_method(params=_params_cls(cursor=cursor))
else:
result = await list_method(cursor=cursor)
except TypeError:
result = await list_method(cursor=cursor)
if cache_meta_out is not None and not items:
_ttl = mcp_field(result, "ttl_ms", "ttlMs")
_scope = mcp_field(result, "cache_scope", "cacheScope")
if _ttl is not None:
cache_meta_out["ttl_ms"] = _ttl
if _scope is not None:
cache_meta_out["cache_scope"] = _scope
items.extend(getattr(result, items_attr, None) or [])
cursor = mcp_field(result, "next_cursor", "nextCursor")
# Cursor is an opaque string; anything else (incl. mocks) = last page.
if not isinstance(cursor, str) or not cursor:
break
else:
logger.warning("MCP server '%s': %s pagination exceeded %d pages; truncating at %d items",
server_name, items_attr, _MCP_LIST_MAX_PAGES, len(items))
return items
def _mcp_types():
"""Late import of ``mcp.types`` (module keeps the SDK import lazy)."""
import mcp.types as _t
return _t
# ---------------------------------------------------------------------------
# Server task -- each MCP server lives in one long-lived asyncio Task
# ---------------------------------------------------------------------------
class MCPServerTask(MCPServerRunMixin, MCPServerTransportMixin, MCPServerHealthMixin):
"""One MCP server connection living in one long-lived asyncio Task, so the transport's
anyio cancel scopes are entered and exited in the same Task. Run state machine in
``MCPServerRunMixin``, transport bring-up in ``MCPServerTransportMixin``, keepalive /
refresh / liveness in ``MCPServerHealthMixin``."""
__slots__ = (
"name", "session", "tool_timeout", "_task", "_ready", "_shutdown_event", "_reconnect_event",
"_tools", "_error", "_config", "_sampling", "_elicitation", "_registered_tool_names",
"_auth_type", "_refresh_lock", "_rpc_lock", "_pending_refresh_tasks", "_pending_call_context",
"_lifecycle_started_at", "_last_tool_call_at", "_idle_timeout_seconds", "_max_lifetime_seconds",
"_recycled_reason", "initialize_result", "_ping_unsupported", "_list_cache_meta",
"_reconnect_retries", "_session_proven", "_was_parked", "_inflight_tasks", "_reconnecting",
"_suspect_reason", "_teardown_race", "_permanent_grace_used", "_stdio_child_pids",
"_ever_connected",
)
def __init__(self, name: str):
self.name = name
self.session: Optional[Any] = None
self.tool_timeout: float = _DEFAULT_TOOL_TIMEOUT
self._task: Optional[asyncio.Task] = None
self._ready = asyncio.Event()
self._shutdown_event = asyncio.Event()
# When set, _run_http/_run_stdio exit their async-with cleanly and run() re-enters
# the transport (auth recovery, manual refresh, ...).
self._reconnect_event = asyncio.Event()
self._tools: list = []
self._error: Optional[Exception] = None
self._config: dict = {}
self._sampling: Optional[SamplingHandler] = None
self._elicitation: Optional[ElicitationHandler] = None
self._registered_tool_names: list[str] = []
self._reconnect_retries: int = 0
# Rapid-drop budget: a (re)established session is UNPROVEN until it survives a full
# keepalive interval or serves a successful call. Only a proven session clears the
# reconnect budget, so a transport that flaps right after the handshake still parks.
self._session_proven: bool = False
# Set once tools were ever registered, never cleared (unlike _ready, which clears every
# reconnect cycle): separates first-connect from reconnect failures in run()'s ladders.
self._ever_connected: bool = False
# True from park until the session proves healthy again; logs the revival once.
self._was_parked: bool = False
# In-flight RPC tasks, so a reconnect/shutdown teardown can fail them fast instead of
# orphaning them on a dying transport.
self._inflight_tasks: set = set()
# True while a deliberate teardown fails in-flight calls; lets _track_inflight_rpc
# turn the cancel into a retryable error.
self._reconnecting: bool = False
# Latched by races (teardown-vs-keepalive, auth-lock corruption); verified lazily by
# ensure_healthy() before the next call.
self._suspect_reason: Optional[str] = None
# A teardown that failed >=1 in-flight call makes the next reconnect a RACE RECOVERY:
# it must not charge the rapid-drop budget.
self._teardown_race: bool = False
# One-time grace: an auth/permanent-classified failure on a previously PROVEN session
# gets one suspect+reconnect cycle before the park ladder applies.
self._permanent_grace_used: bool = False
# Children of the current stdio transport: in-flight calls fail FAST when the child
# dies instead of riding out the tool timeout.
self._stdio_child_pids: Set[int] = set()
self._auth_type: str = ""
self._refresh_lock = asyncio.Lock()
# A stdio session is one JSON-RPC stream: a list_tools issued by the notification
# handler while a tool call is in flight can wedge it. Serialize client-initiated
# RPCs per server (HTTP too, for ordering).
self._rpc_lock = asyncio.Lock()
self._pending_refresh_tasks: set[asyncio.Task] = set()
# contextvars snapshot of the agent task inside session.call_tool(). The SDK dispatches
# elicitation/create on a separate task that does not inherit HERMES_SESSION_PLATFORM;
# replaying this context in the elicitation callback routes the prompt correctly.
self._pending_call_context: Optional[contextvars.Context] = None
now = time.monotonic()
self._lifecycle_started_at: float = now
self._last_tool_call_at: float = now
self._idle_timeout_seconds: Optional[float] = None
self._max_lifetime_seconds: Optional[float] = None
self._recycled_reason: Optional[str] = None
# InitializeResult from the handshake: the server's REAL advertised capabilities.
self.initialize_result: Optional[Any] = None
# SEP-2549 cache hints from the last tools/list (ttl_ms, cache_scope).
self._list_cache_meta: dict = {}
# Latched when keepalive ``ping`` returns -32601 (optional utility not implemented);
# later keepalives use list_tools instead. Reset on every fresh transport connection.
self._ping_unsupported: bool = False
# Content types a real Streamable-HTTP endpoint may return on the initial POST/GET;
# anything else on a 2xx means the URL is not an MCP endpoint.
_MCP_CONTENT_TYPES = ("application/json", "text/event-stream")
# ---------------------------------------------------------------------------
# Module-level state (every mutation under ``_lock``)
# ---------------------------------------------------------------------------
_servers: Dict[str, MCPServerTask] = {}
# Profile registry scope owning each live connection (None outside multiplex): a
# multiplexed /reload-mcp tears down only its own profile's servers.
_server_scope_keys: Dict[str, Optional[str]] = {}
_server_connecting: set[str] = set()
_server_connect_errors: Dict[str, str] = {}
# Lazy startup: servers registered from the on-disk schema cache without connecting;
# popped once a real connection is established on first use.
_lazy_server_configs: Dict[str, dict] = {}
_lazy_server_fingerprints: Dict[str, str] = {}
_lazy_server_tool_names: Dict[str, List[str]] = {}
# Discovery installs a task-local claim around ``_connect_server`` so it can retain a
# recoverable parked task without standalone probe calls publishing failed servers into
# module-global ownership.
_connect_server_claim: contextvars.ContextVar[Optional[Callable[[MCPServerTask], None]]] = (
contextvars.ContextVar("mcp_connect_server_claim", default=None))
# Per-server connect cooldown. A server that fails to spawn never reaches ``_servers``, so
# without this every ``discover_mcp_tools()`` (one per worker session) would respawn it from
# scratch — a restart storm whose unreaped subprocesses destabilise the healthy co-located
# servers. Failed attempts stamp an exponential-backoff deadline that
# ``register_mcp_servers`` honours; a successful connection clears it.
_server_connect_retry_after: Dict[str, float] = {} # name -> monotonic deadline
_server_connect_failures: Dict[str, int] = {} # name -> consecutive failures
_CONNECT_RETRY_BASE_BACKOFF_SEC = 30.0
_CONNECT_RETRY_MAX_BACKOFF_SEC = 600.0
# 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 only via _bump_server_error / _reset_server_error.
_server_error_counts: Dict[str, int] = {}
_server_breaker_opened_at: Dict[str, float] = {}
_CIRCUIT_BREAKER_THRESHOLD = 3
_CIRCUIT_BREAKER_COOLDOWN_SEC = 60.0
# Trust-tier gating (``mcp_servers.<name>.trust: full | untrusted``). On an untrusted server
# every write-capable call needs user approval before the RPC fires; a tool is write-capable
# unless its discovery-time ``annotations.readOnlyHint`` is exactly True (malformed fails
# closed). readOnlyHint is a server-supplied HINT and a hostile server can lie, but on an
# untrusted server a lie can only skip approval for calls the operator was already warned
# about — never widen access. Missing ``trust`` defaults to full (backward compatible); any
# unrecognized value normalizes to untrusted (a typo must never disable the gate).
# Classified at CALL time from DISCOVERY data: no schema mutation, prompt cache intact.
_server_trust_levels: Dict[str, str] = {}
_tool_read_only_hints: Dict[str, Dict[str, bool]] = {}
_TRUST_FULL = "full"
_TRUST_UNTRUSTED = "untrusted"
def _bump_server_error(server_name: str) -> None:
"""Count a failure; at the threshold (re)stamp the breaker-open time."""
n = _server_error_counts.get(server_name, 0) + 1
_server_error_counts[server_name] = n
if n >= _CIRCUIT_BREAKER_THRESHOLD:
_server_breaker_opened_at[server_name] = time.monotonic()
def _reset_server_error(server_name: str) -> None:
"""Close the breaker on any unambiguous success signal."""
_server_error_counts[server_name] = 0
_server_breaker_opened_at.pop(server_name, None)
# Raw server names opted into parallel tool calls. Raw identity matters: ``foo-bar`` and
# ``foo_bar`` both sanitize to ``foo_bar`` but must not share policy.
_parallel_safe_servers: set = set()
# registry tool name -> raw server name, captured at registration. The generated name is
# lossy (punctuation -> ``_``), so never re-parse it.
_mcp_tool_server_names: Dict[str, str] = {}
# Dedicated event loop running in a background daemon thread.
_mcp_loop: Optional[asyncio.AbstractEventLoop] = None
_mcp_thread: Optional[threading.Thread] = None
# Guards the loop handles, _servers, the status maps and the PID ledgers.
_lock = threading.Lock()
def _mcp_registry_scope() -> Optional[str]:
"""Registry scope for MCP registrations: under a profile multiplexer each profile's MCP
tools live in its own registry overlay; single-profile processes stay global (None)."""
from agent.secret_scope import is_multiplex_active
if not is_multiplex_active():
return None
from tools.registry import registry
return registry.current_scope_key()
def _server_registry_scope(name: str) -> Optional[str]:
"""Scope owning server *name*'s tools: recorded at connect, else current. Teardown runs
on the MCP loop without the discovering profile's context, so the scope captured at
adoption into ``_servers`` is authoritative."""
if name in _server_scope_keys:
return _server_scope_keys[name]
return _mcp_registry_scope()
# Cross-process MCP discovery guard: advisory file lock so gateway + CLI + TUI don't all
# run discovery at once.
_LOCK_UNAVAILABLE: Any = object() # sentinel: locking broken/unavailable
_MCP_DISCOVERY_LOCK_PATH: Optional[str] = None # resolved lazily
# Bounded wait when another process holds the lock.
_MCP_DISCOVERY_LOCK_MAX_RETRIES: int = 240
_MCP_DISCOVERY_LOCK_RETRY_DELAY_S: float = 0.5
Reference in New Issue View Git Blame Copy Permalink