8000 refactor: Refactor oauth2_credential_exchanger to exchanger and refre… · DavidSchmidt00/adk-python@9a207cb · GitHub
[go: up one dir, main page]
Skip to content

Commit 9a207cb

Browse files
seanzhougooglecopybara-github
seanzhougoogle
authored and
copybara-github
committed
refactor: Refactor oauth2_credential_exchanger to exchanger and refresher separately
PiperOrigin-RevId: 772979993
1 parent a17ebe6 commit 9a207cb

File tree

16 files changed

+926
-368
lines changed

16 files changed

+926
-368
lines changed

‎src/google/adk/auth/auth_handler.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
from .auth_schemes import AuthSchemeType
2323
from .auth_schemes import OpenIdConnectWithConfig
2424
from .auth_tool import AuthConfig
25-
from .oauth2_credential_fetcher import OAuth2CredentialFetcher
25+
from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
2626

2727
if TYPE_CHECKING:
2828
from ..sessions.state import State
@@ -36,18 +36,23 @@
3636

3737

3838
class AuthHandler:
39+
"""A handler that handles the auth flow in Agent Development Kit to help
40+
orchestrate the credential request and response flow (e.g. OAuth flow)
41+
This class should only be used by Agent Development Kit.
42+
"""
3943

4044
def __init__(self, auth_config: AuthConfig):
4145
self.auth_config = auth_config
4246

43-
def exchange_auth_token(
47+
async def exchange_auth_token(
4448
self,
4549
) -> AuthCredential:
46-
return OAuth2CredentialFetcher(
47-
self.auth_config.auth_scheme, self.auth_config.exchanged_auth_credential
48-
).exchange()
50+
exchanger = OAuth2CredentialExchanger()
51+
return await exchanger.exchange(
52+
self.auth_config.exchanged_auth_credential, self.auth_config.auth_scheme
53+
)
4954

50-
def parse_and_store_auth_response(self, state: State) -> None:
55+
async def parse_and_store_auth_response(self, state: State) -> None:
5156

5257
credential_key = "temp:" + self.auth_config.credential_key
5358

@@ -60,7 +65,7 @@ def parse_and_store_auth_response(self, state: State) -> None:
6065
):
6166
return
6267

63-
state[credential_key] = self.exchange_auth_token()
68+
state[credential_key] = await self.exchange_auth_token()
6469

6570
def _validate(self) -> None:
6671
if not self.auth_scheme:

‎src/google/adk/auth/auth_preprocessor.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,9 @@ async def run_async(
6767
# function call
6868
request_euc_function_call_ids.add(function_call_response.id)
6969
auth_config = AuthConfig.model_validate(function_call_response.response)
70-
AuthHandler(auth_config=auth_config).parse_and_store_auth_response(
71-
state=invocation_context.session.state
72-
)
70+
await AuthHandler(
71+
auth_config=auth_config
72+
).parse_and_store_auth_response(state=invocation_context.session.state)
7373
break
7474

7575
if not request_euc_function_call_ids:

‎src/google/adk/auth/exchanger/oauth2_credential_exchanger.py

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""OAuth2 credential exchanger implementation."""
16+
17+
from __future__ import annotations
18+
19+
import logging
20+
from typing import Optional
21+
22+
from google.adk.auth.auth_credential import AuthCredential
23+
from google.adk.auth.auth_schemes import AuthScheme
24+
from google.adk.auth.auth_schemes import OAuthGrantType
25+
from google.adk.auth.oauth2_credential_util import create_oauth2_session
26+
from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
27+
from google.adk.utils.feature_decorator import experimental
28+
from typing_extensions import override
29+
30+
from .base_credential_exchanger import BaseCredentialExchanger
31+
from .base_credential_exchanger import CredentialExchangError
32+
33+
try:
34+
from authlib.integrations.requests_client import OAuth2Session
35+
36+
AUTHLIB_AVIALABLE = True
37+
except ImportError:
38+
AUTHLIB_AVIALABLE = False
39+
40+
logger = logging.getLogger("google_adk." + __name__)
41+
42+
43+
@experimental
44+
class OAuth2CredentialExchanger(BaseCredentialExchanger):
45+
"""Exchanges OAuth2 credentials from authorization responses."""
46+
47+
@override
48+
async def exchange(
49+
self,
50+
auth_credential: AuthCredential,
51+
auth_scheme: Optional[AuthScheme] = None,
52+
) -> AuthCredential:
53+
"""Exchange OAuth2 credential from authorization response.
54+
if credential exchange failed, the original credential will be returned.
55+
56+
Args:
57+
auth_credential: The OAuth2 credential to exchange.
58+
auth_scheme: The OAuth2 authentication scheme.
59+
60+
Returns:
61+
The exchanged credential with access token.
62+
63+
Raises:
64+
CredentialExchangError: If auth_scheme is missing.
65+
"""
66+
if not auth_scheme:
67+
raise CredentialExchangError(
68+
"auth_scheme is required for OAuth2 credential exchange"
69+
)
70+
71+
if not AUTHLIB_AVIALABLE:
72+
# If authlib is not available, we cannot exchange the credential.
73+
# We return the original credential without exchange.
74+
# The client using this tool can decide to exchange the credential
75+
# themselves using other lib.
76+
logger.warning(
77+
"authlib is not available, skipping OAuth2 credential exchange."
78+
)
79+
return auth_credential
80+
81+
if auth_credential.oauth2 and auth_credential.oauth2.access_token:
82+
return auth_creden CAD4 tial
83+
84+
client, token_endpoint = create_oauth2_session(auth_scheme, auth_credential)
85+
if not client:
86+
logger.warning("Could not create OAuth2 session for token exchange")
87+
return auth_credential
88+
89+
try:
90+
tokens = client.fetch_token(
91+
token_endpoint,
92+
authorization_response=auth_credential.oauth2.auth_response_uri,
93+
code=auth_credential.oauth2.auth_code,
94+
grant_type=OAuthGrantType.AUTHORIZATION_CODE,
95+
)
96+
update_credential_with_tokens(auth_credential, tokens)
97+
logger.debug("Successfully exchanged OAuth2 tokens")
98+
except Exception as e:
99+
# TODO reconsider whether we should raise errors in this case
100+
logger.error("Failed to exchange OAuth2 tokens: %s", e)
101+
# Return original credential on failure
102+
return auth_credential
103+
104+
return auth_credential

‎src/google/adk/auth/oauth2_credential_fetcher.py

Lines changed: 0 additions & 132 deletions
This file was deleted.

0 commit comments

Comments
 (0)
0