mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Support Oauth2 client credentials grant type
PiperOrigin-RevId: 815813477
This commit is contained in:
committed by
Copybara-Service
parent
46d73be41a
commit
5c6cdcd197
@@ -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
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -17,6 +17,9 @@ from unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
|
||||
from authlib.oauth2.rfc6749 import OAuth2Token
|
||||
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 AuthCredential
|
||||
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||
from google.adk.auth.auth_credential import OAuth2Auth
|
||||
@@ -218,3 +221,116 @@ class TestOAuth2CredentialExchanger:
|
||||
# Should return original credential when authlib is not available
|
||||
assert result == credential
|
||||
assert result.oauth2.access_token is None
|
||||
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_client_credentials_success(self, mock_oauth2_session):
|
||||
"""Test successful client credentials exchange."""
|
||||
# Setup mock
|
||||
mock_client = Mock()
|
||||
mock_oauth2_session.return_value = mock_client
|
||||
mock_tokens = OAuth2Token({
|
||||
"access_token": "client_access_token",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
"expires_in": 3600,
|
||||
})
|
||||
mock_client.fetch_token.return_value = mock_tokens
|
||||
|
||||
# Create OAuth2 scheme with client credentials flow
|
||||
flows = OAuthFlows(
|
||||
clientCredentials=OAuthFlowClientCredentials(
|
||||
tokenUrl="https://example.com/token",
|
||||
scopes={"read": "Read access", "write": "Write access"},
|
||||
)
|
||||
)
|
||||
scheme = OAuth2(flows=flows)
|
||||
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Verify client credentials exchange was successful
|
||||
assert result.oauth2.access_token == "client_access_token"
|
||||
mock_client.fetch_token.assert_called_once_with(
|
||||
"https://example.com/token",
|
||||
grant_type="client_credentials",
|
||||
)
|
||||
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_client_credentials_failure(self, mock_oauth2_session):
|
||||
"""Test client credentials exchange failure."""
|
||||
# Setup mock to raise exception during fetch_token
|
||||
mock_client = Mock()
|
||||
mock_oauth2_session.return_value = mock_client
|
||||
mock_client.fetch_token.side_effect = Exception(
|
||||
"Client credentials fetch failed"
|
||||
)
|
||||
|
||||
# Create OAuth2 scheme with client credentials flow
|
||||
flows = OAuthFlows(
|
||||
clientCredentials=OAuthFlowClientCredentials(
|
||||
tokenUrl="https://example.com/token", scopes={"read": "Read access"}
|
||||
)
|
||||
)
|
||||
scheme = OAuth2(flows=flows)
|
||||
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Should return original credential when client credentials exchange fails
|
||||
assert result == credential
|
||||
assert result.oauth2.access_token is None
|
||||
mock_client.fetch_token.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_determine_grant_type_client_credentials(self):
|
||||
"""Test grant type determination for client credentials."""
|
||||
flows = OAuthFlows(
|
||||
clientCredentials=OAuthFlowClientCredentials(
|
||||
tokenUrl="https://example.com/token", scopes={"read": "Read access"}
|
||||
)
|
||||
)
|
||||
scheme = OAuth2(flows=flows)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
grant_type = exchanger._determine_grant_type(scheme)
|
||||
|
||||
from google.adk.auth.auth_schemes import OAuthGrantType
|
||||
|
||||
assert grant_type == OAuthGrantType.CLIENT_CREDENTIALS
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_determine_grant_type_openid_connect(self):
|
||||
"""Test grant type determination for OpenID Connect (defaults to auth code)."""
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
grant_type = exchanger._determine_grant_type(scheme)
|
||||
|
||||
from google.adk.auth.auth_schemes import OAuthGrantType
|
||||
|
||||
assert grant_type == OAuthGrantType.AUTHORIZATION_CODE
|
||||
|
||||
@@ -105,6 +105,9 @@ class TestCredentialManager:
|
||||
auth_config = Mock(spec=AuthConfig)
|
||||
auth_config.raw_auth_credential = None
|
||||
auth_config.exchanged_auth_credential = None
|
||||
# Add auth_scheme for the _is_client_credentials_flow method
|
||||
auth_config.auth_scheme = Mock()
|
||||
auth_config.auth_scheme.flows = None
|
||||
|
||||
callback_context = Mock()
|
||||
|
||||
@@ -562,6 +565,109 @@ class TestCredentialManager:
|
||||
|
||||
assert manager._auth_config.auth_scheme == implicit_oauth2_scheme
|
||||
|
||||
def test_is_client_credentials_flow_oauth2_with_client_credentials(self):
|
||||
"""Test _is_client_credentials_flow returns True for OAuth2 with client credentials."""
|
||||
from fastapi.openapi.models import OAuth2
|
||||
from fastapi.openapi.models import OAuthFlowClientCredentials
|
||||
from fastapi.openapi.models import OAuthFlows
|
||||
|
||||
# Create OAuth2 scheme with client credentials flow
|
||||
auth_scheme = OAuth2(
|
||||
flows=OAuthFlows(
|
||||
clientCredentials=OAuthFlowClientCredentials(
|
||||
tokenUrl="https://example.com/token"
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
auth_config = Mock(spec=AuthConfig)
|
||||
auth_config.auth_scheme = auth_scheme
|
||||
auth_config.raw_auth_credential = None
|
||||
auth_config.exchanged_auth_credential = None
|
||||
|
||||
manager = CredentialManager(auth_config)
|
||||
|
||||
assert manager._is_client_credentials_flow() is True
|
||||
|
||||
def test_is_client_credentials_flow_oauth2_without_client_credentials(self):
|
||||
"""Test _is_client_credentials_flow returns False for OAuth2 without client credentials."""
|
||||
from fastapi.openapi.models import OAuth2
|
||||
from fastapi.openapi.models import OAuthFlowAuthorizationCode
|
||||
from fastapi.openapi.models import OAuthFlows
|
||||
|
||||
# Create OAuth2 scheme with authorization code flow only
|
||||
auth_scheme = OAuth2(
|
||||
flows=OAuthFlows(
|
||||
authorizationCode=OAuthFlowAuthorizationCode(
|
||||
authorizationUrl="https://example.com/auth",
|
||||
tokenUrl="https://example.com/token",
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
auth_config = Mock(spec=AuthConfig)
|
||||
auth_config.auth_scheme = auth_scheme
|
||||
auth_config.raw_auth_credential = None
|
||||
auth_config.exchanged_auth_credential = None
|
||||
|
||||
manager = CredentialManager(auth_config)
|
||||
|
||||
assert manager._is_client_credentials_flow() is False
|
||||
|
||||
def test_is_client_credentials_flow_oidc_with_client_credentials(self):
|
||||
"""Test _is_client_credentials_flow returns True for OIDC with client credentials."""
|
||||
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
|
||||
|
||||
# Create OIDC scheme with client credentials support
|
||||
auth_scheme = OpenIdConnectWithConfig(
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
grant_types_supported=["authorization_code", "client_credentials"],
|
||||
)
|
||||
|
||||
auth_config = Mock(spec=AuthConfig)
|
||||
auth_config.auth_scheme = auth_scheme
|
||||
auth_config.raw_auth_credential = None
|
||||
auth_config.exchanged_auth_credential = None
|
||||
|
||||
manager = CredentialManager(auth_config)
|
||||
|
||||
assert manager._is_client_credentials_flow() is True
|
||||
|
||||
def test_is_client_credentials_flow_oidc_without_client_credentials(self):
|
||||
"""Test _is_client_credentials_flow returns False for OIDC without client credentials."""
|
||||
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
|
||||
|
||||
# Create OIDC scheme without client credentials support
|
||||
auth_scheme = OpenIdConnectWithConfig(
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
grant_types_supported=["authorization_code"],
|
||||
)
|
||||
|
||||
auth_config = Mock(spec=AuthConfig)
|
||||
auth_config.auth_scheme = auth_scheme
|
||||
auth_config.raw_auth_credential = None
|
||||
auth_config.exchanged_auth_credential = None
|
||||
|
||||
manager = CredentialManager(auth_config)
|
||||
|
||||
assert manager._is_client_credentials_flow() is False
|
||||
|
||||
def test_is_client_credentials_flow_other_scheme(self):
|
||||
"""Test _is_client_credentials_flow returns False for other auth schemes."""
|
||||
# Create a non-OAuth2/OIDC scheme
|
||||
auth_scheme = Mock()
|
||||
|
||||
auth_config = Mock(spec=AuthConfig)
|
||||
auth_config.auth_scheme = auth_scheme
|
||||
auth_config.raw_auth_credential = None
|
||||
auth_config.exchanged_auth_credential = None
|
||||
|
||||
manager = CredentialManager(auth_config)
|
||||
|
||||
assert manager._is_client_credentials_flow() is False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def oauth2_auth_scheme():
|
||||
|
||||
Reference in New Issue
Block a user