fix: persist events after plugin on_event_callback modifications

src/google/adk/runners.py

@@ -806,49 +806,44 @@ async def _exec_with_plugin(
                     'Appending transcription finished event: %s', event
                 )
                 if self._should_append_event(event, is_live_call):
-                  await self.session_service.append_event(
-                      session=session, event=event
-                  )
+                  event = await self.__append_event(invocation_context, session, event)
 
                 for buffered_event in buffered_events:
                   logger.debug('Appending buffered event: %s', buffered_event)
-                  await self.session_service.append_event(
-                      session=session, event=buffered_event
-                  )
+                  buffered_event = await self.__append_event(invocation_context, session, buffered_event)
                   yield buffered_event  # yield buffered events to caller
                 buffered_events = []
               else:
                 # non-transcription event or empty transcription event, for
                 # example, event that stores blob reference, should be appended.
                 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
-                  )
+                  event = await self.__append_event(invocation_context, session, event)
+          # Run the on_event callbacks to optionally modify the event.
           else:
-            if event.partial is not True:
-              await self.session_service.append_event(
-                  session=session, event=event
-              )
+            event = await self.__append_event(invocation_context, session, event)
+          yield 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
 
-    # Step 4: Run the after_run callbacks to perform global cleanup tasks or
-    # finalizing logs and metrics data.
-    # This does NOT emit any event.
-    await plugin_manager.run_after_run_callback(
-        invocation_context=invocation_context
+  async def __append_event(self, invocation_context: InvocationContext, session: Session, event: Event) -> None:
+    """Appends an event to the session with plugin callbacks.
+
+    Args:
+      invocation_context: The invocation context.
+      session: The session to append the event to.
+      event: The event to process and append to the session.
+    """
+    plugin_manager = invocation_context.plugin_manager
+    modified_event = await plugin_manager.run_on_event_callback(
+        invocation_context=invocation_context, event=event
     )
+    if modified_event:
+      event = modified_event
+      _apply_run_config_custom_metadata(
+          event, invocation_context.run_config
+      )
+    await self.session_service.append_event(session=session, event=event)
+    return event
 
   async def _append_new_message_to_session(
       self,