diff --git a/src/google/adk/flows/llm_flows/_live_llm_flow.py b/src/google/adk/flows/llm_flows/_live_llm_flow.py index 9bbfafe701..099e0e9e32 100644 --- a/src/google/adk/flows/llm_flows/_live_llm_flow.py +++ b/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( diff --git a/tests/unittests/streaming/test_live_tool_shutdown.py b/tests/unittests/streaming/test_live_tool_shutdown.py index 66102ebea6..21a3ed74a2 100644 --- a/tests/unittests/streaming/test_live_tool_shutdown.py +++ b/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' + )