diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index e0a6c1a221..a9cd6af935 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 0000000000..88bab780d4 --- /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 87f6dd7fa4..a34f915f25 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 # ---------------------------------------------------------------------------