Skip to content

Commit 70e1cbd

Browse files
inamdarmihircopybara-github
authored andcommitted
fix: ask for re-authorization when the OAuth2 refresh token is revoked
Merge #7731 When the token endpoint rejects a refresh with invalid_grant, OAuth2CredentialRefresher.refresh now raises a CredentialRefresherError subclass instead of returning the stale credential, and CredentialManager.get_auth_credential returns None and removes the dead credential from the credential service, so authenticated tools and toolsets ask the user to authorize again. A client-credentials flow exchanges its configured credential again instead. Other refresh failures still keep the stale token. The credential manager now calls BaseCredentialService.save_credential with exchanged_auth_credential=None to remove a stored credential; the save_credential docstring documents that None means remove, so that load_credential returns None afterwards. If a custom credential service raises on None, the error is logged and the call continues. Fixes #7730 PiperOrigin-RevId: 1013682546
1 parent 36cb10d commit 70e1cbd

6 files changed

Lines changed: 605 additions & 8 deletions

File tree

‎src/google/adk/auth/credential_manager.py‎

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
from .exchanger.base_credential_exchanger import BaseCredentialExchanger
3939
from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry
4040
from .oauth2_discovery import OAuth2DiscoveryManager
41+
from .refresher.base_credential_refresher import _CredentialRevokedError
4142
from .refresher.credential_refresher_registry import CredentialRefresherRegistry
4243

