"""find_agent 异步运行与资源清理。""" from __future__ import annotations import asyncio import logging import threading from supply_agent.agent.core import Agent from supply_agent.types import AgentResult logger = logging.getLogger(__name__) _thread_loop = threading.local() async def _drain_event_loop() -> None: """等待 httpx/OpenAI 等异步客户端的收尾任务完成。""" loop = asyncio.get_running_loop() for _ in range(5): pending = [ task for task in asyncio.all_tasks(loop) if task is not asyncio.current_task() and not task.done() ] if not pending: break results = await asyncio.gather(*pending, return_exceptions=True) for result in results: if isinstance(result, Exception) and not isinstance( result, asyncio.CancelledError ): logger.debug("pending async task error during drain: %s", result) await asyncio.sleep(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 async_client.close() 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) -> AgentResult: """在单个事件循环内运行 find_agent 并确保资源释放。""" try: return await agent.arun(user_input) finally: await _close_agent_async_resources(agent) 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 = getattr(_thread_loop, "loop", None) if loop is None or loop.is_closed(): loop = asyncio.new_event_loop() _thread_loop.loop = loop asyncio.set_event_loop(loop) return loop.run_until_complete(coro)