|
22 | 22 | TextPart, |
23 | 23 | ) |
24 | 24 | from google.adk.runners import Runner |
| 25 | +from google.adk.utils.context_utils import Aclosing |
25 | 26 | from opentelemetry import trace |
26 | 27 | from pydantic import BaseModel |
27 | 28 | from typing_extensions import override |
@@ -145,17 +146,6 @@ async def execute( |
145 | 146 | ) |
146 | 147 | except Exception as enqueue_error: |
147 | 148 | logger.error("Failed to publish failure event: %s", enqueue_error, exc_info=True) |
148 | | - finally: |
149 | | - # Shield cleanup from external cancellation so toolsets (e.g., MCP) can |
150 | | - # gracefully close their sessions without being torn down mid-flight. |
151 | | - try: |
152 | | - await asyncio.wait_for(asyncio.shield(runner.close()), timeout=15.0) |
153 | | - except asyncio.CancelledError: |
154 | | - # Suppress cancellation during cleanup to avoid noisy tracebacks |
155 | | - # from libraries that assume non-cancelled close semantics. |
156 | | - logger.warning("Runner.close() was cancelled; suppressing during cleanup") |
157 | | - except Exception as close_error: |
158 | | - logger.error("Error during runner.close(): %s", close_error, exc_info=True) |
159 | 149 |
|
160 | 150 | async def _handle_request( |
161 | 151 | self, |
@@ -203,12 +193,13 @@ async def _handle_request( |
203 | 193 | ) |
204 | 194 |
|
205 | 195 | task_result_aggregator = TaskResultAggregator() |
206 | | - async for adk_event in runner.run_async(**run_args): |
207 | | - for a2a_event in convert_event_to_a2a_events( |
208 | | - adk_event, invocation_context, context.task_id, context.context_id |
209 | | - ): |
210 | | - task_result_aggregator.process_event(a2a_event) |
211 | | - await event_queue.enqueue_event(a2a_event) |
| 196 | + async with Aclosing(runner.run_async(**run_args)) as agen: |
| 197 | + async for adk_event in agen: |
| 198 | + for a2a_event in convert_event_to_a2a_events( |
| 199 | + adk_event, invocation_context, context.task_id, context.context_id |
| 200 | + ): |
| 201 | + task_result_aggregator.process_event(a2a_event) |
| 202 | + await event_queue.enqueue_event(a2a_event) |
212 | 203 |
|
213 | 204 | # publish the task result event - this is final |
214 | 205 | if ( |
|
0 commit comments