fix: execute on_event_callback before append_event to persist plugin modifications

Previously, on_event_callback ran after append_event, so plugin
modifications (e.g. custom_metadata) were only visible in the
yielded stream but never persisted to the session store. This
moves the callback before persistence, matching the documented
contract in base_plugin.py.

Fixes #3990

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 899845848

Author: GWeale

src/google/adk/plugins/base_plugin.py

@@ -155,10 +155,10 @@ async def before_run_callback(
   async def on_event_callback(
       self, *, invocation_context: InvocationContext, event: Event
   ) -> Optional[Event]:
-    """Callback executed after an event is yielded from runner.
+    """Callback executed when the runner produces an event.
 
-    This is the ideal place to make modification to the event before the event
-    is handled by the underlying agent app.
+    This is the ideal place to modify the event before it is persisted to the
+    session service and yielded to the caller.
 
     Args:
       invocation_context: The context for the entire invocation.

src/google/adk/runners.py

@@ -791,6 +791,34 @@ def _should_append_event(self, event: Event, is_live_call: bool) -> bool:
       return False
     return True
 
+  def _get_output_event(
+      self,
+      *,
+      original_event: Event,
+      modified_event: Event | None,
+      run_config: RunConfig | None,
+  ) -> Event:
+    """Returns the event that should be persisted and yielded.
+
+    Plugins may return a replacement event that only overrides a subset of
+    fields. Merge those changes onto the original event so the streamed event
+    and the persisted event stay aligned without losing the original event
+    identity.
+    """
+    if modified_event is None:
+      return original_event
+
+    _apply_run_config_custom_metadata(modified_event, run_config)
+    update = {}
+    for field_name in modified_event.model_fields_set:
+      if field_name in {'id', 'invocation_id', 'timestamp'}:
+        continue
+      update[field_name] = modified_event.__dict__[field_name]
+    output_event = original_event.model_copy(update=update)
+    if not output_event.author:
+      output_event.author = original_event.author
+    return output_event
+
   async def _exec_with_plugin(
       self,
       invocation_context: InvocationContext,
@@ -854,13 +882,24 @@ async def _exec_with_plugin(
           _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,
+          )
+
           if is_live_call:
             if event.partial and _is_transcription(event):
               is_transcribing = True
             if is_transcribing and _is_tool_call_or_response(event):
               # only buffer function call and function response event which is
               # non-partial
-              buffered_events.append(event)
+              buffered_events.append(output_event)
               continue
             # Note for live/bidi: for audio response, it's considered as
             # non-partial event(event.partial=None)
@@ -881,7 +920,7 @@ async def _exec_with_plugin(
                 )
                 if self._should_append_event(event, is_live_call):
                   await self.session_service.append_event(
-                      session=session, event=event
+                      session=session, event=output_event
                   )
 
                 for buffered_event in buffered_events:
@@ -897,25 +936,15 @@ async def _exec_with_plugin(
                 if self._should_append_event(event, is_live_call):
                   logger.debug('Appending non-buffered event: %s', event)
                   await self.session_service.append_event(
-                      session=session, event=event
+                      session=session, event=output_event
                   )
           else:
             if event.partial is not True:
               await self.session_service.append_event(
-                  session=session, event=event
+                  session=session, event=output_event
               )
 
-          # Step 3: Run the on_event callbacks to optionally modify the event.
-          modified_event = await plugin_manager.run_on_event_callback(
-              invocation_context=invocation_context, event=event
-          )
-          if modified_event:
-            _apply_run_config_custom_metadata(
-                modified_event, invocation_context.run_config
-            )
-            yield modified_event
-          else:
-            yield event
+          yield output_event
 
     # Step 4: Run the after_run callbacks to perform global cleanup tasks or
     # finalizing logs and metrics data.

tests/unittests/test_runners.py

@@ -139,6 +139,7 @@ class MockPlugin(BasePlugin):
       "Modified user message ON_USER_CALLBACK_MSG from MockPlugin"
   )
   ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin"
+  ON_EVENT_CALLBACK_METADATA = {"plugin_key": "plugin_value"}
 
   def __init__(self):
     super().__init__(name="mock_plugin")
@@ -184,6 +185,7 @@ async def on_event_callback(
             ],
             role=event.content.role,
         ),
+        custom_metadata=self.ON_EVENT_CALLBACK_METADATA,
     )
 
 
@@ -359,6 +361,60 @@ async def test_run_live_auto_create_session():
   assert session is not None
 
 
+@pytest.mark.asyncio
+async def test_run_live_persists_event_callback_modifications():
+  """run_live should persist the same event it streams after callback changes."""
+  session_service = InMemorySessionService()
+  artifact_service = InMemoryArtifactService()
+  plugin = MockPlugin()
+  plugin.enable_event_callback = True
+  runner = Runner(
+      app_name="live_app",
+      agent=MockLiveAgent("live_agent"),
+      session_service=session_service,
+      artifact_service=artifact_service,
+      plugins=[plugin],
+  )
+  await session_service.create_session(
+      app_name="live_app", user_id="user", session_id="live_session"
+  )
+
+  from google.adk.agents.live_request_queue import LiveRequestQueue
+
+  live_queue = LiveRequestQueue()
+  agen = runner.run_live(
+      user_id="user",
+      session_id="live_session",
+      live_request_queue=live_queue,
+  )
+
+  streamed_event = await agen.__anext__()
+  await agen.aclose()
+
+  session = await session_service.get_session(
+      app_name="live_app", user_id="user", session_id="live_session"
+  )
+  persisted_event = session.events[0]
+
+  assert streamed_event.author == "live_agent"
+  assert streamed_event.invocation_id
+  assert streamed_event.content.parts[0].text == (
+      MockPlugin.ON_EVENT_CALLBACK_MSG
+  )
+  assert streamed_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+
+  assert persisted_event.id == streamed_event.id
+  assert persisted_event.timestamp == streamed_event.timestamp
+  assert persisted_event.author == streamed_event.author
+  assert persisted_event.invocation_id == streamed_event.invocation_id
+  assert persisted_event.content.parts[0].text == (
+      MockPlugin.ON_EVENT_CALLBACK_MSG
+  )
+  assert (
+      persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+  )
+
+
 @pytest.mark.asyncio
 async def test_runner_allows_nested_agent_directories(tmp_path, monkeypatch):
   project_root = tmp_path / "workspace"
@@ -747,6 +803,39 @@ async def test_runner_modifies_event_after_execution(self):
 
     assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG
 
+  @pytest.mark.asyncio
+  async def test_runner_persists_event_callback_modifications(self):
+    """Event callback output should be persisted, not only streamed."""
+    self.plugin.enable_event_callback = True
+
+    events = await self.run_test()
+    streamed_event = events[0]
+
+    session = await self.session_service.get_session(
+        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+    )
+    persisted_event = session.events[1]
+
+    assert streamed_event.author == "test_agent"
+    assert streamed_event.invocation_id
+    assert streamed_event.content.parts[0].text == (
+        MockPlugin.ON_EVENT_CALLBACK_MSG
+    )
+    assert (
+        streamed_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+    )
+
+    assert persisted_event.id == streamed_event.id
+    assert persisted_event.timestamp == streamed_event.timestamp
+    assert persisted_event.author == streamed_event.author
+    assert persisted_event.invocation_id == streamed_event.invocation_id
+    assert persisted_event.content.parts[0].text == (
+        MockPlugin.ON_EVENT_CALLBACK_MSG
+    )
+    assert (
+        persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
+    )
+
   @pytest.mark.asyncio
   async def test_runner_close_calls_plugin_close(self):
     """Test that runner.close() calls plugin manager close."""