"""find_agent 异步运行与资源清理。""" from __future__ import annotations import asyncio import logging import threading from agents.find_agent.completion_guard import ( configure_find_agent_completion_guard, ) from supply_agent.agent.core import Agent from supply_agent.types import AgentResult logger = logging.getLogger(__name__) _ASYNC_CLEANUP_TIMEOUT_SECONDS = 5.0 async def _drain_event_loop() -> None: """Bound cleanup of httpx/OpenAI background tasks.""" loop = asyncio.get_running_loop() pending = { task for task in asyncio.all_tasks(loop) if task is not asyncio.current_task() and not task.done() } if not pending: return done, still_pending = await asyncio.wait( pending, timeout=_ASYNC_CLEANUP_TIMEOUT_SECONDS, ) for task in done: if task.cancelled(): continue error = task.exception() if error is not None: logger.debug("pending async task error during drain: %s", error) if still_pending: logger.warning( "Cancelling %d async cleanup task(s) after %.0fs", len(still_pending), _ASYNC_CLEANUP_TIMEOUT_SECONDS, ) for task in still_pending: task.cancel() await asyncio.wait(still_pending, timeout=1.0) async def _close_agent_async_resources(agent: Agent) -> None: """在事件循环关闭前显式释放异步 HTTP 连接。""" async_client = getattr(agent.llm, "_async_client", None) if async_client is not None: try: await asyncio.wait_for( async_client.close(), timeout=_ASYNC_CLEANUP_TIMEOUT_SECONDS, ) except TimeoutError: logger.warning( "close async llm client timed out after %.0fs", _ASYNC_CLEANUP_TIMEOUT_SECONDS, ) except Exception: logger.debug("close async llm client failed", exc_info=True) await _drain_event_loop() async def arun_find_agent( agent: Agent, user_input: str, *, run_id: str, timeout_seconds: float, ) -> AgentResult: """在单个事件循环内运行 find_agent 核心循环并确保资源释放(不含 OSS 发布)。""" configure_find_agent_completion_guard(agent, run_id) try: return await asyncio.wait_for( agent.arun_core(user_input), timeout=timeout_seconds, ) except TimeoutError: logger.error("find_agent timed out after %.0fs", timeout_seconds) raise TimeoutError( f"find_agent timed out after {timeout_seconds:.0f}s" ) from None finally: await _close_agent_async_resources(agent) def _shutdown_worker_loop(loop: asyncio.AbstractEventLoop) -> None: """Release async generators/executor before closing a worker-thread loop.""" shutdown_timeout = 3.0 try: loop.run_until_complete( asyncio.wait_for(loop.shutdown_asyncgens(), timeout=shutdown_timeout) ) except Exception: logger.debug("shutdown_asyncgens failed", exc_info=True) shutdown_executor = getattr(loop, "shutdown_default_executor", None) if shutdown_executor is not None: try: loop.run_until_complete( asyncio.wait_for( shutdown_executor(), timeout=shutdown_timeout, ) ) except Exception: logger.debug("shutdown_default_executor failed", exc_info=True) def _run_coroutine(coro) -> AgentResult: """同步入口:主线程用 asyncio.run;worker 线程每次新建并关闭独立 loop。""" try: asyncio.get_running_loop() except RuntimeError: pass else: raise RuntimeError("run_find_agent 不能在已运行的事件循环内调用") if threading.current_thread() is threading.main_thread(): return asyncio.run(coro) loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: return loop.run_until_complete(coro) finally: try: _shutdown_worker_loop(loop) finally: loop.close() asyncio.set_event_loop(None)