From 3a8192b39fb0e43304f1f3734933f8d70d592c8f Mon Sep 17 00:00:00 2001 From: Thuongvu Ho Date: Tue, 6 Oct 2026 08:52:37 -0700 Subject: [PATCH 1/2] fix(workflow): keep finished parallel tool results when a sibling raises When the model emits parallel function calls and one FunctionTool raises, the batch is cancelled and the error re-raised, but the response of a sibling that had already finished was thrown away with it. The call was then stripped from the next request, so the model could repeat it. The batch executor now keeps the finished calls' real responses, and once the node has stopped on the error, the node runner persists them before the error event or a retry; the original error propagates unchanged. Nothing is synthesized, and responses that carry control actions or confirmation/credential placeholders are not kept. --- src/google/adk/agents/invocation_context.py | 13 + .../flows/llm_flows/tools/_batch_executor.py | 63 +++- src/google/adk/workflow/_node_runner.py | 28 ++ .../llm_flows/tools/test_batch_executor.py | 71 +++++ tests/unittests/test_runners.py | 282 ++++++++++++++++++ 5 files changed, 456 insertions(+), 1 deletion(-) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 3984d51f0bf..797feb1d61d 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -100,6 +100,12 @@ class _AbortState: ``asyncio.Event`` and can only be awaited on ``loop``, so AgentTool does not pass it when the tool runs on another event loop (e.g. RunConfig's tool thread pool); such a sub-run is not cancelled by a caller abort. + + ``unpersisted_function_responses`` holds the merged response events of + parallel tool calls that finished before their batch ended early because a + sibling call raised. The batch executor cannot persist them itself, and the + error that ends the batch unwinds the agent; keeping them on this shared + instance lets the node runner persist them once the agent has stopped. """ def __init__( @@ -111,6 +117,13 @@ def __init__( self.loop = loop self.aborted = False self.event_synthesized = False + self.unpersisted_function_responses: list[Event] = [] + + def pop_unpersisted_function_responses(self) -> list[Event]: + """Returns and clears the kept function responses not yet persisted.""" + responses = self.unpersisted_function_responses + self.unpersisted_function_responses = [] + return responses def __deepcopy__(self, memo: dict[int, Any] | None) -> _AbortState: # Preserve single-instance sharing across deepcopies and avoid traversing 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 98e15aa6f78..5afc1abaae9 100644 --- a/src/google/adk/flows/llm_flows/tools/_batch_executor.py +++ b/src/google/adk/flows/llm_flows/tools/_batch_executor.py @@ -27,6 +27,7 @@ from google.genai import types +from ....agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from ....events.event import Event from ....events.event_actions import EventActions from ....telemetry.tracing import trace_merged_tool_calls @@ -251,7 +252,14 @@ async def _execute_prepared_function_calls( ] # Wait for all tasks to complete - maybe_function_response_events = await _gather_or_cancel(tasks) + try: + maybe_function_response_events = await _gather_or_cancel(tasks) + except Exception: + # A call raised and its unfinished siblings were cancelled. Calls that had + # already finished may have had side effects, so keep their real results + # for the node runner to persist once the error has unwound the agent. + _keep_completed_responses(invocation_context, tasks) + raise # Filter out None results function_response_events = [ @@ -270,6 +278,59 @@ async def _execute_prepared_function_calls( ) +def _keep_completed_responses( + invocation_context: InvocationContext, + tasks: list[asyncio.Task[Optional[Event]]], +) -> None: + """Keeps the responses of calls that finished before their batch ended early. + + Only plain tool results are kept: a call that was cancelled, raised or + returned no event gets nothing. They are merged as a complete batch would + be and left on the invocation's shared state for the node runner to + persist. Live mode is skipped. + """ + if invocation_context.live_request_queue is not None: + return + finished = [ + t.result() + for t in tasks + if t.done() and not t.cancelled() and t.exception() is None + ] + completed = [ + event + for event in finished + if event is not None and _is_plain_result(event) + ] + if completed: + invocation_context._abort_state.unpersisted_function_responses.append( # pylint: disable=protected-access + _merge_and_trace_function_response_events(invocation_context, completed) + ) + + +def _is_plain_result(event: Event) -> bool: + """True if the event is a tool result with no control-flow actions. + + Besides the tool's state and artifact deltas, any action (a transfer, an + escalation, a confirmation or credential request whose tool has not run) + is an instruction the failed step never carried out, so it is not kept. + finish_task sets no action and is detected by name. + """ + actions = event.actions.model_copy( + update={ + 'state_delta': {}, + 'artifact_delta': {}, + 'skip_summarization': None, + 'render_ui_widgets': None, + } + ) + if actions != EventActions(): + return False + return not any( + response.name == FINISH_TASK_TOOL_NAME + for response in event.get_function_responses() + ) + + async def _execute_prepared_function_calls_async( invocation_context: InvocationContext, prepared_calls: list[_PreparedFunctionCall], diff --git a/src/google/adk/workflow/_node_runner.py b/src/google/adk/workflow/_node_runner.py index a860c34a0a7..b38e13759a4 100644 --- a/src/google/adk/workflow/_node_runner.py +++ b/src/google/adk/workflow/_node_runner.py @@ -150,6 +150,10 @@ async def run( logger.debug("node %s cancelled via signal.", ctx.node_path) raise except Exception as e: + # Before the DynamicNodeFailError return: a tool that runs a failing + # child node raises it out of this node's own tool batch. + await self._enqueue_kept_function_responses(ctx) + if isinstance(e, DynamicNodeFailError): # TODO: consider to retry upon dynamic node failures later. This may # require thorough design to consider a workflow dynamic node and a @@ -185,6 +189,30 @@ async def run( ) attempt_count += 1 + async def _enqueue_kept_function_responses(self, ctx: Context) -> None: + """Enqueues the results kept from a parallel tool batch that failed. + + The node has stopped, so nothing in it can act on them. They go out before + the error event, and before a retry builds its next request. Once the + invocation is aborted, abort sealing decides how its calls are answered. + """ + ic = ctx._invocation_context + if ic.is_aborted: + return + kept = ic._abort_state.pop_unpersisted_function_responses() # pylint: disable=protected-access + if not kept: + return + # ctx.event_author is the author of the node's last yielded event, which + # in a ParallelAgent node can be another branch. Keep each result's own + # author, as BaseAgent._run_impl does for the events it yields. + event_author = ctx.event_author + try: + for event in kept: + ctx.event_author = event.author + await self._enqueue_event(event, ctx) + finally: + ctx.event_author = event_author + async def _attempt_retry(self, e: Exception, attempt_count: int) -> bool: """Checks if node should retry and sleeps if so.""" from ._node_state import NodeState diff --git a/tests/unittests/flows/llm_flows/tools/test_batch_executor.py b/tests/unittests/flows/llm_flows/tools/test_batch_executor.py index b58dcb9b93d..e5daeb442c9 100644 --- a/tests/unittests/flows/llm_flows/tools/test_batch_executor.py +++ b/tests/unittests/flows/llm_flows/tools/test_batch_executor.py @@ -20,11 +20,14 @@ import contextvars from unittest import mock +from google.adk.agents.invocation_context import _AbortState +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from google.adk.events.event import Event from google.adk.events.event_actions import EventActions from google.adk.flows.llm_flows.tools import _batch_executor as _batch_tool_executor from google.adk.flows.llm_flows.tools._caller import _PreparedFunctionCall from google.adk.tools.base_tool import BaseTool +from google.adk.tools.tool_confirmation import ToolConfirmation from google.adk.tools.tool_context import ToolContext from google.genai import types import pytest @@ -227,6 +230,74 @@ async def test_execute_prepared_function_calls_sets_live_session_id_and_merges() assert [p.text for p in merged.content.parts] == ['r1', 'r2'] +def _response_event(name: str, **actions) -> Event: + return Event( + invocation_id='inv-1', + author='agent', + content=types.Content( + role='user', + parts=[types.Part.from_function_response(name=name, response={})], + ), + actions=EventActions(**actions), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'name, actions, kept', + [ + ('plain', {}, True), + ('state_delta', {'state_delta': {'k': 'v'}}, True), + ('transfer', {'transfer_to_agent': 'other'}, False), + ('escalate', {'escalate': True}, False), + ( + 'confirmation', + {'requested_tool_confirmations': {'id': ToolConfirmation()}}, + False, + ), + ('set_model_response', {'set_model_response': {'a': 1}}, False), + (FINISH_TASK_TOOL_NAME, {}, False), + ], +) +async def test_keep_completed_responses_keeps_only_finished_plain_results( + name: str, actions: dict, kept: bool +) -> None: + """Finished plain results are merged and kept; control actions are not.""" + mock_ic = mock.MagicMock() + mock_ic.session.state = {} + mock_ic.live_request_queue = None + mock_ic._abort_state = _AbortState() + + async def returns(event): + return event + + async def raises(): + raise RuntimeError('boom') + + tasks = [ + asyncio.create_task(returns(_response_event('fast'))), + asyncio.create_task(asyncio.Event().wait()), + asyncio.create_task(raises()), + asyncio.create_task(returns(None)), + asyncio.create_task(returns(_response_event(name, **actions))), + ] + await asyncio.sleep(0) + tasks[1].cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + # The cancelled call, the call that raised and the None result get nothing. + _batch_tool_executor._keep_completed_responses(mock_ic, tasks) + [merged] = mock_ic._abort_state.pop_unpersisted_function_responses() + assert [r.name for r in merged.get_function_responses()] == ( + ['fast', name] if kept else ['fast'] + ) + + # Live mode keeps nothing: there is nothing that would persist it. + mock_ic.live_request_queue = mock.MagicMock() + _batch_tool_executor._keep_completed_responses(mock_ic, tasks) + assert not mock_ic._abort_state.unpersisted_function_responses + + @pytest.mark.asyncio async def test_handle_function_calls_live_compat_wrapper() -> None: mock_context = mock.MagicMock() diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 5e312dcc341..820dd3387b6 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -16,6 +16,7 @@ from contextlib import aclosing import importlib import logging +import operator from pathlib import Path import sys import textwrap @@ -36,6 +37,7 @@ from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from google.adk.agents.llm_agent import LlmAgent +from google.adk.agents.parallel_agent import ParallelAgent from google.adk.agents.remote_a2a_agent import RemoteA2aAgent from google.adk.agents.run_config import RunConfig from google.adk.apps.app import App @@ -57,6 +59,8 @@ from google.adk.tools.base_toolset import BaseToolset from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.workflow import BaseNode +from google.adk.workflow import node +from google.adk.workflow import RetryConfig from google.adk.workflow import START from google.adk.workflow._workflow import Workflow from google.genai import types @@ -64,6 +68,7 @@ import pytest from tests.unittests import testing_utils +from tests.unittests._invariants import InvariantPlugin TEST_APP_ID = "test_app" TEST_USER_ID = "test_user" @@ -6253,5 +6258,282 @@ async def _run_async_impl( assert events[0].branch == "root_agent.sub_agent" +def _gated_tools(gate_on=("fast_tool",), before_raise=None): + """Returns fast_tool, failing_tool and an after_tool_callback. + + failing_tool raises once every call named in `gate_on` has finished. + """ + finished = [] + gate = asyncio.Event() + + async def fast_tool() -> dict[str, str]: + return {"record_id": "r1"} + + async def failing_tool() -> dict[str, str]: + await gate.wait() + # The callback runs at the end of a call; one more step lets it finish. + await asyncio.sleep(0) + if before_raise: + before_raise() + raise RuntimeError("downstream 500") from KeyError("upstream") + + def after_tool(tool, args, tool_context, tool_response): + if tool.name in gate_on: + finished.append(tool.name) + if len(finished) == len(gate_on): + gate.set() + + return fast_tool, failing_tool, after_tool + + +def _parallel_calls_agent(tools, after_tool, name="tool_agent", **kwargs): + """An agent whose model calls every tool in parallel, then answers.""" + return LlmAgent( + name=name, + model=testing_utils.MockModel.create( + responses=[ + [ + types.Part.from_function_call(name=tool.__name__, args={}) + for tool in tools + ], + "Recovered", + ] + ), + tools=tools, + after_tool_callback=after_tool, + **kwargs, + ) + + +def _run_async(runner: Runner, **kwargs): + return runner.run_async( + user_id=TEST_USER_ID, + session_id="s", + new_message=types.Content(role="user", parts=[types.Part(text="Run")]), + **kwargs, + ) + + +def _function_responses(contents) -> list[Any]: + return [ + p.function_response.response + for c in contents + if c + for p in c.parts or [] + if p.function_response + ] + + +async def _persisted_events(runner: Runner) -> list[Event]: + session = await runner.session_service.get_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id="s" + ) + return session.events + + +async def _persisted_responses(runner: Runner) -> list[Any]: + return _function_responses(e.content for e in await _persisted_events(runner)) + + +async def _kept_authors(runner: Runner, tool_name: str) -> list[str]: + """Returns the kept results' authors, checking each is filed like its call.""" + events = await _persisted_events(runner) + where = operator.attrgetter( + "author", "branch", "isolation_scope", "node_info.path" + ) + calls = { + fc.id: where(e) + for e in events + for fc in e.get_function_calls() + if fc.name == tool_name + } + kept = { + fr.id: where(e) + for e in events + for fr in e.get_function_responses() + if fr.name == tool_name + } + assert kept == calls + return sorted(author for author, *_ in kept.values()) + + +@pytest.mark.asyncio +async def test_run_async_tool_error_keeps_result_of_completed_parallel_call(): + """A parallel call that raises keeps a finished sibling's real result. + + The error, its cause and the run callbacks are unchanged, and the next + request shows the model that the finished call succeeded. + """ + fast_tool, failing_tool, after_tool = _gated_tools() + agent = _parallel_calls_agent([fast_tool, failing_tool], after_tool) + recorder = BasePlugin(name="recorder") + recorder.on_run_error_callback = AsyncMock(return_value=None) + recorder.after_run_callback = AsyncMock(return_value=None) + runner = _abort_runner(agent, plugins=[InvariantPlugin(), recorder]) + + with pytest.raises(RuntimeError, match="downstream 500") as exc_info: + await _run_turn(runner, "s", "Run") + + assert await _persisted_responses(runner) == [{"record_id": "r1"}] + assert isinstance(exc_info.value.__cause__, KeyError) + recorder.on_run_error_callback.assert_awaited_once_with( + invocation_context=mock.ANY, error=exc_info.value + ) + recorder.after_run_callback.assert_not_awaited() + + await _run_turn(runner, "s", "Next question") + assert _function_responses(agent.model.requests[-1].contents) == [ + {"record_id": "r1"} + ] + + +@pytest.mark.asyncio +async def test_run_async_tool_error_as_abort_lands_leaves_calls_to_abort_sealing(): + """A batch that fails in the step an abort lands is left to abort sealing.""" + abort_signal = asyncio.Event() + fast_tool, failing_tool, after_tool = _gated_tools( + before_raise=abort_signal.set + ) + runner = _abort_runner( + _parallel_calls_agent([fast_tool, failing_tool], after_tool) + ) + + # An aborted run ends without re-raising the node's error. + async for _ in _run_async(runner, abort_signal=abort_signal): + pass + + aborted = {"error": "Invocation was aborted by client."} + assert await _persisted_responses(runner) == [aborted, aborted] + + +@pytest.mark.asyncio +async def test_run_async_tool_error_retried_agent_sees_kept_result(): + """A node retried after the error sees the finished call's result.""" + fast_tool, failing_tool, after_tool = _gated_tools() + agent = _parallel_calls_agent( + [fast_tool, failing_tool], + after_tool, + retry_config=RetryConfig(max_attempts=2, initial_delay=0, jitter=0), + ) + # No InvariantPlugin: the failing call stays unanswered in a completed turn, + # as both calls do on main. + runner = _abort_runner(agent) + + await _run_turn(runner, "s", "Run") + + assert _function_responses(agent.model.requests[1].contents) == [ + {"record_id": "r1"} + ] + + +@pytest.mark.asyncio +async def test_run_async_tool_error_in_parallel_agent_node_keeps_each_author(): + """Each branch's kept result keeps its author, not the node's last one.""" + fast_tool, failing_tool, after_tool = _gated_tools( + gate_on=("fast_tool", "fast_tool") + ) + parallel = ParallelAgent( + name="parallel", + sub_agents=[ + _parallel_calls_agent([fast_tool, failing_tool], after_tool, name) + for name in ("a", "b") + ], + ) + runner = _abort_runner(Workflow(name="wf", edges=[(START, parallel)])) + + with pytest.raises(RuntimeError, match="downstream 500"): + await _run_turn(runner, "s", "Run") + + assert await _kept_authors(runner, "fast_tool") == ["a", "b"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("caught_by_parent", [False, True]) +async def test_run_async_tool_error_from_failing_child_node_keeps_result( + caught_by_parent: bool, +): + """A tool whose child node fails keeps its sibling's result in the agent. + + The child's failure leaves the agent's node as a DynamicNodeFailError. A + parent node that catches it and fails later must not take the result over. + """ + fast_tool, failing_tool, after_tool = _gated_tools() + + async def child_tool(tool_context) -> dict[str, str]: + return await tool_context.run_node(failing_tool) + + agent = _parallel_calls_agent([fast_tool, child_tool], after_tool) + + async def parent(ctx: Context): + try: + await ctx.run_node(agent, node_input="Run") + except Exception: # pylint: disable=broad-exception-caught + pass # The parent carries on after the agent fails, then fails itself. + raise ValueError("parent failed") + + parent_wf = Workflow( + name="wf", + edges=[(START, node(parent, name="parent", rerun_on_resume=True))], + ) + runner = _abort_runner(parent_wf if caught_by_parent else agent) + + with pytest.raises( + Exception, match="parent failed" if caught_by_parent else "downstream 500" + ): + await _run_turn(runner, "s", "Run") + + assert await _kept_authors(runner, "fast_tool") == ["tool_agent"] + + +@pytest.mark.asyncio +async def test_run_async_tool_error_with_transfer_sibling_still_raises(): + """A finished transfer next to a failing call is neither run nor kept.""" + _, failing_tool, after_tool = _gated_tools(gate_on=("transfer_to_agent",)) + sub_agent = LlmAgent( + name="sub_agent", model=testing_utils.MockModel.create(responses=[]) + ) + root_agent = LlmAgent( + name="root_agent", + model=testing_utils.MockModel.create( + responses=[[ + types.Part.from_function_call( + name="transfer_to_agent", args={"agent_name": "sub_agent"} + ), + types.Part.from_function_call(name="failing_tool", args={}), + ]] + ), + tools=[failing_tool], + sub_agents=[sub_agent], + after_tool_callback=after_tool, + ) + runner = _abort_runner(root_agent) + + with pytest.raises(RuntimeError, match="downstream 500"): + await _run_turn(runner, "s", "Run") + + assert not sub_agent.model.requests + assert not await _persisted_responses(runner) + + +@pytest.mark.asyncio +async def test_run_async_tool_error_caller_closing_at_kept_result_returns(): + """A caller that stops reading at the streamed kept result is not held up.""" + fast_tool, failing_tool, after_tool = _gated_tools() + runner = _abort_runner( + _parallel_calls_agent([fast_tool, failing_tool], after_tool) + ) + + async def _read_until_kept_result() -> None: + async with aclosing(_run_async(runner)) as agen: + async for event in agen: + if event.get_function_responses(): + break + + # Bounded, so that a close that never returns fails instead of hanging. + await asyncio.wait_for(_read_until_kept_result(), timeout=5) + + assert await _persisted_responses(runner) == [{"record_id": "r1"}] + + if __name__ == "__main__": pytest.main([__file__]) From 3e309d5d3eb753343590e28baf5b8efc547645c8 Mon Sep 17 00:00:00 2001 From: Thuongvu Ho Date: Wed, 7 Oct 2026 19:56:01 -0700 Subject: [PATCH 2/2] test(workflow): check kept parallel results are saved once across a retry and a resume --- tests/unittests/test_runners.py | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 820dd3387b6..26630baef06 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -6286,7 +6286,9 @@ def after_tool(tool, args, tool_context, tool_response): return fast_tool, failing_tool, after_tool -def _parallel_calls_agent(tools, after_tool, name="tool_agent", **kwargs): +def _parallel_calls_agent( + tools, after_tool, name="tool_agent", answers=("Recovered",), **kwargs +): """An agent whose model calls every tool in parallel, then answers.""" return LlmAgent( name=name, @@ -6296,7 +6298,7 @@ def _parallel_calls_agent(tools, after_tool, name="tool_agent", **kwargs): types.Part.from_function_call(name=tool.__name__, args={}) for tool in tools ], - "Recovered", + *answers, ] ), tools=tools, @@ -6407,23 +6409,40 @@ async def test_run_async_tool_error_as_abort_lands_leaves_calls_to_abort_sealing @pytest.mark.asyncio -async def test_run_async_tool_error_retried_agent_sees_kept_result(): - """A node retried after the error sees the finished call's result.""" +@pytest.mark.parametrize("resumable", [False, True]) +async def test_run_async_tool_error_retried_agent_sees_kept_result( + resumable: bool, +): + """A node retried after the error sees the finished call's result. + + The result is saved once, also when the invocation is resumed afterwards. + """ fast_tool, failing_tool, after_tool = _gated_tools() agent = _parallel_calls_agent( [fast_tool, failing_tool], after_tool, + answers=("Recovered", "Resumed"), retry_config=RetryConfig(max_attempts=2, initial_delay=0, jitter=0), ) # No InvariantPlugin: the failing call stays unanswered in a completed turn, # as both calls do on main. - runner = _abort_runner(agent) + runner = _abort_runner(agent, resumable=resumable) await _run_turn(runner, "s", "Run") + if resumable: + invocation_id = (await _persisted_events(runner))[0].invocation_id + async with aclosing( + runner.run_async( + user_id=TEST_USER_ID, session_id="s", invocation_id=invocation_id + ) + ) as agen: + async for _ in agen: + pass assert _function_responses(agent.model.requests[1].contents) == [ {"record_id": "r1"} ] + assert await _persisted_responses(runner) == [{"record_id": "r1"}] @pytest.mark.asyncio