refactor(mcp): split tool handler into breaker check, server acquire, stdio-death-racing call, pure result renderer and one shared recovery ladder
This commit is contained in:
Teknium
@@ -8,7 +8,7 @@ import json
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Callable, 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
|
||||
@@ -73,6 +73,29 @@ def _trust_gate_check(server_name: str, tool_name: str) -> Optional[str]:
|
||||
)
|
||||
|
||||
|
||||
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).
|
||||
"""
|
||||
failures = _core._server_error_counts.get(server_name, 0)
|
||||
if failures < _core._CIRCUIT_BREAKER_THRESHOLD:
|
||||
return None
|
||||
age = time.monotonic() - _core._server_breaker_opened_at.get(server_name, 0.0)
|
||||
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 "
|
||||
f"{failures} 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."
|
||||
)
|
||||
|
||||
|
||||
def _result_is_error(result) -> bool:
|
||||
"""True only for a JSON payload carrying an ``error`` key (non-JSON = success)."""
|
||||
try:
|
||||
@@ -81,6 +104,35 @@ def _result_is_error(result) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _record_call_outcome(server_name: str, result) -> Any:
|
||||
"""Breaker bookkeeping: 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
|
||||
|
||||
|
||||
def _generic_call_failure(exc: BaseException) -> str:
|
||||
return tool_error(_sanitize_error(
|
||||
f"MCP call failed: {type(exc).__name__}: {_exc_str(exc)}"
|
||||
))
|
||||
|
||||
|
||||
def _lookup_reconnectable_server(server_name: str):
|
||||
"""The registered server object when it can be signalled to reconnect, else None."""
|
||||
with _core._lock:
|
||||
srv = _core._servers.get(server_name)
|
||||
if srv is None or not hasattr(srv, "_reconnect_event"):
|
||||
return None
|
||||
return srv
|
||||
|
||||
|
||||
def _mcp_loop_running() -> bool:
|
||||
loop = _core._mcp_loop
|
||||
return loop is not None and loop.is_running()
|
||||
|
||||
|
||||
def _retry_once(server_name: str, retry_call, op_description: str, what: str):
|
||||
"""Re-run ``retry_call`` after a recovery step.
|
||||
|
||||
@@ -134,22 +186,17 @@ def _handle_auth_error_and_retry(
|
||||
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,
|
||||
)
|
||||
|
||||
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.
|
||||
if reconnected:
|
||||
if srv is not None and _core._signal_reconnect_and_wait(
|
||||
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")
|
||||
@@ -183,14 +230,8 @@ def _handle_session_expired_and_retry(
|
||||
"""
|
||||
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():
|
||||
srv = _lookup_reconnectable_server(server_name)
|
||||
if srv is None or not _mcp_loop_running():
|
||||
return None
|
||||
|
||||
logger.info(
|
||||
@@ -198,7 +239,6 @@ def _handle_session_expired_and_retry(
|
||||
"signalling transport reconnect and retrying once.",
|
||||
server_name, op_description, exc,
|
||||
)
|
||||
|
||||
if not _core._signal_reconnect_and_wait(
|
||||
server_name,
|
||||
srv,
|
||||
@@ -211,7 +251,6 @@ def _handle_session_expired_and_retry(
|
||||
server_name,
|
||||
)
|
||||
return None
|
||||
|
||||
return _retry_once(server_name, retry_call, op_description, "session reconnect")
|
||||
|
||||
|
||||
@@ -240,18 +279,15 @@ def _handle_stdio_child_exited_and_retry(
|
||||
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"):
|
||||
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,
|
||||
)
|
||||
loop = _core._mcp_loop
|
||||
if loop is not None and loop.is_running():
|
||||
if _mcp_loop_running():
|
||||
reconnected = _core._signal_reconnect_and_wait(
|
||||
server_name,
|
||||
srv,
|
||||
@@ -263,49 +299,44 @@ def _handle_stdio_child_exited_and_retry(
|
||||
# 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
|
||||
if not reconnected:
|
||||
_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."
|
||||
)
|
||||
|
||||
_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."
|
||||
)
|
||||
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,
|
||||
)
|
||||
_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)}"
|
||||
))
|
||||
|
||||
|
||||
def _interrupted_call_result() -> str:
|
||||
@@ -348,124 +379,224 @@ async def _track_inflight_rpc(server: Any, server_name: str, op: str):
|
||||
inflight.discard(task)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
try:
|
||||
result = call_once()
|
||||
return _record_call_outcome(server_name, result) if record_outcome else result
|
||||
except InterruptedError:
|
||||
return _interrupted_call_result()
|
||||
except Exception as exc:
|
||||
for recover in recoverers:
|
||||
recovered = recover(server_name, exc, call_once, op)
|
||||
if recovered is not None:
|
||||
return recovered
|
||||
on_final_failure(exc)
|
||||
return _generic_call_failure(exc)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
_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"
|
||||
)
|
||||
_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.
|
||||
return await _call_coro if asyncio.iscoroutine(_call_coro) else _call_coro
|
||||
|
||||
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()
|
||||
raise _StdioChildExited(
|
||||
f"MCP stdio subprocess for "
|
||||
f"'{server_name}' exited mid-call"
|
||||
)
|
||||
return 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)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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."""
|
||||
parts: List[str] = []
|
||||
for block in (result.content or []):
|
||||
if hasattr(block, "text") and block.text:
|
||||
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)
|
||||
)
|
||||
if rendered:
|
||||
parts.append(rendered)
|
||||
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)
|
||||
# Hard-cap pathological payloads; ordinary large results pass to spillover.
|
||||
return _truncate_mcp_text_result("\n".join(parts))
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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"
|
||||
)))
|
||||
|
||||
text_result = _render_content_blocks(result, server_name)
|
||||
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 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
|
||||
if structured is not None:
|
||||
if text_result:
|
||||
payload["structuredContent"] = structured
|
||||
else:
|
||||
payload["result"] = structured
|
||||
if meta is not None:
|
||||
payload["_meta"] = meta
|
||||
payload.setdefault("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)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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)),
|
||||
):
|
||||
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; "
|
||||
f"reconnect requested. Do NOT retry this tool "
|
||||
f"immediately — give it a few seconds to come back."
|
||||
)
|
||||
return None, not_connected
|
||||
|
||||
|
||||
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."""
|
||||
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.
|
||||
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")
|
||||
error = _trust_gate_check(server_name, tool_name) or _check_circuit_breaker(server_name)
|
||||
if error is not None:
|
||||
return error
|
||||
server, error = _acquire_call_server(server_name, tool_timeout)
|
||||
if server is None:
|
||||
return error
|
||||
|
||||
async def _call():
|
||||
_mark_server_call_started(server)
|
||||
async with server._rpc_lock, _track_inflight_rpc(
|
||||
server, server_name, f"tools/call {tool_name}"
|
||||
):
|
||||
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.
|
||||
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
|
||||
)
|
||||
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
|
||||
@@ -473,129 +604,22 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float):
|
||||
_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)
|
||||
return _render_call_tool_result(result, server_name)
|
||||
|
||||
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
|
||||
|
||||
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 tool_error(_sanitize_error(
|
||||
f"MCP call failed: {type(exc).__name__}: {_exc_str(exc)}"
|
||||
))
|
||||
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,
|
||||
)
|
||||
|
||||
return _handler
|
||||
|
||||
@@ -613,7 +637,6 @@ def _make_utility_handler(server_name: str, tool_timeout: float, op: 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
|
||||
@@ -621,19 +644,14 @@ def _make_utility_handler(server_name: str, tool_timeout: float, op: str,
|
||||
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
|
||||
def _on_failure(exc):
|
||||
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 _invoke_with_recovery(
|
||||
server_name, _call_once, op,
|
||||
(_handle_auth_error_and_retry, _handle_session_expired_and_retry),
|
||||
_on_failure,
|
||||
)
|
||||
|
||||
return _handler
|
||||
|
||||
@@ -681,8 +699,7 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float):
|
||||
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:
|
||||
for block in getattr(result, "contents", []):
|
||||
if getattr(block, "text", None) is not None:
|
||||
parts.append(strip_unicode_tags(block.text))
|
||||
elif getattr(block, "blob", None) is not None:
|
||||
@@ -693,7 +710,7 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float):
|
||||
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 json.dumps({"result": "\n".join(parts)}, ensure_ascii=False)
|
||||
return _call
|
||||
|
||||
return _make_utility_handler(server_name, tool_timeout, "resources/read", "read_resource", _build)
|
||||
@@ -746,18 +763,16 @@ def _make_get_prompt_handler(server_name: str, tool_timeout: float):
|
||||
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 []):
|
||||
for msg in getattr(result, "messages", []):
|
||||
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))
|
||||
text = content.text if hasattr(content, "text") else (
|
||||
content if isinstance(content, str) else str(content)
|
||||
)
|
||||
entry["content"] = strip_unicode_tags(text)
|
||||
messages.append(entry)
|
||||
resp = {"messages": messages}
|
||||
if hasattr(result, "description") and result.description:
|
||||
|
||||
Reference in New Issue
Block a user