From 0a38a611ebe0722e06da33d86f3ce88433cac13d Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:57:52 -0700 Subject: [PATCH 1/2] =?UTF-8?q?refactor(mcp):=20handlers=20=E2=80=94=20one?= =?UTF-8?q?=20=5Fstrike=20helper=20for=20breaker-strike+tool=5Ferror=20tai?= =?UTF-8?q?ls,=20reconnectable-server=20lookup=20owns=20the=20loop=20check?= =?UTF-8?q?,=20docstrings=20compacted=20by=20hand=20(783=20->=20602)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tools/mcp_tool_handlers.py | 537 ++++++++++++------------------------- 1 file changed, 178 insertions(+), 359 deletions(-) diff --git a/tools/mcp_tool_handlers.py b/tools/mcp_tool_handlers.py index d25014ffa6..9b121ae3ab 100644 --- a/tools/mcp_tool_handlers.py +++ b/tools/mcp_tool_handlers.py @@ -22,63 +22,36 @@ logger = logging.getLogger("tools.mcp_tool") def _trust_gate_check(server_name: str, tool_name: str) -> Optional[str]: """Approval gate for write-capable tools on ``trust: untrusted`` servers. - - Returns None to proceed, or a ``tool_error`` string when blocked. - Fail-closed: approval-system errors block the call. - """ + None to proceed, else a ``tool_error``. Fail-closed: approval-system errors block.""" trust = _core._server_trust_levels.get(server_name, _core._TRUST_FULL) - if trust != _core._TRUST_UNTRUSTED: + if trust != _core._TRUST_UNTRUSTED or _core._tool_read_only_hints.get(server_name, {}).get(tool_name) is True: return None - if _core._tool_read_only_hints.get(server_name, {}).get(tool_name) is True: - return None - - # Lazy import: tools.approval routes the prompt to whichever surface owns - # the session (CLI, TUI, Telegram, ...) and normalizes the answer. + # Lazy import: tools.approval routes the prompt to whichever surface owns the session. try: from tools.approval import request_elicitation_consent answer = request_elicitation_consent( - ( - f"MCP tool '{tool_name}' on UNTRUSTED server '{server_name}' wants to run. This " - f"tool is write-capable (no readOnlyHint=true annotation) and may modify external " - f"state." - ), - ( - f"Server '{server_name}' is configured 'trust: untrusted'. " - f"Approve to run '{tool_name}' once, or deny to block it." - ), + f"MCP tool '{tool_name}' on UNTRUSTED server '{server_name}' wants to run. This " + f"tool is write-capable (no readOnlyHint=true annotation) and may modify external state.", + f"Server '{server_name}' is configured 'trust: untrusted'. " + f"Approve to run '{tool_name}' once, or deny to block it.", surface=f"mcp-trust/{server_name}", ) except Exception as exc: - logger.error( - "MCP trust gate: approval check failed for %s.%s: %s", - server_name, tool_name, exc, exc_info=True, - ) - return tool_error( - f"MCP tool '{tool_name}' on untrusted server '{server_name}' was blocked: the approval " - f"system was unavailable (fail-closed)." - ) - + logger.error("MCP trust gate: approval check failed for %s.%s: %s", server_name, tool_name, exc, exc_info=True) + return tool_error(f"MCP tool '{tool_name}' on untrusted server '{server_name}' was blocked: the approval " + f"system was unavailable (fail-closed).") if answer == "accept": return None - logger.info( - "MCP trust gate: user %s '%s' on untrusted server '%s'", - "cancelled" if answer == "cancel" else "denied", - tool_name, server_name, - ) - return tool_error( - f"The user did not approve running write-capable MCP tool '{tool_name}' on untrusted server " - f"'{server_name}'. The command was NOT run. Do not retry without explicit user direction." - ) + logger.info("MCP trust gate: user %s '%s' on untrusted server '%s'", + "cancelled" if answer == "cancel" else "denied", tool_name, server_name) + return tool_error(f"The user did not approve running write-capable MCP tool '{tool_name}' on untrusted server " + f"'{server_name}'. The command was NOT run. Do not retry without explicit user direction.") def _check_circuit_breaker(server_name: str) -> Optional[str]: - """Open-breaker error for ``server_name``, or None when calls may proceed. - - After the cooldown the breaker is half-open: the next call goes through as a - probe; success resets it, failure re-bumps (re-stamping the open-time and - re-arming the cooldown). - """ + """Open-breaker error, or None when calls may proceed. After the cooldown the breaker is + half-open: the next call probes; success resets, failure re-bumps and re-arms the cooldown.""" failures = _core._server_error_counts.get(server_name, 0) if failures < _core._CIRCUIT_BREAKER_THRESHOLD: return None @@ -86,38 +59,29 @@ def _check_circuit_breaker(server_name: str) -> Optional[str]: if age >= _core._CIRCUIT_BREAKER_COOLDOWN_SEC: return None remaining = max(1, int(_core._CIRCUIT_BREAKER_COOLDOWN_SEC - age)) - return tool_error( - f"MCP server '{server_name}' is unreachable after {failures} consecutive failures. " - f"Auto-retry available in ~{remaining}s. Do NOT retry this tool yet — use alternative " - f"approaches or ask the user to check the MCP server." - ) + return tool_error(f"MCP server '{server_name}' is unreachable after {failures} consecutive failures. " + f"Auto-retry available in ~{remaining}s. Do NOT retry this tool yet — use alternative " + f"approaches or ask the user to check the MCP server.") def _acquire_call_server(server_name: str, tool_timeout: float): """``(server, None)`` when a call may be dispatched, else ``(None, error)``. - - No session: a reconnect may be completing (fresh session swaps in - asynchronously), so wait briefly before charging a breaker strike. Still - down → reconnecting or parked (e.g. dead stdio child); probing a dead - transport would re-arm the breaker forever, so ask the server task to - rebuild (respawns stdio) and return a clean "reconnecting" error — the - breaker resets once the fresh session initializes. - """ + No session: a reconnect may be completing (fresh session swaps in asynchronously), so wait + briefly before charging a breaker strike. Still down → reconnecting or parked (e.g. dead + stdio child); probing a dead transport would re-arm the breaker forever, so ask the server + task to rebuild and return a clean "reconnecting" error — the breaker resets once the + fresh session initializes.""" not_connected = tool_error(f"MCP server '{server_name}' is not connected") server = _core._get_connected_server_for_call(server_name) if not server: _core._bump_server_error(server_name) return None, not_connected - if server.session or _core._wait_for_server_session_ready( - server, timeout=min(5.0, float(tool_timeout or 5.0)), - ): + if server.session or _core._wait_for_server_session_ready(server, timeout=min(5.0, float(tool_timeout or 5.0))): return server, None _core._bump_server_error(server_name) if _core._signal_reconnect(server): - return None, tool_error( - f"MCP server '{server_name}' transport is down; reconnect requested. Do NOT retry this " - f"tool immediately — give it a few seconds to come back." - ) + return None, tool_error(f"MCP server '{server_name}' transport is down; reconnect requested. Do NOT retry this " + f"tool immediately — give it a few seconds to come back.") return None, not_connected @@ -140,12 +104,21 @@ def _record_call_outcome(server_name: str, result) -> Any: return result -def _lookup_reconnectable_server(server_name: str): - """The registered server object when it can be signalled to reconnect, else None.""" +def _strike(server_name: str, message: str, **extra) -> str: + """Breaker strike + the ``tool_error`` payload for *message*.""" + _core._bump_server_error(server_name) + return tool_error(message, **extra) + + +def _lookup_reconnectable_server(server_name: str, require_loop: bool = False): + """The registered server object when it can be signalled to reconnect, else None. + With *require_loop*, also None unless the MCP loop is running (nothing to wait on).""" with _core._lock: srv = _core._servers.get(server_name) if srv is None or not hasattr(srv, "_reconnect_event"): return None + if require_loop and not _mcp_loop_running(): + return None return srv @@ -155,11 +128,8 @@ def _mcp_loop_running() -> bool: def _retry_once(server_name: str, retry_call, op_description: str, what: str): - """Re-run ``retry_call`` after a recovery step. - - Returns the result (and closes the circuit breaker) when it is not an error - payload; None when the retry raised or errored, so the caller falls through. - """ + """Re-run ``retry_call`` after a recovery step. Returns the result (closing the breaker) + when it is not an error payload; None when the retry raised or errored (caller falls through).""" try: result = retry_call() except Exception as retry_exc: @@ -173,23 +143,14 @@ def _retry_once(server_name: str, retry_call, op_description: str, what: str): # --------------------------------------------------------------- recovery ladder -def _handle_auth_error_and_retry( - server_name: str, - exc: BaseException, - retry_call, - op_description: str, -): - """Attempt OAuth recovery and one retry; return None to fall through. - - Non-auth exceptions return None. Otherwise: ask ``MCPOAuthManager.handle_401`` - whether recovery is viable; if so, signal ``_reconnect_event`` so the server - task rebuilds the session with fresh credentials, wait for ready, retry once. - Any failure returns the structured ``needs_reauth`` error so the model stops - trying to refresh manually. - """ +def _handle_auth_error_and_retry(server_name: str, exc: BaseException, retry_call, op_description: str): + """OAuth recovery + one retry; None when *exc* is not an auth error. + ``MCPOAuthManager.handle_401`` decides whether recovery is viable; if so, signal + ``_reconnect_event`` so the server task rebuilds the session with fresh credentials, wait + for ready, retry once. Any failure returns the structured ``needs_reauth`` error so the + model stops trying to refresh manually.""" if not _core._is_auth_error(exc): return None - from tools.mcp_oauth_manager import get_manager manager = get_manager() @@ -201,155 +162,93 @@ def _handle_auth_error_and_retry( except Exception as rec_exc: logger.warning("MCP OAuth '%s': recovery attempt failed: %s", server_name, rec_exc) recovered = False - if recovered: srv = _lookup_reconnectable_server(server_name) - # OAuth recovery + reconnect is independent evidence the server is viable, - # so close the breaker here, not only on retry success — otherwise a failing - # retry would leave it pinned open forever. A broken server re-trips it via - # _bump_server_error on the retry. + # OAuth recovery + reconnect is independent evidence the server is viable, so close the + # breaker here, not only on retry success — otherwise a failing retry would leave it + # pinned open forever. A broken server re-trips it via _bump_server_error on the retry. if srv is not None and _core._signal_reconnect_and_wait( - server_name, - srv, - op_description=f"{op_description} after OAuth recovery", - timeout=15, - ): + server_name, srv, op_description=f"{op_description} after OAuth recovery", timeout=15): _core._reset_server_error(server_name) - result = _retry_once(server_name, retry_call, op_description, "auth recovery") if result is not None: return result - - # No recovery, or retry failed: structured needs_reauth error + breaker strike. - _core._bump_server_error(server_name) - return tool_error( + return _strike( + server_name, f"MCP server '{server_name}' requires re-authentication. Run `hermes mcp login " f"{server_name}` (or delete the tokens file under ~/.hermes/mcp-tokens/ and restart). Do " f"NOT retry this tool — ask the user to re-authenticate.", - needs_reauth=True, - server=server_name, - ) + needs_reauth=True, server=server_name) -def _handle_session_expired_and_retry( - server_name: str, - exc: BaseException, - retry_call, - op_description: str, -): - """Trigger a transport reconnect and retry once on session expiry. - - Unlike :func:`_handle_auth_error_and_retry` this skips ``handle_401`` — the - token is still valid, only the server-side session is stale. Returns None to - fall through (not session-expired, no server record / loop, reconnect did not - ready in time, or the retry also failed). - """ +def _handle_session_expired_and_retry(server_name: str, exc: BaseException, retry_call, op_description: str): + """Transport reconnect + one retry on session expiry; None to fall through (not + session-expired, no server record / loop, reconnect did not ready in time, retry failed). + Skips ``handle_401`` — the token is still valid, only the server-side session is stale.""" if not _is_session_expired_error(exc): return None - srv = _lookup_reconnectable_server(server_name) - if srv is None or not _mcp_loop_running(): + srv = _lookup_reconnectable_server(server_name, require_loop=True) + if srv is None: return None - - logger.info( - "MCP server '%s': %s failed with session-expired error (%s); " - "signalling transport reconnect and retrying once.", - server_name, op_description, exc, - ) - if not _core._signal_reconnect_and_wait( - server_name, - srv, - op_description=op_description, - timeout=15, - ): - logger.warning( - "MCP server '%s': reconnect did not ready within 15s after " - "session-expired error; falling through to error response.", - server_name, - ) + logger.info("MCP server '%s': %s failed with session-expired error (%s); " + "signalling transport reconnect and retrying once.", server_name, op_description, exc) + if not _core._signal_reconnect_and_wait(server_name, srv, op_description=op_description, timeout=15): + logger.warning("MCP server '%s': reconnect did not ready within 15s after " + "session-expired error; falling through to error response.", server_name) return None return _retry_once(server_name, retry_call, op_description, "session reconnect") class _StdioChildExited(RuntimeError): """A server's stdio subprocess was gone when (or while) a call ran. - - Deliberately NOT a TimeoutError: nothing timed out — the child was already - dead (typically a gateway restart killed it under a live agent session). - Handled by :func:`_handle_stdio_child_exited_and_retry`. - """ + Deliberately NOT a TimeoutError: nothing timed out — the child was already dead + (typically a gateway restart killed it under a live agent session).""" -def _handle_stdio_child_exited_and_retry( - server_name: str, - exc: Exception, - retry_call, - op_description: str, -): - """Respawn a dead stdio child and retry the call once; None if not our error. - - Cannot hot-cycle respawns: this never spawns anything — it sets - ``_reconnect_event`` once and waits for the server task to publish a fresh - session, so spawn frequency stays governed by ``run()``'s rapid-drop budget. - Single-shot: a child that dies again immediately reports and stops. - """ +def _handle_stdio_child_exited_and_retry(server_name: str, exc: Exception, retry_call, op_description: str): + """Respawn a dead stdio child and retry once; None if not our error. + Never spawns anything itself: it sets ``_reconnect_event`` once and waits for the server + task to publish a fresh session, so spawn frequency stays governed by ``run()``'s + rapid-drop budget. Single-shot: a child that dies again immediately reports and stops.""" if not isinstance(exc, _StdioChildExited): return None - reconnected = False srv = _lookup_reconnectable_server(server_name) if srv is not None: - logger.info( - "MCP server '%s': %s found the stdio subprocess dead (%s); " - "respawning and retrying once.", - server_name, op_description, exc, - ) + logger.info("MCP server '%s': %s found the stdio subprocess dead (%s); " + "respawning and retrying once.", server_name, op_description, exc) if _mcp_loop_running(): reconnected = _core._signal_reconnect_and_wait( - server_name, - srv, - op_description=op_description, - timeout=_core._STDIO_RESPAWN_WAIT_SEC, - ) + server_name, srv, op_description=op_description, timeout=_core._STDIO_RESPAWN_WAIT_SEC) else: - # No MCP loop to wait on (non-async adapters, tests) — still request - # the respawn so the next call lands on a live transport. + # No MCP loop to wait on (non-async adapters, tests) — still request the respawn + # so the next call lands on a live transport. _core._signal_reconnect(srv) if not reconnected: - _core._bump_server_error(server_name) - return tool_error( + return _strike( + server_name, f"MCP server '{server_name}' stdio subprocess had exited (this is not a timeout — the " f"call never reached the server). A respawn was requested but no fresh session came " f"back within {_core._STDIO_RESPAWN_WAIT_SEC:.0f}s. Wait a few seconds before retrying; " - f"if it keeps failing the server is not starting and needs the user." - ) - + f"if it keeps failing the server is not starting and needs the user.") try: return _record_call_outcome(server_name, retry_call()) except _StdioChildExited as retry_exc: - # Died again right after respawn: broken server, not a restart - # artifact. Stop here — run()'s budget takes it to the park. - logger.warning( - "MCP server '%s': %s stdio subprocess exited again right " - "after respawn (%s); not retrying further.", - server_name, op_description, retry_exc, - ) - message = ( + # Died again right after respawn: broken server, not a restart artifact. Stop here — + # run()'s budget takes it to the park. + logger.warning("MCP server '%s': %s stdio subprocess exited again right " + "after respawn (%s); not retrying further.", server_name, op_description, retry_exc) + return _strike( + server_name, f"MCP server '{server_name}' respawned its stdio subprocess and it exited again " f"immediately. The server is not starting cleanly — do NOT retry this tool; ask the " - f"user to check the server's command and its stderr log." - ) + f"user to check the server's command and its stderr log.") except Exception as retry_exc: - logger.warning( - "MCP %s/%s retry after stdio respawn failed: %s", - server_name, op_description, retry_exc, - ) - message = _sanitize_error( + logger.warning("MCP %s/%s retry after stdio respawn failed: %s", server_name, op_description, retry_exc) + return _strike(server_name, _sanitize_error( f"MCP call failed after respawning the stdio subprocess for '{server_name}': " - f"{type(retry_exc).__name__}: {_exc_str(retry_exc)}" - ) - _core._bump_server_error(server_name) - return tool_error(message) + f"{type(retry_exc).__name__}: {_exc_str(retry_exc)}")) def _interrupted_call_result() -> str: @@ -360,16 +259,12 @@ def _interrupted_call_result() -> str: def _invoke_with_recovery(server_name: str, call_once: Callable[[], str], op: str, recoverers, on_final_failure: Callable[[BaseException], None], record_outcome: bool = False) -> str: - """Run ``call_once`` on the MCP loop, walking the recovery ladder on failure. - - Each recoverer ``(server_name, exc, retry_call, op) -> Optional[str]`` returns - None when the exception is not its kind. Order matters: dead stdio child - (respawn + retry) → auth (OAuth recovery + retry) → session expiry - (reconnect + retry). Unrecovered exceptions go through ``on_final_failure`` - (breaker strike / logging) and become the generic call-failed error. - ``record_outcome`` applies breaker bookkeeping to the FIRST attempt only; - retries own their own bookkeeping inside the recoverers. - """ + """Run ``call_once``, walking the recovery ladder on failure. Each recoverer + ``(server_name, exc, retry_call, op) -> Optional[str]`` returns None when the exception is + not its kind; order matters: dead stdio child → auth → session expiry. Unrecovered + exceptions go through ``on_final_failure`` (breaker strike / logging) and become the generic + call-failed error. ``record_outcome`` applies breaker bookkeeping to the FIRST attempt only; + retries own their bookkeeping inside the recoverers.""" try: result = call_once() return _record_call_outcome(server_name, result) if record_outcome else result @@ -396,56 +291,40 @@ def _mark_server_call_started(server: Any) -> None: @asynccontextmanager async def _track_inflight_rpc(server: Any, server_name: str, op: str): """Register the running RPC on the server so teardown can fail it fast. - - If a deliberate reconnect/shutdown teardown cancels the task - (``_fail_inflight_calls`` sets ``_reconnecting`` first) the cancel becomes a - clean retryable RuntimeError; external cancels (caller timeout, user - interrupt) propagate unchanged. Test doubles without ``_inflight_tasks`` - simply skip tracking. - """ + A deliberate reconnect/shutdown teardown (``_fail_inflight_calls`` sets ``_reconnecting`` + first) turns the cancel into a clean retryable RuntimeError; external cancels (caller + timeout, user interrupt) propagate unchanged. Doubles without ``_inflight_tasks`` skip tracking.""" inflight = getattr(server, "_inflight_tasks", None) task = asyncio.current_task() - if task is not None and inflight is not None: + tracked = task is not None and inflight is not None + if tracked: inflight.add(task) try: yield except asyncio.CancelledError: if getattr(server, "_reconnecting", False): - raise RuntimeError( - f"MCP {op} on '{server_name}' was aborted by a reconnect " - f"teardown; retry the request on the rebuilt session" - ) from None + raise RuntimeError(f"MCP {op} on '{server_name}' was aborted by a reconnect " + f"teardown; retry the request on the rebuilt session") from None raise finally: - if task is not None and inflight is not None: + if tracked: inflight.discard(task) async def _call_tool_racing_stdio_death(server, server_name: str, tool_name: str, args: dict): """``session.call_tool`` that fails fast when the stdio child is/gets dead. - - Pre-call: an already-dead child must not hold the slot for the full tool - timeout (``server.session`` is stale so the transport-down path never fired). - Mid-call: race the RPC against ``_watch_stdio_children`` so a death fails - immediately. Both raise :class:`_StdioChildExited` for the respawn-and-retry - path, which owns the reconnect signal (nothing clears ``server.session``). - callable()/``is True`` checks because MagicMock attributes return truthy Mocks. - """ + Pre-call: an already-dead child must not hold the slot for the full tool timeout + (``server.session`` is stale so the transport-down path never fired). Mid-call: race the + RPC against ``_watch_stdio_children``. Both raise :class:`_StdioChildExited` for the + respawn-and-retry path, which owns the reconnect signal (nothing clears ``server.session``). + callable()/``is True`` checks because MagicMock attributes return truthy Mocks.""" _stdio_dead = getattr(server, "_stdio_children_dead", None) if callable(_stdio_dead) and _stdio_dead() is True: - raise _StdioChildExited( - f"MCP stdio subprocess for '{server_name}' had " - f"already exited when the call was dispatched" - ) + raise _StdioChildExited(f"MCP stdio subprocess for '{server_name}' had already exited when the call was dispatched") _call_coro = server.session.call_tool(tool_name, arguments=args) _watch_children = getattr(server, "_watch_stdio_children", None) - if not ( - _watch_children is not None - and inspect.iscoroutinefunction(_watch_children) - and asyncio.iscoroutine(_call_coro) - ): - # Stubbed sessions return a non-awaitable, or there is no - # child-watcher to race: plain await. + if not (_watch_children is not None and inspect.iscoroutinefunction(_watch_children) and asyncio.iscoroutine(_call_coro)): + # Stubbed sessions return a non-awaitable, or there is no child-watcher to race: plain await. return await _call_coro if asyncio.iscoroutine(_call_coro) else _call_coro rpc_task = asyncio.ensure_future(_call_coro) @@ -466,33 +345,22 @@ async def _call_tool_racing_stdio_death(server, server_name: str, tool_name: str # ---------------------------------------------------------- result rendering def _error_result_text(result) -> str: - """Concatenated text of an ``isError`` result's blocks (EmbeddedResource - error payloads carry text under ``.resource.text``).""" - error_text = "" - for block in (result.content or []): - if getattr(block, "text", None): - error_text += block.text - continue - res_text = getattr(getattr(block, "resource", None), "text", None) - if res_text: - error_text += str(res_text) - return error_text + """Concatenated text of an ``isError`` result's blocks (EmbeddedResource error payloads + carry text under ``.resource.text``).""" + texts = (getattr(b, "text", None) or getattr(getattr(b, "resource", None), "text", None) for b in (result.content or [])) + return "".join(str(t) for t in texts if t) def _render_content_blocks(result, server_name: str) -> str: - """Text blocks pass through; image/audio blocks are cached via the gateway - image-cache so they flow out as MEDIA: tags; resource blocks (PDFs, docs, ...) - are materialized rather than silently dropped.""" + """Text blocks pass through; image/audio blocks are cached via the gateway image-cache so + they flow out as MEDIA: tags; resource blocks (PDFs, docs, ...) are materialized rather + than silently dropped.""" parts: List[str] = [] for block in (result.content or []): - if hasattr(block, "text") and block.text: + if getattr(block, "text", None): parts.append(strip_unicode_tags(block.text)) continue - rendered = ( - _cache_mcp_image_block(block) - or _cache_mcp_audio_block(block) - or _render_mcp_resource_block(block, server_name) - ) + rendered = _cache_mcp_image_block(block) or _cache_mcp_audio_block(block) or _render_mcp_resource_block(block, server_name) if rendered: parts.append(rendered) continue @@ -507,8 +375,8 @@ def _render_content_blocks(result, server_name: str) -> str: def _capped_structured_content(result): - """``structuredContent`` (or None); over the hard cap it degrades to the - head+tail truncated JSON string (multi-MB JSON flood guard).""" + """``structuredContent`` (or None); over the hard cap it degrades to the head+tail + truncated JSON string (multi-MB JSON flood guard).""" structured = mcp_field(result, "structured_content", "structuredContent") if structured is None: return None @@ -516,41 +384,27 @@ def _capped_structured_content(result): as_json = json.dumps(structured, ensure_ascii=False, default=str) except (TypeError, ValueError): return structured - if len(as_json) > _MCP_HARD_RESULT_CAP_CHARS: - return _truncate_mcp_text_result(as_json) - return structured + return _truncate_mcp_text_result(as_json) if len(as_json) > _MCP_HARD_RESULT_CAP_CHARS else structured def _render_call_tool_result(result, server_name: str) -> str: - """Pure: ``CallToolResult`` → the handler's JSON string. - - ``content`` is the primary (model-oriented) payload; ``structuredContent`` - supplements it (or becomes ``result`` when there is no text). Server-level - ``_meta`` is surfaced minus protocol-reserved keys (``modelcontextprotocol``/ - ``mcp`` label followed by another label; vendor-namespaced keys pass through). - ``.is_error`` is ``.isError`` before mcp 2.0. - """ + """Pure: ``CallToolResult`` → the handler's JSON string. ``content`` is the primary + (model-oriented) payload; ``structuredContent`` supplements it (or becomes ``result`` when + there is no text). Server-level ``_meta`` is surfaced minus protocol-reserved keys. + ``.is_error`` is ``.isError`` before mcp 2.0.""" if mcp_field(result, "is_error", "isError", False): - return tool_error(_sanitize_error(_truncate_mcp_text_result( - _error_result_text(result) or "MCP tool returned an error" - ))) + return tool_error(_sanitize_error(_truncate_mcp_text_result(_error_result_text(result) or "MCP tool returned an error"))) text_result = _render_content_blocks(result, server_name) structured = _capped_structured_content(result) meta = _strip_reserved_meta_keys(mcp_field(result, "meta", "meta")) if structured is None and meta is None: return json.dumps({"result": text_result}, ensure_ascii=False) - - # Key order is part of the output: "result" leads when there is text, - # otherwise "_meta" precedes the (empty) "result". - payload: Dict[str, Any] = {} - if text_result: - payload["result"] = text_result + # Key order is part of the output: "result" leads when there is text, otherwise "_meta" + # precedes the (empty) "result". + payload: Dict[str, Any] = {"result": text_result} if text_result else {} if structured is not None: - if text_result: - payload["structuredContent"] = structured - else: - payload["result"] = structured + payload["structuredContent" if text_result else "result"] = structured if meta is not None: payload["_meta"] = meta payload.setdefault("result", text_result) @@ -564,13 +418,12 @@ def _render_call_tool_result(result, server_name: str) -> str: # ------------------------------------------------------------------- handlers def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): - """Return a sync registry handler (``handler(args_dict, **kwargs) -> str``) - that calls an MCP tool via the background loop.""" + """Sync registry handler (``handler(args_dict, **kwargs) -> str``) calling an MCP tool via the background loop.""" op = f"tools/call {tool_name}" def _handler(args: dict, **kwargs) -> str: - # Security boundary: untrusted-server write tools need approval before - # ANY transport work, including the lazy first-use spawn below. + # Security boundary: untrusted-server write tools need approval before ANY transport + # work, including the lazy first-use spawn below. error = _trust_gate_check(server_name, tool_name) or _check_circuit_breaker(server_name) if error is not None: return error @@ -581,34 +434,28 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): async def _call(): _mark_server_call_started(server) async with server._rpc_lock, _track_inflight_rpc(server, server_name, op): - # Snapshot contextvars so an elicitation callback (fired on the - # MCP recv loop, which doesn't inherit them) can replay them - # for gateway platform / session routing. + # Snapshot contextvars so an elicitation callback (fired on the MCP recv loop, + # which doesn't inherit them) can replay them for gateway platform / session routing. server._pending_call_context = contextvars.copy_context() try: result = await _call_tool_racing_stdio_death(server, server_name, tool_name, args) finally: server._pending_call_context = None - # Round-trip completed: transport is healthy even if the tool - # returned isError. Clear the rapid-drop budget. + # Round-trip completed: transport is healthy even if the tool returned isError. + # Clear the rapid-drop budget. _mark_proven = getattr(server, "_mark_session_proven", None) if _mark_proven is not None: _mark_proven() return _render_call_tool_result(result, server_name) - def _call_once(): - return _core._run_on_mcp_loop(_call, timeout=tool_timeout) - def _on_failure(exc): _core._bump_server_error(server_name) logger.error("MCP tool %s/%s call failed: %s", server_name, tool_name, exc) return _invoke_with_recovery( - server_name, _call_once, op, - (_handle_stdio_child_exited_and_retry, _handle_auth_error_and_retry, - _handle_session_expired_and_retry), - _on_failure, - record_outcome=True, + server_name, lambda: _core._run_on_mcp_loop(_call, timeout=tool_timeout), op, + (_handle_stdio_child_exited_and_retry, _handle_auth_error_and_retry, _handle_session_expired_and_retry), + _on_failure, record_outcome=True, ) return _handler @@ -616,13 +463,10 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): def _make_utility_handler(server_name: str, tool_timeout: float, op: str, log_label: str, rpc, render, required: Optional[str] = None): - """Shared shape of the four utility handlers (resources/prompts). - - ``rpc(session, args)`` is awaited under ``_rpc_lock``; ``render(result, - server_name)`` turns its result into the JSON-able payload. ``required`` names - a parameter that must be present (validated before any transport work). The - wrapper owns the connected check and the auth / session-expired recovery ladder. - """ + """Shared shape of the four utility handlers (resources/prompts): ``rpc(session, args)`` + is awaited under ``_rpc_lock``, ``render(result, server_name)`` builds the JSON-able payload, + ``required`` names a parameter validated before any transport work. The wrapper owns the + connected check and the auth / session-expired recovery ladder.""" def _handler(args: dict, **kwargs) -> str: server = _core._get_connected_server_for_call(server_name) @@ -637,16 +481,12 @@ def _make_utility_handler(server_name: str, tool_timeout: float, op: str, log_la result = await rpc(server.session, args) return json.dumps(render(result, server_name), ensure_ascii=False) - def _call_once(): - return _core._run_on_mcp_loop(_call, timeout=tool_timeout) - def _on_failure(exc): logger.error("MCP %s/%s failed: %s", server_name, log_label, exc) return _invoke_with_recovery( - server_name, _call_once, op, - (_handle_auth_error_and_retry, _handle_session_expired_and_retry), - _on_failure, + server_name, lambda: _core._run_on_mcp_loop(_call, timeout=tool_timeout), op, + (_handle_auth_error_and_retry, _handle_session_expired_and_retry), _on_failure, ) return _handler @@ -654,20 +494,15 @@ def _make_utility_handler(server_name: str, tool_timeout: float, op: str, log_la def _pick(obj, *specs) -> dict: """``{out_key: value}`` for each ``(out_key, attr[, truthy])`` present on *obj*. - - ``hasattr`` (not a default) so SDK models and test stubs behave alike; with - ``truthy`` the field is also skipped when falsy. Output key order = spec order. - """ + ``hasattr`` (not a default) so SDK models and test stubs behave alike; with ``truthy`` + the field is also skipped when falsy. Output key order = spec order.""" entry = {} - for spec in specs: - out_key, attr = spec[0], spec[1] - truthy = len(spec) > 2 and spec[2] + for out_key, attr, *truthy in specs: if not hasattr(obj, attr): continue value = getattr(obj, attr) - if truthy and not value: - continue - entry[out_key] = value + if value or not (truthy and truthy[0]): + entry[out_key] = value return entry @@ -691,11 +526,9 @@ def _render_read_resource(result, server_name: str) -> dict: if getattr(block, "text", None) is not None: parts.append(strip_unicode_tags(block.text)) elif getattr(block, "blob", None) is not None: - # Materialize binary contents into the document cache - # (same contract as EmbeddedResource blocks in tool results). - rendered = _render_mcp_resource_block( - SimpleNamespace(type="resource", resource=block), server_name, - ) + # Materialize binary contents into the document cache (same contract as + # EmbeddedResource blocks in tool results). + rendered = _render_mcp_resource_block(SimpleNamespace(type="resource", resource=block), server_name) parts.append(rendered or f"[binary data, {len(block.blob)} bytes]") return {"result": "\n".join(parts)} @@ -704,7 +537,7 @@ def _render_prompt_list(all_prompts, server_name: str) -> dict: prompts = [] for p in all_prompts: entry = _pick(p, ("name", "name"), ("description", "description", True)) - if hasattr(p, "arguments") and p.arguments: + if getattr(p, "arguments", None): entry["arguments"] = [ {"name": a.name, **_pick(a, ("description", "description", True), ("required", "required"))} for a in p.arguments @@ -719,65 +552,51 @@ def _render_get_prompt(result, server_name: str) -> dict: entry = _pick(msg, ("role", "role")) if hasattr(msg, "content"): content = msg.content - text = content.text if hasattr(content, "text") else ( - content if isinstance(content, str) else str(content) - ) - entry["content"] = strip_unicode_tags(text) + entry["content"] = strip_unicode_tags(content.text if hasattr(content, "text") else str(content)) messages.append(entry) resp = {"messages": messages} - if hasattr(result, "description") and result.description: + if getattr(result, "description", None): resp["description"] = result.description return resp def _make_list_resources_handler(server_name: str, tool_timeout: float): - """Return a sync handler that lists resources from an MCP server.""" - return _make_utility_handler( - server_name, tool_timeout, "resources/list", "list_resources", - lambda session, args: _core._paginate_full_list(session.list_resources, "resources", server_name), - _render_resource_list, - ) + """Sync handler that lists resources from an MCP server.""" + return _make_utility_handler(server_name, tool_timeout, "resources/list", "list_resources", + lambda session, args: _core._paginate_full_list(session.list_resources, "resources", server_name), + _render_resource_list) def _make_read_resource_handler(server_name: str, tool_timeout: float): - """Return a sync handler that reads a resource by URI from an MCP server.""" - return _make_utility_handler( - server_name, tool_timeout, "resources/read", "read_resource", - lambda session, args: session.read_resource(args["uri"]), - _render_read_resource, required="uri", - ) + """Sync handler that reads a resource by URI from an MCP server.""" + return _make_utility_handler(server_name, tool_timeout, "resources/read", "read_resource", + lambda session, args: session.read_resource(args["uri"]), _render_read_resource, required="uri") def _make_list_prompts_handler(server_name: str, tool_timeout: float): - """Return a sync handler that lists prompts from an MCP server.""" - return _make_utility_handler( - server_name, tool_timeout, "prompts/list", "list_prompts", - lambda session, args: _core._paginate_full_list(session.list_prompts, "prompts", server_name), - _render_prompt_list, - ) + """Sync handler that lists prompts from an MCP server.""" + return _make_utility_handler(server_name, tool_timeout, "prompts/list", "list_prompts", + lambda session, args: _core._paginate_full_list(session.list_prompts, "prompts", server_name), + _render_prompt_list) def _make_get_prompt_handler(server_name: str, tool_timeout: float): - """Return a sync handler that gets a prompt by name from an MCP server.""" - return _make_utility_handler( - server_name, tool_timeout, "prompts/get", "get_prompt", - lambda session, args: session.get_prompt(args["name"], arguments=args.get("arguments", {})), - _render_get_prompt, required="name", - ) + """Sync handler that gets a prompt by name from an MCP server.""" + return _make_utility_handler(server_name, tool_timeout, "prompts/get", "get_prompt", + lambda session, args: session.get_prompt(args["name"], arguments=args.get("arguments", {})), + _render_get_prompt, required="name") def _make_check_fn(server_name: str): - """Return a check function that verifies the MCP connection is alive.""" + """Check function that verifies the MCP connection is alive.""" def _check() -> bool: with _core._lock: server = _core._servers.get(server_name) - if server is not None and ( - server.session is not None or server._is_recycled_stdio() - ): + if server is not None and (server.session is not None or server._is_recycled_stdio()): return True - # Lazy (schema-cache registered) servers count as available: the - # first real call spawns/connects them. + # Lazy (schema-cache registered) servers count as available: the first real + # call spawns/connects them. return server_name in _core._lazy_server_configs return _check From 36aefc2638f8c808ace9a2d8f47c774839709429 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:02:34 -0700 Subject: [PATCH 2/2] =?UTF-8?q?refactor(mcp):=20transport=20=E2=80=94=20on?= =?UTF-8?q?e=20guarded-attempt=20helper=20for=20both=20protocol-negotiatio?= =?UTF-8?q?n=20fallbacks,=20kwargs=20tables=20for=20optional=20httpx=20arg?= =?UTF-8?q?s,=20docstrings=20compacted=20by=20hand=20(645=20->=20511)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tools/mcp_tool_transport.py | 484 +++++++++++++----------------------- 1 file changed, 175 insertions(+), 309 deletions(-) diff --git a/tools/mcp_tool_transport.py b/tools/mcp_tool_transport.py index 79ad17a30c..407a30a3e8 100644 --- a/tools/mcp_tool_transport.py +++ b/tools/mcp_tool_transport.py @@ -33,14 +33,13 @@ def _is_2xx(resp) -> bool: def _capture_pgids(pids: Set[int]) -> Dict[int, int]: - """pgid per live pid. Captured while the child is alive — getpgid fails once - it exits, and the sweep needs it to reach reparented descendants.""" + """pgid per live pid. Captured while the child is alive — getpgid fails once it exits, + and the sweep needs it to reach reparented descendants.""" pgids: Dict[int, int] = {} for pid in pids: try: pgids[pid] = os.getpgid(pid) - except (AttributeError, ProcessLookupError, OSError): - # AttributeError: Windows; ProcessLookupError: already exited. + except (AttributeError, ProcessLookupError, OSError): # Windows / already exited pass return pgids @@ -58,22 +57,17 @@ def _pgroup_alive(pgid: Optional[int]) -> bool: async def _osv_malware_preflight(server_name: str, command: str, args: list) -> None: - """OSV malware preflight: off-loop (blocking HTTPS) with a wall-clock bound so - a stalled handshake can't freeze discovery; fail-open on timeout. Must run - against the REAL command/args — the watchdog wrap rewrites argv to the - supervisor, which would turn the check into a no-op.""" + """OSV malware preflight: off-loop (blocking HTTPS) with a wall-clock bound so a stalled + handshake can't freeze discovery; fail-open on timeout. Must run against the REAL + command/args — the watchdog wrap rewrites argv to the supervisor, turning the check into a no-op.""" from tools.osv_check import check_package_for_malware try: malware_error = await asyncio.wait_for( - asyncio.to_thread(check_package_for_malware, command, args), - timeout=_core._OSV_MALWARE_CHECK_TIMEOUT_S, - ) + asyncio.to_thread(check_package_for_malware, command, args), timeout=_core._OSV_MALWARE_CHECK_TIMEOUT_S) except asyncio.TimeoutError: - logger.warning( - "MCP server '%s': OSV malware preflight timed out after %.0fs " - "(network slow/unreachable) — proceeding without the check.", - server_name, _core._OSV_MALWARE_CHECK_TIMEOUT_S, - ) + logger.warning("MCP server '%s': OSV malware preflight timed out after %.0fs " + "(network slow/unreachable) — proceeding without the check.", + server_name, _core._OSV_MALWARE_CHECK_TIMEOUT_S) return if malware_error: raise ValueError(f"MCP server '{server_name}': {malware_error}") @@ -85,17 +79,11 @@ class MCPServerTransportMixin: __slots__ = () def _advertises_tools(self) -> bool: - """Whether the server advertises the ``tools`` capability. - - Prompt-/resource-only servers omit it, and ``tools/list`` against them - raises ``MCPError(-32601)``. True when no capability info was captured - (legacy fallback: keep the old always-call-list_tools behavior). - """ - init_result = self.initialize_result - caps = getattr(init_result, "capabilities", None) if init_result is not None else None - if caps is None: - return True - return getattr(caps, "tools", None) is not None + """Whether the server advertises the ``tools`` capability. Prompt-/resource-only servers + omit it, and ``tools/list`` against them raises ``MCPError(-32601)``. True when no + capability info was captured (legacy fallback: always call list_tools).""" + caps = getattr(self.initialize_result, "capabilities", None) + return caps is None or getattr(caps, "tools", None) is not None def _session_kwargs(self) -> dict: """ClientSession kwargs: sampling, elicitation, notification + logging callbacks.""" @@ -110,66 +98,52 @@ class MCPServerTransportMixin: async def _negotiate_session(self, session, connect_timeout: float): """Negotiate the protocol era (``initialize`` vs ``server/discover``) and return its result. - - Per-server ``protocol`` key: ``auto`` (default) tries the legacy handshake - FIRST and falls back to ``server/discover`` only when the server signals - modern-only (-32022 / initialize -32601) — the reverse of the SDK's - discover-first mode, on purpose: zero extra round-trips for the handshake-era - servers that dominate today. ``stateless`` probes discover first (one legacy - retry on error); ``legacy`` is handshake only, no fallback. Both result - types expose ``.capabilities``, so downstream gates work on either. - A handshake TIMEOUT never triggers a fallback — it propagates. - """ + Per-server ``protocol`` key: ``auto`` (default) tries the legacy handshake FIRST and falls + back to ``server/discover`` only when the server signals modern-only (-32022 / initialize + -32601) — the reverse of the SDK's discover-first mode, on purpose: zero extra round-trips + for the handshake-era servers that dominate today. ``stateless`` probes discover first (one + legacy retry on any error); ``legacy`` is handshake only, no fallback. Both result types + expose ``.capabilities``. A handshake TIMEOUT never triggers a fallback — it propagates.""" def initialize(): return asyncio.wait_for(session.initialize(), timeout=connect_timeout) def discover(): return asyncio.wait_for(session.discover(), timeout=connect_timeout) - mode = str((self._config or {}).get("protocol", "auto")).lower().strip() - if mode in ("stateless", "modern", "2026-07-28"): + async def attempt(primary, fallback, should_fallback, log_fmt, *log_extra): try: - return await discover() + return await primary() except asyncio.TimeoutError: raise except Exception as exc: - logger.info( - "MCP server '%s': server/discover rejected (%s) despite " - "protocol=%s — falling back to the legacy handshake", - self.name, exc, mode, - ) - return await initialize() + if not should_fallback(exc): + raise + logger.info(log_fmt, self.name, exc, *log_extra) + return await fallback() + + mode = str((self._config or {}).get("protocol", "auto")).lower().strip() + if mode in ("stateless", "modern", "2026-07-28"): + return await attempt(discover, initialize, lambda exc: True, + "MCP server '%s': server/discover rejected (%s) despite " + "protocol=%s — falling back to the legacy handshake", mode) if mode in ("legacy", "handshake"): return await initialize() if mode != "auto": - logger.warning( - "MCP server '%s': unknown protocol=%r — treating as 'auto' " - "(valid: auto, stateless, legacy)", self.name, mode, - ) - try: - return await initialize() - except asyncio.TimeoutError: - raise - except Exception as exc: - # mcp 1.x has no server/discover client — nothing to fall back to. - if not _handshake_rejected_as_modern(exc) or not hasattr(session, "discover"): - raise - logger.info( - "MCP server '%s': legacy handshake rejected (%s) — " - "retrying via server/discover (2026-07-28 stateless server)", - self.name, exc, - ) - return await discover() + logger.warning("MCP server '%s': unknown protocol=%r — treating as 'auto' " + "(valid: auto, stateless, legacy)", self.name, mode) + # mcp 1.x has no server/discover client — nothing to fall back to. + return await attempt( + initialize, discover, + lambda exc: _handshake_rejected_as_modern(exc) and hasattr(session, "discover"), + "MCP server '%s': legacy handshake rejected (%s) — " + "retrying via server/discover (2026-07-28 stateless server)") async def _serve_session(self, session, connect_timeout: float, label: str = "", mark_lifecycle: bool = False) -> str: """Handshake, discover, publish readiness, then serve until a lifecycle event. - - Clears stale breaker state from a prior outage, but leaves the session - UNPROVEN: a completed handshake is not proof of health (flapping - transports handshake fine and drop moments later); only keepalive or - tool-call success clears the reconnect budget. - """ + Clears stale breaker state from a prior outage but leaves the session UNPROVEN: a + completed handshake is not proof of health (flapping transports handshake fine and drop + moments later); only keepalive or tool-call success clears the reconnect budget.""" self.initialize_result = await self._negotiate_session(session, connect_timeout) self.session = session if mark_lifecycle: @@ -185,17 +159,12 @@ class MCPServerTransportMixin: return reason async def _serve_transport(self, transport_cm, label: str, connect_timeout: float) -> str: - """Open *transport_cm*, wrap its streams in a ClientSession and serve it. - - Streams are unpacked positionally: mcp 1.x yields ``(read, write, - get_session_id)``, 2.x yields ``(read, write)``. A transport TaskGroup - drop maps to ``"reconnect"`` instead of backoff/park. - """ + """Open *transport_cm*, wrap its streams in a ClientSession and serve it. Streams are + unpacked positionally: mcp 1.x yields ``(read, write, get_session_id)``, 2.x ``(read, write)``. + A transport TaskGroup drop maps to ``"reconnect"`` instead of backoff/park.""" try: async with transport_cm as _streams: - async with _core.ClientSession( - _streams[0], _streams[1], **self._session_kwargs() - ) as session: + async with _core.ClientSession(_streams[0], _streams[1], **self._session_kwargs()) as session: return await self._serve_session(session, connect_timeout, label) except BaseExceptionGroup as _eg: return self._reconnect_or_reraise_group(_eg) @@ -207,10 +176,9 @@ class MCPServerTransportMixin: command = config.get("command") if not command: raise ValueError(f"MCP server '{self.name}' has no 'command' in config") - args = config.get("args", []) safe_env = _core._build_safe_env(config.get("env")) command, safe_env = _core._resolve_stdio_command(command, safe_env) - return command, args, safe_env + return command, config.get("args", []), safe_env def _track_spawned_children(self, new_pids: Set[int]) -> None: """Ledger the freshly spawned stdio children (pids, pgids, machine spawn ledger).""" @@ -219,8 +187,8 @@ class MCPServerTransportMixin: for _pid in new_pids: _stdio_pids[_pid] = self.name _stdio_pgids.update(new_pgids) - # Machine spawn ledger so startup sweeps can reap orphans after an - # unclean parent exit. Best-effort — never break startup. + # Machine spawn ledger so startup sweeps can reap orphans after an unclean parent + # exit. Best-effort — never break startup. for _pid in new_pids: try: from hermes_cli.process_identity import register_child @@ -230,92 +198,71 @@ class MCPServerTransportMixin: logger.debug("spawn-ledger register_child failed for MCP helper pid %s", _pid, exc_info=True) def _release_spawned_children(self, new_pids: Set[int]) -> None: - """Drop the ledger entries; any child (or its pgroup) still alive means SDK - teardown failed (common on cancel mid-way on Linux, where setsid() children - escape the cgroup) — mark it orphaned for the next cleanup sweep.""" + """Drop the ledger entries; any child (or its pgroup) still alive means SDK teardown + failed (common on cancel mid-way on Linux, where setsid() children escape the cgroup) + — mark it orphaned for the next cleanup sweep.""" from gateway.status import _pid_exists with _core._lock: - for _pid in new_pids: - _stdio_pids.pop(_pid, None) for pid in new_pids: - # ``os.kill(pid, 0)`` is NOT a no-op on Windows; use the - # cross-platform check. Child exited but descendants may remain - # in its pgroup. + _stdio_pids.pop(pid, None) + # ``os.kill(pid, 0)`` is NOT a no-op on Windows; use the cross-platform check. + # The child may have exited while descendants remain in its pgroup. if _pid_exists(pid) or _pgroup_alive(_stdio_pgids.get(pid)): _orphan_stdio_pids.add(pid) _orphan_stdio_pid_servers[pid] = self.name else: - # Nothing to reap — drop the pgid so PID reuse can't - # surface stale pgroup state later. + # Nothing to reap — drop the pgid so PID reuse can't surface stale pgroup state. _stdio_pgids.pop(pid, None) async def _run_stdio(self, config: dict): """Run the server using stdio transport.""" if config.get("identity_header") is not None: # No headers on stdio — warn so a copy-pasted HTTP block doesn't mislead. - logger.warning( - "MCP server '%s': identity_header is only supported on " - "HTTP/SSE transports — ignored for stdio servers", self.name, - ) + logger.warning("MCP server '%s': identity_header is only supported on " + "HTTP/SSE transports — ignored for stdio servers", self.name) if not _core._ensure_mcp_sdk(): - raise ImportError( - f"MCP server '{self.name}' requires the 'mcp' Python SDK, but " - "it is not installed. Run `hermes setup` to install MCP support, then retry." - ) + raise ImportError(f"MCP server '{self.name}' requires the 'mcp' Python SDK, but " + "it is not installed. Run `hermes setup` to install MCP support, then retry.") command, args, safe_env = self._resolve_stdio_config(config) await _osv_malware_preflight(self.name, command, args) - # Parent-death watchdog: an ungraceful Hermes exit (kill -9, crash) can't - # leave the child and its descendants running. Clean-exit reaping is - # unchanged. POSIX-only (process groups); no-op elsewhere. AFTER the OSV - # preflight so the check inspects the real package. + # Parent-death watchdog: an ungraceful Hermes exit (kill -9, crash) can't leave the + # child and its descendants running. POSIX-only (process groups); no-op elsewhere. + # AFTER the OSV preflight so the check inspects the real package. command, args = _wrap_command_with_watchdog(command, args) server_params = _core.StdioServerParameters( - command=command, - args=args, - env=safe_env if safe_env else None, - cwd=config.get("cwd"), - # Windows pipes can split non-UTF-8 bytes at chunk boundaries; - # substitute U+FFFD instead of raising UnicodeDecodeError. + command=command, args=args, env=safe_env if safe_env else None, cwd=config.get("cwd"), + # Windows pipes can split non-UTF-8 bytes at chunk boundaries; substitute U+FFFD + # instead of raising UnicodeDecodeError. encoding_error_handler="replace", ) session_kwargs = self._session_kwargs() - - # Reap orphans from prior failed attempts before spawning, else each - # reconnect retry piles up zombie pairs. Unscoped on purpose (also reaps - # orphans of servers that never reconnect). Worker thread: the reaper - # blocks up to 2s (SIGTERM → wait → SIGKILL) and would stall the loop. + # Reap orphans from prior failed attempts before spawning, else each reconnect retry + # piles up zombie pairs. Unscoped on purpose (also reaps orphans of servers that never + # reconnect). Worker thread: the reaper blocks up to 2s (SIGTERM → wait → SIGKILL). await asyncio.to_thread(_core._kill_orphaned_mcp_children) - # Snapshot child PIDs before spawning so the new one can be identified. pids_before = _core._snapshot_child_pids() new_pids: set = set() - # Route subprocess stderr to ~/.hermes/logs/mcp-stderr.log so server - # banners don't land on the user's TTY and corrupt the TUI. + # Route subprocess stderr to ~/.hermes/logs/mcp-stderr.log so server banners don't + # land on the user's TTY and corrupt the TUI. _core._write_stderr_log_header(self.name) _errlog = _core._get_mcp_stderr_log() try: - async with _core.stdio_client(server_params, errlog=_errlog) as ( - read_stream, - write_stream, - ): - # Capture the new PID for force-kill cleanup, filtering non-MCP - # children (slash_worker, LSP servers) that race into the snapshot - # window: they share the TUI parent's pgid, so leaking them into - # _stdio_pgids makes the shutdown killpg() kill the TUI itself. + async with _core.stdio_client(server_params, errlog=_errlog) as (read_stream, write_stream): + # Capture the new PID for force-kill cleanup, filtering non-MCP children + # (slash_worker, LSP servers) that race into the snapshot window: they share + # the TUI parent's pgid, so leaking them into _stdio_pgids makes the shutdown + # killpg() kill the TUI itself. new_pids = _filter_mcp_children(_core._snapshot_child_pids() - pids_before) if new_pids: self._track_spawned_children(new_pids) - # Tracked on the connection so in-flight calls fail fast when - # the subprocess dies. + # Tracked on the connection so in-flight calls fail fast when the subprocess dies. self._stdio_child_pids = set(new_pids) - async with _core.ClientSession( - read_stream, write_stream, **session_kwargs - ) as session: - # Bound the handshake: ``connect_timeout`` only bounds the - # caller's ``.result()`` wait, not this coroutine. A server that - # never answers ``initialize`` would otherwise hang here forever, - # the ``finally`` below would never run, and the child + pipes - # would leak on every retry until EMFILE. + async with _core.ClientSession(read_stream, write_stream, **session_kwargs) as session: + # Bound the handshake: ``connect_timeout`` only bounds the caller's + # ``.result()`` wait, not this coroutine. A server that never answers + # ``initialize`` would otherwise hang here forever, the ``finally`` below + # would never run, and the child + pipes would leak on every retry until EMFILE. connect_timeout = float(config.get("connect_timeout", _core._DEFAULT_CONNECT_TIMEOUT)) return await self._serve_session(session, connect_timeout, mark_lifecycle=True) finally: @@ -325,40 +272,24 @@ class MCPServerTransportMixin: # ------------------------------------------------------------------- HTTP - async def _preflight_content_type( - self, - url: str, - *, - headers: Optional[dict] = None, - ssl_verify: bool = True, - client_cert=None, - timeout: float = 5.0, - ) -> None: - """Probe *url* for an MCP-shaped response before the SDK connects. - - A URL pointing at a plain web page makes the SDK sit out the full - ``connect_timeout`` before an opaque ``CancelledError``; this raises - :class:`NonMcpEndpointError` within ``timeout`` instead. Allow-list based: - only a 2xx with a definite non-MCP content type is rejected, and only - after a JSON-RPC ``initialize`` POST also fails to look like MCP (some - servers serve a UI on GET but speak Streamable HTTP via POST). Missing - content type, non-2xx, or transport errors pass silently — the real - handshake stays the source of truth. Uses its own httpx client OUTSIDE - the SDK's anyio task group so the error isn't wrapped in an ExceptionGroup. - """ + async def _preflight_content_type(self, url: str, *, headers: Optional[dict] = None, + ssl_verify: bool = True, client_cert=None, timeout: float = 5.0) -> None: + """Probe *url* for an MCP-shaped response before the SDK connects. A URL pointing at a + plain web page makes the SDK sit out the full ``connect_timeout`` before an opaque + ``CancelledError``; this raises :class:`NonMcpEndpointError` within ``timeout`` instead. + Allow-list based: only a 2xx with a definite non-MCP content type is rejected, and only + after a JSON-RPC ``initialize`` POST also fails to look like MCP (some servers serve a UI + on GET but speak Streamable HTTP via POST). Missing content type, non-2xx, or transport + errors pass silently — the real handshake stays the source of truth. Uses its own httpx + client OUTSIDE the SDK's anyio task group so the error isn't wrapped in an ExceptionGroup.""" try: import httpx as _httpx except ImportError: return # No httpx → skip probe; SDK import would have failed first. - client_kwargs: dict = { - "verify": ssl_verify, - "follow_redirects": True, - "timeout": _httpx.Timeout(timeout), - } + client_kwargs: dict = {"verify": ssl_verify, "follow_redirects": True, "timeout": _httpx.Timeout(timeout)} if client_cert is not None: client_kwargs["cert"] = client_cert - probe_headers = dict(headers) if headers else {} try: async with _httpx.AsyncClient(**client_kwargs) as client: @@ -366,17 +297,14 @@ class MCPServerTransportMixin: resp = await client.head(url, headers=probe_headers) if resp.status_code in (405, 501): resp = await client.get(url, headers=probe_headers) - # Non-MCP content type on HEAD/GET: try a JSON-RPC POST before - # rejecting, so POST-only servers aren't false positives. + # Non-MCP content type on HEAD/GET: try a JSON-RPC POST before rejecting, so + # POST-only servers aren't false positives. ct = _content_type_base(resp) if ct and ct not in self._MCP_CONTENT_TYPES and _is_2xx(resp): post_resp = await client.post( url, - headers={ - **probe_headers, - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - }, + headers={**probe_headers, "Content-Type": "application/json", + "Accept": "application/json, text/event-stream"}, content=_PROBE_INITIALIZE_BODY, ) if _is_2xx(post_resp) and _content_type_base(post_resp) in self._MCP_CONTENT_TYPES: @@ -384,9 +312,8 @@ class MCPServerTransportMixin: except _httpx.HTTPError: return # DNS/connect/timeout/transport error — let the SDK try. - # Only judge 2xx: a 4xx/5xx may be an auth challenge or transient error - # the real handshake handles correctly. No content type advertised → - # don't second-guess the SDK. + # Only judge 2xx: a 4xx/5xx may be an auth challenge or transient error the real + # handshake handles correctly. No content type advertised → don't second-guess the SDK. if not _is_2xx(resp): return ct_base = _content_type_base(resp) @@ -400,37 +327,27 @@ class MCPServerTransportMixin: ) def _reconnect_or_reraise_group(self, eg: BaseExceptionGroup) -> str: - """Map an SDK transport TaskGroup failure to a clean ``"reconnect"``. - - HTTP/SSE stream pumps run in an anyio TaskGroup, so a transient stream - drop escapes as a ``BaseExceptionGroup``. Unmapped, ``run()`` would back - off and eventually park the server for 300s (deregistering its tools) - over a sub-second glitch; ``"reconnect"`` rebuilds the session at once. - Re-raise when it is not a transient drop: shutdown in progress - (``_shutdown_event`` is set before the task is cancelled), the group - carries KeyboardInterrupt/SystemExit or a real CancelledError (must - propagate), or no live session was reached this attempt (``_ready`` - unset — connect failures must go through backoff, not hot-loop). - """ - if ( - self._shutdown_event.is_set() - or eg.split((KeyboardInterrupt, SystemExit))[0] is not None - or eg.split(asyncio.CancelledError)[0] is not None - or not self._ready.is_set() - ): + """Map an SDK transport TaskGroup failure to a clean ``"reconnect"``. HTTP/SSE stream + pumps run in an anyio TaskGroup, so a transient stream drop escapes as a + ``BaseExceptionGroup``; unmapped, ``run()`` would back off and eventually park the server + for 300s (deregistering its tools) over a sub-second glitch. Re-raise when it is not a + transient drop: shutdown in progress (``_shutdown_event`` is set before the task is + cancelled), the group carries KeyboardInterrupt/SystemExit or a real CancelledError (must + propagate), or no live session was reached this attempt (``_ready`` unset — connect + failures must go through backoff, not hot-loop).""" + if (self._shutdown_event.is_set() + or eg.split((KeyboardInterrupt, SystemExit))[0] is not None + or eg.split(asyncio.CancelledError)[0] is not None + or not self._ready.is_set()): raise eg - logger.debug( - "MCP server '%s': transport TaskGroup exited after a live session " - "(%r) — reconnecting immediately instead of backing off", - self.name, eg, - ) + logger.debug("MCP server '%s': transport TaskGroup exited after a live session " + "(%r) — reconnecting immediately instead of backing off", self.name, eg) return "reconnect" def _build_oauth_auth(self, url: str, config: dict): - """OAuth 2.1 PKCE via the central MCPOAuthManager so one provider is reused - across reconnects and shared with config-time CLI paths. On setup failure - (e.g. non-interactive without cached tokens) re-raise so only this server - is reported failed.""" + """OAuth 2.1 PKCE via the central MCPOAuthManager so one provider is reused across + reconnects and shared with config-time CLI paths. On setup failure (e.g. + non-interactive without cached tokens) re-raise so only this server is reported failed.""" if self._auth_type != "oauth": return None try: @@ -445,48 +362,30 @@ class MCPServerTransportMixin: """``sse_client`` context manager for ``transport: sse`` entries.""" if strict_cfg_headers: # Fail closed: SSE cannot enforce the redirect boundary. - raise ValueError( - f"MCP server '{self.name}': strict_redirect_headers is " - "not supported on the SSE transport." - ) + raise ValueError(f"MCP server '{self.name}': strict_redirect_headers is " + "not supported on the SSE transport.") if _core.sse_client is None: - raise ImportError( - f"MCP server '{self.name}' requires SSE transport but " - "mcp.client.sse.sse_client is not available. " - "Upgrade the mcp package to get SSE support." - ) - # sse_read_timeout bounds the gap between SSE events. SSE servers - # commonly idle for minutes, so tool_timeout (60s) would drop the - # stream; 300s matches the Streamable HTTP read timeout. - sse_kwargs: dict = { - "url": url, - "headers": headers or None, - "timeout": float(connect_timeout), - "sse_read_timeout": 300.0, - } + raise ImportError(f"MCP server '{self.name}' requires SSE transport but " + "mcp.client.sse.sse_client is not available. " + "Upgrade the mcp package to get SSE support.") + # sse_read_timeout bounds the gap between SSE events. SSE servers commonly idle for + # minutes, so tool_timeout (60s) would drop the stream; 300s matches the Streamable + # HTTP read timeout. + sse_kwargs: dict = {"url": url, "headers": headers or None, "timeout": float(connect_timeout), "sse_read_timeout": 300.0} if oauth_auth is not None: # Forward OAuth to sse_client, else OAuth SSE servers 401 silently. sse_kwargs["auth"] = oauth_auth if client_cert is not None or ssl_verify is not True: - # sse_client has no verify/cert kwargs: wrap the SDK defaults - # (follow_redirects=True) in an httpx_client_factory, forwarding - # the SDK's (headers, auth, timeout) and layering TLS on top. The - # client MUST come from the SDK's own httpx module (httpx2 on - # mcp >= 2.0) — see sdk_httpx(). + # sse_client has no verify/cert kwargs: wrap the SDK defaults (follow_redirects=True) + # in an httpx_client_factory, forwarding the SDK's (headers, auth, timeout) and + # layering TLS on top. The client MUST come from the SDK's own httpx module + # (httpx2 on mcp >= 2.0) — see sdk_httpx(). _httpx_mod = _core.sdk_httpx() def _mcp_http_client_factory(headers=None, timeout=None, auth=None): - kwargs: dict = { - "follow_redirects": True, - "verify": ssl_verify, - "timeout": timeout if timeout is not None else _httpx_mod.Timeout(30.0, read=300.0), - } - if headers is not None: - kwargs["headers"] = headers - if auth is not None: - kwargs["auth"] = auth - if client_cert is not None: - kwargs["cert"] = client_cert + kwargs: dict = {"follow_redirects": True, "verify": ssl_verify, + "timeout": timeout if timeout is not None else _httpx_mod.Timeout(30.0, read=300.0)} + kwargs.update({k: v for k, v in (("headers", headers), ("auth", auth), ("cert", client_cert)) if v is not None}) return _httpx_mod.AsyncClient(**kwargs) sse_kwargs["httpx_client_factory"] = _mcp_http_client_factory @@ -497,34 +396,22 @@ class MCPServerTransportMixin: strict_cfg_headers: bool, configured_header_names: set): """Streamable HTTP context manager (mcp >= 1.24.0: caller-owned httpx client).""" if not _core._MCP_NEW_HTTP: - return self._legacy_http_transport(url, headers, connect_timeout, ssl_verify, - oauth_auth, strict_cfg_headers) - # Build an explicit AsyncClient matching the SDK's create_mcp_http_client - # defaults. It MUST come from the SDK's httpx module (httpx2 on mcp >= 2.0) - # since the SDK sends its own Request objects through it — see sdk_httpx(). + return self._legacy_http_transport(url, headers, connect_timeout, ssl_verify, oauth_auth, strict_cfg_headers) + # Build an explicit AsyncClient matching the SDK's create_mcp_http_client defaults. It + # MUST come from the SDK's httpx module (httpx2 on mcp >= 2.0) since the SDK sends its + # own Request objects through it — see sdk_httpx(). httpx = _core.sdk_httpx() _strip_auth_on_cross_origin_redirect = _make_redirect_header_stripper( - httpx.URL(url), - strict=strict_cfg_headers, - configured_header_names=configured_header_names, - ) - client_kwargs: dict = { - "follow_redirects": True, - "timeout": httpx.Timeout(float(connect_timeout), read=300.0), - "verify": ssl_verify, - "event_hooks": {"response": [_strip_auth_on_cross_origin_redirect]}, - } + httpx.URL(url), strict=strict_cfg_headers, configured_header_names=configured_header_names) + client_kwargs: dict = {"follow_redirects": True, "timeout": httpx.Timeout(float(connect_timeout), read=300.0), + "verify": ssl_verify, "event_hooks": {"response": [_strip_auth_on_cross_origin_redirect]}} if headers: client_kwargs["headers"] = headers - if oauth_auth is not None: - client_kwargs["auth"] = oauth_auth - if client_cert is not None: - client_kwargs["cert"] = client_cert + client_kwargs.update({k: v for k, v in (("auth", oauth_auth), ("cert", client_cert)) if v is not None}) @asynccontextmanager async def _owned_client_streams(): - # Caller owns the client lifecycle — the SDK skips cleanup when - # http_client is provided. + # Caller owns the client lifecycle — the SDK skips cleanup when http_client is provided. async with httpx.AsyncClient(**client_kwargs) as http_client: async with _core.streamable_http_client(url, http_client=http_client) as streams: yield streams @@ -535,13 +422,11 @@ class MCPServerTransportMixin: ssl_verify, oauth_auth, strict_cfg_headers: bool): """Deprecated API (mcp < 1.24.0): the SDK owns the httpx client.""" if strict_cfg_headers: - # Fail closed: without an owned client we cannot hook redirects, - # so the cross-origin header boundary cannot be enforced. - raise ImportError( - f"MCP server '{self.name}' requires mcp >= 1.24.0 to " - "enforce the portable redirect-header boundary " - "(strict_redirect_headers). Upgrade the mcp package." - ) + # Fail closed: without an owned client we cannot hook redirects, so the + # cross-origin header boundary cannot be enforced. + raise ImportError(f"MCP server '{self.name}' requires mcp >= 1.24.0 to " + "enforce the portable redirect-header boundary " + "(strict_redirect_headers). Upgrade the mcp package.") http_kwargs: dict = {"headers": headers, "timeout": float(connect_timeout), "verify": ssl_verify} if oauth_auth is not None: http_kwargs["auth"] = oauth_auth @@ -551,26 +436,22 @@ class MCPServerTransportMixin: """Run the server using HTTP/StreamableHTTP (or SSE) transport.""" _core._ensure_mcp_sdk() if not _core._MCP_HTTP_AVAILABLE: - raise ImportError( - f"MCP server '{self.name}' requires HTTP transport but " - "mcp.client.streamable_http is not available. " - "Upgrade the mcp package to get HTTP support." - ) + raise ImportError(f"MCP server '{self.name}' requires HTTP transport but " + "mcp.client.streamable_http is not available. " + "Upgrade the mcp package to get HTTP support.") url = config["url"] headers = dict(config.get("headers") or {}) - # Portable Agent Plugins v1 (strict_redirect_headers): configured - # headers MUST NOT follow a redirect to a different origin. Capture the - # configured names BEFORE client-generated headers are merged in. + # Portable Agent Plugins v1 (strict_redirect_headers): configured headers MUST NOT + # follow a redirect to a different origin. Capture the configured names BEFORE + # client-generated headers are merged in. strict_cfg_headers = bool(config.get("strict_redirect_headers")) configured_header_names = {key.lower() for key in headers} # Optional per-user identity header; explicit headers of the same name win. headers = _apply_identity_header(self.name, config, headers) - # Some servers require MCP-Protocol-Version on the initial request; seed - # it (case-insensitive user override wins). Seeded from the HANDSHAKE - # version, not the latest: the body sent by ``initialize()`` speaks the - # handshake era, and a 2026-07-28 header would route the request onto - # the server's per-request-envelope ladder, which rejects that body. - # The header must agree with what the body actually speaks. + # Some servers require MCP-Protocol-Version on the initial request; seed it (case-insensitive + # user override wins) from the HANDSHAKE version, not the latest: ``initialize()``'s body speaks + # the handshake era, and a 2026-07-28 header would route it onto the server's per-request-envelope + # ladder, which rejects that body. The header must agree with what the body actually speaks. if not any(key.lower() == "mcp-protocol-version" for key in headers): headers["mcp-protocol-version"] = _core.LATEST_HANDSHAKE_VERSION connect_timeout = config.get("connect_timeout", _core._DEFAULT_CONNECT_TIMEOUT) @@ -579,58 +460,43 @@ class MCPServerTransportMixin: oauth_auth = self._build_oauth_auth(url, config) if config.get("transport") == "sse": - transport = self._sse_transport(url, headers, connect_timeout, ssl_verify, - client_cert, oauth_auth, strict_cfg_headers) + transport = self._sse_transport(url, headers, connect_timeout, ssl_verify, client_cert, oauth_auth, strict_cfg_headers) label = "SSE" else: - transport = self._streamable_http_transport( - url, headers, connect_timeout, ssl_verify, client_cert, oauth_auth, - strict_cfg_headers, configured_header_names, - ) + transport = self._streamable_http_transport(url, headers, connect_timeout, ssl_verify, client_cert, + oauth_auth, strict_cfg_headers, configured_header_names) label = "HTTP" if _core._MCP_NEW_HTTP else "legacy HTTP" return await self._serve_transport(transport, label, float(connect_timeout)) # -------------------------------------------------------------- discovery async def _discover_tools(self): - """Discover tools from the connected session. - - Capability-gated: prompt-/resource-only servers raise ``MCPError(-32601)`` - on ``tools/list``, which would abort the connection — skip the call when - ``tools`` isn't advertised. - """ - # Fresh transport: re-probe with cheap ``ping`` in case the server gained - # support across the reconnect. + """Discover tools from the connected session. Capability-gated: prompt-/resource-only + servers raise ``MCPError(-32601)`` on ``tools/list``, which would abort the connection.""" + # Fresh transport: re-probe with cheap ``ping`` in case the server gained support + # across the reconnect. self._ping_unsupported = False if self.session is None: return if not self._advertises_tools(): - logger.info( - "MCP server '%s': does not advertise 'tools' capability — " - "skipping tools/list (prompts/resources remain available)", - self.name, - ) + logger.info("MCP server '%s': does not advertise 'tools' capability — " + "skipping tools/list (prompts/resources remain available)", self.name) self._tools = [] self._register_discovered_tools_if_needed() return async with self._rpc_lock: self._list_cache_meta = {} self._tools = await _core._paginate_full_list( - self.session.list_tools, "tools", self.name, - cache_meta_out=self._list_cache_meta, - ) + self.session.list_tools, "tools", self.name, cache_meta_out=self._list_cache_meta) self._register_discovered_tools_if_needed() def _register_discovered_tools_if_needed(self) -> None: """Publish freshly discovered tools for a registry-owned server if none are registered. - - Initial registration normally happens in ``_discover_and_register_server`` - after ``start()``. On reconnect, outage handling may clear ``_ready`` and - deregister stale tools; ownership via ``_servers`` authorizes publishing - before readiness is restored so a revival never comes back with zero - tools. A server retained after a recoverable initial failure is likewise - owned before its first session, which authorizes its first publication. - """ + Initial registration normally happens in ``_discover_and_register_server`` after + ``start()``. On reconnect, outage handling may clear ``_ready`` and deregister stale + tools; ownership via ``_servers`` authorizes publishing before readiness is restored so a + revival never comes back with zero tools. A server retained after a recoverable initial + failure is likewise owned before its first session, which authorizes its first publication.""" if self._registered_tool_names: return if not self._ready.is_set(): @@ -638,8 +504,8 @@ class MCPServerTransportMixin: if _core._servers.get(self.name) is not self: return self._registered_tool_names = _core._register_server_tools(self.name, self, self._config) - # A retained initial-failure server that just published tools has - # recovered: drop its stale connect error from status surfaces. + # A retained initial-failure server that just published tools has recovered: drop + # its stale connect error from status surfaces. with _core._lock: if _core._servers.get(self.name) is self: _core._server_connect_errors.pop(self.name, None)