diff --git a/src/google/adk/agents/callback_context.py b/src/google/adk/agents/callback_context.py index d89a3af2..c9ca750d 100644 --- a/src/google/adk/agents/callback_context.py +++ b/src/google/adk/agents/callback_context.py @@ -177,6 +177,25 @@ class CallbackContext(ReadonlyContext): auth_config, self ) + def get_auth_response( + self, auth_config: AuthConfig + ) -> Optional[AuthCredential]: + """Gets the auth response credential from session state. + + This method retrieves an authentication credential that was previously + stored in session state after a user completed an OAuth flow or other + authentication process. + + Args: + auth_config: The authentication configuration for the credential. + + Returns: + The auth credential from the auth response, or None if not found. + """ + from ..auth.auth_handler import AuthHandler + + return AuthHandler(auth_config).get_auth_response(self.state) + async def add_session_to_memory(self) -> None: """Triggers memory generation for the current session. diff --git a/src/google/adk/auth/auth_preprocessor.py b/src/google/adk/auth/auth_preprocessor.py index f8885aaf..1c3ac70d 100644 --- a/src/google/adk/auth/auth_preprocessor.py +++ b/src/google/adk/auth/auth_preprocessor.py @@ -33,6 +33,11 @@ from .auth_tool import AuthToolArguments if TYPE_CHECKING: from ..agents.llm_agent import LlmAgent +# Prefix used by toolset auth credential IDs. +# Auth requests with this prefix are for toolset authentication (before tool +# listing) and don't require resuming a function call. +TOOLSET_AUTH_CREDENTIAL_ID_PREFIX = '_adk_toolset_auth_' + class _AuthLlmRequestProcessor(BaseLlmRequestProcessor): """Handles auth information to build the LLM request.""" @@ -96,6 +101,11 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor): continue args = AuthToolArguments.model_validate(function_call.args) + # Skip toolset auth - auth response is already stored in session state + # and we don't need to resume a function call for toolsets + if args.function_call_id.startswith(TOOLSET_AUTH_CREDENTIAL_ID_PREFIX): + continue + tools_to_resume.add(args.function_call_id) if not tools_to_resume: continue diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index bf0ed1e2..21b6328a 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -19,8 +19,8 @@ from typing import Optional from fastapi.openapi.models import OAuth2 +from ..agents.callback_context import CallbackContext from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger -from ..tools.tool_context import ToolContext from ..utils.feature_decorator import experimental from .auth_credential import AuthCredential from .auth_credential import AuthCredentialTypes @@ -124,11 +124,16 @@ class CredentialManager: """ self._exchanger_registry.register(credential_type, exchanger_instance) - async def request_credential(self, tool_context: ToolContext) -> None: - tool_context.request_credential(self._auth_config) + async def request_credential(self, context: CallbackContext) -> None: + if not hasattr(context, "request_credential"): + raise TypeError( + "request_credential requires a ToolContext with request_credential" + " method, not a plain CallbackContext" + ) + context.request_credential(self._auth_config) async def get_auth_credential( - self, tool_context: ToolContext + self, context: CallbackContext ) -> Optional[AuthCredential]: """Load and prepare authentication credential through a structured workflow.""" @@ -140,14 +145,14 @@ class CredentialManager: return self._auth_config.raw_auth_credential # Step 3: Try to load existing processed credential - credential = await self._load_existing_credential(tool_context) + credential = await self._load_existing_credential(context) # Step 4: If no existing credential, load from auth response # TODO instead of load from auth response, we can store auth response in # credential service. was_from_auth_response = False if not credential: - credential = await self._load_from_auth_response(tool_context) + credential = await self._load_from_auth_response(context) was_from_auth_response = True # Step 5: If still no credential available, check if client credentials @@ -169,38 +174,38 @@ class CredentialManager: # Step 8: Save credential if it was modified if was_from_auth_response or was_exchanged or was_refreshed: - await self._save_credential(tool_context, credential) + await self._save_credential(context, credential) return credential async def _load_existing_credential( - self, tool_context: ToolContext + self, context: CallbackContext ) -> Optional[AuthCredential]: """Load existing credential from credential service.""" # Try loading from credential service first - credential = await self._load_from_credential_service(tool_context) + credential = await self._load_from_credential_service(context) if credential: return credential return None async def _load_from_credential_service( - self, tool_context: ToolContext + self, context: CallbackContext ) -> Optional[AuthCredential]: """Load credential from credential service if available.""" - credential_service = tool_context._invocation_context.credential_service + credential_service = context._invocation_context.credential_service if credential_service: # Note: This should be made async in a future refactor # For now, assuming synchronous operation - return await tool_context.load_credential(self._auth_config) + return await context.load_credential(self._auth_config) return None async def _load_from_auth_response( - self, tool_context: ToolContext + self, context: CallbackContext ) -> Optional[AuthCredential]: - """Load credential from auth response in tool context.""" - return tool_context.get_auth_response(self._auth_config) + """Load credential from auth response in context.""" + return context.get_auth_response(self._auth_config) async def _exchange_credential( self, credential: AuthCredential @@ -290,15 +295,15 @@ class CredentialManager: # Additional validation can be added here async def _save_credential( - self, tool_context: ToolContext, credential: AuthCredential + 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 = tool_context._invocation_context.credential_service + credential_service = context._invocation_context.credential_service if credential_service: - await tool_context.save_credential(self._auth_config) + await context.save_credential(self._auth_config) async def _populate_auth_scheme(self) -> bool: """Auto-discover server metadata and populate missing auth scheme info. diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index f9f80e6c..d4af6937 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -37,6 +37,9 @@ from ...agents.live_request_queue import LiveRequestQueue from ...agents.readonly_context import ReadonlyContext from ...agents.run_config import StreamingMode from ...agents.transcription_entry import TranscriptionEntry +from ...auth.auth_handler import AuthHandler +from ...auth.auth_tool import AuthConfig +from ...auth.credential_manager import CredentialManager from ...events.event import Event from ...models.base_llm_connection import BaseLlmConnection from ...models.llm_request import LlmRequest @@ -50,6 +53,11 @@ from ...tools.google_search_tool import google_search from ...tools.tool_context import ToolContext from ...utils.context_utils import Aclosing from .audio_cache_manager import AudioCacheManager +from .functions import build_auth_request_event +from .functions import REQUEST_EUC_FUNCTION_CALL_NAME + +# Prefix used by toolset auth credential IDs +TOOLSET_AUTH_CREDENTIAL_ID_PREFIX = '_adk_toolset_auth_' if TYPE_CHECKING: from ...agents.llm_agent import LlmAgent @@ -528,6 +536,17 @@ class BaseLlmFlow(ABC): async for event in agen: yield event + # Resolve toolset authentication before tool listing. + # This ensures credentials are ready before get_tools() is called. + async with Aclosing( + self._resolve_toolset_auth(invocation_context, agent) + ) as agen: + async for event in agen: + yield event + + if invocation_context.end_invocation: + return + # Run processors for tools. # We may need to wrap some built-in tools if there are other tools @@ -561,6 +580,81 @@ class BaseLlmFlow(ABC): tool_context=tool_context, llm_request=llm_request ) + async def _resolve_toolset_auth( + self, + invocation_context: InvocationContext, + agent: LlmAgent, + ) -> AsyncGenerator[Event, None]: + """Resolves authentication for toolsets before tool listing. + + For each toolset with auth configured via get_auth_config(): + - If credential is available, populate auth_config.exchanged_auth_credential + - If credential is not available, yield auth request event and interrupt + + Args: + invocation_context: The invocation context. + agent: The LLM agent. + + Yields: + Auth request events if any toolset needs authentication. + """ + if not agent.tools: + return + + pending_auth_requests: dict[str, AuthConfig] = {} + callback_context = CallbackContext(invocation_context) + + for tool_union in agent.tools: + if not isinstance(tool_union, BaseToolset): + continue + + auth_config = tool_union.get_auth_config() + if not auth_config: + continue + + try: + credential = await CredentialManager(auth_config).get_auth_credential( + callback_context + ) + except ValueError as e: + # Validation errors from CredentialManager should be logged but not + # block the flow - the toolset may still work without auth + logger.warning( + 'Failed to get auth credential for toolset %s: %s', + type(tool_union).__name__, + e, + ) + credential = None + + if credential: + # Populate in-place for toolset to use in get_tools() + auth_config.exchanged_auth_credential = credential + else: + # Need auth - will interrupt + toolset_id = ( + f'{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}{type(tool_union).__name__}' + ) + pending_auth_requests[toolset_id] = auth_config + + if not pending_auth_requests: + return + + # Build auth requests dict with generated auth requests + auth_requests = { + credential_id: AuthHandler(auth_config).generate_auth_request() + for credential_id, auth_config in pending_auth_requests.items() + } + + # Yield event with auth requests using the shared helper + yield build_auth_request_event( + invocation_context, + auth_requests, + author=agent.name, + ) + + # Interrupt invocation + invocation_context.end_invocation = True + async def _postprocess_async( self, invocation_context: InvocationContext, diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 40412efe..664d0d6c 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -26,6 +26,7 @@ import threading from typing import Any from typing import AsyncGenerator from typing import cast +from typing import Dict from typing import Optional from typing import TYPE_CHECKING import uuid @@ -34,6 +35,7 @@ from google.genai import types from ...agents.active_streaming_tool import ActiveStreamingTool from ...agents.invocation_context import InvocationContext +from ...auth.auth_tool import AuthConfig from ...auth.auth_tool import AuthToolArguments from ...events.event import Event from ...events.event_actions import EventActions @@ -211,41 +213,77 @@ def get_long_running_function_calls( return long_running_tool_ids -def generate_auth_event( +def build_auth_request_event( invocation_context: InvocationContext, - function_response_event: Event, -) -> Optional[Event]: - if not function_response_event.actions.requested_auth_configs: - return None + auth_requests: Dict[str, AuthConfig], + *, + author: Optional[str] = None, + role: Optional[str] = None, +) -> Event: + """Builds an auth request event with function calls for each auth request. + + This is a shared helper used by both tool-level auth (when a tool requests + auth during execution) and toolset-level auth (before tool listing). + + Args: + invocation_context: The invocation context. + auth_requests: Dict mapping function_call_id to AuthConfig. + author: The event author. Defaults to agent name. + role: The content role. Defaults to None. + + Returns: + Event with auth request function calls. + """ parts = [] long_running_tool_ids = set() - for ( - function_call_id, - auth_config, - ) in function_response_event.actions.requested_auth_configs.items(): + for function_call_id, auth_config in auth_requests.items(): request_euc_function_call = types.FunctionCall( name=REQUEST_EUC_FUNCTION_CALL_NAME, + id=generate_client_function_call_id(), args=AuthToolArguments( function_call_id=function_call_id, auth_config=auth_config, ).model_dump(exclude_none=True, by_alias=True), ) - request_euc_function_call.id = generate_client_function_call_id() long_running_tool_ids.add(request_euc_function_call.id) parts.append(types.Part(function_call=request_euc_function_call)) return Event( invocation_id=invocation_context.invocation_id, - author=invocation_context.agent.name, + author=author or invocation_context.agent.name, branch=invocation_context.branch, - content=types.Content( - parts=parts, role=function_response_event.content.role - ), + content=types.Content(parts=parts, role=role), long_running_tool_ids=long_running_tool_ids, ) +def generate_auth_event( + invocation_context: InvocationContext, + function_response_event: Event, +) -> Optional[Event]: + """Generates an auth request event from a function response event. + + This is used for tool-level auth where a tool requests credentials during + execution. + + Args: + invocation_context: The invocation context. + function_response_event: The function response event with auth requests. + + Returns: + Event with auth request function calls, or None if no auth requested. + """ + if not function_response_event.actions.requested_auth_configs: + return None + + return build_auth_request_event( + invocation_context, + function_response_event.actions.requested_auth_configs, + role=function_response_event.content.role, + ) + + def generate_request_confirmation_event( invocation_context: InvocationContext, function_call_event: Event, diff --git a/tests/unittests/auth/test_toolset_auth.py b/tests/unittests/auth/test_toolset_auth.py new file mode 100644 index 00000000..bd4d8f2a --- /dev/null +++ b/tests/unittests/auth/test_toolset_auth.py @@ -0,0 +1,414 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for toolset authentication functionality.""" + +from typing import Optional +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import Mock +from unittest.mock import patch + +from fastapi.openapi.models import OAuth2 +from fastapi.openapi.models import OAuthFlowAuthorizationCode +from fastapi.openapi.models import OAuthFlows +from google.adk.agents.callback_context import CallbackContext +from google.adk.agents.invocation_context import InvocationContext +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_preprocessor import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX +from google.adk.auth.auth_tool import AuthConfig +from google.adk.auth.auth_tool import AuthToolArguments +from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow +from google.adk.flows.llm_flows.base_llm_flow import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX as FLOW_PREFIX +from google.adk.flows.llm_flows.functions import build_auth_request_event +from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.base_toolset import BaseToolset +import pytest + + +class MockToolset(BaseToolset): + """A mock toolset for testing.""" + + def __init__( + self, + auth_config: Optional[AuthConfig] = None, + tools: Optional[list[BaseTool]] = None, + ): + super().__init__() + self._auth_config = auth_config + self._tools = tools or [] + + def get_auth_config(self) -> Optional[AuthConfig]: + return self._auth_config + + async def get_tools(self, readonly_context=None) -> list[BaseTool]: + return self._tools + + async def close(self): + pass + + +def create_oauth2_auth_config() -> AuthConfig: + """Create a sample OAuth2 auth config for testing.""" + return AuthConfig( + auth_scheme=OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={"read": "Read access"}, + ) + ) + ), + raw_auth_credential=AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test_client_id", + client_secret="test_client_secret", + ), + ), + ) + + +class TestToolsetAuthPrefixConstant: + """Test that prefix constants are consistent.""" + + def test_prefix_constants_match(self): + """Ensure auth_preprocessor and base_llm_flow use the same prefix.""" + assert TOOLSET_AUTH_CREDENTIAL_ID_PREFIX == FLOW_PREFIX + assert TOOLSET_AUTH_CREDENTIAL_ID_PREFIX == "_adk_toolset_auth_" + + +class TestResolveToolsetAuth: + """Tests for _resolve_toolset_auth method in BaseLlmFlow.""" + + @pytest.fixture + def mock_invocation_context(self): + """Create a mock invocation context.""" + ctx = Mock(spec=InvocationContext) + ctx.invocation_id = "test-invocation-id" + ctx.end_invocation = False + ctx.branch = None + ctx.session = Mock() + ctx.session.state = {} + ctx.session.id = "test-session-id" + ctx.credential_service = None + ctx.app_name = "test-app" + ctx.user_id = "test-user" + return ctx + + @pytest.fixture + def mock_agent(self): + """Create a mock LLM agent.""" + agent = Mock() + agent.name = "test-agent" + agent.tools = [] + return agent + + @pytest.fixture + def flow(self): + """Create a BaseLlmFlow instance for testing.""" + # BaseLlmFlow is abstract, but we can still test _resolve_toolset_auth + flow = Mock(spec=BaseLlmFlow) + flow._resolve_toolset_auth = BaseLlmFlow._resolve_toolset_auth + return flow + + @pytest.mark.asyncio + async def test_no_tools_returns_no_events( + self, mock_invocation_context, mock_agent + ): + """Test that no events are yielded when agent has no tools.""" + mock_agent.tools = [] + + flow = BaseLlmFlow.__new__(BaseLlmFlow) + events = [] + async for event in flow._resolve_toolset_auth( + mock_invocation_context, mock_agent + ): + events.append(event) + + assert len(events) == 0 + assert mock_invocation_context.end_invocation is False + + @pytest.mark.asyncio + async def test_toolset_without_auth_config_skipped( + self, mock_invocation_context, mock_agent + ): + """Test that toolsets without auth config are skipped.""" + toolset = MockToolset(auth_config=None) + mock_agent.tools = [toolset] + + flow = BaseLlmFlow.__new__(BaseLlmFlow) + events = [] + async for event in flow._resolve_toolset_auth( + mock_invocation_context, mock_agent + ): + events.append(event) + + assert len(events) == 0 + assert mock_invocation_context.end_invocation is False + + @pytest.mark.asyncio + async def test_toolset_with_credential_available_populates_config( + self, mock_invocation_context, mock_agent + ): + """Test that credential is populated in auth_config when available.""" + auth_config = create_oauth2_auth_config() + toolset = MockToolset(auth_config=auth_config) + mock_agent.tools = [toolset] + + # Mock CredentialManager to return a credential + mock_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token="test-token"), + ) + + with patch( + "google.adk.flows.llm_flows.base_llm_flow.CredentialManager" + ) as MockCredentialManager: + mock_manager = AsyncMock() + mock_manager.get_auth_credential = AsyncMock(return_value=mock_credential) + MockCredentialManager.return_value = mock_manager + + flow = BaseLlmFlow.__new__(BaseLlmFlow) + events = [] + async for event in flow._resolve_toolset_auth( + mock_invocation_context, mock_agent + ): + events.append(event) + + # No auth request events - credential was available + assert len(events) == 0 + assert mock_invocation_context.end_invocation is False + # Credential should be populated in auth_config + assert auth_config.exchanged_auth_credential == mock_credential + + @pytest.mark.asyncio + async def test_toolset_without_credential_yields_auth_event( + self, mock_invocation_context, mock_agent + ): + """Test that auth request event is yielded when credential not available.""" + auth_config = create_oauth2_auth_config() + toolset = MockToolset(auth_config=auth_config) + mock_agent.tools = [toolset] + + with patch( + "google.adk.flows.llm_flows.base_llm_flow.CredentialManager" + ) as MockCredentialManager: + mock_manager = AsyncMock() + mock_manager.get_auth_credential = AsyncMock(return_value=None) + MockCredentialManager.return_value = mock_manager + + flow = BaseLlmFlow.__new__(BaseLlmFlow) + events = [] + async for event in flow._resolve_toolset_auth( + mock_invocation_context, mock_agent + ): + events.append(event) + + # Should yield one auth request event + assert len(events) == 1 + assert mock_invocation_context.end_invocation is True + + # Check event structure + event = events[0] + assert event.invocation_id == "test-invocation-id" + assert event.author == "test-agent" + assert event.content is not None + assert len(event.content.parts) == 1 + + # Check function call + fc = event.content.parts[0].function_call + assert fc.name == REQUEST_EUC_FUNCTION_CALL_NAME + # The args use camelCase aliases from the pydantic model + assert fc.args["functionCallId"].startswith( + TOOLSET_AUTH_CREDENTIAL_ID_PREFIX + ) + assert "MockToolset" in fc.args["functionCallId"] + + @pytest.mark.asyncio + async def test_multiple_toolsets_needing_auth( + self, mock_invocation_context, mock_agent + ): + """Test that multiple toolsets needing auth yield multiple function calls.""" + auth_config1 = create_oauth2_auth_config() + auth_config2 = create_oauth2_auth_config() + toolset1 = MockToolset(auth_config=auth_config1) + toolset2 = MockToolset(auth_config=auth_config2) + mock_agent.tools = [toolset1, toolset2] + + with patch( + "google.adk.flows.llm_flows.base_llm_flow.CredentialManager" + ) as MockCredentialManager: + mock_manager = AsyncMock() + mock_manager.get_auth_credential = AsyncMock(return_value=None) + MockCredentialManager.return_value = mock_manager + + flow = BaseLlmFlow.__new__(BaseLlmFlow) + events = [] + async for event in flow._resolve_toolset_auth( + mock_invocation_context, mock_agent + ): + events.append(event) + + # Should yield one event with multiple function calls + # But since both toolsets have same class name, they'll have same ID + # and only one will be in pending_auth_requests (dict overwrites) + assert len(events) == 1 + assert mock_invocation_context.end_invocation is True + + +class TestAuthPreprocessorToolsetAuthSkip: + """Tests for auth preprocessor skipping toolset auth.""" + + def test_toolset_auth_prefix_skipped(self): + """Test that function calls with toolset auth prefix are skipped.""" + from google.adk.auth.auth_preprocessor import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX + + # Verify the prefix is correct + assert TOOLSET_AUTH_CREDENTIAL_ID_PREFIX == "_adk_toolset_auth_" + + # Test that a function_call_id starting with this prefix would be skipped + toolset_function_call_id = f"{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}McpToolset" + assert toolset_function_call_id.startswith( + TOOLSET_AUTH_CREDENTIAL_ID_PREFIX + ) + + # Regular tool auth function_call_id should NOT start with prefix + regular_function_call_id = "call_123" + assert not regular_function_call_id.startswith( + TOOLSET_AUTH_CREDENTIAL_ID_PREFIX + ) + + +class TestCallbackContextGetAuthResponse: + """Tests for CallbackContext.get_auth_response method.""" + + @pytest.fixture + def mock_invocation_context(self): + """Create a mock invocation context.""" + ctx = Mock(spec=InvocationContext) + ctx.session = Mock() + ctx.session.state = {} + return ctx + + def test_get_auth_response_returns_none_when_no_response( + self, mock_invocation_context + ): + """Test that get_auth_response returns None when no auth response in state.""" + callback_context = CallbackContext(mock_invocation_context) + auth_config = create_oauth2_auth_config() + + result = callback_context.get_auth_response(auth_config) + + # Should return None when no auth response is stored + assert result is None + + def test_get_auth_response_delegates_to_auth_handler( + self, mock_invocation_context + ): + """Test that get_auth_response delegates to AuthHandler.""" + callback_context = CallbackContext(mock_invocation_context) + auth_config = create_oauth2_auth_config() + + # AuthHandler is imported inside the method, so we patch the module + with patch("google.adk.auth.auth_handler.AuthHandler") as MockAuthHandler: + mock_handler = Mock() + mock_handler.get_auth_response = Mock(return_value=None) + MockAuthHandler.return_value = mock_handler + + callback_context.get_auth_response(auth_config) + + MockAuthHandler.assert_called_once_with(auth_config) + mock_handler.get_auth_response.assert_called_once() + + +class TestBuildAuthRequestEvent: + """Tests for build_auth_request_event helper function.""" + + @pytest.fixture + def mock_invocation_context(self): + """Create a mock invocation context.""" + ctx = Mock(spec=InvocationContext) + ctx.invocation_id = "test-invocation-id" + ctx.branch = None + ctx.agent = Mock() + ctx.agent.name = "test-agent" + return ctx + + def test_builds_event_with_auth_requests(self, mock_invocation_context): + """Test that build_auth_request_event creates correct event.""" + auth_requests = { + "call_123": create_oauth2_auth_config(), + } + + event = build_auth_request_event(mock_invocation_context, auth_requests) + + assert event.invocation_id == "test-invocation-id" + assert event.author == "test-agent" + assert event.content is not None + assert len(event.content.parts) == 1 + + fc = event.content.parts[0].function_call + assert fc.name == REQUEST_EUC_FUNCTION_CALL_NAME + assert fc.args["functionCallId"] == "call_123" + + def test_multiple_auth_requests_create_multiple_parts( + self, mock_invocation_context + ): + """Test that multiple auth requests create multiple function call parts.""" + auth_requests = { + "call_1": create_oauth2_auth_config(), + "call_2": create_oauth2_auth_config(), + } + + event = build_auth_request_event(mock_invocation_context, auth_requests) + + assert len(event.content.parts) == 2 + function_call_ids = { + p.function_call.args["functionCallId"] for p in event.content.parts + } + assert function_call_ids == {"call_1", "call_2"} + + def test_always_adds_long_running_tool_ids(self, mock_invocation_context): + """Test that long_running_tool_ids is always set.""" + auth_requests = {"call_123": create_oauth2_auth_config()} + + event = build_auth_request_event(mock_invocation_context, auth_requests) + + assert event.long_running_tool_ids is not None + assert len(event.long_running_tool_ids) == 1 + + def test_custom_author_overrides_default(self, mock_invocation_context): + """Test that custom author overrides default agent name.""" + auth_requests = {"call_123": create_oauth2_auth_config()} + + event = build_auth_request_event( + mock_invocation_context, auth_requests, author="custom-author" + ) + + assert event.author == "custom-author" + + def test_role_is_set_in_content(self, mock_invocation_context): + """Test that role is set in content.""" + auth_requests = {"call_123": create_oauth2_auth_config()} + + event = build_auth_request_event( + mock_invocation_context, auth_requests, role="model" + ) + + assert event.content.role == "model"