agentic-mobile-control/apps/cloud-api/cloud_api/app.py

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from fastapi import FastAPI, status
from fastapi import Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.types import Scope

from cloud.auth import (
    ChainedAuthProvider,
    ConfiguredEnrollmentTokenProvider,
    RepositoryHostAuthProvider,
    create_auth_provider,
)
from cloud.config import CloudConfig
from cloud.control_config import (
    CloudConfigurationError,
    CloudControlConfig,
    load_control_config,
    validate_control_config,
)
from cloud.database import CloudDatabase
from cloud.internal_api.api import create_internal_router
from cloud.plugins import PluginRegistry
from cloud.observability import (
    CORRELATION_HEADER,
    bind_correlation_id,
    new_correlation_id,
    normalize_correlation_id,
    reset_correlation_id,
)
from cloud.pool import DevicePool
from cloud.scheduler import TaskScheduler
from cloud.schema import require_current_schema
from cloud.sdk.api import create_cloud_router
from core.models import utc_now


DatabaseFactory = Callable[[CloudControlConfig], CloudDatabase]
logger = logging.getLogger(__name__)


class _RepositoryProxy:
    def __init__(self) -> None:
        self._target: Any | None = None

    def bind(self, target: Any) -> None:
        self._target = target

    def unbind(self) -> None:
        self._target = None

    def __getattr__(self, name: str) -> Any:
        if self._target is None:
            raise RuntimeError("cloud repository is not available outside app lifespan")
        return getattr(self._target, name)


@dataclass(frozen=True)
class CloudApplicationServices:
    repository: _RepositoryProxy
    pool: DevicePool
    scheduler: TaskScheduler
    plugin_registry: PluginRegistry
    auth_provider: Any


class SpaStaticFiles(StaticFiles):
    """``StaticFiles`` variant that falls back to ``index.html`` for SPA routes.

    ``StaticFiles(html=True)`` only serves ``index.html`` for the mount root and
    for directory roots; an unknown path like ``/console/tasks/abc`` raises a
    plain 404, which breaks deep-link refreshes in a browser-running SPA. This
    subclass intercepts 404s for non-asset paths and re-serves ``index.html``
    so the SPA router can take over.
    """

    async def get_response(self, path: str, scope: Scope) -> Any:
        try:
            return await super().get_response(path, scope)
        except StarletteHTTPException as exc:
            if exc.status_code == 404 and path != "index.html":
                return await super().get_response("index.html", scope)
            raise


def create_app(
    *,
    config: CloudControlConfig | None = None,
    database_factory: DatabaseFactory | None = None,
) -> FastAPI:
    """Create the independently deployable cloud API application."""
    control_config = config or load_control_config()
    validate_control_config(control_config)
    build_database = database_factory or _default_database_factory
    configured_auth_provider = create_auth_provider(
        control_config.credentials,
        allow_insecure_anonymous=control_config.allow_insecure_anonymous,
    )
    repository = _RepositoryProxy()
    auth_provider = ChainedAuthProvider(
        (
            configured_auth_provider,
            RepositoryHostAuthProvider(repository),  # type: ignore[arg-type]
        )
    )
    enrollment_auth_provider = ConfiguredEnrollmentTokenProvider(
        control_config.enrollment_credentials
    )
    domain_config = CloudConfig(
        lease_duration_seconds=control_config.lease_duration_seconds,
    )
    pool = DevicePool(repository, domain_config)  # type: ignore[arg-type]
    scheduler = TaskScheduler(pool, repository, domain_config)  # type: ignore[arg-type]
    plugin_registry = PluginRegistry(repository)  # type: ignore[arg-type]
    services = CloudApplicationServices(
        repository=repository,
        pool=pool,
        scheduler=scheduler,
        plugin_registry=plugin_registry,
        auth_provider=auth_provider,
    )

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        database = build_database(control_config)
        repository.bind(database.repository)
        app.state.cloud_config = control_config
        app.state.database = database
        app.state.cloud_services = services
        app.state.startup_complete = False
        app.state.schema_ready = False
        stop_workers = asyncio.Event()
        worker_tasks: list[asyncio.Task[None]] = []
        try:
            repository.health_check()
            if control_config.environment == "production":
                require_current_schema(control_config.database_url)
            app.state.schema_ready = True
            worker_tasks = [
                asyncio.create_task(
                    _run_scheduler_loop(
                        services,
                        stop_workers,
                        interval_seconds=control_config.scheduler_interval_seconds,
                    ),
                    name="cloud-scheduler",
                ),
                asyncio.create_task(
                    _run_lease_reaper_loop(
                        services,
                        stop_workers,
                        interval_seconds=(control_config.lease_reaper_interval_seconds),
                        max_attempts=control_config.max_task_attempts,
                    ),
                    name="cloud-lease-reaper",
                ),
            ]
            app.state.worker_tasks = tuple(worker_tasks)
            app.state.startup_complete = True
            yield
        finally:
            app.state.startup_complete = False
            stop_workers.set()
            try:
                if worker_tasks:
                    await asyncio.gather(*worker_tasks)
            finally:
                try:
                    database.close()
                finally:
                    repository.unbind()

    app = FastAPI(title="Device Cloud API", lifespan=lifespan)

    if control_config.cors_allowed_origins:
        from fastapi.middleware.cors import CORSMiddleware

        app.add_middleware(
            CORSMiddleware,
            allow_origins=list(control_config.cors_allowed_origins),
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

    @app.middleware("http")
    async def correlation_logging(request: Request, call_next):
        correlation_id = normalize_correlation_id(
            request.headers.get(CORRELATION_HEADER)
        )
        correlation_token = bind_correlation_id(correlation_id)
        try:
            response = await call_next(request)
            logger.info(
                "cloud request completed",
                extra={
                    "correlation_id": correlation_id,
                    "method": request.method,
                    "path": request.url.path,
                    "status_code": response.status_code,
                },
            )
            response.headers[CORRELATION_HEADER] = correlation_id
            return response
        finally:
            reset_correlation_id(correlation_token)

    @app.get("/health/live")
    def health_live() -> dict[str, str]:
        return {"status": "live"}

    @app.get("/health/ready")
    def health_ready():
        checks = {
            "configuration": bool(getattr(app.state, "startup_complete", False)),
            "database": False,
            "schema": bool(getattr(app.state, "schema_ready", False)),
            "workers": False,
        }
        if checks["configuration"]:
            try:
                services.repository.health_check()
                checks["database"] = True
            except Exception:
                checks["database"] = False
        worker_state = getattr(app.state, "worker_tasks", ())
        checks["workers"] = bool(worker_state) and all(
            not task.done() for task in worker_state
        )
        ready = all(checks.values())
        return JSONResponse(
            status_code=(
                status.HTTP_200_OK if ready else status.HTTP_503_SERVICE_UNAVAILABLE
            ),
            content={"status": "ready" if ready else "not_ready", "checks": checks},
        )

    app.include_router(
        create_cloud_router(
            pool=pool,
            scheduler=scheduler,
            plugin_registry=plugin_registry,
            auth_provider=auth_provider,
        )
    )
    app.include_router(
        create_internal_router(
            pool=pool,
            auth_provider=auth_provider,
            enrollment_auth_provider=enrollment_auth_provider,
            lease_duration_seconds=control_config.lease_duration_seconds,
        )
    )

    if control_config.console_static_dir:
        dist_dir = Path(control_config.console_static_dir)
        if not dist_dir.is_dir():
            raise CloudConfigurationError(
                f"CLOUD_CONSOLE_STATIC_DIR is not a directory: {dist_dir}"
            )

        @app.get("/", include_in_schema=False)
        async def _redirect_to_console() -> RedirectResponse:
            return RedirectResponse(url="/console/")

        app.mount(
            "/console",
            SpaStaticFiles(directory=str(dist_dir), html=True),
            name="cloud-console",
        )

    return app


def _default_database_factory(config: CloudControlConfig) -> CloudDatabase:
    return CloudDatabase(
        config.database_url,
        create_schema=config.environment != "production",
    )


async def _run_scheduler_loop(
    services: CloudApplicationServices,
    stop: asyncio.Event,
    *,
    interval_seconds: float,
) -> None:
    while not stop.is_set():
        correlation_token = bind_correlation_id(new_correlation_id())
        try:
            services.scheduler.assign()
        except Exception:
            logger.exception(
                "cloud lifecycle iteration failed",
                extra={"worker": "scheduler"},
            )
        finally:
            reset_correlation_id(correlation_token)
        if await _wait_for_stop(stop, interval_seconds):
            return


async def _run_lease_reaper_loop(
    services: CloudApplicationServices,
    stop: asyncio.Event,
    *,
    interval_seconds: float,
    max_attempts: int,
) -> None:
    while not stop.is_set():
        correlation_token = bind_correlation_id(new_correlation_id())
        try:
            services.repository.reap_expired_leases(
                now=utc_now(),
                max_attempts=max_attempts,
            )
        except Exception:
            logger.exception(
                "cloud lifecycle iteration failed",
                extra={"worker": "lease_reaper"},
            )
        finally:
            reset_correlation_id(correlation_token)
        if await _wait_for_stop(stop, interval_seconds):
            return


async def _wait_for_stop(stop: asyncio.Event, interval_seconds: float) -> bool:
    try:
        await asyncio.wait_for(stop.wait(), timeout=interval_seconds)
    except TimeoutError:
        return False
    return True