feat: Modify runner to support resuming an invocation (optionally with a function response)

PiperOrigin-RevId: 813008406

tests/unittests/runners/test_run_tool_confirmation.py

@@ -19,6 +19,8 @@ from unittest import mock
 
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.agents.sequential_agent import SequentialAgentState
 from google.adk.apps.app import App
 from google.adk.apps.app import ResumabilityConfig
 from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
@@ -186,6 +188,7 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
     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(
@@ -209,6 +212,8 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
         ),
         (agent.name, "test llm response after final tool call"),
     ]
+    for event in events:
+      assert event.invocation_id != invocation_id
     assert (
         testing_utils.simplify_events(copy.deepcopy(events))
         == expected_parts_final
@@ -321,6 +326,7 @@ class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest):
     ask_for_confirmation_function_call_id = (
         events[1].content.parts[0].function_call.id
     )
+    invocation_id = events[1].invocation_id
     custom_payload = {
         "test_custom_payload": {
             "int_field": 123,
@@ -358,6 +364,8 @@ class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest):
         ),
         (agent.name, "test llm response after final tool call"),
     ]
+    for event in events:
+      assert event.invocation_id != invocation_id
     assert (
         testing_utils.simplify_events(copy.deepcopy(events))
         == expected_parts_final
@@ -380,9 +388,6 @@ class TestHITLConfirmationFlowWithResumableApp:
     return [
         _create_llm_response_from_tools(tools),
         _create_llm_response_from_text("test llm response after tool call"),
-        _create_llm_response_from_text(
-            "test llm response after final tool call"
-        ),
     ]
 
   @pytest.fixture
@@ -412,7 +417,7 @@ class TestHITLConfirmationFlowWithResumableApp:
     return testing_utils.InMemoryRunner(app=app)
 
   @pytest.mark.asyncio
-  def test_pause_on_request_confirmation(
+  async def test_pause_and_resume_on_request_confirmation(
       self,
       runner: testing_utils.InMemoryRunner,
       agent: LlmAgent,
@@ -449,3 +454,189 @@ class TestHITLConfirmationFlowWithResumableApp:
             ),
         ),
     ]
+    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, invocation_id=invocation_id
+    )
+    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."""
+
+  @pytest.fixture
+  def tools(self) -> list[FunctionTool]:
+    """Provides the tools for the agent."""
+    return [FunctionTool(func=_test_request_confirmation_function)]
+
+  @pytest.fixture
+  def llm_responses(
+      self, tools: list[FunctionTool]
+  ) -> list[GenerateContentResponse]:
+    """Provides mock LLM responses for the tests."""
+    return [
+        _create_llm_response_from_tools(tools),
+        _create_llm_response_from_text("test llm response after tool call"),
+        _create_llm_response_from_text("test llm response from second agent"),
+    ]
+
+  @pytest.fixture
+  def mock_model(
+      self, llm_responses: list[GenerateContentResponse]
+  ) -> testing_utils.MockModel:
+    """Provides a mock model with predefined responses."""
+    return testing_utils.MockModel(responses=llm_responses)
+
+  @pytest.fixture
+  def agent(
+      self, mock_model: testing_utils.MockModel, tools: list[FunctionTool]
+  ) -> SequentialAgent:
+    """Provides a single LlmAgent for the test."""
+    return SequentialAgent(
+        name="root_agent",
+        sub_agents=[
+            LlmAgent(name="agent1", model=mock_model, tools=tools),
+            LlmAgent(name="agent2", model=mock_model, tools=[]),
+        ],
+    )
+
+  @pytest.fixture
+  def runner(self, agent: SequentialAgent) -> testing_utils.InMemoryRunner:
+    """Provides an in-memory runner for the agent."""
+    # Mark the app as resumable. So that the invocation will be paused after the
+    # long running tool call.
+    app = App(
+        name="test_app",
+        resumability_config=ResumabilityConfig(is_resumable=True),
+        root_agent=agent,
+    )
+    return testing_utils.InMemoryRunner(app=app)
+
+  @pytest.mark.asyncio
+  async def test_pause_and_resume_on_request_confirmation(
+      self,
+      runner: testing_utils.InMemoryRunner,
+      agent: SequentialAgent,
+  ):
+    """Tests HITL flow where all tool calls are confirmed."""
+    events = runner.run("test user query")
+    sub_agent1 = agent.sub_agents[0]
+    sub_agent2 = agent.sub_agents[1]
+
+    # Verify that the invocation is paused after the long running tool call.
+    # So that no intermediate function response and llm response is generated.
+    # And the second sub agent is not started.
+    assert testing_utils.simplify_resumable_app_events(
+        copy.deepcopy(events)
+    ) == [
+        (
+            agent.name,
+            SequentialAgentState(current_sub_agent=sub_agent1.name).model_dump(
+                mode="json"
+            ),
+        ),
+        (
+            sub_agent1.name,
+            Part(
+                function_call=FunctionCall(
+                    name=sub_agent1.tools[0].name, args={}
+                )
+            ),
+        ),
+        (
+            sub_agent1.name,
+            Part(
+                function_call=FunctionCall(
+                    name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                    args={
+                        "originalFunctionCall": {
+                            "name": sub_agent1.tools[0].name,
+                            "id": mock.ANY,
+                            "args": {},
+                        },
+                        "toolConfirmation": {
+                            "hint": "test hint for request_confirmation",
+                            "confirmed": False,
+                        },
+                    },
+                )
+            ),
+        ),
+    ]
+    ask_for_confirmation_function_call_id = (
+        events[2].content.parts[0].function_call.id
+    )
+    invocation_id = events[2].invocation_id
+
+    # Resume the invocation and confirm the tool call from sub_agent1, and
+    # sub_agent2 will continue.
+    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, invocation_id=invocation_id
+    )
+    expected_parts_final = [
+        (
+            sub_agent1.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=sub_agent1.tools[0].name,
+                    response={"result": "confirmed=True"},
+                )
+            ),
+        ),
+        (sub_agent1.name, "test llm response after tool call"),
+        (sub_agent1.name, testing_utils.END_OF_AGENT),
+        (
+            agent.name,
+            SequentialAgentState(current_sub_agent=sub_agent2.name).model_dump(
+                mode="json"
+            ),
+        ),
+        (sub_agent2.name, "test llm response from second agent"),
+        (sub_agent2.name, testing_utils.END_OF_AGENT),
+        (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
+    )