fix: rework parallel_agent.py to always aclose async generators

See #1670 (comment)

PiperOrigin-RevId: 795173267

Author: jawoszek

src/google/adk/agents/parallel_agent.py

@@ -17,11 +17,9 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any
+import sys
 from typing import AsyncGenerator
 from typing import ClassVar
-from typing import Dict
-from typing import Type
 
 from typing_extensions import override
 
@@ -49,6 +47,70 @@ def _create_branch_ctx_for_sub_agent(
   return invocation_context
 
 
+# TODO - remove once Python <3.11 is no longer supported.
+async def _merge_agent_run_pre_3_11(
+    agent_runs: list[AsyncGenerator[Event, None]],
+) -> AsyncGenerator[Event, None]:
+  """Merges the agent run event generator.
+  This version works in Python 3.9 and 3.10 and uses custom replacement for
+  asyncio.TaskGroup for tasks cancellation and exception handling.
+
+  This implementation guarantees for each agent, it won't move on until the
+  generated event is processed by upstream runner.
+
+  Args:
+      agent_runs: A list of async generators that yield events from each agent.
+
+  Yields:
+      Event: The next event from the merged generator.
+  """
+  sentinel = object()
+  queue = asyncio.Queue()
+
+  def propagate_exceptions(tasks):
+    # Propagate exceptions and errors from tasks.
+    for task in tasks:
+      if task.done():
+        # Ignore the result (None) of correctly finished tasks and re-raise
+        # exceptions and errors.
+        task.result()
+
+  # Agents are processed in parallel.
+  # Events for each agent are put on queue sequentially.
+  async def process_an_agent(events_for_one_agent):
+    try:
+      async for event in events_for_one_agent:
+        resume_signal = asyncio.Event()
+        await queue.put((event, resume_signal))
+        # Wait for upstream to consume event before generating new events.
+        await resume_signal.wait()
+    finally:
+      # Mark agent as finished.
+      await queue.put((sentinel, None))
+
+  tasks = []
+  try:
+    for events_for_one_agent in agent_runs:
+      tasks.append(asyncio.create_task(process_an_agent(events_for_one_agent)))
+
+    sentinel_count = 0
+    # Run until all agents finished processing.
+    while sentinel_count < len(agent_runs):
+      propagate_exceptions(tasks)
+      event, resume_signal = await queue.get()
+      # Agent finished processing.
+      if event is sentinel:
+        sentinel_count += 1
+      else:
+        yield event
+        # Signal to agent that event has been processed by runner and it can
+        # continue now.
+        resume_signal.set()
+  finally:
+    for task in tasks:
+      task.cancel()
+
+
 async def _merge_agent_run(
     agent_runs: list[AsyncGenerator[Event, None]],
 ) -> AsyncGenerator[Event, None]:
@@ -63,30 +125,37 @@ async def _merge_agent_run(
   Yields:
       Event: The next event from the merged generator.
   """
-  tasks = [
-      asyncio.create_task(events_for_one_agent.__anext__())
-      for events_for_one_agent in agent_runs
-  ]
-  pending_tasks = set(tasks)
-
-  while pending_tasks:
-    done, pending_tasks = await asyncio.wait(
-        pending_tasks, return_when=asyncio.FIRST_COMPLETED
-    )
-    for task in done:
-      try:
-        yield task.result()
-
-        # Find the generator that produced this event and move it on.
-        for i, original_task in enumerate(tasks):
-          if task == original_task:
-            new_task = asyncio.create_task(agent_runs[i].__anext__())
-            tasks[i] = new_task
-            pending_tasks.add(new_task)
-            break  # stop iterating once found
-
-      except StopAsyncIteration:
-        continue
+  sentinel = object()
+  queue = asyncio.Queue()
+
+  # Agents are processed in parallel.
+  # Events for each agent are put on queue sequentially.
+  async def process_an_agent(events_for_one_agent):
+    try:
+      async for event in events_for_one_agent:
+        resume_signal = asyncio.Event()
+        await queue.put((event, resume_signal))
+        # Wait for upstream to consume event before generating new events.
+        await resume_signal.wait()
+    finally:
+      # Mark agent as finished.
+      await queue.put((sentinel, None))
+
+  async with asyncio.TaskGroup() as tg:
+    for events_for_one_agent in agent_runs:
+      tg.create_task(process_an_agent(events_for_one_agent))
+
+    sentinel_count = 0
+    # Run until all agents finished processing.
+    while sentinel_count < len(agent_runs):
+      event, resume_signal = await queue.get()
+      # Agent finished processing.
+      if event is sentinel:
+        sentinel_count += 1
+      else:
+        yield event
+        # Signal to agent that it should generate next event.
+        resume_signal.set()
 
 
 class ParallelAgent(BaseAgent):
@@ -112,10 +181,19 @@ async def _run_async_impl(
         )
         for sub_agent in self.sub_agents
     ]
-
-    async with Aclosing(_merge_agent_run(agent_runs)) as agen:
-      async for event in agen:
-        yield event
+    try:
+      # TODO remove if once Python <3.11 is no longer supported.
+      if sys.version_info >= (3, 11):
+        async with Aclosing(_merge_agent_run(agent_runs)) as agen:
+          async for event in agen:
+            yield event
+      else:
+        async with Aclosing(_merge_agent_run_pre_3_11(agent_runs)) as agen:
+          async for event in agen:
+            yield event
+    finally:
+      for sub_agent_run in agent_runs:
+        await sub_agent_run.aclose()
 
   @override
   async def _run_live_impl(

tests/unittests/agents/test_parallel_agent.py

@@ -33,12 +33,8 @@ class _TestingAgent(BaseAgent):
   delay: float = 0
   """The delay before the agent generates an event."""
 
-  @override
-  async def _run_async_impl(
-      self, ctx: InvocationContext
-  ) -> AsyncGenerator[Event, None]:
-    await asyncio.sleep(self.delay)
-    yield Event(
+  def event(self, ctx: InvocationContext):
+    return Event(
         author=self.name,
         branch=ctx.branch,
         invocation_id=ctx.invocation_id,
@@ -47,6 +43,13 @@ async def _run_async_impl(
         ),
     )
 
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    await asyncio.sleep(self.delay)
+    yield self.event(ctx)
+
 
 async def _create_parent_invocation_context(
     test_name: str, agent: BaseAgent
@@ -135,3 +138,102 @@ async def test_run_async_branches(request: pytest.FixtureRequest):
   # Sub-agents should have different branches.
   assert events[2].branch != events[1].branch
   assert events[2].branch != events[0].branch
+
+
+class _TestingAgentWithMultipleEvents(_TestingAgent):
+  """Mock agent for testing."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    for _ in range(0, 3):
+      event = self.event(ctx)
+      yield event
+      # Check that the event was processed by the consumer.
+      assert event.custom_metadata is not None
+      assert event.custom_metadata['processed']
+
+
+@pytest.mark.asyncio
+async def test_generating_one_event_per_agent_at_once(
+    request: pytest.FixtureRequest,
+):
+  # This test is to verify that the parallel agent won't generate more than one
+  # event per agent at a time.
+  agent1 = _TestingAgentWithMultipleEvents(
+      name=f'{request.function.__name__}_test_agent_1'
+  )
+  agent2 = _TestingAgentWithMultipleEvents(
+      name=f'{request.function.__name__}_test_agent_2'
+  )
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[
+          agent1,
+          agent2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent
+  )
+
+  agen = parallel_agent.run_async(parent_ctx)
+  async for event in agen:
+    event.custom_metadata = {'processed': True}
+    # Asserts on event are done in _TestingAgentWithMultipleEvents.
+
+
+class _TestingAgentWithException(_TestingAgent):
+  """Mock agent for testing."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield self.event(ctx)
+    raise Exception()
+
+
+class _TestingAgentInfiniteEvents(_TestingAgent):
+  """Mock agent for testing."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    while True:
+      yield self.event(ctx)
+
+
+@pytest.mark.asyncio
+async def test_stop_agent_if_sub_agent_fails(
+    request: pytest.FixtureRequest,
+):
+  # This test is to verify that the parallel agent and subagents will all stop
+  # processing and throw exception to top level runner in case of exception.
+  agent1 = _TestingAgentWithException(
+      name=f'{request.function.__name__}_test_agent_1'
+  )
+  agent2 = _TestingAgentInfiniteEvents(
+      name=f'{request.function.__name__}_test_agent_2'
+  )
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[
+          agent1,
+          agent2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent
+  )
+
+  agen = parallel_agent.run_async(parent_ctx)
+  # We expect to receive an exception from one of subagents.
+  # The exception should be propagated to root agent and other subagents.
+  # Otherwise we'll have an infinite loop.
+  with pytest.raises(Exception):
+    async for _ in agen:
+      # The infinite agent could iterate a few times depending on scheduling.
+      pass