Skip to content

Commit 8245a0a

Browse files
authored
fixes [BUG] kagent 100% CPU and not responding well to queries kagent-dev#911 (kagent-dev#956)
fixes kagent-dev#911 The runner.close was tearing down MCP sessions from the wrong task; that triggered cancel-scope errors and a cancellation storm causing the CPU to get to 1000m and stay there. also wrapping the generator with aclosing, so it cleans up the generator safely. Signed-off-by: Peter Jausovec <peter.jausovec@solo.io>
1 parent d55ce27 commit 8245a0a

1 file changed

Lines changed: 8 additions & 17 deletions

File tree

python/packages/kagent-adk/src/kagent/adk/_agent_executor.py

Lines changed: 8 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
TextPart,
2323
)
2424
from google.adk.runners import Runner
25+
from google.adk.utils.context_utils import Aclosing
2526
from opentelemetry import trace
2627
from pydantic import BaseModel
2728
from typing_extensions import override
@@ -145,17 +146,6 @@ async def execute(
145146
)
146147
except Exception as enqueue_error:
147148
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)
159149

160150
async def _handle_request(
161151
self,
@@ -203,12 +193,13 @@ async def _handle_request(
203193
)
204194

205195
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)
212203

213204
# publish the task result event - this is final
214205
if (

0 commit comments

Comments
 (0)