From 59fcbe6875ae6718522dfa86a4bef6c5764f6749 Mon Sep 17 00:00:00 2001 From: Meher Gajula Date: Tue, 20 Oct 2026 21:05:39 -0500 Subject: [PATCH 1/4] feat(workflow): support bounded parallel agent workers Expose max_parallel_workers on Agent and forward it to the existing worker scheduler. Validate worker-only limits and preserve ordered outputs, cancellation, and resume behavior. Fixes #7641 --- docs/guides/workflow/parallel_worker/index.md | 8 +- src/google/adk/agents/llm_agent.py | 14 ++ .../workflow/utils/_workflow_graph_utils.py | 6 +- .../gepa_root_agent_optimizer_test.py | 1 + .../gepa_root_agent_prompt_optimizer_test.py | 1 + .../workflow/test_llm_agent_as_node.py | 46 ++++ .../workflow/test_workflow_parallel_worker.py | 218 ++++++++++++++++++ 7 files changed, 289 insertions(+), 5 deletions(-) diff --git a/docs/guides/workflow/parallel_worker/index.md b/docs/guides/workflow/parallel_worker/index.md index c87fed45312..428fe8b2a59 100644 --- a/docs/guides/workflow/parallel_worker/index.md +++ b/docs/guides/workflow/parallel_worker/index.md @@ -51,7 +51,7 @@ individual parameters. That is the same mechanism described in ### An agent as the worker -An agent used as a workflow node takes the same flag. Each item arrives as the agent's incoming user message, so the instruction refers to it as the user message rather than through a placeholder. +An agent used as a workflow node takes the same flag. Each item arrives as the agent's incoming user message, so the instruction refers to it as the user message rather than through a placeholder. Set `max_parallel_workers` to bound the number of items in flight, for example to stay within an API's concurrent request limit: ```python from google.adk import Agent @@ -60,6 +60,7 @@ analyzer_agent = Agent( name="analyzer", instruction="Analyze the text in the user message.", parallel_worker=True, + max_parallel_workers=4, ) ``` @@ -90,9 +91,9 @@ When several items fail at once, the worker surfaces the failure belonging to th | `parallel_worker` | `bool` | `False` | Turns on parallel worker mode for this node. | | `max_parallel_workers` | `int \| None` | `None` | Upper bound on items in flight at once. `None` means all of them. | -Set both on the same call, as in `@node(parallel_worker=True, max_parallel_workers=4)`, or `Node(parallel_worker=True, max_parallel_workers=4)` for a `Node` subclass. Setting `max_parallel_workers` without `parallel_worker=True` raises `ValueError`, and so does any value below `1`. +Set both on the same call, as in `@node(parallel_worker=True, max_parallel_workers=4)`, `Node(parallel_worker=True, max_parallel_workers=4)` for a `Node` subclass, or `Agent(..., parallel_worker=True, max_parallel_workers=4)`. Setting `max_parallel_workers` without `parallel_worker=True` raises `ValueError`, and so does any value below `1`. -`LlmAgent` accepts `parallel_worker` but has no `max_parallel_workers` field, so `Agent(..., parallel_worker=True, max_parallel_workers=4)` is rejected by Pydantic as an unexpected input. An agent worker cannot be throttled. `max_concurrency` on the enclosing [`Workflow`](../workflow/index.md) does not help either, because it bounds nodes scheduled by graph edges, and the worker's items are dynamic runs, which are exempt. If you need a bound around an agent, drive the fan-out yourself from a node that calls `ctx.run_node()` under your own semaphore. +The limit applies to items within one parallel worker run. It does not bound all model or tool calls across the workflow. `max_concurrency` on the enclosing [`Workflow`](../workflow/index.md) bounds nodes scheduled by graph edges, and the worker's items are dynamic runs, which are exempt. Set `max_parallel_workers` on the worker itself to bound those items. A parallel worker also has `rerun_on_resume=True` forced on, whatever the node was given. It has to: after an interrupt the node runs again in order to collect the results of the items that had already finished. Those completed items are replayed from history rather than re-executed. @@ -114,7 +115,6 @@ Several items can be waiting at once. Each carries its own interrupt id, so each - **The parameter must be `node_input`.** Any other name is resolved from `ctx.state` and fails at bind time. Naming the parameter after the thing it holds, such as `topic` or `document`, is the usual way a parallel worker goes wrong. - **A list goes in and a list comes out.** If the upstream node produces something that is not a list, it is treated as a list of one item. An empty list gives you an empty list back without anything having run. - **Fail-fast: one failure fails everything.** A single item raising fails the whole worker and cancels the rest of the items. There is no "continue on error" setting that would let you collect partial results, so wrap the risky part of the item's own logic in a `try` block when a failed item should not take the batch down with it. -- **Agent workers cannot be throttled.** `max_parallel_workers` does not exist on `LlmAgent`. ## Related samples diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index ed6dd0a1adb..3005acd6243 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -421,6 +421,12 @@ class LlmAgent(BaseAgent, abc.ABC): parallel_worker: bool | None = None """Whether to run the agent in parallel worker mode.""" + max_parallel_workers: int | None = Field(default=None, ge=1) + """Maximum number of input items to run concurrently in parallel worker mode. + + Requires parallel_worker=True. None means no limit on concurrency. + """ + # LLM-based agent transfer configs - Start disallow_transfer_to_parent: bool = False """Disallows LLM-controlled transferring to the parent agent. @@ -1188,6 +1194,14 @@ def __maybe_accumulate_streaming_output( event.actions.state_delta[self.output_key] = accumulator return accumulator + @model_validator(mode='after') + def _validate_parallel_worker_config(self) -> LlmAgent: + if self.max_parallel_workers is not None and not self.parallel_worker: + raise ValueError( + 'max_parallel_workers can only be set when parallel_worker is True.' + ) + return self + @model_validator(mode='before') @classmethod def _reject_misplaced_generate_content_kwargs(cls, data: Any) -> Any: diff --git a/src/google/adk/workflow/utils/_workflow_graph_utils.py b/src/google/adk/workflow/utils/_workflow_graph_utils.py index ac563c3d822..3e4afbb68ab 100644 --- a/src/google/adk/workflow/utils/_workflow_graph_utils.py +++ b/src/google/adk/workflow/utils/_workflow_graph_utils.py @@ -173,8 +173,12 @@ def build_node( if isinstance(agent, LlmAgent) and agent.parallel_worker: from .._parallel_worker import _ParallelWorker + max_parallel_workers = agent.max_parallel_workers + agent.max_parallel_workers = None agent.parallel_worker = False - return _ParallelWorker(node=agent) + return _ParallelWorker( + node=agent, max_parallel_workers=max_parallel_workers + ) return cast(BaseNode, agent) else: if kwargs: diff --git a/tests/unittests/optimization/gepa_root_agent_optimizer_test.py b/tests/unittests/optimization/gepa_root_agent_optimizer_test.py index bb4a23b6c18..dd6ebf3c58f 100644 --- a/tests/unittests/optimization/gepa_root_agent_optimizer_test.py +++ b/tests/unittests/optimization/gepa_root_agent_optimizer_test.py @@ -160,6 +160,7 @@ def mock_sampler(mocker): @pytest.fixture def mock_agent(mocker): agent = mocker.create_autospec(Agent, instance=True) + agent.max_parallel_workers = None agent.instruction = "Initial instruction" agent.sub_agents = {} agent.clone.return_value = agent diff --git a/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py b/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py index e30021a0546..8dde1508858 100644 --- a/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py +++ b/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py @@ -79,6 +79,7 @@ def mock_sampler(mocker): @pytest.fixture def mock_agent(mocker): agent = mocker.MagicMock(spec=Agent) + agent.max_parallel_workers = None agent.instruction = "Initial instruction" agent.sub_agents = {} agent.mode = None diff --git a/tests/unittests/workflow/test_llm_agent_as_node.py b/tests/unittests/workflow/test_llm_agent_as_node.py index 30b3340a662..6a16eb2403c 100644 --- a/tests/unittests/workflow/test_llm_agent_as_node.py +++ b/tests/unittests/workflow/test_llm_agent_as_node.py @@ -233,6 +233,52 @@ def test_rerun_on_resume_defaults_true(self): assert wrapper.rerun_on_resume is True +@pytest.mark.parametrize('parallel_worker', [None, False]) +def test_agent_worker_limit_requires_parallel_worker(parallel_worker): + """An agent rejects a worker limit unless parallel worker mode is enabled.""" + with pytest.raises( + ValidationError, + match='max_parallel_workers can only be set when parallel_worker is True', + ): + LlmAgent( + name='worker', + parallel_worker=parallel_worker, + max_parallel_workers=2, + ) + + +@pytest.mark.parametrize('max_parallel_workers', [0, -1]) +def test_agent_worker_limit_must_be_positive(max_parallel_workers): + """An agent rejects worker limits below one.""" + with pytest.raises(ValidationError, match='greater than or equal to 1'): + LlmAgent( + name='worker', + parallel_worker=True, + max_parallel_workers=max_parallel_workers, + ) + + +@pytest.mark.parametrize('max_parallel_workers', [None, 1, 4]) +def test_agent_worker_limit_survives_clone_and_serialization( + max_parallel_workers, +): + """Cloned and restored agents retain their parallel worker configuration.""" + agent = LlmAgent( + name='worker', + parallel_worker=True, + max_parallel_workers=max_parallel_workers, + ) + + cloned = agent.clone() + restored = LlmAgent.model_validate_json(agent.model_dump_json()) + built = build_node(agent) + + for worker in (cloned, restored, agent): + assert worker.parallel_worker is True + assert worker.max_parallel_workers == max_parallel_workers + assert built.max_parallel_workers == max_parallel_workers + + @pytest.mark.asyncio async def test_single_turn_input_event_inherits_branch_and_scope( request: pytest.FixtureRequest, diff --git a/tests/unittests/workflow/test_workflow_parallel_worker.py b/tests/unittests/workflow/test_workflow_parallel_worker.py index f685a6c86cf..e70c60e5579 100644 --- a/tests/unittests/workflow/test_workflow_parallel_worker.py +++ b/tests/unittests/workflow/test_workflow_parallel_worker.py @@ -19,11 +19,15 @@ 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.apps.app import ResumabilityConfig from google.adk.events._interrupts import extract_event_interrupt_ids from google.adk.events.event import Event from google.adk.events.request_input import RequestInput +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.workflow import BaseNode from google.adk.workflow import START from google.adk.workflow._node import node @@ -47,6 +51,220 @@ _MAX_WAIT_S = 5.0 +@pytest.mark.parametrize('max_parallel_workers', [None, 1, 2]) +async def test_agent_parallel_worker_limits_concurrency_and_preserves_outputs( + max_parallel_workers: int | None, +): + """Agent workers obey their limit and return every result in input order. + + Setup: Five agent inputs wait on separate gates inside their model calls. + Act: Release the newest active item first, allowing another to start. + Assert: In-flight calls never exceed the limit, and all ordered outputs reach + the next node. Without an explicit limit, all five calls start together. + """ + items = ['item1', 'item2', 'item3', 'item4', 'item5'] + limit = max_parallel_workers or len(items) + gates = {item: asyncio.Event() for item in items} + started: asyncio.Queue[str] = asyncio.Queue() + completed: list[str] = [] + active = 0 + peak_active = 0 + + class GatedModel(testing_utils.MockModel): + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + nonlocal active, peak_active + item = llm_request.contents[-1].parts[0].text + active += 1 + peak_active = max(peak_active, active) + started.put_nowait(item) + try: + assert active <= limit + await gates[item].wait() + completed.append(item) + yield LlmResponse( + content=types.Content( + role='model', parts=[types.Part(text=f'{item}_processed')] + ) + ) + finally: + active -= 1 + + kwargs = ( + {'max_parallel_workers': max_parallel_workers} + if max_parallel_workers is not None + else {} + ) + worker = LlmAgent( + name='worker', + model=GatedModel(responses=[]), + parallel_worker=True, + **kwargs, + ) + captured: list[list[str]] = [] + + def produce() -> list[str]: + return items + + def capture(node_input: list[str]) -> list[str]: + captured.append(node_input) + return node_input + + workflow = Workflow( + name='agent_worker', edges=[(START, produce, worker, capture)] + ) + runner = testing_utils.InMemoryRunner(node=workflow) + run_task = asyncio.create_task(runner.run_async('start')) + + try: + in_flight = [ + await asyncio.wait_for(started.get(), timeout=_MAX_WAIT_S) + for _ in range(limit) + ] + remaining = len(items) - limit + while in_flight: + gates[in_flight.pop()].set() + if remaining: + in_flight.append( + await asyncio.wait_for(started.get(), timeout=_MAX_WAIT_S) + ) + remaining -= 1 + await asyncio.wait_for(run_task, timeout=_MAX_WAIT_S) + finally: + run_task.cancel() + await asyncio.gather(run_task, return_exceptions=True) + + assert peak_active == limit + assert active == 0 + assert sorted(completed) == items + if limit > 1: + assert completed != items + assert captured == [[f'{item}_processed' for item in items]] + assert worker.parallel_worker is True + assert worker.max_parallel_workers == max_parallel_workers + + +async def test_agent_parallel_worker_failure_cancels_active_sibling(): + """A failed agent cancels its sibling without starting queued inputs.""" + started: list[str] = [] + cancelled: list[str] = [] + sibling_waiting = asyncio.Event() + never_finish = asyncio.Event() + + class FailingModel(testing_utils.MockModel): + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + item = llm_request.contents[-1].parts[0].text + started.append(item) + if item == 'fail': + await asyncio.wait_for(sibling_waiting.wait(), timeout=_MAX_WAIT_S) + raise RuntimeError('worker failed') + try: + sibling_waiting.set() + await never_finish.wait() + except asyncio.CancelledError: + cancelled.append(item) + raise + yield LlmResponse() + + worker = LlmAgent( + name='worker', + model=FailingModel(responses=[]), + parallel_worker=True, + max_parallel_workers=2, + ) + + def produce() -> list[str]: + return ['fail', 'blocked', 'queued'] + + workflow = Workflow(name='agent_worker', edges=[(START, produce, worker)]) + runner = testing_utils.InMemoryRunner(node=workflow) + + with pytest.raises(RuntimeError, match='worker failed'): + await asyncio.wait_for(runner.run_async('start'), timeout=_MAX_WAIT_S) + + assert started == ['fail', 'blocked'] + assert cancelled == ['blocked'] + + +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_agent_parallel_worker_resumes_without_repeating_finished_items( + is_resumable: bool, +): + """A bounded agent worker resumes its queue without repeating finished work. + + Setup: Three items run with a limit of one. The second requests approval. + Act: Run until interrupted, then answer the pending approval. + Assert: The first result is replayed, the second resumes, and the queued + third item runs exactly once before all outputs reach the next node. + """ + + def approve() -> None: + """Request approval to continue.""" + + model = testing_utils.MockModel.create( + responses=[ + 'first_done', + types.Part.from_function_call(name='approve', args={}), + 'second_done', + 'third_done', + ] + ) + worker = LlmAgent( + name='worker', + model=model, + tools=[LongRunningFunctionTool(approve)], + parallel_worker=True, + max_parallel_workers=1, + ) + captured: list[list[str]] = [] + + def produce() -> list[str]: + return ['first', 'second', 'third'] + + def capture(node_input: list[str]) -> list[str]: + captured.append(node_input) + return node_input + + workflow = Workflow( + name='agent_worker', edges=[(START, produce, worker, capture)] + ) + app = App( + name='agent_worker_resume', + root_agent=workflow, + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + runner = testing_utils.InMemoryRunner(app=app) + + events = await runner.run_async('start') + assert len(model.requests) == 2 + assert captured == [] + interrupt = next(event for event in events if event.long_running_tool_ids) + call_id = next(iter(interrupt.long_running_tool_ids)) + await runner.run_async( + new_message=types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + name='approve', + id=call_id, + response={'result': 'approved'}, + ) + ) + ], + ), + invocation_id=interrupt.invocation_id, + ) + + assert len(model.requests) == 4 + assert captured == [['first_done', 'second_done', 'third_done']] + + class _ProducerNode(BaseNode): """A node that produces a list of items.""" From c0cd43e548e15f8138d2c328ee0843614190e126 Mon Sep 17 00:00:00 2001 From: Meher Gajula Date: Tue, 20 Oct 2026 22:01:12 -0500 Subject: [PATCH 2/4] feat: add owned parallel task collaboration and accepted-result cancellation --- .../workflows/parallel_tasks/README.md | 62 ++ .../samples/workflows/parallel_tasks/agent.py | 73 ++ docs/guides/README.md | 1 + docs/guides/workflow/parallel_tasks/index.md | 192 ++++ docs/guides/workflow/parallel_worker/index.md | 6 + src/google/adk/agents/context.py | 42 +- src/google/adk/agents/invocation_context.py | 3 + .../adk/flows/llm_flows/base_llm_flow.py | 2 + .../core/_function_call_postprocessor.py | 9 + .../flows/llm_flows/tools/_batch_executor.py | 7 + src/google/adk/workflow/__init__.py | 6 + .../adk/workflow/_parallel_task_context.py | 351 +++++++ src/google/adk/workflow/_parallel_tasks.py | 783 ++++++++++++++ src/google/adk/workflow/_parallel_worker.py | 119 +-- .../adk/workflow/utils/_rehydration_utils.py | 53 + .../adk/workflow/utils/_replay_manager.py | 31 +- .../workflow/test_parallel_task_context.py | 893 ++++++++++++++++ .../workflow/test_parallel_task_replay.py | 249 +++++ .../unittests/workflow/test_parallel_tasks.py | 994 ++++++++++++++++++ .../workflow/utils/test_replay_manager.py | 1 + 20 files changed, 3755 insertions(+), 122 deletions(-) create mode 100644 contributing/samples/workflows/parallel_tasks/README.md create mode 100644 contributing/samples/workflows/parallel_tasks/agent.py create mode 100644 docs/guides/workflow/parallel_tasks/index.md create mode 100644 src/google/adk/workflow/_parallel_task_context.py create mode 100644 src/google/adk/workflow/_parallel_tasks.py create mode 100644 tests/unittests/workflow/test_parallel_task_context.py create mode 100644 tests/unittests/workflow/test_parallel_task_replay.py create mode 100644 tests/unittests/workflow/test_parallel_tasks.py diff --git a/contributing/samples/workflows/parallel_tasks/README.md b/contributing/samples/workflows/parallel_tasks/README.md new file mode 100644 index 00000000000..6e8f1115784 --- /dev/null +++ b/contributing/samples/workflows/parallel_tasks/README.md @@ -0,0 +1,62 @@ +# Parallel tasks with a verified winner + +## Overview + +Three agents calculate the same sum using different instructions. The coordinator +accepts the first correct answer, cancels remaining work, and explicitly merges +only that answer into session state. A deterministic check rejects incorrect +answers even if they arrive first. + +Arithmetic keeps the acceptance rule easy to inspect. A production application +should use its own verifier; asking a model to declare its own answer correct is +not independent verification. Computing a sum directly is cheaper than calling +three models when the sum itself is the only objective. + +## Sample Inputs + +- `17, 25, -4, 11` + +- `100, -70, 20, -5` + +## Graph + +```mermaid +graph TD + START --> choose_sum + choose_sum --> left_to_right + choose_sum --> pairwise + choose_sum --> signed_groups +``` + +## How To + +Run the sample with your configured ADK model: + +```sh +uv run adk run contributing/samples/workflows/parallel_tasks +``` + +A function marked `@node(rerun_on_resume=True)` owns concurrent children: + +```python +async with ctx.parallel_tasks( + name="sum_candidates", max_parallel_workers=3 +) as tasks: + for solver in solvers: + tasks.spawn(solver, node_input) + winner = await tasks.first(accept=is_correct) + await tasks.merge_result(winner, keys=["sum"]) +``` + +`first` waits for an acceptable result and drains cancellation before returning. +The agents return proposals without writing an `output_key`; `merge_result` +checks that the selected state key has not changed since the group opened. + +Cancellation stops cooperative local execution. Requests already sent to a +provider can still incur cost, and completed tool side effects are not undone. + +## Related Guides + +- [Parallel tasks](../../../../docs/guides/workflow/parallel_tasks/index.md) - Handles, messaging, progress, cancellation, and explicit result merging. +- [Parallel worker mode](../../../../docs/guides/workflow/parallel_worker/index.md) - Ordered fan-out over a list using the same task scheduler. +- [Dynamic nodes](../../../../docs/guides/workflow/dynamic_nodes/index.md) - Scheduling child nodes inside a coordinator. diff --git a/contributing/samples/workflows/parallel_tasks/agent.py b/contributing/samples/workflows/parallel_tasks/agent.py new file mode 100644 index 00000000000..4986e308517 --- /dev/null +++ b/contributing/samples/workflows/parallel_tasks/agent.py @@ -0,0 +1,73 @@ +# 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 +# +# https://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. + +"""Accept a verified answer from three agents and stop the remaining work.""" + +from collections.abc import AsyncGenerator +from typing import Any + +from google.adk import Agent +from google.adk import Context +from google.adk import Event +from google.adk import Workflow +from google.adk.workflow import node +from pydantic import BaseModel + + +class SumProposal(BaseModel): + sum: int + + +solvers = [ + Agent( + name=name, + instruction=f"Sum the integers in the user message. {strategy}", + output_schema=SumProposal, + ) + for name, strategy in ( + ("left_to_right", "Add them from left to right."), + ("pairwise", "Group them into pairs before adding."), + ("signed_groups", "Add positive and negative groups separately."), + ) +] + + +@node(rerun_on_resume=True) +async def choose_sum( + ctx: Context, node_input: str +) -> AsyncGenerator[Event | str, None]: + numbers = [int(value.strip()) for value in node_input.split(",")] + + def is_correct(output: Any) -> bool: + return ( + isinstance(output, dict) + and type(output.get("sum")) is int + and output["sum"] == sum(numbers) + ) + + async with ctx.parallel_tasks( + name="sum_candidates", max_parallel_workers=3 + ) as tasks: + for solver in solvers: + tasks.spawn(solver, node_input) + yield Event(message="Checking three independent sums.") + winner = await tasks.first(accept=is_correct) + await tasks.merge_result(winner, keys=["sum"]) + yield f"Verified sum: {winner.output['sum']} (from {winner.name})." + + +root_agent = Workflow( + name="root_agent", + edges=[("START", choose_sum)], +) diff --git a/docs/guides/README.md b/docs/guides/README.md index 15b44eb7ecb..92ef6482fd8 100644 --- a/docs/guides/README.md +++ b/docs/guides/README.md @@ -147,4 +147,5 @@ This directory contains specific developer guides for the ADK Python implementat * [JoinNode](workflow/join_node/index.md) - Synchronizing parallel execution paths in workflows. * [RetryConfig](workflow/retry_config/index.md) - Configuring retry policies for resilient workflow nodes. * [ParallelWorker](workflow/parallel_worker/index.md) - Processing lists of items concurrently in workflows. +* [ParallelTasks](workflow/parallel_tasks/index.md) - Owning concurrent agent runs with handles, messages, selective cancellation, and accepted results. * [Dynamic Nodes](workflow/dynamic_nodes/index.md) - Scheduling and executing nodes dynamically at runtime. diff --git a/docs/guides/workflow/parallel_tasks/index.md b/docs/guides/workflow/parallel_tasks/index.md new file mode 100644 index 00000000000..7859a5ea2b8 --- /dev/null +++ b/docs/guides/workflow/parallel_tasks/index.md @@ -0,0 +1,192 @@ +# Parallel tasks + +## Introduction + +`Context.parallel_tasks()` gives a workflow coordinator handles to concurrent +agent or function runs. The coordinator can keep working, inspect results, +message a worker, cancel a selected task, or accept the first satisfactory answer +and stop the remaining work. Children use ADK's existing node scheduler, event +stream, session, and resume machinery. + +Use [parallel worker mode](../parallel_worker/index.md) when every item needs the +same node and all outputs should be collected in input order. It uses the same +task group internally. Use explicit handles when the coordinator needs to decide +which work to keep. + +## Get started + +```python +from google.adk import Agent, Context, Workflow +from google.adk.workflow import node + +researchers = [ + Agent(name="researcher_a", instruction="Investigate the question."), + Agent(name="researcher_b", instruction="Independently check the question."), + Agent(name="researcher_c", instruction="Look for counterexamples."), +] + +@node(rerun_on_resume=True) +async def research(ctx: Context, node_input: str) -> list[str]: + async with ctx.parallel_tasks(name="research", max_parallel_workers=3) as tasks: + handles = [tasks.spawn(agent, node_input) for agent in researchers] + # Each spawn returns immediately; all three can execute concurrently. + handles[2].send_message("Pay particular attention to edge cases.") + return await tasks.gather() + +root_agent = Workflow(name="research_workflow", edges=[("START", research)]) +``` + +`ParallelTask` and `ParallelTasks` are also exported from `google.adk.workflow` +for type annotations. Enter the group as an async context manager. Exiting +normally gathers pending work; an exception or parent cancellation cancels and +drains it. + +## How it works + +`spawn(node, node_input)` reserves a deterministic child node path and returns a +handle without waiting for a concurrency slot. `max_parallel_workers` bounds +active children; extra handles stay queued. Children have separate conversation +branches, while their state remains ordinary shared session state. + +The handle exposes `id`, `path`, `name`, `status`, `output`, and `error`. Status is +one of `queued`, `running`, `completed`, `failed`, `waiting`, or `cancelled`. +`waiting` means a child needs human input. `await handle.result()` propagates the +child's failure or interruption. Cancelling an observer awaiting `result()` does +not cancel the worker; explicitly cancel the worker when that is intended. +IDs identify children within an invocation; use native session and invocation +IDs as well when correlating across runs. Groups opened inside a function tool +also include that tool call's identity, so repeated or concurrent calls remain +independent. + +| Operation | Behavior | +| --- | --- | +| `await tasks.wait(mode="any")` | Returns completed and pending handles in spawn order once any handle is terminal. | +| `await tasks.wait(mode="all")` | Waits for all handles, or until human input prevents further admission and running siblings finish. | +| `await tasks.gather()` | Returns outputs in spawn order; a failure stops admission and cancels siblings. | +| `await tasks.first(accept=predicate)` | Chooses the first acceptable output, cancels siblings, and waits for cleanup. | +| `handle.cancel()` | Requests cancellation without waiting for it to finish. | +| `await tasks.cancel(handle)` | Persists deliberate cancellation and waits for that owned worker to stop. | +| `handle.send_message(text)` | Queues a message and returns a delivery receipt. | +| `await tasks.merge_result(handle, keys=[...])` | Persists selected keys from a completed mapping proposal after conflict and schema checks. | + +`wait` reports the current sets; it does not consume completed handles. Inspect +`status` to distinguish success, failure, cancellation, and human input. + +The coordinator can also call `await ctx.run_node(...)` inside the scope, for +example to ask a planning agent to refine instructions while the workers run. + +### First acceptable result + +```python +async with ctx.parallel_tasks(name="candidates", max_parallel_workers=3) as tasks: + for candidate in candidates: + tasks.spawn(candidate, question) + winner = await tasks.first(accept=verify_answer) + return winner.output +``` + +The predicate receives a candidate's output and can be synchronous or +asynchronous. Rejected answers and failed workers let alternatives continue. +No new queued worker is admitted while a completed candidate is being checked. +The winning decision is recorded before siblings are cancelled. Results observed +in the same completion batch are checked in spawn order. + +If all successful results are rejected, `first` raises `ValueError`; if no answer +is accepted and a worker failed, it propagates a worker failure. Human input +propagates when no remaining active worker can provide an accepted answer. +An exception from the predicate also cancels the group. + +### Progress and messages + +Worker events still flow through the normal Runner event stream. Their +`custom_metadata["parallel_task"]` identifies the owning `task_id` and the task +status at emission. Nested child events keep that correlation. Handle status is +the source of truth after completion; progress events describe when they were +emitted. + +```python +receipt = worker.send_message("Include a source for the numerical estimate.") +# Later, while this group remains open: +delivery = worker.message_receipt(receipt["id"]) +``` + +Messages reach the addressed agent at its next model request boundary, in FIFO +order. A nested agent cannot consume its parent's mailbox. A receipt changes +from `queued` to `delivered` when the message has been persisted and inserted into +the prepared model request. It does not acknowledge provider receipt or promise +that the agent followed the instruction. If there is no later model request, +the message becomes `undelivered` when the task ends. Sending to a finished task +raises `ValueError`. + +Each task accepts at most 128 messages over its lifetime, with a 64 KiB UTF-8 +limit per message. Queued messages and receipts live in memory; delivered +content remains in native session history. A process restart does not restore +the mailbox. A function or workflow without a model request at its root does +not consume messages addressed to that handle. + +Event callbacks may record progress or call the nonblocking `send_message` and +`handle.cancel` methods. Do not await worker results or group operations inside +a Runner event callback: the worker may itself be waiting for that event to be +acknowledged. + +### Explicit result merging + +Have speculative workers return mapping proposals without writing shared state: + +```python +winner = await tasks.first(accept=verify_answer) +await tasks.merge_result(winner, keys=["answer", "sources"]) +``` + +The group snapshots state when it opens. Every selected key must still match +that snapshot before any key is written. Nested values are compared, and an +absent key differs from one holding `None`. An optional `expected` mapping can +supply a different expected value for selected keys. State schema validation +also completes before writes. The state delta and a record of the merged keys +are persisted in the same native event. Repeating a merge of the same task and +keys is a no-op, including when the coordinator resumes; a later external update +is not overwritten by reapplying an old proposal. + +This protects the explicit merge from overwriting concurrent changes. It does +not isolate or roll back writes that workers make directly to shared state, or +undo tool effects. Keep speculative work free of irreversible side effects. + +## Configuration + +| Parameter | Default | Meaning | +| --- | --- | --- | +| `name` | `"parallel"` | Stable group name; must not contain `/` or `@`. | +| `max_parallel_workers` | `None` | Positive active-child limit, or unlimited. | +| `cancellation_timeout` | `5.0` | Positive finite seconds to wait for cooperative cancellation. | +| `spawn(..., run_id=...)` | Generated | Optional nonnumeric child run ID, unique within the coordinator execution and containing neither `/` nor `@`. | + +## Advanced behavior and limitations + +- The coordinating node must use `rerun_on_resume=True`. Keep group names, + creation order, node names, and spawn order stable when resuming. Explicit + cancellations and winner selection are replayed from session events. A parent + abort does not persist cleanup as a deliberate decision to discard children. +- Groups own their workers for the duration of the coordinating node scope. + They do not expose a background task service across arbitrary model turns or + invocations. +- Cancellation is cooperative Python async cancellation. It interrupts pending + awaits and closes cooperative generators; it cannot preempt blocking Python + code, undo side effects, or guarantee a remote provider stops an already sent + request. Completed or in-flight requests can still be billed. +- If an explicit cancellation or winner selection cannot drain a worker before + the timeout, it raises `TimeoutError` instead of reporting successful cleanup. + While unwinding another error or parent cancellation, it logs the timeout and + preserves that original error. A worker that suppresses cancellation may + still be running. Scope exit seals its mailbox and marks remaining messages + undelivered without pretending the worker stopped. The enclosing Workflow has + its own task cleanup, which can continue waiting for uncooperative work; the + group timeout does not guarantee a bounded shutdown of the entire invocation. +- Concurrency helps overlap I/O. It does not make CPU-bound Python work parallel + or establish a universal model latency or token savings claim. + +## Related samples and guides + +- [Verified winner sample](../../../../contributing/samples/workflows/parallel_tasks/README.md) - Three agents, independent validation, sibling cancellation, and selected state merging. +- [Parallel worker mode](../parallel_worker/index.md) - Ordered list fan-out. +- [Dynamic nodes](../dynamic_nodes/index.md) - Native child runs and resume. +- [Runner execution cancellation](../../runners/runner/abort.md) - Cancelling the enclosing invocation. diff --git a/docs/guides/workflow/parallel_worker/index.md b/docs/guides/workflow/parallel_worker/index.md index 428fe8b2a59..eff2e05568b 100644 --- a/docs/guides/workflow/parallel_worker/index.md +++ b/docs/guides/workflow/parallel_worker/index.md @@ -12,6 +12,12 @@ A node with a list to work through, whether that is documents to analyze, querie There is no `ParallelWorker` class to import. You turn the mode on with the `parallel_worker=True` flag on `@node(...)`, on `Node(...)`, or on an `Agent(...)` being used as a node, and the node then runs once per element of its input list rather than once on the list as a whole. +For a coordinator that needs individual task handles, messages, selective +cancellation, or the first acceptable answer, use +[parallel tasks](../parallel_tasks/index.md). Both forms use the same owned task +group and native node scheduler; parallel worker mode retains its ordered list +results and bounded active work. + ## Get started The one thing to get right is the parameter name. Each list element is handed to the worker as `node_input`, so the function's parameter **must** be called `node_input`. Any other name is looked up in `ctx.state` instead, is not found there, and every run fails with `ValueError: Missing value for parameter ""` before the function body executes. diff --git a/src/google/adk/agents/context.py b/src/google/adk/agents/context.py index 8078eb99ff2..090dfa2ac31 100644 --- a/src/google/adk/agents/context.py +++ b/src/google/adk/agents/context.py @@ -47,6 +47,7 @@ from ..workflow._dynamic_node_scheduler import DynamicNodeScheduler from ..workflow._graph import NodeLike from ..workflow._graph import RouteValue + from ..workflow._parallel_tasks import ParallelTasks from .invocation_context import InvocationContext _MAX_PARENT_DEPTH = 50 @@ -220,6 +221,8 @@ def __init__( self._route_value: RouteValue | list[RouteValue] | None = None self._route_emitted: bool = False self._interrupt_ids: set[str] = set() + self._parallel_group_counters: dict[str, int] = {} + self._parallel_task_paths: set[str] = set() # scope tag inherited from parent ctx or invocation_context by default; # NodeRunner / Workflow may override before the node runs. if parent_ctx is not None: @@ -431,6 +434,35 @@ def get_invocation_context(self) -> InvocationContext: ctx_with_proxy = ctx.model_copy(update=update) return ctx_with_proxy + def parallel_tasks( + self, + *, + name: str = 'parallel', + max_parallel_workers: int | None = None, + cancellation_timeout: float = 5.0, + ) -> ParallelTasks: + """Create an owned group of concurrent workflow node runs. + + Use ``async with ctx.parallel_tasks(...) as tasks`` and ``tasks.spawn`` + to obtain stable handles. The group bounds running workers, exposes their + results, and cancels remaining work when the scope exits abnormally. + Worker state remains shared session state; merge returned proposals + explicitly when only selected results should update the parent's state. + + Args: + name: Stable group name, reused in the same order when resuming. + max_parallel_workers: Maximum active children, or None for no limit. + cancellation_timeout: Seconds to wait for cancelled children to stop. + """ + from ..workflow._parallel_tasks import ParallelTasks + + return ParallelTasks( + self, + name=name, + max_parallel_workers=max_parallel_workers, + cancellation_timeout=cancellation_timeout, + ) + async def run_node( self, node: NodeLike, @@ -450,10 +482,8 @@ async def run_node( asynchronously wait for its result. The dynamically executed node becomes a child run of the current node in the workflow. - IMPORTANT: Always ``await`` this method directly. Wrapping it in - ``asyncio.create_task()`` means the task runs unsupervised — errors - are silently swallowed and the task is not cancelled if the parent - node is interrupted (e.g. via HITL). + Await this method directly. For concurrent children, use + ``parallel_tasks()`` so the scope supervises errors and cancellation. Args: node: The node to be executed. This can be a BaseNode instance or a @@ -469,8 +499,8 @@ async def run_node( run_id: An optional custom run ID for the dynamic node execution. If not provided, a default run ID is generated. Useful for correlating events across runs. - use_sub_branch: If True, the dynamic node will be executed in a sub-branch - to isolate its state and events from the main branch. + use_sub_branch: If True, execute in a separate conversation branch. + Session state remains shared. override_branch: An optional branch to use instead of parent's branch. override_isolation_scope: An optional isolation scope to use instead of the parent's scope. diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 338f1e23874..fb9cb018c52 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -51,6 +51,7 @@ from ..tools.base_tool import BaseTool from ..utils._agent_mode import AgentMode from ..workflow._base_node import BaseNode +from ..workflow._parallel_task_context import _tag_parallel_task_event from .base_agent import _agent_state_key from .base_agent import BaseAgent from .base_agent import BaseAgentState @@ -496,6 +497,8 @@ async def _enqueue_event(self, event: Event) -> None: "InvocationContext." ) + event = _tag_parallel_task_event(event) + if event.partial: # Partial events: SSE streaming only, no session append, no blocking. await self._event_queue.put((event, None)) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 365a7cb9ace..3a5100e6f29 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -37,6 +37,7 @@ from ...models.llm_request import LlmRequest from ...models.llm_response import LlmResponse from ...utils.context_utils import Aclosing +from ...workflow._parallel_task_context import _deliver_parallel_task_messages from .core._finalizer import finalize_model_response_event from .core._finalizer import handle_after_model_callback from .core._finalizer import handle_before_model_callback @@ -558,6 +559,7 @@ async def _handle_before_model_callback( llm_request: LlmRequest, model_response_event: Event, ) -> Optional[LlmResponse]: + await _deliver_parallel_task_messages(invocation_context, llm_request) return await handle_before_model_callback( invocation_context, llm_request, model_response_event ) diff --git a/src/google/adk/flows/llm_flows/core/_function_call_postprocessor.py b/src/google/adk/flows/llm_flows/core/_function_call_postprocessor.py index 86000f21dc7..4123e7e4493 100644 --- a/src/google/adk/flows/llm_flows/core/_function_call_postprocessor.py +++ b/src/google/adk/flows/llm_flows/core/_function_call_postprocessor.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio from typing import AsyncGenerator from .. import functions @@ -25,6 +26,7 @@ from ....models.llm_request import LlmRequest from ....utils.context_utils import Aclosing from ....workflow._errors import NodeInterruptedError +from ....workflow._parallel_task_context import _persist_parallel_task_cancelled_tools from ..prompt import _schema as _output_schema_processor from ..tools._batch_executor import _pop_kept_function_response_event from ._utils import require_agent as _require_agent @@ -80,6 +82,13 @@ async def postprocess_handle_function_calls_async( function_response_event = await functions.handle_function_calls_async( invocation_context, function_call_event, llm_request.tools_dict ) + except asyncio.CancelledError as error: + await _persist_parallel_task_cancelled_tools( + invocation_context, + function_call_event, + _pop_kept_function_response_event(error), + ) + raise except (Exception, NodeInterruptedError) as e: kept_event = _pop_kept_function_response_event(e) if kept_event is not None and not invocation_context.is_aborted: diff --git a/src/google/adk/flows/llm_flows/tools/_batch_executor.py b/src/google/adk/flows/llm_flows/tools/_batch_executor.py index 6225bacc845..45ff80da773 100644 --- a/src/google/adk/flows/llm_flows/tools/_batch_executor.py +++ b/src/google/adk/flows/llm_flows/tools/_batch_executor.py @@ -35,6 +35,7 @@ from ....tools.base_tool import BaseTool from ....tools.tool_confirmation import ToolConfirmation from ....workflow._errors import NodeInterruptedError +from ....workflow._parallel_task_context import _parallel_task_cancellation_requested from ..core._utils import as_llm_agent as _as_llm_agent from ._caller import _execute_single_prepared_call_async from ._caller import _prepare_single @@ -272,6 +273,12 @@ async def _execute_prepared_function_calls( # Wait for all tasks to complete try: maybe_function_response_events = await _gather_or_cancel(tasks) + except asyncio.CancelledError as error: + if _parallel_task_cancellation_requested(): + kept_event = _keep_completed_responses(invocation_context, tasks) + if kept_event is not None: + setattr(error, _KEPT_FUNCTION_RESPONSE_EVENT_ATTR, kept_event) + raise except (Exception, NodeInterruptedError) as e: # A call raised and its unfinished siblings were cancelled. Calls that had # already finished may have had side effects, so keep their real results diff --git a/src/google/adk/workflow/__init__.py b/src/google/adk/workflow/__init__.py index b18156f281b..55ae924d206 100644 --- a/src/google/adk/workflow/__init__.py +++ b/src/google/adk/workflow/__init__.py @@ -28,6 +28,8 @@ from ._join_node import JoinNode from ._node import Node from ._node import node + from ._parallel_tasks import ParallelTask + from ._parallel_tasks import ParallelTasks from ._retry_config import RetryConfig from ._workflow import Workflow @@ -39,6 +41,8 @@ 'JoinNode': '._join_node', 'Node': '._node', 'NodeTimeoutError': '._errors', + 'ParallelTask': '._parallel_tasks', + 'ParallelTasks': '._parallel_tasks', 'RetryConfig': '._retry_config', 'START': '._base_node', 'Workflow': '._workflow', @@ -52,6 +56,8 @@ 'JoinNode', 'Node', 'NodeTimeoutError', + 'ParallelTask', + 'ParallelTasks', 'RetryConfig', 'START', 'Workflow', diff --git a/src/google/adk/workflow/_parallel_task_context.py b/src/google/adk/workflow/_parallel_task_context.py new file mode 100644 index 00000000000..71ba5600787 --- /dev/null +++ b/src/google/adk/workflow/_parallel_task_context.py @@ -0,0 +1,351 @@ +# 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 +# +# https://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. + +"""Execution-local messaging and event correlation for parallel task groups. + +The owning group controls task lifetimes. This module does not schedule workers. +A ContextVar carries the ownership lineage into native child execution and is +reset when that execution leaves. +""" + +from __future__ import annotations + +import asyncio +from collections import deque +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING + +from google.genai import types + +from ..events._internal_metadata import INTERNAL_METADATA_PREFIX +from ..platform import uuid as platform_uuid + +if TYPE_CHECKING: + from ..agents.invocation_context import InvocationContext + from ..events.event import Event + from ..models.llm_request import LlmRequest + + +class _ParallelTaskControl: + """One logical task's bounded mailbox and delivery receipts. + + Queued messages and receipts live only as long as their owning task group. + Delivered message content is persisted through the native event stream and + remains part of the addressed node's conversation history. A receipt means + delivery to the model request, not provider acknowledgement or execution of + an instruction. + """ + + def __init__( + self, + *, + task_id: str, + root_node_path: str, + max_messages: int = 128, + max_message_bytes: int = 65_536, + ) -> None: + for name, value in ( + ('max_messages', max_messages), + ('max_message_bytes', max_message_bytes), + ): + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f'{name} must be a positive integer') + if not task_id or not root_node_path: + raise ValueError('Task ID and root node path must be nonempty') + self.task_id = task_id + self.root_node_path = root_node_path + self.status = 'queued' + self.cancellation_requested = False + self._event_delivery_closed = asyncio.Event() + self._max_messages = max_messages + self._max_message_bytes = max_message_bytes + self._messages: deque[tuple[str, str]] = deque() + self._receipts: dict[str, dict[str, str]] = {} + self._finished = False + + def send_message(self, content: str) -> dict[str, str]: + """Queue content for this task's next top-level model boundary. + + Receipt retention counts toward the limit even after delivery. This bounds + the control's lifetime memory rather than only its current queue length. + """ + if ( + self._finished + or not self.event_delivery_open + or self.status not in ('queued', 'running') + ): + raise ValueError('Task is no longer accepting messages') + if not isinstance(content, str) or not content.strip(): + raise ValueError('Message content must be a nonempty string') + if len(content.encode('utf-8')) > self._max_message_bytes: + raise ValueError('Message content exceeds the task message byte limit') + if len(self._receipts) >= self._max_messages: + raise ValueError('Task message limit reached') + message_id = platform_uuid.new_uuid() + receipt = { + 'id': message_id, + 'task_id': self.task_id, + 'status': 'queued', + } + self._messages.append((message_id, content)) + self._receipts[message_id] = receipt + return dict(receipt) + + def message_receipt(self, message_id: str) -> dict[str, str]: + """Return a detached delivery receipt for an accepted message.""" + try: + return dict(self._receipts[message_id]) + except KeyError as error: + raise ValueError('Unknown message receipt for this task') from error + + @property + def receipts(self) -> list[dict[str, str]]: + """Return detached receipts in acceptance order.""" + return [dict(receipt) for receipt in self._receipts.values()] + + def finish(self, status: str) -> None: + """Close admission and mark messages without another boundary undelivered.""" + if self._finished: + return + self._finished = True + self.status = status + while self._messages: + message_id, _ = self._messages.popleft() + self._receipts[message_id]['status'] = 'undelivered' + + @property + def event_delivery_open(self) -> bool: + """Whether the owning scope can still await native event persistence.""" + return not self._event_delivery_closed.is_set() + + def close_event_delivery(self) -> None: + """Seal the mailbox and stop awaiting events after the owning scope exits. + + Closing delivery does not claim that a cancellation-resistant worker has + stopped. Its execution status remains unchanged until it actually finishes. + """ + self._event_delivery_closed.set() + while self._messages: + message_id, _ = self._messages.popleft() + self._receipts[message_id]['status'] = 'undelivered' + + +_current_parallel_task: ContextVar[_ParallelTaskControl | None] = ContextVar( + 'adk_parallel_task_control', default=None +) +_parallel_task_controls: ContextVar[tuple[_ParallelTaskControl, ...]] = ( + ContextVar('adk_parallel_task_controls', default=()) +) + + +@contextmanager +def _bind_parallel_task(control: _ParallelTaskControl) -> Iterator[None]: + """Bind an owned control while the existing scheduler runs its child.""" + token = _current_parallel_task.set(control) + lineage_token = _parallel_task_controls.set( + (*_parallel_task_controls.get(), control) + ) + try: + yield + finally: + _parallel_task_controls.reset(lineage_token) + _current_parallel_task.reset(token) + + +def _tag_parallel_task_event(event: Event) -> Event: + """Correlate native progress without mutating a node's cached event.""" + control = _current_parallel_task.get() + if control is None: + return event + return event.model_copy( + update={ + 'custom_metadata': { + **(event.custom_metadata or {}), + 'parallel_task': { + 'task_id': control.task_id, + 'status': control.status, + }, + INTERNAL_METADATA_PREFIX + + 'parallel_task_owners': [ + owner.task_id for owner in _parallel_task_controls.get() + ], + } + } + ) + + +def _parallel_task_cancellation_requested() -> bool: + """Whether this task or an ancestor deliberately requested cancellation.""" + return any( + control.cancellation_requested + for control in _parallel_task_controls.get() + ) + + +async def _persist_parallel_task_cancelled_tools( + invocation_context: InvocationContext, + function_call_event: Event, + kept_event: Event | None, +) -> None: + """Pair outstanding calls when deliberate task cancellation stops a batch. + + Completed siblings keep their real responses and state deltas. A cancellation + response describes local cancellation, not whether external side effects + occurred. No synthetic response is persisted for a call already answered. + """ + if not _parallel_task_cancellation_requested(): + return + controls = _parallel_task_controls.get() + if not controls or any( + not control.event_delivery_open for control in controls + ): + return + + from ..events._rewind_events import _apply_rewinds + from ..events.event import Event + from ..events.event import NodeInfo + from ..events.event_actions import EventActions + + call_path = function_call_event.node_info.path or invocation_context.node_path + answered = { + response.id + for event in _apply_rewinds(invocation_context.session.events) + if event.invocation_id == invocation_context.invocation_id + and event.node_info.path == call_path + and event.branch == function_call_event.branch + and event.author == function_call_event.author + and event.isolation_scope == function_call_event.isolation_scope + for response in event.get_function_responses() + } + kept = ( + { + part.function_response.id: part + for part in kept_event.content.parts or [] + if part.function_response + } + if kept_event is not None and kept_event.content is not None + else {} + ) + parts = [] + included_kept = False + for call in function_call_event.get_function_calls(): + if call.id in answered: + continue + if call.id in kept: + parts.append(kept[call.id].model_copy(deep=True)) + included_kept = True + else: + parts.append( + types.Part( + function_response=types.FunctionResponse( + id=call.id, + name=call.name, + response={ + 'status': 'cancelled', + 'error': ( + 'The owning parallel task cancelled this local call. ' + 'External side effects may already have occurred.' + ), + }, + ) + ) + ) + if not parts: + return + event = Event( + author=function_call_event.author, + invocation_id=invocation_context.invocation_id, + branch=function_call_event.branch, + isolation_scope=function_call_event.isolation_scope, + node_info=NodeInfo(path=call_path), + content=types.Content(role='user', parts=parts), + actions=kept_event.actions.model_copy(deep=True) + if kept_event is not None and included_kept + else EventActions(), + ) + delivery = asyncio.create_task(invocation_context._enqueue_event(event)) + closures = [ + asyncio.create_task(control._event_delivery_closed.wait()) + for control in controls + ] + try: + done, _ = await asyncio.wait( + [delivery, *closures], return_when=asyncio.FIRST_COMPLETED + ) + if delivery in done: + await delivery + finally: + for waiter in (delivery, *closures): + if not waiter.done(): + waiter.cancel() + await asyncio.gather(delivery, *closures, return_exceptions=True) + + +async def _deliver_parallel_task_messages( + invocation_context: InvocationContext, + llm_request: LlmRequest, +) -> None: + """Persist queued content and add it to the addressed model request. + + Only the task's root node may consume its inbox. Nested agents inherit the + execution binding for progress correlation, but cannot steal messages meant + for the root. Queued messages remain queued if persistence is interrupted. + """ + control = _current_parallel_task.get() + if ( + control is None + or control._finished + or invocation_context.node_path != control.root_node_path + or not control._messages + ): + return + + from ..events.event import Event + from ..events.event import NodeInfo + + messages = list(control._messages) + content = types.Content( + role='user', + parts=[ + types.Part.from_text( + text=f'Message from the coordinating task:\n{message}' + ) + for _, message in messages + ], + ) + event = Event( + author='user', + invocation_id=invocation_context.invocation_id, + branch=invocation_context.branch, + isolation_scope=invocation_context.isolation_scope, + node_info=NodeInfo(path=control.root_node_path), + content=content.model_copy(deep=True), + custom_metadata={ + 'parallel_task_message': { + 'task_id': control.task_id, + 'message_ids': [message_id for message_id, _ in messages], + } + }, + ) + # This is the worker producer, not an on_event callback. The native queue + # waits until the Runner has persisted the event before the worker continues. + await invocation_context._enqueue_event(event) + if control._finished: + return + llm_request.contents.append(content) + for message_id, _ in messages: + control._messages.popleft() + control._receipts[message_id]['status'] = 'delivered' diff --git a/src/google/adk/workflow/_parallel_tasks.py b/src/google/adk/workflow/_parallel_tasks.py new file mode 100644 index 00000000000..a551550d02e --- /dev/null +++ b/src/google/adk/workflow/_parallel_tasks.py @@ -0,0 +1,783 @@ +# 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 +# +# https://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. + +"""Owned parallel node runs using the workflow's existing scheduler.""" + +from __future__ import annotations + +import asyncio +import base64 +from collections.abc import Callable +from collections.abc import Iterable +from collections.abc import Mapping +import copy +import inspect +import logging +import math +from typing import Any +from typing import Literal +from typing import TYPE_CHECKING + +from ..events._internal_metadata import INTERNAL_METADATA_PREFIX +from ..events._node_path_builder import _NodePathBuilder +from ..events._rewind_events import _apply_rewinds +from ..events.event import Event +from ..events.event import NodeInfo +from ..events.event_actions import EventActions +from ._errors import DynamicNodeFailError +from ._errors import NodeInterruptedError +from ._parallel_task_context import _bind_parallel_task +from ._parallel_task_context import _ParallelTaskControl +from .utils._workflow_graph_utils import build_node + +if TYPE_CHECKING: + from ..agents.context import Context + from ._base_node import BaseNode + from ._graph import NodeLike + +logger = logging.getLogger('google_adk.' + __name__) +_CONTROL_KEY = INTERNAL_METADATA_PREFIX + 'parallel_tasks' +_MISSING = object() + + +def _same_value(left: Any, right: Any) -> bool: + """Compare state values without treating booleans as integers.""" + if type(left) is not type(right): + return False + if isinstance(left, dict): + return left.keys() == right.keys() and all( + _same_value(value, right[key]) for key, value in left.items() + ) + if isinstance(left, (list, tuple)): + return len(left) == len(right) and all( + _same_value(a, b) for a, b in zip(left, right) + ) + return bool(left == right) + + +class ParallelTask: + """A stable handle to one node run owned by a ``ParallelTasks`` scope. + + Cancelling an await of ``result()`` only stops that observer. Use ``cancel`` + to request cancellation of the work, or await ``group.cancel(task)`` to also + wait for it to stop. Messages are delivered at the next LLM request boundary. + """ + + def __init__( + self, + group: ParallelTasks, + node: BaseNode, + node_input: Any, + run_id: str, + index: int, + ) -> None: + self._group = group + self._node = node + self._input = node_input + self._run_id = run_id + self._index = index + self._path = str( + _NodePathBuilder.from_string(group._ctx.node_path).append( + node.name, run_id + ) + ) + self._status = 'queued' + self._output: Any = None + self._error: BaseException | None = None + self._runner: asyncio.Task[None] | None = None + self._done: asyncio.Future[None] = ( + asyncio.get_running_loop().create_future() + ) + self._cancel_requested = False + self._cancellation_sent = False + self._cancel_deadline: float | None = None + self._interrupt_ids: set[str] = set() + self._control = _ParallelTaskControl( + task_id=self.id, root_node_path=self.path + ) + + @property + def id(self) -> str: + """Stable node path, including the run ID reserved at spawn time.""" + return self._path + + @property + def path(self) -> str: + """Node path used by native workflow events.""" + return self._path + + @property + def name(self) -> str: + """Name of the scheduled node.""" + return self._node.name + + @property + def status(self) -> str: + """Queued, running, completed, failed, waiting, or cancelled.""" + return self._status + + @property + def output(self) -> Any: + """Completed output, or None when no output is available.""" + return self._output + + @property + def error(self) -> BaseException | None: + """Execution failure, if any.""" + return self._error + + async def result(self) -> Any: + """Wait for this run without transferring cancellation ownership.""" + await asyncio.shield(self._done) + if self.status == 'cancelled': + raise asyncio.CancelledError() + if self._error is not None: + raise self._error + return self._output + + def cancel(self) -> bool: + """Request deliberate cancellation; returns whether it was requested.""" + if self.status in {'completed', 'failed', 'cancelled'}: + return False + if self._group._closed: + raise RuntimeError('The parallel task scope is closed.') + self._cancel_requested = True + self._group._wake.set() + return True + + def send_message(self, content: str) -> dict[str, str]: + """Queue a message for the worker's next LLM request.""" + return self._control.send_message(content) + + def message_receipt(self, message_id: str) -> dict[str, str]: + """Return the delivery status of a queued message.""" + return self._control.message_receipt(message_id) + + @property + def message_receipts(self) -> list[dict[str, str]]: + """Detached snapshots of this worker's message receipts.""" + return self._control.receipts + + def _finish(self, status: str, error: BaseException | None = None) -> None: + self._status = status + self._error = error + self._control.finish(status) + if not self._done.done(): + self._done.set_result(None) + self._group._changed.set() + self._group._wake.set() + + +class ParallelTasks: + """An owned, bounded group of native workflow node runs. + + Use through ``async with ctx.parallel_tasks(...) as tasks``. Worker state + remains normal shared session state. Return mapping proposals and explicitly + call ``merge_result`` when only selected results should update parent state. + Names and spawn order must be stable when a workflow resumes. + """ + + def __init__( + self, + ctx: Context, + *, + name: str = 'parallel', + max_parallel_workers: int | None = None, + cancellation_timeout: float = 5.0, + _legacy_run_ids: bool = False, + ) -> None: + if not name or any(char in name for char in '/@'): + raise ValueError( + 'Group name must be nonempty and contain neither / nor @.' + ) + if max_parallel_workers is not None and ( + isinstance(max_parallel_workers, bool) + or not isinstance(max_parallel_workers, int) + or max_parallel_workers < 1 + ): + raise ValueError( + 'max_parallel_workers must be a positive integer or None.' + ) + if not math.isfinite(cancellation_timeout) or cancellation_timeout <= 0: + raise ValueError('cancellation_timeout must be positive and finite.') + self._ctx = ctx + self._name = name + self._limit = max_parallel_workers + self._timeout = cancellation_timeout + self._legacy_run_ids = _legacy_run_ids + ordinal = ctx._parallel_group_counters.get(name, 0) + 1 + ctx._parallel_group_counters[name] = ordinal + self._prefix = f'{name}-{ordinal}' + self._call_scope: str | None = None + group_run_id = str(ordinal) + if not _legacy_run_ids and ctx.function_call_id: + self._call_scope = ( + base64.urlsafe_b64encode(ctx.function_call_id.encode('utf-8')) + .decode('ascii') + .rstrip('=') + ) + self._prefix += f'-call-{self._call_scope}' + group_run_id += f'-call-{self._call_scope}' + self._id = str( + _NodePathBuilder.from_string(ctx.node_path).append(name, group_run_id) + ) + self._tasks: list[ParallelTask] = [] + self._next_index = 0 + self._wake = asyncio.Event() + self._changed = asyncio.Event() + self._decision_lock = asyncio.Lock() + self._pump_task: asyncio.Task[None] | None = None + self._pump_error: BaseException | None = None + self._entered = False + self._closed = False + self._paused = False + self._fail_fast = False + self._first_mode = False + self._selecting = False + self._decision_pending = False + self._cancelled_ids: set[str] = set() + self._winner_id: str | None = None + self._winner: ParallelTask | None = None + self._state_snapshot: dict[str, Any] = {} + self._merged: set[tuple[str, str]] = set() + + async def __aenter__(self) -> ParallelTasks: + if self._entered: + raise RuntimeError('A parallel task scope can only be entered once.') + if not self._ctx._node_rerun_on_resume: + raise ValueError('Parallel tasks require rerun_on_resume=True.') + self._entered = True + if not self._legacy_run_ids: + self._state_snapshot = copy.deepcopy(self._ctx.state.to_dict()) + for event in _apply_rewinds(self._ctx.session.events): + if event.invocation_id != self._ctx._invocation_context.invocation_id: + continue + record = (event.custom_metadata or {}).get(_CONTROL_KEY) + if not isinstance(record, dict) or record.get('group') != self._id: + continue + self._cancelled_ids.update(record.get('cancelled', [])) + for task_id, keys in record.get('merged', {}).items(): + self._merged.update((task_id, key) for key in keys) + if record.get('winner') is not None: + self._winner_id = record['winner'] + self._pump_task = asyncio.create_task(self._pump()) + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + active_error = exc + try: + if exc is None and self._winner is None: + await self.gather() + except BaseException as error: + active_error = error + raise + finally: + self._closed = True + self._wake.set() + if self._pump_task is not None: + self._pump_task.cancel() + try: + await self._drain( + self._tasks, + deliberate=False, + preserve_error=active_error is not None, + ) + finally: + for task in self._tasks: + task._control.close_event_delivery() + if self._pump_task is not None: + await asyncio.gather(self._pump_task, return_exceptions=True) + + def spawn( + self, + node: NodeLike, + node_input: Any = None, + *, + run_id: str | None = None, + ) -> ParallelTask: + """Reserve a stable identity and queue a child without waiting for a slot.""" + if not self._entered or self._closed: + raise RuntimeError('Spawn requires an open parallel task scope.') + if self._winner is not None: + raise RuntimeError('Cannot spawn after this group selected its result.') + built_node = build_node(node) + index = self._next_index + self._next_index += 1 + if run_id is not None and ( + not run_id or run_id.isdigit() or any(c in run_id for c in '/@') + ): + raise ValueError( + 'Explicit run_id must be nonnumeric and contain neither / nor @.' + ) + if run_id is None: + scheduler = self._ctx._workflow_scheduler + if self._legacy_run_ids and scheduler is not None: + run_id = scheduler._state.next_run_id( + built_node.name, parent_path=self._ctx.node_path + ) + elif self._legacy_run_ids: + run_id = str(index + 1) + else: + run_id = f'{self._prefix}-{index + 1}' + elif self._call_scope is not None: + run_id = f'call-{self._call_scope}-{run_id}' + task = ParallelTask(self, built_node, node_input, run_id, index) + if not self._legacy_run_ids: + if task.id in self._ctx._parallel_task_paths: + raise ValueError(f'Duplicate parallel task ID: {task.id}') + self._ctx._parallel_task_paths.add(task.id) + self._tasks.append(task) + if task.id in self._cancelled_ids or ( + self._winner_id is not None and task.id != self._winner_id + ): + task._finish('cancelled') + self._wake.set() + return task + + async def wait( + self, + *, + mode: Literal['any', 'all'] = 'all', + ) -> tuple[list[ParallelTask], list[ParallelTask]]: + """Return completed and pending handles in spawn order. + + A worker awaiting human input pauses admission. Once running siblings stop, + ``all`` returns with queued handles pending so the parent can suspend. + """ + if mode not in {'any', 'all'}: + raise ValueError('mode must be any or all.') + while True: + self._changed.clear() + self._raise_pump_error() + done = [task for task in self._tasks if task._done.done()] + pending = [task for task in self._tasks if not task._done.done()] + running = any(task.status == 'running' for task in pending) + if ( + not pending + or (mode == 'any' and done) + or (self._paused and not running) + ): + return done, pending + await self._changed.wait() + + async def gather(self) -> list[Any]: + """Gather outputs in spawn order, stopping admission on failure or HITL.""" + self._fail_fast = True + if any(task.status in {'failed', 'waiting'} for task in self._tasks): + self._paused = True + while True: + self._changed.clear() + self._raise_pump_error() + failed = next( + (task for task in self._tasks if task.status == 'failed'), None + ) + if failed is not None: + await self._drain(self._tasks, deliberate=False, preserve_error=True) + await failed.result() + running = any(task.status == 'running' for task in self._tasks) + if self._paused and not running: + waiting = next( + (task for task in self._tasks if task.status == 'waiting'), None + ) + if waiting is not None: + await waiting.result() + if all(task._done.done() for task in self._tasks): + for task in self._tasks: + if task.status == 'cancelled' and task.id not in self._cancelled_ids: + await task.result() + return [task.output for task in self._tasks] + await self._changed.wait() + + async def _gather_items(self, node: NodeLike, items: list[Any]) -> list[Any]: + """Legacy list fan-out retaining only active handles and ordered outputs.""" + if not self._legacy_run_ids or self._tasks: + raise RuntimeError('List fan-out requires an empty legacy worker group.') + self._fail_fast = True + results: list[Any] = [None] * len(items) + next_input = 0 + interrupted: NodeInterruptedError | None = None + while next_input < len(items) or self._tasks: + while ( + not self._paused + and next_input < len(items) + and (self._limit is None or len(self._tasks) < self._limit) + ): + self.spawn(node, items[next_input]) + next_input += 1 + if not self._tasks: + break + done, _ = await self.wait(mode='any') + for task in done: + if task.status == 'waiting': + interrupted = NodeInterruptedError() + elif task.status in {'failed', 'cancelled'}: + await self._drain(self._tasks, deliberate=False, preserve_error=True) + await task.result() + else: + results[task._index] = task.output + self._tasks.remove(task) + if interrupted is not None: + raise interrupted + return results + + async def first( + self, + accept: Callable[[Any], Any] | None = None, + ) -> ParallelTask: + """Return the first acceptable result after cancelling and draining siblings. + + The predicate receives the output and may be asynchronous. Rejections and + worker failures allow remaining candidates to continue. Completion ties are + resolved in spawn order. Raises ValueError when no result is acceptable. + """ + if self._selecting: + raise RuntimeError('Only one first() selection may run at a time.') + if self._winner is not None: + return self._winner + if self._closed: + raise RuntimeError('The parallel task scope is closed.') + self._selecting = True + try: + return await self._first(accept) + finally: + self._selecting = False + self._first_mode = False + self._decision_pending = False + self._wake.set() + + async def _first(self, accept: Callable[[Any], Any] | None) -> ParallelTask: + self._fail_fast = False + self._first_mode = True + self._decision_pending = any( + task.status == 'completed' for task in self._tasks + ) + seen: set[str] = set() + try: + while True: + self._changed.clear() + self._raise_pump_error() + for task in self._tasks: + if task.id in seen or not task._done.done(): + continue + seen.add(task.id) + if task.status != 'completed': + continue + accepted = self._winner_id == task.id + if self._winner_id is None: + self._decision_pending = True + accepted = True if accept is None else accept(task.output) + if inspect.isawaitable(accepted): + accepted = await accepted + if accepted: + # Pause admission before persisting the decision, which yields to + # the Runner and must not admit another candidate in the meantime. + self._paused = True + async with self._decision_lock: + if self._winner_id is None: + losers = [ + other.id for other in self._tasks if other is not task + ] + await self._record(cancelled=losers, winner=task.id) + self._winner_id = task.id + self._cancelled_ids.update(losers) + self._winner = task + self._paused = True + await self._drain( + [other for other in self._tasks if other is not task], + deliberate=True, + ) + return task + self._decision_pending = any( + task.status == 'completed' and task.id not in seen + for task in self._tasks + ) + self._wake.set() + if self._decision_pending: + continue + running = any(task.status == 'running' for task in self._tasks) + if all(task._done.done() for task in self._tasks) or ( + self._paused and not running + ): + waiting = next( + (task for task in self._tasks if task.status == 'waiting'), None + ) + if waiting is not None: + await waiting.result() + failed = next( + (task for task in self._tasks if task.status == 'failed'), None + ) + if failed is not None: + await failed.result() + raise ValueError('No parallel task produced an acceptable result.') + await self._changed.wait() + except BaseException: + await self._drain(self._tasks, deliberate=False, preserve_error=True) + raise + + async def cancel(self, task: ParallelTask) -> None: + """Deliberately cancel one owned task and wait for it to stop.""" + self._require_owned(task) + if task.status in {'completed', 'failed', 'cancelled'}: + return + if self._closed: + raise RuntimeError('The parallel task scope is closed.') + async with self._decision_lock: + await self._record(cancelled=[task.id]) + self._cancelled_ids.add(task.id) + await self._drain([task], deliberate=True) + if self._winner_id is None: + self._paused = any( + other.status == 'waiting' + or (self._fail_fast and other.status == 'failed') + for other in self._tasks + ) + self._wake.set() + + async def merge_result( + self, + task: ParallelTask, + *, + keys: Iterable[str] | None = None, + expected: Mapping[str, Any] | None = None, + ) -> None: + """Persist a completed mapping proposal after validating every key. + + Selected keys must still match their values when the group was entered, + unless ``expected`` explicitly supplies another expected value. The whole + proposal is validated before writing, and a repeated merge of the same task + and keys is a no-op, including after a resume. The merge disposition and + state delta are persisted in one native event. Other worker writes remain + shared; this operation does not roll back their side effects. + """ + async with self._decision_lock: + await self._merge_result(task, keys=keys, expected=expected) + + async def _merge_result( + self, + task: ParallelTask, + *, + keys: Iterable[str] | None, + expected: Mapping[str, Any] | None, + ) -> None: + from ..sessions.state import _validate_state_entry + + if self._closed: + raise RuntimeError('The parallel task scope is closed.') + self._require_owned(task) + if task.status != 'completed' or not isinstance(task.output, Mapping): + raise ValueError('merge_result requires a completed mapping output.') + proposal = copy.deepcopy(dict(task.output)) + if any(not isinstance(key, str) for key in proposal): + raise ValueError('State proposal keys must be strings.') + if isinstance(keys, str): + raise ValueError('keys must be an iterable of state keys, not a string.') + selected = list(proposal) if keys is None else list(dict.fromkeys(keys)) + if any(not isinstance(key, str) or key not in proposal for key in selected): + raise ValueError('Every selected key must exist in the mapping output.') + selected = [key for key in selected if (task.id, key) not in self._merged] + current = self._ctx.state.to_dict() + for key in selected: + original = ( + expected[key] + if expected is not None and key in expected + else self._state_snapshot.get(key, _MISSING) + ) + if not _same_value(current.get(key, _MISSING), original): + raise ValueError(f'State changed before merging key {key!r}.') + schema = self._ctx.state._schema + if schema is not None: + for key in selected: + _validate_state_entry(schema, key, proposal[key]) + for key in selected: + self._ctx.state[key] = proposal[key] + # State's check/write remains synchronous. This event owns these deltas, + # so they must not be flushed again with the parent's later output event. + for key in selected: + self._ctx.actions.state_delta.pop(key, None) + if not selected: + return + await self._record( + cancelled=[], + merged={task.id: selected}, + state_delta={key: proposal[key] for key in selected}, + ) + for key in selected: + self._merged.add((task.id, key)) + + def _require_owned(self, task: ParallelTask) -> None: + if task._group is not self: + raise ValueError('The task belongs to a different parallel scope.') + + def _raise_pump_error(self) -> None: + if self._pump_error is not None: + raise self._pump_error + + async def _record( + self, + *, + cancelled: list[str], + winner: str | None = None, + merged: dict[str, list[str]] | None = None, + state_delta: dict[str, Any] | None = None, + ) -> None: + # A synthetic child path keeps records in the replay manager's subtree + # indexes without making them look like completion events for the owner. + record: dict[str, Any] = {'group': self._id, 'cancelled': cancelled} + if winner is not None: + record['winner'] = winner + if merged is not None: + record['merged'] = merged + await self._ctx._invocation_context._enqueue_event( + Event( + author=self._ctx.node.name + if self._ctx.node is not None + else 'parallel', + invocation_id=self._ctx._invocation_context.invocation_id, + node_info=NodeInfo( + path=str( + _NodePathBuilder.from_string(self._ctx.node_path).append( + '__parallel_tasks_control', self._prefix + ) + ) + ), + actions=EventActions(state_delta=state_delta or {}), + custom_metadata={_CONTROL_KEY: record}, + ) + ) + + async def _pump(self) -> None: + try: + while not self._closed: + self._wake.clear() + for task in self._tasks: + if task._cancel_requested and task.id not in self._cancelled_ids: + await self.cancel(task) + if task.status == 'cancelled': + scheduler = self._ctx._workflow_scheduler + if scheduler is not None: + scheduler._replay_manager.prepare_parent_sequence_barrier( + self._ctx, self._ctx.node_path + ) + await scheduler._replay_manager.advance_sequence( + self._ctx.node_path, f'{task.name}@{task._run_id}' + ) + active = sum(task.status == 'running' for task in self._tasks) + for task in self._tasks: + if ( + self._paused + or self._decision_pending + or (self._limit is not None and active >= self._limit) + ): + break + if task.status != 'queued': + continue + task._status = 'running' + task._control.status = 'running' + task._runner = asyncio.create_task(self._execute(task)) + active += 1 + await self._wake.wait() + except asyncio.CancelledError: + raise + except BaseException as error: + self._pump_error = error + self._changed.set() + + async def _execute(self, task: ParallelTask) -> None: + try: + with _bind_parallel_task(task._control): + child = await self._ctx._run_node_internal( + task._node, + task._input, + run_id=task._run_id, + use_sub_branch=True, + return_ctx=True, + skip_run_id_validation=self._legacy_run_ids, + ) + if task.id in self._cancelled_ids: + self._ctx._interrupt_ids.difference_update(child.interrupt_ids) + scheduler = self._ctx._workflow_scheduler + if scheduler is not None: + scheduler._state.interrupt_ids.difference_update(child.interrupt_ids) + task._finish('cancelled') + return + if child.error is not None: + raise DynamicNodeFailError( + message=f'Dynamic node {task.name} failed', + error=child.error, + error_node_path=child.error_node_path, + ) + if child.interrupt_ids: + task._interrupt_ids = set(child.interrupt_ids) + self._ctx._interrupt_ids.update(task._interrupt_ids) + raise NodeInterruptedError() + task._output = child.output + if self._first_mode: + self._decision_pending = True + task._finish('completed') + except asyncio.CancelledError: + task._finish('cancelled') + raise + except NodeInterruptedError as error: + self._paused = True + task._finish('waiting', error) + except BaseException as error: + if self._fail_fast: + self._paused = True + task._finish('failed', error) + + async def _drain( + self, + tasks: list[ParallelTask], + *, + deliberate: bool, + preserve_error: bool = False, + ) -> None: + running: list[asyncio.Task[None]] = [] + deadlines: list[float] = [] + loop = asyncio.get_running_loop() + for task in tasks: + if deliberate: + task._control.cancellation_requested = True + if deliberate and task._interrupt_ids: + self._ctx._interrupt_ids.difference_update(task._interrupt_ids) + scheduler = self._ctx._workflow_scheduler + if scheduler is not None: + scheduler._state.interrupt_ids.difference_update(task._interrupt_ids) + if task.status == 'queued' or (deliberate and task.status == 'waiting'): + task._finish('cancelled') + if task._runner is not None and not task._runner.done(): + if not task._cancellation_sent: + task._runner.cancel() + task._cancellation_sent = True + if task._cancel_deadline is None: + task._cancel_deadline = loop.time() + self._timeout + deadlines.append(task._cancel_deadline) + running.append(task._runner) + if not running: + return + done, pending = await asyncio.wait( + running, timeout=max(0.0, max(deadlines) - loop.time()) + ) + for runner in done: + if not runner.cancelled(): + runner.exception() + if pending: + message = ( + f'Parallel group {self._id}: {len(pending)} task(s) did not stop' + f' within {self._timeout}s.' + ) + if preserve_error: + logger.warning(message) + else: + raise TimeoutError(message) diff --git a/src/google/adk/workflow/_parallel_worker.py b/src/google/adk/workflow/_parallel_worker.py index 74afd50691c..100b2b87c72 100644 --- a/src/google/adk/workflow/_parallel_worker.py +++ b/src/google/adk/workflow/_parallel_worker.py @@ -16,10 +16,8 @@ from __future__ import annotations -import asyncio from collections.abc import AsyncGenerator from collections.abc import Iterator -import logging from typing import Any from google.genai import types @@ -31,14 +29,12 @@ from ..agents.context import Context from ..utils._schema_utils import _unwrap_content_payload from ._base_node import BaseNode -from ._errors import NodeInterruptedError from ._errors import WorkflowConfigurationError from ._graph import NodeLike +from ._parallel_tasks import ParallelTasks from ._retry_config import RetryConfig from .utils._workflow_graph_utils import build_node -logger = logging.getLogger('google_adk.' + __name__) - # How long to wait for cancelled items to actually stop before giving up on # them, so an item that swallows cancellation cannot hang this node forever. _CANCELLED_ITEM_DRAIN_TIMEOUT_SECONDS = 5.0 @@ -128,108 +124,11 @@ async def _run_impl( yield [] return - results = [None] * len(node_input) - pending_tasks: set[asyncio.Task[Any]] = set() - input_index = 0 - interrupted: NodeInterruptedError | None = None - - try: - while ( - not interrupted and input_index < len(node_input) - ) or pending_tasks: - # Check for any inputs waiting to be processed. Once a worker has - # interrupted, stop scheduling new items and wait for in-flight - # sibling workers to finish. - while ( - not interrupted - and input_index < len(node_input) - and ( - self.max_parallel_workers is None - or len(pending_tasks) < self.max_parallel_workers - ) - ): - item = node_input[input_index] - task = asyncio.create_task( - ctx.run_node( - self._node, - node_input=item, - use_sub_branch=True, - ) - ) - # Store index on task so we can place result correctly when done - setattr(task, '_worker_index', input_index) - pending_tasks.add(task) - input_index += 1 - - # If there are pending tasks, wait for first one to complete. - # We only wait for the first one, because after it completes, we want - # to check if any new items are waiting to be processed. - if pending_tasks: - done, pending_tasks = await asyncio.wait( - pending_tasks, return_when=asyncio.FIRST_COMPLETED - ) - # asyncio.wait returns completed tasks as an unordered set; iterate in - # input order so that, when several tasks fail in the same wake-up, - # the surfaced exception is deterministic across runs and replays. - failure: BaseException | None = None - for task in sorted(done, key=lambda t: getattr(t, '_worker_index')): - # Ask every completed task for its exception, not just the one we - # surface, otherwise asyncio reports the rest as never retrieved. - exc = task.exception() - if exc is not None: - if isinstance(exc, NodeInterruptedError): - if interrupted is None: - interrupted = exc - elif failure is None: - failure = exc - continue - - index = getattr(task, '_worker_index') - results[index] = task.result() - if failure is not None: - # Items still in flight are cancelled and drained below. - raise failure - if interrupted is not None: - raise interrupted - finally: - await self._drain_pending_items(pending_tasks) - - yield results - - async def _drain_pending_items( - self, pending_tasks: set[asyncio.Task[Any]] - ) -> None: - """Cancels the items still in flight and waits for them to stop.""" - if not pending_tasks: - return - # asyncio.wait does not own the tasks it waits on, so an item left running - # when this node is cancelled keeps going and then blocks forever emitting - # into a run nobody consumes any more. - for task in pending_tasks: - task.cancel() - drain = asyncio.gather(*pending_tasks, return_exceptions=True) - try: - # This normally runs while a CancelledError is being unwound, so the - # drain is shielded: a second cancellation would otherwise abandon it - # part-way and leave the remaining exceptions unretrieved. The timeout is - # a separate concern: an item that swallows cancellation would otherwise - # hold this node here forever. - await asyncio.wait_for( - asyncio.shield(drain), - timeout=_CANCELLED_ITEM_DRAIN_TIMEOUT_SECONDS, - ) - except asyncio.TimeoutError: - logger.warning( - 'Node %s: %d item(s) did not stop within %ss of being cancelled;' - ' abandoning them.', - self.name, - sum(1 for task in pending_tasks if not task.done()), - _CANCELLED_ITEM_DRAIN_TIMEOUT_SECONDS, - ) - except asyncio.CancelledError: - logger.warning( - 'Node %s: cancelled again while stopping %d item(s).', - self.name, - sum(1 for task in pending_tasks if not task.done()), - ) - raise + async with ParallelTasks( + ctx, + name=self.name, + max_parallel_workers=self.max_parallel_workers, + cancellation_timeout=_CANCELLED_ITEM_DRAIN_TIMEOUT_SECONDS, + _legacy_run_ids=True, + ) as tasks: + yield await tasks._gather_items(self._node, node_input) diff --git a/src/google/adk/workflow/utils/_rehydration_utils.py b/src/google/adk/workflow/utils/_rehydration_utils.py index 27f66a9d72b..077f9a5f39f 100644 --- a/src/google/adk/workflow/utils/_rehydration_utils.py +++ b/src/google/adk/workflow/utils/_rehydration_utils.py @@ -28,6 +28,7 @@ from pydantic import ValidationError from ...events._branch_path import _BranchPath +from ...events._internal_metadata import INTERNAL_METADATA_PREFIX from ...events._interrupts import extract_event_interrupt_ids from ...events._interrupts import is_interrupt_event from ...events._node_path_builder import _NodePathBuilder @@ -43,6 +44,46 @@ logger = logging.getLogger('google_adk.' + __name__) _RESULT_KEY = 'result' +_PARALLEL_TASKS_CONTROL_KEY = INTERNAL_METADATA_PREFIX + 'parallel_tasks' + + +def _update_parallel_cancellations( + cancelled_by_invocation: dict[str, set[str]], event: Event +) -> None: + """Index deliberate cancellations from one trusted native control event.""" + record = (event.custom_metadata or {}).get(_PARALLEL_TASKS_CONTROL_KEY) + if not isinstance(record, dict): + return + cancelled = record.get('cancelled', []) + if isinstance(cancelled, list) and cancelled: + cancelled_by_invocation.setdefault(event.invocation_id, set()).update( + path for path in cancelled if isinstance(path, str) and path + ) + + +def _is_cancelled_parallel_event( + event: Event, cancelled_by_invocation: dict[str, set[str]] +) -> bool: + """Match cancelled paths and transferred hops sharing their task identity.""" + cancelled_paths = cancelled_by_invocation.get(event.invocation_id) + if not cancelled_paths: + return False + owners = (event.custom_metadata or {}).get( + INTERNAL_METADATA_PREFIX + 'parallel_task_owners' + ) + if isinstance(owners, list) and any( + isinstance(task_id, str) and task_id in cancelled_paths + for task_id in owners + ): + return True + path: _NodePathBuilder | None = _NodePathBuilder.from_string( + event.node_info.path or '' + ) + while path is not None: + if str(path) in cancelled_paths: + return True + path = path.parent + return False @dataclass @@ -277,6 +318,16 @@ def _reconstruct_node_states( interrupt_owner: dict[str, str] = {} schemas_by_id: dict[str, Any] = {} + # Deliberate cancellation is a durable ownership decision. A losing worker's + # old input request must not keep an ancestor waiting after its coordinator + # selected another result. Keep history intact; exclude cancelled subtrees + # only from replay state, scoped to the invocation that recorded the choice. + cancelled_by_invocation: dict[str, set[str]] = {} + for event in events: + if invocation_id and event.invocation_id != invocation_id: + continue + _update_parallel_cancellations(cancelled_by_invocation, event) + base_path_builder = _NodePathBuilder.from_string(base_path) def get_owner_key(event_path_builder: _NodePathBuilder) -> str | None: @@ -297,6 +348,8 @@ def get_owner_key(event_path_builder: _NodePathBuilder) -> str | None: for event in events: if invocation_id and event.invocation_id != invocation_id: continue + if _is_cancelled_parallel_event(event, cancelled_by_invocation): + continue # 1. Match user function responses if event.author == 'user' and event.content and event.content.parts: diff --git a/src/google/adk/workflow/utils/_replay_manager.py b/src/google/adk/workflow/utils/_replay_manager.py index f5cc7f6820f..417dcc2860b 100644 --- a/src/google/adk/workflow/utils/_replay_manager.py +++ b/src/google/adk/workflow/utils/_replay_manager.py @@ -26,7 +26,9 @@ from ...events.event import Event from ._rehydration_utils import _ChildScanState from ._rehydration_utils import _event_has_result +from ._rehydration_utils import _is_cancelled_parallel_event from ._rehydration_utils import _reconstruct_node_states +from ._rehydration_utils import _update_parallel_cancellations from ._rehydration_utils import is_terminal_event from ._replay_sequence_barrier import ReplaySequenceBarrier @@ -45,6 +47,7 @@ def __init__(self) -> None: self._live_events: list[Event] = [] self._indexed_event_count: int = 0 self._indexed_last_event: Event | None = None + self._parallel_cancellations: dict[str, set[str]] = {} @property def sequence_barrier(self) -> ReplaySequenceBarrier | None: @@ -120,6 +123,7 @@ def _build_event_index(self, events: list[Event]) -> None: self._events_by_parent = {} self._transitive_events_by_parent = {} self._fc_to_parent = {} + self._parallel_cancellations = {} self._live_events = list(events) self._index_events(events) self._record_indexed_through(events) @@ -133,6 +137,7 @@ def _index_events(self, events: list[Event]) -> None: from ...events._interrupts import extract_event_interrupt_ids for event in events: + _update_parallel_cancellations(self._parallel_cancellations, event) if event.author == "user": self._index_user_event(event, self._fc_to_parent) continue @@ -200,12 +205,12 @@ def get_events_for_rehydration( path_builder = _NodePathBuilder.from_string(node_path) parent_builder = path_builder.parent if not parent_builder or not str(parent_builder): - return session_events + return self._without_cancelled_parallel_events(session_events) parent_path = str(parent_builder) node_events = self._transitive_events_by_parent.get(parent_path, []) if not node_events: - return session_events + return self._without_cancelled_parallel_events(session_events) # Top-level user text prompts live under root key (""). # Merge them so multi-turn turn inputs remain visible during state reconstruction. @@ -223,11 +228,25 @@ def get_events_for_rehydration( ] if not user_prompts: - return node_events + return self._without_cancelled_parallel_events(node_events) # Retain exact chronological ordering of session events. event_ids = node_event_ids.union(id(e) for e in user_prompts) - return [e for e in session_events if id(e) in event_ids] + return self._without_cancelled_parallel_events( + [e for e in session_events if id(e) in event_ids] + ) + + def _without_cancelled_parallel_events( + self, events: list[Event] + ) -> list[Event]: + """Apply invocation-wide ownership decisions even across transfer parents.""" + if not self._parallel_cancellations: + return events + return [ + event + for event in events + if not _is_cancelled_parallel_event(event, self._parallel_cancellations) + ] def _add_event_to_index(self, parent_path: str, event: Event) -> None: """Indexes an event under its direct parent and all ancestor paths up to root.""" @@ -257,7 +276,7 @@ def _scan_sequence( completed: set[str] = set() invocation_id = ctx._invocation_context.invocation_id - for event in events: + for event in self._without_cancelled_parallel_events(events): # The event index spans the whole session so multi-turn context stays # visible during rehydration. The replay sequence, however, must describe # only the current invocation: a terminal event from an earlier completed @@ -328,7 +347,7 @@ def scan_workflow_events( # Use transitive parent events for static child nodes so deeper descendant events (e.g. delegated outputs/interrupts) are recovered filtered_events = self._transitive_events_by_parent.get(ctx.node_path, []) raw_results = _reconstruct_node_states( - events=filtered_events, + events=self._without_cancelled_parallel_events(filtered_events), base_path=ctx.node_path, group_by_direct_child=True, invocation_id=ic.invocation_id, diff --git a/tests/unittests/workflow/test_parallel_task_context.py b/tests/unittests/workflow/test_parallel_task_context.py new file mode 100644 index 00000000000..7d40cd964fb --- /dev/null +++ b/tests/unittests/workflow/test_parallel_task_context.py @@ -0,0 +1,893 @@ +# 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 +# +# https://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. + +"""Delivery and native event correlation for parallel task messaging.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from typing import Any + +from google.adk.agents.context import Context +from google.adk.agents.llm_agent import LlmAgent +from google.adk.events._internal_metadata import public_event +from google.adk.events.event import Event +from google.adk.events.event import NodeInfo +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.tools.long_running_tool import LongRunningFunctionTool +from google.adk.workflow import node +from google.adk.workflow import START +from google.adk.workflow._parallel_task_context import _bind_parallel_task +from google.adk.workflow._parallel_task_context import _deliver_parallel_task_messages +from google.adk.workflow._parallel_task_context import _parallel_task_cancellation_requested +from google.adk.workflow._parallel_task_context import _ParallelTaskControl +from google.adk.workflow._parallel_task_context import _persist_parallel_task_cancelled_tools +from google.adk.workflow._parallel_task_context import _tag_parallel_task_event +from google.adk.workflow._workflow import Workflow +from google.genai import types +import pytest + +from . import testing_utils + + +def _request_text(request: LlmRequest) -> list[str]: + return [ + part.text + for content in request.contents + for part in content.parts or [] + if part.text + ] + + +def test_receipts_are_detached_and_pending_messages_finish_undelivered(): + """Finishing preserves accepted receipts without promising delivery.""" + control = _ParallelTaskControl(task_id='task', root_node_path='root/worker') + receipt = control.send_message('First instruction') + receipt['status'] = 'caller edit' + + control.finish('completed') + snapshot = control.receipts + snapshot[0]['status'] = 'another caller edit' + + assert control.message_receipt(receipt['id'])['status'] == 'undelivered' + with pytest.raises(ValueError, match='no longer accepting'): + control.send_message('Too late') + with pytest.raises(ValueError, match='Unknown message receipt'): + control.message_receipt('another task receipt') + + +@pytest.mark.parametrize('value', [0, -1, True, 1.5]) +@pytest.mark.parametrize('limit', ['max_messages', 'max_message_bytes']) +def test_mailbox_limits_must_be_positive_integers(limit, value): + """Invalid mailbox limits are rejected before task execution.""" + with pytest.raises(ValueError, match='positive integer'): + _ParallelTaskControl( + task_id='task', root_node_path='root/worker', **{limit: value} + ) + + +def test_mailbox_count_and_utf8_bytes_are_bounded_before_acceptance(): + """Admission counts UTF-8 bytes and retains only accepted receipts.""" + control = _ParallelTaskControl( + task_id='task', + root_node_path='root/worker', + max_messages=1, + max_message_bytes=4, + ) + with pytest.raises(ValueError, match='byte limit'): + control.send_message('ééé') + receipt = control.send_message('éé') + with pytest.raises(ValueError, match='message limit'): + control.send_message('next') + + assert control.receipts == [receipt] + + +def test_progress_tags_copy_events_and_binding_restores_after_nested_scope(): + """Progress metadata never mutates original or previously tagged events.""" + event = Event(author='worker', custom_metadata={'application': 'kept'}) + outer = _ParallelTaskControl(task_id='outer', root_node_path='root/outer') + inner = _ParallelTaskControl( + task_id='inner', root_node_path='root/outer/inner' + ) + outer.status = 'running' + + with _bind_parallel_task(outer): + first = _tag_parallel_task_event(event) + with _bind_parallel_task(inner): + nested = _tag_parallel_task_event(event) + assert not _parallel_task_cancellation_requested() + outer.cancellation_requested = True + assert _parallel_task_cancellation_requested() + outer.finish('completed') + last = _tag_parallel_task_event(event) + + assert event.custom_metadata == {'application': 'kept'} + assert first is not event + assert first.custom_metadata['parallel_task'] == { + 'task_id': 'outer', + 'status': 'running', + } + assert nested.custom_metadata['parallel_task']['task_id'] == 'inner' + assert first.custom_metadata['__adk_internal_parallel_task_owners'] == [ + 'outer' + ] + assert nested.custom_metadata['__adk_internal_parallel_task_owners'] == [ + 'outer', + 'inner', + ] + public = public_event(nested) + assert '__adk_internal_parallel_task_owners' not in public.custom_metadata + assert public.custom_metadata['parallel_task']['task_id'] == 'inner' + assert last.custom_metadata['__adk_internal_parallel_task_owners'] == [ + 'outer' + ] + assert last.custom_metadata['parallel_task']['status'] == 'completed' + assert _tag_parallel_task_event(event) is event + assert not _parallel_task_cancellation_requested() + + +async def test_delivery_waits_for_native_persistence_and_keeps_later_messages(): + """Only the accepted batch is acknowledged after native event persistence. + + Setup: A worker pauses while its message event awaits Runner persistence. + Act: Queue a second message, persist the first event, and finish the worker. + Assert: The first receipt is delivered; the later receipt is undelivered. + """ + agent = LlmAgent(name='worker', model=testing_utils.MockModel.create(['ok'])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.node_path = 'root/worker' + invocation._event_queue = asyncio.Queue() + control = _ParallelTaskControl(task_id='task', root_node_path='root/worker') + first = control.send_message('First instruction') + request = LlmRequest() + + with _bind_parallel_task(control): + delivery = asyncio.create_task( + _deliver_parallel_task_messages(invocation, request) + ) + event, persisted = await asyncio.wait_for(invocation._event_queue.get(), 2) + assert request.contents == [] + assert control.message_receipt(first['id'])['status'] == 'queued' + later = control.send_message('Arrived during persistence') + await invocation.session_service.append_event(invocation.session, event) + persisted.set() + await asyncio.wait_for(delivery, 2) + control.finish('completed') + + assert _request_text(request) == [ + 'Message from the coordinating task:\nFirst instruction' + ] + assert control.message_receipt(first['id'])['status'] == 'delivered' + assert control.message_receipt(later['id'])['status'] == 'undelivered' + assert event.node_info.path == 'root/worker' + assert event.custom_metadata['parallel_task_message']['message_ids'] == [ + first['id'] + ] + assert event in invocation.session.events + + +async def test_cancel_during_persistence_does_not_claim_delivery(): + """A cancelled persistence wait leaves its message explicitly undelivered.""" + agent = LlmAgent(name='worker', model=testing_utils.MockModel.create(['ok'])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.node_path = 'root/worker' + invocation._event_queue = asyncio.Queue() + control = _ParallelTaskControl(task_id='task', root_node_path='root/worker') + receipt = control.send_message('Instruction') + request = LlmRequest() + + with _bind_parallel_task(control): + delivery = asyncio.create_task( + _deliver_parallel_task_messages(invocation, request) + ) + await asyncio.wait_for(invocation._event_queue.get(), 2) + delivery.cancel() + with pytest.raises(asyncio.CancelledError): + await delivery + control.finish('cancelled') + + assert request.contents == [] + assert control.message_receipt(receipt['id'])['status'] == 'undelivered' + + +async def test_finish_during_persistence_does_not_resume_delivery(): + """Persisted mail is not falsely acknowledged after its worker has closed.""" + agent = LlmAgent(name='worker', model=testing_utils.MockModel.create(['ok'])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.node_path = 'root/worker' + invocation._event_queue = asyncio.Queue() + control = _ParallelTaskControl(task_id='task', root_node_path='root/worker') + receipt = control.send_message('Instruction') + request = LlmRequest() + + with _bind_parallel_task(control): + delivery = asyncio.create_task( + _deliver_parallel_task_messages(invocation, request) + ) + event, persisted = await asyncio.wait_for(invocation._event_queue.get(), 2) + await invocation.session_service.append_event(invocation.session, event) + control.finish('cancelled') + persisted.set() + await asyncio.wait_for(delivery, 2) + + assert event in invocation.session.events + assert request.contents == [] + assert control.message_receipt(receipt['id'])['status'] == 'undelivered' + + +async def test_delivered_receipts_still_count_toward_lifetime_admission_limit(): + """A caller cannot grow retained receipts by draining and refilling inboxes.""" + agent = LlmAgent(name='worker', model=testing_utils.MockModel.create(['ok'])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.node_path = 'root/worker' + invocation._event_queue = asyncio.Queue() + control = _ParallelTaskControl( + task_id='task', root_node_path='root/worker', max_messages=1 + ) + receipt = control.send_message('Instruction') + + with _bind_parallel_task(control): + delivery = asyncio.create_task( + _deliver_parallel_task_messages(invocation, LlmRequest()) + ) + _, persisted = await asyncio.wait_for(invocation._event_queue.get(), 2) + persisted.set() + await asyncio.wait_for(delivery, 2) + + assert control.message_receipt(receipt['id'])['status'] == 'delivered' + with pytest.raises(ValueError, match='message limit'): + control.send_message('Another instruction') + + +async def test_native_runner_delivers_fifo_only_to_root_and_retains_history(): + """Nested agents cannot consume messages addressed to their task root. + + Setup: The task root pauses in a tool before invoking a nested agent. + Act: Its coordinator sends two messages, then releases the nested agent. + Assert: Only later root requests contain both messages in FIFO order. Their + persisted native events retain task identity, and root progress is live. + """ + before_nested = asyncio.Event() + release_nested = asyncio.Event() + observations: dict[str, Any] = {} + progress: list[Event] = [] + nested_model = testing_utils.MockModel.create(['nested complete']) + nested = LlmAgent(name='nested', model=nested_model) + + class Observer(BasePlugin): + + async def on_event_callback(self, *, invocation_context, event): + progress.append(event.model_copy(deep=True)) + return None + + async def consult(ctx: Context) -> str: + before_nested.set() + await release_nested.wait() + return await ctx.run_node(nested, 'Nested task', use_sub_branch=True) + + def another_step() -> str: + return 'One more model boundary' + + root_model = testing_utils.MockModel.create([ + types.Part.from_function_call(name='consult', args={}), + types.Part.from_function_call(name='another_step', args={}), + 'root complete', + ]) + worker = LlmAgent( + name='worker', model=root_model, tools=[consult, another_step] + ) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as group: + task = group.spawn(worker, 'Original task') + await asyncio.wait_for(before_nested.wait(), 2) + observations['live_progress'] = any( + (event.custom_metadata or {}).get('parallel_task', {}).get('task_id') + == task.id + for event in progress + ) + first = task.send_message('Root-only first instruction') + second = task.send_message('Root-only second instruction') + release_nested.set() + output = await task.result() + observations['task_id'] = task.id + observations['receipts'] = [ + task.message_receipt(first['id']), + task.message_receipt(second['id']), + ] + return output + + workflow = Workflow(name='team', edges=[(START, coordinate)]) + runner = testing_utils.InMemoryRunner( + node=workflow, plugins=[Observer(name='parallel_observer')] + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 5) + finally: + await runner.runner.close() + + assert observations['live_progress'] + assert len(nested_model.requests) == 1 + assert not any( + 'Root-only' in text for text in _request_text(nested_model.requests[0]) + ) + assert len(root_model.requests) == 3 + expected = [ + 'Message from the coordinating task:\nRoot-only first instruction', + 'Message from the coordinating task:\nRoot-only second instruction', + ] + for request in root_model.requests[1:]: + assert [ + text for text in _request_text(request) if 'Root-only' in text + ] == expected + assert [receipt['status'] for receipt in observations['receipts']] == [ + 'delivered', + 'delivered', + ] + message_events = [ + event + for event in events + if 'parallel_task_message' in (event.custom_metadata or {}) + ] + assert len(message_events) == 1 + assert message_events[0].node_info.path == observations['task_id'] + + +async def test_native_runner_final_model_race_reports_undelivered_message(): + """Messages accepted after the final boundary never silently disappear.""" + entered = asyncio.Event() + release = asyncio.Event() + observations: dict[str, Any] = {} + + class GatedModel(testing_utils.MockModel): + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + self.requests.append(llm_request.model_copy(deep=True)) + entered.set() + await release.wait() + yield LlmResponse( + content=types.Content( + role='model', parts=[types.Part.from_text(text='complete')] + ) + ) + + model = GatedModel(responses=[]) + worker = LlmAgent(name='worker', model=model) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as group: + task = group.spawn(worker, 'Task') + await asyncio.wait_for(entered.wait(), 2) + receipt = task.send_message('After the final model boundary') + release.set() + result = await task.result() + observations['receipt'] = task.message_receipt(receipt['id']) + return result + + workflow = Workflow(name='team', edges=[(START, coordinate)]) + runner = testing_utils.InMemoryRunner(node=workflow) + try: + await asyncio.wait_for(runner.run_async('start'), 5) + finally: + await runner.runner.close() + + assert observations['receipt']['status'] == 'undelivered' + assert len(model.requests) == 1 + assert not any( + 'After the final' in text for text in _request_text(model.requests[0]) + ) + + +async def test_native_runner_rebuilds_delivered_message_history_after_pause(): + """A resumed native task reads persisted mail without retaining its mailbox. + + Setup: A task receives mail before pausing on a long-running tool. + Act: Resume the workflow with the tool result, recreating the task scope. + Assert: The resumed model request contains exactly one copy of the message. + """ + coordinator_runs = 0 + handles = [] + + def await_approval() -> None: + return None + + model = testing_utils.MockModel.create([ + types.Part.from_function_call(name='await_approval', args={}), + 'Approved output', + ]) + worker = LlmAgent( + name='worker', + model=model, + tools=[LongRunningFunctionTool(func=await_approval)], + ) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + nonlocal coordinator_runs + coordinator_runs += 1 + async with ctx.parallel_tasks() as group: + task = group.spawn(worker, 'Task', run_id='owned') + handles.append(task) + if coordinator_runs == 1: + task.send_message('Remember this after approval') + return await task.result() + + workflow = Workflow(name='team', edges=[(START, coordinate)]) + runner = testing_utils.InMemoryRunner(node=workflow) + try: + first_events = await asyncio.wait_for(runner.run_async('start'), 5) + function_calls = [ + part.function_call + for event in first_events + if event.content + for part in event.content.parts or [] + if part.function_call and part.function_call.name == 'await_approval' + ] + assert len(function_calls) == 1 + assert handles[0].status == 'waiting' + assert handles[0].message_receipts[0]['status'] == 'delivered' + await asyncio.wait_for( + runner.run_async( + types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + name='await_approval', + id=function_calls[0].id, + response={'approved': True}, + ) + ) + ], + ) + ), + 5, + ) + finally: + await runner.runner.close() + + assert coordinator_runs == 2 + assert handles[0] is not handles[1] + assert handles[1].message_receipts == [] + assert handles[1].status == 'completed' + assert len(model.requests) == 2 + assert [ + text + for text in _request_text(model.requests[1]) + if 'Remember this after approval' in text + ] == ['Message from the coordinating task:\nRemember this after approval'] + + +@pytest.mark.parametrize('boundary', ['model', 'tool']) +@pytest.mark.parametrize('nested', [False, True]) +async def test_native_cancellation_closes_generators_while_parent_can_work( + boundary, + nested, +): + """An active agent does not block root work or outlive deliberate cancel.""" + entered = asyncio.Event() + released = asyncio.Event() + closed = asyncio.Event() + observations = {} + + class GatedModel(testing_utils.MockModel): + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + try: + entered.set() + await released.wait() + yield LlmResponse( + content=types.Content( + role='model', parts=[types.Part.from_text(text='complete')] + ) + ) + finally: + closed.set() + + async def gated_tool() -> AsyncGenerator[str, None]: + try: + entered.set() + await released.wait() + yield 'complete' + finally: + closed.set() + + worker = LlmAgent( + name='worker', + model=GatedModel(responses=[]) + if boundary == 'model' + else testing_utils.MockModel.create([ + types.Part.from_function_call(name='gated_tool', args={}), + 'complete', + ]), + tools=[] if boundary == 'model' else [gated_tool], + ) + + async def parent_work() -> str: + return 'parent made progress' + + @node(rerun_on_resume=True) + async def nested_team(ctx: Context) -> str: + async with ctx.parallel_tasks() as group: + return await group.spawn(worker, 'Nested task').result() + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as group: + task = group.spawn(nested_team if nested else worker, 'Task') + await asyncio.wait_for(entered.wait(), 2) + observations['parent_output'] = await ctx.run_node(parent_work) + observations['worker_status_before_cancel'] = task.status + observations['closed_before_cancel'] = closed.is_set() + receipt = task.send_message('Queued before cancellation') + await group.cancel(task) + observations['closed_after_cancel'] = closed.is_set() + observations['worker_status_after_cancel'] = task.status + observations['receipt'] = task.message_receipt(receipt['id']) + with pytest.raises(asyncio.CancelledError): + await task.result() + return 'parent complete' + + runner = testing_utils.InMemoryRunner( + node=Workflow(name='team', edges=[(START, coordinate)]) + ) + try: + await asyncio.wait_for(runner.run_async('start'), 5) + history = runner.session.events + finally: + released.set() + await runner.runner.close() + + assert observations['parent_output'] == 'parent made progress' + assert observations['worker_status_before_cancel'] == 'running' + assert not observations['closed_before_cancel'] + assert observations['closed_after_cancel'] + assert observations['worker_status_after_cancel'] == 'cancelled' + assert observations['receipt']['status'] == 'undelivered' + if boundary == 'tool': + responses = [event for event in history if event.get_function_responses()] + assert len(responses) == 1 + metadata = responses[0].custom_metadata + owners = metadata['__adk_internal_parallel_task_owners'] + assert len(owners) == (2 if nested else 1) + assert owners[-1] == metadata['parallel_task']['task_id'] + + +async def test_first_pauses_queued_admission_for_already_completed_candidate(): + """Evaluating an existing result must not admit another queued worker.""" + predicate_entered = asyncio.Event() + release_predicate = asyncio.Event() + second_entered = asyncio.Event() + observations = {} + + async def first_worker() -> str: + return 'accepted' + + async def second_worker() -> str: + second_entered.set() + return 'unnecessary' + + async def parent_work() -> str: + return 'parent progress while evaluating' + + async def accept(value: str) -> bool: + predicate_entered.set() + await release_predicate.wait() + return value == 'accepted' + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks(max_parallel_workers=1) as group: + first = group.spawn(first_worker) + await first.result() + second = group.spawn(second_worker) + + async def observe_pending_decision() -> None: + try: + await asyncio.wait_for(predicate_entered.wait(), 2) + await ctx.run_node(parent_work) + observations['second_started'] = second_entered.is_set() + observations['second_status'] = second.status + finally: + release_predicate.set() + + observer = asyncio.create_task(observe_pending_decision()) + winner = await group.first(accept) + await observer + return winner.output + + runner = testing_utils.InMemoryRunner( + node=Workflow(name='team', edges=[(START, coordinate)]) + ) + try: + await asyncio.wait_for(runner.run_async('start'), 5) + finally: + release_predicate.set() + await runner.runner.close() + + assert not observations['second_started'] + assert observations['second_status'] == 'queued' + + +async def test_cancelled_tool_batch_preserves_completed_sibling_and_lineage(): + """Cancellation retains completed work and pairs the unfinished sibling.""" + completed = asyncio.Event() + entered = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + + async def completed_tool(ctx: Context) -> dict[str, int]: + ctx.state['completed_tool_effect'] = 7 + completed.set() + return {'value': 7} + + async def blocked_tool() -> AsyncGenerator[str, None]: + try: + entered.set() + await release.wait() + yield 'complete' + finally: + closed.set() + + async def parent_work() -> str: + return 'continue coordinating' + + model = testing_utils.MockModel.create([ + [ + types.Part.from_function_call(name='completed_tool', args={}), + types.Part.from_function_call(name='blocked_tool', args={}), + ], + 'Should not call the model again', + ]) + worker = LlmAgent( + name='worker', model=model, tools=[completed_tool, blocked_tool] + ) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as group: + task = group.spawn(worker, 'Task') + await asyncio.wait_for( + asyncio.gather(completed.wait(), entered.wait()), 2 + ) + await ctx.run_node(parent_work) + await group.cancel(task) + return 'complete' + + runner = testing_utils.InMemoryRunner( + node=Workflow(name='team', edges=[(START, coordinate)]) + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 5) + finally: + release.set() + await runner.runner.close() + + assert closed.is_set() + assert len(model.requests) == 1 + calls = { + call.id: (event, call) + for event in events + for call in event.get_function_calls() + } + responses = [ + (event, response) + for event in events + for response in event.get_function_responses() + ] + assert len(calls) == len(responses) == 2 + assert {response.id for _, response in responses} == calls.keys() + for response_event, response in responses: + call_event, call = calls[response.id] + assert response.name == call.name + assert response_event.author == call_event.author + assert response_event.branch == call_event.branch + assert response_event.node_info.path == call_event.node_info.path + assert response_event.invocation_id == call_event.invocation_id + values = {response.name: response.response for _, response in responses} + assert values['completed_tool'] == {'value': 7} + assert values['blocked_tool']['status'] == 'cancelled' + assert 'External side effects may already have occurred' in ( + values['blocked_tool']['error'] + ) + assert any( + event.actions.state_delta.get('completed_tool_effect') == 7 + for event, _ in responses + ) + + +@pytest.mark.parametrize('response_path', ['root/worker', 'root/other']) +async def test_cancellation_response_deduplicates_only_its_own_node( + response_path, +): + """Another node reusing a provider call ID cannot suppress this response.""" + agent = LlmAgent(name='worker', model=testing_utils.MockModel.create(['ok'])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.node_path = 'root/worker' + invocation._event_queue = asyncio.Queue() + call_event = Event( + author='worker', + invocation_id=invocation.invocation_id, + node_info=NodeInfo(path=invocation.node_path), + content=types.Content( + role='model', + parts=[ + types.Part( + function_call=types.FunctionCall( + id='finished', name='completed_tool', args={} + ) + ), + types.Part( + function_call=types.FunctionCall( + id='unfinished', name='blocked_tool', args={} + ) + ), + ], + ), + ) + kept_event = Event( + author='worker', + invocation_id=invocation.invocation_id, + node_info=NodeInfo(path=response_path), + content=types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id='finished', + name='completed_tool', + response={'value': 7}, + ) + ), + ], + ), + ) + await invocation.session_service.append_event(invocation.session, kept_event) + control = _ParallelTaskControl(task_id='task', root_node_path='root/worker') + with _bind_parallel_task(control): + await _persist_parallel_task_cancelled_tools( + invocation, call_event, kept_event + ) + assert invocation._event_queue.empty() + control.cancellation_requested = True + cleanup = asyncio.create_task( + _persist_parallel_task_cancelled_tools( + invocation, call_event, kept_event + ) + ) + event, persisted = await asyncio.wait_for(invocation._event_queue.get(), 2) + persisted.set() + await asyncio.wait_for(cleanup, 2) + + expected_ids = ( + ['unfinished'] + if response_path == 'root/worker' + else ['finished', 'unfinished'] + ) + assert [ + response.id for response in event.get_function_responses() + ] == expected_ids + + +@pytest.mark.parametrize('close_ancestor', [False, True]) +async def test_late_cancellation_cleanup_stops_when_scope_delivery_closes( + close_ancestor, +): + """Cleanup cannot wait forever after a timed-out owner's consumer exits.""" + agent = LlmAgent(name='worker', model=testing_utils.MockModel.create(['ok'])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.node_path = 'root/worker' + invocation._event_queue = asyncio.Queue() + call_event = Event( + author='worker', + invocation_id=invocation.invocation_id, + content=types.Content( + role='model', + parts=[ + types.Part( + function_call=types.FunctionCall( + id='unfinished', name='blocked_tool', args={} + ) + ), + ], + ), + ) + control = _ParallelTaskControl(task_id='task', root_node_path='root/worker') + inner = _ParallelTaskControl( + task_id='inner', root_node_path='root/worker/nested' + ) + control.cancellation_requested = True + with _bind_parallel_task(control), _bind_parallel_task(inner): + cleanup = asyncio.create_task( + _persist_parallel_task_cancelled_tools(invocation, call_event, None) + ) + await asyncio.wait_for(invocation._event_queue.get(), 2) + (control if close_ancestor else inner).close_event_delivery() + await asyncio.wait_for(cleanup, 2) + await _persist_parallel_task_cancelled_tools(invocation, call_event, None) + + assert invocation._event_queue.empty() + + +async def test_timeout_seals_mailbox_and_allows_late_tool_cleanup_to_finish(): + """A timed-out scope rejects mail without pretending its worker has stopped.""" + entered = asyncio.Event() + ignored_cancellation = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + observations = {} + + async def resistant_tool() -> str: + try: + entered.set() + try: + await release.wait() + except asyncio.CancelledError: + ignored_cancellation.set() + await release.wait() + return 'Finished external work after cancellation' + finally: + closed.set() + + worker = LlmAgent( + name='worker', + model=testing_utils.MockModel.create([ + types.Part.from_function_call(name='resistant_tool', args={}), + 'Should not request another model response', + ]), + tools=[resistant_tool], + ) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> None: + async with ctx.parallel_tasks(cancellation_timeout=0.01) as group: + task = group.spawn(worker, 'Task') + observations['task'] = task + await asyncio.wait_for(entered.wait(), 2) + observations['receipt'] = task.send_message('Too late to deliver') + await group.cancel(task) + + # A plain root node isolates the group's bounded ownership contract from a + # Workflow's independent cleanup of every dynamic task in that workflow. + runner = testing_utils.InMemoryRunner(node=coordinate) + try: + with pytest.raises(TimeoutError, match='did not stop'): + await asyncio.wait_for(runner.run_async('start'), 2) + task = observations['task'] + assert ignored_cancellation.is_set() + assert not closed.is_set() + assert task.status == 'running' + assert task.message_receipt(observations['receipt']['id'])['status'] == ( + 'undelivered' + ) + with pytest.raises(ValueError, match='no longer accepting'): + task.send_message('The owning scope has already closed') + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task.result(), 2) + assert closed.is_set() + assert task.status == 'cancelled' + finally: + release.set() + await runner.runner.close() diff --git a/tests/unittests/workflow/test_parallel_task_replay.py b/tests/unittests/workflow/test_parallel_task_replay.py new file mode 100644 index 00000000000..4291bd53028 --- /dev/null +++ b/tests/unittests/workflow/test_parallel_task_replay.py @@ -0,0 +1,249 @@ +# 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 +# +# https://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. + +"""Replay of owned tasks whose agent transfers cross workflow schedulers.""" + +from __future__ import annotations + +import asyncio +from collections.abc 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.apps.app import ResumabilityConfig +from google.adk.events.event import Event +from google.adk.events.event import NodeInfo +from google.adk.events.event_actions import EventActions +from google.adk.events.request_input import RequestInput +from google.adk.tools.long_running_tool import LongRunningFunctionTool +from google.adk.tools.tool_context import ToolContext +from google.adk.workflow import node +from google.adk.workflow import START +from google.adk.workflow import Workflow +from google.adk.workflow._errors import NodeInterruptedError +from google.adk.workflow.utils._replay_manager import ReplayManager +from google.genai import types +import pytest + +from . import testing_utils + + +@pytest.mark.parametrize('is_resumable', [False, True]) +@pytest.mark.parametrize('nested_group', [False, True]) +async def test_cancelled_transfer_to_foreign_scheduler_does_not_block_resume( + is_resumable: bool, + nested_group: bool, +): + """A cancelled parent transfer cannot retain an unrelated replay barrier. + + The outer agent calls an inner workflow through a tool. Its child worker + transfers back to that outer agent, crossing from the inner scheduler to the + outer scheduler, and pauses for input. Another worker wins. Later resuming + an unrelated input request must skip the cancelled hop and completed work. + """ + chosen = [] + cancelled_ids = [] + + def abandoned_approval() -> None: + """An approval requested by the losing transferred worker.""" + + child_model = testing_utils.MockModel.create([ + types.Part.from_function_call( + name='transfer_to_agent', args={'agent_name': 'outer_agent'} + ), + ]) + child_agent = LlmAgent(name='child_agent', model=child_model) + + @node(rerun_on_resume=True) + async def nested(ctx: Context, node_input: str) -> str: + async with ctx.parallel_tasks() as tasks: + return await tasks.spawn(child_agent, node_input).result() + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + losing = tasks.spawn(nested if nested_group else child_agent, 'work') + + async def accepted() -> str: + try: + await losing.result() + except NodeInterruptedError: + return 'winner' + raise AssertionError('The transferred worker should request approval.') + + tasks.spawn(accepted) + winner = await tasks.first() + chosen.append(winner.output) + cancelled_ids.append(losing.id) + assert losing.status == 'cancelled' + return winner.output + + inner = Workflow(name='inner', edges=[(START, coordinate)]) + + async def run_inner(tool_context: ToolContext) -> str: + """Run the inner collaboration workflow.""" + return await tool_context.run_node(inner, node_input='work') + + outer_model = testing_utils.MockModel.create([ + types.Part.from_function_call(name='run_inner', args={}), + types.Part.from_function_call(name='abandoned_approval', args={}), + 'outer complete', + ]) + outer_agent = LlmAgent( + name='outer_agent', + model=outer_model, + sub_agents=[child_agent], + tools=[run_inner, LongRunningFunctionTool(abandoned_approval)], + ) + + @node(rerun_on_resume=True) + async def approval(ctx: Context) -> AsyncGenerator[Event, None]: + if 'keep_approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='keep_approval', message='Continue?') + return + yield Event(output='finished') + + workflow = Workflow(name='outer', edges=[(START, outer_agent, approval)]) + runner = testing_utils.InMemoryRunner( + app=App( + name='transfer_tasks', + root_agent=workflow, + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 5) + assert chosen == ['winner'] + assert len(outer_model.requests) == 3 + cancelled_hops = [ + event + for event in runner.session.events + if cancelled_ids[0] + in (event.custom_metadata or {}).get( + '__adk_internal_parallel_task_owners', [] + ) + and event.long_running_tool_ids + ] + assert cancelled_hops + assert all( + not event.node_info.path.startswith(cancelled_ids[0] + '/') + and event.node_info.path != cancelled_ids[0] + for event in cancelled_hops + ) + response = types.Content( + role='user', + parts=[ + types.Part.from_function_response( + name='adk_request_input', + response={'result': 'approved'}, + ) + ], + ) + response.parts[0].function_response.id = 'keep_approval' + resumed = await asyncio.wait_for( + runner.run_async( + new_message=response, + invocation_id=events[0].invocation_id, + ), + 5, + ) + assert any(event.output == 'finished' for event in resumed) + assert chosen == ['winner'] + assert len(outer_model.requests) == 3 + assert len(child_model.requests) == 1 + finally: + await runner.runner.close() + + +async def test_cancellation_index_rebuilds_after_rewind(): + """Rewinding the decision permits a new execution of the same task path.""" + agent = LlmAgent(name='wf', model=testing_utils.MockModel.create([])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.invocation_id = 'run' + ctx = Context(invocation, node=agent, run_id='1') + worker_path = 'wf@1/worker@one' + invocation.session.events.extend([ + Event( + author='worker', + invocation_id='run', + output='discarded', + node_info=NodeInfo(path=worker_path), + ), + Event( + author='wf', + invocation_id='run', + node_info=NodeInfo(path='wf@1/__parallel_tasks_control@parallel-1'), + custom_metadata={ + '__adk_internal_parallel_tasks': { + 'group': 'wf@1/parallel@1', + 'cancelled': [worker_path], + } + }, + ), + ]) + replay = ReplayManager() + recovered, sequence = replay.scan_workflow_events(ctx) + assert 'worker@one' not in recovered + assert 'worker@one' not in sequence + + invocation.session.events.extend([ + Event( + author='user', + invocation_id='rewind', + actions=EventActions(rewind_before_invocation_id='run'), + ), + Event( + author='worker', + invocation_id='run', + output='new execution', + node_info=NodeInfo(path=worker_path), + custom_metadata={'parallel_task': {'task_id': []}}, + ), + ]) + recovered, sequence = replay.scan_workflow_events(ctx) + assert recovered['worker@one'].output == 'new execution' + assert sequence == ['worker@one'] + + +@pytest.mark.parametrize('task_id', [{'arbitrary': True}, [], 'wf@1/loser@one']) +async def test_cancellation_filter_ignores_unrelated_public_metadata(task_id): + """Arbitrary user metadata cannot break replay after a cancellation.""" + agent = LlmAgent(name='wf', model=testing_utils.MockModel.create([])) + invocation = await testing_utils.create_invocation_context(agent) + invocation.invocation_id = 'run' + ctx = Context(invocation, node=agent, run_id='1') + invocation.session.events.extend([ + Event( + author='wf', + invocation_id='run', + node_info=NodeInfo(path='wf@1/__parallel_tasks_control@parallel-1'), + custom_metadata={ + '__adk_internal_parallel_tasks': { + 'group': 'wf@1/parallel@1', + 'cancelled': ['wf@1/loser@one'], + } + }, + ), + Event( + author='worker', + invocation_id='run', + output='kept', + node_info=NodeInfo(path='wf@1/worker@one'), + custom_metadata={'parallel_task': {'task_id': task_id}}, + ), + ]) + recovered, sequence = ReplayManager().scan_workflow_events(ctx) + assert recovered['worker@one'].output == 'kept' + assert sequence == ['worker@one'] diff --git a/tests/unittests/workflow/test_parallel_tasks.py b/tests/unittests/workflow/test_parallel_tasks.py new file mode 100644 index 00000000000..e4f9e241278 --- /dev/null +++ b/tests/unittests/workflow/test_parallel_tasks.py @@ -0,0 +1,994 @@ +# 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 +# +# https://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. + +"""Runner-backed contracts for owned parallel workflow tasks.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from typing import Any + +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.apps.app import ResumabilityConfig +from google.adk.events.event import Event +from google.adk.events.event_actions import EventActions +from google.adk.events.request_input import RequestInput +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.sessions.state import StateSchemaError +from google.adk.workflow import node +from google.adk.workflow import START +from google.adk.workflow._workflow import Workflow +from google.genai import types +from pydantic import BaseModel +import pytest + +from . import testing_utils + +_MAX_WAIT = 5 + + +async def _run(coordinator, **workflow_options): + runner = testing_utils.InMemoryRunner( + node=Workflow( + name='team', edges=[(START, coordinator)], **workflow_options + ) + ) + try: + return await asyncio.wait_for(runner.run_async('start'), _MAX_WAIT) + finally: + await runner.runner.close() + + +def _response(text: str) -> LlmResponse: + return LlmResponse( + content=types.Content(role='model', parts=[types.Part(text=text)]) + ) + + +async def test_four_agents_overlap_while_coordinator_does_independent_work(): + """Spawning is nonblocking and all four real Agent model calls overlap.""" + entered: asyncio.Queue[str] = asyncio.Queue() + release = asyncio.Event() + active = 0 + peak = 0 + observations = {} + + class GatedModel(testing_utils.MockModel): + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + nonlocal active, peak + item = llm_request.contents[-1].parts[0].text + active += 1 + peak = max(peak, active) + entered.put_nowait(item) + try: + await release.wait() + yield _response(item) + finally: + active -= 1 + + agents = [ + LlmAgent(name=f'worker_{i}', model=GatedModel(responses=[])) + for i in range(4) + ] + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> list[str]: + async with ctx.parallel_tasks() as tasks: + handles = [tasks.spawn(agent, str(i)) for i, agent in enumerate(agents)] + observations['started'] = [await entered.get() for _ in handles] + observations['coordinator_work'] = sum(range(10)) + observations['statuses'] = [handle.status for handle in handles] + release.set() + done, pending = await tasks.wait(mode='all') + assert done == handles + assert pending == [] + observations['outputs'] = [await handle.result() for handle in handles] + return observations['outputs'] + + await _run(coordinate) + + assert set(observations['started']) == {'0', '1', '2', '3'} + assert observations['coordinator_work'] == 45 + assert observations['statuses'] == ['running'] * 4 + assert observations['outputs'] == ['0', '1', '2', '3'] + assert peak == 4 + assert active == 0 + + +async def test_first_drains_model_and_tool_losers_before_returning(): + """A selected result closes losing model/tool work before parent proceeds.""" + model_entered = asyncio.Event() + tool_entered = asyncio.Event() + model_closed = asyncio.Event() + tool_closed = asyncio.Event() + never = asyncio.Event() + calls = [] + observations = {} + + class BlockedModel(testing_utils.MockModel): + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + calls.append('blocked model') + model_entered.set() + try: + await never.wait() + yield _response('must not complete') + finally: + model_closed.set() + + async def blocked_tool() -> str: + calls.append('blocked tool') + tool_entered.set() + try: + await never.wait() + return 'must not complete' + finally: + tool_closed.set() + + tool_model = testing_utils.MockModel.create([ + types.Part.from_function_call(name='blocked_tool', args={}), + 'must not make another model call', + ]) + model_worker = LlmAgent(name='model_worker', model=BlockedModel(responses=[])) + tool_worker = LlmAgent( + name='tool_worker', model=tool_model, tools=[blocked_tool] + ) + + async def accepted() -> str: + await model_entered.wait() + await tool_entered.wait() + return 'accepted' + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + losers = [ + tasks.spawn(model_worker, 'work'), + tasks.spawn(tool_worker, 'work'), + ] + expected = tasks.spawn(accepted) + winner = await tasks.first() + assert winner is expected + assert model_closed.is_set() + assert tool_closed.is_set() + observations['statuses'] = [handle.status for handle in losers] + observations['calls_at_selection'] = list(calls) + return await winner.result() + + await _run(coordinate) + + assert observations['statuses'] == ['cancelled', 'cancelled'] + assert calls == observations['calls_at_selection'] + assert len(tool_model.requests) == 1 + + +@pytest.mark.parametrize('async_predicate', [False, True]) +async def test_first_rejects_fastest_and_selects_next_acceptable( + async_predicate, +): + """Rejected outputs leave other candidates running; both predicate forms work.""" + release_acceptable = asyncio.Event() + blocked_entered = asyncio.Event() + blocked_closed = asyncio.Event() + seen = [] + + async def candidate(node_input: str) -> str: + if node_input == 'rejected': + await blocked_entered.wait() + elif node_input == 'accepted': + await release_acceptable.wait() + else: + blocked_entered.set() + try: + await asyncio.Event().wait() + finally: + blocked_closed.set() + return node_input + + def accept(output): + seen.append(output) + if output == 'rejected': + release_acceptable.set() + return output == 'accepted' + + async def accept_async(output): + return accept(output) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + handles = [ + tasks.spawn(candidate, value) + for value in ['rejected', 'accepted', 'blocked'] + ] + winner = await tasks.first( + accept=accept_async if async_predicate else accept + ) + assert winner is handles[1] + assert blocked_closed.is_set() + return winner.output + + await _run(coordinate) + assert seen == ['rejected', 'accepted'] + + +async def test_first_with_cap_one_never_starts_queued_losers(): + """The winner decision pauses admission before a freed slot starts more work.""" + started = [] + + async def candidate(node_input: str) -> str: + started.append(node_input) + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + handles = [ + tasks.spawn(candidate, value) + for value in ['winner', 'queued1', 'queued2'] + ] + winner = await tasks.first() + assert winner is handles[0] + assert [handle.status for handle in handles] == [ + 'completed', + 'cancelled', + 'cancelled', + ] + return winner.output + + await _run(coordinate) + assert started == ['winner'] + + +async def test_cancelled_result_observer_does_not_cancel_worker(): + """Cancellation of an observer is separate from cancellation of owned work.""" + entered = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + + async def candidate() -> str: + entered.set() + try: + await release.wait() + return 'complete' + finally: + closed.set() + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + handle = tasks.spawn(candidate) + await entered.wait() + observer = asyncio.create_task(handle.result()) + observer.cancel() + with pytest.raises(asyncio.CancelledError): + await observer + assert handle.status == 'running' + assert not closed.is_set() + release.set() + return await handle.result() + + await _run(coordinate) + assert closed.is_set() + + +async def test_completion_ties_select_spawn_order(): + """When every candidate is already complete, first picks the first spawned.""" + + async def candidate(node_input: str) -> str: + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + handles = [tasks.spawn(candidate, value) for value in ['first', 'second']] + await tasks.wait(mode='all') + winner = await tasks.first() + assert winner is handles[0] + assert await tasks.first() is winner + return winner.output + + await _run(coordinate) + + +@pytest.mark.parametrize('all_fail', [False, True]) +async def test_first_tolerates_partial_failure_and_surfaces_all_failure( + all_fail, +): + """A failed candidate cannot prevent a viable candidate from winning.""" + + async def candidate(node_input: str) -> str: + if node_input == 'failure' or all_fail: + raise ValueError(f'{node_input} failed') + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + tasks.spawn(candidate, 'failure') + success = tasks.spawn(candidate, 'success') + winner = await tasks.first() + assert winner is success + return winner.output + + if all_fail: + with pytest.raises(ValueError, match='failure failed'): + await _run(coordinate) + else: + await _run(coordinate) + + +async def test_predicate_error_closes_in_flight_workers(): + """An acceptance callback error propagates only after sibling cleanup.""" + entered = asyncio.Event() + closed = asyncio.Event() + + async def candidate(node_input: str) -> str: + if node_input == 'blocked': + entered.set() + try: + await asyncio.Event().wait() + finally: + closed.set() + await entered.wait() + return node_input + + def accept(output): + raise LookupError('predicate failed') + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + tasks.spawn(candidate, 'blocked') + tasks.spawn(candidate, 'complete') + await tasks.first(accept=accept) + return 'unreachable' + + with pytest.raises(LookupError, match='predicate failed'): + await _run(coordinate) + assert closed.is_set() + + +async def test_parent_cancellation_drains_workers_without_durable_cancel_record(): + """Aborting the invocation cleans up work without recording a user decision.""" + entered: asyncio.Queue[str] = asyncio.Queue() + closed = set() + + async def candidate(node_input: str) -> str: + entered.put_nowait(node_input) + try: + await asyncio.Event().wait() + finally: + closed.add(node_input) + return 'unreachable' + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> list[str]: + async with ctx.parallel_tasks() as tasks: + tasks.spawn(candidate, 'one') + tasks.spawn(candidate, 'two') + return await tasks.gather() + + runner = testing_utils.InMemoryRunner( + node=Workflow(name='team', edges=[(START, coordinate)]) + ) + running = asyncio.create_task(runner.run_async('start')) + try: + assert { + await asyncio.wait_for(entered.get(), _MAX_WAIT) for _ in range(2) + } == {'one', 'two'} + running.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(running, _MAX_WAIT) + assert closed == {'one', 'two'} + assert not any( + key.endswith('parallel_tasks') + for event in runner.session.events + for key in event.custom_metadata or {} + ) + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + await runner.runner.close() + + +async def test_uncooperative_loser_reports_cleanup_timeout(): + """First does not claim clean shutdown when a worker ignores cancellation.""" + entered = asyncio.Event() + cancelled = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + observations = {} + + async def stubborn() -> str: + entered.set() + try: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + await release.wait() + return 'late' + finally: + closed.set() + + async def accepted() -> str: + await entered.wait() + return 'winner' + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks(cancellation_timeout=0.01) as tasks: + loser = tasks.spawn(stubborn) + tasks.spawn(accepted) + try: + with pytest.raises(TimeoutError, match='did not stop within'): + await tasks.first() + observations['cancelled'] = cancelled.is_set() + observations['closed_at_timeout'] = closed.is_set() + finally: + release.set() + await closed.wait() + with pytest.raises(asyncio.CancelledError): + await loser.result() + return 'timeout observed' + + await _run(coordinate) + assert observations == {'cancelled': True, 'closed_at_timeout': False} + assert closed.is_set() + + +async def test_all_rejected_results_raise_clear_error(): + """Rejecting every successful output is an explicit no-winner outcome.""" + + async def candidate(node_input: str) -> str: + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + tasks.spawn(candidate, 'one') + tasks.spawn(candidate, 'two') + await tasks.first(accept=lambda output: False) + return 'unreachable' + + with pytest.raises( + ValueError, match='No parallel task produced an acceptable' + ): + await _run(coordinate) + + +async def test_explicit_cancel_stops_running_worker_and_releases_queued_slot(): + """Awaited cancellation drains a worker and permits the next queued task.""" + entered = asyncio.Event() + closed = asyncio.Event() + started = [] + + async def candidate(node_input: str) -> str: + started.append(node_input) + if node_input == 'blocked': + entered.set() + try: + await asyncio.Event().wait() + finally: + closed.set() + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + blocked = tasks.spawn(candidate, 'blocked') + queued = tasks.spawn(candidate, 'queued') + await entered.wait() + await tasks.cancel(blocked) + assert blocked.status == 'cancelled' + assert closed.is_set() + assert blocked.cancel() is False + with pytest.raises(asyncio.CancelledError): + await blocked.result() + return await queued.result() + + await _run(coordinate) + assert started == ['blocked', 'queued'] + + +def _resumable_runner(coordinate, is_resumable): + return testing_utils.InMemoryRunner( + app=App( + name='team', + root_agent=Workflow(name='team', edges=[(START, coordinate)]), + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + + +async def _resume(runner, events, interrupt_id): + return await asyncio.wait_for( + runner.run_async( + new_message=types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=interrupt_id, + name='user_input', + response={'text': 'approved'}, + ) + ) + ], + ), + invocation_id=events[0].invocation_id, + ), + _MAX_WAIT, + ) + + +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_winner_and_cancelled_losers_survive_downstream_hitl( + is_resumable, +): + """Resume reuses the chosen result without rerunning predicate or losers.""" + calls = [] + predicate_calls = [] + winners = [] + + async def candidate(node_input: str) -> str: + calls.append(node_input) + return node_input + + def accept(output): + predicate_calls.append(output) + return True + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + first = tasks.spawn(candidate, 'winner') + loser = tasks.spawn(candidate, 'queued') + winner = await tasks.first(accept=accept) + assert winner is first + assert loser.status == 'cancelled' + winners.append(winner.id) + if 'approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='approval', message='Approve winner?') + return + yield Event(output=winner.output) + + runner = _resumable_runner(coordinate, is_resumable) + try: + first_events = await asyncio.wait_for(runner.run_async('start'), _MAX_WAIT) + resumed = await _resume(runner, first_events, 'approval') + assert any(event.output == 'winner' for event in resumed) + finally: + await runner.runner.close() + + assert len(winners) == 2 + assert winners[0] == winners[1] + assert calls == ['winner'] + assert predicate_calls == ['winner'] + + +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_merge_disposition_survives_resume_and_preserves_external_edit( + is_resumable, +): + """A committed patch is not reapplied over edits made while awaiting input.""" + calls = [] + + async def proposal() -> dict[str, str]: + calls.append('proposal') + return {'answer': 'selected', 'unselected': 'must not merge'} + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks() as tasks: + task = tasks.spawn(proposal) + await task.result() + await tasks.merge_result(task, keys=['answer']) + if 'approval' not in ctx.resume_inputs: + yield RequestInput( + interrupt_id='approval', message='Approve merged answer?' + ) + return + yield Event(output=ctx.state['answer']) + + runner = _resumable_runner(coordinate, is_resumable) + try: + first_events = await asyncio.wait_for(runner.run_async('start'), _MAX_WAIT) + assert runner.session.state['answer'] == 'selected' + assert 'unselected' not in runner.session.state + await runner.runner.session_service.append_event( + runner.session, + Event( + author='user', + actions=EventActions(state_delta={'answer': 'external edit'}), + ), + ) + resumed = await _resume(runner, first_events, 'approval') + assert any(event.output == 'external edit' for event in resumed) + assert runner.session.state['answer'] == 'external edit' + finally: + await runner.runner.close() + assert calls == ['proposal'] + + +@pytest.mark.parametrize( + 'initial,changed', + [(False, 0), ({'enabled': False}, {'enabled': 0}), ('before', 'after')], +) +async def test_merge_conflicts_are_type_sensitive_and_atomic(initial, changed): + """A conflict rejects every proposed key, including bool/number changes.""" + + async def proposal() -> dict[str, Any]: + return {'new_key': 'must not leak', 'conflict': 'proposal'} + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + ctx.state['conflict'] = initial + async with ctx.parallel_tasks() as tasks: + task = tasks.spawn(proposal) + await task.result() + ctx.state['conflict'] = changed + with pytest.raises(ValueError, match='State changed before merging'): + await tasks.merge_result(task) + assert 'new_key' not in ctx.state + assert ctx.state['conflict'] == changed + await tasks.merge_result( + task, keys=['conflict'], expected={'conflict': changed} + ) + assert ctx.state['conflict'] == 'proposal' + return 'checked' + + await _run(coordinate) + + +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_gather_pauses_queue_for_hitl_and_replays_completed_workers( + is_resumable, +): + """An approval request pauses bounded admission and resumes only pending work.""" + calls = [] + outputs = [] + + @node(rerun_on_resume=True) + async def candidate( + ctx: Context, node_input: str + ) -> AsyncGenerator[Event, None]: + calls.append(node_input) + if node_input == 'approval' and 'worker_approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='worker_approval', message='Continue?') + return + yield Event(output=node_input) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> list[str]: + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + for value in ['complete', 'approval', 'queued']: + tasks.spawn(candidate, value) + result = await tasks.gather() + outputs.append(result) + return result + + runner = _resumable_runner(coordinate, is_resumable) + try: + first_events = await asyncio.wait_for(runner.run_async('start'), _MAX_WAIT) + assert calls == ['complete', 'approval'] + assert outputs == [] + await _resume(runner, first_events, 'worker_approval') + finally: + await runner.runner.close() + + assert calls == ['complete', 'approval', 'approval', 'queued'] + assert outputs == [['complete', 'approval', 'queued']] + + +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_cancelled_waiting_worker_cannot_block_unrelated_resume( + is_resumable, +): + """Cancelling a waiting child clears its interruption across later resumes.""" + calls = [] + results = [] + + @node(rerun_on_resume=True) + async def candidate( + ctx: Context, node_input: str + ) -> AsyncGenerator[Event, None]: + calls.append(node_input) + if node_input == 'cancelled_approval': + yield RequestInput(interrupt_id='unused_approval', message='Not required') + return + yield Event(output=node_input) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + cancelled = tasks.spawn(candidate, 'cancelled_approval') + survivor = tasks.spawn(candidate, 'survivor') + if cancelled.status != 'cancelled': + done, pending = await tasks.wait(mode='all') + assert done == [cancelled] + assert pending == [survivor] + assert cancelled.status == 'waiting' + await tasks.cancel(cancelled) + result = await survivor.result() + results.append(result) + if 'parent_approval' not in ctx.resume_inputs: + yield RequestInput( + interrupt_id='parent_approval', message='Approve survivor?' + ) + return + yield Event(output=result) + + runner = _resumable_runner(coordinate, is_resumable) + try: + first_events = await asyncio.wait_for(runner.run_async('start'), _MAX_WAIT) + assert calls == ['cancelled_approval', 'survivor'] + resumed = await _resume(runner, first_events, 'parent_approval') + assert any(event.output == 'survivor' for event in resumed) + finally: + await runner.runner.close() + + assert calls == ['cancelled_approval', 'survivor'] + assert results == ['survivor', 'survivor'] + + +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_new_invocation_does_not_reuse_previous_winner_or_patch( + is_resumable, +): + """Replay dispositions are scoped to a continuation, not future user turns.""" + calls = [] + predicates = [] + answers = [] + + async def proposal() -> dict[str, int]: + calls.append('run') + return {'answer': len(calls)} + + def accept(output): + predicates.append(output['answer']) + return True + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> int: + async with ctx.parallel_tasks() as tasks: + tasks.spawn(proposal) + winner = await tasks.first(accept=accept) + await tasks.merge_result(winner) + answers.append(ctx.state['answer']) + return ctx.state['answer'] + + runner = _resumable_runner(coordinate, is_resumable) + try: + await asyncio.wait_for(runner.run_async('first task'), _MAX_WAIT) + await asyncio.wait_for(runner.run_async('new task'), _MAX_WAIT) + assert runner.session.state['answer'] == 2 + finally: + await runner.runner.close() + + assert calls == ['run', 'run'] + assert predicates == [1, 2] + assert answers == [1, 2] + + +async def test_merge_validates_all_selected_schema_fields_before_writing(): + """An invalid later field cannot leave a valid earlier field partially merged.""" + + class Schema(BaseModel): + good: str + count: int + + async def proposal() -> dict[str, Any]: + return {'good': 'must not leak', 'count': 'not an integer'} + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + handle = tasks.spawn(proposal) + await handle.result() + with pytest.raises(StateSchemaError, match='count'): + await tasks.merge_result(handle) + assert 'good' not in ctx.state + assert 'count' not in ctx.state + await tasks.merge_result(handle, keys=['good']) + assert ctx.state['good'] == 'must not leak' + return 'checked' + + await _run(coordinate, state_schema=Schema) + + +async def test_merge_rejects_bad_selection_and_foreign_handle(): + """Selection and ownership errors cannot accidentally modify session state.""" + + async def proposal() -> dict[str, int]: + return {'answer': 42} + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks(name='outer') as outer: + handle = outer.spawn(proposal) + await handle.result() + with pytest.raises(ValueError, match='not a string'): + await outer.merge_result(handle, keys='answer') + with pytest.raises(ValueError, match='Every selected key'): + await outer.merge_result(handle, keys=['missing']) + async with ctx.parallel_tasks(name='inner') as inner: + with pytest.raises(ValueError, match='different parallel scope'): + await inner.merge_result(handle) + with pytest.raises(ValueError, match='different parallel scope'): + await inner.cancel(handle) + assert 'answer' not in ctx.state + await outer.merge_result(handle) + ctx.state['answer'] = 99 + await outer.merge_result(handle) + assert ctx.state['answer'] == 99 + return 'checked' + + await _run(coordinate) + + +async def test_first_closes_losing_node_generator_after_progress_event(): + """A losing generator is finalized even after it emitted native progress.""" + progressed = asyncio.Event() + closed = asyncio.Event() + outputs = [] + + @node + async def streaming() -> AsyncGenerator[Event, None]: + try: + yield Event(custom_metadata={'progress': 'started'}) + progressed.set() + await asyncio.Event().wait() + outputs.append('late output') + yield Event(output='late output') + finally: + closed.set() + + async def accepted() -> str: + await progressed.wait() + return 'winner' + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks() as tasks: + loser = tasks.spawn(streaming) + winner = tasks.spawn(accepted) + assert await tasks.first() is winner + assert closed.is_set() + assert loser.status == 'cancelled' + return winner.output + + events = await _run(coordinate) + assert any( + (event.custom_metadata or {}).get('progress') == 'started' + for event in events + ) + assert outputs == [] + + +async def test_handle_cancel_request_prevents_queued_worker_from_starting(): + """Synchronous cancellation requests are persisted before queued admission.""" + entered = asyncio.Event() + release = asyncio.Event() + started = [] + + async def candidate(node_input: str) -> str: + started.append(node_input) + entered.set() + await release.wait() + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> str: + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + running = tasks.spawn(candidate, 'running') + queued = tasks.spawn(candidate, 'cancelled') + await entered.wait() + assert queued.cancel() is True + release.set() + done, pending = await tasks.wait(mode='all') + assert done == [running, queued] + assert pending == [] + assert queued.status == 'cancelled' + return await running.result() + + await _run(coordinate) + assert started == ['running'] + + +@pytest.mark.parametrize('concurrent', [False, True]) +@pytest.mark.parametrize('explicit_run_ids', [False, True]) +async def test_tool_calls_own_distinct_groups_and_current_selection( + concurrent, explicit_run_ids +): + """Each tool invocation owns its task IDs, selection, and merge decisions. + + The same real Agent calls one coordinating FunctionTool twice in a turn. + Each tool uses the default group name and the same candidate names/run IDs, + but must select the output requested by that call's current arguments. + """ + both_entered = asyncio.Event() + entered = [] + ids = [] + selected = {} + predicate_calls = {'a': [], 'b': []} + + async def candidate(node_input: str) -> dict[str, str]: + return {'answer': node_input} + + async def choose(desired: str, ctx: Context) -> str: + """Choose the requested candidate and merge its answer.""" + async with ctx.parallel_tasks(max_parallel_workers=1) as tasks: + handles = [ + tasks.spawn( + candidate, + value, + **({'run_id': f'candidate-{value}'} if explicit_run_ids else {}), + ) + for value in ['a', 'b'] + ] + ids.extend(handle.id for handle in handles) + entered.append(desired) + if len(entered) == 2: + both_entered.set() + if concurrent: + await both_entered.wait() + + def accept(output): + predicate_calls[desired].append(output['answer']) + return output['answer'] == desired + + winner = await tasks.first(accept=accept) + expected = ( + {'answer': ctx.state['answer']} if 'answer' in ctx.state else None + ) + await tasks.merge_result(winner, expected=expected) + selected[desired] = winner.output['answer'] + return winner.output['answer'] + + calls = [ + types.Part.from_function_call(name='choose', args={'desired': desired}) + for desired in ['a', 'b'] + ] + model = testing_utils.MockModel.create( + [calls, 'complete'] if concurrent else [*calls, 'complete'] + ) + agent = LlmAgent(name='coordinator', model=model, tools=[choose]) + runner = testing_utils.InMemoryRunner(root_agent=agent) + try: + events = await asyncio.wait_for(runner.run_async('start'), _MAX_WAIT) + assert selected == {'a': 'a', 'b': 'b'} + assert predicate_calls == {'a': ['a'], 'b': ['a', 'b']} + assert len(ids) == len(set(ids)) == 4 + merged_answers = [ + event.actions.state_delta['answer'] + for event in events + if 'answer' in event.actions.state_delta + ] + assert sorted(merged_answers) == ['a', 'b'] + if not concurrent: + assert runner.session.state['answer'] == 'b' + finally: + both_entered.set() + await runner.runner.close() diff --git a/tests/unittests/workflow/utils/test_replay_manager.py b/tests/unittests/workflow/utils/test_replay_manager.py index 05b8896ac37..28e516dc4e8 100644 --- a/tests/unittests/workflow/utils/test_replay_manager.py +++ b/tests/unittests/workflow/utils/test_replay_manager.py @@ -56,6 +56,7 @@ def _make_event( event.long_running_tool_ids = set(interrupt_ids) if interrupt_ids else None event.content = None event.actions = None + event.custom_metadata = None return event From ac99c09f4265f391bece4da817aeb91d44c89038 Mon Sep 17 00:00:00 2001 From: Meher Gajula Date: Tue, 20 Oct 2026 22:17:35 -0500 Subject: [PATCH 3/4] fix: preserve legacy worker events and hermetic test fixtures --- src/google/adk/workflow/_parallel_tasks.py | 9 ++- tests/unittests/test_samples.py | 12 +++- .../workflow/test_dynamic_node_scheduler.py | 2 + .../workflow/test_parallel_task_replay.py | 57 +++++++++++++++++++ 4 files changed, 78 insertions(+), 2 deletions(-) diff --git a/src/google/adk/workflow/_parallel_tasks.py b/src/google/adk/workflow/_parallel_tasks.py index a551550d02e..39a7f8032c4 100644 --- a/src/google/adk/workflow/_parallel_tasks.py +++ b/src/google/adk/workflow/_parallel_tasks.py @@ -21,6 +21,7 @@ from collections.abc import Callable from collections.abc import Iterable from collections.abc import Mapping +from contextlib import nullcontext import copy import inspect import logging @@ -695,7 +696,13 @@ async def _pump(self) -> None: async def _execute(self, task: ParallelTask) -> None: try: - with _bind_parallel_task(task._control): + # Legacy workers preserve their event format while inheriting any outer + # public task's ownership, including cancellation and transfer tracking. + with ( + nullcontext() + if self._legacy_run_ids + else _bind_parallel_task(task._control) + ): child = await self._ctx._run_node_internal( task._node, task._input, diff --git a/tests/unittests/test_samples.py b/tests/unittests/test_samples.py index 5a6c4c789bb..af41cb33ded 100644 --- a/tests/unittests/test_samples.py +++ b/tests/unittests/test_samples.py @@ -17,7 +17,6 @@ import asyncio import contextlib import importlib -import json import logging import os from pathlib import Path @@ -243,6 +242,8 @@ def test_sample_loads(sample_dir: Path, monkeypatch): """Smoke test: every sample's agent imports and constructs a root agent.""" for key, value in _DUMMY_ENV.items(): monkeypatch.setenv(key, value) + from google.adk.integrations.parameter_manager import parameter_client + from google.adk.integrations.secret_manager import secret_client import google.auth import google.auth.credentials import google.auth.transport @@ -265,6 +266,15 @@ def refresh(self, request: google.auth.transport.Request) -> None: "dummy-project", ), ) + # Integration tests may already have imported these module-level aliases. + # Import them before patching google.auth.default so monkeypatch also restores + # the original aliases when this test was their first importer. + monkeypatch.setattr( + parameter_client, "default_service_credential", google.auth.default + ) + monkeypatch.setattr( + secret_client, "default_service_credential", google.auth.default + ) monkeypatch.setattr( mtls, "has_default_client_cert_source", diff --git a/tests/unittests/workflow/test_dynamic_node_scheduler.py b/tests/unittests/workflow/test_dynamic_node_scheduler.py index b28eea74372..0a35ee35762 100644 --- a/tests/unittests/workflow/test_dynamic_node_scheduler.py +++ b/tests/unittests/workflow/test_dynamic_node_scheduler.py @@ -92,6 +92,7 @@ def _make_event( """Create a minimal Event for session event lists.""" event = MagicMock(spec=Event) event.invocation_id = invocation_id + event.custom_metadata = None event.author = author event.output = output event.error_code = None @@ -112,6 +113,7 @@ def _make_fr_event(fc_id, response, invocation_id='inv-1'): """Create a user FR event.""" event = MagicMock(spec=Event) event.invocation_id = invocation_id + event.custom_metadata = None event.author = 'user' event.output = None event.error_code = None diff --git a/tests/unittests/workflow/test_parallel_task_replay.py b/tests/unittests/workflow/test_parallel_task_replay.py index 4291bd53028..6562cdb0d2d 100644 --- a/tests/unittests/workflow/test_parallel_task_replay.py +++ b/tests/unittests/workflow/test_parallel_task_replay.py @@ -40,6 +40,63 @@ from . import testing_utils +@pytest.mark.parametrize('explicit_group', [False, True]) +async def test_legacy_worker_preserves_event_metadata_and_outer_ownership( + explicit_group: bool, +): + """Legacy workers add no task identity and inherit any public task owner.""" + owner_ids = [] + captured = [] + + @node(parallel_worker=True, max_parallel_workers=1) + async def worker(node_input: str) -> AsyncGenerator[Event, None]: + yield Event( + output=f'{node_input}_processed', + custom_metadata={'source': node_input}, + ) + + def produce() -> list[str]: + return ['first', 'second'] + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> list[str]: + async with ctx.parallel_tasks() as tasks: + task = tasks.spawn(worker, produce()) + owner_ids.append(task.id) + return await task.result() + + def capture(node_input: list[str]) -> list[str]: + captured.append(node_input) + return node_input + + workflow = Workflow( + name='legacy_metadata', + edges=[ + (START, coordinate, capture) + if explicit_group + else (START, produce, worker, capture) + ], + ) + runner = testing_utils.InMemoryRunner(node=workflow) + try: + events = await runner.run_async('start') + worker_events = [ + event for event in events if 'source' in (event.custom_metadata or {}) + ] + assert len(worker_events) == 2 + assert captured == [['first_processed', 'second_processed']] + for event in worker_events: + expected = {'source': event.output.removesuffix('_processed')} + if explicit_group: + expected.update({ + 'parallel_task': {'task_id': owner_ids[0], 'status': 'running'}, + '__adk_internal_parallel_task_owners': owner_ids, + }) + assert event.custom_metadata == expected + finally: + await runner.runner.close() + + @pytest.mark.parametrize('is_resumable', [False, True]) @pytest.mark.parametrize('nested_group', [False, True]) async def test_cancelled_transfer_to_foreign_scheduler_does_not_block_resume( From d36495faae14b7514bf05e9d41fe32fdf5e1298c Mon Sep 17 00:00:00 2001 From: Meher Gajula Date: Tue, 20 Oct 2026 22:38:17 -0500 Subject: [PATCH 4/4] fix(workflow): keep parallel decisions consistent across resume Serialize cancellation persistence and terminal status. Retain completed sibling identities through native replay, free saved-winner capacity from unneeded retries, and require cleanup before repeated winner selection succeeds. --- docs/guides/workflow/parallel_tasks/index.md | 6 + src/google/adk/workflow/_parallel_tasks.py | 243 ++++++--- .../workflow/test_parallel_task_replay.py | 461 ++++++++++++++++++ .../unittests/workflow/test_parallel_tasks.py | 37 +- 4 files changed, 661 insertions(+), 86 deletions(-) diff --git a/docs/guides/workflow/parallel_tasks/index.md b/docs/guides/workflow/parallel_tasks/index.md index 7859a5ea2b8..2ee3e22b46a 100644 --- a/docs/guides/workflow/parallel_tasks/index.md +++ b/docs/guides/workflow/parallel_tasks/index.md @@ -91,6 +91,10 @@ No new queued worker is admitted while a completed candidate is being checked. The winning decision is recorded before siblings are cancelled. Results observed in the same completion batch are checked in spawn order. +Cancellation applies to unfinished siblings. Completed results stay available, +including after resume. Failed nodes keep ADK's native retry behavior; a retry +that is still running when a saved winner is restored is cancelled and drained. + If all successful results are rejected, `first` raises `ValueError`; if no answer is accepted and a worker failed, it propagates a worker failure. Human input propagates when no remaining active worker can provide an accepted answer. @@ -175,6 +179,8 @@ undo tool effects. Keep speculative work free of irreversible side effects. request. Completed or in-flight requests can still be billed. - If an explicit cancellation or winner selection cannot drain a worker before the timeout, it raises `TimeoutError` instead of reporting successful cleanup. + Retrying `first` keeps the selected winner and still requires cleanup to finish; + it does not reset the cancellation deadline. While unwinding another error or parent cancellation, it logs the timeout and preserves that original error. A worker that suppresses cancellation may still be running. Scope exit seals its mailbox and marks remaining messages diff --git a/src/google/adk/workflow/_parallel_tasks.py b/src/google/adk/workflow/_parallel_tasks.py index 39a7f8032c4..a9fc4d828aa 100644 --- a/src/google/adk/workflow/_parallel_tasks.py +++ b/src/google/adk/workflow/_parallel_tasks.py @@ -250,6 +250,7 @@ def __init__( self._cancelled_ids: set[str] = set() self._winner_id: str | None = None self._winner: ParallelTask | None = None + self._completed_ids: set[str] = set() self._state_snapshot: dict[str, Any] = {} self._merged: set[tuple[str, str]] = set() @@ -268,6 +269,7 @@ async def __aenter__(self) -> ParallelTasks: if not isinstance(record, dict) or record.get('group') != self._id: continue self._cancelled_ids.update(record.get('cancelled', [])) + self._completed_ids.update(record.get('completed', [])) for task_id, keys in record.get('merged', {}).items(): self._merged.update((task_id, key) for key in keys) if record.get('winner') is not None: @@ -339,9 +341,7 @@ def spawn( raise ValueError(f'Duplicate parallel task ID: {task.id}') self._ctx._parallel_task_paths.add(task.id) self._tasks.append(task) - if task.id in self._cancelled_ids or ( - self._winner_id is not None and task.id != self._winner_id - ): + if task.id in self._cancelled_ids: task._finish('cancelled') self._wake.set() return task @@ -436,18 +436,25 @@ async def first( self, accept: Callable[[Any], Any] | None = None, ) -> ParallelTask: - """Return the first acceptable result after cancelling and draining siblings. + """Return an acceptable result after cancelling unfinished siblings. The predicate receives the output and may be asynchronous. Rejections and worker failures allow remaining candidates to continue. Completion ties are resolved in spawn order. Raises ValueError when no result is acceptable. + Already completed or failed siblings retain their native replay behavior. + Retrying after a cleanup timeout still requires unfinished siblings to stop. """ if self._selecting: raise RuntimeError('Only one first() selection may run at a time.') - if self._winner is not None: - return self._winner if self._closed: raise RuntimeError('The parallel task scope is closed.') + if self._winner is not None: + await self._drain( + [task for task in self._tasks if task.id in self._cancelled_ids], + deliberate=True, + ) + await self._restore_completed() + return self._winner self._selecting = True try: return await self._first(accept) @@ -465,6 +472,15 @@ async def _first(self, accept: Callable[[Any], Any] | None) -> ParallelTask: ) seen: set[str] = set() try: + if self._winner_id is not None: + # A failed sibling can retry before the saved winner is admitted. Stop + # that retry first so it cannot occupy the winner's bounded slot. + await self._record_selection(self._winner_id) + await self._drain( + [task for task in self._tasks if task.id in self._cancelled_ids], + deliberate=True, + ) + self._wake.set() while True: self._changed.clear() self._raise_pump_error() @@ -481,23 +497,17 @@ async def _first(self, accept: Callable[[Any], Any] | None) -> ParallelTask: if inspect.isawaitable(accepted): accepted = await accepted if accepted: - # Pause admission before persisting the decision, which yields to - # the Runner and must not admit another candidate in the meantime. - self._paused = True - async with self._decision_lock: - if self._winner_id is None: - losers = [ - other.id for other in self._tasks if other is not task - ] - await self._record(cancelled=losers, winner=task.id) - self._winner_id = task.id - self._cancelled_ids.update(losers) - self._winner = task - self._paused = True + await self._record_selection(task.id) + self._winner = task await self._drain( - [other for other in self._tasks if other is not task], + [ + other + for other in self._tasks + if other.id in self._cancelled_ids + ], deliberate=True, ) + await self._restore_completed() return task self._decision_pending = any( task.status == 'completed' and task.id not in seen @@ -507,8 +517,13 @@ async def _first(self, accept: Callable[[Any], Any] | None) -> ParallelTask: if self._decision_pending: continue running = any(task.status == 'running' for task in self._tasks) + restoring = self._winner_id is not None and any( + not task._done.done() + and (task.id == self._winner_id or task.id in self._completed_ids) + for task in self._tasks + ) if all(task._done.done() for task in self._tasks) or ( - self._paused and not running + self._paused and not running and not restoring ): waiting = next( (task for task in self._tasks if task.status == 'waiting'), None @@ -523,9 +538,63 @@ async def _first(self, accept: Callable[[Any], Any] | None) -> ParallelTask: raise ValueError('No parallel task produced an acceptable result.') await self._changed.wait() except BaseException: - await self._drain(self._tasks, deliberate=False, preserve_error=True) + await self._drain( + [ + task + for task in self._tasks + if self._winner_id is None or task.id in self._cancelled_ids + ], + deliberate=False, + preserve_error=True, + ) raise + async def _record_selection(self, winner_id: str) -> None: + # Pause admission before persistence yields to the Runner. Saved completed + # siblings may still hydrate through the same bounded scheduler. + self._paused = True + async with self._decision_lock: + losers = [ + task.id + for task in self._tasks + if task.id != winner_id + and task.status in {'queued', 'running', 'waiting'} + and task.id not in self._completed_ids + ] + completed = [ + task.id + for task in self._tasks + if task.id != winner_id and task.status == 'completed' + ] + if ( + self._winner_id is None + or any(task_id not in self._cancelled_ids for task_id in losers) + or any(task_id not in self._completed_ids for task_id in completed) + ): + await self._record( + cancelled=losers, winner=winner_id, completed=completed + ) + self._winner_id = winner_id + self._cancelled_ids.update(losers) + self._completed_ids.update(completed) + + async def _restore_completed(self) -> None: + """Finish native replay of siblings completed before the saved selection.""" + while True: + self._changed.clear() + self._raise_pump_error() + for task in self._tasks: + if task.id in self._completed_ids and task.status == 'cancelled': + await task.result() + if all( + task._done.done() + for task in self._tasks + if task.id in self._completed_ids + ): + return + self._wake.set() + await self._changed.wait() + async def cancel(self, task: ParallelTask) -> None: """Deliberately cancel one owned task and wait for it to stop.""" self._require_owned(task) @@ -534,16 +603,21 @@ async def cancel(self, task: ParallelTask) -> None: if self._closed: raise RuntimeError('The parallel task scope is closed.') async with self._decision_lock: - await self._record(cancelled=[task.id]) - self._cancelled_ids.add(task.id) - await self._drain([task], deliberate=True) - if self._winner_id is None: - self._paused = any( - other.status == 'waiting' - or (self._fail_fast and other.status == 'failed') - for other in self._tasks - ) - self._wake.set() + # Completion may have won while this caller waited for another decision. + if task.status in {'completed', 'failed', 'cancelled'}: + return + if task.id not in self._cancelled_ids: + await self._record(cancelled=[task.id]) + self._cancelled_ids.add(task.id) + # A finishing worker needs the decision lock to publish its terminal state. + await self._drain([task], deliberate=True) + if self._winner_id is None: + self._paused = any( + other.status == 'waiting' + or (self._fail_fast and other.status == 'failed') + for other in self._tasks + ) + self._wake.set() async def merge_result( self, @@ -629,6 +703,7 @@ async def _record( *, cancelled: list[str], winner: str | None = None, + completed: list[str] | None = None, merged: dict[str, list[str]] | None = None, state_delta: dict[str, Any] | None = None, ) -> None: @@ -637,6 +712,8 @@ async def _record( record: dict[str, Any] = {'group': self._id, 'cancelled': cancelled} if winner is not None: record['winner'] = winner + if completed: + record['completed'] = completed if merged is not None: record['merged'] = merged await self._ctx._invocation_context._enqueue_event( @@ -675,14 +752,18 @@ async def _pump(self) -> None: ) active = sum(task.status == 'running' for task in self._tasks) for task in self._tasks: - if ( - self._paused - or self._decision_pending - or (self._limit is not None and active >= self._limit) - ): + if self._limit is not None and active >= self._limit: break if task.status != 'queued': continue + if self._paused or self._decision_pending: + if self._winner_id is None: + break + if ( + task.id != self._winner_id + and task.id not in self._completed_ids + ): + continue task._status = 'running' task._control.status = 'running' task._runner = asyncio.create_task(self._execute(task)) @@ -695,43 +776,61 @@ async def _pump(self) -> None: self._changed.set() async def _execute(self, task: ParallelTask) -> None: + child: Context | None = None + execution_error: BaseException | None = None try: - # Legacy workers preserve their event format while inheriting any outer - # public task's ownership, including cancellation and transfer tracking. - with ( - nullcontext() - if self._legacy_run_ids - else _bind_parallel_task(task._control) - ): - child = await self._ctx._run_node_internal( - task._node, - task._input, - run_id=task._run_id, - use_sub_branch=True, - return_ctx=True, - skip_run_id_validation=self._legacy_run_ids, - ) - if task.id in self._cancelled_ids: - self._ctx._interrupt_ids.difference_update(child.interrupt_ids) - scheduler = self._ctx._workflow_scheduler - if scheduler is not None: - scheduler._state.interrupt_ids.difference_update(child.interrupt_ids) - task._finish('cancelled') - return - if child.error is not None: - raise DynamicNodeFailError( - message=f'Dynamic node {task.name} failed', - error=child.error, - error_node_path=child.error_node_path, - ) - if child.interrupt_ids: + try: + # Legacy workers preserve their event format while inheriting any outer + # public task's ownership, including cancellation and transfer tracking. + with ( + nullcontext() + if self._legacy_run_ids + else _bind_parallel_task(task._control) + ): + child = await self._ctx._run_node_internal( + task._node, + task._input, + run_id=task._run_id, + use_sub_branch=True, + return_ctx=True, + skip_run_id_validation=self._legacy_run_ids, + ) + # Cancellation can interrupt the following lock acquisition. Expose + # returned HITL IDs first so deliberate drain can always remove them. task._interrupt_ids = set(child.interrupt_ids) - self._ctx._interrupt_ids.update(task._interrupt_ids) - raise NodeInterruptedError() - task._output = child.output - if self._first_mode: - self._decision_pending = True - task._finish('completed') + except asyncio.CancelledError: + raise + except BaseException as error: + execution_error = error + # Persistence yields to the Runner. Serialize terminal publication with + # cancellation decisions so observers and replay see the same outcome. + async with self._decision_lock: + if task.id in self._cancelled_ids: + if child is not None: + self._ctx._interrupt_ids.difference_update(child.interrupt_ids) + scheduler = self._ctx._workflow_scheduler + if scheduler is not None: + scheduler._state.interrupt_ids.difference_update( + child.interrupt_ids + ) + task._finish('cancelled') + return + if execution_error is not None: + raise execution_error + assert child is not None + if child.error is not None: + raise DynamicNodeFailError( + message=f'Dynamic node {task.name} failed', + error=child.error, + error_node_path=child.error_node_path, + ) + if child.interrupt_ids: + self._ctx._interrupt_ids.update(task._interrupt_ids) + raise NodeInterruptedError() + task._output = child.output + if self._first_mode: + self._decision_pending = True + task._finish('completed') except asyncio.CancelledError: task._finish('cancelled') raise diff --git a/tests/unittests/workflow/test_parallel_task_replay.py b/tests/unittests/workflow/test_parallel_task_replay.py index 6562cdb0d2d..6957fb304ce 100644 --- a/tests/unittests/workflow/test_parallel_task_replay.py +++ b/tests/unittests/workflow/test_parallel_task_replay.py @@ -27,11 +27,13 @@ from google.adk.events.event import NodeInfo from google.adk.events.event_actions import EventActions from google.adk.events.request_input import RequestInput +from google.adk.plugins.base_plugin import BasePlugin from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.tools.tool_context import ToolContext from google.adk.workflow import node from google.adk.workflow import START from google.adk.workflow import Workflow +from google.adk.workflow._errors import DynamicNodeFailError from google.adk.workflow._errors import NodeInterruptedError from google.adk.workflow.utils._replay_manager import ReplayManager from google.genai import types @@ -40,6 +42,116 @@ from . import testing_utils +@pytest.mark.parametrize('is_resumable', [False, True]) +@pytest.mark.parametrize('cancel_from_callback', [False, True]) +@pytest.mark.parametrize('final_interruption', [False, True]) +async def test_cancel_at_final_output_agrees_with_replayed_status( + is_resumable: bool, + cancel_from_callback: bool, + final_interruption: bool, +): + """A cancellation racing final output has one live and replay disposition. + + Setup: Hold the worker's final event until cancellation can be requested. + Act: Cancel from the progress callback or await group cancellation, then + pause for human input and resume. + Assert: Both outcomes agree with the durable decision; the worker runs once. + """ + observations = {'calls': [], 'results': [], 'statuses': [], 'requests': []} + final_output_seen = asyncio.Event() + release_final_output = asyncio.Event() + + class CancelAtOutput(BasePlugin): + + async def on_event_callback(self, *, invocation_context, event): + at_boundary = event.output == 'done' or 'unused_approval' in ( + event.long_running_tool_ids or set() + ) + if at_boundary and 'parallel_task' in (event.custom_metadata or {}): + if cancel_from_callback: + observations['requests'].append(observations['task'].cancel()) + else: + final_output_seen.set() + await release_final_output.wait() + return None + + @node(rerun_on_resume=True) + async def worker() -> AsyncGenerator[Event, None]: + observations['calls'].append('worker') + if final_interruption: + yield RequestInput(interrupt_id='unused_approval', message='Not needed') + return + yield Event(output='done') + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks() as group: + task = group.spawn(worker) + observations['task'] = task + if not cancel_from_callback and 'approval' not in ctx.resume_inputs: + await final_output_seen.wait() + release_final_output.set() + await group.cancel(task) + try: + outcome = await task.result() + except asyncio.CancelledError: + outcome = 'cancelled' + observations['results'].append(outcome) + observations['statuses'].append(task.status) + assert not ctx.interrupt_ids + if 'approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='approval', message='Continue?') + return + yield Event(output=outcome) + + runner = testing_utils.InMemoryRunner( + app=App( + name='cancellation_race', + root_agent=Workflow(name='team', edges=[(START, coordinate)]), + plugins=[CancelAtOutput(name='cancel_at_output')], + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 3) + assert observations['requests'] == ([True] if cancel_from_callback else []) + cancelled = any( + observations['task'].id + in (event.custom_metadata or {}) + .get('__adk_internal_parallel_tasks', {}) + .get('cancelled', []) + for event in runner.session.events + ) + expected_result = 'cancelled' if cancelled else 'done' + expected_status = 'cancelled' if cancelled else 'completed' + if final_interruption: + assert cancelled + assert observations['results'] == [expected_result] + assert observations['statuses'] == [expected_status] + response = types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id='approval', + name='adk_request_input', + response={'result': 'yes'}, + ) + ) + ], + ) + resumed = await asyncio.wait_for( + runner.run_async(response, invocation_id=events[0].invocation_id), 3 + ) + assert any(event.output == expected_result for event in resumed) + finally: + await runner.runner.close() + + assert observations['calls'] == ['worker'] + assert observations['results'] == [expected_result, expected_result] + assert observations['statuses'] == [expected_status, expected_status] + + @pytest.mark.parametrize('explicit_group', [False, True]) async def test_legacy_worker_preserves_event_metadata_and_outer_ownership( explicit_group: bool, @@ -97,6 +209,355 @@ def capture(node_input: list[str]) -> list[str]: await runner.runner.close() +@pytest.mark.parametrize('is_resumable', [False, True]) +@pytest.mark.parametrize('other_fails', [False, True]) +@pytest.mark.parametrize('observe_after_first', [False, True]) +async def test_first_preserves_previously_observed_nonwinner_on_resume( + is_resumable: bool, + other_fails: bool, + observe_after_first: bool, +): + """Selection preserves observed terminal results and failures across replay. + + Setup: Inspect a terminal candidate before or after selecting another winner. + Act: Select, pause for human input, and resume the same coordinator. + Assert: Terminal status stays consistent; completed work is reused and + failed work retains the native scheduler's retry behavior. + """ + calls = [] + observed = [] + + async def worker(node_input: str) -> str: + calls.append(node_input) + if other_fails and node_input == 'other': + raise ValueError('Rejected candidate') + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks(max_parallel_workers=1) as group: + group.spawn(worker, 'winner') + other = group.spawn(worker, 'other') + await group.wait(mode='all') + if observe_after_first: + winner = await group.first(lambda output: output == 'winner') + try: + result = await other.result() + except DynamicNodeFailError: + result = 'failed' + observed.append((result, other.status)) + if not observe_after_first: + winner = await group.first(lambda output: output == 'winner') + assert other.status == ('failed' if other_fails else 'completed') + if 'approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='approval', message='Continue?') + return + yield Event(output=winner.output) + + runner = testing_utils.InMemoryRunner( + app=App( + name='completed_candidate', + root_agent=Workflow(name='team', edges=[(START, coordinate)]), + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + expected = ('failed', 'failed') if other_fails else ('other', 'completed') + try: + events = await asyncio.wait_for(runner.run_async('start'), 3) + assert observed == [expected] + response = types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id='approval', + name='adk_request_input', + response={'result': 'yes'}, + ) + ) + ], + ) + resumed = await asyncio.wait_for( + runner.run_async(response, invocation_id=events[0].invocation_id), 3 + ) + assert any(event.output == 'winner' for event in resumed) + finally: + await runner.runner.close() + + # Native replay retries failed nodes; completed nodes reuse their output. + assert calls == ['winner', 'other'] + (['other'] if other_fails else []) + assert observed == [expected, expected] + + +@pytest.mark.parametrize('is_resumable', [False, True]) +@pytest.mark.parametrize('retry_before_winner', [False, True]) +async def test_restored_winner_drains_a_failed_siblings_unfinished_retry( + is_resumable: bool, + retry_before_winner: bool, +): + """Restoring a winner cancels an unfinished retry before returning it.""" + calls = [] + retry_started = asyncio.Event() + retry_closed = asyncio.Event() + + async def worker(node_input: str) -> str: + calls.append(node_input) + if node_input == 'other': + if calls.count('other') == 1: + raise ValueError('Candidate failed on its first attempt') + retry_started.set() + try: + await asyncio.Event().wait() + finally: + retry_closed.set() + return node_input + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks( + max_parallel_workers=1 if retry_before_winner else 2 + ) as group: + if retry_before_winner: + other = group.spawn(worker, 'other') + group.spawn(worker, 'winner') + else: + group.spawn(worker, 'winner') + other = group.spawn(worker, 'other') + if 'approval' in ctx.resume_inputs: + await retry_started.wait() + else: + await group.wait(mode='all') + assert other.status == 'failed' + winner = await group.first(lambda output: output == 'winner') + if 'approval' in ctx.resume_inputs: + assert retry_closed.is_set() + assert other.status == 'cancelled' + if 'approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='approval', message='Continue?') + return + yield Event(output=winner.output) + + runner = testing_utils.InMemoryRunner( + app=App( + name='retried_candidate', + root_agent=Workflow(name='team', edges=[(START, coordinate)]), + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 3) + response = types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id='approval', + name='adk_request_input', + response={'result': 'yes'}, + ) + ) + ], + ) + resumed = await asyncio.wait_for( + runner.run_async(response, invocation_id=events[0].invocation_id), 3 + ) + assert any(event.output == 'winner' for event in resumed) + finally: + await runner.runner.close() + + assert calls == ( + ['other', 'winner', 'other'] + if retry_before_winner + else ['winner', 'other', 'other'] + ) + assert retry_closed.is_set() + + +@pytest.mark.parametrize('is_resumable', [False, True]) +@pytest.mark.parametrize('workflow_sibling', [False, True]) +async def test_restored_winner_retains_completed_sibling_with_bounded_admission( + is_resumable: bool, + workflow_sibling: bool, +): + """A retained sibling queued behind the restored winner hydrates normally.""" + calls = [] + observed = [] + + async def worker(node_input: str) -> str: + calls.append(node_input) + return node_input + + other_node = ( + Workflow(name='other_workflow', edges=[(START, worker)]) + if workflow_sibling + else worker + ) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks(max_parallel_workers=1) as group: + group.spawn(worker, 'winner') + other = group.spawn(other_node, 'other') + if 'approval' not in ctx.resume_inputs: + await group.wait(mode='all') + winner = await group.first(lambda output: output == 'winner') + # Resume enters first() before the second handle can hydrate. Selection + # must preserve this known completed sibling and admit it through the cap. + assert other.status == 'completed' + observed.append(await other.result()) + if 'approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='approval', message='Continue?') + return + yield Event(output=winner.output) + + runner = testing_utils.InMemoryRunner( + app=App( + name='retained_candidate', + root_agent=Workflow(name='team', edges=[(START, coordinate)]), + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 3) + response = types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id='approval', + name='adk_request_input', + response={'result': 'yes'}, + ) + ) + ], + ) + resumed = await asyncio.wait_for( + runner.run_async(response, invocation_id=events[0].invocation_id), 3 + ) + assert any(event.output == 'winner' for event in resumed) + finally: + await runner.runner.close() + + assert calls == ['winner', 'other'] + assert observed == ['other', 'other'] + + +async def test_selection_timeout_preserves_retained_workflow_until_retry(): + """Timing out a losing retry must not cancel retained native hydration.""" + calls = [] + resuming = asyncio.Event() + retry_started = asyncio.Event() + retry_closed = asyncio.Event() + release_retry = asyncio.Event() + hydration_started = asyncio.Event() + hydration_cancelled = asyncio.Event() + release_hydration = asyncio.Event() + + async def worker(node_input: str) -> str: + calls.append(node_input) + if node_input == 'failed': + if calls.count('failed') == 1: + raise ValueError('Candidate failed on its first attempt') + retry_started.set() + try: + while not release_retry.is_set(): + try: + await release_retry.wait() + except asyncio.CancelledError: + pass + finally: + retry_closed.set() + return node_input + + class GatedWorkflow(Workflow): + + async def _run_impl(self, *, ctx, node_input): + if resuming.is_set(): + hydration_started.set() + try: + await release_hydration.wait() + except asyncio.CancelledError: + hydration_cancelled.set() + raise + async for event in super()._run_impl(ctx=ctx, node_input=node_input): + yield event + + retained_node = GatedWorkflow(name='retained', edges=[(START, worker)]) + + @node(rerun_on_resume=True) + async def coordinate(ctx: Context) -> AsyncGenerator[Event, None]: + async with ctx.parallel_tasks(cancellation_timeout=0.01) as group: + group.spawn(worker, 'winner') + failed = group.spawn(worker, 'failed') + retained = group.spawn(retained_node, 'retained') + if not resuming.is_set(): + await group.wait(mode='all') + winner = await group.first(lambda output: output == 'winner') + else: + selection = None + try: + await retry_started.wait() + await hydration_started.wait() + with pytest.raises(TimeoutError, match='did not stop within'): + await group.first() + assert retained.status == 'running' + assert not hydration_cancelled.is_set() + release_retry.set() + await retry_closed.wait() + with pytest.raises(asyncio.CancelledError): + await failed.result() + selection = asyncio.create_task(group.first()) + await asyncio.sleep(0) + assert not selection.done() + release_hydration.set() + winner = await selection + assert await retained.result() == 'retained' + assert retained.status == 'completed' + assert not hydration_cancelled.is_set() + finally: + release_retry.set() + release_hydration.set() + if selection is not None: + await asyncio.gather(selection, return_exceptions=True) + if 'approval' not in ctx.resume_inputs: + yield RequestInput(interrupt_id='approval', message='Continue?') + return + yield Event(output=winner.output) + + runner = testing_utils.InMemoryRunner( + app=App( + name='retry_retained', + root_agent=Workflow(name='team', edges=[(START, coordinate)]), + resumability_config=ResumabilityConfig(is_resumable=True), + ) + ) + try: + events = await asyncio.wait_for(runner.run_async('start'), 3) + resuming.set() + response = types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id='approval', + name='adk_request_input', + response={'result': 'yes'}, + ) + ) + ], + ) + resumed = await asyncio.wait_for( + runner.run_async(response, invocation_id=events[0].invocation_id), 3 + ) + assert any(event.output == 'winner' for event in resumed) + finally: + release_retry.set() + release_hydration.set() + await runner.runner.close() + + assert calls == ['winner', 'failed', 'retained', 'failed'] + + @pytest.mark.parametrize('is_resumable', [False, True]) @pytest.mark.parametrize('nested_group', [False, True]) async def test_cancelled_transfer_to_foreign_scheduler_does_not_block_resume( diff --git a/tests/unittests/workflow/test_parallel_tasks.py b/tests/unittests/workflow/test_parallel_tasks.py index e4f9e241278..f3fe8852e19 100644 --- a/tests/unittests/workflow/test_parallel_tasks.py +++ b/tests/unittests/workflow/test_parallel_tasks.py @@ -412,22 +412,21 @@ async def coordinate(ctx: Context) -> list[str]: await runner.runner.close() -async def test_uncooperative_loser_reports_cleanup_timeout(): - """First does not claim clean shutdown when a worker ignores cancellation.""" +async def test_first_retry_waits_for_cleanup_after_timeout(): + """A cached winner cannot bypass incomplete cleanup or renew its deadline.""" entered = asyncio.Event() - cancelled = asyncio.Event() release = asyncio.Event() closed = asyncio.Event() - observations = {} + cancellations = [] async def stubborn() -> str: entered.set() try: - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - cancelled.set() - await release.wait() + while not release.is_set(): + try: + await release.wait() + except asyncio.CancelledError: + cancellations.append('cancelled') return 'late' finally: closed.set() @@ -440,21 +439,31 @@ async def accepted() -> str: async def coordinate(ctx: Context) -> str: async with ctx.parallel_tasks(cancellation_timeout=0.01) as tasks: loser = tasks.spawn(stubborn) - tasks.spawn(accepted) + winner = tasks.spawn(accepted) try: with pytest.raises(TimeoutError, match='did not stop within'): await tasks.first() - observations['cancelled'] = cancelled.is_set() - observations['closed_at_timeout'] = closed.is_set() + assert loser.status == 'running' + assert not closed.is_set() + deadline = loser._cancel_deadline + assert deadline is not None + with pytest.raises(TimeoutError, match='did not stop within'): + await tasks.first() + assert loser.status == 'running' + assert not closed.is_set() + assert loser._cancel_deadline == deadline + assert cancellations == ['cancelled'] finally: release.set() await closed.wait() with pytest.raises(asyncio.CancelledError): await loser.result() - return 'timeout observed' + assert await tasks.first() is winner + assert loser.status == 'cancelled' + assert cancellations == ['cancelled'] + return winner.output await _run(coordinate) - assert observations == {'cancelled': True, 'closed_at_timeout': False} assert closed.is_set()