Skip to content

Commit 847b851

Browse files
author
Daniel Partatus
committed
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 google#6032 by @garyzava (rebased onto current main and extended with a root-LlmAgent regression test). Fixes google#6828
1 parent 4599a52 commit 847b851

2 files changed

Lines changed: 116 additions & 1 deletion

File tree

‎src/google/adk/runners.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -673,11 +673,44 @@ async def _run() -> AsyncGenerator[Event, None]:
673673
yield user_event
674674

675675
# Run before_run callbacks
676-
await ic.plugin_manager.run_before_run_callback(invocation_context=ic)
676+
early_exit_result = await ic.plugin_manager.run_before_run_callback(
677+
invocation_context=ic
678+
)
677679
except Exception as e:
678680
await _notify_run_error(ic.plugin_manager, ic, e)
679681
raise
680682

683+
# A Content returned by before_run halts the run and becomes the
684+
# final response, mirroring the early-exit contract of
685+
# _exec_with_plugin. The success-only cleanup below (after_run and
686+
# compaction) is run explicitly: the finally that normally runs it
687+
# belongs to the main loop, which a halted run never enters.
688+
if isinstance(early_exit_result, types.Content):
689+
early_exit_event = Event(
690+
invocation_id=ic.invocation_id,
691+
author='model',
692+
content=early_exit_result,
693+
)
694+
_apply_run_config_custom_metadata(early_exit_event, ic.run_config)
695+
if self._should_append_event(early_exit_event, is_live_call=False):
696+
await self.session_service.append_event(
697+
session=ic.session,
698+
event=early_exit_event,
699+
)
700+
yield early_exit_event
701+
try:
702+
await ic.plugin_manager.run_after_run_callback(
703+
invocation_context=ic
704+
)
705+
await self._run_post_invocation_compaction(
706+
session=session,
707+
skip_token_compaction=ic.token_compaction_checked,
708+
)
709+
except Exception as e:
710+
await _notify_run_error(ic.plugin_manager, ic, e)
711+
raise
712+
return
713+
681714
# 3. Start root node in background
682715
from .agents.context import Context
683716
from .workflow._dynamic_node_scheduler import DynamicNodeScheduler

‎tests/unittests/workflow/test_workflow_failures.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,11 @@
2121

2222
from google.adk import platform as adk_platform
2323
from google.adk.agents.context import Context
24+
from google.adk.agents.invocation_context import InvocationContext
25+
from google.adk.agents.llm_agent import Agent
2426
from google.adk.apps.app import App
2527
from google.adk.events.event import Event
28+
from google.adk.plugins.base_plugin import BasePlugin
2629
# Added for the moved test
2730
from google.adk.runners import Runner
2831
from google.adk.sessions.in_memory_session_service import InMemorySessionService
@@ -1162,3 +1165,82 @@ async def failing_node_2(ctx: Context):
11621165

11631166
with pytest.raises(ValueError, match='Fail 1'):
11641167
await runner.run_async(testing_utils.get_user_content('start'))
1168+
1169+
1170+
class _HaltingPlugin(BasePlugin):
1171+
"""Plugin whose before_run_callback halts the run with a Content."""
1172+
1173+
def __init__(self):
1174+
super().__init__(name='halting_plugin')
1175+
self.after_run_called = False
1176+
1177+
async def before_run_callback(
1178+
self, *, invocation_context: InvocationContext
1179+
) -> types.Content:
1180+
return types.Content(
1181+
role='model', parts=[types.Part(text='halted by plugin')]
1182+
)
1183+
1184+
async def after_run_callback(
1185+
self, *, invocation_context: InvocationContext
1186+
) -> None:
1187+
self.after_run_called = True
1188+
1189+
1190+
def _texts(events: list[Event]) -> list[str]:
1191+
return [
1192+
part.text
1193+
for event in events
1194+
if event.content and event.content.parts
1195+
for part in event.content.parts
1196+
if part.text
1197+
]
1198+
1199+
1200+
@pytest.mark.asyncio
1201+
async def test_workflow_halts_when_before_run_callback_returns_content(
1202+
request: pytest.FixtureRequest,
1203+
):
1204+
"""Regression for #6013: a Content returned by before_run_callback must
1205+
halt the run with that content and skip node execution."""
1206+
node_a = TestingNode(name='NodeA', output='should not run')
1207+
graph = Graph(edges=[Edge(from_node=START, to_node=node_a)])
1208+
workflow = Workflow(name='halt_workflow', graph=graph)
1209+
1210+
plugin = _HaltingPlugin()
1211+
app = App(
1212+
name=request.function.__name__,
1213+
root_agent=workflow,
1214+
plugins=[plugin],
1215+
)
1216+
runner = testing_utils.InMemoryRunner(app=app)
1217+
events = await runner.run_async(testing_utils.get_user_content('start'))
1218+
1219+
assert node_a.received_inputs == []
1220+
assert 'halted by plugin' in _texts(events)
1221+
# A halted run is a completed run: after_run must still fire, matching
1222+
# the _exec_with_plugin early-exit behavior.
1223+
assert plugin.after_run_called
1224+
1225+
1226+
@pytest.mark.asyncio
1227+
async def test_llm_agent_root_halts_when_before_run_callback_returns_content(
1228+
request: pytest.FixtureRequest,
1229+
):
1230+
"""Same regression for the other node-path shape: a root LlmAgent. The
1231+
model must never be called on a halted run."""
1232+
mock_model = testing_utils.MockModel.create(responses=['should not run'])
1233+
agent = Agent(name='root_agent', model=mock_model)
1234+
1235+
plugin = _HaltingPlugin()
1236+
app = App(
1237+
name=request.function.__name__,
1238+
root_agent=agent,
1239+
plugins=[plugin],
1240+
)
1241+
runner = testing_utils.InMemoryRunner(app=app)
1242+
events = await runner.run_async(testing_utils.get_user_content('hello'))
1243+
1244+
assert not mock_model.requests
1245+
assert 'halted by plugin' in _texts(events)
1246+
assert plugin.after_run_called

0 commit comments

Comments
 (0)