From 2bfc0f2c5d420f902e2d3b4cfa4de0c76421beb4 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Sat, 4 Apr 2026 16:06:42 -0700 Subject: [PATCH] fix: harden MCP client against misbehaving servers (#296) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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). --- docs/architecture.md | 15 + docs/diagrams/20-mcp-architecture.puml | 29 +- docs/diagrams/png/20-mcp-architecture.png | 4 +- tests/test_mcp_client.py | 408 +++++++++++++++++++++- turnstone/core/mcp_client.py | 340 +++++++++++++++++- 5 files changed, 782 insertions(+), 14 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index cf831359..1f211516 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -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. diff --git a/docs/diagrams/20-mcp-architecture.puml b/docs/diagrams/20-mcp-architecture.puml index d098d240..ca39d9d9 100644 --- a/docs/diagrams/20-mcp-architecture.puml +++ b/docs/diagrams/20-mcp-architecture.puml @@ -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 diff --git a/docs/diagrams/png/20-mcp-architecture.png b/docs/diagrams/png/20-mcp-architecture.png index 2193e9eb..34cdf4df 100644 --- a/docs/diagrams/png/20-mcp-architecture.png +++ b/docs/diagrams/png/20-mcp-architecture.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a6b7769aa7e732ffbeb1eb7f5b65273a135fb3a78d9802ec36d3b92801c34f6b -size 427745 +oid sha256:7623df33be9baf7647ca1c2450640df57e1cd73e8be1f8168aae16e546ad683c +size 459941 diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 459f7b20..fb275da8 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -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 diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index b0fb50e8..37c16c3e 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -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: