From 847b85114f15741f42e0ea23b293fc25fe2baaa8 Mon Sep 17 00:00:00 2001 From: Daniel Partatus Date: Thu, 20 Aug 2026 12:59:59 +0200 Subject: [PATCH] fix(runners): honor before_run_callback early-exit on the node execution path The node execution path (_run_node_async) invoked plugin_manager.run_before_run_callback but discarded its return value, so a plugin returning types.Content (the documented signal to halt the run) was ignored and execution continued. The legacy path (_exec_with_plugin) already honors this contract. This affects every root that dispatches through the node path: a Workflow root and a root LlmAgent (chat/task mode), which means plugin-based guardrails (e.g. safety filters that block a turn in before_run_callback) are silently bypassed for those shapes. Fix mirrors the _exec_with_plugin early-exit contract: the returned Content becomes the final response event (with RunConfig custom_metadata applied), is appended to the session, and the run ends. after_run callbacks and post-invocation compaction are run explicitly (the finally that normally runs them belongs to the main loop, which a halted run never enters), matching the success path and the legacy early-exit behavior. Adapted from the stale PR #6032 by @garyzava (rebased onto current main and extended with a root-LlmAgent regression test). Fixes #6828 --- src/google/adk/runners.py | 35 +++++++- .../workflow/test_workflow_failures.py | 82 +++++++++++++++++++ 2 files changed, 116 insertions(+), 1 deletion(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 626b6ee5d1..6f3fb01d46 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -673,11 +673,44 @@ async def _run() -> AsyncGenerator[Event, None]: yield user_event # Run before_run callbacks - await ic.plugin_manager.run_before_run_callback(invocation_context=ic) + early_exit_result = await ic.plugin_manager.run_before_run_callback( + invocation_context=ic + ) except Exception as e: await _notify_run_error(ic.plugin_manager, ic, e) raise + # A Content returned by before_run halts the run and becomes the + # final response, mirroring the early-exit contract of + # _exec_with_plugin. The success-only cleanup below (after_run and + # compaction) is run explicitly: the finally that normally runs it + # belongs to the main loop, which a halted run never enters. + if isinstance(early_exit_result, types.Content): + early_exit_event = Event( + invocation_id=ic.invocation_id, + author='model', + content=early_exit_result, + ) + _apply_run_config_custom_metadata(early_exit_event, ic.run_config) + if self._should_append_event(early_exit_event, is_live_call=False): + await self.session_service.append_event( + session=ic.session, + event=early_exit_event, + ) + yield early_exit_event + try: + await ic.plugin_manager.run_after_run_callback( + invocation_context=ic + ) + await self._run_post_invocation_compaction( + session=session, + skip_token_compaction=ic.token_compaction_checked, + ) + except Exception as e: + await _notify_run_error(ic.plugin_manager, ic, e) + raise + return + # 3. Start root node in background from .agents.context import Context from .workflow._dynamic_node_scheduler import DynamicNodeScheduler diff --git a/tests/unittests/workflow/test_workflow_failures.py b/tests/unittests/workflow/test_workflow_failures.py index 8721acd883..c7f2835ee7 100644 --- a/tests/unittests/workflow/test_workflow_failures.py +++ b/tests/unittests/workflow/test_workflow_failures.py @@ -21,8 +21,11 @@ from google.adk import platform as adk_platform from google.adk.agents.context import Context +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm_agent import Agent from google.adk.apps.app import App from google.adk.events.event import Event +from google.adk.plugins.base_plugin import BasePlugin # Added for the moved test from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService @@ -1162,3 +1165,82 @@ async def failing_node_2(ctx: Context): with pytest.raises(ValueError, match='Fail 1'): await runner.run_async(testing_utils.get_user_content('start')) + + +class _HaltingPlugin(BasePlugin): + """Plugin whose before_run_callback halts the run with a Content.""" + + def __init__(self): + super().__init__(name='halting_plugin') + self.after_run_called = False + + async def before_run_callback( + self, *, invocation_context: InvocationContext + ) -> types.Content: + return types.Content( + role='model', parts=[types.Part(text='halted by plugin')] + ) + + async def after_run_callback( + self, *, invocation_context: InvocationContext + ) -> None: + self.after_run_called = True + + +def _texts(events: list[Event]) -> list[str]: + return [ + part.text + for event in events + if event.content and event.content.parts + for part in event.content.parts + if part.text + ] + + +@pytest.mark.asyncio +async def test_workflow_halts_when_before_run_callback_returns_content( + request: pytest.FixtureRequest, +): + """Regression for #6013: a Content returned by before_run_callback must + halt the run with that content and skip node execution.""" + node_a = TestingNode(name='NodeA', output='should not run') + graph = Graph(edges=[Edge(from_node=START, to_node=node_a)]) + workflow = Workflow(name='halt_workflow', graph=graph) + + plugin = _HaltingPlugin() + app = App( + name=request.function.__name__, + root_agent=workflow, + plugins=[plugin], + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('start')) + + assert node_a.received_inputs == [] + assert 'halted by plugin' in _texts(events) + # A halted run is a completed run: after_run must still fire, matching + # the _exec_with_plugin early-exit behavior. + assert plugin.after_run_called + + +@pytest.mark.asyncio +async def test_llm_agent_root_halts_when_before_run_callback_returns_content( + request: pytest.FixtureRequest, +): + """Same regression for the other node-path shape: a root LlmAgent. The + model must never be called on a halted run.""" + mock_model = testing_utils.MockModel.create(responses=['should not run']) + agent = Agent(name='root_agent', model=mock_model) + + plugin = _HaltingPlugin() + app = App( + name=request.function.__name__, + root_agent=agent, + plugins=[plugin], + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('hello')) + + assert not mock_model.requests + assert 'halted by plugin' in _texts(events) + assert plugin.after_run_called