diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 1a00d41b..467fdb66 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -66,10 +66,13 @@ class AuthHandler: return exchange_result.credential async def parse_and_store_auth_response(self, state: State) -> None: + credential_key = self.auth_config.credential_key + if not credential_key: + raise ValueError("credential_key is empty.") - credential_key = "temp:" + self.auth_config.credential_key + temp_credential_key = "temp:" + credential_key - state[credential_key] = self.auth_config.exchanged_auth_credential + state[temp_credential_key] = self.auth_config.exchanged_auth_credential if not isinstance( self.auth_config.auth_scheme, SecurityBase ) or self.auth_config.auth_scheme.type_ not in ( @@ -78,15 +81,78 @@ class AuthHandler: ): return - state[credential_key] = await self.exchange_auth_token() + state[temp_credential_key] = await self.exchange_auth_token() def _validate(self) -> None: if not self.auth_config.auth_scheme: raise ValueError("auth_scheme is empty.") - def get_auth_response(self, state: State) -> AuthCredential: - credential_key = "temp:" + self.auth_config.credential_key - return state.get(credential_key, None) + def get_auth_response(self, state: State) -> AuthCredential | None: + # 1. Try reading the temp credential key (standard ADK flow) + credential_key = self.auth_config.credential_key + if not credential_key: + return None + + temp_credential_key = "temp:" + credential_key + val = state.get(temp_credential_key, None) + if val is not None: + if isinstance(val, AuthCredential): + return val + if isinstance(val, dict): + return AuthCredential.model_validate(val) + if isinstance(val, str) and val: + return self._build_credential_from_string(val) + + # 2. Try reading the credential key without the 'temp:' prefix + val = state.get(credential_key, None) + if val is not None: + if isinstance(val, AuthCredential): + return val + if isinstance(val, dict): + return AuthCredential.model_validate(val) + if isinstance(val, str) and val: + return self._build_credential_from_string(val) + + return None + + def _build_credential_from_string(self, val: str) -> AuthCredential: + from .auth_credential import AuthCredentialTypes + from .auth_credential import HttpAuth + from .auth_credential import HttpCredentials + from .auth_credential import OAuth2Auth + + auth_scheme = self.auth_config.auth_scheme + if not auth_scheme: + return AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token=val), + ) + + scheme_type = auth_scheme.type_ + if scheme_type == AuthSchemeType.apiKey: + return AuthCredential( + auth_type=AuthCredentialTypes.API_KEY, + api_key=val, + ) + elif scheme_type == AuthSchemeType.http: + scheme = getattr(auth_scheme, "scheme", "bearer") + return AuthCredential( + auth_type=AuthCredentialTypes.HTTP, + http=HttpAuth( + scheme=scheme, + credentials=HttpCredentials(token=val), + ), + ) + elif scheme_type in (AuthSchemeType.oauth2, AuthSchemeType.openIdConnect): + return AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token=val), + ) + else: + return AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token=val), + ) def generate_auth_request(self) -> AuthConfig: if not isinstance( diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index f63c6170..2217821f 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -607,6 +607,112 @@ class TestGetAuthResponse: result = handler.get_auth_response(state) assert result is None + def test_get_auth_response_temp_prefix_str_token(self, auth_config): + """Test retrieving a string token stored under temp prefix in state.""" + handler = AuthHandler(auth_config) + state = MockState() + credential_key = auth_config.credential_key + state["temp:" + credential_key] = "ya29.mock_token" + + result = handler.get_auth_response(state) + + assert result is not None + assert result.auth_type == AuthCredentialTypes.OAUTH2 + assert result.oauth2.access_token == "ya29.mock_token" + + def test_get_auth_response_no_prefix_credential( + self, auth_config, oauth2_credentials_with_auth_uri + ): + """Test retrieving a credential stored under the key without prefix.""" + handler = AuthHandler(auth_config) + state = MockState() + credential_key = auth_config.credential_key + state[credential_key] = oauth2_credentials_with_auth_uri + + result = handler.get_auth_response(state) + + assert result == oauth2_credentials_with_auth_uri + + def test_get_auth_response_no_prefix_str_token(self, auth_config): + """Test retrieving a string token stored under the key without prefix.""" + handler = AuthHandler(auth_config) + state = MockState() + credential_key = auth_config.credential_key + state[credential_key] = "ya29.mock_token_no_prefix" + + result = handler.get_auth_response(state) + + assert result is not None + assert result.auth_type == AuthCredentialTypes.OAUTH2 + assert result.oauth2.access_token == "ya29.mock_token_no_prefix" + + def test_get_auth_response_temp_prefix_dict(self, auth_config): + """Test retrieving a credential dictionary stored under temp prefix.""" + handler = AuthHandler(auth_config) + state = MockState() + credential_key = auth_config.credential_key + # Store dict in state representing an AuthCredential + state["temp:" + credential_key] = { + "auth_type": "oauth2", + "oauth2": {"access_token": "ya29.mock_token_from_dict"}, + } + + result = handler.get_auth_response(state) + + assert result is not None + assert result.auth_type == AuthCredentialTypes.OAUTH2 + assert result.oauth2.access_token == "ya29.mock_token_from_dict" + + def test_get_auth_response_no_prefix_dict(self, auth_config): + """Test retrieving a credential dictionary stored under the key without prefix.""" + handler = AuthHandler(auth_config) + state = MockState() + credential_key = auth_config.credential_key + state[credential_key] = { + "auth_type": "oauth2", + "oauth2": {"access_token": "ya29.mock_token_from_dict_no_prefix"}, + } + + result = handler.get_auth_response(state) + + assert result is not None + assert result.auth_type == AuthCredentialTypes.OAUTH2 + assert result.oauth2.access_token == "ya29.mock_token_from_dict_no_prefix" + + def test_get_auth_response_api_key_str(self): + """Test retrieving a string token under apiKey scheme wraps it as APIKey.""" + auth_scheme = APIKey(**{"name": "X-API-Key", "in": APIKeyIn.header}) + config = AuthConfig(auth_scheme=auth_scheme) + handler = AuthHandler(config) + state = MockState() + credential_key = config.credential_key + state["temp:" + credential_key] = "my_api_key_value" + + result = handler.get_auth_response(state) + + assert result is not None + assert result.auth_type == AuthCredentialTypes.API_KEY + assert result.api_key == "my_api_key_value" + + def test_get_auth_response_http_str(self): + """Test retrieving a string token under http bearer scheme wraps it as HTTP Bearer.""" + from fastapi.openapi.models import HTTPBearer + + auth_scheme = HTTPBearer() + config = AuthConfig(auth_scheme=auth_scheme) + handler = AuthHandler(config) + state = MockState() + credential_key = config.credential_key + state["temp:" + credential_key] = "my_http_bearer_token" + + result = handler.get_auth_response(state) + + assert result is not None + assert result.auth_type == AuthCredentialTypes.HTTP + assert result.http is not None + assert result.http.scheme == "bearer" + assert result.http.credentials.token == "my_http_bearer_token" + class TestParseAndStoreAuthResponse: """Tests for the parse_and_store_auth_response method.""" @@ -650,6 +756,19 @@ class TestParseAndStoreAuthResponse: assert state["temp:" + credential_key] == mock_exchange_token.return_value assert mock_exchange_token.called + @pytest.mark.asyncio + async def test_empty_credential_key_raises_error(self, oauth2_auth_scheme): + """Test that ValueError is raised when credential_key is empty.""" + config = AuthConfig( + auth_scheme=oauth2_auth_scheme, + ) + config.credential_key = "" # Bypass init logic that sets it + handler = AuthHandler(config) + state = MockState() + + with pytest.raises(ValueError, match="credential_key is empty."): + await handler.parse_and_store_auth_response(state) + class TestExchangeAuthToken: """Tests for the exchange_auth_token method."""