786 lines
34 KiB
Python
786 lines
34 KiB
Python
"""Registry-facing sync handlers for MCP tools and utility tools (resources/prompts), plus the per-call recovery ladder: trust gating, circuit breaker, auth (401) refresh, session-expired reconnect and dead-stdio respawn retry. Split from tools/mcp_tool.py."""
|
|
|
|
import logging
|
|
import asyncio
|
|
import contextvars
|
|
import inspect
|
|
import json
|
|
import time
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict, List, Optional
|
|
from tools.registry import tool_error
|
|
from tools.ansi_strip import strip_unicode_tags
|
|
from tools.mcp_tool_common import _exc_str, _sanitize_error, mcp_field, _core
|
|
from tools.mcp_tool_content import _MCP_HARD_RESULT_CAP_CHARS, _cache_mcp_audio_block, _cache_mcp_image_block, _render_mcp_resource_block, _strip_reserved_meta_keys, _truncate_mcp_text_result
|
|
from tools.mcp_tool_errors import _is_session_expired_error
|
|
|
|
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.
|
|
"""
|
|
trust = _core._server_trust_levels.get(server_name, _core._TRUST_FULL)
|
|
if trust != _core._TRUST_UNTRUSTED:
|
|
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.
|
|
try:
|
|
from tools.approval import request_elicitation_consent
|
|
|
|
answer = request_elicitation_consent(
|
|
(
|
|
f"MCP tool '{tool_name}' on UNTRUSTED server "
|
|
f"'{server_name}' wants to run. This tool is write-capable "
|
|
f"(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."
|
|
),
|
|
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}' "
|
|
f"was blocked: the approval system was unavailable "
|
|
f"(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 "
|
|
f"'{tool_name}' on untrusted server '{server_name}'. The command "
|
|
f"was NOT run. Do not retry without explicit user direction."
|
|
)
|
|
|
|
|
|
def _result_is_error(result) -> bool:
|
|
"""True only for a JSON payload carrying an ``error`` key (non-JSON = success)."""
|
|
try:
|
|
return "error" in json.loads(result)
|
|
except (json.JSONDecodeError, TypeError):
|
|
return False
|
|
|
|
|
|
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.
|
|
"""
|
|
try:
|
|
result = retry_call()
|
|
except Exception as retry_exc:
|
|
logger.warning(
|
|
"MCP %s/%s retry after %s failed: %s",
|
|
server_name, op_description, what, retry_exc,
|
|
)
|
|
return None
|
|
if _result_is_error(result):
|
|
return None
|
|
_core._reset_server_error(server_name)
|
|
return result
|
|
|
|
|
|
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.
|
|
"""
|
|
if not _core._is_auth_error(exc):
|
|
return None
|
|
|
|
from tools.mcp_oauth_manager import get_manager
|
|
manager = get_manager()
|
|
|
|
async def _recover():
|
|
return await manager.handle_401(server_name, None)
|
|
|
|
try:
|
|
recovered = _core._run_on_mcp_loop(_recover, timeout=10)
|
|
except Exception as rec_exc:
|
|
logger.warning(
|
|
"MCP OAuth '%s': recovery attempt failed: %s",
|
|
server_name, rec_exc,
|
|
)
|
|
recovered = False
|
|
|
|
if recovered:
|
|
with _core._lock:
|
|
srv = _core._servers.get(server_name)
|
|
reconnected = False
|
|
if srv is not None and hasattr(srv, "_reconnect_event"):
|
|
reconnected = _core._signal_reconnect_and_wait(
|
|
server_name,
|
|
srv,
|
|
op_description=f"{op_description} after OAuth recovery",
|
|
timeout=15,
|
|
)
|
|
|
|
# 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 reconnected:
|
|
_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(
|
|
f"MCP server '{server_name}' requires re-authentication. "
|
|
f"Run `hermes mcp login {server_name}` (or delete the tokens "
|
|
f"file under ~/.hermes/mcp-tokens/ and restart). Do NOT retry "
|
|
f"this tool — ask the user to re-authenticate.",
|
|
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).
|
|
"""
|
|
if not _is_session_expired_error(exc):
|
|
return None
|
|
|
|
with _core._lock:
|
|
srv = _core._servers.get(server_name)
|
|
if srv is None or not hasattr(srv, "_reconnect_event"):
|
|
return None
|
|
|
|
loop = _core._mcp_loop
|
|
if loop is None or not loop.is_running():
|
|
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,
|
|
)
|
|
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`.
|
|
"""
|
|
|
|
|
|
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.
|
|
"""
|
|
if not isinstance(exc, _StdioChildExited):
|
|
return None
|
|
|
|
with _core._lock:
|
|
srv = _core._servers.get(server_name)
|
|
|
|
reconnected = False
|
|
if srv is not None and hasattr(srv, "_reconnect_event"):
|
|
logger.info(
|
|
"MCP server '%s': %s found the stdio subprocess dead (%s); "
|
|
"respawning and retrying once.",
|
|
server_name, op_description, exc,
|
|
)
|
|
loop = _core._mcp_loop
|
|
if loop is not None and loop.is_running():
|
|
reconnected = _core._signal_reconnect_and_wait(
|
|
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.
|
|
_core._signal_reconnect(srv)
|
|
|
|
if reconnected:
|
|
try:
|
|
result = 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,
|
|
)
|
|
_core._bump_server_error(server_name)
|
|
return tool_error(
|
|
f"MCP server '{server_name}' respawned its stdio subprocess "
|
|
f"and it exited again immediately. The server is not "
|
|
f"starting cleanly — do NOT retry this tool; ask the user to "
|
|
f"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,
|
|
)
|
|
_core._bump_server_error(server_name)
|
|
return tool_error(_sanitize_error(
|
|
f"MCP call failed after respawning the stdio subprocess for "
|
|
f"'{server_name}': {type(retry_exc).__name__}: "
|
|
f"{_exc_str(retry_exc)}"
|
|
))
|
|
if _result_is_error(result):
|
|
_core._bump_server_error(server_name)
|
|
else:
|
|
_core._reset_server_error(server_name)
|
|
return result
|
|
|
|
_core._bump_server_error(server_name)
|
|
return tool_error(
|
|
f"MCP server '{server_name}' stdio subprocess had exited (this is "
|
|
f"not a timeout — the call never reached the server). A respawn was "
|
|
f"requested but no fresh session came back within "
|
|
f"{_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."
|
|
)
|
|
|
|
|
|
def _interrupted_call_result() -> str:
|
|
"""Standardized JSON error for a user-interrupted MCP tool call."""
|
|
return tool_error("MCP call interrupted: user sent a new message")
|
|
|
|
|
|
def _mark_server_call_started(server: Any) -> None:
|
|
"""Record a user-visible MCP operation when the server supports it."""
|
|
mark_tool_call = getattr(server, "mark_tool_call", None)
|
|
if callable(mark_tool_call):
|
|
mark_tool_call()
|
|
|
|
|
|
@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.
|
|
"""
|
|
inflight = getattr(server, "_inflight_tasks", None)
|
|
task = asyncio.current_task()
|
|
if task is not None and inflight is not None:
|
|
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
|
|
finally:
|
|
if task is not None and inflight is not None:
|
|
inflight.discard(task)
|
|
|
|
|
|
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."""
|
|
|
|
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.
|
|
gate_error = _trust_gate_check(server_name, tool_name)
|
|
if gate_error is not None:
|
|
return gate_error
|
|
|
|
# Circuit breaker. After the cooldown the breaker is half-open: the next
|
|
# call goes through as a probe; success resets it, failure re-bumps (which
|
|
# re-stamps the open-time and re-arms the cooldown).
|
|
if _core._server_error_counts.get(server_name, 0) >= _core._CIRCUIT_BREAKER_THRESHOLD:
|
|
opened_at = _core._server_breaker_opened_at.get(server_name, 0.0)
|
|
age = time.monotonic() - opened_at
|
|
if age < _core._CIRCUIT_BREAKER_COOLDOWN_SEC:
|
|
remaining = max(1, int(_core._CIRCUIT_BREAKER_COOLDOWN_SEC - age))
|
|
return tool_error(
|
|
f"MCP server '{server_name}' is unreachable after "
|
|
f"{_core._server_error_counts[server_name]} consecutive "
|
|
f"failures. Auto-retry available in ~{remaining}s. "
|
|
f"Do NOT retry this tool yet — use alternative "
|
|
f"approaches or ask the user to check the MCP server."
|
|
)
|
|
|
|
server = _core._get_connected_server_for_call(server_name)
|
|
if not server:
|
|
_core._bump_server_error(server_name)
|
|
return tool_error(f"MCP server '{server_name}' is not connected")
|
|
|
|
# No session: a reconnect may be completing (fresh session swaps in
|
|
# asynchronously), so wait briefly before charging a breaker strike.
|
|
if not server.session and not _core._wait_for_server_session_ready(
|
|
server, timeout=min(5.0, float(tool_timeout or 5.0)),
|
|
):
|
|
# 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.
|
|
_core._bump_server_error(server_name)
|
|
if _core._signal_reconnect(server):
|
|
return tool_error(
|
|
f"MCP server '{server_name}' transport is down; "
|
|
f"reconnect requested. Do NOT retry this tool "
|
|
f"immediately — give it a few seconds to come back."
|
|
)
|
|
return tool_error(f"MCP server '{server_name}' is not connected")
|
|
|
|
async def _call():
|
|
_mark_server_call_started(server)
|
|
async with server._rpc_lock, _track_inflight_rpc(
|
|
server, server_name, f"tools/call {tool_name}"
|
|
):
|
|
# 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:
|
|
# Fast-fail: an already-dead stdio child must not hold this
|
|
# slot for the full tool timeout. callable() + real-bool check
|
|
# because MagicMock attributes return truthy Mocks.
|
|
_stdio_dead = getattr(server, "_stdio_children_dead", None)
|
|
if (
|
|
callable(_stdio_dead)
|
|
and isinstance(_stdio_dead_result := _stdio_dead(), bool)
|
|
and _stdio_dead_result
|
|
):
|
|
# server.session is stale so the transport-down path above
|
|
# never fired; hand this to the respawn-and-retry path.
|
|
raise _StdioChildExited(
|
|
f"MCP stdio subprocess for '{server_name}' had "
|
|
f"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)
|
|
_watch_ok = (
|
|
_watch_children is not None
|
|
and inspect.iscoroutinefunction(_watch_children)
|
|
and asyncio.iscoroutine(_call_coro)
|
|
)
|
|
if not _watch_ok:
|
|
# Stubbed sessions return a non-awaitable, or there is no
|
|
# child-watcher to race: plain await.
|
|
result = (
|
|
await _call_coro
|
|
if asyncio.iscoroutine(_call_coro)
|
|
else _call_coro
|
|
)
|
|
else:
|
|
# Race the RPC against the stdio-children watcher so a
|
|
# mid-call death fails immediately.
|
|
rpc_task = asyncio.ensure_future(_call_coro)
|
|
watch_task = asyncio.ensure_future(_watch_children())
|
|
try:
|
|
done, _pending = await asyncio.wait(
|
|
{rpc_task, watch_task},
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
if watch_task in done and not rpc_task.done():
|
|
rpc_task.cancel()
|
|
# Nothing clears server.session on a mid-call
|
|
# death; the respawn-and-retry path owns the
|
|
# reconnect signal.
|
|
raise _StdioChildExited(
|
|
f"MCP stdio subprocess for "
|
|
f"'{server_name}' exited mid-call"
|
|
)
|
|
result = await rpc_task
|
|
finally:
|
|
watch_task.cancel()
|
|
if not rpc_task.done():
|
|
rpc_task.cancel()
|
|
await asyncio.gather(
|
|
rpc_task, watch_task, return_exceptions=True
|
|
)
|
|
finally:
|
|
server._pending_call_context = None
|
|
# 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()
|
|
# CallToolResult: .content (blocks) and .is_error (.isError before mcp 2.0).
|
|
if mcp_field(result, "is_error", "isError", False):
|
|
error_text = ""
|
|
for block in (result.content or []):
|
|
if getattr(block, "text", None):
|
|
error_text += block.text
|
|
continue
|
|
# EmbeddedResource error payloads carry text under .resource.text.
|
|
res_text = getattr(getattr(block, "resource", None), "text", None)
|
|
if res_text:
|
|
error_text += str(res_text)
|
|
return tool_error(_sanitize_error(
|
|
_truncate_mcp_text_result(
|
|
error_text or "MCP tool returned an error"
|
|
)
|
|
))
|
|
|
|
# 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:
|
|
parts.append(strip_unicode_tags(block.text))
|
|
continue
|
|
image_tag = _cache_mcp_image_block(block)
|
|
if image_tag:
|
|
parts.append(image_tag)
|
|
continue
|
|
audio_tag = _cache_mcp_audio_block(block)
|
|
if audio_tag:
|
|
parts.append(audio_tag)
|
|
continue
|
|
resource_text = _render_mcp_resource_block(block, server_name)
|
|
if resource_text:
|
|
parts.append(resource_text)
|
|
continue
|
|
# Benign empty renders log at debug; warn only for unknown shapes.
|
|
block_type = getattr(block, "type", None) or type(block).__name__
|
|
if block_type in {"text", "resource", "audio", "image"}:
|
|
logger.debug(
|
|
"MCP %s: content block type %r rendered empty",
|
|
server_name, block_type,
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"MCP %s: dropping unsupported content block type %r",
|
|
server_name, block_type,
|
|
)
|
|
text_result = "\n".join(parts) if parts else ""
|
|
|
|
# Hard-cap pathological payloads; ordinary large results pass to spillover.
|
|
text_result = _truncate_mcp_text_result(text_result)
|
|
|
|
# content is the primary (model-oriented) payload; structuredContent
|
|
# supplements it. Server-level `_meta` is surfaced too, minus
|
|
# protocol-reserved keys (`modelcontextprotocol`/`mcp` label followed
|
|
# by another label); vendor-namespaced keys pass through.
|
|
structured = mcp_field(result, "structured_content", "structuredContent")
|
|
# Cap structuredContent too (multi-MB JSON flood); over the hard cap it
|
|
# degrades to the head+tail truncated string.
|
|
if structured is not None:
|
|
try:
|
|
_structured_json = json.dumps(structured, ensure_ascii=False, default=str)
|
|
except (TypeError, ValueError):
|
|
_structured_json = None
|
|
if _structured_json is not None and len(_structured_json) > _MCP_HARD_RESULT_CAP_CHARS:
|
|
structured = _truncate_mcp_text_result(_structured_json)
|
|
meta = _strip_reserved_meta_keys(mcp_field(result, "meta", "meta"))
|
|
if structured is not None or meta is not None:
|
|
payload: Dict[str, Any] = {}
|
|
if text_result:
|
|
payload["result"] = text_result
|
|
if structured is not None:
|
|
if text_result:
|
|
payload["structuredContent"] = structured
|
|
else:
|
|
payload["result"] = structured
|
|
if meta is not None:
|
|
payload["_meta"] = meta
|
|
if "result" not in payload:
|
|
payload["result"] = text_result
|
|
try:
|
|
return json.dumps(payload, ensure_ascii=False)
|
|
except (TypeError, ValueError):
|
|
# Non-serializable metadata: drop the extras, keep the call.
|
|
return json.dumps({"result": text_result}, ensure_ascii=False)
|
|
return json.dumps({"result": text_result}, ensure_ascii=False)
|
|
|
|
def _call_once():
|
|
return _core._run_on_mcp_loop(_call, timeout=tool_timeout)
|
|
|
|
try:
|
|
result = _call_once()
|
|
# An error payload from the tool itself still counts as a strike.
|
|
if _result_is_error(result):
|
|
_core._bump_server_error(server_name)
|
|
else:
|
|
_core._reset_server_error(server_name)
|
|
return result
|
|
except InterruptedError:
|
|
return _interrupted_call_result()
|
|
except Exception as exc:
|
|
# Recovery ladder, in order: dead stdio child (respawn + retry),
|
|
# auth (OAuth recovery + retry), session expiry (reconnect + retry).
|
|
# Each returns None when the exception is not its kind.
|
|
for recover in (
|
|
_handle_stdio_child_exited_and_retry,
|
|
_handle_auth_error_and_retry,
|
|
_handle_session_expired_and_retry,
|
|
):
|
|
recovered = recover(server_name, exc, _call_once, f"tools/call {tool_name}")
|
|
if recovered is not None:
|
|
return recovered
|
|
|
|
_core._bump_server_error(server_name)
|
|
logger.error(
|
|
"MCP tool %s/%s call failed: %s",
|
|
server_name, tool_name, exc,
|
|
)
|
|
return tool_error(_sanitize_error(
|
|
f"MCP call failed: {type(exc).__name__}: {_exc_str(exc)}"
|
|
))
|
|
|
|
return _handler
|
|
|
|
|
|
def _make_utility_handler(server_name: str, tool_timeout: float, op: str,
|
|
log_label: str, build_call):
|
|
"""Shared shape of the four utility handlers (resources/prompts).
|
|
|
|
``build_call(server, args)`` returns an error string (validation failed) or a
|
|
zero-arg coroutine function doing the RPC under ``_rpc_lock``. 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)
|
|
if not server or not server.session:
|
|
return tool_error(f"MCP server '{server_name}' is not connected")
|
|
|
|
call = build_call(server, args)
|
|
if isinstance(call, str):
|
|
return call
|
|
|
|
def _call_once():
|
|
return _core._run_on_mcp_loop(call, timeout=tool_timeout)
|
|
|
|
try:
|
|
return _call_once()
|
|
except InterruptedError:
|
|
return _interrupted_call_result()
|
|
except Exception as exc:
|
|
for recover in (_handle_auth_error_and_retry, _handle_session_expired_and_retry):
|
|
recovered = recover(server_name, exc, _call_once, op)
|
|
if recovered is not None:
|
|
return recovered
|
|
logger.error("MCP %s/%s failed: %s", server_name, log_label, exc)
|
|
return tool_error(_sanitize_error(
|
|
f"MCP call failed: {type(exc).__name__}: {_exc_str(exc)}"
|
|
))
|
|
|
|
return _handler
|
|
|
|
|
|
def _make_list_resources_handler(server_name: str, tool_timeout: float):
|
|
"""Return a sync handler that lists resources from an MCP server."""
|
|
|
|
def _build(server, args):
|
|
async def _call():
|
|
_mark_server_call_started(server)
|
|
async with server._rpc_lock:
|
|
all_resources = await _core._paginate_full_list(
|
|
server.session.list_resources, "resources", server_name
|
|
)
|
|
resources = []
|
|
for r in all_resources:
|
|
entry = {}
|
|
if hasattr(r, "uri"):
|
|
entry["uri"] = str(r.uri)
|
|
if hasattr(r, "name"):
|
|
entry["name"] = r.name
|
|
if hasattr(r, "description") and r.description:
|
|
entry["description"] = r.description
|
|
# Key stays camelCase — this is the tool's own JSON output shape.
|
|
_mime = mcp_field(r, "mime_type", "mimeType")
|
|
if _mime:
|
|
entry["mimeType"] = _mime
|
|
resources.append(entry)
|
|
return json.dumps({"resources": resources}, ensure_ascii=False)
|
|
return _call
|
|
|
|
return _make_utility_handler(server_name, tool_timeout, "resources/list", "list_resources", _build)
|
|
|
|
|
|
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."""
|
|
|
|
def _build(server, args):
|
|
uri = args.get("uri")
|
|
if not uri:
|
|
return tool_error("Missing required parameter 'uri'")
|
|
|
|
async def _call():
|
|
_mark_server_call_started(server)
|
|
async with server._rpc_lock:
|
|
result = await server.session.read_resource(uri)
|
|
parts: List[str] = []
|
|
contents = result.contents if hasattr(result, "contents") else []
|
|
for block in contents:
|
|
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,
|
|
)
|
|
parts.append(rendered or f"[binary data, {len(block.blob)} bytes]")
|
|
return json.dumps({"result": "\n".join(parts) if parts else ""}, ensure_ascii=False)
|
|
return _call
|
|
|
|
return _make_utility_handler(server_name, tool_timeout, "resources/read", "read_resource", _build)
|
|
|
|
|
|
def _make_list_prompts_handler(server_name: str, tool_timeout: float):
|
|
"""Return a sync handler that lists prompts from an MCP server."""
|
|
|
|
def _build(server, args):
|
|
async def _call():
|
|
_mark_server_call_started(server)
|
|
async with server._rpc_lock:
|
|
all_prompts = await _core._paginate_full_list(
|
|
server.session.list_prompts, "prompts", server_name
|
|
)
|
|
prompts = []
|
|
for p in all_prompts:
|
|
entry = {}
|
|
if hasattr(p, "name"):
|
|
entry["name"] = p.name
|
|
if hasattr(p, "description") and p.description:
|
|
entry["description"] = p.description
|
|
if hasattr(p, "arguments") and p.arguments:
|
|
entry["arguments"] = [
|
|
{
|
|
"name": a.name,
|
|
**({"description": a.description} if hasattr(a, "description") and a.description else {}),
|
|
**({"required": a.required} if hasattr(a, "required") else {}),
|
|
}
|
|
for a in p.arguments
|
|
]
|
|
prompts.append(entry)
|
|
return json.dumps({"prompts": prompts}, ensure_ascii=False)
|
|
return _call
|
|
|
|
return _make_utility_handler(server_name, tool_timeout, "prompts/list", "list_prompts", _build)
|
|
|
|
|
|
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."""
|
|
|
|
def _build(server, args):
|
|
name = args.get("name")
|
|
if not name:
|
|
return tool_error("Missing required parameter 'name'")
|
|
arguments = args.get("arguments", {})
|
|
|
|
async def _call():
|
|
_mark_server_call_started(server)
|
|
async with server._rpc_lock:
|
|
result = await server.session.get_prompt(name, arguments=arguments)
|
|
messages = []
|
|
for msg in (result.messages if hasattr(result, "messages") else []):
|
|
entry = {}
|
|
if hasattr(msg, "role"):
|
|
entry["role"] = msg.role
|
|
if hasattr(msg, "content"):
|
|
content = msg.content
|
|
if hasattr(content, "text"):
|
|
entry["content"] = strip_unicode_tags(content.text)
|
|
elif isinstance(content, str):
|
|
entry["content"] = strip_unicode_tags(content)
|
|
else:
|
|
entry["content"] = strip_unicode_tags(str(content))
|
|
messages.append(entry)
|
|
resp = {"messages": messages}
|
|
if hasattr(result, "description") and result.description:
|
|
resp["description"] = result.description
|
|
return json.dumps(resp, ensure_ascii=False)
|
|
return _call
|
|
|
|
return _make_utility_handler(server_name, tool_timeout, "prompts/get", "get_prompt", _build)
|
|
|
|
|
|
def _make_check_fn(server_name: str):
|
|
"""Return a 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()
|
|
):
|
|
return True
|
|
# 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
|