From 4f7b13a3f13f3312e9c64c9afd78e41883628a3a Mon Sep 17 00:00:00 2001 From: Atul Joshi <120785343+AtulJoshi1206@users.noreply.github.com> Date: Sat, 3 Oct 2026 16:55:22 +0530 Subject: [PATCH 1/2] fix(plugins): return None from MultimodalToolResultsPlugin for non-part results For a tool result that is not a Part or a list of Parts, after_tool_callback returned the result itself. The plugin manager stops at the first non-None plugin result, and the agent's after_tool_callbacks only run when the plugin result is None, so registering this plugin silently disabled every agent after_tool_callback, and every later plugin, for ordinary dict results. Returning None is the BasePlugin contract for leaving a result unchanged. --- .../plugins/multimodal_tool_results_plugin.py | 4 +- .../test_multimodal_tool_results_plugin.py | 39 ++++++++++++++++++- 2 files changed, 41 insertions(+), 2 deletions(-) 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/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}] From aaeeab97428269b666ae227c66dc136a3ca9230e Mon Sep 17 00:00:00 2001 From: Atul Joshi <120785343+AtulJoshi1206@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:28:00 +0530 Subject: [PATCH 2/2] fix(plugins): return None for unchanged results in the other built-in plugins ReflectAndRetryModelPlugin and SaveFilesAsArtifactsPlugin had the same pattern as MultimodalToolResultsPlugin, returning their input to mean "unchanged". The plugin manager stops at the first non-None result, so: - ReflectAndRetryModelPlugin, once the retry limit was exceeded with throw_exception_if_retry_exceeded=False, returned the error response itself, and the flow then skipped the remaining plugins and the agent's canonical after_model_callbacks for it. - SaveFilesAsArtifactsPlugin, with no artifact service configured, returned the user message itself, which stopped every later plugin's on_user_message_callback. Both now return None, the BasePlugin contract for leaving a value unchanged. --- .../plugins/_reflect_retry_model_plugin.py | 9 ++-- .../plugins/save_files_as_artifacts_plugin.py | 3 +- .../test_reflect_retry_model_plugin.py | 41 ++++++++++++++++--- .../plugins/test_save_files_as_artifacts.py | 41 ++++++++++++++++++- 4 files changed, 81 insertions(+), 13 deletions(-) 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/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_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):