mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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";
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user