vent_agent_2.0/main.py

"""deep_agent 2.0 服务入口。

启动顺序（lifespan）：
1. 日志
2. 业务扩展发现与校验；建表（幂等）
3. 计费单价表种子
4. checkpointer 初始化
5. **后台**预热 agent（图编译耗时百毫秒级，不能挡住 uvicorn 就绪）
6. 调度器启动并装载定时任务
7. 单进程约束自检

退出时逆序释放：调度器 → checkpointer。
"""

from __future__ import annotations

import asyncio
import os

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles

import config
from api import router as api_router
from api import concurrency
from billing import pricing, 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 已停止")


async def lifespan(app: FastAPI):
    await _lifespan_startup()
    try:
        yield
    finally:
        await _lifespan_shutdown()


app = FastAPI(
    title="deep_agent 2.0",
    description="基于 deepagents 的深度智能体服务骨架；主入口 POST /api/chat（SSE）",
    version=config.APP_VERSION,
    lifespan=lifespan,
)

# 业务路由统一挂在 /api 下（api/__init__.py 里声明了 prefix）
app.include_router(api_router)


@app.exception_handler(guard.QuotaExceededError)
async def quota_error(request: Request, exc: guard.QuotaExceededError):
    return JSONResponse(status_code=429, content={"detail": {"code": "quota_exceeded", "message": str(exc)}})


@app.exception_handler(guard.BillingConfigError)
async def billing_error(request: Request, exc: guard.BillingConfigError):
    logger.error("计费准入失败：%s", exc)
    return JSONResponse(status_code=503, content={"detail": {"code": "billing_unavailable", "message": "计费暂不可用"}})


@app.exception_handler(concurrency.TooManyTasksError)
async def concurrency_error(request: Request, exc: concurrency.TooManyTasksError):
    return JSONResponse(status_code=429, headers={"Retry-After": str(exc.retry_after)},
                        content={"detail": {"code": "too_many_tasks", "message": str(exc),
                                            "retry_after": exc.retry_after}})


@app.exception_handler(concurrency.SessionBusyError)
async def session_busy(request: Request, exc: concurrency.SessionBusyError):
    return JSONResponse(status_code=409,
                        content={"detail": {"code": "session_busy", "message": str(exc)}})


@app.exception_handler(concurrency.SessionMissingError)
async def session_missing(request: Request, exc: concurrency.SessionMissingError):
    return JSONResponse(status_code=404,
                        content={"detail": {"code": "session_not_found", "message": str(exc)}})


if config.TRACE_MIDDLEWARE:

    @app.middleware("http")
    async def trace_middleware(request: Request, call_next):
        """把 ``X-Trace-Id`` 透传进来或新生成一个，回写响应头。

        注意：这里绑定的 ContextVar 在 StreamingResponse 的生成器里**不保证**
        可见（响应体迭代发生在中间件调用栈之外），所以 SSE 生成器内部会自己
        再绑一次。这个中间件的作用是让普通 HTTP 接口与访问日志带上 trace_id。
        """
        trace_id = request.headers.get("X-Trace-Id") or new_trace_id()
        request.state.trace_id = trace_id
        token = bind_trace_id(trace_id)
        try:
            response = await call_next(request)
        finally:
            unbind_trace_id(token)
        response.headers["X-Trace-Id"] = trace_id
        return response


@app.exception_handler(Exception)
async def unhandled_exception_handler(request: Request, exc: Exception):
    """兜底异常处理：统一成 ``{code, message}``，与业务错误同形状。

    不返回堆栈——它可能含连接串、内网地址。堆栈只进日志，用 trace_id 关联。
    """
    trace_id = getattr(request.state, "trace_id", "")
    logger.exception("未处理异常 path=%s trace=%s", request.url.path, trace_id)
    return JSONResponse(
        status_code=500,
        content={"code": "internal_error", "message": f"服务内部错误：{type(exc).__name__}"},
    )


@app.get("/")
async def index():
    """重定向到测试页（与 1.5 的行为一致：``GET /`` → 静态页）。"""
    return RedirectResponse(url="/static/index.html")


@app.get("/health")
async def health():
    """健康检查：逐子系统报状态。

    整体 ``status`` 只由**关键路径**决定（存储 + checkpoint）。调度器、
    计费单价缺失这类是 ``degraded`` 而不是 ``down``：对话还能用，不该被
    负载均衡摘掉。
    """
    from agents import agent_memory

    detail: dict = {
        "version": config.APP_VERSION,
        "pid": os.getpid(),
        "dialect": config.DB_DIALECT,
        "storage": db_conn.health(),
        "checkpointer": {
            "kind": agent_memory.saver_kind(),
            "initialized": agent_memory.is_initialized(),
        },
        "concurrency": concurrency.stats(),
        "model": config.MODEL,
        "port": config.PORT,
        # 环境变量冲突要能在健康检查里看到，不必去翻日志
        "config_conflicts": dict(config.CONFLICTS),
    }

    try:
        from scheduler import scheduler as sched

        detail["scheduler"] = {
            "running": sched.is_running(),
            "jobs": len(sched.jobs()),
        }
    except Exception as exc:  # pragma: no cover
        detail["scheduler"] = {"running": False, "error": f"{type(exc).__name__}: {exc}"}

    try:
        from tools import registry

        detail["tools"] = len(registry.get_tool_names())
    except Exception as exc:  # pragma: no cover
        detail["tools"] = {"error": f"{type(exc).__name__}: {exc}"}

    storage_ok = detail["storage"].get("status") == "ok"
    checkpointer_ok = bool(detail["checkpointer"]["initialized"])
    status = "ok" if (storage_ok and checkpointer_ok) else "down"
    detail["status"] = status
    detail["degraded"] = status == "ok" and not detail.get("scheduler", {}).get("running", True)
    return detail


try:
    config.resolve_path("static/.keep")
    app.mount("/static", StaticFiles(directory=str(config.ROOT / "static")), name="static")
except Exception as exc:  # pragma: no cover - 静态目录缺失不该阻止服务启动
    logger.warning("挂载静态目录失败：%s: %s", type(exc).__name__, exc)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "main:app",
        host=config.HOST,
        port=config.PORT,
        workers=1,  # 单进程是硬约束，见 _check_single_process
        reload=False,  # 开 reload 会起子进程，同样破坏进程内状态
        log_level=config.LOG_LEVEL.lower(),
    )