fix(mcp): auto-reconnect + retry once when the transport session expires (#13383)

Streamable HTTP MCP servers may garbage-collect their server-side
session state while the OAuth token remains valid — idle TTL, server
restart, pod rotation, etc.  Before this fix, the tool-call handler
treated the resulting "Invalid or expired session" error as a plain
tool failure with no recovery path, so **every subsequent call on
the affected server failed until the gateway was manually
restarted**.  Reporter: #13383.

The OAuth-based recovery path (``_handle_auth_error_and_retry``)
already exists for 401s, but it only fires on auth errors.  Session
expiry slipped through because the access token is still valid —
nothing 401'd, so the existing recovery branch was skipped.

Fix
---
Add a sibling function ``_handle_session_expired_and_retry`` that
detects MCP session-expiry via ``_is_session_expired_error`` (a
narrow allow-list of known-stable substrings: ``"invalid or expired
session"``, ``"session expired"``, ``"session not found"``,
``"unknown session"``, etc.) and then uses the existing transport
reconnect mechanism:

* Sets ``MCPServerTask._reconnect_event`` — the server task's
  lifecycle loop already interprets this as "tear down the current
  ``streamablehttp_client`` + ``ClientSession`` and rebuild them,
  reusing the existing OAuth provider instance".
* Waits up to 15 s for the new session to come back ready.
* Retries the original call once.  If the retry succeeds, returns
  its result and resets the circuit-breaker error count.  If the
  retry raises, or if the reconnect doesn't ready in time, falls
  through to the caller's generic error path.

Unlike the 401 path, this does **not** call ``handle_401`` — the
access token is already valid and running an OAuth refresh would be
a pointless round-trip.

All 5 MCP handlers (``call_tool``, ``list_resources``, ``read_resource``,
``list_prompts``, ``get_prompt``) now consult both recovery paths
before falling through:

    recovered = _handle_auth_error_and_retry(...)          # 401 path
    if recovered is not None: return recovered
    recovered = _handle_session_expired_and_retry(...)     # new
    if recovered is not None: return recovered
    # generic error response

Narrow scope — explicitly not changed
-------------------------------------
* **Detection is string-based on a 5-entry allow-list.**  The MCP
  SDK wraps JSON-RPC errors in ``McpError`` whose exception type +
  attributes vary across SDK versions, so matching on message
  substrings is the durable path.  Kept narrow to avoid false
  positives — a regular ``RuntimeError("Tool failed")`` will NOT
  trigger spurious reconnects (pinned by
  ``test_is_session_expired_rejects_unrelated_errors``).
* **No change to the existing 401 recovery flow.**  The new path is
  consulted only after the auth path declines (returns ``None``).
* **Retry count stays at 1.**  If the reconnect-then-retry also
  fails, we don't loop — the error surfaces normally so the model
  sees a failed tool call rather than a hang.
* **``InterruptedError`` is explicitly excluded** from session-expired
  detection so user-cancel signals always short-circuit the same
  way they did before (pinned by
  ``test_is_session_expired_rejects_interrupted_error``).

Regression coverage
-------------------
``tests/tools/test_mcp_tool_session_expired.py`` (new, 16 cases):

Unit tests for ``_is_session_expired_error``:
* ``test_is_session_expired_detects_invalid_or_expired_session`` —
  reporter's exact wpcom-mcp text.
* ``test_is_session_expired_detects_expired_session_variant`` —
  "Session expired" / "expired session" variants.
* ``test_is_session_expired_detects_session_not_found`` — server GC
  variant ("session not found", "unknown session").
* ``test_is_session_expired_is_case_insensitive``.
* ``test_is_session_expired_rejects_unrelated_errors`` — narrow-scope
  canary: random RuntimeError / ValueError / 401 don't trigger.
* ``test_is_session_expired_rejects_interrupted_error`` — user cancel
  must never route through reconnect.
* ``test_is_session_expired_rejects_empty_message``.

Handler integration tests:
* ``test_call_tool_handler_reconnects_on_session_expired`` — reporter's
  full repro: first call raises "Invalid or expired session", handler
  signals ``_reconnect_event``, retries once, returns the retry's
  success result with no ``error`` key.
* ``test_call_tool_handler_non_session_expired_error_falls_through``
  — preserved-behaviour canary: random tool failures do NOT trigger
  reconnect.
* ``test_session_expired_handler_returns_none_without_loop`` —
  defensive: cold-start / shutdown race.
* ``test_session_expired_handler_returns_none_without_server_record``
  — torn-down server falls through cleanly.
* ``test_session_expired_handler_returns_none_when_retry_also_fails``
  — no retry loop on repeated failure.

Parametrised across all 4 non-``tools/call`` handlers:
* ``test_non_tool_handlers_also_reconnect_on_session_expired``
  [list_resources / read_resource / list_prompts / get_prompt].

**15 of 16 fail on clean ``origin/main`` (``6fb69229``)** with
``ImportError: cannot import name '_is_session_expired_error'``
— the fix's surface symbols don't exist there yet.  The 1 passing
test is an ordering artefact of pytest-xdist worker collection.

Validation
----------
``source venv/bin/activate && python -m pytest
tests/tools/test_mcp_tool_session_expired.py -q`` → **16 passed**.

Broader MCP suite (5 files:
``test_mcp_tool.py``, ``test_mcp_tool_401_handling.py``,
``test_mcp_tool_session_expired.py``, ``test_mcp_reconnect_signal.py``,
``test_mcp_oauth.py``) → **230 passed, 0 regressions**.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Brian D. Evans
2026-04-21 09:05:46 +01:00
committed by Teknium
parent 8c2732a9f9
commit e87a2100f6
2 changed files with 512 additions and 0 deletions

View File

@@ -1582,6 +1582,129 @@ def _handle_auth_error_and_retry(
"server": server_name,
}, ensure_ascii=False)
# Substrings (lower-cased match) that indicate the MCP server rejected
# the request because its server-side transport session expired /
# was garbage-collected. The caller's OAuth token is still valid —
# only the transport-layer session state needs rebuilding. See #13383.
_SESSION_EXPIRED_MARKERS: tuple = (
"invalid or expired session",
"expired session",
"session expired",
"session not found",
"unknown session",
)
def _is_session_expired_error(exc: BaseException) -> bool:
"""Return True if ``exc`` looks like an MCP transport session expiry.
Streamable HTTP MCP servers may garbage-collect server-side session
state while the OAuth token remains valid — idle TTL, server
restart, horizontal-scaling pod rotation, etc. The SDK surfaces
this as a JSON-RPC error whose message contains phrases like
``"Invalid or expired session"``. This class of failure is
distinct from :func:`_is_auth_error`: re-running the OAuth refresh
flow would be pointless because the access token is fine. What's
needed is a transport reconnect — tear down and rebuild the
``streamablehttp_client`` + ``ClientSession`` pair, which is
exactly what ``MCPServerTask._reconnect_event`` triggers.
"""
if isinstance(exc, InterruptedError):
return False
# Exception messages vary across SDK versions + server
# implementations, so match on a small allow-list of stable
# substrings rather than exception type. Kept narrow to avoid
# false positives on unrelated server errors.
msg = str(exc).lower()
if not msg:
return False
return any(marker in msg for marker in _SESSION_EXPIRED_MARKERS)
def _handle_session_expired_and_retry(
server_name: str,
exc: BaseException,
retry_call,
op_description: str,
):
"""Trigger a transport reconnect and retry once on session expiry.
Unlike :func:`_handle_auth_error_and_retry`, this does **not** call
the OAuth manager's ``handle_401`` — the access token is still
valid, only the server-side session state is stale. Setting
``_reconnect_event`` causes the server task's lifecycle loop to
tear down the current ``streamablehttp_client`` + ``ClientSession``
and rebuild them, reusing the existing OAuth provider instance.
See #13383.
Args:
server_name: Name of the MCP server that raised.
exc: The exception from the failed call.
retry_call: Zero-arg callable that re-runs the operation,
returning the same JSON string format as the handler.
op_description: Human-readable name of the operation (logs).
Returns:
A JSON string if reconnect + retry was attempted and produced
a response, or ``None`` to fall through to the caller's
generic error path (not a session-expired error, no server
record, reconnect didn't ready in time, or retry also failed).
"""
if not _is_session_expired_error(exc):
return None
with _lock:
srv = _servers.get(server_name)
if srv is None or not hasattr(srv, "_reconnect_event"):
return None
loop = _mcp_loop
if loop is None or not loop.is_running():
return None
logger.info(
"MCP server '%s': %s failed with session-expired error (%s); "
"signalling transport reconnect and retrying once.",
server_name, op_description, exc,
)
# 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:
logger.warning(
"MCP server '%s': reconnect did not ready within 15s after "
"session-expired error; falling through to error response.",
server_name,
)
return None
try:
result = retry_call()
try:
parsed = json.loads(result)
if "error" not in parsed:
_server_error_counts[server_name] = 0
return result
except (json.JSONDecodeError, TypeError):
_server_error_counts[server_name] = 0
return result
except Exception as retry_exc:
logger.warning(
"MCP %s/%s retry after session reconnect failed: %s",
server_name, op_description, retry_exc,
)
return None
# Dedicated event loop running in a background daemon thread.
_mcp_loop: Optional[asyncio.AbstractEventLoop] = None
_mcp_thread: Optional[threading.Thread] = None
@@ -1868,6 +1991,16 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float):
if recovered is not None:
return recovered
# Transport session expiry (#13383): same reconnect flow
# but skips OAuth recovery because the access token is
# still valid — only the server-side session is stale.
recovered = _handle_session_expired_and_retry(
server_name, exc, _call_once,
f"tools/call {tool_name}",
)
if recovered is not None:
return recovered
_bump_server_error(server_name)
logger.error(
"MCP tool %s/%s call failed: %s",
@@ -1920,6 +2053,11 @@ def _make_list_resources_handler(server_name: str, tool_timeout: float):
recovered = _handle_auth_error_and_retry(
server_name, exc, _call_once, "resources/list",
)
if recovered is not None:
return recovered
recovered = _handle_session_expired_and_retry(
server_name, exc, _call_once, "resources/list",
)
if recovered is not None:
return recovered
logger.error(
@@ -1974,6 +2112,11 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float):
recovered = _handle_auth_error_and_retry(
server_name, exc, _call_once, "resources/read",
)
if recovered is not None:
return recovered
recovered = _handle_session_expired_and_retry(
server_name, exc, _call_once, "resources/read",
)
if recovered is not None:
return recovered
logger.error(
@@ -2031,6 +2174,11 @@ def _make_list_prompts_handler(server_name: str, tool_timeout: float):
recovered = _handle_auth_error_and_retry(
server_name, exc, _call_once, "prompts/list",
)
if recovered is not None:
return recovered
recovered = _handle_session_expired_and_retry(
server_name, exc, _call_once, "prompts/list",
)
if recovered is not None:
return recovered
logger.error(
@@ -2096,6 +2244,11 @@ def _make_get_prompt_handler(server_name: str, tool_timeout: float):
recovered = _handle_auth_error_and_retry(
server_name, exc, _call_once, "prompts/get",
)
if recovered is not None:
return recovered
recovered = _handle_session_expired_and_retry(
server_name, exc, _call_once, "prompts/get",
)
if recovered is not None:
return recovered
logger.error(