mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-05 16:13:51 +00:00
arc53-machine
f882ef49a7
About 85 call sites read a setting as getattr(settings, "NAME", fallback), each carrying its own copy of the default. Every one of those names is a field with a default on the model, so the fallback could never apply to the real settings object; it only masked drift. Two had drifted: - OPENAI_PROMPT_CACHE_KEY defaults to True on the model but the reader fell back to False, and two test stubs relied on that. - SharePoint's MICROSOFT_AUTHORITY fallback to https://login.microsoftonline.com/<tenant> never fired, because the attribute always exists (as None), so MSAL got authority=None. The connector now derives the tenant authority when the setting is unset, as its test always assumed. Four places read EMBEDDINGS_KEY straight from os.environ, skipping the "None"/"" normalisation the model applies; they read the setting now. Test stubs that replaced a module's settings with a SimpleNamespace list every setting the code under test reads.
675 lines
26 KiB
Python
675 lines
26 KiB
Python
"""Redis-backed broker that routes invocations to active device sessions.
|
|
|
|
The device CLI only ever talks to the web process (poll / SSE / output
|
|
POST), but an agent may dispatch a command from *any* process — notably a
|
|
Celery worker during a scheduled run. An in-memory broker can't bridge
|
|
that gap: the worker and the web process don't share Python memory, so a
|
|
worker-side dispatch never reaches the web-side SSE session and the tool
|
|
times out. Routing through Redis makes every hop cross-process.
|
|
|
|
What lives in Redis (shared) vs. in the process (ephemeral):
|
|
|
|
* ``dev:cmd:{device_id}`` — list of queued command envelopes (JSON).
|
|
* ``dev:ticket:{device_id}``— poll-issued SSE upgrade ticket (string, TTL).
|
|
* ``dev:inv:{invocation_id}``— invocation metadata hash (status, result).
|
|
* ``dev:out:{invocation_id}``— stream of stdout/stderr/control chunks.
|
|
* ``SessionState`` is per-connection state for the one SSE handler that
|
|
owns the live socket; it stays in that process's memory.
|
|
|
|
Lifecycle:
|
|
1. Agent (any process) calls ``dispatch_invocation`` → metadata hash + an
|
|
RPUSH onto the device's command list.
|
|
2. The web SSE handler blocks on that list (``next_command``) and emits
|
|
each envelope to the wire; an offline device leaves the envelope queued
|
|
until its next ``/poll`` issues a ticket and upgrades to SSE.
|
|
3. The CLI POSTs ack + chunked output back; ``submit_output_chunk`` XADDs
|
|
chunks to the invocation's output stream and updates the hash.
|
|
4. A ``control`` chunk closes the invocation; ``drain_output`` (in the
|
|
dispatching process) reads the stream from the start and stops on it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, Iterator, Optional
|
|
|
|
import anyio
|
|
|
|
from docsgpt.cache import get_redis_instance
|
|
from docsgpt.core.settings import settings
|
|
from docsgpt.streaming.async_redis import get_async_redis_instance
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Kept for backwards compatibility with callers/tests importing the
|
|
# sentinel; the Redis path signals end-of-output with a ``control`` chunk.
|
|
INVOCATION_DONE = object()
|
|
|
|
|
|
def _cmd_key(device_id: str) -> str:
|
|
return f"dev:cmd:{device_id}"
|
|
|
|
|
|
def _ticket_key(device_id: str) -> str:
|
|
return f"dev:ticket:{device_id}"
|
|
|
|
|
|
# Deletes the device's ticket iff it still holds the presented one. Runs
|
|
# server-side so two requests racing with one ticket can't both redeem it.
|
|
_REDEEM_TICKET_LUA = """
|
|
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
|
return redis.call('DEL', KEYS[1])
|
|
end
|
|
return 0
|
|
"""
|
|
|
|
|
|
def _inv_key(invocation_id: str) -> str:
|
|
return f"dev:inv:{invocation_id}"
|
|
|
|
|
|
def _out_key(invocation_id: str) -> str:
|
|
return f"dev:out:{invocation_id}"
|
|
|
|
|
|
def _as_str(value: Any) -> str:
|
|
if isinstance(value, (bytes, bytearray)):
|
|
return value.decode("utf-8", "replace")
|
|
return str(value)
|
|
|
|
|
|
@dataclass
|
|
class Invocation:
|
|
"""Snapshot of an invocation's cross-process state.
|
|
|
|
Returned by ``dispatch_invocation`` (fresh) and ``get_invocation`` (read
|
|
from Redis). Carries enough for ownership checks and audit without the
|
|
caller touching Redis directly.
|
|
"""
|
|
|
|
invocation_id: str
|
|
device_id: str
|
|
completed: bool = False
|
|
exit_code: Optional[int] = None
|
|
duration_ms: Optional[int] = None
|
|
error: Optional[str] = None
|
|
started_at: Optional[float] = None
|
|
finished_at: Optional[float] = None
|
|
stdout_bytes: int = 0
|
|
stderr_bytes: int = 0
|
|
|
|
|
|
@dataclass
|
|
class SessionState:
|
|
"""Per-connection state for the SSE handler that owns the live socket."""
|
|
|
|
session_id: str
|
|
device_id: str
|
|
user_id: str
|
|
last_event_id: int = 0
|
|
last_activity_at: float = field(default_factory=time.time)
|
|
closed: threading.Event = field(default_factory=threading.Event)
|
|
|
|
|
|
class DeviceBroker:
|
|
"""Cross-process device registry backed by Redis.
|
|
|
|
Redis is the source of truth for queued commands, output, and tickets,
|
|
so dispatch and drain work regardless of which process they run in. A
|
|
small in-memory map of live SSE sessions stays local to the web process
|
|
that holds each socket (sessions are inherently per-connection).
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._sessions_by_id: Dict[str, SessionState] = {}
|
|
self._sessions_by_device: Dict[str, SessionState] = {}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Session lifecycle (web process / CLI side)
|
|
# ------------------------------------------------------------------
|
|
def register_session(self, device_id: str, user_id: str) -> SessionState:
|
|
"""Open the SSE session for ``device_id``, adopting its poll ticket.
|
|
|
|
The poll-issued ticket becomes the ``session_id`` so the URL the CLI
|
|
opens matches the live session. A previous local session for the
|
|
same device is closed and replaced (the CLI reconnected).
|
|
"""
|
|
redis = get_redis_instance()
|
|
issued = None
|
|
if redis is not None:
|
|
try:
|
|
issued = redis.get(_ticket_key(device_id))
|
|
if issued is not None:
|
|
redis.delete(_ticket_key(device_id))
|
|
except Exception:
|
|
logger.exception("ticket lookup failed for %s", device_id)
|
|
session_id = _as_str(issued) if issued else f"st_{uuid.uuid4().hex}"
|
|
return self._adopt_session(session_id, device_id, user_id)
|
|
|
|
def redeem_ticket(self, device_id: str, user_id: str, ticket: str) -> Optional[SessionState]:
|
|
"""Consume the device's poll-issued ``ticket`` and open its session under it.
|
|
|
|
Returns ``None``, and opens nothing, unless ``ticket`` is still the
|
|
device's unexpired ticket. The compare and delete are one Redis step,
|
|
so a replay racing the first redeem can't replace its session.
|
|
"""
|
|
if not ticket:
|
|
return None
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
redeemed = redis.eval(_REDEEM_TICKET_LUA, 1, _ticket_key(device_id), ticket)
|
|
except Exception:
|
|
logger.exception("ticket redeem failed for %s", device_id)
|
|
return None
|
|
if not redeemed:
|
|
return None
|
|
return self._adopt_session(ticket, device_id, user_id)
|
|
|
|
def _adopt_session(self, session_id: str, device_id: str, user_id: str) -> SessionState:
|
|
"""Make a new session the device's live one, closing the one it replaces."""
|
|
sess = SessionState(
|
|
session_id=session_id, device_id=device_id, user_id=user_id
|
|
)
|
|
with self._lock:
|
|
prior = self._sessions_by_device.get(device_id)
|
|
if prior is not None:
|
|
prior.closed.set()
|
|
self._sessions_by_id.pop(prior.session_id, None)
|
|
self._sessions_by_device[device_id] = sess
|
|
self._sessions_by_id[session_id] = sess
|
|
return sess
|
|
|
|
def close_session(self, session_id: str, *, reason: str = "closed") -> None:
|
|
"""Close the local session by id.
|
|
|
|
Queued-but-undelivered commands stay on the device's Redis list and
|
|
are picked up by the next session; an in-flight command whose socket
|
|
drops falls back to the tool's own drain deadline.
|
|
"""
|
|
with self._lock:
|
|
sess = self._sessions_by_id.pop(session_id, None)
|
|
if sess is None:
|
|
return
|
|
sess.closed.set()
|
|
if self._sessions_by_device.get(sess.device_id) is sess:
|
|
self._sessions_by_device.pop(sess.device_id, None)
|
|
logger.debug("device session closed: %s (%s)", session_id, reason)
|
|
|
|
def get_session(self, session_id: str) -> Optional[SessionState]:
|
|
with self._lock:
|
|
return self._sessions_by_id.get(session_id)
|
|
|
|
def next_command(
|
|
self, session: SessionState, timeout: float = 1.0
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Block up to ``timeout`` for the next queued command envelope.
|
|
|
|
Returns the decoded envelope, or ``None`` on timeout so the SSE
|
|
handler can emit a keepalive and re-check the session lifecycle.
|
|
"""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
time.sleep(timeout)
|
|
return None
|
|
try:
|
|
popped = redis.blpop(_cmd_key(session.device_id), timeout=timeout)
|
|
except Exception:
|
|
logger.exception("blpop failed for %s", session.device_id)
|
|
time.sleep(timeout)
|
|
return None
|
|
if not popped:
|
|
return None
|
|
envelope = self._decode_envelope(popped[1])
|
|
if envelope is None:
|
|
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
|
|
|
|
async def next_command_async(
|
|
self, session: SessionState, timeout: float = 1.0
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Await up to ``timeout`` for the next queued command envelope.
|
|
|
|
Event-loop twin of :meth:`next_command` for the native-async SSE
|
|
stream: the same envelope checks over the async Redis client, so an
|
|
idle device session holds no thread.
|
|
"""
|
|
redis = await get_async_redis_instance()
|
|
if redis is None:
|
|
await anyio.sleep(timeout)
|
|
return None
|
|
try:
|
|
popped = await redis.blpop(_cmd_key(session.device_id), timeout=timeout)
|
|
except Exception:
|
|
logger.exception("async blpop failed for %s", session.device_id)
|
|
await anyio.sleep(timeout)
|
|
return None
|
|
if not popped:
|
|
return None
|
|
envelope = self._decode_envelope(popped[1])
|
|
if envelope is None:
|
|
return None
|
|
# Same reaped-invocation guard as ``next_command``; a failed lookup
|
|
# drops the envelope too, as ``get_invocation`` returning None does.
|
|
inv_id = envelope.get("invocation_id")
|
|
if inv_id:
|
|
try:
|
|
reaped = not await redis.exists(_inv_key(inv_id))
|
|
except Exception:
|
|
logger.exception("invocation lookup failed for %s", inv_id)
|
|
reaped = True
|
|
if reaped:
|
|
logger.debug("dropping reaped invocation %s", inv_id)
|
|
return None
|
|
return envelope
|
|
|
|
@staticmethod
|
|
def _decode_envelope(raw: Any) -> Optional[Dict[str, Any]]:
|
|
"""Parse a queued command envelope; ``None`` for anything malformed."""
|
|
try:
|
|
envelope = json.loads(_as_str(raw))
|
|
except (TypeError, ValueError):
|
|
logger.warning("dropping malformed command envelope")
|
|
return None
|
|
return envelope if isinstance(envelope, dict) else None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Polling / tickets
|
|
# ------------------------------------------------------------------
|
|
def claim_ticket(self, device_id: str, ttl_seconds: float) -> Optional[str]:
|
|
"""Return an SSE upgrade ticket iff the device has queued work.
|
|
|
|
Reuses an unexpired ticket so repeated polls don't churn it; the
|
|
ticket's Redis TTL doubles as the advertised ``expires_in`` window.
|
|
"""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
if redis.llen(_cmd_key(device_id)) <= 0:
|
|
return None
|
|
existing = redis.get(_ticket_key(device_id))
|
|
if existing:
|
|
return _as_str(existing)
|
|
ticket = f"st_{uuid.uuid4().hex}"
|
|
redis.set(_ticket_key(device_id), ticket, ex=int(ttl_seconds))
|
|
return ticket
|
|
except Exception:
|
|
logger.exception("claim_ticket failed for %s", device_id)
|
|
return None
|
|
|
|
def validate_ticket(self, device_id: str, session_id: str) -> bool:
|
|
"""True iff ``session_id`` is the unexpired ticket issued to the device.
|
|
|
|
An absent, mismatched, or expired ticket is rejected; expiry is
|
|
enforced by Redis's TTL (a ``GET`` of an expired key returns nil).
|
|
"""
|
|
if not session_id:
|
|
return False
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return False
|
|
try:
|
|
issued = redis.get(_ticket_key(device_id))
|
|
except Exception:
|
|
logger.exception("validate_ticket failed for %s", device_id)
|
|
return False
|
|
return issued is not None and _as_str(issued) == session_id
|
|
|
|
# ------------------------------------------------------------------
|
|
# Dispatch (server-issued, any process)
|
|
# ------------------------------------------------------------------
|
|
def dispatch_invocation(
|
|
self,
|
|
device_id: str,
|
|
user_id: str,
|
|
envelope: Dict[str, Any],
|
|
) -> Invocation:
|
|
"""Queue an invocation for ``device_id`` and record its metadata.
|
|
|
|
Writes the metadata hash, then RPUSHes the envelope onto the
|
|
device's command list. A live SSE session draining the list picks it
|
|
up immediately; otherwise it waits for the next poll-issued ticket.
|
|
"""
|
|
invocation_id = envelope["invocation_id"]
|
|
inv = Invocation(invocation_id=invocation_id, device_id=device_id)
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
inv.error = "device broker unavailable"
|
|
inv.completed = True
|
|
return inv
|
|
envelope_json = json.dumps(envelope)
|
|
try:
|
|
redis.hset(
|
|
_inv_key(invocation_id),
|
|
mapping={
|
|
"device_id": device_id,
|
|
"user_id": user_id,
|
|
"envelope": envelope_json,
|
|
"completed": "0",
|
|
"stdout_bytes": "0",
|
|
"stderr_bytes": "0",
|
|
},
|
|
)
|
|
redis.expire(_inv_key(invocation_id), self._inv_ttl())
|
|
redis.rpush(_cmd_key(device_id), envelope_json)
|
|
redis.expire(_cmd_key(device_id), self._cmd_ttl())
|
|
except Exception:
|
|
logger.exception("dispatch_invocation failed for %s", invocation_id)
|
|
# Don't strand the metadata hash (it holds the plaintext command)
|
|
# if the queue write failed partway through.
|
|
try:
|
|
redis.delete(_inv_key(invocation_id))
|
|
except Exception:
|
|
logger.debug(
|
|
"cleanup after failed dispatch failed for %s", invocation_id
|
|
)
|
|
inv.error = "device broker dispatch failed"
|
|
inv.completed = True
|
|
return inv
|
|
|
|
def get_invocation(self, invocation_id: str) -> Optional[Invocation]:
|
|
"""Read an invocation's current metadata snapshot from Redis."""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
raw = redis.hgetall(_inv_key(invocation_id))
|
|
except Exception:
|
|
logger.exception("get_invocation failed for %s", invocation_id)
|
|
return None
|
|
if not raw:
|
|
return None
|
|
h = {_as_str(k): _as_str(v) for k, v in raw.items()}
|
|
return Invocation(
|
|
invocation_id=invocation_id,
|
|
device_id=h.get("device_id", ""),
|
|
completed=h.get("completed") == "1",
|
|
exit_code=_to_int(h.get("exit_code")),
|
|
duration_ms=_to_int(h.get("duration_ms")),
|
|
error=h.get("error") or None,
|
|
started_at=_to_float(h.get("started_at")),
|
|
finished_at=_to_float(h.get("finished_at")),
|
|
stdout_bytes=_to_int(h.get("stdout_bytes")) or 0,
|
|
stderr_bytes=_to_int(h.get("stderr_bytes")) or 0,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Output streaming (web process / CLI side)
|
|
# ------------------------------------------------------------------
|
|
def submit_output_chunk(
|
|
self, invocation_id: str, chunk: Dict[str, Any]
|
|
) -> bool:
|
|
"""Forward one CLI output chunk to the dispatching process's drain.
|
|
|
|
Updates the metadata hash (byte counts; result fields on the closing
|
|
``control`` chunk) and XADDs the chunk to the invocation's output
|
|
stream. Returns ``False`` for an unknown invocation.
|
|
"""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return False
|
|
key = _inv_key(invocation_id)
|
|
try:
|
|
device_id = redis.hget(key, "device_id")
|
|
except Exception:
|
|
logger.exception("submit_output_chunk read failed for %s", invocation_id)
|
|
return False
|
|
if device_id is None:
|
|
return False
|
|
device_id = _as_str(device_id)
|
|
now = time.time()
|
|
stream = chunk.get("stream")
|
|
try:
|
|
# Append to the output stream BEFORE flipping completion state, so a
|
|
# reader that observes completed=1 is guaranteed the chunk (and every
|
|
# earlier one) is already on the stream — otherwise drain could see
|
|
# completion and stop before the control chunk lands.
|
|
redis.xadd(
|
|
_out_key(invocation_id),
|
|
{"c": json.dumps(chunk)},
|
|
maxlen=self._out_maxlen(),
|
|
approximate=True,
|
|
)
|
|
redis.expire(_out_key(invocation_id), self._inv_ttl())
|
|
if stream in ("stdout", "stderr"):
|
|
text = chunk.get("chunk")
|
|
if isinstance(text, str):
|
|
field = "stdout_bytes" if stream == "stdout" else "stderr_bytes"
|
|
redis.hincrby(key, field, len(text.encode("utf-8")))
|
|
elif stream == "control":
|
|
redis.hsetnx(key, "started_at", repr(now))
|
|
mapping = {"completed": "1", "finished_at": repr(now)}
|
|
if chunk.get("exit_code") is not None:
|
|
mapping["exit_code"] = str(int(chunk["exit_code"]))
|
|
if chunk.get("duration_ms") is not None:
|
|
mapping["duration_ms"] = str(int(chunk["duration_ms"]))
|
|
if chunk.get("error"):
|
|
mapping["error"] = _as_str(chunk["error"])
|
|
redis.hset(key, mapping=mapping)
|
|
redis.expire(key, self._inv_ttl())
|
|
except Exception:
|
|
logger.exception("submit_output_chunk failed for %s", invocation_id)
|
|
return False
|
|
# Keep a co-located SSE session alive while output flows (single-worker
|
|
# web tier); a cross-worker session relies on its own keepalive.
|
|
with self._lock:
|
|
sess = self._sessions_by_device.get(device_id)
|
|
if sess is not None:
|
|
sess.last_activity_at = now
|
|
return True
|
|
|
|
def submit_ack(
|
|
self,
|
|
invocation_id: str,
|
|
decision: str,
|
|
reason: Optional[str] = None,
|
|
) -> bool:
|
|
"""Record the CLI's accept/deny decision for an invocation."""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return False
|
|
key = _inv_key(invocation_id)
|
|
try:
|
|
if not redis.exists(key):
|
|
return False
|
|
now = time.time()
|
|
mapping = {"decision": decision}
|
|
if reason:
|
|
mapping["decision_reason"] = reason
|
|
redis.hset(key, mapping=mapping)
|
|
redis.hsetnx(key, "started_at", repr(now))
|
|
redis.expire(key, self._inv_ttl())
|
|
if decision == "denied":
|
|
# XADD the synthetic control chunk BEFORE marking completed, so a
|
|
# racing drain that observes completed=1 always finds the chunk
|
|
# and reports "denied" rather than a false timeout.
|
|
redis.xadd(
|
|
_out_key(invocation_id),
|
|
{"c": json.dumps(
|
|
{"stream": "control", "exit_code": None, "error": "denied"}
|
|
)},
|
|
maxlen=self._out_maxlen(),
|
|
approximate=True,
|
|
)
|
|
redis.expire(_out_key(invocation_id), self._inv_ttl())
|
|
redis.hset(
|
|
key,
|
|
mapping={
|
|
"completed": "1",
|
|
"error": "denied",
|
|
"finished_at": repr(now),
|
|
},
|
|
)
|
|
except Exception:
|
|
logger.exception("submit_ack failed for %s", invocation_id)
|
|
return False
|
|
return True
|
|
|
|
def drain_output(
|
|
self,
|
|
invocation_id: str,
|
|
timeout: float = 0.5,
|
|
deadline: Optional[float] = None,
|
|
) -> Iterator[Dict[str, Any]]:
|
|
"""Yield queued output chunks for an invocation; stop on ``control``.
|
|
|
|
Reads the invocation's output stream from the start (so chunks
|
|
XADDed before draining began are never missed) and blocks up to
|
|
``timeout`` per read. ``deadline`` is an absolute ``time.time()``;
|
|
once it passes with no closing ``control`` chunk the generator
|
|
returns so a device that never responds can't loop forever.
|
|
"""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
yield {
|
|
"stream": "control",
|
|
"exit_code": None,
|
|
"error": "device broker unavailable",
|
|
}
|
|
return
|
|
last_id = "0-0"
|
|
block_ms = max(1, int(timeout * 1000))
|
|
out_key = _out_key(invocation_id)
|
|
# ``resp`` is parsed as the RESP2 list shape ``[[key, [(id, fields)]]]``;
|
|
# the client must stay protocol=2 (redis-py default — see cache.py).
|
|
tail_flush = False
|
|
while True:
|
|
try:
|
|
# On the final sweep, read non-blocking (block=None) so any
|
|
# entries that landed after the prior empty read are drained
|
|
# before returning — closing the completion-vs-control race.
|
|
resp = redis.xread(
|
|
{out_key: last_id},
|
|
count=200,
|
|
block=None if tail_flush else block_ms,
|
|
)
|
|
except Exception:
|
|
logger.exception("xread failed for %s", invocation_id)
|
|
return
|
|
if not resp:
|
|
if tail_flush:
|
|
return
|
|
done = self._is_completed(redis, invocation_id)
|
|
timed_out = deadline is not None and time.time() >= deadline
|
|
if done or timed_out:
|
|
# One final non-blocking sweep from last_id, then stop.
|
|
tail_flush = True
|
|
continue
|
|
for _stream_key, entries in resp:
|
|
for entry_id, fields in entries:
|
|
last_id = _as_str(entry_id)
|
|
raw = fields.get(b"c")
|
|
if raw is None:
|
|
raw = fields.get("c")
|
|
if raw is None:
|
|
continue
|
|
try:
|
|
chunk = json.loads(_as_str(raw))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if not isinstance(chunk, dict):
|
|
continue
|
|
if chunk.get("stream") in ("stdout", "stderr", "control"):
|
|
yield chunk
|
|
if chunk.get("stream") == "control":
|
|
return
|
|
|
|
def cleanup_invocation(self, invocation_id: str) -> None:
|
|
"""Drop an invocation's Redis state, including any undelivered command.
|
|
|
|
Deletes the metadata hash first so a concurrent ``next_command`` that
|
|
just BLPOPped this envelope re-checks ``get_invocation``, sees it gone,
|
|
and drops it instead of delivering a command the user already saw fail;
|
|
the ``LREM`` then clears it for the still-queued (offline-device) case.
|
|
Best-effort: a delivery that wins a tight race with the delete can still
|
|
reach the device once, and that run's output is discarded.
|
|
"""
|
|
redis = get_redis_instance()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
raw = redis.hgetall(_inv_key(invocation_id))
|
|
device_id = envelope_json = None
|
|
if raw:
|
|
h = {_as_str(k): _as_str(v) for k, v in raw.items()}
|
|
device_id = h.get("device_id")
|
|
envelope_json = h.get("envelope")
|
|
redis.delete(_inv_key(invocation_id))
|
|
redis.delete(_out_key(invocation_id))
|
|
if device_id and envelope_json:
|
|
redis.lrem(_cmd_key(device_id), 0, envelope_json)
|
|
except Exception:
|
|
logger.exception("cleanup_invocation failed for %s", invocation_id)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internals
|
|
# ------------------------------------------------------------------
|
|
@staticmethod
|
|
def _is_completed(redis, invocation_id: str) -> bool:
|
|
try:
|
|
return _as_str(redis.hget(_inv_key(invocation_id), "completed")) == "1"
|
|
except Exception:
|
|
return False
|
|
|
|
@staticmethod
|
|
def _inv_ttl() -> int:
|
|
return int(settings.REMOTE_DEVICE_INVOCATION_TTL_SECONDS)
|
|
|
|
@staticmethod
|
|
def _cmd_ttl() -> int:
|
|
return int(settings.REMOTE_DEVICE_CMD_QUEUE_TTL_SECONDS)
|
|
|
|
@staticmethod
|
|
def _out_maxlen() -> int:
|
|
return int(settings.REMOTE_DEVICE_OUTPUT_STREAM_MAXLEN)
|
|
|
|
|
|
def _to_int(value: Optional[str]) -> Optional[int]:
|
|
if value is None or value == "":
|
|
return None
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _to_float(value: Optional[str]) -> Optional[float]:
|
|
if value is None or value == "":
|
|
return None
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
_broker_instance: Optional[DeviceBroker] = None
|
|
_broker_lock = threading.Lock()
|
|
|
|
|
|
def get_broker() -> DeviceBroker:
|
|
"""Return the process-wide ``DeviceBroker`` instance."""
|
|
global _broker_instance
|
|
if _broker_instance is None:
|
|
with _broker_lock:
|
|
if _broker_instance is None:
|
|
_broker_instance = DeviceBroker()
|
|
return _broker_instance
|