Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 25 additions & 8 deletions src/google/adk/auth/exchanger/oauth2_credential_exchanger.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
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 normalize_oauth2_tokens
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
Expand Down Expand Up @@ -193,19 +194,35 @@ async def _exchange_authorization_code(
return ExchangeResult(auth_credential, False)

try:
tokens = client.fetch_token(
token_endpoint,
authorization_response=self._normalize_auth_uri(
token_auth_method = (
auth_credential.oauth2.token_endpoint_auth_method
if auth_credential.oauth2
else None
)
fetch_token_kwargs = {
'authorization_response': self._normalize_auth_uri(
auth_credential.oauth2.auth_response_uri
),
code=auth_credential.oauth2.auth_code,
grant_type=OAuthGrantType.AUTHORIZATION_CODE,
client_id=auth_credential.oauth2.client_id,
)
'code': auth_credential.oauth2.auth_code,
'grant_type': OAuthGrantType.AUTHORIZATION_CODE,
}
# For client_secret_post, Authlib already includes client_id in POST body
# from OAuth2Session. Passing client_id again can duplicate parameters.
if token_auth_method != 'client_secret_post':
fetch_token_kwargs['client_id'] = auth_credential.oauth2.client_id

tokens = client.fetch_token(token_endpoint, **fetch_token_kwargs)
tokens = normalize_oauth2_tokens(tokens)
update_credential_with_tokens(auth_credential, tokens)
logger.debug("Successfully exchanged authorization code for access token")
except Exception as e:
logger.error("Failed to exchange authorization code: %s", e)
logger.error(
"Failed to exchange authorization code (token_endpoint_auth_method=%s): %s",
auth_credential.oauth2.token_endpoint_auth_method
if auth_credential.oauth2
else None,
Comment on lines +221 to +223
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To improve maintainability and avoid code duplication, you can reuse the token_auth_method variable that is already defined within the try block. This avoids re-evaluating the same expression in the except block.

          token_auth_method,

e,
)
return ExchangeResult(auth_credential, False)

return ExchangeResult(auth_credential, True)
13 changes: 13 additions & 0 deletions src/google/adk/auth/oauth2_credential_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

from __future__ import annotations

from collections.abc import Mapping
import logging
from typing import Optional
from typing import Tuple
Expand Down Expand Up @@ -97,6 +98,17 @@ def create_oauth2_session(
)


@experimental
def normalize_oauth2_tokens(tokens: object) -> OAuth2Token:
"""Validates and normalizes token payload returned by OAuth libraries."""
if not isinstance(tokens, Mapping):
raise ValueError(
"OAuth2 token response must be a dict-like object, got "
f"{type(tokens).__name__}."
)
return tokens # type: ignore[return-value]


@experimental
def update_credential_with_tokens(
auth_credential: AuthCredential, tokens: OAuth2Token
Expand All @@ -107,6 +119,7 @@ def update_credential_with_tokens(
auth_credential: The authentication credential to update.
tokens: The OAuth2Token object containing new token information.
"""
tokens = normalize_oauth2_tokens(tokens)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The normalize_oauth2_tokens function is already being called in oauth2_credential_exchanger.py and oauth2_credential_refresher.py before they call update_credential_with_tokens. Calling it again here is redundant. It's best to have the callers be responsible for normalization and have this function expect a valid token object.

auth_credential.oauth2.access_token = tokens.get("access_token")
auth_credential.oauth2.refresh_token = tokens.get("refresh_token")
auth_credential.oauth2.expires_at = (
Expand Down
2 changes: 2 additions & 0 deletions src/google/adk/auth/refresher/oauth2_credential_refresher.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_schemes import AuthScheme
from google.adk.auth.oauth2_credential_util import create_oauth2_session
from google.adk.auth.oauth2_credential_util import normalize_oauth2_tokens
from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
from google.adk.utils.feature_decorator import experimental
from google.auth.transport.requests import Request
Expand Down Expand Up @@ -115,6 +116,7 @@ async def refresh(
url=token_endpoint,
refresh_token=auth_credential.oauth2.refresh_token,
)
tokens = normalize_oauth2_tokens(tokens)
update_credential_with_tokens(auth_credential, tokens)
logger.debug("Successfully refreshed OAuth2 tokens")
except Exception as e:
Expand Down