Coverage for src/local_deep_research/web/routes/globals.py: 99%
124 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-19 23:35 +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
12from sqlalchemy.orm import Session
14# ---------------------------------------------------------------------------
15# Internal state — never import these directly from other modules
16# ---------------------------------------------------------------------------
17_active_research: dict[int, dict] = {}
18_termination_flags: dict[int, bool] = {}
20# A single lock protects both dicts. They have strongly correlated lifecycles
21# (entries are created/destroyed together) and operations are fast dict lookups,
22# so a single lock is simpler and eliminates any deadlock risk from lock ordering.
23_lock = threading.RLock()
26# ===================================================================
27# active_research accessors
28# ===================================================================
31def is_research_active(research_id):
32 """Return True if *research_id* is in the active-research dict."""
33 with _lock:
34 return research_id in _active_research
37def get_active_research_ids():
38 """Return a list of all active research IDs (snapshot)."""
39 with _lock:
40 return list(_active_research.keys())
43def get_active_research_snapshot(research_id):
44 """Return a safe snapshot of an active-research entry, or ``None``.
46 The returned dict contains only serialisable fields (``thread`` is
47 excluded). The ``log`` list is shallow-copied — individual entries
48 are never mutated after creation, so this is safe.
49 """
50 with _lock:
51 entry = _active_research.get(research_id)
52 if entry is None:
53 return None
54 return {
55 "progress": entry.get("progress", 0),
56 "status": entry.get("status"),
57 "log": list(entry.get("log", [])),
58 "settings": dict(s)
59 if (s := entry.get("settings")) is not None
60 else None,
61 }
64def get_research_field(research_id, field, default=None):
65 """Return a single field from an active-research entry.
67 For mutable fields (``list``, ``dict``) a shallow copy is returned so
68 callers cannot accidentally mutate shared state.
69 """
70 with _lock:
71 entry = _active_research.get(research_id)
72 if entry is None:
73 return default
74 value = entry.get(field, default)
75 # We explicitly check for list/dict rather than using copy.copy()
76 # because entries may contain threading.Thread objects —
77 # copy.copy(Thread) creates a broken shallow copy sharing the same
78 # OS thread ident. Scalars (int, str, bool) are immutable and safe.
79 # For Thread objects, access them via iter_active_research() snapshots.
80 if isinstance(value, list):
81 return list(value)
82 if isinstance(value, dict):
83 return dict(value)
84 return value
87def set_active_research(research_id, data):
88 """Insert or replace the active-research entry for *research_id*."""
89 with _lock:
90 _active_research[research_id] = data
93def check_and_start_research(research_id, data) -> bool:
94 """Atomically register a research entry iff no live thread exists.
96 If ``_active_research[research_id]`` already holds an entry whose
97 ``thread`` is alive, returns ``False`` without mutating state and
98 without calling ``.start()``. Otherwise, starts ``data['thread']``
99 and writes *data* into the active-research dict, then returns ``True``.
101 The entire check-and-start is done under the single shared lock, so
102 two concurrent callers with the same *research_id* cannot both pass
103 the liveness check and end up with two live threads for the same ID.
104 """
105 thread = data.get("thread") if isinstance(data, dict) else None
106 if thread is None:
107 raise ValueError("data must contain a 'thread' entry")
108 with _lock:
109 entry = _active_research.get(research_id)
110 if entry is not None:
111 existing = entry.get("thread")
112 if existing is not None and existing.is_alive():
113 return False
114 thread.start()
115 _active_research[research_id] = data
116 return True
119def update_active_research(research_id, **fields):
120 """Update one or more fields on an existing entry.
122 Silently does nothing if *research_id* is not active.
123 """
124 with _lock:
125 entry = _active_research.get(research_id)
126 if entry is not None:
127 entry.update(fields)
130def append_research_log(research_id, log_entry):
131 """Append *log_entry* to the ``log`` list for *research_id*.
133 Silently does nothing if *research_id* is not active.
134 """
135 with _lock:
136 entry = _active_research.get(research_id)
137 if entry is not None:
138 entry.setdefault("log", []).append(log_entry)
141def update_progress_if_higher(research_id, new_progress):
142 """Atomically update progress only if *new_progress* exceeds current.
144 Returns the resulting progress value, or ``None`` if the research is
145 not active.
146 """
147 with _lock:
148 entry = _active_research.get(research_id)
149 if entry is None:
150 return None
151 current = entry.get("progress", 0)
152 if new_progress is not None and new_progress > current:
153 entry["progress"] = new_progress
154 return new_progress
155 return current
158def remove_active_research(research_id):
159 """Remove *research_id* from the active-research dict (if present)."""
160 with _lock:
161 _active_research.pop(research_id, None)
164def iter_active_research():
165 """Yield ``(research_id, snapshot)`` pairs for all active research.
167 Each *snapshot* is a shallow copy of the entry dict, safe to read
168 outside the lock.
169 """
170 with _lock:
171 items = [
172 (
173 rid,
174 {
175 "progress": entry.get("progress", 0),
176 "status": entry.get("status"),
177 "log": list(entry.get("log", [])),
178 "settings": dict(s)
179 if (s := entry.get("settings")) is not None
180 else None,
181 },
182 )
183 for rid, entry in _active_research.items()
184 ]
185 for rid, data in items:
186 yield rid, data
189def get_active_research_count():
190 """Return the number of active research entries."""
191 with _lock:
192 return len(_active_research)
195def get_usernames_with_active_research() -> set:
196 """Return set of usernames that have research currently running."""
197 with _lock:
198 return {
199 entry.get("settings", {}).get("username")
200 for entry in _active_research.values()
201 if entry.get("settings", {}).get("username")
202 }
205# ===================================================================
206# termination_flags accessors
207# ===================================================================
210def is_termination_requested(research_id):
211 """Return ``True`` if termination was requested for *research_id*."""
212 with _lock:
213 return _termination_flags.get(research_id, False)
216def set_termination_flag(research_id):
217 """Signal that *research_id* should be terminated."""
218 with _lock:
219 _termination_flags[research_id] = True
222def clear_termination_flag(research_id):
223 """Remove the termination flag for *research_id* (if present)."""
224 with _lock:
225 _termination_flags.pop(research_id, None)
228# ===================================================================
229# Compound / cleanup helpers
230# ===================================================================
233def is_research_thread_alive(research_id):
234 """Return ``True`` if the research thread for *research_id* is alive.
236 Returns ``False`` if the research is not active or has no thread.
237 """
238 with _lock:
239 entry = _active_research.get(research_id)
240 if entry is None:
241 return False
242 thread = entry.get("thread")
243 return thread is not None and thread.is_alive()
246def update_progress_and_check_active(research_id, new_progress):
247 """Atomically update progress (if higher) and check if research is active.
249 Returns ``(progress_value, is_active)`` where *progress_value* is the
250 resulting progress (or ``None`` if not active) and *is_active* indicates
251 whether *research_id* is still in the active-research dict.
252 """
253 with _lock:
254 entry = _active_research.get(research_id)
255 if entry is None:
256 return (None, False)
257 current = entry.get("progress", 0)
258 if new_progress is not None and new_progress > current:
259 entry["progress"] = new_progress
260 return (new_progress, True)
261 return (current, True)
264def cleanup_research(research_id):
265 """Remove *research_id* from both active_research and termination_flags
266 atomically under the single shared lock.
267 """
268 with _lock:
269 _active_research.pop(research_id, None)
270 _termination_flags.pop(research_id, None)
273def reclaim_stale_user_active_research(
274 db_session: Session, username, *, grace_cutoff_dt=None, logger=None
275):
276 """Flip ``UserActiveResearch`` rows whose worker thread is dead.
278 Shared between ``research_routes.start_research`` (no grace window)
279 and ``chat.routes.send_message`` (30-second grace window) — both
280 sites historically iterated the same query / status flip / cleanup
281 pattern inline; consolidating here makes the difference (grace vs.
282 no grace) explicit instead of two near-duplicate copies drifting.
284 The helper does NOT commit — callers compose this with surrounding
285 DB writes and commit once at the end. Returns ``True`` if any rows
286 were reclaimed so the caller can decide whether a commit is needed.
288 Args:
289 db_session: open SQLAlchemy session against the user's DB.
290 username: the per-user DB owner (used to scope the query).
291 grace_cutoff_dt: optional ``datetime`` boundary; rows whose
292 ``started_at`` is at-or-after this are skipped (avoids
293 killing a sibling request's just-spawned thread). ``None``
294 disables the grace filter — matches the
295 ``research_routes.start_research`` original behaviour.
296 logger: optional ``loguru``-style logger for the reclaim
297 audit line. Both call sites log at WARNING so operators can
298 trace why an active-research cap was released.
300 Returns:
301 ``True`` if any row was flipped (caller should commit), else
302 ``False``.
303 """
304 # Lazy import to avoid pulling SQLAlchemy at module import time;
305 # this helper is only called from the route paths that already have
306 # the ORM models in scope.
307 from ...constants import ResearchStatus
308 from ...database.models import UserActiveResearch
310 query = db_session.query(UserActiveResearch).filter(
311 UserActiveResearch.username == username,
312 UserActiveResearch.status == ResearchStatus.IN_PROGRESS,
313 )
314 if grace_cutoff_dt is not None:
315 query = query.filter(UserActiveResearch.started_at < grace_cutoff_dt)
317 reclaimed = False
318 for row in query.all():
319 if is_research_thread_alive(row.research_id):
320 continue
321 if logger is not None: 321 ↛ 322line 321 didn't jump to line 322 because the condition on line 321 was never true
322 logger.warning(
323 "Reclaiming stale UserActiveResearch {short_id}... "
324 "(thread dead) for user {user}",
325 short_id=row.research_id[:8],
326 user=username,
327 )
328 row.status = ResearchStatus.FAILED
329 cleanup_research(row.research_id)
330 reclaimed = True
331 return reclaimed