fix: use OAuth2 client-credentials scheme for OpenAPI SA helpers

Merge #6660

Change service-account OpenAPI helpers to return an OAuth2 client-credentials scheme so CredentialManager can perform token exchange.

Also, bypass the credential service caching for all SERVICE_ACCOUNT credentials. This ensures we don't cache exchanged tokens that cannot be refreshed, but means token exchange will run on each tool execution if the manager/exchanger is not reused.

Fixes #6656

PiperOrigin-RevId: 966305100

Author: a2105z

scripts/compliance_checks.py

@@ -40,7 +40,6 @@
     'src/google/adk/tools/_google_credentials.py',
     'src/google/adk/tools/apihub_tool/clients/apihub_client.py',
     'src/google/adk/tools/google_api_tool/google_api_toolset.py',
-    'src/google/adk/tools/openapi_tool/auth/auth_helpers.py',
     'tests/unittests/auth/test_credential_manager.py',
     'tests/unittests/flows/llm_flows/test_functions_request_euc.py',
     'tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py',

src/google/adk/auth/credential_manager.py

@@ -238,7 +238,12 @@ async def get_auth_credential(
       return raw_auth_credential.model_copy(deep=True)
 
     # Step 3: Try to load existing processed credential
-    credential = await self._load_existing_credential(context)
+    credential = None
+    if not (
+        raw_auth_credential
+        and raw_auth_credential.auth_type == AuthCredentialTypes.SERVICE_ACCOUNT
+    ):
+      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
@@ -269,7 +274,12 @@ async def get_auth_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)
+      if not (
+          raw_auth_credential
+          and raw_auth_credential.auth_type
+          == AuthCredentialTypes.SERVICE_ACCOUNT
+      ):
+        await self._save_credential(context, credential)
 
     return credential
 

src/google/adk/tools/openapi_tool/auth/auth_helpers.py

@@ -26,6 +26,8 @@
 from fastapi.openapi.models import HTTPBase
 from fastapi.openapi.models import HTTPBearer
 from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowClientCredentials
+from fastapi.openapi.models import OAuthFlows
 from fastapi.openapi.models import OpenIdConnect
 from fastapi.openapi.models import Schema
 import httpx
