From ddeb2ba00d614003c21c7e219f23da8690422b7e Mon Sep 17 00:00:00 2001 From: Atul Joshi <120785343+AtulJoshi1206@users.noreply.github.com> Date: Sat, 19 Sep 2026 09:22:32 +0530 Subject: [PATCH 1/2] fix(agents): close the inner run_async when an agent node stops early BaseAgent._run_impl iterated run_async with a bare async for, so when BaseNode.run's Aclosing closed it early the inner generator was dropped rather than closed. Its cleanup then fell to the asyncgen finalizer hook, which resumes it in a different contextvars context, and the OTel span it is suspended in fails to detach: "ValueError: Token was created in a different Context". after_agent_callback and the exit-stack teardown ran out of band with it. Wrap it in Aclosing, which this module already imports and which both BaseNode.run and the LlmAgent._run_impl override already use. --- src/google/adk/agents/base_agent.py | 22 ++++++++----- tests/unittests/agents/test_base_agent.py | 40 +++++++++++++++++++++++ 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index 80a2dfa95cc..847626eb1aa 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 a consumer that stops early closes run_async here + # rather than leaving it to the asyncgen finalizer hook, which resumes it + # in a different contextvars context and breaks its OTel span teardown. + 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/tests/unittests/agents/test_base_agent.py b/tests/unittests/agents/test_base_agent.py index f13136c8fc0..296c2ff2ac8 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,45 @@ 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. + + Leaving it to the asyncgen finalizer hook resumes the generator in a + different contextvars context, which breaks the OTel span teardown and the + after_agent_callback unwinding it is suspended in. + """ + 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') From 1c7fd498a7e874d566b17205864e6bc313b8c42b Mon Sep 17 00:00:00 2001 From: Atul Joshi <120785343+AtulJoshi1206@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:22:56 +0530 Subject: [PATCH 2/2] fix(agents): close run_async in the ManagedAgent and AntigravityAgent node paths ManagedAgent._run_impl and AntigravityAgent._run_impl mirror BaseAgent._run_impl so they can thread node_input into the invocation context, and they iterated run_async without Aclosing the same way, so an early close left run_async's finally blocks and after_agent_callback unwinding to the asyncgen finalizer hook. Wrap both in Aclosing, and describe the invariant in BaseAgent without the OTel detach symptom, which no longer occurs now that the agent span is attached per step. --- src/google/adk/agents/_managed_agent.py | 15 +++++---- src/google/adk/agents/base_agent.py | 6 ++-- .../labs/antigravity/_antigravity_agent.py | 20 ++++++------ tests/unittests/agents/test_base_agent.py | 5 ++- tests/unittests/agents/test_managed_agent.py | 31 +++++++++++++++++++ .../antigravity/test_antigravity_agent.py | 29 +++++++++++++++++ 6 files changed, 85 insertions(+), 21 deletions(-) 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 847626eb1aa..6bcb047b338 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -392,9 +392,9 @@ async def _run_impl( node_input: Any, ) -> AsyncGenerator[Any, None]: """Runs the agent as a node.""" - # Aclosing, so that a consumer that stops early closes run_async here - # rather than leaving it to the asyncgen finalizer hook, which resumes it - # in a different contextvars context and breaks its OTel span teardown. + # 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: 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 296c2ff2ac8..df0937105ca 100644 --- a/tests/unittests/agents/test_base_agent.py +++ b/tests/unittests/agents/test_base_agent.py @@ -177,9 +177,8 @@ async def _create_parent_invocation_context( async def test_run_impl_closes_run_async_when_consumer_stops_early(): """A consumer that stops early closes run_async before aclose() returns. - Leaving it to the asyncgen finalizer hook resumes the generator in a - different contextvars context, which breaks the OTel span teardown and the - after_agent_callback unwinding it is suspended in. + 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 = [] 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]