From 4ac2cf85df644ec74f4a9cd69afd80cbb62dd228 Mon Sep 17 00:00:00 2001 From: Aarav Mittal <137450929+a2105z@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:18:51 -0500 Subject: [PATCH] feat(agents): let nested LoopAgents exit independently escalate=True still stops every enclosing loop so exit_loop stays backward compatible. EscalationContext(type='parent') is consumed by the nearest LoopAgent so an outer retry loop can continue. exit_current_loop is the parent-scoped tool. Fixes #2808 --- src/google/adk/agents/loop_agent.py | 9 +- src/google/adk/agents/parallel_agent.py | 2 +- src/google/adk/events/__init__.py | 2 + src/google/adk/events/event_actions.py | 71 +++++ src/google/adk/tools/__init__.py | 2 + src/google/adk/tools/exit_loop_tool.py | 13 +- tests/unittests/a2a/converters/test_to_adk.py | 6 + tests/unittests/agents/test_loop_agent.py | 260 ++++++++++++++++++ tests/unittests/events/test_event_actions.py | 45 +++ tests/unittests/tools/test_exit_loop_tool.py | 42 +++ 10 files changed, 448 insertions(+), 4 deletions(-) create mode 100644 tests/unittests/tools/test_exit_loop_tool.py diff --git a/src/google/adk/agents/loop_agent.py b/src/google/adk/agents/loop_agent.py index aefe076a115..aead52fceb7 100644 --- a/src/google/adk/agents/loop_agent.py +++ b/src/google/adk/agents/loop_agent.py @@ -60,6 +60,11 @@ class LoopAgent(BaseAgent): When sub-agent generates an event with escalate or max_iterations are reached, the loop agent will stop. + ``escalate=True`` with no ``escalation_context`` (including ``exit_loop``) + exits every enclosing LoopAgent. ``EscalationContext(type='parent')`` + (including ``exit_current_loop``) exits only this loop so an outer loop + can continue. + .. deprecated:: LoopAgent is deprecated in favor of Workflow and will be removed in a future version. Workflow cannot yet be used as an LlmAgent sub-agent. @@ -112,9 +117,9 @@ async def _run_async_impl( async with Aclosing(sub_agent.run_async(ctx)) as agen: async for event in agen: - yield event - if event.actions.escalate: + if event.actions.consume_workflow_escalation(self.name): should_exit = True + yield event if ctx.should_pause_invocation(event): pause_invocation = True diff --git a/src/google/adk/agents/parallel_agent.py b/src/google/adk/agents/parallel_agent.py index 908b0e9893a..32ba7c157ed 100644 --- a/src/google/adk/agents/parallel_agent.py +++ b/src/google/adk/agents/parallel_agent.py @@ -57,7 +57,7 @@ def _create_branch_ctx_for_sub_agent( def _has_escalate_action(event: Event) -> bool: """Returns whether the event asks the parent workflow to exit early.""" - return bool(event.actions.escalate) + return event.actions.is_active_escalation() def _cancel_tasks(tasks: list[asyncio.Task[None]]) -> None: diff --git a/src/google/adk/events/__init__.py b/src/google/adk/events/__init__.py index d027ffbe424..3ac5dc41971 100644 --- a/src/google/adk/events/__init__.py +++ b/src/google/adk/events/__init__.py @@ -13,10 +13,12 @@ # limitations under the License. from .event import Event +from .event_actions import EscalationContext from .event_actions import EventActions from .request_input import RequestInput __all__ = [ + 'EscalationContext', 'Event', 'EventActions', 'RequestInput', diff --git a/src/google/adk/events/event_actions.py b/src/google/adk/events/event_actions.py index dd0ada936bc..7bd05d5aebc 100644 --- a/src/google/adk/events/event_actions.py +++ b/src/google/adk/events/event_actions.py @@ -17,6 +17,7 @@ import logging from typing import Any from typing import cast +from typing import Literal from typing import Optional from typing import TYPE_CHECKING from typing import Union @@ -75,6 +76,34 @@ class EventCompaction(BaseModel): # type: ignore[misc] """The compacted content of the events.""" +class EscalationContext(BaseModel): # type: ignore[misc] + """Controls how far an ``escalate`` action travels in nested workflows. + + ``escalate=True`` with no context is treated as ``type='root'`` so existing + ``exit_loop`` callers keep stopping every enclosing LoopAgent. + """ + + model_config = ConfigDict( + extra='forbid', + alias_generator=alias_generators.to_camel, + populate_by_name=True, + ) + """The pydantic model config.""" + + type: Literal['parent', 'root'] = 'root' + """Which ancestor should stop. + + - ``root``: every enclosing LoopAgent exits (same as escalate=True alone). + - ``parent``: only the nearest LoopAgent exits; outer loops continue. + """ + + handled_by: list[str] = Field(default_factory=list) + """LoopAgents that already consumed a parent-scoped escalation.""" + + target_agent: Optional[str] = None + """If set, only this named agent treats the event as a loop exit.""" + + class EventActions(BaseModel): # type: ignore[misc] """Represents the actions attached to an event.""" @@ -125,6 +154,13 @@ def _serialize_state_delta( escalate: Optional[bool] = None """The agent is escalating to a higher level agent.""" + escalation_context: Optional[EscalationContext] = None + """Optional scope for ``escalate``. + + When omitted, ``escalate=True`` stops every enclosing LoopAgent (root). + ``EscalationContext(type='parent')`` stops only the nearest LoopAgent. + """ + requested_auth_configs: dict[str, AuthConfig] = Field(default_factory=dict) """Authentication configurations requested by tool responses. @@ -200,3 +236,38 @@ def _serialize_agent_state( set_model_response: Optional[Any] = None """The model response structured output.""" + + def is_active_escalation(self) -> bool: + """Whether this event should still stop ancestor workflows. + + Parent-scoped escalations already consumed by a LoopAgent are inactive so + outer loops and ParallelAgent siblings keep running. + """ + context = self.escalation_context + if context is None: + return bool(self.escalate) + if context.target_agent is not None: + return False + if context.type == 'parent' and context.handled_by: + return False + return bool(self.escalate) or context is not None + + def consume_workflow_escalation(self, agent_name: str) -> bool: + """Returns True if ``agent_name`` should exit because of this event. + + Parent-scoped escalations are consumed by the nearest LoopAgent, which + records itself on ``escalation_context.handled_by``. Root-scoped + escalations (including bare ``escalate=True``) stay active for every + ancestor. + """ + context = self.escalation_context + if context is None: + return bool(self.escalate) + if context.target_agent is not None: + return context.target_agent == agent_name + if context.type == 'root': + return True + if context.handled_by: + return False + context.handled_by.append(agent_name) + return True diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index 9cd22f305ff..d86fd8f0dad 100644 --- a/src/google/adk/tools/__init__.py +++ b/src/google/adk/tools/__init__.py @@ -30,6 +30,7 @@ from .discovery_engine_search_tool import SearchResultMode from .enterprise_search_tool import enterprise_web_search_tool as enterprise_web_search from .example_tool import ExampleTool + from .exit_loop_tool import exit_current_loop from .exit_loop_tool import exit_loop from .function_tool import FunctionTool from .get_user_choice_tool import get_user_choice_tool as get_user_choice @@ -68,6 +69,7 @@ 'enterprise_web_search_tool', ), 'ExampleTool': ('.example_tool', 'ExampleTool'), + 'exit_current_loop': ('.exit_loop_tool', 'exit_current_loop'), 'exit_loop': ('.exit_loop_tool', 'exit_loop'), 'FunctionTool': ('.function_tool', 'FunctionTool'), 'get_user_choice': ('.get_user_choice_tool', 'get_user_choice_tool'), diff --git a/src/google/adk/tools/exit_loop_tool.py b/src/google/adk/tools/exit_loop_tool.py index 07a925eeae8..0f82fa4f308 100644 --- a/src/google/adk/tools/exit_loop_tool.py +++ b/src/google/adk/tools/exit_loop_tool.py @@ -14,13 +14,24 @@ from __future__ import annotations +from ..events.event_actions import EscalationContext from .tool_context import ToolContext def exit_loop(tool_context: ToolContext) -> None: - """Exits the loop. + """Exits every enclosing loop. Call this function only when you are instructed to do so. """ tool_context.actions.escalate = True tool_context.actions.skip_summarization = True + + +def exit_current_loop(tool_context: ToolContext) -> None: + """Exits only the nearest enclosing LoopAgent. + + Outer loops keep running. Use ``exit_loop`` to stop every enclosing loop. + """ + tool_context.actions.escalate = True + tool_context.actions.escalation_context = EscalationContext(type='parent') + tool_context.actions.skip_summarization = True diff --git a/tests/unittests/a2a/converters/test_to_adk.py b/tests/unittests/a2a/converters/test_to_adk.py index ac1f05836b7..85f003ba46a 100644 --- a/tests/unittests/a2a/converters/test_to_adk.py +++ b/tests/unittests/a2a/converters/test_to_adk.py @@ -305,6 +305,11 @@ def test_peer_supplied_actions_cannot_mutate_caller_session(self): metadata = { _get_adk_metadata_key("actions"): { "escalate": True, + "escalationContext": { + "type": "parent", + "handledBy": ["attacker-loop"], + "targetAgent": "attacker-agent", + }, "stateDelta": {"app:is_admin": True, "user:persona": "attacker"}, "artifactDelta": {"report.pdf": 7}, "transferToAgent": "attacker-agent", @@ -427,6 +432,7 @@ def test_peer_supplied_actions_cannot_mutate_caller_session(self): assert event.actions.route is None assert event.actions.render_ui_widgets is None assert event.actions.set_model_response is None + assert event.actions.escalation_context is None # Inert fields a peer may set are still honored. assert event.actions.escalate is True diff --git a/tests/unittests/agents/test_loop_agent.py b/tests/unittests/agents/test_loop_agent.py index e4465883890..dbe1d3269a0 100644 --- a/tests/unittests/agents/test_loop_agent.py +++ b/tests/unittests/agents/test_loop_agent.py @@ -21,8 +21,10 @@ from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.loop_agent import LoopAgent from google.adk.agents.loop_agent import LoopAgentState +from google.adk.agents.sequential_agent import SequentialAgent from google.adk.apps import ResumabilityConfig from google.adk.events.event import Event +from google.adk.events.event_actions import EscalationContext from google.adk.events.event_actions import EventActions from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.genai import types @@ -290,3 +292,261 @@ def mock_should_pause(event): def test_deprecation_mentions_sub_agent_limitation(): with pytest.warns(DeprecationWarning, match='sub-agent'): LoopAgent(name='deprecated_loop', sub_agents=[]) + + +class _CountingAgent(BaseAgent): + + def __init__(self, name: str, bucket: list[str]): + super().__init__(name=name) + object.__setattr__(self, '_bucket', bucket) + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + self._bucket.append(self.name) + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text=self.name)]), + ) + + +class _ParentEscalateAgent(BaseAgent): + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text='exit inner')]), + actions=EventActions( + escalate=True, + escalation_context=EscalationContext(type='parent'), + ), + ) + + +class _TargetedEscalateAgent(BaseAgent): + + def __init__(self, name: str, target_agent: str): + super().__init__(name=name) + object.__setattr__(self, '_target_agent', target_agent) + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text='exit named')]), + actions=EventActions( + escalate=True, + escalation_context=EscalationContext( + type='parent', target_agent=self._target_agent + ), + ), + ) + + +@pytest.mark.asyncio +async def test_nested_loop_root_escalate_stops_every_loop( + request: pytest.FixtureRequest, +): + """Bare escalate=True still exits inner and outer loops (issue #2808 default).""" + inner_runs: list[str] = [] + outer_runs: list[str] = [] + inner = LoopAgent( + name=f'{request.function.__name__}_inner', + sub_agents=[ + _CountingAgent(f'{request.function.__name__}_inner_work', inner_runs), + _TestingAgentWithEscalateAction( + name=f'{request.function.__name__}_root_exit' + ), + ], + ) + outer = LoopAgent( + name=f'{request.function.__name__}_outer', + max_iterations=3, + sub_agents=[ + inner, + _CountingAgent(f'{request.function.__name__}_outer_work', outer_runs), + ], + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, outer + ) + _ = [e async for e in outer.run_async(parent_ctx)] + assert inner_runs == [f'{request.function.__name__}_inner_work'] + assert outer_runs == [] + + +@pytest.mark.asyncio +async def test_nested_loop_parent_escalate_exits_only_inner( + request: pytest.FixtureRequest, +): + """EscalationContext(type='parent') lets the outer loop keep iterating.""" + inner_runs: list[str] = [] + outer_runs: list[str] = [] + inner_name = f'{request.function.__name__}_inner' + inner = LoopAgent( + name=inner_name, + sub_agents=[ + _CountingAgent(f'{request.function.__name__}_inner_work', inner_runs), + _ParentEscalateAgent(name=f'{request.function.__name__}_parent_exit'), + ], + ) + outer = LoopAgent( + name=f'{request.function.__name__}_outer', + max_iterations=3, + sub_agents=[ + inner, + _CountingAgent(f'{request.function.__name__}_outer_work', outer_runs), + ], + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, outer + ) + events = [e async for e in outer.run_async(parent_ctx)] + + assert inner_runs == [f'{request.function.__name__}_inner_work'] * 3 + assert outer_runs == [f'{request.function.__name__}_outer_work'] * 3 + parent_events = [ + event + for event in events + if event.actions.escalation_context + and event.actions.escalation_context.type == 'parent' + ] + assert parent_events + assert parent_events[0].actions.escalation_context.handled_by == [inner_name] + + +class _EscalateAfterNAgent(BaseAgent): + """Escalates with parent scope after this instance has run ``limit`` times.""" + + def __init__(self, name: str, limit: int): + super().__init__(name=name) + object.__setattr__(self, '_limit', limit) + object.__setattr__(self, '_calls', 0) + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + self._calls += 1 + if self._calls >= self._limit: + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text='break')]), + actions=EventActions( + escalate=True, + escalation_context=EscalationContext(type='parent'), + ), + ) + return + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text='continue')]), + ) + + +class _EscalateEveryNCountsAgent(BaseAgent): + """Parent-escalates whenever ``bucket`` length is a multiple of ``n``.""" + + def __init__(self, name: str, n: int, bucket: list[str]): + super().__init__(name=name) + object.__setattr__(self, '_n', n) + object.__setattr__(self, '_bucket', bucket) + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + if self._bucket and len(self._bucket) % self._n == 0: + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text='break inner')]), + actions=EventActions( + escalate=True, + escalation_context=EscalationContext(type='parent'), + ), + ) + return + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text='continue')]), + ) + + +@pytest.mark.asyncio +async def test_nested_loop_parent_escalate_runs_inner_times_outer( + request: pytest.FixtureRequest, +): + """Issue #2808: inner 5 × outer 5 should be 25 inner steps, not 5.""" + inner_runs: list[str] = [] + inner = LoopAgent( + name=f'{request.function.__name__}_inner', + sub_agents=[ + _CountingAgent(f'{request.function.__name__}_inner_work', inner_runs), + _EscalateEveryNCountsAgent( + name=f'{request.function.__name__}_inner_break', + n=5, + bucket=inner_runs, + ), + ], + ) + outer = LoopAgent( + name=f'{request.function.__name__}_outer', + sub_agents=[ + inner, + _EscalateAfterNAgent( + name=f'{request.function.__name__}_outer_break', limit=5 + ), + ], + ) + root = SequentialAgent( + name=f'{request.function.__name__}_root', + sub_agents=[outer], + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, root + ) + _ = [e async for e in root.run_async(parent_ctx)] + assert inner_runs == [f'{request.function.__name__}_inner_work'] * 25 + + +@pytest.mark.asyncio +async def test_nested_loop_target_agent_skips_inner( + request: pytest.FixtureRequest, +): + """target_agent on the outer loop leaves the inner loop running until max.""" + inner_runs: list[str] = [] + outer_name = f'{request.function.__name__}_outer' + inner = LoopAgent( + name=f'{request.function.__name__}_inner', + max_iterations=2, + sub_agents=[ + _CountingAgent(f'{request.function.__name__}_inner_work', inner_runs), + _TargetedEscalateAgent( + name=f'{request.function.__name__}_aim_outer', + target_agent=outer_name, + ), + ], + ) + outer = LoopAgent( + name=outer_name, + max_iterations=5, + sub_agents=[inner], + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, outer + ) + _ = [e async for e in outer.run_async(parent_ctx)] + assert inner_runs == [f'{request.function.__name__}_inner_work'] * 2 diff --git a/tests/unittests/events/test_event_actions.py b/tests/unittests/events/test_event_actions.py index ef864060db9..e3ea6c1a2cb 100644 --- a/tests/unittests/events/test_event_actions.py +++ b/tests/unittests/events/test_event_actions.py @@ -20,6 +20,7 @@ import logging from google.adk.events.event_actions import _make_json_serializable +from google.adk.events.event_actions import EscalationContext from google.adk.events.event_actions import EventActions from pydantic import BaseModel @@ -148,3 +149,47 @@ def test_non_serializable_agent_state_logs_warning(self, caplog): 'Failed to serialize `agent_state`' in record.message for record in caplog.records ) + + +class TestEscalationContext: + """Tests for scoped loop escalation.""" + + def test_bare_escalate_is_root_and_stays_active(self): + actions = EventActions(escalate=True) + assert actions.consume_workflow_escalation('inner') + assert actions.consume_workflow_escalation('outer') + assert actions.is_active_escalation() + assert actions.escalation_context is None + + def test_parent_scope_is_consumed_by_nearest_agent(self): + actions = EventActions( + escalate=True, + escalation_context=EscalationContext(type='parent'), + ) + assert actions.consume_workflow_escalation('inner') + assert actions.escalation_context.handled_by == ['inner'] + assert not actions.consume_workflow_escalation('outer') + assert not actions.is_active_escalation() + + def test_target_agent_only_matches_named_loop(self): + actions = EventActions( + escalate=True, + escalation_context=EscalationContext( + type='parent', target_agent='outer' + ), + ) + assert not actions.consume_workflow_escalation('inner') + assert actions.consume_workflow_escalation('outer') + assert not actions.is_active_escalation() + + def test_escalation_context_round_trips(self): + actions = EventActions( + escalation_context=EscalationContext( + type='parent', handled_by=['inner'], target_agent='outer' + ) + ) + dumped = actions.model_dump(mode='json', by_alias=True) + restored = EventActions.model_validate(dumped) + assert restored.escalation_context.type == 'parent' + assert restored.escalation_context.handled_by == ['inner'] + assert restored.escalation_context.target_agent == 'outer' diff --git a/tests/unittests/tools/test_exit_loop_tool.py b/tests/unittests/tools/test_exit_loop_tool.py new file mode 100644 index 00000000000..022c94b740b --- /dev/null +++ b/tests/unittests/tools/test_exit_loop_tool.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from google.adk.events.event_actions import EventActions +from google.adk.tools.exit_loop_tool import exit_current_loop +from google.adk.tools.exit_loop_tool import exit_loop + + +class _FakeToolContext: + + def __init__(self): + self.actions = EventActions() + + +def test_exit_loop_is_root_scoped(): + ctx = _FakeToolContext() + exit_loop(ctx) + assert ctx.actions.escalate is True + assert ctx.actions.skip_summarization is True + assert ctx.actions.escalation_context is None + + +def test_exit_current_loop_is_parent_scoped(): + ctx = _FakeToolContext() + exit_current_loop(ctx) + assert ctx.actions.escalate is True + assert ctx.actions.skip_summarization is True + assert ctx.actions.escalation_context is not None + assert ctx.actions.escalation_context.type == 'parent'