fix: aclose all async generators to fix OTel tracing context

See #1670 (comment) PiperOrigin-RevId: 794659547

Author: jawoszek

contributing/samples/telemetry/main.py

@@ -46,13 +46,19 @@ async def run_prompt(session: Session, new_message: str):
  role='user', parts=[types.Part.from_text(text=new_message)]
  )
  print('** User says:', content.model_dump(exclude_none=True))
- async for event in runner.run_async(
+ # TODO - migrate try...finally to contextlib.aclosing after Python 3.9 is
+ # no longer supported.
+ agen = runner.run_async(
  user_id=user_id_1,
  session_id=session.id,
  new_message=content,
- ):
- if event.content.parts and event.content.parts[0].text:
- print(f'** {event.author}: {event.content.parts[0].text}')
+ )
+ try:
+ async for event in agen:
+ if event.content.parts and event.content.parts[0].text:
+ print(f'** {event.author}: {event.content.parts[0].text}')
+ finally:
+ await agen.aclose()
 
  async def run_prompt_bytes(session: Session, new_message: str):
  content = types.Content(
@@ -64,14 +70,20 @@ async def run_prompt_bytes(session: Session, new_message: str):
  ],
  )
  print('** User says:', content.model_dump(exclude_none=True))
- async for event in runner.run_async(
+ # TODO - migrate try...finally to contextlib.aclosing after Python 3.9 is
+ # no longer supported.
+ agen = runner.run_async(
  user_id=user_id_1,
  session_id=session.id,
  new_message=content,
  run_config=RunConfig(save_input_blobs_as_artifacts=True),
- ):
- if event.content.parts and event.content.parts[0].text:
- print(f'** {event.author}: {event.content.parts[0].text}')
+ )
+ try:
+ async for event in agen:
+ if event.content.parts and event.content.parts[0].text:
+ print(f'** {event.author}: {event.content.parts[0].text}')
+ finally:
+ await agen.aclose()
 
  start_time = time.time()
  print('Start time:', start_time)

src/google/adk/a2a/executor/a2a_agent_executor.py

@@ -24,6 +24,8 @@
 from typing import Optional
 import uuid
 
+from ...utils.context_utils import Aclosing
+
 try:
  from a2a.server.agent_execution import AgentExecutor
  from a2a.server.agent_execution.context import RequestContext
@@ -212,12 +214,13 @@ async def _handle_request(
  )
 
  task_result_aggregator = TaskResultAggregator()
