feat: Support Oauth2 client credentials grant type

PiperOrigin-RevId: 815813477
This commit is contained in:
Xiang (Sean) Zhou
2025-10-06 11:28:17 -07:00
committed by Copybara-Service
parent 46d73be41a
commit 5c6cdcd197
5 changed files with 381 additions and 22 deletions
+40 -3
View File
@@ -25,11 +25,11 @@ from .auth_credential import AuthCredential
from .auth_credential import AuthCredentialTypes
from .auth_schemes import AuthSchemeType
from .auth_schemes import ExtendedOAuth2
from .auth_schemes import OpenIdConnectWithConfig
from .auth_tool import AuthConfig
from .exchanger.base_credential_exchanger import BaseCredentialExchanger
from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry
from .oauth2_discovery import OAuth2DiscoveryManager
from .refresher.base_credential_refresher import BaseCredentialRefresher
from .refresher.credential_refresher_registry import CredentialRefresherRegistry
logger = logging.getLogger("google_adk." + __name__)
@@ -85,8 +85,17 @@ class CredentialManager:
# Register default exchangers and refreshers
# TODO: support service account credential exchanger
from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
from .refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
oauth2_exchanger = OAuth2CredentialExchanger()
self._exchanger_registry.register(
AuthCredentialTypes.OAUTH2, oauth2_exchanger
)
self._exchanger_registry.register(
AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_exchanger
)
oauth2_refresher = OAuth2CredentialRefresher()
self._refresher_registry.register(
AuthCredentialTypes.OAUTH2, oauth2_refresher
@@ -134,9 +143,14 @@ class CredentialManager:
credential = await self._load_from_auth_response(callback_context)
was_from_auth_response = True
# Step 5: If still no credential available, return None
# Step 5: If still no credential available, check if client credentials
if not credential:
return None
# For client credentials flow, use raw credentials directly
if self._is_client_credentials_flow():
credential = self._auth_config.raw_auth_credential
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)
@@ -328,3 +342,26 @@ class CredentialManager:
and not flows.authorizationCode.tokenUrl
)
return False
def _is_client_credentials_flow(self) -> bool:
"""Check if the auth scheme uses client credentials flow.
Supports both OAuth2 and OIDC schemes.
Returns:
True if using client credentials flow, False otherwise.
"""
auth_scheme = self._auth_config.auth_scheme
# Check OAuth2 schemes
if isinstance(auth_scheme, OAuth2) and auth_scheme.flows:
return auth_scheme.flows.clientCredentials is not None
# Check OIDC schemes
if isinstance(auth_scheme, OpenIdConnectWithConfig):
return (
auth_scheme.grant_types_supported is not None
and "client_credentials" in auth_scheme.grant_types_supported
)
return False
@@ -19,9 +19,11 @@ from __future__ import annotations
import logging
from typing import Optional
from fastapi.openapi.models import OAuth2
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.auth_schemes import OpenIdConnectWithConfig
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
@@ -81,9 +83,100 @@ class OAuth2CredentialExchanger(BaseCredentialExchanger):
if auth_credential.oauth2 and auth_credential.oauth2.access_token:
return auth_credential
# Determine grant type from auth_scheme
grant_type = self._determine_grant_type(auth_scheme)
if grant_type == OAuthGrantType.CLIENT_CREDENTIALS:
return await self._exchange_client_credentials(
auth_credential, auth_scheme
)
elif grant_type == OAuthGrantType.AUTHORIZATION_CODE:
return await self._exchange_authorization_code(
auth_credential, auth_scheme
)
else:
logger.warning("Unsupported OAuth2 grant type: %s", grant_type)
return auth_credential
def _determine_grant_type(
self, auth_scheme: AuthScheme
) -> Optional[OAuthGrantType]:
"""Determine the OAuth2 grant type from the auth scheme.
Args:
auth_scheme: The OAuth2 authentication scheme.
Returns:
The OAuth2 grant type or None if cannot be determined.
"""
if isinstance(auth_scheme, OAuth2) and auth_scheme.flows:
return OAuthGrantType.from_flow(auth_scheme.flows)
elif isinstance(auth_scheme, OpenIdConnectWithConfig):
# Check supported grant types for OIDC
if (
auth_scheme.grant_types_supported
and "client_credentials" in auth_scheme.grant_types_supported
):
return OAuthGrantType.CLIENT_CREDENTIALS
else:
# Default to authorization code if client credentials not supported
return OAuthGrantType.AUTHORIZATION_CODE
return None
async def _exchange_client_credentials(
self,
auth_credential: AuthCredential,
auth_scheme: AuthScheme,
) -> AuthCredential:
"""Exchange client credentials for access token.
Args:
auth_credential: The OAuth2 credential to exchange.
auth_scheme: The OAuth2 authentication scheme.
Returns:
The credential with access token.
"""
client, token_endpoint = create_oauth2_session(auth_scheme, auth_credential)
if not client:
logger.warning("Could not create OAuth2 session for token exchange")
logger.warning(
"Could not create OAuth2 session for client credentials exchange"
)
return auth_credential
try:
tokens = client.fetch_token(
token_endpoint,
grant_type=OAuthGrantType.CLIENT_CREDENTIALS,
)
update_credential_with_tokens(auth_credential, tokens)
logger.debug("Successfully exchanged client credentials for access token")
except Exception as e:
logger.error("Failed to exchange client credentials: %s", e)
return auth_credential
return auth_credential
async def _exchange_authorization_code(
self,
auth_credential: AuthCredential,
auth_scheme: AuthScheme,
) -> AuthCredential:
"""Exchange authorization code for access token.
Args:
auth_credential: The OAuth2 credential to exchange.
auth_scheme: The OAuth2 authentication scheme.
Returns:
The credential with access token.
"""
client, token_endpoint = create_oauth2_session(auth_scheme, auth_credential)
if not client:
logger.warning(
"Could not create OAuth2 session for authorization code exchange"
)
return auth_credential
try:
@@ -94,11 +187,9 @@ class OAuth2CredentialExchanger(BaseCredentialExchanger):
grant_type=OAuthGrantType.AUTHORIZATION_CODE,
)
update_credential_with_tokens(auth_credential, tokens)
logger.debug("Successfully exchanged OAuth2 tokens")
logger.debug("Successfully exchanged authorization code for access token")
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
logger.error("Failed to exchange authorization code: %s", e)
return auth_credential
return auth_credential
+23 -14
View File
@@ -18,6 +18,8 @@ import logging
from typing import Optional
from typing import Tuple
from authlib.integrations.requests_client import OAuth2Session
from authlib.oauth2.rfc6749 import OAuth2Token
from fastapi.openapi.models import OAuth2
from ..utils.feature_decorator import experimental
@@ -25,15 +27,6 @@ from .auth_credential import AuthCredential
from .auth_schemes import AuthScheme
from .auth_schemes import OpenIdConnectWithConfig
try:
from authlib.integrations.requests_client import OAuth2Session
from authlib.oauth2.rfc6749 import OAuth2Token
AUTHLIB_AVAILABLE = True
except ImportError:
AUTHLIB_AVAILABLE = False
logger = logging.getLogger("google_adk." + __name__)
@@ -53,18 +46,34 @@ def create_oauth2_session(
"""
if isinstance(auth_scheme, OpenIdConnectWithConfig):
if not hasattr(auth_scheme, "token_endpoint"):
logger.warning("OpenIdConnect scheme missing token_endpoint")
return None, None
token_endpoint = auth_scheme.token_endpoint
scopes = auth_scheme.scopes
scopes = auth_scheme.scopes or []
elif isinstance(auth_scheme, OAuth2):
# Support both authorization code and client credentials flows
if (
not auth_scheme.flows.authorizationCode
or not auth_scheme.flows.authorizationCode.tokenUrl
auth_scheme.flows.authorizationCode
and auth_scheme.flows.authorizationCode.tokenUrl
):
token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl
scopes = list(auth_scheme.flows.authorizationCode.scopes.keys())
elif (
auth_scheme.flows.clientCredentials
and auth_scheme.flows.clientCredentials.tokenUrl
):
token_endpoint = auth_scheme.flows.clientCredentials.tokenUrl
scopes = list(auth_scheme.flows.clientCredentials.scopes.keys())
else:
logger.warning(
"OAuth2 scheme missing required flow configuration. Expected either"
" authorizationCode.tokenUrl or clientCredentials.tokenUrl. Auth"
" scheme: %s",
auth_scheme,
)
return None, None
token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl
scopes = list(auth_scheme.flows.authorizationCode.scopes.keys())
else:
logger.warning(f"Unsupported auth_scheme type: {type(auth_scheme)}")
return None, None
if (