Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/google/adk/workflow/_llm_agent_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,19 @@ def _extract_task_delegation_fcs(
"""Return task-delegation FCs from this event.

A task-delegation FC is one whose tool is a ``_TaskAgentTool`` instance.

Partial progressive-SSE chunks are ignored. Under
``PROGRESSIVE_SSE_STREAMING``, intermediate chunks that carry a task FC are
marked ``partial=True``; the Runner only persists non-partial events. If we
dispatch and ``break`` on a partial chunk, the non-partial aggregate that
carries the FC is never yielded, so the session keeps a synthesized task FR
with no matching FC. Dispatch must wait for the final aggregate, which also
carries complete streamed arguments.
"""
# Mirror process_llm_agent_output: never act on partial model chunks.
if event.partial:
return []

from ..tools.agent_tool import _TaskAgentTool

return [
Expand Down
32 changes: 32 additions & 0 deletions tests/unittests/workflow/test_llm_agent_as_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -1545,6 +1545,38 @@ async def _gen():
assert drained[1].get_function_responses()[0].name == 'echo'


def test_extract_task_delegation_fcs_skips_partial_events():
"""Progressive-SSE partial chunks must not trigger task dispatch (#6583)."""

def _fc(name: str, call_id: str) -> types.Part:
return types.Part(
function_call=types.FunctionCall(
name=name, args={'request': 'x'}, id=call_id
)
)

task_agent = LlmAgent(name='specialist', mode='task', model='unused')
tools_dict = {'specialist': _TaskAgentTool(task_agent)}
partial = Event(
author='coordinator',
content=types.Content(role='model', parts=[_fc('specialist', '1')]),
partial=True,
)
final = Event(
author='coordinator',
content=types.Content(role='model', parts=[_fc('specialist', '1')]),
partial=False,
)

assert not agent_wrapper._extract_task_delegation_fcs( # pylint: disable=protected-access
partial, tools_dict
)
extracted = agent_wrapper._extract_task_delegation_fcs( # pylint: disable=protected-access
final, tools_dict
)
assert [fc.id for fc in extracted] == ['1']


# --- process_llm_agent_output ---


Expand Down
147 changes: 147 additions & 0 deletions tests/unittests/workflow/test_task_api_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,15 +28,20 @@

from __future__ import annotations

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

from google.adk.agents.context import Context
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.run_config import RunConfig
from google.adk.agents.run_config import StreamingMode
from google.adk.apps.app import App
from google.adk.apps.app import ResumabilityConfig
from google.adk.events.event import Event
from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
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
Expand Down Expand Up @@ -829,3 +834,145 @@ def my_long_run(value: str) -> None:
]
assert len(model_events) == 1
assert "fc-lro-001" in model_events[0].long_running_tool_ids


# ---------------------------------------------------------------------------
# Progressive SSE: do not dispatch task FCs from partial=True chunks (#6583)
# ---------------------------------------------------------------------------


class _ProgressiveSseDispatchCoordinatorLlm(BaseLlm):
"""Emits the progressive-SSE task-FC shape: partial chunk, then aggregate."""

model: str = "progressive_sse_dispatch_stub"
calls: int = 0
specialist_name: str = "specialist"
full_request: str = "Top landing pages for June 2026"

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

async def generate_content_async( # type: ignore[override]
self, llm_request: Any, stream: bool = False
) -> AsyncIterator[LlmResponse]:
del llm_request, stream
self.calls += 1
if self.calls == 1:
# Intermediate progressive-SSE chunk: incomplete args, partial=True.
yield LlmResponse(
content=types.Content(
role="model",
parts=[
types.Part(
function_call=types.FunctionCall(
name=self.specialist_name,
args={"request": ""},
id="fc-dispatch-1",
)
)
],
),
partial=True,
)
# Non-partial aggregate: complete args; Runner persists this event.
yield LlmResponse(
content=types.Content(
role="model",
parts=[
types.Part(
function_call=types.FunctionCall(
name=self.specialist_name,
args={"request": self.full_request},
id="fc-dispatch-1",
)
)
],
),
partial=False,
turn_complete=True,
)
return

yield LlmResponse(
content=types.Content(
role="model",
parts=[
types.Part.from_text(text="Here are your top landing pages.")
],
),
partial=False,
turn_complete=True,
)


@pytest.mark.asyncio
async def test_chat_root_dispatches_task_fc_only_from_non_partial_sse_chunk(
request: pytest.FixtureRequest,
):
"""Task dispatch must wait for the non-partial progressive-SSE aggregate.

Regression for #6583: extracting from ``partial=True`` closes the generator
before the persisted FC event is yielded, leaving an orphaned task FR.
"""
child = _make_task_agent(
name="specialist",
responses=[_finish_part({"result": "1. /pricing 2. /blog 3. /home"})],
)
coordinator_llm = _ProgressiveSseDispatchCoordinatorLlm()
root = LlmAgent(
name="coordinator",
model=coordinator_llm,
mode="chat",
sub_agents=[child],
)

app = App(name=request.function.__name__, root_agent=root)
runner = testing_utils.InMemoryRunner(app=app)
# Partial chunks are only yielded under SSE streaming; without this the
# regression path never reaches the chat wrapper.
run_config = RunConfig(streaming_mode=StreamingMode.SSE)

events = []
async for event in runner.runner.run_async(
user_id=runner.session.user_id,
session_id=runner.session.id,
new_message=testing_utils.get_user_content(
"Top landing pages June 2026?"
),
run_config=run_config,
):
events.append(event)

assert _collect_finish_outputs(events) == [
{"result": "1. /pricing 2. /blog 3. /home"}
]
assert coordinator_llm.calls == 2

persisted = list(runner.session.events)
task_fc_ids = [
fc.id
for e in persisted
for fc in e.get_function_calls()
if fc.name == "specialist"
]
task_fr_ids = [
fr.id
for e in persisted
for fr in e.get_function_responses()
if fr.name == "specialist"
]
assert task_fc_ids == ["fc-dispatch-1"]
assert task_fr_ids == ["fc-dispatch-1"]

task_fc_args = [
dict(fc.args or {})
for e in persisted
for fc in e.get_function_calls()
if fc.name == "specialist"
]
assert task_fc_args == [{"request": "Top landing pages for June 2026"}]
assert any(
"Here are your top landing pages." in t
for t in _get_text_responses(events)
)