@@ -152,13 +154,40 @@ def token_to_scheme_credential(
     raise ValueError(f"Invalid security scheme type: {type}")
 
 
+def _service_account_auth_scheme() -> OAuth2:
+  """Auth scheme for Google Service Account credentials.
+
+  CredentialManager only auto-loads raw non-interactive credentials when the
+  scheme is an OAuth2/OIDC client-credentials flow. An HTTPBearer scheme makes
+  ``_is_client_credentials_flow`` return False, so ``get_auth_credential``
+  returns None and the tool falls back to ``adk_request_credential`` instead of
+  exchanging the service account for a token.
+
+  The token URL is unused by ServiceAccountCredentialExchanger (ADC / JWT
+  assertion), but is required by the OAuth2 client-credentials model.
+  """
+  return OAuth2(
+      flows=OAuthFlows(
+          clientCredentials=OAuthFlowClientCredentials(
+              # Placeholder only; SA exchange does not call this endpoint.
+              # Use the mTLS host form for compliance with Google API endpoint
+              # requirements.
+              tokenUrl="https://oauth2.mtls.googleapis.com/token",
+              scopes={},
+          )
+      )
+  )
+
+
 def service_account_dict_to_scheme_credential(
     config: Dict[str, Any],
     scopes: List[str],
 ) -> Tuple[AuthScheme, AuthCredential]:
   """Creates AuthScheme and AuthCredential for Google Service Account.
 
-  Returns a bearer token scheme, and a service account credential.
+  Returns an OAuth2 client-credentials scheme (so CredentialManager can
+  exchange the service account) and a service account credential. After
+  exchange the credential is an HTTP bearer token.
 
   Args:
       config: A ServiceAccount object containing the Google Service Account
@@ -168,7 +197,6 @@ def service_account_dict_to_scheme_credential(
   Returns:
       Tuple: (AuthScheme, AuthCredential)
   """
-  auth_scheme = HTTPBearer(bearerFormat="JWT")
   service_account = ServiceAccount(
       service_account_credential=ServiceAccountCredential.model_construct(
           **config
@@ -179,15 +207,17 @@ def service_account_dict_to_scheme_credential(
       auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
       service_account=service_account,
   )
-  return auth_scheme, auth_credential
+  return _service_account_auth_scheme(), auth_credential
 
 
 def service_account_scheme_credential(
     config: ServiceAccount,
 ) -> Tuple[AuthScheme, AuthCredential]:
   """Creates AuthScheme and AuthCredential for Google Service Account.
 
-  Returns a bearer token scheme, and a service account credential.
+  Returns an OAuth2 client-credentials scheme (so CredentialManager can
+  exchange the service account) and a service account credential. After
+  exchange the credential is an HTTP bearer token.
 
   Args:
       config: A ServiceAccount object containing the Google Service Account
@@ -196,11 +226,10 @@ def service_account_scheme_credential(
   Returns:
       Tuple: (AuthScheme, AuthCredential)
   """
-  auth_scheme = HTTPBearer(bearerFormat="JWT")
   auth_credential = AuthCredential(
       auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, service_account=config
   )
-  return auth_scheme, auth_credential
+  return _service_account_auth_scheme(), auth_credential
 
 
 def openid_dict_to_scheme_credential(

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

@@ -16,10 +16,14 @@
 
 from __future__ import annotations
 
+import calendar
+import time
+from typing import Any
 from typing import Optional
 
 import google.auth
 from google.auth import exceptions as google_auth_exceptions
+from google.auth import jwt
 from google.auth.transport.requests import Request
 from google.oauth2 import service_account
 import google.oauth2.credentials
@@ -33,6 +37,38 @@
 from .base_credential_exchanger import AuthCredentialMissingError
 from .base_credential_exchanger import BaseAuthCredentialExchanger
 
+_access_token_cache: dict[tuple[Any, ...], tuple[AuthCredential, float]] = {}
+_id_token_cache: dict[tuple[Any, ...], tuple[AuthCredential, float]] = {}
+
+
+def _get_cache_key(sa_config: ServiceAccount) -> tuple[Any, ...]:
+  scopes_tuple = tuple(sa_config.scopes) if sa_config.scopes else ()
+  if sa_config.use_default_credential:
+    return (
+        True,
+        scopes_tuple,
+        sa_config.use_id_token,
+        sa_config.audience,
+    )
+  else:
+    cred = sa_config.service_account_credential
+    cred_id = cred.private_key_id if cred else None
+    client_email = cred.client_email if cred else None
+    return (
+        False,
+        cred_id,
+        client_email,
+        scopes_tuple,
+        sa_config.use_id_token,
+        sa_config.audience,
+    )
+
+
+def _reset_cache():
+  global _access_token_cache, _id_token_cache
+  _access_token_cache.clear()
+  _id_token_cache.clear()
+
 
 class ServiceAccountCredentialExchanger(BaseAuthCredentialExchanger):
   """Fetches credentials for Google Service Account.
@@ -95,6 +131,13 @@ def _exchange_for_id_token(self, sa_config: ServiceAccount) -> AuthCredential:
     Raises:
         AuthCredentialMissingError: If token exchange fails.
     """
+    cache_key = _get_cache_key(sa_config)
+    cached_val = _id_token_cache.get(cache_key)
+    if cached_val:
+      token, expires_at = cached_val
+      if time.time() < expires_at - 300:
+        return token
+
     # audience and credential presence are validated by the ServiceAccount
     # model_validator at construction time.
     try:
@@ -103,6 +146,11 @@ def _exchange_for_id_token(self, sa_config: ServiceAccount) -> AuthCredential:
 
         request = Request()
         token = oauth2_id_token.fetch_id_token(request, sa_config.audience)
+        try:
+          decoded = jwt.decode(token, verify=False)
+          expires_at = decoded.get("exp") or int(time.time() + 3600)
+        except Exception:  # pylint: disable=broad-except
+          expires_at = int(time.time() + 3600)
       else:
         # Guaranteed non-None by ServiceAccount model_validator.
         assert sa_config.service_account_credential is not None
@@ -114,14 +162,24 @@ def _exchange_for_id_token(self, sa_config: ServiceAccount) -> AuthCredential:
         )
         credentials.refresh(Request())
         token = credentials.token
-
-      return AuthCredential(
+        try:
+          expires_at = (
+              calendar.timegm(credentials.expiry.utctimetuple())
+              if credentials.expiry
+              else int(time.time() + 3600)
+          )
+        except (AttributeError, TypeError, ValueError):
+          expires_at = int(time.time() + 3600)
+
+      res = AuthCredential(
           auth_type=AuthCredentialTypes.HTTP,
           http=HttpAuth(
               scheme="bearer",
               credentials=HttpCredentials(token=token),
           ),
       )
+      _id_token_cache[cache_key] = (res, expires_at)
+      return res
 
     # ValueError is raised by google-auth when service account JSON is
     # missing required fields (e.g. client_email, private_key), or when
@@ -146,6 +204,13 @@ def _exchange_for_access_token(
         AuthCredentialMissingError: If scopes are missing for explicit
             credentials or token exchange fails.
     """
+    cache_key = _get_cache_key(sa_config)
+    cached_val = _access_token_cache.get(cache_key)
+    if cached_val:
+      token, expires_at = cached_val
+      if time.time() < expires_at - 300:
+        return token
+
     if not sa_config.use_default_credential and not sa_config.scopes:
       raise AuthCredentialMissingError(
           "scopes are required when using explicit service account credentials"
@@ -173,8 +238,16 @@ def _exchange_for_access_token(
         quota_project_id = None
 
       credentials.refresh(Request())
+      try:
+        expires_at = (
+            calendar.timegm(credentials.expiry.utctimetuple())
+            if credentials.expiry
+            else int(time.time() + 3600)
+        )
+      except (AttributeError, TypeError, ValueError):
+        expires_at = int(time.time() + 3600)
 
-      return AuthCredential(
+      res = AuthCredential(
           auth_type=AuthCredentialTypes.HTTP,
           http=HttpAuth(
               scheme="bearer",
@@ -186,6 +259,8 @@ def _exchange_for_access_token(
               else None,
           ),
       )
+      _access_token_cache[cache_key] = (res, expires_at)
+      return res
 
     # ValueError is raised by google-auth when service account JSON is
     # missing required fields (e.g. client_email, private_key).

tests/unittests/auth/test_credential_manager.py

@@ -329,6 +329,70 @@ async def test_load_auth_credentials_no_credential(self):
 
     assert result is None
 
+  @pytest.mark.asyncio
+  async def test_get_auth_credential_service_account_skips_cache(
+      self, mocker, service_account_credential
+  ):
+    """Test that Service Account credentials bypass the load/save cache."""
+    from fastapi.openapi.models import OAuth2
+    from fastapi.openapi.models import OAuthFlowClientCredentials
+    from fastapi.openapi.models import OAuthFlows
+    from google.adk.auth.auth_credential import HttpAuth
+    from google.adk.auth.auth_credential import HttpCredentials
+
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            clientCredentials=OAuthFlowClientCredentials(
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    auth_config = AuthConfig(
+        auth_scheme=auth_scheme,
+        raw_auth_credential=service_account_credential,
+    )
+
+    exchanged_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        http=HttpAuth(
+            scheme="bearer",
+            credentials=HttpCredentials(token="sa-access-token"),
+        ),
+    )
+
+    tool_context = mocker.Mock(spec=CallbackContext)
+
+    manager = CredentialManager(auth_config)
+
+    # Mock the private methods
+    manager._validate_credential = mocker.AsyncMock()
+    manager._is_credential_ready = mocker.Mock(return_value=False)
+    manager._load_existing_credential = mocker.AsyncMock()
+    manager._load_from_auth_response = mocker.AsyncMock(return_value=None)
+    manager._exchange_credential = mocker.AsyncMock(
+        return_value=(exchanged_credential, True)
+    )
+    manager._refresh_credential = mocker.AsyncMock(
+        return_value=(exchanged_credential, False)
+    )
+    manager._save_credential = mocker.AsyncMock()
+    manager._is_client_credentials_flow = mocker.Mock(return_value=True)
+
+    result = await manager.get_auth_credential(tool_context)
+
+    # Verify load and save were NOT called
+    manager._load_existing_credential.assert_not_called()
+    manager._save_credential.assert_not_called()
+
+    # Verify exchange WAS called
+    manager._exchange_credential.assert_called_once()
+    called_arg = manager._exchange_credential.call_args[0][0]
+    assert called_arg.auth_type == AuthCredentialTypes.SERVICE_ACCOUNT
+
+    assert result == exchanged_credential
+
   @pytest.mark.asyncio
   async def test_load_existing_credential_already_exchanged(self):
     """Test _load_existing_credential ignores shared config cache."""