diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 802278a9..2ecef140 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -2535,7 +2535,7 @@ class MCPClientManager: """ if not user_id or self._loop is None: return - if not (self._oauth_user_server_names or self._obo_server_names): + if not self._pool_server_names: return if self._app_state is None or self._storage is None: return @@ -2649,7 +2649,7 @@ class MCPClientManager: finally: self._priming_keys.discard(key) - prime_names = self._oauth_user_server_names | self._obo_server_names + prime_names = self._pool_server_names if self._obo_server_names: # One raw existence SELECT (no decrypt) decides ALL obo servers for # this user: a user with no captured credential — a local-auth @@ -4675,7 +4675,7 @@ class MCPClientManager: result: dict[str, dict[str, Any]] = {} for name in list(self._server_configs): result[name] = self.get_server_status(name, user_id, aggregate=aggregate) - for name in list(self._oauth_user_server_names | self._obo_server_names): + for name in list(self._pool_server_names): if name not in result: result[name] = self.get_server_status(name, user_id, aggregate=aggregate) return result @@ -4935,6 +4935,17 @@ class MCPClientManager: """ return self.server_auth_type(server_name) is not None + @property + def _pool_server_names(self) -> set[str]: + """Union of the per-auth-type pool registries (oauth_user + oauth_obo). + + The set-level counterpart to :meth:`_is_pool_server`, so the iteration + sites (priming, keep-hot sweep, status) share ONE definition of "which + servers are pool-backed" — a future third pool-backed auth type is added + in one place instead of being missed at an ad-hoc inline union. + """ + return self._oauth_user_server_names | self._obo_server_names + @property def server_count(self) -> int: return sum(1 for s in self._static_servers.values() if s.session is not None) @@ -7221,77 +7232,74 @@ def _build_consent_url( return f"/v1/api/mcp/oauth/start?{qs}" -def _consent_missing_detail(server_row: dict[str, Any]) -> str: - """User-facing detail for a ``missing`` token lookup, per auth model. +# User-facing remediation copy for every (auth model, failure situation) a +# pool-backed dispatch can surface. Consolidated into ONE table so the +# oauth_user-vs-oauth_obo split lives in a single place instead of a branch +# fanned across four helpers: a lookup situation that must show obo re-login / +# admin guidance can't silently keep oauth_user's "re-consent at a per-server +# flow that doesn't exist for obo" copy — the exact wrong-remediation bug this +# feature's review caught repeatedly. ``{kind}`` / ``{kind_cap}`` are filled per +# call (only the insufficient_scope rows use them; other rows ignore them). +_POOL_ERROR_DETAIL: dict[tuple[str, str], str] = { + ("oauth_user", "missing"): "No token for user. Consent flow required.", + ("oauth_obo", "missing"): ( + "No sign-in credential for this account. Sign in to Turnstone again to reconnect." + ), + ("oauth_user", "refresh_failed"): "Refresh token rejected. Re-consent required.", + ("oauth_obo", "refresh_failed"): ( + "Sign-in credential was rejected for this server. Sign in to Turnstone " + "again; if it keeps failing, your administrator may need to grant access." + ), + ("oauth_user", "insufficient_scope"): ( + "{kind_cap} requires elevated scopes. Re-consent flow with new scopes required." + ), + ("oauth_obo", "insufficient_scope"): ( + "This {kind} needs additional permissions your sign-in token does not " + "carry. Ask your administrator to grant the required access (add the " + "scope to the server, or widen your delegated permissions at the " + "identity provider)." + ), + ("oauth_user", "token_rejected"): "Refreshed token still rejected. Re-consent required.", + ("oauth_obo", "token_rejected"): ( + "Server rejected a freshly issued sign-in token. This usually means the " + "server's audience setting doesn't match what the server expects — ask " + "your administrator to check it." + ), +} - ``oauth_obo`` has no per-server consent flow — the missing thing is the - user's captured sign-in credential, so the affordance is a re-login - (``_build_consent_url`` already returns None for these servers, so no - per-server Connect button is advertised). + +def _pool_error_detail(server_row: dict[str, Any], situation: str, *, kind: str = "") -> str: + """Remediation copy for *situation* on *server_row*, chosen by auth model. + + Single lookup into :data:`_POOL_ERROR_DETAIL` — the ONE place the + oauth_user-vs-oauth_obo decision is made — so no dispatch site can pair a + situation with the wrong auth model's copy. ``oauth_obo`` rows never have a + per-server consent flow (``_build_consent_url`` returns None), so their copy + points at a re-login / administrator remedy rather than a dead-end consent + card; everything else is treated as ``oauth_user``. """ - if server_row.get("auth_type") == "oauth_obo": - return "No sign-in credential for this account. Sign in to Turnstone again to reconnect." - return "No token for user. Consent flow required." + model = "oauth_obo" if server_row.get("auth_type") == "oauth_obo" else "oauth_user" + return _POOL_ERROR_DETAIL[(model, situation)].format(kind=kind, kind_cap=kind.capitalize()) + + +def _consent_missing_detail(server_row: dict[str, Any]) -> str: + """Detail for a ``missing`` token lookup (see :func:`_pool_error_detail`).""" + return _pool_error_detail(server_row, "missing") def _refresh_failed_detail(server_row: dict[str, Any]) -> str: - """User-facing detail for a ``refresh_failed`` (permanent) lookup, per auth model. - - For ``oauth_obo`` there is no per-server consent to redo and - ``_build_consent_url`` returns None, so "re-consent required" is a dead end — - the remedy is an admin tenant-grant fix or a re-login. Point the user there - instead of at a consent flow that does not exist for this auth type. - """ - if server_row.get("auth_type") == "oauth_obo": - return ( - "Sign-in credential was rejected for this server. Sign in to Turnstone " - "again; if it keeps failing, your administrator may need to grant access." - ) - return "Refresh token rejected. Re-consent required." + """Detail for a permanent ``refresh_failed`` lookup (see :func:`_pool_error_detail`).""" + return _pool_error_detail(server_row, "refresh_failed") def _insufficient_scope_detail(server_row: dict[str, Any], kind: str) -> str: - """User-facing detail for a 403 insufficient_scope, per auth model. - - ``oauth_user``: a per-server step-up re-consent with the widened scope set - is the remedy (``_build_consent_url`` attaches the affordance). ``oauth_obo``: - there is no per-server consent flow — the minted token's scopes come from - the user's captured sign-in (entra: the app's delegated permissions; - rfc8693: the server's configured ``oauth_scopes``), so a user-driven - step-up is impossible (``_build_consent_url`` returns None). Point at the - real fix — an administrator widening access — instead of a dead-end - re-consent card. - """ - if server_row.get("auth_type") == "oauth_obo": - return ( - f"This {kind} needs additional permissions your sign-in token does not " - "carry. Ask your administrator to grant the required access (add the " - "scope to the server, or widen your delegated permissions at the " - "identity provider)." - ) - return ( - f"{kind.capitalize()} requires elevated scopes. Re-consent flow with new scopes required." - ) + """Detail for a 403 insufficient_scope (see :func:`_pool_error_detail`).""" + return _pool_error_detail(server_row, "insufficient_scope", kind=kind) def _token_rejected_detail(server_row: dict[str, Any]) -> str: - """User-facing detail for a 401 that survived one forced refresh (retry ceiling). - - ``oauth_user``: the freshly refreshed bearer was rejected — per-server - re-consent is the actionable remedy (``_build_consent_url`` attaches the - Connect affordance). ``oauth_obo``: the bearer was JUST minted from the - captured credential and there is no per-server consent flow - (``_build_consent_url`` returns None), so "re-consent required" is a dead - end — the realistic causes are server-side (audience mismatch, upstream - auth misconfig, clock skew), so point at the administrator instead. - """ - if server_row.get("auth_type") == "oauth_obo": - return ( - "Server rejected a freshly issued sign-in token. This usually means the " - "server's audience setting doesn't match what the server expects — ask " - "your administrator to check it." - ) - return "Refreshed token still rejected. Re-consent required." + """Detail for a 401 that survived one forced refresh (see :func:`_pool_error_detail`).""" + return _pool_error_detail(server_row, "token_rejected") def _pool_lookup_error( diff --git a/turnstone/core/mcp_oauth.py b/turnstone/core/mcp_oauth.py index 840bc60e..920eef9b 100644 --- a/turnstone/core/mcp_oauth.py +++ b/turnstone/core/mcp_oauth.py @@ -1435,6 +1435,21 @@ def _refresh_backoff_state(app_state: Any, user_id: str, server_name: str) -> _R return state +def _arm_cooldown(app_state: Any, user_id: str, server_name: str) -> _RefreshBackoffState: + """Stamp the per-(user, server) transient-failure cooldown clock to now. + + Single definition of the "back off this pair" operation (previously written + inline at every failure site) so a change to how the cooldown is armed — + jitter, a min-interval, a second timestamp — is one edit, not four, and a + missed site can't silently keep hammering the AS/IdP on that path. Returns + the backoff state so a caller that also mutates the ambiguous streak reuses + the same object instead of re-fetching it. + """ + state = _refresh_backoff_state(app_state, user_id, server_name) + state.last_failure_monotonic = time.monotonic() + 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)``. @@ -1687,15 +1702,12 @@ async def _handle_refresh_failure( if arm_cooldown_on_permanent: # _revoke_after_refresh_failure cleared the backoff; re-arm the # cooldown as the terminal backstop for the surviving credential. - _refresh_backoff_state( - app_state, user_id, server_name - ).last_failure_monotonic = time.monotonic() + _arm_cooldown(app_state, user_id, server_name) return result # 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() + backoff = _arm_cooldown(app_state, user_id, server_name) if exc.failure_class is _RefreshFailureClass.AMBIGUOUS: backoff.ambiguous_streak += 1 if backoff.ambiguous_streak >= _AMBIGUOUS_ESCALATION_THRESHOLD: @@ -2383,9 +2395,7 @@ async def get_obo_access_token_classified( # on the next tick after the window lapses. (The write path also rejects # audience-less oauth_obo rows, so this branch is normally a typo'd # grant profile, not a common state.) - _refresh_backoff_state( - app_state, user_id, server_name - ).last_failure_monotonic = time.monotonic() + _arm_cooldown(app_state, user_id, server_name) log.error( "mcp_server.oauth.obo_misconfigured", server_name=server_name, @@ -2396,12 +2406,16 @@ async def get_obo_access_token_classified( return TokenLookupResult(kind="refresh_failed_transient") issuer = str(getattr(oidc_config, "issuer", "")) - credential_result = await _read_obo_credential( - app_state, token_store, user_id, server_name, issuer - ) - if isinstance(credential_result, TokenLookupResult): - return credential_result # missing or decrypt_failure - # credential present (unlocked pre-check; re-read authoritatively under lock) + # Cheap pre-lock presence check: a raw existence read (NO decrypt) is enough + # to short-circuit the common "no captured credential" case before taking + # the pg advisory lock. The authoritative decrypt happens exactly once under + # the lock (credential2 below), where decrypt_failure is already classified — + # so the refresh token is never Fernet-decrypted twice per mint (which, at + # session-start priming across N obo servers, was N redundant decrypts per + # user). Mirrors the raw existence guard the priming path already uses. + if await asyncio.to_thread(storage.get_oidc_user_credential, user_id, issuer) is None: + # No captured credential → the consent affordance is a re-login. + return _no_token_result(app_state, user_id, server_name, TokenLookupResult(kind="missing")) lock = _refresh_lock_for(app_state, user_id, server_name) credential_key = f"__obo__:{issuer}" @@ -2507,8 +2521,7 @@ async def get_obo_access_token_classified( access_token = tokens.get("access_token") if not isinstance(access_token, str) or not access_token: - backoff = _refresh_backoff_state(app_state, user_id, server_name) - backoff.last_failure_monotonic = time.monotonic() + backoff = _arm_cooldown(app_state, user_id, server_name) # A malformed 200 is a clean transient (not a dead grant): reset the # ambiguous streak, matching the sibling's _refresh_and_persist path. backoff.ambiguous_streak = 0