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 typing import Any

from fastapi import FastAPI

from cloud.auth import create_auth_provider
from cloud.config import CloudConfig
from cloud.control_config import (
    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.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


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
    auth_provider = create_auth_provider(
        control_config.credentials,
        allow_insecure_anonymous=control_config.allow_insecure_anonymous,
    )
    repository = _RepositoryProxy()
    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
        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)
            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)
            yield
        finally:
            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)
    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,
            lease_duration_seconds=control_config.lease_duration_seconds,
        )
    )
    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():
        try:
            services.scheduler.assign()
        except Exception:
            logger.exception(
                "cloud lifecycle iteration failed",
                extra={"worker": "scheduler"},
            )
        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():
        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"},
            )
        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