Fix remote-device tool timing out on scheduled runs (Redis-backed broker) (#2511)

* fix: route remote-device tool through Redis so scheduled runs reach the device

The remote-device tool worked interactively but timed out on every scheduled
run. DeviceBroker was an in-process, in-memory singleton, but scheduled runs
execute in the Celery worker — a different process from the gunicorn web tier
that holds the device's SSE session — so a worker-side dispatch never reached
the device and the tool always hit its deadline.

Make the broker Redis-backed so every hop crosses the process boundary:
- queued commands       -> Redis list   dev:cmd:{device_id}
- output chunks         -> Redis stream dev:out:{invocation_id}
- invocation metadata   -> Redis hash   dev:inv:{invocation_id}
- SSE upgrade tickets    -> Redis key    dev🎫{device_id}
Per-connection SSE session state stays in the web process. Reuses the existing
get_redis_instance()/CACHE_REDIS_URL; no new infrastructure. Also makes the web
tier safe to scale past one worker.

Concurrency hardening (from adversarial review + real-Redis e2e):
- XADD the output/control chunk before flipping completed=1, and have
  drain_output do a final non-blocking flush after observing completion, so a
  reader can't see completion and stop before the control chunk lands (this had
  reintroduced the false "device did not respond (timed out)" under a race).
- _collect_result builds the result from drained chunks, checks the deadline
  only after capturing a chunk, and falls back to the authoritative snapshot
  (before cleanup) when no control chunk was observed.
- Audit outcome is written from locally-known fields so it survives the worker
  racing to delete the invocation; a denied command now records a terminal
  "denied" outcome instead of staying "dispatched".
- cmd-queue TTL raised to 900s (>= max drain deadline); dispatch-failure and
  reaped-invocation cleanup; UTF-8 byte counts.

Tests: new tests/devices/{conftest (FakeRedis double), test_broker_cross_process,
test_broker_race, test_submit_output_audit}; drain/cleanup/ticket tests rewritten
for the Redis contract. The race tests fail against the pre-fix code. ruff clean;
device + tool-executor suites green.

* fix: log instead of silently passing on failed-dispatch cleanup

Addresses the code-quality lint on the best-effort hash delete in
dispatch_invocation's failure path: replace the bare `except: pass` with a
logger.debug carrying the invocation_id. No behavior change — cleanup stays
best-effort and still returns a failed Invocation.

application/agents/tools/remote_device.py

@@ -246,37 +246,72 @@ class RemoteDeviceTool(Tool):
             return False
 
     def _collect_result(self, broker, inv, device: dict, timeout_ms: int) -> Dict[str, Any]:
-        """Drain output from the broker until the control chunk arrives."""
+        """Drain output from the broker until the control chunk arrives.
+
+        Result fields come from the drained chunks, not from ``inv``: the
+        invocation runs and reports back in a different process (the web
+        tier), so the dispatching process never sees ``inv`` mutated.
+        """
+        # Dispatch already failed (e.g. broker/Redis unavailable): report it.
+        if inv.completed and inv.error:
+            return {
+                "exit_code": None,
+                "stdout": "",
+                "stderr": "",
+                "duration_ms": None,
+                "device_name": device.get("name"),
+                "error": inv.error,
+            }
+
         deadline = time.time() + (timeout_ms / 1000.0) + 5.0
         stdout = []
         stderr = []
+        exit_code = None
+        duration_ms = None
+        error = None
+        saw_control = False
         try:
             for chunk in broker.drain_output(
                 inv.invocation_id, timeout=1.0, deadline=deadline
             ):
-                if time.time() > deadline:
-                    break
                 stream = chunk.get("stream")
                 if stream == "stdout":
                     stdout.append(chunk.get("chunk", ""))
                 elif stream == "stderr":
                     stderr.append(chunk.get("chunk", ""))
                 elif stream == "control":
-                    # control chunks include exit_code; drain loop will stop next iter
-                    pass
+                    saw_control = True
+                    exit_code = chunk.get("exit_code")
+                    duration_ms = chunk.get("duration_ms")
+                    error = chunk.get("error") or error
+                # Stop once past the deadline — but only AFTER capturing a chunk
+                # the drain already yielded, so a near-deadline control chunk
+                # isn't dropped and misreported as a timeout.
+                if time.time() > deadline:
+                    break
+            # No control chunk observed: consult the authoritative completion
+            # state (before cleanup deletes it) so a late or dropped control
+            # chunk isn't misreported as "device did not respond".
+            if not saw_control:
+                final = broker.get_invocation(inv.invocation_id)
+                if final is not None and final.completed:
+                    saw_control = True
+                    exit_code = final.exit_code
+                    duration_ms = final.duration_ms
+                    error = final.error or error
         finally:
             broker.cleanup_invocation(inv.invocation_id)
 
-        # Deadline hit with no control chunk: the device never connected or
+        # Deadline hit with no completion at all: the device never connected or
         # never finished. Surface a clear timeout instead of empty success.
-        if not inv.completed.is_set() and inv.exit_code is None and not inv.error:
-            inv.error = "device did not respond (timed out)"
+        if not saw_control and exit_code is None and not error:
+            error = "device did not respond (timed out)"
 
         return {
-            "exit_code": inv.exit_code,
-            "stdout": "".join(stdout) if stdout else "".join(inv.stdout_parts),
-            "stderr": "".join(stderr) if stderr else "".join(inv.stderr_parts),
-            "duration_ms": inv.duration_ms,
+            "exit_code": exit_code,
+            "stdout": "".join(stdout),
+            "stderr": "".join(stderr),
+            "duration_ms": duration_ms,
             "device_name": device.get("name"),
-            "error": inv.error,
+            "error": error,
         }

application/api/devices/session.py

@@ -7,7 +7,6 @@ import io
 import json
 import logging
 import time
-from queue import Empty
 from typing import Iterator
 
 from flask import Response, jsonify, make_response, request, stream_with_context
@@ -118,16 +117,15 @@ def session_events(session_id: str) -> Response:
                     sess.last_event_id += 1
                     broker.close_session(sess.session_id, reason="idle")
                     return
-                try:
-                    inv = sess.invocation_queue.get(timeout=1.0)
-                except Empty:
+                envelope = broker.next_command(sess, timeout=1.0)
+                if envelope is None:
                     if time.time() - last_keepalive >= keepalive_interval:
                         last_keepalive = time.time()
                         yield ": heartbeat\n\n"
                     continue
                 sess.last_event_id += 1
                 sess.last_activity_at = time.time()
-                yield _sse_event("invocation", inv.envelope, sess.last_event_id)
+                yield _sse_event("invocation", envelope, sess.last_event_id)
                 last_keepalive = time.time()
         except GeneratorExit:
             logger.debug("device SSE generator exiting for session %s", sess.session_id)
@@ -164,6 +162,22 @@ def ack_invocation(session_id: str, invocation_id: str) -> Response:
             jsonify({"success": False, "error": "invocation_not_found"}), 404
         )
     broker.submit_ack(invocation_id, decision, reason)
+    if decision == "denied":
+        # A denial is terminal and produces no device output, so submit_output's
+        # audit write is never reached. Record the outcome here from locally
+        # known facts (not re-read Redis state the agent's drain races to clean
+        # up), so the audit row reflects the denial instead of staying
+        # "dispatched". Accepted/auto_approved runs record via submit_output.
+        from datetime import datetime, timezone
+        try:
+            with db_session() as conn:
+                DeviceAuditLogRepository(conn).record_result(
+                    invocation_id,
+                    finished_at=datetime.now(timezone.utc),
+                    error="denied",
+                )
+        except Exception:
+            logger.exception("audit record_result (denied) failed for %s", invocation_id)
     return make_response(jsonify({"success": True}), 200)
 
 
@@ -189,6 +203,7 @@ def submit_output(session_id: str, invocation_id: str) -> Response:
             )
 
     received = 0
+    control_chunk = None
     for line in io.BytesIO(body):
         line = line.strip()
         if not line:
@@ -199,33 +214,33 @@ def submit_output(session_id: str, invocation_id: str) -> Response:
             continue
         if not isinstance(chunk, dict):
             continue
+        if chunk.get("stream") == "control":
+            control_chunk = chunk
         broker.submit_output_chunk(invocation_id, chunk)
         received += 1
 
