diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 9ac294332d..3d2217b372 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -752,14 +752,9 @@ async def _consume_event_queue( event = event.model_copy() event.output = None - _apply_run_config_custom_metadata(event, ic.run_config) - modified_event = await ic.plugin_manager.run_on_event_callback( - invocation_context=ic, event=event - ) - output_event = self._get_output_event( - original_event=event, - modified_event=modified_event, - run_config=ic.run_config, + output_event = await self._process_event_with_plugin_callbacks( + invocation_context=ic, + event=event, ) if not event.partial: @@ -1346,6 +1341,26 @@ def _get_output_event( output_event.author = original_event.author return output_event + async def _process_event_with_plugin_callbacks( + self, + *, + invocation_context: InvocationContext, + event: Event, + ) -> Event: + """Applies runner metadata and plugin callbacks to an output event.""" + _apply_run_config_custom_metadata(event, invocation_context.run_config) + modified_event = ( + await invocation_context.plugin_manager.run_on_event_callback( + invocation_context=invocation_context, + event=event, + ) + ) + return self._get_output_event( + original_event=event, + modified_event=modified_event, + run_config=invocation_context.run_config, + ) + async def _exec_with_plugin( self, invocation_context: InvocationContext, @@ -1378,31 +1393,25 @@ async def _exec_with_plugin( author='model', content=early_exit_result, ) - _apply_run_config_custom_metadata( - early_exit_event, invocation_context.run_config + output_event = await self._process_event_with_plugin_callbacks( + invocation_context=invocation_context, + event=early_exit_event, ) if self._should_append_event(early_exit_event, is_live_call): await self.session_service.append_event( session=invocation_context.session, - event=early_exit_event, + event=output_event, ) - yield early_exit_event + yield output_event else: # Step 2: Otherwise continue with normal execution async with aclosing(execute_fn(invocation_context)) as agen: async for event in agen: - _apply_run_config_custom_metadata( - event, invocation_context.run_config - ) # Step 3: Run the on_event callbacks before persisting so callback # changes are stored in the session and match the streamed event. - modified_event = await plugin_manager.run_on_event_callback( - invocation_context=invocation_context, event=event - ) - output_event = self._get_output_event( - original_event=event, - modified_event=modified_event, - run_config=invocation_context.run_config, + output_event = await self._process_event_with_plugin_callbacks( + invocation_context=invocation_context, + event=event, ) if is_live_call: diff --git a/src/google/adk/workflow/_node_runner_utils.py b/src/google/adk/workflow/_node_runner_utils.py index 4f6ca4887e..7a6c63200a 100644 --- a/src/google/adk/workflow/_node_runner_utils.py +++ b/src/google/adk/workflow/_node_runner_utils.py @@ -61,7 +61,6 @@ async def run_node_async( session: Optional[Session] = None, ) -> AsyncGenerator[Event, None]: """Runs a BaseNode or Workflow in async mode.""" - from ..runners import _apply_run_config_custom_metadata from ..runners import _find_active_task_scope caller_ctx = context.get_current() @@ -187,15 +186,18 @@ async def _run() -> AsyncGenerator[Event, None]: author="model", content=early_exit_result, ) - _apply_run_config_custom_metadata(early_exit_event, ic.run_config) + output_event = await runner._process_event_with_plugin_callbacks( # pylint: disable=protected-access + invocation_context=ic, + event=early_exit_event, + ) if runner._should_append_event( # pylint: disable=protected-access early_exit_event, is_live_call=False ): await runner.session_service.append_event( session=ic.session, - event=early_exit_event, + event=output_event, ) - yield early_exit_event + yield output_event else: # 3. Start root node in background root_ctx = Context(ic) diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 4fa4df522b..a02056a9cc 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -162,6 +162,7 @@ def __init__(self): super().__init__(name="mock_plugin") self.enable_user_message_callback = False self.enable_event_callback = False + self.before_run_response: Optional[types.Content] = None self.user_content_seen_in_before_run_callback = None async def on_user_message_callback( @@ -181,10 +182,11 @@ async def before_run_callback( self, *, invocation_context: InvocationContext, - ) -> None: + ) -> Optional[types.Content]: self.user_content_seen_in_before_run_callback = ( invocation_context.user_content ) + return self.before_run_response async def on_event_callback( self, *, invocation_context: InvocationContext, event: Event @@ -1218,6 +1220,68 @@ async def test_runner_persists_event_callback_modifications(self): persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("agent_cls", "is_live"), + [ + (MockAgent, False), + (MockLlmAgent, False), + (MockLiveAgent, True), + ], + ids=("legacy", "node", "live"), + ) + async def test_runner_processes_before_run_early_exit_with_event_callback( + self, agent_cls, is_live + ): + """Before-run early exits still pass through on-event hooks.""" + from google.adk.live import LiveRequestQueue + + plugin = MockPlugin() + plugin.before_run_response = types.Content( + role="model", parts=[types.Part(text="blocked by before_run")] + ) + plugin.enable_event_callback = True + session_service = InMemorySessionService() + runner = Runner( + app=App( + name=TEST_APP_ID, + root_agent=agent_cls("test_agent"), + plugins=[plugin], + ), + session_service=session_service, + ) + session = await session_service.create_session( + app_name=TEST_APP_ID, + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + ) + + if is_live: + event_stream = runner.run_live( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + live_request_queue=LiveRequestQueue(), + ) + else: + event_stream = runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=types.Content( + role="user", parts=[types.Part(text="hello")] + ), + ) + events = [event async for event in event_stream] + persisted_session = await session_service.get_session( + app_name=TEST_APP_ID, + user_id=TEST_USER_ID, + session_id=session.id, + ) + + assert len(events) == 1 + assert events[0].content.parts[0].text == MockPlugin.ON_EVENT_CALLBACK_MSG + assert events[0].custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA + assert persisted_session.events[-1] == events[0] + @pytest.mark.asyncio async def test_runner_close_calls_plugin_close(self): """Test that runner.close() calls plugin manager close."""