src/api/main.py - master - azathd/mutiagent

src/api/main.py
140 строк4 KB
Your Name
init
15 май 2026, 19:11
15 май 2026, 19:11f0c1163
"""FastAPI orchestrator: REST, UI, WebSocket, background NATS consumer."""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from pathlib import Path

import src.common.telemetry  # noqa: F401
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from prometheus_client import CONTENT_TYPE_LATEST, REGISTRY, generate_latest
from starlette.responses import Response

from src.api.auth import Role, create_token
from src.api.routes import actions, hitl, incidents, reports_api
from src.api.ws import feed

SERVICE_NAME = "orchestrator"
TEMPLATES = Jinja2Templates(directory=str(Path(__file__).resolve().parents[1] / "ui" / "templates"))


@asynccontextmanager
async def lifespan(app: FastAPI):  # noqa: ARG001
    from src.api.orchestrator_svc import OrchestratorBackground
    from src.common.logging import configure_logging, get_logger
    from src.common.migrate import ensure_database_schema, incidents_table_exists

    configure_logging()
    log = get_logger("orchestrator")

    schema_ok = False
    for attempt in range(1, 16):
        try:
            await ensure_database_schema()
            schema_ok = await incidents_table_exists()
            if schema_ok:
                log.info("schema_bootstrap_ok", attempt=attempt)
                break
        except Exception as exc:
            log.error("schema_bootstrap_failed", attempt=attempt, error=str(exc))
        await asyncio.sleep(2)

    task = asyncio.create_task(_idle())
    if schema_ok:
        bg = OrchestratorBackground()
        await bg.start()
    else:
        log.error(
            "orchestrator_consumer_disabled",
            hint="Run: make init-db && docker compose restart orchestrator",
        )
    yield
    task.cancel()


async def _idle() -> None:
    await asyncio.Future()


app = FastAPI(title="Xray Orchestrator", lifespan=lifespan)
app.include_router(incidents.router)
app.include_router(actions.router)
app.include_router(hitl.router)
app.include_router(reports_api.router)
app.include_router(feed.router)


@app.middleware("http")
async def trace(request: Request, call_next):  # type: ignore[type-arg]
    response = await call_next(request)
    response.headers.setdefault("X-Service", SERVICE_NAME)
    return response


@app.get("/healthz")
async def healthz() -> dict[str, str]:
    return {"status": "ok", "service": SERVICE_NAME}


@app.get("/metrics")
async def metrics() -> Response:
    return Response(generate_latest(REGISTRY), media_type=CONTENT_TYPE_LATEST)


@app.post("/api/auth/dev-token")
async def dev_token(role: str = "admin") -> dict[str, str]:
    """Development-only token issuer."""
    try:
        r = Role(role)
    except ValueError:
        r = Role.ADMIN
    return {"token": create_token("dev", r)}


@app.get("/api/auth/dev-token")
async def dev_token_get(role: str = "admin") -> dict[str, str]:
    return await dev_token(role)


@app.get("/", response_class=HTMLResponse)
async def overview(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse(
        "overview.html",
        {"request": request, "title": "Overview"},
    )


@app.get("/security", response_class=HTMLResponse)
async def security_page(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse("security.html", {"request": request})


@app.get("/traffic", response_class=HTMLResponse)
async def traffic_page(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse("traffic.html", {"request": request})


@app.get("/users", response_class=HTMLResponse)
async def users_page(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse("users.html", {"request": request})


@app.get("/system", response_class=HTMLResponse)
async def system_page(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse(
        "system.html",
        {"request": request, "grafana_url": "http://localhost:3000"},
    )


@app.get("/settings", response_class=HTMLResponse)
async def settings_page(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse("settings.html", {"request": request})


@app.get("/reports-ui", response_class=HTMLResponse)
async def reports_page(request: Request) -> HTMLResponse:
    return TEMPLATES.TemplateResponse("reports.html", {"request": request})