- async for adk_event in runner.run_async(**run_args):
- for a2a_event in convert_event_to_a2a_events(
- adk_event, invocation_context, context.task_id, context.context_id
- ):
- task_result_aggregator.process_event(a2a_event)
- await event_queue.enqueue_event(a2a_event)
+ async with Aclosing(runner.run_async(**run_args)) as agen:
+ async for adk_event in agen:
+ for a2a_event in convert_event_to_a2a_events(
+ adk_event, invocation_context, context.task_id, context.context_id
+ ):
+ task_result_aggregator.process_event(a2a_event)
+ await event_queue.enqueue_event(a2a_event)
 
  # publish the task result event - this is final
  if (

src/google/adk/agents/base_agent.py

@@ -39,6 +39,7 @@
 from typing_extensions import TypeAlias
 
 from ..events.event import Event
+from ..utils.context_utils import Aclosing
 from ..utils.feature_decorator import experimental
 from .base_agent_config import BaseAgentConfig
 from .callback_context import CallbackContext
@@ -212,21 +213,27 @@ async def run_async(
  Event: the events generated by the agent.
  """
 
- with tracer.start_as_current_span(f'agent_run [{self.name}]'):
- ctx = self._create_invocation_context(parent_context)
+ async def _run_with_trace() -> AsyncGenerator[Event, None]:
+ with tracer.start_as_current_span(f'agent_run [{self.name}]'):
+ ctx = self._create_invocation_context(parent_context)
 
- if event := await self.__handle_before_agent_callback(ctx):
- yield event
- if ctx.end_invocation:
- return
+  if event := await self.__handle_before_agent_callback(ctx):
+  yield event
+  if ctx.end_invocation:
+  return
 
- async for event in self._run_async_impl(ctx):
- yield event
+ async with Aclosing(self._run_async_impl(ctx)) as agen:
+ async for event in agen:
+ yield event
 
- if ctx.end_invocation:
- return
+  if ctx.end_invocation:
+  return
 
- if event := await self.__handle_after_agent_callback(ctx):
+ if event := await self.__handle_after_agent_callback(ctx):
+ yield event
+
+ async with Aclosing(_run_with_trace()) as agen:
+ async for event in agen:
  yield event
 
  @final
@@ -243,18 +250,25 @@ async def run_live(
  Yields:
  Event: the events generated by the agent.
  """
- with tracer.start_as_current_span(f'agent_run [{self.name}]'):
- ctx = self._create_invocation_context(parent_context)
 
- if event := await self.__handle_before_agent_callback(ctx):
- yield event
- if ctx.end_invocation:
- return
+ async def _run_with_trace() -> AsyncGenerator[Event, None]:
+ with tracer.start_as_current_span(f'agent_run [{self.name}]'):
+ ctx = self._create_invocation_context(parent_context)
 
- async for event in self._run_live_impl(ctx):
- yield event
+ if event := await self.__handle_before_agent_callback(ctx):
+ yield event
+ if ctx.end_invocation:
+ return
+
+ async with Aclosing(self._run_live_impl(ctx)) as agen:
+ async for event in agen:
+ yield event
+
+ if event := await self.__handle_after_agent_callback(ctx):
+ yield event
 
- if event := await self.__handle_after_agent_callback(ctx):
+ async with Aclosing(_run_with_trace()) as agen:
+ async for event in agen:
  yield event
 
  async def _run_async_impl(

src/google/adk/agents/llm_agent.py

@@ -51,6 +51,7 @@
 from ..tools.function_tool import FunctionTool
 from ..tools.tool_configs import ToolConfig
 from ..tools.tool_context import ToolContext
+from ..utils.context_utils import Aclosing
 from ..utils.feature_decorator import experimental
 from .base_agent import BaseAgent
 from .base_agent_config import BaseAgentConfig
@@ -283,19 +284,21 @@ class LlmAgent(BaseAgent):
  async def _run_async_impl(
  self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
- async for event in self._llm_flow.run_async(ctx):
- self.__maybe_save_output_to_state(event)
- yield event
+ async with Aclosing(self._llm_flow.run_async(ctx)) as agen:
+ async for event in agen:
+ self.__maybe_save_output_to_state(event)
+ yield event
 
  @override
  async def _run_live_impl(
  self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
- async for event in self._llm_flow.run_live(ctx):
- self.__maybe_save_output_to_state(event)
- yield event
- if ctx.end_invocation:
- return
+ async with Aclosing(self._llm_flow.run_live(ctx)) as agen:
+ async for event in agen:
+ self.__maybe_save_output_to_state(event)
+ yield event
+ if ctx.end_invocation:
+ return
 
  @property
  def canonical_model(self) -> BaseLlm:

src/google/adk/agents/loop_agent.py

@@ -27,6 +27,7 @@
 
 from ..agents.invocation_context import InvocationContext
 from ..events.event import Event
+from ..utils.context_utils import Aclosing
 from ..utils.feature_decorator import experimental
 from .base_agent import BaseAgent
 from .base_agent_config import BaseAgentConfig
@@ -58,10 +59,11 @@ async def _run_async_impl(
  while not self.max_iterations or times_looped < self.max_iterations:
  for sub_agent in self.sub_agents:
  should_exit = False
- async for event in sub_agent.run_async(ctx):
- yield event
- if event.actions.escalate:
- should_exit = True
+ async with Aclosing(sub_agent.run_async(ctx)) as agen:
+ async for event in agen:
+ yield event
+ if event.actions.escalate:
+ should_exit = True
 
  if should_exit:
  return

src/google/adk/agents/parallel_agent.py

@@ -26,6 +26,7 @@
 from typing_extensions import override
 
 from ..events.event import Event
+from ..utils.context_utils import Aclosing
 from .base_agent import BaseAgent
 from .base_agent_config import BaseAgentConfig
 from .invocation_context import InvocationContext
@@ -111,8 +112,10 @@ async def _run_async_impl(
  )
  for sub_agent in self.sub_agents
  ]
- async for event in _merge_agent_run(agent_runs):
- yield event
+
+ async with Aclosing(_merge_agent_run(agent_runs)) as agen:
+ async for event in agen:
+ yield event
 
  @override
  async def _run_live_impl(

src/google/adk/agents/sequential_agent.py

@@ -22,6 +22,7 @@
 from typing_extensions import override
 
 from ..events.event import Event
+from ..utils.context_utils import Aclosing
 from .base_agent import BaseAgent
 from .base_agent import BaseAgentConfig
 from .invocation_context import InvocationContext
@@ -40,8 +41,9 @@ async def _run_async_impl(
  self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
  for sub_agent in self.sub_agents:
- async for event in sub_agent.run_async(ctx):
- yield event
+ async with Aclosing(sub_agent.run_async(ctx)) as agen:
+ async for event in agen:
+ yield event
 
  @override
  async def _run_live_impl(
@@ -78,5 +80,6 @@ def task_completed():
  do not generate any text other than the function call."""
 
  for sub_agent in self.sub_agents:
- async for event in sub_agent.run_live(ctx):
- yield event
+ async with Aclosing(sub_agent.run_live(ctx)) as agen:
+ async for event in agen:
+ yield event

src/google/adk/cli/adk_web_server.py

@@ -73,6 +73,7 @@
 from ..runners import Runner
 from ..sessions.base_session_service import BaseSessionService
 from ..sessions.session import Session
+from ..utils.context_utils import Aclosing
 from .cli_eval import EVAL_SESSION_ID_PREFIX
 from .cli_eval import EvalStatus
 from .utils import cleanup
@@ -828,14 +829,16 @@ async def agent_run(req: AgentRunRequest) -> list[Event]:
  if not session:
  raise HTTPException(status_code=404, detail="Session not found")
  runner = await self.get_runner_async(req.app_name)
- events = [
- event
- async for event in runner.run_async(
+
+ events = []
+ async with Aclosing(
+ runner.run_async(
  user_id=req.user_id,
  session_id=req.session_id,
  new_message=req.new_message,
  )
- ]
+ ) as agen:
+ events = [event async for event in agen]
  logger.info("Generated %s events in agent run", len(events))
  logger.debug("Events generated: %s", events)
  return events
@@ -856,19 +859,24 @@ async def event_generator():
  StreamingMode.SSE if req.streaming else StreamingMode.NONE
  )
  runner = await self.get_runner_async(req.app_name)
- async for event in runner.run_async(
- user_id=req.user_id,
- session_id=req.session_id,
- new_message=req.new_message,
- state_delta=req.state_delta,
- run_config=RunConfig(streaming_mode=stream_mode),
- ):
- # Format as SSE data
- sse_event = event.model_dump_json(exclude_none=True, by_alias=True)
- logger.debug(
- "Generated event in agent run streaming: %s", sse_event
- )
- yield f"data: {sse_event}\n\n"
+ async with Aclosing(
+ runner.run_async(
+ user_id=req.user_id,
+ session_id=req.session_id,
+ new_message=req.new_message,
+ state_delta=req.state_delta,
+ run_config=RunConfig(streaming_mode=stream_mode),
+ )
+ ) as agen:
+ async for event in agen:
+ # Format as SSE data
+ sse_event = event.model_dump_json(
+ exclude_none=True, by_alias=True
+ )
+ logger.debug(
+ "Generated event in agent run streaming: %s", sse_event
+ )
+ yield f"data: {sse_event}\n\n"
  except Exception as e:
  logger.exception("Error in event_generator: %s", e)
  # You might want to yield an error event here
@@ -954,12 +962,15 @@ async def agent_live_run(
 
  async def forward_events():
  runner = await self.get_runner_async(app_name)
- async for event in runner.run_live(
- session=session, live_request_queue=live_request_queue
- ):
- await websocket.send_text(
- event.model_dump_json(exclude_none=True, by_alias=True)
- )
+ async with Aclosing(
+ runner.run_live(
+ session=session, live_request_queue=live_request_queue
+ )
+ ) as agen:
+ async for event in agen:
+ await websocket.send_text(
+ event.model_dump_json(exclude_none=True, by_alias=True)
+ )
 
  async def process_messages():
  try: