From 772658fd81fac71be1828707ff893ce362540944 Mon Sep 17 00:00:00 2001
From: "Xinran (Sherry) Tang" <txinran@google.com>
Date: Mon, 29 Sep 2025 12:21:10 -0700
Subject: [PATCH] chore: Refactor runner run_async flow to extract out
 execution context setup logic

PiperOrigin-RevId: 812894540
---
 src/google/adk/runners.py | 100 ++++++++++++++++++++++++++++++--------
 1 file changed, 79 insertions(+), 21 deletions(-)

diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py
index 853368b2..12ffdb17 100644
--- a/src/google/adk/runners.py
+++ b/src/google/adk/runners.py
@@ -306,30 +306,12 @@ class Runner:
         if not session:
           raise ValueError(f'Session not found: {session_id}')
 
-        invocation_context = self._new_invocation_context(
-            session,
+        invocation_context = await self._setup_context_for_new_invocation(
+            session=session,
             new_message=new_message,
             run_config=run_config,
+            state_delta=state_delta,
         )
-        root_agent = self.agent
-
-        # Modify user message before execution.
-        modified_user_message = await invocation_context.plugin_manager.run_on_user_message_callback(
-            invocation_context=invocation_context, user_message=new_message
-        )
-        if modified_user_message is not None:
-          new_message = modified_user_message
-
-        if new_message:
-          await self._append_new_message_to_session(
-              session,
-              new_message,
-              invocation_context,
-              run_config.save_input_blobs_as_artifacts,
-              state_delta,
-          )
-
-        invocation_context.agent = self._find_agent_to_run(session, root_agent)
 
         async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
           async with Aclosing(ctx.agent.run_async(ctx)) as agen:
@@ -420,6 +402,7 @@ class Runner:
 
   async def _append_new_message_to_session(
       self,
+      *,
       session: Session,
       new_message: types.Content,
       invocation_context: InvocationContext,
@@ -433,6 +416,7 @@ class Runner:
         new_message: The new message to append.
         invocation_context: The invocation context for the message.
         save_input_blobs_as_artifacts: Whether to save input blobs as artifacts.
+        state_delta: Optional state changes to apply to the session.
     """
     if not new_message.parts:
       raise ValueError('No parts in the new_message.')
@@ -661,6 +645,44 @@ class Runner:
       agent = agent.parent_agent
     return True
 
+  async def _setup_context_for_new_invocation(
+      self,
+      *,
+      session: Session,
+      new_message: types.Content,
+      run_config: RunConfig,
+      state_delta: Optional[dict[str, Any]],
+  ) -> InvocationContext:
+    """Sets up the context for a new invocation.
+
+    Args:
+      session: The session to setup the invocation context for.
+      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.
+
+    Returns:
+      The invocation context for the new invocation.
+    """
+    # Step 1: Create invocation context in memory.
+    invocation_context = self._new_invocation_context(
+        session,
+        new_message=new_message,
+        run_config=run_config,
+    )
+    # Step 2: Handle new message, by running callbacks and appending to
+    # session.
+    await self._handle_new_message(
+        session=session,
+        new_message=new_message,
+        invocation_context=invocation_context,
+        run_config=run_config,
+        state_delta=state_delta,
+    )
+    # Step 3: Set agent to run for the invocation.
+    invocation_context.agent = self._find_agent_to_run(session, self.agent)
+    return invocation_context
+
   def _new_invocation_context(
       self,
       session: Session,
@@ -743,6 +765,42 @@ class Runner:
         run_config=run_config,
     )
 
+  async def _handle_new_message(
+      self,
+      *,
+      session: Session,
+      new_message: types.Content,
+      invocation_context: InvocationContext,
+      run_config: RunConfig,
+      state_delta: Optional[dict[str, Any]],
+  ) -> None:
+    """Handles a new message by running callbacks and appending to session.
+
+    Args:
+      session: The session of the new message.
+      new_message: The new message to process and append to the session.
+      invocation_context: The invocation context to use for the message
+        handling.
+      run_config: The run config of the agent.
+      state_delta: Optional state changes to apply to the session.
+    """
+    modified_user_message = (
+        await invocation_context.plugin_manager.run_on_user_message_callback(
+            invocation_context=invocation_context, user_message=new_message
+        )
+    )
+    if modified_user_message is not None:
+      new_message = modified_user_message
+
+    if new_message:
+      await self._append_new_message_to_session(
+          session=session,
+          new_message=new_message,
+          invocation_context=invocation_context,
+          save_input_blobs_as_artifacts=run_config.save_input_blobs_as_artifacts,
+          state_delta=state_delta,
+      )
+
   def _collect_toolset(self, agent: BaseAgent) -> set[BaseToolset]:
     toolsets = set()
     if isinstance(agent, LlmAgent):