fix: end a ParallelAgent early only when a direct sub-agent escalates (cherry-pick to release/candidate) (#6918)

Co-authored-by: George Weale <gweale@google.com>

Author: wukath

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

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',
+  ]