-    # Persist outcome on the audit row when a control chunk closed the stream.
-    if inv.completed.is_set():
+    # Persist the outcome when the closing control chunk arrived in this POST.
+    # Its fields are captured locally so the audit write survives the draining
+    # (worker) process racing to delete the invocation's Redis state. Byte
+    # totals / started_at live in the hash, read best-effort (the functional
+    # exit_code/error/duration still land even if the hash is already gone).
+    if control_chunk is not None:
+        from datetime import datetime, timezone
+        snap = broker.get_invocation(invocation_id)
         try:
-            from datetime import datetime, timezone
             with db_session() as conn:
                 DeviceAuditLogRepository(conn).record_result(
                     invocation_id,
                     started_at=(
-                        datetime.fromtimestamp(inv.started_at, tz=timezone.utc)
-                        if inv.started_at else None
+                        datetime.fromtimestamp(snap.started_at, tz=timezone.utc)
+                        if snap is not None and snap.started_at else None
                     ),
-                    finished_at=(
-                        datetime.fromtimestamp(inv.finished_at, tz=timezone.utc)
-                        if inv.finished_at else None
-                    ),
-                    exit_code=inv.exit_code,
-                    duration_ms=inv.duration_ms,
-                    stdout_bytes=sum(
-                        len(c) for c in inv.stdout_parts
-                    ),
-                    stderr_bytes=sum(
-                        len(c) for c in inv.stderr_parts
-                    ),
-                    error=inv.error,
+                    finished_at=datetime.now(timezone.utc),
+                    exit_code=_as_opt_int(control_chunk.get("exit_code")),
+                    duration_ms=_as_opt_int(control_chunk.get("duration_ms")),
+                    stdout_bytes=(snap.stdout_bytes if snap is not None else 0),
+                    stderr_bytes=(snap.stderr_bytes if snap is not None else 0),
+                    error=control_chunk.get("error"),
                 )
         except Exception:
             logger.exception("audit record_result failed for %s", invocation_id)
@@ -235,6 +250,16 @@ def submit_output(session_id: str, invocation_id: str) -> Response:
     )
 
 
+def _as_opt_int(value) -> int | None:
+    """Coerce a CLI-supplied JSON value to int for an INTEGER audit column."""
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
 def _sse_event(name: str, payload: dict, event_id: int) -> str:
     return (
         f"event: {name}\n"

application/core/settings.py

@@ -241,6 +241,14 @@ class Settings(BaseSettings):
     REMOTE_DEVICE_SESSION_IDLE_SECONDS: int = 60
     REMOTE_DEVICE_REQUIRE_SIGNATURE: bool = False
     REMOTE_DEVICE_PAIRING_TTL_SECONDS: int = 600
+    # Redis-backed broker tunables (route invocations cross-process so a
+    # scheduled/Celery run reaches the web-held device session). The command
+    # queue TTL must exceed the max command drain deadline (the tool caps
+    # timeout_ms at 600s, drained with a +5s margin = 605s) so a queued command
+    # for a briefly-offline device isn't evicted before its own drain gives up.
+    REMOTE_DEVICE_CMD_QUEUE_TTL_SECONDS: int = 900
+    REMOTE_DEVICE_INVOCATION_TTL_SECONDS: int = 900
+    REMOTE_DEVICE_OUTPUT_STREAM_MAXLEN: int = 10_000
 
     # Scheduler (see scheduler.md).
     SCHEDULE_DISPATCHER_INTERVAL: int = 30

application/devices/broker.py

@@ -1,257 +1,398 @@
-"""In-process broker that routes invocations to active device sessions.
+"""Redis-backed broker that routes invocations to active device sessions.
 
-Single-worker assumption (per spec 13.1). Each device has at most one
-active SSE session; concurrent invocations queue and drain in FIFO order.
+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 calls ``RemoteDeviceTool.execute_action`` → ``broker.dispatch_invocation``.
-2. If a session is active, the envelope goes on the session's queue; the
-   SSE handler reads it and emits to the wire. Otherwise it sits on the
-   pending-invocation queue until ``register_session`` is called by the
-   next ``GET /api/devices/poll`` -> SSE upgrade.
-3. CLI POSTs ack and chunked output back. Output chunks land in the
-   invocation's stdout/stderr queue, which the tool's generator drains
-   to yield streaming events.
-4. Final ``control`` chunk closes the invocation. The tool resolves with
-   the aggregated result.
+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 queue import Empty, Queue
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Iterator, Optional
+
+from application.cache import get_redis_instance
+from application.core.settings import settings
 
 
 logger = logging.getLogger(__name__)
 
 
-# Sentinel object signalling the end of an invocation's output stream.
+# 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 InvocationState:
-    """In-flight invocation: envelope + back-streaming queues."""
+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
-    session_id: str
-    envelope: Dict[str, Any]
-    output_queue: Queue = field(default_factory=Queue)
-    stdout_parts: List[str] = field(default_factory=list)
-    stderr_parts: List[str] = field(default_factory=list)
+    completed: bool = False
     exit_code: Optional[int] = None
     duration_ms: Optional[int] = None
     error: Optional[str] = None
-    decision: Optional[str] = None
-    decision_reason: Optional[str] = None
     started_at: Optional[float] = None
     finished_at: Optional[float] = None
-    completed: threading.Event = field(default_factory=threading.Event)
+    stdout_bytes: int = 0
+    stderr_bytes: int = 0
 
 
 @dataclass
 class SessionState:
-    """Active SSE session for one device."""
+    """Per-connection state for the SSE handler that owns the live socket."""
 
     session_id: str
     device_id: str
     user_id: str
-    invocation_queue: Queue = field(default_factory=Queue)
     last_event_id: int = 0
-    invocations: Dict[str, InvocationState] = field(default_factory=dict)
     last_activity_at: float = field(default_factory=time.time)
     closed: threading.Event = field(default_factory=threading.Event)
 
 
 class DeviceBroker:
-    """In-process device registry; thread-safe."""
+    """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()
-        # device_id -> SessionState (active SSE).
-        self._sessions: Dict[str, SessionState] = {}
-        # session_id -> device_id (reverse map for fast session-by-id lookup).
-        self._session_by_id: Dict[str, str] = {}
-        # device_id -> pending invocations (no active session yet).
-        self._pending: Dict[str, List[InvocationState]] = {}
-        # invocation_id -> InvocationState for ack/output lookup.
-        self._invocations: Dict[str, InvocationState] = {}
-        # device_id -> (session ticket, monotonic expiry) for the next poll.
-        self._tickets: Dict[str, tuple[str, float]] = {}
+        self._sessions_by_id: Dict[str, SessionState] = {}
+        self._sessions_by_device: Dict[str, SessionState] = {}
 
     # ------------------------------------------------------------------
-    # Session lifecycle (CLI side)
+    # Session lifecycle (web process / CLI side)
     # ------------------------------------------------------------------
     def register_session(self, device_id: str, user_id: str) -> SessionState:
-        """Open or refresh the active SSE session for ``device_id``.
+        """Open the SSE session for ``device_id``, adopting its poll ticket.
 
-        Drains any pending invocations into the session's queue so the
-        next SSE read sees them in order. If a previous session is still
-        registered, it is closed and replaced (the CLI reconnected).
+        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:
-            existing = self._sessions.get(device_id)
-            if existing is not None:
-                existing.closed.set()
-                self._session_by_id.pop(existing.session_id, None)
-            issued = self._tickets.pop(device_id, None)
-            session_id = issued[0] if issued else f"st_{uuid.uuid4().hex}"
-            sess = SessionState(
-                session_id=session_id,
-                device_id=device_id,
-                user_id=user_id,
-            )
-            for inv in self._pending.pop(device_id, []):
-                inv.session_id = session_id
-                sess.invocations[inv.invocation_id] = inv
-                sess.invocation_queue.put(inv)
-            self._sessions[device_id] = sess
-            self._session_by_id[session_id] = device_id
-            return sess
+            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 session by id; outstanding invocations get an error."""
+        """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:
-            device_id = self._session_by_id.pop(session_id, None)
-            if device_id is None:
-                return
-            sess = self._sessions.pop(device_id, None)
+            sess = self._sessions_by_id.pop(session_id, None)
             if sess is None:
                 return
             sess.closed.set()
-            for inv in sess.invocations.values():
-                if not inv.completed.is_set():
-                    inv.error = inv.error or f"session_{reason}"
-                    inv.output_queue.put(INVOCATION_DONE)
-                    inv.completed.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:
-            device_id = self._session_by_id.get(session_id)
-            if device_id is None:
-                return None
-            return self._sessions.get(device_id)
+            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
+    # Polling / tickets
     # ------------------------------------------------------------------
     def claim_ticket(self, device_id: str, ttl_seconds: float) -> Optional[str]:
-        """Return a session ticket if the device has work waiting, else None.
+        """Return an SSE upgrade ticket iff the device has queued work.
 
-        A pending invocation queued before the CLI polls allocates a ticket
-        eagerly; the CLI then upgrades to SSE with that ticket within
-        ``ttl_seconds``. The ticket is remembered (with its expiry) so the SSE
-        ``session_id`` can be validated against the one this poll issued.
+        Reuses an unexpired ticket so repeated polls don't churn it; the
+        ticket's Redis TTL doubles as the advertised ``expires_in`` window.
         """
-        now = time.monotonic()
-        with self._lock:
-            if not self._pending.get(device_id):
+        redis = get_redis_instance()
+        if redis is None:
+            return None
+        try:
+            if redis.llen(_cmd_key(device_id)) <= 0:
                 return None
-            existing = self._tickets.get(device_id)
-            if existing and existing[1] > now:
-                return existing[0]
+            existing = redis.get(_ticket_key(device_id))
+            if existing:
+                return _as_str(existing)
             ticket = f"st_{uuid.uuid4().hex}"
-            self._tickets[device_id] = (ticket, now + ttl_seconds)
+            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 ``device_id``.
+        """True iff ``session_id`` is the unexpired ticket issued to the device.
 
-        Gates the SSE upgrade: the CLI must open the events stream with the
-        exact ``session_ticket`` returned by its ``/poll`` (which became the
-        ``session_url``). An absent, mismatched, or expired ticket is rejected.
-        Expired entries are evicted so a stale ticket can't be reused.
+        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
-        now = time.monotonic()
-        with self._lock:
-            issued = self._tickets.get(device_id)
-            if issued is None:
-                return False
-            ticket, expires_at = issued
-            if expires_at <= now:
-                self._tickets.pop(device_id, None)
-                return False
-            return session_id == ticket
+        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)
+    # Dispatch (server-issued, any process)
     # ------------------------------------------------------------------
     def dispatch_invocation(
         self,
         device_id: str,
         user_id: str,
         envelope: Dict[str, Any],
-    ) -> InvocationState:
-        """Enqueue an invocation for ``device_id``.
+    ) -> Invocation:
+        """Queue an invocation for ``device_id`` and record its metadata.
 
-        If there's an active session, drops it onto the session's
-        invocation queue. Otherwise parks it in the pending list and
-        the next poll will see a ticket.
+        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"]
-        with self._lock:
-            sess = self._sessions.get(device_id)
-            session_id = sess.session_id if sess else ""
-            inv = InvocationState(
-                invocation_id=invocation_id,
-                device_id=device_id,
-                session_id=session_id,
-                envelope=envelope,
+        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",
+                },
             )
-            self._invocations[invocation_id] = inv
-            if sess is not None:
-                sess.invocations[invocation_id] = inv
-                sess.invocation_queue.put(inv)
-                sess.last_activity_at = time.time()
-            else:
-                self._pending.setdefault(device_id, []).append(inv)
+            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[InvocationState]:
-        with self._lock:
-            return self._invocations.get(invocation_id)
+    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 (CLI side)
+    # Output streaming (web process / CLI side)
     # ------------------------------------------------------------------
     def submit_output_chunk(
-        self,
-        invocation_id: str,
-        chunk: Dict[str, Any],
+        self, invocation_id: str, chunk: Dict[str, Any]
     ) -> bool:
-        """Forward a chunk from the CLI to the agent stream's generator."""
-        inv = self.get_invocation(invocation_id)
-        if inv is None:
+        """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
-        with self._lock:
-            sess = self._sessions.get(inv.device_id)
-            if sess is not None:
-                sess.last_activity_at = time.time()
-            stream = chunk.get("stream")
-            if stream == "stdout":
-                inv.stdout_parts.append(chunk.get("chunk", ""))
-            elif stream == "stderr":
-                inv.stderr_parts.append(chunk.get("chunk", ""))
+        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":
-                inv.exit_code = chunk.get("exit_code")
-                inv.duration_ms = chunk.get("duration_ms")
-                inv.error = chunk.get("error") or inv.error
-                if inv.started_at is None:
-                    inv.started_at = time.time()
-                inv.finished_at = time.time()
-        inv.output_queue.put(chunk)
-        if chunk.get("stream") == "control":
-            inv.output_queue.put(INVOCATION_DONE)
-            inv.completed.set()
+                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(
@@ -260,17 +401,45 @@ class DeviceBroker:
         decision: str,
         reason: Optional[str] = None,
     ) -> bool:
-        inv = self.get_invocation(invocation_id)
-        if inv is None:
+        """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
-        with self._lock:
-            inv.decision = decision
-            inv.decision_reason = reason
-            inv.started_at = inv.started_at or time.time()
-        if decision == "denied":
-            inv.error = inv.error or "denied"
-            inv.output_queue.put(INVOCATION_DONE)
-            inv.completed.set()
         return True
 
     def drain_output(
@@ -278,56 +447,136 @@ class DeviceBroker:
         invocation_id: str,
         timeout: float = 0.5,
         deadline: Optional[float] = None,
-    ):
-        """Yield queued chunks for an invocation; stops on INVOCATION_DONE.
+    ) -> Iterator[Dict[str, Any]]:
+        """Yield queued output chunks for an invocation; stop on ``control``.
 
