fix(agent/registry): Improve tool execution timeout and cancellation handling

This commit is contained in:
2026-06-15 15:24:39 +02:00
parent d3ddefdc27
commit 8a0d825274

View File

@@ -214,13 +214,26 @@ class NativeToolProvider:
raise asyncio.CancelledError()
if definition.llm_call and context.consume_llm_call:
context.consume_llm_call()
handler_task = asyncio.create_task(definition.handler(arguments, context))
cancellation_task = asyncio.create_task(context.cancellation_event.wait())
try:
result = await asyncio.wait_for(
definition.handler(arguments, context),
done, _pending = await asyncio.wait(
{handler_task, cancellation_task},
timeout=definition.timeout_seconds,
return_when=asyncio.FIRST_COMPLETED,
)
except asyncio.TimeoutError as exc:
raise TimeoutError(f"Tool {name} exceeded {definition.timeout_seconds:g}s") from exc
if cancellation_task in done and cancellation_task.result():
handler_task.cancel()
await asyncio.gather(handler_task, return_exceptions=True)
raise asyncio.CancelledError()
if handler_task not in done:
handler_task.cancel()
await asyncio.gather(handler_task, return_exceptions=True)
raise TimeoutError(f"Tool {name} exceeded {definition.timeout_seconds:g}s")
result = handler_task.result()
finally:
cancellation_task.cancel()
await asyncio.gather(cancellation_task, return_exceptions=True)
validate_json_schema(result, definition.output_schema)
size = len(json.dumps(result, ensure_ascii=False, default=str))
if size > definition.result_size_limit: