fix(mcp): classify OAuth refresh failures so neither a blip revokes consent nor a dead grant strands the user

Follow-up to #714 (Entra OBO, #682). A refresh failure deleted the user token +
emitted token_revoked regardless of cause, so a transient AS/network blip during
a forced refresh (the live 401-retry path) permanently revoked consent
cluster-wide. Fixing only that, though, opens the dual failure: a genuinely-dead
grant the AS reports in a non-standard shape would now be kept forever and the
user stranded on a retryable error with no re-consent path. This classifies the
failure three ways so each is handled correctly.

Classification (_classify_refresh_failure): MCPOAuthRefreshFailed carries a
_RefreshFailureClass instead of a bool —
- PERMANENT (revoke + re-consent): an explicit dead-grant / re-consent signal —
  invalid_grant at any 4xx (400/401/403), invalid_scope, or an OIDC
  interaction-required code (interaction_required / login_required /
  consent_required / account_selection_required) the AS surfaces.
- TRANSIENT (keep, retry, never escalate): infrastructure (network, 5xx, 429,
  malformed body) and operator-fixable codes (invalid_client, invalid_request,
  unauthorized_client, unsupported_grant_type, temporarily_unavailable) —
  re-consenting the user can't fix a bad client_secret, and an outage must not
  revoke consent however long it lasts.
- AMBIGUOUS (keep, but escalate after a run): a 400/401 we can't map to a
  standard code. A one-off can't revoke, but an uninterrupted streak past a
  threshold escalates to re-consent so a dead grant in a non-standard shape
  can't strand the user. Infra transients reset the streak, so an outage never
  escalates.

Concurrency: do NOT drop the per-(user,server) refresh lock on the keep-the-token
path. Evicting it while the token is still live let a second concurrent caller
mint a fresh lock and refresh the same token in parallel; with refresh-token
rotation the second send reuses the consumed token, gets invalid_grant, and
spuriously revokes — the exact bug this commit prevents. The async-with still
releases the lock on return; the registry entry is pruned only when the token is
actually refreshed or revoked. Bit SQLite single-node hardest, where the pg
advisory lock is a no-op.

perf: a per-(user,server) cooldown short-circuits the token-endpoint round-trip
for a brief window after a transient failure, so a down AS isn't hit once per
tool call; self-heals when the window expires. Plus the lock-free in-flight key
set that collapses concurrent session-start pool primes (single mcp-loop thread).

dispatch/FE: the transient kind maps to a retryable mcp_refresh_unavailable
structured error (not mcp_consent_required); the FE titles it "Temporarily
unavailable" under a new soft "transient" category (amber, not the red hard-error
styling) in both stylesheets, with no wrong re-consent button.

