Nexus/server/main.py

"""Nexus — FastAPI Application Entry Point
Clean Architecture: Presentation Layer (API routes + middleware + lifespan)

Lifespan startup (when .env exists — normal mode):
1. Verify SECRET_KEY is set (fatal if missing)
2. init_db() — create tables
3. init_redis() — ADR-009 mandatory validation (exit if unavailable)
4. load_settings_from_db() — override mutable settings from MySQL
5. Start Redis Pub/Sub subscriber (ADR-010)
6. Launch background tasks (heartbeat flush, self-monitor)

Lifespan startup (no .env — install mode):
- Only the /api/install/ routes are functional
- All other API routes return 503 "System not configured"
- User completes the install wizard at /app/install.html

D2: install.php → install.html migration (conditional lifespan)
D7: DB session leak fix — middleware auto-manages session lifecycle per request.
"""

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from starlette.middleware.base import BaseHTTPMiddleware

from server.config import settings
from server.infrastructure.database.session import init_db, AsyncSessionLocal

# Detect install mode: no .env = first-time setup
ROOT_DIR = Path(__file__).resolve().parent.parent
INSTALL_MODE = not (ROOT_DIR / ".env").exists()

# Install wizard (no JWT required — runs before system is configured)
from server.api.install import router as install_router

# API routes (require JWT + full app init)
from server.api.servers import router as servers_router
from server.api.auth import router as auth_router
from server.api.agent import router as agent_router
from server.api.scripts import router as scripts_router
from server.api.settings import (
    router as settings_router,
    schedule_router,
    preset_router,
    audit_router,
    retry_router,
)
from server.api.websocket import router as websocket_router
from server.api.health import router as health_router
from server.api.assets import router as assets_router
from server.api.webssh import router as webssh_router
from server.api.sync_v2 import router as sync_v2_router
from server.api.search import router as search_router

# Background tasks
from server.background.heartbeat_flush import heartbeat_flush_loop
from server.background.self_monitor import self_monitor_loop
from server.background.schedule_runner import schedule_runner_loop
from server.background.retry_runner import retry_runner_loop

logger = logging.getLogger("nexus")

# Track background tasks for clean shutdown
_background_tasks: list[asyncio.Task] = []

# Redis-based primary worker lock (prevents duplicate background tasks with multi-worker uvicorn)
_PRIMARY_LOCK_KEY = "nexus:primary_worker"
_PRIMARY_LOCK_TTL = 30  # seconds — must renew periodically


async def _acquire_primary_lock() -> bool:
    """Try to become the primary worker via Redis SETNX.

    Only one uvicorn worker will hold this lock at a time.
    The lock auto-expires after PRIMARY_LOCK_TTL seconds (safety net for crashes).
    Primary worker renews the lock periodically via _renew_primary_lock().
    """
    try:
        from server.infrastructure.redis.client import get_redis
        import os
        redis = get_redis()
        worker_id = str(os.getpid())
        # SET with NX (only if not exists) + EX (TTL)
        acquired = await redis.set(_PRIMARY_LOCK_KEY, worker_id, nx=True, ex=_PRIMARY_LOCK_TTL)
        if acquired:
            # Start lock renewal background task
            asyncio.create_task(_renew_primary_lock(), name="primary_lock_renewal")
            logger.info(f"Primary worker lock acquired (pid={worker_id})")
            return True
        return False
    except Exception as e:
        # If Redis fails, assume primary (single-worker mode)
        logger.warning(f"Primary lock check failed, assuming primary: {e}")
        return True


async def _renew_primary_lock():
    """Periodically renew the primary worker lock in Redis."""
    import os
    worker_id = str(os.getpid())
    while True:
        await asyncio.sleep(_PRIMARY_LOCK_TTL // 2)  # Renew at half the TTL
        try:
            from server.infrastructure.redis.client import get_redis
            redis = get_redis()
            current = await redis.get(_PRIMARY_LOCK_KEY)
            if current == worker_id:
                await redis.expire(_PRIMARY_LOCK_KEY, _PRIMARY_LOCK_TTL)
            else:
                logger.warning("Primary worker lock lost — another worker took over")
                break
        except Exception as e:
            logger.error(f"Primary lock renewal failed: {e}")


# ── D7: DB Session Middleware (fixes session leak in 4 service factories) ──

class DbSessionMiddleware(BaseHTTPMiddleware):
    """Auto-manage DB session lifecycle per HTTP request.

    Problem: 4 service factories (get_server_service, get_script_service,
    get_auth_service, get_sync_service) create AsyncSessionLocal() but never
    close them — causing connection pool exhaustion (P0 bug).

    Solution: Middleware opens session at request start, stores in request.state.db,
    and closes it after response. Service factories read from request.state.db
    instead of creating new sessions.

    Pattern: Similar to Django's ATOMIC_REQUESTS — session lifetime = request lifetime.
    """

    async def dispatch(self, request: Request, call_next):
        # Skip non-API routes (WebSocket, health, static files)
        # WebSocket connections are long-lived — never wrap them in a DB session
        path = request.url.path
        if not path.startswith("/api/"):
            return await call_next(request)

        # Skip install API (it manages its own DB connections)
        if path.startswith("/api/install/"):
            return await call_next(request)

        # Skip in install mode (no global engine available)
        if INSTALL_MODE:
            return await call_next(request)

        # Open session for this request
        async with AsyncSessionLocal() as session:
            request.state.db = session
            try:
                response = await call_next(request)
                return response
            except Exception:
                await session.rollback()
                raise


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifecycle: startup → background tasks → shutdown

    Install mode (no .env): only install API + static files available.
    Normal mode (.env exists): full initialization with DB, Redis, background tasks.
    """

    # ── Install mode: skip everything ──
    if INSTALL_MODE:
        logger.info("⚠ INSTALL MODE — .env not found. Visit /app/install.html to configure.")
        yield
        return

    # ── Normal mode: full startup ──
    if not settings.SECRET_KEY:
        logger.error("SECRET_KEY is empty! Set it in .env or MySQL settings table.")
        raise SystemExit("SECRET_KEY is required for Nexus to start.")

    if not settings.API_KEY:
        logger.error("API_KEY is empty! Set it in .env (generated by install wizard).")
        raise SystemExit("API_KEY is required for Agent authentication and encryption fallback.")

    if not settings.ENCRYPTION_KEY:
        logger.error("ENCRYPTION_KEY is empty! Set it in .env (generated by install wizard).")
        raise SystemExit("ENCRYPTION_KEY is required for credential encryption.")

    # 1. Initialize database tables
    await init_db()
    logger.info("Database tables initialized")

    # 1b. D8: Run data migrations (category → Node, backfill platform_id)
    from server.infrastructure.database.migrations import run_migrations
    await run_migrations()

    # 2. Initialize Redis (ADR-009: mandatory — exits if unavailable)
    from server.infrastructure.redis.client import init_redis, close_redis
    await init_redis(app)

    # 3. Load mutable settings from MySQL (override .env values)
    async with AsyncSessionLocal() as session:
        overridden = await settings.load_settings_from_db(session)
        logger.info(f"Settings from DB: {overridden} overrides applied")

    # 4. Start Redis Pub/Sub subscriber (ADR-010)
    from server.api.websocket import start_redis_subscriber, stop_redis_subscriber
    await start_redis_subscriber()

    # 5. Launch background tasks (only on primary worker to avoid duplicate execution)
    # When running with --workers N, each worker gets its own lifespan.
    # Use Redis-based leader election to ensure only one worker runs background tasks.
    is_primary = await _acquire_primary_lock()
    if is_primary:
        task_flush = asyncio.create_task(heartbeat_flush_loop(), name="heartbeat_flush")
        task_monitor = asyncio.create_task(self_monitor_loop(), name="self_monitor")
        task_schedule = asyncio.create_task(schedule_runner_loop(), name="schedule_runner")
        task_retry = asyncio.create_task(retry_runner_loop(), name="retry_runner")
        _background_tasks.extend([task_flush, task_monitor, task_schedule, task_retry])
        logger.info("Primary worker — background tasks launched")
    else:
        logger.info("Secondary worker — background tasks skipped (primary worker handles them)")

    # 6. Start asyncssh connection pool (W1)
    from server.infrastructure.ssh.asyncssh_pool import ssh_pool
    await ssh_pool.start()

    logger.info(
        f"{settings.SYSTEM_NAME} v6.0.0 started — "
        f"background tasks: heartbeat_flush(10min), self_monitor(30s), "
        f"schedule_runner(60s), retry_runner(5min)"
    )

    yield

    # ── Shutdown ──
    for task in _background_tasks:
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass
    _background_tasks.clear()

    # Stop Redis Pub/Sub subscriber
    await stop_redis_subscriber()

    # Close Redis connection
    await close_redis()

    # Stop asyncssh connection pool
    from server.infrastructure.ssh.asyncssh_pool import ssh_pool
    await ssh_pool.stop()

    # Close shared Telegram httpx client
    from server.infrastructure.telegram import close_client as close_telegram_client
    await close_telegram_client()

    logger.info(f"{settings.SYSTEM_NAME} shutting down — background tasks cancelled")


app = FastAPI(
    title=settings.SYSTEM_NAME if not INSTALL_MODE else "Nexus Installer",
    description=f"{settings.SYSTEM_TITLE}" if not INSTALL_MODE else "Nexus Installation Wizard",
    version="6.0.0",
    lifespan=lifespan,
)


# ── Install mode middleware: block non-install routes when not configured ──

class InstallModeMiddleware(BaseHTTPMiddleware):
    """In install mode, only /api/install/, /app/, /health are accessible.
    All other API routes return 503 Service Unavailable."""

    async def dispatch(self, request: Request, call_next):
        if INSTALL_MODE:
            path = request.url.path
            # Redirect root to install wizard
            if path == "/":
                from fastapi.responses import RedirectResponse
                return RedirectResponse(url="/app/install.html")
            # Allow: install API, static files, health check
            if path.startswith("/api/install/") or path.startswith("/app/") or path == "/health":
                return await call_next(request)
            if path.startswith("/ws/"):
                return await call_next(request)
            # Block everything else
            from fastapi.responses import JSONResponse
            return JSONResponse(
                status_code=503,
                content={"detail": "系统尚未配置，请先访问 /app/install.html 完成安装"},
            )
        return await call_next(request)


# Install mode middleware (must be first — before DB session middleware)
app.add_middleware(InstallModeMiddleware)

# Security headers middleware (before CORS so headers appear on all responses)
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """Inject security-related response headers on every response.

    - X-Content-Type-Options: nosniff — prevents MIME-type sniffing
    - X-Frame-Options: DENY — prevents clickjacking via iframes
    - Referrer-Policy: strict-origin-when-cross-origin — limits referrer leakage
    - Permissions-Policy: restricts browser features (camera, mic, geolocation)

    Not included:
    - Content-Security-Policy: too many inline scripts/styles to enforce easily
    - Strict-Transport-Security: handled by Nginx in production
    """
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        response.headers["Permissions-Policy"] = "camera=(), microphone=(), geolocation=()"
        return response

app.add_middleware(SecurityHeadersMiddleware)

# D7: DB session leak fix middleware (must be added BEFORE CORS so it wraps all requests)
app.add_middleware(DbSessionMiddleware)

# CORS — restrict to actual frontend origin (ECC security fix: was allow_origins=["*"])
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.CORS_ORIGINS.split(",") if settings.CORS_ORIGINS else ["http://localhost:8600"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Register all API routers

# D2: Install wizard (always registered — works in both modes)
app.include_router(install_router)

# Normal API routers (require full app init)
app.include_router(servers_router)
app.include_router(auth_router)
app.include_router(agent_router)
app.include_router(scripts_router)
app.include_router(settings_router)
app.include_router(schedule_router)
app.include_router(preset_router)
app.include_router(audit_router)
app.include_router(retry_router)

# WebSocket + Health check routers (new in v6.0)
app.include_router(websocket_router)
app.include_router(health_router)

# Asset management (Platform, Node, SSH Sessions, Command Logs)
app.include_router(assets_router)

# Web SSH Terminal
app.include_router(webssh_router)

# Sync Engine v2 (Step 4)
app.include_router(sync_v2_router)

# Global Search
app.include_router(search_router)

# ── Static files: serve web/app/ at /app/ (HTML pages + JS/CSS assets) ──
# In production, Nginx serves these directly; this is for dev/standalone mode.
WEB_APP_DIR = ROOT_DIR / "web" / "app"
if WEB_APP_DIR.is_dir():
    app.mount("/app", StaticFiles(directory=str(WEB_APP_DIR), html=True), name="static_app")