From 218ea76e30ced48898a46ca48a014f7dffd266a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= Date: Mon, 11 May 2026 15:50:45 -0700 Subject: [PATCH] fix(auth): persist refreshed OAuth2 credentials to store Merge https://github.com/google/adk-python/pull/5350 Fixes #5329 Co-authored-by: Xuan Yang COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5350 from voidborne-d:fix/persist-refreshed-oauth2-credential 0ef3d8c19616b45f932b15ee2a200bc74c8334b3 PiperOrigin-RevId: 913922074 --- .../openapi_spec_parser/tool_auth_handler.py | 6 ++ .../test_tool_auth_handler.py | 67 ++++++++++++++++++- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py index 21337cfa..0d78a575 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py @@ -242,6 +242,12 @@ class ToolAuthHandler: existing_credential = await refresher.refresh( existing_credential, self.auth_scheme ) + # Persist the refreshed credential so the next invocation + # reads the new tokens instead of the stale pre-refresh ones. + # Without this, providers that rotate refresh_tokens on each + # refresh (e.g. Salesforce, many OIDC providers) will fail + # because the old refresh_token has already been invalidated. + self._store_credential(existing_credential) return existing_credential return None diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py index a6babce6..d32fc132 100644 --- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py +++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py @@ -13,6 +13,7 @@ # limitations under the License. from typing import Optional +from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import patch @@ -29,6 +30,7 @@ from google.adk.sessions.session import Session from google.adk.tools.openapi_tool.auth.auth_helpers import openid_dict_to_scheme_credential from google.adk.tools.openapi_tool.auth.auth_helpers import token_to_scheme_credential from google.adk.tools.openapi_tool.auth.credential_exchangers.auto_auth_credential_exchanger import OAuth2CredentialExchanger +from google.adk.tools.openapi_tool.openapi_spec_parser import tool_auth_handler from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolContextCredentialStore from google.adk.tools.tool_context import ToolContext @@ -223,9 +225,7 @@ async def test_openid_connect_existing_token( assert result.auth_credential == existing_credential -@patch( - 'google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler.OAuth2CredentialRefresher' -) +@patch.object(tool_auth_handler, 'OAuth2CredentialRefresher') @pytest.mark.asyncio async def test_openid_connect_existing_oauth2_token_refresh( mock_oauth2_refresher, openid_connect_scheme, openid_connect_credential @@ -292,3 +292,64 @@ async def test_openid_connect_existing_oauth2_token_refresh( assert result.state == 'done' # The result should contain the refreshed credential after exchange assert result.auth_credential is not None + + +@patch.object(tool_auth_handler, 'OAuth2CredentialRefresher') +@pytest.mark.asyncio +async def test_refreshed_credential_is_persisted_to_store( + mock_oauth2_refresher, openid_connect_scheme, openid_connect_credential +): + """Test that refreshed OAuth2 credentials are persisted back to the store.""" + # Create existing OAuth2 credential with an "old" refresh token. + existing_credential = AuthCredential( + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, + oauth2=OAuth2Auth( + client_id='test_client_id', + client_secret='test_client_secret', + access_token='old_access_token', + refresh_token='old_refresh_token', + ), + ) + + # The refresher will return a credential with rotated tokens. + refreshed_credential = AuthCredential( + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, + oauth2=OAuth2Auth( + client_id='test_client_id', + client_secret='test_client_secret', + access_token='new_access_token', + refresh_token='new_refresh_token', + ), + ) + + mock_refresher_instance = MagicMock() + mock_refresher_instance.is_refresh_needed = AsyncMock(return_value=True) + mock_refresher_instance.refresh = AsyncMock(return_value=refreshed_credential) + mock_oauth2_refresher.return_value = mock_refresher_instance + + tool_context = create_mock_tool_context() + credential_store = ToolContextCredentialStore(tool_context=tool_context) + + # Store the existing (stale) credential. + key = credential_store.get_credential_key( + openid_connect_scheme, openid_connect_credential + ) + credential_store.store_credential(key, existing_credential) + + handler = ToolAuthHandler( + tool_context, + openid_connect_scheme, + openid_connect_credential, + credential_store=credential_store, + ) + + await handler.prepare_auth_credentials() + + # The critical assertion: the *refreshed* credential must now be in the + # store so that the next invocation reads the new tokens, not the old ones. + persisted = credential_store.get_credential( + openid_connect_scheme, openid_connect_credential + ) + assert persisted is not None + assert persisted.oauth2.access_token == 'new_access_token' + assert persisted.oauth2.refresh_token == 'new_refresh_token'