chore: import upstream snapshot with attribution

tests/unittests/workflow/test_task_api_e2e.py

@@ -0,0 +1,517 @@
+# 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.
+
+"""End-to-end tests for the Task Delegation API matrix.
+
+Covers the complete cross-product of dispatch-shape × hierarchy-depth so
+the chat-coordinator wrapper, the workflow-node task path, and the
+nested-delegation path are all exercised:
+
+* LlmAgent root → single task sub-agent (basic FC delegation).
+* LlmAgent root → multiple task sub-agents (sequential delegation).
+* LlmAgent root → task sub-agent → nested task sub-agent (chained).
+* Workflow with a task-mode node (no FC delegation).
+* Workflow with a task-mode node that itself has a task sub-agent.
+* Dynamic node case (task agent dispatched via ``ctx.run_node``).
+"""
+
+from __future__ import annotations
+
+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.apps.app import App
+from google.adk.events.event import Event
+from google.adk.workflow import node
+from google.adk.workflow import START
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from tests.unittests import testing_utils
+
+# ---------------------------------------------------------------------------
+# Fixture helpers
+# ---------------------------------------------------------------------------
+
+
+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}
+  )
+
+
+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)
+
+
+def _text_part(text: str) -> types.Part:
+  return types.Part.from_text(text=text)
+
+
+def _make_task_agent(
+    name: str,
+    responses: list,
+    *,
+    sub_agents: list[LlmAgent] | None = None,
+) -> LlmAgent:
+  return LlmAgent(
+      name=name,
+      model=testing_utils.MockModel.create(responses=responses),
+      mode='task',
+      sub_agents=sub_agents or [],
+  )
+
+
+def _collect_finish_outputs(events: list[Event]) -> list[Any]:
+  """Pull out finish_task FC arg dicts in chronological order."""
+  out = []
+  for e in events:
+    for fc in e.get_function_calls():
+      if fc.name == 'finish_task':
+        out.append(dict(fc.args or {}))
+  return out
+
+
+def _get_text_responses(events: list[Event]) -> list[str]:
+  """Concatenate text responses from all model events."""
+  texts = []
+  for e in events:
+    if not e.content or not e.content.parts:
+      continue
+    for p in e.content.parts:
+      if p.text and not p.thought:
+        texts.append(p.text)
+  return texts
+
+
+# ---------------------------------------------------------------------------
+# 1. LlmAgent root → single task sub-agent
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_chat_root_with_single_task_sub_agent(
+    request: pytest.FixtureRequest,
+):
+  """Chat coordinator delegates to one task sub-agent and reports its output."""
+  child = _make_task_agent(
+      name='child',
+      responses=[_finish_part({'result': 'child output'})],
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'do the thing'),
+              'All done: child output.',
+          ]
+      ),
+      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('hi'))
+
+  finish_args = _collect_finish_outputs(events)
+  assert finish_args == [{'result': 'child output'}]
+  assert any(
+      'All done: child output.' in t for t in _get_text_responses(events)
+  )
+
+
+# ---------------------------------------------------------------------------
+# 2. LlmAgent root → multiple task sub-agents (sequential)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_chat_root_with_two_task_sub_agents_sequential(
+    request: pytest.FixtureRequest,
+):
+  """Chat coordinator delegates to two task sub-agents in one turn."""
+  collector = _make_task_agent(
+      name='collector',
+      responses=[_finish_part({'result': 'collected'})],
+  )
+  payer = _make_task_agent(
+      name='payer',
+      responses=[_finish_part({'result': 'paid'})],
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('collector', 'collect'),
+              _delegate_part('payer', 'pay'),
+              'Order placed.',
+          ]
+      ),
+      sub_agents=[collector, payer],
+  )
+
+  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'))
+
+  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))
+
+
+# ---------------------------------------------------------------------------
+# 3. LlmAgent root → task sub-agent → nested task sub-agent
+# ---------------------------------------------------------------------------
+
+
+@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.'
+    ),
+    strict=True,
+)
+@pytest.mark.asyncio
+async def test_chat_root_with_nested_task_delegation(
+    request: pytest.FixtureRequest,
+):
+  """Task agent itself has a task sub-agent and delegates further."""
+  grandchild = _make_task_agent(
+      name='grandchild',
+      responses=[_finish_part({'result': 'leaf'})],
+  )
+
+  child = LlmAgent(
+      name='child',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('grandchild', 'leaf work'),
+              _finish_part({'result': 'middle wraps leaf'}),
+          ]
+      ),
+      mode='task',
+      sub_agents=[grandchild],
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'do the thing'),
+              'Top-level done.',
+          ]
+      ),
+      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('hi'))
+
+  finish_args = _collect_finish_outputs(events)
+  # grandchild fires first (deepest), then child.
+  assert finish_args == [
+      {'result': 'leaf'},
+      {'result': 'middle wraps leaf'},
+  ]
+  assert any('Top-level done.' in t for t in _get_text_responses(events))
+
+
+# ---------------------------------------------------------------------------
+# 4. Workflow with a single task-mode node (no FC delegation)
+# ---------------------------------------------------------------------------
+
+
+class _CaptureNode(BaseNode):
+  """Records its node_input for assertion."""
+
+  received: list[Any] = []
+
+  async def _run_impl(self, *, ctx, node_input):
+    type(self).received.append(node_input)
+    yield Event(output=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')
+
+  wf = Workflow(name='wf', edges=[(START, intake), (intake, capture)])
+  assert wf is not None
+
+
+# ---------------------------------------------------------------------------
+# 6. Dynamic node: function node that dispatches a task agent via ctx.run_node
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_dynamic_dispatch_of_task_agent(
+    request: pytest.FixtureRequest,
+):
+  """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'})],
+  )
+
+  @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}')
+
+  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'))
+
+  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}'
+
+
+# ---------------------------------------------------------------------------
+# 7. Validation error -> retry: wrapper yields the error FR and lets the LLM
+#    emit a corrected finish_task on the next round.
+# ---------------------------------------------------------------------------
+
+
+class _StrictOutput(BaseModel):
+  name: str
+  age: int
+
+
+@pytest.mark.asyncio
+async def test_task_validation_error_drives_retry(
+    request: pytest.FixtureRequest,
+):
+  """Bad finish_task args produce an error FR; the LLM gets a 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}),
+      ]
+  )
+  child = LlmAgent(
+      name='child',
+      model=child_model,
+      mode='task',
+      output_schema=_StrictOutput,
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'gather identity'),
+              'All set.',
+          ]
+      ),
+      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('hi'))
+
+  # The mock LLM was called twice for the child (the bad attempt + the
+  # corrected one), proving the wrapper looped instead of terminating
+  # on the first finish_task.
+  assert child_model.response_index == 1
+  finish_args = _collect_finish_outputs(events)
+  assert finish_args == [
+      {'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.
+  error_frs = [
+      fr.response
+      for e in events
+      for fr in e.get_function_responses()
+      if fr.name == 'finish_task'
+      and isinstance(fr.response, dict)
+      and 'error' in fr.response
+  ]
+  assert len(error_frs) == 1, f'expected one error FR, got {error_frs}'
+
+
+# ---------------------------------------------------------------------------
+# 8. Cross-turn resumption: an unresolved task FC from a prior turn is
+#    re-dispatched by the chat coordinator on the next user turn, before
+#    the LLM is called.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_chat_coordinator_resumes_unresolved_task_fc(
+    request: pytest.FixtureRequest,
+):
+  """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'})]
+  )
+  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.',
+      ]
+  )
+  root = LlmAgent(
+      name='root',
+      model=root_model,
+      sub_agents=[child],
+  )
+
+  # Seed the session with an unresolved task delegation FC authored by
+  # root from a "prior turn".  No matching FR exists.
+  from google.adk.sessions.in_memory_session_service import InMemorySessionService
+
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name=request.function.__name__,
+      user_id='u',
+  )
+  await session_service.append_event(
+      session=session,
+      event=Event(
+          invocation_id='prior-inv',
+          author='root',
+          content=types.Content(
+              role='model',
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          id='fc-pending',
+                          name='child',
+                          args={'request': 'leftover work'},
+                      )
+                  )
+              ],
+          ),
+      ),
+  )
+
+  from google.adk.runners import Runner
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = Runner(app=app, session_service=session_service)
+
+  events = []
+  async for ev in runner.run_async(
+      user_id='u',
+      session_id=session.id,
+      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'
+  finish_args = _collect_finish_outputs(events)
+  assert {
+      'result': 'finished after resume'
+  } in finish_args, f'expected resumed task to finish; got {finish_args}'
+
+
+# ---------------------------------------------------------------------------
+# 9. Strict isolation filtering: a stranger event with a foreign
+#    isolation_scope must NOT appear in the task agent's LLM context.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_strict_isolation_filter_excludes_foreign_scope(
+    request: pytest.FixtureRequest,
+):
+  """Garbage-scoped events are excluded from the task agent's view."""
+  child_model = testing_utils.MockModel.create(
+      responses=[_finish_part({'result': 'ok'})]
+  )
+  child = LlmAgent(name='child', model=child_model, mode='task')
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'do the thing'),
+              'Done.',
+          ]
+      ),
+      sub_agents=[child],
+  )
+
+  from google.adk.sessions.in_memory_session_service import InMemorySessionService
+
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name=request.function.__name__,
+      user_id='u',
+  )
+  # Seed a stranger event with a different scope.
+  stranger = Event(
+      invocation_id='stranger-inv',
+      author='someone_else',
+      content=types.Content(
+          role='user',
+          parts=[types.Part(text='SECRET-SHOULD-NOT-LEAK')],
+      ),
+  )
+  stranger.isolation_scope = 'garbage-scope'
+  session.events.append(stranger)
+
+  from google.adk.runners import Runner
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = Runner(app=app, session_service=session_service)
+
+  async for _ in runner.run_async(
+      user_id='u',
+      session_id=session.id,
+      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 []
+  )
+  assert (
+      'SECRET-SHOULD-NOT-LEAK' not in rendered
+  ), 'stranger event leaked across isolation_scope filter'