feat: Pass invocation_id to _setup_invocation_context

PiperOrigin-RevId: 934183388

Author: chengxin-wang

src/google/adk/runners.py

@@ -192,8 +192,8 @@ def __init__(
     Args:
         app: An `App` instance. Mutually exclusive with `agent` and `node`.
         app_name: The application name. Required when `agent` is provided.
-          Optional override for `app.name` when `app` is provided. Defaults
-          to `node.name` when only `node` is provided.
+          Optional override for `app.name` when `app` is provided. Defaults to
+          `node.name` when only `node` is provided.
         agent: The root agent to run. Mutually exclusive with `app` and `node`.
         node: The root node to run. Mutually exclusive with `app` and `agent`.
         plugins: Deprecated. A list of plugins for the runner. Please use the
@@ -203,8 +203,8 @@ def __init__(
         memory_service: The memory service for the runner.
         credential_service: The credential service for the runner.
         plugin_close_timeout: The timeout in seconds for plugin close methods.
-        auto_create_session: Whether to automatically create a session when
-          not found. Defaults to False. If False, a missing session raises
+        auto_create_session: Whether to automatically create a session when not
+          found. Defaults to False. If False, a missing session raises
           ValueError with a helpful message.
 
     Raises:
@@ -1071,6 +1071,7 @@ async def _run_with_trace(
               new_message=new_message,
               run_config=run_config,
               state_delta=state_delta,
+              invocation_id=invocation_id,
           )
         else:
           invocation_id = self._resolve_invocation_id(
@@ -1833,6 +1834,7 @@ async def _setup_context_for_new_invocation(
       new_message: types.Content,
       run_config: RunConfig,
       state_delta: Optional[dict[str, Any]],
+      invocation_id: Optional[str] = None,
   ) -> InvocationContext:
     """Sets up the context for a new invocation.
 
@@ -1841,6 +1843,7 @@ async def _setup_context_for_new_invocation(
       new_message: The new message to process and append to the session.
       run_config: The run config of the agent.
       state_delta: Optional state changes to apply to the session.
+      invocation_id: Optional invocation identifier.
 
     Returns:
       The invocation context for the new invocation.
@@ -1850,6 +1853,7 @@ async def _setup_context_for_new_invocation(
         session,
         new_message=new_message,
         run_config=run_config,
+        invocation_id=invocation_id,
     )
     # Step 2: Handle new message, by running callbacks and appending to
     # session.

tests/unittests/test_runners.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from contextlib import aclosing
 import importlib
 from pathlib import Path
 import sys
@@ -763,6 +764,48 @@ def test_run_passes_state_delta():
   assert user_event.actions.state_delta == state_delta
 
 
+@pytest.mark.asyncio
+async def test_run_async_propagates_invocation_id():
+  """run_async should propagate invocation_id to the invocation context and events."""
+
+  session_service = InMemorySessionService()
+  runner = Runner(
+      app_name=TEST_APP_ID,
+      agent=MockAgent("test_agent"),
+      session_service=session_service,
+      artifact_service=InMemoryArtifactService(),
+      auto_create_session=True,
+  )
+
+  custom_invocation_id = "my_custom_invocation_id"
+
+  agen = runner.run_async(
+      user_id=TEST_USER_ID,
+      session_id=TEST_SESSION_ID,
+      new_message=types.Content(role="user", parts=[types.Part(text="hello")]),
+      invocation_id=custom_invocation_id,
+  )
+
+  events = []
+  async with aclosing(agen) as a:
+    async for event in a:
+      events.append(event)
+
+  assert len(events) >= 1
+  # Verify yielded events have the custom invocation ID
+  for event in events:
+    assert event.invocation_id == custom_invocation_id
+
+  # Verify the session has the custom invocation ID in its events
+  session = await session_service.get_session(
+      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+  )
+  assert session is not None
+  assert len(session.events) == 2
+  for event in session.events:
+    assert event.invocation_id == custom_invocation_id
+
+
 @pytest.mark.asyncio
 async def test_run_live_persists_event_callback_modifications():
   """run_live should persist the same event it streams after callback changes."""