diff --git a/tests/tools/test_mcp_circuit_breaker.py b/tests/tools/test_mcp_circuit_breaker.py index 6fdfd08ce3c..fb876d8d65a 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 b17e6484aab..6a693d5aaf8 100644 --- a/tests/tools/test_mcp_tool_session_expired.py +++ b/tests/tools/test_mcp_tool_session_expired.py @@ -112,23 +112,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``). @@ -188,6 +212,74 @@ def test_call_tool_handler_reconnects_on_session_expired(monkeypatch, tmp_path): 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 98edbd0e00c..50c5dcc57e9 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -2539,12 +2539,12 @@ class MCPServerTask: # permanently kill an otherwise-healthy server. self._reconnect_retries = 0 backoff = 1.0 - # 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 @@ -2816,6 +2816,84 @@ def _signal_reconnect(server: Any) -> bool: return True +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) # --------------------------------------------------------------------------- @@ -2941,41 +3019,24 @@ def _handle_auth_error_and_retry( 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) + reconnected = _signal_reconnect_and_wait( + server_name, + srv, + op_description=f"{op_description} after OAuth recovery", + timeout=15, + ) - # 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. The async helper - # runs on the MCP event loop via _run_on_mcp_loop so it - # does NOT block the event loop during the poll interval. - async def _await_ready() -> bool: - deadline = time.monotonic() + 15 - while time.monotonic() < deadline: - if srv.session is not None and srv._ready.is_set(): - return True - await asyncio.sleep(0.25) - return False - - try: - _run_on_mcp_loop(_await_ready(), timeout=15) - except Exception as exc: - logger.warning( - "MCP OAuth '%s': ready poll failed: %s", - server_name, exc, - ) - - # 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) + # 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() @@ -3104,15 +3165,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.", @@ -3561,27 +3619,38 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): }, ensure_ascii=False) if not server.session: - # No live session — the server task is reconnecting, or it has - # exhausted its retry budget and parked (e.g. a dead stdio - # subprocess). Probing here would write into a dead/absent - # transport and re-arm the breaker forever (#16788). Instead, - # ask the (always-present) server task to rebuild the transport - # — which respawns a dead stdio subprocess — and return a clean - # "reconnecting" error so the model backs off without burning - # iterations. The breaker resets once the fresh session - # initializes (_run_stdio/_run_http call _reset_server_error). - _bump_server_error(server_name) - if _signal_reconnect(server): + # No live session. A reconnect may already be completing (the + # transport swaps in a fresh session object asynchronously) — + # wait briefly before treating this as a failure, so a + # transient reconnect window doesn't burn a circuit-breaker + # strike (#26892). + if _wait_for_server_session_ready( + server, timeout=min(5.0, float(tool_timeout or 5.0)), + ): + pass # Fresh session arrived; proceed below. + else: + # Still down — the server task is reconnecting, or it has + # exhausted its retry budget and parked (e.g. a dead stdio + # subprocess). Probing here would write into a dead/absent + # transport and re-arm the breaker forever (#16788). Instead, + # ask the (always-present) server task to rebuild the + # transport — which respawns a dead stdio subprocess — and + # return a clean "reconnecting" error so the model backs off + # without burning iterations. The breaker resets once the + # fresh session initializes (_run_stdio/_run_http call + # _reset_server_error). + _bump_server_error(server_name) + if _signal_reconnect(server): + return json.dumps({ + "error": ( + f"MCP server '{server_name}' transport is down; " + f"reconnect requested. Do NOT retry this tool " + f"immediately — give it a few seconds to come back." + ) + }, ensure_ascii=False) return json.dumps({ - "error": ( - f"MCP server '{server_name}' transport is down; " - f"reconnect requested. Do NOT retry this tool " - f"immediately — give it a few seconds to come back." - ) + "error": f"MCP server '{server_name}' is not connected" }, ensure_ascii=False) - return json.dumps({ - "error": f"MCP server '{server_name}' is not connected" - }, ensure_ascii=False) async def _call(): async with server._rpc_lock: