vent_agent_2.0/main.py

"""deep_agent 2.0 HTTP 入口；公共资源生命周期由 RuntimeEnvironment 管理。"""

from __future__ import annotations

import os
from contextlib import asynccontextmanager

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 runtime import concurrency
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__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    async with RuntimeEnvironment(scheduler=True):
        yield


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(),
    )