chat-mode wrapper (_llm_agent_wrapper.py) silently drops regular-tool FR when model emits regular FC + task FC in same turn → session poisoning

[kenwilly]

Environment

• Python 3.13
• google-adk==2.0.0
• Affected models: gemini-3.1-pro-preview (routinely); gemini-2.5-pro (rarely)

Bug description

When an LlmAgent(mode='chat') coordinator emits a model turn containing both a
regular function call (non-task tool, e.g. a state-write tool) and one or more
task-delegation function calls (to LlmAgent(mode='task') sub-agents), the wrapper in
_llm_agent_wrapper.py breaks out of the event generator after handling the task FCs.
The pending function-response event for the regular tool is discarded (never yielded to
the session).

Two effects:

1. The regular tool never executes. Its side-effect is silently lost; the model is
   told the call was "interrupted" on the next turn.
2. Session poisoning. The session now holds a model turn with N function calls and
   fewer than N function responses. Every subsequent request that replays this history is
   rejected by Gemini:
   400 INVALID_ARGUMENT: Please ensure that the number of function response parts is equal to the number of function call parts
   The session cannot recover without external repair.

Reproduction

A self-contained Python repro is below. It uses a hand-authored model response to prove
the bug is in the wrapper, not in any particular model's parallel-calling behaviour.

"""
Minimal repro: ADK 2.0.0 mixed-turn function-call drop.

Run:  cd <repo with google-adk==2.0.0> && python repro_mixed_turn.py
Expected: prints "FAIL: set_todo_list FR absent from session events (bug confirmed)"
"""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from typing import Any

from google.adk.agents import LlmAgent
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.models.registry import LLMRegistry
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools.function_tool import FunctionTool
from google.genai import types

_REGULAR_TOOL_NAME = "set_todo_list"
_TASK_AGENT_NAME = "ga_specialist"
_COORDINATOR_MODEL_ID = "stub-coordinator-mixed-turn"
_TASK_MODEL_ID = "stub-task-specialist"
_CALL_COUNT: list[int] = [0]


class _CoordinatorStubLlm(BaseLlm):
    """Emits a hand-authored model response with a regular FC + task-dispatch FC."""

    @classmethod
    def supported_models(cls) -> list[str]:
        return [_COORDINATOR_MODEL_ID]

    async def generate_content_async(
        self,
        llm_request: LlmRequest,
        stream: bool = False,
    ) -> AsyncIterator[LlmResponse]:
        _CALL_COUNT[0] += 1
        if _CALL_COUNT[0] == 1:
            yield LlmResponse(
                content=types.Content(
                    role="model",
                    parts=[
                        types.Part(
                            function_call=types.FunctionCall(
                                name=_REGULAR_TOOL_NAME,
                                args={"items": ["write upstream bug report"]},
                                id="fc-todo-001",
                            )
                        ),
                        types.Part(
                            function_call=types.FunctionCall(
                                name=_TASK_AGENT_NAME,
                                args={"request": "analyse GA traffic"},
                                id="fc-ga-001",
                            )
                        ),
                    ],
                ),
            )
        else:
            yield LlmResponse(
                content=types.Content(
                    role="model",
                    parts=[types.Part(text="Done. Todo written and GA analysed.")],
                ),
                turn_complete=True,
            )


class _TaskStubLlm(BaseLlm):
    @classmethod
    def supported_models(cls) -> list[str]:
        return [_TASK_MODEL_ID]

    async def generate_content_async(
        self,
        llm_request: LlmRequest,
        stream: bool = False,
    ) -> AsyncIterator[LlmResponse]:
        yield LlmResponse(
            content=types.Content(
                role="model",
                parts=[types.Part(text="GA analysis complete.")],
            ),
            turn_complete=True,
        )


LLMRegistry.register(_CoordinatorStubLlm)
LLMRegistry.register(_TaskStubLlm)


def set_todo_list(items: list[str]) -> dict[str, Any]:
    """Regular function tool — writes the todo list to session state."""
    return {"status": "ok", "items_written": items}


def _build_agents() -> LlmAgent:
    task_specialist = LlmAgent(
        name=_TASK_AGENT_NAME,
        mode="task",
        model=_TASK_MODEL_ID,
        instruction="You are a GA analyst.",
    )
    coordinator = LlmAgent(
        name="coordinator",
        mode="chat",
        model=_COORDINATOR_MODEL_ID,
        instruction=(
            "You coordinate tasks. Use set_todo_list to record work, "
            "then dispatch to ga_specialist."
        ),
        tools=[FunctionTool(set_todo_list)],
        sub_agents=[task_specialist],
    )
    return coordinator


