diff --git a/src/google/adk/agents/_managed_agent.py b/src/google/adk/agents/_managed_agent.py index a1e89d5fa85..4024735436d 100644 --- a/src/google/adk/agents/_managed_agent.py +++ b/src/google/adk/agents/_managed_agent.py @@ -366,12 +366,15 @@ async def _run_impl( parent_context = parent_context.model_copy( update={'user_content': to_user_content(node_input)} ) - async for event in self.run_async(parent_context=parent_context): - if event.author: - ctx.event_author = event.author - if not event.node_info.path and event.author == self.name: - event.node_info.path = ctx.node_path - yield event + # Aclosing, as in BaseAgent._run_impl: closing the node early must close + # run_async too. + async with Aclosing(self.run_async(parent_context=parent_context)) as agen: + async for event in agen: + if event.author: + ctx.event_author = event.author + if not event.node_info.path and event.author == self.name: + event.node_info.path = ctx.node_path + yield event async def _run_async_impl( self, ctx: InvocationContext diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index 80a2dfa95cc..6bcb047b338 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -392,16 +392,20 @@ async def _run_impl( node_input: Any, ) -> AsyncGenerator[Any, None]: """Runs the agent as a node.""" - async for event in self.run_async( - parent_context=ctx.get_invocation_context() - ): - # Preserve author by setting it in context for NodeRunner - if event.author: - ctx.event_author = event.author + # Aclosing, so that closing the node early closes run_async with it: its + # finally blocks and after_agent_callback unwinding run before aclose() + # returns instead of being deferred to the asyncgen finalizer hook. + async with Aclosing( + self.run_async(parent_context=ctx.get_invocation_context()) + ) as agen: + async for event in agen: + # Preserve author by setting it in context for NodeRunner + if event.author: + ctx.event_author = event.author - if not event.node_info.path and event.author == self.name: - event.node_info.path = ctx.node_path - yield event + if not event.node_info.path and event.author == self.name: + event.node_info.path = ctx.node_path + yield event @final async def run_live( diff --git a/src/google/adk/labs/antigravity/_antigravity_agent.py b/src/google/adk/labs/antigravity/_antigravity_agent.py index a051acc54e5..46d385fee08 100644 --- a/src/google/adk/labs/antigravity/_antigravity_agent.py +++ b/src/google/adk/labs/antigravity/_antigravity_agent.py @@ -57,6 +57,7 @@ from ...utils._agent_mode import AgentMode from ...utils._agent_mode import SingleTurnAgentMode from ...utils.content_utils import to_user_content +from ...utils.context_utils import Aclosing from ._event_converter import convert_step_to_events from ._event_converter import drain_tool_results from ._event_converter import final_model_text @@ -802,15 +803,16 @@ async def _run_impl( last_text: str | None = None # Keep in sync with BaseAgent._run_impl: super() cannot be delegated to, # since it re-derives the invocation context and would drop node_input. - async for event in self.run_async(parent_context=parent_context): - if event.author: - ctx.event_author = event.author - if not event.node_info.path and event.author == self.name: - event.node_info.path = ctx.node_path - text = final_model_text(event, self.name) - if text is not None: - last_text = text - yield event + async with Aclosing(self.run_async(parent_context=parent_context)) as agen: + async for event in agen: + if event.author: + ctx.event_author = event.author + if not event.node_info.path and event.author == self.name: + event.node_info.path = ctx.node_path + text = final_model_text(event, self.name) + if text is not None: + last_text = text + yield event # Both assignments are needed: NodeRunner._enrich_event reads # ctx.event_author, and a direct consumer reads author=. diff --git a/tests/unittests/agents/test_base_agent.py b/tests/unittests/agents/test_base_agent.py index f13136c8fc0..df0937105ca 100644 --- a/tests/unittests/agents/test_base_agent.py +++ b/tests/unittests/agents/test_base_agent.py @@ -29,6 +29,7 @@ from google.adk.agents.base_agent import BaseAgent from google.adk.agents.base_agent import BaseAgentState from google.adk.agents.callback_context import CallbackContext +from google.adk.agents.context import Context from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.llm_agent import LlmAgent from google.adk.apps.app import ResumabilityConfig @@ -172,6 +173,44 @@ async def _create_parent_invocation_context( ) +@pytest.mark.asyncio +async def test_run_impl_closes_run_async_when_consumer_stops_early(): + """A consumer that stops early closes run_async before aclose() returns. + + Otherwise run_async is left to the asyncgen finalizer hook, and its finally + blocks and after_agent_callback unwinding run at some later point. + """ + cleaned_up = [] + + class _CleanupAgent(BaseAgent): + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + try: + for index in range(3): + yield Event( + author=self.name, + branch=ctx.branch, + invocation_id=ctx.invocation_id, + content=types.Content(parts=[types.Part(text=f'e{index}')]), + ) + finally: + cleaned_up.append(True) + + agent = _CleanupAgent(name='cleanup_agent') + parent_ctx = await _create_parent_invocation_context( + 'test_run_impl_closes_run_async_when_consumer_stops_early', agent + ) + + agen = agent._run_impl(ctx=Context(parent_ctx), node_input=None) + await agen.__anext__() + await agen.aclose() + + assert cleaned_up == [True] + + def test_invalid_agent_name(): with pytest.raises(ValueError): _ = _TestingAgent(name='not an identifier') diff --git a/tests/unittests/agents/test_managed_agent.py b/tests/unittests/agents/test_managed_agent.py index 98c4dcc9e9c..2b42a0e3186 100644 --- a/tests/unittests/agents/test_managed_agent.py +++ b/tests/unittests/agents/test_managed_agent.py @@ -17,6 +17,7 @@ import sys from typing import Any from unittest.mock import MagicMock +from unittest.mock import patch from google.adk.agents._managed_agent import ManagedAgent from google.adk.agents.run_config import RunConfig @@ -1207,3 +1208,33 @@ def test_string_instruction_injects_session_state(): client.aio.interactions.calls[0]['system_instruction'] == 'Discuss volcanoes.' ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('node_input', [None, 'Fix bug 42.']) +async def test_run_impl_closes_run_async_when_consumer_stops_early(node_input): + cleaned_up = [] + + async def _run_async(self, parent_context): + try: + for index in range(3): + yield Event( + author=self.name, + invocation_id='inv', + content=types.Content(parts=[types.Part(text=f'e{index}')]), + ) + finally: + cleaned_up.append(True) + + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + ctx = MagicMock() + ctx.node_path = 'root/mgr' + + with patch.object(ManagedAgent, 'run_async', _run_async): + agen = agent._run_impl(ctx=ctx, node_input=node_input) + await agen.__anext__() + await agen.aclose() + + assert cleaned_up == [True] diff --git a/tests/unittests/labs/antigravity/test_antigravity_agent.py b/tests/unittests/labs/antigravity/test_antigravity_agent.py index 72785ebb1ad..857d16b6545 100644 --- a/tests/unittests/labs/antigravity/test_antigravity_agent.py +++ b/tests/unittests/labs/antigravity/test_antigravity_agent.py @@ -1851,3 +1851,32 @@ async def test_absent_last_turn_usage_leaves_usage_metadata_none(): assert len(events) == 1 assert events[-1].usage_metadata is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize('node_input', [None, 'Fix bug 42.']) +async def test_run_impl_closes_run_async_when_consumer_stops_early(node_input): + cleaned_up = [] + + async def _run_async(self, parent_context): + try: + for index in range(3): + yield Event( + author=self.name, + invocation_id='inv', + content=genai_types.Content( + parts=[genai_types.Part(text=f'e{index}')] + ), + ) + finally: + cleaned_up.append(True) + + agent = AntigravityAgent(name='agy', config=_make_config()) + ctx = await _node_ctx(agent=agent) + + with patch.object(AntigravityAgent, 'run_async', _run_async): + agen = agent._run_impl(ctx=ctx, node_input=node_input) + await agen.__anext__() + await agen.aclose() + + assert cleaned_up == [True]