Coverage for src/local_deep_research/web/research_state.py: 98%
175 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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 progress only if *new_progress* exceeds current.
275 Returns the resulting progress value, or ``None`` if the research is
276 not active.
277 """
278 with _lock:
279 entry = _active_research.get(research_id)
280 if entry is None:
281 return None
282 current = entry.get("progress", 0)
283 if new_progress is not None and new_progress > current:
284 entry["progress"] = new_progress
285 return new_progress
286 return current
289def remove_active_research(research_id):
290 """Remove *research_id* from the active-research dict (if present)."""
291 with _lock:
292 _active_research.pop(research_id, None)
295def iter_active_research():
296 """Yield ``(research_id, snapshot)`` pairs for all active research.
298 Each *snapshot* is a shallow copy of the entry dict, safe to read
299 outside the lock.
300 """
301 with _lock:
302 items = [
303 (
304 rid,
305 {
306 "progress": entry.get("progress", 0),
307 "status": entry.get("status"),
308 "log": list(entry.get("log", [])),
309 "settings": dict(s)
310 if (s := entry.get("settings")) is not None
311 else None,
312 },
313 )
314 for rid, entry in _active_research.items()
315 ]
316 for rid, data in items:
317 yield rid, data
320def get_active_research_count():
321 """Return the number of active research entries."""
322 with _lock:
323 return len(_active_research)
326def get_usernames_with_active_research() -> set:
327 """Return set of usernames that have research currently running."""
328 with _lock:
329 return {
330 entry.get("settings", {}).get("username")
331 for entry in _active_research.values()
332 if entry.get("settings", {}).get("username")
333 }
336# ===================================================================
337# termination_flags accessors
338# ===================================================================
341def is_termination_requested(research_id):
342 """Return ``True`` if termination was requested for *research_id*."""
343 with _lock:
344 return _termination_flags.get(research_id, False)
347def set_termination_flag(research_id):
348 """Signal that *research_id* should be terminated."""
349 with _lock:
350 _termination_flags[research_id] = True
353def clear_termination_flag(research_id):
354 """Remove the termination flag for *research_id* (if present)."""
355 with _lock:
356 _termination_flags.pop(research_id, None)
359# ===================================================================
360# Compound / cleanup helpers
361# ===================================================================
364def is_research_thread_alive(research_id):
365 """Return ``True`` if the research thread for *research_id* is alive.
367 Returns ``False`` if the research is not active or has no thread.
368 """
369 with _lock:
370 entry = _active_research.get(research_id)
371 if entry is None:
372 return False
373 thread = entry.get("thread")
374 return thread is not None and thread.is_alive()
377def update_progress_and_check_active(research_id, new_progress):
378 """Atomically update progress (if higher) and check if research is active.
380 Returns ``(progress_value, is_active)`` where *progress_value* is the
381 resulting progress (or ``None`` if not active) and *is_active* indicates
382 whether *research_id* is still in the active-research dict.
383 """
384 with _lock:
385 entry = _active_research.get(research_id)
386 if entry is None:
387 return (None, False)
388 current = entry.get("progress", 0)
389 if new_progress is not None and new_progress > current:
390 entry["progress"] = new_progress
391 return (new_progress, True)
392 return (current, True)
395def cleanup_research(research_id, *, preserve_termination_flag=False):
396 """Remove a research's process state atomically under the shared lock.
398 The normal worker cleanup removes both registries. The HTTP cancellation
399 path may preserve the flag after removing the active entry so the worker
400 that is still unwinding can observe the stop request at its next
401 checkpoint and perform the final cleanup itself.
402 """
403 with _lock:
404 _active_research.pop(research_id, None)
405 if not preserve_termination_flag:
406 _termination_flags.pop(research_id, None)
409def reclaim_stale_user_active_research(
410 db_session, username, *, grace_cutoff_dt=None, logger=None
411):
412 """Flip ``UserActiveResearch`` rows whose worker thread is dead.
414 Shared between ``research_routes.start_research`` (no grace window)
415 and ``chat.routes.send_message`` (30-second grace window) — both
416 sites historically iterated the same query / status flip / cleanup
417 pattern inline; consolidating here makes the difference (grace vs.
418 no grace) explicit instead of two near-duplicate copies drifting.
420 The helper does NOT commit — callers compose this with surrounding
421 DB writes and commit once at the end. Returns ``True`` if any rows
422 were reclaimed so the caller can decide whether a commit is needed.
424 Args:
425 db_session: open SQLAlchemy session against the user's DB.
426 username: the per-user DB owner (used to scope the query).
427 grace_cutoff_dt: optional ``datetime`` boundary; rows whose
428 ``started_at`` is at-or-after this are skipped (avoids
429 killing a sibling request's just-spawned thread). ``None``
430 disables the grace filter — matches the
431 ``research_routes.start_research`` original behaviour.
432 logger: optional ``loguru``-style logger for the reclaim
433 audit line. Both call sites log at WARNING so operators can
434 trace why an active-research cap was released.
436 Returns:
437 ``True`` if any row was flipped (caller should commit), else
438 ``False``.
439 """
440 # Lazy import to avoid pulling SQLAlchemy at module import time;
441 # this helper is only called from the route paths that already have
442 # the ORM models in scope.
443 from ..constants import ResearchStatus
444 from ..database.models import UserActiveResearch
446 query = db_session.query(UserActiveResearch).filter(
447 UserActiveResearch.username == username,
448 UserActiveResearch.status == ResearchStatus.IN_PROGRESS,
449 )
450 if grace_cutoff_dt is not None:
451 query = query.filter(UserActiveResearch.started_at < grace_cutoff_dt)
453 reclaimed = False
454 for row in query.all():
455 if is_research_thread_alive(row.research_id):
456 continue
457 if logger is not None:
458 logger.warning(
459 "Reclaiming stale UserActiveResearch {short_id}... "
460 "(thread dead) for user {user}",
461 short_id=row.research_id[:8],
462 user=username,
463 )
464 row.status = ResearchStatus.FAILED
465 cleanup_research(row.research_id)
466 reclaimed = True
467 return reclaimed