google-adk 1.2.1__py3-none-any.whl → 1.4.1__py3-none-any.whl

google/adk/auth/credential_manager.py

@@ -0,0 +1,261 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import Optional
+
+from ..tools.tool_context import ToolContext
+from ..utils.feature_decorator import experimental
+from .auth_credential import AuthCredential
+from .auth_credential import AuthCredentialTypes
+from .auth_schemes import AuthSchemeType
+from .auth_tool import AuthConfig
+from .exchanger.base_credential_exchanger import BaseCredentialExchanger
+from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry
+from .refresher.base_credential_refresher import BaseCredentialRefresher
+from .refresher.credential_refresher_registry import CredentialRefresherRegistry
+
+
+@experimental
+class CredentialManager:
+  """Manages authentication credentials through a structured workflow.
+
+  The CredentialManager orchestrates the complete lifecycle of authentication
+  credentials, from initial loading to final preparation for use. It provides
+  a centralized interface for handling various credential types and authentication
+  schemes while maintaining proper credential hygiene (refresh, exchange, caching).
+
+  This class is only for use by Agent Development Kit.
+
+  Args:
+      auth_config: Configuration containing authentication scheme and credentials
+
+  Example:
+      ```python
+      auth_config = AuthConfig(
+          auth_scheme=oauth2_scheme,
+          raw_auth_credential=service_account_credential
+      )
+      manager = CredentialManager(auth_config)
+
+      # Register custom exchanger if needed
+      manager.register_credential_exchanger(
+          AuthCredentialTypes.CUSTOM_TYPE,
+          CustomCredentialExchanger()
+      )
+
+      # Register custom refresher if needed
+      manager.register_credential_refresher(
+          AuthCredentialTypes.CUSTOM_TYPE,
+          CustomCredentialRefresher()
+      )
+
+      # Load and prepare credential
+      credential = await manager.load_auth_credential(tool_context)
+      ```
+  """
+
+  def __init__(
+      self,
+      auth_config: AuthConfig,
+  ):
+    self._auth_config = auth_config
+    self._exchanger_registry = CredentialExchangerRegistry()
+    self._refresher_registry = CredentialRefresherRegistry()
+
+    # Register default exchangers and refreshers
+    # TODO: support service account credential exchanger
+    from .refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
+
+    oauth2_refresher = OAuth2CredentialRefresher()
+    self._refresher_registry.register(
+        AuthCredentialTypes.OAUTH2, oauth2_refresher
+    )
+    self._refresher_registry.register(
+        AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_refresher
+    )
+
+  def register_credential_exchanger(
+      self,
+      credential_type: AuthCredentialTypes,
+      exchanger_instance: BaseCredentialExchanger,
+  ) -> None:
+    """Register a credential exchanger for a credential type.
+
+    Args:
+        credential_type: The credential type to register for.
+        exchanger_instance: The exchanger instance to register.
+    """
+    self._exchanger_registry.register(credential_type, exchanger_instance)
+
+  async def request_credential(self, tool_context: ToolContext) -> None:
+    tool_context.request_credential(self._auth_config)
+
+  async def get_auth_credential(
+      self, tool_context: ToolContext
+  ) -> Optional[AuthCredential]:
+    """Load and prepare authentication credential through a structured workflow."""
+
+    # Step 1: Validate credential configuration
+    await self._validate_credential()
+
+    # Step 2: Check if credential is already ready (no processing needed)
+    if self._is_credential_ready():
+      return self._auth_config.raw_auth_credential
+
+    # Step 3: Try to load existing processed credential
+    credential = await self._load_existing_credential(tool_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(tool_context)
+      was_from_auth_response = True
+
+    # Step 5: If still no credential available, return None
+    if not credential:
+      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
+    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(tool_context, credential)
+
+    return credential
+
+  async def _load_existing_credential(
+      self, tool_context: ToolContext
+  ) -> Optional[AuthCredential]:
+    """Load existing credential from credential service or cached exchanged credential."""
+
+    # Try loading from credential service first
+    credential = await self._load_from_credential_service(tool_context)
+    if credential:
+      return credential
+
+    # Check if we have a cached exchanged credential
+    if self._auth_config.exchanged_auth_credential:
+      return self._auth_config.exchanged_auth_credential
+
+    return None
+
+  async def _load_from_credential_service(
+      self, tool_context: ToolContext
+  ) -> Optional[AuthCredential]:
+    """Load credential from credential service if available."""
+    credential_service = tool_context._invocation_context.credential_service
+    if credential_service:
+      # Note: This should be made async in a future refactor
+      # For now, assuming synchronous operation
+      return await credential_service.load_credential(
+          self._auth_config, tool_context
+      )
+    return None
+
+  async def _load_from_auth_response(
+      self, tool_context: ToolContext
+  ) -> Optional[AuthCredential]:
+    """Load credential from auth response in tool context."""
+    return tool_context.get_auth_response(self._auth_config)
+
+  async def _exchange_credential(
+      self, credential: AuthCredential
+  ) -> tuple[AuthCredential, bool]:
+    """Exchange credential if needed and return the credential and whether it was exchanged."""
+    exchanger = self._exchanger_registry.get_exchanger(credential.auth_type)
+    if not exchanger:
+      return credential, False
+
+    exchanged_credential = await exchanger.exchange(
+        credential, self._auth_config.auth_scheme
+    )
+    return exchanged_credential, True
+
+  async def _refresh_credential(
+      self, credential: AuthCredential
+  ) -> tuple[AuthCredential, bool]:
+    """Refresh credential if expired and return the credential and whether it was refreshed."""
+    refresher = self._refresher_registry.get_refresher(credential.auth_type)
+    if not refresher:
+      return credential, False
+
+    if await refresher.is_refresh_needed(
+        credential, self._auth_config.auth_scheme
+    ):
+      refreshed_credential = await refresher.refresh(
+          credential, self._auth_config.auth_scheme
+      )
+      return refreshed_credential, True
+
+    return credential, False
+
+  def _is_credential_ready(self) -> bool:
+    """Check if credential is ready to use without further processing."""
+    raw_credential = self._auth_config.raw_auth_credential
+    if not raw_credential:
+      return False
+
+    # Simple credentials that don't need exchange or refresh
+    return raw_credential.auth_type in (
+        AuthCredentialTypes.API_KEY,
+        AuthCredentialTypes.HTTP,
+        # Add other simple auth types as needed
+    )
+
+  async def _validate_credential(self) -> None:
+    """Validate credential configuration and raise errors if invalid."""
+    if not self._auth_config.raw_auth_credential:
+      if self._auth_config.auth_scheme.type_ in (
+          AuthSchemeType.oauth2,
+          AuthSchemeType.openIdConnect,
+      ):
+        raise ValueError(
+            "raw_auth_credential is required for auth_scheme type "
+            f"{self._auth_config.auth_scheme.type_}"
+        )
+
+    raw_credential = self._auth_config.raw_auth_credential
+    if raw_credential:
+      if (
+          raw_credential.auth_type
+          in (
+              AuthCredentialTypes.OAUTH2,
+              AuthCredentialTypes.OPEN_ID_CONNECT,
+          )
+          and not raw_credential.oauth2
+      ):
+        raise ValueError(
+            "auth_config.raw_credential.oauth2 required for credential type "
+            f"{raw_credential.auth_type}"
+        )
+        # Additional validation can be added here
+
+  async def _save_credential(
+      self, tool_context: ToolContext, credential: AuthCredential
+  ) -> None:
+    """Save credential to credential service if available."""
+    credential_service = tool_context._invocation_context.credential_service
+    if credential_service:
+      # Update the exchanged credential in config
+      self._auth_config.exchanged_auth_credential = credential
+      await credential_service.save_credential(self._auth_config, tool_context)

google/adk/auth/credential_service/__init__.py

@@ -0,0 +1,13 @@
+# 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.

google/adk/auth/credential_service/base_credential_service.py

@@ -0,0 +1,75 @@
+# 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.
+
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from typing import Optional
+
+from ...tools.tool_context import ToolContext
+from ...utils.feature_decorator import experimental
+from ..auth_credential import AuthCredential
+from ..auth_tool import AuthConfig
+
+
+@experimental
+class BaseCredentialService(ABC):
+  """Abstract class for Service that loads / saves tool credentials from / to
+  the backend credential store."""
+
+  @abstractmethod
+  async def load_credential(
+      self,
+      auth_config: AuthConfig,
+      tool_context: ToolContext,
+  ) -> Optional[AuthCredential]:
+    """
+    Loads the credential by auth config and current tool context from the
+    backend credential store.
+
+    Args:
+        auth_config: The auth config which contains the auth scheme and auth
+        credential information. auth_config.get_credential_key will be used to
+        build the key to load the credential.
+
+        tool_context: The context of the current invocation when the tool is
+        trying to load the credential.
+
+    Returns:
+        Optional[AuthCredential]: the credential saved in the store.
+
+    """
+
+  @abstractmethod
+  async def save_credential(
+      self,
+      auth_config: AuthConfig,
+      tool_context: ToolContext,
+  ) -> None:
+    """
+    Saves the exchanged_auth_credential in auth config to the backend credential
+    store.
+
+    Args:
+        auth_config: The auth config which contains the auth scheme and auth
+        credential information. auth_config.get_credential_key will be used to
+        build the key to save the credential.
+
+        tool_context: The context of the current invocation when the tool is
+        trying to save the credential.
+
+    Returns:
+        None
+    """

google/adk/auth/credential_service/in_memory_credential_service.py

@@ -0,0 +1,64 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import Optional
+
+from typing_extensions import override
+
+from ...tools.tool_context import ToolContext
+from ...utils.feature_decorator import experimental
+from ..auth_credential import AuthCredential
+from ..auth_tool import AuthConfig
+from .base_credential_service import BaseCredentialService
+
+
+@experimental
+class InMemoryCredentialService(BaseCredentialService):
+  """Class for in memory implementation of credential service(Experimental)"""
+
+  def __init__(self):
+    super().__init__()
+    self._credentials = {}
+
+  @override
+  async def load_credential(
+      self,
+      auth_config: AuthConfig,
+      tool_context: ToolContext,
+  ) -> Optional[AuthCredential]:
+    credential_bucket = self._get_bucket_for_current_context(tool_context)
+    return credential_bucket.get(auth_config.credential_key)
+
+  @override
+  async def save_credential(
+      self,
+      auth_config: AuthConfig,
+      tool_context: ToolContext,
+  ) -> None:
+    credential_bucket = self._get_bucket_for_current_context(tool_context)
+    credential_bucket[auth_config.credential_key] = (
+        auth_config.exchanged_auth_credential
+    )
+
+  def _get_bucket_for_current_context(self, tool_context: ToolContext) -> str:
+    app_name = tool_context._invocation_context.app_name
+    user_id = tool_context._invocation_context.user_id
+
+    if app_name not in self._credentials:
+      self._credentials[app_name] = {}
+    if user_id not in self._credentials[app_name]:
+      self._credentials[app_name][user_id] = {}
+    return self._credentials[app_name][user_id]

google/adk/auth/exchanger/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+
+"""Credential exchanger module."""
+
+from .base_credential_exchanger import BaseCredentialExchanger
+
+__all__ = [
+    "BaseCredentialExchanger",
+]

google/adk/auth/exchanger/base_credential_exchanger.py

@@ -0,0 +1,57 @@
+# 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.
+
+"""Base credential exchanger interface."""
+
+from __future__ import annotations
+
+import abc
+from typing import Optional
+
+from ...utils.feature_decorator import experimental
+from ..auth_credential import AuthCredential
+from ..auth_schemes import AuthScheme
+
+
+class CredentialExchangError(Exception):
+  """Base exception for credential exchange errors."""
+
+
+@experimental
+class BaseCredentialExchanger(abc.ABC):
+  """Base interface for credential exchangers.
+
+  Credential exchangers are responsible for exchanging credentials from
+  one format or scheme to another.
+  """
+
+  @abc.abstractmethod
+  async def exchange(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> AuthCredential:
+    """Exchange credential if needed.
+
+    Args:
+        auth_credential: The credential to exchange.
+        auth_scheme: The authentication scheme (optional, some exchangers don't need it).
+
+    Returns:
+        The exchanged credential.
+
+    Raises:
+        CredentialExchangError: If credential exchange fails.
+    """
+    pass

google/adk/auth/exchanger/credential_exchanger_registry.py

@@ -0,0 +1,58 @@
+# 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.
+
+"""Credential exchanger registry."""
+
+from __future__ import annotations
+
+from typing import Dict
+from typing import Optional
+
+from ...utils.feature_decorator import experimental
+from ..auth_credential import AuthCredentialTypes
+from .base_credential_exchanger import BaseCredentialExchanger
+
+
+@experimental
+class CredentialExchangerRegistry:
+  """Registry for credential exchanger instances."""
+
+  def __init__(self):
+    self._exchangers: Dict[AuthCredentialTypes, BaseCredentialExchanger] = {}
+
+  def register(
+      self,
+      credential_type: AuthCredentialTypes,
+      exchanger_instance: BaseCredentialExchanger,
+  ) -> None:
+    """Register an exchanger instance for a credential type.
+
+    Args:
+        credential_type: The credential type to register for.
+        exchanger_instance: The exchanger instance to register.
+    """
+    self._exchangers[credential_type] = exchanger_instance
+
+  def get_exchanger(
+      self, credential_type: AuthCredentialTypes
+  ) -> Optional[BaseCredentialExchanger]:
+    """Get the exchanger instance for a credential type.
+
+    Args:
+        credential_type: The credential type to get exchanger for.
+
+    Returns:
+        The exchanger instance if registered, None otherwise.
+    """
+    return self._exchangers.get(credential_type)

google/adk/auth/exchanger/oauth2_credential_exchanger.py

@@ -0,0 +1,104 @@
+# 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.
+
+"""OAuth2 credential exchanger implementation."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_schemes import AuthScheme
+from google.adk.auth.auth_schemes import OAuthGrantType
+from google.adk.auth.oauth2_credential_util import create_oauth2_session
+from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
+from google.adk.utils.feature_decorator import experimental
+from typing_extensions import override
+
+from .base_credential_exchanger import BaseCredentialExchanger
+from .base_credential_exchanger import CredentialExchangError
+
+try:
+  from authlib.integrations.requests_client import OAuth2Session
+
+  AUTHLIB_AVIALABLE = True
+except ImportError:
+  AUTHLIB_AVIALABLE = False
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class OAuth2CredentialExchanger(BaseCredentialExchanger):
+  """Exchanges OAuth2 credentials from authorization responses."""
+
+  @override
+  async def exchange(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> AuthCredential:
+    """Exchange OAuth2 credential from authorization response.
+    if credential exchange failed, the original credential will be returned.
+
+    Args:
+        auth_credential: The OAuth2 credential to exchange.
+        auth_scheme: The OAuth2 authentication scheme.
+
+    Returns:
+        The exchanged credential with access token.
+
+    Raises:
+        CredentialExchangError: If auth_scheme is missing.
+    """
+    if not auth_scheme:
+      raise CredentialExchangError(
+          "auth_scheme is required for OAuth2 credential exchange"
+      )
+
+    if not AUTHLIB_AVIALABLE:
+      # If authlib is not available, we cannot exchange the credential.
+      # We return the original credential without exchange.
+      # The client using this tool can decide to exchange the credential
+      # themselves using other lib.
+      logger.warning(
+          "authlib is not available, skipping OAuth2 credential exchange."
+      )
+      return auth_credential
+
+    if auth_credential.oauth2 and auth_credential.oauth2.access_token:
+      return auth_credential
+
+    client, token_endpoint = create_oauth2_session(auth_scheme, auth_credential)
+    if not client:
+      logger.warning("Could not create OAuth2 session for token exchange")
+      return auth_credential
+
+    try:
+      tokens = client.fetch_token(
+          token_endpoint,
+          authorization_response=auth_credential.oauth2.auth_response_uri,
+          code=auth_credential.oauth2.auth_code,
+          grant_type=OAuthGrantType.AUTHORIZATION_CODE,
+      )
+      update_credential_with_tokens(auth_credential, tokens)
+      logger.debug("Successfully exchanged OAuth2 tokens")
+    except Exception as e:
+      # TODO reconsider whether we should raise errors in this case
+      logger.error("Failed to exchange OAuth2 tokens: %s", e)
+      # Return original credential on failure
+      return auth_credential
+
+    return auth_credential