diff --git a/tests/test_mcp_oauth_refresh.py b/tests/test_mcp_oauth_refresh.py index 7faa51a9..5774fce4 100644 --- a/tests/test_mcp_oauth_refresh.py +++ b/tests/test_mcp_oauth_refresh.py @@ -378,11 +378,12 @@ class TestRefreshFailureClassification: # 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.""" + def test_backoff_and_lock_cleared_when_token_vanishes(self, storage: SQLiteBackend) -> None: + """A transient failure retains BOTH sibling per-(user,server) entries — the + refresh lock (for serialization) and the backoff (for the cooldown). If + the token is then deleted cluster-wide (another node's permanent revoke), + the next lookup returns ``missing`` AND prunes both, so neither in-process + dict grows unboundedly on the missing path.""" _seed_server(storage) client = MagicMock(spec=httpx.AsyncClient) client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc())) @@ -392,16 +393,19 @@ class TestRefreshFailureClassification: state = _make_app_state(storage, http_client=client) _seed_token(state, expires_in_seconds=-1000) - # First lookup: a transient 503 records a backoff entry. + # First lookup: a transient 503 records a backoff entry AND retains the + # refresh lock (the keep-path must not drop it — bug-1). assert self._lookup(state).kind == "refresh_failed_transient" assert ("user-1", "srv-oauth") in state.mcp_oauth_refresh_backoff + assert ("user-1", "srv-oauth") in state.mcp_oauth_refresh_locks # 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. + # Next lookup sees the row gone -> missing -> both stale entries cleared. assert self._lookup(state).kind == "missing" assert ("user-1", "srv-oauth") not in state.mcp_oauth_refresh_backoff + assert ("user-1", "srv-oauth") not in state.mcp_oauth_refresh_locks # --------------------------------------------------------------------------- diff --git a/turnstone/core/mcp_oauth.py b/turnstone/core/mcp_oauth.py index 1ddcc7b9..cead45c9 100644 --- a/turnstone/core/mcp_oauth.py +++ b/turnstone/core/mcp_oauth.py @@ -1429,15 +1429,19 @@ def _token_result( 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*. + """Drop per-(user, server) refresh state, 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. + on this node (token deleted cluster-wide, key rotated), so BOTH sibling + per-(user, server) dicts are pruned — the refresh lock and the transient + backoff — keeping each bounded to live pairs (the mirror of + :func:`_token_result` on success and :func:`_revoke_after_refresh_failure` + on revoke). The transient keep-path deliberately retains the lock (so + concurrent same-key refreshes stay serialized) and the backoff (for the + cooldown), so without this prune a token that vanishes after a transient + failure would strand both entries. """ + _drop_refresh_lock(app_state, user_id, server_name) _clear_refresh_backoff(app_state, user_id, server_name) return result @@ -1615,7 +1619,6 @@ async def get_user_access_token_classified( server_name=server_name, exc_info=True, ) - _drop_refresh_lock(app_state, user_id, server_name) return _no_token_result( app_state, user_id,