async def _run_repro() -> None:
    svc = InMemorySessionService()
    session = await svc.create_session(app_name="repro", user_id="tester")
    runner = Runner(agent=_build_agents(), app_name="repro", session_service=svc)

    async for _ in runner.run_async(
        user_id=session.user_id,
        session_id=session.id,
        new_message=types.Content(
            role="user",
            parts=[types.Part(text="Write my todo list and analyse GA traffic.")],
        ),
    ):
        pass

    final = await svc.get_session(
        app_name="repro", user_id=session.user_id, session_id=session.id
    )
    events = final.events if final else []

    todo_fr_found = any(
        any(
            p.function_response is not None
            and p.function_response.name == _REGULAR_TOOL_NAME
            for p in (e.content.parts if e.content else [])
        )
        for e in events
    )

    print("\n── Event dump ──")
    for i, e in enumerate(events):
        role = getattr(e.content, "role", "?") if e.content else "?"
        parts_summary = []
        for p in (e.content.parts if e.content else []):
            if p.function_call:
                parts_summary.append(f"FC:{p.function_call.name}")
            elif p.function_response:
                parts_summary.append(f"FR:{p.function_response.name}")
            elif p.text:
                parts_summary.append(f"text:{p.text[:40]!r}")
        print(f"  [{i}] role={role}  parts={parts_summary}")

    print()
    if not todo_fr_found:
        print(
            "FAIL: set_todo_list FR absent from session events (bug confirmed).\n"
            "The regular tool's function-response was dropped by _llm_agent_wrapper.py:388.\n"
            "FC count in history for the mixed turn: 2; FR count: 1 → 400 on next request."
        )
    else:
        print(
            "PASS: set_todo_list FR present in session events.\n"
            "If this is ADK 2.0.0 the bug may have been patched upstream — re-check."
        )


if __name__ == "__main__":
    asyncio.run(_run_repro())

Observed output (run against google-adk==2.0.0, Python 3.13):

── Event dump ──
  [0] role=user  parts=["text:'Write my todo list and analyse GA traffi'"]
  [1] role=model  parts=['FC:set_todo_list', 'FC:ga_specialist']
  [2] role=model  parts=["text:'GA analysis complete.'"]

FAIL: set_todo_list FR absent from session events (bug confirmed).
The regular tool's function-response was dropped by _llm_agent_wrapper.py:388.
FC count in history for the mixed turn: 2; FR count: 1 → 400 on next request.

Event [1] persisted both function calls, but event [2] is the task specialist's
text completion — there is no FR:set_todo_list event anywhere. The regular tool's
function-response was built by flows/llm_flows/functions.py and queued in the
coordinator's event generator, but the wrapper's break closed the generator before it
was ever read. The next turn replays a history with 2 FCs and 1 FR.

Root cause

google/adk/workflow/_llm_agent_wrapper.py, lines 375–388:

while True:
  had_task_fc = False
  transferred = False
  run_method = agent.run_live(ic) if is_live else agent.run_async(ic)
  async with aclosing(run_method) as run_iter:
    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:
        had_task_fc = True
        break  # ← closes run_iter; pending FR events are discarded

The break closes the aclosing() context. The regular tool's FR event is produced by
flows/llm_flows/functions.py (because _defers_response is False for regular tools)
and is sitting in the generator's pending output, but the wrapper never reads it.

_TaskAgentTool._defers_response = True (agent_tool.py) so task FCs correctly skip
the auto-FR build in functions.py. The problem is that the break discards
unprocessed events for any non-deferred tools in the same turn.

The relevant functions.py block (flows/llm_flows/functions.py, lines 579–589):

if (
    tool.is_long_running or tool._defers_response
) and not function_response:
  # The tool either runs long (FR will arrive later via session
  # injection) or defers its response by design (e.g., the LlmAgent
  # wrapper for task delegation synthesizes the FR after the
  # sub-agent completes).  Either way, skip the auto-FR build when
  # the tool returned nothing.
  return None

Regular tools (_defers_response=False) do NOT return early here, so their FR event IS
built and queued — only to be discarded by the wrapper's break.

Expected behaviour

All regular-tool FRs from a mixed-model-turn should be yielded before breaking out of
run_iter. One fix sketch: before the break, drain remaining events from run_iter
until a model event is reached (non-model events after a mixed turn are the FRs for
non-deferred tools in that same turn).

Observed behaviour

The FR for the regular tool is dropped. The session holds an unbalanced FC/FR history that
Gemini permanently rejects.

Related

• #3984 — AgentTool.run_async
  discards inner-stream events (distinct issue: different code path, symptom is billing/
  token gaps rather than session poisoning).
• Declaration hint _TaskAgentTool._get_declaration appends "Do NOT call this tool in
  parallel with any other tools." gemini-3.1-pro-preview routinely ignores this hint,
  making the bug systematic on that model. The fix should be in the wrapper, not in
  model-instruction workarounds.
• A separate but related defect — an orphaned function-response poisoning the session
  from the inverse direction (a persisted FR with no matching FC) — is documented in
  the companion draft at docs/upstream-reports/02-orphaned-function-response-poisoning.md.
  The mixed-turn drop here can produce that orphaned-FR state when the discarded FR
  later lands without its call surviving the branch filter.