diff --git a/src/google/adk/auth/auth_preprocessor.py b/src/google/adk/auth/auth_preprocessor.py index b0677497..133b456b 100644 --- a/src/google/adk/auth/auth_preprocessor.py +++ b/src/google/adk/auth/auth_preprocessor.py @@ -51,26 +51,33 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor): return request_euc_function_call_ids = set() - for k in range(len(events) - 1, -1, -1): - event = events[k] - # look for first event authored by user - if not event.author or event.author != 'user': - continue - responses = event.get_function_responses() - if not responses: - return + # find the last event with non-None content + last_event_with_content = None + for i in range(len(events) - 1, -1, -1): + event = events[i] + if event.content is not None: + last_event_with_content = event + break - for function_call_response in responses: - if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME: - continue - # found the function call response for the system long running request euc - # function call - request_euc_function_call_ids.add(function_call_response.id) - auth_config = AuthConfig.model_validate(function_call_response.response) - await AuthHandler( - auth_config=auth_config - ).parse_and_store_auth_response(state=invocation_context.session.state) - break + # check if the last event with content is authored by user + if not last_event_with_content or last_event_with_content.author != 'user': + return + + responses = last_event_with_content.get_function_responses() + if not responses: + return + + # look for auth response + for function_call_response in responses: + if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME: + continue + # found the function call response for the system long running request euc + # function call + request_euc_function_call_ids.add(function_call_response.id) + auth_config = AuthConfig.model_validate(function_call_response.response) + await AuthHandler(auth_config=auth_config).parse_and_store_auth_response( + state=invocation_context.session.state + ) if not request_euc_function_call_ids: return diff --git a/tests/unittests/auth/test_auth_preprocessor.py b/tests/unittests/auth/test_auth_preprocessor.py new file mode 100644 index 00000000..9b589db2 --- /dev/null +++ b/tests/unittests/auth/test_auth_preprocessor.py @@ -0,0 +1,541 @@ +# Copyright 2025 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. + +"""Unit tests for auth_preprocessor module.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock +from unittest.mock import Mock +from unittest.mock import patch + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.auth.auth_handler import AuthHandler +from google.adk.auth.auth_preprocessor import _AuthLlmRequestProcessor +from google.adk.auth.auth_tool import AuthConfig +from google.adk.auth.auth_tool import AuthToolArguments +from google.adk.events.event import Event +from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME +from google.adk.models.llm_request import LlmRequest +import pytest + + +class TestAuthLlmRequestProcessor: + """Tests for _AuthLlmRequestProcessor class.""" + + @pytest.fixture + def processor(self): + """Create an _AuthLlmRequestProcessor instance.""" + return _AuthLlmRequestProcessor() + + @pytest.fixture + def mock_llm_agent(self): + """Create a mock LlmAgent.""" + from google.adk.agents.llm_agent import LlmAgent + + agent = Mock(spec=LlmAgent) + agent.canonical_tools = AsyncMock(return_value=[]) + return agent + + @pytest.fixture + def mock_non_llm_agent(self): + """Create a mock non-LLM agent.""" + agent = Mock() + agent.__class__.__name__ = 'BaseAgent' + return agent + + @pytest.fixture + def mock_session(self): + """Create a mock session.""" + session = Mock() + session.state = {} + session.events = [] + return session + + @pytest.fixture + def mock_invocation_context(self, mock_llm_agent, mock_session): + """Create a mock invocation context.""" + context = Mock(spec=InvocationContext) + context.agent = mock_llm_agent + context.session = mock_session + return context + + @pytest.fixture + def mock_llm_request(self): + """Create a mock LlmRequest.""" + return Mock(spec=LlmRequest) + + @pytest.fixture + def mock_auth_config(self): + """Create a mock AuthConfig.""" + return Mock(spec=AuthConfig) + + @pytest.fixture + def mock_function_response_with_auth(self, mock_auth_config): + """Create a mock function response with auth data.""" + function_response = Mock() + function_response.name = REQUEST_EUC_FUNCTION_CALL_NAME + function_response.id = 'auth_response_id' + function_response.response = mock_auth_config + return function_response + + @pytest.fixture + def mock_function_response_without_auth(self): + """Create a mock function response without auth data.""" + function_response = Mock() + function_response.name = 'some_other_function' + function_response.id = 'other_response_id' + return function_response + + @pytest.fixture + def mock_user_event_with_auth_response( + self, mock_function_response_with_auth + ): + """Create a mock user event with auth response.""" + event = Mock(spec=Event) + event.author = 'user' + event.content = Mock() # Non-None content + event.get_function_responses.return_value = [ + mock_function_response_with_auth + ] + return event + + @pytest.fixture + def mock_user_event_without_auth_response( + self, mock_function_response_without_auth + ): + """Create a mock user event without auth response.""" + event = Mock(spec=Event) + event.author = 'user' + event.content = Mock() # Non-None content + event.get_function_responses.return_value = [ + mock_function_response_without_auth + ] + return event + + @pytest.fixture + def mock_user_event_no_responses(self): + """Create a mock user event with no responses.""" + event = Mock(spec=Event) + event.author = 'user' + event.content = Mock() # Non-None content + event.get_function_responses.return_value = [] + return event + + @pytest.fixture + def mock_agent_event(self): + """Create a mock agent-authored event.""" + event = Mock(spec=Event) + event.author = 'test_agent' + event.content = Mock() # Non-None content + return event + + @pytest.fixture + def mock_event_no_content(self): + """Create a mock event with no content.""" + event = Mock(spec=Event) + event.author = 'user' + event.content = None + return event + + @pytest.fixture + def mock_agent_event_with_content(self): + """Create a mock agent event with content.""" + event = Mock(spec=Event) + event.author = 'test_agent' + event.content = Mock() # Non-None content + return event + + @pytest.mark.asyncio + async def test_non_llm_agent_returns_early( + self, processor, mock_llm_request, mock_session + ): + """Test that non-LLM agents return early.""" + mock_context = Mock(spec=InvocationContext) + mock_context.agent = Mock() + mock_context.agent.__class__.__name__ = 'BaseAgent' + mock_context.session = mock_session + + result = [] + async for event in processor.run_async(mock_context, mock_llm_request): + result.append(event) + + assert result == [] + + @pytest.mark.asyncio + async def test_empty_events_returns_early( + self, processor, mock_invocation_context, mock_llm_request + ): + """Test that empty events list returns early.""" + mock_invocation_context.session.events = [] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + assert result == [] + + @pytest.mark.asyncio + async def test_no_events_with_content_returns_early( + self, + processor, + mock_invocation_context, + mock_llm_request, + mock_event_no_content, + ): + """Test that no events with content returns early.""" + mock_invocation_context.session.events = [mock_event_no_content] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + assert result == [] + + @pytest.mark.asyncio + async def test_last_event_with_content_not_user_authored_returns_early( + self, + processor, + mock_invocation_context, + mock_llm_request, + mock_event_no_content, + mock_agent_event_with_content, + ): + """Test that last event with content not user-authored returns early.""" + # Mix of events: user event with no content, then agent event with content + mock_invocation_context.session.events = [ + mock_event_no_content, + mock_agent_event_with_content, + ] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + assert result == [] + + @pytest.mark.asyncio + async def test_last_event_no_responses_returns_early( + self, + processor, + mock_invocation_context, + mock_llm_request, + mock_user_event_no_responses, + ): + """Test that user event with no responses returns early.""" + mock_invocation_context.session.events = [mock_user_event_no_responses] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + assert result == [] + + @pytest.mark.asyncio + async def test_last_event_no_auth_responses_returns_early( + self, + processor, + mock_invocation_context, + mock_llm_request, + mock_user_event_without_auth_response, + ): + """Test that user event with non-auth responses returns early.""" + mock_invocation_context.session.events = [ + mock_user_event_without_auth_response + ] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + assert result == [] + + @pytest.mark.asyncio + @patch('google.adk.auth.auth_preprocessor.AuthHandler') + @patch('google.adk.auth.auth_tool.AuthConfig.model_validate') + async def test_processes_auth_response_successfully( + self, + mock_auth_config_validate, + mock_auth_handler_class, + processor, + mock_invocation_context, + mock_llm_request, + mock_user_event_with_auth_response, + mock_auth_config, + ): + """Test successful processing of auth response in last event.""" + # Setup mocks + mock_auth_config_validate.return_value = mock_auth_config + mock_auth_handler = Mock(spec=AuthHandler) + mock_auth_handler.parse_and_store_auth_response = AsyncMock() + mock_auth_handler_class.return_value = mock_auth_handler + + mock_invocation_context.session.events = [ + mock_user_event_with_auth_response + ] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + # Verify auth config validation was called + mock_auth_config_validate.assert_called_once() + + # Verify auth handler was created with the config + mock_auth_handler_class.assert_called_once_with( + auth_config=mock_auth_config + ) + + # Verify parse_and_store_auth_response was called + mock_auth_handler.parse_and_store_auth_response.assert_called_once_with( + state=mock_invocation_context.session.state + ) + + @pytest.mark.asyncio + @patch('google.adk.auth.auth_preprocessor.AuthHandler') + @patch('google.adk.auth.auth_tool.AuthConfig.model_validate') + @patch('google.adk.flows.llm_flows.functions.handle_function_calls_async') + async def test_processes_multiple_auth_responses_and_resumes_tools( + self, + mock_handle_function_calls, + mock_auth_config_validate, + mock_auth_handler_class, + processor, + mock_invocation_context, + mock_llm_request, + mock_auth_config, + ): + """Test processing multiple auth responses and resuming tools.""" + # Create multiple auth responses + auth_response_1 = Mock() + auth_response_1.name = REQUEST_EUC_FUNCTION_CALL_NAME + auth_response_1.id = 'auth_id_1' + auth_response_1.response = mock_auth_config + + auth_response_2 = Mock() + auth_response_2.name = REQUEST_EUC_FUNCTION_CALL_NAME + auth_response_2.id = 'auth_id_2' + auth_response_2.response = mock_auth_config + + user_event_with_multiple_responses = Mock(spec=Event) + user_event_with_multiple_responses.author = 'user' + user_event_with_multiple_responses.content = Mock() # Non-None content + user_event_with_multiple_responses.get_function_responses.return_value = [ + auth_response_1, + auth_response_2, + ] + + # Create system function call events + system_function_call_1 = Mock() + system_function_call_1.id = 'auth_id_1' + system_function_call_1.args = { + 'function_call_id': 'tool_id_1', + 'auth_config': mock_auth_config, + } + + system_function_call_2 = Mock() + system_function_call_2.id = 'auth_id_2' + system_function_call_2.args = { + 'function_call_id': 'tool_id_2', + 'auth_config': mock_auth_config, + } + + system_event = Mock(spec=Event) + system_event.content = Mock() # Non-None content + system_event.get_function_calls.return_value = [ + system_function_call_1, + system_function_call_2, + ] + + # Create original function call event + original_function_call_1 = Mock() + original_function_call_1.id = 'tool_id_1' + + original_function_call_2 = Mock() + original_function_call_2.id = 'tool_id_2' + + original_event = Mock(spec=Event) + original_event.content = Mock() # Non-None content + original_event.get_function_calls.return_value = [ + original_function_call_1, + original_function_call_2, + ] + + # Setup events in order: original -> system -> user_with_responses + mock_invocation_context.session.events = [ + original_event, + system_event, + user_event_with_multiple_responses, + ] + + # Setup mocks + mock_auth_config_validate.return_value = mock_auth_config + mock_auth_handler = Mock(spec=AuthHandler) + mock_auth_handler.parse_and_store_auth_response = AsyncMock() + mock_auth_handler_class.return_value = mock_auth_handler + + mock_function_response_event = Mock(spec=Event) + mock_handle_function_calls.return_value = mock_function_response_event + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + # Verify auth responses were processed + assert mock_auth_handler.parse_and_store_auth_response.call_count == 2 + + # Verify function calls were resumed + mock_handle_function_calls.assert_called_once() + call_args = mock_handle_function_calls.call_args + assert call_args[0][1] == original_event # The original event + assert call_args[0][3] == {'tool_id_1', 'tool_id_2'} # Tools to resume + + # Verify the function response event was yielded + assert result == [mock_function_response_event] + + @pytest.mark.asyncio + @patch('google.adk.auth.auth_preprocessor.AuthHandler') + @patch('google.adk.auth.auth_tool.AuthConfig.model_validate') + async def test_no_matching_system_function_calls_returns_early( + self, + mock_auth_config_validate, + mock_auth_handler_class, + processor, + mock_invocation_context, + mock_llm_request, + mock_user_event_with_auth_response, + mock_auth_config, + ): + """Test that missing matching system function calls returns early.""" + # Setup mocks + mock_auth_config_validate.return_value = mock_auth_config + mock_auth_handler = Mock(spec=AuthHandler) + mock_auth_handler.parse_and_store_auth_response = AsyncMock() + mock_auth_handler_class.return_value = mock_auth_handler + + # Create a non-matching system event + non_matching_function_call = Mock() + non_matching_function_call.id = ( # Different from 'auth_response_id' + 'different_id' + ) + + system_event = Mock(spec=Event) + system_event.content = Mock() # Non-None content + system_event.get_function_calls.return_value = [non_matching_function_call] + + mock_invocation_context.session.events = [ + system_event, + mock_user_event_with_auth_response, + ] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + # Should process auth response but not resume any tools + mock_auth_handler.parse_and_store_auth_response.assert_called_once() + assert result == [] + + @pytest.mark.asyncio + @patch('google.adk.auth.auth_preprocessor.AuthHandler') + @patch('google.adk.auth.auth_tool.AuthConfig.model_validate') + @patch('google.adk.auth.auth_tool.AuthToolArguments.model_validate') + async def test_handles_missing_original_function_calls( + self, + mock_auth_tool_args_validate, + mock_auth_config_validate, + mock_auth_handler_class, + processor, + mock_invocation_context, + mock_llm_request, + mock_user_event_with_auth_response, + mock_auth_config, + ): + """Test handling when original function calls are not found.""" + # Setup mocks + mock_auth_config_validate.return_value = mock_auth_config + mock_auth_handler = Mock(spec=AuthHandler) + mock_auth_handler.parse_and_store_auth_response = AsyncMock() + mock_auth_handler_class.return_value = mock_auth_handler + + # Create matching system function call + auth_tool_args = Mock(spec=AuthToolArguments) + auth_tool_args.function_call_id = 'tool_id_1' + mock_auth_tool_args_validate.return_value = auth_tool_args + + system_function_call = Mock() + system_function_call.id = 'auth_response_id' # Matches the response ID + system_function_call.args = { + 'function_call_id': 'tool_id_1', + 'auth_config': mock_auth_config, + } + + system_event = Mock(spec=Event) + system_event.content = Mock() # Non-None content + system_event.get_function_calls.return_value = [system_function_call] + + # Create event with no function calls (original function calls missing) + empty_event = Mock(spec=Event) + empty_event.content = Mock() # Non-None content + empty_event.get_function_calls.return_value = [] + + mock_invocation_context.session.events = [ + empty_event, + system_event, + mock_user_event_with_auth_response, + ] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + # Should process auth response but not find original function calls + mock_auth_handler.parse_and_store_auth_response.assert_called_once() + assert result == [] + + @pytest.mark.asyncio + async def test_isinstance_check_for_llm_agent( + self, processor, mock_llm_request, mock_session + ): + """Test that isinstance check works correctly for LlmAgent.""" + # This test ensures the isinstance check work as expected + + # Create a mock that fails isinstance check + mock_context = Mock(spec=InvocationContext) + mock_context.agent = Mock() # This will fail isinstance(agent, LlmAgent) + mock_context.session = mock_session + + result = [] + async for event in processor.run_async(mock_context, mock_llm_request): + result.append(event) + + assert result == []