fix(live): transfer to the target agent regardless of function response order

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

@@ -642,43 +642,44 @@ async def run_live(
  # the same function response. By handling agent transfer here,
  # we ensure that only child agent processes its own function
  # responses after the transfer.
- if (
- event.content
- and event.content.parts
- and event.content.parts[0].function_response
- and event.content.parts[0].function_response.name
- == 'transfer_to_agent'
- ):
+ #
+ # The transfer is gated on the `transfer_to_agent` action
+ # rather than on the position of the `transfer_to_agent`
+ # function response: the model may issue the transfer alongside
+ # other function calls, whose responses are merged into a
+ # single event in call order, so the transfer response is not
+ # necessarily `parts[0]`. Gating on the action matches
+ # `_postprocess_handle_function_calls_async`, and also covers
+ # tools that request a transfer by setting the action directly
+ # instead of calling `transfer_to_agent`.
+ transfer_to_agent = event.actions.transfer_to_agent
+ if transfer_to_agent:
  await asyncio.sleep(DEFAULT_TRANSFER_AGENT_DELAY)
  # cancel the tasks that belongs to the closed connection.
  send_task.cancel()
  logger.debug('Closing live connection')
  await llm_connection.close()
  logger.debug('Live connection closed.')
  # transfer to the sub agent.
- transfer_to_agent = event.actions.transfer_to_agent
- if transfer_to_agent:
- logger.debug('Transferring to agent: %s', transfer_to_agent)
- agent_to_run = self._get_agent_to_run(
- invocation_context, transfer_to_agent
+ logger.debug('Transferring to agent: %s', transfer_to_agent)
+ agent_to_run = self._get_agent_to_run(
+ invocation_context, transfer_to_agent
+ )
+ child_ctx = invocation_context.model_copy()
+ # Child Live agent should start a new Live session.
+ # Do not reuse the parent session's resumption handle.
+ child_ctx.live_session_resumption_handle = None
+
+ if child_ctx.run_config:
+ child_ctx.run_config = child_ctx.run_config.model_copy(
+ deep=True
  )
- child_ctx = invocation_context.model_copy()
- # Child Live agent should start a new Live session.
- # Do not reuse the parent session's resumption handle.
- child_ctx.live_session_resumption_handle = None
-
- if child_ctx.run_config:
- child_ctx.run_config = child_ctx.run_config.model_copy(
- deep=True
- )
- if child_ctx.run_config.session_resumption:
- child_ctx.run_config.session_resumption.handle = None
-
- async with Aclosing(
- agent_to_run.run_live(child_ctx)
- ) as agen:
- async for item in agen:
- yield item
+ if child_ctx.run_config.session_resumption:
+ child_ctx.run_config.session_resumption.handle = None
+
+ async with Aclosing(agent_to_run.run_live(child_ctx)) as agen:
+ async for item in agen:
+ yield item
  if (
  event.content
  and event.content.parts

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -14,6 +14,7 @@
 
 """Unit tests for BaseLlmFlow toolset integration."""
 
+from typing import Optional
 from unittest import mock
 from unittest.mock import AsyncMock
 
@@ -1081,6 +1082,121 @@ async def mock_run_live_sub_agent(child_ctx, *args, **kwargs):
  )
 
 
+@pytest.mark.parametrize(
+ ('function_response_names', 'transfer_action', 'expect_transfer'),
+ [
+ # A lone transfer call.
+ (('transfer_to_agent',), 'sub_agent', True),
+ # Parallel calls whose transfer response is merged first.
+ (('transfer_to_agent', 'set_state'), 'sub_agent', True),
+ # Parallel calls whose transfer response is merged after another
+ # tool's response, so it is not `parts[0]`.
+ (('set_state', 'transfer_to_agent'), 'sub_agent', True),
+ (('set_state', 'log_event', 'transfer_to_agent'), 'sub_agent', True),
+ # A tool that requests the transfer by setting the action directly
+ # instead of calling `transfer_to_agent`.
+ (('escalate',), 'sub_agent', True),
+ # Parallel calls that do not transfer.
+ (('set_state', 'other_tool'), None, False),
+ # A transfer response whose action was suppressed, e.g. by a
+ # `before_tool_callback` overriding the transfer tool. The parent
+ # connection must stay open because no child agent takes over.
+ (('transfer_to_agent',), None, False),
+ (('set_state', 'transfer_to_agent'), None, False),
+ ],
+)
+@pytest.mark.asyncio
+async def test_run_live_transfer_is_independent_of_response_order(
+ function_response_names: tuple[str, ...],
+ transfer_action: Optional[str],
+ expect_transfer: bool,
+):
+ """Live transfer keys off the action, not the transfer response's position."""
+
+ agent = Agent(name='test_agent')
+ invocation_context = await testing_utils.create_invocation_context(
+ agent=agent
+ )
+ invocation_context.live_request_queue = LiveRequestQueue()
+ invocation_context.run_config = RunConfig()
+
+ flow = BaseLlmFlowForTesting()
+
+ # Parallel function responses are merged into a single event in call order
+ # by `merge_parallel_function_response_events`, so the transfer response may
+ # land at any index.
+ function_response_event = Event(
+ id=Event.new_id(),
+ invocation_id=invocation_context.invocation_id,
+ author=agent.name,
+ content=types.Content(
+ role='user',
+ parts=[
+ types.Part(
+ function_response=types.FunctionResponse(name=name),
+ )
+ for name in function_response_names
+ ],
+ ),
+ )
+ function_response_event.actions.transfer_to_agent = transfer_action
+
+ # A follow-up model turn, used to tell a live parent connection that is still
+ # usable apart from one that was torn down without a child taking over.
+ follow_up_event = Event(
+ id=Event.new_id(),
+ invocation_id=invocation_context.invocation_id,
+ author=agent.name,
+ content=types.Content(role='model', parts=[types.Part(text='follow up')]),
+ )
+
+ async def mock_receive_from_model(*args, **kwargs):
+ yield function_response_event
+ yield follow_up_event
+
+ flow._receive_from_model = mock.Mock(side_effect=mock_receive_from_model)
+
+ mock_sub_agent = mock.Mock()
+
+ async def mock_run_live_sub_agent(child_ctx, *args, **kwargs):
+ for item in []:
+ yield item
+
+ mock_sub_agent.run_live = mock.Mock(side_effect=mock_run_live_sub_agent)
+ flow._get_agent_to_run = mock.Mock(return_value=mock_sub_agent)
+
+ # Mock _send_to_model to prevent it from running indefinitely
+ flow._send_to_model = mock.AsyncMock()
+
+ with (
+ mock.patch('google.adk.models.google_llm.Gemini.connect') as mock_connect,
+ mock.patch(
+ 'google.adk.flows.llm_flows.base_llm_flow.DEFAULT_TRANSFER_AGENT_DELAY',
+ 0,
+ ),
+ ):
+ mock_connection = mock.AsyncMock()
+ mock_connect.return_value.__aenter__.return_value = mock_connection
+
+ events = [event async for event in flow.run_live(invocation_context)]
+
+ # The merged function response is always forwarded back to the model.
+ assert events[0] is function_response_event
+
+ if expect_transfer:
+ # The child agent takes over exactly once, and the parent connection is
+ # closed first so that only the child processes subsequent responses.
+ mock_sub_agent.run_live.assert_called_once()
+ assert flow._get_agent_to_run.call_args[0][1] == transfer_action
+ assert mock_connection.close.await_count == 1
+ else:
+ # No child agent takes over, so the parent connection must stay open and
+ # keep processing the live session.
+ mock_sub_agent.run_live.assert_not_called()
+ assert mock_connection.close.await_count == 0
+ assert follow_up_event in events
+
+
 @pytest.mark.asyncio
 async def test_postprocess_live_yields_grounding_metadata_only():
  """Test that _postprocess_live yields LlmResponse with only grounding_metadata."""