From 8083d536d815ced02e2f2de59917942bbc322685 Mon Sep 17 00:00:00 2001 From: Tyler Gibbs Date: Sat, 25 Jul 2026 10:12:14 -0500 Subject: [PATCH] Serialize credential lifecycle by key --- src/google/adk/auth/credential_manager.py | 119 +++++++++++------- .../unittests/auth/test_credential_manager.py | 97 ++++++++++++++ 2 files changed, 172 insertions(+), 44 deletions(-) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index d693f67262d..a8878904fce 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -14,10 +14,13 @@ from __future__ import annotations +import asyncio from collections.abc import Sequence import logging import threading from typing import Optional +from weakref import WeakKeyDictionary +from weakref import WeakValueDictionary from fastapi.openapi.models import OAuth2 @@ -102,6 +105,27 @@ class CredentialManager: _auth_provider_registry = AuthProviderRegistry() _registry_lock = threading.Lock() + _credential_locks_lock = threading.Lock() + _credential_locks: WeakKeyDictionary[ + asyncio.AbstractEventLoop, WeakValueDictionary[str, asyncio.Lock] + ] = WeakKeyDictionary() + + def _get_credential_lock(self) -> asyncio.Lock: + """Returns the event-loop-local lock for this credential.""" + loop = asyncio.get_running_loop() + credential_key = getattr(self._auth_config, "credential_key", None) or ( + f"auth_config:{id(self._auth_config)}" + ) + with self._credential_locks_lock: + locks_for_loop = self._credential_locks.get(loop) + if locks_for_loop is None: + locks_for_loop = WeakValueDictionary() + self._credential_locks[loop] = locks_for_loop + lock = locks_for_loop.get(credential_key) + if lock is None: + lock = asyncio.Lock() + locks_for_loop[credential_key] = lock + return lock @classmethod def register_auth_provider(cls, provider: BaseAuthProvider) -> None: @@ -227,51 +251,58 @@ async def get_auth_credential( return None return provided_credential - # Step 1: Validate credential configuration - await self._validate_credential() - - # Step 2: Check if credential is already ready (no processing needed) - raw_auth_credential = self._auth_config.raw_auth_credential - if self._is_credential_ready() and raw_auth_credential is not None: - # Return a copy to avoid leaking mutations across invocations/users when - # tools share a long-lived AuthConfig instance. - return raw_auth_credential.model_copy(deep=True) - - # Step 3: Try to load existing processed credential - credential = await self._load_existing_credential(context) - - # Step 4: If no existing credential, load from auth response - # TODO instead of load from auth response, we can store auth response in - # credential service. - was_from_auth_response = False - if not credential: - credential = await self._load_from_auth_response(context) - was_from_auth_response = True - - # Step 5: If still no credential available, check if client credentials - if not credential: - # For client credentials flow, use raw credentials directly - if self._is_client_credentials_flow(): - # 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 - - # Step 6: Exchange credential if needed (e.g., service account to access token) - credential, was_exchanged = await self._exchange_credential(credential) + async with self._get_credential_lock(): + # Step 1: Validate credential configuration + await self._validate_credential() + + # Step 2: Check if credential is already ready (no processing needed) + raw_auth_credential = self._auth_config.raw_auth_credential + if self._is_credential_ready() and raw_auth_credential is not None: + # Return a copy to avoid leaking mutations across invocations/users when + # tools share a long-lived AuthConfig instance. + return raw_auth_credential.model_copy(deep=True) + + # Step 3: Try to load existing processed credential + # This is intentionally inside the lock so an invocation that waited can + # observe a credential saved by the previous holder. + credential = await self._load_existing_credential(context) + + # Step 4: If no existing credential, load from auth response + # TODO instead of load from auth response, we can store auth response in + # credential service. + was_from_auth_response = False + if not credential: + credential = await self._load_from_auth_response(context) + was_from_auth_response = True + + # Step 5: If still no credential available, check if client credentials + if not credential: + # For client credentials flow, use raw credentials directly + if self._is_client_credentials_flow(): + # Exchange/refresh steps may mutate the credential object in-place, + # so do not operate on the shared tool config. + raw_credential = self._auth_config.raw_auth_credential + assert raw_credential is not None + credential = raw_credential.model_copy(deep=True) + else: + # For authorization code flow, return None to trigger user + # authorization. + return None + + # Step 6: Exchange credential if needed (e.g., service account to access + # token) + credential, was_exchanged = await self._exchange_credential(credential) + + # Step 7: Refresh credential if expired + was_refreshed = False + if not was_exchanged: + credential, was_refreshed = await self._refresh_credential(credential) + + # Step 8: Save credential if it was modified + if was_from_auth_response or was_exchanged or was_refreshed: + await self._save_credential(context, credential) - # Step 7: Refresh credential if expired - was_refreshed = False - if not was_exchanged: - credential, was_refreshed = await self._refresh_credential(credential) - - # Step 8: Save credential if it was modified - if was_from_auth_response or was_exchanged or was_refreshed: - await self._save_credential(context, credential) - - return credential + return credential async def _load_existing_credential( self, context: CallbackContext diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index 0ba94908978..db4c69835eb 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio from unittest.mock import ANY from unittest.mock import AsyncMock from unittest.mock import Mock @@ -329,6 +330,102 @@ async def test_load_auth_credentials_no_credential(self): assert result is None + @pytest.mark.asyncio + async def test_get_auth_credential_serializes_same_credential_key(self): + """Concurrent managers should process a credential only once.""" + raw_credential = Mock(spec=AuthCredential) + processed_credential = Mock(spec=AuthCredential) + auth_config = Mock( + spec=AuthConfig, + auth_scheme=Mock(spec=AuthScheme), + raw_auth_credential=raw_credential, + credential_key="shared-key", + ) + managers = [CredentialManager(auth_config), CredentialManager(auth_config)] + context = Mock(spec=CallbackContext) + stored_credential = None + exchange_count = 0 + + async def load_existing_credential(_): + return stored_credential + + async def exchange_credential(credential): + nonlocal exchange_count + if credential is processed_credential: + return credential, False + exchange_count += 1 + await asyncio.sleep(0) + return processed_credential, True + + async def save_credential(_, credential): + nonlocal stored_credential + stored_credential = credential + + for manager in managers: + manager._validate_credential = AsyncMock() + manager._is_credential_ready = Mock(return_value=False) + manager._load_existing_credential = AsyncMock( + side_effect=load_existing_credential + ) + manager._load_from_auth_response = AsyncMock(return_value=raw_credential) + manager._exchange_credential = AsyncMock(side_effect=exchange_credential) + manager._refresh_credential = AsyncMock( + side_effect=lambda credential: (credential, False) + ) + manager._save_credential = AsyncMock(side_effect=save_credential) + + results = await asyncio.gather( + *(manager.get_auth_credential(context) for manager in managers) + ) + + assert results == [processed_credential, processed_credential] + assert exchange_count == 1 + + @pytest.mark.asyncio + async def test_get_auth_credential_keeps_different_keys_concurrent(self): + """Credential processing for different keys should not block.""" + both_started = asyncio.Event() + active_count = 0 + max_active_count = 0 + + async def exchange_credential(credential): + nonlocal active_count, max_active_count + active_count += 1 + max_active_count = max(max_active_count, active_count) + if active_count == 2: + both_started.set() + await asyncio.wait_for(both_started.wait(), timeout=1) + active_count -= 1 + return credential, False + + managers = [] + for credential_key in ("key-a", "key-b"): + raw_credential = Mock(spec=AuthCredential) + auth_config = Mock( + spec=AuthConfig, + auth_scheme=Mock(spec=AuthScheme), + raw_auth_credential=raw_credential, + credential_key=credential_key, + ) + manager = CredentialManager(auth_config) + manager._validate_credential = AsyncMock() + manager._is_credential_ready = Mock(return_value=False) + manager._load_existing_credential = AsyncMock(return_value=None) + manager._load_from_auth_response = AsyncMock(return_value=raw_credential) + manager._exchange_credential = AsyncMock(side_effect=exchange_credential) + manager._refresh_credential = AsyncMock( + side_effect=lambda credential: (credential, False) + ) + manager._save_credential = AsyncMock() + managers.append(manager) + + await asyncio.gather(*( + manager.get_auth_credential(Mock(spec=CallbackContext)) + for manager in managers + )) + + assert max_active_count == 2 + @pytest.mark.asyncio async def test_load_existing_credential_already_exchanged(self): """Test _load_existing_credential ignores shared config cache."""