feat(agents): annotate events with the LoopAgent iteration

src/google/adk/agents/loop_agent.py

@@ -79,6 +79,27 @@ class LoopAgent(BaseAgent):
   escalates.
   """
 
+  LOOP_ITERATION_KEY: ClassVar[str] = 'loop_iteration'
+  """Key under ``Event.custom_metadata`` recording the 0-based iteration of the
+  nearest enclosing ``LoopAgent`` during which the event was produced.
+
+  Without it, events produced in different iterations are indistinguishable
+  (same ``author``, ``branch`` and ``invocation_id``), so consumers cannot tell
+  loop iterations apart — e.g. a ``ParallelAgent`` nested in a ``LoopAgent``,
+  whose branch ids repeat every iteration. See
+  https://github.com/google/adk-python/issues/6266.
+  """
+
+  def _annotate_loop_iteration(self, event: Event, iteration: int) -> None:
+    """Record the current loop iteration on ``event.custom_metadata``.
+
+    Uses ``setdefault`` so that when ``LoopAgent``s are nested the nearest
+    enclosing loop (which annotates first, as the event bubbles up) wins.
+    """
+    metadata = dict(event.custom_metadata) if event.custom_metadata else {}
+    metadata.setdefault(self.LOOP_ITERATION_KEY, iteration)
+    event.custom_metadata = metadata
+
   @override
   async def _run_async_impl(
       self, ctx: InvocationContext
@@ -112,6 +133,7 @@ async def _run_async_impl(
 
         async with Aclosing(sub_agent.run_async(ctx)) as agen:
           async for event in agen:
+            self._annotate_loop_iteration(event, times_looped)
             yield event
             if event.actions.escalate:
               should_exit = True

tests/unittests/agents/test_loop_agent.py

@@ -139,6 +139,61 @@ async def test_run_async(request: pytest.FixtureRequest, resumable: bool):
   assert simplified_events == expected_events
 
 
+@pytest.mark.asyncio
+async def test_run_async_annotates_loop_iteration(
+    request: pytest.FixtureRequest,
+):
+  agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
+  loop_agent = LoopAgent(
+      name=f'{request.function.__name__}_test_loop_agent',
+      max_iterations=3,
+      sub_agents=[agent],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, loop_agent, resumable=False
+  )
+
+  events = [e async for e in loop_agent.run_async(parent_ctx)]
+
+  sub_agent_events = [e for e in events if e.author == agent.name]
+  assert len(sub_agent_events) == 3
+  assert [
+      e.custom_metadata[LoopAgent.LOOP_ITERATION_KEY]
+      for e in sub_agent_events
+  ] == [0, 1, 2]
+
+
+@pytest.mark.asyncio
+async def test_run_async_loop_iteration_survives_nested_loops(
+    request: pytest.FixtureRequest,
+):
+  # An inner LoopAgent annotates first (as events bubble up), so its iteration
+  # wins over the outer loop's for events it produced.
+  agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
+  inner = LoopAgent(
+      name=f'{request.function.__name__}_inner_loop',
+      max_iterations=2,
+      sub_agents=[agent],
+  )
+  outer = LoopAgent(
+      name=f'{request.function.__name__}_outer_loop',
+      max_iterations=2,
+      sub_agents=[inner],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, outer, resumable=False
+  )
+
+  events = [e async for e in outer.run_async(parent_ctx)]
+
+  sub_agent_events = [e for e in events if e.author == agent.name]
+  # outer x2 * inner x2 = 4 events; each carries the inner loop's iteration.
+  assert [
+      e.custom_metadata[LoopAgent.LOOP_ITERATION_KEY]
+      for e in sub_agent_events
+  ] == [0, 1, 0, 1]
+
+
 @pytest.mark.asyncio
 async def test_resume_async(request: pytest.FixtureRequest):
   agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')