diff --git a/application/agents/tools/remote_device.py b/application/agents/tools/remote_device.py index f728a0b7..f4b86c7f 100644 --- a/application/agents/tools/remote_device.py +++ b/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, } diff --git a/application/api/devices/session.py b/application/api/devices/session.py index 02530d15..8f15aabc 100644 --- a/application/api/devices/session.py +++ b/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" diff --git a/application/core/settings.py b/application/core/settings.py index b93dbbf0..0dee0db6 100644 --- a/application/core/settings.py +++ b/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 diff --git a/application/devices/broker.py b/application/devices/broker.py index 52305ee8..6f542162 100644 --- a/application/devices/broker.py +++ b/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 diff --git a/tests/devices/conftest.py b/tests/devices/conftest.py new file mode 100644 index 00000000..20aa6c55 --- /dev/null +++ b/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 diff --git a/tests/devices/test_broker_cleanup.py b/tests/devices/test_broker_cleanup.py index 561cc744..d0ced0dd 100644 --- a/tests/devices/test_broker_cleanup.py +++ b/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" diff --git a/tests/devices/test_broker_cross_process.py b/tests/devices/test_broker_cross_process.py new file mode 100644 index 00000000..5ced19e0 --- /dev/null +++ b/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" diff --git a/tests/devices/test_broker_drain.py b/tests/devices/test_broker_drain.py index 16dc784f..00778620 100644 --- a/tests/devices/test_broker_drain.py +++ b/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"] diff --git a/tests/devices/test_broker_race.py b/tests/devices/test_broker_race.py new file mode 100644 index 00000000..6fbc5e9d --- /dev/null +++ b/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 diff --git a/tests/devices/test_session_ticket.py b/tests/devices/test_session_ticket.py index 9a235af0..41a56c5e 100644 --- a/tests/devices/test_session_ticket.py +++ b/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" diff --git a/tests/devices/test_submit_output_audit.py b/tests/devices/test_submit_output_audit.py new file mode 100644 index 00000000..2484dd6f --- /dev/null +++ b/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 == []