mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-08 13:12:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

fix: reconnect stale MCP sessions before retry

Browse source
This commit is contained in:
Marin Pesa 2026-05-16 14:44:35 -04:00 • committed by Teknium
parent a124d16764
commit 27beeb1830
3 changed files with 261 additions and 77 deletions
Download patch file Download diff file Expand all files Collapse all files

29
tests/tools/test_mcp_circuit_breaker.py
View file

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

108
tests/tools/test_mcp_tool_session_expired.py
View file

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

201
tools/mcp_tool.py
View file

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