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 api.request_logging import RequestLogMiddleware
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, 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("计费准入失败", exc_info=(type(exc), exc, exc.__traceback__),
                 extra={"kv": {"event": "billing.admission_failed"}})
    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)}})


app.add_middleware(RequestLogMiddleware, trace_enabled=config.TRACE_MIDDLEWARE)


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

    不返回堆栈——它可能含连接串、内网地址。堆栈只进日志，用 trace_id 关联。
    """
    trace_id = getattr(request.state, "trace_id", "")
    token = bind_trace_id(trace_id)
    try:
        logger.exception("请求处理异常", extra={"kv": {"event": "http.request.failed",
                         "method": request.method, "path": request.url.path,
                         "error_type": type(exc).__name__}})
    finally:
        unbind_trace_id(token)
    return JSONResponse(
        status_code=500,
        content={"code": "internal_error", "message": f"服务内部错误：{type(exc).__name__}"},
        headers={"X-Trace-Id": trace_id} if trace_id else None,
    )


@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),
        # 能力开关：上线后要能一眼确认「资料目录挂上了没有」「附件上限是多少」，
        # 不必靠"猜哪个进程在服务"。
        "capabilities": {
            "max_upload_files": config.MAX_UPLOAD_FILES,
            "airflow_monitoring_dir": config.AIRFLOW_MONITORING_DIR,
            "airflow_output_dir": config.AIRFLOW_OUTPUT_DIR or f"data/{config.AIRFLOW_OUTPUT_DIRNAME}",
        },
    }

    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(),
        log_config=None, access_log=False,
    )