Skip to content

Commit cbe786b

Browse files
AkshayShah03copybara-github
authored andcommitted
fix: handle an expired OAuth2 token that has no refresh token
Merge #7729 An expired OAuth2 token without a refresh token no longer raises AttributeError. For the client credentials flow a new token is fetched with the configured client credentials; other flows keep the stale token and log a warning, as for a failed refresh. Fixes #7728 PiperOrigin-RevId: 1013172404
1 parent d118089 commit cbe786b

4 files changed

Lines changed: 226 additions & 0 deletions

File tree

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -416,6 +416,22 @@ async def _refresh_credential(
416416
if await refresher.is_refresh_needed(
417417
credential, self._auth_config.auth_scheme
418418
):
419+
raw_auth_credential = self._auth_config.raw_auth_credential
420+
if (
421+
raw_auth_credential is not None
422+
and self._is_client_credentials_flow()
423+
and credential.oauth2
424+
and not credential.oauth2.refresh_token
425+
):
426+
# A client credentials token has no refresh token; get a new one from
427+
# the client credentials instead. Copy them so the exchange does not
428+
# mutate the shared tool config.
429+
new_credential, was_exchanged = await self._exchange_credential(
430+
raw_auth_credential.model_copy(deep=True)
431+
)
432+
if was_exchanged:
433+
return new_credential, True
434+
return credential, False
419435
refreshed_credential = await refresher.refresh(
420436
credential, self._auth_config.auth_scheme
421437
)

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,14 @@ async def refresh(
103103
"expires_at": auth_credential.oauth2.expires_at,
104104
"expires_in": auth_credential.oauth2.expires_in,
105105
}).is_expired():
106+
if not auth_credential.oauth2.refresh_token:
107+
# Client credentials tokens, and some providers' user tokens, come
108+
# without a refresh token. Keep the stale token, as for a failed
109+
# refresh below.
110+
logger.warning(
111+
"OAuth2 token expired and has no refresh token; cannot refresh"
112+
)
113+
return auth_credential
106114
client, token_endpoint = create_oauth2_session(
107115
auth_scheme, auth_credential
108116
)

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,3 +308,35 @@ async def test_refresh_unexpected_error_propagates(
308308
refresher = OAuth2CredentialRefresher()
309309
with pytest.raises(ValueError, match="unexpected bug"):
310310
await refresher.refresh(credential, scheme)
311+
312+
@patch("google.adk.auth.refresher.oauth2_credential_refresher.logger")
313+
@pytest.mark.asyncio
314+
async def test_refresh_without_refresh_token_returns_original_and_logs(
315+
self, mock_logger
316+
):
317+
"""An expired token with no refresh token is returned as is, not a crash."""
318+
scheme = OpenIdConnectWithConfig(
319+
type_="openIdConnect",
320+
openId_connect_url=(
321+
"https://example.com/.well-known/openid_configuration"
322+
),
323+
authorization_endpoint="https://example.com/auth",
324+
token_endpoint="https://example.com/token",
325+
scopes=["openid"],
326+
)
327+
credential = AuthCredential(
328+
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
329+
oauth2=OAuth2Auth(
330+
client_id="test_client_id",
331+
client_secret="test_client_secret",
332+
access_token="old_token",
333+
expires_at=int(time.time()) - 3600, # Expired
334+
),
335+
)
336+
337+
refresher = OAuth2CredentialRefresher()
338+
result = await refresher.refresh(credential, scheme)
339+
340+
assert result is credential
341+
assert result.oauth2.access_token == "old_token"
342+
mock_logger.warning.assert_called_once()

‎tests/unittests/auth/test_credential_manager.py‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import asyncio
16+
import time
1617
from unittest.mock import ANY
1718
from unittest.mock import AsyncMock
1819
from unittest.mock import Mock
@@ -1180,6 +1181,175 @@ async def test_refresh_credential_no_refresher(self):
11801181
assert result == mock_credential
11811182
assert was_refreshed is False
11821183

1184+
@pytest.mark.asyncio
1185+
@pytest.mark.parametrize(
1186+
"auth_scheme, auth_type",
1187+
[
1188+
(
1189+
OAuth2(
1190+
flows=OAuthFlows(
1191+
clientCredentials=OAuthFlowClientCredentials(
1192+
tokenUrl="https://example.com/token",
1193+
scopes={},
1194+
)
1195+
)
1196+
),
1197+
AuthCredentialTypes.OAUTH2,
1198+
),
1199+
(
1200+
OpenIdConnectWithConfig(
1201+
authorization_endpoint="https://example.com/auth",
1202+
token_endpoint="https://example.com/token",
1203+
grant_types_supported=["client_credentials"],
1204+
),
1205+
AuthCredentialTypes.OPEN_ID_CONNECT,
1206+
),
1207+
],
1208+
ids=["oauth2", "oidc"],
1209+
)
1210+
async def test_get_auth_credential_refetches_expired_client_credentials_token(
1211+
self, auth_scheme, auth_type
1212+
):
1213+
"""Client credentials tokens have no refresh token, so fetch a new one."""
1214+
raw_credential = AuthCredential(
1215+
auth_type=auth_type,
1216+
oauth2=OAuth2Auth(
1217+
client_id="mock_client_id",
1218+
client_secret="mock_client_secret",
1219+
),
1220+
)
1221+
stored_credential = raw_credential.model_copy(deep=True)
1222+
stored_credential.oauth2.access_token = "expired_token"
1223+
stored_credential.oauth2.expires_at = int(time.time()) - 60
1224+
manager = CredentialManager(
1225+
AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=raw_credential)
1226+
)
1227+
manager._load_existing_credential = AsyncMock(
1228+
return_value=stored_credential
1229+
)
1230+
manager._save_credential = AsyncMock()
1231+
context = Mock()
1232+
1233+
with patch.object(
1234+
oauth2_credential_util, "OAuth2Session", autospec=True
1235+
) as mock_oauth2_session:
1236+
mock_client = mock_oauth2_session.return_value
1237+
mock_client.fetch_token.return_value = {
1238+
"access_token": "new_token",
1239+
"expires_at": int(time.time()) + 3600,
1240+
}
1241+
# There is no refresh token to send.
1242+
mock_client.refresh_token.side_effect = AssertionError(
1243+
"refresh_token() called without a refresh token"
1244+
)
1245+
result = await manager.get_auth_credential(context)
1246+
1247+
assert result.oauth2.access_token == "new_token"
1248+
mock_client.fetch_token.assert_called_once()
1249+
manager._save_credential.assert_awaited_once_with(context, result)
1250+
# The shared tool config is not mutated.
1251+
assert raw_credential.oauth2.access_token is None
1252+
1253+
@pytest.mark.asyncio
1254+
async def test_get_auth_credential_refreshes_token_that_has_refresh_token(
1255+
self,
1256+
):
1257+
"""A client credentials token that has a refresh token is refreshed."""
1258+
auth_scheme = OAuth2(
1259+
flows=OAuthFlows(
1260+
clientCredentials=OAuthFlowClientCredentials(
1261+
tokenUrl="https://example.com/token",
1262+
scopes={},
1263+
)
1264+
)
1265+
)
1266+
raw_credential = AuthCredential(
1267+
auth_type=AuthCredentialTypes.OAUTH2,
1268+
oauth2=OAuth2Auth(
1269+
client_id="mock_client_id",
1270+
client_secret="mock_client_secret",
1271+
),
1272+
)
1273+
stored_credential = raw_credential.model_copy(deep=True)
1274+
stored_credential.oauth2.access_token = "expired_token"
1275+
stored_credential.oauth2.refresh_token = "refresh_token"
1276+
stored_credential.oauth2.expires_at = int(time.time()) - 60
1277+
manager = CredentialManager(
1278+
AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=raw_credential)
1279+
)
1280+
manager._load_existing_credential = AsyncMock(
1281+
return_value=stored_credential
1282+
)
1283+
manager._save_credential = AsyncMock()
1284+
context = Mock()
1285+
1286+
with patch.object(
1287+
oauth2_credential_util, "OAuth2Session", autospec=True
1288+
) as mock_oauth2_session:
1289+
mock_client = mock_oauth2_session.return_value
1290+
mock_client.refresh_token.return_value = {
1291+
"access_token": "refreshed_token",
1292+
"refresh_token": "refresh_token",
1293+
"expires_at": int(time.time()) + 3600,
1294+
}
1295+
result = await manager.get_auth_credential(context)
1296+
1297+
assert result.oauth2.access_token == "refreshed_token"
1298+
mock_client.refresh_token.assert_called_once()
1299+
assert (
1300+
mock_client.refresh_token.call_args.kwargs["refresh_token"]
1301+
== "refresh_token"
1302+
)
1303+
mock_client.fetch_token.assert_not_called()
1304+
1305+
@pytest.mark.asyncio
1306+
async def test_get_auth_credential_keeps_stored_token_when_refetch_fails(
1307+
self,
1308+
):
1309+
"""A failed client credentials refetch returns the stored token unsaved."""
1310+
auth_scheme = OAuth2(
1311+
flows=OAuthFlows(
1312+
clientCredentials=OAuthFlowClientCredentials(
1313+
tokenUrl="https://example.com/token",
1314+
scopes={},
1315+
)
1316+
)
1317+
)
1318+
raw_credential = AuthCredential(
1319+
auth_type=AuthCredentialTypes.OAUTH2,
1320+
oauth2=OAuth2Auth(
1321+
client_id="mock_client_id",
1322+
client_secret="mock_client_secret",
1323+
),
1324+
)
1325+
stored_credential = raw_credential.model_copy(deep=True)
1326+
stored_credential.oauth2.access_token = "expired_token"
1327+
stored_credential.oauth2.expires_at = int(time.time()) - 60
1328+
manager = CredentialManager(
1329+
AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=raw_credential)
1330+
)
1331+
manager._load_existing_credential = AsyncMock(
1332+
return_value=stored_credential
1333+
)
1334+
manager._save_credential = AsyncMock()
1335+
context = Mock()
1336+
1337+
with patch.object(
1338+
oauth2_credential_util, "OAuth2Session", autospec=True
1339+
) as mock_oauth2_session:
1340+
mock_client = mock_oauth2_session.return_value
1341+
mock_client.fetch_token.side_effect = OAuth2Error(error="invalid_client")
1342+
mock_client.refresh_token.side_effect = AssertionError(
1343+
"refresh_token() called without a refresh token"
1344+
)
1345+
result = await manager.get_auth_credential(context)
1346+
1347+
assert result is stored_credential
1348+
assert result.oauth2.access_token == "expired_token"
1349+
mock_client.fetch_token.assert_called_once()
1350+
manager._save_credential.assert_not_awaited()
1351+
assert raw_credential.oauth2.access_token is None
1352+
11831353
@pytest.mark.asyncio
11841354
async def test_is_credential_ready_api_key(self):
11851355
"""Test _is_credential_ready with API key credential."""

0 commit comments

Comments
 (0)