Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion src/google/adk/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
82 changes: 82 additions & 0 deletions tests/unittests/workflow/test_workflow_failures.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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