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..26630baef06 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,301 @@ 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", answers=("Recovered",), **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 + ], + *answers, + ] + ), + 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 +@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, 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 +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__])