From baedd19136a7ad39d229d30cb431e4c2a31280f9 Mon Sep 17 00:00:00 2001 From: Aarav Mittal <137450929+a2105z@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:03:03 -0700 Subject: [PATCH 1/2] fix(workflow): Keep regular-tool FRs on mixed task turns When a chat coordinator emits a regular tool call and a task-delegation call in the same model turn, the wrapper broke out of run_async before draining the regular-tool function response. That left unbalanced FC/FR history and caused Gemini to reject later turns. Fixes #6581 --- src/google/adk/workflow/_llm_agent_wrapper.py | 69 +++++++- .../test_llm_agent_wrapper_mixed_turn.py | 115 +++++++++++++ tests/unittests/workflow/test_task_api_e2e.py | 159 ++++++++++++++++++ 3 files changed, 340 insertions(+), 3 deletions(-) create mode 100644 tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index e0a6c1a2217..a9cd6af9354 100644 --- a/src/google/adk/workflow/_llm_agent_wrapper.py +++ b/src/google/adk/workflow/_llm_agent_wrapper.py @@ -77,6 +77,60 @@ def _extract_task_delegation_fcs( ] +def _event_has_eager_tool_calls( + event: Event, tools_dict: Mapping[str, ToolUnion] +) -> bool: + """True if this event has FCs that produce FR events in the current step. + + Task-delegation tools (``_TaskAgentTool``) and other deferred / long-running + tools do not emit an FR from ``handle_function_calls_async``; the chat + wrapper synthesizes task FRs itself. Regular tools do emit FRs in the same + LLM step, after the model FC event. The wrapper must drain those FR events + before closing the generator, or they are lost and the session history + becomes unbalanced for Gemini. + """ + from ..tools.agent_tool import _TaskAgentTool + + for fc in event.get_function_calls(): + if not fc.name: + continue + tool = tools_dict.get(fc.name) + if tool is None or isinstance(tool, _TaskAgentTool): + continue + if getattr(tool, 'is_long_running', False): + continue + if getattr(tool, '_defers_response', False): + continue + return True + return False + + +async def _drain_pending_tool_response_events( + run_iter: AsyncGenerator[Event, None], +) -> AsyncGenerator[Event, None]: + """Yield remaining non-model events from the current LLM step. + + After a mixed model turn (regular tools + task delegation), the LLM flow + still has pending function-response events. Closing the generator before + reading them drops regular-tool FRs. + + Stops after the first event that carries function responses, or before the + next model-role event (which would start another LLM round without + synthesized task FRs). + """ + async for pending_event in run_iter: + if ( + pending_event.content is not None + and pending_event.content.role == 'model' + ): + # Next LLM round already started; abandon it by stopping iteration. + # Closing the outer generator cancels further work. + return + yield pending_event + if pending_event.get_function_responses(): + return + + def _find_unresolved_task_delegations( session: Session, owner: str, @@ -392,10 +446,19 @@ async def run_llm_agent_as_node( async for event in run_iter: yield event task_fcs = _extract_task_delegation_fcs(event, tools_dict) - for fc in task_fcs: - output = await _dispatch_task_fc(agent, fc, ctx) - yield _synthesize_task_fr_event(fc, output) if task_fcs: + # Mixed turns (regular tool FC + task FC) still have pending + # regular-tool FR events in this generator. Drain them before + # breaking, otherwise aclosing drops them and the session is + # left with unbalanced FC/FR history that Gemini rejects. + if _event_has_eager_tool_calls(event, tools_dict): + async for pending_event in _drain_pending_tool_response_events( + run_iter + ): + yield pending_event + for fc in task_fcs: + output = await _dispatch_task_fc(agent, fc, ctx) + yield _synthesize_task_fr_event(fc, output) had_task_fc = True break # close this run_iter; outer loop re-enters if event.actions.transfer_to_agent: diff --git a/tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py b/tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py new file mode 100644 index 00000000000..88bab780d47 --- /dev/null +++ b/tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py @@ -0,0 +1,115 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for chat-wrapper mixed-turn FR draining helpers. + +Verifies that the wrapper can detect eager (non-deferred) tool calls that +must be drained before breaking out of ``run_async`` on task delegation. +""" + +from __future__ import annotations + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.events.event import Event +from google.adk.tools.agent_tool import _TaskAgentTool +from google.adk.tools.function_tool import FunctionTool +from google.adk.workflow import _llm_agent_wrapper as wrapper +from google.genai import types +import pytest + + +def _echo(value: str) -> dict[str, str]: + """Return the provided value.""" + return {'value': value} + + +def _model_event(*parts: types.Part) -> Event: + return Event( + author='coordinator', + content=types.Content(role='model', parts=list(parts)), + ) + + +def _fc(name: str, call_id: str) -> types.Part: + return types.Part( + function_call=types.FunctionCall(name=name, args={}, id=call_id) + ) + + +def _fr(name: str, call_id: str) -> types.Part: + return types.Part( + function_response=types.FunctionResponse( + name=name, response={'ok': True}, id=call_id + ) + ) + + +def test_event_has_eager_tool_calls_true_for_regular_plus_task(): + """A mixed turn with a FunctionTool and task tool reports eager calls.""" + task_agent = LlmAgent(name='specialist', mode='task', model='unused') + tools_dict = { + 'echo': FunctionTool(_echo), + 'specialist': _TaskAgentTool(task_agent), + } + event = _model_event(_fc('echo', '1'), _fc('specialist', '2')) + + assert wrapper._event_has_eager_tool_calls(event, tools_dict) is True + + +def test_event_has_eager_tool_calls_false_for_task_only(): + """Task-only turns should not drain (no FR is produced by the flow).""" + task_agent = LlmAgent(name='specialist', mode='task', model='unused') + tools_dict = {'specialist': _TaskAgentTool(task_agent)} + event = _model_event(_fc('specialist', '1')) + + assert wrapper._event_has_eager_tool_calls(event, tools_dict) is False + + +@pytest.mark.asyncio +async def test_drain_pending_tool_response_events_yields_fr_then_stops(): + """Drain yields the FR event and stops before a following model event.""" + + async def _gen(): + yield Event( + author='coordinator', + content=types.Content(role='user', parts=[_fr('echo', '1')]), + ) + yield _model_event(types.Part.from_text(text='should not be drained')) + + drained = [ + event + async for event in wrapper._drain_pending_tool_response_events(_gen()) + ] + + assert len(drained) == 1 + assert drained[0].get_function_responses()[0].name == 'echo' + + +@pytest.mark.asyncio +async def test_drain_pending_tool_response_events_stops_on_model_role(): + """Drain stops immediately when the next event is already a model turn.""" + + async def _gen(): + yield _model_event(types.Part.from_text(text='next round')) + yield Event( + author='coordinator', + content=types.Content(role='user', parts=[_fr('echo', '1')]), + ) + + drained = [ + event + async for event in wrapper._drain_pending_tool_response_events(_gen()) + ] + + assert drained == [] diff --git a/tests/unittests/workflow/test_task_api_e2e.py b/tests/unittests/workflow/test_task_api_e2e.py index 87f6dd7fa45..a34f915f257 100644 --- a/tests/unittests/workflow/test_task_api_e2e.py +++ b/tests/unittests/workflow/test_task_api_e2e.py @@ -190,6 +190,165 @@ async def test_chat_root_with_two_task_sub_agents_sequential( assert any('Order placed.' in t for t in _get_text_responses(events)) +# --------------------------------------------------------------------------- +# 2b. Mixed turn: regular tool FC + task FC in the same model response +# --------------------------------------------------------------------------- + + +def _function_call_part( + name: str, args: dict[str, Any], *, call_id: str +) -> types.Part: + """Build a function-call Part with a stable id for FC/FR matching.""" + return types.Part( + function_call=types.FunctionCall(name=name, args=args, id=call_id) + ) + + +def _fr_names(events: list[Event]) -> list[str]: + names: list[str] = [] + for event in events: + for fr in event.get_function_responses(): + if fr.name: + names.append(fr.name) + return names + + +def _fc_names(events: list[Event], *, author: str) -> list[str]: + names: list[str] = [] + for event in events: + if event.author != author: + continue + for fc in event.get_function_calls(): + if fc.name: + names.append(fc.name) + return names + + +@pytest.mark.asyncio +async def test_chat_root_mixed_regular_tool_and_task_keeps_regular_fr( + request: pytest.FixtureRequest, +): + """Regular-tool FR is persisted when emitted with a task FC in one turn. + + Regression for github.com/google/adk-python/issues/6581: the chat wrapper + used to break out of ``run_async`` after dispatching task FCs, dropping the + pending regular-tool FR and poisoning the session for Gemini. + """ + tool_calls: list[list[str]] = [] + + def set_todo_list(items: list[str]) -> dict[str, Any]: + """Record a todo list in session-visible tool output.""" + tool_calls.append(list(items)) + return {'status': 'ok', 'items_written': items} + + child = _make_task_agent( + name='specialist', + responses=[_finish_part({'result': 'specialist done'})], + ) + root = LlmAgent( + name='coordinator', + model=testing_utils.MockModel.create( + responses=[ + [ + _function_call_part( + 'set_todo_list', + {'items': ['write report']}, + call_id='fc-todo-001', + ), + _function_call_part( + 'specialist', + {'request': 'analyse'}, + call_id='fc-task-001', + ), + ], + 'Todos saved and analysis complete.', + ] + ), + tools=[FunctionTool(set_todo_list)], + sub_agents=[child], + ) + + app = App(name=request.function.__name__, root_agent=root) + runner = testing_utils.InMemoryRunner(app=app) + + events = await runner.run_async(testing_utils.get_user_content('go')) + + assert tool_calls == [['write report']] + assert 'set_todo_list' in _fr_names(events) + assert 'specialist' in _fr_names(events) + assert _collect_finish_outputs(events) == [{'result': 'specialist done'}] + assert any( + 'Todos saved and analysis complete.' in t + for t in _get_text_responses(events) + ) + + # Persisted session must keep FC/FR pairs balanced for the mixed turn. + session_events = runner.session.events + assert 'set_todo_list' in _fr_names(session_events) + assert 'specialist' in _fr_names(session_events) + coordinator_fcs = _fc_names(session_events, author='coordinator') + assert coordinator_fcs.count('set_todo_list') == 1 + assert coordinator_fcs.count('specialist') == 1 + + +@pytest.mark.asyncio +async def test_chat_root_mixed_turn_with_two_regular_tools_and_task( + request: pytest.FixtureRequest, +): + """All regular-tool FRs survive when two tools share a turn with a task FC.""" + seen: list[str] = [] + + def note_a(value: str) -> dict[str, str]: + """Record note A.""" + seen.append(f'a:{value}') + return {'note': value} + + def note_b(value: str) -> dict[str, str]: + """Record note B.""" + seen.append(f'b:{value}') + return {'note': value} + + child = _make_task_agent( + name='worker', + responses=[_finish_part({'result': 'worked'})], + ) + root = LlmAgent( + name='coordinator', + model=testing_utils.MockModel.create( + responses=[ + [ + _function_call_part( + 'note_a', {'value': 'one'}, call_id='fc-a' + ), + _function_call_part( + 'note_b', {'value': 'two'}, call_id='fc-b' + ), + _function_call_part( + 'worker', {'request': 'run'}, call_id='fc-w' + ), + ], + 'Combined turn complete.', + ] + ), + tools=[FunctionTool(note_a), FunctionTool(note_b)], + sub_agents=[child], + ) + + app = App(name=request.function.__name__, root_agent=root) + runner = testing_utils.InMemoryRunner(app=app) + + events = await runner.run_async(testing_utils.get_user_content('go')) + + assert sorted(seen) == ['a:one', 'b:two'] + fr_names = _fr_names(events) + assert 'note_a' in fr_names + assert 'note_b' in fr_names + assert 'worker' in fr_names + assert any( + 'Combined turn complete.' in t for t in _get_text_responses(events) + ) + + # --------------------------------------------------------------------------- # 3. LlmAgent root → task sub-agent → nested task sub-agent # --------------------------------------------------------------------------- From 3c802611a288955f32a791c0162dbf9a074e851c Mon Sep 17 00:00:00 2001 From: Aarav Mittal <137450929+a2105z@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:15:33 -0700 Subject: [PATCH 2/2] chore(workflow): Align mixed-turn fix with landed main Copybara already merged #6586 onto main. Keep main's wrapper and tests so the open PR can merge cleanly without reintroducing a divergent patch. --- src/google/adk/workflow/_llm_agent_wrapper.py | 50 ++- .../workflow/test_llm_agent_as_node.py | 297 +++++++++++++- .../test_llm_agent_wrapper_mixed_turn.py | 115 ------ tests/unittests/workflow/test_task_api_e2e.py | 362 +++++++++++------- 4 files changed, 534 insertions(+), 290 deletions(-) delete mode 100644 tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index a9cd6af9354..f1209781ef3 100644 --- a/src/google/adk/workflow/_llm_agent_wrapper.py +++ b/src/google/adk/workflow/_llm_agent_wrapper.py @@ -29,6 +29,7 @@ from ..agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT from ..agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME as _FINISH_TASK_FC_NAME from ..events.event import Event +from ..flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME from ..utils._schema_utils import validate_schema from ..utils.content_utils import to_user_content @@ -84,12 +85,20 @@ def _event_has_eager_tool_calls( Task-delegation tools (``_TaskAgentTool``) and other deferred / long-running tools do not emit an FR from ``handle_function_calls_async``; the chat - wrapper synthesizes task FRs itself. Regular tools do emit FRs in the same - LLM step, after the model FC event. The wrapper must drain those FR events - before closing the generator, or they are lost and the session history - becomes unbalanced for Gemini. + wrapper synthesizes task FRs itself. Regular tools (including long-running or + deferred tools that return a value) do emit FRs in the same LLM step, after + the model FC event. The wrapper must drain those FR events before closing the + generator, or they are lost and the session history becomes unbalanced for + Gemini. + + Args: + event: The event containing function calls. + tools_dict: Map of tool names to Tool objects. + + Returns: + True if the event has eager tool calls. """ - from ..tools.agent_tool import _TaskAgentTool + from ..tools.agent_tool import _TaskAgentTool # pylint: disable=g-import-not-at-top for fc in event.get_function_calls(): if not fc.name: @@ -97,10 +106,6 @@ def _event_has_eager_tool_calls( tool = tools_dict.get(fc.name) if tool is None or isinstance(tool, _TaskAgentTool): continue - if getattr(tool, 'is_long_running', False): - continue - if getattr(tool, '_defers_response', False): - continue return True return False @@ -117,12 +122,29 @@ async def _drain_pending_tool_response_events( Stops after the first event that carries function responses, or before the next model-role event (which would start another LLM round without synthesized task FRs). + + Args: + run_iter: The generator to drain events from. + + Yields: + Events from the current LLM step. """ async for pending_event in run_iter: if ( pending_event.content is not None and pending_event.content.role == 'model' ): + # Tool confirmation events have role 'model' but they are part of the + # current step (asking for confirmation before executing the tool). + # We must yield them and continue draining the actual FR. + is_confirmation = any( + fc.name == REQUEST_CONFIRMATION_FUNCTION_CALL_NAME + for fc in pending_event.get_function_calls() + ) + if is_confirmation: + yield pending_event + continue + # Next LLM round already started; abandon it by stopping iteration. # Closing the outer generator cancels further work. return @@ -452,10 +474,12 @@ async def run_llm_agent_as_node( # breaking, otherwise aclosing drops them and the session is # left with unbalanced FC/FR history that Gemini rejects. if _event_has_eager_tool_calls(event, tools_dict): - async for pending_event in _drain_pending_tool_response_events( - run_iter - ): - yield pending_event + async with aclosing( + _drain_pending_tool_response_events(run_iter) + ) as drain_iter: + async for pending_event in drain_iter: + yield pending_event + for fc in task_fcs: output = await _dispatch_task_fc(agent, fc, ctx) yield _synthesize_task_fr_event(fc, output) diff --git a/tests/unittests/workflow/test_llm_agent_as_node.py b/tests/unittests/workflow/test_llm_agent_as_node.py index 71cf7cee3b5..f0f9ad6b431 100644 --- a/tests/unittests/workflow/test_llm_agent_as_node.py +++ b/tests/unittests/workflow/test_llm_agent_as_node.py @@ -26,10 +26,17 @@ from google.adk.agents.context import Context from google.adk.agents.llm.task._task_models import TaskResult from google.adk.agents.llm_agent import LlmAgent +from google.adk.apps.app import App +from google.adk.apps.app import ResumabilityConfig from google.adk.events.event import Event from google.adk.events.event_actions import EventActions from google.adk.features import FeatureName from google.adk.features import override_feature_enabled +from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME +from google.adk.tools.agent_tool import _TaskAgentTool +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.long_running_tool import LongRunningFunctionTool +from google.adk.workflow import _llm_agent_wrapper as agent_wrapper from google.adk.workflow import START from google.adk.workflow._workflow import Workflow from google.adk.workflow.utils._workflow_graph_utils import build_node @@ -158,8 +165,6 @@ def __exit__(self, *args): def _new_workflow_runner(wf, test_name): """Creates an InMemoryRunner for the new Workflow (root_agent path).""" - from google.adk.apps.app import App - from . import testing_utils app = App(name=test_name, root_agent=wf) @@ -290,8 +295,6 @@ async def test_single_turn_defaults_include_contents_only_when_unset( """Single-turn workflow nodes preserve explicit content inclusion.""" from unittest.mock import MagicMock - from google.adk.workflow import _llm_agent_wrapper - agent = LlmAgent( name='test_agent', model='gemini-2.5-flash', @@ -311,12 +314,12 @@ async def mock_run_async(*args, **kwargs): object.__setattr__(wrapper, 'run_async', mock_run_async) monkeypatch.setattr( - _llm_agent_wrapper, + agent_wrapper, 'prepare_llm_agent_context', lambda agent, ctx: ctx, ) monkeypatch.setattr( - _llm_agent_wrapper, + agent_wrapper, 'prepare_llm_agent_input', lambda agent, ctx, node_input: None, ) @@ -805,7 +808,6 @@ async def test_long_running_tool_interrupts_workflow( request: pytest.FixtureRequest, ): """Long-running tool stops the workflow after one LLM call.""" - from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.workflow._workflow import Workflow as NewWorkflow from . import testing_utils @@ -841,9 +843,6 @@ async def test_resume_after_interrupt_completes_workflow( request: pytest.FixtureRequest, ): """Resuming after interrupt calls the LLM once more to complete.""" - from google.adk.apps.app import App - from google.adk.apps.app import ResumabilityConfig - from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.workflow._workflow import Workflow as NewWorkflow from . import testing_utils @@ -923,9 +922,6 @@ async def test_multiple_sequential_interrupts_in_workflow( request: pytest.FixtureRequest, ): """Two interrupts in sequence each resume and complete in a workflow.""" - from google.adk.apps.app import App - from google.adk.apps.app import ResumabilityConfig - from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.workflow._workflow import Workflow as NewWorkflow from . import testing_utils @@ -1209,9 +1205,6 @@ async def test_three_layer_llm_agent_transfer_round_trip( request: pytest.FixtureRequest, ): """Verify 3-layer LlmAgent transfers end-to-end (Root -> Child -> Grandchild -> Child -> Root).""" - from google.adk.apps.app import App - from google.adk.apps.app import ResumabilityConfig - from . import testing_utils # Prepare the transfer function call parts @@ -1382,3 +1375,275 @@ class InputSchema(BaseModel): with _mock_agent_run(agent_clone, content_text='hi'): with pytest.raises(ValidationError): await runner.run_async('{"wrong_field": "hello"}') + + +# --- Tests for chat-wrapper mixed-turn FR draining helpers --- + + +def _model_event(*parts: types.Part) -> Event: + return Event( + author='coordinator', + content=types.Content(role='model', parts=list(parts)), + ) + + +def test_event_has_eager_tool_calls_true_for_regular_plus_task(): + """A mixed turn with a FunctionTool and task tool reports eager calls.""" + + def _echo(value: str) -> dict[str, str]: + return {'value': value} + + def _fc(name: str, call_id: str) -> types.Part: + return types.Part( + function_call=types.FunctionCall(name=name, args={}, id=call_id) + ) + + task_agent = LlmAgent(name='specialist', mode='task', model='unused') + tools_dict = { + 'echo': FunctionTool(_echo), + 'specialist': _TaskAgentTool(task_agent), + } + event = _model_event(_fc('echo', '1'), _fc('specialist', '2')) + + assert agent_wrapper._event_has_eager_tool_calls(event, tools_dict) # pylint: disable=protected-access + + +def test_event_has_eager_tool_calls_false_for_task_only(): + """Task-only turns should not drain (no FR is produced by the flow).""" + + def _fc(name: str, call_id: str) -> types.Part: + return types.Part( + function_call=types.FunctionCall(name=name, args={}, id=call_id) + ) + + task_agent = LlmAgent(name='specialist', mode='task', model='unused') + tools_dict = {'specialist': _TaskAgentTool(task_agent)} + event = _model_event(_fc('specialist', '1')) + + assert not agent_wrapper._event_has_eager_tool_calls(event, tools_dict) # pylint: disable=protected-access + + +@pytest.mark.asyncio +async def test_drain_pending_tool_response_events_yields_fr_then_stops(): + """Drain yields the FR event and stops before a following model event.""" + + def _fr(name: str, call_id: str) -> types.Part: + return types.Part( + function_response=types.FunctionResponse( + name=name, response={'ok': True}, id=call_id + ) + ) + + async def _gen(): + yield Event( + author='coordinator', + content=types.Content(role='user', parts=[_fr('echo', '1')]), + ) + yield _model_event(types.Part.from_text(text='should not be drained')) + + drained = [ + event + async for event in agent_wrapper._drain_pending_tool_response_events( # pylint: disable=protected-access + _gen() + ) + ] + + assert len(drained) == 1 + assert drained[0].get_function_responses()[0].name == 'echo' + + +@pytest.mark.asyncio +async def test_drain_pending_tool_response_events_stops_on_model_role(): + """Drain stops immediately when the next event is already a model turn.""" + + def _fr(name: str, call_id: str) -> types.Part: + return types.Part( + function_response=types.FunctionResponse( + name=name, response={'ok': True}, id=call_id + ) + ) + + async def _gen(): + yield _model_event(types.Part.from_text(text='next round')) + yield Event( + author='coordinator', + content=types.Content(role='user', parts=[_fr('echo', '1')]), + ) + + drained = [ + event + async for event in agent_wrapper._drain_pending_tool_response_events( # pylint: disable=protected-access + _gen() + ) + ] + + assert not drained + + +def test_event_has_eager_tool_calls_true_for_long_running_tool(): + """A mixed turn with a LongRunningFunctionTool and task tool reports eager calls.""" + + def _long_run(value: str) -> None: + del value + + def _fc(name: str, call_id: str) -> types.Part: + return types.Part( + function_call=types.FunctionCall(name=name, args={}, id=call_id) + ) + + task_agent = LlmAgent(name='specialist', mode='task', model='unused') + tools_dict = { + 'long_run': LongRunningFunctionTool(_long_run), + 'specialist': _TaskAgentTool(task_agent), + } + event = _model_event(_fc('long_run', '1'), _fc('specialist', '2')) + + assert agent_wrapper._event_has_eager_tool_calls(event, tools_dict) # pylint: disable=protected-access + + +@pytest.mark.asyncio +async def test_drain_pending_tool_response_events_yields_confirmation_then_fr(): + """Drain yields confirmation event (role model) AND following FR, then stops.""" + + def _fr(name: str, call_id: str) -> types.Part: + return types.Part( + function_response=types.FunctionResponse( + name=name, response={'ok': True}, id=call_id + ) + ) + + def _confirmation_fc(call_id: str) -> types.Part: + return types.Part( + function_call=types.FunctionCall( + name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, args={}, id=call_id + ) + ) + + async def _gen(): + yield Event( + author='coordinator', + content=types.Content(role='model', parts=[_confirmation_fc('conf-1')]), + ) + yield Event( + author='coordinator', + content=types.Content(role='user', parts=[_fr('echo', '1')]), + ) + yield _model_event(types.Part.from_text(text='should not be drained')) + + drained = [ + event + async for event in agent_wrapper._drain_pending_tool_response_events( # pylint: disable=protected-access + _gen() + ) + ] + + assert len(drained) == 2 + assert ( + drained[0].get_function_calls()[0].name + == REQUEST_CONFIRMATION_FUNCTION_CALL_NAME + ) + assert drained[1].get_function_responses()[0].name == 'echo' + + +# --- process_llm_agent_output --- + + +def _output_model_event(*parts: types.Part, **kwargs: Any) -> Event: + return Event( + invocation_id='inv', + author='test_agent', + content=types.Content(role='model', parts=list(parts)), + **kwargs, + ) + + +def _bare_ctx() -> Context: + """A Context that only needs to carry actions for output processing.""" + from unittest.mock import MagicMock + + ctx = MagicMock(spec=Context) + ctx.actions = EventActions() + return ctx + + +def test_process_llm_agent_output_drops_thought_parts_from_the_output(): + """Thought parts are model reasoning, not part of the node's answer.""" + from google.adk.workflow._llm_agent_wrapper import process_llm_agent_output + + agent = _make_agent(output_key='answer') + ctx = _bare_ctx() + event = _output_model_event( + types.Part(text='thinking out loud', thought=True), + types.Part(text='the '), + types.Part(text='answer'), + ) + + process_llm_agent_output(agent, ctx, event) + + assert event.output == 'the answer' + assert event.node_info.message_as_output is True + assert ctx.actions.state_delta == {'answer': 'the answer'} + + +def test_process_llm_agent_output_skips_events_carrying_function_calls(): + """A tool call is mid-turn work, not the agent's output.""" + from google.adk.workflow._llm_agent_wrapper import process_llm_agent_output + + agent = _make_agent(output_key='answer') + ctx = _bare_ctx() + event = _output_model_event( + types.Part( + function_call=types.FunctionCall(name='some_tool', args={}, id='fc-1') + ) + ) + + process_llm_agent_output(agent, ctx, event) + + assert event.output is None + assert not event.node_info.message_as_output + assert ctx.actions.state_delta == {} + + +def test_process_llm_agent_output_skips_partial_events(): + """Streaming chunks must not each be treated as the finished output.""" + from google.adk.workflow._llm_agent_wrapper import process_llm_agent_output + + agent = _make_agent(output_key='answer') + ctx = _bare_ctx() + event = _output_model_event(types.Part(text='half of an ans'), partial=True) + + process_llm_agent_output(agent, ctx, event) + + assert event.output is None + assert not event.node_info.message_as_output + assert ctx.actions.state_delta == {} + + +def test_process_llm_agent_output_parses_text_against_the_output_schema(): + """With an output_schema the text is parsed, not stored as a raw string.""" + from google.adk.workflow._llm_agent_wrapper import process_llm_agent_output + + agent = _make_agent(output_schema=StoryOutput, output_key='story') + ctx = _bare_ctx() + event = _output_model_event( + types.Part(text='{"title": "T", "content": "C"}'), + ) + + process_llm_agent_output(agent, ctx, event) + + assert event.output == {'title': 'T', 'content': 'C'} + assert ctx.actions.state_delta == {'story': {'title': 'T', 'content': 'C'}} + + +def test_process_llm_agent_output_blank_schema_response_writes_no_state(): + """An empty response cannot satisfy the schema, so nothing is stored.""" + from google.adk.workflow._llm_agent_wrapper import process_llm_agent_output + + agent = _make_agent(output_schema=StoryOutput, output_key='story') + ctx = _bare_ctx() + event = _output_model_event(types.Part(text=' ')) + + process_llm_agent_output(agent, ctx, event) + + assert event.output is None + assert ctx.actions.state_delta == {} diff --git a/tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py b/tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py deleted file mode 100644 index 88bab780d47..00000000000 --- a/tests/unittests/workflow/test_llm_agent_wrapper_mixed_turn.py +++ /dev/null @@ -1,115 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for chat-wrapper mixed-turn FR draining helpers. - -Verifies that the wrapper can detect eager (non-deferred) tool calls that -must be drained before breaking out of ``run_async`` on task delegation. -""" - -from __future__ import annotations - -from google.adk.agents.llm_agent import LlmAgent -from google.adk.events.event import Event -from google.adk.tools.agent_tool import _TaskAgentTool -from google.adk.tools.function_tool import FunctionTool -from google.adk.workflow import _llm_agent_wrapper as wrapper -from google.genai import types -import pytest - - -def _echo(value: str) -> dict[str, str]: - """Return the provided value.""" - return {'value': value} - - -def _model_event(*parts: types.Part) -> Event: - return Event( - author='coordinator', - content=types.Content(role='model', parts=list(parts)), - ) - - -def _fc(name: str, call_id: str) -> types.Part: - return types.Part( - function_call=types.FunctionCall(name=name, args={}, id=call_id) - ) - - -def _fr(name: str, call_id: str) -> types.Part: - return types.Part( - function_response=types.FunctionResponse( - name=name, response={'ok': True}, id=call_id - ) - ) - - -def test_event_has_eager_tool_calls_true_for_regular_plus_task(): - """A mixed turn with a FunctionTool and task tool reports eager calls.""" - task_agent = LlmAgent(name='specialist', mode='task', model='unused') - tools_dict = { - 'echo': FunctionTool(_echo), - 'specialist': _TaskAgentTool(task_agent), - } - event = _model_event(_fc('echo', '1'), _fc('specialist', '2')) - - assert wrapper._event_has_eager_tool_calls(event, tools_dict) is True - - -def test_event_has_eager_tool_calls_false_for_task_only(): - """Task-only turns should not drain (no FR is produced by the flow).""" - task_agent = LlmAgent(name='specialist', mode='task', model='unused') - tools_dict = {'specialist': _TaskAgentTool(task_agent)} - event = _model_event(_fc('specialist', '1')) - - assert wrapper._event_has_eager_tool_calls(event, tools_dict) is False - - -@pytest.mark.asyncio -async def test_drain_pending_tool_response_events_yields_fr_then_stops(): - """Drain yields the FR event and stops before a following model event.""" - - async def _gen(): - yield Event( - author='coordinator', - content=types.Content(role='user', parts=[_fr('echo', '1')]), - ) - yield _model_event(types.Part.from_text(text='should not be drained')) - - drained = [ - event - async for event in wrapper._drain_pending_tool_response_events(_gen()) - ] - - assert len(drained) == 1 - assert drained[0].get_function_responses()[0].name == 'echo' - - -@pytest.mark.asyncio -async def test_drain_pending_tool_response_events_stops_on_model_role(): - """Drain stops immediately when the next event is already a model turn.""" - - async def _gen(): - yield _model_event(types.Part.from_text(text='next round')) - yield Event( - author='coordinator', - content=types.Content(role='user', parts=[_fr('echo', '1')]), - ) - - drained = [ - event - async for event in wrapper._drain_pending_tool_response_events(_gen()) - ] - - assert drained == [] diff --git a/tests/unittests/workflow/test_task_api_e2e.py b/tests/unittests/workflow/test_task_api_e2e.py index a34f915f257..f2f6716d07a 100644 --- a/tests/unittests/workflow/test_task_api_e2e.py +++ b/tests/unittests/workflow/test_task_api_e2e.py @@ -38,6 +38,7 @@ from google.adk.events.event import Event from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.tools.tool_context import ToolContext from google.adk.workflow import node from google.adk.workflow import START @@ -57,13 +58,13 @@ def _delegate_part(target_name: str, request_text: str) -> types.Part: """LLM response calling a task sub-agent (the _TaskAgentTool FC).""" return types.Part.from_function_call( - name=target_name, args={'request': request_text} + name=target_name, args={"request": request_text} ) def _finish_part(args: dict[str, Any]) -> types.Part: """LLM response calling finish_task with the given args.""" - return types.Part.from_function_call(name='finish_task', args=args) + return types.Part.from_function_call(name="finish_task", args=args) def _text_part(text: str) -> types.Part: @@ -72,7 +73,7 @@ def _text_part(text: str) -> types.Part: def _confirmed_task_step(tool_context: ToolContext) -> dict[str, bool]: """Return whether the resumable task step was confirmed.""" - return {'confirmed': tool_context.tool_confirmation.confirmed} + return {"confirmed": tool_context.tool_confirmation.confirmed} def _make_task_agent( @@ -84,7 +85,7 @@ def _make_task_agent( return LlmAgent( name=name, model=testing_utils.MockModel.create(responses=responses), - mode='task', + mode="task", sub_agents=sub_agents or [], ) @@ -94,7 +95,7 @@ def _collect_finish_outputs(events: list[Event]) -> list[Any]: out = [] for e in events: for fc in e.get_function_calls(): - if fc.name == 'finish_task': + if fc.name == "finish_task": out.append(dict(fc.args or {})) return out @@ -122,16 +123,16 @@ async def test_chat_root_with_single_task_sub_agent( ): """Chat coordinator delegates to one task sub-agent and reports its output.""" child = _make_task_agent( - name='child', - responses=[_finish_part({'result': 'child output'})], + name="child", + responses=[_finish_part({"result": "child output"})], ) root = LlmAgent( - name='root', + name="root", model=testing_utils.MockModel.create( responses=[ - _delegate_part('child', 'do the thing'), - 'All done: child output.', + _delegate_part("child", "do the thing"), + "All done: child output.", ] ), sub_agents=[child], @@ -140,12 +141,12 @@ async def test_chat_root_with_single_task_sub_agent( app = App(name=request.function.__name__, root_agent=root) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('hi')) + events = await runner.run_async(testing_utils.get_user_content("hi")) finish_args = _collect_finish_outputs(events) - assert finish_args == [{'result': 'child output'}] + assert finish_args == [{"result": "child output"}] assert any( - 'All done: child output.' in t for t in _get_text_responses(events) + "All done: child output." in t for t in _get_text_responses(events) ) @@ -160,21 +161,21 @@ async def test_chat_root_with_two_task_sub_agents_sequential( ): """Chat coordinator delegates to two task sub-agents in one turn.""" collector = _make_task_agent( - name='collector', - responses=[_finish_part({'result': 'collected'})], + name="collector", + responses=[_finish_part({"result": "collected"})], ) payer = _make_task_agent( - name='payer', - responses=[_finish_part({'result': 'paid'})], + name="payer", + responses=[_finish_part({"result": "paid"})], ) root = LlmAgent( - name='root', + name="root", model=testing_utils.MockModel.create( responses=[ - _delegate_part('collector', 'collect'), - _delegate_part('payer', 'pay'), - 'Order placed.', + _delegate_part("collector", "collect"), + _delegate_part("payer", "pay"), + "Order placed.", ] ), sub_agents=[collector, payer], @@ -183,11 +184,11 @@ async def test_chat_root_with_two_task_sub_agents_sequential( app = App(name=request.function.__name__, root_agent=root) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('place order')) + events = await runner.run_async(testing_utils.get_user_content("place order")) finish_args = _collect_finish_outputs(events) - assert finish_args == [{'result': 'collected'}, {'result': 'paid'}] - assert any('Order placed.' in t for t in _get_text_responses(events)) + assert finish_args == [{"result": "collected"}, {"result": "paid"}] + assert any("Order placed." in t for t in _get_text_responses(events)) # --------------------------------------------------------------------------- @@ -239,29 +240,29 @@ async def test_chat_root_mixed_regular_tool_and_task_keeps_regular_fr( def set_todo_list(items: list[str]) -> dict[str, Any]: """Record a todo list in session-visible tool output.""" tool_calls.append(list(items)) - return {'status': 'ok', 'items_written': items} + return {"status": "ok", "items_written": items} child = _make_task_agent( - name='specialist', - responses=[_finish_part({'result': 'specialist done'})], + name="specialist", + responses=[_finish_part({"result": "specialist done"})], ) root = LlmAgent( - name='coordinator', + name="coordinator", model=testing_utils.MockModel.create( responses=[ [ _function_call_part( - 'set_todo_list', - {'items': ['write report']}, - call_id='fc-todo-001', + "set_todo_list", + {"items": ["write report"]}, + call_id="fc-todo-001", ), _function_call_part( - 'specialist', - {'request': 'analyse'}, - call_id='fc-task-001', + "specialist", + {"request": "analyse"}, + call_id="fc-task-001", ), ], - 'Todos saved and analysis complete.', + "Todos saved and analysis complete.", ] ), tools=[FunctionTool(set_todo_list)], @@ -271,24 +272,24 @@ def set_todo_list(items: list[str]) -> dict[str, Any]: app = App(name=request.function.__name__, root_agent=root) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('go')) + events = await runner.run_async(testing_utils.get_user_content("go")) - assert tool_calls == [['write report']] - assert 'set_todo_list' in _fr_names(events) - assert 'specialist' in _fr_names(events) - assert _collect_finish_outputs(events) == [{'result': 'specialist done'}] + assert tool_calls == [["write report"]] + assert "set_todo_list" in _fr_names(events) + assert "specialist" in _fr_names(events) + assert _collect_finish_outputs(events) == [{"result": "specialist done"}] assert any( - 'Todos saved and analysis complete.' in t + "Todos saved and analysis complete." in t for t in _get_text_responses(events) ) # Persisted session must keep FC/FR pairs balanced for the mixed turn. session_events = runner.session.events - assert 'set_todo_list' in _fr_names(session_events) - assert 'specialist' in _fr_names(session_events) - coordinator_fcs = _fc_names(session_events, author='coordinator') - assert coordinator_fcs.count('set_todo_list') == 1 - assert coordinator_fcs.count('specialist') == 1 + assert "set_todo_list" in _fr_names(session_events) + assert "specialist" in _fr_names(session_events) + coordinator_fcs = _fc_names(session_events, author="coordinator") + assert coordinator_fcs.count("set_todo_list") == 1 + assert coordinator_fcs.count("specialist") == 1 @pytest.mark.asyncio @@ -300,34 +301,34 @@ async def test_chat_root_mixed_turn_with_two_regular_tools_and_task( def note_a(value: str) -> dict[str, str]: """Record note A.""" - seen.append(f'a:{value}') - return {'note': value} + seen.append(f"a:{value}") + return {"note": value} def note_b(value: str) -> dict[str, str]: """Record note B.""" - seen.append(f'b:{value}') - return {'note': value} + seen.append(f"b:{value}") + return {"note": value} child = _make_task_agent( - name='worker', - responses=[_finish_part({'result': 'worked'})], + name="worker", + responses=[_finish_part({"result": "worked"})], ) root = LlmAgent( - name='coordinator', + name="coordinator", model=testing_utils.MockModel.create( responses=[ [ _function_call_part( - 'note_a', {'value': 'one'}, call_id='fc-a' + "note_a", {"value": "one"}, call_id="fc-a" ), _function_call_part( - 'note_b', {'value': 'two'}, call_id='fc-b' + "note_b", {"value": "two"}, call_id="fc-b" ), _function_call_part( - 'worker', {'request': 'run'}, call_id='fc-w' + "worker", {"request": "run"}, call_id="fc-w" ), ], - 'Combined turn complete.', + "Combined turn complete.", ] ), tools=[FunctionTool(note_a), FunctionTool(note_b)], @@ -337,15 +338,15 @@ def note_b(value: str) -> dict[str, str]: app = App(name=request.function.__name__, root_agent=root) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('go')) + events = await runner.run_async(testing_utils.get_user_content("go")) - assert sorted(seen) == ['a:one', 'b:two'] + assert sorted(seen) == ["a:one", "b:two"] fr_names = _fr_names(events) - assert 'note_a' in fr_names - assert 'note_b' in fr_names - assert 'worker' in fr_names + assert "note_a" in fr_names + assert "note_b" in fr_names + assert "worker" in fr_names assert any( - 'Combined turn complete.' in t for t in _get_text_responses(events) + "Combined turn complete." in t for t in _get_text_responses(events) ) @@ -356,9 +357,9 @@ def note_b(value: str) -> dict[str, str]: @pytest.mark.xfail( reason=( - 'Task-mode wrapper does not dispatch task-delegation FCs (only the ' - 'chat-mode wrapper does), so a task-mode middle agent cannot delegate ' - 'to its task sub-agent. Documented limitation.' + "Task-mode wrapper does not dispatch task-delegation FCs (only the " + "chat-mode wrapper does), so a task-mode middle agent cannot delegate " + "to its task sub-agent. Documented limitation." ), strict=True, ) @@ -368,28 +369,28 @@ async def test_chat_root_with_nested_task_delegation( ): """Task agent itself has a task sub-agent and delegates further.""" grandchild = _make_task_agent( - name='grandchild', - responses=[_finish_part({'result': 'leaf'})], + name="grandchild", + responses=[_finish_part({"result": "leaf"})], ) child = LlmAgent( - name='child', + name="child", model=testing_utils.MockModel.create( responses=[ - _delegate_part('grandchild', 'leaf work'), - _finish_part({'result': 'middle wraps leaf'}), + _delegate_part("grandchild", "leaf work"), + _finish_part({"result": "middle wraps leaf"}), ] ), - mode='task', + mode="task", sub_agents=[grandchild], ) root = LlmAgent( - name='root', + name="root", model=testing_utils.MockModel.create( responses=[ - _delegate_part('child', 'do the thing'), - 'Top-level done.', + _delegate_part("child", "do the thing"), + "Top-level done.", ] ), sub_agents=[child], @@ -398,15 +399,15 @@ async def test_chat_root_with_nested_task_delegation( app = App(name=request.function.__name__, root_agent=root) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('hi')) + events = await runner.run_async(testing_utils.get_user_content("hi")) finish_args = _collect_finish_outputs(events) # grandchild fires first (deepest), then child. assert finish_args == [ - {'result': 'leaf'}, - {'result': 'middle wraps leaf'}, + {"result": "leaf"}, + {"result": "middle wraps leaf"}, ] - assert any('Top-level done.' in t for t in _get_text_responses(events)) + assert any("Top-level done." in t for t in _get_text_responses(events)) # --------------------------------------------------------------------------- @@ -427,10 +428,10 @@ async def _run_impl(self, *, ctx, node_input): @pytest.mark.asyncio async def test_workflow_accepts_task_mode_graph_node(): """A mode='task' LlmAgent can be used as a static workflow graph node.""" - intake = _make_task_agent(name='intake', responses=[]) - capture = _CaptureNode(name='capture') + intake = _make_task_agent(name="intake", responses=[]) + capture = _CaptureNode(name="capture") - wf = Workflow(name='wf', edges=[(START, intake), (intake, capture)]) + wf = Workflow(name="wf", edges=[(START, intake), (intake, capture)]) assert wf is not None @@ -445,26 +446,26 @@ async def test_dynamic_dispatch_of_task_agent( ): """A custom function node can dispatch a task agent and consume its output.""" task_agent = _make_task_agent( - name='task_agent', - responses=[_finish_part({'result': 'dynamic output'})], + name="task_agent", + responses=[_finish_part({"result": "dynamic output"})], ) @node(rerun_on_resume=True) async def driver(*, ctx: Context, node_input: Any): - output = await ctx.run_node(task_agent, node_input='go') - yield Event(output=f'wrapped: {output}') + output = await ctx.run_node(task_agent, node_input="go") + yield Event(output=f"wrapped: {output}") - wf = Workflow(name='wf', edges=[(START, driver)]) + wf = Workflow(name="wf", edges=[(START, driver)]) app = App(name=request.function.__name__, root_agent=wf) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('start')) + events = await runner.run_async(testing_utils.get_user_content("start")) outputs = [e.output for e in events if e.output] assert any( - isinstance(o, str) and 'dynamic output' in o for o in outputs - ), f'expected wrapped dynamic output, got: {outputs}' + isinstance(o, str) and "dynamic output" in o for o in outputs + ), f"expected wrapped dynamic output, got: {outputs}" # --------------------------------------------------------------------------- @@ -486,23 +487,23 @@ async def test_task_validation_error_drives_retry( # First finish_task call has wrong types (age as string), second is correct. child_model = testing_utils.MockModel.create( responses=[ - _finish_part({'name': 'Jane', 'age': 'thirty'}), - _finish_part({'name': 'Jane', 'age': 30}), + _finish_part({"name": "Jane", "age": "thirty"}), + _finish_part({"name": "Jane", "age": 30}), ] ) child = LlmAgent( - name='child', + name="child", model=child_model, - mode='task', + mode="task", output_schema=_StrictOutput, ) root = LlmAgent( - name='root', + name="root", model=testing_utils.MockModel.create( responses=[ - _delegate_part('child', 'gather identity'), - 'All set.', + _delegate_part("child", "gather identity"), + "All set.", ] ), sub_agents=[child], @@ -511,7 +512,7 @@ async def test_task_validation_error_drives_retry( app = App(name=request.function.__name__, root_agent=root) runner = testing_utils.InMemoryRunner(app=app) - events = await runner.run_async(testing_utils.get_user_content('hi')) + events = await runner.run_async(testing_utils.get_user_content("hi")) # The mock LLM was called twice for the child (the bad attempt + the # corrected one), proving the wrapper looped instead of terminating @@ -519,8 +520,8 @@ async def test_task_validation_error_drives_retry( assert child_model.response_index == 1 finish_args = _collect_finish_outputs(events) assert finish_args == [ - {'name': 'Jane', 'age': 'thirty'}, - {'name': 'Jane', 'age': 30}, + {"name": "Jane", "age": "thirty"}, + {"name": "Jane", "age": 30}, ] # The validation-error FR should be present in session for the LLM # to see on its retry round. @@ -528,11 +529,11 @@ async def test_task_validation_error_drives_retry( fr.response for e in events for fr in e.get_function_responses() - if fr.name == 'finish_task' + if fr.name == "finish_task" and isinstance(fr.response, dict) - and 'error' in fr.response + and "error" in fr.response ] - assert len(error_frs) == 1, f'expected one error FR, got {error_frs}' + assert len(error_frs) == 1, f"expected one error FR, got {error_frs}" # --------------------------------------------------------------------------- @@ -548,19 +549,19 @@ async def test_chat_coordinator_resumes_unresolved_task_fc( ): """Pending task FC from a prior turn is dispatched before the new LLM call.""" child_model = testing_utils.MockModel.create( - responses=[_finish_part({'result': 'finished after resume'})] + responses=[_finish_part({"result": "finished after resume"})] ) - child = LlmAgent(name='child', model=child_model, mode='task') + child = LlmAgent(name="child", model=child_model, mode="task") root_model = testing_utils.MockModel.create( responses=[ # Only response needed: post-resume continuation after the # pre-LLM scan dispatches the pending task and synthesizes its FR. - 'Resumed and done.', + "Resumed and done.", ] ) root = LlmAgent( - name='root', + name="root", model=root_model, sub_agents=[child], ) @@ -572,21 +573,21 @@ async def test_chat_coordinator_resumes_unresolved_task_fc( session_service = InMemorySessionService() session = await session_service.create_session( app_name=request.function.__name__, - user_id='u', + user_id="u", ) await session_service.append_event( session=session, event=Event( - invocation_id='prior-inv', - author='root', + invocation_id="prior-inv", + author="root", content=types.Content( - role='model', + role="model", parts=[ types.Part( function_call=types.FunctionCall( - id='fc-pending', - name='child', - args={'request': 'leftover work'}, + id="fc-pending", + name="child", + args={"request": "leftover work"}, ) ) ], @@ -601,20 +602,20 @@ async def test_chat_coordinator_resumes_unresolved_task_fc( events = [] async for ev in runner.run_async( - user_id='u', + user_id="u", session_id=session.id, - new_message=testing_utils.get_user_content('continue'), + new_message=testing_utils.get_user_content("continue"), ): events.append(ev) # The child must have been dispatched once (resuming the pending FC). assert ( child_model.response_index == 0 - ), 'child LLM should have been called exactly once for the resumed task' + ), "child LLM should have been called exactly once for the resumed task" finish_args = _collect_finish_outputs(events) assert { - 'result': 'finished after resume' - } in finish_args, f'expected resumed task to finish; got {finish_args}' + "result": "finished after resume" + } in finish_args, f"expected resumed task to finish; got {finish_args}" # --------------------------------------------------------------------------- @@ -633,23 +634,23 @@ async def test_task_sub_agent_resumes_without_parent_delegation_fc( require_confirmation=True, ) child = _make_task_agent( - name='child', + name="child", responses=[ types.Part.from_function_call( name=confirmation_tool.name, args={}, ), - _finish_part({'result': 'confirmed'}), + _finish_part({"result": "confirmed"}), ], ) child.tools.append(confirmation_tool) root = LlmAgent( - name='root', + name="root", model=testing_utils.MockModel.create( responses=[ - _delegate_part('child', 'perform a confirmed step'), - 'Task confirmed.', + _delegate_part("child", "perform a confirmed step"), + "Task confirmed.", ] ), sub_agents=[child], @@ -661,7 +662,7 @@ async def test_task_sub_agent_resumes_without_parent_delegation_fc( ) runner = testing_utils.InMemoryRunner(app=app) - first_events = await runner.run_async(testing_utils.get_user_content('start')) + first_events = await runner.run_async(testing_utils.get_user_content("start")) confirmation_fc = next( fc for event in first_events @@ -680,16 +681,16 @@ async def test_task_sub_agent_resumes_without_parent_delegation_fc( function_response=types.FunctionResponse( id=confirmation_fc.id, name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, - response={'confirmed': True}, + response={"confirmed": True}, ) ) ), invocation_id=invocation_id, ) - assert {'result': 'confirmed'} in _collect_finish_outputs(resumed_events) + assert {"result": "confirmed"} in _collect_finish_outputs(resumed_events) assert any( - 'Task confirmed.' in text for text in _get_text_responses(resumed_events) + "Task confirmed." in text for text in _get_text_responses(resumed_events) ) @@ -705,16 +706,16 @@ async def test_strict_isolation_filter_excludes_foreign_scope( ): """Garbage-scoped events are excluded from the task agent's view.""" child_model = testing_utils.MockModel.create( - responses=[_finish_part({'result': 'ok'})] + responses=[_finish_part({"result": "ok"})] ) - child = LlmAgent(name='child', model=child_model, mode='task') + child = LlmAgent(name="child", model=child_model, mode="task") root = LlmAgent( - name='root', + name="root", model=testing_utils.MockModel.create( responses=[ - _delegate_part('child', 'do the thing'), - 'Done.', + _delegate_part("child", "do the thing"), + "Done.", ] ), sub_agents=[child], @@ -725,18 +726,18 @@ async def test_strict_isolation_filter_excludes_foreign_scope( session_service = InMemorySessionService() session = await session_service.create_session( app_name=request.function.__name__, - user_id='u', + user_id="u", ) # Seed a stranger event with a different scope. stranger = Event( - invocation_id='stranger-inv', - author='someone_else', + invocation_id="stranger-inv", + author="someone_else", content=types.Content( - role='user', - parts=[types.Part(text='SECRET-SHOULD-NOT-LEAK')], + role="user", + parts=[types.Part(text="SECRET-SHOULD-NOT-LEAK")], ), ) - stranger.isolation_scope = 'garbage-scope' + stranger.isolation_scope = "garbage-scope" session.events.append(stranger) from google.adk.runners import Runner @@ -745,17 +746,86 @@ async def test_strict_isolation_filter_excludes_foreign_scope( runner = Runner(app=app, session_service=session_service) async for _ in runner.run_async( - user_id='u', + user_id="u", session_id=session.id, - new_message=testing_utils.get_user_content('go'), + new_message=testing_utils.get_user_content("go"), ): pass # Inspect the child's LLM request: SECRET text must not appear. child_request = child_model.requests[0] - rendered = '\n'.join( - p.text or '' for c in child_request.contents or [] for p in c.parts or [] - ) + parts = [] + for c in child_request.contents or []: + for p in c.parts or []: + parts.append(p.text or "") + rendered = "\n".join(parts) assert ( - 'SECRET-SHOULD-NOT-LEAK' not in rendered - ), 'stranger event leaked across isolation_scope filter' + "SECRET-SHOULD-NOT-LEAK" not in rendered + ), "stranger event leaked across isolation_scope filter" + + +@pytest.mark.asyncio +async def test_chat_root_mixed_turn_with_long_running_tool_and_task_pauses( + request: pytest.FixtureRequest, +): + """Mixed turn with a task FC and a long-running tool (which returns None) pauses.""" + + long_run_called = [] + + def my_long_run(value: str) -> None: + long_run_called.append(value) + return None + + child = _make_task_agent( + name="specialist", + responses=[_finish_part({"result": "specialist done"})], + ) + root = LlmAgent( + name="coordinator", + model=testing_utils.MockModel.create( + responses=[ + [ + _function_call_part( + "my_long_run", + {"value": "hello"}, + call_id="fc-lro-001", + ), + _function_call_part( + "specialist", + {"request": "analyse"}, + call_id="fc-task-001", + ), + ], + "Resume complete.", + ] + ), + tools=[LongRunningFunctionTool(my_long_run)], + sub_agents=[child], + ) + + app = App( + name=request.function.__name__, + root_agent=root, + resumability_config=ResumabilityConfig(is_resumable=True), + ) + runner = testing_utils.InMemoryRunner(app=app) + + events = await runner.run_async(testing_utils.get_user_content("go")) + + assert long_run_called == ["hello"] + assert _collect_finish_outputs(events) == [{"result": "specialist done"}] + + fr_names = _fr_names(events) + assert "specialist" in fr_names + assert "my_long_run" not in fr_names + + assert not any("Resume complete." in t for t in _get_text_responses(events)) + + assert runner.session.events + model_events = [ + e + for e in runner.session.events + if e.author == "coordinator" and e.get_function_calls() + ] + assert len(model_events) == 1 + assert "fc-lro-001" in model_events[0].long_running_tool_ids