diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index abcfbe54ed9..5c7a6deab32 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -336,16 +336,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 - - if not event.node_info.path and event.author == self.name: - event.node_info.path = ctx.node_path - yield event + # 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 @final async def run_live( diff --git a/src/google/adk/cli/browser/assets/config/runtime-config.json b/src/google/adk/cli/browser/assets/config/runtime-config.json index 888614ad561..873e88b1f1d 100644 --- a/src/google/adk/cli/browser/assets/config/runtime-config.json +++ b/src/google/adk/cli/browser/assets/config/runtime-config.json @@ -1,4 +1,4 @@ { "backendUrl": "", - "telemetry": false + "telemetry": null } diff --git a/tests/unittests/agents/test_base_agent.py b/tests/unittests/agents/test_base_agent.py index 7a8e0ba09aa..d1eab509655 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 @@ -171,6 +172,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')