fix(live): stop background tools before the transfer delay

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

@@ -813,6 +813,17 @@ async def run_live_flow(
                 # instead of calling `transfer_to_agent`.
                 transfer_to_agent = event.actions.transfer_to_agent
                 if transfer_to_agent:
+                  # The sub agent takes over the live request queue, so this
+                  # agent's background tools have to stop here rather than
+                  # when this run_live eventually returns: it does not return
+                  # until the sub agent is done, and until then a tool of this
+                  # agent would keep feeding function responses to a model
+                  # that never made those calls. They stop before the transfer
+                  # delay below rather than after it: for the length of that
+                  # delay the connection is still open and the send task is
+                  # still draining the live request queue, so a tool that
+                  # finished inside the window would be forwarded after all.
+                  await flow._stop_background_tool_tasks(invocation_context)
                   await asyncio.sleep(
                       base_llm_flow.DEFAULT_TRANSFER_AGENT_DELAY
                   )
@@ -821,13 +832,6 @@ async def run_live_flow(
                   logger.debug('Closing live connection')
                   await llm_connection.close()
                   logger.debug('Live connection closed.')
-                  # The sub agent takes over the live request queue, so this
-                  # agent's background tools have to stop here rather than
-                  # when this run_live eventually returns: it does not return
-                  # until the sub agent is done, and until then a tool of this
-                  # agent would keep feeding function responses to a model
-                  # that never made those calls.
-                  await flow._stop_background_tool_tasks(invocation_context)
                   # transfer to the sub agent.
                   logger.debug('Transferring to agent: %s', transfer_to_agent)
                   agent_to_run = flow._get_agent_to_run(

tests/unittests/streaming/test_live_tool_shutdown.py

@@ -28,6 +28,7 @@
 from typing import Any
 from typing import AsyncGenerator
 from typing import Callable
+from unittest import mock
 
 from google.adk.agents.active_streaming_tool import ActiveStreamingTool
 from google.adk.agents.invocation_context import InvocationContext
@@ -469,3 +470,179 @@ async def task_completed() -> str:
 
   assert ended
   assert cancelled.is_set()
+
+
+def _late_tool_agents(tool: Any, call_name: str) -> Agent:
+  """A root agent that calls ``tool`` and then hands off to a sub agent."""
+
+  def report() -> str:
+    """The sub agent's own tool, proving it took over the queue."""
+    return 'sub agent is live'
+
+  sub_agent = Agent(
+      name='sub_agent',
+      model=testing_utils.MockModel.create([_call('report')]),
+      tools=[report],
+  )
+  root_agent = Agent(
+      name='root_agent',
+      model=testing_utils.MockModel.create([
+          _call(call_name),
+          LlmResponse(
+              content=types.Content(
+                  role='model',
+                  parts=[
+                      types.Part.from_function_call(
+                          name='transfer_to_agent',
+                          args={'agent_name': 'sub_agent'},
+                      )
+                  ],
+              ),
+              turn_complete=False,
+          ),
+      ]),
+      tools=[tool],
+      sub_agents=[sub_agent],
+  )
+  return root_agent
+
+
+async def _handoff_with_background_tool(
+    *,
+    transfer_delay: float,
+    completes_at: float,
+    streaming: bool = False,
+) -> tuple[bool, list[str]]:
+  """Hands off while a background tool of the root agent is still in flight.
+
+  Returns whether the tool ever started, and the names of the function
+  responses that reached the model connection after the transfer event was
+  yielded. A streaming tool legitimately streams to its own agent's model
+  before the handoff, so only what follows the handoff is of interest.
+  """
+  started = asyncio.Event()
+  transferred = asyncio.Event()
+  sent: list[tuple[bool, types.Content]] = []
+
+  async def _record_send_content(self, content, *, partial=False) -> None:
+    sent.append((transferred.is_set(), content))
+
+  async def late_lookup() -> str:
+    started.set()
+    await asyncio.sleep(completes_at)
+    return 'late'
+
+  async def late_stream() -> AsyncGenerator[Any, None]:
+    started.set()
+    while True:
+      await asyncio.sleep(completes_at)
+      yield {'late': True}
+
+  if streaming:
+    tool: Any = late_stream
+    call_name = 'late_stream'
+  else:
+    scheduled = FunctionTool(func=late_lookup)
+    scheduled.response_scheduling = types.FunctionResponseScheduling.SILENT
+    tool = scheduled
+    call_name = 'late_lookup'
+
+  root_agent = _late_tool_agents(tool, call_name)
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(app_name='app', user_id='u')
+  runner = Runner(
+      app_name='app', agent=root_agent, session_service=session_service
+  )
+  live_request_queue = LiveRequestQueue()
+  live_request_queue.send_realtime(
+      types.Blob(data=b'question', mime_type='audio/pcm')
+  )
+
+  async def _consume() -> None:
+    async with aclosing(
+        runner.run_live(
+            user_id='u',
+            session_id=session.id,
+            live_request_queue=live_request_queue,
+            run_config=RunConfig(response_modalities=['TEXT']),
+        )
+    ) as agen:
+      seen = 0
+      async for event in agen:
+        seen += 1
+        if event.actions and event.actions.transfer_to_agent:
+          transferred.set()
+        if seen >= _MAX_EVENTS:
+          return
+
+  with (
+      mock.patch.object(
+          testing_utils.MockLlmConnection, '_send_content', _record_send_content
+      ),
+      mock.patch.object(
+          base_llm_flow, 'DEFAULT_TRANSFER_AGENT_DELAY', transfer_delay
+      ),
+  ):
+    try:
+      await asyncio.wait_for(_consume(), timeout=10.0)
+    except asyncio.TimeoutError:
+      pass
+    # Give a tool that outlived the handoff time to report in.
+    await asyncio.sleep(max(completes_at, transfer_delay))
+
+  forwarded = [
+      part.function_response.name
+      for after_transfer, content in sent
+      if after_transfer
+      for part in content.parts or []
+      if part.function_response
+  ]
+  return started.is_set(), forwarded
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'transfer_delay, completes_at',
+    [
+        (0.5, 0.05),  # finishes early in the delay window
+        (0.5, 0.4),  # finishes just before the window closes
+        (0.5, 2.0),  # would finish long after the handoff
+        (1.0, 0.5),  # the shipped delay, tool halfway through it
+        (0.0, 0.05),  # no window at all
+    ],
+)
+async def test_handoff_stops_tools_before_the_transfer_delay(
+    transfer_delay: float, completes_at: float
+):
+  """A tool finishing during the transfer delay never reaches the model.
+
+  The handoff waits ``DEFAULT_TRANSFER_AGENT_DELAY`` before it cancels the send
+  task and closes the connection. The connection is open for the whole of that
+  wait and the send task is still draining the live request queue, so a tool of
+  the handing-off agent that completes inside the window would have its
+  function response forwarded to a model that never called it.
+  """
+  started, forwarded = await _handoff_with_background_tool(
+      transfer_delay=transfer_delay, completes_at=completes_at
+  )
+
+  # Without this the assertion below would hold for a tool that never ran.
+  assert started
+  assert 'late_lookup' not in forwarded, (
+      "the handing-off agent's tool finished during the transfer delay and its"
+      ' response was forwarded to a model that never called it'
+  )
+
+
+@pytest.mark.asyncio
+async def test_handoff_stops_streaming_tools_before_the_transfer_delay():
+  """The same window closes for streaming tools, which yield repeatedly."""
+  started, forwarded = await _handoff_with_background_tool(
+      transfer_delay=0.5, completes_at=0.05, streaming=True
+  )
+
+  assert started
+  assert 'late_stream' not in forwarded, (
+      "the handing-off agent's streaming tool yielded during the transfer"
+      ' delay and its response was forwarded to a model that never called it'
+  )