mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
fix: harden MCP client against misbehaving servers (#296)
* fix: harden MCP client against misbehaving servers Misbehaving/failed/misconfigured MCP servers could peg CPU at 100% due to anyio cancel-scope busy-loops (SDK #2147), uncancelled orphaned futures, and missing application-layer resilience. Five fixes: 1. Cancel orphaned futures on timeout — future.cancel() in all sync bridge methods prevents coroutine accumulation on the event loop 2. Per-server circuit breaker — 3-failure threshold with exponential cooldown (30s–5min), per-server jitter, auto-reconnect on half-open probe, McpError excluded (protocol errors from healthy servers) 3. Safe transport stream pre-close — store stream refs and close them before stack teardown in all error/shutdown paths, preventing the anyio zero-buffer CPU busy-loop 4. Notification debounce — 5s per-server rate limit on list_changed refresh storms from buggy servers 5. Periodic refresh backoff with auto-reconnect — disconnected servers get reconnection attempts with exponential backoff (60s–1hr) instead of being silently skipped forever * docs: add MCP resilience section to architecture docs and diagram Document the circuit breaker, future cancellation, stream pre-close, notification debounce, and periodic refresh backoff in the architecture guide and the MCP architecture PlantUML diagram. * fix: address review — stack leak on transport error, half-open comment - Widen _connect_one guard to check _per_server_stacks too, not just _sessions. Transport errors in sync dispatch methods evict the session but left the stack behind, leaking anyio tasks on reconnect. - Clarify half-open design: multiple callers are intentionally allowed through (reconnects serialize on the event loop, first failure re-trips).
This commit is contained in:
@@ -547,6 +547,21 @@ expanded tools).
|
||||
**Tool naming:** `mcp__{server}__{tool}` — double underscore delimiter, validated
|
||||
at connection time (server names with `__` are rejected).
|
||||
|
||||
**Resilience:** Each MCP server has an independent circuit breaker that opens
|
||||
after 3 consecutive transport failures (timeouts, broken pipes, connection
|
||||
resets). Cooldown uses capped exponential backoff (30 s base, 5 min max) with
|
||||
per-server jitter to avoid thundering herd. Protocol-level errors (`McpError`)
|
||||
from a healthy connection do not trip the breaker. When the cooldown expires
|
||||
(half-open), the next operation attempt triggers automatic reconnection. Manual
|
||||
`/mcp refresh` also clears the circuit on success. All sync bridge methods
|
||||
(`call_tool_sync`, `read_resource_sync`, `get_prompt_sync`, `refresh_sync`)
|
||||
cancel orphaned futures on timeout to prevent coroutine accumulation on the
|
||||
background event loop. Push notification refreshes are debounced (5 s per
|
||||
server) to protect against notification storms. The periodic refresh loop
|
||||
attempts reconnection for disconnected servers with exponential backoff
|
||||
(60 s–1 h). Transport stream references are pre-closed before stack teardown to
|
||||
work around the MCP SDK's anyio cancel-scope CPU busy-loop (SDK #2147).
|
||||
|
||||
**Error isolation:** Per-server connection/refresh failures are caught and logged; other
|
||||
servers are unaffected. Tool execution errors return error strings to the LLM
|
||||
rather than crashing the session.
|
||||
|
||||
@@ -152,10 +152,34 @@ MCPMgr -> MCPSrv : prompts/get
|
||||
MCPSrv --> MCPMgr : GetPromptResult
|
||||
MCPMgr --> Session : messages [{role, content}]
|
||||
|
||||
== Resilience: Circuit Breaker & Stream Safety ==
|
||||
|
||||
note over MCPMgr
|
||||
**Per-server circuit breaker**
|
||||
CLOSED --(3 failures)--> OPEN
|
||||
OPEN --(cooldown expires)--> half-open probe
|
||||
Probe success --> CLOSED (trip_count decays by 1)
|
||||
Probe failure --> OPEN (cooldown doubles, max 5 min)
|
||||
|
||||
McpError (protocol) does NOT trip breaker.
|
||||
BrokenPipeError / EOFError evicts dead session.
|
||||
All sync methods cancel orphaned futures on timeout.
|
||||
Transport streams pre-closed before stack teardown
|
||||
to avoid anyio cancel-scope CPU busy-loop (SDK #2147).
|
||||
end note
|
||||
|
||||
Session -> MCPMgr : call_tool_sync()
|
||||
MCPMgr -> MCPMgr : _cb_gate(server)\n[reject if circuit open]
|
||||
MCPMgr -> MCPMgr : _cb_auto_reconnect()\n[if session gone + cooldown expired]
|
||||
MCPMgr -> MCPSrv : tools/call
|
||||
MCPSrv --> MCPMgr : result or error
|
||||
MCPMgr -> MCPMgr : _cb_record_success()\nor _cb_record_failure()
|
||||
|
||||
== Three-Tier Refresh ==
|
||||
|
||||
group Push Notifications
|
||||
group Push Notifications (debounced 5s per server)
|
||||
MCPSrv -> MCPMgr : ToolListChangedNotification
|
||||
MCPMgr -> MCPMgr : debounce check\n(skip if < 5s since last)
|
||||
MCPMgr -> MCPMgr : _refresh_server_tools()
|
||||
|
||||
MCPSrv -> MCPMgr : ResourceListChangedNotification
|
||||
@@ -172,6 +196,9 @@ group Periodic Polling (default 4h)
|
||||
Only polls capabilities
|
||||
without push support.
|
||||
Staggered per-server.
|
||||
Disconnected servers get
|
||||
reconnect attempts with
|
||||
exponential backoff (60s-1h).
|
||||
end note
|
||||
end
|
||||
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a6b7769aa7e732ffbeb1eb7f5b65273a135fb3a78d9802ec36d3b92801c34f6b
|
||||
size 427745
|
||||
oid sha256:7623df33be9baf7647ca1c2450640df57e1cd73e8be1f8168aae16e546ad683c
|
||||
size 459941
|
||||
|
||||
+407
-1
@@ -3,7 +3,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import time
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
@@ -304,7 +306,7 @@ class TestMCPClientManager:
|
||||
def test_call_tool_sync_disconnected_server(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tool_map["mcp__dead__ping"] = ("dead", "ping")
|
||||
# No session registered for "dead"
|
||||
# No session registered for "dead", no config/loop → reconnect fails
|
||||
with pytest.raises(RuntimeError, match="not connected"):
|
||||
mgr.call_tool_sync("mcp__dead__ping", {})
|
||||
|
||||
@@ -1553,3 +1555,407 @@ class TestSafeCloseStack:
|
||||
await MCPClientManager._safe_close_stack(stack)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 1: Cancel orphaned futures on timeout
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFutureCancellation:
|
||||
"""Verify future.cancel() is called when sync bridge methods time out."""
|
||||
|
||||
def _make_manager_with_session(self) -> MCPClientManager:
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
# Prevent auto-spec from creating async coroutines that trigger warnings
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mock_session.read_resource = MagicMock(return_value="sentinel")
|
||||
mock_session.get_prompt = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__search"] = ("test", "search")
|
||||
mgr._resource_map["file:///a.txt"] = ("test", "file:///a.txt")
|
||||
mgr._prompt_map["mcp__test__review"] = ("test", "review")
|
||||
return mgr
|
||||
|
||||
def test_call_tool_sync_cancels_future_on_timeout(self):
|
||||
mgr = self._make_manager_with_session()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__search", {"query": "x"}, timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
def test_read_resource_sync_cancels_future_on_timeout(self):
|
||||
mgr = self._make_manager_with_session()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.read_resource_sync("file:///a.txt", timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
def test_get_prompt_sync_cancels_future_on_timeout(self):
|
||||
mgr = self._make_manager_with_session()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.get_prompt_sync("mcp__test__review", timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
def test_refresh_sync_cancels_future_on_timeout(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._loop = MagicMock()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch.object(mgr, "_refresh_all", return_value=MagicMock()),
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.refresh_sync(timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 2: Per-server circuit breaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCircuitBreaker:
|
||||
"""Verify per-server circuit breaker behavior."""
|
||||
|
||||
def test_circuit_stays_closed_below_threshold(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._cb_record_failure("srv")
|
||||
mgr._cb_record_failure("srv")
|
||||
is_open, _ = mgr._cb_check("srv")
|
||||
assert not is_open
|
||||
|
||||
def test_circuit_opens_at_threshold(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
is_open, cooldown_expired = mgr._cb_check("srv")
|
||||
assert is_open
|
||||
assert not cooldown_expired # just opened, cooldown not expired
|
||||
|
||||
def test_circuit_half_open_after_cooldown(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
# Simulate cooldown expiry
|
||||
mgr._circuit_open_until["srv"] = time.monotonic() - 1
|
||||
is_open, cooldown_expired = mgr._cb_check("srv")
|
||||
assert is_open
|
||||
assert cooldown_expired
|
||||
|
||||
def test_circuit_resets_on_success(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
assert "srv" in mgr._circuit_open_until
|
||||
mgr._cb_record_success("srv")
|
||||
is_open, _ = mgr._cb_check("srv")
|
||||
assert not is_open
|
||||
assert mgr._consecutive_failures.get("srv") is None
|
||||
|
||||
def test_success_decays_trip_count(self):
|
||||
"""Success decays trip_count by 1 so flapping servers escalate backoff."""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._circuit_trip_count["srv"] = 3
|
||||
mgr._cb_record_success("srv")
|
||||
assert mgr._circuit_trip_count["srv"] == 2
|
||||
mgr._cb_record_success("srv")
|
||||
assert mgr._circuit_trip_count["srv"] == 1
|
||||
mgr._cb_record_success("srv")
|
||||
assert "srv" not in mgr._circuit_trip_count
|
||||
|
||||
def test_cooldown_is_exponential(self):
|
||||
mgr = MCPClientManager({})
|
||||
# First trip (trip_count starts at 0)
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
deadline1 = mgr._circuit_open_until["srv"]
|
||||
base1 = deadline1 - time.monotonic()
|
||||
# Reset circuit but keep trip_count at 1 (set by first trip)
|
||||
mgr._cb_record_success("srv")
|
||||
# trip_count decayed from 1 to 0 — manually set to 1 for test
|
||||
mgr._circuit_trip_count["srv"] = 1
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
deadline2 = mgr._circuit_open_until["srv"]
|
||||
base2 = deadline2 - time.monotonic()
|
||||
# Second trip should have longer cooldown (roughly 2x, within jitter)
|
||||
assert base2 > base1 * 1.5
|
||||
|
||||
def test_cooldown_capped_at_max(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._circuit_trip_count["srv"] = 100 # very high trip count
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
deadline = mgr._circuit_open_until["srv"]
|
||||
cooldown = deadline - time.monotonic()
|
||||
# Should not exceed max (300s) + 10% jitter = 330s
|
||||
assert cooldown <= mgr._CB_MAX_COOLDOWN * 1.11
|
||||
|
||||
def test_cb_gate_rejects_when_open(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
with pytest.raises(RuntimeError, match="circuit open"):
|
||||
mgr._cb_gate("srv")
|
||||
|
||||
def test_cb_gate_allows_after_cooldown(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
mgr._circuit_open_until["srv"] = time.monotonic() - 1
|
||||
# Should not raise
|
||||
mgr._cb_gate("srv")
|
||||
# Deadline should be removed (half-open probe allowed)
|
||||
assert "srv" not in mgr._circuit_open_until
|
||||
|
||||
def test_cb_clear_removes_all_state(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
mgr._cb_clear("srv")
|
||||
assert "srv" not in mgr._consecutive_failures
|
||||
assert "srv" not in mgr._circuit_open_until
|
||||
assert "srv" not in mgr._circuit_trip_count
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
|
||||
@pytest.mark.filterwarnings("ignore:coroutine.*was never awaited:RuntimeWarning")
|
||||
def test_call_tool_sync_records_failure_on_timeout(self):
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=1)
|
||||
assert mgr._consecutive_failures.get("test", 0) == 1
|
||||
|
||||
def test_call_tool_sync_records_success(self):
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
# Pre-set a failure
|
||||
mgr._consecutive_failures["test"] = 2
|
||||
mock_result = MagicMock()
|
||||
mock_result.content = []
|
||||
mock_result.isError = False
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.return_value = mock_result
|
||||
with patch("asyncio.run_coroutine_threadsafe", return_value=mock_future):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
||||
assert mgr._consecutive_failures.get("test") is None
|
||||
|
||||
def test_connection_error_evicts_session(self):
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = BrokenPipeError("dead")
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(BrokenPipeError),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
||||
assert "test" not in mgr._sessions
|
||||
|
||||
def test_independent_circuits_per_server(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("a")
|
||||
is_open_a, _ = mgr._cb_check("a")
|
||||
is_open_b, _ = mgr._cb_check("b")
|
||||
assert is_open_a
|
||||
assert not is_open_b
|
||||
|
||||
def test_mcp_error_does_not_trip_circuit(self):
|
||||
"""Protocol errors (McpError) should not count as transport failures."""
|
||||
from mcp import McpError
|
||||
from mcp.types import ErrorData
|
||||
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = McpError(ErrorData(code=-32601, message="tool not found"))
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(McpError),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
||||
# Circuit should NOT have recorded a failure
|
||||
assert mgr._consecutive_failures.get("test", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 3: Safe transport stream pre-close
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSafeTransportStreams:
|
||||
"""Verify stream references are stored and pre-closed."""
|
||||
|
||||
def test_pre_close_streams_closes_both(self):
|
||||
mgr = MCPClientManager({})
|
||||
stream_a = MagicMock()
|
||||
stream_b = MagicMock()
|
||||
mgr._server_streams["srv"] = (stream_a, stream_b)
|
||||
|
||||
async def _run():
|
||||
await mgr._pre_close_streams("srv")
|
||||
|
||||
asyncio.run(_run())
|
||||
stream_a.aclose.assert_called_once()
|
||||
stream_b.aclose.assert_called_once()
|
||||
assert "srv" not in mgr._server_streams
|
||||
|
||||
def test_pre_close_streams_ignores_missing(self):
|
||||
mgr = MCPClientManager({})
|
||||
|
||||
async def _run():
|
||||
await mgr._pre_close_streams("nonexistent")
|
||||
|
||||
asyncio.run(_run()) # should not raise
|
||||
|
||||
def test_pre_close_streams_suppresses_errors(self):
|
||||
mgr = MCPClientManager({})
|
||||
stream_a = MagicMock()
|
||||
stream_a.aclose.side_effect = RuntimeError("boom")
|
||||
stream_b = MagicMock()
|
||||
mgr._server_streams["srv"] = (stream_a, stream_b)
|
||||
|
||||
async def _run():
|
||||
await mgr._pre_close_streams("srv")
|
||||
|
||||
asyncio.run(_run()) # should not raise despite stream_a error
|
||||
stream_b.aclose.assert_called_once()
|
||||
|
||||
def test_shutdown_clears_stream_refs(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._server_streams["srv"] = (MagicMock(), MagicMock())
|
||||
mgr.shutdown()
|
||||
assert len(mgr._server_streams) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 4: Notification debounce
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestNotificationDebounce:
|
||||
"""Verify notification-triggered refreshes are debounced."""
|
||||
|
||||
def test_debounce_within_window(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._last_notification_refresh["srv"] = time.monotonic()
|
||||
# We can't easily call _on_notification (it's a closure), so test
|
||||
# the debounce logic directly via the timestamp check
|
||||
now = time.monotonic()
|
||||
last = mgr._last_notification_refresh.get("srv", 0.0)
|
||||
assert now - last < mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
def test_debounce_passes_after_window(self):
|
||||
mgr = MCPClientManager({})
|
||||
# Set timestamp well in the past
|
||||
mgr._last_notification_refresh["srv"] = time.monotonic() - 10
|
||||
now = time.monotonic()
|
||||
last = mgr._last_notification_refresh.get("srv", 0.0)
|
||||
assert now - last >= mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
def test_debounce_is_per_server(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._last_notification_refresh["srv_a"] = time.monotonic()
|
||||
# srv_b has no timestamp — should pass debounce
|
||||
now = time.monotonic()
|
||||
last_b = mgr._last_notification_refresh.get("srv_b", 0.0)
|
||||
assert now - last_b >= mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 5: Periodic refresh backoff
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPeriodicRefreshBackoff:
|
||||
"""Verify periodic refresh backoff and auto-reconnect."""
|
||||
|
||||
def test_backoff_set_on_failure(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._refresh_failures["srv"] = 1
|
||||
# Simulate what _periodic_refresh does on failure
|
||||
failures = mgr._refresh_failures.get("srv", 0) + 1
|
||||
mgr._refresh_failures["srv"] = failures
|
||||
backoff = min(mgr._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)), mgr._REFRESH_BACKOFF_MAX)
|
||||
mgr._refresh_backoff_until["srv"] = time.monotonic() + backoff
|
||||
assert mgr._refresh_backoff_until["srv"] > time.monotonic()
|
||||
assert failures == 2
|
||||
|
||||
def test_backoff_doubles(self):
|
||||
mgr = MCPClientManager({})
|
||||
b1 = min(mgr._REFRESH_BACKOFF_BASE * (2**0), mgr._REFRESH_BACKOFF_MAX)
|
||||
b2 = min(mgr._REFRESH_BACKOFF_BASE * (2**1), mgr._REFRESH_BACKOFF_MAX)
|
||||
b3 = min(mgr._REFRESH_BACKOFF_BASE * (2**2), mgr._REFRESH_BACKOFF_MAX)
|
||||
assert b1 == 60
|
||||
assert b2 == 120
|
||||
assert b3 == 240
|
||||
|
||||
def test_backoff_capped(self):
|
||||
mgr = MCPClientManager({})
|
||||
b = min(mgr._REFRESH_BACKOFF_BASE * (2**20), mgr._REFRESH_BACKOFF_MAX)
|
||||
assert b == mgr._REFRESH_BACKOFF_MAX
|
||||
|
||||
def test_backoff_clears_on_success(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._refresh_failures["srv"] = 3
|
||||
mgr._refresh_backoff_until["srv"] = time.monotonic() + 1000
|
||||
# Simulate success
|
||||
mgr._refresh_failures.pop("srv", None)
|
||||
mgr._refresh_backoff_until.pop("srv", None)
|
||||
assert "srv" not in mgr._refresh_failures
|
||||
assert "srv" not in mgr._refresh_backoff_until
|
||||
|
||||
def test_server_status_includes_circuit_info(self):
|
||||
mgr = MCPClientManager({"srv": {"type": "stdio", "command": "echo"}})
|
||||
status = mgr.get_server_status("srv")
|
||||
assert "circuit_open" in status
|
||||
assert "consecutive_failures" in status
|
||||
assert status["circuit_open"] is False
|
||||
assert status["consecutive_failures"] == 0
|
||||
|
||||
def test_server_status_shows_open_circuit(self):
|
||||
mgr = MCPClientManager({"srv": {"type": "stdio", "command": "echo"}})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
status = mgr.get_server_status("srv")
|
||||
assert status["circuit_open"] is True
|
||||
assert status["consecutive_failures"] == 3
|
||||
|
||||
+330
-10
@@ -35,7 +35,7 @@ if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
import mcp.types as mcp_types
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp import ClientSession, McpError, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
|
||||
@@ -151,6 +151,22 @@ class MCPClientManager:
|
||||
self._refresh_interval = refresh_interval
|
||||
self._refresh_task: asyncio.Task[None] | None = None
|
||||
|
||||
# Circuit breaker (per-server) — prevents repeated calls to broken servers
|
||||
self._consecutive_failures: dict[str, int] = {}
|
||||
self._circuit_open_until: dict[str, float] = {} # monotonic timestamp
|
||||
self._circuit_trip_count: dict[str, int] = {} # backoff exponent
|
||||
|
||||
# Safe transport stream refs (pre-close before stack teardown to avoid
|
||||
# the anyio cancel-scope CPU busy-loop — MCP SDK #2147)
|
||||
self._server_streams: dict[str, tuple[Any, Any]] = {}
|
||||
|
||||
# Notification debounce (per-server)
|
||||
self._last_notification_refresh: dict[str, float] = {}
|
||||
|
||||
# Periodic refresh backoff (per-server)
|
||||
self._refresh_failures: dict[str, int] = {}
|
||||
self._refresh_backoff_until: dict[str, float] = {} # monotonic timestamp
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
@@ -181,6 +197,7 @@ class MCPClientManager:
|
||||
except Exception as exc:
|
||||
log.warning("Failed to connect MCP server '%s'", name, exc_info=True)
|
||||
self._set_error(name, f"{type(exc).__name__}: {exc}")
|
||||
self._cb_record_failure(name)
|
||||
|
||||
self._connected.set()
|
||||
|
||||
@@ -203,6 +220,94 @@ class MCPClientManager:
|
||||
_CONNECT_TIMEOUT = 30 # seconds — prevents hung connections on broken remotes
|
||||
_TCP_PROBE_TIMEOUT = 5 # seconds — fast TCP pre-flight for HTTP transports
|
||||
|
||||
# Circuit breaker constants
|
||||
_CB_FAILURE_THRESHOLD = 3
|
||||
_CB_BASE_COOLDOWN = 30.0 # seconds
|
||||
_CB_MAX_COOLDOWN = 300.0 # 5 minutes
|
||||
|
||||
# Notification debounce
|
||||
_NOTIFICATION_DEBOUNCE = 5.0 # seconds between refreshes per server
|
||||
|
||||
# Periodic refresh backoff
|
||||
_REFRESH_BACKOFF_BASE = 60.0 # seconds
|
||||
_REFRESH_BACKOFF_MAX = 3600.0 # 1 hour
|
||||
|
||||
# -- circuit breaker (per-server) -----------------------------------------
|
||||
|
||||
def _cb_check(self, name: str) -> tuple[bool, bool]:
|
||||
"""Check circuit breaker state for *name*.
|
||||
|
||||
Returns ``(is_open, cooldown_expired)``. When the circuit is closed
|
||||
both values are False. When open, *cooldown_expired* indicates
|
||||
whether a probe attempt is allowed.
|
||||
"""
|
||||
deadline = self._circuit_open_until.get(name)
|
||||
if deadline is None:
|
||||
return False, False
|
||||
now = time.monotonic()
|
||||
if now >= deadline:
|
||||
return True, True # half-open: allow one probe
|
||||
return True, False # still in cooldown
|
||||
|
||||
def _cb_record_failure(self, name: str) -> None:
|
||||
"""Record a failure against *name*, potentially opening the circuit."""
|
||||
count = self._consecutive_failures.get(name, 0) + 1
|
||||
self._consecutive_failures[name] = count
|
||||
# Guard: don't extend an already-open deadline. Additional failures
|
||||
# while open still accumulate in _consecutive_failures, so the circuit
|
||||
# re-opens immediately after the next half-open probe fails (count is
|
||||
# already >= threshold).
|
||||
if count >= self._CB_FAILURE_THRESHOLD and name not in self._circuit_open_until:
|
||||
trips = self._circuit_trip_count.get(name, 0)
|
||||
cooldown = min(self._CB_BASE_COOLDOWN * (2**trips), self._CB_MAX_COOLDOWN)
|
||||
# Per-server jitter seeded from server name (varies across process
|
||||
# restarts via PYTHONHASHSEED, which is desirable — each cluster
|
||||
# node gets different jitter to avoid thundering herd).
|
||||
jitter = random.Random(hash(name)).random() * cooldown * 0.1
|
||||
self._circuit_open_until[name] = time.monotonic() + cooldown + jitter
|
||||
self._circuit_trip_count[name] = trips + 1
|
||||
log.warning(
|
||||
"MCP circuit open for '%s': %d consecutive failures, cooldown %.0fs",
|
||||
name,
|
||||
count,
|
||||
cooldown + jitter,
|
||||
)
|
||||
|
||||
def _cb_record_success(self, name: str) -> None:
|
||||
"""Record a successful operation for *name*, decaying circuit state.
|
||||
|
||||
Decays trip count by 1 rather than resetting to 0, so a chronically
|
||||
flapping server escalates its backoff over time instead of always
|
||||
restarting at the minimum cooldown.
|
||||
"""
|
||||
self._consecutive_failures.pop(name, None)
|
||||
self._circuit_open_until.pop(name, None)
|
||||
trips = self._circuit_trip_count.get(name, 0)
|
||||
if trips > 1:
|
||||
self._circuit_trip_count[name] = trips - 1
|
||||
else:
|
||||
self._circuit_trip_count.pop(name, None)
|
||||
|
||||
def _cb_clear(self, name: str) -> None:
|
||||
"""Remove all circuit breaker state for *name*."""
|
||||
self._consecutive_failures.pop(name, None)
|
||||
self._circuit_open_until.pop(name, None)
|
||||
self._circuit_trip_count.pop(name, None)
|
||||
|
||||
# -- safe transport helpers ------------------------------------------------
|
||||
|
||||
async def _pre_close_streams(self, name: str) -> None:
|
||||
"""Close MCP transport streams before stack teardown.
|
||||
|
||||
Pre-closing unblocks anyio transport tasks stuck on zero-buffer
|
||||
``send()`` calls, preventing the CPU busy-loop from SDK #2147.
|
||||
"""
|
||||
streams = self._server_streams.pop(name, None)
|
||||
if streams:
|
||||
for s in streams:
|
||||
with contextlib.suppress(Exception):
|
||||
await s.aclose()
|
||||
|
||||
async def _tcp_probe(self, name: str, url: str) -> None:
|
||||
"""Fast TCP connect check before entering the MCP transport context.
|
||||
|
||||
@@ -251,6 +356,16 @@ class MCPClientManager:
|
||||
log.error("MCP server name '%s' contains '__' (reserved delimiter), skipping", name)
|
||||
return
|
||||
|
||||
# Guard: tear down stale session/stack so we don't leak. Checks both
|
||||
# _sessions and _per_server_stacks because transport errors in the sync
|
||||
# dispatch methods evict the session but leave the stack behind.
|
||||
if name in self._sessions or name in self._per_server_stacks:
|
||||
self._sessions.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
old_stack = self._per_server_stacks.pop(name, None)
|
||||
if old_stack:
|
||||
await self._safe_close_stack(old_stack)
|
||||
|
||||
# Per-server exit stack for clean per-server lifecycle management
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
@@ -271,6 +386,9 @@ class MCPClientManager:
|
||||
),
|
||||
timeout=self._CONNECT_TIMEOUT,
|
||||
)
|
||||
# Stash stream refs so _pre_close_streams can unblock anyio
|
||||
# transport tasks before the cancel scope fires (SDK #2147).
|
||||
self._server_streams[name] = (read, write)
|
||||
else:
|
||||
# Default: stdio transport
|
||||
command = cfg.get("command", "")
|
||||
@@ -287,24 +405,29 @@ class MCPClientManager:
|
||||
env=env,
|
||||
)
|
||||
read, write = await stack.enter_async_context(stdio_client(params))
|
||||
self._server_streams[name] = (read, write)
|
||||
except asyncio.CancelledError:
|
||||
# Stray CancelledError from broken anyio cancel scope — treat as
|
||||
# Stray CancelledError from broken anyio cancel scope -- treat as
|
||||
# connection failure. But if the task is genuinely being cancelled
|
||||
# (shutdown), re-raise so we don't block teardown.
|
||||
task = asyncio.current_task()
|
||||
if task is not None and task.cancelling():
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
log.warning("MCP server '%s' connection failed (anyio cancel)", name)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"Connection failed for '{name}'") from None
|
||||
except TimeoutError:
|
||||
log.warning(
|
||||
"MCP server '%s' connection timed out after %ds", name, self._CONNECT_TIMEOUT
|
||||
)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"Connection timed out after {self._CONNECT_TIMEOUT}s") from None
|
||||
except Exception:
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
|
||||
@@ -316,15 +439,30 @@ class MCPClientManager:
|
||||
if not isinstance(msg, mcp_types.ServerNotification):
|
||||
return
|
||||
root = msg.root
|
||||
|
||||
# Debounce: skip if we refreshed this server very recently
|
||||
now = time.monotonic()
|
||||
last = self._last_notification_refresh.get(name, 0.0)
|
||||
if now - last < self._NOTIFICATION_DEBOUNCE:
|
||||
log.debug(
|
||||
"Debouncing notification from '%s' (%.1fs since last refresh)",
|
||||
name,
|
||||
now - last,
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
if isinstance(root, mcp_types.ToolListChangedNotification):
|
||||
log.info("Received tools/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_tools(name)
|
||||
elif isinstance(root, mcp_types.ResourceListChangedNotification):
|
||||
log.info("Received resources/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_resources(name)
|
||||
elif isinstance(root, mcp_types.PromptListChangedNotification):
|
||||
log.info("Received prompts/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_prompts(name)
|
||||
self._last_error.pop(name, None)
|
||||
except Exception as exc:
|
||||
@@ -336,6 +474,7 @@ class MCPClientManager:
|
||||
ClientSession(read, write, message_handler=_on_notification) # type: ignore[arg-type]
|
||||
)
|
||||
except Exception:
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
|
||||
@@ -346,16 +485,20 @@ class MCPClientManager:
|
||||
self._per_server_stacks.pop(name, None)
|
||||
task = asyncio.current_task()
|
||||
if task is not None and task.cancelling():
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"MCP handshake failed for '{name}'") from None
|
||||
except TimeoutError:
|
||||
self._per_server_stacks.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"MCP handshake timed out after {self._CONNECT_TIMEOUT}s") from None
|
||||
except Exception:
|
||||
self._per_server_stacks.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
self._sessions[name] = session
|
||||
@@ -548,12 +691,14 @@ class MCPClientManager:
|
||||
if cfg:
|
||||
log.info("Reconnecting MCP server '%s'", name)
|
||||
await self._connect_one(name, cfg)
|
||||
self._cb_record_success(name)
|
||||
new_names = [
|
||||
t["function"]["name"] for t in self._per_server_tools.get(name, [])
|
||||
]
|
||||
results[name] = (new_names, [])
|
||||
continue
|
||||
added, removed = await self._refresh_server(name)
|
||||
self._cb_record_success(name)
|
||||
results[name] = (added, removed)
|
||||
except Exception as exc:
|
||||
log.warning("Refresh failed for MCP server '%s'", name, exc_info=True)
|
||||
@@ -577,10 +722,18 @@ class MCPClientManager:
|
||||
"""
|
||||
assert self._loop is not None
|
||||
future = asyncio.run_coroutine_threadsafe(self._refresh_all(server_name), self._loop)
|
||||
return future.result(timeout=timeout)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
raise TimeoutError(f"MCP refresh timed out after {timeout}s") from None
|
||||
|
||||
async def _periodic_refresh(self) -> None:
|
||||
"""Periodically refresh servers that lack push notifications."""
|
||||
"""Periodically refresh servers that lack push notifications.
|
||||
|
||||
Applies per-server exponential backoff on failure and attempts
|
||||
reconnection for disconnected servers.
|
||||
"""
|
||||
# Stagger start using a launch-time seed so cluster nodes don't
|
||||
# all hit MCP servers simultaneously.
|
||||
seed = random.Random(time.monotonic_ns() ^ os.getpid()).random()
|
||||
@@ -588,8 +741,42 @@ class MCPClientManager:
|
||||
await asyncio.sleep(initial_delay)
|
||||
while True:
|
||||
for name in list(self._server_configs):
|
||||
now = time.monotonic()
|
||||
|
||||
# Check per-server backoff
|
||||
backoff_until = self._refresh_backoff_until.get(name, 0.0)
|
||||
if now < backoff_until:
|
||||
continue # still in backoff
|
||||
|
||||
if name not in self._sessions:
|
||||
continue # not connected — skip (reconnect on manual refresh)
|
||||
# Attempt reconnection for disconnected servers
|
||||
cfg = self._server_configs.get(name)
|
||||
if cfg:
|
||||
try:
|
||||
log.info("Periodic reconnect attempt for '%s'", name)
|
||||
await self._connect_one(name, cfg)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
self._cb_record_success(name)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
failures = self._refresh_failures.get(name, 0) + 1
|
||||
self._refresh_failures[name] = failures
|
||||
backoff = min(
|
||||
self._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)),
|
||||
self._REFRESH_BACKOFF_MAX,
|
||||
)
|
||||
self._refresh_backoff_until[name] = time.monotonic() + backoff
|
||||
log.warning(
|
||||
"Periodic reconnect failed for '%s' (attempt %d, backoff %.0fs)",
|
||||
name,
|
||||
failures,
|
||||
backoff,
|
||||
)
|
||||
self._set_error(name, f"Reconnect failed: {exc}")
|
||||
continue
|
||||
|
||||
try:
|
||||
if not self._supports_list_changed.get(name, False):
|
||||
await self._refresh_server_tools(name)
|
||||
@@ -598,9 +785,26 @@ class MCPClientManager:
|
||||
if not self._supports_prompt_list_changed.get(name, False):
|
||||
await self._refresh_server_prompts(name)
|
||||
self._last_error.pop(name, None)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
except Exception as exc:
|
||||
log.warning("Periodic refresh failed for '%s'", name, exc_info=True)
|
||||
failures = self._refresh_failures.get(name, 0) + 1
|
||||
self._refresh_failures[name] = failures
|
||||
backoff = min(
|
||||
self._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)),
|
||||
self._REFRESH_BACKOFF_MAX,
|
||||
)
|
||||
self._refresh_backoff_until[name] = time.monotonic() + backoff
|
||||
log.warning(
|
||||
"Periodic refresh failed for '%s' (attempt %d, backoff %.0fs)",
|
||||
name,
|
||||
failures,
|
||||
backoff,
|
||||
)
|
||||
self._set_error(name, f"Periodic refresh failed: {exc}")
|
||||
# Note: per-server backoff (max 1h) is only meaningful when
|
||||
# refresh_interval is shorter than _REFRESH_BACKOFF_MAX. With
|
||||
# the default 4h interval this sleep already bounds retry frequency.
|
||||
await asyncio.sleep(self._refresh_interval)
|
||||
|
||||
# -- resource refresh ----------------------------------------------------
|
||||
@@ -940,6 +1144,9 @@ class MCPClientManager:
|
||||
if self._loop and self._per_server_stacks:
|
||||
|
||||
async def _close_all_stacks() -> None:
|
||||
# Pre-close streams to prevent anyio CPU busy-loop during teardown
|
||||
for srv_name in list(self._server_streams):
|
||||
await self._pre_close_streams(srv_name)
|
||||
for stack in self._per_server_stacks.values():
|
||||
await self._safe_close_stack(stack)
|
||||
|
||||
@@ -985,6 +1192,14 @@ class MCPClientManager:
|
||||
self._listeners.clear()
|
||||
self._resource_listeners.clear()
|
||||
self._prompt_listeners.clear()
|
||||
# Clear resilience state
|
||||
self._consecutive_failures.clear()
|
||||
self._circuit_open_until.clear()
|
||||
self._circuit_trip_count.clear()
|
||||
self._server_streams.clear()
|
||||
self._last_notification_refresh.clear()
|
||||
self._refresh_failures.clear()
|
||||
self._refresh_backoff_until.clear()
|
||||
|
||||
log.info("MCP client shut down")
|
||||
|
||||
@@ -1049,6 +1264,7 @@ class MCPClientManager:
|
||||
async def _remove() -> None:
|
||||
# Close session + transport via per-server stack
|
||||
self._sessions.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
stack = self._per_server_stacks.pop(name, None)
|
||||
if stack is not None:
|
||||
await self._safe_close_stack(stack)
|
||||
@@ -1062,6 +1278,10 @@ class MCPClientManager:
|
||||
self._supports_prompts.pop(name, None)
|
||||
self._supports_prompt_list_changed.pop(name, None)
|
||||
self._last_error.pop(name, None)
|
||||
self._last_notification_refresh.pop(name, None)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
self._cb_clear(name)
|
||||
# Rebuild merged state (serialized with notification handlers)
|
||||
self._rebuild_tools()
|
||||
self._rebuild_resources()
|
||||
@@ -1075,6 +1295,7 @@ class MCPClientManager:
|
||||
else:
|
||||
# No event loop (tests / pre-start) — mutate directly
|
||||
self._sessions.pop(name, None)
|
||||
self._server_streams.pop(name, None)
|
||||
self._per_server_tools.pop(name, None)
|
||||
self._per_server_resources.pop(name, None)
|
||||
self._per_server_prompts.pop(name, None)
|
||||
@@ -1084,6 +1305,10 @@ class MCPClientManager:
|
||||
self._supports_prompts.pop(name, None)
|
||||
self._supports_prompt_list_changed.pop(name, None)
|
||||
self._last_error.pop(name, None)
|
||||
self._last_notification_refresh.pop(name, None)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
self._cb_clear(name)
|
||||
self._rebuild_tools()
|
||||
self._rebuild_resources()
|
||||
self._rebuild_prompts()
|
||||
@@ -1107,6 +1332,8 @@ class MCPClientManager:
|
||||
connected = name in self._sessions
|
||||
cfg = self._server_configs.get(name, {})
|
||||
transport = cfg.get("type", "stdio")
|
||||
cb_deadline = self._circuit_open_until.get(name)
|
||||
cb_open = cb_deadline is not None and time.monotonic() < cb_deadline
|
||||
return {
|
||||
"connected": connected,
|
||||
"tools": len(self._per_server_tools.get(name, [])) if connected else 0,
|
||||
@@ -1116,6 +1343,8 @@ class MCPClientManager:
|
||||
"transport": transport,
|
||||
"command": cfg.get("command", "") if transport == "stdio" else "",
|
||||
"url": cfg.get("url", "") if transport != "stdio" else "",
|
||||
"circuit_open": cb_open,
|
||||
"consecutive_failures": self._consecutive_failures.get(name, 0),
|
||||
}
|
||||
|
||||
def get_all_server_status(self) -> dict[str, dict[str, Any]]:
|
||||
@@ -1245,6 +1474,55 @@ class MCPClientManager:
|
||||
|
||||
# -- tool invocation -----------------------------------------------------
|
||||
|
||||
def _cb_gate(self, server_name: str) -> None:
|
||||
"""Check circuit breaker before dispatching to *server_name*.
|
||||
|
||||
Raises ``RuntimeError`` if the circuit is open and cooldown has not
|
||||
expired. When the cooldown has expired (half-open), clears the
|
||||
deadline so the probe attempt is allowed through.
|
||||
"""
|
||||
is_open, cooldown_expired = self._cb_check(server_name)
|
||||
if is_open and not cooldown_expired:
|
||||
remaining = self._circuit_open_until.get(server_name, 0) - time.monotonic()
|
||||
raise RuntimeError(
|
||||
f"MCP server '{server_name}' circuit open "
|
||||
f"(cooldown {remaining:.0f}s remaining). "
|
||||
f"Use '/mcp refresh {server_name}' to retry manually."
|
||||
)
|
||||
if cooldown_expired:
|
||||
# Remove deadline so concurrent callers aren't rejected while the
|
||||
# probe is in-flight. This intentionally allows multiple callers
|
||||
# through rather than a single probe: reconnects serialize on the
|
||||
# event loop via _connect_one's guard, and if the server is truly
|
||||
# broken the first failure re-trips the circuit immediately.
|
||||
self._circuit_open_until.pop(server_name, None)
|
||||
|
||||
def _cb_auto_reconnect(self, server_name: str) -> Any:
|
||||
"""Attempt reconnection for a disconnected server during half-open probe.
|
||||
|
||||
Returns the new session on success, or raises on failure.
|
||||
"""
|
||||
cfg = self._server_configs.get(server_name)
|
||||
if not cfg or self._loop is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
reconnect_future = asyncio.run_coroutine_threadsafe(
|
||||
self._connect_one(server_name, cfg), self._loop
|
||||
)
|
||||
try:
|
||||
reconnect_future.result(timeout=self._CONNECT_TIMEOUT)
|
||||
except concurrent.futures.TimeoutError:
|
||||
reconnect_future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise RuntimeError(f"MCP server '{server_name}' reconnect timed out") from None
|
||||
except Exception as exc:
|
||||
self._cb_record_failure(server_name)
|
||||
raise RuntimeError(f"MCP server '{server_name}' reconnect failed: {exc}") from None
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
self._cb_record_failure(server_name)
|
||||
raise RuntimeError(f"MCP server '{server_name}' reconnect produced no session")
|
||||
return session
|
||||
|
||||
def call_tool_sync(
|
||||
self,
|
||||
func_name: str,
|
||||
@@ -1254,15 +1532,19 @@ class MCPClientManager:
|
||||
"""Execute an MCP tool call synchronously (blocks the calling thread).
|
||||
|
||||
Dispatches an async ``tools/call`` to the background event loop and
|
||||
waits for the result.
|
||||
waits for the result. Includes circuit-breaker gating and automatic
|
||||
reconnection for servers recovering from failure.
|
||||
"""
|
||||
mapping = self._tool_map.get(func_name)
|
||||
if mapping is None:
|
||||
raise ValueError(f"Unknown MCP tool: {func_name}")
|
||||
server_name, original_name = mapping
|
||||
|
||||
self._cb_gate(server_name)
|
||||
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
session = self._cb_auto_reconnect(server_name)
|
||||
assert self._loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
@@ -1271,7 +1553,19 @@ class MCPClientManager:
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise TimeoutError(f"MCP tool call timed out after {timeout}s") from None
|
||||
except Exception as exc:
|
||||
# Protocol errors (McpError) come from a healthy connection that
|
||||
# rejected the request — only transport errors trip the breaker.
|
||||
if not isinstance(exc, McpError):
|
||||
self._cb_record_failure(server_name)
|
||||
if isinstance(exc, (BrokenPipeError, ConnectionResetError, EOFError)):
|
||||
self._sessions.pop(server_name, None)
|
||||
raise
|
||||
|
||||
self._cb_record_success(server_name)
|
||||
|
||||
# Extract text from the content array
|
||||
texts: list[str] = []
|
||||
@@ -1320,16 +1614,29 @@ class MCPClientManager:
|
||||
if mapping is None:
|
||||
raise ValueError(f"Unknown MCP resource: {uri}")
|
||||
server_name, _ = mapping
|
||||
|
||||
self._cb_gate(server_name)
|
||||
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
session = self._cb_auto_reconnect(server_name)
|
||||
assert self._loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(session.read_resource(uri), self._loop)
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise TimeoutError(f"MCP resource read timed out after {timeout}s") from None
|
||||
except Exception as exc:
|
||||
if not isinstance(exc, McpError):
|
||||
self._cb_record_failure(server_name)
|
||||
if isinstance(exc, (BrokenPipeError, ConnectionResetError, EOFError)):
|
||||
self._sessions.pop(server_name, None)
|
||||
raise
|
||||
|
||||
self._cb_record_success(server_name)
|
||||
|
||||
parts: list[str] = []
|
||||
for item in result.contents:
|
||||
@@ -1357,9 +1664,12 @@ class MCPClientManager:
|
||||
if mapping is None:
|
||||
raise ValueError(f"Unknown MCP prompt: {prefixed_name}")
|
||||
server_name, original_name = mapping
|
||||
|
||||
self._cb_gate(server_name)
|
||||
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
session = self._cb_auto_reconnect(server_name)
|
||||
assert self._loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
@@ -1368,7 +1678,17 @@ class MCPClientManager:
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise TimeoutError(f"MCP prompt retrieval timed out after {timeout}s") from None
|
||||
except Exception as exc:
|
||||
if not isinstance(exc, McpError):
|
||||
self._cb_record_failure(server_name)
|
||||
if isinstance(exc, (BrokenPipeError, ConnectionResetError, EOFError)):
|
||||
self._sessions.pop(server_name, None)
|
||||
raise
|
||||
|
||||
self._cb_record_success(server_name)
|
||||
|
||||
messages: list[dict[str, Any]] = []
|
||||
for msg in result.messages:
|
||||
|
||||
Reference in New Issue
Block a user