hermes-agent/hermes_cli/web_routers/display.py

"""``/api/display/ws`` — raw RFB over WebSocket for the Bot Desktop viewer.

The Desktop renderer calls ``display.observe`` on its authenticated ``/api/ws`` connection, gets a
single-use 30 s ticket pinned to that profile's RFB socket, then opens this route with
``?display_ticket=``. No websockify, no new port: the bridge splices the profile's 0600 Unix socket
into the WebSocket as binary frames with backpressure both ways, and runs the client stream through
:class:`tools.bot_desktop.rfb_filter.RfbClientFilter` so keyboard, pointer and clipboard reach Xvnc
only from the viewer that currently holds the lease. noVNC's ``viewOnly`` is UX; this is the gate.

A lease change closes the evicted viewer's socket with 4000 ``control-taken`` so its UI drops back to
Watch mode and reconnects.

Design notes. The ticket rides in the query string on purpose: noVNC's Websock owns the socket and
cannot negotiate a subprotocol or add a header, and the ticket is single-use and expires in 30 s, so
a logged URL is spent by the time anyone reads it. The bridge's input gate re-reads ``lease.json`` at
most every ``_LEASE_REFRESH_S`` (250 ms) between in-process ``on_change`` callbacks: that bounds how
long another PROCESS's takeover can go unnoticed, and is the price of not stat-ing the lease file on
the event loop for every pointer move.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from hermes_cli.web_server_chat import _ws_request_is_allowed

_log = logging.getLogger(__name__)
router = APIRouter()

_ACTIVITY_STAMP_S = 60.0
_READ_CHUNK = 64 * 1024
_CLOSE_CONTROL_TAKEN = 4000
_CLEAN_CLOSE = frozenset({1000, 1001})
_CLOSE_DESKTOP_GONE = 4001
_CLOSE_BAD_TICKET = 4401
_CLOSE_NOT_ALLOWED = 4403
_CLOSE_PROTOCOL = 1003
_LEASE_REFRESH_S = 0.25


def _should_evict(held: dict, lease, viewer_id: str) -> bool:
    """A viewer that held control during this connection and lost it to ANOTHER human is kicked so its
    UI repaints; a plain hand-back to the agent, pure watchers and the new holder stay connected. The
    hand-back also forgets that this viewer ever held: after it, they are a plain watcher again and a
    later takeover by someone else must not evict them. ``held`` is the per-connection memory."""
    from tools.bot_desktop import lease as _lease
    if lease.holder != _lease.HUMAN:
        held["ever"] = False
        return False
    if lease.viewer_id == viewer_id:
        held["ever"] = True
        return False
    return bool(held["ever"])


def _consume_display_ticket(ws: WebSocket) -> Optional[dict]:
    from hermes_cli.dashboard_auth.ws_tickets import TicketInvalid, consume_ticket
    ticket = ws.query_params.get("display_ticket", "")
    if not ticket:
        return None
    try:
        info = consume_ticket(ticket)
    except TicketInvalid:
        return None
    if info.get("provider") != "bot-desktop" or not info.get("hermes_home"):
        return None
    return info


@router.websocket("/api/display/ws")
async def display_ws(ws: WebSocket) -> None:
    # Host/Origin/client-IP policy is refused BEFORE accept, like every other dashboard route: a
    # cross-origin page never gets a completed handshake. Everything after that (ticket, desktop
    # state) is refused AFTER accept so the code + reason arrive in a close frame — a close before
    # accept surfaces in the browser as an opaque HTTP 403 and the renderer cannot tell "re-observe"
    # (4401) from "screen is gone" (4001).
    if not _ws_request_is_allowed(ws):
        await ws.close(code=_CLOSE_NOT_ALLOWED)
        return
    await ws.accept()
    info = _consume_display_ticket(ws)
    if info is None:
        await ws.close(code=_CLOSE_BAD_TICKET, reason="display ticket missing, expired or used")
        return
    await _bridge(ws, info)


async def _bridge(ws: WebSocket, info: dict) -> None:
    """Pump RFB bytes between the viewer socket (already accepted) and THIS profile's Xvnc, gated by
    the lease."""
    from hermes_constants import hermes_home_key
    from tools.bot_desktop import lease as _lease
    from tools.bot_desktop.rfb_filter import RfbClientFilter
    from pathlib import Path

    sock = Path(info["hermes_home"]) / "bot-desktop" / "rfb.sock"
    profile_home = str(info["hermes_home"])
    profile_key = hermes_home_key(profile_home)
    viewer_id = str(info.get("viewer_id") or info.get("user_id") or "viewer")
    if not sock.exists():
        await ws.close(code=_CLOSE_DESKTOP_GONE, reason="Bot Desktop is not running")
        return
    try:
        reader, writer = await asyncio.open_unix_connection(str(sock))
    except OSError as exc:
        _log.warning("display ws: cannot reach RFB socket %s: %s", sock, exc)
        await ws.close(code=_CLOSE_DESKTOP_GONE, reason="Bot Desktop socket unreachable")
        return

    loop = asyncio.get_running_loop()
    evicted = asyncio.Event()
    held = {"ever": _lease.viewer_may_send_input(viewer_id, profile_key=profile_home)}
    # Input gate cache: reading lease.json per client message (a stat + read on the event loop for
    # every pointer move) is replaced by a decision refreshed on this process's on_change callback
    # and by a file re-read at most every _LEASE_REFRESH_S, so another process's takeover still lands.
    allowed = {"input": held["ever"], "at": loop.time()}

    def _refresh_allowed(lease=None) -> None:
        if lease is None:
            lease = _lease.get(profile_key=profile_home)
        allowed["input"] = lease.holder == _lease.HUMAN and lease.viewer_id == viewer_id
        allowed["at"] = loop.time()

    def _may_send_input() -> bool:
        if loop.time() - allowed["at"] > _LEASE_REFRESH_S:
            _refresh_allowed()
        return allowed["input"]

    def _on_lease(key: str, lease) -> None:
        if key != profile_key:
            return
        loop.call_soon_threadsafe(_refresh_allowed, lease)
        if _should_evict(held, lease, viewer_id):
            loop.call_soon_threadsafe(evicted.set)
    unsubscribe = _lease.on_change(_on_lease)

    rfb_filter = RfbClientFilter(_may_send_input)

    viewer_closed = asyncio.Event()
    activity_file = Path(profile_home) / "bot-desktop" / "activity"
    stamped = {"at": 0.0}

    def _stamp_activity() -> None:
        # An attached viewer is use: the idle auto-stop must not take a screen someone is watching.
        if loop.time() - stamped["at"] < _ACTIVITY_STAMP_S:
            return
        stamped["at"] = loop.time()
        try:
            activity_file.touch()
        except OSError:
            pass

    async def rfb_to_ws() -> None:
        while True:
            chunk = await reader.read(_READ_CHUNK)
            if not chunk:
                return
            _stamp_activity()
            await ws.send_bytes(chunk)  # awaiting the send is the backpressure toward Xvnc

    async def ws_to_rfb() -> None:
        while True:
            message = await ws.receive()
            if message.get("type") == "websocket.disconnect":
                # 1000/1001 = the viewer closed the window; anything else is a dropped link.
                if message.get("code") in _CLEAN_CLOSE:
                    viewer_closed.set()
                return
            data = message.get("bytes")
            if data is None:
                await ws.close(code=_CLOSE_PROTOCOL, reason="RFB is binary")
                return
            try:
                allowed = rfb_filter.feed(data)
            except ValueError as exc:
                await ws.close(code=_CLOSE_PROTOCOL, reason=str(exc)[:100])
                return
            if allowed:
                writer.write(allowed)
                await writer.drain()  # backpressure toward the browser

    async def watch_eviction() -> None:
        await evicted.wait()
        await ws.close(code=_CLOSE_CONTROL_TAKEN, reason="control-taken")

    tasks = [asyncio.create_task(rfb_to_ws()), asyncio.create_task(ws_to_rfb()),
             asyncio.create_task(watch_eviction())]
    try:
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for t in pending:
            t.cancel()
        # Cancelled pumps must finish before we tear down the socket they hold, or they outlive the
        # bridge on the loop (a viewer reconnecting in a loop piled them up).
        await asyncio.gather(*pending, return_exceptions=True)
        for t in done:
            exc = t.exception()
            if exc and not isinstance(exc, (WebSocketDisconnect, ConnectionError)):
                _log.debug("display ws ended: %r", exc)
    finally:
        unsubscribe()
        writer.close()
        try:
            await writer.wait_closed()
        except OSError:  # Xvnc already gone (ECONNRESET / EPIPE on the FIN)
            pass
        # Closing the viewer window hands control back. A DROPPED link (laptop lid, Wi-Fi, 1006)
        # keeps the human's exclusion: they may be mid-login on that screen and the agent must not
        # resume into it. The Desktop reconnects into the same lease, or the human hands back.
        if viewer_closed.is_set() and _lease.viewer_may_send_input(viewer_id, profile_key=profile_home):
            _lease.release(viewer_id, profile_key=profile_home)
        try:
            await ws.close()
        except Exception:  # already closed by the peer or by an eviction
            pass