-        Used by the tool's generator. Blocks up to ``timeout`` per chunk
-        so the caller can interleave with other work. ``deadline`` is an
-        absolute ``time.time()`` value; once it passes with no output and
-        the invocation never completed, the generator returns so a device
-        that never connects can't loop forever.
+        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.
         """
-        inv = self.get_invocation(invocation_id)
-        if inv is None:
+        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:
-                item = inv.output_queue.get(timeout=timeout)
-            except Empty:
-                if inv.completed.is_set():
-                    return
-                if deadline is not None and time.time() >= deadline:
-                    return
-                continue
-            if item is INVOCATION_DONE:
+                # 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
-            yield item
+            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 from all registries.
+        """Drop an invocation's Redis state, including any undelivered command.
 
-        Removes it from ``_invocations``, the active session's
-        ``invocations`` map, and — if it was still parked for an offline
-        device — ``_pending`` too. Dropping it from the active session keeps
-        a long-lived SSE session from accumulating completed invocations
-        (unbounded memory). Without the pending removal a timed-out
-        invocation would be drained and executed by a later
-        ``register_session`` — running a command the user already saw fail.
+        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.
         """
-        with self._lock:
-            inv = self._invocations.pop(invocation_id, None)
-            if inv is None:
-                return
-            sess = self._sessions.get(inv.device_id)
-            if sess is not None:
-                sess.invocations.pop(invocation_id, None)
-            pending = self._pending.get(inv.device_id)
-            if pending:
-                remaining = [p for p in pending if p.invocation_id != invocation_id]
-                if remaining:
-                    self._pending[inv.device_id] = remaining
-                else:
-                    self._pending.pop(inv.device_id, None)
+        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

tests/devices/conftest.py

@@ -0,0 +1,201 @@
+"""Shared fixtures for device-broker tests.
+
+The repo has no ``fakeredis`` dependency, so ``FakeRedis`` is a tiny
+in-memory double covering only the commands ``DeviceBroker`` issues. It is
+deliberately *not* faithful in several ways, so tests must not lean on them:
+- blocking ops (``blpop`` / ``xread``) don't truly block — they return
+  immediately (or after a tiny sleep) so tests stay fast;
+- TTL is not modeled — ``set(ex=)`` and ``expire`` are no-ops, so expiry must
+  be simulated by an explicit ``delete``;
+- ``xadd`` trims to an EXACT ``maxlen`` and ignores ``approximate``, whereas
+  real Redis with ``approximate=True`` only trims past a slack.
+One ``FakeRedis`` shared by two ``DeviceBroker`` instances models two
+processes (e.g. a Celery worker and the web tier) talking through one Redis.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+
+import pytest
+
+from application.devices.broker import DeviceBroker
+
+
+def _b(value) -> bytes:
+    if isinstance(value, bytes):
+        return value
+    return str(value).encode("utf-8")
+
+
+def _seq_of(entry_id) -> int:
+    text = entry_id.decode() if isinstance(entry_id, bytes) else str(entry_id)
+    try:
+        return int(text.split("-", 1)[0])
+    except (ValueError, IndexError):
+        return 0
+
+
+class FakeRedis:
+    """In-memory stand-in for the redis-py commands the broker uses."""
+
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self.kv: dict = {}
+        self.lists: dict = {}
+        self.hashes: dict = {}
+        self.streams: dict = {}
+        self._seq = 0
+
+    # -- strings ------------------------------------------------------
+    def get(self, key):
+        with self._lock:
+            return self.kv.get(key)
+
+    def set(self, key, value, ex=None, nx=False):
+        with self._lock:
+            if nx and key in self.kv:
+                return None
+            self.kv[key] = _b(value)
+            return True
+
+    def delete(self, *keys):
+        removed = 0
+        with self._lock:
+            for key in keys:
+                for store in (self.kv, self.lists, self.hashes, self.streams):
+                    if key in store:
+                        del store[key]
+                        removed += 1
+        return removed
+
+    def exists(self, key):
+        with self._lock:
+            present = (
+                key in self.kv
+                or key in self.hashes
+                or key in self.lists
+                or key in self.streams
+            )
+            return 1 if present else 0
+
+    def expire(self, key, ttl):  # TTL is not simulated.
+        return True
+
+    # -- lists --------------------------------------------------------
+    def llen(self, key):
+        with self._lock:
+            return len(self.lists.get(key, []))
+
+    def rpush(self, key, *values):
+        with self._lock:
+            lst = self.lists.setdefault(key, [])
+            lst.extend(_b(v) for v in values)
+            return len(lst)
+
+    def blpop(self, key, timeout=0):
+        with self._lock:
+            lst = self.lists.get(key)
+            if lst:
+                value = lst.pop(0)
+                if not lst:
+                    del self.lists[key]
+                return (_b(key), value)
+        return None
+
+    def lrem(self, key, count, value):
+        with self._lock:
+            lst = self.lists.get(key)
+            if not lst:
+                return 0
+            target = _b(value)
+            kept = [item for item in lst if item != target]
+            removed = len(lst) - len(kept)
+            if kept:
+                self.lists[key] = kept
+            else:
+                self.lists.pop(key, None)
+            return removed
+
+    # -- hashes -------------------------------------------------------
+    def hset(self, key, field=None, value=None, mapping=None):
+        with self._lock:
+            h = self.hashes.setdefault(key, {})
+            count = 0
+            if mapping:
+                for f, v in mapping.items():
+                    h[f] = _b(v)
+                    count += 1
+            if field is not None:
+                h[field] = _b(value)
+                count += 1
+            return count
+
+    def hsetnx(self, key, field, value):
+        with self._lock:
+            h = self.hashes.setdefault(key, {})
+            if field in h:
+                return 0
+            h[field] = _b(value)
+            return 1
+
+    def hget(self, key, field):
+        with self._lock:
+            return self.hashes.get(key, {}).get(field)
+
+    def hgetall(self, key):
+        with self._lock:
+            return {_b(f): v for f, v in self.hashes.get(key, {}).items()}
+
+    def hincrby(self, key, field, amount):
+        with self._lock:
+            h = self.hashes.setdefault(key, {})
+            current = int(h.get(field, b"0"))
+            current += int(amount)
+            h[field] = _b(str(current))
+            return current
+
+    # -- streams ------------------------------------------------------
+    def xadd(self, key, fields, maxlen=None, approximate=True):
+        with self._lock:
+            self._seq += 1
+            entry_id = f"{self._seq}-0"
+            entries = self.streams.setdefault(key, [])
+            entries.append((entry_id, {_b(f): _b(v) for f, v in fields.items()}))
+            if maxlen is not None and len(entries) > maxlen:
+                del entries[: len(entries) - maxlen]
+            return _b(entry_id)
+
+    def xread(self, streams, count=None, block=None):
+        if block:
+            time.sleep(min(block / 1000.0, 0.02))
+        out = []
+        with self._lock:
+            for key, last_id in streams.items():
+                entries = self.streams.get(key, [])
+                last_seq = _seq_of(last_id)
+                fresh = [
+                    (eid, fields)
+                    for (eid, fields) in entries
+                    if _seq_of(eid) > last_seq
+                ]
+                if count:
+                    fresh = fresh[:count]
+                if fresh:
+                    out.append([_b(key), [(_b(eid), f) for eid, f in fresh]])
+        return out or None
+
+
+@pytest.fixture
+def fake_redis() -> FakeRedis:
+    return FakeRedis()
+
+
+@pytest.fixture
+def broker_env(monkeypatch, fake_redis):
+    """A ``DeviceBroker`` wired to a fresh ``FakeRedis``; returns (broker, fake)."""
+    monkeypatch.setattr(
+        "application.devices.broker.get_redis_instance", lambda: fake_redis
+    )
+    return DeviceBroker(), fake_redis

