From 177980106b2f7be9a8c0a02f395ff0f85faa0c5a Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Thu, 12 Jun 2025 20:07:06 -0700 Subject: [PATCH] feat: Support refresh access token automatically for rest_api_tool 1. let auth_handler.py to utilize the oauth2 credential fetcher to exchange token 2. restructure tool_auth_handler.py to support refresh token PiperOrigin-RevId: 770901469 --- src/google/adk/auth/auth_handler.py | 82 +++--------------- .../openapi_spec_parser/tool_auth_handler.py | 83 ++++++++++++------- tests/unittests/auth/test_auth_handler.py | 7 +- .../test_tool_auth_handler.py | 69 ++++++++++++++- 4 files changed, 137 insertions(+), 104 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 1607dcad..5f80ee3f 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -16,16 +16,13 @@ from __future__ import annotations from typing import TYPE_CHECKING -from fastapi.openapi.models import OAuth2 from fastapi.openapi.models import SecurityBase from .auth_credential import AuthCredential -from .auth_credential import AuthCredentialTypes -from .auth_credential import OAuth2Auth from .auth_schemes import AuthSchemeType -from .auth_schemes import OAuthGrantType from .auth_schemes import OpenIdConnectWithConfig from .auth_tool import AuthConfig +from .oauth2_credential_fetcher import OAuth2CredentialFetcher if TYPE_CHECKING: from ..sessions.state import State @@ -33,9 +30,9 @@ if TYPE_CHECKING: try: from authlib.integrations.requests_client import OAuth2Session - SUPPORT_TOKEN_EXCHANGE = True + AUTHLIB_AVIALABLE = True except ImportError: - SUPPORT_TOKEN_EXCHANGE = False + AUTHLIB_AVIALABLE = False class AuthHandler: @@ -46,69 +43,9 @@ class AuthHandler: def exchange_auth_token( self, ) -> AuthCredential: - """Generates an auth token from the authorization response. - - Returns: - An AuthCredential object containing the access token. - - Raises: - ValueError: If the token endpoint is not configured in the auth - scheme. - AuthCredentialMissingError: If the access token cannot be retrieved - from the token endpoint. - """ - auth_scheme = self.auth_config.auth_scheme - auth_credential = self.auth_config.exchanged_auth_credential - if not SUPPORT_TOKEN_EXCHANGE: - return auth_credential - if isinstance(auth_scheme, OpenIdConnectWithConfig): - if not hasattr(auth_scheme, "token_endpoint"): - return self.auth_config.exchanged_auth_credential - token_endpoint = auth_scheme.token_endpoint - scopes = auth_scheme.scopes - elif isinstance(auth_scheme, OAuth2): - if ( - not auth_scheme.flows.authorizationCode - or not auth_scheme.flows.authorizationCode.tokenUrl - ): - return self.auth_config.exchanged_auth_credential - token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl - scopes = list(auth_scheme.flows.authorizationCode.scopes.keys()) - else: - return self.auth_config.exchanged_auth_credential - - if ( - not auth_credential - or not auth_credential.oauth2 - or not auth_credential.oauth2.client_id - or not auth_credential.oauth2.client_secret - or auth_credential.oauth2.access_token - or auth_credential.oauth2.refresh_token - ): - return self.auth_config.exchanged_auth_credential - - client = OAuth2Session( - auth_credential.oauth2.client_id, - auth_credential.oauth2.client_secret, - scope=" ".join(scopes), - redirect_uri=auth_credential.oauth2.redirect_uri, - state=auth_credential.oauth2.state, - ) - tokens = client.fetch_token( - token_endpoint, - authorization_response=auth_credential.oauth2.auth_response_uri, - code=auth_credential.oauth2.auth_code, - grant_type=OAuthGrantType.AUTHORIZATION_CODE, - ) - - updated_credential = AuthCredential( - auth_type=AuthCredentialTypes.OAUTH2, - oauth2=OAuth2Auth( - access_token=tokens.get("access_token"), - refresh_token=tokens.get("refresh_token"), - ), - ) - return updated_credential + return OAuth2CredentialFetcher( + self.auth_config.auth_scheme, self.auth_config.exchanged_auth_credential + ).exchange() def parse_and_store_auth_response(self, state: State) -> None: @@ -204,6 +141,13 @@ class AuthHandler: ValueError: If the authorization endpoint is not configured in the auth scheme. """ + if not AUTHLIB_AVIALABLE: + return ( + self.auth_config.raw_auth_credential.model_copy(deep=True) + if self.auth_config.raw_auth_credential + else None + ) + auth_scheme = self.auth_config.auth_scheme auth_credential = self.auth_config.raw_auth_credential 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 71e760e3..c36793fd 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 @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import logging from typing import Literal from typing import Optional -from fastapi.encoders import jsonable_encoder from pydantic import BaseModel from ....auth.auth_credential import AuthCredential @@ -25,6 +25,7 @@ from ....auth.auth_credential import AuthCredentialTypes from ....auth.auth_schemes import AuthScheme from ....auth.auth_schemes import AuthSchemeType from ....auth.auth_tool import AuthConfig +from ....auth.oauth2_credential_fetcher import OAuth2CredentialFetcher from ...tool_context import ToolContext from ..auth.credential_exchangers.auto_auth_credential_exchanger import AutoAuthCredentialExchanger from ..auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError @@ -95,10 +96,9 @@ class ToolContextCredentialStore: auth_credential: Optional[AuthCredential], ): if self.tool_context: - serializable_credential = jsonable_encoder( - auth_credential, exclude_none=True + self.tool_context.state[key] = auth_credential.model_dump( + exclude_none=True ) - self.tool_context.state[key] = serializable_credential def remove_credential(self, key: str): del self.tool_context.state[key] @@ -146,20 +146,20 @@ class ToolAuthHandler: credential_store, ) - def _handle_existing_credential( + def _get_existing_credential( self, - ) -> Optional[AuthPreparationResult]: + ) -> Optional[AuthCredential]: """Checks for and returns an existing, exchanged credential.""" if self.credential_store: existing_credential = self.credential_store.get_credential( self.auth_scheme, self.auth_credential ) if existing_credential: - return AuthPreparationResult( - state="done", - auth_scheme=self.auth_scheme, - auth_credential=existing_credential, - ) + if existing_credential.oauth2: + existing_credential = OAuth2CredentialFetcher( + self.auth_scheme, existing_credential + ).refresh() + return existing_credential return None def _exchange_credential( @@ -223,6 +223,17 @@ class ToolAuthHandler: ) ) + def _external_exchange_required(self, credential) -> bool: + return ( + credential.auth_type + in ( + AuthCredentialTypes.OAUTH2, + AuthCredentialTypes.OPEN_ID_CONNECT, + ) + and not credential.oauth2.access_token + and not credential.google_oauth2_json + ) + def prepare_auth_credentials( self, ) -> AuthPreparationResult: @@ -233,31 +244,41 @@ class ToolAuthHandler: return AuthPreparationResult(state="done") # Check for existing credential. - existing_result = self._handle_existing_credential() - if existing_result: - return existing_result + existing_credential = self._get_existing_credential() + credential = existing_credential or self.auth_credential # fetch credential from adk framework # Some auth scheme like OAuth2 AuthCode & OpenIDConnect may require # multi-step exchange: # client_id , client_secret -> auth_uri -> auth_code -> access_token - # -> bearer token # adk framework supports exchange access_token already - fetched_credential = self._get_auth_response() or self.auth_credential + # for other credential, adk can also get back the credential directly + if not credential or self._external_exchange_required(credential): + credential = self._get_auth_response() + # store fetched credential + if credential: + self._store_credential(credential) + else: + self._request_credential() + return AuthPreparationResult( + state="pending", + auth_scheme=self.auth_scheme, + auth_credential=self.auth_credential, + ) - exchanged_credential = self._exchange_credential(fetched_credential) + # here exchangers are doing two different thing: + # for service account the exchanger is doing actualy token exchange + # while for oauth2 it's actually doing the credentail conversion + # from OAuth2 credential to HTTP credentails for setting credential in + # http header + # TODO cleanup the logic: + # 1. service account token exchanger should happen before we store them in + # the token store + # 2. blow line should only do credential conversion - if exchanged_credential: - self._store_credential(exchanged_credential) - return AuthPreparationResult( - state="done", - auth_scheme=self.auth_scheme, - auth_credential=exchanged_credential, - ) - else: - self._request_credential() - return AuthPreparationResult( - state="pending", - auth_scheme=self.auth_scheme, - auth_credential=self.auth_credential, - ) + exchanged_credential = self._exchange_credential(credential) + return AuthPreparationResult( + state="done", + auth_scheme=self.auth_scheme, + auth_credential=exchanged_credential, + ) diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index c717682a..402e4d0c 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -449,7 +449,7 @@ class TestExchangeAuthToken: ): """Test when token exchange is not supported.""" monkeypatch.setattr( - "google.adk.auth.auth_handler.SUPPORT_TOKEN_EXCHANGE", False + "google.adk.auth.oauth2_credential_fetcher.AUTHLIB_AVIALABLE", False ) handler = AuthHandler(auth_config_with_auth_code) @@ -537,7 +537,10 @@ class TestExchangeAuthToken: assert result == oauth2_credentials_with_token - @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session) + @patch( + "google.adk.auth.oauth2_credential_fetcher.OAuth2Session", + MockOAuth2Session, + ) def test_successful_token_exchange(self, auth_config_with_auth_code): """Test a successful token exchange.""" handler = AuthHandler(auth_config_with_auth_code) 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 0a3b8ccc..8db151fc 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 @@ -14,6 +14,7 @@ from typing import Optional from unittest.mock import MagicMock +from unittest.mock import patch from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.llm_agent import LlmAgent @@ -147,10 +148,11 @@ def test_openid_connect_with_auth_response( tool_context = create_mock_tool_context() mock_auth_handler = MagicMock() - mock_auth_handler.get_auth_response.return_value = AuthCredential( + returned_credentail = AuthCredential( auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, oauth2=OAuth2Auth(auth_response_uri='test_auth_response_uri'), ) + mock_auth_handler.get_auth_response.return_value = returned_credentail mock_auth_handler_path = 'google.adk.tools.tool_context.AuthHandler' monkeypatch.setattr( mock_auth_handler_path, lambda *args, **kwargs: mock_auth_handler @@ -172,7 +174,7 @@ def test_openid_connect_with_auth_response( stored_credential = credential_store.get_credential( openid_connect_scheme, openid_connect_credential ) - assert stored_credential == result.auth_credential + assert stored_credential == returned_credentail mock_auth_handler.get_auth_response.assert_called_once() @@ -199,3 +201,66 @@ def test_openid_connect_existing_token( result = handler.prepare_auth_credentials() assert result.state == 'done' assert result.auth_credential == existing_credential + + +@patch( + 'google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler.OAuth2CredentialFetcher' +) +def test_openid_connect_existing_oauth2_token_refresh( + mock_oauth2_fetcher, openid_connect_scheme, openid_connect_credential +): + """Test that OAuth2 tokens are refreshed when existing credentials are found.""" + # Create existing OAuth2 credential + existing_credential = AuthCredential( + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, + oauth2=OAuth2Auth( + client_id='test_client_id', + client_secret='test_client_secret', + access_token='existing_token', + refresh_token='refresh_token', + ), + ) + + # Mock the refreshed credential + refreshed_credential = AuthCredential( + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, + oauth2=OAuth2Auth( + client_id='test_client_id', + client_secret='test_client_secret', + access_token='refreshed_token', + refresh_token='new_refresh_token', + ), + ) + + # Setup mock OAuth2CredentialFetcher + mock_fetcher_instance = MagicMock() + mock_fetcher_instance.refresh.return_value = refreshed_credential + mock_oauth2_fetcher.return_value = mock_fetcher_instance + + tool_context = create_mock_tool_context() + credential_store = ToolContextCredentialStore(tool_context=tool_context) + + # Store the existing 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, + ) + + result = handler.prepare_auth_credentials() + + # Verify OAuth2CredentialFetcher was called for refresh + mock_oauth2_fetcher.assert_called_once_with( + openid_connect_scheme, existing_credential + ) + mock_fetcher_instance.refresh.assert_called_once() + + assert result.state == 'done' + # The result should contain the refreshed credential after exchange + assert result.auth_credential is not None