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 fa8466b7..6675d4db 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -757,12 +757,15 @@ class BaseLlmFlow(ABC): 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 - and event.content.parts[0].function_response - and event.content.parts[0].function_response.name - == 'task_completed' + # `task_completed` is an ordinary tool, so the model may call + # it alongside others. Their responses are merged into a single + # event in call order, so scan every response rather than only + # `parts[0]`. Unlike agent transfer there is no corresponding + # action to key off, since `task_completed` only signals + # completion through its function response. + if any( + function_response.name == 'task_completed' + for function_response in event.get_function_responses() ): # this is used for sequential agent to signal the end of the agent. await asyncio.sleep(DEFAULT_TASK_COMPLETION_DELAY) 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 2a6fde1d..5f0f0a70 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -1501,6 +1501,88 @@ async def test_run_live_transfer_is_independent_of_response_order( assert follow_up_event in events +@pytest.mark.parametrize( + ('function_response_names', 'expect_completion'), + [ + # A lone task_completed call. + (('task_completed',), True), + # Parallel calls whose task_completed response is merged first. + (('task_completed', 'set_state'), True), + # Parallel calls whose task_completed response is merged after another + # tool's response, so it is not `parts[0]`. + (('set_state', 'task_completed'), True), + (('set_state', 'log_event', 'task_completed'), True), + # Parallel calls that do not signal completion. + (('set_state', 'other_tool'), False), + ], +) +@pytest.mark.asyncio +async def test_run_live_task_completion_is_independent_of_response_order( + function_response_names: tuple[str, ...], expect_completion: bool +): + """`task_completed` ends the live agent from any position in the event.""" + + 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, + # so the `task_completed` 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 + ], + ), + ) + + # A follow-up model turn. `task_completed` must end the agent before this is + # processed, so that the next sub-agent of a SequentialAgent can take 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='more')]), + ) + + 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 _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_TASK_COMPLETION_DELAY', + 0, + ), + ): + mock_connect.return_value.__aenter__.return_value = mock.AsyncMock() + + events = [event async for event in flow.run_live(invocation_context)] + + assert events[0] is function_response_event + # The agent stops right after signaling completion, so the follow-up turn is + # only reached when completion was not signaled. + assert (follow_up_event not in events) == expect_completion + + @pytest.mark.asyncio async def test_postprocess_live_yields_grounding_metadata_only(): """Test that _postprocess_live yields LlmResponse with only grounding_metadata."""