tests/devices/test_broker_cleanup.py

@@ -1,77 +1,59 @@
-"""Tests for ``DeviceBroker.cleanup_invocation`` pending removal."""
+"""Tests for ``DeviceBroker.cleanup_invocation`` queued-command removal."""
 
 from __future__ import annotations
 
-from queue import Empty
 
-from application.devices.broker import DeviceBroker
+def _dispatch_offline(broker, invocation_id="inv_stale", device_id="dev_x"):
+    # No session draining -> the envelope stays queued on the device's list.
+    envelope = {"invocation_id": invocation_id, "action": "run_command"}
+    return broker.dispatch_invocation(device_id, "user_x", envelope)
 
 
-def _dispatch_offline(broker: DeviceBroker):
-    # No session registered -> the invocation lands in _pending.
-    envelope = {"invocation_id": "inv_stale", "action": "run_command"}
-    return broker.dispatch_invocation("dev_x", "user_x", envelope)
-
-
-def test_cleanup_removes_pending_so_register_session_doesnt_replay():
-    # A timed-out invocation cleaned up while still pending must NOT be
-    # drained/executed by a later session registration.
-    broker = DeviceBroker()
+def test_cleanup_removes_queued_command_so_it_doesnt_replay(broker_env):
+    # A timed-out invocation cleaned up while still queued must NOT later be
+    # delivered to (and run by) a freshly connected session.
+    broker, fake = broker_env
     inv = _dispatch_offline(broker)
-    assert broker._pending.get("dev_x") == [inv]
+    assert fake.llen("dev:cmd:dev_x") == 1
 
     broker.cleanup_invocation(inv.invocation_id)
-    # Pending key gone entirely once the list empties.
-    assert "dev_x" not in broker._pending
+
+    assert fake.llen("dev:cmd:dev_x") == 0
     assert broker.get_invocation(inv.invocation_id) is None
 
+    # A session that connects now finds nothing queued.
     sess = broker.register_session("dev_x", "user_x")
-    # The stale invocation must not have been re-queued.
-    assert inv.invocation_id not in sess.invocations
-    try:
-        sess.invocation_queue.get_nowait()
-        raise AssertionError("stale invocation was dispatched after cleanup")
-    except Empty:
-        pass
+    assert broker.next_command(sess, timeout=0.05) is None
 
 
-def test_cleanup_removes_from_active_session():
-    # A completed invocation must not linger in the active session's map,
-    # or a long-lived SSE session accumulates them unboundedly.
-    broker = DeviceBroker()
-    sess = broker.register_session("dev_live", "user_live")
-    inv = broker.dispatch_invocation(
-        "dev_live", "user_live",
-        {"invocation_id": "inv_live", "action": "run_command"},
-    )
-    # Dispatched into the active session.
-    assert "inv_live" in sess.invocations
-    assert broker.get_invocation("inv_live") is inv
-
-    # Drive it to completion via a control chunk, then clean up.
+def test_cleanup_deletes_invocation_and_output(broker_env):
+    # The metadata hash and output stream are removed on cleanup.
+    broker, fake = broker_env
+    _dispatch_offline(broker, invocation_id="inv_live", device_id="dev_live")
     broker.submit_output_chunk(
         "inv_live", {"stream": "control", "exit_code": 0, "duration_ms": 1}
     )
