mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-20 15:33:54 +00:00
fix: reconnect stale MCP sessions before retry
This commit is contained in:
Marin Pesa
Teknium
parent
a124d16764
commit
27beeb1830
3 changed files with 261 additions and 77 deletions
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue