Coverage for src/local_deep_research/web/research_state.py: 98%
167 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Thread-safe global state management.
4Wraps two module-level dicts (`_active_research`, `_termination_flags`)
5with accessor functions protected by a single ``threading.RLock``.
6All external code should use the accessor functions instead of touching the
7dicts directly.
8"""
10import threading
11from contextlib import contextmanager
13# ---------------------------------------------------------------------------
14# Internal state — never import these directly from other modules
15# ---------------------------------------------------------------------------
16_active_research: dict[int, dict] = {}
17_termination_flags: dict[int, bool] = {}
19# A single lock protects both dicts. They have strongly correlated lifecycles
20# (entries are created/destroyed together) and operations are fast dict lookups,
21# so a single lock is simpler and eliminates any deadlock risk from lock ordering.
22_lock = threading.RLock()
24# Per-user gate serialising a user's password-change rekey against that user's
25# own newly-starting research threads. Deliberately NOT the process-global
26# ``_lock`` above: the SQLCipher ``PRAGMA rekey`` can take seconds, and holding
27# ``_lock`` across it would freeze EVERY user's research (every log-append /
28# progress-update / status read takes ``_lock``). Keyed by username;
29# ``check_and_start_research`` acquires the gate for the research's owner, and
30# ``change_password`` (web/routers/auth.py) holds it across its "is research
31# active?" check + rekey, so a rekey blocks only the SAME user's new-research
32# registration. A dedicated small lock guards the dict so gate lookups never
33# touch ``_lock`` — keeping the two consistently orderable: a caller always
34# takes a gate BEFORE ``_lock``, never the reverse.
35#
36# INTENTIONALLY NEVER POPPED. This is a per-user mutual-exclusion primitive
37# that may be HELD across a multi-second SQLCipher ``PRAGMA rekey``. A gate
38# that may be held across a long operation cannot be safely removed from its
39# registry: if a logout / idle sweep popped a gate while one tab holds it
40# across the rekey, a concurrent same-user ``check_and_start_research`` would
41# find no entry, CREATE A SECOND ``threading.Lock`` instance, acquire it
42# uncontended, and start a worker that writes the user's DB concurrently with
43# the in-flight rekey — defeating the exact exclusion this gate exists to
44# provide. Any lookup-then-acquire scheme has this lookup->acquire gap, so the
45# only correct policy is to never remove the gate. It also serialises the
46# per-user capacity count -> DB claim used by fresh requests and queue replay;
47# an ``RLock`` is required because the outer admission section can reach
48# ``check_and_start_research``, which takes the same gate before spawning.
49# The dict is bounded by the user population (one small lock per distinct
50# username), so it does not grow unbounded.
53class _ReentrantUserLock:
54 """Small ``RLock`` adapter with ``locked()`` on supported Python 3.12.
56 Native ``threading.RLock.locked`` was only added in Python 3.14, while the
57 project and CI support 3.12. Tracking acquisition depth preserves the
58 existing lock-introspection contract without relying on private runtime
59 methods.
60 """
62 def __init__(self) -> None:
63 self._lock = threading.RLock()
64 self._state_lock = threading.Lock()
65 self._depth = 0
67 def acquire(self, blocking: bool = True, timeout: float = -1) -> bool:
68 if timeout == -1: 68 ↛ 71line 68 didn't jump to line 71 because the condition on line 68 was always true
69 acquired = self._lock.acquire(blocking)
70 else:
71 acquired = self._lock.acquire(blocking, timeout)
72 if acquired:
73 with self._state_lock:
74 self._depth += 1
75 return acquired
77 def release(self) -> None:
78 self._lock.release()
79 with self._state_lock:
80 self._depth -= 1
82 def locked(self) -> bool:
83 with self._state_lock:
84 return self._depth > 0
86 def __enter__(self):
87 self.acquire()
88 return self
90 def __exit__(self, exc_type, exc_value, traceback) -> None:
91 del exc_type, exc_value, traceback
92 self.release()
95_user_research_start_gates: dict[str, _ReentrantUserLock] = {}
96_user_research_start_gates_lock = threading.Lock()
99# ===================================================================
100# active_research accessors
101# ===================================================================
104def is_research_active(research_id):
105 """Return True if *research_id* is in the active-research dict."""
106 with _lock:
107 return research_id in _active_research
110def get_active_research_ids():
111 """Return a list of all active research IDs (snapshot)."""
112 with _lock:
113 return list(_active_research.keys())
116def get_active_research_snapshot(research_id):
117 """Return a safe snapshot of an active-research entry, or ``None``.
119 The returned dict contains only serialisable fields (``thread`` is
120 excluded). The ``log`` list is shallow-copied — individual entries
121 are never mutated after creation, so this is safe.
122 """
123 with _lock:
124 entry = _active_research.get(research_id)
125 if entry is None:
126 return None
127 return {
128 "progress": entry.get("progress", 0),
129 "status": entry.get("status"),
130 "log": list(entry.get("log", [])),
131 "settings": dict(s)
132 if (s := entry.get("settings")) is not None
133 else None,
134 }
137def get_research_field(research_id, field, default=None):
138 """Return a single field from an active-research entry.
140 For mutable fields (``list``, ``dict``) a shallow copy is returned so
141 callers cannot accidentally mutate shared state.
142 """
143 with _lock:
144 entry = _active_research.get(research_id)
145 if entry is None:
146 return default
147 value = entry.get(field, default)
148 # We explicitly check for list/dict rather than using copy.copy()
149 # because entries may contain threading.Thread objects —
150 # copy.copy(Thread) creates a broken shallow copy sharing the same
151 # OS thread ident. Scalars (int, str, bool) are immutable and safe.
152 # For Thread objects, access them via iter_active_research() snapshots.
153 if isinstance(value, list):
154 return list(value)
155 if isinstance(value, dict):
156 return dict(value)
157 return value
160def set_active_research(research_id, data):
161 """Insert or replace the active-research entry for *research_id*."""
162 with _lock:
163 _active_research[research_id] = data
166def _get_user_research_start_gate(username: str) -> _ReentrantUserLock:
167 """Return the per-user research-start gate, creating it on first use."""
168 with _user_research_start_gates_lock:
169 gate = _user_research_start_gates.get(username)
170 if gate is None:
171 gate = _ReentrantUserLock()
172 _user_research_start_gates[username] = gate
173 return gate
176def get_user_research_start_lock(username: str):
177 """Return the stable lock shared by admission and worker registration."""
178 return _get_user_research_start_gate(username)
181@contextmanager
182def user_research_start_gate(username):
183 """Hold the per-user research-start gate across a check-then-act sequence.
185 ``check_and_start_research`` acquires this same per-user gate before it
186 registers and starts a new research thread for a user. Fresh HTTP starts
187 and persisted queue replay also hold it across their live-count -> DB-claim
188 sections, so neither can consume the same last user slot. A caller that
189 holds the gate across its own "is this user researching?" check and a
190 password-change rekey likewise closes that TOCTOU window for that user.
192 Unlike the process-global ``_lock`` this does NOT block other users'
193 research — essential because the rekey it guards can take seconds and
194 ``_lock`` is taken on every progress/log update.
196 ``username`` of ``None`` yields without acquiring anything: there is no user
197 to gate, and real research always carries an owning username.
198 """
199 if username is None:
200 yield
201 return
202 with _get_user_research_start_gate(username):
203 yield
206def check_and_start_research(research_id, data) -> bool:
207 """Atomically register a research entry iff no live thread exists.
209 If ``_active_research[research_id]`` already holds an entry whose
210 ``thread`` is alive, returns ``False`` without mutating state and
211 without calling ``.start()``. Otherwise, starts ``data['thread']``
212 and writes *data* into the active-research dict, then returns ``True``.
214 The entire check-and-start is done under the single shared lock, so
215 two concurrent callers with the same *research_id* cannot both pass
216 the liveness check and end up with two live threads for the same ID.
218 Registration is additionally wrapped in the research owner's per-user
219 ``user_research_start_gate``, so a new thread cannot register (and start
220 writing to the pre-rekey engine) while that user's password-change rekey
221 is in flight — see ``change_password`` in web/routers/auth.py.
222 """
223 thread = data.get("thread") if isinstance(data, dict) else None
224 if thread is None:
225 raise ValueError("data must contain a 'thread' entry")
226 # Owning username drives the per-user rekey gate below. Every real research
227 # start passes settings["username"] (via start_research_process); a
228 # missing/None username degrades to no gating (nothing to serialise
229 # against, since the rekey gate is keyed by username).
230 username = None
231 if isinstance(data, dict): 231 ↛ 238line 231 didn't jump to line 238 because the condition on line 231 was always true
232 settings = data.get("settings")
233 if isinstance(settings, dict): 233 ↛ 238line 233 didn't jump to line 238 because the condition on line 233 was always true
234 username = settings.get("username")
235 # Serialise against an in-flight password-change rekey for this user. The
236 # gate is per-user, so this never stalls other users' research; the inner
237 # ``_lock`` still provides the atomic check-and-start for research_id.
238 with user_research_start_gate(username):
239 with _lock:
240 entry = _active_research.get(research_id)
241 if entry is not None:
242 existing = entry.get("thread")
243 if existing is not None and existing.is_alive():
244 return False
245 thread.start()
246 _active_research[research_id] = data
247 return True
250def update_active_research(research_id, **fields):
251 """Update one or more fields on an existing entry.
253 Silently does nothing if *research_id* is not active.
254 """
255 with _lock:
256 entry = _active_research.get(research_id)
257 if entry is not None:
258 entry.update(fields)
261def append_research_log(research_id, log_entry):
262 """Append *log_entry* to the ``log`` list for *research_id*.
264 Silently does nothing if *research_id* is not active.
265 """
266 with _lock:
267 entry = _active_research.get(research_id)
268 if entry is not None:
269 entry.setdefault("log", []).append(log_entry)
272def update_progress_if_higher(research_id, new_progress):
273 """Atomically update bounded progress only if it exceeds current.
275 Returns the resulting progress value, or ``None`` if the research is
276 not active.
277 """
278 return update_progress_and_check_active(research_id, new_progress)[0]
281def remove_active_research(research_id):
282 """Remove *research_id* from the active-research dict (if present)."""
283 with _lock:
284 _active_research.pop(research_id, None)
287def iter_active_research():
288 """Yield ``(research_id, snapshot)`` pairs for all active research.
290 Each *snapshot* is a shallow copy of the entry dict, safe to read
291 outside the lock.
292 """
293 with _lock:
294 items = [
295 (
296 rid,
297 {
298 "progress": entry.get("progress", 0),
299 "status": entry.get("status"),
300 "log": list(entry.get("log", [])),
301 "settings": dict(s)
302 if (s := entry.get("settings")) is not None
303 else None,
304 },
305 )
306 for rid, entry in _active_research.items()
307 ]
308 for rid, data in items:
309 yield rid, data
312def get_active_research_count():
313 """Return the number of active research entries."""
314 with _lock:
315 return len(_active_research)
318def get_usernames_with_active_research() -> set:
319 """Return set of usernames that have research currently running."""
320 with _lock:
321 return {
322 entry.get("settings", {}).get("username")
323 for entry in _active_research.values()
324 if entry.get("settings", {}).get("username")
325 }
328# ===================================================================
329# termination_flags accessors
330# ===================================================================
333def is_termination_requested(research_id):
334 """Return ``True`` if termination was requested for *research_id*."""
335 with _lock:
336 return _termination_flags.get(research_id, False)
339def set_termination_flag(research_id):
340 """Signal that *research_id* should be terminated."""
341 with _lock:
342 _termination_flags[research_id] = True
345def clear_termination_flag(research_id):
346 """Remove the termination flag for *research_id* (if present)."""
347 with _lock:
348 _termination_flags.pop(research_id, None)
351# ===================================================================
352# Compound / cleanup helpers
353# ===================================================================
356def is_research_thread_alive(research_id):
357 """Return ``True`` if the research thread for *research_id* is alive.
359 Returns ``False`` if the research is not active or has no thread.
360 """
361 with _lock:
362 entry = _active_research.get(research_id)
363 if entry is None:
364 return False
365 thread = entry.get("thread")
366 return thread is not None and thread.is_alive()
369def update_progress_and_check_active(research_id, new_progress):
370 """Atomically bound progress to 0..100 and advance it monotonically.
372 Returns ``(progress_value, is_active)`` where *progress_value* is the
373 resulting progress (or ``None`` if not active) and *is_active* indicates
374 whether *research_id* is still in the active-research dict.
375 """
376 with _lock:
377 entry = _active_research.get(research_id)
378 if entry is None:
379 return (None, False)
380 # A strategy can overshoot its phase allocation. Clamp before the
381 # monotonic comparison: otherwise an overshoot becomes a permanent
382 # floor and even the worker's final 100% update cannot correct it.
383 current = min(100, max(0, entry.get("progress", 0)))
384 if new_progress is not None:
385 current = max(current, min(100, max(0, new_progress)))
386 entry["progress"] = current
387 return (current, True)
390def cleanup_research(research_id, *, preserve_termination_flag=False):
391 """Remove a research's process state atomically under the shared lock.
393 The normal worker cleanup removes both registries. The HTTP cancellation
394 path may preserve the flag after removing the active entry so the worker
395 that is still unwinding can observe the stop request at its next
396 checkpoint and perform the final cleanup itself.
397 """
398 with _lock:
399 _active_research.pop(research_id, None)
400 if not preserve_termination_flag:
401 _termination_flags.pop(research_id, None)
404def reclaim_stale_user_active_research(
405 db_session, username, *, grace_cutoff_dt=None, logger=None
406):
407 """Flip ``UserActiveResearch`` rows whose worker thread is dead.
409 Shared between ``research_routes.start_research`` (no grace window)
410 and ``chat.routes.send_message`` (30-second grace window) — both
411 sites historically iterated the same query / status flip / cleanup
412 pattern inline; consolidating here makes the difference (grace vs.
413 no grace) explicit instead of two near-duplicate copies drifting.
415 The helper does NOT commit — callers compose this with surrounding
416 DB writes and commit once at the end. Returns ``True`` if any rows
417 were reclaimed so the caller can decide whether a commit is needed.
419 Args:
420 db_session: open SQLAlchemy session against the user's DB.
421 username: the per-user DB owner (used to scope the query).
422 grace_cutoff_dt: optional ``datetime`` boundary; rows whose
423 ``started_at`` is at-or-after this are skipped (avoids
424 killing a sibling request's just-spawned thread). ``None``
425 disables the grace filter — matches the
426 ``research_routes.start_research`` original behaviour.
427 logger: optional ``loguru``-style logger for the reclaim
428 audit line. Both call sites log at WARNING so operators can
429 trace why an active-research cap was released.
431 Returns:
432 ``True`` if any row was flipped (caller should commit), else
433 ``False``.
434 """
435 # Lazy import to avoid pulling SQLAlchemy at module import time;
436 # this helper is only called from the route paths that already have
437 # the ORM models in scope.
438 from ..constants import ResearchStatus
439 from ..database.models import UserActiveResearch
441 query = db_session.query(UserActiveResearch).filter(
442 UserActiveResearch.username == username,
443 UserActiveResearch.status == ResearchStatus.IN_PROGRESS,
444 )
445 if grace_cutoff_dt is not None:
446 query = query.filter(UserActiveResearch.started_at < grace_cutoff_dt)
448 reclaimed = False
449 for row in query.all():
450 if is_research_thread_alive(row.research_id):
451 continue
452 if logger is not None:
453 logger.warning(
454 "Reclaiming stale UserActiveResearch {short_id}... "
455 "(thread dead) for user {user}",
456 short_id=row.research_id[:8],
457 user=username,
458 )
459 row.status = ResearchStatus.FAILED
460 cleanup_research(row.research_id)
461 reclaimed = True
462 return reclaimed