diff --git a/src/google/adk/agents/parallel_agent.py b/src/google/adk/agents/parallel_agent.py index 82d0446bcd9..0a31d989b73 100644 --- a/src/google/adk/agents/parallel_agent.py +++ b/src/google/adk/agents/parallel_agent.py @@ -55,9 +55,21 @@ def _create_branch_ctx_for_sub_agent( return invocation_context -def _has_escalate_action(event: Event) -> bool: - """Returns whether the event asks the parent workflow to exit early.""" - return bool(event.actions.escalate) +def _asks_this_agent_to_exit(event: Event, sub_agent_names: set[str]) -> bool: + """Returns whether the event asks this parallel agent to exit early. + + An escalation ends the workflow that directly encloses the escalating agent, + and that workflow re-yields the event while unwinding, so only an escalation + authored by a direct sub-agent is addressed to this one. + + Args: + event: The event to inspect. + sub_agent_names: Names of this agent's direct sub-agents. + + Returns: + Whether this parallel agent should stop its remaining branches. + """ + return bool(event.actions.escalate) and event.author in sub_agent_names def _cancel_tasks(tasks: list[asyncio.Task[None]]) -> None: @@ -69,6 +81,7 @@ def _cancel_tasks(tasks: list[asyncio.Task[None]]) -> None: async def _merge_agent_run( agent_runs: list[AsyncGenerator[Event, None]], + sub_agent_names: set[str], ) -> AsyncGenerator[Event, None]: """Merges agent runs using asyncio.TaskGroup on Python 3.11+.""" sentinel = _AgentRunComplete() @@ -121,7 +134,7 @@ async def process_an_agent( raise payload else: yield event - if _has_escalate_action(event): + if _asks_this_agent_to_exit(event, sub_agent_names): _cancel_tasks(tasks) return # Signal to agent that it should generate next event. @@ -135,6 +148,7 @@ async def process_an_agent( # TODO - remove once Python <3.11 is no longer supported. async def _merge_agent_run_pre_3_11( agent_runs: list[AsyncGenerator[Event, None]], + sub_agent_names: set[str], ) -> AsyncGenerator[Event, None]: """Merges agent runs for Python 3.10 without asyncio.TaskGroup. @@ -143,6 +157,7 @@ async def _merge_agent_run_pre_3_11( Args: agent_runs: Async generators that yield events from each agent. + sub_agent_names: Names of the parallel agent's direct sub-agents. Yields: Event: The next event from the merged generator. @@ -190,7 +205,7 @@ async def process_an_agent( raise payload else: yield event - if _has_escalate_action(event): + if _asks_this_agent_to_exit(event, sub_agent_names): _cancel_tasks(tasks) return # Signal to agent that event has been processed by runner and it can @@ -246,6 +261,7 @@ async def _run_async_impl( yield self._create_agent_state_event(ctx) agent_runs = [] + sub_agent_names = {sub_agent.name for sub_agent in self.sub_agents} # Prepare and collect async generators for each sub-agent. for sub_agent in self.sub_agents: sub_agent_ctx = _create_branch_ctx_for_sub_agent(self, sub_agent, ctx) @@ -261,10 +277,10 @@ async def _run_async_impl( if sys.version_info >= (3, 11) else _merge_agent_run_pre_3_11 ) - async with Aclosing(merge_func(agent_runs)) as agen: + async with Aclosing(merge_func(agent_runs, sub_agent_names)) as agen: async for event in agen: yield event - if _has_escalate_action(event): + if _asks_this_agent_to_exit(event, sub_agent_names): escalated = True if ctx.should_pause_invocation(event): pause_invocation = True diff --git a/tests/unittests/agents/test_parallel_agent.py b/tests/unittests/agents/test_parallel_agent.py index 924cc3f87c5..5d9dac30dc7 100644 --- a/tests/unittests/agents/test_parallel_agent.py +++ b/tests/unittests/agents/test_parallel_agent.py @@ -23,6 +23,7 @@ from google.adk.agents.base_agent import BaseAgent from google.adk.agents.base_agent import BaseAgentState from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.loop_agent import LoopAgent from google.adk.agents.parallel_agent import _merge_agent_run_pre_3_11 from google.adk.agents.parallel_agent import ParallelAgent from google.adk.agents.sequential_agent import SequentialAgent @@ -461,7 +462,7 @@ async def test_sub_agent_failure_reaches_busy_caller( async def test_merge_agent_run_pre_3_11_surfaces_failure_without_events(): """A branch failing before it emits anything must not look successful.""" with pytest.raises(ValueError, match='simulated sub-agent failure'): - async for _ in _merge_agent_run_pre_3_11([_failing_agent()]): + async for _ in _merge_agent_run_pre_3_11([_failing_agent()], set()): pass @@ -475,7 +476,7 @@ async def test_merge_agent_run_pre_3_11_no_aclose_error_on_failure(): agent_runs = [_slow_agent_with_cleanup_delay(), _failing_agent()] with pytest.raises(ValueError, match='simulated sub-agent failure'): - async for _ in _merge_agent_run_pre_3_11(agent_runs): + async for _ in _merge_agent_run_pre_3_11(agent_runs, set()): pass # If tasks were not properly awaited, aclose() on a still-running generator @@ -705,3 +706,77 @@ def mock_should_pause(event: Event) -> bool: assert events[0].author == fast_agent.name assert events[1].author == escalating_agent.name assert events[1].actions.escalate + + +@pytest.mark.asyncio +@pytest.mark.parametrize('use_pre_3_11_merge', [False, True]) +async def test_run_async_keeps_siblings_when_a_nested_loop_ends_itself( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, + use_pre_3_11_merge: bool, +): + """A sub-agent ending its own LoopAgent must not cancel sibling branches.""" + + ticks: dict[str, int] = {} + + class _LoopingAgent(_TestingAgent): + """Escalates on its `escalate_on`-th run to end its enclosing loop.""" + + escalate_on: int = 1 + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + await asyncio.sleep(self.delay) + ticks[self.name] = ticks.get(self.name, 0) + 1 + escalating = ticks[self.name] >= self.escalate_on + yield self.event( + ctx, + text=f'{self.name}#{ticks[self.name]}', + actions=EventActions(escalate=True) if escalating else EventActions(), + ) + + if use_pre_3_11_merge: + monkeypatch.setattr( + parallel_agent_module, + 'sys', + SimpleNamespace(version_info=(3, 10)), + ) + + fast_agent = _LoopingAgent( + name=f'{request.function.__name__}_test_fast_agent', + escalate_on=1, + ) + slow_agent = _LoopingAgent( + name=f'{request.function.__name__}_test_slow_agent', + delay=0.05, + escalate_on=3, + ) + parallel_agent = ParallelAgent( + name=f'{request.function.__name__}_test_parallel_agent', + sub_agents=[ + LoopAgent( + name=f'{request.function.__name__}_test_fast_loop', + sub_agents=[fast_agent], + max_iterations=5, + ), + LoopAgent( + name=f'{request.function.__name__}_test_slow_loop', + sub_agents=[slow_agent], + max_iterations=5, + ), + ], + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, parallel_agent + ) + + events = [e async for e in parallel_agent.run_async(parent_ctx)] + + assert [event.content.parts[0].text for event in events] == [ + f'{fast_agent.name}#1', + f'{slow_agent.name}#1', + f'{slow_agent.name}#2', + f'{slow_agent.name}#3', + ]