app/backend/fluksio/cloud/connector.py

"""The link this installation opens to its portal.

Same shape as the remote-worker agent, with the direction of the favour
reversed: there, a GPU box dials the engine so the engine can give it work;
here, the engine dials the portal so the portal can pass on what a browser
asked for. Both exist because the interesting machine is the one nobody can
route to.

A proxied request is dispatched straight into this process's own ASGI app
rather than over localhost HTTP. That matters beyond saving a hop: HTTP
trigger nodes install themselves into ``app.routes`` while flows run, and an
in-process transport sees that table live, with no port, TLS or self-address
to discover.
"""

from __future__ import annotations

import asyncio
import base64
import contextlib
import logging
import random
import time
import uuid
from dataclasses import replace
from datetime import timedelta
from typing import Any

import httpx
from fastapi import FastAPI

import fluksio
from fluksio.cloud import config as cloud_config

logger = logging.getLogger(__name__)

PROTOCOL = 1
HEARTBEAT_S = 20.0
STATUS_S = 60.0
BASE_BACKOFF_S = 1.0
MAX_BACKOFF_S = 30.0
#: How long a silent socket is given before it is treated as gone. A laptop
#: that suspended, or a wifi that changed underneath us, leaves a connection
#: the operating system still believes in — nothing arrives and nothing fails.
#: Only an unanswered ping tells us, so this is the worst case for noticing.
PING_INTERVAL_S = 20.0
PING_TIMEOUT_S = 20.0
CHUNK_BYTES = 256 * 1024
MAX_WS_MESSAGE = 1024 * 1024
#: Only the versioned API is served over the tunnel. The MCP mount and the
#: OAuth endpoints live outside it and stay local-only.
ALLOWED_PREFIX = "/api/v1/"
#: How often to look for a config that appeared while the engine was up.
#: Enrolment is something a person does and then waits on, so this is the
#: delay they sit through — short enough not to be worth a restart.
ENROL_POLL_S = 3.0


def start(app: FastAPI) -> None:
    """Dial the portal, replacing any link already up."""
    existing = getattr(app.state, "cloud_task", None)
    if existing is not None:
        existing.cancel()
    connector = CloudConnector(app)
    app.state.cloud_connector = connector
    app.state.cloud_task = asyncio.create_task(
        connector.serve_forever(), name="cloud-connector"
    )


async def watch_enrolment(app: FastAPI) -> None:
    """Notice an enrolment that happened outside this process.

    ``fluksio enroll`` writes the config against the database directly, with no
    idea whether an engine is running — so without this, pairing an installation
    that is already serving would take a restart to come into effect. Enrolling
    through the API starts the link itself and this sees a task already there.
    """
    while True:
        await asyncio.sleep(ENROL_POLL_S)
        task = getattr(app.state, "cloud_task", None)
        if task is not None and task.done():
            # It returns of its own accord when the config goes away, which is
            # what `fluksio disconnect` and the portal's own Disconnect do.
            app.state.cloud_task = None
            app.state.cloud_connector = None
            task = None
        # `load()`, not `exists()`: a file that cannot be read is not an
        # enrolment, and the connector would give up on it the instant it
        # started — which, started from here, is a restart every few seconds.
        # A config that is fine but unreachable keeps its task, and the
        # retrying belongs to the connector rather than to this.
        if task is None and cloud_config.load() is not None:
            logger.info("Enrolled while running; dialling the portal")
            start(app)


class CloudConnector:
    """Holds one websocket to the portal open, and answers over it."""

    def __init__(self, app: FastAPI) -> None:
        self._app = app
        self._calls: dict[str, asyncio.Task[None]] = {}
        self._streams: dict[str, asyncio.Task[None]] = {}
        self._connected = False
        self._connected_since: float | None = None
        self._last_error: str | None = None

    # -------------------------------------------------------------------------
    # Status, for the local settings screen
    # -------------------------------------------------------------------------

    def status(self) -> dict[str, Any]:
        config = cloud_config.load()
        return {
            "enrolled": config is not None,
            "connected": self._connected,
            "portal_url": config.portal_url if config else None,
            # Where a browser reaches the portal, which is not always where
            # this process does: enrolment may have named a container.
            "issuer": config.issuer if config else None,
            "portal_account": config.portal_account if config else None,
            "installation_id": config.installation_id if config else None,
            "last_error": self._last_error,
            "connected_since": self._connected_since,
        }

    # -------------------------------------------------------------------------
    # The connection
    # -------------------------------------------------------------------------

    async def serve_forever(self) -> None:
        """Hold the link up, and put it back up when it goes down.

        Every ending is the same ending here: a portal that closed the socket,
        a network that changed under it, a laptop that woke up somewhere else.
        The link is dialled again in all of them — the only question is how
        soon, and that turns on whether this attempt got as far as being
        attached. One that did and then dropped is a network event, so it
        retries at once; one that never stood up is met with a longer wait
        each time, because whatever is refusing is unlikely to stop within a
        second.
        """
        delay = BASE_BACKOFF_S
        while True:
            config = cloud_config.load()
            if config is None:
                # Disconnected locally while we were running.
                return
            attached = False
            try:
                attached = await self._session(config)
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                self._last_error = str(exc)
                logger.warning("Portal link down: %s", exc)
            finally:
                self._connected = False
                self._connected_since = None

            delay = BASE_BACKOFF_S if attached else min(MAX_BACKOFF_S, delay * 2)
            # Jittered, so a portal coming back up is not met by every
            # installation it serves in the same instant.
            wait = delay * (0.75 + random.random() * 0.5)
            logger.info("Reconnecting to the portal in %.0fs", wait)
            await asyncio.sleep(wait)

    async def _session(self, config: cloud_config.CloudConfig) -> bool:
        """One connection, from dial to close. True if it ever attached."""
        import websockets

        attached = False
        url = f"{config.ws_url}?token={config.token}"
        async with websockets.connect(
            url,
            max_size=MAX_WS_MESSAGE,
            ping_interval=PING_INTERVAL_S,
            ping_timeout=PING_TIMEOUT_S,
        ) as socket:
            await socket.send(
                _dump({"op": "hello", "protocol": PROTOCOL, "app_version": _version()})
            )
            welcome = _load(await socket.recv())
            if welcome.get("op") != "welcome":
                raise RuntimeError(
                    str(welcome.get("reason") or "refused by the portal")
                )

            self._adopt_owner(config, welcome.get("owner"))

            attached = True
            self._connected = True
            self._connected_since = time.time()
            self._last_error = None
            logger.info(
                "Attached to the portal as installation %s", config.installation_id
            )

            keepalive = asyncio.create_task(self._keepalive(socket, config))
            try:
                async for raw in socket:
                    frame = _load(raw)
                    op = frame.get("op")
                    if op == "req":
                        self._start_call(socket, frame)
                    elif op == "cancel":
                        self._cancel(str(frame.get("id") or ""))
                    elif op == "ws_open":
                        self._start_stream(socket, frame)
                    elif op == "ws_close":
                        self._cancel_stream(str(frame.get("id") or ""))
            finally:
                keepalive.cancel()
                # Awaited, not just cancelled: a keepalive that died of its own
                # accord holds the reason the link went, and an un-awaited task
                # takes it to the garbage collector instead of the log.
                with contextlib.suppress(asyncio.CancelledError, Exception):
                    await keepalive
                self._connected = False
                self._connected_since = None
                for task in list(self._calls.values()) + list(self._streams.values()):
                    task.cancel()
                self._calls.clear()
                self._streams.clear()
        return attached

    @staticmethod
    def _adopt_owner(config: cloud_config.CloudConfig, owner: Any) -> None:
        """Map the enrolling account to the portal account that owns us.

        Enrolment does this itself. This is for the enrolments that predate
        per-user mapping: without it their owner would be refused here until
        somebody enrolled the machine again — which, for a machine reached only
        through the portal, means standing in front of it. Runs on every attach
        because it is a no-op once the mapping is there.

        Never moves a mapping somebody else holds: that would be this code
        guessing at something enrolment was told.
        """
        from sqlmodel import Session, select

        from fluksio.core.db import engine
        from fluksio.models import User

        if not owner:
            return
        owner_id = str(owner)
        try:
            with Session(engine) as session:
                if session.exec(
                    select(User).where(User.portal_sub == owner_id)
                ).first():
                    return
                user = CloudConnector._local_account(session, config)
                if user is None or user.portal_sub:
                    return
                user.portal_sub = owner_id
                session.add(user)
                session.commit()
                logger.info("Mapped %s to the portal account that owns us", user.email)
        except Exception:
            # A mapping that could not be written is not a reason to drop the
            # link: everything else this connection does still works.
            logger.exception("Could not map the enrolling account to the portal owner")

    @staticmethod
    def _local_account(session: Any, config: cloud_config.CloudConfig) -> Any:
        """The account this installation acts as, or the one that now stands for it.

        ``local_user_id`` names whoever redeemed the claim code. A database
        restored from a backup, or rebuilt under an enrolment that outlived it,
        has the same operator behind a different row — and then this resolves to
        nobody, which would refuse the owner on a machine they may only be able
        to reach through the portal. That is the lockout the enrolment handshake
        exists to prevent, so it is closed here too: one superuser is not a guess,
        it is the account enrolment would have used. Several is a guess, and this
        code does not make it.

        The new id is written back, because the same field is what a screen
        paired through the portal borrows — a panel would otherwise be refused
        for the same reason and with no way to say so.
        """
        from fluksio.models import User

        try:
            user = session.get(User, uuid.UUID(config.local_user_id))
        except ValueError:
            user = None
        if user is not None:
            return user

        from sqlmodel import select

        superusers = list(
            session.exec(select(User).where(User.is_superuser == True))  # noqa: E712
        )
        if len(superusers) != 1:
            logger.warning(
                "The account this installation enrolled as is gone, and there "
                "%s to adopt unambiguously — enrol again from Settings.",
                "is no superuser" if not superusers else "are several superusers",
            )
            return None
        adopted = superusers[0]
        cloud_config.save(replace(config, local_user_id=str(adopted.id)))
        logger.warning(
            "The account this installation enrolled as is gone; acting as %s instead",
            adopted.email,
        )
        return adopted

    async def _keepalive(self, socket: Any, config: cloud_config.CloudConfig) -> None:
        """Heartbeats, plus a health snapshot the portal can show while offline."""
        last_status = 0.0
        while True:
            now = time.time()
            if now - last_status >= STATUS_S or last_status == 0.0:
                summary = await self._health_summary(config)
                await socket.send(
                    _dump(
                        {"op": "status", "summary": summary, "app_version": _version()}
                    )
                )
                last_status = now
            else:
                await socket.send(_dump({"op": "heartbeat"}))
            await asyncio.sleep(HEARTBEAT_S)

    async def _health_summary(
        self, config: cloud_config.CloudConfig
    ) -> dict[str, Any] | None:
        """This installation's own health, as the dashboard reads it.

        Fetched through the same in-process transport as everything else, with
        a short-lived local token: the observability API is authenticated, and
        this process is entitled to mint one for the enrolling user.

        ``failures_24h`` is added here rather than by ``/summary``: it is the
        portal's tile and nothing on this installation reads it, and summing
        the same rollups over the same window the app's own Home tile sums
        keeps the two from drifting apart the way they already did once.
        """
        from fluksio.core import security

        token = security.create_access_token(config.local_user_id, timedelta(minutes=5))
        headers = {"Authorization": f"Bearer {token}"}
        try:
            async with self._client() as client:
                response = await client.get(
                    "/api/v1/observability/summary", headers=headers
                )
                if response.status_code != 200:
                    return None
                summary: dict[str, Any] = response.json()
                rollups = await client.get(
                    "/api/v1/observability/flows",
                    params={"hours": 24},
                    headers=headers,
                )
            if rollups.status_code == 200:
                summary["failures_24h"] = sum(row["errors"] for row in rollups.json())
            return summary
        except Exception:
            logger.debug("Could not collect a health summary for the portal")
        return None

    # -------------------------------------------------------------------------
    # Proxied requests
    # -------------------------------------------------------------------------

    def _client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(
            transport=httpx.ASGITransport(app=self._app),
            base_url="http://installation.local",
            timeout=httpx.Timeout(300.0, connect=5.0),
        )

    def _start_call(self, socket: Any, frame: dict[str, Any]) -> None:
        call_id = str(frame.get("id") or "")
        if not call_id:
            return
        task = asyncio.create_task(self._serve(socket, frame))
        self._calls[call_id] = task
        task.add_done_callback(lambda _t: self._calls.pop(call_id, None))

    async def _serve(self, socket: Any, frame: dict[str, Any]) -> None:
        call_id = str(frame["id"])
        path = str(frame.get("path") or "")
        try:
            if not path.startswith(ALLOWED_PREFIX):
                # Defence in depth: the hub only ever forwards API paths, and
                # an installation that trusted anything else would be an SSRF
                # gadget aimed at its own process.
                await socket.send(
                    _dump({"op": "err", "id": call_id, "message": "path not allowed"})
                )
                return
            body = frame.get("body")
            content = base64.b64decode(body) if body else None
            query = str(frame.get("query") or "")

            headers = dict(frame.get("headers") or {})
            # Set here rather than trusted from the frame: a browser can send
            # any header it likes through the proxy, and this one decides where
            # a pairing device's credential is minted. Arriving on this socket
            # is the only thing that makes it true, and the value is this
            # process's own so nothing off the network can imitate it.
            headers[cloud_config.VIA_HEADER] = cloud_config.VIA_PORTAL

            async with self._client() as client:
                request = client.build_request(
                    str(frame.get("method") or "GET"),
                    f"{path}?{query}" if query else path,
                    headers=headers,
                    content=content,
                )
                response = await client.send(request, stream=True)
                try:
                    await socket.send(
                        _dump(
                            {
                                "op": "res",
                                "id": call_id,
                                "status": response.status_code,
                                "headers": dict(response.headers),
                            }
                        )
                    )
                    async for chunk in response.aiter_bytes(CHUNK_BYTES):
                        await socket.send(
                            _dump(
                                {
                                    "op": "res_body",
                                    "id": call_id,
                                    "body": base64.b64encode(chunk).decode("ascii"),
                                    "done": False,
                                }
                            )
                        )
                    await socket.send(
                        _dump({"op": "res_body", "id": call_id, "done": True})
                    )
                finally:
                    await response.aclose()
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("Proxied %s failed: %s", path, exc)
            with contextlib.suppress(Exception):
                await socket.send(
                    _dump({"op": "err", "id": call_id, "message": str(exc)})
                )

    def _cancel(self, call_id: str) -> None:
        task = self._calls.pop(call_id, None)
        if task is not None:
            task.cancel()

    # -------------------------------------------------------------------------
    # Bridged live stream
    # -------------------------------------------------------------------------

    def _start_stream(self, socket: Any, frame: dict[str, Any]) -> None:
        stream_id = str(frame.get("id") or "")
        if not stream_id:
            return
        task = asyncio.create_task(self._stream(socket, frame))
        self._streams[stream_id] = task
        task.add_done_callback(lambda _t: self._streams.pop(stream_id, None))

    async def _stream(self, socket: Any, frame: dict[str, Any]) -> None:
        """Feed the engine's live events down the tunnel.

        The flows websocket is subscribed to directly rather than dialled:
        httpx cannot speak websocket, and everything that endpoint does — check
        the token, send a snapshot, forward the bus — is a few lines against
        objects this process already holds. The snapshot is built by the route
        itself, so the two cannot drift apart behind the portal's back.
        """
        from sqlmodel import Session

        from fluksio.api.deps import user_from_token
        from fluksio.api.routes.flows import (
            event_for_panel,
            panel_scope,
            snapshot_payload,
        )
        from fluksio.core.db import engine
        from fluksio.flow.events import event_bus

        stream_id = str(frame["id"])
        path = str(frame.get("path") or "")
        query = str(frame.get("query") or "")
        token = _query_value(query, "token")

        if not path.startswith(ALLOWED_PREFIX) or not path.endswith("/flows/ws"):
            await socket.send(_dump({"op": "ws_close", "id": stream_id, "code": 1008}))
            return
        with Session(engine) as session:
            if user_from_token(session, token) is None:
                await socket.send(
                    _dump({"op": "ws_close", "id": stream_id, "code": 1008})
                )
                return
        only = panel_scope(token, self._app)

        await socket.send(_dump({"op": "ws_open_ok", "id": stream_id}))

        controller = getattr(self._app.state, "flow_controller", None)
        if controller is not None:
            await socket.send(
                _dump(
                    {
                        "op": "ws_msg",
                        "id": stream_id,
                        "text": _dump(snapshot_payload(controller, only)),
                    }
                )
            )

        try:
            async with event_bus.subscribe() as queue:
                while True:
                    event = await queue.get()
                    if only is not None and event.get("type") == "dashboard_changed":
                        # What this screen shows may have just changed under
                        # it. Rescope, then resend the snapshot so a dashboard
                        # it has only now been assigned draws values instead of
                        # blanks. ``or set()`` is the point: a panel that was
                        # deleted scopes to nothing, and None would widen this
                        # socket to everything on the bus.
                        only = panel_scope(token, self._app) or set()
                        if controller is not None:
                            catch_up = snapshot_payload(controller, only)
                            await socket.send(
                                _dump(
                                    {
                                        "op": "ws_msg",
                                        "id": stream_id,
                                        "text": _dump(catch_up),
                                    }
                                )
                            )
                    elif only is not None and not event_for_panel(event, only):
                        continue
                    await socket.send(
                        _dump({"op": "ws_msg", "id": stream_id, "text": _dump(event)})
                    )
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.debug("Live stream %s ended: %s", stream_id, exc)
            with contextlib.suppress(Exception):
                await socket.send(
                    _dump({"op": "ws_close", "id": stream_id, "code": 1011})
                )

    def _cancel_stream(self, stream_id: str) -> None:
        task = self._streams.pop(stream_id, None)
        if task is not None:
            task.cancel()


def _dump(payload: Any) -> str:
    import json

    return json.dumps(payload, default=str)


def _load(raw: Any) -> dict[str, Any]:
    import json

    data: dict[str, Any] = json.loads(raw)
    return data


def _query_value(query: str, key: str) -> str:
    from urllib.parse import parse_qs

    values = parse_qs(query).get(key) or [""]
    return values[0]


def _version() -> str:
    return fluksio.__version__