fix(agents): close the inner run_async when an agent node stops early

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

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(

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=.

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')

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]

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]