Coverage for src/local_deep_research/web/queue/processor_v2.py: 93%
666 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2Queue processor v2 - uses encrypted user databases instead of service.db
3Supports both direct execution and queue modes.
4"""
6import threading
7import time
8import uuid
9from dataclasses import dataclass
10from enum import Enum, auto
11from typing import Any, Dict, Optional
13from loguru import logger
14from sqlalchemy.orm import Session
16from ...constants import ResearchStatus
17from ...database.encrypted_db import db_manager
18from ...database.models import (
19 QueuedResearch,
20 ResearchHistory,
21 TaskMetadata,
22 UserActiveResearch,
23)
24from ...database.queue_service import UserQueueService
25from ...database.session_context import get_user_db_session
26from ...database.session_passwords import session_password_store
27from ...exceptions import (
28 DuplicateResearchError,
29 InvalidQueuedResearchOverridesError,
30 SystemAtCapacityError,
31)
32from ...security.log_sanitizer import redact_secrets
33from ...notifications.queue_helpers import (
34 send_research_completed_notification_from_session,
35 send_research_failed_notification_from_session,
36)
37from .lifecycle_cleanup import (
38 cleanup_queued_research_state,
39 reconcile_research_queue_status,
40)
41from ..routes.research_validation import (
42 validate_research_query_length,
43 validate_search_overrides,
44)
45from ..research_state import get_user_research_start_lock
46from ..services.research_service import (
47 clamp_user_max_concurrent,
48 run_research_process,
49 start_research_process,
50)
52# Retry configuration constants for notification database queries
53MAX_RESEARCH_LOOKUP_RETRIES = 3
54INITIAL_RESEARCH_LOOKUP_DELAY = 0.5 # seconds
55RETRY_BACKOFF_MULTIPLIER = 2
57# Give up on a queued research after this many consecutive spawn failures.
58# Each failure leaves is_processing=False so the next loop tick retries.
59SPAWN_RETRY_LIMIT = 3
62@dataclass(frozen=True, slots=True)
63class QueueSweepResult:
64 can_dispatch: bool
65 cleaned_ids: frozenset[str]
68class _DirectStartOutcome(Enum):
69 """Outcome of the direct-start attempt used by the queue handoff.
71 ``notify_research_queued`` must know whether it should return because the
72 research is running/terminal, or fall through to create queue metadata and
73 register the user for a later dispatch. A bare ``None`` return cannot make
74 that distinction and previously stranded rows on setup and global-capacity
75 failures.
76 """
78 STARTED = auto()
79 ALREADY_RUNNING = auto()
80 QUEUE_FALLBACK = auto()
81 TERMINAL_FAILURE = auto()
84class QueueProcessorV2:
85 """
86 Processes queued researches using encrypted user databases.
87 This replaces the service.db approach.
88 """
90 def __init__(self, check_interval=10):
91 """
92 Initialize the queue processor.
94 Args:
95 check_interval: How often to check for work (seconds)
96 """
97 self.check_interval = check_interval
98 self.running = False
99 self.thread = None
100 self._loop_iteration = 0
101 # Wakes the loop out of its inter-iteration wait so stop() returns
102 # in milliseconds instead of blocking for up to check_interval.
103 # The test suite stops the processor after every app test
104 # (tests/conftest.py reset_singletons), so a blocking stop() adds
105 # ~check_interval seconds of teardown to each of those tests.
106 self._stop_event = threading.Event()
108 # Per-user settings will be retrieved from each user's database
109 # when processing their queue using SettingsManager
110 logger.info(
111 "Queue processor v2 initialized - will use per-user settings from SettingsManager"
112 )
114 # Track which users we should check
115 self._users_to_check: set[tuple[str, str]] = set()
116 self._users_lock = threading.Lock()
118 # Track pending operations from background threads
119 self.pending_operations = {}
120 self._pending_operations_lock = threading.Lock()
122 # Compatibility view of the shared per-user admission locks. The
123 # canonical registry lives in research_state so fresh HTTP starts and
124 # queue replay use the SAME lock for count -> DB claim; a queue-local
125 # lock would still let those two entry points race each other.
126 #
127 # Size is bounded by distinct-usernames-ever-seen, not by session
128 # count. Entries are tiny (~40 bytes) and deliberately retain stable
129 # identity for the process lifetime; see pop_user_critical_lock().
130 self._user_critical_locks: Dict[str, threading.Lock] = {}
131 self._user_critical_locks_lock = threading.Lock()
133 # Count consecutive spawn failures per research_id. Entries are
134 # popped on success or after hitting SPAWN_RETRY_LIMIT (then the
135 # research is marked FAILED). In-memory is sufficient: a restart
136 # resets the counter and the research gets a fresh N retries,
137 # which is the desired behavior if the underlying system issue
138 # (thread pool, memory) cleared.
139 # Access is guarded by _spawn_retry_counts_lock because the
140 # increment path is a read-modify-write and the loop and direct
141 # request paths can interleave.
142 self._spawn_retry_counts: dict[str, int] = {}
143 self._spawn_retry_counts_lock = threading.Lock()
145 def _get_user_critical_lock(self, username: str) -> threading.Lock:
146 """Get (or lazily create) the per-user lock used to serialise the
147 count-active-and-start-direct critical section for a given user.
148 """
149 lock = get_user_research_start_lock(username)
150 with self._user_critical_locks_lock:
151 # Keep the compatibility/introspection view synchronized with the
152 # canonical research_state registry (tests may snapshot/restore
153 # that registry between cases).
154 self._user_critical_locks[username] = lock
155 return lock
157 def pop_user_critical_lock(self, username: str) -> None:
158 """Compatibility no-op: per-user critical locks keep stable identity.
160 A caller receives the lock from ``_get_user_critical_lock`` before it
161 acquires it. Removing even an apparently-unlocked entry therefore has
162 an unavoidable lookup-to-acquire race: another request can retain the
163 old lock while a later request creates and acquires a second one. If a
164 held lock is removed, the race is immediate and the count-to-start
165 critical section is no longer serialised.
167 Keep one tiny lock per distinct username for the process lifetime,
168 matching the deliberately persistent rekey gate in ``research_state``.
169 The method remains as a no-op so existing user-close cleanup callers do
170 not need a special-case import path.
171 """
172 del username
174 def _bump_spawn_retry_count(self, research_id: str) -> int:
175 """Atomically increment and return the spawn-retry counter for
176 ``research_id``. Extracted so tests can exercise the real
177 locked increment path instead of duplicating the lock in the
178 test worker (which would be a tautology).
179 """
180 with self._spawn_retry_counts_lock:
181 attempts = self._spawn_retry_counts.get(research_id, 0) + 1
182 self._spawn_retry_counts[research_id] = attempts
183 return attempts
185 @staticmethod
186 def _commit_with_safe_rollback(db_session: Session, context: str) -> bool:
187 """Commit ``db_session`` with best-effort rollback on failure.
189 Returns ``True`` on success, ``False`` if the commit raised.
190 The failure path logs via ``logger.exception`` and attempts a
191 rollback, itself guarded so a subsequent rollback failure is
192 logged at debug level rather than propagated.
194 Extracted because the ``try: commit / except: log + try:
195 rollback`` idiom repeats at ≥5 sites in this module; inlining
196 hides the defensive structure behind nested ``try`` blocks and
197 makes each callsite longer than the work it describes.
198 """
199 try:
200 db_session.commit()
201 return True
202 except Exception:
203 logger.exception(f"Commit failed: {context}")
204 try:
205 db_session.rollback()
206 except Exception:
207 logger.debug(
208 f"Rollback after commit failure ({context})",
209 exc_info=True,
210 )
211 return False
213 def _delete_queue_row_safely(
214 self, db_session: Session, username: str, research_id: str
215 ) -> None:
216 """Best-effort delete of the ``QueuedResearch`` row for
217 ``(username, research_id)``.
219 Rolls back any pending state first (the session may be in
220 ``PendingRollbackError`` from a failed commit inside
221 ``_start_research``), re-queries the row fresh, deletes it if
222 present, and commits via ``_commit_with_safe_rollback``.
224 Use this for ``DuplicateResearchError`` cleanup where the goal
225 is "drop the queue row regardless of session state." Do NOT
226 use it for paths that need the delete to be atomic with other
227 writes (e.g. the terminal FAILED path bundles
228 ``ResearchHistory.status = FAILED`` with the queue-row delete
229 in a single commit — that stays inline).
230 """
231 try:
232 db_session.rollback()
233 except Exception:
234 logger.debug(
235 f"Rollback before queue-row delete for {research_id}",
236 exc_info=True,
237 )
238 try:
239 fresh_queued = (
240 db_session.query(QueuedResearch)
241 .filter_by(username=username, research_id=research_id)
242 .first()
243 )
244 if fresh_queued:
245 db_session.delete(fresh_queued)
246 self._commit_with_safe_rollback(
247 db_session,
248 f"queue-row delete for research {research_id}",
249 )
250 except Exception:
251 logger.exception(
252 f"Failed to query/delete queue row for {research_id}"
253 )
254 try:
255 db_session.rollback()
256 except Exception:
257 logger.debug(
258 f"Rollback after queue-row delete failure for {research_id}",
259 exc_info=True,
260 )
262 def start(self):
263 """Start the queue processor thread."""
264 if self.running:
265 logger.warning("Queue processor already running")
266 return
268 self.running = True
269 # Re-arm the wait for restart: stop() leaves the event set, and
270 # create_app() restarts this singleton after the test-suite
271 # teardown stops it.
272 self._stop_event.clear()
273 self.thread = threading.Thread(
274 target=self._process_queue_loop, daemon=True
275 )
276 self.thread.start()
277 logger.info("Queue processor v2 started")
279 def stop(self):
280 """Stop the queue processor thread."""
281 self.running = False
282 # Wake the loop immediately rather than waiting for the next
283 # check_interval tick (otherwise shutdown blocks for up to 10s).
284 self._stop_event.set()
285 if self.thread:
286 self.thread.join(timeout=10)
287 logger.info("Queue processor v2 stopped")
289 def notify_user_activity(self, username: str, session_id: str):
290 """
291 Notify that a user has activity and their queue should be checked.
293 Args:
294 username: The username
295 session_id: The Flask session ID (for password access)
296 """
297 with self._users_lock:
298 self._users_to_check.add((username, session_id))
299 logger.debug(f"User {username} added to queue check list")
301 def notify_research_queued(self, username: str, research_id: str, **kwargs):
302 """
303 Notify that a research was queued.
304 In direct mode, this immediately starts the research if slots are available.
305 In queue mode, it adds to the queue.
307 Args:
308 username: The username
309 research_id: The research ID
310 **kwargs: Additional parameters for direct execution (query, mode, etc.)
311 """
312 # Pre-declared so the except handlers below can pass it to
313 # redact_secrets even on paths where it is never assigned.
314 password = None
315 # Check user's queue_mode setting when we have database access
316 if kwargs:
317 session_id = kwargs.get("session_id")
318 if session_id:
319 # Check if we can start it directly
320 password = session_password_store.get_session_password(
321 username, session_id
322 )
323 if password:
324 # TOCTOU guard against post-logout DB resurrection.
325 # Between the read above and open_user_database() below the
326 # user can log out: logout clears the session password store
327 # (auth/routes.py step 2) and THEN closes the DB. Re-opening
328 # here with the captured password would resurrect the user's
329 # decrypted database — and even resume a queued research —
330 # after logout. Re-read the store immediately before the
331 # open; if the credential is gone the user has logged out,
332 # so abort the direct start (``password`` becomes None and
333 # the block below is skipped). Keeping this re-check adjacent
334 # to the open shrinks the residual window to a negligible
335 # remainder — the store clear and the DB open live in
336 # different subsystems and are not jointly lockable without a
337 # fragile cross-subsystem lock. Otherwise adopt the freshest
338 # credential (never the stale captured one) for both the open
339 # and _start_research_directly. On abort we fall through to
340 # the queue-mode path below, which re-reads fresh and no-ops
341 # cleanly when the user is logged out.
342 current_password = (
343 session_password_store.get_session_password(
344 username, session_id
345 )
346 )
347 if current_password is None:
348 logger.info(
349 f"Aborting direct-start DB open for {username}: "
350 "session credential cleared since read (logout?); "
351 "leaving research queued"
352 )
353 password = current_password
354 if password:
355 # The canonical ordering is admission gate -> database.
356 # Fresh starts, logout/rekey, and library initialisation
357 # all follow that order. Opening/querying a pooled session
358 # first and then waiting for the gate can deadlock under
359 # pool exhaustion with a gate holder waiting for a
360 # connection. The inner admission section below is
361 # intentionally retained: the shared gate is reentrant and
362 # that narrower block documents the count -> claim scope.
363 admission_lock = self._get_user_critical_lock(username)
364 admission_lock.acquire()
365 try:
366 # Open database and check settings + active count
367 engine = db_manager.open_user_database(
368 username, password
369 )
370 if engine:
371 with get_user_db_session(username) as db_session:
372 # Get user's settings using SettingsManager
373 from ...settings.manager import SettingsManager
375 settings_manager = SettingsManager(db_session)
377 # Get user's queue_mode setting (env > DB > default)
378 queue_mode = settings_manager.get_setting(
379 "app.queue_mode", "direct"
380 )
382 # Get user's max concurrent setting (env > DB > default),
383 # clamped to the global semaphore ceiling so a stale or
384 # user-inflated value can't monopolize it.
385 max_concurrent = clamp_user_max_concurrent(
386 settings_manager.get_setting(
387 "app.max_concurrent_researches", 3
388 )
389 )
391 logger.debug(
392 f"User {username} settings: queue_mode={queue_mode}, "
393 f"max_concurrent={max_concurrent}"
394 )
396 # Only try direct execution if user has queue_mode="direct"
397 if queue_mode == "direct":
398 # Serialise the count→check→start critical
399 # section at the application layer. Two
400 # concurrent submissions for the same user
401 # must not both observe the same active
402 # count and both start — that would exceed
403 # max_concurrent. A per-user Python lock
404 # gives us that atomicity independent of
405 # the DB isolation level.
406 with self._get_user_critical_lock(username):
407 active_count = (
408 db_session.query(UserActiveResearch)
409 .filter_by(
410 username=username,
411 status=ResearchStatus.IN_PROGRESS,
412 )
413 .count()
414 )
416 if active_count < max_concurrent:
417 # We have slots - start directly!
418 logger.info(
419 f"Direct mode: Starting research {research_id} immediately "
420 f"(active: {active_count}/{max_concurrent})"
421 )
423 # Start the research directly
424 outcome = (
425 self._start_research_directly(
426 username,
427 research_id,
428 password,
429 **kwargs,
430 )
431 )
432 if (
433 outcome
434 is not _DirectStartOutcome.QUEUE_FALLBACK
435 ):
436 return
437 logger.info(
438 "Direct start deferred for "
439 f"{research_id}; queueing for "
440 "a later dispatch"
441 )
442 else:
443 logger.info(
444 "Direct mode: Max concurrent "
445 f"reached ({active_count}/"
446 f"{max_concurrent}), queueing "
447 f"{research_id}"
448 )
449 else:
450 logger.info(
451 f"User {username} has queue_mode={queue_mode}, "
452 f"queueing research {research_id}"
453 )
454 except Exception as e:
455 # ``password`` is in scope — drop the traceback
456 # chain and redact str(e) so the SQLCipher master
457 # password can't leak via diagnose=True frame
458 # locals (see the generic handler in
459 # _start_research_directly for the full rationale).
460 safe_msg = redact_secrets(str(e), password)
461 logger.warning(
462 f"Error in direct execution for {username}: {safe_msg}"
463 )
464 finally:
465 admission_lock.release()
467 # Fall back to queue mode (or if direct mode failed)
468 try:
469 with get_user_db_session(username) as session:
470 queue_service = UserQueueService(session)
471 queue_service.add_task_metadata(
472 task_id=research_id,
473 task_type="research",
474 priority=0,
475 )
476 logger.info(
477 f"Research {research_id} queued for user {username}"
478 )
480 # Register the user with the processing loop. The loop only
481 # dispatches users present in _users_to_check; under Flask that
482 # set was fed by a before_request hook on every authenticated
483 # request (deleted in the FastAPI migration), so without this
484 # call a queued research would never start. The loop keeps
485 # re-checking a registered user every tick until their queue
486 # drains, so this single registration also covers the
487 # at-capacity case: when a running research finishes, the next
488 # tick dispatches the queued one.
489 session_id = kwargs.get("session_id")
490 if session_id:
491 self.notify_user_activity(username, session_id)
492 else:
493 logger.warning(
494 f"Research {research_id} queued for {username} without "
495 "a session_id — the processing loop cannot pick it up "
496 "until the user logs in again"
497 )
498 except Exception as e:
499 # ``password`` may be bound above — same redaction rationale.
500 safe_msg = redact_secrets(str(e), password)
501 logger.warning(
502 f"Failed to update queue status for {username}: {safe_msg}"
503 )
505 def _start_research_directly(
506 self, username: str, research_id: str, password: str, **kwargs
507 ) -> _DirectStartOutcome:
508 """
509 Start a research directly without queueing.
511 Args:
512 username: The username
513 research_id: The research ID
514 password: The user's password
515 **kwargs: Research parameters (query, mode, settings, etc.)
516 """
517 query = kwargs.get("query")
518 mode = kwargs.get("mode")
519 settings_snapshot = kwargs.get("settings_snapshot", {})
521 # Create active research record
522 active_record_persisted = False
523 try:
524 with get_user_db_session(username) as db_session:
525 active_record = UserActiveResearch(
526 username=username,
527 research_id=research_id,
528 status=ResearchStatus.IN_PROGRESS,
529 thread_id="pending",
530 settings_snapshot=settings_snapshot,
531 )
532 db_session.add(active_record)
533 db_session.commit()
534 active_record_persisted = True
536 # Update task status if it exists
537 queue_service = UserQueueService(db_session)
538 queue_service.update_task_status(research_id, "processing")
539 except Exception as e:
540 # ``password`` is a parameter of this method — drop the
541 # traceback chain and redact str(e) (full rationale at the
542 # generic handler below).
543 safe_msg = redact_secrets(str(e), password)
544 logger.warning(
545 f"Failed to create active research record for {research_id}: {safe_msg}"
546 )
547 # The first commit lands before TaskMetadata is transitioned. If
548 # that later update failed, remove the now-stale active row before
549 # asking the caller to use the normal queue fallback.
550 if active_record_persisted:
551 try:
552 with get_user_db_session(username) as db_session:
553 active_record = (
554 db_session.query(UserActiveResearch)
555 .filter_by(
556 username=username, research_id=research_id
557 )
558 .first()
559 )
560 if active_record: 560 ↛ 562line 560 didn't jump to line 562 because the condition on line 560 was always true
561 db_session.delete(active_record)
562 research_row = (
563 db_session.query(ResearchHistory)
564 .filter_by(id=research_id)
565 .first()
566 )
567 if research_row: 567 ↛ 569line 567 didn't jump to line 569 because the condition on line 567 was always true
568 research_row.status = ResearchStatus.QUEUED
569 db_session.commit()
570 except Exception as cleanup_error:
571 safe_msg = redact_secrets(str(cleanup_error), password)
572 logger.warning(
573 "Cleanup after direct-start setup failure failed for "
574 f"{research_id}: {safe_msg}"
575 )
576 return _DirectStartOutcome.QUEUE_FALLBACK
578 # Extract parameters from kwargs
579 model_provider = kwargs.get("model_provider")
580 model = kwargs.get("model")
581 custom_endpoint = kwargs.get("custom_endpoint")
582 search_engine = kwargs.get("search_engine")
584 # Start the research process
585 try:
586 research_thread = start_research_process(
587 research_id,
588 query,
589 mode,
590 run_research_process,
591 username=username,
592 user_password=password,
593 model_provider=model_provider,
594 model=model,
595 custom_endpoint=custom_endpoint,
596 search_engine=search_engine,
597 max_results=kwargs.get("max_results"),
598 time_period=kwargs.get("time_period"),
599 iterations=kwargs.get("iterations"),
600 questions_per_iteration=kwargs.get("questions_per_iteration"),
601 strategy=kwargs.get("strategy", "source-based"),
602 settings_snapshot=settings_snapshot,
603 )
605 # Update thread ID
606 try:
607 with get_user_db_session(username) as db_session:
608 active_record = (
609 db_session.query(UserActiveResearch)
610 .filter_by(username=username, research_id=research_id)
611 .first()
612 )
613 if active_record: 613 ↛ 623line 613 didn't jump to line 623
614 active_record.thread_id = str(research_thread.ident)
615 db_session.commit()
616 except Exception as e:
617 # ``password`` is in scope — same redaction rationale.
618 safe_msg = redact_secrets(str(e), password)
619 logger.warning(
620 f"Failed to update thread ID for {research_id}: {safe_msg}"
621 )
623 logger.info(
624 f"Direct execution: Started research {research_id} for user {username} "
625 f"in thread {research_thread.ident}"
626 )
627 return _DirectStartOutcome.STARTED
629 except DuplicateResearchError:
630 # A live thread already owns this research_id. Do NOT delete
631 # the UserActiveResearch row or mark ResearchHistory FAILED —
632 # that state belongs to the live thread, and mutating it
633 # would terminate a running research from the user's
634 # perspective while it keeps executing. Same contract as the
635 # queue processor's dedicated dup branch (#3506).
636 logger.warning(
637 f"Duplicate live thread detected for {research_id} "
638 "in direct mode; leaving state intact"
639 )
640 return _DirectStartOutcome.ALREADY_RUNNING
641 except SystemAtCapacityError:
642 # System at concurrent-research capacity in the direct-execution
643 # path. Roll back the IN_PROGRESS active row and mark history
644 # back to QUEUED so the queue processor can pick it up later.
645 logger.info(
646 f"Direct execution hit capacity for {research_id}; re-queueing"
647 )
648 try:
649 with get_user_db_session(username) as db_session:
650 active_record = (
651 db_session.query(UserActiveResearch)
652 .filter_by(username=username, research_id=research_id)
653 .first()
654 )
655 if active_record: 655 ↛ 657line 655 didn't jump to line 657 because the condition on line 655 was always true
656 db_session.delete(active_record)
657 research_row = (
658 db_session.query(ResearchHistory)
659 .filter_by(id=research_id)
660 .first()
661 )
662 if research_row: 662 ↛ 664line 662 didn't jump to line 664 because the condition on line 662 was always true
663 research_row.status = ResearchStatus.QUEUED
664 db_session.commit()
665 except Exception as e:
666 # ``password`` is in scope — same redaction rationale.
667 safe_msg = redact_secrets(str(e), password)
668 logger.warning(
669 f"Cleanup after capacity reject failed for "
670 f"{research_id}; the stale UserActiveResearch row is "
671 f"recovered by reclaim_stale_user_active_research: {safe_msg}"
672 )
673 # Let notify_research_queued create TaskMetadata through its one
674 # canonical queue-fallback path and register (username, session_id)
675 # with the processing loop. Handling metadata here used to make
676 # the row look queued while silently omitting that registration.
677 return _DirectStartOutcome.QUEUE_FALLBACK
678 except Exception as e:
679 # ``password`` is in lexical scope (function parameter,
680 # passed through to ``start_research_process``). A
681 # SQLAlchemy / requests exception from anywhere in
682 # ``start_research_process`` could carry frame locals that
683 # include the SQLCipher master password (which is
684 # unrecoverable — see TRUST.md §5). Drop the traceback chain
685 # and redact str(e) defensively.
686 safe_msg = redact_secrets(str(e), password)
687 logger.warning(
688 f"Failed to start research {research_id} directly: {safe_msg}"
689 )
690 # Clean up the active record AND mark the research terminal
691 # FAILED so the user-visible state matches reality (no running
692 # thread, not IN_PROGRESS). Same contract as the queue
693 # processor's terminal-failure branch (#3481).
694 try:
695 with get_user_db_session(username) as db_session:
696 active_record = (
697 db_session.query(UserActiveResearch)
698 .filter_by(username=username, research_id=research_id)
699 .first()
700 )
701 if active_record:
702 db_session.delete(active_record)
703 research_row = (
704 db_session.query(ResearchHistory)
705 .filter_by(id=research_id)
706 .first()
707 )
708 if research_row:
709 research_row.status = ResearchStatus.FAILED
710 # This direct attempt originated from an already-persisted
711 # queue row. A terminal spawn failure must consume that row
712 # and its TaskMetadata atomically with FAILED, otherwise the
713 # non-QUEUED parent makes it permanently undispatchable
714 # while queue counters continue to report it.
715 cleanup_queued_research_state(
716 db_session, [research_id], include_claimed=True
717 )
718 db_session.commit()
719 except Exception as e2:
720 # ``password`` is in scope — same redaction rationale.
721 safe_msg = redact_secrets(str(e2), password)
722 logger.warning(
723 f"Failed to clean up active research record for {research_id}: {safe_msg}"
724 )
725 return _DirectStartOutcome.TERMINAL_FAILURE
727 @staticmethod
728 def _drop_active_research_row(session, username: str, research_id: str):
729 """Delete the per-user concurrency-cap row for a finished research.
731 ``UserActiveResearch`` exists to count what a user currently has
732 running. Main deleted finished rows from a ``before_request`` hook
733 (``web/auth/cleanup_middleware.cleanup_completed_research``) that
734 sampled ~1% of requests, walked the user's rows, and dropped any
735 whose thread was no longer active. That hook has no successor
736 under FastAPI, and nothing replaced its *delete*: the spawn-failure
737 branches drop their own row, but neither terminal notification
738 below did, so every normally-completed research left its row
739 behind permanently.
741 The visible effect is not a wrong concurrency count —
742 ``reclaim_stale_user_active_research`` re-derives that at the
743 user's next start, so the cap self-heals — but the rows are never
744 removed, only flipped to FAILED, and only for a user who starts
745 another research. A user who finishes and stops accumulates rows
746 with no purge path anywhere in the codebase, and the ones that do
747 get reclaimed are mislabelled FAILED despite having succeeded.
749 Deleting at the terminal notification is strictly better than
750 main's sampled sweep: it is event-driven, so it cannot miss, and
751 it costs one statement on a path that already holds an open
752 session — no extra connection, no per-request overhead.
754 The caller's session context does not commit on exit, so commit
755 here rather than relying on a later write to flush this one.
756 """
757 from ...database.models import UserActiveResearch
759 try:
760 deleted = (
761 session.query(UserActiveResearch)
762 .filter_by(username=username, research_id=research_id)
763 .delete(synchronize_session=False)
764 )
765 if deleted:
766 session.commit()
767 logger.debug(
768 f"Dropped UserActiveResearch row for {research_id}"
769 )
770 except Exception:
771 # Never let bookkeeping break the completion notification the
772 # user is actually waiting on; the row stays and the next
773 # reclaim sweep will pick it up.
774 logger.exception(
775 f"Could not drop UserActiveResearch row for {research_id}"
776 )
777 # Swallowing is not enough: a failed statement leaves the
778 # session raising PendingRollbackError on everything that
779 # follows, and the caller runs
780 # ``send_research_completed_notification_from_session`` on
781 # *this* session on the very next line. Without the rollback
782 # the swallow would hand that notification a dead session --
783 # exactly what the comment above promises not to do. Main did
784 # the same (web/auth/cleanup_middleware.py, five rollbacks).
785 try:
786 session.rollback()
787 except Exception:
788 # A genuinely exhausted pool can fail the rollback too;
789 # this path is user-visible completion, so still no raise.
790 logger.warning(
791 "Could not roll back after a failed UserActiveResearch "
792 f"delete for {research_id}"
793 )
795 def notify_research_completed(
796 self, username: str, research_id: str, user_password: str | None = None
797 ):
798 """
799 Notify that a research completed.
800 Updates the user's queue status in their database.
802 Args:
803 username: The username
804 research_id: The research ID
805 user_password: User password for database access. Required for queue
806 updates and database lookups during notification sending.
807 Optional only because some callers may not have it
808 available, in which case only basic updates occur.
809 """
810 try:
811 # get_user_db_session is already imported at module level (line 19)
812 # It accepts optional password parameter and returns a context manager
813 with get_user_db_session(username, user_password) as session:
814 queue_service = UserQueueService(session)
815 queue_service.update_task_status(
816 research_id, ResearchStatus.COMPLETED
817 )
818 self._drop_active_research_row(session, username, research_id)
819 logger.info(
820 f"Research {research_id} completed for user {username}"
821 )
823 # Send notification using helper from notification module
824 send_research_completed_notification_from_session(
825 username=username,
826 research_id=research_id,
827 db_session=session,
828 )
830 except Exception as e:
831 # ``user_password`` is a parameter of this method — drop the
832 # traceback chain and redact str(e) so the SQLCipher master
833 # password can't leak via diagnose=True frame locals.
834 safe_msg = redact_secrets(str(e), user_password)
835 logger.warning(
836 f"Failed to update completion status for {username}: {safe_msg}"
837 )
839 # Auto-convert research to document in History collection.
840 # Documents only — FAISS indexing is triggered separately by the user
841 # via "Index All" on the History page.
842 from ...research_library.search.services.research_history_indexer import (
843 auto_convert_research,
844 )
846 auto_convert_research(username, research_id, db_password=user_password)
848 def notify_research_failed(
849 self,
850 username: str,
851 research_id: str,
852 error_message: str | None = None,
853 user_password: str | None = None,
854 ):
855 """
856 Notify that a research failed.
857 Updates the user's queue status in their database and sends notification.
859 Args:
860 username: The username
861 research_id: The research ID
862 error_message: Optional error message
863 user_password: User password for database access. Required for queue
864 updates and database lookups during notification sending.
865 Optional only because some callers may not have it
866 available, in which case only basic updates occur.
867 """
868 try:
869 # get_user_db_session is already imported at module level (line 19)
870 # It accepts optional password parameter and returns a context manager
871 with get_user_db_session(username, user_password) as session:
872 queue_service = UserQueueService(session)
873 queue_service.update_task_status(
874 research_id,
875 ResearchStatus.FAILED,
876 error_message=error_message,
877 )
878 self._drop_active_research_row(session, username, research_id)
879 logger.info(
880 f"Research {research_id} failed for user {username}: "
881 f"{error_message}"
882 )
884 # Send notification using helper from notification module
885 send_research_failed_notification_from_session(
886 username=username,
887 research_id=research_id,
888 error_message=error_message or "Unknown error",
889 db_session=session,
890 )
892 except Exception as e:
893 # ``user_password`` is a parameter of this method — same
894 # redaction rationale as notify_research_completed.
895 safe_msg = redact_secrets(str(e), user_password)
896 logger.warning(
897 f"Failed to update failure status for {username}: {safe_msg}"
898 )
900 def _process_queue_loop(self):
901 """Main loop that processes the queue."""
902 while self.running:
903 try:
904 # Get list of users to check (don't clear immediately)
905 with self._users_lock:
906 users_to_check = list(self._users_to_check)
908 # Process each user's queue
909 users_to_remove = []
910 for user_session in users_to_check:
911 try:
912 username, session_id = user_session
913 # _process_user_queue returns True if queue is empty
914 queue_empty = self._process_user_queue(
915 username, session_id
916 )
917 if queue_empty:
918 users_to_remove.append(user_session)
919 except Exception:
920 logger.exception(
921 f"Error processing queue for user {username}"
922 )
923 # Don't remove on error - the _process_user_queue method
924 # determines whether to keep checking based on error type
926 # Only remove users whose queues are now empty
927 with self._users_lock:
928 for user_session in users_to_remove:
929 self._users_to_check.discard(user_session)
931 # Drain pending operations (progress / error updates queued
932 # by background threads via queue_progress_update /
933 # queue_error_update). Without this, FAILED status is never
934 # persisted unless the user happens to make another HTTP
935 # request that drains it from a request handler.
936 self._drain_pending_operations()
938 except Exception:
939 logger.exception("Error in queue processor loop")
940 finally:
941 # Clean up thread-local database session after each iteration.
942 # The loop opens a new session each iteration via get_user_db_session();
943 # closing it returns the connection to the shared QueuePool promptly.
944 try:
945 from ...database.thread_local_session import (
946 cleanup_current_thread,
947 cleanup_dead_threads,
948 )
950 cleanup_current_thread()
951 except Exception:
952 logger.debug(
953 "thread-local cleanup on shutdown", exc_info=True
954 )
956 # Periodic dead-thread credential sweep (every ~60s).
957 # One of three sweep trigger points (fastapi_app's
958 # post-request cleanup, connection_cleanup scheduler,
959 # and here).
960 self._loop_iteration += 1
961 if self._loop_iteration % 6 == 0: # Every ~60s (10s × 6)
962 try:
963 cleanup_dead_threads()
964 except Exception:
965 logger.debug(
966 "periodic dead-thread sweep", exc_info=True
967 )
969 # Event.wait, not time.sleep: stop() must be able to interrupt
970 # this pause, otherwise shutdown blocks for up to
971 # check_interval seconds.
972 self._stop_event.wait(self.check_interval)
974 def _drain_pending_operations(self) -> None:
975 """Per-iteration drain of pending_operations across users.
977 Background threads (research workers) write into pending_operations
978 when they need to mark a research FAILED or update progress but
979 don't have direct DB access (e.g. password lookup failed). We
980 opportunistically drain by username here using the queue
981 processor's own DB access.
982 """
983 with self._pending_operations_lock:
984 usernames = {
985 op["username"]
986 for op in self.pending_operations.values()
987 if op.get("username")
988 }
989 if not usernames:
990 return
992 for username in usernames:
993 try:
994 with get_user_db_session(username) as db_session:
995 n = self.process_pending_operations_for_user(
996 username, db_session
997 )
998 if n: 998 ↛ 992line 998 didn't jump to line 992
999 logger.info(
1000 f"Drained {n} pending operation(s) for {username}"
1001 )
1002 except Exception:
1003 # Likely no password available for this user yet — leave the
1004 # ops queued; eviction will reap them after _PENDING_OPS_TTL.
1005 logger.debug(
1006 f"Could not drain pending ops for {username} this tick",
1007 exc_info=True,
1008 )
1010 def _process_user_queue(self, username: str, session_id: str) -> bool:
1011 """
1012 Process the queue for a specific user.
1014 Args:
1015 username: The username
1016 session_id: The Flask session ID
1018 Returns:
1019 True if the queue is empty, False if there are still items
1020 """
1021 # Get the user's password from session store
1022 password = session_password_store.get_session_password(
1023 username, session_id
1024 )
1025 if not password:
1026 logger.debug(
1027 f"No password available for user {username}, skipping queue check"
1028 )
1029 return True # Remove from checking - session expired
1031 # Keep the process-wide lock order consistent with fresh starts and
1032 # logout/rekey: per-user admission gate before database/session work.
1033 # Taking a pooled session first and then blocking on the gate creates a
1034 # resource cycle when the gate holder itself needs a connection.
1035 admission_lock = self._get_user_critical_lock(username)
1036 admission_lock.acquire()
1038 # Open the user's encrypted database
1039 try:
1040 # First ensure the database is open
1041 engine = db_manager.open_user_database(username, password)
1042 if not engine:
1043 logger.error(f"Failed to open database for user {username}")
1044 return False # Keep checking - could be temporary DB issue
1046 # Get a session and process the queue
1047 with get_user_db_session(username, password) as db_session:
1048 queue_service = UserQueueService(db_session)
1050 sweep_result = self._sweep_missing_parent_queue_rows(
1051 db_session, username
1052 )
1053 if not sweep_result.can_dispatch: 1053 ↛ 1054line 1053 didn't jump to line 1054 because the condition on line 1053 was never true
1054 return False
1056 # Get user's settings using SettingsManager
1057 from ...settings.manager import SettingsManager
1059 settings_manager = SettingsManager(db_session)
1061 # Get user's max concurrent setting (env > DB > default),
1062 # clamped to the global semaphore ceiling so a stale or
1063 # user-inflated value can't monopolize it.
1064 max_concurrent = clamp_user_max_concurrent(
1065 settings_manager.get_setting(
1066 "app.max_concurrent_researches", 3
1067 )
1068 )
1070 # Serialize the live-count -> dispatch handoff with direct
1071 # starts made through this processor. QueueStatus.active_tasks
1072 # only counts tasks that previously travelled through the
1073 # queue, so it misses already-running direct researches and
1074 # cannot enforce the user's actual concurrency limit.
1075 with self._get_user_critical_lock(username):
1076 # A crashed worker can leave an IN_PROGRESS row behind.
1077 # The fresh-submit route performs this same liveness
1078 # reconciliation before counting; queue replay must do so
1079 # as well or a dead row can block the queue forever.
1080 from ..routes.globals import (
1081 reclaim_stale_user_active_research,
1082 )
1084 if reclaim_stale_user_active_research(
1085 db_session, username, logger=logger
1086 ):
1087 db_session.commit()
1089 active_count = (
1090 db_session.query(UserActiveResearch)
1091 .filter_by(
1092 username=username,
1093 status=ResearchStatus.IN_PROGRESS,
1094 )
1095 .count()
1096 )
1097 queue_status = queue_service.get_queue_status() or {
1098 "queued_tasks": 0,
1099 }
1100 available_slots = max_concurrent - active_count
1102 if available_slots <= 0:
1103 # No slots available, but queue might not be empty.
1104 return False # Keep checking
1106 if queue_status["queued_tasks"] == 0:
1107 # Queue is empty.
1108 return True # Remove from checking
1110 logger.info(
1111 f"Processing queue for {username}: "
1112 f"{active_count} active, "
1113 f"{queue_status['queued_tasks']} queued, "
1114 f"{available_slots} slots available"
1115 )
1117 # Process queued researches while retaining the same lock
1118 # used for the capacity decision. Otherwise a direct
1119 # notification could consume the last slot between the
1120 # count and the queued spawn.
1121 self._start_queued_researches(
1122 db_session,
1123 queue_service,
1124 username,
1125 password,
1126 available_slots,
1127 )
1129 # Check if there are still items in queue.
1130 updated_status = queue_service.get_queue_status() or {
1131 "queued_tasks": 0
1132 }
1133 return bool(updated_status["queued_tasks"] == 0)
1135 except Exception as e:
1136 # ``password`` (from the session store above) is in scope —
1137 # drop the traceback chain and redact str(e).
1138 safe_msg = redact_secrets(str(e), password)
1139 logger.warning(
1140 f"Error processing queue for user {username}: {safe_msg}"
1141 )
1142 return False # Keep checking - errors might be temporary
1143 finally:
1144 admission_lock.release()
1146 def _reclaim_stranded_queue_rows(
1147 self, db_session: Session, username: str
1148 ) -> int:
1149 """Reclaim queue rows stranded by a crash or restart.
1151 A row is stranded when ``is_processing=True`` but no live thread
1152 exists in ``_active_research`` for its ``research_id``. This can
1153 happen after a crash/restart between the pre-spawn IN_PROGRESS
1154 commit and the queue-row deletion in ``_start_queued_researches``
1155 — the row is invisible to the normal ``is_processing=False``
1156 query and would never be retried.
1158 Reverts only rows whose parent is not database-backed IN_PROGRESS,
1159 preserving the worker spawn-grace window. Returns the number of
1160 rows reclaimed.
1161 """
1162 from ..routes.globals import is_research_active
1164 stranded = (
1165 db_session.query(QueuedResearch)
1166 .filter_by(username=username, is_processing=True)
1167 .all()
1168 )
1169 reclaimed = 0
1170 for row in stranded:
1171 if is_research_active(row.research_id):
1172 # A legitimate in-flight claim; don't touch.
1173 continue
1174 row.is_processing = False
1175 research = (
1176 db_session.query(ResearchHistory)
1177 .filter_by(id=row.research_id)
1178 .first()
1179 )
1180 status_changed = (
1181 research is not None
1182 and research.status == ResearchStatus.IN_PROGRESS
1183 )
1184 if status_changed:
1185 research.status = ResearchStatus.QUEUED
1186 reclaimed += 1
1187 logger.warning(
1188 f"Reclaimed stranded queue row for research "
1189 f"{row.research_id} (user {username}): no live thread, "
1190 "resetting is_processing=False"
1191 + (" and status=QUEUED" if status_changed else "")
1192 )
1193 if reclaimed:
1194 if not self._commit_with_safe_rollback( 1194 ↛ 1198line 1194 didn't jump to line 1198 because the condition on line 1194 was never true
1195 db_session,
1196 f"reclaim of stranded rows for user {username}",
1197 ):
1198 return 0
1199 return reclaimed
1201 def _sweep_missing_parent_queue_rows(
1202 self, db_session: Session, username: str
1203 ) -> QueueSweepResult:
1204 queued_orphan_ids = {
1205 research_id
1206 for (research_id,) in (
1207 db_session.query(QueuedResearch.research_id)
1208 .outerjoin(
1209 ResearchHistory,
1210 QueuedResearch.research_id == ResearchHistory.id,
1211 )
1212 .filter(
1213 QueuedResearch.username == username,
1214 ResearchHistory.id.is_(None),
1215 )
1216 .all()
1217 )
1218 }
1219 metadata_orphan_ids = {
1220 task_id
1221 for (task_id,) in (
1222 db_session.query(TaskMetadata.task_id)
1223 .outerjoin(
1224 ResearchHistory,
1225 TaskMetadata.task_id == ResearchHistory.id,
1226 )
1227 .filter(
1228 TaskMetadata.task_type == "research",
1229 ResearchHistory.id.is_(None),
1230 )
1231 .all()
1232 )
1233 }
1234 orphaned_research_ids = queued_orphan_ids | metadata_orphan_ids
1235 if not orphaned_research_ids:
1236 if reconcile_research_queue_status(db_session):
1237 if not self._commit_with_safe_rollback( 1237 ↛ 1241line 1237 didn't jump to line 1241 because the condition on line 1237 was never true
1238 db_session,
1239 f"queue status reconciliation for user {username}",
1240 ):
1241 return QueueSweepResult(False, frozenset())
1242 return QueueSweepResult(True, frozenset())
1244 cleanup_result = cleanup_queued_research_state(
1245 db_session,
1246 orphaned_research_ids,
1247 include_claimed=True,
1248 )
1249 if not self._commit_with_safe_rollback(
1250 db_session,
1251 f"missing-parent queue sweep for user {username}",
1252 ):
1253 return QueueSweepResult(False, frozenset())
1255 with self._spawn_retry_counts_lock:
1256 for research_id in cleanup_result.cleaned_ids:
1257 self._spawn_retry_counts.pop(research_id, None)
1258 return QueueSweepResult(True, cleanup_result.cleaned_ids)
1260 def _start_queued_researches(
1261 self,
1262 db_session: Session,
1263 queue_service: UserQueueService,
1264 username: str,
1265 password: str,
1266 available_slots: int,
1267 ):
1268 """Start queued researches up to available slots."""
1269 sweep_result = self._sweep_missing_parent_queue_rows(
1270 db_session, username
1271 )
1272 if not sweep_result.can_dispatch:
1273 return
1275 # Before picking work, reclaim any rows stranded by a prior
1276 # crash — otherwise they are invisible to the is_processing=False
1277 # filter below and would never retry.
1278 self._reclaim_stranded_queue_rows(db_session, username)
1280 # Get queued researches
1281 queued = (
1282 db_session.query(QueuedResearch)
1283 .filter_by(username=username, is_processing=False)
1284 .order_by(QueuedResearch.position, QueuedResearch.id)
1285 .limit(available_slots)
1286 .all()
1287 )
1289 for queued_research in queued:
1290 research_id = queued_research.research_id
1291 try:
1292 # Atomically claim this item by flipping is_processing from
1293 # False to True in a single UPDATE. If another worker has
1294 # already claimed it since our SELECT above, the UPDATE will
1295 # match zero rows and we skip. Under non-IMMEDIATE isolation
1296 # the previous SELECT+assign pattern would race and two
1297 # workers could both process the same queued item.
1298 claimed = (
1299 db_session.query(QueuedResearch)
1300 .filter(
1301 QueuedResearch.id == queued_research.id,
1302 QueuedResearch.is_processing.is_(False),
1303 db_session.query(ResearchHistory.id)
1304 .filter(
1305 ResearchHistory.id == research_id,
1306 ResearchHistory.status == ResearchStatus.QUEUED,
1307 )
1308 .exists(),
1309 )
1310 .update(
1311 {QueuedResearch.is_processing: True},
1312 synchronize_session=False,
1313 )
1314 )
1315 db_session.commit()
1316 if not claimed:
1317 logger.debug(
1318 f"Queued research {research_id} "
1319 f"already claimed by another worker; skipping"
1320 )
1321 continue
1322 # Refresh local object state now that we hold the claim
1323 db_session.refresh(queued_research)
1325 # Update task status
1326 queue_service.update_task_status(research_id, "processing")
1328 # Start the research
1329 self._start_research(
1330 db_session,
1331 username,
1332 password,
1333 queued_research,
1334 )
1336 # Success — clear any prior spawn-failure count and
1337 # remove the queue row.
1338 with self._spawn_retry_counts_lock:
1339 self._spawn_retry_counts.pop(research_id, None)
1340 db_session.delete(queued_research)
1341 db_session.commit()
1343 logger.info(
1344 f"Started queued research {research_id} for user {username}"
1345 )
1347 except DuplicateResearchError:
1348 # Raised by _start_research when a prior attempt's thread
1349 # is still live, OR when the ResearchHistory row is in a
1350 # non-QUEUED state (IN_PROGRESS from a prior attempt's
1351 # successful pre-spawn commit; terminal COMPLETED /
1352 # FAILED / SUSPENDED from a thread that already finished
1353 # and cleaned up). In every case the correct behavior is
1354 # the same: clear the stale queue row and the retry
1355 # counter, and do NOT fall through to the FAILED/notify
1356 # path — that would terminate-status a live thread or
1357 # emit a false failure for a completed one.
1358 logger.warning(
1359 f"Research {research_id} is already started "
1360 "(live thread or non-QUEUED status); clearing stale "
1361 "queue row"
1362 )
1363 with self._spawn_retry_counts_lock:
1364 self._spawn_retry_counts.pop(research_id, None)
1365 self._delete_queue_row_safely(db_session, username, research_id)
1366 continue
1368 except SystemAtCapacityError:
1369 # System hit the global concurrent-research capacity while
1370 # dispatching this queued item. _start_research already
1371 # reset the ResearchHistory row back to QUEUED before
1372 # re-raising. This is a transient condition, NOT a spawn
1373 # failure, so it must NOT count toward SPAWN_RETRY_LIMIT —
1374 # otherwise a busy system would wrongly mark a perfectly
1375 # valid queued research FAILED after a few ticks. Just
1376 # release the processing claim so the next tick retries.
1377 # Mirrors the dedicated handler in _start_research_directly.
1378 logger.info(
1379 f"System at capacity dispatching queued research "
1380 f"{research_id}; leaving queued for next tick"
1381 )
1382 # Revert the queued->processing claim from
1383 # update_task_status("processing") above. The research stays
1384 # queued for the next tick, so its slot must return to
1385 # queued_tasks rather than leaking into active_tasks on
1386 # every capacity-rejected retry.
1387 queue_service.update_task_status(research_id, "queued")
1388 fresh_queued = (
1389 db_session.query(QueuedResearch)
1390 .filter_by(username=username, research_id=research_id)
1391 .first()
1392 )
1393 if fresh_queued: 1393 ↛ 1400line 1393 didn't jump to line 1400 because the condition on line 1393 was always true
1394 fresh_queued.is_processing = False
1395 self._commit_with_safe_rollback(
1396 db_session,
1397 "is_processing reset after capacity reject for "
1398 f"research {research_id}",
1399 )
1400 continue
1402 except InvalidQueuedResearchOverridesError as error:
1403 logger.warning(
1404 "Invalid persisted search overrides for queued research "
1405 f"{research_id}; marking FAILED"
1406 )
1407 with self._spawn_retry_counts_lock:
1408 self._spawn_retry_counts.pop(research_id, None)
1410 research = (
1411 db_session.query(ResearchHistory)
1412 .filter_by(id=research_id)
1413 .first()
1414 )
1415 fresh_queued = (
1416 db_session.query(QueuedResearch)
1417 .filter_by(username=username, research_id=research_id)
1418 .first()
1419 )
1420 if research: 1420 ↛ 1422line 1420 didn't jump to line 1422 because the condition on line 1420 was always true
1421 research.status = ResearchStatus.FAILED
1422 if fresh_queued: 1422 ↛ 1424line 1422 didn't jump to line 1424 because the condition on line 1422 was always true
1423 db_session.delete(fresh_queued)
1424 self._commit_with_safe_rollback(
1425 db_session,
1426 "terminal invalid persisted overrides for research "
1427 f"{research_id}",
1428 )
1429 self.notify_research_failed(
1430 username=username,
1431 research_id=research_id,
1432 error_message=str(error),
1433 user_password=password,
1434 )
1435 continue
1437 except Exception as e:
1438 # ``password`` is a parameter of this method — drop the
1439 # traceback chain and redact str(e).
1440 safe_msg = redact_secrets(str(e), password)
1441 logger.warning(
1442 f"Error starting queued research {research_id}: {safe_msg}"
1443 )
1444 # Session may be in PendingRollbackError state after a
1445 # failed commit inside _start_research.
1446 try:
1447 db_session.rollback()
1448 except Exception as rb_err:
1449 # No exc_info: ``password`` is in this frame and a
1450 # rendered traceback could expose it via
1451 # diagnose=True frame locals.
1452 logger.debug(
1453 "Rollback after start failure: "
1454 f"{redact_secrets(str(rb_err), password)}"
1455 )
1457 attempts = self._bump_spawn_retry_count(research_id)
1459 # Re-query in case rollback expired the ORM object.
1460 fresh_queued = (
1461 db_session.query(QueuedResearch)
1462 .filter_by(username=username, research_id=research_id)
1463 .first()
1464 )
1466 if attempts < SPAWN_RETRY_LIMIT:
1467 # Transient failure — allow the next loop tick to
1468 # retry. _start_research rolls back its own
1469 # IN_PROGRESS write on spawn failure, so the only
1470 # fix-up needed here is resetting is_processing.
1471 logger.warning(
1472 f"Spawn failed for research {research_id} "
1473 f"(attempt {attempts}/{SPAWN_RETRY_LIMIT}), "
1474 "leaving queued for retry"
1475 )
1476 if fresh_queued: 1476 ↛ 1482line 1476 didn't jump to line 1482 because the condition on line 1476 was always true
1477 fresh_queued.is_processing = False
1478 self._commit_with_safe_rollback(
1479 db_session,
1480 f"is_processing reset for research {research_id}",
1481 )
1482 continue
1484 # Exhausted retries — mark terminal FAILED, delete the
1485 # queue row to stop re-dispatch, and notify the user.
1486 # The spawn failure was already logged (redacted) at the
1487 # top of this except block; no need to repeat it here.
1488 logger.warning(
1489 f"Spawn failed for research {research_id} "
1490 f"after {attempts} attempts; marking FAILED"
1491 )
1492 with self._spawn_retry_counts_lock:
1493 self._spawn_retry_counts.pop(research_id, None)
1494 try:
1495 research = (
1496 db_session.query(ResearchHistory)
1497 .filter_by(id=research_id)
1498 .first()
1499 )
1500 if research: 1500 ↛ 1502line 1500 didn't jump to line 1502 because the condition on line 1500 was always true
1501 research.status = ResearchStatus.FAILED
1502 if fresh_queued: 1502 ↛ 1504line 1502 didn't jump to line 1504 because the condition on line 1502 was always true
1503 db_session.delete(fresh_queued)
1504 db_session.commit()
1505 except Exception as e2:
1506 # ``password`` is in scope — same redaction rationale.
1507 safe_msg = redact_secrets(str(e2), password)
1508 logger.warning(
1509 "Failed to persist terminal FAILED state for "
1510 f"research {research_id}: {safe_msg}"
1511 )
1512 try:
1513 db_session.rollback()
1514 except Exception as rb_err:
1515 # No exc_info: same frame-locals rationale as the
1516 # rollback handler above.
1517 logger.debug(
1518 "Rollback after terminal update failure: "
1519 f"{redact_secrets(str(rb_err), password)}"
1520 )
1522 # notify_research_failed opens its own session and
1523 # sends the user notification. Called exactly once
1524 # per research_id because the counter is popped above.
1525 self.notify_research_failed(
1526 username=username,
1527 research_id=research_id,
1528 error_message=(
1529 f"Failed to start research after {attempts} attempts"
1530 ),
1531 user_password=password,
1532 )
1534 def _start_research(
1535 self,
1536 db_session: Session,
1537 username: str,
1538 password: str,
1539 queued_research,
1540 ):
1541 """Start a queued research.
1543 Commits ``ResearchHistory.status = IN_PROGRESS`` BEFORE spawning
1544 the thread. If we did this after, a fast-completing thread
1545 (which opens its own DB session) could write ``COMPLETED`` and
1546 then our post-spawn commit would overwrite that with
1547 ``IN_PROGRESS``, stranding the research as stuck IN_PROGRESS
1548 after it had already finished.
1550 If ``start_research_process`` raises, reset status back to
1551 ``QUEUED`` and re-raise so the caller's 3-strike retry logic
1552 handles it. ``DuplicateResearchError`` is re-raised as-is
1553 because a thread is already running for this research; mutating
1554 status further would be wrong.
1555 """
1556 research_id = queued_research.research_id
1557 research = (
1558 db_session.query(ResearchHistory).filter_by(id=research_id).first()
1559 )
1561 if not research:
1562 raise ValueError(f"Research {research_id} not found")
1564 # Guard against re-entering _start_research on a retry when a
1565 # prior attempt's post-spawn UserActiveResearch commit failed:
1566 # - IN_PROGRESS means the prior thread is (or was) running.
1567 # - COMPLETED/FAILED means the prior thread already finished
1568 # and cleaned itself up out of _active_research, so a bare
1569 # retry would both overwrite the terminal status with
1570 # IN_PROGRESS and then spawn a *second* thread (because
1571 # check_and_start_research sees no live entry), re-running
1572 # the whole research.
1573 # In all three cases the correct behavior is the same: raise
1574 # DuplicateResearchError so the caller's existing dup branch
1575 # deletes the queue row without mutating status or notifying.
1576 if research.status != ResearchStatus.QUEUED:
1577 raise DuplicateResearchError(
1578 f"Research {research_id} is already started "
1579 f"(status={research.status})"
1580 )
1582 # Extract settings
1583 settings_snapshot = queued_research.settings_snapshot or {}
1585 # Handle new vs legacy structure
1586 if (
1587 isinstance(settings_snapshot, dict)
1588 and "submission" in settings_snapshot
1589 ):
1590 submission_params = settings_snapshot.get("submission", {})
1591 complete_settings = settings_snapshot.get("settings_snapshot", {})
1592 allowed_override_keys = (
1593 "model_provider",
1594 "model",
1595 "custom_endpoint",
1596 "search_engine",
1597 "max_results",
1598 "time_period",
1599 "iterations",
1600 "questions_per_iteration",
1601 "strategy",
1602 )
1603 explicit_override_keys = settings_snapshot.get(
1604 "submission_overrides"
1605 )
1606 override_keys = (
1607 explicit_override_keys
1608 if isinstance(explicit_override_keys, list)
1609 else allowed_override_keys
1610 )
1611 runtime_overrides = {
1612 key: submission_params[key]
1613 for key in allowed_override_keys
1614 if key in override_keys and key in submission_params
1615 }
1616 else:
1617 submission_params = settings_snapshot
1618 # A legacy-flat queued row (enqueued by an older version, pre-
1619 # "submission" wrapper) carries no settings_snapshot. Seed the run's
1620 # primary engine from the submitted search_engine so the worker's
1621 # egress build (resolve_run_primary_engine) doesn't fail closed on
1622 # an empty snapshot and refuse the run.
1623 _legacy_engine = submission_params.get("search_engine")
1624 complete_settings = (
1625 {"search.tool": _legacy_engine} if _legacy_engine else {}
1626 )
1627 runtime_overrides = {
1628 "model_provider": submission_params.get("model_provider"),
1629 "model": submission_params.get("model"),
1630 "custom_endpoint": submission_params.get("custom_endpoint"),
1631 "search_engine": submission_params.get("search_engine"),
1632 "max_results": submission_params.get("max_results"),
1633 "time_period": submission_params.get("time_period"),
1634 "iterations": submission_params.get("iterations"),
1635 "questions_per_iteration": submission_params.get(
1636 "questions_per_iteration"
1637 ),
1638 "strategy": submission_params.get("strategy", "source-based"),
1639 }
1641 query_validation_error = validate_research_query_length(
1642 queued_research.query
1643 )
1644 if query_validation_error is not None:
1645 raise InvalidQueuedResearchOverridesError(query_validation_error)
1647 validation_error = validate_search_overrides(runtime_overrides)
1648 if validation_error is not None:
1649 raise InvalidQueuedResearchOverridesError(validation_error)
1651 # Claim IN_PROGRESS before spawn to close the
1652 # thread-completes-before-parent-commits race.
1653 research.status = ResearchStatus.IN_PROGRESS
1654 db_session.commit()
1656 try:
1657 research_thread = start_research_process(
1658 research_id,
1659 queued_research.query,
1660 queued_research.mode,
1661 run_research_process,
1662 username=username,
1663 user_password=password, # Pass password for metrics
1664 settings_snapshot=complete_settings,
1665 **runtime_overrides,
1666 )
1667 except DuplicateResearchError:
1668 # A live thread already exists for this research_id (e.g.
1669 # previous attempt's post-spawn commit failed). Do NOT
1670 # reset status — that would contradict the running thread.
1671 raise
1672 except SystemAtCapacityError:
1673 # System at concurrent-research capacity. No thread was
1674 # spawned. Reset to QUEUED so the next dispatch tick can try
1675 # again — this is not a permanent spawn failure and should
1676 # NOT count toward SPAWN_RETRY_LIMIT.
1677 logger.info(
1678 f"System at capacity when dispatching {research_id}; "
1679 "re-queueing for next tick"
1680 )
1681 research.status = ResearchStatus.QUEUED
1682 self._commit_with_safe_rollback(
1683 db_session,
1684 f"status reset to QUEUED after capacity reject for research {research_id}",
1685 )
1686 raise
1687 except Exception:
1688 # Genuine spawn failure: no thread exists. Roll back the
1689 # IN_PROGRESS claim so the retry sees a clean QUEUED row.
1690 research.status = ResearchStatus.QUEUED
1691 self._commit_with_safe_rollback(
1692 db_session,
1693 f"status reset to QUEUED after spawn failure for research {research_id}",
1694 )
1695 raise
1697 # Thread is running. Record the active-research row. If this
1698 # commit fails the live thread is unrecorded but still running.
1699 # Raise DuplicateResearchError instead of letting a generic
1700 # exception propagate, so the caller's dup branch cleans up the
1701 # queue row without bumping the retry counter — if we let this
1702 # count as a spawn failure, three consecutive post-spawn commit
1703 # failures (or one at LIMIT-1) would push the counter to
1704 # SPAWN_RETRY_LIMIT and mark a LIVE thread as terminal FAILED.
1705 active_record = UserActiveResearch(
1706 username=username,
1707 research_id=research_id,
1708 status=ResearchStatus.IN_PROGRESS,
1709 thread_id=str(research_thread.ident),
1710 settings_snapshot=queued_research.settings_snapshot,
1711 )
1712 db_session.add(active_record)
1713 if not self._commit_with_safe_rollback(
1714 db_session,
1715 f"UserActiveResearch persist after spawn for research {research_id}",
1716 ):
1717 # Thread is live; the commit failing leaves the UAR row
1718 # unrecorded but the thread running. Raise
1719 # DuplicateResearchError so the caller's dup branch deletes
1720 # the queue row without bumping the retry counter — if we
1721 # let a plain exception count as a spawn failure, a commit
1722 # failure at SPAWN_RETRY_LIMIT - 1 would mark a LIVE thread
1723 # as terminal FAILED.
1724 raise DuplicateResearchError(
1725 f"Research {research_id} thread is live; "
1726 "UserActiveResearch commit failed"
1727 )
1729 # Keep pending_operations bounded: if no request comes in to drain
1730 # them (user never navigates after a crash), they would otherwise
1731 # accumulate forever in memory. We drop oldest entries past this cap
1732 # and entries older than this TTL.
1733 _PENDING_OPS_MAX = 10_000
1734 _PENDING_OPS_TTL_SECONDS = 24 * 60 * 60 # 24 hours
1736 def reconcile_orphan_active_research(
1737 self, username: str, db_session
1738 ) -> int:
1739 """Mark stale UserActiveResearch rows as FAILED.
1741 Called after a user logs in (when their DB is open). If the server
1742 was restarted mid-research, rows with status=IN_PROGRESS have a
1743 thread_id pointing to a now-dead thread. Leaving them as IN_PROGRESS
1744 means the dashboard shows "running forever" with no way to cancel.
1746 Only rows whose worker thread is actually DEAD are reconciled. This
1747 runs on EVERY login, so without the liveness guard a user opening a
1748 second tab / logging in on another device while a research is still
1749 running would wrongly flag that live research FAILED — and
1750 delete_attempt could then hard-delete it out from under the running
1751 worker. Mirrors the is_research_thread_alive() guard chat.py and the
1752 research-start reclaim already use.
1754 Returns the number of rows reconciled.
1755 """
1756 from datetime import UTC, datetime
1758 from ..research_state import cleanup_research, is_research_thread_alive
1760 try:
1761 stale = (
1762 db_session.query(UserActiveResearch)
1763 .filter_by(username=username, status=ResearchStatus.IN_PROGRESS)
1764 .all()
1765 )
1766 if not stale:
1767 return 0
1769 count = 0
1770 for record in stale:
1771 # Skip rows whose worker thread is still alive — those are
1772 # genuinely running (e.g. a concurrent login), not orphans.
1773 if is_research_thread_alive(record.research_id):
1774 continue
1775 # Also mark the ResearchHistory row as failed
1776 research = (
1777 db_session.query(ResearchHistory)
1778 .filter_by(id=record.research_id)
1779 .first()
1780 )
1781 if research and research.status == ResearchStatus.IN_PROGRESS: 1781 ↛ 1789line 1781 didn't jump to line 1789 because the condition on line 1781 was always true
1782 research.status = ResearchStatus.FAILED
1783 research.completed_at = datetime.now(UTC).isoformat()
1784 meta = dict(research.research_meta or {})
1785 meta["failure_reason"] = (
1786 "Server restarted while research was in progress"
1787 )
1788 research.research_meta = meta
1789 record.status = ResearchStatus.FAILED
1790 # Drop the dead thread's in-memory entry too, mirroring
1791 # reclaim_stale_user_active_research. Without this the stale
1792 # _active_research entry outlives the DB reconciliation, so
1793 # capacity checks keep counting a research that is already
1794 # FAILED and its termination flag never clears.
1795 cleanup_research(record.research_id)
1796 count += 1
1798 db_session.commit()
1799 logger.info(
1800 f"Reconciled {count} orphan IN_PROGRESS research records for {username}"
1801 )
1802 return count
1803 except Exception:
1804 logger.exception(
1805 f"Failed to reconcile orphan research records for {username}"
1806 )
1807 return 0
1809 def _evict_stale_pending_operations(self) -> None:
1810 """Drop expired and over-capacity entries. Caller must hold the lock."""
1811 now = time.time()
1812 # TTL eviction
1813 expired = [
1814 op_id
1815 for op_id, op in self.pending_operations.items()
1816 if now - op.get("timestamp", now) > self._PENDING_OPS_TTL_SECONDS
1817 ]
1818 for op_id in expired:
1819 del self.pending_operations[op_id]
1820 # Size cap — drop oldest entries if over
1821 overflow = len(self.pending_operations) - self._PENDING_OPS_MAX
1822 if overflow > 0: 1822 ↛ 1823line 1822 didn't jump to line 1823 because the condition on line 1822 was never true
1823 oldest = sorted(
1824 self.pending_operations.items(),
1825 key=lambda kv: kv[1].get("timestamp", 0),
1826 )[:overflow]
1827 for op_id, _ in oldest:
1828 del self.pending_operations[op_id]
1830 def queue_progress_update(
1831 self, username: str, research_id: str, progress: float
1832 ):
1833 """
1834 Queue a progress update that needs database access.
1835 For compatibility with old processor during migration.
1837 Args:
1838 username: The username
1839 research_id: The research ID
1840 progress: The progress value (0-100)
1841 """
1842 # In processor_v2, we can update directly if we have database access
1843 # or queue it for later processing
1844 operation_id = str(uuid.uuid4())
1845 with self._pending_operations_lock:
1846 self.pending_operations[operation_id] = {
1847 "username": username,
1848 "operation_type": "progress_update",
1849 "research_id": research_id,
1850 "progress": progress,
1851 "timestamp": time.time(),
1852 }
1853 self._evict_stale_pending_operations()
1854 logger.debug(
1855 f"Queued progress update for research {research_id}: {progress}%"
1856 )
1858 def queue_error_update(
1859 self,
1860 username: str,
1861 research_id: str,
1862 status: str,
1863 error_message: str,
1864 metadata: Dict[str, Any],
1865 completed_at: str,
1866 report_path: Optional[str] = None,
1867 ):
1868 """
1869 Queue an error status update that needs database access.
1870 For compatibility with old processor during migration.
1872 Args:
1873 username: The username
1874 research_id: The research ID
1875 status: The status to set (failed, suspended, etc.)
1876 error_message: The error message
1877 metadata: Research metadata
1878 completed_at: Completion timestamp
1879 report_path: Optional path to error report
1880 """
1881 operation_id = str(uuid.uuid4())
1882 with self._pending_operations_lock:
1883 self.pending_operations[operation_id] = {
1884 "username": username,
1885 "operation_type": "error_update",
1886 "research_id": research_id,
1887 "status": status,
1888 "error_message": error_message,
1889 "metadata": metadata,
1890 "completed_at": completed_at,
1891 "report_path": report_path,
1892 "timestamp": time.time(),
1893 }
1894 self._evict_stale_pending_operations()
1895 logger.info(
1896 f"Queued error update for research {research_id} with status {status}"
1897 )
1899 def process_pending_operations_for_user(
1900 self, username: str, db_session: Session
1901 ) -> int:
1902 """
1903 Process pending operations for a user when we have database access.
1904 Called from request context where encrypted database is accessible.
1905 For compatibility with old processor during migration.
1907 Args:
1908 username: Username to process operations for
1909 db_session: Active database session for the user
1911 Returns:
1912 Number of operations cleared from the queue
1913 """
1914 # Find pending operations for this user (with lock)
1915 raw_operations = []
1916 with self._pending_operations_lock:
1917 for op_id, op_data in list(self.pending_operations.items()):
1918 if op_data["username"] == username:
1919 raw_operations.append((op_id, op_data))
1920 # Remove immediately to prevent duplicate processing
1921 del self.pending_operations[op_id]
1923 if not raw_operations:
1924 return 0
1926 # Deduplicate and consolidate operations by research_id to minimize DB locks
1927 latest_progress = {}
1928 latest_errors = {}
1929 other_ops = []
1931 for op_id, op_data in raw_operations:
1932 op_type = op_data.get("operation_type")
1933 rid = op_data.get("research_id")
1934 ts = op_data.get("timestamp", 0)
1936 if op_type == "progress_update" and rid:
1937 if rid not in latest_progress or ts >= latest_progress[rid][
1938 1
1939 ].get("timestamp", 0):
1940 latest_progress[rid] = (op_id, op_data)
1941 elif op_type == "error_update" and rid: 1941 ↛ 1947line 1941 didn't jump to line 1947 because the condition on line 1941 was always true
1942 if rid not in latest_errors or ts >= latest_errors[rid][1].get( 1942 ↛ 1931line 1942 didn't jump to line 1931 because the condition on line 1942 was always true
1943 "timestamp", 0
1944 ):
1945 latest_errors[rid] = (op_id, op_data)
1946 else:
1947 other_ops.append((op_id, op_data))
1949 operations_to_process = (
1950 list(latest_progress.values())
1951 + list(latest_errors.values())
1952 + other_ops
1953 )
1955 from sqlalchemy.exc import (
1956 OperationalError,
1957 PendingRollbackError,
1958 TimeoutError,
1959 )
1961 from ...database.models import ResearchHistory
1963 max_retries = 3
1964 backoff = 0.05
1965 max_requeue_attempts = 10
1966 # Track the raw count of operations being flushed so the returned
1967 # count reflects "operations cleared from the queue" rather than
1968 # "deduplicated batches applied". Set after a successful commit
1969 # so retries don't double-count.
1970 processed_count = 0
1972 for attempt in range(max_retries):
1973 try:
1974 # Consolidate DB queries: fetch all relevant ResearchHistory rows in a single batch
1975 rids = list(
1976 {
1977 op_data["research_id"]
1978 for _, op_data in operations_to_process
1979 if op_data.get("research_id")
1980 }
1981 )
1982 if rids: 1982 ↛ 1990line 1982 didn't jump to line 1990 because the condition on line 1982 was always true
1983 researches = (
1984 db_session.query(ResearchHistory)
1985 .filter(ResearchHistory.id.in_(rids))
1986 .all()
1987 )
1988 research_map = {r.id: r for r in researches}
1989 else:
1990 research_map = {}
1992 for op_id, op_data in operations_to_process:
1993 op_type = op_data.get("operation_type")
1994 rid = op_data.get("research_id")
1995 research = research_map.get(rid)
1997 if not research:
1998 continue
2000 if op_type == "progress_update":
2001 research.progress = op_data.get("progress")
2002 elif op_type == "error_update": 2002 ↛ 1992line 2002 didn't jump to line 1992 because the condition on line 2002 was always true
2003 research.status = op_data.get("status", "failed")
2004 research.error_message = op_data.get("error_message")
2005 research.research_meta = op_data.get("metadata")
2006 research.completed_at = op_data.get("completed_at")
2007 report_path = op_data.get("report_path")
2008 if report_path:
2009 research.report_path = report_path
2011 # Single commit for the entire batch of consolidated updates.
2012 # Count is set AFTER the commit succeeds so retries don't
2013 # double-count the raw operations.
2014 db_session.commit()
2015 processed_count = len(raw_operations)
2016 break
2018 except (OperationalError, PendingRollbackError, TimeoutError) as e:
2019 try:
2020 db_session.rollback()
2021 except Exception as rollback_err:
2022 logger.debug(f"Queue flush rollback failed: {rollback_err}")
2024 if attempt < max_retries - 1:
2025 logger.debug(
2026 f"Database locked during queue flush for {username} (attempt {attempt + 1}/{max_retries}); retrying in {backoff * 1000:.0f}ms: {type(e).__name__}"
2027 )
2028 time.sleep(backoff)
2029 backoff *= 2
2030 else:
2031 logger.warning(
2032 f"Failed to commit queue operations for {username} after {max_retries} retries: {type(e).__name__}"
2033 )
2034 # Re-queue raw operations so data is not lost, up to max_requeue_attempts
2035 with self._pending_operations_lock:
2036 for op_id, op_data in raw_operations:
2037 retry_count = op_data.get("retry_count", 0) + 1
2038 if retry_count > max_requeue_attempts:
2039 logger.exception(
2040 f"Dropping queued operation {op_id} ({op_data.get('operation_type')}) for user {username} "
2041 f"after exceeding maximum requeue attempts ({max_requeue_attempts})"
2042 )
2043 elif op_id not in self.pending_operations: 2043 ↛ 2036line 2043 didn't jump to line 2036 because the condition on line 2043 was always true
2044 op_data["retry_count"] = retry_count
2045 self.pending_operations[op_id] = op_data
2046 processed_count = 0
2048 except Exception:
2049 logger.exception(
2050 f"Unexpected error committing queue operations for {username}"
2051 )
2052 try:
2053 db_session.rollback()
2054 except Exception as rollback_err:
2055 logger.debug(f"Queue flush rollback failed: {rollback_err}")
2056 # Re-queue raw operations up to max_requeue_attempts
2057 with self._pending_operations_lock:
2058 for op_id, op_data in raw_operations:
2059 retry_count = op_data.get("retry_count", 0) + 1
2060 if retry_count > max_requeue_attempts: 2060 ↛ 2061line 2060 didn't jump to line 2061 because the condition on line 2060 was never true
2061 logger.exception(
2062 f"Dropping queued operation {op_id} ({op_data.get('operation_type')}) for user {username} "
2063 f"after exceeding maximum requeue attempts ({max_requeue_attempts})"
2064 )
2065 elif op_id not in self.pending_operations: 2065 ↛ 2058line 2065 didn't jump to line 2058 because the condition on line 2065 was always true
2066 op_data["retry_count"] = retry_count
2067 self.pending_operations[op_id] = op_data
2068 processed_count = 0
2069 break
2071 return processed_count
2074# Global queue processor instance
2075queue_processor = QueueProcessorV2()