diff --git a/src/google/adk/plugins/_reflect_retry_model_plugin.py b/src/google/adk/plugins/_reflect_retry_model_plugin.py index 0436e6a43b8..ab23414a634 100644 --- a/src/google/adk/plugins/_reflect_retry_model_plugin.py +++ b/src/google/adk/plugins/_reflect_retry_model_plugin.py @@ -219,17 +219,16 @@ async def _handle_model_error( llm_response: LlmResponse, ) -> LlmResponse | None: """Handles detected model errors by initiating retry logic.""" - retry_response = await self._handle_model_retry( + # Without a retry this returns None, which leaves the error response + # unchanged. Returning the error response itself would stop the remaining + # plugins and the agent's after_model_callbacks. + return await self._handle_model_retry( callback_context=callback_context, llm_response=llm_response, error_type=llm_response.error_code, error_details=llm_response.error_message, finish_reason=llm_response.finish_reason, ) - if retry_response is not None: - return retry_response - - return llm_response async def _handle_model_retry( self, diff --git a/src/google/adk/plugins/multimodal_tool_results_plugin.py b/src/google/adk/plugins/multimodal_tool_results_plugin.py index 893bb0a0848..a2bbb1ad046 100644 --- a/src/google/adk/plugins/multimodal_tool_results_plugin.py +++ b/src/google/adk/plugins/multimodal_tool_results_plugin.py @@ -110,7 +110,9 @@ async def after_tool_callback( and result and isinstance(result[0], types.Part) ): - return result + # None leaves the result unchanged. Returning the result itself would + # stop the remaining plugins and the agent's after_tool_callbacks. + return None parts = [result] if isinstance(result, types.Part) else result[:] diff --git a/src/google/adk/plugins/save_files_as_artifacts_plugin.py b/src/google/adk/plugins/save_files_as_artifacts_plugin.py index ff751cfd369..ee876be5921 100644 --- a/src/google/adk/plugins/save_files_as_artifacts_plugin.py +++ b/src/google/adk/plugins/save_files_as_artifacts_plugin.py @@ -83,7 +83,8 @@ async def on_user_message_callback( "Artifact service is not set. SaveFilesAsArtifactsPlugin" " will not be enabled." ) - return user_message + # None leaves the message unchanged without stopping later plugins. + return None if not user_message.parts: return None diff --git a/tests/unittests/plugins/test_multimodal_tool_results_plugin.py b/tests/unittests/plugins/test_multimodal_tool_results_plugin.py index 16f92fc7995..1cb5b1e561f 100644 --- a/tests/unittests/plugins/test_multimodal_tool_results_plugin.py +++ b/tests/unittests/plugins/test_multimodal_tool_results_plugin.py @@ -103,7 +103,7 @@ async def test_tool_returning_non_list_of_parts_is_unchanged( result=original_result, ) - assert result == original_result + assert result is None assert PARTS_RETURNED_BY_TOOLS_ID not in tool_context.state callback_context = Mock(spec=CallbackContext) @@ -708,3 +708,40 @@ async def test_session_retention_retains_parts_on_subsequent_model_calls_in_same callback_context=callback_context, llm_request=llm_request_2 ) assert llm_request_2.contents[-1].parts == [file_part] + + +def test_agent_after_tool_callback_runs_for_non_part_results(): + """The plugin does not short-circuit the agent's after_tool_callback.""" + + def get_weather(city: str) -> dict[str, Any]: + return {"city": city, "temp": 20} + + seen = [] + + def after_tool_callback(tool, args, tool_context, tool_response): + seen.append(tool_response) + return {"redacted": True} + + mock_model = testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name="get_weather", args={"city": "Paris"} + ), + "done", + ] + ) + agent = Agent( + name="root_agent", + model=mock_model, + tools=[get_weather], + after_tool_callback=after_tool_callback, + ) + runner = testing_utils.InMemoryRunner( + agent, plugins=[MultimodalToolResultsPlugin()] + ) + + events = runner.run("weather?") + + assert seen == [{"city": "Paris", "temp": 20}] + responses = [fr.response for e in events for fr in e.get_function_responses()] + assert responses == [{"redacted": True}] diff --git a/tests/unittests/plugins/test_reflect_retry_model_plugin.py b/tests/unittests/plugins/test_reflect_retry_model_plugin.py index f0931680027..cf4b54d5ff2 100644 --- a/tests/unittests/plugins/test_reflect_retry_model_plugin.py +++ b/tests/unittests/plugins/test_reflect_retry_model_plugin.py @@ -28,6 +28,8 @@ from google.adk.tools.function_tool import FunctionTool from google.genai import types +from .. import testing_utils + class TestReflectAndRetryModelPlugin(IsolatedAsyncioTestCase): """Tests for model error handling in the ReflectAndRetryModelPlugin.""" @@ -266,10 +268,10 @@ async def test_after_model_callback_can_perform_multiple_retries(self): response3.content.parts[0].function_call.args["retry_count"], 3 ) - async def test_after_model_callback_returns_response_when_retry_limit_reached( + async def test_after_model_callback_leaves_response_unchanged_when_retry_limit_reached( self, ): - """Checks that after_model_callback returns the failed response when retry limit is reached and throw_exception_if_retry_exceeded is False.""" + """Checks that after_model_callback leaves the failed response unchanged when retry limit is reached and throw_exception_if_retry_exceeded is False.""" mock_agent = Mock(spec=LlmAgent) mock_agent.canonical_model = Mock() mock_agent.canonical_model.model = "TEST_MODEL_NAME" @@ -303,9 +305,7 @@ async def test_after_model_callback_returns_response_when_retry_limit_reached( callback_context=mock_callback_context, llm_response=llm_response, ) - self.assertEqual(response2.error_code, llm_response.error_code) - self.assertEqual(response2.error_message, llm_response.error_message) - self.assertEqual(response2.finish_reason, llm_response.finish_reason) + self.assertIsNone(response2) async def test_after_model_callback_throws_when_retry_limit_reached(self): """Checks that after_model_callback raises an Exception when retry limit is reached and throw_exception_if_retry_exceeded is True.""" @@ -673,3 +673,34 @@ async def test_global_tracking_scope_for_models(self): self.assertEqual( resp2.content.parts[0].function_call.args["retry_count"], 2 ) + + +def test_agent_after_model_callback_runs_when_retry_limit_is_exceeded(): + """The plugin does not short-circuit the agent's after_model_callback.""" + seen = [] + + def after_model_callback(callback_context, llm_response): + seen.append(llm_response.error_message) + + mock_model = testing_utils.MockModel.create( + responses=[ + LlmResponse( + error_code=types.FinishReason.MALFORMED_FUNCTION_CALL, + error_message="TEST_ERROR_MESSAGE", + finish_reason=types.FinishReason.MALFORMED_FUNCTION_CALL, + ) + ] + ) + agent = LlmAgent( + name="root_agent", + model=mock_model, + after_model_callback=after_model_callback, + ) + plugin = ReflectAndRetryModelPlugin( + max_retries=0, throw_exception_if_retry_exceeded=False + ) + runner = testing_utils.InMemoryRunner(agent, plugins=[plugin]) + + runner.run("hi") + + assert seen == ["TEST_ERROR_MESSAGE"] diff --git a/tests/unittests/plugins/test_save_files_as_artifacts.py b/tests/unittests/plugins/test_save_files_as_artifacts.py index d46b004eedd..673709b838b 100644 --- a/tests/unittests/plugins/test_save_files_as_artifacts.py +++ b/tests/unittests/plugins/test_save_files_as_artifacts.py @@ -20,6 +20,8 @@ from google.adk.agents.invocation_context import InvocationContext from google.adk.artifacts.base_artifact_service import ArtifactVersion +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.plugins.plugin_manager import PluginManager from google.adk.plugins.save_files_as_artifacts_plugin import SaveFilesAsArtifactsPlugin from google.genai import types import pytest @@ -220,8 +222,43 @@ async def test_no_artifact_service(self): invocation_context=self.mock_context, user_message=user_message ) - assert result == user_message - assert result.parts[0].inline_data == inline_data + assert result is None + assert user_message.parts[0].inline_data == inline_data + + @pytest.mark.asyncio + async def test_no_artifact_service_does_not_stop_later_plugins(self): + """Without an artifact service, later plugins still see the message.""" + self.mock_context.artifact_service = None + seen = [] + + class _RecordingPlugin(BasePlugin): + + async def on_user_message_callback( + self, *, invocation_context, user_message + ): + seen.append(user_message) + + plugin_manager = PluginManager( + plugins=[self.plugin, _RecordingPlugin(name="recording")] + ) + user_message = types.Content( + parts=[ + types.Part( + inline_data=types.Blob( + display_name="test.pdf", + data=b"test data", + mime_type="application/pdf", + ) + ) + ] + ) + + result = await plugin_manager.run_on_user_message_callback( + invocation_context=self.mock_context, user_message=user_message + ) + + assert result is None + assert seen == [user_message] @pytest.mark.asyncio async def test_no_parts_in_message(self):