+    assert fake.exists("dev:inv:inv_live") == 1
+    assert fake.exists("dev:out:inv_live") == 1
+
     broker.cleanup_invocation("inv_live")
 
-    assert "inv_live" not in sess.invocations
+    assert fake.exists("dev:inv:inv_live") == 0
+    assert fake.exists("dev:out:inv_live") == 0
     assert broker.get_invocation("inv_live") is None
 
 
-def test_cleanup_keeps_other_pending_invocations():
-    # Cleaning one invocation must leave a sibling pending invocation intact.
-    broker = DeviceBroker()
-    inv1 = broker.dispatch_invocation(
-        "dev_y", "user_y", {"invocation_id": "inv_1", "action": "run_command"}
-    )
-    inv2 = broker.dispatch_invocation(
-        "dev_y", "user_y", {"invocation_id": "inv_2", "action": "run_command"}
-    )
-    assert broker._pending.get("dev_y") == [inv1, inv2]
+def test_cleanup_keeps_other_queued_commands(broker_env):
+    # Cleaning one invocation must leave a sibling queued command intact.
+    broker, fake = broker_env
+    inv1 = _dispatch_offline(broker, invocation_id="inv_1", device_id="dev_y")
+    _dispatch_offline(broker, invocation_id="inv_2", device_id="dev_y")
+    assert fake.llen("dev:cmd:dev_y") == 2
 
     broker.cleanup_invocation(inv1.invocation_id)
-    assert broker._pending.get("dev_y") == [inv2]
+    assert fake.llen("dev:cmd:dev_y") == 1
 
     sess = broker.register_session("dev_y", "user_y")
-    assert "inv_2" in sess.invocations
-    assert "inv_1" not in sess.invocations
+    envelope = broker.next_command(sess, timeout=0.05)
+    assert envelope is not None
+    assert envelope["invocation_id"] == "inv_2"

tests/devices/test_broker_cross_process.py

@@ -0,0 +1,83 @@
+"""Regression: a command dispatched in one process reaches a session in another.
+
+This reproduces the scheduled-run failure. The agent runs inside a Celery
+worker (``worker_broker``) while the device's SSE session lives in the web
+process (``web_broker``). With the old in-process broker these two never
+shared state, so every scheduled invocation timed out. Backed by one Redis
+(here a shared ``FakeRedis``), dispatch and drain cross the process line.
+"""
+
+from __future__ import annotations
+
+import time
+
+from application.devices.broker import DeviceBroker
+
+
+def test_dispatch_in_worker_reaches_session_in_web(monkeypatch, fake_redis):
+    monkeypatch.setattr(
+        "application.devices.broker.get_redis_instance", lambda: fake_redis
+    )
+    worker_broker = DeviceBroker()  # e.g. Celery scheduled run
+    web_broker = DeviceBroker()  # e.g. gunicorn web tier holding the SSE socket
+
+    # 1) Agent (worker process) dispatches a command.
+    envelope = {
+        "invocation_id": "inv_xproc",
+        "action": "run_command",
+        "params": {"command": "echo hi"},
+    }
+    worker_broker.dispatch_invocation("dev_1", "user_1", envelope)
+
+    # 2) Device polls/upgrades on the web process and receives the command.
+    assert web_broker.claim_ticket("dev_1", 30) is not None
+    sess = web_broker.register_session("dev_1", "user_1")
+    delivered = web_broker.next_command(sess, timeout=0.1)
+    assert delivered is not None
+    assert delivered["invocation_id"] == "inv_xproc"
+    assert delivered["params"]["command"] == "echo hi"
+
+    # 3) Device streams output back through the web process.
+    web_broker.submit_output_chunk("inv_xproc", {"stream": "stdout", "chunk": "hi\n"})
+    web_broker.submit_output_chunk(
+        "inv_xproc",
+        {"stream": "control", "exit_code": 0, "duration_ms": 7},
+    )
+
+    # 4) The worker's drain (what the tool does) sees the full result.
+    deadline = time.time() + 2.0
+    chunks = list(
+        worker_broker.drain_output("inv_xproc", timeout=0.05, deadline=deadline)
+    )
+    stdout = "".join(c.get("chunk", "") for c in chunks if c.get("stream") == "stdout")
+    control = [c for c in chunks if c.get("stream") == "control"]
+    assert stdout == "hi\n"
+    assert control and control[0]["exit_code"] == 0
+
+    # And the metadata snapshot reflects completion across instances.
+    final = worker_broker.get_invocation("inv_xproc")
+    assert final is not None
+    assert final.completed is True
+    assert final.exit_code == 0
+    assert final.stdout_bytes == len("hi\n")
+
+
+def test_denied_ack_unblocks_drain(monkeypatch, fake_redis):
+    # A denial on the web side must promptly stop a worker-side drain.
+    monkeypatch.setattr(
+        "application.devices.broker.get_redis_instance", lambda: fake_redis
+    )
+    worker_broker = DeviceBroker()
+    web_broker = DeviceBroker()
+
+    worker_broker.dispatch_invocation(
+        "dev_2", "user_2", {"invocation_id": "inv_deny", "action": "run_command"}
+    )
+    web_broker.submit_ack("inv_deny", "denied", reason="user rejected")
+
+    deadline = time.time() + 2.0
+    chunks = list(
+        worker_broker.drain_output("inv_deny", timeout=0.05, deadline=deadline)
+    )
+    assert chunks and chunks[-1]["stream"] == "control"
+    assert chunks[-1]["error"] == "denied"

tests/devices/test_broker_drain.py

@@ -1,21 +1,19 @@
-"""Tests for ``DeviceBroker.drain_output`` deadline handling."""
+"""Tests for ``DeviceBroker.drain_output`` deadline + control handling."""
 
 from __future__ import annotations
 
 import time
 
-from application.devices.broker import DeviceBroker, INVOCATION_DONE
 
-
-def _dispatch(broker: DeviceBroker):
+def _dispatch(broker):
     envelope = {"invocation_id": "inv_test", "action": "run_command"}
     return broker.dispatch_invocation("dev_x", "user_x", envelope)
 
 
-def test_drain_output_returns_after_deadline_with_no_output():
-    # Device never connects / posts output and ``completed`` is never set:
-    # the generator must stop once the deadline passes instead of looping.
-    broker = DeviceBroker()
+def test_drain_output_returns_after_deadline_with_no_output(broker_env):
+    # Device never connects / posts output: the generator must stop once the
+    # deadline passes instead of looping forever.
+    broker, _fake = broker_env
     inv = _dispatch(broker)
     deadline = time.time() + 0.2
     start = time.time()
@@ -24,28 +22,45 @@ def test_drain_output_returns_after_deadline_with_no_output():
     )
     elapsed = time.time() - start
     assert chunks == []
-    assert not inv.completed.is_set()
     assert elapsed < 2.0
 
 
-def test_drain_output_yields_then_stops_on_done():
-    # Sanity: a normal stream still drains and stops on INVOCATION_DONE,
-    # with a generous deadline that should not fire.
-    broker = DeviceBroker()
+def test_drain_output_yields_then_stops_on_control(broker_env):
+    # A normal stream drains stdout then stops on the closing control chunk.
+    broker, _fake = broker_env
     inv = _dispatch(broker)
-    inv.output_queue.put({"stream": "stdout", "chunk": "hello"})
-    inv.output_queue.put(INVOCATION_DONE)
+    broker.submit_output_chunk(inv.invocation_id, {"stream": "stdout", "chunk": "hello"})
+    broker.submit_output_chunk(
+        inv.invocation_id, {"stream": "control", "exit_code": 0, "duration_ms": 1}
+    )
     deadline = time.time() + 5.0
     chunks = list(
         broker.drain_output(inv.invocation_id, timeout=0.05, deadline=deadline)
     )
-    assert chunks == [{"stream": "stdout", "chunk": "hello"}]
+    assert chunks[0] == {"stream": "stdout", "chunk": "hello"}
+    assert chunks[-1]["stream"] == "control"
+    assert chunks[-1]["exit_code"] == 0
 
 
-def test_drain_output_no_deadline_stops_when_completed():
-    # Without a deadline, the completed event still terminates the loop.
-    broker = DeviceBroker()
+def test_drain_output_stops_when_completed_without_control(broker_env):
+    # If the invocation is marked completed but no control chunk is on the
+    # stream, the completed flag still terminates the drain loop.
+    broker, fake = broker_env
     inv = _dispatch(broker)
