arc53--docsgpt/application/devices/broker.py

"""Redis-backed broker that routes invocations to active device sessions.

The device CLI only ever talks to the web process (poll / SSE / output
POST), but an agent may dispatch a command from *any* process — notably a
Celery worker during a scheduled run. An in-memory broker can't bridge
that gap: the worker and the web process don't share Python memory, so a
worker-side dispatch never reaches the web-side SSE session and the tool
times out. Routing through Redis makes every hop cross-process.

What lives in Redis (shared) vs. in the process (ephemeral):

* ``dev:cmd:{device_id}``   — list of queued command envelopes (JSON).
* ``dev:ticket:{device_id}``— poll-issued SSE upgrade ticket (string, TTL).
* ``dev:inv:{invocation_id}``— invocation metadata hash (status, result).
* ``dev:out:{invocation_id}``— stream of stdout/stderr/control chunks.
* ``SessionState`` is per-connection state for the one SSE handler that
  owns the live socket; it stays in that process's memory.

Lifecycle:
1. Agent (any process) calls ``dispatch_invocation`` → metadata hash + an
   RPUSH onto the device's command list.
2. The web SSE handler blocks on that list (``next_command``) and emits
   each envelope to the wire; an offline device leaves the envelope queued
   until its next ``/poll`` issues a ticket and upgrades to SSE.
3. The CLI POSTs ack + chunked output back; ``submit_output_chunk`` XADDs
   chunks to the invocation's output stream and updates the hash.
4. A ``control`` chunk closes the invocation; ``drain_output`` (in the
   dispatching process) reads the stream from the start and stops on it.
"""

from __future__ import annotations

import json
import logging
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, Iterator, Optional

from application.cache import get_redis_instance
from application.core.settings import settings


logger = logging.getLogger(__name__)


# Kept for backwards compatibility with callers/tests importing the
# sentinel; the Redis path signals end-of-output with a ``control`` chunk.
INVOCATION_DONE = object()


def _cmd_key(device_id: str) -> str:
    return f"dev:cmd:{device_id}"


def _ticket_key(device_id: str) -> str:
    return f"dev:ticket:{device_id}"


def _inv_key(invocation_id: str) -> str:
    return f"dev:inv:{invocation_id}"


def _out_key(invocation_id: str) -> str:
    return f"dev:out:{invocation_id}"


def _as_str(value: Any) -> str:
    if isinstance(value, (bytes, bytearray)):
        return value.decode("utf-8", "replace")
    return str(value)


@dataclass
class Invocation:
    """Snapshot of an invocation's cross-process state.

    Returned by ``dispatch_invocation`` (fresh) and ``get_invocation`` (read
    from Redis). Carries enough for ownership checks and audit without the
    caller touching Redis directly.
    """

    invocation_id: str
    device_id: str
    completed: bool = False
    exit_code: Optional[int] = None
    duration_ms: Optional[int] = None
    error: Optional[str] = None
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    stdout_bytes: int = 0
    stderr_bytes: int = 0


@dataclass
class SessionState:
    """Per-connection state for the SSE handler that owns the live socket."""

    session_id: str
    device_id: str
    user_id: str
    last_event_id: int = 0
    last_activity_at: float = field(default_factory=time.time)
    closed: threading.Event = field(default_factory=threading.Event)


class DeviceBroker:
    """Cross-process device registry backed by Redis.

    Redis is the source of truth for queued commands, output, and tickets,
    so dispatch and drain work regardless of which process they run in. A
    small in-memory map of live SSE sessions stays local to the web process
    that holds each socket (sessions are inherently per-connection).
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._sessions_by_id: Dict[str, SessionState] = {}
        self._sessions_by_device: Dict[str, SessionState] = {}

    # ------------------------------------------------------------------
    # Session lifecycle (web process / CLI side)
    # ------------------------------------------------------------------
    def register_session(self, device_id: str, user_id: str) -> SessionState:
        """Open the SSE session for ``device_id``, adopting its poll ticket.

        The poll-issued ticket becomes the ``session_id`` so the URL the CLI
        opens matches the live session. A previous local session for the
        same device is closed and replaced (the CLI reconnected).
        """
        redis = get_redis_instance()
        issued = None
        if redis is not None:
            try:
                issued = redis.get(_ticket_key(device_id))
                if issued is not None:
                    redis.delete(_ticket_key(device_id))
            except Exception:
                logger.exception("ticket lookup failed for %s", device_id)
        session_id = _as_str(issued) if issued else f"st_{uuid.uuid4().hex}"
        sess = SessionState(
            session_id=session_id, device_id=device_id, user_id=user_id
        )
        with self._lock:
            prior = self._sessions_by_device.get(device_id)
            if prior is not None:
                prior.closed.set()
                self._sessions_by_id.pop(prior.session_id, None)
            self._sessions_by_device[device_id] = sess
            self._sessions_by_id[session_id] = sess
        return sess

    def close_session(self, session_id: str, *, reason: str = "closed") -> None:
        """Close the local session by id.

        Queued-but-undelivered commands stay on the device's Redis list and
        are picked up by the next session; an in-flight command whose socket
        drops falls back to the tool's own drain deadline.
        """
        with self._lock:
            sess = self._sessions_by_id.pop(session_id, None)
            if sess is None:
                return
            sess.closed.set()
            if self._sessions_by_device.get(sess.device_id) is sess:
                self._sessions_by_device.pop(sess.device_id, None)
        logger.debug("device session closed: %s (%s)", session_id, reason)

    def get_session(self, session_id: str) -> Optional[SessionState]:
        with self._lock:
            return self._sessions_by_id.get(session_id)

    def next_command(
        self, session: SessionState, timeout: float = 1.0
    ) -> Optional[Dict[str, Any]]:
        """Block up to ``timeout`` for the next queued command envelope.

        Returns the decoded envelope, or ``None`` on timeout so the SSE
        handler can emit a keepalive and re-check the session lifecycle.
        """
        redis = get_redis_instance()
        if redis is None:
            time.sleep(timeout)
            return None
        try:
            popped = redis.blpop(_cmd_key(session.device_id), timeout=timeout)
        except Exception:
            logger.exception("blpop failed for %s", session.device_id)
            time.sleep(timeout)
            return None
        if not popped:
            return None
        _key, raw = popped
        try:
            envelope = json.loads(_as_str(raw))
        except (TypeError, ValueError):
            logger.warning("dropping malformed command envelope")
            return None
        if not isinstance(envelope, dict):
            return None
        # Drop an envelope whose invocation was already reaped (timed out /
        # cleaned up) after it was queued, so a command the user already saw
        # fail can't still run on the device. Best-effort — it narrows but does
        # not fully close the BLPOP-vs-cleanup window (see cleanup_invocation).
        inv_id = envelope.get("invocation_id")
        if inv_id and self.get_invocation(inv_id) is None:
            logger.debug("dropping reaped invocation %s", inv_id)
            return None
        return envelope

    # ------------------------------------------------------------------
    # Polling / tickets
    # ------------------------------------------------------------------
    def claim_ticket(self, device_id: str, ttl_seconds: float) -> Optional[str]:
        """Return an SSE upgrade ticket iff the device has queued work.

        Reuses an unexpired ticket so repeated polls don't churn it; the
        ticket's Redis TTL doubles as the advertised ``expires_in`` window.
        """
        redis = get_redis_instance()
        if redis is None:
            return None
        try:
            if redis.llen(_cmd_key(device_id)) <= 0:
                return None
            existing = redis.get(_ticket_key(device_id))
            if existing:
                return _as_str(existing)
            ticket = f"st_{uuid.uuid4().hex}"
            redis.set(_ticket_key(device_id), ticket, ex=int(ttl_seconds))
            return ticket
        except Exception:
            logger.exception("claim_ticket failed for %s", device_id)
            return None

    def validate_ticket(self, device_id: str, session_id: str) -> bool:
        """True iff ``session_id`` is the unexpired ticket issued to the device.

        An absent, mismatched, or expired ticket is rejected; expiry is
        enforced by Redis's TTL (a ``GET`` of an expired key returns nil).
        """
        if not session_id:
            return False
        redis = get_redis_instance()
        if redis is None:
            return False
        try:
            issued = redis.get(_ticket_key(device_id))
        except Exception:
            logger.exception("validate_ticket failed for %s", device_id)
            return False
        return issued is not None and _as_str(issued) == session_id

    # ------------------------------------------------------------------
    # Dispatch (server-issued, any process)
    # ------------------------------------------------------------------
    def dispatch_invocation(
        self,
        device_id: str,
        user_id: str,
        envelope: Dict[str, Any],
    ) -> Invocation:
        """Queue an invocation for ``device_id`` and record its metadata.

        Writes the metadata hash, then RPUSHes the envelope onto the
        device's command list. A live SSE session draining the list picks it
        up immediately; otherwise it waits for the next poll-issued ticket.
        """
        invocation_id = envelope["invocation_id"]
        inv = Invocation(invocation_id=invocation_id, device_id=device_id)
        redis = get_redis_instance()
        if redis is None:
            inv.error = "device broker unavailable"
            inv.completed = True
            return inv
        envelope_json = json.dumps(envelope)
        try:
            redis.hset(
                _inv_key(invocation_id),
                mapping={
                    "device_id": device_id,
                    "user_id": user_id,
                    "envelope": envelope_json,
                    "completed": "0",
                    "stdout_bytes": "0",
                    "stderr_bytes": "0",
                },
            )
            redis.expire(_inv_key(invocation_id), self._inv_ttl())
            redis.rpush(_cmd_key(device_id), envelope_json)
            redis.expire(_cmd_key(device_id), self._cmd_ttl())
        except Exception:
            logger.exception("dispatch_invocation failed for %s", invocation_id)
            # Don't strand the metadata hash (it holds the plaintext command)
            # if the queue write failed partway through.
            try:
                redis.delete(_inv_key(invocation_id))
            except Exception:
                logger.debug(
                    "cleanup after failed dispatch failed for %s", invocation_id
                )
            inv.error = "device broker dispatch failed"
            inv.completed = True
        return inv

    def get_invocation(self, invocation_id: str) -> Optional[Invocation]:
        """Read an invocation's current metadata snapshot from Redis."""
        redis = get_redis_instance()
        if redis is None:
            return None
        try:
            raw = redis.hgetall(_inv_key(invocation_id))
        except Exception:
            logger.exception("get_invocation failed for %s", invocation_id)
            return None
        if not raw:
            return None
        h = {_as_str(k): _as_str(v) for k, v in raw.items()}
        return Invocation(
            invocation_id=invocation_id,
            device_id=h.get("device_id", ""),
            completed=h.get("completed") == "1",
            exit_code=_to_int(h.get("exit_code")),
            duration_ms=_to_int(h.get("duration_ms")),
            error=h.get("error") or None,
            started_at=_to_float(h.get("started_at")),
            finished_at=_to_float(h.get("finished_at")),
            stdout_bytes=_to_int(h.get("stdout_bytes")) or 0,
            stderr_bytes=_to_int(h.get("stderr_bytes")) or 0,
        )

    # ------------------------------------------------------------------
    # Output streaming (web process / CLI side)
    # ------------------------------------------------------------------
    def submit_output_chunk(
        self, invocation_id: str, chunk: Dict[str, Any]
    ) -> bool:
        """Forward one CLI output chunk to the dispatching process's drain.

        Updates the metadata hash (byte counts; result fields on the closing
        ``control`` chunk) and XADDs the chunk to the invocation's output
        stream. Returns ``False`` for an unknown invocation.
        """
        redis = get_redis_instance()
        if redis is None:
            return False
        key = _inv_key(invocation_id)
        try:
            device_id = redis.hget(key, "device_id")
        except Exception:
            logger.exception("submit_output_chunk read failed for %s", invocation_id)
            return False
        if device_id is None:
            return False
        device_id = _as_str(device_id)
        now = time.time()
        stream = chunk.get("stream")
        try:
            # Append to the output stream BEFORE flipping completion state, so a
            # reader that observes completed=1 is guaranteed the chunk (and every
            # earlier one) is already on the stream — otherwise drain could see
            # completion and stop before the control chunk lands.
            redis.xadd(
                _out_key(invocation_id),
                {"c": json.dumps(chunk)},
                maxlen=self._out_maxlen(),
                approximate=True,
            )
            redis.expire(_out_key(invocation_id), self._inv_ttl())
            if stream in ("stdout", "stderr"):
                text = chunk.get("chunk")
                if isinstance(text, str):
                    field = "stdout_bytes" if stream == "stdout" else "stderr_bytes"
                    redis.hincrby(key, field, len(text.encode("utf-8")))
            elif stream == "control":
                redis.hsetnx(key, "started_at", repr(now))
                mapping = {"completed": "1", "finished_at": repr(now)}
                if chunk.get("exit_code") is not None:
                    mapping["exit_code"] = str(int(chunk["exit_code"]))
                if chunk.get("duration_ms") is not None:
                    mapping["duration_ms"] = str(int(chunk["duration_ms"]))
                if chunk.get("error"):
                    mapping["error"] = _as_str(chunk["error"])
                redis.hset(key, mapping=mapping)
            redis.expire(key, self._inv_ttl())
        except Exception:
            logger.exception("submit_output_chunk failed for %s", invocation_id)
            return False
        # Keep a co-located SSE session alive while output flows (single-worker
        # web tier); a cross-worker session relies on its own keepalive.
        with self._lock:
            sess = self._sessions_by_device.get(device_id)
        if sess is not None:
            sess.last_activity_at = now
        return True

    def submit_ack(
        self,
        invocation_id: str,
        decision: str,
        reason: Optional[str] = None,
    ) -> bool:
        """Record the CLI's accept/deny decision for an invocation."""
        redis = get_redis_instance()
        if redis is None:
            return False
        key = _inv_key(invocation_id)
        try:
            if not redis.exists(key):
                return False
            now = time.time()
            mapping = {"decision": decision}
            if reason:
                mapping["decision_reason"] = reason
            redis.hset(key, mapping=mapping)
            redis.hsetnx(key, "started_at", repr(now))
            redis.expire(key, self._inv_ttl())
            if decision == "denied":
                # XADD the synthetic control chunk BEFORE marking completed, so a
                # racing drain that observes completed=1 always finds the chunk
                # and reports "denied" rather than a false timeout.
                redis.xadd(
                    _out_key(invocation_id),
                    {"c": json.dumps(
                        {"stream": "control", "exit_code": None, "error": "denied"}
                    )},
                    maxlen=self._out_maxlen(),
                    approximate=True,
                )
                redis.expire(_out_key(invocation_id), self._inv_ttl())
                redis.hset(
                    key,
                    mapping={
                        "completed": "1",
                        "error": "denied",
                        "finished_at": repr(now),
                    },
                )
        except Exception:
            logger.exception("submit_ack failed for %s", invocation_id)
            return False
        return True

    def drain_output(
        self,
        invocation_id: str,
        timeout: float = 0.5,
        deadline: Optional[float] = None,
    ) -> Iterator[Dict[str, Any]]:
        """Yield queued output chunks for an invocation; stop on ``control``.

        Reads the invocation's output stream from the start (so chunks
        XADDed before draining began are never missed) and blocks up to
        ``timeout`` per read. ``deadline`` is an absolute ``time.time()``;
        once it passes with no closing ``control`` chunk the generator
        returns so a device that never responds can't loop forever.
        """
        redis = get_redis_instance()
        if redis is None:
            yield {
                "stream": "control",
                "exit_code": None,
                "error": "device broker unavailable",
            }
            return
        last_id = "0-0"
        block_ms = max(1, int(timeout * 1000))
        out_key = _out_key(invocation_id)
        # ``resp`` is parsed as the RESP2 list shape ``[[key, [(id, fields)]]]``;
        # the client must stay protocol=2 (redis-py default — see cache.py).
        tail_flush = False
        while True:
            try:
                # On the final sweep, read non-blocking (block=None) so any
                # entries that landed after the prior empty read are drained
                # before returning — closing the completion-vs-control race.
                resp = redis.xread(
                    {out_key: last_id},
                    count=200,
                    block=None if tail_flush else block_ms,
                )
            except Exception:
                logger.exception("xread failed for %s", invocation_id)
                return
            if not resp:
                if tail_flush:
                    return
                done = self._is_completed(redis, invocation_id)
                timed_out = deadline is not None and time.time() >= deadline
                if done or timed_out:
                    # One final non-blocking sweep from last_id, then stop.
                    tail_flush = True
                continue
            for _stream_key, entries in resp:
                for entry_id, fields in entries:
                    last_id = _as_str(entry_id)
                    raw = fields.get(b"c")
                    if raw is None:
                        raw = fields.get("c")
                    if raw is None:
                        continue
                    try:
                        chunk = json.loads(_as_str(raw))
                    except (TypeError, ValueError):
                        continue
                    if not isinstance(chunk, dict):
                        continue
                    if chunk.get("stream") in ("stdout", "stderr", "control"):
                        yield chunk
                    if chunk.get("stream") == "control":
                        return

    def cleanup_invocation(self, invocation_id: str) -> None:
        """Drop an invocation's Redis state, including any undelivered command.

        Deletes the metadata hash first so a concurrent ``next_command`` that
        just BLPOPped this envelope re-checks ``get_invocation``, sees it gone,
        and drops it instead of delivering a command the user already saw fail;
        the ``LREM`` then clears it for the still-queued (offline-device) case.
        Best-effort: a delivery that wins a tight race with the delete can still
        reach the device once, and that run's output is discarded.
        """
        redis = get_redis_instance()
        if redis is None:
            return
        try:
            raw = redis.hgetall(_inv_key(invocation_id))
            device_id = envelope_json = None
            if raw:
                h = {_as_str(k): _as_str(v) for k, v in raw.items()}
                device_id = h.get("device_id")
                envelope_json = h.get("envelope")
            redis.delete(_inv_key(invocation_id))
            redis.delete(_out_key(invocation_id))
            if device_id and envelope_json:
                redis.lrem(_cmd_key(device_id), 0, envelope_json)
        except Exception:
            logger.exception("cleanup_invocation failed for %s", invocation_id)

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------
    @staticmethod
    def _is_completed(redis, invocation_id: str) -> bool:
        try:
            return _as_str(redis.hget(_inv_key(invocation_id), "completed")) == "1"
        except Exception:
            return False

    @staticmethod
    def _inv_ttl() -> int:
        return int(getattr(settings, "REMOTE_DEVICE_INVOCATION_TTL_SECONDS", 900))

    @staticmethod
    def _cmd_ttl() -> int:
        return int(getattr(settings, "REMOTE_DEVICE_CMD_QUEUE_TTL_SECONDS", 900))

    @staticmethod
    def _out_maxlen() -> int:
        return int(getattr(settings, "REMOTE_DEVICE_OUTPUT_STREAM_MAXLEN", 10_000))


def _to_int(value: Optional[str]) -> Optional[int]:
    if value is None or value == "":
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _to_float(value: Optional[str]) -> Optional[float]:
    if value is None or value == "":
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


_broker_instance: Optional[DeviceBroker] = None
_broker_lock = threading.Lock()


def get_broker() -> DeviceBroker:
    """Return the process-wide ``DeviceBroker`` instance."""
    global _broker_instance
    if _broker_instance is None:
        with _broker_lock:
            if _broker_instance is None:
                _broker_instance = DeviceBroker()
    return _broker_instance