fix: Propagate RunConfig custom metadata to all events

Adds a method to merge custom metadata from the RunConfig into each Event. This metadata is applied to events generated by the agent, early exit events, and the initial user message event.

Close #3953

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

Author: GWeale

src/google/adk/runners.py

@@ -84,6 +84,19 @@ def _has_non_empty_transcription_text(transcription) -> bool:
   )
 
 
+def _apply_run_config_custom_metadata(
+    event: Event, run_config: RunConfig | None
+) -> None:
+  """Merges run-level custom metadata into the event, if present."""
+  if not run_config or not run_config.custom_metadata:
+    return
+
+  event.custom_metadata = {
+      **run_config.custom_metadata,
+      **(event.custom_metadata or {}),
+  }
+
+
 class Runner:
   """The Runner class is used to run agents.
 
@@ -695,6 +708,9 @@ async def _exec_with_plugin(
           author='model',
           content=early_exit_result,
       )
+      _apply_run_config_custom_metadata(
+          early_exit_event, invocation_context.run_config
+      )
       if self._should_append_event(early_exit_event, is_live_call):
         await self.session_service.append_event(
             session=session,
@@ -721,6 +737,9 @@ async def _exec_with_plugin(
 
       async with Aclosing(execute_fn(invocation_context)) as agen:
         async for event in agen:
+          _apply_run_config_custom_metadata(
+              event, invocation_context.run_config
+          )
           if is_live_call:
             if event.partial and _is_transcription(event):
               is_transcribing = True
@@ -775,7 +794,13 @@ async def _exec_with_plugin(
           modified_event = await plugin_manager.run_on_event_callback(
               invocation_context=invocation_context, event=event
           )
-          yield (modified_event if modified_event else 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.
@@ -846,6 +871,7 @@ async def _append_new_message_to_session(
           author='user',
           content=new_message,
       )
+    _apply_run_config_custom_metadata(event, invocation_context.run_config)
     # If new_message is a function response, find the matching function call
     # and use its branch as the new event's branch.
     if function_call := invocation_context._find_matching_function_call(event):

tests/unittests/test_runners.py

@@ -16,13 +16,15 @@
 from pathlib import Path
 import sys
 import textwrap
+from typing import AsyncGenerator
 from typing import Optional
 from unittest.mock import AsyncMock
 
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
 from google.adk.apps.app import App
 from google.adk.apps.app import ResumabilityConfig
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
@@ -54,7 +56,9 @@ def __init__(
     if parent_agent:
       self.parent_agent = parent_agent
 
-  async def _run_async_impl(self, invocation_context):
+  async def _run_async_impl(
+      self, invocation_context: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
     yield Event(
         invocation_id=invocation_context.invocation_id,
         author=self.name,
@@ -78,7 +82,9 @@ def __init__(
     self.disallow_transfer_to_parent = disallow_transfer_to_parent
     self.parent_agent = parent_agent
 
-  async def _run_async_impl(self, invocation_context):
+  async def _run_async_impl(
+      self, invocation_context: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
     yield Event(
         invocation_id=invocation_context.invocation_id,
         author=self.name,
@@ -88,6 +94,25 @@ async def _run_async_impl(self, invocation_context):
     )
 
 
+class MockAgentWithMetadata(BaseAgent):
+  """Mock agent that returns event-level custom metadata."""
+
+  def __init__(self, name: str):
+    super().__init__(name=name, sub_agents=[])
+
+  async def _run_async_impl(
+      self, invocation_context: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        invocation_id=invocation_context.invocation_id,
+        author=self.name,
+        content=types.Content(
+            role="model", parts=[types.Part(text="Test response")]
+        ),
+        custom_metadata={"event_key": "event_value"},
+    )
+
+
 class MockPlugin(BasePlugin):
   """Mock plugin for unit testing."""
 
@@ -495,6 +520,41 @@ def test_is_transferable_across_agent_tree_with_non_llm_agent(self):
     assert result is False
 
 
+@pytest.mark.asyncio
+async def test_run_config_custom_metadata_propagates_to_events():
+  session_service = InMemorySessionService()
+  runner = Runner(
+      app_name=TEST_APP_ID,
+      agent=MockAgentWithMetadata("metadata_agent"),
+      session_service=session_service,
+      artifact_service=InMemoryArtifactService(),
+  )
+  await session_service.create_session(
+      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+  )
+
+  run_config = RunConfig(custom_metadata={"request_id": "req-1"})
+  events = [
+      event
+      async for event in runner.run_async(
+          user_id=TEST_USER_ID,
+          session_id=TEST_SESSION_ID,
+          new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
+          run_config=run_config,
+      )
+  ]
+
+  assert events[0].custom_metadata is not None
+  assert events[0].custom_metadata["request_id"] == "req-1"
+  assert events[0].custom_metadata["event_key"] == "event_value"
+
+  session = await session_service.get_session(
+      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+  )
+  user_event = next(event for event in session.events if event.author == "user")
+  assert user_event.custom_metadata == {"request_id": "req-1"}
+
+
 class TestRunnerWithPlugins:
   """Tests for Runner with plugins."""