From 513f9dc211ce545fffa0eeef483e9602b0e45aa3 Mon Sep 17 00:00:00 2001 From: Marin Pesa Date: Sat, 16 May 2026 14:44:35 -0400 Subject: [PATCH] fix: reconnect stale MCP sessions before retry --- tests/tools/test_mcp_circuit_breaker.py | 29 +++- tests/tools/test_mcp_tool_session_expired.py | 108 ++++++++++++- tools/mcp_tool.py | 160 ++++++++++++++----- 3 files changed, 247 insertions(+), 50 deletions(-) diff --git a/tests/tools/test_mcp_circuit_breaker.py b/tests/tools/test_mcp_circuit_breaker.py index 0173fa52afe26..6d39881b06628 100644 --- a/tests/tools/test_mcp_circuit_breaker.py +++ b/tests/tools/test_mcp_circuit_breaker.py @@ -31,14 +31,37 @@ def _install_stub_server(mcp_tool_module, name: str, call_tool_impl): ``call_tool_impl`` is an async function stored at ``session.call_tool`` (it's what the tool handler invokes). """ + import threading + server = MagicMock() server.name = name session = MagicMock() session.call_tool = call_tool_impl server.session = session - server._reconnect_event = MagicMock() - server._ready = MagicMock() - server._ready.is_set.return_value = True + + ready_flag = threading.Event() + ready_flag.set() + + class _ReadyAdapter: + def is_set(self): + return ready_flag.is_set() + + def clear(self): + ready_flag.clear() + + def set(self): + ready_flag.set() + + class _ReconnectAdapter: + def set(self): + old_session = server.session + new_session = MagicMock() + new_session.call_tool = old_session.call_tool + server.session = new_session + ready_flag.set() + + server._reconnect_event = _ReconnectAdapter() + server._ready = _ReadyAdapter() mcp_tool_module._servers[name] = server mcp_tool_module._server_error_counts.pop(name, None) diff --git a/tests/tools/test_mcp_tool_session_expired.py b/tests/tools/test_mcp_tool_session_expired.py index 59601ba1c3d7c..fc4b108b806e4 100644 --- a/tests/tools/test_mcp_tool_session_expired.py +++ b/tests/tools/test_mcp_tool_session_expired.py @@ -113,23 +113,47 @@ def _install_stub_server(name: str = "wpcom"): server = MagicMock() server.name = name + + ready_flag = threading.Event() + ready_flag.set() + + class _ReadyAdapter: + def is_set(self): + return ready_flag.is_set() + + def clear(self): + ready_flag.clear() + + def set(self): + ready_flag.set() + + server._ready = _ReadyAdapter() + # _reconnect_event is called via loop.call_soon_threadsafe(…set); use - # a threading-safe substitute. + # a threading-safe substitute. The production reconnect path must not + # treat the old stale session as fresh, so this test double swaps in a + # distinct session object when reconnect is requested. reconnect_flag = threading.Event() class _EventAdapter: def set(self): reconnect_flag.set() + old_session = server.session + new_session = MagicMock() + for method_name in ( + "call_tool", + "list_resources", + "read_resource", + "list_prompts", + "get_prompt", + ): + if hasattr(old_session, method_name): + setattr(new_session, method_name, getattr(old_session, method_name)) + server.session = new_session + ready_flag.set() server._reconnect_event = _EventAdapter() - # Immediately "ready" — simulates a fast reconnect (_ready.is_set() - # is polled by _handle_session_expired_and_retry until the timeout). - ready_flag = threading.Event() - ready_flag.set() - server._ready = MagicMock() - server._ready.is_set = ready_flag.is_set - # session attr must be truthy for the handler's initial check # (``if not server or not server.session``) and for the post- # reconnect readiness probe (``srv.session is not None``). @@ -189,6 +213,74 @@ async def _call_sequence(*a, **kw): mcp_tool._server_error_counts.pop("wpcom", None) +def test_session_expired_retry_waits_for_new_session(monkeypatch, tmp_path): + """Regression for long-lived HTTP/stream MCP sessions. + + If the reconnect helper only checks ``_ready.is_set()`` and + ``session is not None``, it can return immediately while ``session`` still + points at the stale transport. The retry then hits the same dead session + and the circuit breaker eventually reports the server as unreachable. The + handler must wait for a distinct session object before retrying. + """ + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + + from tools import mcp_tool + from tools.mcp_tool import _make_tool_handler + + mcp_tool._ensure_mcp_loop() + server = MagicMock() + server.name = "hindsight" + ready_flag = threading.Event() + ready_flag.set() + + class _ReadyAdapter: + def is_set(self): + return ready_flag.is_set() + + def clear(self): + ready_flag.clear() + + def set(self): + ready_flag.set() + + old_session = MagicMock() + + async def _old_call(*a, **kw): + raise RuntimeError("Session terminated") + + old_session.call_tool = _old_call + new_session = MagicMock() + + async def _new_call(*a, **kw): + result = MagicMock() + result.isError = False + result.content = [MagicMock(type="text", text="bank ok")] + result.structuredContent = None + return result + + new_session.call_tool = _new_call + server.session = old_session + server._ready = _ReadyAdapter() + + class _ReconnectAdapter: + def set(self): + server.session = new_session + ready_flag.set() + + server._reconnect_event = _ReconnectAdapter() + mcp_tool._servers["hindsight"] = server + mcp_tool._server_error_counts.pop("hindsight", None) + + try: + handler = _make_tool_handler("hindsight", "get_bank", 10.0) + parsed = json.loads(handler({})) + assert parsed.get("result") == "bank ok", parsed + assert mcp_tool._server_error_counts.get("hindsight", 0) == 0 + finally: + mcp_tool._servers.pop("hindsight", None) + mcp_tool._server_error_counts.pop("hindsight", None) + + def test_call_tool_handler_non_session_expired_error_falls_through( monkeypatch, tmp_path ): diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index b24bb9705adfb..b3136ebc60342 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -1480,12 +1480,12 @@ async def run(self, config: dict): "manual refresh)", self.name, ) - # Reset the session reference; _run_http/_run_stdio will - # repopulate it on successful re-entry. + # Reset the session reference and readiness; _run_http/_run_stdio + # will repopulate both on successful re-entry. Leaving + # _ready set here lets handler-side recovery mistake the stale + # pre-reconnect session for a fresh one and retry too early. + self._ready.clear() self.session = None - # Keep _ready set across reconnects so tool handlers can - # still detect a transient in-flight state — it'll be - # re-set after the fresh session initializes. continue except asyncio.CancelledError: # Task was cancelled (shutdown, gateway restart, explicit @@ -1670,6 +1670,85 @@ def _reset_server_error(server_name: str) -> None: _server_error_counts[server_name] = 0 _server_breaker_opened_at.pop(server_name, None) + +def _wait_for_server_session_ready( + srv: "MCPServerTask", + *, + old_session: Any = None, + timeout: float = 15.0, +) -> bool: + """Wait for an MCP server to expose a usable session. + + Tool handlers run in normal worker threads while the MCP transport lives on + the module's background asyncio loop. During a reconnect there is a short + window where ``srv.session`` is ``None`` (or still points at the stale + session until the lifecycle coroutine has left the transport context). A + handler that blindly retries in that window can burn circuit-breaker strikes + and return ``not connected`` even though the reconnect is already in + progress. + + When ``old_session`` is supplied, require the observed session object to be + different so callers do not mistake the pre-reconnect, stale session for a + fresh one. + """ + deadline = time.monotonic() + max(float(timeout), 0.0) + while time.monotonic() < deadline: + session = getattr(srv, "session", None) + ready = getattr(srv, "_ready", None) + is_ready = True + if ready is not None and hasattr(ready, "is_set"): + try: + is_ready = bool(ready.is_set()) + except Exception: + is_ready = True + if session is not None and session is not old_session and is_ready: + return True + time.sleep(0.25) + return False + + +def _signal_reconnect_and_wait( + server_name: str, + srv: "MCPServerTask", + *, + op_description: str, + timeout: float = 15.0, +) -> bool: + """Ask a live MCP server task to rebuild its transport session. + + The important detail is clearing ``_ready`` on the MCP event loop before + setting ``_reconnect_event``. Older code left ``_ready`` set across + reconnects, so the caller's readiness poll could return immediately and + retry against the same dead HTTP/stream session. That was observed as + repeated ``Session terminated`` / ``not connected`` / circuit-breaker + failures in long-lived gateway sessions even though a fresh CLI process + could connect successfully. + """ + loop = _mcp_loop + if loop is None or not loop.is_running(): + return False + + old_session = getattr(srv, "session", None) + + def _request_reconnect() -> None: + ready = getattr(srv, "_ready", None) + if ready is not None and hasattr(ready, "clear"): + ready.clear() + reconnect_event = getattr(srv, "_reconnect_event", None) + if reconnect_event is not None and hasattr(reconnect_event, "set"): + reconnect_event.set() + + logger.info( + "MCP server '%s': %s requesting transport reconnect", + server_name, op_description, + ) + loop.call_soon_threadsafe(_request_reconnect) + return _wait_for_server_session_ready( + srv, + old_session=old_session, + timeout=timeout, + ) + # --------------------------------------------------------------------------- # Auth-failure detection helpers (Task 6 of MCP OAuth consolidation) # --------------------------------------------------------------------------- @@ -1795,28 +1874,24 @@ async def _recover(): if recovered: with _lock: srv = _servers.get(server_name) + reconnected = False if srv is not None and hasattr(srv, "_reconnect_event"): - loop = _mcp_loop - if loop is not None and loop.is_running(): - loop.call_soon_threadsafe(srv._reconnect_event.set) - # Wait briefly for the session to come back ready. Bounded - # so that a stuck reconnect falls through to the error - # path rather than hanging the caller. - deadline = time.monotonic() + 15 - while time.monotonic() < deadline: - if srv.session is not None and srv._ready.is_set(): - break - time.sleep(0.25) - - # A successful OAuth recovery is independent evidence that the - # server is viable again, so close the circuit breaker here — - # not only on retry success. Without this, a reconnect - # followed by a failing retry would leave the breaker pinned - # above threshold forever (the retry-exception branch below - # bumps the count again). The post-reset retry still goes - # through _bump_server_error on failure, so a genuinely broken - # server will re-trip the breaker as normal. - _reset_server_error(server_name) + reconnected = _signal_reconnect_and_wait( + server_name, + srv, + op_description=f"{op_description} after OAuth recovery", + timeout=15, + ) + + # A successful OAuth recovery + transport reconnect is independent + # evidence that the server is viable again, so close the circuit + # breaker here — not only on retry success. Without this, a reconnect + # followed by a failing retry would leave the breaker pinned above + # threshold forever. The post-reset retry still goes through + # _bump_server_error on failure, so a genuinely broken server will + # re-trip the breaker as normal. + if reconnected: + _reset_server_error(server_name) try: result = retry_call() @@ -1945,15 +2020,12 @@ def _handle_session_expired_and_retry( # Trigger the same reconnect mechanism the OAuth recovery path # uses, then wait briefly for the new session to come back ready. - loop.call_soon_threadsafe(srv._reconnect_event.set) - deadline = time.monotonic() + 15 - ready = False - while time.monotonic() < deadline: - if srv.session is not None and srv._ready.is_set(): - ready = True - break - time.sleep(0.25) - if not ready: + if not _signal_reconnect_and_wait( + server_name, + srv, + op_description=op_description, + timeout=15, + ): logger.warning( "MCP server '%s': reconnect did not ready within 15s after " "session-expired error; falling through to error response.", @@ -2231,10 +2303,20 @@ def _handler(args: dict, **kwargs) -> str: with _lock: server = _servers.get(server_name) if not server or not server.session: - _bump_server_error(server_name) - return json.dumps({ - "error": f"MCP server '{server_name}' is not connected" - }, ensure_ascii=False) + if server is not None and _wait_for_server_session_ready( + server, + timeout=min(5.0, float(tool_timeout or 5.0)), + ): + # A reconnect completed while this handler was starting. + # Re-read the server below and proceed with the fresh session + # rather than counting a transient reconnect window as a + # circuit-breaker failure. + pass + else: + _bump_server_error(server_name) + return json.dumps({ + "error": f"MCP server '{server_name}' is not connected" + }, ensure_ascii=False) async def _call(): async with server._rpc_lock: