chore: Refactor runner to infer invocation_id from FunctionResponse Event for HITL resuming

invocation_id is no longer required in resuming case, unless no new_message is provided.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 875432024

Author: DeanChensj

src/google/adk/agents/invocation_context.py

@@ -396,23 +396,20 @@ def should_pause_invocation(self, event: Event) -> bool:
     return False
 
   # TODO: Move this method from invocation_context to a dedicated module.
-  # TODO: Converge this method with find_matching_function_call in llm_flows.
   def _find_matching_function_call(
       self, function_response_event: Event
   ) -> Optional[Event]:
     """Finds the function call event in the current invocation that matches the function response id."""
+    from ..flows.llm_flows.functions import find_event_by_function_call_id
+
     function_responses = function_response_event.get_function_responses()
     if not function_responses:
       return None
-    function_call_id = function_responses[0].id
-
-    events = self._get_events(current_invocation=True)
-    # The last event is function_response_event, so we search backwards from the
-    # one before it.
-    for event in reversed(events[:-1]):
-      if any(fc.id == function_call_id for fc in event.get_function_calls()):
-        return event
-    return None
+
+    # Search backwards from the event before the current response event.
+    return find_event_by_function_call_id(
+        self._get_events(current_invocation=True)[:-1], function_responses[0].id
+    )
 
 
 def new_invocation_context_id() -> str:

src/google/adk/flows/llm_flows/functions.py

@@ -37,7 +37,6 @@
 from google.genai import types
 
 from ...agents.active_streaming_tool import ActiveStreamingTool
-from ...agents.invocation_context import InvocationContext
 from ...agents.live_request_queue import LiveRequestQueue
 from ...auth.auth_tool import AuthConfig
 from ...auth.auth_tool import AuthToolArguments
@@ -52,6 +51,7 @@
 from ...utils.context_utils import Aclosing
 
 if TYPE_CHECKING:
+  from ...agents.invocation_context import InvocationContext
   from ...agents.llm_agent import LlmAgent
 
 AF_FUNCTION_CALL_ID_PREFIX = 'adk-'
@@ -1157,6 +1157,18 @@ def merge_parallel_function_response_events(
   return merged_event
 
 
+def find_event_by_function_call_id(
+    events: list[Event],
+    function_call_id: str,
+) -> Optional[Event]:
+  """Finds the function call event that matches the function call id."""
+  for event in reversed(events):
+    for function_call in event.get_function_calls():
+      if function_call.id == function_call_id:
+        return event
+  return None
+
+
 def find_matching_function_call(
     events: list[Event],
 ) -> Optional[Event]:
@@ -1165,25 +1177,8 @@ def find_matching_function_call(
     return None
 
   last_event = events[-1]
-  if (
-      last_event.content
-      and last_event.content.parts
-      and any(part.function_response for part in last_event.content.parts)
-  ):
+  function_responses = last_event.get_function_responses()
+  if not function_responses:
+    return None
 
-    function_call_id = next(
-        part.function_response.id
-        for part in last_event.content.parts
-        if part.function_response
-    )
-    for i in range(len(events) - 2, -1, -1):
-      event = events[i]
-      # looking for the system long-running request euc function call
-      function_calls = event.get_function_calls()
-      if not function_calls:
-        continue
-
-      for function_call in function_calls:
-        if function_call.id == function_call_id:
-          return event
-  return None
+  return find_event_by_function_call_id(events[:-1], function_responses[0].id)

src/google/adk/runners.py

@@ -49,6 +49,7 @@
 from .events.event import Event
 from .events.event import EventActions
 from .flows.llm_flows import contents
+from .flows.llm_flows.functions import find_event_by_function_call_id
 from .flows.llm_flows.functions import find_matching_function_call
 from .memory.base_memory_service import BaseMemoryService
 from .memory.in_memory_memory_service import InMemoryMemoryService
@@ -70,6 +71,16 @@ def _is_tool_call_or_response(event: Event) -> bool:
   return bool(event.get_function_calls() or event.get_function_responses())
 
 
+def _get_function_responses_from_content(
+    content: types.Content,
+) -> list[types.FunctionResponse]:
+  if not content:
+    return []
+  return [
+      part.function_response for part in content.parts if part.function_response
+  ]
+
+
 def _is_transcription(event: Event) -> bool:
   return (
       event.input_transcription is not None
@@ -341,6 +352,35 @@ def _enforce_app_name_alignment(self) -> None:
     self._app_name_alignment_hint = f'{mismatch_details} {resolution}'
     logger.warning('App name mismatch detected. %s', mismatch_details)
 
+  def _resolve_invocation_id(
+      self,
+      session: Session,
+      new_message: Optional[types.Content],
+      invocation_id: Optional[str],
+  ) -> Optional[str]:
+    """Infers invocation_id from new_message if it is a function response."""
+    function_responses = _get_function_responses_from_content(new_message)
+    if not function_responses:
+      return invocation_id
+
+    fc_event = find_event_by_function_call_id(
+        session.events, function_responses[0].id
+    )
+    if not fc_event:
+      raise ValueError(
+          'Function call event not found for function response id:'
+          f' {function_responses[0].id}'
+      )
+
+    if invocation_id and invocation_id != fc_event.invocation_id:
+      logger.warning(
+          'Provided invocation_id %s is ignored because new_message has a '
+          'function response with invocation_id %s.',
+          invocation_id,
+          fc_event.invocation_id,
+      )
+    return fc_event.invocation_id
+
   def _format_session_not_found_message(self, session_id: str) -> str:
     message = f'Session not found: {session_id}'
     if not self._app_name_alignment_hint:
@@ -497,42 +537,57 @@ async def _run_with_trace(
         session = await self._get_or_create_session(
             user_id=user_id, session_id=session_id
         )
+
         if not invocation_id and not new_message:
           raise ValueError(
               'Running an agent requires either a new_message or an '
               'invocation_id to resume a previous invocation. '
               f'Session: {session_id}, User: {user_id}'
           )
 
-        if invocation_id:
-          if (
-              not self.resumability_config
-              or not self.resumability_config.is_resumable
-          ):
-            raise ValueError(
-                f'invocation_id: {invocation_id} is provided but the app is not'
-                ' resumable.'
-            )
-          invocation_context = await self._setup_context_for_resumed_invocation(
+        is_resumable = (
+            self.resumability_config and self.resumability_config.is_resumable
+        )
+        if not is_resumable and not new_message:
+          raise ValueError(
+              'Running an agent requires a new_message or a resumable app. '
+              f'Session: {session_id}, User: {user_id}'
+          )
+
+        if not is_resumable:
+          invocation_context = await self._setup_context_for_new_invocation(
               session=session,
               new_message=new_message,
-              invocation_id=invocation_id,
               run_config=run_config,
               state_delta=state_delta,
           )
-          if invocation_context.end_of_agents.get(
-              invocation_context.agent.name
-          ):
-            # Directly return if the current agent in invocation context is
-            # already final.
-            return
         else:
-          invocation_context = await self._setup_context_for_new_invocation(
-              session=session,
-              new_message=new_message,  # new_message is not None.
-              run_config=run_config,
-              state_delta=state_delta,
+          invocation_id = self._resolve_invocation_id(
+              session, new_message, invocation_id
           )
+          if not invocation_id:
+            invocation_context = await self._setup_context_for_new_invocation(
+                session=session,
+                new_message=new_message,
+                run_config=run_config,
+                state_delta=state_delta,
+            )
+          else:
+            invocation_context = (
+                await self._setup_context_for_resumed_invocation(
+                    session=session,
+                    new_message=new_message,
+                    invocation_id=invocation_id,
+                    run_config=run_config,
+                    state_delta=state_delta,
+                )
+            )
+            if invocation_context.end_of_agents.get(
+                invocation_context.agent.name
+            ):
+              # Directly return if the current agent in invocation context is
+              # already final.
+              return
 
         async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
           async with Aclosing(ctx.agent.run_async(ctx)) as agen:

tests/unittests/runners/test_run_tool_confirmation.py

@@ -502,6 +502,86 @@ async def test_pause_and_resume_on_request_confirmation(
         == expected_parts_final
     )
 
+  @pytest.mark.asyncio
+  async def test_pause_and_resume_on_request_confirmation_without_invocation_id(
+      self,
+      runner: testing_utils.InMemoryRunner,
+      agent: LlmAgent,
+  ):
+    """Tests HITL flow where all tool calls are confirmed."""
+    events = runner.run("test user query")
+
+    # Verify that the invocation is paused when tool confirmation is requested.
+    # The tool call returns error response, and summarization was skipped.
+    assert testing_utils.simplify_resumable_app_events(
+        copy.deepcopy(events)
+    ) == [
+        (
+            agent.name,
+            Part(function_call=FunctionCall(name=agent.tools[0].name, args={})),
+        ),
+        (
+            agent.name,
+            Part(
+                function_call=FunctionCall(
+                    name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                    args={
+                        "originalFunctionCall": {
+                            "name": agent.tools[0].name,
+                            "id": mock.ANY,
+                            "args": {},
+                        },
+                        "toolConfirmation": {
+                            "hint": HINT_TEXT,
+                            "confirmed": False,
+                        },
+                    },
+                )
+            ),
+        ),
+        (
+            agent.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=agent.tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
+                )
+            ),
+        ),
+    ]
+    ask_for_confirmation_function_call_id = (
+        events[1].content.parts[0].function_call.id
+    )
+    invocation_id = events[1].invocation_id
+    user_confirmation = testing_utils.UserContent(
+        Part(
+            function_response=FunctionResponse(
+                id=ask_for_confirmation_function_call_id,
+                name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                response={"confirmed": True},
+            )
+        )
+    )
+    events = await runner.run_async(user_confirmation)
+    expected_parts_final = [
+        (
+            agent.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=agent.tools[0].name,
+                    response={"result": "confirmed=True"},
+                )
+            ),
+        ),
+        (agent.name, "test llm response after tool call"),
+        (agent.name, testing_utils.END_OF_AGENT),
+    ]
+    for event in events:
+      assert event.invocation_id == invocation_id
+    assert (
+        testing_utils.simplify_resumable_app_events(copy.deepcopy(events))
+        == expected_parts_final
+    )
+
 
 class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
   """Tests the HITL confirmation flow with a resumable sequential agent app."""