encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Make credential key generation stable and prevent cross-user credential leaks

Browse Source 
This change updates the credential key generation to use a stable hash (SHA256) instead of Python's built-in hash, which can vary based on PYTHONHASHSEED. It also makes sure that temporary or exchanged OAuth2 fields are excluded from the key calculation. I also added when saving credentials, a copy of the AuthConfig is used to avoid modifying the original shared AuthConfig instance with user-specific exchanged credentials.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 864599326
This commit is contained in:
George Weale
2026-02-02 17:51:19 -08:00
committed by Copybara-Service
parent 666cebe369
commit 33012e6dda
11 changed files with 448 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/auth/auth_handler.py
+2
View File
@@ -115,6 +115,7 @@ class AuthHandler:
exchanged_auth_credential=self.auth_config.raw_auth_credential.model_copy(
deep=True
),
credential_key=self.auth_config.credential_key,
)
# Check for client_id and client_secret
@@ -133,6 +134,7 @@ class AuthHandler:
auth_scheme=self.auth_config.auth_scheme,
raw_auth_credential=self.auth_config.raw_auth_credential,
exchanged_auth_credential=exchanged_credential,
credential_key=self.auth_config.credential_key,
)
def generate_auth_uri(
src/google/adk/auth/auth_preprocessor.py
+32
View File
@@ -72,6 +72,7 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor):
if not responses:
return
requested_auth_config_by_request_id = {}
# look for auth response
for function_call_response in responses:
if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME:
@@ -79,7 +80,38 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor):
# found the function call response for the system long running request euc
# function call
request_euc_function_call_ids.add(function_call_response.id)
if request_euc_function_call_ids:
for event in events:
function_calls = event.get_function_calls()
if not function_calls:
continue
try:
for function_call in function_calls:
if (
function_call.id in request_euc_function_call_ids
and function_call.name == REQUEST_EUC_FUNCTION_CALL_NAME
):
args = AuthToolArguments.model_validate(function_call.args)
requested_auth_config_by_request_id[function_call.id] = (
args.auth_config
)
except TypeError:
continue
for function_call_response in responses:
if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME:
continue
auth_config = AuthConfig.model_validate(function_call_response.response)
requested_auth_config = requested_auth_config_by_request_id.get(
function_call_response.id
)
if (
requested_auth_config
and requested_auth_config.credential_key is not None
):
auth_config.credential_key = requested_auth_config.credential_key
await AuthHandler(auth_config=auth_config).parse_and_store_auth_response(
state=invocation_context.session.state
)
src/google/adk/auth/auth_tool.py
+49 -4
View File
@@ -14,8 +14,11 @@
from __future__ import annotations
import hashlib
import json
from typing import Optional
from pydantic import BaseModel
from typing_extensions import deprecated
from .auth_credential import AuthCredential
@@ -23,6 +26,28 @@ from .auth_credential import BaseModelWithConfig
from .auth_schemes import AuthScheme
def _stable_model_digest(model: BaseModel) -> str:
"""Returns a stable digest for a pydantic model.
The digest is stable across:
- Python hash seeds (does not use `hash()`).
- Dict insertion ordering differences (canonicalizes via `sort_keys=True`).
- Pydantic `model_extra` values (ignored).
"""
if getattr(model, "model_extra", None):
model = model.model_copy(deep=True)
model.model_extra.clear()
dumped = model.model_dump(by_alias=True, exclude_none=True, mode="json")
canonical_json = json.dumps(
dumped,
sort_keys=True,
ensure_ascii=False,
separators=(",", ":"),
)
return hashlib.sha256(canonical_json.encode("utf-8")).hexdigest()[:16]
class AuthConfig(BaseModelWithConfig):
"""The auth config sent by tool asking client to collect auth credentials and
@@ -58,12 +83,22 @@ class AuthConfig(BaseModelWithConfig):
super().__init__(**data)
if self.credential_key:
return
for obj in (self.raw_auth_credential, self.auth_scheme):
if not obj or not getattr(obj, "model_extra", None):
continue
for key in ("credential_key", "credentialKey"):
value = obj.model_extra.get(key)
if isinstance(value, str) and value:
self.credential_key = value
return
self.credential_key = self.get_credential_key()
@deprecated("This method is deprecated. Use credential_key instead.")
def get_credential_key(self):
"""Builds a hash key based on auth_scheme and raw_auth_credential used to
save / load this credential to / from a credentials service.
"""Builds a stable key based on auth_scheme and raw_auth_credential.
This is used to save/load credentials to/from a credential service when
`credential_key` is not explicitly provided.
"""
auth_scheme = self.auth_scheme
@@ -72,7 +107,7 @@ class AuthConfig(BaseModelWithConfig):
auth_scheme = auth_scheme.model_copy(deep=True)
auth_scheme.model_extra.clear()
scheme_name = (
f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}"
f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}"
if auth_scheme
else ""
)
@@ -81,8 +116,18 @@ class AuthConfig(BaseModelWithConfig):
if auth_credential and auth_credential.model_extra:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.model_extra.clear()
if auth_credential and auth_credential.oauth2:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.oauth2.auth_uri = None
auth_credential.oauth2.state = None
auth_credential.oauth2.auth_response_uri = None
auth_credential.oauth2.auth_code = None
auth_credential.oauth2.access_token = None
auth_credential.oauth2.refresh_token = None
auth_credential.oauth2.expires_at = None
auth_credential.oauth2.expires_in = None
credential_name = (
f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}"
f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}"
if auth_credential
else ""
)
src/google/adk/auth/credential_manager.py
+9 -6
View File
@@ -142,7 +142,9 @@ class CredentialManager:
# Step 2: Check if credential is already ready (no processing needed)
if self._is_credential_ready():
return self._auth_config.raw_auth_credential
# Return a copy to avoid leaking mutations across invocations/users when
# tools share a long-lived AuthConfig instance.
return self._auth_config.raw_auth_credential.model_copy(deep=True)
# Step 3: Try to load existing processed credential
credential = await self._load_existing_credential(context)
@@ -159,7 +161,9 @@ class CredentialManager:
if not credential:
# For client credentials flow, use raw credentials directly
if self._is_client_credentials_flow():
credential = self._auth_config.raw_auth_credential
# Exchange/refresh steps may mutate the credential object in-place, so
# do not operate on the shared tool config.
credential = self._auth_config.raw_auth_credential.model_copy(deep=True)
else:
# For authorization code flow, return None to trigger user authorization
return None
@@ -298,12 +302,11 @@ class CredentialManager:
self, context: CallbackContext, credential: AuthCredential
) -> None:
"""Save credential to credential service if available."""
# Update the exchanged credential in config
self._auth_config.exchanged_auth_credential = credential
credential_service = context._invocation_context.credential_service
if credential_service:
await context.save_credential(self._auth_config)
auth_config_to_save = self._auth_config.model_copy(deep=True)
auth_config_to_save.exchanged_auth_credential = credential
await context.save_credential(auth_config_to_save)
async def _populate_auth_scheme(self) -> bool:
"""Auto-discover server metadata and populate missing auth scheme info.
src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py
+16 -4
View File
@@ -70,6 +70,7 @@ class OpenAPIToolset(BaseToolset):
spec_str_type: Literal["json", "yaml"] = "json",
auth_scheme: Optional[AuthScheme] = None,
auth_credential: Optional[AuthCredential] = None,
credential_key: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
tool_name_prefix: Optional[str] = None,
ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None,
@@ -108,6 +109,8 @@ class OpenAPIToolset(BaseToolset):
auth_credential: The auth credential to use for all tools. Use
AuthCredential or use helpers in
``google.adk.tools.openapi_tool.auth.auth_helpers``
credential_key: Optional stable key used for interactive auth and
credential caching across all tools in this toolset.
tool_filter: The filter used to filter the tools in the toolset. It can be
either a tool predicate or a list of tool names of the tools to expose.
tool_name_prefix: The prefix to prepend to the names of the tools returned
@@ -137,6 +140,7 @@ class OpenAPIToolset(BaseToolset):
AuthConfig(
auth_scheme=auth_scheme,
raw_auth_credential=auth_credential,
credential_key=credential_key,
)
if auth_scheme
else None
@@ -147,6 +151,8 @@ class OpenAPIToolset(BaseToolset):
self._tools: Final[List[RestApiTool]] = list(self._parse(spec_dict))
if auth_scheme or auth_credential:
self._configure_auth_all(auth_scheme, auth_credential)
if credential_key:
self._configure_credential_key_all(credential_key)
def _configure_auth_all(
self, auth_scheme: AuthScheme, auth_credential: AuthCredential
@@ -159,6 +165,11 @@ class OpenAPIToolset(BaseToolset):
if auth_credential:
tool.configure_auth_credential(auth_credential)
def _configure_credential_key_all(self, credential_key: str):
"""Configure credential key for all tools."""
for tool in self._tools:
tool.configure_credential_key(credential_key)
def configure_ssl_verify_all(
self, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None
):
@@ -229,8 +240,9 @@ class OpenAPIToolset(BaseToolset):
def get_auth_config(self) -> Optional[AuthConfig]:
"""Returns the auth config for this toolset.
ADK will populate exchanged_auth_credential on this config before calling
get_tools(). The toolset can then access the ready-to-use credential via
self._auth_config.exchanged_auth_credential.
Note: This returns a copy so any exchanged credentials populated by the ADK
framework do not persist on the toolset instance across invocations.
"""
return self._auth_config
return (
self._auth_config.model_copy(deep=True) if self._auth_config else None
)
src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py
+13 -1
View File
@@ -100,6 +100,8 @@ class RestApiTool(BaseTool):
header_provider: Optional[
Callable[[ReadonlyContext], Dict[str, str]]
] = None,
*,
credential_key: Optional[str] = None,
):
"""Initializes the RestApiTool with the given parameters.
@@ -137,6 +139,8 @@ class RestApiTool(BaseTool):
an argument, allowing dynamic header generation based on the current
context. Useful for adding custom headers like correlation IDs,
authentication tokens, or other request metadata.
credential_key: Optional stable key used for interactive auth and
credential caching.
"""
# Gemini restrict the length of function name to be less than 64 characters
self.name = name[:60]
@@ -152,6 +156,7 @@ class RestApiTool(BaseTool):
else operation
)
self.auth_credential, self.auth_scheme = None, None
self.credential_key = credential_key
self.configure_auth_credential(auth_credential)
self.configure_auth_scheme(auth_scheme)
@@ -266,6 +271,10 @@ class RestApiTool(BaseTool):
auth_credential = AuthCredential.model_validate_json(auth_credential)
self.auth_credential = auth_credential
def configure_credential_key(self, credential_key: Optional[str] = None):
"""Configures the credential key for interactive auth / caching."""
self.credential_key = credential_key
def configure_ssl_verify(
self, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None
):
@@ -449,7 +458,10 @@ class RestApiTool(BaseTool):
"""
# Prepare auth credentials for the API call
tool_auth_handler = ToolAuthHandler.from_tool_context(
tool_context, self.auth_scheme, self.auth_credential
tool_context,
self.auth_scheme,
self.auth_credential,
credential_key=self.credential_key,
)
auth_result = await tool_auth_handler.prepare_auth_credentials()
auth_state, auth_scheme, auth_credential = (
src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py
+89 -18
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import hashlib
import logging
from typing import Literal
from typing import Optional
@@ -24,6 +25,7 @@ from ....auth.auth_credential import AuthCredential
from ....auth.auth_credential import AuthCredentialTypes
from ....auth.auth_schemes import AuthScheme
from ....auth.auth_schemes import AuthSchemeType
from ....auth.auth_tool import _stable_model_digest
from ....auth.auth_tool import AuthConfig
from ....auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
from ...tool_context import ToolContext
@@ -50,19 +52,60 @@ class ToolContextCredentialStore:
def __init__(self, tool_context: ToolContext):
self.tool_context = tool_context
def _legacy_stable_digest(self, text: str) -> str:
return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]
def _get_legacy_credential_key(
self,
auth_scheme: Optional[AuthScheme],
auth_credential: Optional[AuthCredential],
) -> str:
if auth_credential and auth_credential.oauth2:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.oauth2.auth_uri = None
auth_credential.oauth2.state = None
auth_credential.oauth2.auth_response_uri = None
auth_credential.oauth2.auth_code = None
auth_credential.oauth2.access_token = None
auth_credential.oauth2.refresh_token = None
auth_credential.oauth2.expires_at = None
auth_credential.oauth2.expires_in = None
scheme_name = (
f"{auth_scheme.type_.name}_{self._legacy_stable_digest(auth_scheme.model_dump_json())}"
if auth_scheme
else ""
)
credential_name = (
f"{auth_credential.auth_type.value}_{self._legacy_stable_digest(auth_credential.model_dump_json())}"
if auth_credential
else ""
)
return f"{scheme_name}_{credential_name}_existing_exchanged_credential"
def get_credential_key(
self,
auth_scheme: Optional[AuthScheme],
auth_credential: Optional[AuthCredential],
) -> str:
"""Generates a unique key for the given auth scheme and credential."""
if auth_credential and auth_credential.oauth2:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.oauth2.auth_uri = None
auth_credential.oauth2.state = None
auth_credential.oauth2.auth_response_uri = None
auth_credential.oauth2.auth_code = None
auth_credential.oauth2.access_token = None
auth_credential.oauth2.refresh_token = None
auth_credential.oauth2.expires_at = None
auth_credential.oauth2.expires_in = None
scheme_name = (
f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}"
f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}"
if auth_scheme
else ""
)
credential_name = (
f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}"
f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}"
if auth_credential
else ""
)
@@ -86,9 +129,19 @@ class ToolContextCredentialStore:
# session implementation, we don't want session to persist the token,
# meanwhile we want the token shared across runs.
serialized_credential = self.tool_context.state.get(token_key)
if not serialized_credential:
if serialized_credential:
return AuthCredential.model_validate(serialized_credential)
legacy_key = self._get_legacy_credential_key(auth_scheme, auth_credential)
if legacy_key == token_key:
return None
return AuthCredential.model_validate(serialized_credential)
serialized_legacy_credential = self.tool_context.state.get(legacy_key)
if not serialized_legacy_credential:
return None
# Migrate to the current key for future lookups.
self.tool_context.state[token_key] = serialized_legacy_credential
return AuthCredential.model_validate(serialized_legacy_credential)
def store_credential(
self,
@@ -114,6 +167,8 @@ class ToolAuthHandler:
auth_credential: Optional[AuthCredential],
credential_exchanger: Optional[BaseAuthCredentialExchanger] = None,
credential_store: Optional["ToolContextCredentialStore"] = None,
*,
credential_key: Optional[str] = None,
):
self.tool_context = tool_context
self.auth_scheme = (
@@ -122,12 +177,35 @@ class ToolAuthHandler:
self.auth_credential = (
auth_credential.model_copy(deep=True) if auth_credential else None
)
self._credential_key = credential_key
self.credential_exchanger = (
credential_exchanger or AutoAuthCredentialExchanger()
)
self.credential_store = credential_store
self.should_store_credential = True
def _get_credential_key_override(self) -> Optional[str]:
"""Returns a user-provided credential_key if available."""
if self._credential_key:
return self._credential_key
for obj in (self.auth_credential, self.auth_scheme):
if not obj or not obj.model_extra:
continue
for key in ("credential_key", "credentialKey"):
value = obj.model_extra.get(key)
if isinstance(value, str) and value:
return value
return None
def _build_auth_config(self) -> AuthConfig:
return AuthConfig(
auth_scheme=self.auth_scheme,
raw_auth_credential=self.auth_credential,
credential_key=self._get_credential_key_override(),
)
@classmethod
def from_tool_context(
cls,
@@ -135,6 +213,8 @@ class ToolAuthHandler:
auth_scheme: Optional[AuthScheme],
auth_credential: Optional[AuthCredential],
credential_exchanger: Optional[BaseAuthCredentialExchanger] = None,
*,
credential_key: Optional[str] = None,
) -> "ToolAuthHandler":
"""Creates a ToolAuthHandler instance from a ToolContext."""
credential_store = ToolContextCredentialStore(tool_context)
@@ -142,8 +222,9 @@ class ToolAuthHandler:
tool_context,
auth_scheme,
auth_credential,
credential_exchanger,
credential_store,
credential_key=credential_key,
credential_exchanger=credential_exchanger,
credential_store=credential_store,
)
async def _get_existing_credential(
@@ -209,21 +290,11 @@ class ToolAuthHandler:
"OAuth2 credentials client_secret is missing."
)
self.tool_context.request_credential(
AuthConfig(
auth_scheme=self.auth_scheme,
raw_auth_credential=self.auth_credential,
)
)
self.tool_context.request_credential(self._build_auth_config())
return None
def _get_auth_response(self) -> AuthCredential:
return self.tool_context.get_auth_response(
AuthConfig(
auth_scheme=self.auth_scheme,
raw_auth_credential=self.auth_credential,
)
)
return self.tool_context.get_auth_response(self._build_auth_config())
def _external_exchange_required(self, credential) -> bool:
return (
tests/unittests/auth/test_auth_config.py
+55
View File
@@ -12,6 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from pathlib import Path
import subprocess
import sys
from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from fastapi.openapi.models import OAuthFlows
@@ -107,3 +112,53 @@ def test_get_credential_key_with_extras(auth_config):
assert original_key == key
assert "extra_field" in auth_config.auth_scheme.model_extra
assert "extra_field" in auth_config.raw_auth_credential.model_extra
def test_credential_key_is_stable_across_python_hash_seed():
"""Test AuthConfig key generation does not depend on PYTHONHASHSEED."""
repo_root = Path(__file__).resolve().parents[3]
pythonpath = str(repo_root / "src")
code = "\n".join([
"from fastapi.openapi.models import OAuth2",
"from fastapi.openapi.models import OAuthFlowAuthorizationCode",
"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",
"from google.adk.auth.auth_tool import AuthConfig",
"",
"auth_scheme = OAuth2(",
" flows=OAuthFlows(",
" authorizationCode=OAuthFlowAuthorizationCode(",
" authorizationUrl='https://example.com/oauth2/authorize',",
" tokenUrl='https://example.com/oauth2/token',",
" scopes={'read': 'Read access'},",
" )",
" )",
")",
"auth_cred = AuthCredential(",
" auth_type=AuthCredentialTypes.OAUTH2,",
" oauth2=OAuth2Auth(",
" client_id='mock_client_id',",
" client_secret='mock_client_secret',",
" ),",
")",
"print(AuthConfig(",
" auth_scheme=auth_scheme,",
" raw_auth_credential=auth_cred,",
").credential_key)",
])
def _run_with_seed(seed: str) -> str:
env = os.environ.copy()
env["PYTHONHASHSEED"] = seed
env["PYTHONPATH"] = os.pathsep.join(
[pythonpath, env.get("PYTHONPATH", "")]
).strip(os.pathsep)
return subprocess.check_output(
[sys.executable, "-c", code],
env=env,
text=True,
).strip()
assert _run_with_seed("0") == _run_with_seed("1")
tests/unittests/auth/test_auth_handler.py
+27
View File
@@ -348,10 +348,12 @@ class TestGenerateAuthRequest:
exchanged_auth_credential=oauth2_credentials_with_auth_uri.model_copy(
deep=True
),
credential_key="my_tool_tokens",
)
handler = AuthHandler(config)
result = handler.generate_auth_request()
assert result.credential_key == "my_tool_tokens"
assert (
result.exchanged_auth_credential.oauth2.auth_uri
== oauth2_credentials_with_auth_uri.oauth2.auth_uri
@@ -400,6 +402,31 @@ class TestGenerateAuthRequest:
assert mock_generate_auth_uri.called
assert result.exchanged_auth_credential == mock_credential
@patch("google.adk.auth.auth_handler.AuthHandler.generate_auth_uri")
def test_preserves_credential_key_on_generated_request(
self, mock_generate_auth_uri, oauth2_auth_scheme, oauth2_credentials
):
"""Test that AuthHandler preserves an explicit credential_key."""
mock_generate_auth_uri.return_value = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="mock_client_id",
client_secret="mock_client_secret",
auth_uri="https://example.com/generated",
state="generated_state",
),
)
config = AuthConfig(
auth_scheme=oauth2_auth_scheme,
raw_auth_credential=oauth2_credentials,
credential_key="my_tool_tokens",
)
handler = AuthHandler(config)
result = handler.generate_auth_request()
assert result.credential_key == "my_tool_tokens"
class TestGetAuthResponse:
"""Tests for the get_auth_response method."""
tests/unittests/auth/test_credential_manager.py
+139 -10
View File
@@ -21,6 +21,8 @@ from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from fastapi.openapi.models import OAuthFlowImplicit
from fastapi.openapi.models import OAuthFlows
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import Agent
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import OAuth2Auth
@@ -33,8 +35,12 @@ from google.adk.auth.auth_tool import AuthConfig
from google.adk.auth.credential_manager import CredentialManager
from google.adk.auth.credential_manager import ServiceAccountCredentialExchanger
from google.adk.auth.oauth2_discovery import AuthorizationServerMetadata
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools.tool_context import ToolContext
import pytest
from .. import testing_utils
class TestCredentialManager:
"""Test suite for CredentialManager."""
@@ -211,8 +217,30 @@ class TestCredentialManager:
@pytest.mark.asyncio
async def test_save_credential_with_service(self):
"""Test _save_credential with credential service."""
auth_config = Mock(spec=AuthConfig)
mock_credential = Mock(spec=AuthCredential)
auth_scheme = OAuth2(
flows=OAuthFlows(
authorizationCode=OAuthFlowAuthorizationCode(
authorizationUrl="https://example.com/oauth2/authorize",
tokenUrl="https://example.com/oauth2/token",
scopes={"read": "Read access"},
)
)
)
auth_config = AuthConfig(
auth_scheme=auth_scheme,
raw_auth_credential=AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="mock_client_id",
client_secret="mock_client_secret",
),
),
credential_key="test_key",
)
mock_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(access_token="mock_access_token"),
)
# Mock credential service
credential_service = AsyncMock()
@@ -228,15 +256,39 @@ class TestCredentialManager:
manager = CredentialManager(auth_config)
await manager._save_credential(tool_context, mock_credential)
tool_context.save_credential.assert_called_once_with(auth_config)
assert auth_config.exchanged_auth_credential == mock_credential
tool_context.save_credential.assert_called_once()
saved_auth_config = tool_context.save_credential.call_args.args[0]
assert saved_auth_config.credential_key == "test_key"
assert saved_auth_config.exchanged_auth_credential == mock_credential
assert auth_config.exchanged_auth_credential is None
@pytest.mark.asyncio
async def test_save_credential_no_service(self):
"""Test _save_credential when no credential service is available."""
auth_config = Mock(spec=AuthConfig)
auth_config.exchanged_auth_credential = None
mock_credential = Mock(spec=AuthCredential)
auth_scheme = OAuth2(
flows=OAuthFlows(
authorizationCode=OAuthFlowAuthorizationCode(
authorizationUrl="https://example.com/oauth2/authorize",
tokenUrl="https://example.com/oauth2/token",
scopes={"read": "Read access"},
)
)
)
auth_config = AuthConfig(
auth_scheme=auth_scheme,
raw_auth_credential=AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="mock_client_id",
client_secret="mock_client_secret",
),
),
credential_key="test_key",
)
mock_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(access_token="mock_access_token"),
)
# Mock invocation context with no credential service
invocation_context = Mock()
@@ -248,9 +300,86 @@ class TestCredentialManager:
manager = CredentialManager(auth_config)
await manager._save_credential(tool_context, mock_credential)
# Should not raise an error, and credential should be set in auth_config
# even when there's no credential service (config is updated regardless)
assert auth_config.exchanged_auth_credential == mock_credential
assert auth_config.exchanged_auth_credential is None
@pytest.mark.asyncio
async def test_request_credential_does_not_leak_across_users(self):
"""Test that user-specific tokens are not cached on shared tool configs."""
auth_scheme = OAuth2(
flows=OAuthFlows(
authorizationCode=OAuthFlowAuthorizationCode(
authorizationUrl="https://example.com/oauth2/authorize",
tokenUrl="https://example.com/oauth2/token",
scopes={"read": "Read access"},
)
)
)
auth_config = AuthConfig(
auth_scheme=auth_scheme,
raw_auth_credential=AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="mock_client_id",
client_secret="mock_client_secret",
),
),
credential_key="shared_key",
)
manager = CredentialManager(auth_config)
agent = Agent(
name="root_agent",
model=testing_utils.MockModel.create(responses=[]),
tools=[],
)
session_service = InMemorySessionService()
session_a = await session_service.create_session(
app_name="test_app",
user_id="user_a",
session_id="session_a",
)
session_b = await session_service.create_session(
app_name="test_app",
user_id="user_b",
session_id="session_b",
)
invocation_context_a = InvocationContext(
session_service=session_service,
invocation_id="invocation_a",
agent=agent,
session=session_a,
credential_service=None,
)
invocation_context_b = InvocationContext(
session_service=session_service,
invocation_id="invocation_b",
agent=agent,
session=session_b,
credential_service=None,
)
tool_context_a = ToolContext(
invocation_context_a, function_call_id="call_a"
)
tool_context_b = ToolContext(
invocation_context_b, function_call_id="call_b"
)
tool_context_a.state["temp:" + auth_config.credential_key] = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
auth_uri="https://example.com/oauth2/authorize?x=y",
state="state_a",
access_token="token_a",
expires_at=9_999_999_999,
),
)
await manager.get_auth_credential(tool_context_a)
await manager.request_credential(tool_context_b)
requested = tool_context_b.actions.requested_auth_configs["call_b"]
assert requested.exchanged_auth_credential.oauth2.access_token is None
@pytest.mark.asyncio
async def test_refresh_credential_oauth2(self):
tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py
+17
View File
@@ -138,6 +138,23 @@ async def test_openid_connect_no_auth_response(
assert result.auth_credential == openid_connect_credential
@pytest.mark.asyncio
async def test_openid_connect_uses_explicit_credential_key(
openid_connect_scheme, openid_connect_credential
):
tool_context = create_mock_tool_context()
handler = ToolAuthHandler(
tool_context,
openid_connect_scheme,
openid_connect_credential,
credential_key='my_tool_tokens',
)
result = await handler.prepare_auth_credentials()
assert result.state == 'pending'
requested = tool_context.actions.requested_auth_configs['test-fc-id']
assert requested.credential_key == 'my_tool_tokens'
@pytest.mark.asyncio
async def test_openid_connect_with_auth_response(
openid_connect_scheme, openid_connect_credential, monkeypatch