-    inv.completed.set()
+    fake.hset(f"dev:inv:{inv.invocation_id}", mapping={"completed": "1"})
     chunks = list(broker.drain_output(inv.invocation_id, timeout=0.05))
     assert chunks == []
+
+
+def test_drain_output_reports_error_when_redis_unavailable(monkeypatch):
+    # No Redis: the tool must get a clear control/error chunk, not hang.
+    from application.devices.broker import DeviceBroker
+
+    monkeypatch.setattr(
+        "application.devices.broker.get_redis_instance", lambda: None
+    )
+    broker = DeviceBroker()
+    chunks = list(broker.drain_output("inv_missing", timeout=0.05))
+    assert len(chunks) == 1
+    assert chunks[0]["stream"] == "control"
+    assert "unavailable" in chunks[0]["error"]

tests/devices/test_broker_race.py

@@ -0,0 +1,194 @@
+"""Race-condition regressions for the Redis-backed broker.
+
+These cover the concurrency paths the happy-path suite cannot reach: the
+completion-flag-vs-control-chunk ordering, drain's final flush, the tool's
+near-deadline capture + authoritative fallback, and the cleanup/dispatch
+edge paths. Each fails against the pre-fix code and passes after.
+"""
+
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+
+from application.agents.tools.remote_device import RemoteDeviceTool, _MAX_TIMEOUT_MS
+from application.core.settings import settings
+from application.devices.broker import DeviceBroker, Invocation
+
+from .conftest import FakeRedis
+
+
+class RacyFakeRedis(FakeRedis):
+    """FakeRedis that fires a one-shot hook on the FIRST ``xread`` and returns
+    empty for that call — modelling a drain observing completion before the
+    chunks are visible to its read cursor.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.on_first_empty = None
+        self._fired = False
+
+    def xread(self, streams, count=None, block=None):
+        if not self._fired and self.on_first_empty is not None:
+            self._fired = True
+            self.on_first_empty()  # web side posts output + control here
+            return None  # this read still sees nothing
+        return super().xread(streams, count=count, block=block)
+
+
+class _StubBroker:
+    """Minimal broker for exercising RemoteDeviceTool._collect_result."""
+
+    def __init__(self, chunks, snapshot=None):
+        self._chunks = chunks
+        self._snapshot = snapshot
+        self.cleaned = False
+
+    def drain_output(self, invocation_id, timeout=1.0, deadline=None):
+        for chunk in self._chunks:
+            yield chunk
+
+    def get_invocation(self, invocation_id):
+        return self._snapshot
+
+    def cleanup_invocation(self, invocation_id):
+        self.cleaned = True
+
+
+def _tool():
+    # Bypass __init__ (which loads the device from the DB); _collect_result
+    # uses only its arguments, no instance state.
+    return RemoteDeviceTool.__new__(RemoteDeviceTool)
+
+
+# ---------------------------------------------------------------------------
+# drain_output final-flush (HIGH): completion observed before control read
+# ---------------------------------------------------------------------------
+def test_drain_flushes_chunks_posted_after_first_empty_read(monkeypatch):
+    fake = RacyFakeRedis()
+    monkeypatch.setattr(
+        "application.devices.broker.get_redis_instance", lambda: fake
+    )
+    worker = DeviceBroker()  # Celery side: dispatch + drain
+    web = DeviceBroker()  # web side: posts output
+    worker.dispatch_invocation(
+        "d_race", "u_race",
+        {"invocation_id": "inv_race", "action": "run_command"},
+    )
+
+    def hook():
+        web.submit_output_chunk("inv_race", {"stream": "stdout", "chunk": "out"})
+        web.submit_output_chunk(
+            "inv_race", {"stream": "control", "exit_code": 0, "duration_ms": 1}
+        )
+
+    fake.on_first_empty = hook
+    chunks = list(
+        worker.drain_output("inv_race", timeout=0.05, deadline=time.time() + 5)
+    )
+    stdout = "".join(c.get("chunk", "") for c in chunks if c.get("stream") == "stdout")
+    control = [c for c in chunks if c.get("stream") == "control"]
+    # Pre-fix: drain returned [] on the _is_completed early-return, dropping both.
+    assert stdout == "out"
+    assert control and control[0]["exit_code"] == 0
+
+
+# ---------------------------------------------------------------------------
+# _collect_result authoritative fallback + near-deadline capture (MEDIUM)
+# ---------------------------------------------------------------------------
+def test_collect_result_uses_completion_snapshot_when_no_control():
+    inv = SimpleNamespace(invocation_id="i", device_id="d", completed=False, error=None)
+    snap = Invocation("i", "d", completed=True, exit_code=0, duration_ms=9)
+    broker = _StubBroker([], snapshot=snap)  # drain yields nothing
+    res = _tool()._collect_result(broker, inv, {"name": "dev"}, 1000)
+    assert res["exit_code"] == 0
+    assert not res["error"]
+    assert broker.cleaned
+
+
+def test_collect_result_times_out_when_snapshot_incomplete():
+    inv = SimpleNamespace(invocation_id="i", device_id="d", completed=False, error=None)
+    broker = _StubBroker([], snapshot=None)
+    res = _tool()._collect_result(broker, inv, {"name": "dev"}, 1000)
+    assert "did not respond" in (res["error"] or "")
+
+
+def test_collect_result_captures_control_chunk_past_deadline(monkeypatch):
+    from application.agents.tools import remote_device as rd
+
+    # First time.time() seeds the deadline; later calls are far past it, so the
+    # post-capture break fires — the control chunk must still be captured.
+    seq = iter([100.0])
+
+    def fake_time():
+        try:
+            return next(seq)
+        except StopIteration:
+            return 1e9
+
+    monkeypatch.setattr(rd.time, "time", fake_time)
+    inv = SimpleNamespace(invocation_id="i", device_id="d", completed=False, error=None)
+    broker = _StubBroker([{"stream": "control", "exit_code": 7, "duration_ms": 3}])
+    res = _tool()._collect_result(broker, inv, {"name": "dev"}, 30000)
+    assert res["exit_code"] == 7
+    assert not res["error"]
+
+
+# ---------------------------------------------------------------------------
+# cleanup / dispatch / next_command edge paths
+# ---------------------------------------------------------------------------
+def test_dispatch_failure_cleans_inv_hash(monkeypatch):
+    class RpushFailRedis(FakeRedis):
+        def rpush(self, key, *values):
+            raise RuntimeError("queue write failed")
+
+    fake = RpushFailRedis()
+    monkeypatch.setattr(
+        "application.devices.broker.get_redis_instance", lambda: fake
+    )
+    broker = DeviceBroker()
+    inv = broker.dispatch_invocation(
+        "d_fail", "u_fail",
+        {
+            "invocation_id": "inv_fail",
+            "action": "run_command",
+            "params": {"command": "echo secret"},
+        },
+    )
+    assert inv.completed is True
+    assert inv.error
+    # The plaintext command must not be stranded in the orphaned hash.
+    assert broker.get_invocation("inv_fail") is None
+    assert fake.exists("dev:inv:inv_fail") == 0
+
+
+def test_next_command_drops_reaped_invocation(broker_env):
+    broker, fake = broker_env
+    broker.dispatch_invocation(
+        "d_reap", "u_reap",
+        {"invocation_id": "inv_reap", "action": "run_command"},
+    )
+    # Invocation reaped (timed out / cleaned up) after it was queued.
+    fake.delete("dev:inv:inv_reap")
+    sess = broker.register_session("d_reap", "u_reap")
+    assert broker.next_command(sess, timeout=0.05) is None
+
+
+def test_byte_counts_are_utf8(broker_env):
+    broker, _fake = broker_env
+    broker.dispatch_invocation(
+        "d_utf", "u_utf",
+        {"invocation_id": "inv_utf", "action": "run_command"},
+    )
+    text = "héllo 世界 🚀"
+    broker.submit_output_chunk("inv_utf", {"stream": "stdout", "chunk": text})
+    snap = broker.get_invocation("inv_utf")
+    assert snap.stdout_bytes == len(text.encode("utf-8"))
+    assert snap.stdout_bytes != len(text)  # bytes != characters for this string
+
+
+def test_cmd_queue_ttl_covers_max_command_timeout():
+    # A queued command must outlive its own drain deadline so a briefly-offline
+    # device that reconnects late still receives it.
+    assert settings.REMOTE_DEVICE_CMD_QUEUE_TTL_SECONDS >= _MAX_TIMEOUT_MS / 1000 + 5

tests/devices/test_session_ticket.py

@@ -3,7 +3,7 @@
 Two layers:
 
 * Broker unit tests for ``claim_ticket`` / ``validate_ticket`` (issue, match,
-  mismatch, expiry, absence).
+  mismatch, eviction, absence).
 * Route tests proving the real CLI loop still works: ``/poll`` issues a
   ticket and ``session_events`` accepts *that* ticket, while a mismatched
   ticket is rejected with ``410`` before any stream opens.
@@ -20,72 +20,80 @@ from application.api.devices import auth as auth_module
 from application.api.devices import session as session_module
 from application.devices.broker import DeviceBroker
 
+from .conftest import FakeRedis
+
 
 # ---------------------------------------------------------------------------
 # Broker unit tests
 # ---------------------------------------------------------------------------
 def _queue_work(broker: DeviceBroker, device_id: str = "dev_x") -> None:
-    # No active session -> the invocation parks in _pending, which is what
-    # makes claim_ticket hand out a ticket.
+    # Queuing a command (no draining session) is what makes claim_ticket
+    # hand out a ticket: the device's command list is non-empty.
     broker.dispatch_invocation(
         device_id, "user_x", {"invocation_id": "inv_1", "action": "run_command"}
     )
 
 
-def test_claim_ticket_none_when_no_work():
-    broker = DeviceBroker()
+def test_claim_ticket_none_when_no_work(broker_env):
+    broker, _fake = broker_env
     assert broker.claim_ticket("dev_x", 30) is None
 
 
-def test_validate_accepts_issued_ticket():
-    broker = DeviceBroker()
+def test_validate_accepts_issued_ticket(broker_env):
+    broker, _fake = broker_env
     _queue_work(broker)
     ticket = broker.claim_ticket("dev_x", 30)
     assert ticket is not None
     assert broker.validate_ticket("dev_x", ticket) is True
 
 
-def test_validate_rejects_wrong_ticket():
-    broker = DeviceBroker()
+def test_validate_rejects_wrong_ticket(broker_env):
+    broker, _fake = broker_env
     _queue_work(broker)
     broker.claim_ticket("dev_x", 30)
     assert broker.validate_ticket("dev_x", "st_not_the_one") is False
 
 
-def test_validate_rejects_when_no_ticket_issued():
-    broker = DeviceBroker()
-    # Never polled -> nothing issued.
+def test_validate_rejects_when_no_ticket_issued(broker_env):
+    broker, _fake = broker_env
     assert broker.validate_ticket("dev_x", "st_anything") is False
 
 
-def test_validate_rejects_empty_session_id():
-    broker = DeviceBroker()
+def test_validate_rejects_empty_session_id(broker_env):
+    broker, _fake = broker_env
     _queue_work(broker)
     broker.claim_ticket("dev_x", 30)
     assert broker.validate_ticket("dev_x", "") is False
 
 
-def test_validate_rejects_expired_ticket_and_evicts():
-    broker = DeviceBroker()
+def test_validate_rejects_after_ticket_evicted(broker_env):
+    # Redis enforces the TTL; once the ticket key is gone (expired/evicted),
+    # validate can't resurrect it.
+    broker, fake = broker_env
     _queue_work(broker)
-    with patch("application.devices.broker.time.monotonic", return_value=1000.0):
-        ticket = broker.claim_ticket("dev_x", 30)
-    # 31s later the ticket is past its 30s TTL.
-    with patch("application.devices.broker.time.monotonic", return_value=1031.0):
-        assert broker.validate_ticket("dev_x", ticket) is False
-    # Evicted: a second check can't resurrect it.
+    ticket = broker.claim_ticket("dev_x", 30)
+    fake.delete("dev:ticket:dev_x")  # simulate TTL expiry
     assert broker.validate_ticket("dev_x", ticket) is False
 
 
-def test_register_session_consumes_issued_ticket_as_session_id():
+def test_claim_ticket_reused_while_unexpired(broker_env):
+    broker, _fake = broker_env
+    _queue_work(broker)
+    first = broker.claim_ticket("dev_x", 30)
+    second = broker.claim_ticket("dev_x", 30)
+    assert first == second
+
+
+def test_register_session_consumes_issued_ticket_as_session_id(broker_env):
     # The issued ticket becomes the session_id, so the URL the CLI opens
-    # (= the ticket) matches the live session. This is the invariant the
-    # route enforcement relies on.
-    broker = DeviceBroker()
+    # (= the ticket) matches the live session.
+    broker, fake = broker_env
     _queue_work(broker)
     ticket = broker.claim_ticket("dev_x", 30)
     sess = broker.register_session("dev_x", "user_x")
     assert sess.session_id == ticket
+    # Ticket is consumed on registration.
+    assert fake.get("dev:ticket:dev_x") is None
 
 
 # ---------------------------------------------------------------------------
@@ -151,67 +159,70 @@ def _call(app: Flask, view, path: str, *args):
 
 def test_poll_to_sse_upgrade_with_issued_ticket_works(app):
     """The legitimate loop: /poll issues a ticket, session_events accepts it."""
+    fake = FakeRedis()
     broker = DeviceBroker()
-    # Queue work so /poll returns a ticket rather than 202.
-    broker.dispatch_invocation(
-        "dev_route", "user_route",
-        {"invocation_id": "inv_a", "action": "run_command"},
-    )
-    with patch.object(session_module, "get_broker", return_value=broker):
-        poll_resp = _call(app, session_module.poll, "/api/devices/poll")
-        assert poll_resp.status_code == 200
-        payload = poll_resp.get_json()
-        ticket = payload["session_ticket"]
-        assert payload["session_url"] == f"/api/devices/sessions/{ticket}/events"
-        assert payload["expires_in"] == 30
-
-        # CLI opens the exact session_url it was handed.
-        sse_resp = _call(
-            app,
-            session_module.session_events,
-            f"/api/devices/sessions/{ticket}/events",
-            ticket,
+    with patch("application.devices.broker.get_redis_instance", return_value=fake):
+        # Queue work so /poll returns a ticket rather than 202.
+        broker.dispatch_invocation(
+            "dev_route", "user_route",
+            {"invocation_id": "inv_a", "action": "run_command"},
         )
-        try:
-            assert sse_resp.status_code == 200
-            assert sse_resp.mimetype == "text/event-stream"
-        finally:
-            # Close the streamed (stream_with_context) generator so its
-            # deferred app-context teardown runs now instead of leaking into
-            # the next test as a "popped wrong app context" error.
-            sse_resp.close()
+        with patch.object(session_module, "get_broker", return_value=broker):
+            poll_resp = _call(app, session_module.poll, "/api/devices/poll")
+            assert poll_resp.status_code == 200
+            payload = poll_resp.get_json()
+            ticket = payload["session_ticket"]
+            assert payload["session_url"] == f"/api/devices/sessions/{ticket}/events"
+            assert payload["expires_in"] == 30
+
+            # CLI opens the exact session_url it was handed.
+            sse_resp = _call(
+                app,
+                session_module.session_events,
+                f"/api/devices/sessions/{ticket}/events",
+                ticket,
+            )
+            try:
+                assert sse_resp.status_code == 200
+                assert sse_resp.mimetype == "text/event-stream"
+            finally:
+                sse_resp.close()
 
 
 def test_session_events_rejects_mismatched_ticket(app):
     """A fabricated/mismatched session_id is 410 Gone, no stream opened."""
+    fake = FakeRedis()
     broker = DeviceBroker()
-    broker.dispatch_invocation(
-        "dev_route", "user_route",
-        {"invocation_id": "inv_b", "action": "run_command"},
-    )
-    with patch.object(session_module, "get_broker", return_value=broker):
-        # Poll issues the real ticket...
-        _call(app, session_module.poll, "/api/devices/poll")
-        # ...but the client opens a different one.
-        resp = _call(
-            app,
-            session_module.session_events,
-            "/api/devices/sessions/st_bogus/events",
-            "st_bogus",
+    with patch("application.devices.broker.get_redis_instance", return_value=fake):
+        broker.dispatch_invocation(
+            "dev_route", "user_route",
+            {"invocation_id": "inv_b", "action": "run_command"},
         )
+        with patch.object(session_module, "get_broker", return_value=broker):
+            # Poll issues the real ticket...
+            _call(app, session_module.poll, "/api/devices/poll")
+            # ...but the client opens a different one.
+            resp = _call(
+                app,
+                session_module.session_events,
+                "/api/devices/sessions/st_bogus/events",
+                "st_bogus",
+            )
     assert resp.status_code == 410
     assert resp.get_json()["error"] == "session_ticket_invalid"
 
 
 def test_session_events_rejects_when_never_polled(app):
     """Opening the SSE stream without a prior poll is rejected (410)."""
+    fake = FakeRedis()
     broker = DeviceBroker()
-    with patch.object(session_module, "get_broker", return_value=broker):
-        resp = _call(
-            app,
-            session_module.session_events,
-            "/api/devices/sessions/st_anything/events",
-            "st_anything",
-        )
+    with patch("application.devices.broker.get_redis_instance", return_value=fake):
+        with patch.object(session_module, "get_broker", return_value=broker):
+            resp = _call(
+                app,
+                session_module.session_events,
+                "/api/devices/sessions/st_anything/events",
+                "st_anything",
+            )
     assert resp.status_code == 410
     assert resp.get_json()["error"] == "session_ticket_invalid"

tests/devices/test_submit_output_audit.py

@@ -0,0 +1,206 @@
+"""Route-level coverage for submit_output's audit write (fix 4).
+
+The audit outcome must come from the control chunk captured locally during
+the parse loop, so it survives the draining (worker) process racing to delete
+the invocation's Redis state — modelled here by the post-submit snapshot
+returning None.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+from flask import Flask
+
+from application.api.devices import auth as auth_module
+from application.api.devices import session as session_module
+
+
+@pytest.fixture
+def app():
+    return Flask(__name__)
+
+
+def _device_row() -> dict:
+    return {
+        "id": "dev_route",
+        "user_id": "user_route",
+        "name": "laptop",
+        "status": "active",
+        "approval_mode": "ask",
+        "machine_pubkey_fingerprint": "fp",
+        "token_hash": "tokhash",
+    }
+
+
+class _Repo:
+    def __init__(self, _conn):
+        pass
+
+    def find_by_token_hash(self, _token_hash):
+        return _device_row()
+
+    def touch_last_seen(self, _device_id):
+        pass
+
+
+class _Ctx:
+    def __enter__(self):
+        return None
+
+    def __exit__(self, *a):
+        return False
+
+
+class _AuditSpy:
+    calls: list = []
+
+    def __init__(self, _conn):
+        pass
+
+    def record_result(self, invocation_id, **kwargs):
+        _AuditSpy.calls.append((invocation_id, kwargs))
+
+
+class _RaceBroker:
+    """First get_invocation (ownership check) returns a live inv; the second
+    (post-submit snapshot) returns ``snap`` — None models the worker having
+    already deleted the hash.
+    """
+
+    def __init__(self, snap):
+        self._snap = snap
+        self.calls = 0
+
+    def get_invocation(self, _invocation_id):
+        self.calls += 1
+        if self.calls == 1:
+            return SimpleNamespace(device_id="dev_route", completed=True)
+        return self._snap
+
+    def submit_output_chunk(self, _invocation_id, _chunk):
+        return True
+
+    def submit_ack(self, _invocation_id, _decision, _reason=None):
+        return True
+
+
+CONTROL_BODY = b'{"stream":"control","exit_code":9,"duration_ms":42}\n'
+
+
+def _post(app: Flask, broker, body: bytes, invocation_id: str = "inv_audit"):
+    path = f"/api/devices/sessions/s/invocations/{invocation_id}/output"
+    with app.test_request_context(
+        path, method="POST", data=body, headers={"Authorization": "Bearer tok"}
+    ):
+        patches = [
+            patch.object(auth_module, "DevicesRepository", _Repo),
+            patch.object(auth_module, "db_readonly", _Ctx),
+            patch.object(auth_module, "db_session", _Ctx),
+            patch.object(session_module, "get_broker", return_value=broker),
+            patch.object(session_module, "DeviceAuditLogRepository", _AuditSpy),
+            patch.object(session_module, "db_session", _Ctx),
+        ]
+        for p in patches:
+            p.start()
+        try:
+            return session_module.submit_output("s", invocation_id)
+        finally:
+            for p in patches:
+                p.stop()
+
+
+def test_audit_written_from_control_chunk_when_snapshot_present(app):
+    _AuditSpy.calls.clear()
+    snap = SimpleNamespace(started_at=None, stdout_bytes=3, stderr_bytes=0)
+    resp = _post(app, _RaceBroker(snap), CONTROL_BODY)
+    assert resp.status_code == 200
+    assert len(_AuditSpy.calls) == 1
+    inv_id, kw = _AuditSpy.calls[0]
+    assert inv_id == "inv_audit"
+    assert kw["exit_code"] == 9
+    assert kw["duration_ms"] == 42
+    assert kw["stdout_bytes"] == 3
+
+
+def test_audit_written_even_when_snapshot_deleted_by_race(app):
+    # snap is None: the worker already cleaned up the hash. The functional
+    # outcome fields must still land, sourced from the local control chunk.
+    _AuditSpy.calls.clear()
+    resp = _post(app, _RaceBroker(None), CONTROL_BODY)
+    assert resp.status_code == 200
+    assert len(_AuditSpy.calls) == 1
+    _inv_id, kw = _AuditSpy.calls[0]
+    assert kw["exit_code"] == 9
+    assert kw["duration_ms"] == 42
+    assert kw["started_at"] is None
+    assert kw["stdout_bytes"] == 0
+
+
+def test_no_audit_without_control_chunk(app):
+    _AuditSpy.calls.clear()
+    resp = _post(app, _RaceBroker(None), b'{"stream":"stdout","chunk":"hi"}\n')
+    assert resp.status_code == 200
+    assert _AuditSpy.calls == []
+
+
+# ---------------------------------------------------------------------------
+# ack_invocation: a denial is a terminal outcome with no subsequent output,
+# so it must record its own audit row (the device never POSTs output).
+# ---------------------------------------------------------------------------
+def _ack(app: Flask, broker, decision: str, invocation_id: str = "inv_audit",
+         reason=None):
+    path = f"/api/devices/sessions/s/invocations/{invocation_id}/ack"
+    with app.test_request_context(
+        path,
+        method="POST",
+        json={"decision": decision, "reason": reason},
+        headers={"Authorization": "Bearer tok"},
+    ):
+        patches = [
+            patch.object(auth_module, "DevicesRepository", _Repo),
+            patch.object(auth_module, "db_readonly", _Ctx),
+            patch.object(auth_module, "db_session", _Ctx),
+            patch.object(session_module, "get_broker", return_value=broker),
+            patch.object(session_module, "DeviceAuditLogRepository", _AuditSpy),
+            patch.object(session_module, "db_session", _Ctx),
+        ]
+        for p in patches:
+            p.start()
+        try:
+            return session_module.ack_invocation("s", invocation_id)
+        finally:
+            for p in patches:
+                p.stop()
+
+
+def test_denied_ack_records_terminal_audit_outcome(app):
+    _AuditSpy.calls.clear()
+    resp = _ack(app, _RaceBroker(None), "denied", reason="user rejected")
+    assert resp.status_code == 200
+    assert len(_AuditSpy.calls) == 1
+    inv_id, kw = _AuditSpy.calls[0]
+    assert inv_id == "inv_audit"
+    assert kw["error"] == "denied"
+    assert kw["finished_at"] is not None
+    # exit_code is intentionally not passed (a denied command never ran), so it
+    # stays NULL via record_result's COALESCE update.
+    assert kw.get("exit_code") is None
+
+
+def test_accepted_ack_does_not_record_outcome(app):
+    # An accepted command will run and POST output; submit_output records the
+    # outcome then, so ack must not write a (premature) terminal row.
+    _AuditSpy.calls.clear()
+    resp = _ack(app, _RaceBroker(None), "accepted")
+    assert resp.status_code == 200
+    assert _AuditSpy.calls == []
+
+
+def test_auto_approved_ack_does_not_record_outcome(app):
+    _AuditSpy.calls.clear()
+    resp = _ack(app, _RaceBroker(None), "auto_approved")
+    assert resp.status_code == 200
+    assert _AuditSpy.calls == []