mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 20:13:04 +00:00
* fix: route remote-device tool through Redis so scheduled runs reach the device
The remote-device tool worked interactively but timed out on every scheduled
run. DeviceBroker was an in-process, in-memory singleton, but scheduled runs
execute in the Celery worker — a different process from the gunicorn web tier
that holds the device's SSE session — so a worker-side dispatch never reached
the device and the tool always hit its deadline.
Make the broker Redis-backed so every hop crosses the process boundary:
- queued commands -> Redis list dev:cmd:{device_id}
- output chunks -> Redis stream dev:out:{invocation_id}
- invocation metadata -> Redis hash dev:inv:{invocation_id}
- SSE upgrade tickets -> Redis key dev🎫{device_id}
Per-connection SSE session state stays in the web process. Reuses the existing
get_redis_instance()/CACHE_REDIS_URL; no new infrastructure. Also makes the web
tier safe to scale past one worker.
Concurrency hardening (from adversarial review + real-Redis e2e):
- XADD the output/control chunk before flipping completed=1, and have
drain_output do a final non-blocking flush after observing completion, so a
reader can't see completion and stop before the control chunk lands (this had
reintroduced the false "device did not respond (timed out)" under a race).
- _collect_result builds the result from drained chunks, checks the deadline
only after capturing a chunk, and falls back to the authoritative snapshot
(before cleanup) when no control chunk was observed.
- Audit outcome is written from locally-known fields so it survives the worker
racing to delete the invocation; a denied command now records a terminal
"denied" outcome instead of staying "dispatched".
- cmd-queue TTL raised to 900s (>= max drain deadline); dispatch-failure and
reaped-invocation cleanup; UTF-8 byte counts.
Tests: new tests/devices/{conftest (FakeRedis double), test_broker_cross_process,
test_broker_race, test_submit_output_audit}; drain/cleanup/ticket tests rewritten
for the Redis contract. The race tests fail against the pre-fix code. ruff clean;
device + tool-executor suites green.
* fix: log instead of silently passing on failed-dispatch cleanup
Addresses the code-quality lint on the best-effort hash delete in
dispatch_invocation's failure path: replace the bare `except: pass` with a
logger.debug carrying the invocation_id. No behavior change — cleanup stays
best-effort and still returns a failed Invocation.
202 lines
6.3 KiB
Python
202 lines
6.3 KiB
Python
"""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
|