tests: invalid_client kept (pins the discriminator on the error code, not the 4xx
status), single ambiguous 400 kept, 403 invalid_grant revokes, interaction_required
revokes, ambiguous streak escalates at the threshold, sustained 5xx never escalates
(outage safety), and the cooldown skips the second AS round-trip — all through the
real AS HTTP boundary.
This commit is contained in:
Patrick Buckley
2026-06-25 21:52:11 -07:00
parent de271dc2f9
commit 800b561f56
7 changed files with 754 additions and 85 deletions
+232
View File
@@ -172,6 +172,238 @@ def _public_addr_patch():
return patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("93.184.216.34", 0))])
# ---------------------------------------------------------------------------
# Refresh-failure classification (#714 follow-up + hardening): a TRANSIENT
# failure (network / 5xx / 429 / operator-fixable code) keeps the token and
# returns a retryable kind; an explicit dead-grant / re-consent signal
# (``invalid_grant`` at any 4xx, ``invalid_scope``, an OIDC interaction-required
# code) revokes consent; and an unclassifiable 400/401 is AMBIGUOUS — kept until
# a sustained run escalates to re-consent. A per-(user,server) cooldown
# short-circuits the AS round-trip during an outage. All exercised through the
# real AS HTTP boundary so an AS/network blip on the live 401-retry path can
# never revoke a user, while a genuinely dead grant can't strand one forever.
# ---------------------------------------------------------------------------
class TestRefreshFailureClassification:
def _lookup(self, state: SimpleNamespace) -> Any:
from turnstone.core.mcp_oauth import get_user_access_token_classified
async def _run() -> Any:
with _public_addr_patch():
return await get_user_access_token_classified(
app_state=state,
user_id="user-1",
server_name="srv-oauth",
force_refresh=True,
)
return asyncio.run(_run())
def test_transient_503_keeps_token(self, storage: SQLiteBackend) -> None:
"""A 503 from the token endpoint is transient: keep the token, retryable kind."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(503, {"error": "temporarily_unavailable"})
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed_transient"
# Token survives a transient failure — no cluster-wide revoke; self-heals.
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is not None
def test_transient_network_error_keeps_token(self, storage: SQLiteBackend) -> None:
"""A network error (httpx.HTTPError) is transient: keep the token, retryable kind."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed_transient"
# Token survives a transient failure — no cluster-wide revoke; self-heals.
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is not None
def test_permanent_invalid_grant_revokes(self, storage: SQLiteBackend) -> None:
"""Contrast: 400 invalid_grant IS permanent — deletion is correct and the
eventual fix MUST preserve it."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(return_value=_mk_response(400, {"error": "invalid_grant"}))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is None
def test_400_invalid_client_keeps_token(self, storage: SQLiteBackend) -> None:
"""A 400 ``invalid_client`` is operator-fixable, NOT a dead grant: keep
the token. Pins the discriminator on the *error code*, not the 4xx
status — broadening ``permanent`` to "any 400" would silently revoke
consent on a config blip (the regression this guards)."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(return_value=_mk_response(400, {"error": "invalid_client"}))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed_transient"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is not None
def test_400_unrecognised_body_is_ambiguous_keeps_token(self, storage: SQLiteBackend) -> None:
"""A single 400 with a non-JSON / no-``error`` body is ambiguous: keep
the token — one oddity must not revoke. Escalation only bites after a
sustained run (see ``test_ambiguous_streak_escalates_to_revoke``)."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(return_value=_mk_response(400, None))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed_transient"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is not None
def test_403_invalid_grant_revokes(self, storage: SQLiteBackend) -> None:
"""``invalid_grant`` is a dead grant at ANY client-error status, not just
400/401 — a 403 invalid_grant must still revoke + re-consent."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(return_value=_mk_response(403, {"error": "invalid_grant"}))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is None
def test_interaction_required_revokes(self, storage: SQLiteBackend) -> None:
"""An OIDC interaction-required code (Entra surfaces these) means the user
must re-consent / re-auth — treat as permanent, revoke."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(return_value=_mk_response(401, {"error": "interaction_required"}))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
result = self._lookup(state)
assert result.kind == "refresh_failed"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is None
def test_ambiguous_streak_escalates_to_revoke(self, storage: SQLiteBackend) -> None:
"""A *sustained* run of unclassifiable 400s is treated as a dead grant in
a non-standard shape: the token survives below the threshold, then the
threshold-crossing attempt escalates to re-consent so the user isn't
stranded on a retryable error forever."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(return_value=_mk_response(400, None))
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
with (
patch("turnstone.core.mcp_oauth._AMBIGUOUS_ESCALATION_THRESHOLD", 3),
patch("turnstone.core.mcp_oauth._REFRESH_TRANSIENT_COOLDOWN_SECONDS", 0.0),
):
# Below threshold: the token survives each attempt.
for _ in range(2):
assert self._lookup(state).kind == "refresh_failed_transient"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is not None
# The threshold-crossing attempt escalates to a revoke.
result = self._lookup(state)
assert result.kind == "refresh_failed"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is None
def test_sustained_5xx_never_escalates(self, storage: SQLiteBackend) -> None:
"""Outage safety: infra failures (5xx) never feed the escalation counter,
so even a long AS outage — far past the ambiguous threshold — keeps the
token. A blip must never revoke consent, however long it lasts."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(503, {"error": "temporarily_unavailable"})
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
with (
patch("turnstone.core.mcp_oauth._AMBIGUOUS_ESCALATION_THRESHOLD", 2),
patch("turnstone.core.mcp_oauth._REFRESH_TRANSIENT_COOLDOWN_SECONDS", 0.0),
):
for _ in range(5):
assert self._lookup(state).kind == "refresh_failed_transient"
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is not None
def test_transient_cooldown_skips_as_roundtrip(self, storage: SQLiteBackend) -> None:
"""After a transient failure, a follow-up lookup inside the cooldown
window returns the retryable kind WITHOUT a second token-endpoint
round-trip — so a down AS isn't hammered once per tool call."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(503, {"error": "temporarily_unavailable"})
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
first = self._lookup(state)
second = self._lookup(state)
assert first.kind == "refresh_failed_transient"
assert second.kind == "refresh_failed_transient"
# The cooldown short-circuited the second attempt: exactly one AS POST.
assert client.post.call_count == 1
def test_backoff_cleared_when_token_vanishes(self, storage: SQLiteBackend) -> None:
"""A transient failure records per-(user,server) backoff; if the token is
then deleted cluster-wide (another node's permanent revoke), the next
lookup returns ``missing`` AND clears this node's now-stale backoff entry,
so the dict stays bounded to live pairs."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(503, {"error": "temporarily_unavailable"})
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
# First lookup: a transient 503 records a backoff entry.
assert self._lookup(state).kind == "refresh_failed_transient"
assert ("user-1", "srv-oauth") in state.mcp_oauth_refresh_backoff
# Another node revokes the token cluster-wide (shared Postgres store).
state.mcp_token_store.delete_user_token("user-1", "srv-oauth")
# Next lookup sees the row gone -> missing -> the stale entry is cleared.
assert self._lookup(state).kind == "missing"
assert ("user-1", "srv-oauth") not in state.mcp_oauth_refresh_backoff
# ---------------------------------------------------------------------------
# Happy paths
# ---------------------------------------------------------------------------
+90
View File
@@ -701,6 +701,47 @@ class TestDispatcherAuthFlows:
assert payload["error"]["server"] == "pool-srv"
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
def test_dispatch_pool_transient_refresh_emits_retryable_not_consent(
self, running_loop_mgr, storage: SQLiteBackend
) -> None:
"""A TRANSIENT refresh failure on the 401-retry surfaces a retryable
``mcp_refresh_unavailable`` error — NOT a re-consent prompt — and does
not tick the breaker."""
from unittest.mock import patch
mgr, loop, _ = running_loop_mgr
cipher = make_mcp_token_cipher()
_seed_oauth_server(storage, name="pool-srv")
_seed_user_token(storage, cipher)
self._wire_pool(mgr, storage, cipher)
from turnstone.core.mcp_oauth import TokenLookupResult
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
if kwargs.get("force_refresh"):
return TokenLookupResult(kind="refresh_failed_transient")
return TokenLookupResult(kind="token", token="access-aaa")
async def _call_tool(name: str, args: dict[str, Any]) -> Any:
_populate_active_capture(mgr, status=401, header='Bearer error="invalid_token"')
raise RuntimeError("upstream 401")
self._seed_pool_entry_with_call_tool(mgr, loop, _call_tool)
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_refresh_unavailable"
assert payload["error"]["server"] == "pool-srv"
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
def test_dispatch_pool_401_retry_ceiling_caps_at_one(
self, running_loop_mgr, storage: SQLiteBackend
) -> None:
@@ -1749,6 +1790,55 @@ class TestPoolPrimingAndTokenRotation:
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
assert entry.tools == [{"name": "do_thing"}]
def test_prime_user_pools_skips_when_already_in_flight(
self, running_loop_mgr, storage: SQLiteBackend
) -> None:
"""A concurrent prime already in flight for (user, server) collapses the
duplicate before the redundant DB reads."""
mgr, loop, _ = running_loop_mgr
cipher = make_mcp_token_cipher()
_seed_oauth_server(storage, name="pool-srv")
_seed_user_token(storage, cipher, expires_in_seconds=3600)
self._wire(mgr, storage, cipher)
mgr._priming_keys.add(("user-1", "pool-srv")) # simulate an in-flight prime
primed: list[tuple[str, str]] = []
async def _fake_prime(
self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str
) -> int:
primed.append(key)
return 0
mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign]
_run_on_loop(loop, mgr._prime_user_pools("user-1"))
assert primed == [], "an in-flight prime must collapse the duplicate"
# The marker belongs to the other (still-running) prime — left intact.
assert ("user-1", "pool-srv") in mgr._priming_keys
def test_prime_user_pools_clears_in_flight_marker_after(
self, running_loop_mgr, storage: SQLiteBackend
) -> None:
"""The in-flight marker is cleared in ``finally`` once a prime completes."""
mgr, loop, _ = running_loop_mgr
cipher = make_mcp_token_cipher()
_seed_oauth_server(storage, name="pool-srv")
_seed_user_token(storage, cipher, expires_in_seconds=3600)
self._wire(mgr, storage, cipher)
async def _fake_prime(
self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str
) -> int:
return 1
mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign]
_run_on_loop(loop, mgr._prime_user_pools("user-1"))
assert mgr._priming_keys == set(), "in-flight marker must be cleared in finally"
# Suppress unused-import warning for AsyncMock.
_ = AsyncMock
+71 -34
View File
@@ -542,6 +542,11 @@ class MCPClientManager:
# mutated only on the mcp-loop. Sync threads interact via
# ``asyncio.run_coroutine_threadsafe``.
self._user_pool_entries: dict[tuple[str, str], PoolEntryState] = {}
# (user, server) keys with a session-start prime in flight — collapses
# concurrent primes (multiple ChatSession starts) before the redundant
# token/server-row DB reads. Mutated only on the single mcp-loop thread,
# so no lock is needed.
self._priming_keys: set[tuple[str, str]] = set()
# LRU tracking: monotonic last-access timestamp per pool key.
self._user_pool_last_used: dict[tuple[str, str], float] = {}
# Per-key lock guarding open / dispatch / close. Allocated lazily
@@ -1733,40 +1738,48 @@ class MCPClientManager:
sem = asyncio.Semaphore(_PRIME_MAX_CONCURRENCY)
async def _prime_one(server_name: str) -> None:
async with sem:
try:
key = (user_id, server_name)
entry = self._user_pool_entries.get(key)
if entry is not None and entry.session is not None:
return # already connected — nothing to do
# Non-refreshing read — priming must not drive a refresh whose
# transient failure would revoke the token (see docstring).
plain = await asyncio.to_thread(
token_store.get_user_token, user_id, server_name
)
if plain is None or not plain.get("access_token"):
return # no usable token (not consented) — lazy paths handle it
if _token_needs_refresh(plain.get("expires_at")):
return # near expiry — let lazy dispatch refresh on actual use
server_row = await asyncio.to_thread(
self._storage.get_mcp_server_by_name, server_name
)
if not server_row:
return
cfg = _pool_cfg_from_row(server_row)
await self._prime_user_server(key, cfg, plain["access_token"])
log.info(
"mcp pool auto-primed at session start user=%s server=%s",
user_id,
server_name,
)
except Exception:
log.debug(
"mcp pool auto-prime failed user=%s server=%s",
user_id,
server_name,
exc_info=True,
)
key = (user_id, server_name)
entry = self._user_pool_entries.get(key)
if entry is not None and entry.session is not None:
return # already connected — nothing to do
if key in self._priming_keys:
return # a concurrent prime for this (user, server) is in flight
# Claim synchronously before any await — the mcp-loop is single-
# threaded, so check-then-add can't interleave with another coroutine.
self._priming_keys.add(key)
try:
async with sem:
try:
# Non-refreshing read — priming must not drive a refresh
# whose transient failure would revoke the token.
plain = await asyncio.to_thread(
token_store.get_user_token, user_id, server_name
)
if plain is None or not plain.get("access_token"):
return # no usable token (not consented) — lazy paths handle it
if _token_needs_refresh(plain.get("expires_at")):
return # near expiry — let lazy dispatch refresh on actual use
server_row = await asyncio.to_thread(
self._storage.get_mcp_server_by_name, server_name
)
if not server_row:
return
cfg = _pool_cfg_from_row(server_row)
await self._prime_user_server(key, cfg, plain["access_token"])
log.info(
"mcp pool auto-primed at session start user=%s server=%s",
user_id,
server_name,
)
except Exception:
log.debug(
"mcp pool auto-prime failed user=%s server=%s",
user_id,
server_name,
exc_info=True,
)
finally:
self._priming_keys.discard(key)
await asyncio.gather(*(_prime_one(s) for s in list(self._oauth_user_server_names)))
@@ -4269,6 +4282,14 @@ class MCPClientManager:
"Operator action required."
),
)
if lookup.kind == "refresh_failed_transient":
# Transient refresh failure (AS/network blip) — the token was kept;
# a retry may succeed once the AS recovers. Retryable, NOT re-consent.
return _structured_error(
code="mcp_refresh_unavailable",
server=server_name,
detail="Token refresh temporarily failed; please retry.",
)
if lookup.kind == "refresh_failed":
# ``mcp_server.oauth.token_revoked`` audit was already emitted
# by ``get_user_access_token_classified`` when it deleted the
@@ -4446,6 +4467,14 @@ class MCPClientManager:
"Operator action required."
),
)
if lookup.kind == "refresh_failed_transient":
# Transient refresh failure (AS/network blip) — the token was kept;
# a retry may succeed once the AS recovers. Retryable, NOT re-consent.
return _structured_error(
code="mcp_refresh_unavailable",
server=server_name,
detail="Token refresh temporarily failed; please retry.",
)
if lookup.kind == "refresh_failed":
return _structured_error(
code="mcp_consent_required",
@@ -4597,6 +4626,14 @@ class MCPClientManager:
"Operator action required."
),
)
if lookup.kind == "refresh_failed_transient":
# Transient refresh failure (AS/network blip) — the token was kept;
# a retry may succeed once the AS recovers. Retryable, NOT re-consent.
return _structured_error(
code="mcp_refresh_unavailable",
server=server_name,
detail="Token refresh temporarily failed; please retry.",
)
if lookup.kind == "refresh_failed":
return _structured_error(
code="mcp_consent_required",
+348 -51
View File
@@ -24,6 +24,7 @@ import asyncio
import base64
import concurrent.futures
import contextlib
import enum
import hashlib
import json
import secrets
@@ -94,13 +95,45 @@ class MCPOAuthExchangeError(MCPOAuthError):
"""Authorization-code exchange failed."""
class MCPOAuthRefreshFailed(MCPOAuthError): # noqa: N818 — name reflects domain semantics
"""Refresh-token grant failed.
class _RefreshFailureClass(enum.Enum):
"""How the caller should react to a failed refresh-token grant.
Caller should treat this as a re-consent trigger: delete the user
token row and emit ``mcp_consent_required``.
- ``PERMANENT`` the AS rejected the grant as dead (``invalid_grant`` /
``invalid_scope`` / an OIDC interaction-required code): revoke the stored
token and trigger re-consent.
- ``TRANSIENT`` an infrastructure or operator-fixable blip (network, 5xx,
429, ``invalid_client``, malformed body): keep the token and surface a
retryable error so a blip can never revoke a user's consent. Never
escalates, so even a sustained AS outage can't strand consent.
- ``AMBIGUOUS`` a 400/401 token-endpoint rejection we couldn't pin to a
standard code: keep the token, but the caller counts consecutive
occurrences and escalates to re-consent past a threshold so a dead grant
delivered in a non-standard shape can't strand the user forever, while a
one-off oddity still can't revoke consent.
"""
PERMANENT = "permanent"
TRANSIENT = "transient"
AMBIGUOUS = "ambiguous"
class MCPOAuthRefreshFailed(MCPOAuthError): # noqa: N818 — name reflects domain semantics
"""Refresh-token grant failed; ``failure_class`` tells the caller how to react.
See :class:`_RefreshFailureClass` for the three handling classes. Defaults to
``TRANSIENT`` the safe direction, since the caller then keeps the token
rather than revoking a user's consent on an unclassified failure.
"""
def __init__(
self,
message: str = "",
*,
failure_class: _RefreshFailureClass = _RefreshFailureClass.TRANSIENT,
) -> None:
super().__init__(message)
self.failure_class = failure_class
# ---------------------------------------------------------------------------
# Authorization-server metadata
@@ -606,6 +639,25 @@ async def pop_pending_state(*, storage: StorageBackend, state: str) -> dict[str,
_AS_ERROR_FIELD_MAX = 80
def _as_error_code(resp: httpx.Response) -> str | None:
"""Return the RFC 6749 ``error`` code from a token-endpoint JSON error body.
Feeds :func:`_classify_refresh_failure`, which maps the code to a handling
class. Returns ``None`` when the body isn't JSON or carries no ``error``
field an absent code is treated as an *ambiguous* rejection, not a
permanent one, so a non-standard error shape can't revoke consent outright.
"""
try:
doc = resp.json()
except ValueError:
return None
if isinstance(doc, dict):
code = doc.get("error")
if isinstance(code, str) and code:
return code
return None
def _format_as_error(resp: httpx.Response) -> str:
"""Build a safe, redacted summary of an AS error response.
@@ -638,6 +690,61 @@ def _format_as_error(resp: httpx.Response) -> str:
return sanitize_log_text(redact_credentials(composite), 200)
# RFC 6749 / OIDC token-endpoint error codes that mean the grant is genuinely
# dead and the user must re-consent — the only PERMANENT (revoke) signals.
_PERMANENT_AS_ERRORS = frozenset(
{
"invalid_grant", # refresh token expired/revoked (RFC 6749 §5.2)
"invalid_scope", # requested scope no longer grantable → re-consent
}
)
# OIDC interaction-required family: the AS needs the user back in the loop
# (consent / login / account selection) — also a re-consent (PERMANENT) signal.
_INTERACTION_AS_ERRORS = frozenset(
{
"interaction_required",
"login_required",
"consent_required",
"account_selection_required",
}
)
# Operator-fixable or RFC-transient codes: keep the token and never escalate —
# re-consenting the user won't fix a bad client_secret, and
# ``temporarily_unavailable`` is explicitly retryable.
_TRANSIENT_AS_ERRORS = frozenset(
{
"invalid_client",
"invalid_request",
"unauthorized_client",
"unsupported_grant_type",
"temporarily_unavailable",
}
)
def _classify_refresh_failure(resp: httpx.Response) -> _RefreshFailureClass:
"""Classify a non-200 refresh response into a handling class.
Conservative by construction the only path to a PERMANENT
(consent-revoking) outcome is an explicit dead-grant / re-consent error code
at a client-error status. Everything infrastructural (5xx, 429) or
operator-fixable (``invalid_client`` ) is TRANSIENT and never escalates, so
a sustained AS outage can't revoke consent. A 400/401 carrying an error code
we don't recognise (or none at all) is AMBIGUOUS: the caller keeps the token
but escalates to re-consent after an uninterrupted run, so a dead grant in a
non-standard shape can't strand the user while a one-off can't revoke.
"""
status = resp.status_code
code = _as_error_code(resp)
if status in (400, 401, 403) and (
code in _PERMANENT_AS_ERRORS or code in _INTERACTION_AS_ERRORS
):
return _RefreshFailureClass.PERMANENT
if status in (400, 401) and code not in _TRANSIENT_AS_ERRORS:
return _RefreshFailureClass.AMBIGUOUS
return _RefreshFailureClass.TRANSIENT
# ---------------------------------------------------------------------------
# DCR (RFC 7591 minimal one-shot)
# ---------------------------------------------------------------------------
@@ -813,8 +920,13 @@ async def refresh_token(
raise MCPOAuthRefreshFailed("refresh endpoint response body exceeds size limit")
if resp.status_code != 200:
# Classify the rejection (see _classify_refresh_failure): an explicit
# dead-grant / re-consent code revokes consent; infra (5xx/429) and
# operator-fixable codes keep the token; an unrecognised 400/401 is
# ambiguous and the caller escalates only after a sustained run.
raise MCPOAuthRefreshFailed(
f"refresh endpoint returned HTTP {resp.status_code}: {_format_as_error(resp)}"
f"refresh endpoint returned HTTP {resp.status_code}: {_format_as_error(resp)}",
failure_class=_classify_refresh_failure(resp),
)
try:
@@ -1239,11 +1351,128 @@ class TokenLookupResult:
would be wrong (the user can't fix it; only an operator can).
"""
kind: Literal["token", "missing", "decrypt_failure", "refresh_failed"] = "missing"
kind: Literal[
"token", "missing", "decrypt_failure", "refresh_failed", "refresh_failed_transient"
] = "missing"
token: str | None = None
decrypt_fingerprints: tuple[str, ...] = field(default_factory=tuple)
# In-process (per-node) backoff bookkeeping for transient refresh failures,
# keyed ``(user_id, server_name)`` on ``app_state.mcp_oauth_refresh_backoff``.
# The cooldown timer short-circuits the token-endpoint round-trip during a
# sustained AS outage (perf); the ambiguous streak escalates an
# unclassifiable-but-persistent rejection to re-consent so a dead grant in a
# non-standard shape can't strand the user forever.
_REFRESH_TRANSIENT_COOLDOWN_SECONDS = 30.0
_AMBIGUOUS_ESCALATION_THRESHOLD = 5
@dataclass
class _RefreshBackoffState:
"""Per-(user, server) transient-refresh backoff state (see the helpers below)."""
last_failure_monotonic: float = 0.0
ambiguous_streak: int = 0
def _refresh_backoff_state(app_state: Any, user_id: str, server_name: str) -> _RefreshBackoffState:
"""Return (creating if absent) the backoff state for ``(user_id, server_name)``."""
states = getattr(app_state, "mcp_oauth_refresh_backoff", None)
if states is None:
states = {}
app_state.mcp_oauth_refresh_backoff = states
key = (user_id, server_name)
state = states.get(key)
if state is None:
state = _RefreshBackoffState()
states[key] = state
return state
def _clear_refresh_backoff(app_state: Any, user_id: str, server_name: str) -> None:
"""Drop the backoff state for ``(user_id, server_name)``.
Called whenever a usable token is returned or the token is revoked, so a
healthy grant resets the cooldown timer + ambiguous streak and the dict
stays bounded to live ``(user, server)`` pairs.
"""
states = getattr(app_state, "mcp_oauth_refresh_backoff", None)
if isinstance(states, dict):
states.pop((user_id, server_name), None)
def _refresh_in_cooldown(app_state: Any, user_id: str, server_name: str) -> bool:
"""Return True while within the post-transient-failure cooldown window."""
states = getattr(app_state, "mcp_oauth_refresh_backoff", None)
if not isinstance(states, dict):
return False
state: _RefreshBackoffState | None = states.get((user_id, server_name))
if state is None or not state.last_failure_monotonic:
return False
elapsed = time.monotonic() - state.last_failure_monotonic
return elapsed < _REFRESH_TRANSIENT_COOLDOWN_SECONDS
def _token_result(
app_state: Any, user_id: str, server_name: str, token: str | None
) -> TokenLookupResult:
"""Return a ``kind="token"`` result, resetting any transient-refresh backoff.
A usable token means the grant is healthy, so the cooldown timer and the
ambiguous-failure streak are cleared here at the single success choke point.
"""
_clear_refresh_backoff(app_state, user_id, server_name)
return TokenLookupResult(kind="token", token=token)
def _no_token_result(
app_state: Any, user_id: str, server_name: str, result: TokenLookupResult
) -> TokenLookupResult:
"""Clear any transient-refresh backoff, then return a non-token *result*.
A ``missing`` / ``decrypt_failure`` outcome means the grant is no longer live
on this node (token deleted cluster-wide, key rotated), so the
per-(user, server) backoff is dropped to keep the dict bounded to live pairs
the mirror of :func:`_token_result` (success) and
:func:`_revoke_after_refresh_failure` (revoke). Backoff then survives only
the two intentional keep-paths: the transient handler and the cooldown gate.
"""
_clear_refresh_backoff(app_state, user_id, server_name)
return result
async def _revoke_after_refresh_failure(
app_state: Any,
token_store: MCPTokenStore,
user_id: str,
server_name: str,
server_id_for_audit: str,
*,
reason: str,
) -> TokenLookupResult:
"""Delete the stored token, emit a ``token_revoked`` audit, and drop locks.
The single revoke choke point for every "the grant is dead" outcome
(permanent rejection, ambiguous-streak escalation, expired-with-no-refresh).
Drops the per-key refresh lock and backoff state both safe to call when no
entry exists and returns ``refresh_failed`` so the dispatcher surfaces
re-consent.
"""
await asyncio.to_thread(token_store.delete_user_token, user_id, server_name)
await _audit_event(
app_state,
server_id=server_id_for_audit,
user_id=user_id,
action="mcp_server.oauth.token_revoked",
server_name=server_name,
detail={"reason": reason},
)
_drop_refresh_lock(app_state, user_id, server_name)
_clear_refresh_backoff(app_state, user_id, server_name)
return TokenLookupResult(kind="refresh_failed")
async def get_user_access_token(*, app_state: Any, user_id: str, server_name: str) -> str | None:
"""Return a valid plaintext access token, refreshing if needed.
@@ -1272,9 +1501,12 @@ async def get_user_access_token_classified(
"""Tagged token lookup with refresh-on-expiry.
Walks the token-lookup state machine (token / missing /
decrypt_failure / refresh_failed) and returns a tagged result so
the dispatcher can map each failure mode to the right user-facing
error.
decrypt_failure / refresh_failed / refresh_failed_transient) and
returns a tagged result so the dispatcher can map each failure mode
to the right user-facing error. A transient refresh failure keeps the
token and returns ``refresh_failed_transient`` (retryable, no revoke);
only a permanent rejection or a sustained run of ambiguous ones
revokes the stored token and returns ``refresh_failed``.
Multi-node correctness: the refresh path is serialized at two
layers an outer ``asyncio.Lock`` from :func:`_refresh_lock_for`
@@ -1318,42 +1550,54 @@ async def get_user_access_token_classified(
server_name=server_name,
exc_info=True,
)
return TokenLookupResult(
kind="decrypt_failure",
decrypt_fingerprints=tuple(exc.key_fingerprints_attempted),
return _no_token_result(
app_state,
user_id,
server_name,
TokenLookupResult(
kind="decrypt_failure",
decrypt_fingerprints=tuple(exc.key_fingerprints_attempted),
),
)
if plain is None:
return TokenLookupResult(kind="missing")
return _no_token_result(app_state, user_id, server_name, TokenLookupResult(kind="missing"))
expires_at = plain.get("expires_at")
if not force_refresh and not _token_needs_refresh(expires_at):
return TokenLookupResult(kind="token", token=plain["access_token"])
needs_refresh = _token_needs_refresh(expires_at)
if not force_refresh and not needs_refresh:
return _token_result(app_state, user_id, server_name, plain["access_token"])
# perf: during a sustained AS outage, short-circuit the token-endpoint
# round-trip for a brief window after a transient failure rather than
# re-attempting on every dispatch. Gated on the locally-read token being
# itself expired — a force_refresh on a still-fresh-looking token (the 401
# retry) falls through so the in-lock race check can still pick up a token a
# cluster-mate just refreshed.
if needs_refresh and _refresh_in_cooldown(app_state, user_id, server_name):
return TokenLookupResult(kind="refresh_failed_transient")
storage = _get_storage(app_state)
if storage is None:
return TokenLookupResult(kind="missing")
return _no_token_result(app_state, user_id, server_name, TokenLookupResult(kind="missing"))
# bug-6: load server_row BEFORE the no-refresh-token branch so the
# pre-lock audit event carries the immutable server_id rather than
# falling back to a name-keyed lookup that races admin renames.
server_row = await asyncio.to_thread(storage.get_mcp_server_by_name, server_name)
if server_row is None:
return TokenLookupResult(kind="missing")
return _no_token_result(app_state, user_id, server_name, TokenLookupResult(kind="missing"))
server_id_for_audit = str(server_row.get("server_id") or "")
refresh_value = plain.get("refresh_token")
if not refresh_value:
await asyncio.to_thread(token_store.delete_user_token, user_id, server_name)
_drop_refresh_lock(app_state, user_id, server_name)
await _audit_event(
return await _revoke_after_refresh_failure(
app_state,
server_id=server_id_for_audit,
user_id=user_id,
action="mcp_server.oauth.token_revoked",
server_name=server_name,
detail={"reason": "expired_no_refresh"},
token_store,
user_id,
server_name,
server_id_for_audit,
reason="expired_no_refresh",
)
return TokenLookupResult(kind="refresh_failed")
lock = _refresh_lock_for(app_state, user_id, server_name)
pg_lock = await _acquire_pg_refresh_lock(storage, user_id, server_name)
@@ -1372,12 +1616,19 @@ async def get_user_access_token_classified(
exc_info=True,
)
_drop_refresh_lock(app_state, user_id, server_name)
return TokenLookupResult(
kind="decrypt_failure",
decrypt_fingerprints=tuple(exc.key_fingerprints_attempted),
return _no_token_result(
app_state,
user_id,
server_name,
TokenLookupResult(
kind="decrypt_failure",
decrypt_fingerprints=tuple(exc.key_fingerprints_attempted),
),
)
if plain2 is None:
return TokenLookupResult(kind="missing")
return _no_token_result(
app_state, user_id, server_name, TokenLookupResult(kind="missing")
)
expires_at2 = plain2.get("expires_at")
# Reuse the freshly-refreshed token under two conditions:
# 1. ``force_refresh=False`` and the cached token is still fresh
@@ -1389,23 +1640,20 @@ async def get_user_access_token_classified(
# we lost the race and should reuse rather than refresh again.
if not _token_needs_refresh(expires_at2):
if not force_refresh:
return TokenLookupResult(kind="token", token=plain2["access_token"])
return _token_result(app_state, user_id, server_name, plain2["access_token"])
last_refreshed = _parse_iso_to_utc(plain2.get("last_refreshed") or "")
if last_refreshed is not None and last_refreshed >= t_lock_request_started:
return TokenLookupResult(kind="token", token=plain2["access_token"])
return _token_result(app_state, user_id, server_name, plain2["access_token"])
refresh_value2 = plain2.get("refresh_token")
if not refresh_value2:
await asyncio.to_thread(token_store.delete_user_token, user_id, server_name)
await _audit_event(
return await _revoke_after_refresh_failure(
app_state,
server_id=server_id_for_audit,
user_id=user_id,
action="mcp_server.oauth.token_revoked",
server_name=server_name,
detail={"reason": "expired_no_refresh"},
token_store,
user_id,
server_name,
server_id_for_audit,
reason="expired_no_refresh",
)
_drop_refresh_lock(app_state, user_id, server_name)
return TokenLookupResult(kind="refresh_failed")
try:
new_access, _new_refresh, _new_expires_at = await _refresh_and_persist(
@@ -1418,19 +1666,68 @@ async def get_user_access_token_classified(
refresh_value=refresh_value2,
existing_scopes=plain2.get("scopes") or "",
)
except MCPOAuthRefreshFailed:
await asyncio.to_thread(token_store.delete_user_token, user_id, server_name)
await _audit_event(
app_state,
server_id=server_id_for_audit,
except MCPOAuthRefreshFailed as exc:
if exc.failure_class is _RefreshFailureClass.PERMANENT:
# The AS rejected the grant as dead (invalid_grant / invalid_scope
# / an OIDC interaction-required code) — revoke and re-consent.
return await _revoke_after_refresh_failure(
app_state,
token_store,
user_id,
server_name,
server_id_for_audit,
reason="refresh_failed",
)
# Transient or ambiguous: keep the token (a blip must never revoke a
# user's consent) and arm the cooldown so a down AS isn't hit on
# every later dispatch.
backoff = _refresh_backoff_state(app_state, user_id, server_name)
backoff.last_failure_monotonic = time.monotonic()
if exc.failure_class is _RefreshFailureClass.AMBIGUOUS:
backoff.ambiguous_streak += 1
if backoff.ambiguous_streak >= _AMBIGUOUS_ESCALATION_THRESHOLD:
# A persistent 400/401 rejection we can't map to a standard
# code most likely IS a dead grant the AS reports in a
# non-standard shape. Escalate to re-consent so the user
# isn't stranded on a retryable error forever. (Infra
# transients never reach here, so an outage can't escalate.)
log.warning(
"mcp_server.oauth.refresh_ambiguous_escalated",
user_id=user_id,
server_name=server_name,
streak=backoff.ambiguous_streak,
error=str(exc),
)
return await _revoke_after_refresh_failure(
app_state,
token_store,
user_id,
server_name,
server_id_for_audit,
reason="refresh_failed_ambiguous_escalated",
)
else:
# A clean infra/operator-fixable transient breaks any ambiguous
# run — only an uninterrupted streak escalates.
backoff.ambiguous_streak = 0
log.warning(
"mcp_server.oauth.refresh_transient_failure",
user_id=user_id,
action="mcp_server.oauth.token_revoked",
server_name=server_name,
detail={"reason": "refresh_failed"},
failure_class=exc.failure_class.value,
ambiguous_streak=backoff.ambiguous_streak,
error=str(exc),
)
_drop_refresh_lock(app_state, user_id, server_name)
return TokenLookupResult(kind="refresh_failed")
return TokenLookupResult(kind="token", token=new_access)
# Do NOT drop the refresh lock here: the token is kept, so the
# per-key asyncio.Lock must stay registered to keep serializing
# concurrent refreshes. Dropping it would let a second concurrent
# caller mint a fresh lock and refresh the same token in parallel —
# with refresh-token rotation that races to invalid_grant and a
# spurious revoke (the exact bug this path prevents). The async-with
# still releases the lock on return; the entry is pruned when the
# token is later refreshed or revoked.
return TokenLookupResult(kind="refresh_failed_transient")
return _token_result(app_state, user_id, server_name, new_access)
def _token_needs_refresh(expires_at: str | None) -> bool:
+4
View File
@@ -372,6 +372,10 @@ audio.media-player {
color: var(--red);
}
.mcp-error-card.mcp-error-transient .mcp-error-icon {
color: var(--yellow);
}
.mcp-error-body {
display: flex;
flex-direction: column;
+6
View File
@@ -3091,6 +3091,10 @@ function _mcpErrorCategory(code) {
) {
return "operator";
}
if (code === "mcp_refresh_unavailable") {
// Soft, retryable state — a transient refresh failure, not a hard denial.
return "transient";
}
// Default for any other mcp_*_forbidden / unrecognised mcp_ code.
return "forbidden";
}
@@ -3104,6 +3108,8 @@ function _mcpErrorTitle(err) {
case "mcp_token_undecryptable_key_unknown":
case "mcp_oauth_url_insecure":
return "Operator action required";
case "mcp_refresh_unavailable":
return "Temporarily unavailable";
default:
return "Forbidden";
}
+3
View File
@@ -1699,6 +1699,9 @@ audio.media-player {
.mcp-error-card.mcp-error-operator .mcp-error-icon {
color: var(--red);
}
.mcp-error-card.mcp-error-transient .mcp-error-icon {
color: var(--yellow);
}
.mcp-error-body {
display: flex;
flex-direction: column;