diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index e2259cf3669..e359cc56259 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -642,13 +642,18 @@ 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() @@ -656,29 +661,25 @@ async def run_live( 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 diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 35e2d8b8193..6cc86066410 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/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."""