diff --git a/src/google/adk/agents/loop_agent.py b/src/google/adk/agents/loop_agent.py index 6bf5221404..14a1c2fd41 100644 --- a/src/google/adk/agents/loop_agent.py +++ b/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 diff --git a/tests/unittests/agents/test_loop_agent.py b/tests/unittests/agents/test_loop_agent.py index 0e23d9d42c..be93e7b8ef 100644 --- a/tests/unittests/agents/test_loop_agent.py +++ b/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')