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_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,
    *,
    timeout_seconds: float,
) -> AgentResult:
    """在单个事件循环内运行 find_agent 并确保资源释放。"""
    try:
        return await asyncio.wait_for(
            agent.arun(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 _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)