vent_agent_2.0/runtime/environment.py

"""网页、独立脚本和集成测试共用的进程级运行环境。"""
from __future__ import annotations

import asyncio
from contextlib import AsyncExitStack
import os
import threading

import config
from agents import agent_memory, agent_registry
from billing import pricing
from db import db_conn, run_store
from extensions.registry import approval_names, load_extensions
from obs import get_logger, setup_logging

logger = get_logger(__name__)
# 底层缓存与连接仍是进程级单例；拒绝嵌套/并行环境，防止一方关闭另一方的资源。
_owner_lock = threading.Lock()


async def _retry_billing(stop: asyncio.Event) -> None:
    from billing.meter import retry_pending
    while not stop.is_set():
        try:
            await asyncio.to_thread(retry_pending)
        except Exception:
            logger.exception("重试待结算流水失败")
        try:
            await asyncio.wait_for(stop.wait(), 60)
        except asyncio.TimeoutError:
            pass


async def _stop_billing(stop: asyncio.Event, task: asyncio.Task) -> None:
    stop.set()
    # 不取消 to_thread：必须等数据库写入完成才能关闭连接。
    await task


async def _prewarm() -> None:
    ok = await agent_registry.prewarm()
    logger.info("agent 预热%s", "完成" if ok else "失败（将在首次请求时重试）")


async def _join(task: asyncio.Task) -> None:
    # 初始化和编译可能在线程中进行，取消协程并不能停止底层工作。
    await task


def _check_single_process() -> None:
    workers = os.getenv("WEB_CONCURRENCY", "")
    try:
        count = int(workers) if workers else 1
    except ValueError:
        count = 1
    if count > 1:
        logger.warning(
            "检测到 WEB_CONCURRENCY=%s：checkpoint / 并发闸 / 调度器 "
            "使用进程级状态，生产请用 --workers 1。", workers,
        )


class RuntimeEnvironment:
    """使用 ``async with`` 管理一次完整生命周期。

    默认加载扩展、建表、初始化定价/checkpoint、恢复中断记录、后台预热和重试计费。
    独立脚本默认不启动调度器；网页入口显式传 ``scheduler=True``。
    同一进程只允许一个活动环境，可在退出后重新进入；退出前调用方必须等待
    自己创建的执行任务结束（或取消并等待收尾），再释放共享资源。
    """

    def __init__(self, *, scheduler: bool = False, prewarm: bool = True,
                 billing_retry: bool = True, recover_interrupted: bool = True) -> None:
        self.scheduler = scheduler
        self.prewarm = prewarm
        self.billing_retry = billing_retry
        self.recover_interrupted = recover_interrupted
        self._stack: AsyncExitStack | None = None

    async def __aenter__(self) -> RuntimeEnvironment:
        if not _owner_lock.acquire(blocking=False):
            raise RuntimeError("已有活动的 RuntimeEnvironment，请复用外层环境")
        stack = self._stack = AsyncExitStack()
        stack.callback(_owner_lock.release)
        try:
            setup_logging(config.LOG_LEVEL, config.LOG_FORMAT)
            config.log_config_sources()
            _check_single_process()
            load_extensions()
            approval_names()

            stack.callback(db_conn.close)
            db_conn.init_schema()
            logger.info("存储就绪 %s", db_conn.health())
            pricing.seed_default_pricing()
            if self.recover_interrupted:
                run_store.recover_interrupted()

            # 注册在初始化前，部分初始化失败也会清理。
            stack.push_async_callback(agent_memory.close_checkpointer)
            agent_registry.invalidate()
            stack.callback(agent_registry.invalidate)
            init_task = asyncio.create_task(agent_memory.init_checkpointer(), name="runtime-checkpoint-init")
            stack.push_async_callback(_join, init_task)
            # aiosqlite 在线程中连接；取消调用方时先等初始化结束，再关闭连接。
            await asyncio.shield(init_task)

            if self.prewarm:
                task = asyncio.create_task(_prewarm(), name="runtime-prewarm")
                stack.push_async_callback(_join, task)
            if self.billing_retry:
                stop = asyncio.Event()
                task = asyncio.create_task(_retry_billing(stop), name="runtime-billing")
                stack.push_async_callback(_stop_billing, stop, task)
            if self.scheduler:
                from scheduler import scheduler
                stack.push_async_callback(scheduler.shutdown)
                task = asyncio.create_task(scheduler.start(), name="runtime-scheduler-init")
                stack.push_async_callback(_join, task)
                await asyncio.shield(task)

            if not config.OPENAI_API_KEY:
                logger.warning("OPENAI_API_KEY 未配置，对话会失败——请在 .env 里设置")
            logger.info("运行环境就绪 scheduler=%s", self.scheduler)
            return self
        except BaseException:
            await self._close()
            raise

    async def _close(self) -> None:
        stack = self._stack
        if stack is None:
            return
        # shield 保证取消（含二次取消）时仍完成整条资源清理链。
        task = asyncio.create_task(stack.aclose(), name="runtime-close")
        cancelled = False
        try:
            while not task.done():
                try:
                    await asyncio.shield(task)
                except asyncio.CancelledError:
                    cancelled = True
            task.result()
        finally:
            self._stack = None
        logger.info("运行环境已关闭")
        if cancelled:
            raise asyncio.CancelledError

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self._close()