SupplyAgent/agents/find_agent/async_runner.py

"""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)