ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

refactor(mcp): dedup obo credential decrypt, cooldown arming, pool set, error copy

Browse Source 
Round-4 review follow-up — no correctness findings; these are the four
cleanups it surfaced.

- The obo mint path decrypted the captured IdP refresh token twice per
  mint: once pre-lock only to test presence, then again under the lock.
  The pre-lock presence check now uses the raw existence read (no
  decrypt), mirroring the priming path; the single authoritative decrypt
  happens under the lock. Removes N throwaway decrypts per user at
  session-start priming across N obo servers.

- The "arm the per-(user,server) cooldown" idiom was written inline at
  four failure sites. Extracted _arm_cooldown (returns the backoff state
  so the streak-mutating callers reuse it), so a change to how backoff
  works is one edit.

- The oauth_user|obo pool-membership union was rebuilt inline at three
  iteration sites. Added a _pool_server_names property, the set-level
  counterpart to _is_pool_server, so a future third pool-backed auth type
  is registered in one place.

- The four per-situation remediation-copy helpers each repeated the
  oauth_user-vs-obo branch. Consolidated the copy into one
  (auth_model, situation) table behind _pool_error_detail — the single
  place the auth-model decision is made — so a dispatch site can't pair a
  situation with the wrong auth model's copy (the wrong-remediation bug
  class this review caught repeatedly). The named helpers remain as thin,
  tested wrappers.
This commit is contained in:
Patrick Buckley
2026-07-12 06:29:25 -07:00
parent 6d80051925
commit c53bd464d0
2 changed files with 100 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
turnstone/core/mcp_client.py
+71 -63
View File
@@ -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(
turnstone/core/mcp_oauth.py
+29 -16
View File
@@ -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