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

Commit 618dc01

Browse files
seanzhougooglecopybara-github
seanzhougoogle
authored and
copybara-github
committed
refactor: Refactor oauth2_credential_exchanger to exchanger and refresher separately
PiperOrigin-RevId: 771637958
1 parent a4d432a commit 618dc01

31 files changed

+2308
-618
lines changed

‎src/google/adk/agents/invocation_context.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from pydantic import ConfigDict
2323

2424
from ..artifacts.base_artifact_service import BaseArtifactService
25+
from ..auth.credential_service.base_credential_service import BaseCredentialService
2526
from ..memory.base_memory_service import BaseMemoryService
2627
from ..sessions.base_session_service import BaseSessionService
2728
from ..sessions.session import Session
@@ -115,6 +116,7 @@ class InvocationContext(BaseModel):
115116
artifact_service: Optional[BaseArtifactService] = None
116117
session_service: BaseSessionService
117118
memory_service: Optional[BaseMemoryService] = None
119+
credential_service: Optional[BaseCredentialService] = None
118120

119121
invocation_id: str
120122
"""The id of this invocation context. Readonly."""

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

Lines changed: 5 additions & 4 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
@@ -43,9 +43,10 @@ def __init__(self, auth_config: AuthConfig):
4343
def exchange_auth_token(
4444
self,
4545
) -> AuthCredential:
46-
return OAuth2CredentialFetcher(
47-
self.auth_config.auth_scheme, self.auth_config.exchanged_auth_credential
48-
).exchange()
46+
exchanger = OAuth2CredentialExchanger()
47+
return exchanger.exchange(
48+
self.auth_config.exchanged_auth_credential, self.auth_config.auth_scheme
49+
)
4950

5051
def parse_and_store_auth_response(self, state: State) -> None:
5152

‎src/google/adk/auth/credential_service/base_credential_service.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,12 @@
1919
from typing import Optional
2020

2121
from ...tools.tool_context import ToolContext
22-
from ...utils.feature_decorator import working_in_progress
22+
from ...utils.feature_decorator import experimental
2323
from ..auth_credential import AuthCredential
2424
from ..auth_tool import AuthConfig
2525

2626

27-
@working_in_progress("Implementation are in progress. Don't use it for now.")
27+
@experimental
2828
class BaseCredentialService(ABC):
2929
"""Abstract class for Service that loads / saves tool credentials from / to
3030
the backend credential store."""

‎src/google/adk/auth/credential_service/in_memory_credential_service.py

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
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+
from __future__ import annotations
16+
17+
from typing import Optional
18+
19+
from typing_extensions import override
20+
21+
from ...tools.tool_context import ToolContext
22+
from ...utils.feature_decorator import experimental
23+
from ..auth_credential import AuthCredential
24+
from ..auth_tool import AuthConfig
25+
from .base_credential_service import BaseCredentialService
26+
27+
28+
@experimental
29+
class InMemoryCredentialService(BaseCredentialService):
30+
"""Class for in memory implementation of credential service(Experimental)"""
31+
32+
def __init__(self):
33+
super().__init__()
34+
self._store: dict[str, AuthCredential] = {}
35+
36+
@override
37+
async def load_credential(
38+
self,
39+
auth_config: AuthConfig,
40+
tool_context: ToolContext,
41+
) -> Optional[AuthCredential]:
42+
"""
43+
Loads the credential by auth config and current tool context from the
44+
backend credential store.
45+
46+
Args:
47+
auth_config: The auth config which contains the auth scheme and auth
48+
credential information. auth_config.get_credential_key will be used to
49+
build the key to load the credential.
50+
51+
tool_context: The context of the current invocation when the tool is
52+
trying to load the credential.
53+
54+
Returns:
55+
Optional[AuthCredential]: the credential saved in the store.
56+
57+
"""
58+
storage = self._get_storage_for_current_context(tool_context)
59+
return storage.get(auth_config.credential_key)
60+
61+
@override
62+
async def save_credential(
63+
self,
64+
auth_config: AuthConfig,
65+
tool_context: ToolContext,
66+
) -> None:
67+
"""
68+
Saves the exchanged_auth_credential in auth config to the backend credential
69+
store.
70+
71+
Args:
72+
auth_config: The auth config which contains the auth scheme and auth
73+
credential information. auth_config.get_credential_key will be used to
74+
build the key to save the credential.
75+
76+
tool_context: The context of the current invocation when the tool is
77+
trying to save the credential.
78+
79+
Returns:
80+
None
81+
"""
82+
storage = self._get_storage_for_current_context(tool_context)
83+
storage[auth_config.credential_key] = auth_config.exchanged_auth_credential
84+
85+
def _get_storage_for_current_context(self, tool_context: ToolContext) -> str:
86+
app_name = tool_context._invocation_context.app_name
87+
user_id = tool_context._invocation_context.user_id
88+
89+
if app_name not in self._store:
90+
self._store[app_name] = {}
91+
if user_id not in self._store[app_name]:
92+
self._store[app_name][user_id] = {}
93+
return self._store[app_name][user_id]

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
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+
"""Credential exchanger module."""
16+
17+
from .base_credential_exchanger import BaseCredentialExchanger
18+
from .service_account_credential_exchanger import ServiceAccountCredentialExchanger
19+
20+
__all__ = [
21+
"BaseCredentialExchanger",
22+
"ServiceAccountCredentialExchanger",
23+
]

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
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+
"""Base credential exchanger interface."""
16+
17+
from __future__ import annotations
18+
19+
import abc
20+
from typing import Optional
21+
22+
from ...utils.feature_decorator import experimental
23+
from ..auth_credential import AuthCredential
24+
from ..auth_schemes import AuthScheme
25+
26+
27+
@experimental
28+
class BaseCredentialExchanger(abc.ABC):
29+
"""Base interface for credential exchangers."""
30+
31+
@abc.abstractmethod
32+
def exchange(
33+
self,
34+
auth_credential: AuthCredential,
35+
auth_scheme: Optional[AuthScheme] = None,
36+
) -> AuthCredential:
37+
"""Exchange credential if needed.
38+
39+
Args:
40+
auth_credential: The credential to exchange.
41+
auth_scheme: The authentication scheme (optional, some exchangers don't need it).
42+
43+
Returns:
44+
The exchanged credential.
45+
46+
Raises:
47+
ValueError: If credential exchange fails.
48+
"""
49+
pass

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
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+
"""Credential exchanger registry."""
16+
17+
from __future__ import annotations
18+
19+
from typing import Dict
20+
from typing import Optional
21+
22+
from ...utils.feature_decorator import experimental
23+
from ..auth_credential import AuthCredentialTypes
24+
from .base_credential_exchanger import BaseCredentialExchanger
25+
26+
27+
@experimental
28+
class CredentialExchangerRegistry:
29+
"""Registry for credential exchanger instances."""
30+
31+
def __init__(self):
32+
self._exchangers: Dict[AuthCredentialTypes, BaseCredentialExchanger] = {}
33+
34+
def register(
35+
self,
36+
credential_type: AuthCredentialTypes,
37+
exchanger_instance: BaseCredentialExchanger,
38+
) -> None:
39+
"""Register an exchanger instance for a credential type.
40+
41+
Args:
42+
credential_type: The credential type to register for.
43+
exchanger_instance: The exchanger instance to register.
44+
"""
45+
self._exchangers[credential_type] = exchanger_instance
46+
47+
def get_exchanger(
48+
self, credential_type: AuthCredentialTypes
49+
) -> Optional[BaseCredentialExchanger]:
50+
"""Get the exchanger instance for a credential type.
51+
52+
Args:
53+
credential_type: The credential type to get exchanger for.
54+
55+
Returns:
56+
The exchanger instance if registered, None otherwise.
57+
"""
58+
return self._exchangers.get(credential_type)

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

