Răsfoiți Sursa

refactor: 统一运行环境初始化与资源清理

gaoshutao 1 lună în urmă
părinte
comite
e7c5168cbe
10 a modificat fișierele cu 591 adăugiri și 157 ștergeri
  1. 1 1
      README.md
  2. 9 5
      agents/agent_memory.py
  3. 42 0
      db/db_conn.py
  4. 44 17
      docs/DEVELOPMENT.md
  5. 6 120
      main.py
  6. 151 0
      runtime/environment.py
  7. 13 13
      scripts/selfcheck.py
  8. 322 0
      tests/test_environment.py
  9. 1 0
      tests/test_execution.py
  10. 2 1
      tests/test_extensions.py

+ 1 - 1
README.md

@@ -132,7 +132,7 @@ DEEP_AGENT_ENV_PRIORITY=process LOG_LEVEL=DEBUG ~/vent-agent/current/release/ser
 
 ## 7. 业务扩展与开发验证
 
-聊天、定时任务和批处理共用 `runtime/` 执行层;HTTP 层负责 SSE 编码。非 HTTP 调用方式见 [执行层与入口](docs/DEVELOPMENT.md#执行层与入口)。
+聊天、定时任务和批处理共用 `runtime/` 执行层;HTTP 层负责 SSE 编码。独立脚本使用 `RuntimeEnvironment` 自动初始化和清理资源,默认不启动调度器。调用方式见 [执行层与入口](docs/DEVELOPMENT.md#执行层与入口)。
 
 新增业务时,在 `extensions/<业务名>/extension.py` 编写 `setup(registry)`,统一声明工具、子 Agent、提示词和审批要求。启动自动发现并校验,无需修改工厂函数;修改后重启生效。示例和完整步骤见 [扩展方式](docs/DEVELOPMENT.md#扩展方式)。
 

+ 9 - 5
agents/agent_memory.py

@@ -7,8 +7,7 @@ agents/agent_memory.py 的注释)。2.0 的单进程约束下(uvicorn worker
 只有一个事件循环,不需要这层桥接:启动时创建一次 saver,全进程复用。
 
 代价是 **不能多 worker**:多进程各自持有自己的 sqlite 文件句柄与图状态,
-同一会话的续跑会落到不同进程而看不到对方的状态。这条约束由 main.py 的
-启动检查兜住。
+并发锁与调度器也仍使用进程级状态。RuntimeEnvironment 启动时提示单进程约束。
 """
 
 from __future__ import annotations
@@ -55,7 +54,7 @@ def get_checkpointer() -> Any:
         logger.info("checkpointer=memory 惰性创建(图状态不落盘)")
     if _saver is None:
         raise CheckpointerConfigError(
-            "checkpointer 尚未初始化:请确认已在 lifespan 中 await init_checkpointer()"
+            "checkpointer 尚未初始化:请在 async with RuntimeEnvironment() 内执行任务"
         )
     return _saver
 
@@ -97,11 +96,16 @@ async def init_checkpointer() -> Any:
             ) from exc
 
         path = config.resolve_path(config.CHECKPOINT_DB_PATH)
+        conn = None
         try:
             conn = await aiosqlite.connect(str(path))
             saver = AsyncSqliteSaver(conn)
             await saver.setup()
-        except Exception as exc:
+        except BaseException as exc:
+            if conn is not None:
+                await conn.close()
+            if not isinstance(exc, Exception):
+                raise
             raise CheckpointerConfigError(
                 f"打开 checkpoint 库失败 path={path}:{type(exc).__name__}: {exc}"
             ) from exc
@@ -112,7 +116,7 @@ async def init_checkpointer() -> Any:
 
 
 async def close_checkpointer() -> None:
-    """关闭底层连接(lifespan 退出时调用)。"""
+    """关闭底层连接(RuntimeEnvironment 退出时调用)。"""
     global _saver
     saver = _saver
     _saver = None

+ 42 - 0
db/db_conn.py

@@ -157,6 +157,13 @@ class _SqliteAdapter:
     def ping(self) -> None:
         self._connection().execute("SELECT 1").fetchone()
 
+    def close(self) -> None:
+        with self._lock:
+            if self._conn is not None:
+                self._conn.close()
+                self._conn = None
+            self._transaction_depth = 0
+
 
 # ============================================================
 # MySQL 适配
@@ -172,6 +179,8 @@ class _MySqlAdapter:
     def __init__(self) -> None:
         self._local = threading.local()
         self._last_used: dict[int, float] = {}
+        self._connections: dict[int, Any] = {}
+        self._connections_lock = threading.Lock()
 
     def _connect(self, database: str | None):
         """建一个 pymysql 连接。``database=None`` 表示不选库。"""
@@ -222,12 +231,21 @@ class _MySqlAdapter:
                 else:
                     raise
             self._local.conn = conn
+            with self._connections_lock:
+                self._connections[id(conn)] = conn
         else:
             ident = id(conn)
             if not getattr(self._local, "transaction_depth", 0) and time.time() - self._last_used.get(ident, 0) > config.DB_PING_IDLE_SECONDS:
                 try:
                     conn.ping(reconnect=True)
                 except Exception:
+                    with self._connections_lock:
+                        self._connections.pop(id(conn), None)
+                    self._last_used.pop(id(conn), None)
+                    try:
+                        conn.close()
+                    except Exception:
+                        pass
                     self._local.conn = None
                     return self._connection()
         self._last_used[id(conn)] = time.time()
@@ -297,6 +315,19 @@ class _MySqlAdapter:
     def ping(self) -> None:
         self._connection().ping(reconnect=True)
 
+    def close(self) -> None:
+        """在所有数据库任务结束后关闭各线程连接,下一轮生命周期重新建立。"""
+        with self._connections_lock:
+            connections = list(self._connections.values())
+            self._connections.clear()
+            self._local = threading.local()
+            self._last_used.clear()
+        for conn in connections:
+            try:
+                conn.close()
+            except Exception as exc:
+                logger.warning("关闭 MySQL 连接失败:%s", type(exc).__name__)
+
 
 def _build_adapter() -> _SqliteAdapter | _MySqlAdapter:
     if DIALECT == "sqlite":
@@ -426,8 +457,19 @@ def health() -> dict[str, Any]:
         return {"status": "error", "dialect": DIALECT, "error": f"{type(exc).__name__}: {exc}"}
 
 
+def close() -> None:
+    """关闭业务库连接;只能在执行任务和后台维护全部结束后调用。"""
+    global _schema_ready
+    with _schema_lock:
+        try:
+            _adapter.close()
+        finally:
+            _schema_ready = False
+
+
 def reset_for_tests() -> None:
     """仅供测试:丢弃缓存连接与建表标记,让下一次调用重新初始化。"""
     global _adapter, _schema_ready
+    close()
     _adapter = _build_adapter()
     _schema_ready = False

+ 44 - 17
docs/DEVELOPMENT.md

@@ -171,25 +171,37 @@ HTTP 错误语义:额度超限 429;并发超限 429(带 `retry_after`)
 
 ### 非 HTTP 调用
 
-在已完成数据库建表、模型定价和 checkpoint 初始化的服务生命周期中,对归属当前用户的已有会话执行:
+独立脚本使用公共运行环境;以下示例保存为项目根目录的 Python 文件后即可运行:
 
 ```python
+import asyncio
+
+from db import chat_store
+from runtime.environment import RuntimeEnvironment
 from runtime.execution import execute
 
-result = await execute(
-    user_name="report-user",
-    session_id=session_id,
-    message="生成本次报告",
-    kind="batch",
-    graph_config={"recursion_limit": 30},
-)
-if result.status == "success":
-    print(result.output)
-elif result.status == "waiting_input":
-    # 保存任务状态,等待用户确认,不作为已完成报告发布。
-    print(result.interrupts)
-else:
-    print(result.error)
+
+async def main():
+    async with RuntimeEnvironment():
+        session = await asyncio.to_thread(chat_store.create_session, "report-user")
+        result = await execute(
+            user_name="report-user",
+            session_id=session["id"],
+            message="生成本次报告",
+            kind="batch",
+            graph_config={"recursion_limit": 30},
+        )
+        if result.status == "success":
+            print(result.output)
+        elif result.status == "waiting_input":
+            # 保存任务状态,等待用户确认。
+            print(result.interrupts)
+        else:
+            print(result.error)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
 ```
 
 - `execute()` 复用会话归属校验、额度及并发闸,最后释放资源。已有审批未处理时拒绝开始新任务。
@@ -199,7 +211,22 @@ else:
 - 审批恢复使用 `runtime.interrupts.ResumeDecision` / `resume_value()`,并在持有会话锁期间读取和校验最新图状态,再将结果作为 `resume_payload` 交给 `build_execution()`。
 - 每个 `Execution` 只能消费一次;`run()` 消费事件并清理后返回结果,`events()` 供逐条消费,不能对同一个对象重复调用。
 
-独立脚本需自行初始化数据库、定价和 `agents.agent_memory.init_checkpointer()`,退出时关闭 checkpoint。执行层不会自动启动 HTTP 或调度器;并发锁仍在单进程内,拆分模块不代表支持多进程共享会话。
+### 公共运行环境
+
+`RuntimeEnvironment` 负责加载和校验扩展、初始化业务库与计费单价、标记上次意外中断的执行、打开 checkpoint,并在后台预热 Agent、重试待结算流水。退出时停止可选调度器,等待维护和编译完成,清除 Agent 缓存,关闭 checkpoint 与业务库连接。启动失败或调用方取消时同样清理。
+
+| 参数 | 默认值 | 用途 |
+|---|---|---|
+| `scheduler` | `False` | 网页服务传 `True`;脚本默认不触发定时任务 |
+| `prewarm` | `True` | 后台编译 Agent;关闭后首次执行时编译 |
+| `billing_retry` | `True` | 后台重试待结算记录;测试可关闭,本轮执行仍正常计费 |
+| `recover_interrupted` | `True` | 将上次残留的运行中记录标为中断;自检时关闭 |
+
+- `main.py` 与自检脚本已使用这一入口。网页请求内部直接调用执行层,不要为每个请求创建环境。
+- 同一进程仅允许一个活动环境;嵌套或同时启动会报错。退出后可以重新进入,会重新建立连接和 Agent 缓存。
+- 并行批处理应在环境内部使用 `asyncio.TaskGroup()`;退出前必须等待自己的任务完成,或取消并等待其收尾。
+- 沿用单进程部署约束。独立脚本使用单独的数据配置,或先停止共享同一套数据的服务;关闭恢复选项也不代表支持多进程共享并发锁。
+- 集成测试可以使用 `RuntimeEnvironment(prewarm=False, billing_retry=False, recover_interrupted=False)`,并在导入模块前设置测试数据库与 checkpoint 路径。配置仍按进程加载,环境重入不会重新读取 `.env`。
 
 ---
 
@@ -332,7 +359,7 @@ config.py                  配置唯一入口(.env 读取与冲突检测)
 api/                       HTTP 参数、路由与协议适配
   sse_events.py            SSE 编码(保留前端字段)
   sse_core.py              事件转 SSE、断连关闭内层生成器
-runtime/                   与传输无关的执行层
+runtime/                   公共生命周期与传输无关的执行层
   execution.py            准入、装配、execute() 非 HTTP 入口
   engine.py               ★ 事件流、运行结果、落库与收尾
   events.py               ★ 统一事件字段

+ 6 - 120
main.py

@@ -1,21 +1,9 @@
-"""deep_agent 2.0 服务入口。
-
-启动顺序(lifespan):
-1. 日志
-2. 业务扩展发现与校验;建表(幂等)
-3. 计费单价表种子
-4. checkpointer 初始化
-5. **后台**预热 agent(图编译耗时百毫秒级,不能挡住 uvicorn 就绪)
-6. 调度器启动并装载定时任务
-7. 单进程约束自检
-
-退出时逆序释放:调度器 → checkpointer。
-"""
+"""deep_agent 2.0 HTTP 入口;公共资源生命周期由 RuntimeEnvironment 管理。"""
 
 from __future__ import annotations
 
-import asyncio
 import os
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, RedirectResponse
@@ -24,121 +12,19 @@ from fastapi.staticfiles import StaticFiles
 import config
 from api import router as api_router
 from runtime import concurrency
-from billing import pricing, guard
+from runtime.environment import RuntimeEnvironment
+from billing import guard
 from db import db_conn
 from obs import bind_trace_id, get_logger, new_trace_id, setup_logging, unbind_trace_id
 
 setup_logging(config.LOG_LEVEL, config.LOG_FORMAT)
 logger = get_logger(__name__)
-_maintenance_stop: asyncio.Event | None = None
-_maintenance_task: asyncio.Task | None = None
-_prewarm_task: asyncio.Task | None = None
-
-
-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
-
-
-def _check_single_process() -> None:
-    """多 worker 自检。
-
-    2.0 有三处进程内状态:SQLite checkpointer、并发计数、调度器。多 worker
-    下它们各自独立——同一会话的续跑会落到不同进程导致图状态读不到,并发闸
-    形同虚设,定时任务会被每个 worker 各触发一次。所以这里只警告不阻止:
-    本地调试时多 worker 有时是有意的。
-    """
-    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:SQLite checkpointer / 并发闸 / 调度器 "
-            "都是进程内的,多 worker 会导致会话状态分裂、定时任务重复触发。"
-            "生产请用 --workers 1,或把这三处状态换成 Redis/MySQL 后再扩。",
-            workers,
-        )
-
-
-async def _prewarm() -> None:
-    """后台预热 agent。失败不阻断启动——首次真实请求会重试并暴露问题。"""
-    from agents import agent_registry
-
-    ok = await agent_registry.prewarm()
-    logger.info("agent 预热%s", "完成" if ok else "失败(将在首次请求时重试)")
-
-
-async def _lifespan_startup() -> None:
-    global _maintenance_stop, _maintenance_task, _prewarm_task
-    # 配置来源与冲突在日志装配后就打出来:排查「配置不生效」时它是第一手信息。
-    # 放在 lifespan 而不是模块顶层,是因为 `python main.py` 会让本模块被导入
-    # 两次(`__main__` 与 `main`),模块级副作用会重复执行。
-    config.log_config_sources()
-    logger.info(
-        "deep_agent 2.0 启动中 dialect=%s model=%s port=%s",
-        config.DB_DIALECT, config.MODEL, config.PORT,
-    )
-    _check_single_process()
-    # 扩展定义错误先阻断启动,避免启动调度器后才在后台预热时发现。
-    from extensions.registry import load_extensions, approval_names
-    load_extensions()
-    approval_names()
-
-    # 建表与单价:失败即拒绝启动。半可用的服务比直接失败更难排查——
-    # 用户会遇到「对话能开但一落库就报错」这类现象。
-    db_conn.init_schema()
-    logger.info("存储就绪 %s", db_conn.health())
-    pricing.seed_default_pricing()
-    from db import run_store
-    run_store.recover_interrupted()
-
-    from agents import agent_memory
-
-    await agent_memory.init_checkpointer()
-
-    from scheduler import scheduler
-
-    await scheduler.start()
-
-    _prewarm_task = asyncio.create_task(_prewarm())
-    _maintenance_stop = asyncio.Event()
-    _maintenance_task = asyncio.create_task(_retry_billing(_maintenance_stop))
-
-    if not config.OPENAI_API_KEY:
-        logger.warning("OPENAI_API_KEY 未配置,对话会失败——请在 .env 里设置")
-
-
-async def _lifespan_shutdown() -> None:
-    from agents import agent_memory
-    from scheduler import scheduler
-
-    await scheduler.shutdown()
-    if _maintenance_stop:
-        _maintenance_stop.set()
-    if _maintenance_task:
-        await _maintenance_task
-    if _prewarm_task:
-        await _prewarm_task
-    await agent_memory.close_checkpointer()
-    logger.info("deep_agent 2.0 已停止")
 
 
+@asynccontextmanager
 async def lifespan(app: FastAPI):
-    await _lifespan_startup()
-    try:
+    async with RuntimeEnvironment(scheduler=True):
         yield
-    finally:
-        await _lifespan_shutdown()
 
 
 app = FastAPI(

+ 151 - 0
runtime/environment.py

@@ -0,0 +1,151 @@
+"""网页、独立脚本和集成测试共用的进程级运行环境。"""
+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()

+ 13 - 13
scripts/selfcheck.py

@@ -47,7 +47,7 @@ MODULES = [
     "agents.agent_memory", "agents.agent_factory", "agents.agent_registry",
     "api.sse_events", "api.subagent_tracker", "api.context_meter",
     "api.source_tracker", "api.concurrency", "api.sse_core", "api.models",
-    "runtime.engine", "runtime.execution", "runtime.events", "runtime.interrupts",
+    "runtime.environment", "runtime.engine", "runtime.execution", "runtime.events", "runtime.interrupts",
     "runtime.concurrency", "runtime.context_meter", "runtime.source_tracker", "runtime.subagent_tracker",
     "api.chat_routes", "api.session_routes", "api.model_routes",
     "api.usage_routes", "api.schedule_routes", "api",
@@ -151,18 +151,18 @@ def check_import_boundaries() -> str:
 
 
 async def check_agent() -> str:
-    from agents import agent_factory, agent_memory, agent_registry
-
-    await agent_memory.init_checkpointer()
-    agent = await asyncio.to_thread(agent_registry.get_agent)
-    if agent is None:
-        raise RuntimeError("agent 编译返回 None")
-    info = agent_factory.describe()
-    if not info["tools"]["count"]:
-        raise RuntimeError("agent 没有挂任何工具")
-    if not info["interrupt_on"]:
-        raise RuntimeError("interrupt_on 为空:中断/恢复协议会失效")
-    await agent_memory.close_checkpointer()
+    from agents import agent_factory, agent_registry
+    from runtime.environment import RuntimeEnvironment
+
+    async with RuntimeEnvironment(prewarm=False, billing_retry=False, recover_interrupted=False):
+        agent = await asyncio.to_thread(agent_registry.get_agent)
+        if agent is None:
+            raise RuntimeError("agent 编译返回 None")
+        info = agent_factory.describe()
+        if not info["tools"]["count"]:
+            raise RuntimeError("agent 没有挂任何工具")
+        if not info["interrupt_on"]:
+            raise RuntimeError("interrupt_on 为空:中断/恢复协议会失效")
     return (
         f"{info['name']} · 工具 {info['tools']['count']} · "
         f"子智能体 {info['subagents']} · 审批 {info['interrupt_on']} · "

+ 322 - 0
tests/test_environment.py

@@ -0,0 +1,322 @@
+"""公共生命周期:启动失败、取消、重复进入以及真实 SQLite 资源释放。"""
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+import threading
+from types import SimpleNamespace
+
+import pytest
+from langchain_core.messages import AIMessageChunk
+
+import config
+from agents import agent_memory, agent_registry
+from db import db_conn, chat_store, billing_store
+from runtime import environment
+from runtime.environment import RuntimeEnvironment
+from runtime.execution import execute
+from scheduler import scheduler
+
+
+@pytest.fixture
+def lifecycle(monkeypatch):
+    calls = []
+    def mark(name):
+        def call(*args, **kwargs):
+            calls.append(name)
+        return call
+    def amark(name):
+        async def call(*args, **kwargs):
+            calls.append(name)
+        return call
+    for name in ('setup_logging', 'load_extensions', 'approval_names'):
+        monkeypatch.setattr(environment, name, mark(name))
+    monkeypatch.setattr(config, 'log_config_sources', mark('config'))
+    monkeypatch.setattr(db_conn, 'init_schema', mark('schema'))
+    monkeypatch.setattr(db_conn, 'health', lambda: {'status': 'ok'})
+    monkeypatch.setattr(db_conn, 'close', mark('db_close'))
+    monkeypatch.setattr(environment.pricing, 'seed_default_pricing', mark('pricing'))
+    monkeypatch.setattr(environment.run_store, 'recover_interrupted', mark('recover'))
+    monkeypatch.setattr(agent_registry, 'invalidate', mark('invalidate'))
+    monkeypatch.setattr(agent_memory, 'init_checkpointer', amark('checkpoint'))
+    monkeypatch.setattr(agent_memory, 'close_checkpointer', amark('checkpoint_close'))
+    monkeypatch.setattr(scheduler, 'start', amark('scheduler_start'))
+    monkeypatch.setattr(scheduler, 'shutdown', amark('scheduler_stop'))
+    return calls
+
+
+@pytest.mark.asyncio
+async def test_default_environment_and_sequential_reuse(lifecycle):
+    env = RuntimeEnvironment(prewarm=False, billing_retry=False)
+    for _ in range(2):
+        lifecycle.clear()
+        async with env as current:
+            assert current is env
+            assert lifecycle[-4:] == ['pricing', 'recover', 'invalidate', 'checkpoint']
+        assert lifecycle[-3:] == ['invalidate', 'checkpoint_close', 'db_close']
+        assert 'scheduler_start' not in lifecycle
+
+
+@pytest.mark.asyncio
+async def test_nested_environment_does_not_close_outer(lifecycle):
+    env = RuntimeEnvironment(prewarm=False, billing_retry=False)
+    async with env:
+        for other in (env, RuntimeEnvironment()):
+            with pytest.raises(RuntimeError, match='已有活动'):
+                async with other:
+                    pytest.fail('不应启动嵌套环境')
+        assert 'checkpoint_close' not in lifecycle
+    assert lifecycle.count('db_close') == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('stage', ['pricing', 'checkpoint', 'scheduler'])
+async def test_partial_startup_failure_cleans_and_allows_restart(stage, lifecycle, monkeypatch):
+    def fail():
+        raise ValueError(stage)
+    async def afail():
+        raise ValueError(stage)
+    with monkeypatch.context() as patch:
+        if stage == 'pricing':
+            patch.setattr(environment.pricing, 'seed_default_pricing', fail)
+        elif stage == 'checkpoint':
+            patch.setattr(agent_memory, 'init_checkpointer', afail)
+        else:
+            patch.setattr(scheduler, 'start', afail)
+        with pytest.raises(ValueError, match=stage):
+            async with RuntimeEnvironment(scheduler=True, prewarm=False, billing_retry=False):
+                pytest.fail('启动失败后不应进入执行体')
+    assert lifecycle[-1] == 'db_close'
+    if stage != 'pricing':
+        assert 'checkpoint_close' in lifecycle
+    if stage == 'scheduler':
+        assert lifecycle.index('scheduler_stop') < lifecycle.index('checkpoint_close')
+    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_body_exception_and_cleanup_error_still_release_resources(lifecycle, monkeypatch):
+    async def fail_close():
+        lifecycle.append('checkpoint_close')
+        raise RuntimeError('close failed')
+    with monkeypatch.context() as patch:
+        patch.setattr(agent_memory, 'close_checkpointer', fail_close)
+        with pytest.raises(RuntimeError, match='close failed'):
+            async with RuntimeEnvironment(prewarm=False, billing_retry=False):
+                raise ValueError('body failed')
+    assert lifecycle[-1] == 'db_close'
+    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_cancel_during_startup_waits_for_initialization(lifecycle, monkeypatch):
+    started, finish = asyncio.Event(), asyncio.Event()
+    async def init():
+        started.set()
+        await finish.wait()
+        lifecycle.append('checkpoint')
+    monkeypatch.setattr(agent_memory, 'init_checkpointer', init)
+    async def run():
+        async with RuntimeEnvironment(prewarm=False, billing_retry=False):
+            pytest.fail('取消后不应进入执行体')
+    task = asyncio.create_task(run())
+    await started.wait()
+    task.cancel()
+    await asyncio.sleep(0)
+    task.cancel()  # 清理期间再次取消,也要等待底层初始化收尾。
+    assert 'db_close' not in lifecycle
+    finish.set()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    assert lifecycle[-4:] == ['checkpoint', 'invalidate', 'checkpoint_close', 'db_close']
+
+
+@pytest.mark.asyncio
+async def test_shutdown_waits_for_workers_before_closing_connections(lifecycle, monkeypatch):
+    from billing import meter
+    prewarm_started, billing_started = threading.Event(), threading.Event()
+    release = threading.Event()
+    def compile_agent():
+        prewarm_started.set()
+        assert release.wait(5)
+        lifecycle.append('compiled')
+    def retry():
+        billing_started.set()
+        assert release.wait(5)
+        lifecycle.append('settled')
+    monkeypatch.setattr(agent_registry, 'get_agent', compile_agent)
+    monkeypatch.setattr(meter, 'retry_pending', retry)
+    running = asyncio.Event()
+    async def run():
+        async with RuntimeEnvironment(scheduler=True):
+            running.set()
+            await asyncio.Event().wait()
+    task = asyncio.create_task(run())
+    try:
+        await running.wait()
+        assert await asyncio.to_thread(prewarm_started.wait, 2)
+        assert await asyncio.to_thread(billing_started.wait, 2)
+        task.cancel()
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        assert 'checkpoint_close' not in lifecycle
+    finally:
+        release.set()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    assert lifecycle.index('scheduler_stop') < lifecycle.index('checkpoint_close')
+    assert lifecycle.index('settled') < lifecycle.index('db_close')
+    assert lifecycle.index('compiled') < lifecycle.index('checkpoint_close')
+    assert not [t for t in asyncio.all_tasks() if t.get_name().startswith('runtime-')]
+
+
+class FakeAgent:
+    def __init__(self):
+        self.checkpointer = agent_memory.get_checkpointer()
+
+    async def aget_state(self, options):
+        return SimpleNamespace(tasks=[])
+
+    async def astream(self, *args, **kwargs):
+        yield {'type': 'messages', 'data': (AIMessageChunk(content='报告完成', usage_metadata={
+            'input_tokens': 8, 'output_tokens': 2, 'total_tokens': 10}), {})}
+
+
+@pytest.mark.asyncio
+async def test_sqlite_script_executes_bills_and_reopens(monkeypatch, tmp_path):
+    # 独立业务库与 checkpoint,不使用线上配置,也不调用模型网络。
+    adapter = db_conn._SqliteAdapter(str(tmp_path / 'business.sqlite'))
+    monkeypatch.setattr(db_conn, '_adapter', adapter)
+    monkeypatch.setattr(db_conn, '_schema_ready', False)
+    monkeypatch.setattr(agent_memory, '_saver', None)
+    monkeypatch.setattr(agent_memory, '_kind', 'sqlite')
+    monkeypatch.setattr(config, 'CHECKPOINT_DB_PATH', str(tmp_path / 'checkpoint.sqlite'))
+    monkeypatch.setattr(agent_registry._slot, '_compile', FakeAgent)
+    schema = config.ROOT / 'db' / 'schema_sqlite.sql'
+    monkeypatch.setattr(adapter, 'schema_file', lambda: schema)
+    monkeypatch.setattr(config, 'ROOT', tmp_path)
+    previous = None
+    env = RuntimeEnvironment(prewarm=False, billing_retry=False)
+    for _ in range(2):
+        async with env:
+            session = chat_store.create_session('runtime-test')
+            agent = agent_registry.get_agent()
+            assert agent is not previous
+            assert agent.checkpointer is agent_memory.get_checkpointer()
+            result = await execute(user_name='runtime-test', session_id=session['id'], message='生成报告')
+            assert result.status == 'success'
+            assert result.output == '报告完成'
+            assert result.usage['total_tokens'] == 10
+            assert billing_store.list_pricing()
+            assert db_conn.scalar('SELECT total_tokens FROM billing_usage WHERE session_id = ?',
+                                  (session['id'],)) == 10
+            conn = adapter._conn
+            saver_conn = agent.checkpointer.conn
+            previous = agent
+        assert adapter._conn is None
+        with pytest.raises(Exception, match='closed'):
+            conn.execute('SELECT 1')
+        with pytest.raises(ValueError, match='no active connection'):
+            await saver_conn.execute('SELECT 1')
+        assert not agent_memory.is_initialized()
+        assert agent_registry.peek_agent() is None
+
+
+def test_mysql_closes_connections_from_all_threads_and_reopens():
+    adapter = db_conn._MySqlAdapter()
+    class Connection:
+        def __init__(self):
+            self.closed = False
+        def close(self):
+            self.closed = True
+        def ping(self, **kwargs):
+            assert not self.closed
+    adapter._connect = lambda database: Connection()
+    barrier = threading.Barrier(3)
+    def connect():
+        conn = adapter._connection()
+        barrier.wait(timeout=3)
+        return conn
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        tasks = [pool.submit(connect) for _ in range(2)]
+        main_conn = adapter._connection()
+        barrier.wait(timeout=3)
+        connections = [main_conn, *(task.result() for task in tasks)]
+        assert len({id(conn) for conn in connections}) == 3
+        adapter.close()
+        assert all(conn.closed for conn in connections)
+        assert pool.submit(adapter._connection).result() not in connections
+        assert adapter._connection() not in connections
+        adapter.close()
+
+
+@pytest.mark.asyncio
+async def test_checkpoint_setup_failure_closes_partial_connection(monkeypatch, tmp_path):
+    import aiosqlite
+    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+    monkeypatch.setattr(agent_memory, '_saver', None)
+    monkeypatch.setattr(agent_memory, '_kind', 'sqlite')
+    monkeypatch.setattr(config, 'CHECKPOINT_DB_PATH', str(tmp_path / 'bad.sqlite'))
+    connected = []
+    original = aiosqlite.connect
+    async def connect(*args, **kwargs):
+        conn = await original(*args, **kwargs)
+        connected.append(conn)
+        return conn
+    async def fail(self):
+        raise ValueError('setup failed')
+    monkeypatch.setattr(aiosqlite, 'connect', connect)
+    monkeypatch.setattr(AsyncSqliteSaver, 'setup', fail)
+    with pytest.raises(agent_memory.CheckpointerConfigError, match='setup failed'):
+        await agent_memory.init_checkpointer()
+    assert not agent_memory.is_initialized()
+    with pytest.raises(ValueError, match='no active connection'):
+        await connected[0].execute('SELECT 1')
+
+
+@pytest.mark.asyncio
+async def test_scheduler_start_cancellation_waits_for_database_read(lifecycle, monkeypatch):
+    started, finish = asyncio.Event(), asyncio.Event()
+    async def start():
+        started.set()
+        await finish.wait()
+        lifecycle.append('scheduler_loaded')
+    monkeypatch.setattr(scheduler, 'start', start)
+    async def run():
+        async with RuntimeEnvironment(scheduler=True, prewarm=False, billing_retry=False):
+            pytest.fail('取消后不应进入执行体')
+    task = asyncio.create_task(run())
+    await started.wait()
+    task.cancel()
+    await asyncio.sleep(0)
+    assert 'db_close' not in lifecycle
+    finish.set()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    assert lifecycle.index('scheduler_loaded') < lifecycle.index('scheduler_stop')
+    assert lifecycle.index('scheduler_stop') < lifecycle.index('db_close')
+
+
+@pytest.mark.asyncio
+async def test_http_lifespan_uses_common_environment(lifecycle, monkeypatch):
+    import main
+    def env(**kwargs):
+        assert kwargs == {'scheduler': True}
+        return RuntimeEnvironment(**kwargs, prewarm=False, billing_retry=False)
+    monkeypatch.setattr(main, 'RuntimeEnvironment', env)
+    async with main.app.router.lifespan_context(main.app):
+        assert 'scheduler_start' in lifecycle
+        assert 'db_close' not in lifecycle
+    assert lifecycle[-4:] == ['scheduler_stop', 'invalidate', 'checkpoint_close', 'db_close']
+
+
+def test_environment_reopens_across_event_loops(lifecycle):
+    env = RuntimeEnvironment(prewarm=False, billing_retry=False, recover_interrupted=False)
+    async def run():
+        async with env:
+            pass
+    asyncio.run(run())
+    asyncio.run(run())
+    assert lifecycle.count('checkpoint') == 2
+    assert lifecycle.count('db_close') == 2
+    assert 'recover' not in lifecycle

+ 1 - 0
tests/test_execution.py

@@ -215,6 +215,7 @@ def test_runtime_and_scheduler_do_not_import_http_layer():
                 assert not any(m.split(".")[0] in {"api", "fastapi", "starlette"} for m in modules), path
     # 新进程防止测试会话里提前导入 API 掩盖反向依赖。
     script = """
+import runtime.environment
 import runtime.engine
 import runtime.execution
 import runtime.interrupts

+ 2 - 1
tests/test_extensions.py

@@ -169,7 +169,8 @@ async def test_invalid_extension_blocks_startup_before_database(monkeypatch):
         pytest.fail("扩展校验失败后不应初始化数据库")
     monkeypatch.setattr(main.db_conn, "init_schema", forbidden)
     with pytest.raises(ValueError, match="未注册工具"):
-        await main._lifespan_startup()
+        async with main.lifespan(main.app):
+            pytest.fail("无效扩展不应启动服务")
 
 
 def test_release_includes_extension_code_and_prompt_files():