stackpilot/backend/routers/ws.py

"""WebSocket endpoints for real-time log streaming and Docker events."""
from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import threading
from collections.abc import Awaitable, Callable

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
from jose import JWTError
from sqlmodel import Session

from auth import decode_token, resolve_token_user
from database import engine
from models.setting import EVENT_PULL_FAILED, EVENT_STACK_ERROR, EVENT_STACK_START
from models.stack import Stack
from services import (
    audit_service,
    compose_run_service,
    compose_service,
    exec_service,
    notify_service,
    update_service,
)

logger = logging.getLogger("stackpilot.ws")

router = APIRouter(tags=["ws"])


def _user_for(token: str | None):
    """The live user behind a socket's token, or None.

    Resolves against the database rather than reading the role straight off the
    JWT: a socket can outlive a demotion, a disabled account or a password
    reset, and the exec endpoint below is root-equivalent on the host. Same
    check the HTTP routes make.
    """
    if not token:
        return None
    try:
        payload = decode_token(token, "access")
    except (JWTError, Exception):  # noqa: BLE001
        return None
    with Session(engine) as session:
        user = resolve_token_user(session, payload)
        if user:
            session.expunge(user)
        return user


async def _authorize(websocket: WebSocket, token: str | None):
    """The user behind the token, or None after closing the socket with 4401."""
    user = _user_for(token)
    if user is None:
        await websocket.close(code=4401)
    return user


async def _authorize_admin(websocket: WebSocket, token: str | None):
    """Like _authorize but also requires the admin role.

    Anything that changes what runs on the host is admin-only over HTTP, so it
    has to be over a socket too — the deploy console runs the same
    ``compose up`` the REST ``/start`` does. Closes 4401 on a missing/invalid
    token, 4403 on a valid non-admin one."""
    user = _user_for(token)
    if user is None:
        await websocket.close(code=4401)
        return None
    if user.role != "admin":
        await websocket.close(code=4403)
        return None
    return user


async def _wait_disconnect(websocket: WebSocket) -> None:
    """Return once the client has gone. Whatever it sends is ignored.

    A stream that only ever *sends* notices a closed socket on its next send —
    which for a quiet log or an idle Docker host can be hours away, and until
    then everything behind the socket (a ``compose logs -f``, an event reader
    thread) stays alive for nobody.
    """
    try:
        while True:
            message = await websocket.receive()
            if message.get("type") == "websocket.disconnect":
                return
    except Exception:  # noqa: BLE001 - a broken socket is a gone client
        return


async def _until_gone(websocket: WebSocket, work: Awaitable[None]) -> None:
    """Run ``work`` until it finishes or the client disconnects, whichever is first.

    Cancelling ``work`` is what stops whatever it was reading from; see
    :func:`compose_service.stream_compose`.
    """
    task = asyncio.ensure_future(work)
    gone = asyncio.create_task(_wait_disconnect(websocket))
    try:
        await asyncio.wait({task, gone}, return_when=asyncio.FIRST_COMPLETED)
    finally:
        for pending in (task, gone):
            pending.cancel()
        await asyncio.gather(task, gone, return_exceptions=True)
    if task.done() and not task.cancelled() and task.exception() is not None:
        exc = task.exception()
        if not isinstance(exc, WebSocketDisconnect):
            with contextlib.suppress(Exception):
                await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))


async def _stream_logs(websocket: WebSocket, stack_id: str, service: str | None):
    """Stream `docker compose logs -f` output to the client."""
    args = ["logs", "--no-color", "--tail", "200", "--timestamps", "-f"]
    if service:
        # "--" so a service name can never be read as a compose option.
        args += ["--", service]

    async def pump() -> None:
        # aclosing: cancelled while *sending*, the loop is left by an exception
        # outside the generator, which would then only be closed — and its
        # `compose logs -f` killed — whenever garbage collection got to it.
        async with contextlib.aclosing(
            compose_service.stream_compose(stack_id, args)
        ) as lines:
            async for line in lines:
                await _send_log(line)

    async def _send_log(line: str) -> None:
        await websocket.send_text(
            json.dumps({"type": "log", "stack_id": stack_id, "service": service, "line": line})
        )

    await _until_gone(websocket, pump())


@router.websocket("/ws/logs/{stack_id}")
async def ws_stack_logs(
    websocket: WebSocket,
    stack_id: str,
    token: str | None = Query(default=None),
):
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    await _stream_logs(websocket, stack_id, None)


@router.websocket("/ws/logs/{stack_id}/{service}")
async def ws_service_logs(
    websocket: WebSocket,
    stack_id: str,
    service: str,
    token: str | None = Query(default=None),
):
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    await _stream_logs(websocket, stack_id, service)


# --------------------------------------------------------------------------- #
# Deploy / update: a compose run the socket only watches
# --------------------------------------------------------------------------- #


def _stack_known(stack_id: str) -> bool:
    with Session(engine) as session:
        return session.get(Stack, stack_id) is not None


async def _watch_run(
    websocket: WebSocket,
    stack_id: str,
    action: str,
    username: str,
    stream: Callable,
    on_finish: Callable[[int | None], Awaitable[None]],
) -> None:
    """Start (or attach to) the compose run and relay it until it ends.

    The run belongs to :mod:`services.compose_run_service`, not to this socket:
    closing the deploy console no longer releases the stack lock under a
    running compose, and the audit entry carries the real exit code.
    """
    if not _stack_known(stack_id):
        with contextlib.suppress(Exception):
            await websocket.send_text(
                json.dumps({"type": "error", "detail": f"Stack '{stack_id}' not found"})
            )
            await websocket.close(code=4404)
        return
    try:
        run, queue, _attached = compose_run_service.start(
            stack_id, action, username, stream, on_finish
        )
    except compose_run_service.RunBusy as exc:
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
            await websocket.close(code=4409)
        return

    async def relay() -> None:
        while True:
            item = await queue.get()
            if item is None:
                return
            kind, payload = item
            if kind == "log":
                await websocket.send_text(json.dumps({"type": "log", "line": payload}))
            else:
                await websocket.send_text(json.dumps({"type": "done", "returncode": payload}))

    try:
        await _until_gone(websocket, relay())
    finally:
        run.unsubscribe(queue)
    with contextlib.suppress(Exception):
        await websocket.close()


@router.websocket("/ws/deploy/{stack_id}")
async def ws_deploy(
    websocket: WebSocket,
    stack_id: str,
    token: str | None = Query(default=None),
):
    """Run `docker compose up -d` and stream its output (image pulls, container
    creation) to the browser so the user sees deploy progress live. Records the
    same audit entry and notification as the REST `/start` endpoint. Admin-only,
    like that endpoint."""
    await websocket.accept()
    user = await _authorize_admin(websocket, token)
    if not user:
        return
    username = user.username

    async def finish(rc: int | None) -> None:
        with Session(engine) as session:
            audit_service.record(
                session, user=username, action="stack.start", target=stack_id,
                detail=f"rc={rc} (deploy console)", ip="ws",
            )
            if rc == 0:
                await notify_service.notify(
                    EVENT_STACK_START, f"Stack '{stack_id}' started",
                    "compose up completed successfully.", session,
                )
            else:
                await notify_service.notify(
                    EVENT_STACK_ERROR, f"Stack '{stack_id}' start failed",
                    "compose up returned a non-zero exit code.", session,
                )

    await _watch_run(
        websocket, stack_id, "start", username,
        lambda: compose_service.stream_up(stack_id), finish,
    )


@router.websocket("/ws/update/{stack_id}")
async def ws_update(
    websocket: WebSocket,
    stack_id: str,
    token: str | None = Query(default=None),
):
    """Run `docker compose pull && up -d` and stream its output, so the stacks
    list can render real update progress. Same audit/notify contract as the
    REST `/update` endpoint, which stays for non-interactive callers."""
    await websocket.accept()
    user = await _authorize_admin(websocket, token)
    if not user:
        return
    username = user.username

    async def finish(rc: int | None) -> None:
        with Session(engine) as session:
            audit_service.record(
                session, user=username, action="stack.update", target=stack_id,
                detail=f"rc={rc} (update stream)", ip="ws",
            )
            if rc == 0:
                await notify_service.notify(
                    EVENT_STACK_START, f"Stack '{stack_id}' updated",
                    "compose pull + up completed successfully.", session,
                )
            else:
                await notify_service.notify(
                    EVENT_PULL_FAILED, f"Stack '{stack_id}' update failed",
                    "compose pull/up returned a non-zero exit code.", session,
                )
        if rc == 0:
            await asyncio.to_thread(update_service.refresh_stack_local, stack_id)

    await _watch_run(
        websocket, stack_id, "update", username,
        lambda: compose_service.stream_update(stack_id), finish,
    )


# --------------------------------------------------------------------------- #
# Docker events
# --------------------------------------------------------------------------- #

#: Docker event types worth telling the UI about. Filtered daemon-side, so the
#: bulk of the firehose never crosses the socket.
_EVENT_TYPES = ["container", "image", "network", "volume"]

#: Container actions that say nothing about state a page renders. exec_* alone
#: is three events per web-terminal keystroke session, and `top`/`attach` fire
#: whenever something inspects a container — invalidating queries on those would
#: make the stream noisier than the polling it replaces.
_IGNORED_ACTIONS = {
    "exec_create", "exec_start", "exec_die", "exec_detach",
    "attach", "top", "resize", "archive-path", "extract-to-dir",
}


@router.websocket("/ws/events")
async def ws_events(
    websocket: WebSocket,
    token: str | None = Query(default=None),
):
    """Stream Docker events so the UI can refresh on change instead of polling.

    Every page used to poll its own endpoint every few seconds. Almost all of
    that state only changes when Docker does something, which is exactly what
    this reports — so the client refreshes on an event and keeps a slow poll as
    a safety net.

    Payload per event::

        {"type": "event", "resource": "container", "action": "start",
         "container": "jellyfin", "stack": "jellyfin"}

    ``resource`` is what the client needs to decide which queries to drop.

    The blocking read runs on a thread of its own, not on the event loop's
    shared executor: every open tab holds one reader for as long as it is open,
    and on the executor a handful of tabs used up every worker — after which
    each ``asyncio.to_thread`` (backups, the dashboard, NFS destinations) waited
    forever.
    """
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    from docker_client import get_client

    loop = asyncio.get_running_loop()
    queue: asyncio.Queue = asyncio.Queue()
    guard = threading.Lock()
    state: dict = {"stream": None, "stop": False}

    def post(item) -> None:
        try:
            loop.call_soon_threadsafe(queue.put_nowait, item)
        except RuntimeError:  # the loop is gone (shutdown)
            pass

    def reader() -> None:
        """Blocking read of the event stream, handed to the loop thread-safely."""
        try:
            stream = get_client().events(decode=True, filters={"type": _EVENT_TYPES})
            with guard:
                if state["stop"]:
                    # The client left while we were connecting.
                    with contextlib.suppress(Exception):
                        stream.close()
                    return
                state["stream"] = stream
            for event in stream:
                post(event)
        except Exception:  # noqa: BLE001 - a closed stream lands here on teardown
            pass
        finally:
            post(None)

    threading.Thread(target=reader, name="docker-events", daemon=True).start()
    gone = asyncio.create_task(_wait_disconnect(websocket))
    try:
        await websocket.send_text(json.dumps({"type": "ready"}))
        while True:
            getter = asyncio.ensure_future(queue.get())
            done, _ = await asyncio.wait({getter, gone}, return_when=asyncio.FIRST_COMPLETED)
            if getter not in done:
                getter.cancel()
                break  # the client disconnected
            event = getter.result()
            if event is None:  # reader finished — daemon gone or stream closed
                await websocket.send_text(
                    json.dumps({"type": "error", "detail": "Docker event stream ended"})
                )
                break
            action = (event.get("Action") or "").split(":")[0]
            if action in _IGNORED_ACTIONS:
                continue
            actor = event.get("Actor", {}) or {}
            attrs = actor.get("Attributes", {}) or {}
            await websocket.send_text(
                json.dumps(
                    {
                        "type": "event",
                        "resource": event.get("Type"),
                        "action": event.get("Action"),
                        "container": attrs.get("name"),
                        "stack": attrs.get("com.docker.compose.project"),
                    }
                )
            )
    except WebSocketDisconnect:
        pass
    except Exception as exc:  # noqa: BLE001
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
    finally:
        gone.cancel()
        # Closing the stream is what actually unblocks the reader thread.
        with guard:
            state["stop"] = True
            stream = state["stream"]
        if stream is not None:
            with contextlib.suppress(Exception):
                stream.close()


@router.websocket("/ws/exec/{container_id}")
async def ws_exec(
    websocket: WebSocket,
    container_id: str,
    token: str | None = Query(default=None),
    cmd: str | None = Query(default=None),
):
    """Interactive shell into a compose-managed container (admin only)."""
    await websocket.accept()
    user = await _authorize_admin(websocket, token)
    if not user:
        return
    username = user.username
    shell = cmd or exec_service.DEFAULT_SHELL

    try:
        exec_id = exec_service.create_exec(container_id, [shell])
        holder, raw = exec_service.start_exec(exec_id)
    except Exception as exc:  # noqa: BLE001
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
        with contextlib.suppress(Exception):
            await websocket.close()
        return

    with contextlib.suppress(Exception):
        with Session(engine) as session:
            audit_service.record(
                session, user=username, action="container.exec",
                target=container_id[:12], detail=shell, ip="ws",
            )

    try:
        await exec_service.pump_exec(websocket, exec_id, holder, raw)
    except WebSocketDisconnect:
        pass
    finally:
        with contextlib.suppress(Exception):
            await websocket.close()