"""网页、独立脚本和集成测试共用的进程级运行环境。""" from __future__ import annotations import asyncio from contextlib import AsyncExitStack import os import threading import time 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 extensions.services import ServiceContainer 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.debug("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, business_jobs: bool = False) -> None: self.scheduler = scheduler self.prewarm = prewarm self.billing_retry = billing_retry self.recover_interrupted = recover_interrupted self.business_jobs = business_jobs 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) started = time.perf_counter() try: setup_logging(config.LOG_LEVEL, config.LOG_FORMAT) logger.info("运行环境启动", extra={"kv": {"event": "runtime.starting", "version": config.APP_VERSION, "dialect": config.DB_DIALECT, "model": config.MODEL, "scheduler": self.scheduler}}) config.log_config_sources() _check_single_process() catalog = load_extensions() approval_names() logger.info("扩展加载完成", extra={"kv": {"event": "extensions.loaded", "modules": catalog.modules, "disabled": catalog.disabled_modules, "tools": len(catalog.tools), "agents": len(catalog.agents)}}) 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() 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) services = ServiceContainer(catalog.services) stack.push_async_callback(services.close) task = asyncio.create_task(services.start(), name="runtime-services-init") stack.push_async_callback(_join, task) await asyncio.shield(task) stack.callback(agent_registry.invalidate) catalog.check_middleware_factories() 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 self.business_jobs: from extensions.services import get_service from runtime.business_tasks import BusinessTaskController gateway = get_service("shaanxi_gateway") if "shaanxi_gateway" in catalog.services else None publisher = get_service("result_platform") if "result_platform" in catalog.services else None controller = BusinessTaskController(gateway, publisher=publisher) stack.push_async_callback(controller.close) task = asyncio.create_task(controller.start(), name="runtime-business-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("运行环境就绪", extra={"kv": {"event": "runtime.ready", "scheduler": self.scheduler, "duration_ms": round((time.perf_counter() - started) * 1000, 1)}}) return self except BaseException as exc: if isinstance(exc, asyncio.CancelledError): logger.warning("运行环境启动已取消", extra={"kv": {"event": "runtime.cancelled"}}) else: logger.exception("运行环境启动失败", extra={"kv": {"event": "runtime.startup_failed"}}) 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("运行环境已关闭", extra={"kv": {"event": "runtime.closed"}}) if cancelled: raise asyncio.CancelledError async def __aexit__(self, exc_type, exc, tb) -> None: await self._close()