Merge branch 'main' into fix/handle-circular-refs

Author: foivos-all

src/google/adk/auth/auth_credential.py

@@ -25,6 +25,7 @@
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import model_validator
 
 
 class BaseModelWithConfig(BaseModel):
@@ -145,11 +146,45 @@ class ServiceAccountCredential(BaseModelWithConfig):
 
 
 class ServiceAccount(BaseModelWithConfig):
-  """Represents Google Service Account configuration."""
+  """Represents Google Service Account configuration.
+
+  Attributes:
+    service_account_credential: The service account credential (JSON key).
+    scopes: The OAuth2 scopes to request. Optional; when omitted with
+        ``use_default_credential=True``, defaults to the cloud-platform scope.
+    use_default_credential: Whether to use Application Default Credentials.
+    use_id_token: Whether to exchange for an ID token instead of an access
+        token. Required for service-to-service authentication with Cloud Run,
+        Cloud Functions, and other Google Cloud services that require identity
+        verification. When True, ``audience`` must also be set.
+    audience: The target audience for the ID token, typically the URL of the
+        receiving service (e.g. ``https://my-service-xyz.run.app``). Required
+        when ``use_id_token`` is True.
+  """
 
   service_account_credential: Optional[ServiceAccountCredential] = None
-  scopes: List[str]
+  scopes: Optional[List[str]] = None
   use_default_credential: Optional[bool] = False
+  use_id_token: Optional[bool] = False
+  audience: Optional[str] = None
+
+  @model_validator(mode="after")
+  def _validate_config(self) -> ServiceAccount:
+    if (
+        not self.use_default_credential
+        and self.service_account_credential is None
+    ):
+      raise ValueError(
+          "service_account_credential is required when"
+          " use_default_credential is False."
+      )
+    if self.use_id_token and not self.audience:
+      raise ValueError(
+          "audience is required when use_id_token is True. Set it to the"
+          " URL of the target service"
+          " (e.g. 'https://my-service.run.app')."
+      )
+    return self
 
 
 class AuthCredentialTypes(str, Enum):

src/google/adk/errors/session_not_found_error.py

@@ -14,14 +14,11 @@
 
 from __future__ import annotations
 
-from .not_found_error import NotFoundError
 
-
-class SessionNotFoundError(ValueError, NotFoundError):
+class SessionNotFoundError(ValueError):
   """Raised when a session cannot be found.
 
-  Inherits from both ValueError (for backward compatibility) and NotFoundError
-  (for semantic consistency with the project's error hierarchy).
+  Inherits from ValueError (for backward compatibility).
   """
 
   def __init__(self, message="Session not found."):

src/google/adk/models/google_llm.py

@@ -85,6 +85,23 @@ class Gemini(BaseLlm):
 
   Attributes:
     model: The name of the Gemini model.
+    client: An optional preconfigured ``google.genai.Client`` instance.
+      When provided, ADK uses this client for all API calls instead of
+      creating one internally from environment variables or ADC.  This
+      allows fine-grained control over authentication, project, location,
+      and other client-level settings — and enables running agents that
+      target different Vertex AI regions within the same process.
+
+      Example::
+
+        from google import genai
+        from google.adk.models import Gemini
+
+        client = genai.Client(
+            vertexai=True, project="my-project", location="us-central1"
+        )
+        model = Gemini(model="gemini-2.5-flash", client=client)
+
     use_interactions_api: Whether to use the interactions API for model
       invocation.
   """
@@ -131,6 +148,35 @@ class Gemini(BaseLlm):
   ```
   """
 
+  def __init__(self, *, client: Optional[Client] = None, **kwargs: Any):
+    """Initialises a Gemini model wrapper.
+
+    Args:
+      client: An optional preconfigured ``google.genai.Client``.  When
+        provided, ADK uses this client for **all** Gemini API calls
+        (including the Live API) instead of creating one internally.
+
+        .. note::
+          When a custom client is supplied it is used as-is for Live API
+          connections.  ADK will **not** override the client's
+          ``api_version``; you are responsible for setting the correct
+          version (``v1beta1`` for Vertex AI, ``v1alpha`` for the
+          Gemini developer API) on the client yourself.
+
+        .. warning::
+          ``google.genai.Client`` contains threading primitives that
+          cannot be pickled.  If you are deploying to Agent Engine (or
+          any environment that serialises the model), do **not** pass a
+          custom client — let ADK create one from the environment
+          instead.
+
+      **kwargs: Forwarded to the Pydantic ``BaseLlm`` constructor
+        (``model``, ``base_url``, ``retry_options``, etc.).
+    """
+    super().__init__(**kwargs)
+    # Store after super().__init__ so Pydantic validation runs first.
+    object.__setattr__(self, '_client', client)
+
   @classmethod
   @override
   def supported_models(cls) -> list[str]:
