fix: populate user_content in resumed invocations

Change-Id: I85fac02e9a926ff02ee9e95b03e6d31fb51926a4

Author: GWeale

src/google/adk/runners.py

@@ -459,6 +459,7 @@ async def _run_node_async(
       *,
       user_id: str,
       session_id: str,
+      invocation_id: Optional[str] = None,
       new_message: Optional[types.Content] = None,
       state_delta: Optional[dict[str, Any]] = None,
       run_config: Optional[RunConfig] = None,
@@ -481,11 +482,10 @@ async def _run_node_async(
       resume_inputs = self._extract_resume_inputs(new_message)
       self._validate_new_message(new_message, resume_inputs)
 
-      invocation_id = (
-          self._resolve_invocation_id_from_fr(session, new_message)
-          if new_message
-          else None
-      )
+      if not invocation_id and new_message:
+        invocation_id = self._resolve_invocation_id_from_fr(
+            session, new_message
+        )
 
       ic = self._new_invocation_context(
           session,
@@ -496,12 +496,16 @@ async def _run_node_async(
       ic._event_queue = asyncio.Queue()
 
       # 2. Append user message to session and resolve node_input
-      if resume_inputs:
-        # Resume: find original user message, use as node_input
+      node_input = None
+      if resume_inputs or invocation_id:
+        # Resume: recover the original user content. new_message here is a
+        # function response (or None), so it can't populate user_content.
         node_input = self._find_original_user_content(
             ic.session, ic.invocation_id
         )
-      else:
+        if node_input:
+          ic.user_content = node_input
+      if not node_input:
         # Fresh: use user message as node_input
         node_input = new_message
 
@@ -1019,6 +1023,7 @@ async def run_async(
           self._run_node_async(
               user_id=user_id,
               session_id=session_id,
+              invocation_id=invocation_id,
               new_message=new_message,
               state_delta=state_delta,
               run_config=run_config,
@@ -1039,6 +1044,7 @@ async def run_async(
           self._run_node_async(
               user_id=user_id,
               session_id=session_id,
+              invocation_id=invocation_id,
               new_message=new_message,
               state_delta=state_delta,
               run_config=run_config,

tests/unittests/runners/test_runner_node.py

@@ -476,6 +476,72 @@ async def _run_impl(
   assert 'original:my original input' in outputs
 
 
+@pytest.mark.asyncio
+async def test_resume_populates_invocation_user_content():
+  """On resume via a function response, ic.user_content is the original turn."""
+  seen: list[Any] = []
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      if ctx.resume_inputs and 'fc-1' in ctx.resume_inputs:
+        user_content = ctx.get_invocation_context().user_content
+        seen.append(user_content.parts[0].text if user_content else None)
+        yield 'resumed'
+        return
+      yield _make_interrupt_event(fc_name='tool')
+
+  await _run_two_turns(
+      _Node(name='node'),
+      'remember me',
+      _make_resume_message(fc_name='tool', response={'v': 1}),
+  )
+
+  assert seen == ['remember me']
+
+
+@pytest.mark.asyncio
+async def test_resume_by_invocation_id_populates_user_content():
+  """Resuming by invocation_id alone recovers the original user_content."""
+  seen: list[Any] = []
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      user_content = ctx.get_invocation_context().user_content
+      seen.append(user_content.parts[0].text if user_content else None)
+      yield _make_interrupt_event(fc_name='tool')
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=_Node(name='node'), session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  async for _ in runner.run_async(
+      user_id='u',
+      session_id=session.id,
+      new_message=types.Content(
+          parts=[types.Part(text='original text')], role='user'
+      ),
+  ):
+    pass
+
+  updated = await ss.get_session(
+      app_name='test', user_id='u', session_id=session.id
+  )
+  invocation_id = updated.events[0].invocation_id
+
+  async for _ in runner.run_async(
+      user_id='u', session_id=session.id, invocation_id=invocation_id
+  ):
+    pass
+
+  assert seen == ['original text', 'original text']
+
+
 @pytest.mark.asyncio
 async def test_plain_text_does_not_trigger_resume():
   """Sending plain text (no FR) starts fresh, does not enter resume path."""