chore: import upstream snapshot with attribution

tests/unittests/agents/test_llm_agent_interruptions.py

@@ -0,0 +1,481 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.agents import LlmAgent
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.long_running_tool import LongRunningFunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from tests.unittests import testing_utils
+from tests.unittests.agents.llm.event_utils import text_parts
+
+_USER_ID = 'test_user'
+_SESSION_ID = 'test_session'
+
+
+async def _setup_runner(mock_model, tools=None, **agent_kwargs):
+  """Setup runner with LlmAgent directly."""
+  llm_agent = LlmAgent(
+      name='test_agent',
+      model=mock_model,
+      tools=tools or [],
+      **agent_kwargs,
+  )
+  session_service = InMemorySessionService()
+  await session_service.create_session(
+      app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+  )
+  runner = Runner(
+      app_name='test',
+      agent=llm_agent,
+      session_service=session_service,
+  )
+  return runner
+
+
+async def _run_turn(runner, user_message):
+  """Run a single turn."""
+  return [
+      e
+      async for e in runner.run_async(
+          user_id=_USER_ID,
+          session_id=_SESSION_ID,
+          new_message=types.Content(
+              role='user', parts=[types.Part(text=user_message)]
+          ),
+      )
+  ]
+
+
+async def _resume_turn(
+    runner, prev_events, tool_name, tool_response_value='done'
+):
+  """Resume after an interrupt."""
+  fc_ids = []
+  for e in prev_events:
+    if e.content and e.content.parts:
+      for p in e.content.parts:
+        if (
+            p.function_call
+            and p.function_call.name == tool_name
+            and p.function_call.id
+        ):
+          fc_ids.append(p.function_call.id)
+    if getattr(e.output, 'function_calls', None):
+      for fc in e.output.function_calls:
+        if fc.name == tool_name and fc.id:
+          fc_ids.append(fc.id)
+
+  if not fc_ids:
+    for e in prev_events:
+      if e.long_running_tool_ids:
+        fc_ids = list(e.long_running_tool_ids)
+        break
+
+  invocation_id = prev_events[0].invocation_id
+
+  fr_parts = [
+      types.Part(
+          function_response=types.FunctionResponse(
+              name=tool_name,
+              id=fc_id,
+              response={'result': tool_response_value},
+          )
+      )
+      for fc_id in fc_ids
+  ]
+  resume_msg = types.Content(role='user', parts=fr_parts)
+
+  return [
+      e
+      async for e in runner.run_async(
+          user_id=_USER_ID,
+          session_id=_SESSION_ID,
+          invocation_id=invocation_id,
+          new_message=resume_msg,
+      )
+  ]
+
+
+def create_lro_tool(name: str = 'long_running_op') -> LongRunningFunctionTool:
+  """Creates a minimal LRO tool for testing."""
+
+  def _impl() -> None:
+    return None
+
+  _impl.__name__ = name
+  return LongRunningFunctionTool(_impl)
+
+
+# ---------------------------------------------------------------------------
+# Tests: Single Agent
+# ---------------------------------------------------------------------------
+
+
+class TestSingleAgentInterruptions:
+  """Tests for single agent triggering interruptions."""
+
+  async def test_single_agent_yields_on_long_running_tool(self):
+    """Single agent yields on Long Running Tool.
+
+    Arrange: Set up a single agent with a long running tool.
+    Act: Run the agent with a prompt that triggers the tool.
+    Assert: Verify that the execution yields a long running tool interrupt.
+    """
+
+    fc = types.Part.from_function_call(name='long_running_op', args={})
+    mock_model = testing_utils.MockModel.create(responses=[fc, 'Final answer'])
+
+    lro_tool = create_lro_tool()
+    runner = await _setup_runner(mock_model, tools=[lro_tool])
+
+    # Act: Run first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Assert: Should have triggered function call
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'long_running_op'
+            for p in e.content.parts or []
+        )
+        for e in events
+    )
+    assert len(mock_model.requests) == 1
+
+    # Act: Resume
+    resume_events = await _resume_turn(runner, events, 'long_running_op')
+
+    # Assert: Should have completed
+    assert any('Final answer' in t for t in text_parts(resume_events))
+    assert len(mock_model.requests) == 2
+
+  async def test_single_agent_request_input_tool_interrupt_and_resume(self):
+    """Test that using RequestInputTool successfully triggers an interrupt and resumes with user input."""
+    from google.adk.tools import request_input
+
+    fc = types.Part.from_function_call(
+        name='adk_request_input',
+        args={'message': 'Which file?', 'response_schema': {'type': 'string'}},
+    )
+    mock_model = testing_utils.MockModel.create(
+        responses=[fc, 'Continuing with file: file_a.txt']
+    )
+
+    runner = await _setup_runner(mock_model, tools=[request_input])
+
+    # Act: Run first turn
+    events = await _run_turn(runner, 'Start')
+
+    # Assert: Verify the interrupt event is produced
+    assert any(e.long_running_tool_ids for e in events)
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in e.content.parts or []
+        )
+        for e in events
+    )
+
+    # Act: Resume with the response
+    resume_events = await _resume_turn(
+        runner, events, 'adk_request_input', tool_response_value='file_a.txt'
+    )
+
+    # Assert: Execution should continue with user response in the prompt history
+    assert len(mock_model.requests) == 2
+
+    # Assert: Verify the second request contains the FunctionCall & FunctionResponse pair
+    second_req_contents = mock_model.requests[1].contents
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+    assert any(
+        any(
+            p.function_response
+            and p.function_response.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+
+  async def test_single_agent_request_input_tool_structured_schema(self):
+    """Test that using RequestInputTool with a structured object schema successfully interrupts and resumes with a dictionary response."""
+    from google.adk.tools import request_input
+
+    schema = {
+        'type': 'object',
+        'properties': {
+            'host': {'type': 'string'},
+            'port': {'type': 'integer'},
+        },
+        'required': ['host'],
+    }
+    fc = types.Part.from_function_call(
+        name='adk_request_input',
+        args={
+            'message': 'Provide DB connection details:',
+            'response_schema': schema,
+        },
+    )
+    mock_model = testing_utils.MockModel.create(
+        responses=[fc, 'Connected to localhost:3306']
+    )
+
+    runner = await _setup_runner(mock_model, tools=[request_input])
+
+    # Act: Run first turn
+    events = await _run_turn(runner, 'Start')
+
+    # Assert: Verify the interrupt event is produced with the schema args
+    assert any(e.long_running_tool_ids for e in events)
+    fc_event = next(
+        e
+        for e in events
+        if e.content
+        and any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in e.content.parts or []
+        )
+    )
+    fc_part = next(p for p in fc_event.content.parts if p.function_call)
+    assert fc_part.function_call.args['response_schema'] == schema
+
+    # Act: Resume with a structured dict response
+    db_details = {'host': 'localhost', 'port': 3306}
+    resume_events = await _resume_turn(
+        runner, events, 'adk_request_input', tool_response_value=db_details
+    )
+
+    # Assert: Execution should continue with the structured user response
+    assert len(mock_model.requests) == 2
+
+    # Assert: Verify the second request contains the FunctionCall & FunctionResponse pair
+    second_req_contents = mock_model.requests[1].contents
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+    assert any(
+        any(
+            p.function_response
+            and p.function_response.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+
+
+class TestNestedAgentInterruptions:
+  """Tests for multi-agent setups with interruptions."""
+
+  async def test_child_agent_interrupt_and_resume(self):
+    """Child agent yields on LRO and resumes successfully.
+
+    Arrange: Parent agent with Child agent. Parent transfers to Child.
+      Child calls LRO tool.
+    Act: Run, expect LRO interrupt. Then resume.
+    Assert: Should complete successfully.
+    """
+
+    def transfer_to_child(tool_context: ToolContext) -> str:
+      tool_context.actions.transfer_to_agent = 'child_agent'
+      return 'transferring'
+
+    # Child agent
+    fc_child = types.Part.from_function_call(name='child_lro', args={})
+    child_mock_model = testing_utils.MockModel.create(
+        responses=[fc_child, 'Child final answer']
+    )
+
+    lro_tool = create_lro_tool('child_lro')
+
+    child_agent = LlmAgent(
+        name='child_agent',
+        model=child_mock_model,
+        tools=[lro_tool],
+    )
+
+    # Parent agent
+    fc_parent = types.Part.from_function_call(name='transfer_to_child', args={})
+    parent_mock_model = testing_utils.MockModel.create(
+        responses=[fc_parent, fc_parent, 'Parent final answer']
+    )
+
+    parent_agent = LlmAgent(
+        name='parent_agent',
+        model=parent_mock_model,
+        tools=[transfer_to_child],
+        sub_agents=[child_agent],
+    )
+
+    # Setup runner
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+    )
+    runner = Runner(
+        app_name='test', agent=parent_agent, session_service=session_service
+    )
+
+    # When Parent runs the first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Then it should trigger child LRO interrupt
+    assert any(e.long_running_tool_ids for e in events)
+
+    # When Parent resumes the turn
+    resume_events = await _resume_turn(runner, events, 'child_lro')
+
+    # Then it should complete successfully
+    assert any('Child final answer' in t for t in text_parts(resume_events))
+
+  @pytest.mark.xfail(reason='Task agent as subagent not supported yet.')
+  async def test_task_child_agent_interrupt_and_resume(self):
+    """Task child agent yields on LRO and resumes successfully.
+
+    Arrange: Parent agent with Task Child agent. Parent transfers to Child.
+      Child calls LRO tool.
+    Act: Run, expect LRO interrupt. Then resume.
+    Assert: Should complete successfully.
+    """
+
+    def transfer_to_child(tool_context: ToolContext) -> str:
+      tool_context.actions.transfer_to_agent = 'child_agent'
+      return 'transferring'
+
+    # Child agent (Task mode)
+    fc_child = types.Part.from_function_call(name='child_lro', args={})
+    fc_finish = types.Part.from_function_call(
+        name='finish_task', args={'result': 'Task done'}
+    )
+    child_mock_model = testing_utils.MockModel.create(
+        responses=[fc_child, fc_finish, 'Child final answer']
+    )
+
+    lro_tool = create_lro_tool('child_lro')
+
+    child_agent = LlmAgent(
+        name='child_agent',
+        model=child_mock_model,
+        tools=[lro_tool],
+        mode='task',
+    )
+
+    # Parent agent
+    fc_parent = types.Part.from_function_call(name='transfer_to_child', args={})
+    parent_mock_model = testing_utils.MockModel.create(
+        responses=[fc_parent, 'Parent final answer']
+    )
+
+    parent_agent = LlmAgent(
+        name='parent_agent',
+        model=parent_mock_model,
+        tools=[transfer_to_child],
+        sub_agents=[child_agent],
+    )
+
+    # Setup runner
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+    )
+    runner = Runner(
+        app_name='test', agent=parent_agent, session_service=session_service
+    )
+
+    # When Parent runs the first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Then it should trigger child LRO interrupt
+    assert any(e.long_running_tool_ids for e in events)
+
+    # When Parent resumes the turn
+    resume_events = await _resume_turn(runner, events, 'child_lro')
+
+    # Then it should complete successfully
+    assert any('Parent final answer' in t for t in text_parts(resume_events))
+
+  @pytest.mark.xfail(reason='Single-turn agent as subagent not supported yet.')
+  async def test_single_turn_child_agent_interrupt_and_resume(self):
+    """Single-turn child agent yields on LRO and resumes successfully.
+
+    Arrange: Parent agent with Single-turn Child agent.
+      Parent calls Child via tool.
+      Child calls LRO tool.
+    Act: Run, expect LRO interrupt. Then resume.
+    Assert: Should complete successfully.
+    """
+
+    # Child agent (Single-turn)
+    fc_child = types.Part.from_function_call(name='child_lro', args={})
+    child_mock_model = testing_utils.MockModel.create(
+        responses=[fc_child, 'Child final answer']
+    )
+
+    lro_tool = create_lro_tool('child_lro')
+
+    child_agent = LlmAgent(
+        name='child_agent',
+        model=child_mock_model,
+        tools=[lro_tool],
+        mode='single_turn',
+    )
+
+    # Parent agent
+    fc_call_child = types.Part.from_function_call(
+        name='child_agent', args={'request': 'Go to child'}
+    )
+    parent_mock_model = testing_utils.MockModel.create(
+        responses=[fc_call_child, 'Parent final answer']
+    )
+
+    parent_agent = LlmAgent(
+        name='parent_agent',
+        model=parent_mock_model,
+        sub_agents=[child_agent],
+    )
+
+    # Setup runner
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+    )
+    runner = Runner(
+        app_name='test', agent=parent_agent, session_service=session_service
+    )
+
+    # When Parent runs the first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Then it should trigger child LRO interrupt
+    assert any(e.long_running_tool_ids for e in events)
+
+    # When Parent resumes the turn
+    resume_events = await _resume_turn(runner, events, 'child_lro')
+
+    # Then it should complete successfully
+    assert any('Parent final answer' in t for t in text_parts(resume_events))