@@ -299,9 +345,16 @@ async def _generate_content_via_interactions(
   def api_client(self) -> Client:
     """Provides the api client.
 
+    If a preconfigured ``client`` was passed to the constructor it is
+    returned directly; otherwise a new ``Client`` is created using the
+    default environment/ADC configuration.
+
     Returns:
       The api client.
     """
+    if self._client is not None:
+      return self._client
+
     from google.genai import Client
 
     return Client(
@@ -334,6 +387,9 @@ def _live_api_version(self) -> str:
 
   @cached_property
   def _live_api_client(self) -> Client:
+    if self._client is not None:
+      return self._client
+
     from google.genai import Client
 
     return Client(

src/google/adk/tools/openapi_tool/auth/credential_exchangers/service_account_exchanger.py

@@ -19,6 +19,7 @@
 from typing import Optional
 
 import google.auth
+from google.auth import exceptions as google_auth_exceptions
 from google.auth.transport.requests import Request
 from google.oauth2 import service_account
 import google.oauth2.credentials
@@ -27,6 +28,7 @@
 from .....auth.auth_credential import AuthCredentialTypes
 from .....auth.auth_credential import HttpAuth
 from .....auth.auth_credential import HttpCredentials
+from .....auth.auth_credential import ServiceAccount
 from .....auth.auth_schemes import AuthScheme
 from .base_credential_exchanger import AuthCredentialMissingError
 from .base_credential_exchanger import BaseAuthCredentialExchanger
@@ -38,59 +40,142 @@ class ServiceAccountCredentialExchanger(BaseAuthCredentialExchanger):
   Uses the default service credential if `use_default_credential = True`.
   Otherwise, uses the service account credential provided in the auth
   credential.
+
+  Supports exchanging for either an access token (default) or an ID token
+  when ``ServiceAccount.use_id_token`` is True. ID tokens are required for
+  service-to-service authentication with Cloud Run, Cloud Functions, and
+  other services that verify caller identity.
   """
 
   def exchange_credential(
       self,
       auth_scheme: AuthScheme,
       auth_credential: Optional[AuthCredential] = None,
   ) -> AuthCredential:
-    """Exchanges the service account auth credential for an access token.
+    """Exchanges the service account auth credential for a token.
 
     If auth_credential contains a service account credential, it will be used
-    to fetch an access token. Otherwise, the default service credential will be
-    used for fetching an access token.
+    to fetch a token. Otherwise, the default service credential will be
+    used for fetching a token.
+
+    When ``service_account.use_id_token`` is True, an ID token is fetched
+    using the configured ``audience``. This is required for authenticating
+    to Cloud Run, Cloud Functions, and similar services.
 
     Args:
         auth_scheme: The auth scheme.
         auth_credential: The auth credential.
 
     Returns:
-        An AuthCredential in HTTPBearer format, containing the access token.
+        An AuthCredential in HTTPBearer format, containing the token.
     """
-    if (
-        auth_credential is None
-        or auth_credential.service_account is None
-        or (
-            auth_credential.service_account.service_account_credential is None
-            and not auth_credential.service_account.use_default_credential
-        )
-    ):
+    if auth_credential is None or auth_credential.service_account is None:
       raise AuthCredentialMissingError(
-          "Service account credentials are missing. Please provide them, or set"
-          " `use_default_credential = True` to use application default"
+          "Service account credentials are missing. Please provide them, or"
+          " set `use_default_credential = True` to use application default"
           " credential in a hosted service like Cloud Run."
       )
 
+    sa_config = auth_credential.service_account
+
+    if sa_config.use_id_token:
+      return self._exchange_for_id_token(sa_config)
+
+    return self._exchange_for_access_token(sa_config)
+
+  def _exchange_for_id_token(self, sa_config: ServiceAccount) -> AuthCredential:
+    """Exchanges the service account credential for an ID token.
+
+    Args:
+        sa_config: The service account configuration.
+
+    Returns:
+        An AuthCredential in HTTPBearer format containing the ID token.
+
+    Raises:
+        AuthCredentialMissingError: If token exchange fails.
+    """
+    # audience and credential presence are validated by the ServiceAccount
+    # model_validator at construction time.
     try:
-      if auth_credential.service_account.use_default_credential:
-        credentials, project_id = google.auth.default(
-            scopes=["https://www.googleapis.com/auth/cloud-platform"],
+      if sa_config.use_default_credential:
+        from google.oauth2 import id_token as oauth2_id_token
+
+        request = Request()
+        token = oauth2_id_token.fetch_id_token(request, sa_config.audience)
+      else:
+        # Guaranteed non-None by ServiceAccount model_validator.
+        assert sa_config.service_account_credential is not None
+        credentials = (
+            service_account.IDTokenCredentials.from_service_account_info(
+                sa_config.service_account_credential.model_dump(),
+                target_audience=sa_config.audience,
+            )
         )
-        quota_project_id = (
-            getattr(credentials, "quota_project_id", None) or project_id
+        credentials.refresh(Request())
+        token = credentials.token
+
+      return AuthCredential(
+          auth_type=AuthCredentialTypes.HTTP,
+          http=HttpAuth(
+              scheme="bearer",
+              credentials=HttpCredentials(token=token),
+          ),
+      )
+
+    # ValueError is raised by google-auth when service account JSON is
+    # missing required fields (e.g. client_email, private_key), or when
+    # fetch_id_token cannot determine credentials from the environment.
+    except (google_auth_exceptions.GoogleAuthError, ValueError) as e:
+      raise AuthCredentialMissingError(
+          f"Failed to exchange service account for ID token: {e}"
+      ) from e
+
+  def _exchange_for_access_token(
+      self, sa_config: ServiceAccount
+  ) -> AuthCredential:
+    """Exchanges the service account credential for an access token.
+
+    Args:
+        sa_config: The service account configuration.
+
+    Returns:
+        An AuthCredential in HTTPBearer format containing the access token.
+
+    Raises:
+        AuthCredentialMissingError: If scopes are missing for explicit
+            credentials or token exchange fails.
+    """
+    if not sa_config.use_default_credential and not sa_config.scopes:
+      raise AuthCredentialMissingError(
+          "scopes are required when using explicit service account credentials"
+          " for access token exchange."
+      )
+
+    try:
+      if sa_config.use_default_credential:
+        scopes = (
+            sa_config.scopes
+            if sa_config.scopes
+            else ["https://www.googleapis.com/auth/cloud-platform"]
+        )
+        credentials, project_id = google.auth.default(
+            scopes=scopes,
         )
+        quota_project_id = credentials.quota_project_id or project_id
       else:
-        config = auth_credential.service_account
+        # Guaranteed non-None by ServiceAccount model_validator.
+        assert sa_config.service_account_credential is not None
         credentials = service_account.Credentials.from_service_account_info(
-            config.service_account_credential.model_dump(), scopes=config.scopes
+            sa_config.service_account_credential.model_dump(),
+            scopes=sa_config.scopes,
         )
         quota_project_id = None
 
       credentials.refresh(Request())
 
-      updated_credential = AuthCredential(
-          auth_type=AuthCredentialTypes.HTTP,  # Store as a bearer token
+      return AuthCredential(
+          auth_type=AuthCredentialTypes.HTTP,
           http=HttpAuth(
               scheme="bearer",
               credentials=HttpCredentials(token=credentials.token),
@@ -101,9 +186,10 @@ def exchange_credential(
               else None,
           ),
       )
-      return updated_credential
 
-    except Exception as e:
+    # ValueError is raised by google-auth when service account JSON is
+    # missing required fields (e.g. client_email, private_key).
+    except (google_auth_exceptions.GoogleAuthError, ValueError) as e:
       raise AuthCredentialMissingError(
           f"Failed to exchange service account token: {e}"
       ) from e