From 33012e6dda9ef20c7a1dae66a84717de5d782097 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 2 Feb 2026 17:50:50 -0800 Subject: [PATCH] fix: Make credential key generation stable and prevent cross-user credential leaks This change updates the credential key generation to use a stable hash (SHA256) instead of Python's built-in hash, which can vary based on PYTHONHASHSEED. It also makes sure that temporary or exchanged OAuth2 fields are excluded from the key calculation. I also added when saving credentials, a copy of the AuthConfig is used to avoid modifying the original shared AuthConfig instance with user-specific exchanged credentials. Co-authored-by: George Weale PiperOrigin-RevId: 864599326 --- src/google/adk/auth/auth_handler.py | 2 + src/google/adk/auth/auth_preprocessor.py | 32 ++++ src/google/adk/auth/auth_tool.py | 53 ++++++- src/google/adk/auth/credential_manager.py | 15 +- .../openapi_spec_parser/openapi_toolset.py | 20 ++- .../openapi_spec_parser/rest_api_tool.py | 14 +- .../openapi_spec_parser/tool_auth_handler.py | 107 ++++++++++--- tests/unittests/auth/test_auth_config.py | 55 +++++++ tests/unittests/auth/test_auth_handler.py | 27 ++++ .../unittests/auth/test_credential_manager.py | 149 ++++++++++++++++-- .../test_tool_auth_handler.py | 17 ++ 11 files changed, 448 insertions(+), 43 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 96413bf2..ec7c7571 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -115,6 +115,7 @@ class AuthHandler: exchanged_auth_credential=self.auth_config.raw_auth_credential.model_copy( deep=True ), + credential_key=self.auth_config.credential_key, ) # Check for client_id and client_secret @@ -133,6 +134,7 @@ class AuthHandler: auth_scheme=self.auth_config.auth_scheme, raw_auth_credential=self.auth_config.raw_auth_credential, exchanged_auth_credential=exchanged_credential, + credential_key=self.auth_config.credential_key, ) def generate_auth_uri( diff --git a/src/google/adk/auth/auth_preprocessor.py b/src/google/adk/auth/auth_preprocessor.py index 1c3ac70d..fc50cdff 100644 --- a/src/google/adk/auth/auth_preprocessor.py +++ b/src/google/adk/auth/auth_preprocessor.py @@ -72,6 +72,7 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor): if not responses: return + requested_auth_config_by_request_id = {} # look for auth response for function_call_response in responses: if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME: @@ -79,7 +80,38 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor): # found the function call response for the system long running request euc # function call request_euc_function_call_ids.add(function_call_response.id) + + if request_euc_function_call_ids: + for event in events: + function_calls = event.get_function_calls() + if not function_calls: + continue + try: + for function_call in function_calls: + if ( + function_call.id in request_euc_function_call_ids + and function_call.name == REQUEST_EUC_FUNCTION_CALL_NAME + ): + args = AuthToolArguments.model_validate(function_call.args) + requested_auth_config_by_request_id[function_call.id] = ( + args.auth_config + ) + except TypeError: + continue + + for function_call_response in responses: + if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME: + continue + auth_config = AuthConfig.model_validate(function_call_response.response) + requested_auth_config = requested_auth_config_by_request_id.get( + function_call_response.id + ) + if ( + requested_auth_config + and requested_auth_config.credential_key is not None + ): + auth_config.credential_key = requested_auth_config.credential_key await AuthHandler(auth_config=auth_config).parse_and_store_auth_response( state=invocation_context.session.state ) diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index 3b09b7e3..cc8a244e 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -14,8 +14,11 @@ from __future__ import annotations +import hashlib +import json from typing import Optional +from pydantic import BaseModel from typing_extensions import deprecated from .auth_credential import AuthCredential @@ -23,6 +26,28 @@ from .auth_credential import BaseModelWithConfig from .auth_schemes import AuthScheme +def _stable_model_digest(model: BaseModel) -> str: + """Returns a stable digest for a pydantic model. + + The digest is stable across: + - Python hash seeds (does not use `hash()`). + - Dict insertion ordering differences (canonicalizes via `sort_keys=True`). + - Pydantic `model_extra` values (ignored). + """ + if getattr(model, "model_extra", None): + model = model.model_copy(deep=True) + model.model_extra.clear() + + dumped = model.model_dump(by_alias=True, exclude_none=True, mode="json") + canonical_json = json.dumps( + dumped, + sort_keys=True, + ensure_ascii=False, + separators=(",", ":"), + ) + return hashlib.sha256(canonical_json.encode("utf-8")).hexdigest()[:16] + + class AuthConfig(BaseModelWithConfig): """The auth config sent by tool asking client to collect auth credentials and @@ -58,12 +83,22 @@ class AuthConfig(BaseModelWithConfig): super().__init__(**data) if self.credential_key: return + for obj in (self.raw_auth_credential, self.auth_scheme): + if not obj or not getattr(obj, "model_extra", None): + continue + for key in ("credential_key", "credentialKey"): + value = obj.model_extra.get(key) + if isinstance(value, str) and value: + self.credential_key = value + return self.credential_key = self.get_credential_key() @deprecated("This method is deprecated. Use credential_key instead.") def get_credential_key(self): - """Builds a hash key based on auth_scheme and raw_auth_credential used to - save / load this credential to / from a credentials service. + """Builds a stable key based on auth_scheme and raw_auth_credential. + + This is used to save/load credentials to/from a credential service when + `credential_key` is not explicitly provided. """ auth_scheme = self.auth_scheme @@ -72,7 +107,7 @@ class AuthConfig(BaseModelWithConfig): auth_scheme = auth_scheme.model_copy(deep=True) auth_scheme.model_extra.clear() scheme_name = ( - f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}" + f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}" if auth_scheme else "" ) @@ -81,8 +116,18 @@ class AuthConfig(BaseModelWithConfig): if auth_credential and auth_credential.model_extra: auth_credential = auth_credential.model_copy(deep=True) auth_credential.model_extra.clear() + if auth_credential and auth_credential.oauth2: + auth_credential = auth_credential.model_copy(deep=True) + auth_credential.oauth2.auth_uri = None + auth_credential.oauth2.state = None + auth_credential.oauth2.auth_response_uri = None + auth_credential.oauth2.auth_code = None + auth_credential.oauth2.access_token = None + auth_credential.oauth2.refresh_token = None + auth_credential.oauth2.expires_at = None + auth_credential.oauth2.expires_in = None credential_name = ( - f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}" + f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}" if auth_credential else "" ) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 21b6328a..d037a43e 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -142,7 +142,9 @@ class CredentialManager: # Step 2: Check if credential is already ready (no processing needed) if self._is_credential_ready(): - return self._auth_config.raw_auth_credential + # Return a copy to avoid leaking mutations across invocations/users when + # tools share a long-lived AuthConfig instance. + return self._auth_config.raw_auth_credential.model_copy(deep=True) # Step 3: Try to load existing processed credential credential = await self._load_existing_credential(context) @@ -159,7 +161,9 @@ class CredentialManager: if not credential: # For client credentials flow, use raw credentials directly if self._is_client_credentials_flow(): - credential = self._auth_config.raw_auth_credential + # Exchange/refresh steps may mutate the credential object in-place, so + # do not operate on the shared tool config. + credential = self._auth_config.raw_auth_credential.model_copy(deep=True) else: # For authorization code flow, return None to trigger user authorization return None @@ -298,12 +302,11 @@ class CredentialManager: self, context: CallbackContext, credential: AuthCredential ) -> None: """Save credential to credential service if available.""" - # Update the exchanged credential in config - self._auth_config.exchanged_auth_credential = credential - credential_service = context._invocation_context.credential_service if credential_service: - await context.save_credential(self._auth_config) + auth_config_to_save = self._auth_config.model_copy(deep=True) + auth_config_to_save.exchanged_auth_credential = credential + await context.save_credential(auth_config_to_save) async def _populate_auth_scheme(self) -> bool: """Auto-discover server metadata and populate missing auth scheme info. diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py index 15317ff3..9e5dd8dd 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py @@ -70,6 +70,7 @@ class OpenAPIToolset(BaseToolset): spec_str_type: Literal["json", "yaml"] = "json", auth_scheme: Optional[AuthScheme] = None, auth_credential: Optional[AuthCredential] = None, + credential_key: Optional[str] = None, tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, tool_name_prefix: Optional[str] = None, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None, @@ -108,6 +109,8 @@ class OpenAPIToolset(BaseToolset): auth_credential: The auth credential to use for all tools. Use AuthCredential or use helpers in ``google.adk.tools.openapi_tool.auth.auth_helpers`` + credential_key: Optional stable key used for interactive auth and + credential caching across all tools in this toolset. tool_filter: The filter used to filter the tools in the toolset. It can be either a tool predicate or a list of tool names of the tools to expose. tool_name_prefix: The prefix to prepend to the names of the tools returned @@ -137,6 +140,7 @@ class OpenAPIToolset(BaseToolset): AuthConfig( auth_scheme=auth_scheme, raw_auth_credential=auth_credential, + credential_key=credential_key, ) if auth_scheme else None @@ -147,6 +151,8 @@ class OpenAPIToolset(BaseToolset): self._tools: Final[List[RestApiTool]] = list(self._parse(spec_dict)) if auth_scheme or auth_credential: self._configure_auth_all(auth_scheme, auth_credential) + if credential_key: + self._configure_credential_key_all(credential_key) def _configure_auth_all( self, auth_scheme: AuthScheme, auth_credential: AuthCredential @@ -159,6 +165,11 @@ class OpenAPIToolset(BaseToolset): if auth_credential: tool.configure_auth_credential(auth_credential) + def _configure_credential_key_all(self, credential_key: str): + """Configure credential key for all tools.""" + for tool in self._tools: + tool.configure_credential_key(credential_key) + def configure_ssl_verify_all( self, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None ): @@ -229,8 +240,9 @@ class OpenAPIToolset(BaseToolset): def get_auth_config(self) -> Optional[AuthConfig]: """Returns the auth config for this toolset. - ADK will populate exchanged_auth_credential on this config before calling - get_tools(). The toolset can then access the ready-to-use credential via - self._auth_config.exchanged_auth_credential. + Note: This returns a copy so any exchanged credentials populated by the ADK + framework do not persist on the toolset instance across invocations. """ - return self._auth_config + return ( + self._auth_config.model_copy(deep=True) if self._auth_config else None + ) diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py index 70a53a9e..300c47e1 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py @@ -100,6 +100,8 @@ class RestApiTool(BaseTool): header_provider: Optional[ Callable[[ReadonlyContext], Dict[str, str]] ] = None, + *, + credential_key: Optional[str] = None, ): """Initializes the RestApiTool with the given parameters. @@ -137,6 +139,8 @@ class RestApiTool(BaseTool): an argument, allowing dynamic header generation based on the current context. Useful for adding custom headers like correlation IDs, authentication tokens, or other request metadata. + credential_key: Optional stable key used for interactive auth and + credential caching. """ # Gemini restrict the length of function name to be less than 64 characters self.name = name[:60] @@ -152,6 +156,7 @@ class RestApiTool(BaseTool): else operation ) self.auth_credential, self.auth_scheme = None, None + self.credential_key = credential_key self.configure_auth_credential(auth_credential) self.configure_auth_scheme(auth_scheme) @@ -266,6 +271,10 @@ class RestApiTool(BaseTool): auth_credential = AuthCredential.model_validate_json(auth_credential) self.auth_credential = auth_credential + def configure_credential_key(self, credential_key: Optional[str] = None): + """Configures the credential key for interactive auth / caching.""" + self.credential_key = credential_key + def configure_ssl_verify( self, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None ): @@ -449,7 +458,10 @@ class RestApiTool(BaseTool): """ # Prepare auth credentials for the API call tool_auth_handler = ToolAuthHandler.from_tool_context( - tool_context, self.auth_scheme, self.auth_credential + tool_context, + self.auth_scheme, + self.auth_credential, + credential_key=self.credential_key, ) auth_result = await tool_auth_handler.prepare_auth_credentials() auth_state, auth_scheme, 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 e8d9738b..1c534b44 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 @@ -14,6 +14,7 @@ from __future__ import annotations +import hashlib import logging from typing import Literal from typing import Optional @@ -24,6 +25,7 @@ from ....auth.auth_credential import AuthCredential from ....auth.auth_credential import AuthCredentialTypes from ....auth.auth_schemes import AuthScheme from ....auth.auth_schemes import AuthSchemeType +from ....auth.auth_tool import _stable_model_digest from ....auth.auth_tool import AuthConfig from ....auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher from ...tool_context import ToolContext @@ -50,19 +52,60 @@ class ToolContextCredentialStore: def __init__(self, tool_context: ToolContext): self.tool_context = tool_context + def _legacy_stable_digest(self, text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16] + + def _get_legacy_credential_key( + self, + auth_scheme: Optional[AuthScheme], + auth_credential: Optional[AuthCredential], + ) -> str: + if auth_credential and auth_credential.oauth2: + auth_credential = auth_credential.model_copy(deep=True) + auth_credential.oauth2.auth_uri = None + auth_credential.oauth2.state = None + auth_credential.oauth2.auth_response_uri = None + auth_credential.oauth2.auth_code = None + auth_credential.oauth2.access_token = None + auth_credential.oauth2.refresh_token = None + auth_credential.oauth2.expires_at = None + auth_credential.oauth2.expires_in = None + scheme_name = ( + f"{auth_scheme.type_.name}_{self._legacy_stable_digest(auth_scheme.model_dump_json())}" + if auth_scheme + else "" + ) + credential_name = ( + f"{auth_credential.auth_type.value}_{self._legacy_stable_digest(auth_credential.model_dump_json())}" + if auth_credential + else "" + ) + return f"{scheme_name}_{credential_name}_existing_exchanged_credential" + def get_credential_key( self, auth_scheme: Optional[AuthScheme], auth_credential: Optional[AuthCredential], ) -> str: """Generates a unique key for the given auth scheme and credential.""" + + if auth_credential and auth_credential.oauth2: + auth_credential = auth_credential.model_copy(deep=True) + auth_credential.oauth2.auth_uri = None + auth_credential.oauth2.state = None + auth_credential.oauth2.auth_response_uri = None + auth_credential.oauth2.auth_code = None + auth_credential.oauth2.access_token = None + auth_credential.oauth2.refresh_token = None + auth_credential.oauth2.expires_at = None + auth_credential.oauth2.expires_in = None scheme_name = ( - f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}" + f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}" if auth_scheme else "" ) credential_name = ( - f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}" + f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}" if auth_credential else "" ) @@ -86,9 +129,19 @@ class ToolContextCredentialStore: # session implementation, we don't want session to persist the token, # meanwhile we want the token shared across runs. serialized_credential = self.tool_context.state.get(token_key) - if not serialized_credential: + if serialized_credential: + return AuthCredential.model_validate(serialized_credential) + + legacy_key = self._get_legacy_credential_key(auth_scheme, auth_credential) + if legacy_key == token_key: return None - return AuthCredential.model_validate(serialized_credential) + serialized_legacy_credential = self.tool_context.state.get(legacy_key) + if not serialized_legacy_credential: + return None + + # Migrate to the current key for future lookups. + self.tool_context.state[token_key] = serialized_legacy_credential + return AuthCredential.model_validate(serialized_legacy_credential) def store_credential( self, @@ -114,6 +167,8 @@ class ToolAuthHandler: auth_credential: Optional[AuthCredential], credential_exchanger: Optional[BaseAuthCredentialExchanger] = None, credential_store: Optional["ToolContextCredentialStore"] = None, + *, + credential_key: Optional[str] = None, ): self.tool_context = tool_context self.auth_scheme = ( @@ -122,12 +177,35 @@ class ToolAuthHandler: self.auth_credential = ( auth_credential.model_copy(deep=True) if auth_credential else None ) + self._credential_key = credential_key self.credential_exchanger = ( credential_exchanger or AutoAuthCredentialExchanger() ) self.credential_store = credential_store self.should_store_credential = True + def _get_credential_key_override(self) -> Optional[str]: + """Returns a user-provided credential_key if available.""" + if self._credential_key: + return self._credential_key + + for obj in (self.auth_credential, self.auth_scheme): + if not obj or not obj.model_extra: + continue + for key in ("credential_key", "credentialKey"): + value = obj.model_extra.get(key) + if isinstance(value, str) and value: + return value + + return None + + def _build_auth_config(self) -> AuthConfig: + return AuthConfig( + auth_scheme=self.auth_scheme, + raw_auth_credential=self.auth_credential, + credential_key=self._get_credential_key_override(), + ) + @classmethod def from_tool_context( cls, @@ -135,6 +213,8 @@ class ToolAuthHandler: auth_scheme: Optional[AuthScheme], auth_credential: Optional[AuthCredential], credential_exchanger: Optional[BaseAuthCredentialExchanger] = None, + *, + credential_key: Optional[str] = None, ) -> "ToolAuthHandler": """Creates a ToolAuthHandler instance from a ToolContext.""" credential_store = ToolContextCredentialStore(tool_context) @@ -142,8 +222,9 @@ class ToolAuthHandler: tool_context, auth_scheme, auth_credential, - credential_exchanger, - credential_store, + credential_key=credential_key, + credential_exchanger=credential_exchanger, + credential_store=credential_store, ) async def _get_existing_credential( @@ -209,21 +290,11 @@ class ToolAuthHandler: "OAuth2 credentials client_secret is missing." ) - self.tool_context.request_credential( - AuthConfig( - auth_scheme=self.auth_scheme, - raw_auth_credential=self.auth_credential, - ) - ) + self.tool_context.request_credential(self._build_auth_config()) return None def _get_auth_response(self) -> AuthCredential: - return self.tool_context.get_auth_response( - AuthConfig( - auth_scheme=self.auth_scheme, - raw_auth_credential=self.auth_credential, - ) - ) + return self.tool_context.get_auth_response(self._build_auth_config()) def _external_exchange_required(self, credential) -> bool: return ( diff --git a/tests/unittests/auth/test_auth_config.py b/tests/unittests/auth/test_auth_config.py index 3547b7c7..ba87954b 100644 --- a/tests/unittests/auth/test_auth_config.py +++ b/tests/unittests/auth/test_auth_config.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +from pathlib import Path +import subprocess +import sys + from fastapi.openapi.models import OAuth2 from fastapi.openapi.models import OAuthFlowAuthorizationCode from fastapi.openapi.models import OAuthFlows @@ -107,3 +112,53 @@ def test_get_credential_key_with_extras(auth_config): assert original_key == key assert "extra_field" in auth_config.auth_scheme.model_extra assert "extra_field" in auth_config.raw_auth_credential.model_extra + + +def test_credential_key_is_stable_across_python_hash_seed(): + """Test AuthConfig key generation does not depend on PYTHONHASHSEED.""" + repo_root = Path(__file__).resolve().parents[3] + pythonpath = str(repo_root / "src") + code = "\n".join([ + "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", + "", + "auth_scheme = OAuth2(", + " flows=OAuthFlows(", + " authorizationCode=OAuthFlowAuthorizationCode(", + " authorizationUrl='https://example.com/oauth2/authorize',", + " tokenUrl='https://example.com/oauth2/token',", + " scopes={'read': 'Read access'},", + " )", + " )", + ")", + "auth_cred = AuthCredential(", + " auth_type=AuthCredentialTypes.OAUTH2,", + " oauth2=OAuth2Auth(", + " client_id='mock_client_id',", + " client_secret='mock_client_secret',", + " ),", + ")", + "print(AuthConfig(", + " auth_scheme=auth_scheme,", + " raw_auth_credential=auth_cred,", + ").credential_key)", + ]) + + def _run_with_seed(seed: str) -> str: + env = os.environ.copy() + env["PYTHONHASHSEED"] = seed + env["PYTHONPATH"] = os.pathsep.join( + [pythonpath, env.get("PYTHONPATH", "")] + ).strip(os.pathsep) + return subprocess.check_output( + [sys.executable, "-c", code], + env=env, + text=True, + ).strip() + + assert _run_with_seed("0") == _run_with_seed("1") diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index 4261de0e..2faeeb15 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -348,10 +348,12 @@ class TestGenerateAuthRequest: exchanged_auth_credential=oauth2_credentials_with_auth_uri.model_copy( deep=True ), + credential_key="my_tool_tokens", ) handler = AuthHandler(config) result = handler.generate_auth_request() + assert result.credential_key == "my_tool_tokens" assert ( result.exchanged_auth_credential.oauth2.auth_uri == oauth2_credentials_with_auth_uri.oauth2.auth_uri @@ -400,6 +402,31 @@ class TestGenerateAuthRequest: assert mock_generate_auth_uri.called assert result.exchanged_auth_credential == mock_credential + @patch("google.adk.auth.auth_handler.AuthHandler.generate_auth_uri") + def test_preserves_credential_key_on_generated_request( + self, mock_generate_auth_uri, oauth2_auth_scheme, oauth2_credentials + ): + """Test that AuthHandler preserves an explicit credential_key.""" + mock_generate_auth_uri.return_value = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="mock_client_id", + client_secret="mock_client_secret", + auth_uri="https://example.com/generated", + state="generated_state", + ), + ) + + config = AuthConfig( + auth_scheme=oauth2_auth_scheme, + raw_auth_credential=oauth2_credentials, + credential_key="my_tool_tokens", + ) + handler = AuthHandler(config) + result = handler.generate_auth_request() + + assert result.credential_key == "my_tool_tokens" + class TestGetAuthResponse: """Tests for the get_auth_response method.""" diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index 95e01046..7000c9b8 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -21,6 +21,8 @@ from fastapi.openapi.models import OAuth2 from fastapi.openapi.models import OAuthFlowAuthorizationCode from fastapi.openapi.models import OAuthFlowImplicit from fastapi.openapi.models import OAuthFlows +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm_agent import Agent from google.adk.auth.auth_credential import AuthCredential from google.adk.auth.auth_credential import AuthCredentialTypes from google.adk.auth.auth_credential import OAuth2Auth @@ -33,8 +35,12 @@ 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.auth.oauth2_discovery import AuthorizationServerMetadata +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.tools.tool_context import ToolContext import pytest +from .. import testing_utils + class TestCredentialManager: """Test suite for CredentialManager.""" @@ -211,8 +217,30 @@ class TestCredentialManager: @pytest.mark.asyncio async def test_save_credential_with_service(self): """Test _save_credential with credential service.""" - auth_config = Mock(spec=AuthConfig) - mock_credential = Mock(spec=AuthCredential) + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/oauth2/authorize", + tokenUrl="https://example.com/oauth2/token", + scopes={"read": "Read access"}, + ) + ) + ) + auth_config = AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="mock_client_id", + client_secret="mock_client_secret", + ), + ), + credential_key="test_key", + ) + mock_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token="mock_access_token"), + ) # Mock credential service credential_service = AsyncMock() @@ -228,15 +256,39 @@ class TestCredentialManager: manager = CredentialManager(auth_config) await manager._save_credential(tool_context, mock_credential) - tool_context.save_credential.assert_called_once_with(auth_config) - assert auth_config.exchanged_auth_credential == mock_credential + tool_context.save_credential.assert_called_once() + saved_auth_config = tool_context.save_credential.call_args.args[0] + assert saved_auth_config.credential_key == "test_key" + assert saved_auth_config.exchanged_auth_credential == mock_credential + assert auth_config.exchanged_auth_credential is None @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.exchanged_auth_credential = None - mock_credential = Mock(spec=AuthCredential) + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/oauth2/authorize", + tokenUrl="https://example.com/oauth2/token", + scopes={"read": "Read access"}, + ) + ) + ) + auth_config = AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="mock_client_id", + client_secret="mock_client_secret", + ), + ), + credential_key="test_key", + ) + mock_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token="mock_access_token"), + ) # Mock invocation context with no credential service invocation_context = Mock() @@ -248,9 +300,86 @@ class TestCredentialManager: manager = CredentialManager(auth_config) await manager._save_credential(tool_context, mock_credential) - # Should not raise an error, and credential should be set in auth_config - # even when there's no credential service (config is updated regardless) - assert auth_config.exchanged_auth_credential == mock_credential + assert auth_config.exchanged_auth_credential is None + + @pytest.mark.asyncio + async def test_request_credential_does_not_leak_across_users(self): + """Test that user-specific tokens are not cached on shared tool configs.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/oauth2/authorize", + tokenUrl="https://example.com/oauth2/token", + scopes={"read": "Read access"}, + ) + ) + ) + auth_config = AuthConfig( + auth_scheme=auth_scheme, + raw_auth_credential=AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="mock_client_id", + client_secret="mock_client_secret", + ), + ), + credential_key="shared_key", + ) + manager = CredentialManager(auth_config) + agent = Agent( + name="root_agent", + model=testing_utils.MockModel.create(responses=[]), + tools=[], + ) + + session_service = InMemorySessionService() + session_a = await session_service.create_session( + app_name="test_app", + user_id="user_a", + session_id="session_a", + ) + session_b = await session_service.create_session( + app_name="test_app", + user_id="user_b", + session_id="session_b", + ) + + invocation_context_a = InvocationContext( + session_service=session_service, + invocation_id="invocation_a", + agent=agent, + session=session_a, + credential_service=None, + ) + invocation_context_b = InvocationContext( + session_service=session_service, + invocation_id="invocation_b", + agent=agent, + session=session_b, + credential_service=None, + ) + + tool_context_a = ToolContext( + invocation_context_a, function_call_id="call_a" + ) + tool_context_b = ToolContext( + invocation_context_b, function_call_id="call_b" + ) + + tool_context_a.state["temp:" + auth_config.credential_key] = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + auth_uri="https://example.com/oauth2/authorize?x=y", + state="state_a", + access_token="token_a", + expires_at=9_999_999_999, + ), + ) + await manager.get_auth_credential(tool_context_a) + await manager.request_credential(tool_context_b) + + requested = tool_context_b.actions.requested_auth_configs["call_b"] + assert requested.exchanged_auth_credential.oauth2.access_token is None @pytest.mark.asyncio async def test_refresh_credential_oauth2(self): 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 50eb7d11..35914c46 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 @@ -138,6 +138,23 @@ async def test_openid_connect_no_auth_response( assert result.auth_credential == openid_connect_credential +@pytest.mark.asyncio +async def test_openid_connect_uses_explicit_credential_key( + openid_connect_scheme, openid_connect_credential +): + tool_context = create_mock_tool_context() + handler = ToolAuthHandler( + tool_context, + openid_connect_scheme, + openid_connect_credential, + credential_key='my_tool_tokens', + ) + result = await handler.prepare_auth_credentials() + assert result.state == 'pending' + requested = tool_context.actions.requested_auth_configs['test-fc-id'] + assert requested.credential_key == 'my_tool_tokens' + + @pytest.mark.asyncio async def test_openid_connect_with_auth_response( openid_connect_scheme, openid_connect_credential, monkeypatch