4344
logger = logging.getLogger("google_adk." + __name__)
@@ -331,7 +332,53 @@ async def get_auth_credential(
331332
# Step 7: Refresh credential if expired
332333
was_refreshed = False
333334
if not was_exchanged:
334-
credential, was_refreshed = await self._refresh_credential(credential)
335+
try:
336+
credential, was_refreshed = await self._refresh_credential(credential)
337+
except _CredentialRevokedError:
338+
logger.warning(
339+
"The refresh token for auth scheme %s was rejected.",
340+
type(self._auth_config.auth_scheme).__name__,
341+
)
342+
if was_from_auth_response:
343+
return None
344+
revoked_credential = credential
345+
# Evict the dead stored credential so later turns load the user's
346+
# new authorization instead of refreshing it again.
347+
try:
348+
await self._save_credential(context, None)
349+
except Exception: # pylint: disable=broad-exception-caught
350+
# A failed eviction only costs one more rejected refresh next turn.
351+
logger.warning(
352+
"Failed to remove the revoked credential for auth scheme %s.",
353+
type(self._auth_config.auth_scheme).__name__,
354+
exc_info=True,
355+
)
356+
# A stored dead credential must not shadow a fresh auth response.
357+
credential = await self._load_from_auth_response(context)
358+
revoked_oauth2 = revoked_credential.oauth2
359+
new_oauth2 = credential.oauth2 if credential else None
360+
if (
361+
not credential
362+
or credential == revoked_credential
363+
or (
364+
revoked_oauth2
365+
and new_oauth2
366+
and revoked_oauth2.refresh_token
367+
and new_oauth2.refresh_token == revoked_oauth2.refresh_token
368+
)
369+
):
370+
return None
371+
was_from_auth_response = True
372+
credential, was_exchanged = await self._exchange_credential(
373+
credential
374+
)
375+
if not was_exchanged:
376+
try:
377+
credential, was_refreshed = await self._refresh_credential(
378+
credential
379+
)
380+
except _CredentialRevokedError:
381+
return None
335382

336383
if _has_unredeemed_auth_code(credential):
337384
logger.warning(
@@ -432,9 +479,24 @@ async def _refresh_credential(
432479
if was_exchanged:
433480
return new_credential, True
434481
return credential, False
435-
refreshed_credential = await refresher.refresh(
436-
credential, self._auth_config.auth_scheme
437-
)
482+
try:
483+
refreshed_credential = await refresher.refresh(
484+
credential, self._auth_config.auth_scheme
485+
)
486+
except _CredentialRevokedError:
487+
if (
488+
raw_auth_credential is None
489+
or not self._is_client_credentials_flow()
490+
):
491+
raise
492+
# A client credentials flow can get a new token without the user. Copy
493+
# the raw credential so the exchange does not mutate the shared config.
494+
new_credential, was_exchanged = await self._exchange_credential(
495+
raw_auth_credential.model_copy(deep=True)
496+
)
497+
if not was_exchanged:
498+
raise
499+
return new_credential, True
438500
return refreshed_credential, True
439501

440502
return credential, False
@@ -488,9 +550,9 @@ async def _validate_credential(self) -> None:
488550
# Additional validation can be added here
489551

490552
async def _save_credential(
491-
self, context: CallbackContext, credential: AuthCredential
553+
self, context: CallbackContext, credential: AuthCredential | None
492554
) -> None:
493-
"""Save credential to credential service if available."""
555+
"""Save credential to credential service if available; None removes it."""
494556
credential_service = context._invocation_context.credential_service
495557
if credential_service:
496558
auth_config_to_save = self._auth_config.model_copy(deep=True)

‎src/google/adk/auth/credential_service/base_credential_service.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@ async def save_credential(
6363
Saves the exchanged_auth_credential in auth config to the backend credential
6464
store.
6565
66+
An exchanged_auth_credential of None means the stored credential must be
67+
removed, so that load_credential returns None for that key afterwards.
68+
6669
Args:
6770
auth_config: The auth config which contains the auth scheme and auth
6871
credential information. auth_config.get_credential_key will be used to

‎src/google/adk/auth/refresher/base_credential_refresher.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,14 @@ class CredentialRefresherError(Exception):
2727
"""Base exception for credential refresh errors."""
2828

2929

30+
class _CredentialRevokedError(CredentialRefresherError):
31+
"""Raised when the refresh token was permanently rejected.
32+
33+
The credential can never be refreshed again, so only a new user
34+
authorization can replace it.
35+
"""
36+
37+
3038
class BaseCredentialRefresher(abc.ABC):
3139
"""Base interface for credential refreshers.
3240

‎src/google/adk/auth/refresher/oauth2_credential_refresher.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import requests
2828
from typing_extensions import override
2929

30+
from .base_credential_refresher import _CredentialRevokedError
3031
from .base_credential_refresher import BaseCredentialRefresher
3132

3233
try:
@@ -80,7 +81,8 @@ async def refresh(
8081
auth_scheme: Optional[AuthScheme] = None,
8182
) -> AuthCredential:
8283
"""Refresh the OAuth2 credential.
83-
If refresh failed, return the original credential.
84+
85+
If refresh failed transiently, return the original credential.
8486
8587
Args:
8688
auth_credential: The OAuth2 credential to refresh.
@@ -89,6 +91,9 @@ async def refresh(
8991
Returns:
9092
The refreshed credential.
9193
94+
Raises:
95+
CredentialRefresherError: If the authorization server rejected the
96+
refresh token with `invalid_grant` (revoked or expired).
9297
"""
9398

9499
# Handle regular OAuth2 credentials
@@ -128,6 +133,12 @@ async def refresh(
128133
update_credential_with_tokens(auth_credential, tokens)
129134
logger.debug("Successfully refreshed OAuth2 tokens")
130135
except (AuthlibBaseError, requests.RequestException) as e:
136+
if getattr(e, "error", None) == "invalid_grant":
137+
# Permanent: the refresh token is dead, so a stale access token
138+
# would only fail later; the caller must ask for re-authorization.
139+
raise _CredentialRevokedError(
140+
f"OAuth2 refresh token was rejected: {e}"
141+
) from e
131142
# Non-fatal: keep the stale token so its eventual 401
132143
# re-triggers auth.
133144
logger.error("Failed to refresh OAuth2 tokens: %s", e)

‎tests/unittests/auth/refresher/test_oauth2_credential_refresher.py‎

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,15 @@
1616
from unittest.mock import Mock
1717
from unittest.mock import patch
1818

19+
from authlib.integrations.base_client.errors import OAuthError
1920
from authlib.oauth2.rfc6749 import OAuth2Token
21+
from authlib.oauth2.rfc6749.errors import InvalidGrantError
2022
from authlib.oauth2.rfc6749.errors import OAuth2Error
2123
from google.adk.auth.auth_credential import AuthCredential
2224
from google.adk.auth.auth_credential import AuthCredentialTypes
2325
from google.adk.auth.auth_credential import OAuth2Auth
2426
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
27+
from google.adk.auth.refresher import base_credential_refresher
2528
from google.adk.auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
2629
import pytest
2730
import requests
@@ -195,7 +198,7 @@ async def test_refresh_oauth2_error_returns_original_and_logs(
195198
mock_client = Mock()
196199
mock_oauth2_session.return_value = mock_client
197200
mock_client.refresh_token.side_effect = OAuth2Error(
198-
description="invalid_grant"
201+
description="temporarily_unavailable"
199202
)
200203

201204
scheme = OpenIdConnectWithConfig(
@@ -309,6 +312,94 @@ async def test_refresh_unexpected_error_propagates(
309312
with pytest.raises(ValueError, match="unexpected bug"):
310313
await refresher.refresh(credential, scheme)
311314

315+
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
316+
@patch("google.adk.auth.oauth2_credential_util.OAuth2Token")
317+
@pytest.mark.parametrize(
318+
"error",
319+
[
320+
InvalidGrantError(description="Token has been revoked."),
321+
OAuth2Error(error="invalid_grant"),
322+
OAuthError(error="invalid_grant", description="Bad Request"),
323+
],
324+
)
325+
@pytest.mark.asyncio
326+
async def test_refresh_invalid_grant_raises_revoked(
327+
self, mock_oauth2_token, mock_oauth2_session, error
328+
):
329+
"""invalid_grant means the refresh token is dead and must not be kept."""
330+
mock_token_instance = Mock()
331+
mock_token_instance.is_expired.return_value = True
332+
mock_oauth2_token.return_value = mock_token_instance
333+
334+
mock_client = Mock()
335+
mock_oauth2_session.return_value = mock_client
336+
mock_client.refresh_token.side_effect = error
337+
338+
scheme = OpenIdConnectWithConfig(
339+
type_="openIdConnect",
340+
openId_connect_url=(
341+
"https://example.com/.well-known/openid_configuration"
342+
),
343+
authorization_endpoint="https://example.com/auth",
344+
token_endpoint="https://example.com/token",
345+
scopes=["openid"],
346+
)
347+
credential = AuthCredential(
348+
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
349+
oauth2=OAuth2Auth(
350+
client_id="test_client_id",
351+
client_secret="test_client_secret",
352+
access_token="old_token",
353+
refresh_token="old_refresh_token",
354+
expires_at=int(time.time()) - 3600, # Expired
355+
),
356+
)
357+
358+
refresher = OAuth2CredentialRefresher()
359+
with pytest.raises(base_credential_refresher._CredentialRevokedError):
360+
await refresher.refresh(credential, scheme)
361+
362+
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
363+
@patch("google.adk.auth.oauth2_credential_util.OAuth2Token")
364+
@pytest.mark.asyncio
365+
async def test_refresh_other_oauth2_error_returns_original(
366+
self, mock_oauth2_token, mock_oauth2_session
367+
):
368+
"""An OAuth2 error other than invalid_grant is transient: keep the token."""
369+
mock_token_instance = Mock()
370+
mock_token_instance.is_expired.return_value = True
371+
mock_oauth2_token.return_value = mock_token_instance
372+
373+
mock_client = Mock()
374+
mock_oauth2_session.return_value = mock_client
375+
mock_client.refresh_token.side_effect = OAuth2Error(error="server_error")
376+
377+
scheme = OpenIdConnectWithConfig(
378+
type_="openIdConnect",
379+
openId_connect_url=(
380+
"https://example.com/.well-known/openid_configuration"
381+
),
382+
authorization_endpoint="https://example.com/auth",
383+
token_endpoint="https://example.com/token",
384+
scopes=["openid"],
385+
)
386+
credential = AuthCredential(
387+
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
388+
oauth2=OAuth2Auth(
389+
client_id="test_client_id",
390+
client_secret="test_client_secret",
391+
access_token="old_token",
392+
refresh_token="old_refresh_token",
393+
expires_at=int(time.time()) - 3600, # Expired
394+
),
395+
)
396+
397+
refresher = OAuth2CredentialRefresher()
398+
result = await refresher.refresh(credential, scheme)
399+
400+
assert result is credential
401+
assert result.oauth2.access_token == "old_token"
402+
312403
@patch("google.adk.auth.refresher.oauth2_credential_refresher.logger")
313404
@pytest.mark.asyncio
314405
async def test_refresh_without_refresh_token_returns_original_and_logs(

0 commit comments

Comments
 (0)