fix: run on_event callbacks for before_run early exits

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:

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)

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."""