Copy file name to clipboard
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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_session_helper import OAuth2SessionHelper
26+
from google.adk.utils.feature_decorator import experimental
27+
28+
from .base_credential_exchanger import BaseCredentialExchanger
29+
30+
try:
31+
from authlib.integrations.requests_client import OAuth2Session
32+
33+
AUTHLIB_AVIALABLE = True
34+
except ImportError:
35+
AUTHLIB_AVIALABLE = False
36+
37+
logger = logging.getLogger("google_adk." + __name__)
38+
39+
40+
@experimental
41+
class OAuth2CredentialExchanger(BaseCredentialExchanger):
42+
"""Exchanges OAuth2 credentials from authorization responses."""
43+
44+
def exchange(
45+
self,
46+
auth_credential: AuthCredential,
47+
auth_scheme: Optional[AuthScheme] = None,
48+
) -> AuthCredential:
49+
"""Exchange OAuth2 credential from authorization response.
50+
51+
Args:
52+
auth_credential: The OAuth2 credential to exchange.
53+
auth_scheme: The OAuth2 authentication scheme.
54+
55+
Returns:
56+
The exchanged credential with access token.
57+
58+
Raises:
59+
ValueError: If credential exchange fails or auth_scheme is missing.
60+
"""
61+
if not auth_scheme:
62+
raise ValueError("auth_scheme is required for OAuth2 credential exchange")
63+
64+
if not AUTHLIB_AVIALABLE:
65+
return auth_credential
66+
67+
if auth_credential.oauth2 and auth_credential.oauth2.access_token:
68+
return auth_credential
69+
70+
helper = OAuth2SessionHelper(auth_scheme, auth_credential)
71+
client, token_endpoint = helper.create_oauth2_session()
72+
if not client:
73+
logger.warning("Could not create OAuth2 session for token exchange")
74+
return auth_credential
75+
76+
try:
77+
tokens = client.fetch_token(
78+
token_endpoint,
79+
authorization_response=auth_credential.oauth2.auth_response_uri,
80+
code=auth_credential.oauth2.auth_code,
81+
grant_type=OAuthGrantType.AUTHORIZATION_CODE,
82+
)
83+
helper.update_credential_with_tokens(tokens)
84+
logger.info("Successfully exchanged OAuth2 tokens")
85+
except Exception as e:
86+
logger.error("Failed to exchange OAuth2 tokens: %s", e)
87+
# Return original credential on failure
88+
return auth_credential
89+
90+
return auth_credential

0 commit comments

Comments
 (0)
0