fix(workflow): Keep regular-tool FRs on mixed task turns

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:

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 == []

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
 # ---------------------------------------------------------------------------