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: