app/backend/fluksio/main.py

import asyncio
import contextlib
import inspect
import logging
from collections.abc import AsyncIterator, Callable, Coroutine
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from typing import Any

from fastapi import FastAPI, Request
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute
from fluksio_worker.worker_main import ARTIFACT_DIR_ENV
from starlette.middleware.cors import CORSMiddleware

from fluksio import __version__
from fluksio.api.main import api_router
from fluksio.api.routes.alerts import read_config as read_alerts_config
from fluksio.cloud import config as cloud_config
from fluksio.core import security
from fluksio.core.config import settings
from fluksio.core.db import engine as db_engine
from fluksio.core.db import prepare
from fluksio.flow import logs, modules
from fluksio.flow.alerts import AlertManager
from fluksio.flow.artifacts import ArtifactStore
from fluksio.flow.controller import FlowController, RebuildBusy
from fluksio.flow.dashboards import DashboardStore
from fluksio.flow.events import event_bus
from fluksio.flow.executor import ExecutionService
from fluksio.flow.metrics import MetricsCollector
from fluksio.flow.nodes.http import close_shared_client
from fluksio.flow.pipeline import ValueSource
from fluksio.flow.placement import Placer
from fluksio.flow.plugins import load_plugins
from fluksio.flow.provision import load_provisioners
from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue
from fluksio.flow.remote import RemoteWorkerHub
from fluksio.flow.resources import ResourceAccountant, fair_share_env
from fluksio.flow.runs import RUN_STATE_TTL, RunService, sweep_artifacts
from fluksio.flow.secrets import init_secrets
from fluksio.flow.state import MemoryState, RedisState, StateBackend
from fluksio.flow.store import FlowStore
from fluksio.flow.watchdog import LoopWatchdog
from fluksio.flow.workers import PythonWorkerPool

logger = logging.getLogger(__name__)


def custom_generate_unique_id(route: APIRoute) -> str:
    return f"{route.tags[0]}-{route.name}"


if settings.SENTRY_DSN and settings.ENVIRONMENT != "local":
    # Imported here rather than at the top: it is a `fluksio[server]` extra, so
    # a pip install without one has no sentry to import — and no DSN either.
    import sentry_sdk

    # `enable_tracing` was removed in sentry-sdk 2.x; this is what it meant.
    sentry_sdk.init(dsn=str(settings.SENTRY_DSN), traces_sample_rate=1.0)


def _state_backend() -> StateBackend:
    if settings.REDIS_HOST:
        return RedisState(host=settings.REDIS_HOST, port=settings.REDIS_PORT)
    return MemoryState()


def _work_queue(namespace: str = "queue") -> WorkQueue:
    """Redis makes queued work survive the process; memory does not pretend to."""
    if settings.REDIS_HOST:
        return RedisWorkQueue(
            host=settings.REDIS_HOST, port=settings.REDIS_PORT, namespace=namespace
        )
    return MemoryWorkQueue()


def _run_state(namespace: str) -> StateBackend:
    """A state backend of a run's own, which is what isolates it.

    It expires: a finished run's messages are read out into its result, and
    what is left is only worth keeping while someone might look at it.
    """
    if settings.REDIS_HOST:
        return RedisState(
            host=settings.REDIS_HOST,
            port=settings.REDIS_PORT,
            namespace=namespace,
            ttl=RUN_STATE_TTL,
        )
    return MemoryState()


async def _sweep_artifacts(store: ArtifactStore, controller: FlowController) -> None:
    """Take unreferenced artifact bytes off the disk, on a slow loop.

    A flow streaming media writes one artifact per frame, so a store nothing
    prunes only grows. Runs in a thread: it walks a directory and reads state.
    """
    interval = settings.ARTIFACT_GC_INTERVAL_S
    if interval <= 0:
        return
    while True:
        await asyncio.sleep(interval)
        try:
            await run_in_threadpool(
                sweep_artifacts,
                store,
                controller.state,
                settings.ARTIFACT_GC_GRACE_S,
            )
        except Exception:
            logger.exception("Artifact sweep failed")


def _mcp_sessions() -> AbstractAsyncContextManager[None]:
    """The MCP session manager's run scope, or nothing when MCP is off."""
    if not settings.MCP_ENABLED:
        return contextlib.nullcontext()
    from fluksio.mcp.server import mcp as mcp_server

    return mcp_server.session_manager.run()


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Start the flow engine alongside the API."""
    # The schema and the first superuser, before anything reads either. It is
    # idempotent, so a deployment that ran this from its own prestart step
    # pays a version check for it and nothing else.
    await run_in_threadpool(prepare, db_engine)
    event_bus.bind(asyncio.get_running_loop())

    # Everything acquired below registers how to release it, and the `finally`
    # at the end walks the list backwards. It used to be a fixed block of
    # shutdown steps *after* every acquisition — so a failure part-way through
    # startup reached none of them and left the worker pool's subprocesses and
    # every background task behind. Under `--reload` that is once per bad edit.
    started: list[Callable[[], Any]] = []

    async def _release() -> None:
        for close in reversed(started):
            try:
                result = close()
                if inspect.isawaitable(result):
                    await result
            except Exception:
                logger.exception("Error while shutting down")

    def _background(
        coro: Coroutine[Any, Any, None], name: str, once: bool = False
    ) -> asyncio.Task[None]:
        """Start a background task that says something if it ever stops.

        The loops catch their own exceptions *inside* the loop, so one raised
        anywhere else simply ended the task — an engine that went on serving
        with no metrics, no alerts or no artifact sweep and nothing anywhere
        saying so. ``once`` is for a task that is meant to finish.
        """
        task = asyncio.create_task(coro, name=name)

        def _finished(done: asyncio.Task[None]) -> None:
            if done.cancelled():
                return
            exc = done.exception()
            if exc is None:
                if not once:
                    logger.warning("Background task '%s' stopped on its own", name)
                return
            logger.error("Background task '%s' died: %s", name, exc, exc_info=exc)
            event_bus.publish(
                {"type": "engine_degraded", "detail": f"{name} stopped: {exc}"}
            )

        task.add_done_callback(_finished)
        started.append(task.cancel)
        return task

    # The startup itself is inside the try, so a step that raises part-way
    # through still reaches the release below.
    try:
        # Node code is user code, and `print` is how it says things.
        logs.install()
        init_secrets(settings.SECRETS_FILE, settings.SECRET_KEY)
        # Connectors register their node types before any flow is built with them.
        load_plugins()

        alerts = AlertManager(event_bus, config=read_alerts_config())
        execution = ExecutionService(
            queue=_work_queue(),
            max_workers=settings.FLOW_MAX_WORKERS,
            events=event_bus,
            max_cascades=settings.FLOW_MAX_CASCADES,
        )
        store = FlowStore(settings.FLOWS_DIR)
        # The packages node code imports. In the background, because a first
        # install of anything substantial takes minutes and `uv` is given five
        # of them twice over — held here, the container never answers its
        # healthcheck, autoheal restarts it, and it never finishes installing at
        # all. A node whose import is not there yet fails and is retried, which
        # is recoverable; a restart loop is not.
        _background(
            run_in_threadpool(modules.reconcile, store), "module-reconcile", once=True
        )
        # Beside the flows rather than in them: an artifact is what a run produced,
        # not something anyone wrote, so it has no business in the git repository.
        artifacts = ArtifactStore(settings.FLOWS_DIR.parent / "artifacts")
        app.state.artifact_store = artifacts
        accountant = ResourceAccountant(
            cpus=settings.FLOW_CPUS, gpus=settings.FLOW_GPUS
        )
        app.state.resources = accountant
        # Every machine a node could run on: this one, and whatever attaches.
        placer = Placer(local=accountant, events=event_bus)
        # Where more machines can be asked for when nothing attached will do.
        placer.provisioners = load_provisioners(settings.PROVISIONERS_FILE, event_bus)
        app.state.placer = placer
        pool = PythonWorkerPool(
            python=modules.venv_python(),
            size=settings.FLOW_MAX_WORKERS,
            events=event_bus,
            # A worker in this container writes to the store directly; a remote one
            # is given a URL instead. Node code calls the same two functions.
            env={
                ARTIFACT_DIR_ENV: str(artifacts.root),
                # Every slot can be busy at once, so a worker left to size its own
                # thread pool to the machine means as many processes as there are
                # slots, each believing it has the whole of it. A node that says
                # what it needs overrides this; one that says nothing gets a share.
                **fair_share_env(accountant.cpus, settings.FLOW_MAX_WORKERS),
            },
        )
        pool.start()
        # On a thread: stopping a worker waits up to five seconds on each child,
        # and with a pool per declared environment that is a shutdown the event
        # loop should not be holding.
        started.append(lambda: run_in_threadpool(pool.stop))
        app.state.worker_pool = pool
        # Assigned rather than passed both ways: the hub tells the placer when a
        # machine comes or goes, and the placer needs the hub to know what is there.
        worker_hub = RemoteWorkerHub(on_change=placer.wake)
        placer.hub = worker_hub
        app.state.worker_hub = worker_hub
        controller = FlowController(
            store=store,
            state=_state_backend(),
            events=event_bus,
            max_workers=settings.FLOW_MAX_WORKERS,
            fastapi_app=app,
            execution=execution,
            alerts=alerts,
            workers=pool,
            remote=worker_hub,
            resources=accountant,
            placer=placer,
        )
        app.state.flow_controller = controller
        # A "dashboard" alert channel puts its alert into the graph. Bound here
        # rather than passed in: the manager is built before the controller is.
        alerts.publish = lambda name, value: controller.publish_message(
            name, value, ValueSource(kind="api", id="alerts", label="Alerts")
        )
        dashboards = DashboardStore(controller.store)
        app.state.dashboard_store = dashboards
        controller.dashboards = dashboards
        # Charts need a deeper series than the default; tell the engine
        # before it starts recording.
        controller.set_history_limits(dashboards.history_requirements())
        # Runs read from a stream of their own: a burst of sweep runs must not
        # stand between the automations and their work, and a run that takes an
        # hour must not be judged by the cascade reaper's timings.
        run_service = RunService(
            controller=controller,
            queue=_work_queue("run"),
            state_factory=_run_state,
            artifacts=artifacts,
            parallel=settings.FLOW_MAX_RUNS,
        )
        app.state.run_service = run_service
        # Said out loud because it is the only way to tell that a settings file was
        # read at all — the numbers are what someone raising them is looking for.
        logger.info(
            "Engine limits: workers=%s cascades=%s runs=%s",
            settings.FLOW_MAX_WORKERS,
            execution.max_cascades,
            run_service.parallel,
        )
        watchdog = LoopWatchdog(event_bus)
        app.state.watchdog = watchdog

        _background(watchdog.run(), "loop-watchdog")
        _background(alerts.run(), "alert-manager")
        _background(MetricsCollector(event_bus).run(), "metrics-collector")
        _background(_sweep_artifacts(artifacts, controller), "artifact-gc")
        await controller.start()
        started.append(controller.stop)
        run_service.start()
        started.append(lambda: run_in_threadpool(run_service.stop))
        # A machine asked for and not yet arrived would hold an allocation nobody
        # is going to use.
        started.append(
            lambda: asyncio.gather(
                *(
                    run_in_threadpool(provisioner.shutdown)
                    for provisioner in placer.provisioners
                )
            )
        )
        started.append(lambda: run_in_threadpool(close_shared_client))
        # Optional, and off unless someone enrolled this installation: the
        # connector dials the portal, nothing dials in.
        cloud_task: asyncio.Task[None] | None = None
        app.state.cloud_connector = None
        app.state.cloud_task = None
        from fluksio.cloud import connector as cloud_connector

        if cloud_config.exists():
            cloud_connector.start(app)
            cloud_task = app.state.cloud_task

        def _stop_cloud() -> None:
            # Re-read from app.state: enrolling at runtime replaces this.
            running = getattr(app.state, "cloud_task", None) or cloud_task
            if running is not None:
                running.cancel()

        started.append(_stop_cloud)
        # Watched whether or not one exists now: enrolling from the CLI writes the
        # config from another process entirely, and an engine already serving
        # should pick it up rather than need restarting.
        _background(cloud_connector.watch_enrolment(app), "cloud-enrolment-watch")
        # A mounted sub-app gets no lifespan of its own, so the MCP session
        # manager is entered here; without it every /mcp request fails.
        async with _mcp_sessions():
            yield
    except Exception:
        # A startup that fails part-way lands here too, because everything
        # since `event_bus.bind` registered its own release.
        logger.exception("The engine could not start")
        raise
    finally:
        await _release()
        if settings.MCP_ENABLED:
            from fluksio.mcp.http import aclose

            await aclose()


# The schema enumerates every endpoint this installation serves, including the
# paths trigger nodes mount at runtime. That is exactly what a developer wants
# and exactly what an internet-facing deployment should not hand out, so it
# follows the environment — the same rule the portal's backend uses. The
# generated client is built from a local run, not from the deployed host.
_docs_enabled = settings.ENVIRONMENT != "production"

app = FastAPI(
    title=settings.PROJECT_NAME,
    version=__version__,
    openapi_url=f"{settings.API_V1_STR}/openapi.json" if _docs_enabled else None,
    docs_url="/docs" if _docs_enabled else None,
    redoc_url="/redoc" if _docs_enabled else None,
    generate_unique_id_function=custom_generate_unique_id,
    lifespan=lifespan,
)

# Set all CORS enabled origins
if settings.all_cors_origins:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=settings.all_cors_origins,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

app.include_router(api_router, prefix=settings.API_V1_STR)


@app.exception_handler(RebuildBusy)
async def rebuild_busy(request: Request, exc: RebuildBusy) -> JSONResponse:  # noqa: ARG001
    """Every route that deploys something answers a wedged rebuild the same way.

    The request was fine and retrying it may well work, so this is the engine
    saying it is busy rather than the request having gone wrong.
    """
    return JSONResponse(status_code=503, content={"detail": str(exc)})


# Tagged because the operation-id builder reads the first tag; the route
# itself stays out of the schema.
@app.get(
    "/.well-known/oauth-authorization-server",
    include_in_schema=False,
    tags=["oauth"],
)
def oauth_authorization_server() -> JSONResponse:
    """RFC 8414 metadata, so an agent can find its way in unaided.

    The authorization endpoint is the dashboard rather than the API: approving
    a client needs a signed-in human, and the browser session lives there.
    """
    issuer = settings.oauth_issuer
    return JSONResponse(
        content={
            "issuer": issuer,
            "authorization_endpoint": (
                f"{settings.FRONTEND_HOST.rstrip('/')}/oauth/authorize"
            ),
            "token_endpoint": f"{issuer}{settings.API_V1_STR}/oauth/token",
            "registration_endpoint": f"{issuer}{settings.API_V1_STR}/oauth/register",
            "jwks_uri": f"{issuer}/.well-known/jwks.json",
            "response_types_supported": ["code"],
            "grant_types_supported": ["authorization_code", "refresh_token"],
            "code_challenge_methods_supported": ["S256"],
            "token_endpoint_auth_methods_supported": ["none"],
            "scopes_supported": [security.MCP_SCOPE],
        },
        headers={"Cache-Control": "public, max-age=3600"},
    )


@app.get("/.well-known/jwks.json", include_in_schema=False, tags=["oauth"])
def jwks() -> JSONResponse:
    """The public half of the MCP signing key."""
    return JSONResponse(
        content=security.public_jwks(),
        headers={"Cache-Control": "public, max-age=3600"},
    )


# Mounted last, and at the root: the SDK serves both /mcp and the protected
# resource metadata that has to sit beside it, so mounting under /mcp would put
# that metadata somewhere no client looks for it.
if settings.MCP_ENABLED:
    from fluksio.mcp.http import build_http_app

    app.mount("/", build_http_app(app))