Skip to content

Commit a2e2417

Browse files
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 compaction still execute via the enclosing finally, matching the success path. Adapted from the stale PR #6032 by @garyzava (rebased onto current main and extended with a root-LlmAgent regression test). Fixes #6828
1 parent 4599a52 commit a2e2417

2 files changed

Lines changed: 92 additions & 1 deletion

File tree

src/google/adk/runners.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -673,11 +673,32 @@ 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. Returning here still runs after_run and
686+
# compaction via the enclosing finally, same as a successful run.
687+
if isinstance(early_exit_result, types.Content):
688+
early_exit_event = Event(
689+
invocation_id=ic.invocation_id,
690+
author='model',
691+
content=early_exit_result,
692+
)
693+
_apply_run_config_custom_metadata(early_exit_event, ic.run_config)
694+
if self._should_append_event(early_exit_event, is_live_call=False):
695+
await self.session_service.append_event(
696+
session=ic.session,
697+
event=early_exit_event,
698+
)
699+
yield early_exit_event
700+
return
701+
681702
# 3. Start root node in background
682703
from .agents.context import Context
683704
from .workflow._dynamic_node_scheduler import DynamicNodeScheduler

tests/unittests/workflow/test_workflow_failures.py

Lines changed: 70 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,70 @@ 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+
1176+
async def before_run_callback(
1177+
self, *, invocation_context: InvocationContext
1178+
) -> types.Content:
1179+
return types.Content(
1180+
role='model', parts=[types.Part(text='halted by plugin')]
1181+
)
1182+
1183+
1184+
def _texts(events: list[Event]) -> list[str]:
1185+
return [
1186+
part.text
1187+
for event in events
1188+
if event.content and event.content.parts
1189+
for part in event.content.parts
1190+
if part.text
1191+
]
1192+
1193+
1194+
@pytest.mark.asyncio
1195+
async def test_workflow_halts_when_before_run_callback_returns_content(
1196+
request: pytest.FixtureRequest,
1197+
):
1198+
"""Regression for #6013: a Content returned by before_run_callback must
1199+
halt the run with that content and skip node execution."""
1200+
node_a = TestingNode(name='NodeA', output='should not run')
1201+
graph = Graph(edges=[Edge(from_node=START, to_node=node_a)])
1202+
workflow = Workflow(name='halt_workflow', graph=graph)
1203+
1204+
app = App(
1205+
name=request.function.__name__,
1206+
root_agent=workflow,
1207+
plugins=[_HaltingPlugin()],
1208+
)
1209+
runner = testing_utils.InMemoryRunner(app=app)
1210+
events = await runner.run_async(testing_utils.get_user_content('start'))
1211+
1212+
assert node_a.received_inputs == []
1213+
assert 'halted by plugin' in _texts(events)
1214+
1215+
1216+
@pytest.mark.asyncio
1217+
async def test_llm_agent_root_halts_when_before_run_callback_returns_content(
1218+
request: pytest.FixtureRequest,
1219+
):
1220+
"""Same regression for the other node-path shape: a root LlmAgent. The
1221+
model must never be called on a halted run."""
1222+
mock_model = testing_utils.MockModel.create(responses=['should not run'])
1223+
agent = Agent(name='root_agent', model=mock_model)
1224+
1225+
app = App(
1226+
name=request.function.__name__,
1227+
root_agent=agent,
1228+
plugins=[_HaltingPlugin()],
1229+
)
1230+
runner = testing_utils.InMemoryRunner(app=app)
1231+
events = await runner.run_async(testing_utils.get_user_content('hello'))
1232+
1233+
assert not mock_model.requests
1234+
assert 'halted by plugin' in _texts(events)

0 commit comments

Comments
 (0)