MCP session reconnect and tool rediscovery fixes

Fixes on record

Merged to main. A stamp counts the issues a fix is linked to.

• #59331 fix(mcp): session-expired retry waits for a fresh session; wake parked servers at session startup (salvage #27052)

  merged to main

  files: scripts/release.py, tests/tools/test_mcp_circuit_breaker.py, tests/tools/test_mcp_register_wakes_stale.py, tests/tools/test_mcp_tool_session_expired.py, tools/mcp_tool.py

  Show the diff

  • scripts/release.py

    @@ -1211,6 +1211,7 @@
         "[redacted:email]": "BlackishGreen33",
         "[email protected]": "gnodet",  # PR #37598 salvage (MCP preflight POST probe fallback)
         "[email protected]": "kaishi00",  # PR #55203 salvage (skip_preflight opt-out)
    +    "[redacted:email]": "setclock",  # PR #27052 salvage (MCP session-expired retry waits for a distinct fresh session)
         "[redacted:email]": "sberan",  # PR #54494 salvage (--connect-timeout flag on hermes mcp add)
         "[email protected]": "labsobsidian",  # PR #56699 salvage (MCP OAuth login connect_timeout floor)
         "[redacted:email]": "SaulJWu",

  • tests/tools/test_mcp_circuit_breaker.py

    @@ -31,14 +31,49 @@ 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 __init__(self):
    +            self.set_calls = 0
    +
    +        def set(self):
    +            self.set_calls += 1
    +            old_session = server.session
    +            new_session = MagicMock()
    +            if old_session is not None:
    +                new_session.call_tool = old_session.call_tool
    +            elif call_tool_impl is not None:
    +                new_session.call_tool = call_tool_impl
    +            server.session = new_session
    +            ready_flag.set()
    +
    +        # MagicMock-compat shim: the dead-session half-open test asserts the
    +        # reconnect signal was delivered exactly once.
    +        def assert_called_once(self):
    +            assert self.set_calls == 1, f"set() called {self.set_calls} times"
    +
    +    server._reconnect_event = _ReconnectAdapter()
    +    server._ready = _ReadyAdapter()
     
         mcp_tool_module._servers[name] = server
         mcp_tool_module._server_error_counts.pop(name, None)
    @@ -220,7 +255,7 @@ def _fake_monotonic():
     
             # Clean "reconnecting" error, and a reconnect was actually signalled.
             assert "reconnect" in parsed.get("error", "").lower(), parsed
    -        server._reconnect_event.set.assert_called_once()
    +        server._reconnect_event.assert_called_once()
         finally:
             _cleanup(mcp_tool, "srv")
     

  • tests/tools/test_mcp_register_wakes_stale.py

    @@ -0,0 +1,62 @@
    +"""New sessions must wake parked/stale cached MCP servers immediately.
    +
    +Regression for #50170: after a keepalive failure parks a server, its tools
    +are deregistered — so a NEW agent session starting up saw the tools silently
    +absent and had no way to trigger recovery until the next timed self-probe
    +(up to _PARKED_RETRY_INTERVAL later). register_mcp_servers now nudges any
    +cached entry whose session is None via _signal_reconnect.
    +"""
    +
    +import pytest
    +
    +
    [redacted:email]_isolate
    +def test_register_wakes_stale_cached_server(monkeypatch, tmp_path):
    +    monkeypatch.setenv("HERMES_HOME", str(tmp_path))
    +
    +    from tools import mcp_tool
    +
    +    woken: list[str] = []
    +
    +    class _Event:
    +        def __init__(self, name):
    +            self._name = name
    +
    +        def set(self):
    +            woken.append(self._name)
    +
    +    class _Stale:
    +        session = None
    +
    +        def __init__(self, name):
    +            self.name = name
    +            self._reconnect_event = _Event(name)
    +            self._registered_tool_names: list[str] = []
    +
    +    class _Alive:
    +        session = object()
    +
    +        def __init__(self, name):
    +            self.name = name
    +            self._reconnect_event = _Event(name)
    +            self._registered_tool_names = [f"{name}__tool"]
    +
    +    monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True)
    +    stale = _Stale("parked-srv")
    +    alive = _Alive("healthy-srv")
    +    monkeypatch.setitem(mcp_tool._servers, "parked-srv", stale)
    +    monkeypatch.setitem(mcp_tool._servers, "healthy-srv", alive)
    +
    +    try:
    +        result = mcp_tool.register_mcp_servers({
    +            "parked-srv": {"url": "http://127.0.0.1:9/mcp"},
    +            "healthy-srv": {"url": "http://127.0.0.1:9/mcp"},
    +        })
    +        # Both cached → no new connections attempted; existing names returned.
    +        assert "healthy-srv__tool" in result
    +        # The parked (session=None) entry got a reconnect nudge; the healthy
    +        # one was left alone.
    +        assert woken == ["parked-srv"]
    +    finally:
    +        mcp_tool._servers.pop("parked-srv", None)
    +        mcp_tool._servers.pop("healthy-srv", None)

  • 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 @@ 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
     ):

  • tools/mcp_tool.py

    @@ -2539,12 +2539,12 @@ async def run(self, config: dict):
                     # 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,89 @@ 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.
    +    """
    +    # Iteration-bounded rather than deadline-bounded: several tests (and the
    +    # circuit-breaker cooldown logic) monkeypatch time.monotonic to a frozen
    +    # clock, which would make a monotonic-deadline loop spin forever.
    +    poll_interval = 0.25
    +    iterations = max(1, int(max(float(timeout), 0.0) / poll_interval))
    +    for i in range(iterations):
    +        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
    +        if i < iterations - 1:
    +            time.sleep(poll_interval)
    +    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 +3024,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.  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,
    -                    )
    +            reconnected = _signal_reconnect_and_wait(
    +                server_name,
    +                srv,
    +                op_description=f"{op_description} after OAuth recovery",
    +                timeout=15,
    +            )
     
    -        # 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 +3170,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 +3624,38 @@ def _handler(args: dict, **kwargs) -> str:
                 }, 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:
    @@ -4529,6 +4603,16 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
                 for k, v in servers.items()
                 if k not in _servers and _parse_boolish(v.get("enabled", True), default=True)
             }
    +        # Cached entries with no live session are parked or mid-reconnect.
    +        # Their tools are deregistered, so nothing else can reach
    +        # _signal_reconnect — without this nudge a new session silently
    +        # waits up to _PARKED_RETRY_INTERVAL for the next self-probe
    +        # (#50170). Wake them now so their tools come back promptly.
    +        stale_cached = [
    +            _servers[k]
    +            for k in servers
    +            if k in _servers and getattr(_servers[k], "session", None) is None
    +        ]
             _server_connecting.update(new_servers)
             for srv_name in new_servers:
                 _server_connect_errors.pop(srv_name, None)
    @@ -4539,6 +4623,9 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
                 else:
                     _parallel_safe_servers.discard(sanitize_mcp_name_component(srv_name))
     
    +    for srv in stale_cached:
    +        _signal_reconnect(srv)
    +
         if not new_servers:
             return _existing_tool_names()