From 006d547ffb5bea79875657c1a16545a4b3b5fd8c Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Fri, 5 Dec 2025 23:46:20 +0100 Subject: [PATCH 01/12] fix: solve #3845 --- src/google/adk/auth/auth_handler.py | 46 ++++++++--- src/google/adk/auth/auth_tool.py | 6 ++ src/google/adk/auth/credential_manager.py | 76 ++++++++++++++++++- src/google/adk/auth/oauth2_credential_util.py | 2 + .../unittests/auth/test_credential_manager.py | 75 ++++++++++-------- 5 files changed, 164 insertions(+), 41 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 07515ab2e8..c8ae77c9a6 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -7,11 +7,6 @@ # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - from __future__ import annotations from typing import TYPE_CHECKING @@ -48,9 +43,30 @@ async def exchange_auth_token( self, ) -> AuthCredential: exchanger = OAuth2CredentialExchanger() - return await exchanger.exchange( - self.auth_config.exchanged_auth_credential, self.auth_config.auth_scheme - ) + + # Restore secret if needed + credential = self.auth_config.exchanged_auth_credential + redacted = False + original_secret = None + + if credential and credential.oauth2 and credential.oauth2.client_id: + # Check if secret needs restoration + from .credential_manager import CredentialManager + secret = CredentialManager.get_client_secret(credential.oauth2.client_id) + if secret and credential.oauth2.client_secret == "": + original_secret = credential.oauth2.client_secret + credential.oauth2.client_secret = secret + redacted = True + + try: + res = await exchanger.exchange( + credential, self.auth_config.auth_scheme + ) + return res + finally: + # Always re-redact if we restored it + if redacted and credential and credential.oauth2: + credential.oauth2.client_secret = "" async def parse_and_store_auth_response(self, state: State) -> None: @@ -182,9 +198,19 @@ def generate_auth_uri( ) scopes = list(scopes.keys()) + client_id = auth_credential.oauth2.client_id + client_secret = auth_credential.oauth2.client_secret + + # Check if secret is redacted and restore it from manager + if client_secret == "" and client_id: + from .credential_manager import CredentialManager + secret = CredentialManager.get_client_secret(client_id) + if secret: + client_secret = secret + client = OAuth2Session( - auth_credential.oauth2.client_id, - auth_credential.oauth2.client_secret, + client_id, + client_secret, scope=" ".join(scopes), redirect_uri=auth_credential.oauth2.redirect_uri, ) diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index 0316e5258e..0ab605c5a7 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -81,6 +81,12 @@ def get_credential_key(self): if auth_credential and auth_credential.model_extra: auth_credential = auth_credential.model_copy(deep=True) auth_credential.model_extra.clear() + + # Normalize secret to ensure stable key regardless of redaction + if auth_credential and auth_credential.oauth2: + auth_credential = auth_credential.model_copy(deep=True) + auth_credential.oauth2.client_secret = None + credential_name = ( f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}" if auth_credential diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index c022ab694c..b34ad6ef97 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -75,11 +75,30 @@ class CredentialManager: ``` """ + # A map to store client secrets in memory. Key is client_id, value is client_secret + _CLIENT_SECRETS: dict[str, str] = {} + def __init__( self, auth_config: AuthConfig, ): - self._auth_config = auth_config + # We deep copy the auth_config to avoid modifying the original one passed by user + # and to ensure we can safely redact sensitive information + # However we cannot rely on copy.deepcopy because AuthConfig is a pydantic model + # and deepcopy on pydantic model is not always reliable? No, deepcopy works. + # But better to use model_copy if possible. AuthConfig inherits BaseModelWithConfig which is Pydantic. + # auth_config.model_copy(deep=True) is available in Pydantic V2. + # For safe side, we use the passed instance but we redact sensitive info immediately. + # Wait, modifying passed instance is bad practice if user reuses it. + # But CredentialManager usually takes ownership? + # Let's perform redaction on `self._auth_config` which we assign. + # And we should clone it first. + self._auth_config = auth_config.model_copy(deep=True) + + # Secure the client secret + self._secure_client_secret(self._auth_config.raw_auth_credential) + self._secure_client_secret(self._auth_config.exchanged_auth_credential) + self._exchanger_registry = CredentialExchangerRegistry() self._refresher_registry = CredentialRefresherRegistry() self._discovery_manager = OAuth2DiscoveryManager() @@ -110,6 +129,31 @@ def __init__( AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_refresher ) + def _secure_client_secret(self, credential: Optional[AuthCredential]): + """Extracts client secret to memory and redacts it from the credential.""" + if ( + credential + and credential.oauth2 + and credential.oauth2.client_id + and credential.oauth2.client_secret + and credential.oauth2.client_secret != "" + ): + logger.info(f"Securing client secret for client_id: {credential.oauth2.client_id}") + # Store in memory map + self._CLIENT_SECRETS[credential.oauth2.client_id] = ( + credential.oauth2.client_secret + ) + # Redact from config + credential.oauth2.client_secret = "" + else: + if credential and credential.oauth2: + logger.debug(f"Not securing secret for client_id {credential.oauth2.client_id}: secret is {credential.oauth2.client_secret}") + + @staticmethod + def get_client_secret(client_id: str) -> Optional[str]: + """Retrieves the client secret for a given client_id.""" + return CredentialManager._CLIENT_SECRETS.get(client_id) + def register_credential_exchanger( self, credential_type: AuthCredentialTypes, @@ -124,6 +168,9 @@ def register_credential_exchanger( self._exchanger_registry.register(credential_type, exchanger_instance) async def request_credential(self, callback_context: CallbackContext) -> None: + # We send the auth_config (which is already redacted in __init__) to the client + # Note: we need to ensure we don't send any stale exchanged credentials if they are not valid + # But usually CredentialManager manages that. callback_context.request_credential(self._auth_config) async def get_auth_credential( @@ -218,10 +265,37 @@ async def _exchange_credential( self._auth_config.auth_scheme, credential ) else: + # Restore client secret from memory map for exchange + restored = False + if ( + credential.oauth2 + and credential.oauth2.client_id + and credential.oauth2.client_id in self._CLIENT_SECRETS + ): + credential.oauth2.client_secret = self._CLIENT_SECRETS[ + credential.oauth2.client_id + ] + restored = True + elif ( + self._auth_config.raw_auth_credential + and self._auth_config.raw_auth_credential.oauth2 + and self._auth_config.raw_auth_credential.oauth2.client_id + in self._CLIENT_SECRETS + ): + # Fallback to look up using raw credential client id if credential client id is missing (unlikely for valid flow) + credential.oauth2.client_secret = self._CLIENT_SECRETS[ + self._auth_config.raw_auth_credential.oauth2.client_id + ] + restored = True + exchanged_credential = await exchanger.exchange( credential, self._auth_config.auth_scheme ) + # Redact client secret again after exchange to prevent leakage + if exchanged_credential.oauth2: + exchanged_credential.oauth2.client_secret = "" + return exchanged_credential, True async def _refresh_credential( diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py index d90103a88c..954f000fab 100644 --- a/src/google/adk/auth/oauth2_credential_util.py +++ b/src/google/adk/auth/oauth2_credential_util.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging +import sys from typing import Optional from typing import Tuple @@ -107,6 +108,7 @@ def update_credential_with_tokens( tokens: The OAuth2Token object containing new token information. """ auth_credential.oauth2.access_token = tokens.get("access_token") + sys.stderr.write(f"[DEBUG] Assigned access_token: {auth_credential.oauth2.access_token}\n") auth_credential.oauth2.refresh_token = tokens.get("refresh_token") auth_credential.oauth2.expires_at = ( int(tokens.get("expires_at")) if tokens.get("expires_at") else None diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index ab021d1eaa..92874c6cb7 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -36,19 +36,30 @@ import pytest + +def create_auth_config_mock(): + """Creates a mock AuthConfig that returns itself on model_copy.""" + # We remove spec=AuthConfig because accessing Pydantic fields on a spec-ed mock + # can fail if they are not seen as class attributes or if we need dynamic attributes. + m = Mock() + m.spec = AuthConfig # Optional: if we want isinstance to work, but Mock(spec=X) enforces attributes. + # Let's just use a plain Mock and configure what we need. + m.model_copy.side_effect = lambda **kwargs: m + return m + class TestCredentialManager: """Test suite for CredentialManager.""" def test_init(self): """Test CredentialManager initialization.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() manager = CredentialManager(auth_config) assert manager._auth_config == auth_config @pytest.mark.asyncio async def test_request_credential(self): """Test request_credential method.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() callback_context = Mock() callback_context.request_credential = Mock() @@ -61,7 +72,7 @@ async def test_request_credential(self): async def test_load_auth_credentials_success(self): """Test load_auth_credential with successful flow.""" # Create mocks - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None @@ -104,7 +115,7 @@ async def test_load_auth_credentials_success(self): @pytest.mark.asyncio async def test_load_auth_credentials_no_credential(self): """Test load_auth_credential when no credential is available.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None # Add auth_scheme for the _is_client_credentials_flow method @@ -134,8 +145,9 @@ async def test_load_auth_credentials_no_credential(self): @pytest.mark.asyncio async def test_load_existing_credential_already_exchanged(self): """Test _load_existing_credential when credential is already exchanged.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() mock_credential = Mock(spec=AuthCredential) + mock_credential.oauth2 = Mock() auth_config.exchanged_auth_credential = mock_credential callback_context = Mock() @@ -150,7 +162,7 @@ async def test_load_existing_credential_already_exchanged(self): @pytest.mark.asyncio async def test_load_existing_credential_with_credential_service(self): """Test _load_existing_credential with credential service.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.exchanged_auth_credential = None mock_credential = Mock(spec=AuthCredential) @@ -172,7 +184,7 @@ async def test_load_existing_credential_with_credential_service(self): @pytest.mark.asyncio async def test_load_from_credential_service_with_service(self): """Test _load_from_credential_service from callback context when credential service is available.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() mock_credential = Mock(spec=AuthCredential) @@ -196,7 +208,7 @@ async def test_load_from_credential_service_with_service(self): @pytest.mark.asyncio async def test_load_from_credential_service_no_service(self): """Test _load_from_credential_service when no credential service is available.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() # Mock invocation context with no credential service invocation_context = Mock() @@ -213,7 +225,7 @@ async def test_load_from_credential_service_no_service(self): @pytest.mark.asyncio async def test_save_credential_with_service(self): """Test _save_credential with credential service.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() mock_credential = Mock(spec=AuthCredential) # Mock credential service @@ -236,7 +248,7 @@ async def test_save_credential_with_service(self): @pytest.mark.asyncio async def test_save_credential_no_service(self): """Test _save_credential when no credential service is available.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.exchanged_auth_credential = None mock_credential = Mock(spec=AuthCredential) @@ -260,9 +272,10 @@ async def test_refresh_credential_oauth2(self): mock_oauth2_auth = Mock(spec=OAuth2Auth) mock_credential = Mock(spec=AuthCredential) + mock_credential.oauth2 = Mock() mock_credential.auth_type = AuthCredentialTypes.OAUTH2 - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = Mock() # Mock refresher @@ -297,7 +310,7 @@ async def test_refresh_credential_no_refresher(self): mock_credential = Mock(spec=AuthCredential) mock_credential.auth_type = AuthCredentialTypes.API_KEY - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() manager = CredentialManager(auth_config) @@ -316,9 +329,10 @@ async def test_refresh_credential_no_refresher(self): async def test_is_credential_ready_api_key(self): """Test _is_credential_ready with API key credential.""" mock_raw_credential = Mock(spec=AuthCredential) + mock_raw_credential.oauth2 = Mock() mock_raw_credential.auth_type = AuthCredentialTypes.API_KEY - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = mock_raw_credential manager = CredentialManager(auth_config) @@ -330,9 +344,10 @@ async def test_is_credential_ready_api_key(self): async def test_is_credential_ready_oauth2(self): """Test _is_credential_ready with OAuth2 credential (needs processing).""" mock_raw_credential = Mock(spec=AuthCredential) + mock_raw_credential.oauth2 = Mock() mock_raw_credential.auth_type = AuthCredentialTypes.OAUTH2 - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = mock_raw_credential manager = CredentialManager(auth_config) @@ -346,7 +361,7 @@ async def test_validate_credential_no_raw_credential_oauth2(self): auth_scheme = Mock() auth_scheme.type_ = AuthSchemeType.oauth2 - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = None auth_config.auth_scheme = auth_scheme @@ -361,7 +376,7 @@ async def test_validate_credential_no_raw_credential_openid(self): auth_scheme = Mock() auth_scheme.type_ = AuthSchemeType.openIdConnect - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = None auth_config.auth_scheme = auth_scheme @@ -376,7 +391,7 @@ async def test_validate_credential_no_raw_credential_other_scheme(self): auth_scheme = Mock() auth_scheme.type_ = AuthSchemeType.apiKey - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = None auth_config.auth_scheme = auth_scheme @@ -392,7 +407,7 @@ async def test_validate_credential_oauth2_missing_oauth2_field(self): mock_raw_credential.auth_type = AuthCredentialTypes.OAUTH2 mock_raw_credential.oauth2 = None - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = mock_raw_credential auth_config.auth_scheme = Mock() @@ -407,10 +422,10 @@ async def test_validate_credential_oauth2_missing_scheme_info( ): """Test _validate_credential with OAuth2 missing scheme info.""" mock_raw_credential = Mock(spec=AuthCredential) + mock_raw_credential.oauth2 = Mock() mock_raw_credential.auth_type = AuthCredentialTypes.OAUTH2 - mock_raw_credential.oauth2 = Mock(spec=OAuth2Auth) - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.raw_auth_credential = mock_raw_credential auth_config.auth_scheme = extended_oauth2_scheme @@ -428,7 +443,7 @@ async def test_exchange_credentials_service_account( self, service_account_credential, oauth2_auth_scheme ): """Test _exchange_credential with service account credential.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = oauth2_auth_scheme exchanged_credential = Mock(spec=AuthCredential) @@ -457,7 +472,7 @@ async def test_exchange_credential_no_exchanger(self): mock_credential = Mock(spec=AuthCredential) mock_credential.auth_type = AuthCredentialTypes.API_KEY - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() manager = CredentialManager(auth_config) @@ -513,7 +528,7 @@ async def test_populate_auth_scheme_success( self, auth_server_metadata, extended_oauth2_scheme ): """Test _populate_auth_scheme successfully populates missing info.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = extended_oauth2_scheme manager = CredentialManager(auth_config) @@ -536,7 +551,7 @@ async def test_populate_auth_scheme_success( @pytest.mark.asyncio async def test_populate_auth_scheme_fail(self, extended_oauth2_scheme): """Test _populate_auth_scheme when auto-discovery fails.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = extended_oauth2_scheme manager = CredentialManager(auth_config) @@ -555,7 +570,7 @@ async def test_populate_auth_scheme_fail(self, extended_oauth2_scheme): @pytest.mark.asyncio async def test_populate_auth_scheme_noop(self, implicit_oauth2_scheme): """Test _populate_auth_scheme when auth scheme info not missing.""" - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = implicit_oauth2_scheme manager = CredentialManager(auth_config) @@ -578,7 +593,7 @@ def test_is_client_credentials_flow_oauth2_with_client_credentials(self): ) ) - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = auth_scheme auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None @@ -603,7 +618,7 @@ def test_is_client_credentials_flow_oauth2_without_client_credentials(self): ) ) - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = auth_scheme auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None @@ -623,7 +638,7 @@ def test_is_client_credentials_flow_oidc_with_client_credentials(self): grant_types_supported=["authorization_code", "client_credentials"], ) - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = auth_scheme auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None @@ -643,7 +658,7 @@ def test_is_client_credentials_flow_oidc_without_client_credentials(self): grant_types_supported=["authorization_code"], ) - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = auth_scheme auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None @@ -657,7 +672,7 @@ def test_is_client_credentials_flow_other_scheme(self): # Create a non-OAuth2/OIDC scheme auth_scheme = Mock() - auth_config = Mock(spec=AuthConfig) + auth_config = create_auth_config_mock() auth_config.auth_scheme = auth_scheme auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None From b6cb922e5a8b457d94c60d93c959da2a808dcf94 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Fri, 5 Dec 2025 23:55:24 +0100 Subject: [PATCH 02/12] chore: autoformat --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - src/google/adk/auth/auth_handler.py | 42 +++++++++---------- src/google/adk/auth/auth_tool.py | 6 +-- src/google/adk/auth/credential_manager.py | 11 +++-- src/google/adk/auth/oauth2_credential_util.py | 4 +- .../unittests/auth/test_credential_manager.py | 8 ++-- 7 files changed, 39 insertions(+), 34 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index 2f5d03a772..f68b349d9c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index cfd850b3a3..1bc4ee58c8 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index c8ae77c9a6..b80d2295ab 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -43,30 +43,29 @@ async def exchange_auth_token( self, ) -> AuthCredential: exchanger = OAuth2CredentialExchanger() - + # Restore secret if needed credential = self.auth_config.exchanged_auth_credential redacted = False original_secret = None - + if credential and credential.oauth2 and credential.oauth2.client_id: - # Check if secret needs restoration - from .credential_manager import CredentialManager - secret = CredentialManager.get_client_secret(credential.oauth2.client_id) - if secret and credential.oauth2.client_secret == "": - original_secret = credential.oauth2.client_secret - credential.oauth2.client_secret = secret - redacted = True + # Check if secret needs restoration + from .credential_manager import CredentialManager + + secret = CredentialManager.get_client_secret(credential.oauth2.client_id) + if secret and credential.oauth2.client_secret == "": + original_secret = credential.oauth2.client_secret + credential.oauth2.client_secret = secret + redacted = True try: - res = await exchanger.exchange( - credential, self.auth_config.auth_scheme - ) - return res + res = await exchanger.exchange(credential, self.auth_config.auth_scheme) + return res finally: - # Always re-redact if we restored it - if redacted and credential and credential.oauth2: - credential.oauth2.client_secret = "" + # Always re-redact if we restored it + if redacted and credential and credential.oauth2: + credential.oauth2.client_secret = "" async def parse_and_store_auth_response(self, state: State) -> None: @@ -200,13 +199,14 @@ def generate_auth_uri( client_id = auth_credential.oauth2.client_id client_secret = auth_credential.oauth2.client_secret - + # Check if secret is redacted and restore it from manager if client_secret == "" and client_id: - from .credential_manager import CredentialManager - secret = CredentialManager.get_client_secret(client_id) - if secret: - client_secret = secret + from .credential_manager import CredentialManager + + secret = CredentialManager.get_client_secret(client_id) + if secret: + client_secret = secret client = OAuth2Session( client_id, diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index 0ab605c5a7..471a4c0df6 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -81,11 +81,11 @@ def get_credential_key(self): if auth_credential and auth_credential.model_extra: auth_credential = auth_credential.model_copy(deep=True) auth_credential.model_extra.clear() - + # Normalize secret to ensure stable key regardless of redaction if auth_credential and auth_credential.oauth2: - auth_credential = auth_credential.model_copy(deep=True) - auth_credential.oauth2.client_secret = None + auth_credential = auth_credential.model_copy(deep=True) + auth_credential.oauth2.client_secret = None credential_name = ( f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}" diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index b34ad6ef97..15c93930d2 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -138,7 +138,9 @@ def _secure_client_secret(self, credential: Optional[AuthCredential]): and credential.oauth2.client_secret and credential.oauth2.client_secret != "" ): - logger.info(f"Securing client secret for client_id: {credential.oauth2.client_id}") + logger.info( + f"Securing client secret for client_id: {credential.oauth2.client_id}" + ) # Store in memory map self._CLIENT_SECRETS[credential.oauth2.client_id] = ( credential.oauth2.client_secret @@ -146,8 +148,11 @@ def _secure_client_secret(self, credential: Optional[AuthCredential]): # Redact from config credential.oauth2.client_secret = "" else: - if credential and credential.oauth2: - logger.debug(f"Not securing secret for client_id {credential.oauth2.client_id}: secret is {credential.oauth2.client_secret}") + if credential and credential.oauth2: + logger.debug( + f"Not securing secret for client_id {credential.oauth2.client_id}:" + f" secret is {credential.oauth2.client_secret}" + ) @staticmethod def get_client_secret(client_id: str) -> Optional[str]: diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py index 954f000fab..9723926611 100644 --- a/src/google/adk/auth/oauth2_credential_util.py +++ b/src/google/adk/auth/oauth2_credential_util.py @@ -108,7 +108,9 @@ def update_credential_with_tokens( tokens: The OAuth2Token object containing new token information. """ auth_credential.oauth2.access_token = tokens.get("access_token") - sys.stderr.write(f"[DEBUG] Assigned access_token: {auth_credential.oauth2.access_token}\n") + sys.stderr.write( + f"[DEBUG] Assigned access_token: {auth_credential.oauth2.access_token}\n" + ) auth_credential.oauth2.refresh_token = tokens.get("refresh_token") auth_credential.oauth2.expires_at = ( int(tokens.get("expires_at")) if tokens.get("expires_at") else None diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index 92874c6cb7..aa30f7be40 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -36,17 +36,17 @@ import pytest - def create_auth_config_mock(): """Creates a mock AuthConfig that returns itself on model_copy.""" - # We remove spec=AuthConfig because accessing Pydantic fields on a spec-ed mock + # We remove spec=AuthConfig because accessing Pydantic fields on a spec-ed mock # can fail if they are not seen as class attributes or if we need dynamic attributes. - m = Mock() - m.spec = AuthConfig # Optional: if we want isinstance to work, but Mock(spec=X) enforces attributes. + m = Mock() + m.spec = AuthConfig # Optional: if we want isinstance to work, but Mock(spec=X) enforces attributes. # Let's just use a plain Mock and configure what we need. m.model_copy.side_effect = lambda **kwargs: m return m + class TestCredentialManager: """Test suite for CredentialManager.""" From a4fb713244641fb030cbf79abf67671ede9c6d49 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 18:46:11 +0100 Subject: [PATCH 03/12] Update src/google/adk/auth/auth_tool.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/auth/auth_tool.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index 471a4c0df6..ae1022b8ab 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -78,14 +78,13 @@ def get_credential_key(self): ) auth_credential = self.raw_auth_credential - if auth_credential and auth_credential.model_extra: + if auth_credential and (auth_credential.model_extra or auth_credential.oauth2): auth_credential = auth_credential.model_copy(deep=True) - auth_credential.model_extra.clear() - - # Normalize secret to ensure stable key regardless of redaction - if auth_credential and auth_credential.oauth2: - auth_credential = auth_credential.model_copy(deep=True) - auth_credential.oauth2.client_secret = None + if auth_credential.model_extra: + auth_credential.model_extra.clear() + # Normalize secret to ensure stable key regardless of redaction + if auth_credential.oauth2: + auth_credential.oauth2.client_secret = None credential_name = ( f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}" From f339c5023c8ea66f34faa02981a1df41b6279767 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 18:46:21 +0100 Subject: [PATCH 04/12] Update src/google/adk/auth/credential_manager.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/auth/credential_manager.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 15c93930d2..71c5e839d3 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -82,17 +82,10 @@ def __init__( self, auth_config: AuthConfig, ): - # We deep copy the auth_config to avoid modifying the original one passed by user - # and to ensure we can safely redact sensitive information - # However we cannot rely on copy.deepcopy because AuthConfig is a pydantic model - # and deepcopy on pydantic model is not always reliable? No, deepcopy works. - # But better to use model_copy if possible. AuthConfig inherits BaseModelWithConfig which is Pydantic. - # auth_config.model_copy(deep=True) is available in Pydantic V2. - # For safe side, we use the passed instance but we redact sensitive info immediately. - # Wait, modifying passed instance is bad practice if user reuses it. - # But CredentialManager usually takes ownership? - # Let's perform redaction on `self._auth_config` which we assign. - # And we should clone it first. + # We deep copy the auth_config to avoid modifying the original object passed + # by the user. This allows for safe redaction of sensitive information without + # causing side effects. + self._auth_config = auth_config.model_copy(deep=True) # Secure the client secret From 4636a5b9a3945ec1397e985870de06729159c7c9 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 18:59:42 +0100 Subject: [PATCH 05/12] fix: with gemini recommendation --- src/google/adk/auth/auth_handler.py | 2 -- src/google/adk/auth/oauth2_credential_util.py | 4 ---- 2 files changed, 6 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index b80d2295ab..e1607b07cf 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -47,7 +47,6 @@ async def exchange_auth_token( # Restore secret if needed credential = self.auth_config.exchanged_auth_credential redacted = False - original_secret = None if credential and credential.oauth2 and credential.oauth2.client_id: # Check if secret needs restoration @@ -55,7 +54,6 @@ async def exchange_auth_token( secret = CredentialManager.get_client_secret(credential.oauth2.client_id) if secret and credential.oauth2.client_secret == "": - original_secret = credential.oauth2.client_secret credential.oauth2.client_secret = secret redacted = True diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py index 9723926611..d90103a88c 100644 --- a/src/google/adk/auth/oauth2_credential_util.py +++ b/src/google/adk/auth/oauth2_credential_util.py @@ -15,7 +15,6 @@ from __future__ import annotations import logging -import sys from typing import Optional from typing import Tuple @@ -108,9 +107,6 @@ def update_credential_with_tokens( tokens: The OAuth2Token object containing new token information. """ auth_credential.oauth2.access_token = tokens.get("access_token") - sys.stderr.write( - f"[DEBUG] Assigned access_token: {auth_credential.oauth2.access_token}\n" - ) auth_credential.oauth2.refresh_token = tokens.get("refresh_token") auth_credential.oauth2.expires_at = ( int(tokens.get("expires_at")) if tokens.get("expires_at") else None From 9e8d85ddbc77adc20fdd80e149580b61bb06f6a2 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 19:01:02 +0100 Subject: [PATCH 06/12] test: add new test --- .../auth/test_auth_handler_secrets.py | 127 +++++++++++++ .../auth/test_auth_tool_key_stability.py | 73 ++++++++ .../auth/test_credential_manager_secrets.py | 176 ++++++++++++++++++ 3 files changed, 376 insertions(+) create mode 100644 tests/unittests/auth/test_auth_handler_secrets.py create mode 100644 tests/unittests/auth/test_auth_tool_key_stability.py create mode 100644 tests/unittests/auth/test_credential_manager_secrets.py diff --git a/tests/unittests/auth/test_auth_handler_secrets.py b/tests/unittests/auth/test_auth_handler_secrets.py new file mode 100644 index 0000000000..3f2fc916c0 --- /dev/null +++ b/tests/unittests/auth/test_auth_handler_secrets.py @@ -0,0 +1,127 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import Mock +from unittest.mock import patch + +from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_credential import OAuth2Auth +from google.adk.auth.auth_handler import AuthHandler +from google.adk.auth.auth_tool import AuthConfig +from google.adk.auth.credential_manager import CredentialManager +import pytest + + +class TestAuthHandlerSecrets: + + def setUp(self): + # Clear secret store + CredentialManager._CLIENT_SECRETS = {} + + @pytest.mark.asyncio + async def test_exchange_auth_token_restores_and_reredacts_secret(self): + client_id = "test_client_id" + secret = "super_secret_value" + + # Setup secure storage + CredentialManager._CLIENT_SECRETS[client_id] = secret + + # Create credential with redacted secret + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(client_id=client_id, client_secret=""), + ) + + auth_config = Mock(spec=AuthConfig) + auth_config.exchanged_auth_credential = credential + auth_config.auth_scheme = Mock() + + handler = AuthHandler(auth_config) + + # Mock exchanger + mock_exchanger = AsyncMock() + + # Check secret inside exchange + def check_secret(cred, scheme): + assert cred.oauth2.client_secret == secret + return cred + + mock_exchanger.exchange.side_effect = check_secret + + with patch( + "google.adk.auth.auth_handler.OAuth2CredentialExchanger", + return_value=mock_exchanger, + ): + await handler.exchange_auth_token() + + # Verify secret is re-redacted + assert credential.oauth2.client_secret == "" + + def test_generate_auth_uri_uses_restored_secret(self): + client_id = "test_client_id" + secret = "super_secret_value" + + # Setup secure storage + CredentialManager._CLIENT_SECRETS[client_id] = secret + + # Create credential with redacted secret + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id=client_id, + client_secret="", + redirect_uri="http://localhost/callback", + ), + ) + + auth_config = Mock(spec=AuthConfig) + auth_config.raw_auth_credential = credential + auth_config.auth_scheme = Mock() + # Mock flows for scopes + auth_config.auth_scheme.flows.implicit = None + auth_config.auth_scheme.flows.clientCredentials = None + auth_config.auth_scheme.flows.password = None + auth_config.auth_scheme.flows.authorizationCode.scopes = {"scope": "desc"} + auth_config.auth_scheme.flows.authorizationCode.authorizationUrl = ( + "http://auth" + ) + + handler = AuthHandler(auth_config) + + # Mock OAuth2Session + with ( + patch("google.adk.auth.auth_handler.OAuth2Session") as mock_session_cls, + patch("google.adk.auth.auth_handler.AUTHLIB_AVAILABLE", True), + ): + + mock_session = Mock() + mock_session.create_authorization_url.return_value = ( + "http://auth?param=1", + "state", + ) + mock_session_cls.return_value = mock_session + + handler.generate_auth_uri() + + # Verify session was created with the REAL secret, not redacted one + mock_session_cls.assert_called_with( + client_id, + secret, + scope="scope", + redirect_uri="http://localhost/callback", + ) diff --git a/tests/unittests/auth/test_auth_tool_key_stability.py b/tests/unittests/auth/test_auth_tool_key_stability.py new file mode 100644 index 0000000000..827ffbb98e --- /dev/null +++ b/tests/unittests/auth/test_auth_tool_key_stability.py @@ -0,0 +1,73 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest.mock import Mock + +from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_credential import OAuth2Auth +from google.adk.auth.auth_tool import AuthConfig + +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +class TestAuthToolKeyStability(unittest.TestCase): + + def test_key_stability_with_different_secrets(self): + from google.adk.auth.auth_schemes import AuthSchemeType + from google.adk.auth.auth_schemes import OAuth2 + + # Consistent scheme for both + auth_scheme = OAuth2(type=AuthSchemeType.oauth2, flows={}) + + # Config 1: Real secret + auth_credential_1 = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="client_id", client_secret="real_secret", auth_uri="uri" + ), + ) + config1 = AuthConfig( + auth_scheme=auth_scheme, raw_auth_credential=auth_credential_1 + ) + + # Config 2: Redacted secret + auth_credential_2 = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="client_id", client_secret="", auth_uri="uri" + ), + ) + config2 = AuthConfig( + auth_scheme=auth_scheme, raw_auth_credential=auth_credential_2 + ) + + # Keys should be identical + key1 = config1.credential_key + key2 = config2.credential_key + + self.assertEqual(key1, key2, f"Keys should match! {key1} vs {key2}") diff --git a/tests/unittests/auth/test_credential_manager_secrets.py b/tests/unittests/auth/test_credential_manager_secrets.py new file mode 100644 index 0000000000..5825b44825 --- /dev/null +++ b/tests/unittests/auth/test_credential_manager_secrets.py @@ -0,0 +1,176 @@ +from unittest.mock import AsyncMock +from unittest.mock import Mock +from unittest.mock import patch + +from fastapi.openapi.models import OAuth2 +from fastapi.openapi.models import OAuthFlowAuthorizationCode +from fastapi.openapi.models import OAuthFlows +from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_credential import OAuth2Auth +from google.adk.auth.auth_tool import AuthConfig +from google.adk.auth.credential_manager import CredentialManager +import pytest + + +@pytest.mark.asyncio +async def test_credential_manager_redacts_secrets_in_raw_credential(): + """Test that CredentialManager redacts client_secret from raw_auth_credential upon initialization.""" + + # Setup + client_id = "test_client_id" + client_secret = "test_client_secret" + + oauth_auth = OAuth2Auth(client_id=client_id, client_secret=client_secret) + + auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, oauth2=oauth_auth + ) + + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + ) + ) + ) + + auth_config = AuthConfig( + auth_scheme=auth_scheme, raw_auth_credential=auth_credential + ) + + # Act + manager = CredentialManager(auth_config) + + # Assert + # 1. Check if secret is in memory map + assert client_id in manager._CLIENT_SECRETS + assert manager._CLIENT_SECRETS[client_id] == client_secret + + # 2. Check if secret is redacted in the manager's config + assert ( + manager._auth_config.raw_auth_credential.oauth2.client_secret + == "" + ) + + # 3. Check original config is NOT modified (AuthConfig copy behavior) + # Since we used model_copy(deep=True), calling on Pydantic model copies it. + assert auth_config.raw_auth_credential.oauth2.client_secret == client_secret + + +@pytest.mark.asyncio +async def test_credential_manager_redacts_secrets_in_exchanged_credential(): + """Test that CredentialManager redacts client_secret from exchanged_auth_credential if present.""" + + # Setup + client_id = "test_client_id_exchanged" + client_secret = "test_client_secret_exchanged" + + oauth_auth = OAuth2Auth( + client_id=client_id, + client_secret=client_secret, + access_token="some_token", + ) + + exchanged_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, oauth2=oauth_auth + ) + + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + ) + ) + ) + + auth_config = AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=None, + exchanged_auth_credential=exchanged_credential, + ) + + # Act + manager = CredentialManager(auth_config) + + # Assert + assert client_id in manager._CLIENT_SECRETS + assert manager._CLIENT_SECRETS[client_id] == client_secret + + assert ( + manager._auth_config.exchanged_auth_credential.oauth2.client_secret + == "" + ) + + +@pytest.mark.asyncio +async def test_exchange_credential_restores_secret(): + """Test that _exchange_credential restores the secret before calling exchanger.""" + + # Setup + client_id = "test_client_id_exchange" + client_secret = "test_client_secret_exchange" + + oauth_auth = OAuth2Auth(client_id=client_id, client_secret=client_secret) + + raw_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, oauth2=oauth_auth + ) + + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + ) + ) + ) + + auth_config = AuthConfig( + auth_scheme=auth_scheme, raw_auth_credential=raw_credential + ) + + manager = CredentialManager(auth_config) + + # Secret should be redacted now + assert ( + manager._auth_config.raw_auth_credential.oauth2.client_secret + == "" + ) + + # Prepare a credential to be exchanged (e.g. from client response, has no secret or redacted) + credential_to_exchange = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id=client_id, + client_secret="", # or None + auth_code="some_code", + ), + ) + + # Mock exchanger + mock_exchanger = AsyncMock() + + # We use side_effect to verify the secret at the moment of call, because the object is mutated later + def check_secret(cred, scheme): + assert cred.oauth2.client_secret == client_secret + return credential_to_exchange + + mock_exchanger.exchange.side_effect = check_secret + + with patch.object( + manager._exchanger_registry, "get_exchanger", return_value=mock_exchanger + ): + # Act + result_credential, exchanged = await manager._exchange_credential( + credential_to_exchange + ) + + # Assert + # Verification happened in side_effect + assert mock_exchanger.exchange.called + + # Check that the result credential (modified in place or returned) has secret REDACTED again + assert result_credential.oauth2.client_secret == "" From 428cd57cb09081d39e1afc197940446337475f42 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 22:01:07 +0100 Subject: [PATCH 07/12] Update src/google/adk/auth/credential_manager.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/auth/credential_manager.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 71c5e839d3..b726adc808 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -286,14 +286,14 @@ async def _exchange_credential( ] restored = True - exchanged_credential = await exchanger.exchange( - credential, self._auth_config.auth_scheme - ) - - # Redact client secret again after exchange to prevent leakage - if exchanged_credential.oauth2: - exchanged_credential.oauth2.client_secret = "" - + try: + exchanged_credential = await exchanger.exchange( + credential, self._auth_config.auth_scheme + ) + finally: + # Redact client secret again after exchange to prevent leakage + if restored and credential.oauth2: + credential.oauth2.client_secret = "" return exchanged_credential, True async def _refresh_credential( From facc6f49fbe02c3e4be15ee20ea63b9635e6ab1b Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 22:04:30 +0100 Subject: [PATCH 08/12] fix: move import to the top --- src/google/adk/auth/auth_handler.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index e1607b07cf..ff277124a0 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -18,6 +18,7 @@ from .auth_schemes import OpenIdConnectWithConfig from .auth_tool import AuthConfig from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger +from .credential_manager import CredentialManager if TYPE_CHECKING: from ..sessions.state import State @@ -49,9 +50,6 @@ async def exchange_auth_token( redacted = False if credential and credential.oauth2 and credential.oauth2.client_id: - # Check if secret needs restoration - from .credential_manager import CredentialManager - secret = CredentialManager.get_client_secret(credential.oauth2.client_id) if secret and credential.oauth2.client_secret == "": credential.oauth2.client_secret = secret @@ -200,8 +198,6 @@ def generate_auth_uri( # Check if secret is redacted and restore it from manager if client_secret == "" and client_id: - from .credential_manager import CredentialManager - secret = CredentialManager.get_client_secret(client_id) if secret: client_secret = secret From 962dcde21039c3ff9e684efb3585f6510a7f2246 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 6 Dec 2025 22:42:07 +0100 Subject: [PATCH 09/12] fix: follow gemini advices --- src/google/adk/auth/auth_handler.py | 55 ++++++---------- src/google/adk/auth/credential_manager.py | 63 +++++++++++++------ .../auth/test_auth_handler_secrets.py | 8 ++- .../unittests/auth/test_credential_manager.py | 2 +- .../auth/test_credential_manager_secrets.py | 8 +++ 5 files changed, 79 insertions(+), 57 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index ff277124a0..d282d16d41 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -47,21 +47,10 @@ async def exchange_auth_token( # Restore secret if needed credential = self.auth_config.exchanged_auth_credential - redacted = False - if credential and credential.oauth2 and credential.oauth2.client_id: - secret = CredentialManager.get_client_secret(credential.oauth2.client_id) - if secret and credential.oauth2.client_secret == "": - credential.oauth2.client_secret = secret - redacted = True - - try: + with CredentialManager.restore_client_secret(credential): res = await exchanger.exchange(credential, self.auth_config.auth_scheme) return res - finally: - # Always re-redact if we restored it - if redacted and credential and credential.oauth2: - credential.oauth2.client_secret = "" async def parse_and_store_auth_response(self, state: State) -> None: @@ -194,29 +183,25 @@ def generate_auth_uri( scopes = list(scopes.keys()) client_id = auth_credential.oauth2.client_id - client_secret = auth_credential.oauth2.client_secret - - # Check if secret is redacted and restore it from manager - if client_secret == "" and client_id: - secret = CredentialManager.get_client_secret(client_id) - if secret: - client_secret = secret - - client = OAuth2Session( - client_id, - client_secret, - scope=" ".join(scopes), - redirect_uri=auth_credential.oauth2.redirect_uri, - ) - params = { - "access_type": "offline", - "prompt": "consent", - } - if auth_credential.oauth2.audience: - params["audience"] = auth_credential.oauth2.audience - uri, state = client.create_authorization_url( - url=authorization_endpoint, **params - ) + + + with CredentialManager.restore_client_secret(auth_credential): + client_secret = auth_credential.oauth2.client_secret + client = OAuth2Session( + client_id, + client_secret, + scope=" ".join(scopes), + redirect_uri=auth_credential.oauth2.redirect_uri, + ) + params = { + "access_type": "offline", + "prompt": "consent", + } + if auth_credential.oauth2.audience: + params["audience"] = auth_credential.oauth2.audience + uri, state = client.create_authorization_url( + url=authorization_endpoint, **params + ) exchanged_auth_credential = auth_credential.model_copy(deep=True) exchanged_auth_credential.oauth2.auth_uri = uri diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index b726adc808..88eaec07c1 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -14,6 +14,7 @@ from __future__ import annotations +import contextlib import logging from typing import Optional @@ -135,7 +136,7 @@ def _secure_client_secret(self, credential: Optional[AuthCredential]): f"Securing client secret for client_id: {credential.oauth2.client_id}" ) # Store in memory map - self._CLIENT_SECRETS[credential.oauth2.client_id] = ( + CredentialManager._CLIENT_SECRETS[credential.oauth2.client_id] = ( credential.oauth2.client_secret ) # Redact from config @@ -250,6 +251,40 @@ async def _load_from_auth_response( """Load credential from auth response in callback context.""" return callback_context.get_auth_response(self._auth_config) + @staticmethod + @contextlib.contextmanager + def restore_client_secret(credential: AuthCredential, secret: str = None): + """Context manager to temporarily restore client secret in a credential. + + Args: + credential: The credential to restore secret for. + secret: Optional secret to use. If not provided, looks up by client_id. + """ + if not credential or not credential.oauth2: + yield + return + + restored = False + if secret: + credential.oauth2.client_secret = secret + restored = True + elif ( + credential.oauth2.client_id + and credential.oauth2.client_secret == "" + ): + stored_secret = CredentialManager.get_client_secret( + credential.oauth2.client_id + ) + if stored_secret: + credential.oauth2.client_secret = stored_secret + restored = True + + try: + yield + finally: + if restored: + credential.oauth2.client_secret = "" + async def _exchange_credential( self, credential: AuthCredential ) -> tuple[AuthCredential, bool]: @@ -263,37 +298,27 @@ async def _exchange_credential( self._auth_config.auth_scheme, credential ) else: - # Restore client secret from memory map for exchange - restored = False + # Determine if we need to fallback/lookup secret from raw credential + secret_to_use = None if ( credential.oauth2 and credential.oauth2.client_id - and credential.oauth2.client_id in self._CLIENT_SECRETS - ): - credential.oauth2.client_secret = self._CLIENT_SECRETS[ - credential.oauth2.client_id - ] - restored = True - elif ( - self._auth_config.raw_auth_credential + and credential.oauth2.client_id not in self._CLIENT_SECRETS + and self._auth_config.raw_auth_credential and self._auth_config.raw_auth_credential.oauth2 and self._auth_config.raw_auth_credential.oauth2.client_id in self._CLIENT_SECRETS ): - # Fallback to look up using raw credential client id if credential client id is missing (unlikely for valid flow) - credential.oauth2.client_secret = self._CLIENT_SECRETS[ + # Fallback to look up using raw credential client id + secret_to_use = self._CLIENT_SECRETS[ self._auth_config.raw_auth_credential.oauth2.client_id ] - restored = True - try: + with self.restore_client_secret(credential, secret=secret_to_use): exchanged_credential = await exchanger.exchange( credential, self._auth_config.auth_scheme ) - finally: - # Redact client secret again after exchange to prevent leakage - if restored and credential.oauth2: - credential.oauth2.client_secret = "" + return exchanged_credential, True async def _refresh_credential( diff --git a/tests/unittests/auth/test_auth_handler_secrets.py b/tests/unittests/auth/test_auth_handler_secrets.py index 3f2fc916c0..a912fbb5d1 100644 --- a/tests/unittests/auth/test_auth_handler_secrets.py +++ b/tests/unittests/auth/test_auth_handler_secrets.py @@ -29,9 +29,13 @@ class TestAuthHandlerSecrets: - def setUp(self): - # Clear secret store + @pytest.fixture(autouse=True) + def clear_credential_manager_secrets(self): + """Clear CredentialManager secrets buffer before/after each test.""" CredentialManager._CLIENT_SECRETS = {} + yield + CredentialManager._CLIENT_SECRETS = {} + @pytest.mark.asyncio async def test_exchange_auth_token_restores_and_reredacts_secret(self): diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index aa30f7be40..f59b00a601 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -31,7 +31,7 @@ from google.adk.auth.auth_schemes import ExtendedOAuth2 from google.adk.auth.auth_tool import AuthConfig from google.adk.auth.credential_manager import CredentialManager -from google.adk.auth.credential_manager import ServiceAccountCredentialExchanger +from google.adk.tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger from google.adk.auth.oauth2_discovery import AuthorizationServerMetadata import pytest diff --git a/tests/unittests/auth/test_credential_manager_secrets.py b/tests/unittests/auth/test_credential_manager_secrets.py index 5825b44825..ea5d210818 100644 --- a/tests/unittests/auth/test_credential_manager_secrets.py +++ b/tests/unittests/auth/test_credential_manager_secrets.py @@ -13,6 +13,14 @@ import pytest +@pytest.fixture(autouse=True) +def clear_credential_manager_secrets(): + """Clear CredentialManager secrets buffer before/after each test.""" + CredentialManager._CLIENT_SECRETS = {} + yield + CredentialManager._CLIENT_SECRETS = {} + + @pytest.mark.asyncio async def test_credential_manager_redacts_secrets_in_raw_credential(): """Test that CredentialManager redacts client_secret from raw_auth_credential upon initialization.""" From 2e015655749d65c7304d8c8da7e1dcf787daa3bb Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sun, 7 Dec 2025 09:30:56 +0100 Subject: [PATCH 10/12] Update src/google/adk/auth/credential_manager.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/auth/credential_manager.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 88eaec07c1..fe804066c6 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -298,23 +298,7 @@ async def _exchange_credential( self._auth_config.auth_scheme, credential ) else: - # Determine if we need to fallback/lookup secret from raw credential - secret_to_use = None - if ( - credential.oauth2 - and credential.oauth2.client_id - and credential.oauth2.client_id not in self._CLIENT_SECRETS - and self._auth_config.raw_auth_credential - and self._auth_config.raw_auth_credential.oauth2 - and self._auth_config.raw_auth_credential.oauth2.client_id - in self._CLIENT_SECRETS - ): - # Fallback to look up using raw credential client id - secret_to_use = self._CLIENT_SECRETS[ - self._auth_config.raw_auth_credential.oauth2.client_id - ] - - with self.restore_client_secret(credential, secret=secret_to_use): + with self.restore_client_secret(credential): exchanged_credential = await exchanger.exchange( credential, self._auth_config.auth_scheme ) From 65c40e52a67ad65840ea7bbbd9696dfc06b96cab Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sun, 7 Dec 2025 09:39:58 +0100 Subject: [PATCH 11/12] fix: local import to avoid circular dependencies --- src/google/adk/auth/credential_manager.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index fe804066c6..b1a3e28d9c 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -21,7 +21,6 @@ from fastapi.openapi.models import OAuth2 from ..agents.callback_context import CallbackContext -from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger from ..utils.feature_decorator import experimental from .auth_credential import AuthCredential from .auth_credential import AuthCredentialTypes @@ -110,6 +109,7 @@ def __init__( ) # TODO: Move ServiceAccountCredentialExchanger to the auth module + from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger self._exchanger_registry.register( AuthCredentialTypes.SERVICE_ACCOUNT, ServiceAccountCredentialExchanger(), @@ -293,6 +293,7 @@ async def _exchange_credential( if not exchanger: return credential, False + from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger if isinstance(exchanger, ServiceAccountCredentialExchanger): exchanged_credential = exchanger.exchange_credential( self._auth_config.auth_scheme, credential From d6a049da382db776d1bbd41cb3bf84456991c19e Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Thu, 11 Dec 2025 17:06:04 +0100 Subject: [PATCH 12/12] fix: test chore: isort & pyink --- src/google/adk/auth/auth_handler.py | 3 +-- src/google/adk/auth/auth_tool.py | 4 +++- src/google/adk/auth/credential_manager.py | 2 ++ tests/unittests/auth/test_auth_handler_secrets.py | 1 - tests/unittests/auth/test_credential_manager.py | 2 +- tests/unittests/tools/mcp_tool/test_mcp_tool.py | 14 +++++++++++++- 6 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index d282d16d41..55d6cddfa3 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -17,8 +17,8 @@ from .auth_schemes import AuthSchemeType from .auth_schemes import OpenIdConnectWithConfig from .auth_tool import AuthConfig -from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger from .credential_manager import CredentialManager +from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger if TYPE_CHECKING: from ..sessions.state import State @@ -184,7 +184,6 @@ def generate_auth_uri( client_id = auth_credential.oauth2.client_id - with CredentialManager.restore_client_secret(auth_credential): client_secret = auth_credential.oauth2.client_secret client = OAuth2Session( diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index ae1022b8ab..e77c43f6af 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -78,7 +78,9 @@ def get_credential_key(self): ) auth_credential = self.raw_auth_credential - if auth_credential and (auth_credential.model_extra or auth_credential.oauth2): + if auth_credential and ( + auth_credential.model_extra or auth_credential.oauth2 + ): auth_credential = auth_credential.model_copy(deep=True) if auth_credential.model_extra: auth_credential.model_extra.clear() diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index b1a3e28d9c..dd39a0b44a 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -110,6 +110,7 @@ def __init__( # TODO: Move ServiceAccountCredentialExchanger to the auth module from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger + self._exchanger_registry.register( AuthCredentialTypes.SERVICE_ACCOUNT, ServiceAccountCredentialExchanger(), @@ -294,6 +295,7 @@ async def _exchange_credential( return credential, False from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger + if isinstance(exchanger, ServiceAccountCredentialExchanger): exchanged_credential = exchanger.exchange_credential( self._auth_config.auth_scheme, credential diff --git a/tests/unittests/auth/test_auth_handler_secrets.py b/tests/unittests/auth/test_auth_handler_secrets.py index a912fbb5d1..e52658597d 100644 --- a/tests/unittests/auth/test_auth_handler_secrets.py +++ b/tests/unittests/auth/test_auth_handler_secrets.py @@ -36,7 +36,6 @@ def clear_credential_manager_secrets(self): yield CredentialManager._CLIENT_SECRETS = {} - @pytest.mark.asyncio async def test_exchange_auth_token_restores_and_reredacts_secret(self): client_id = "test_client_id" diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index f59b00a601..53ff1cd2cc 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -31,8 +31,8 @@ from google.adk.auth.auth_schemes import ExtendedOAuth2 from google.adk.auth.auth_tool import AuthConfig from google.adk.auth.credential_manager import CredentialManager -from google.adk.tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger from google.adk.auth.oauth2_discovery import AuthorizationServerMetadata +from google.adk.tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger import pytest diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 1284e73bce..bdfaa7444e 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -84,10 +84,13 @@ def test_init_with_auth(self): # Create real auth scheme instances instead of mocks from fastapi.openapi.models import OAuth2 + test_client_secret = "test_secret" auth_scheme = OAuth2(flows={}) auth_credential = AuthCredential( auth_type=AuthCredentialTypes.OAUTH2, - oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"), + oauth2=OAuth2Auth( + client_id="test_id", client_secret=test_client_secret + ), ) tool = MCPTool( @@ -100,6 +103,15 @@ def test_init_with_auth(self): # The auth config is stored in the parent class _credentials_manager assert tool._credentials_manager is not None assert tool._credentials_manager._auth_config.auth_scheme == auth_scheme + assert ( + tool._credentials_manager._auth_config.raw_auth_credential.oauth2.client_secret + == "" + ) + + # Restore the client secret and validate it's the same credential in the end. + tool._credentials_manager._auth_config.raw_auth_credential.oauth2.client_secret = ( + test_client_secret + ) assert ( tool._credentials_manager._auth_config.raw_auth_credential == auth_credential