Coverage for src/local_deep_research/web/services/research_service.py: 92%
1075 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1import contextvars
2import hashlib
3import json
4import re
5import threading
6import time
7import uuid
8from datetime import datetime, UTC
9from pathlib import Path
11from loguru import logger
13from ...exceptions import DuplicateResearchError, ResearchTerminatedException
14from ...config.llm_config import get_llm, LLMNotConfiguredError
15from ...settings.manager import (
16 SnapshotSettingsContext,
17 apply_environment_overrides_to_snapshot,
18)
20# Output directory for research results
21from ...config.paths import get_research_outputs_directory
22from ...config.search_config import get_search
23from ...security.log_sanitizer import redact_and_bound_for_log
24from ...constants import ResearchStatus
25from ...database.models import ResearchHistory, ResearchStrategy
26from ...database.session_context import get_user_db_session
27from ...database.thread_local_session import thread_cleanup
28from ...error_handling.error_messages import TYPED_RESEARCH_ERROR_MESSAGES
29from ...error_handling.openai_compat_errors import (
30 friendly_openai_compatible_error,
31 is_openai_compat_runtime_error,
32)
33from ...error_handling.report_generator import ErrorReportGenerator
34from ...llm.providers.base import normalize_provider
35from ...utilities.thread_context import set_search_context
36from ...report_generator import IntegratedReportGenerator
37from ...search_system import AdvancedSearchSystem
38from ...text_optimization import CitationFormatter, CitationMode
39from ...text_optimization.citation_formatter import (
40 strip_appended_sources_sentinel,
41)
42from ...utilities.log_utils import log_for_research
43from ...utilities.search_utilities import extract_links_from_search_results
44from ...utilities.type_utils import overlay_runtime_settings
45from ..models.database import calculate_duration
46from ...settings.env_registry import get_env_setting
47from .socketio_asgi import (
48 emit_to_subscribers as _sio_emit,
49 remove_subscriptions_for_research as _sio_remove,
50)
52OUTPUT_DIR = get_research_outputs_directory()
54# Providers whose runtime failures are OpenAI-compatible client errors and
55# get rewritten by ``friendly_openai_compatible_error`` to name the provider,
56# base URL, and model (#3878). Pinned by
57# tests/error_handling/test_openai_compat_errors.py.
58_OPENAI_COMPAT_RUNTIME_ERROR_PROVIDERS = frozenset(
59 {
60 "openai_endpoint",
61 "lmstudio",
62 "llamacpp",
63 "openai",
64 "openrouter",
65 "requesty",
66 "orcarouter",
67 "atlascloud",
68 "qwen",
69 "deepseek",
70 "google",
71 "ionos",
72 "xai",
73 }
74)
77def _resolve_error_rewrite_base_url(
78 provider, custom_endpoint, settings_context=None
79):
80 """Resolve the base URL for the friendly error rewrite.
82 ``custom_endpoint`` is the ``llm.openai_endpoint.url`` value and is only
83 the URL the client uses when the provider is ``openai_endpoint``. Some
84 callers pass it for every provider (it defaults to the OpenRouter URL),
85 so it is ignored for any other provider — otherwise QwenCloud, DeepSeek,
86 LM Studio, ... would be reported against a URL they never contacted.
88 Providers with a hardcoded endpoint (QwenCloud, DeepSeek, Atlas Cloud,
89 ...) fall back to the provider class default so the rewritten message
90 names the real base URL instead of ``<unknown>``. Providers whose URL is
91 configurable via a settings key (LM Studio, llama.cpp, custom endpoints)
92 resolve their configured URL through the settings context — their class
93 default is just a localhost guess that would misdirect troubleshooting
94 when an operator overrode it.
95 """
96 provider = normalize_provider(provider)
97 if not provider:
98 return None
99 if custom_endpoint and provider == "openai_endpoint":
100 return custom_endpoint
101 from ...llm.providers.auto_discovery import get_provider_class
103 provider_cls = get_provider_class(provider)
104 if provider_cls is None:
105 return None
106 url_setting = getattr(provider_cls, "url_setting", None)
107 if url_setting:
108 if settings_context is None:
109 # No context to read the operator's configured URL — the class
110 # default is just a localhost guess; report "<unknown>" rather
111 # than risk naming a URL the client is not using.
112 return None
113 url = settings_context.get_setting(url_setting, None)
114 return url or None
115 if provider == "openai" and settings_context is not None:
116 # OpenAIProvider's create_llm honors an operator-configured
117 # llm.openai.api_base gateway; naming the class default would
118 # misdirect troubleshooting the same way a wrong localhost would.
119 api_base = settings_context.get_setting("llm.openai.api_base", None)
120 if api_base: 120 ↛ 122line 120 didn't jump to line 122 because the condition on line 120 was always true
121 return api_base
122 return getattr(provider_cls, "default_base_url", None)
125class _SocketEmitter:
126 """Object-style adapter over ``socketio_asgi.emit_to_subscribers``.
128 Main's Chat Mode code (merged from #2953) instantiated the old
129 Flask ``SocketIOService`` and passed it around as an object with an
130 ``.emit_to_subscribers`` method. Under FastAPI the real emit path is
131 the module-level ``socketio_asgi`` function (aliased ``_sio_emit``);
132 this thin adapter preserves the object API those call sites expect.
133 """
135 @staticmethod
136 def emit_to_subscribers(event, research_id, data, *, owner):
137 # ``owner`` is forwarded, not defaulted. Subscriptions are keyed by
138 # (owner, research_id), so a default here would send every event
139 # through the adapter to the wrong key and deliver nothing --
140 # silently, because every call site sits inside a swallowing
141 # ``except``. Keeping it required means a caller that forgets it
142 # fails loudly at the call.
143 return _sio_emit(event, research_id, data, owner=owner)
146# Module-level singleton — the adapter is stateless.
147_socket_emitter = _SocketEmitter()
150# Global concurrent research limit (server-wide, across all users)
151_MAX_GLOBAL_CONCURRENT = get_env_setting(
152 "server.max_concurrent_research", default=10
153)
154_global_research_semaphore = threading.Semaphore(_MAX_GLOBAL_CONCURRENT)
156# Fallback used by clamp_user_max_concurrent when the stored
157# `app.max_concurrent_researches` value can't be parsed as an int. Mirrors
158# that setting's own JSON-schema default (see default_settings.json).
159_DEFAULT_USER_MAX_CONCURRENT = 3
162def clamp_user_max_concurrent(raw_value) -> int:
163 """Clamp a user-supplied ``app.max_concurrent_researches`` value to the
164 server-wide semaphore ceiling.
166 ``app.max_concurrent_researches`` is a per-user, user-editable setting
167 (default 3) gating how many of a user's OWN researches may start
168 immediately instead of queueing. Its JSON schema now caps it at 20
169 (previously uncapped), but a value written to a user's DB row before
170 that cap existed keeps its raw stored value -- schema metadata is only
171 reconciled (value is preserved) on the next login after a version bump,
172 see ``SettingsManager.import_settings``. Without this clamp, a user
173 could set an arbitrarily high per-user limit and flood the global
174 research semaphore (``server.max_concurrent_research``, default 10)
175 with their own researches, starving every other user. Clamping at each
176 read site is defense-in-depth on top of the JSON-schema ``max_value``.
178 ACCEPTED RESIDUAL: this clamp fixes accounting against the global
179 ceiling, not per-user fairness. A single user whose clamped value equals
180 the full global ceiling (the default is uncapped up to 20, clamped to
181 10) can still legitimately occupy every global semaphore slot at once,
182 queueing every other user's researches until one of theirs finishes.
183 Per-user fairness / anti-monopoly scheduling (e.g. a lower default
184 per-user share, or reserving slots across users) is intentionally out of
185 scope here -- tracked as a known limitation, not a bug in this function.
187 Args:
188 raw_value: The raw setting value (any type -- callers pass whatever
189 ``SettingsManager.get_setting`` returned).
191 Returns:
192 An int between 1 and the live global semaphore ceiling
193 (``_MAX_GLOBAL_CONCURRENT``), inclusive.
194 """
195 try:
196 value = int(raw_value)
197 except (TypeError, ValueError):
198 value = _DEFAULT_USER_MAX_CONCURRENT
199 value = max(1, value)
200 return min(value, _MAX_GLOBAL_CONCURRENT)
203# Progress allocation for detailed mode — report generation is the bulk of the work
204_DETAILED_SEARCH_PROGRESS_CAP = 8 # Search/output phases capped here
205_DETAILED_REPORT_PROGRESS_START = 10 # Report generation starts here
206_DETAILED_REPORT_PROGRESS_END = 100 # Report generation ends here
207# Quick mode: the search strategy reports up to here, then the answer is
208# generated in the 85-95% band (see progress_callback).
209_QUICK_OUTPUT_PROGRESS_START = 85
210# 100% is reserved for the run's own successful completion. The frontend
211# treats ``progress >= 100`` as "research finished" (chat.js, progress.js),
212# so every in-progress value is capped below it -- a strategy that reports
213# its own completion at 100%, or an overshooting phase mapping, must not
214# end the run in the UI while the answer is still being generated.
215_IN_PROGRESS_CAP = 99
216# Metadata flag carried only by run_research_process's final
217# ``phase="complete"`` emit. It is the explicit completion signal: strategies
218# also emit ``phase="complete"`` when their analyze_topic returns, which is
219# mid-run (quick: before answer generation; detailed: once per subsection).
220_RUN_COMPLETE_FLAG = "_run_complete"
221# Phases that belong to the report-generation stage. "report_generation" is
222# emitted in this file; the other four are emitted by report_generator.py. If
223# you add or rename a phase, update both this set and the emitter.
224_REPORT_PHASES = frozenset(
225 {
226 "report_generation",
227 "report_section_research",
228 "report_formatting",
229 "report_structure",
230 "report_complete",
231 }
232)
234# Phases that produce user-visible step messages in chat mode.
235# "complete" is excluded — it fires AFTER the response message is written,
236# which would create a step with a higher sequence_number than the response.
237_STEP_PHASES = frozenset(
238 {
239 "init",
240 "setup",
241 "search_planning",
242 "search",
243 "observation",
244 "output_generation",
245 "synthesis_error",
246 "synthesis_fallback",
247 "report_generation",
248 "report_complete",
249 "error",
250 }
251)
253# Socket.IO emission throttling: minimum interval between progress emissions per research
254_EMIT_THROTTLE_SECONDS = 0.2 # 200ms
255_EMIT_TTL_SECONDS = 3600 # 1 hour — evict stale entries from orphaned research
257# Cap on the partial-content buffer kept server-side so chat-mode termination
258# can persist whatever was already streamed. Bounded to keep memory predictable
259# under pathologically long answers (typical answers are a few KB).
260_MAX_PARTIAL_BUFFER_BYTES = 256 * 1024 # 256 KB
261_emit_cleanup_counter = 0
262_last_emit_times: dict[str, float] = {}
263_last_emit_lock = threading.Lock()
266def _chat_step_decision(
267 phase: str | None,
268 last_step_phase: str | None,
269 is_final: bool,
270) -> tuple[bool, bool]:
271 """Decide whether to persist + emit a chat-mode step event.
273 Encodes the symmetry invariant the progress_callback in
274 run_research_process enforces: for chat sessions, what the live UI
275 surfaces over the socket must equal what `loadSession` reconstructs
276 from chat_progress_steps on reload. The repeat-phase dedup must
277 therefore block BOTH writes, not just the DB write.
279 Returns:
280 (persist, suppress_emit)
281 - persist: True iff add_progress_step should be called for this event.
282 - suppress_emit: True iff the caller should null the socket payload
283 so the emit is dropped. Only suppressed when this is a non-final
284 repeat — final events (error/report_complete, or the run's
285 explicit completion) always emit so the client completion
286 handler fires.
288 Args:
289 phase: the phase tag from this event (e.g. "search", "observation")
290 last_step_phase: phase tag of the previously persisted chat step
291 (None until the first persist of this research).
292 is_final: True if this is a "final" event the client must see to
293 fire its completion handler (error | report_complete | the
294 run's own completion emit, flagged with _RUN_COMPLETE_FLAG).
295 A progress value alone never makes an event final.
296 """
297 if phase not in _STEP_PHASES:
298 # Not a chat-step phase at all (e.g. "complete"). Persist no,
299 # and let the emit through unchanged — completion / control
300 # events still need to reach the client.
301 return False, False
302 dedup_ok = phase != last_step_phase or phase == "observation"
303 if dedup_ok:
304 return True, False
305 return False, not is_final
308def _compose_chat_step_content(
309 message: str, phase: str | None, metadata: dict
310) -> str:
311 """Build the content persisted for a chat progress step.
313 Observation events carry the (bounded) full tool output in
314 ``metadata["content"]`` alongside the one-line ``message``. Persist it
315 beneath the message so the chat step row can expand to show what the
316 tool actually returned — and so a session reload reconstructs exactly
317 what the live accordion showed (the same composition happens client-side
318 in chat.js handleProgressUpdate; keep the two in sync).
320 Only observation events compose: other phases reuse ``metadata["content"]``
321 for large payloads (e.g. full synthesized answers) that must not be
322 duplicated into step rows.
323 """
324 detail = metadata.get("content")
325 if phase == "observation" and detail:
326 return f"{message}\n\n{detail}"
327 return message
330def _parse_research_metadata(research_meta) -> dict:
331 """Parse research_meta into a dict, handling both dict and JSON string types."""
332 if isinstance(research_meta, dict):
333 return dict(research_meta)
334 if isinstance(research_meta, str):
335 try:
336 parsed = json.loads(research_meta)
337 return dict(parsed) if isinstance(parsed, dict) else {}
338 except (ValueError, RecursionError):
339 # JSONDecodeError is a ValueError, but the decoder's integer
340 # digit limit raises a plain ValueError. Keep the same fallback
341 # for unreadable legacy metadata, including nesting beyond the
342 # decoder's recursion limit.
343 logger.exception("Failed to parse research_meta as JSON")
344 return {}
345 return {}
348def _ensure_settings_snapshot(metadata: dict, settings_snapshot) -> dict:
349 """Back-fill ``settings_snapshot`` into completion metadata when missing.
351 Chat-created (``web/routers/chat.py``) and follow-up-created
352 (``web/routers/followup.py``) rows store only ``{"submission": ...}``
353 at creation, while main-form and scheduler runs persist the full
354 snapshot. Without this, ``report_assembly_service`` view/export
355 rendering (which resolves ``report.uncited_sources_mode`` from the
356 row's snapshot) always falls back for those rows, diverging from
357 what the generation-time render used.
359 Only fills when the row lacks a dict snapshot; an existing snapshot
360 is preserved (creation-time is the faithful preference). No-op when
361 the worker has no usable snapshot.
362 """
363 if not isinstance(metadata, dict): 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 return metadata
365 existing = metadata.get("settings_snapshot")
366 if isinstance(existing, dict):
367 return metadata
368 if isinstance(settings_snapshot, dict) and settings_snapshot:
369 metadata["settings_snapshot"] = settings_snapshot
370 return metadata
373def _merge_completion_metadata(
374 research_meta, extra: dict, settings_snapshot
375) -> dict:
376 """Build the completion-time ``research_meta`` for quick + detailed runs.
378 Single choke point for both completion merges in
379 ``run_research_process``: preserves the row's existing metadata,
380 overlays the run's result values, and back-fills the run's settings
381 snapshot for creation paths (chat/follow-up) that store none — so
382 later view/export renders the same uncited-sources mode the
383 generation-time render used. Covered directly by
384 ``TestMergeCompletionMetadata``; removing the back-fill here fails
385 those tests.
386 """
387 metadata = _parse_research_metadata(research_meta)
388 if isinstance(extra, dict): 388 ↛ 390line 388 didn't jump to line 390 because the condition on line 388 was always true
389 metadata.update(extra)
390 _ensure_settings_snapshot(metadata, settings_snapshot)
391 return metadata
394def _is_error_shaped_content(content: object) -> bool:
395 """Recognize strategy error text that must not become report content."""
396 return isinstance(content, str) and content.lstrip().lower().startswith(
397 "error:"
398 )
401def _extract_synthesized_answer(results: dict) -> str:
402 """Pull the LLM-synthesized answer out of a strategy result dict.
404 ``report_content`` must store ONLY the synthesized answer (LLM
405 prose with [N] inline citations). The strategy's
406 ``formatted_findings`` is the full ``format_findings`` blob —
407 answer + ``format_links_to_markdown`` source list +
408 ``## SEARCH QUESTIONS BY ITERATION`` + ``## DETAILED FINDINGS``
409 + ``## ALL SOURCES`` — and ``format_document_split`` only knows
410 how to strip ``## Sources`` headers, not those other sections.
411 Saving the blob would leak sources/findings into ``report_content``.
413 Resolution order:
414 1. ``Final synthesis`` finding (set by source_based, parallel,
415 rapid, focused_iteration, iterdrag).
416 2. ``current_knowledge`` (other strategies expose the answer
417 there).
418 3. Empty string — caller decides whether to fall back further.
420 Error-shaped content is not an answer and is skipped at both steps.
421 """
422 for finding in results.get("findings") or []:
423 if finding.get("phase") == "Final synthesis":
424 content = finding.get("content") or ""
425 if content and not _is_error_shaped_content(content):
426 return content
427 knowledge = results.get("current_knowledge") or ""
428 return "" if _is_error_shaped_content(knowledge) else knowledge
431def get_citation_formatter():
432 """Get citation formatter with settings from thread context."""
433 # Import here to avoid circular imports
434 from ...config.search_config import get_setting_from_snapshot
436 citation_format = get_setting_from_snapshot(
437 "report.citation_format", "number_hyperlinks"
438 )
439 mode_map = {
440 "number_hyperlinks": CitationMode.NUMBER_HYPERLINKS,
441 "domain_hyperlinks": CitationMode.DOMAIN_HYPERLINKS,
442 "domain_id_hyperlinks": CitationMode.DOMAIN_ID_HYPERLINKS,
443 "domain_id_always_hyperlinks": CitationMode.DOMAIN_ID_ALWAYS_HYPERLINKS,
444 "source_tagged_hyperlinks": CitationMode.SOURCE_TAGGED_HYPERLINKS,
445 "no_hyperlinks": CitationMode.NO_HYPERLINKS,
446 }
447 mode = mode_map.get(citation_format, CitationMode.NUMBER_HYPERLINKS)
448 return CitationFormatter(mode=mode)
451def export_report_to_memory(
452 markdown_content: str, format: str, title: str | None = None
453):
454 """
455 Export a markdown report to different formats in memory.
457 Uses the modular exporter registry to support multiple formats.
458 Available formats can be queried with ExporterRegistry.get_available_formats().
460 Args:
461 markdown_content: The markdown content to export
462 format: Export format (e.g., 'pdf', 'odt', 'latex', 'quarto', 'ris')
463 title: Optional title for the document
465 Returns:
466 Tuple of (content_bytes, filename, mimetype)
467 """
468 from ...exporters import ExporterRegistry, ExportOptions
470 # Normalize format
471 format_lower = format.lower()
473 # Get exporter from registry
474 exporter = ExporterRegistry.get_exporter(format_lower)
476 if exporter is None:
477 available = ExporterRegistry.get_available_formats()
478 raise ValueError(
479 f"Unsupported export format: {format}. "
480 f"Available formats: {', '.join(available)}"
481 )
483 # Title prepending is now handled by each exporter via _prepend_title_if_needed()
484 # PDF and ODT exporters prepend titles; RIS and other formats ignore them
486 # Create options
487 options = ExportOptions(title=title)
489 # Export
490 result = exporter.export(markdown_content, options)
492 logger.info(
493 f"Generated {format_lower} in memory, size: {len(result.content)} bytes"
494 )
496 return result.content, result.filename, result.mimetype
499def save_research_strategy(research_id, strategy_name, *, username):
500 """
501 Save the strategy used for a research to the database.
503 Args:
504 research_id: The ID of the research
505 strategy_name: The name of the strategy used
506 username: The username whose encrypted DB to write. Required —
507 without it get_user_db_session would silently fall back to
508 the Flask session user (or fail off-request-context), so
509 callers must state whose DB they mean.
510 """
511 # ``strategy_name`` is the raw request-body ``strategy`` (the research
512 # router passes it through unvalidated), so it goes through the bounded
513 # preview (#6191) here and on the info line below instead of an eager
514 # f-string (#6305). Imported at call time, not at module level, so this
515 # service does not pull the notes router (and FastAPI) into the CLI /
516 # MCP import tree; same pattern as the rag router import in
517 # research_library/services/download_service.py. Inside the ``try``
518 # (#6939): a transient import failure joins the function's own handled
519 # path instead of escaping it.
520 try:
521 from ..routers.notes import _log_value_preview
523 logger.opt(lazy=True).debug(
524 "save_research_strategy called with research_id={}, strategy_name={}",
525 lambda: research_id,
526 lambda: _log_value_preview(strategy_name),
527 )
528 with get_user_db_session(username) as session:
529 # Check if a strategy already exists for this research
530 existing_strategy = (
531 session.query(ResearchStrategy)
532 .filter_by(research_id=research_id)
533 .first()
534 )
536 if existing_strategy:
537 # Update existing strategy
538 existing_strategy.strategy_name = strategy_name
539 logger.debug(
540 f"Updating existing strategy for research {research_id}"
541 )
542 else:
543 # Create new strategy record
544 new_strategy = ResearchStrategy(
545 research_id=research_id, strategy_name=strategy_name
546 )
547 session.add(new_strategy)
548 logger.debug(
549 f"Creating new strategy record for research {research_id}"
550 )
552 session.commit()
553 logger.info(
554 "Saved strategy {} for research {}",
555 _log_value_preview(strategy_name),
556 research_id,
557 )
558 except Exception:
559 logger.exception("Error saving research strategy")
562def get_research_strategy(research_id, *, username):
563 """
564 Get the strategy used for a research.
566 Args:
567 research_id: The ID of the research
568 username: The username whose encrypted DB to read. Required —
569 without it get_user_db_session would silently fall back to
570 the Flask session user (or fail off-request-context, which
571 the except below would swallow into a None return), so
572 callers must state whose DB they mean.
574 Returns:
575 str: The strategy name or None if not found
576 """
577 try:
578 with get_user_db_session(username) as session:
579 strategy = (
580 session.query(ResearchStrategy)
581 .filter_by(research_id=research_id)
582 .first()
583 )
585 return strategy.strategy_name if strategy else None
586 except Exception:
587 logger.exception("Error getting research strategy")
588 return None
591def start_research_process(
592 research_id,
593 query,
594 mode,
595 run_research_callback,
596 **kwargs,
597):
598 """
599 Start a research process in a background thread.
601 Args:
602 research_id: The ID of the research
603 query: The research query
604 mode: The research mode (quick/detailed)
605 run_research_callback: The callback function to run the research
606 **kwargs: Additional parameters to pass to the research process (model, search_engine, etc.)
608 Returns:
609 threading.Thread: The thread running the research
610 """
611 from ..research_state import check_and_start_research
612 from ...exceptions import SystemAtCapacityError
614 # Acquire the global concurrency semaphore SYNCHRONOUSLY in the caller's
615 # thread. Previously this happened inside the worker after the HTTP route
616 # had already returned 200 — at capacity, the worker parked and the user
617 # saw an infinite thinking spinner with the partial unique in-progress
618 # index blocking retries. Surfacing capacity as an exception lets the
619 # route return HTTP 429 before committing any DB state.
620 if not _global_research_semaphore.acquire(blocking=False):
621 raise SystemAtCapacityError(
622 f"At research capacity (max {_MAX_GLOBAL_CONCURRENT} concurrent)"
623 )
625 # Wrap callback so the worker releases the already-held semaphore on exit.
626 original_callback = run_research_callback
628 def _release_semaphore_on_exit(*args, **kw):
629 try:
630 return original_callback(*args, **kw)
631 finally:
632 _global_research_semaphore.release()
634 # Copy the calling request's contextvars (username, session_id) so
635 # service code inside the worker thread sees the authenticated user
636 # when it calls `get_current_username()`. Bare `threading.Thread`
637 # starts with an empty context; without this, metrics/token attribution
638 # and any code path using request_context silently fails to the
639 # Flask-compat fallback (which returns None under FastAPI).
640 _ctx = contextvars.copy_context()
642 def _ctx_wrapped(*a, **kw):
643 return _ctx.run(_release_semaphore_on_exit, *a, **kw)
645 # Prepare (but do not start) the background thread.
646 thread = threading.Thread(
647 target=_ctx_wrapped,
648 args=(
649 research_id,
650 query,
651 mode,
652 ),
653 kwargs=kwargs,
654 )
655 thread.daemon = True
657 # Atomic check-and-start: refuses to spawn a second live thread
658 # for the same research_id. Guards against the double-spawn window
659 # where a post-spawn commit failure in the queue processor could
660 # otherwise cause the retry loop to dispatch the same research twice.
661 try:
662 started = check_and_start_research(
663 research_id,
664 {
665 "thread": thread,
666 "progress": 0,
667 "status": ResearchStatus.IN_PROGRESS,
668 "log": [],
669 "settings": kwargs,
670 },
671 )
672 except Exception:
673 # check_and_start_research raised before the thread ran — the
674 # semaphore won't be released by the worker, so release it here.
675 _global_research_semaphore.release()
676 raise
677 if not started:
678 # No thread will run → no _release_semaphore_on_exit → release here.
679 _global_research_semaphore.release()
680 raise DuplicateResearchError(
681 f"Research {research_id} already has a live thread"
682 )
684 return thread
687def _generate_report_path(query: str) -> Path:
688 """
689 Generates a path for a new report file based on the query.
691 Args:
692 query: The query used for the report.
694 Returns:
695 The path that it generated.
697 """
698 # Generate a unique filename that does not contain
699 # non-alphanumeric characters.
700 query_hash = hashlib.md5( # DevSkim: ignore DS126858
701 query.encode("utf-8"), usedforsecurity=False
702 ).hexdigest()[:10]
703 # Preventive hardening (this helper has no production caller today):
704 # the directory is shared install-wide, and the hash+timestamp prefix
705 # alone is deterministic across users, so the trailing random segment
706 # keeps same-query, same-second calls from colliding.
707 return OUTPUT_DIR / (
708 f"research_report_{query_hash}_"
709 f"{int(datetime.now(UTC).timestamp())}_{uuid.uuid4().hex[:8]}.md"
710 )
713def _save_chat_message_and_context(
714 chat_session_id,
715 research_id,
716 username,
717 report_content,
718 streaming_enabled,
719 streaming_state,
720 socket_service,
721 settings_snapshot=None,
722):
723 """Save assistant message to chat and update accumulated context."""
724 from ...chat.service import ChatService
725 from ...chat.context import ChatContextManager
727 chat_service = ChatService(username)
728 # chat_messages.content is NOT NULL. Write report_content
729 # inline (snapshot pattern). Falls back to a placeholder marker only
730 # if report_content is itself empty — the same pattern used by the
731 # terminate path's _STOPPED_BEFORE_OUTPUT_MARKER.
732 snapshot_content = report_content or _NO_OUTPUT_MARKER
733 # allow_archived=True: a multi-tab race can flip the session to
734 # archived between research.status=COMPLETED commit and this write.
735 # Losing the final assistant answer is worse than violating the
736 # "no writes to archived sessions" rule for a system-generated row;
737 # user-message writes from chat/routes.py still keep the default.
738 chat_service.add_message(
739 session_id=chat_session_id,
740 role="assistant",
741 content=snapshot_content,
742 message_type="response",
743 research_id=research_id,
744 allow_archived=True,
745 )
746 # Mark the response row as persisted so the trailing
747 # progress_callback("Research completed successfully", 100) — which
748 # runs the termination check and could fire if the user clicks Stop
749 # in the small window between this write and the final emit — does
750 # NOT write a duplicate row via _save_partial_chat_message_on_terminate.
751 if streaming_state is not None: 751 ↛ 753line 751 didn't jump to line 753 because the condition on line 751 was always true
752 streaming_state["_persisted"] = True
753 logger.info(f"Added research result to chat {chat_session_id[:8]}...")
755 try:
756 # report_content is the answer-only string; no extraction needed.
757 chat_content = report_content or ""
758 ctx_manager = ChatContextManager(
759 session_id=chat_session_id,
760 messages=[],
761 settings_snapshot=settings_snapshot,
762 )
763 updates = ctx_manager.extract_context_updates(new_content=chat_content)
764 chat_service.update_accumulated_context(
765 session_id=chat_session_id, **updates
766 )
767 logger.info(
768 f"Updated accumulated context for chat {chat_session_id[:8]}..."
769 )
770 except Exception:
771 # Bumped from debug to warning: a failed accumulated-context
772 # update silently degrades multi-turn context for the next
773 # follow-up in this chat (entities/topics/source counts are
774 # missing). Ops need visibility to catch widespread breakage
775 # before users notice "the AI forgot what we were talking
776 # about" — debug-level was invisible in production.
777 logger.opt(exception=True).warning(
778 "Could not update accumulated context"
779 )
781 if streaming_enabled and streaming_state.get("chunks_sent", 0) > 0:
782 # Flush any partial-bracket fragment held in the citation carry
783 # buffer before sending the final empty sentinel. Without this,
784 # a stream that ends mid-token (LLM emits "[12" as its last
785 # bytes before closing) silently drops the leading "[" from
786 # what the client renders — the carry would be discarded when
787 # the callback's closure goes out of scope.
788 flush = streaming_state.get("_flush_carry")
789 if flush:
790 leftover = flush()
791 if leftover:
792 try:
793 socket_service.emit_to_subscribers(
794 "response_chunk",
795 research_id,
796 {
797 "chunk": leftover,
798 "is_streaming": True,
799 "is_final": False,
800 },
801 owner=username,
802 )
803 except Exception:
804 logger.debug(
805 "Carry-buffer flush emit failed (non-critical)"
806 )
807 socket_service.emit_to_subscribers(
808 "response_chunk",
809 research_id,
810 {"chunk": "", "is_streaming": True, "is_final": True},
811 owner=username,
812 )
815_STOPPED_BEFORE_OUTPUT_MARKER = "[Stopped before any output was generated.]"
816_STOPPED_FOOTER = "\n\n_— Stopped by user._"
817_NO_OUTPUT_MARKER = "_(Research completed without producing output.)_"
820# Match a trailing incomplete citation opener at the chunk boundary.
821# Covers both ASCII "[" and the lenticular "【" (U+3010) — some LLMs
822# emit Chinese-style brackets and the citation formatter accepts them,
823# so we have to hold those back the same way to avoid breaking a token
824# across the next chunk.
825_PARTIAL_BRACKET_RE = re.compile(r"[\[【]\d*$")
827# Upper bound on the inline-citation carry buffer. A real citation token
828# is a handful of bytes ("[123"); anything longer means the "[" wasn't a
829# citation opener or the stream is pathological. Flush raw past this so a
830# never-closing "[" + endless digits can't grow the buffer without limit.
831_MAX_CARRY_BYTES = 64
834def _make_chat_stream_callback(
835 research_id,
836 streaming_state,
837 socket_service,
838 username,
839 source_resolver=None,
840 formatter=None,
841):
842 """Build the chat-mode streaming callback.
844 The callback:
845 * Counts chunks (``chunks_sent``).
846 * Buffers RAW chunks in ``streaming_state['chunks']`` so partial
847 content survives termination — the citation handler's local list
848 is discarded on raise. Capped at ``_MAX_PARTIAL_BUFFER_BYTES``;
849 once capped, ``streaming_state['_truncated']`` flips to True and
850 further chunks aren't accumulated server-side.
851 * Raises ``ResearchTerminatedException`` if the user clicked Stop
852 mid-stream — fails fast instead of letting the LLM finish.
853 ``ResearchTerminatedException`` inherits from ``BaseException``,
854 so the citation handler's ``except Exception`` blocks naturally
855 propagate it.
856 * Emits ``response_chunk`` over Socket.IO for live display, with
857 inline citation hyperlinks applied per-chunk when both
858 ``source_resolver`` and ``formatter`` are provided — so the
859 client sees ``[[arxiv.org-1]](url)`` appearing live rather than
860 plain ``[1]`` brackets that only get hyperlinked after the full
861 response saves. Bracket tokens split across chunk boundaries are
862 held in a small carry buffer until the closing ``]`` arrives.
864 ``source_resolver`` — optional ``() -> list[dict]`` returning the
865 current ``all_links_of_system`` (so we read it lazily; the list
866 grows as the agent collects sources).
867 ``formatter`` — optional :class:`CitationFormatter` instance. Its
868 ``mode`` controls the inline format (``[[arxiv.org-1]](url)``
869 etc.), matching what the final-save formatter produces so the
870 live display doesn't mode-shift when ``handleResearchComplete``
871 swaps in the DB-saved version.
873 Extracted to module level so it can be unit-tested without spinning
874 up the full ``run_research_process``.
875 """
877 # Per-call closure state for the streaming-substitution carry
878 # buffer (holds a trailing incomplete bracket like "[12" so the
879 # closing "]3]" on the next chunk completes the citation token).
880 carry = [""]
882 def _flush_carry() -> str:
883 """Release and clear any held partial-bracket fragment.
885 The completion finalizer (``_save_chat_message_and_context``)
886 lives in a different scope and can't reach ``carry`` directly,
887 so it calls this through ``streaming_state['_flush_carry']`` to
888 avoid silently dropping the tail of a stream that ends mid-token
889 like ``"[12"``.
890 """
891 released, carry[0] = carry[0], ""
892 return released
894 # Expose to the completion path via the shared state dict.
895 streaming_state["_flush_carry"] = _flush_carry
897 def _hyperlink_chunk(chunk: str) -> str:
898 """Apply inline citation hyperlinks to a single chunk.
900 Maintains the closure-level ``carry`` buffer for incomplete
901 ``[N`` tokens straddling chunk boundaries. The carry contains
902 the trailing partial-bracket fragment from the previous chunk
903 that we haven't been able to substitute yet — it's prepended
904 to the next chunk so the regex can see the full token. The
905 DELTA returned is what the client should append next: it
906 includes the just-completed carry (now hyperlinked) plus the
907 new chunk's safe portion, MINUS the new chunk's own trailing
908 partial bracket (which becomes the new carry).
909 """
910 if source_resolver is None or formatter is None:
911 return chunk
912 # Keep the complete raw delta before updating carry for this chunk.
913 # The source resolver runs before that update; only formatting
914 # can fail after it.
915 text = carry[0] + chunk
916 try:
917 sources = source_resolver() or []
918 if not sources:
919 # Reset carry so the leading "[" we held onto doesn't
920 # disappear from the client's accumulated text.
921 carry[0] = ""
922 return text
923 pending = _PARTIAL_BRACKET_RE.search(text)
924 if pending:
925 safe = text[: pending.start()]
926 new_carry = text[pending.start() :]
927 # Bound the carry. A well-formed citation token is a few
928 # bytes (`[12`); if the held fragment grows past this, the
929 # "[" was not actually opening a citation (or a hostile /
930 # misbehaving LLM is streaming `[` + endless digits with no
931 # closing `]`). Flush it raw rather than buffering without
932 # limit — preserves the text, just doesn't hyperlink it.
933 if len(new_carry) > _MAX_CARRY_BYTES:
934 safe = text
935 carry[0] = ""
936 else:
937 carry[0] = new_carry
938 else:
939 safe = text
940 carry[0] = ""
941 if not safe:
942 return ""
943 return formatter.apply_inline_hyperlinks(safe, sources)
944 except Exception:
945 # Include the previous carry exactly once even if formatting
946 # failed after replacing it with a new incomplete bracket.
947 logger.debug(
948 "Inline citation hyperlinking failed; emitting raw chunk",
949 exc_info=True,
950 )
951 carry[0] = ""
952 return text
954 def stream_callback(chunk: str):
955 # Resolve through the module namespace each call so tests can
956 # ``patch("local_deep_research.web.routes.globals.is_termination_requested")``.
957 # Cached in sys.modules — negligible cost.
958 from ..routes.globals import is_termination_requested
960 if not chunk:
961 return
962 streaming_state["chunks_sent"] += 1
963 if not streaming_state["_truncated"]: 963 ↛ 984line 963 didn't jump to line 984 because the condition on line 963 was always true
964 chunk_bytes = len(chunk.encode("utf-8"))
965 if (
966 streaming_state["_bytes"] + chunk_bytes
967 <= _MAX_PARTIAL_BUFFER_BYTES
968 ):
969 # IMPORTANT: store the RAW chunk in the partial-content
970 # buffer (terminate handler joins these and saves them as
971 # the partial assistant message). If we stored the
972 # hyperlinked version, the saved-on-terminate text would
973 # be double-formatted when the user resumes.
974 streaming_state["chunks"].append(chunk)
975 streaming_state["_bytes"] += chunk_bytes
976 else:
977 streaming_state["_truncated"] = True
978 logger.warning(
979 f"Partial-content buffer hit {_MAX_PARTIAL_BUFFER_BYTES} bytes "
980 f"for research {research_id}; further chunks won't be persisted on terminate"
981 )
982 # Mid-stream interrupt — fail fast if the user clicked Stop while
983 # the LLM is still streaming.
984 if is_termination_requested(research_id):
985 raise ResearchTerminatedException( # noqa: TRY301 — propagated through citation handler
986 "Research was terminated by user during streaming"
987 )
988 # Apply citation hyperlinks for the client emit only — the
989 # client accumulates substituted text into the streaming bubble
990 # so the user sees [[arxiv.org-1]](url) appearing live as the
991 # model writes "According to [1]…".
992 display_chunk = _hyperlink_chunk(chunk)
993 if not display_chunk:
994 return # nothing safe to emit yet (all held in carry buffer)
995 try:
996 socket_service.emit_to_subscribers(
997 "response_chunk",
998 research_id,
999 {
1000 "chunk": display_chunk,
1001 "is_streaming": True,
1002 "is_final": False,
1003 },
1004 owner=username,
1005 )
1006 except Exception:
1007 logger.debug("Stream chunk emit failed (non-critical)")
1009 return stream_callback
1012def _save_partial_chat_message_on_terminate(
1013 chat_session_id,
1014 research_id,
1015 username,
1016 partial_content,
1017 truncated=False,
1018 streaming_state=None,
1019):
1020 """Persist a chat 'response' row capturing whatever was streamed before
1021 termination, and emit a final ``response_chunk`` so the client strips
1022 the streaming class from the bubble.
1024 Must run BEFORE ``handle_termination()`` because that path runs
1025 ``cleanup_research_resources()`` which removes the Socket.IO room
1026 subscriptions — anything emitted afterwards goes nowhere.
1028 Idempotent: when ``streaming_state`` is supplied, sets a ``_persisted``
1029 flag so duplicate calls (one from the progress callback, one from the
1030 outer except handler when a stream_callback raises mid-stream) only
1031 write a single row.
1033 Skips silently when ``chat_session_id`` is falsy (single-turn case).
1034 All failures are swallowed — termination cleanup must never crash the
1035 worker.
1036 """
1037 if not chat_session_id:
1038 return
1039 if streaming_state is not None and streaming_state.get("_persisted"):
1040 return
1041 try:
1042 from ...chat.service import ChatService
1044 if partial_content:
1045 content = partial_content + _STOPPED_FOOTER
1046 if truncated:
1047 content += " _(output was very long; truncated.)_"
1048 else:
1049 content = _STOPPED_BEFORE_OUTPUT_MARKER
1051 # allow_archived=True: same rationale as the completion path —
1052 # the partial response on Stop is system-generated and must
1053 # survive a concurrent archive (see _save_chat_message_and_context).
1054 ChatService(username).add_message(
1055 session_id=chat_session_id,
1056 role="assistant",
1057 content=content,
1058 message_type="response",
1059 research_id=research_id,
1060 allow_archived=True,
1061 )
1062 # Set the idempotency flag ONLY after the write succeeds. If
1063 # `add_message` raises (DB lock, encryption error, archived
1064 # session), the outer ResearchTerminatedException handler retries
1065 # this helper — flipping the flag pre-write would short-circuit
1066 # the retry and silently lose the partial response.
1067 if streaming_state is not None:
1068 streaming_state["_persisted"] = True
1069 logger.info(
1070 f"Persisted partial chat response for terminated research "
1071 f"{research_id} ({len(content)} chars)"
1072 )
1073 except Exception:
1074 logger.opt(exception=True).warning(
1075 "Failed to persist partial chat message on terminate"
1076 )
1077 try:
1078 flush = streaming_state.get("_flush_carry") if streaming_state else None
1079 if flush:
1080 leftover = flush()
1081 if leftover:
1082 try:
1083 _socket_emitter.emit_to_subscribers(
1084 "response_chunk",
1085 research_id,
1086 {
1087 "chunk": leftover,
1088 "is_streaming": True,
1089 "is_final": False,
1090 },
1091 owner=username,
1092 )
1093 except Exception:
1094 logger.debug(
1095 "Carry-buffer flush emit on terminate failed (non-critical)"
1096 )
1098 _socket_emitter.emit_to_subscribers(
1099 "response_chunk",
1100 research_id,
1101 {"chunk": "", "is_streaming": True, "is_final": True},
1102 owner=username,
1103 )
1104 except Exception:
1105 logger.debug("Final-chunk emit on terminate failed (non-critical)")
1108@log_for_research
1109@thread_cleanup
1110def run_research_process(research_id, query, mode, **kwargs):
1111 """
1112 Run the research process in the background for a given research ID.
1114 Args:
1115 research_id: The ID of the research
1116 query: The research query
1117 mode: The research mode (quick/detailed)
1118 **kwargs: Additional parameters for the research (model_provider, model, search_engine, etc.)
1119 MUST include 'username' for database access
1120 """
1121 from ..research_state import (
1122 is_research_active,
1123 is_termination_requested,
1124 update_progress_and_check_active,
1125 )
1127 # Extract username - required for database access
1128 username = kwargs.get("username")
1129 if not username:
1130 logger.error("No username provided to research thread")
1131 raise ValueError("Username is required for research process")
1132 # Extract user_password early so it's available for all cleanup paths
1133 user_password = kwargs.get("user_password")
1135 # Establish thread context FIRST so every subsequent log line in this
1136 # thread can be attributed to the correct user/research and persisted
1137 # to the user's encrypted ResearchLog. Otherwise the early INFO logs
1138 # below ("Research thread started", "Research strategy", "Research
1139 # parameters") fire before start_research_process gets to its own
1140 # set_search_context call (~line 417) and the daemon can't open the
1141 # encrypted DB to write them — silently dropped via the bare-except.
1142 set_search_context(
1143 {
1144 "research_id": research_id,
1145 "username": username,
1146 "user_password": user_password,
1147 # Settings snapshot may be absent here (early init); the
1148 # shared context built later overwrites this with the
1149 # populated snapshot. Defaults to {} so cache reads in any
1150 # intermediate code path see a non-None dict.
1151 "settings_snapshot": kwargs.get("settings_snapshot") or {},
1152 }
1153 )
1155 logger.info(f"Research thread started with username: {username}")
1157 try:
1158 # Check if this research has been terminated before we even start
1159 if is_termination_requested(research_id):
1160 logger.info(
1161 f"Research {research_id} was terminated before starting"
1162 )
1163 cleanup_research_resources(
1164 research_id,
1165 username,
1166 user_password=user_password,
1167 final_status=ResearchStatus.SUSPENDED,
1168 )
1169 return
1171 # (#6938: the query reaches this thread unvalidated from the
1172 # request body, so it previews instead of f-stringing — same
1173 # conversion as the parameter lines below. Imported at call time
1174 # for the same reason as in save_research_strategy above.)
1175 from ..routers.notes import _log_value_preview
1176 from ...security.ssrf_validator import redact_url_for_log
1178 logger.info(
1179 "Starting research process for ID {}, query: {}",
1180 research_id,
1181 _log_value_preview(query),
1182 )
1184 # Extract key parameters
1185 model_provider = kwargs.get("model_provider")
1186 model = kwargs.get("model")
1187 custom_endpoint = kwargs.get("custom_endpoint")
1188 search_engine = kwargs.get("search_engine")
1189 max_results = kwargs.get("max_results")
1190 time_period = kwargs.get("time_period")
1191 iterations = kwargs.get("iterations")
1192 questions_per_iteration = kwargs.get("questions_per_iteration")
1193 strategy = kwargs.get("strategy")
1194 settings_snapshot = kwargs.get(
1195 "settings_snapshot", {}
1196 ) # Complete settings snapshot
1197 # NOTE: the "_username" injection that lets the agent register the
1198 # user's document collections is done downstream in
1199 # AdvancedSearchSystem.__init__ (the narrowest common consumer of the
1200 # strategy-running paths — web run and the programmatic API), not here.
1201 # Run-start egress (below) doesn't need it: it passes username
1202 # explicitly. See ensure_snapshot_username in search_system.py.
1204 # Log settings snapshot to debug
1205 from ...settings.logger import log_settings
1207 log_settings(settings_snapshot, "Settings snapshot received in thread")
1209 # Strategy should already be saved in the database before thread starts
1210 # (#6938: thread-side parameters reach here unvalidated from the
1211 # request body, so both lines preview them instead of f-stringing —
1212 # same conversion #6756 made on the save path; the helper is
1213 # imported at function entry above.)
1214 logger.info("Research strategy: {}", _log_value_preview(strategy))
1216 # Log all parameters for debugging
1217 logger.info(
1218 "Research parameters: provider={}, model={}, "
1219 "search_engine={}, max_results={}, "
1220 "time_period={}, iterations={}, "
1221 "questions_per_iteration={}, custom_endpoint={}, strategy={}",
1222 _log_value_preview(model_provider),
1223 _log_value_preview(model),
1224 _log_value_preview(search_engine),
1225 _log_value_preview(max_results),
1226 _log_value_preview(time_period),
1227 _log_value_preview(iterations),
1228 _log_value_preview(questions_per_iteration),
1229 # The endpoint URL can carry a credential (``?key=``, userinfo),
1230 # so log only its scheme://host:port; the preview's cut must
1231 # never run on the raw URL.
1232 redact_url_for_log(custom_endpoint)
1233 if isinstance(custom_endpoint, str)
1234 else _log_value_preview(custom_endpoint),
1235 _log_value_preview(strategy),
1236 )
1238 # Reapply current operator environment policy at actual dispatch time.
1239 # Queued snapshots may predate a restart or configuration change; every
1240 # downstream consumer must see this same effective copy.
1241 settings_snapshot = apply_environment_overrides_to_snapshot(
1242 settings_snapshot
1243 )
1244 overlay_runtime_settings(
1245 settings_snapshot,
1246 search_engine=search_engine,
1247 model_provider=model_provider,
1248 model=model,
1249 custom_endpoint=custom_endpoint,
1250 iterations=iterations,
1251 questions_per_iteration=questions_per_iteration,
1252 strategy=strategy,
1253 max_results=max_results,
1254 time_period=time_period,
1255 )
1257 # Create a settings context that uses snapshot - no database access in threads
1258 settings_context = SnapshotSettingsContext(
1259 settings_snapshot, username=username
1260 )
1262 if "model_provider" not in kwargs:
1263 model_provider = settings_context.get_setting("llm.provider")
1264 if "model" not in kwargs:
1265 model = settings_context.get_setting("llm.model")
1266 if (
1267 "custom_endpoint" not in kwargs
1268 and normalize_provider(model_provider) == "openai_endpoint"
1269 ):
1270 custom_endpoint = settings_context.get_setting(
1271 "llm.openai_endpoint.url"
1272 )
1273 if "search_engine" not in kwargs:
1274 search_engine = settings_context.get_setting("search.tool")
1275 if "iterations" not in kwargs:
1276 iterations = settings_context.get_setting("search.iterations")
1277 if "questions_per_iteration" not in kwargs:
1278 questions_per_iteration = settings_context.get_setting(
1279 "search.questions_per_iteration"
1280 )
1281 if "strategy" not in kwargs:
1282 strategy = settings_context.get_setting(
1283 "search.search_strategy", "source-based"
1284 )
1286 # Only log settings if explicitly enabled via LDR_LOG_SETTINGS env var
1287 from ...settings.logger import log_settings
1289 log_settings(
1290 settings_context.values, "SettingsContext values extracted"
1291 )
1293 # Set the settings context for this thread
1294 from ...config.thread_settings import (
1295 set_settings_context,
1296 )
1298 set_settings_context(settings_context)
1300 # Defense-in-depth: register an EgressContext on the audit-hook
1301 # thread-local so socket.connect calls that bypass our explicit
1302 # PEPs (third-party libraries, future contributors reaching for
1303 # raw requests.get, prompt-injection-steered tools) are gated
1304 # against the same scope. Cleared by the @thread_cleanup exit
1305 # handler in database/thread_local_session.py when the worker
1306 # winds down. If snapshot construction errors here, we log and
1307 # continue — the explicit PEPs still gate the known callers, so
1308 # the run is not less secure than before this hook existed.
1309 try:
1310 from ...security import set_active_context as _set_egress_ctx
1311 from ...security.egress.policy import (
1312 PolicyDeniedError,
1313 build_run_egress_context,
1314 )
1316 try:
1317 _egress_ctx = build_run_egress_context(
1318 settings_snapshot, search_engine, username=username
1319 )
1320 _egress_primary = _egress_ctx.primary_engine
1321 _set_egress_ctx(_egress_ctx)
1322 except (PolicyDeniedError, ValueError) as _ctx_err:
1323 # Corrupted scope or invalid policy config. The run-start
1324 # precheck already rejects this at the API boundary;
1325 # if we get here it is because the run was launched
1326 # without going through the precheck (CLI, scheduler).
1327 # Re-raise so the worker fails fast with a clear reason
1328 # rather than running unprotected.
1329 logger.bind(policy_audit=True).warning(
1330 "research worker refused: egress policy unevaluable",
1331 research_id=research_id,
1332 reason=str(_ctx_err),
1333 )
1334 raise
1336 # Stage C (ADR-0007): enforce the two-axis (sensitivity x exposure)
1337 # rule at this shared chokepoint too. The /api/start_research route
1338 # applies it before the run starts, but follow-up / chat / queue
1339 # runs reach the worker without that precheck — so evaluate it here,
1340 # where every entry point flows through, reusing the context
1341 # resolved just above. UNPROTECTED evaluates permissively (never
1342 # blocks); on an internal error audit_run fails CLOSED (a denying
1343 # decision) so the run is refused, not silently degraded. A refusal
1344 # raises PolicyDeniedError, failing the run like the
1345 # unevaluable-config path.
1346 from ...security.egress.run_classification import (
1347 audit_run_from_snapshot,
1348 )
1350 _two_axis = audit_run_from_snapshot(
1351 settings_snapshot,
1352 _egress_ctx,
1353 _egress_primary,
1354 llm_provider=model_provider,
1355 )
1356 if not _two_axis.allowed: 1356 ↛ 1357line 1356 didn't jump to line 1357 because the condition on line 1356 was never true
1357 logger.bind(policy_audit=True).warning(
1358 "research worker refused: two-axis egress",
1359 research_id=research_id,
1360 reason=_two_axis.reason,
1361 offending=_two_axis.offending,
1362 )
1363 raise PolicyDeniedError(_two_axis, target=_egress_primary)
1364 except ImportError:
1365 # security module unavailable — preserve legacy behaviour.
1366 logger.debug("egress audit hook unavailable for research worker")
1368 # user_password already extracted above (before termination check)
1370 # Create shared research context that can be updated during research
1371 shared_research_context = {
1372 "research_id": research_id,
1373 "research_query": query,
1374 "research_mode": mode,
1375 "research_phase": "init",
1376 "search_iteration": 0,
1377 "search_engines_planned": None,
1378 "search_engine_selected": search_engine,
1379 "username": username, # Add username for queue operations
1380 "user_password": user_password, # Add password for metrics access
1381 # Thread-safe settings snapshot propagated to background search
1382 # threads (engine config, per-user resolution, egress scope).
1383 "settings_snapshot": settings_snapshot,
1384 "chat_session_id": kwargs.get("chat_session_id"),
1385 }
1387 # If this is a follow-up research, include the parent context. It is
1388 # read back from JSON, so the context and its findings may be of any
1389 # JSON type.
1390 parent_context = kwargs.get("research_context")
1391 if isinstance(parent_context, dict) and parent_context:
1392 past_findings = parent_context.get("past_findings", "")
1393 findings_chars = (
1394 len(past_findings) if isinstance(past_findings, str) else 0
1395 )
1396 logger.info(
1397 f"Adding parent research context with {findings_chars} chars of findings"
1398 )
1399 shared_research_context.update(parent_context)
1400 elif parent_context:
1401 logger.warning(
1402 f"Ignoring non-dict research_context "
1403 f"{type(parent_context).__name__}"
1404 )
1406 # Do not log context keys as they may contain sensitive information
1407 logger.info(f"Created shared_research_context for user: {username}")
1409 # Set search context for search tracking
1410 set_search_context(shared_research_context)
1412 # Per-research dedup state for step message persistence
1413 last_step_phase = None
1415 # Pre-bind streaming_state so progress_callback's closure cell has a
1416 # value even if an exception fires between this point and the full
1417 # initialization later in this function. Without this, the except handler's
1418 # call to progress_callback during a concurrent termination raises
1419 # UnboundLocalError, silently skipping the DB FAILED update and the
1420 # error socket emit. The full streaming_state is reassigned below
1421 # before any real streaming starts; keys must match the canonical shape.
1422 streaming_state: dict = {
1423 "chunks_sent": 0,
1424 "chunks": [],
1425 "_bytes": 0,
1426 "_truncated": False,
1427 }
1428 streaming_enabled = False
1430 # Set up progress callback
1431 def progress_callback(message, progress_percent, metadata):
1432 nonlocal last_step_phase
1433 # Frequent termination check
1434 if is_termination_requested(research_id):
1435 # Persist the partial chat row + emit final chunk BEFORE
1436 # handle_termination — afterwards the Socket.IO room is gone.
1437 _save_partial_chat_message_on_terminate(
1438 shared_research_context.get("chat_session_id"),
1439 research_id,
1440 username,
1441 "".join(streaming_state.get("chunks", [])),
1442 truncated=streaming_state.get("_truncated", False),
1443 streaming_state=streaming_state,
1444 )
1445 handle_termination(research_id, username)
1446 streaming_state["_termination_handled"] = True
1447 raise ResearchTerminatedException( # noqa: TRY301 — inside nested callback, not caught by enclosing try
1448 "Research was terminated by user"
1449 )
1451 # Silent phase — no UI logging or socket emission needed
1452 if metadata.get("phase") == "termination_check":
1453 return
1455 # Bind research_id AND username so the database_sink + queue
1456 # daemon can resolve the per-user encrypted DB. Without username
1457 # the daemon's _write_log_to_database hits "No authenticated
1458 # user", silently swallows the error, and ResearchLog ends up
1459 # with zero milestone rows — leaving /api/research/<id>/status
1460 # without a log_entry to render and the frontend stuck on the
1461 # "Performing research..." fallback.
1462 bound_logger = logger.bind(
1463 research_id=research_id, username=username
1464 )
1465 bound_logger.log("MILESTONE", message)
1467 if "SEARCH_PLAN:" in message:
1468 engines = message.split("SEARCH_PLAN:")[1].strip()
1469 metadata["planned_engines"] = engines
1470 metadata["phase"] = "search_planning" # Use existing phase
1471 # Update shared context for token tracking
1472 shared_research_context["search_engines_planned"] = engines
1473 shared_research_context["research_phase"] = "search_planning"
1475 if "ENGINE_SELECTED:" in message:
1476 engine = message.split("ENGINE_SELECTED:")[1].strip()
1477 metadata["selected_engine"] = engine
1478 metadata["phase"] = "search" # Use existing 'search' phase
1479 # Update shared context for token tracking
1480 shared_research_context["search_engine_selected"] = engine
1481 shared_research_context["research_phase"] = "search"
1483 # Capture other research phases for better context tracking
1484 if metadata.get("phase"): 1484 ↛ 1488line 1484 didn't jump to line 1488 because the condition on line 1484 was always true
1485 shared_research_context["research_phase"] = metadata["phase"]
1487 # Update search iteration if available
1488 if "iteration" in metadata:
1489 shared_research_context["search_iteration"] = metadata[
1490 "iteration"
1491 ]
1493 # Only run_research_process's own final emit carries this flag;
1494 # pop it so it never reaches the socket payload or chat step.
1495 run_complete = metadata.pop(_RUN_COMPLETE_FLAG, False) is True
1497 # Adjust progress based on research mode
1498 adjusted_progress = progress_percent
1499 phase = metadata.get("phase", "")
1501 if run_complete:
1502 # The run's own completion emit is never remapped or capped.
1503 pass
1504 elif mode == "detailed":
1505 # Report phases pass through (already mapped by wrapper).
1506 # All other phases — including "complete" emitted by each
1507 # strategy when its analyze_topic finishes (report
1508 # generation runs analyze_topic per subsection, so a
1509 # strategy "complete" fires mid-report and must NOT be
1510 # treated as the end of the whole run) — are capped.
1511 # update_progress_and_check_active enforces global
1512 # monotonicity, so backwards jumps from non-monotonic
1513 # strategy updates are absorbed there.
1514 # None values (error path, sub-search relays) pass through
1515 # unchanged.
1516 if phase not in _REPORT_PHASES and progress_percent is not None:
1517 adjusted_progress = min(
1518 _DETAILED_SEARCH_PROGRESS_CAP, progress_percent
1519 )
1520 elif (
1521 mode == "quick"
1522 and phase == "output_generation"
1523 and progress_percent is not None
1524 ):
1525 # For quick mode, scale output_generation to 85-95% range
1526 if progress_percent > 0:
1527 adjusted_progress = (
1528 _QUICK_OUTPUT_PROGRESS_START
1529 + (progress_percent / 100) * 10
1530 )
1531 else:
1532 adjusted_progress = _QUICK_OUTPUT_PROGRESS_START
1533 elif (
1534 mode == "quick"
1535 and phase == "complete"
1536 and progress_percent is not None
1537 ):
1538 # The strategy finished searching; answer generation (the
1539 # 85-95% band) is still ahead.
1540 adjusted_progress = min(
1541 _QUICK_OUTPUT_PROGRESS_START, progress_percent
1542 )
1544 # Reserve 100 for the run's completion signal. Written as a
1545 # comparison (not min()) so a NaN passes through unchanged and
1546 # is rejected by the registry instead of becoming 99.
1547 if (
1548 not run_complete
1549 and adjusted_progress is not None
1550 and adjusted_progress > _IN_PROGRESS_CAP
1551 ):
1552 adjusted_progress = _IN_PROGRESS_CAP
1554 # Atomically update progress and check if research is still active
1555 if adjusted_progress is not None:
1556 adjusted_progress, still_active = (
1557 update_progress_and_check_active(
1558 research_id, adjusted_progress
1559 )
1560 )
1561 else:
1562 still_active = is_research_active(research_id)
1564 if still_active:
1565 # Queue the progress update to be processed in main thread
1566 if adjusted_progress is not None:
1567 from ..queue.processor_v2 import queue_processor
1569 if username: 1569 ↛ 1574line 1569 didn't jump to line 1574 because the condition on line 1569 was always true
1570 queue_processor.queue_progress_update(
1571 username, research_id, adjusted_progress
1572 )
1573 else:
1574 logger.warning(
1575 f"Cannot queue progress update for research {research_id} - no username available"
1576 )
1578 # Determine socket emit throttling. Completion is the
1579 # explicit run_complete signal -- never a progress value
1580 # (a clamped overshoot used to reach 100 mid-run) and never
1581 # a bare phase="complete" (strategies emit that mid-run).
1582 phase = metadata.get("phase", "")
1583 is_final = run_complete or phase in (
1584 "error",
1585 "report_complete",
1586 )
1588 should_emit = is_final
1589 if not is_final:
1590 now = time.monotonic()
1591 with _last_emit_lock:
1592 last = _last_emit_times.get(research_id, 0)
1593 if now - last >= _EMIT_THROTTLE_SECONDS:
1594 _last_emit_times[research_id] = now
1595 should_emit = True
1596 # Periodic TTL cleanup for orphaned entries
1597 global _emit_cleanup_counter # noqa: PLW0603
1598 _emit_cleanup_counter += 1
1599 if _emit_cleanup_counter % 100 == 0:
1600 stale = [
1601 rid
1602 for rid, t in _last_emit_times.items()
1603 if now - t > _EMIT_TTL_SECONDS
1604 ]
1605 for rid in stale: 1605 ↛ 1606line 1605 didn't jump to line 1606 because the loop on line 1605 never started
1606 del _last_emit_times[rid]
1608 # Decide chat-step persistence BEFORE building the event:
1609 # a step that will be persisted must always emit, even
1610 # inside the throttle window above — otherwise a burst of
1611 # back-to-back observations (the agent issuing several tool
1612 # calls in one turn) writes rows the live UI never showed,
1613 # breaking the symmetry invariant in the live direction.
1614 #
1615 # Symmetry invariant: for chat sessions, what the user sees
1616 # live must equal what `loadSession` reconstructs on reload.
1617 # If dedup blocks persistence of a repeat step phase, drop
1618 # the socket emit too (unless this is a final-phase event
1619 # the completion handler depends on) so live UI doesn't
1620 # surface events that vanish on reload.
1621 _chat_session_id = shared_research_context.get(
1622 "chat_session_id"
1623 )
1624 _persist = False
1625 _suppress_emit = False
1626 if _chat_session_id:
1627 _persist, _suppress_emit = _chat_step_decision(
1628 phase, last_step_phase, is_final
1629 )
1630 if _persist:
1631 should_emit = True
1633 # Build event data (before emit, shared scope)
1634 event_data = None
1635 if should_emit:
1636 event_data = {
1637 "progress": adjusted_progress,
1638 "message": message,
1639 "phase": phase,
1640 }
1641 # Include additional metadata for MCP/ReAct strategy display
1642 if metadata.get("thought"): 1642 ↛ 1643line 1642 didn't jump to line 1643 because the condition on line 1642 was never true
1643 event_data["thought"] = metadata["thought"]
1644 if metadata.get("tool"):
1645 event_data["tool"] = metadata["tool"]
1646 if metadata.get("arguments"): 1646 ↛ 1647line 1646 didn't jump to line 1647 because the condition on line 1646 was never true
1647 event_data["arguments"] = metadata["arguments"]
1648 if metadata.get("iteration"): 1648 ↛ 1649line 1648 didn't jump to line 1649 because the condition on line 1648 was never true
1649 event_data["iteration"] = metadata["iteration"]
1650 if metadata.get("error"):
1651 event_data["error"] = metadata["error"]
1652 if metadata.get("content"):
1653 event_data["content"] = metadata["content"]
1655 # Persist step message to chat session BEFORE emitting socket
1656 # (ensures DB has the step when clients receive the event).
1657 if _persist:
1658 try:
1659 from ...chat.service import ChatService
1661 ChatService(username).add_progress_step(
1662 session_id=_chat_session_id,
1663 research_id=research_id,
1664 content=_compose_chat_step_content(
1665 message, phase, metadata
1666 ),
1667 phase=phase,
1668 )
1669 last_step_phase = phase
1670 except Exception:
1671 logger.opt(exception=True).warning(
1672 "Failed to persist progress step"
1673 )
1674 # Symmetry invariant: if persistence failed
1675 # (e.g. OperationalError under DB contention),
1676 # suppress the live emit too — otherwise the
1677 # client sees a step that vanishes on reload.
1678 event_data = None
1679 if _suppress_emit:
1680 event_data = None
1682 # Emit socket event AFTER DB persistence
1683 if event_data is not None:
1684 try:
1685 _sio_emit(
1686 "research_progress",
1687 research_id,
1688 event_data,
1689 owner=username,
1690 )
1691 except Exception:
1692 logger.exception("Socket emit error (non-critical)")
1694 # Function to check termination during long-running operations
1695 def check_termination():
1696 if is_termination_requested(research_id):
1697 _save_partial_chat_message_on_terminate(
1698 shared_research_context.get("chat_session_id"),
1699 research_id,
1700 username,
1701 "".join(streaming_state.get("chunks", [])),
1702 truncated=streaming_state.get("_truncated", False),
1703 streaming_state=streaming_state,
1704 )
1705 handle_termination(research_id, username)
1706 streaming_state["_termination_handled"] = True
1707 raise ResearchTerminatedException( # noqa: TRY301 — inside nested callback, not caught by enclosing try
1708 "Research was terminated by user during long-running operation"
1709 )
1710 return False # Not terminated
1712 # Configure the system with the specified parameters
1713 use_llm = None
1714 if model or search_engine or model_provider: 1714 ↛ 1727line 1714 didn't jump to line 1727 because the condition on line 1714 was always true
1715 # Log that we're overriding system settings
1716 # (#6938: same request-derived values as the parameter lines
1717 # above, so they preview instead of f-stringing.)
1718 logger.info(
1719 "Overriding system settings with: provider={}, "
1720 "model={}, search_engine={}",
1721 _log_value_preview(model_provider),
1722 _log_value_preview(model),
1723 _log_value_preview(search_engine),
1724 )
1726 # Override LLM if model or model_provider specified
1727 if model or model_provider:
1728 try:
1729 # Get LLM with the overridden settings.
1730 # Pass settings_snapshot explicitly — without it the LLM
1731 # PEP (llm_config.py:295) skips evaluate_llm_endpoint
1732 # because the gate is "if settings_snapshot is not None
1733 # and provider:". The snapshot is already stored on
1734 # shared_research_context as "settings_snapshot".
1735 use_llm = get_llm(
1736 model_name=model,
1737 provider=model_provider,
1738 openai_endpoint_url=custom_endpoint,
1739 research_id=research_id,
1740 research_context=shared_research_context,
1741 settings_snapshot=shared_research_context.get(
1742 "settings_snapshot"
1743 ),
1744 # Thread the run's user so the LLM PEP can resolve a
1745 # per-user registered provider and enforce local-only
1746 # policy. The shared_research_context snapshot is captured
1747 # before AdvancedSearchSystem injects `_username`, so
1748 # without this the whole overridden run (question gen,
1749 # summaries, final report) could fail open to a cloud LLM
1750 # for a user whose primary is a private retriever.
1751 username=username,
1752 )
1754 logger.info(
1755 "Successfully set LLM to: provider={}, model={}",
1756 _log_value_preview(model_provider),
1757 _log_value_preview(model),
1758 )
1759 except Exception as e:
1760 logger.exception(
1761 "Error setting LLM provider={}, model={}",
1762 _log_value_preview(model_provider),
1763 _log_value_preview(model),
1764 )
1765 error_msg = str(e)
1766 # Surface configuration errors to user instead of silently continuing
1767 config_error_keywords = [
1768 "model path",
1769 "llamacpp",
1770 "cannot connect",
1771 "server",
1772 "not configured",
1773 "not responding",
1774 "directory",
1775 ".gguf",
1776 ]
1777 if any(
1778 keyword in error_msg.lower()
1779 for keyword in config_error_keywords
1780 ):
1781 # This is a configuration error the user can fix.
1782 #
1783 # Preserve ``LLMNotConfiguredError`` when that is what
1784 # ``get_llm`` raised. The keyword match below is a
1785 # string test over ``str(e)``, and re-raising a plain
1786 # ``ValueError`` threw the subclass away: a caller that
1787 # distinguishes "no model configured" from "bad request"
1788 # BY TYPE (see web/routers/notes.py's
1789 # ``_llm_not_configured_response``) could no longer tell
1790 # them apart downstream of this wrapper, and a fresh
1791 # install's missing ``llm.model`` -- which matches the
1792 # "not configured" keyword -- was exactly the case that
1793 # lost it.
1794 #
1795 # Only *preserved*, never *added*: several keywords here
1796 # ("cannot connect", "server", "not responding") match
1797 # runtime failures of a provider that IS configured, and
1798 # promoting those to ``LLMNotConfiguredError`` would make
1799 # a boundary answer "open Settings and choose a model" to
1800 # a user whose model is already chosen. Both branches
1801 # stay ``ValueError`` instances, so every existing
1802 # ``except ValueError`` caller is unaffected.
1803 error_type = (
1804 LLMNotConfiguredError
1805 if isinstance(e, LLMNotConfiguredError)
1806 else ValueError
1807 )
1808 raise error_type(
1809 f"LLM Configuration Error: {error_msg}"
1810 ) from e
1811 # For other errors, re-raise to avoid silent failures
1812 raise
1814 # Create search engine first if specified, to avoid default creation without username
1815 use_search = None
1816 if search_engine: 1816 ↛ 1862line 1816 didn't jump to line 1862 because the condition on line 1816 was always true
1817 try:
1818 # Create a new search object with these settings
1819 use_search = get_search(
1820 search_tool=search_engine,
1821 llm_instance=use_llm,
1822 username=username,
1823 settings_snapshot=settings_snapshot,
1824 )
1825 logger.info(
1826 "Successfully created search engine: {}",
1827 _log_value_preview(search_engine),
1828 )
1829 except Exception as e:
1830 logger.exception(
1831 "Error creating search engine {}",
1832 _log_value_preview(search_engine),
1833 )
1834 error_msg = str(e)
1835 # Surface configuration errors to user instead of silently continuing
1836 config_error_keywords = [
1837 "searxng",
1838 "instance_url",
1839 "api_key",
1840 "cannot connect",
1841 "connection",
1842 "timeout",
1843 "not configured",
1844 ]
1845 if any(
1846 keyword in error_msg.lower()
1847 for keyword in config_error_keywords
1848 ):
1849 # This is a configuration error the user can fix.
1850 # (#6938: the engine name is request-derived, so the
1851 # raised message carries its bounded preview — the
1852 # ValueError type and the keyword gate above it are
1853 # unchanged.)
1854 raise ValueError(
1855 f"Search Engine Configuration Error "
1856 f"({_log_value_preview(search_engine)}): {error_msg}"
1857 ) from e
1858 # For other errors, re-raise to avoid silent failures
1859 raise
1861 # Set the progress callback in the system
1862 system = AdvancedSearchSystem(
1863 llm=use_llm, # type: ignore[arg-type]
1864 search=use_search, # type: ignore[arg-type]
1865 strategy_name=strategy,
1866 max_iterations=iterations,
1867 questions_per_iteration=questions_per_iteration,
1868 username=username,
1869 settings_snapshot=settings_snapshot,
1870 research_id=research_id,
1871 research_context=shared_research_context,
1872 )
1873 system.set_progress_callback(progress_callback)
1875 # Chat mode: set up LLM streaming callback for real-time response chunks
1876 streaming_enabled = False
1877 # chunks: server-side buffer so partial content survives termination
1878 # (the citation handler's local list is discarded on raise).
1879 # _bytes / _truncated cap the buffer to bound memory on pathologically
1880 # long answers; once capped we still forward to the frontend but stop
1881 # accumulating server-side.
1882 streaming_state = {
1883 "chunks_sent": 0,
1884 "chunks": [],
1885 "_bytes": 0,
1886 "_truncated": False,
1887 }
1888 chat_session_id = shared_research_context.get("chat_session_id")
1890 if chat_session_id:
1891 try:
1892 socket_service = _socket_emitter
1894 # Source resolver returns the strategy's currently-collected
1895 # source list. Late-bound so the streaming callback can
1896 # apply inline hyperlinks as the agent finishes adding to
1897 # all_links_of_system (sources may still be growing when
1898 # synthesis starts; reading via the closure picks up the
1899 # final list at chunk-emit time, not callback-build time).
1900 def _resolve_sources():
1901 if not hasattr(system, "all_links_of_system"): 1901 ↛ 1902line 1901 didn't jump to line 1902 because the condition on line 1901 was never true
1902 return []
1903 return list(system.all_links_of_system or [])
1905 # Build a formatter matching the user's report.citation_format
1906 # so live-display brackets ([[arxiv.org-1]] / [[arxiv-1]] /
1907 # [[1]] etc.) match what the final-save formatter will emit
1908 # — avoids a visible format-flip when handleResearchComplete
1909 # swaps in the DB-saved version.
1910 live_formatter = get_citation_formatter()
1912 stream_callback = _make_chat_stream_callback(
1913 research_id,
1914 streaming_state,
1915 socket_service,
1916 username,
1917 source_resolver=_resolve_sources,
1918 formatter=live_formatter,
1919 )
1921 # Hook into the citation handler's streaming
1922 if hasattr(system, "strategy") and hasattr(
1923 system.strategy, "citation_handler"
1924 ):
1925 handler = system.strategy.citation_handler
1926 if hasattr(handler, "set_stream_callback"): 1926 ↛ 1941line 1926 didn't jump to line 1941 because the condition on line 1926 was always true
1927 handler.set_stream_callback(stream_callback)
1928 streaming_enabled = True
1929 logger.info(
1930 f"Streaming enabled for chat {chat_session_id[:8]}..."
1931 )
1932 except Exception:
1933 # exception=True so the traceback is visible: streaming is
1934 # non-critical (research still completes without it), but a
1935 # silent warning would hide real bugs in the setup above.
1936 logger.opt(exception=True).warning(
1937 "Could not set up streaming (non-critical)"
1938 )
1940 # Helper to save chat message (closes over outer scope vars)
1941 def _maybe_save_chat_message(content):
1942 _chat_sid = shared_research_context.get("chat_session_id")
1943 if not _chat_sid:
1944 return
1945 try:
1946 _save_chat_message_and_context(
1947 _chat_sid,
1948 research_id,
1949 username,
1950 content,
1951 streaming_enabled,
1952 streaming_state,
1953 _socket_emitter,
1954 settings_snapshot=settings_snapshot,
1955 )
1956 except Exception:
1957 # Promoted from debug→warning: at debug level this is invisible
1958 # in production and the user sees "Research completed but no
1959 # report available" with no operator signal. Same rationale as
1960 # the accumulated_context handler in _save_chat_message_and_context.
1961 logger.opt(exception=True).warning(
1962 "Could not add message to chat session — assistant "
1963 "response NOT persisted; user will see 'no report available'"
1964 )
1965 # The DB write failed, so _save_chat_message_and_context never
1966 # emitted its is_final response_chunk. Emit one here so the
1967 # streaming UI clears its 'thinking' state instead of stalling.
1968 try:
1969 _socket_emitter.emit_to_subscribers(
1970 "response_chunk",
1971 research_id,
1972 {"chunk": "", "is_streaming": True, "is_final": True},
1973 owner=username,
1974 )
1975 except Exception:
1976 logger.opt(exception=True).debug(
1977 "Failed to emit final chunk after chat-persist error"
1978 )
1980 # Run the search
1981 progress_callback("Starting research process", 5, {"phase": "init"})
1983 try:
1984 results = system.analyze_topic(query)
1985 if mode == "quick":
1986 progress_callback(
1987 "Search complete, preparing to generate summary...",
1988 85,
1989 {"phase": "output_generation"},
1990 )
1991 else:
1992 progress_callback(
1993 "Search complete, generating output",
1994 80,
1995 {"phase": "output_generation"},
1996 )
1997 except Exception as search_error:
1998 # Better handling of specific search errors
1999 error_message = str(search_error)
2000 error_type = "unknown"
2002 # OpenAI-compatible runtime failures (LM Studio / vLLM / llama.cpp
2003 # server / OpenRouter / custom endpoint) -- rewrite to a message
2004 # that names the provider, base URL, and model (#3878).
2005 # Normalise first: a legacy upper-case stored provider
2006 # ("QWEN") must still get the rewrite.
2007 rewrite_provider = normalize_provider(model_provider)
2008 if (
2009 rewrite_provider in _OPENAI_COMPAT_RUNTIME_ERROR_PROVIDERS
2010 and is_openai_compat_runtime_error(search_error)
2011 ):
2012 rewritten = friendly_openai_compatible_error(
2013 search_error,
2014 provider=rewrite_provider,
2015 base_url=_resolve_error_rewrite_base_url(
2016 rewrite_provider, custom_endpoint, settings_context
2017 ),
2018 model=model,
2019 )
2020 raise RuntimeError(rewritten) from search_error
2022 # Extract error details for common issues
2023 if "status code: 503" in error_message:
2024 error_message = "Ollama AI service is unavailable (HTTP 503). Please check that Ollama is running properly on your system."
2025 error_type = "ollama_unavailable"
2026 elif "status code: 404" in error_message:
2027 error_message = "Ollama model not found (HTTP 404). Please check that you have pulled the required model."
2028 error_type = "model_not_found"
2029 elif "status code:" in error_message:
2030 # Extract the status code for other HTTP errors
2031 status_code = error_message.split("status code:")[1].strip()
2032 error_message = f"API request failed with status code {status_code}. Please check your configuration."
2033 error_type = "api_error"
2034 elif "connection" in error_message.lower():
2035 error_message = "Connection error. Please check that your LLM service (Ollama/API) is running and accessible."
2036 error_type = "connection_error"
2038 # Raise with improved error message
2039 raise RuntimeError(
2040 f"{error_message} (Error type: {error_type})"
2041 ) from search_error
2043 # Generate output based on mode
2044 if mode == "quick":
2045 # Quick Summary
2046 if results.get("findings") or results.get("formatted_findings"):
2047 raw_formatted_findings = results["formatted_findings"]
2049 # Check if formatted_findings contains an error message
2050 if _is_error_shaped_content(raw_formatted_findings):
2051 logger.error("Detected error in formatted findings")
2053 # Determine error type for better user feedback
2054 error_type = "unknown"
2055 error_message = raw_formatted_findings.lower()
2057 if (
2058 "token limit" in error_message
2059 or "context length" in error_message
2060 ):
2061 error_type = "token_limit"
2062 # Log specific error type
2063 logger.warning(
2064 "Detected token limit error in synthesis"
2065 )
2067 # Update progress with specific error type
2068 progress_callback(
2069 "Synthesis hit token limits. Attempting fallback...",
2070 87,
2071 {
2072 "phase": "synthesis_error",
2073 "error_type": error_type,
2074 },
2075 )
2076 elif (
2077 "timeout" in error_message
2078 or "timed out" in error_message
2079 ):
2080 error_type = "timeout"
2081 logger.warning("Detected timeout error in synthesis")
2082 progress_callback(
2083 "Synthesis timed out. Attempting fallback...",
2084 87,
2085 {
2086 "phase": "synthesis_error",
2087 "error_type": error_type,
2088 },
2089 )
2090 elif "rate limit" in error_message:
2091 error_type = "rate_limit"
2092 logger.warning("Detected rate limit error in synthesis")
2093 progress_callback(
2094 "LLM rate limit reached. Attempting fallback...",
2095 87,
2096 {
2097 "phase": "synthesis_error",
2098 "error_type": error_type,
2099 },
2100 )
2101 elif (
2102 "connection" in error_message
2103 or "network" in error_message
2104 ):
2105 error_type = "connection"
2106 logger.warning("Detected connection error in synthesis")
2107 progress_callback(
2108 "Connection issue with LLM. Attempting fallback...",
2109 87,
2110 {
2111 "phase": "synthesis_error",
2112 "error_type": error_type,
2113 },
2114 )
2115 elif (
2116 "llm error" in error_message
2117 or "final answer synthesis fail" in error_message
2118 ):
2119 error_type = "llm_error"
2120 logger.warning(
2121 "Detected general LLM error in synthesis"
2122 )
2123 progress_callback(
2124 "LLM error during synthesis. Attempting fallback...",
2125 87,
2126 {
2127 "phase": "synthesis_error",
2128 "error_type": error_type,
2129 },
2130 )
2131 else:
2132 # Generic error
2133 logger.warning("Detected unknown error in synthesis")
2134 progress_callback(
2135 "Error during synthesis. Attempting fallback...",
2136 87,
2137 {
2138 "phase": "synthesis_error",
2139 "error_type": "unknown",
2140 },
2141 )
2143 # Extract synthesized content from findings if available
2144 synthesized_content = ""
2145 for finding in results.get("findings", []):
2146 if finding.get("phase") == "Final synthesis":
2147 synthesized_content = finding.get("content", "")
2148 break
2150 # Use synthesized content as fallback
2151 if synthesized_content and not _is_error_shaped_content(
2152 synthesized_content
2153 ):
2154 logger.info(
2155 "Using existing synthesized content as fallback"
2156 )
2157 raw_formatted_findings = synthesized_content
2159 # Or use current_knowledge as another fallback
2160 elif results.get("current_knowledge") and not (
2161 _is_error_shaped_content(results["current_knowledge"])
2162 ):
2163 logger.info("Using current_knowledge as fallback")
2164 raw_formatted_findings = results["current_knowledge"]
2166 # Or combine all finding contents as last resort
2167 elif results.get("findings"):
2168 logger.info("Combining all findings as fallback")
2169 # First try to use any findings that are not errors
2170 valid_findings = [
2171 f"## {finding.get('phase', 'Finding')}\n\n{finding.get('content', '')}"
2172 for finding in results.get("findings", [])
2173 if finding.get("content")
2174 and not _is_error_shaped_content(
2175 finding.get("content")
2176 )
2177 ]
2179 if valid_findings:
2180 raw_formatted_findings = (
2181 "# Research Results (Fallback Mode)\n\n"
2182 )
2183 raw_formatted_findings += "\n\n".join(
2184 valid_findings
2185 )
2186 raw_formatted_findings += "\n\n*Final synthesis failed; showing partial findings.*"
2187 # If every finding is error-shaped, keep the failure
2188 # as private input to ErrorReportGenerator below. Its
2189 # authored replacements retain useful category hints
2190 # without echoing raw provider details into the report.
2192 progress_callback(
2193 f"Using fallback synthesis due to {error_type} error",
2194 88,
2195 {
2196 "phase": "synthesis_fallback",
2197 "error_type": error_type,
2198 },
2199 )
2201 logger.info(
2202 "Found formatted_findings of length: {}",
2203 len(str(raw_formatted_findings)),
2204 )
2206 try:
2207 # Check if we have an error in the findings and use enhanced error handling
2208 if _is_error_shaped_content(raw_formatted_findings):
2209 logger.info(
2210 "Generating enhanced error report using ErrorReportGenerator"
2211 )
2213 # Generate comprehensive error report
2214 # ErrorReportGenerator does not use LLM (kept for compat)
2215 error_generator = ErrorReportGenerator()
2216 clean_markdown = error_generator.generate_error_report(
2217 error_message=raw_formatted_findings,
2218 query=query,
2219 partial_results=results,
2220 search_iterations=results.get("iterations", 0),
2221 research_id=research_id,
2222 )
2224 logger.info(
2225 "Generated enhanced error report with {} characters",
2226 len(clean_markdown),
2227 )
2228 else:
2229 # report_content stores the synthesized answer
2230 # only (see _extract_synthesized_answer for the
2231 # full rationale). Fall back to the formatted
2232 # blob only when neither Final synthesis nor
2233 # current_knowledge is populated — leaks
2234 # sources, but at least we save *something*.
2235 clean_markdown = (
2236 _extract_synthesized_answer(results)
2237 or raw_formatted_findings
2238 )
2240 # Pull sources from the search-system's accumulated link
2241 # buffer first — same source the detailed-report path
2242 # uses at the equivalent point below. Wrapper
2243 # strategies (e.g. EnhancedContextualFollowUpStrategy,
2244 # IterativeRefinementStrategy) delegate the actual
2245 # search to an inner strategy that populates
2246 # `self.all_links_of_system`, but they don't bubble
2247 # that list back into the result dict's `findings`. So
2248 # the legacy `findings[*].search_results` extraction
2249 # below stays empty for chat follow-ups, leaving the
2250 # citation formatter with no urls to hyperlink. Prefer
2251 # the system-level accumulator; fall back to the
2252 # legacy extraction so direct strategies that bypass
2253 # the system buffer still work.
2254 all_links = list(
2255 getattr(system, "all_links_of_system", None) or []
2256 )
2257 if not all_links:
2258 for finding in results.get("findings", []):
2259 search_results = finding.get("search_results", [])
2260 if search_results:
2261 try:
2262 links = extract_links_from_search_results(
2263 search_results
2264 )
2265 all_links.extend(links)
2266 except Exception:
2267 logger.exception(
2268 "Error processing search results/links"
2269 )
2271 logger.info(
2272 "Successfully converted to clean markdown of length: {}",
2273 len(clean_markdown),
2274 )
2276 # First send a progress update for generating the summary
2277 progress_callback(
2278 "Generating clean summary from research data...",
2279 90,
2280 {"phase": "output_generation"},
2281 )
2283 # Send progress update for saving report
2284 progress_callback(
2285 "Saving research report to database...",
2286 95,
2287 {"phase": "report_complete"},
2288 )
2290 # Format citations in the markdown content. The
2291 # split returns the answer-with-hyperlinks half
2292 # separately from the trailing sources section the
2293 # LLM may have emitted (which is discarded — sources
2294 # live in research_resources, the canonical store).
2295 # When no Sources section is found, fall back to
2296 # structured-source hyperlinking — never re-parse
2297 # concatenated formatter output downstream.
2298 formatter = get_citation_formatter()
2299 try:
2300 # Quick-mode (chat summary) input is raw LLM
2301 # output; the report generator never wraps it
2302 # in the appended-sources sentinel, so disable
2303 # sentinel trust. Otherwise an LLM that quotes
2304 # the marker verbatim would silently trip the
2305 # over-strip safety bypass on a path where the
2306 # bypass can only ever fire as a false-positive.
2307 answer_with_links, llm_sources, on_sentinel = (
2308 formatter.format_document_split(
2309 clean_markdown, trust_sentinel=False
2310 )
2311 )
2312 if not llm_sources:
2313 answer_with_links = (
2314 formatter.apply_inline_hyperlinks(
2315 clean_markdown, all_links
2316 )
2317 )
2318 # Safety check: a >50% strip on a long input
2319 # likely means the regex over-stripped on a
2320 # "Sources:" header inside the answer body.
2321 # Fall back to structured-source hyperlinking
2322 # on the full text. Min-length floor prevents
2323 # false-fires on legitimately short answers.
2324 # Skip the check when the boundary came from
2325 # the unique appended-sources sentinel emitted
2326 # by IntegratedReportGenerator — that boundary
2327 # is correct by construction (e.g. a 1380-source
2328 # detailed-mode run legitimately pushes the
2329 # answer/sources ratio below 50% without any
2330 # over-strip having occurred). Quick-mode
2331 # always passes trust_sentinel=False above, so
2332 # ``on_sentinel`` can never be True here.
2333 SAFETY_MIN_LEN = 800
2334 if ( 2334 ↛ 2341line 2334 didn't jump to line 2341 because the condition on line 2334 was never true
2335 not on_sentinel
2336 and llm_sources
2337 and len(clean_markdown) > SAFETY_MIN_LEN
2338 and len(answer_with_links)
2339 < len(clean_markdown) * 0.5
2340 ):
2341 logger.warning(
2342 "format_document_split appears to have "
2343 "over-stripped (answer={} chars, "
2344 "original={} chars) for research {}. "
2345 "Falling back to structured-source "
2346 "hyperlinking on full input.",
2347 len(answer_with_links),
2348 len(clean_markdown),
2349 research_id,
2350 )
2351 answer_with_links = (
2352 formatter.apply_inline_hyperlinks(
2353 clean_markdown, all_links
2354 )
2355 )
2356 except Exception:
2357 # Hyperlinking is quality-of-life, not a hard
2358 # requirement. If anything blows up, save the
2359 # raw LLM text rather than fail the research.
2360 logger.exception(
2361 "Citation formatter failed for research {}; saving raw answer",
2362 research_id,
2363 )
2364 answer_with_links = clean_markdown
2366 # report_content stores ONLY the synthesized answer.
2367 # The legacy "answer + ## Sources + ## Research
2368 # Metrics" view is reconstructed at render time by
2369 # report_assembly_service.assemble_full_report.
2370 full_report_content = answer_with_links
2372 # Save report FIRST, then sources:
2373 # a chat read between commits sees a report with no
2374 # sources (assembler renders just the answer) — better
2375 # failure mode than partial assembly with sources but
2376 # no answer body.
2377 from ...storage import get_report_storage
2379 with get_user_db_session(username) as db_session:
2380 storage = get_report_storage(session=db_session)
2382 # Prepare metadata
2383 metadata = {
2384 "iterations": results["iterations"],
2385 "generated_at": datetime.now(UTC).isoformat(),
2386 }
2388 # Save report using storage abstraction
2389 success = storage.save_report(
2390 research_id=research_id,
2391 content=full_report_content,
2392 metadata=metadata,
2393 username=username,
2394 )
2396 if not success:
2397 raise RuntimeError("Failed to save research report") # noqa: TRY301 — triggers research failure handling in outer except
2399 # Save sources to database (non-fatal - report
2400 # already saved; sources missing is recoverable
2401 # because the assembler omits empty Sources blocks)
2402 try:
2403 from .research_sources_service import (
2404 ResearchSourcesService,
2405 )
2407 sources_service = ResearchSourcesService()
2408 if all_links:
2409 logger.info(
2410 f"Quick summary: Saving {len(all_links)} sources to database"
2411 )
2412 sources_saved = (
2413 sources_service.save_research_sources(
2414 research_id=research_id,
2415 sources=all_links,
2416 username=username,
2417 )
2418 )
2419 logger.info(
2420 f"Quick summary: Saved {sources_saved} sources for research {research_id}"
2421 )
2422 except Exception:
2423 logger.exception(
2424 f"Failed to save sources for research {research_id} (continuing with report save)"
2425 )
2427 logger.info(f"Report saved for research_id: {research_id}")
2429 # Skip export to additional formats - we're storing in database only
2431 # Update research status in database
2432 completed_at = datetime.now(UTC).isoformat()
2434 with get_user_db_session(username) as db_session:
2435 research = (
2436 db_session.query(ResearchHistory)
2437 .filter_by(id=research_id)
2438 .first()
2439 )
2441 # Preserve existing metadata and update with new values
2442 # (snapshot back-fill included — see helper).
2443 metadata = _merge_completion_metadata(
2444 research.research_meta,
2445 {
2446 "iterations": results["iterations"],
2447 "generated_at": datetime.now(UTC).isoformat(),
2448 },
2449 settings_snapshot,
2450 )
2452 # Use the helper function for consistent duration calculation
2453 duration_seconds = calculate_duration(
2454 research.created_at, completed_at
2455 )
2457 research.status = ResearchStatus.COMPLETED
2458 research.progress = 100
2459 research.completed_at = completed_at
2460 research.duration_seconds = duration_seconds
2461 # report_content was already saved above via the report
2462 # storage abstraction; this block only updates
2463 # status/metadata. report_path is not used in the
2464 # encrypted-database version.
2466 # Generate headline and topics only for news searches
2467 if (
2468 metadata.get("is_news_search")
2469 or metadata.get("search_type") == "news_analysis"
2470 ):
2471 try:
2472 from ...news.utils.headline_generator import (
2473 generate_headline,
2474 )
2475 from ...news.utils.topic_generator import (
2476 generate_topics,
2477 )
2478 from ...search_system import (
2479 ensure_snapshot_username,
2480 )
2482 # Scope the snapshot to this run's owner so
2483 # the LLM PEP resolves a per-user provider
2484 # and honors local-only policy, mirroring
2485 # the scheduler's _store_research_result fix
2486 # (background.py). Without it, headline/topic
2487 # generation silently no-ops for cloud-
2488 # provider users here (fails closed, not a
2489 # security hole, but inconsistent).
2490 headline_topic_snapshot = (
2491 ensure_snapshot_username(
2492 settings_snapshot, username
2493 )
2494 )
2496 # Get the report content from database for better headline/topic generation
2497 report_content = ""
2498 try:
2499 research = (
2500 db_session.query(ResearchHistory)
2501 .filter_by(id=research_id)
2502 .first()
2503 )
2504 if research and research.report_content: 2504 ↛ 2510line 2504 didn't jump to line 2510 because the condition on line 2504 was always true
2505 report_content = research.report_content
2506 logger.info(
2507 f"Retrieved {len(report_content)} chars from database for headline generation"
2508 )
2509 else:
2510 logger.warning(
2511 f"No report content found in database for research_id: {research_id}"
2512 )
2513 except Exception:
2514 logger.warning(
2515 "Could not retrieve report content from database"
2516 )
2518 # Generate headline
2519 logger.info(
2520 "Generating headline for query: {}",
2521 # (#6938: redacted before it is cut.)
2522 redact_and_bound_for_log(query, 100),
2523 )
2524 headline = generate_headline(
2525 query,
2526 report_content,
2527 settings_snapshot=headline_topic_snapshot,
2528 )
2529 metadata["generated_headline"] = headline
2531 # Generate topics
2532 # (#6938: the category comes from the
2533 # request body's metadata.)
2534 logger.info(
2535 "Generating topics with category: {}",
2536 _log_value_preview(
2537 metadata.get("category", "News")
2538 ),
2539 )
2540 topics = generate_topics(
2541 query=query,
2542 findings=report_content,
2543 category=metadata.get("category", "News"),
2544 max_topics=6,
2545 settings_snapshot=headline_topic_snapshot,
2546 )
2547 metadata["generated_topics"] = topics
2549 logger.info(f"Generated headline: {headline}")
2550 logger.info(f"Generated topics: {topics}")
2552 except Exception:
2553 logger.warning(
2554 "Could not generate headline/topics"
2555 )
2557 research.research_meta = metadata
2559 db_session.commit()
2560 logger.info(
2561 f"Database commit completed for research_id: {research_id}"
2562 )
2564 # Update subscription if this was triggered by a subscription
2565 if metadata.get("subscription_id"):
2566 try:
2567 from ...news.subscription_runner import (
2568 advance_refresh_schedule_by_id,
2569 )
2571 subscription_id = metadata["subscription_id"]
2572 if advance_refresh_schedule_by_id( 2572 ↛ 2584line 2572 didn't jump to line 2584
2573 db_session, subscription_id
2574 ):
2575 db_session.commit()
2576 logger.info(
2577 f"Updated subscription {subscription_id} refresh times"
2578 )
2579 except Exception:
2580 logger.warning(
2581 "Could not update subscription refresh time"
2582 )
2584 logger.info(
2585 f"Database updated successfully for research_id: {research_id}"
2586 )
2588 _maybe_save_chat_message(full_report_content)
2590 # Send the final completion message
2591 progress_callback(
2592 "Research completed successfully",
2593 100,
2594 {"phase": "complete", _RUN_COMPLETE_FLAG: True},
2595 )
2597 # Clean up resources
2598 logger.info(
2599 "Cleaning up resources for research_id: {}", research_id
2600 )
2601 cleanup_research_resources(
2602 research_id, username, user_password=user_password
2603 )
2604 logger.info(
2605 "Resources cleaned up for research_id: {}", research_id
2606 )
2608 except Exception as inner_e:
2609 logger.exception("Error during quick summary generation")
2610 raise RuntimeError(
2611 f"Error generating quick summary: {inner_e!s}"
2612 )
2613 else:
2614 raise RuntimeError( # noqa: TRY301 — triggers research failure handling in outer except
2615 "No research findings were generated. Please try again."
2616 )
2617 else:
2618 # Full Report
2619 progress_callback(
2620 "Generating detailed report...",
2621 _DETAILED_REPORT_PROGRESS_START,
2622 {"phase": "report_generation"},
2623 )
2625 # Extract the search system from the results if available
2626 search_system = results.get("search_system", None)
2628 # Wrapper that maps report generator's 0-100% to the configured
2629 # detailed-mode range and relays cancellation checks through the
2630 # outer progress_callback
2631 _report_range = (
2632 _DETAILED_REPORT_PROGRESS_END - _DETAILED_REPORT_PROGRESS_START
2633 )
2635 def report_progress_callback(message, progress_percent, metadata):
2636 if progress_percent is not None:
2637 adjusted = (
2638 _DETAILED_REPORT_PROGRESS_START
2639 + (progress_percent / 100) * _report_range
2640 )
2641 else:
2642 adjusted = progress_percent
2643 progress_callback(message, adjusted, metadata)
2645 # Pass the existing search system to maintain citation indices
2646 report_generator = IntegratedReportGenerator(
2647 search_system=search_system,
2648 settings_snapshot=settings_snapshot,
2649 )
2650 final_report = report_generator.generate_report(
2651 results, query, progress_callback=report_progress_callback
2652 )
2654 progress_callback(
2655 "Report generation complete",
2656 _DETAILED_REPORT_PROGRESS_END,
2657 {"phase": "report_complete"},
2658 )
2660 # Format citations and split off the trailing Sources
2661 # section. Save only the answer half — sources are
2662 # persisted structurally to research_resources.
2663 all_links = (
2664 getattr(search_system, "all_links_of_system", None) or []
2665 )
2666 formatter = get_citation_formatter()
2667 try:
2668 # Detailed-mode input comes from
2669 # ``IntegratedReportGenerator._format_final_report``,
2670 # which wraps the appended ``## Sources`` block in the
2671 # unique sentinel. Trust the sentinel so a legitimate
2672 # large-tail run (e.g. 1380 sources) doesn't trip the
2673 # over-strip safety check on the right answer half.
2674 answer_with_links, llm_sources, on_sentinel = (
2675 formatter.format_document_split(final_report["content"])
2676 )
2677 if not llm_sources:
2678 answer_with_links = formatter.apply_inline_hyperlinks(
2679 final_report["content"], all_links
2680 )
2681 # Skip the over-strip safety check when the boundary
2682 # came from the unique appended-sources sentinel emitted
2683 # by IntegratedReportGenerator — that boundary is
2684 # correct by construction (e.g. a 1380-source detailed
2685 # run legitimately pushes the answer/sources ratio
2686 # below 50% without any over-strip having occurred).
2687 SAFETY_MIN_LEN = 800
2688 if (
2689 not on_sentinel
2690 and llm_sources
2691 and len(final_report["content"]) > SAFETY_MIN_LEN
2692 and len(answer_with_links)
2693 < len(final_report["content"]) * 0.5
2694 ):
2695 logger.warning(
2696 "format_document_split appears to have over-stripped "
2697 "(answer={} chars, original={} chars) for research {}.",
2698 len(answer_with_links),
2699 len(final_report["content"]),
2700 research_id,
2701 )
2702 answer_with_links = formatter.apply_inline_hyperlinks(
2703 final_report["content"], all_links
2704 )
2705 except Exception:
2706 logger.exception(
2707 "Citation formatter failed for research {}; saving raw answer",
2708 research_id,
2709 )
2710 answer_with_links = final_report["content"]
2711 # The answer keeps the generator's appended-sources sentinel
2712 # whenever the split above kept the Sources block with it
2713 # (NO_HYPERLINKS mode, or the formatter raised). It is an HTML
2714 # comment, which renderMarkdown() shows as literal text.
2715 formatted_content = strip_appended_sources_sentinel(
2716 answer_with_links
2717 )
2719 # Save report FIRST, sources after.
2720 # See quick-summary path for rationale.
2721 from ...storage import get_report_storage
2723 with get_user_db_session(username) as db_session:
2724 storage = get_report_storage(session=db_session)
2726 # Update metadata. Include generated_at like the quick path so
2727 # the detailed file backup's _metadata.json has parity with
2728 # quick-mode (the detailed path previously omitted it).
2729 metadata = final_report["metadata"]
2730 metadata["iterations"] = results["iterations"]
2731 metadata["generated_at"] = datetime.now(UTC).isoformat()
2733 # Save the report through the storage abstraction, exactly as
2734 # the quick-summary branch does, so detailed reports also honor
2735 # the report.enable_file_backup setting. This previously did a
2736 # raw ORM write of report_content that silently skipped the
2737 # file backup, so a user who enabled file backup got it for
2738 # quick research but never for detailed research.
2739 success = storage.save_report(
2740 research_id=research_id,
2741 content=formatted_content,
2742 metadata=metadata,
2743 username=username,
2744 )
2746 if not success: 2746 ↛ 2747line 2746 didn't jump to line 2747 because the condition on line 2746 was never true
2747 raise RuntimeError("Failed to save research report") # noqa: TRY301 — triggers research failure handling in outer except
2749 logger.info(
2750 f"Report saved to database for research_id: {research_id}"
2751 )
2753 # Save sources AFTER report (non-fatal; assembler omits
2754 # empty Sources blocks if this fails).
2755 try:
2756 from .research_sources_service import ResearchSourcesService
2758 sources_service = ResearchSourcesService()
2759 if all_links:
2760 logger.info(f"Saving {len(all_links)} sources to database")
2761 sources_saved = sources_service.save_research_sources(
2762 research_id=research_id,
2763 sources=all_links,
2764 username=username,
2765 )
2766 logger.info(
2767 f"Saved {sources_saved} sources for research {research_id}"
2768 )
2769 except Exception:
2770 logger.exception(
2771 f"Failed to save sources for research {research_id} (continuing)"
2772 )
2774 # Update research status in database
2775 completed_at = datetime.now(UTC).isoformat()
2777 with get_user_db_session(username) as db_session:
2778 research = (
2779 db_session.query(ResearchHistory)
2780 .filter_by(id=research_id)
2781 .first()
2782 )
2784 # Preserve existing metadata and merge with report metadata
2785 # (snapshot back-fill included — see helper).
2786 completion_extra = dict(final_report["metadata"])
2787 completion_extra["iterations"] = results["iterations"]
2788 metadata = _merge_completion_metadata(
2789 research.research_meta, completion_extra, settings_snapshot
2790 )
2792 # Use the helper function for consistent duration calculation
2793 duration_seconds = calculate_duration(
2794 research.created_at, completed_at
2795 )
2797 research.status = ResearchStatus.COMPLETED
2798 research.progress = 100
2799 research.completed_at = completed_at
2800 research.duration_seconds = duration_seconds
2801 # report_content was already saved above via the report storage
2802 # abstraction; this block only updates status/metadata.
2803 # report_path is not used in the encrypted-database version.
2805 # Generate headline and topics only for news searches
2806 if (
2807 metadata.get("is_news_search")
2808 or metadata.get("search_type") == "news_analysis"
2809 ):
2810 try:
2811 from ...news.utils.headline_generator import (
2812 generate_headline, # type: ignore[no-redef]
2813 )
2814 from ...news.utils.topic_generator import (
2815 generate_topics, # type: ignore[no-redef]
2816 )
2817 from ...search_system import (
2818 ensure_snapshot_username,
2819 )
2821 # Scope the snapshot to this run's owner so the LLM
2822 # PEP resolves a per-user provider and honors
2823 # local-only policy. Mirrors the quick-summary
2824 # branch above and the scheduler's
2825 # _store_research_result fix (background.py).
2826 headline_topic_snapshot = ensure_snapshot_username(
2827 settings_snapshot, username
2828 )
2830 # Get the report content from database for better headline/topic generation
2831 report_content = ""
2832 try:
2833 research = (
2834 db_session.query(ResearchHistory)
2835 .filter_by(id=research_id)
2836 .first()
2837 )
2838 if research and research.report_content: 2838 ↛ 2841line 2838 didn't jump to line 2841 because the condition on line 2838 was always true
2839 report_content = research.report_content
2840 else:
2841 logger.warning(
2842 f"No report content found in database for research_id: {research_id}"
2843 )
2844 except Exception:
2845 logger.warning(
2846 "Could not retrieve report content from database"
2847 )
2849 # Generate headline
2850 headline = generate_headline(
2851 query,
2852 report_content,
2853 settings_snapshot=headline_topic_snapshot,
2854 )
2855 metadata["generated_headline"] = headline
2857 # Generate topics
2858 topics = generate_topics(
2859 query=query,
2860 findings=report_content,
2861 category=metadata.get("category", "News"),
2862 max_topics=6,
2863 settings_snapshot=headline_topic_snapshot,
2864 )
2865 metadata["generated_topics"] = topics
2867 logger.info(f"Generated headline: {headline}")
2868 logger.info(f"Generated topics: {topics}")
2870 except Exception:
2871 logger.warning("Could not generate headline/topics")
2873 research.research_meta = metadata
2875 db_session.commit()
2877 # Update subscription if this was triggered by a subscription
2878 if metadata.get("subscription_id"): 2878 ↛ 2879line 2878 didn't jump to line 2879 because the condition on line 2878 was never true
2879 try:
2880 from ...news.subscription_runner import (
2881 advance_refresh_schedule_by_id,
2882 )
2884 subscription_id = metadata["subscription_id"]
2885 if advance_refresh_schedule_by_id(
2886 db_session, subscription_id
2887 ):
2888 db_session.commit()
2889 logger.info(
2890 f"Updated subscription {subscription_id} refresh times"
2891 )
2892 except Exception:
2893 logger.warning(
2894 "Could not update subscription refresh time"
2895 )
2897 _maybe_save_chat_message(formatted_content)
2899 progress_callback(
2900 "Research completed successfully",
2901 100,
2902 {"phase": "complete", _RUN_COMPLETE_FLAG: True},
2903 )
2905 # Clean up resources
2906 cleanup_research_resources(
2907 research_id, username, user_password=user_password
2908 )
2910 except ResearchTerminatedException:
2911 logger.info(f"Research {research_id} terminated by user")
2912 # Fallback path: when termination was raised from the streaming
2913 # callback (mid-stream interrupt), progress_callback hasn't run
2914 # since the flag was set, so handle_termination() was NOT called
2915 # and the partial row hasn't been persisted yet. The helper is
2916 # idempotent via streaming_state["_persisted"]; if the in-callback
2917 # path already ran, both calls are no-ops.
2918 _save_partial_chat_message_on_terminate(
2919 shared_research_context.get("chat_session_id"),
2920 research_id,
2921 username,
2922 "".join(streaming_state.get("chunks", [])),
2923 truncated=streaming_state.get("_truncated", False),
2924 streaming_state=streaming_state,
2925 )
2926 # Ensure the SUSPENDED status update + cleanup runs even when the
2927 # exception was raised mid-stream. The in-callback termination paths
2928 # set "_termination_handled"; only run here when they did NOT, so a
2929 # single termination doesn't queue two SUSPENDED updates, emit two
2930 # final socket messages, and (in test mode) sleep twice.
2931 if not streaming_state.get("_termination_handled"):
2932 try:
2933 handle_termination(research_id, username)
2934 except Exception:
2935 logger.opt(exception=True).debug(
2936 "handle_termination in except block failed"
2937 )
2939 except Exception as e:
2940 # Handle error
2941 error_message = f"Research failed: {e!s}"
2942 logger.exception(error_message)
2944 try:
2945 # Check for common Ollama error patterns in the exception and provide more user-friendly errors
2946 user_friendly_error = str(e)
2947 error_context = {}
2949 # Egress policy refusal (two-axis rule, or a fail-closed audit_error
2950 # from the worker chokepoint) — surface the reason and the override
2951 # instead of a generic "unexpected error".
2952 from ...security.egress.policy import PolicyDeniedError
2954 if isinstance(e, PolicyDeniedError): 2954 ↛ 2955line 2954 didn't jump to line 2955 because the condition on line 2954 was never true
2955 _reason = getattr(e.decision, "reason", "") or "egress_denied"
2956 if _reason == "audit_error":
2957 user_friendly_error = (
2958 "Egress policy could not verify this run, so it was "
2959 "refused (fail-closed)."
2960 )
2961 elif _reason.startswith("sensitive_to_exposing"):
2962 # The two-axis rule specifically: a sensitive source would
2963 # reach an exposing sink.
2964 user_friendly_error = (
2965 "Egress policy refused this run: a sensitive source "
2966 f"would reach an exposing destination ({_reason})."
2967 )
2968 else:
2969 # Other egress PEPs (scope mismatch, require-local
2970 # inference, etc.) — no "sensitive source" is necessarily
2971 # involved, so don't claim one.
2972 # nosemgrep: semgrep.rules.sql-string-concatenation, reason: Client error message
2973 user_friendly_error = (
2974 f"Egress policy refused this run ({_reason})."
2975 )
2976 error_context = {
2977 "solution": "Set Egress Scope to 'Unprotected' to override, or use local inference / non-sensitive sources."
2978 }
2979 elif "Error type: ollama_unavailable" in user_friendly_error:
2980 user_friendly_error = "Ollama AI service is unavailable. Please check that Ollama is running properly on your system."
2981 error_context = {
2982 "solution": "Start Ollama with 'ollama serve' or check if it's installed correctly."
2983 }
2984 elif "Error type: model_not_found" in user_friendly_error:
2985 user_friendly_error = "Required Ollama model not found. Please pull the model first."
2986 error_context = {
2987 "solution": "Run 'ollama pull mistral' to download the required model."
2988 }
2989 elif "Error type: connection_error" in user_friendly_error:
2990 user_friendly_error = "Connection error with LLM service. Please check that your AI service is running."
2991 error_context = {
2992 "solution": "Ensure Ollama or your API service is running and accessible."
2993 }
2994 elif "Error type: api_error" in user_friendly_error:
2995 user_friendly_error = (
2996 "The language model API rejected the request."
2997 )
2998 error_context = {
2999 "solution": "Check API configuration and credentials."
3000 }
3001 # OpenAI-compatible runtime tokens (#3878). The friendly message
3002 # built by friendly_openai_compatible_error() names the provider,
3003 # base URL and model and appends the raw provider error (only
3004 # credential-scrubbed). The base URL can be a server-level endpoint
3005 # (settings_snapshot bakes in LDR_* env overrides) and the appended
3006 # detail can carry internal hosts/paths, so it must not reach the
3007 # client (CWE-209). Replace it with a safe category message; full
3008 # detail stays in the logs above and the actionable hint is in
3009 # ``solution``.
3010 elif "Error type: openai_connection_refused" in user_friendly_error:
3011 user_friendly_error = TYPED_RESEARCH_ERROR_MESSAGES[
3012 "openai_connection_refused"
3013 ]
3014 error_context = {
3015 "solution": "Check that the configured LLM endpoint is reachable (a local server such as LM Studio, vLLM or llama.cpp must be running) and verify its base URL in Settings -> LLM Providers."
3016 }
3017 elif "Error type: openai_timeout" in user_friendly_error:
3018 user_friendly_error = TYPED_RESEARCH_ERROR_MESSAGES[
3019 "openai_timeout"
3020 ]
3021 error_context = {
3022 "solution": "The endpoint is reachable but slow (the model may still be loading). Retry, or increase the request timeout."
3023 }
3024 elif "Error type: openai_auth" in user_friendly_error:
3025 user_friendly_error = (
3026 "Authentication with the configured LLM provider failed."
3027 )
3028 error_context = {
3029 "solution": "Check that the API key for this provider is valid and set in Settings -> LLM Providers."
3030 }
3031 elif "Error type: openai_permission_denied" in user_friendly_error:
3032 user_friendly_error = (
3033 "The configured LLM provider denied access to the model."
3034 )
3035 error_context = {
3036 "solution": "Your API key is valid but lacks access to this model. Pick a model your account is permitted to use."
3037 }
3038 elif "Error type: openai_model_not_found" in user_friendly_error:
3039 user_friendly_error = TYPED_RESEARCH_ERROR_MESSAGES[
3040 "openai_model_not_found"
3041 ]
3042 error_context = {
3043 "solution": "The provider does not offer (or has not loaded) this model id. Pick a model from the provider's model list."
3044 }
3045 elif "Error type: openai_bad_request" in user_friendly_error:
3046 user_friendly_error = TYPED_RESEARCH_ERROR_MESSAGES[
3047 "openai_bad_request"
3048 ]
3049 error_context = {
3050 "solution": "The provider rejected the request. Check the model id and any provider-specific parameters."
3051 }
3052 elif "Error type: openai_unknown" in user_friendly_error:
3053 user_friendly_error = (
3054 "The configured LLM provider returned an error."
3055 )
3056 error_context = {
3057 "solution": "Check the provider's logs for the full error and verify the base URL / model id."
3058 }
3059 elif "Error type: openai_rate_limit" in user_friendly_error:
3060 user_friendly_error = (
3061 "The LLM provider rate-limited the request."
3062 )
3063 error_context = {
3064 "solution": "The provider rate-limited the request. Wait a moment and retry, or enable LLM Rate Limiting in Settings."
3065 }
3066 elif "LLM Configuration Error:" in user_friendly_error:
3067 # The raw text here is str(e) from LLM setup and can carry
3068 # server-level endpoints/paths (settings_snapshot includes
3069 # LDR_* env overrides), so it must not be surfaced to the client
3070 # (CWE-209) -- it is captured in the logs above. Keep only the
3071 # safe category + actionable hint.
3072 user_friendly_error = (
3073 "There was a problem with the LLM configuration."
3074 )
3075 error_context = {
3076 "solution": "Review your LLM model settings (or, on a shared server, contact your administrator) and ensure they are correct."
3077 }
3078 elif "Search Engine Configuration Error" in user_friendly_error:
3079 # Same rationale as the LLM config branch above.
3080 user_friendly_error = (
3081 "There was a problem with the search engine configuration."
3082 )
3083 error_context = {
3084 "solution": "Review your search engine settings (or, on a shared server, contact your administrator) and ensure they are correct."
3085 }
3086 else:
3087 # Unrecognized exception. The raw str(e) here is server-side
3088 # internal detail (file paths, DB/driver text, Python tracebacks)
3089 # with no curated, user-actionable form, so it must not be
3090 # surfaced to the client (CWE-209) — it is already captured by
3091 # the logger.exception above. The branches above classify known
3092 # errors and replace the message with a safe category string +
3093 # hint; this branch handles everything that wasn't classified.
3094 user_friendly_error = (
3095 "Research failed due to an unexpected error. Contact your "
3096 "administrator or check the server logs for details."
3097 )
3099 # Generate enhanced error report for failed research
3100 enhanced_report_content = None
3101 try:
3102 # Get partial results if they exist
3103 partial_results = results if "results" in locals() else None
3104 search_iterations = (
3105 results.get("iterations", 0) if partial_results else 0
3106 )
3108 # Generate comprehensive error report
3109 # ErrorReportGenerator does not use LLM (kept for compat)
3110 error_generator = ErrorReportGenerator()
3111 enhanced_report_content = error_generator.generate_error_report(
3112 # Use the sanitized user_friendly_error (curated for known
3113 # errors, generic for unexpected ones) instead of raw {e!s}:
3114 # this report is persisted and retrievable via the report
3115 # routes, so embedding raw exception text would leak server
3116 # internals (CWE-209). Full detail stays in the logs.
3117 error_message=f"Research failed: {user_friendly_error}",
3118 query=query,
3119 partial_results=partial_results,
3120 search_iterations=search_iterations,
3121 research_id=research_id,
3122 )
3124 logger.info(
3125 "Generated enhanced error report for failed research (length: {})",
3126 len(enhanced_report_content),
3127 )
3129 # Save enhanced error report to encrypted database
3130 try:
3131 # username already available from function scope (line 281)
3132 if username: 3132 ↛ 3154line 3132 didn't jump to line 3154 because the condition on line 3132 was always true
3133 from ...storage import get_report_storage
3135 with get_user_db_session(username) as db_session:
3136 storage = get_report_storage(session=db_session)
3137 success = storage.save_report(
3138 research_id=research_id,
3139 content=enhanced_report_content,
3140 metadata={"error_report": True},
3141 username=username,
3142 )
3143 if success:
3144 logger.info(
3145 "Saved enhanced error report to encrypted database for research {}",
3146 research_id,
3147 )
3148 else:
3149 logger.warning(
3150 "Failed to save enhanced error report to database for research {}",
3151 research_id,
3152 )
3153 else:
3154 logger.warning(
3155 "Cannot save error report: username not available"
3156 )
3158 except Exception as report_error:
3159 logger.exception(
3160 "Failed to save enhanced error report: {}", report_error
3161 )
3163 except Exception as error_gen_error:
3164 logger.exception(
3165 "Failed to generate enhanced error report: {}",
3166 error_gen_error,
3167 )
3168 enhanced_report_content = None
3170 # Get existing metadata from database first
3171 existing_metadata = {}
3172 try:
3173 # username already available from function scope (line 281)
3174 if username: 3174 ↛ 3187line 3174 didn't jump to line 3187 because the condition on line 3174 was always true
3175 with get_user_db_session(username) as db_session:
3176 research = (
3177 db_session.query(ResearchHistory)
3178 .filter_by(id=research_id)
3179 .first()
3180 )
3181 if research and research.research_meta:
3182 existing_metadata = dict(research.research_meta)
3183 except Exception:
3184 logger.exception("Failed to get existing metadata")
3186 # Update metadata with more context about the error while preserving existing values
3187 metadata = existing_metadata
3188 metadata.update({"phase": "error", "error": user_friendly_error})
3189 if error_context:
3190 metadata.update(error_context)
3191 if enhanced_report_content: 3191 ↛ 3195line 3191 didn't jump to line 3195 because the condition on line 3191 was always true
3192 metadata["has_enhanced_report"] = True
3194 # If we still have an active research record, update its log
3195 if is_research_active(research_id):
3196 progress_callback(user_friendly_error, None, metadata)
3198 # If termination was requested, mark as suspended instead of failed
3199 status = (
3200 ResearchStatus.SUSPENDED
3201 if is_termination_requested(research_id)
3202 else ResearchStatus.FAILED
3203 )
3204 message = (
3205 "Research was terminated by user"
3206 if status == ResearchStatus.SUSPENDED
3207 else user_friendly_error
3208 )
3210 # A subscription-triggered run that FAILED must be made due again.
3211 # run_subscription_now / the overdue sweep advance next_refresh at
3212 # spawn time (to avoid the scheduler double-running an in-flight
3213 # subscription), but the completion-time advance only runs on
3214 # success. Without this reset a failed run would leave next_refresh
3215 # pushed a full interval out, silently hiding the subscription from
3216 # the scheduler. Skipped for SUSPENDED (user-terminated) runs.
3217 if (
3218 status == ResearchStatus.FAILED
3219 and username
3220 and metadata.get("subscription_id")
3221 ):
3222 try:
3223 from ...news.subscription_runner import (
3224 mark_subscription_due_by_id,
3225 )
3227 with get_user_db_session(username) as sub_db:
3228 if mark_subscription_due_by_id( 3228 ↛ 3243line 3228 didn't jump to line 3243
3229 sub_db, metadata["subscription_id"]
3230 ):
3231 sub_db.commit()
3232 logger.info(
3233 f"Reset subscription {metadata['subscription_id']} "
3234 "to due after failed run"
3235 )
3236 except Exception:
3237 logger.warning(
3238 "Could not reset subscription refresh time after "
3239 "failed run"
3240 )
3242 # Calculate duration up to termination point - using UTC consistently
3243 now = datetime.now(UTC)
3244 completed_at = now.isoformat()
3246 # NOTE: Database updates from threads are handled by queue processor
3247 # The queue_processor.queue_error_update() method is already being used below
3248 # to safely update the database from the main thread
3250 # Queue the error update to be processed in main thread
3251 # Using the queue processor v2 system
3252 from ..queue.processor_v2 import queue_processor
3254 if username: 3254 ↛ 3268line 3254 didn't jump to line 3268 because the condition on line 3254 was always true
3255 queue_processor.queue_error_update(
3256 username=username,
3257 research_id=research_id,
3258 status=status,
3259 error_message=message,
3260 metadata=metadata,
3261 completed_at=completed_at,
3262 report_path=None,
3263 )
3264 logger.info(
3265 f"Queued error update for research {research_id} with status '{status}'"
3266 )
3267 else:
3268 logger.error(
3269 f"Cannot queue error update for research {research_id} - no username provided. "
3270 f"Status: '{status}', Message: {message}"
3271 )
3273 try:
3274 _sio_emit(
3275 "research_progress",
3276 research_id,
3277 {"status": status, "error": message},
3278 owner=username,
3279 )
3280 except Exception:
3281 logger.exception("Failed to emit error via socket")
3283 # Add error message to chat session if applicable.
3284 # Read chat_session_id from shared_research_context (the canonical
3285 # source — kept consistent with the six other reads in this file).
3286 chat_session_id = shared_research_context.get("chat_session_id")
3287 if chat_session_id and username:
3288 try:
3289 from ...chat.service import ChatService
3291 chat_service = ChatService(username)
3292 # allow_archived=True: same multi-tab race rationale as
3293 # the completion / stop-and-partial paths — if the user
3294 # archived (or deleted) the session between research
3295 # start and failure, the error message would otherwise
3296 # be silently dropped by the active-only insert guard
3297 # and the user sees nothing in the chat.
3298 chat_service.add_message(
3299 session_id=chat_session_id,
3300 role="assistant",
3301 content=f"Sorry, the research failed: {message}",
3302 message_type="response",
3303 allow_archived=True,
3304 )
3305 except Exception:
3306 # Promoted from debug → warning to match the success-path
3307 # rationale: if this write fails the user never sees the
3308 # error and operators get no signal at debug-off level.
3309 logger.opt(exception=True).warning(
3310 "Could not add error message to chat session"
3311 )
3313 except Exception:
3314 logger.exception("Error in error handler")
3316 # Clean up resources. This is the error path, so report FAILED on
3317 # the final socket message rather than a spurious "completed".
3318 cleanup_research_resources(
3319 research_id,
3320 username,
3321 user_password=user_password,
3322 final_status=ResearchStatus.FAILED,
3323 )
3325 finally:
3326 # RESOURCE CLEANUP: Close search engine HTTP sessions.
3327 #
3328 # Search engines (created via get_search()) may hold HTTP connection
3329 # pools. Currently only SemanticScholarSearchEngine creates a
3330 # persistent SafeSession; other engines use stateless safe_get()/
3331 # safe_post() utility functions. However, BaseSearchEngine.close()
3332 # is safe to call on any engine — it checks for a 'session'
3333 # attribute and is fully idempotent (SemanticScholar sets
3334 # self.session = None after close).
3335 #
3336 # Neither @thread_cleanup nor cleanup_research_resources() close
3337 # the search engine — @thread_cleanup only handles database sessions
3338 # and context cleanup, and cleanup_research_resources() only handles
3339 # status updates, notifications, and tracking dict removal.
3340 #
3341 # Without this explicit close, search engine sessions rely on
3342 # Python's non-deterministic garbage collection (__del__) for
3343 # cleanup, which can cause file descriptor exhaustion under
3344 # sustained load.
3345 from ...utilities.resource_utils import safe_close
3347 if "use_search" in locals():
3348 safe_close(use_search, "research search engine")
3349 # Close search system (cascades to strategy thread pools).
3350 # See AdvancedSearchSystem.close() for details.
3351 if "system" in locals():
3352 safe_close(system, "research system")
3353 # Close the LLM instance created for model/provider overrides.
3354 # system.close() does NOT close the LLM passed to it via system.model,
3355 # so we must close it explicitly here.
3356 if "use_llm" in locals():
3357 safe_close(use_llm, "research LLM")
3360def cleanup_research_resources(
3361 research_id,
3362 username=None,
3363 user_password=None,
3364 final_status=ResearchStatus.COMPLETED,
3365 preserve_termination_flag=False,
3366):
3367 """
3368 Clean up resources for a completed research.
3370 Args:
3371 research_id: The ID of the research
3372 username: The username for database access (required for thread context)
3373 final_status: The terminal status to report on the final socket
3374 message. Callers that end a research for a reason other than
3375 normal completion MUST pass the real status (e.g. SUSPENDED on
3376 user termination, FAILED on error) so the final ``progress``
3377 event matches reality. Defaulting this to COMPLETED — and
3378 previously hard-coding it — caused the stop/error paths to emit
3379 a spurious "completed" signal to subscribers.
3380 preserve_termination_flag: Keep a cancellation signal visible after
3381 removing the active registry entry. Used only by the HTTP cancel
3382 handoff; the worker's own termination cleanup uses the default and
3383 removes the flag.
3384 """
3385 from ..research_state import cleanup_research
3387 logger.info("Cleaning up resources for research {}", research_id)
3389 # For testing: Add a small delay to simulate research taking time
3390 # This helps test concurrent research limits
3391 from ...settings.env_registry import is_test_mode
3393 if is_test_mode():
3394 import time
3396 logger.info(
3397 f"Test mode: Adding 5 second delay before cleanup for {research_id}"
3398 )
3399 time.sleep(5)
3401 # The terminal status to report on the final socket message. This comes
3402 # from the caller (which knows why the research ended) rather than a
3403 # hard-coded COMPLETED, so termination (SUSPENDED) and error (FAILED)
3404 # paths no longer emit a false "completed" signal to subscribers.
3405 current_status = final_status
3407 # NOTE: Queue processor already handles database updates from the main thread
3408 # The notify_research_completed() method is called at the end of this function
3409 # which safely updates the database status
3411 # Notify queue processor that research completed
3412 # This uses processor_v2 which handles database updates in the main thread
3413 # avoiding the Flask request context issues that occur in background threads
3414 from ..queue.processor_v2 import queue_processor
3416 if username:
3417 queue_processor.notify_research_completed(
3418 username, research_id, user_password=user_password
3419 )
3420 logger.info(
3421 f"Notified queue processor of completion for research {research_id} (user: {username})"
3422 )
3423 else:
3424 logger.warning(
3425 f"Cannot notify completion for research {research_id} - no username provided"
3426 )
3428 # Remove from active research and termination flags atomically
3429 if preserve_termination_flag:
3430 cleanup_research(research_id, preserve_termination_flag=True)
3431 else:
3432 cleanup_research(research_id)
3434 # Clean up throttle state for this research
3435 with _last_emit_lock:
3436 _last_emit_times.pop(research_id, None)
3438 # Send a final message to subscribers
3439 try:
3440 # Send a final message to any remaining subscribers with explicit status
3441 # Use the proper status message based on database status
3442 if current_status in (
3443 ResearchStatus.SUSPENDED,
3444 ResearchStatus.FAILED,
3445 ):
3446 final_message = {
3447 "status": current_status,
3448 "message": f"Research was {current_status}",
3449 "progress": 0, # For suspended research, show 0% not 100%
3450 }
3451 else:
3452 final_message = {
3453 "status": ResearchStatus.COMPLETED,
3454 "message": "Research process has ended and resources have been cleaned up",
3455 "progress": 100,
3456 }
3458 logger.info(
3459 "Sending final {} socket message for research {}",
3460 current_status,
3461 research_id,
3462 )
3464 _sio_emit(
3465 "research_progress", research_id, final_message, owner=username
3466 )
3468 # Clean up socket subscriptions for this research
3469 _sio_remove(research_id, username)
3471 except Exception:
3472 logger.exception("Error sending final cleanup message")
3475def handle_termination(
3476 research_id, username=None, *, preserve_termination_flag=False
3477):
3478 """
3479 Handle the termination of a research process.
3481 Args:
3482 research_id: The ID of the research
3483 username: The username for database access (required for thread context)
3484 """
3485 logger.info(f"Handling termination for research {research_id}")
3487 # Queue the status update to be processed in the main thread
3488 # This avoids Flask request context errors in background threads
3489 try:
3490 from ..queue.processor_v2 import queue_processor
3492 now = datetime.now(UTC)
3493 completed_at = now.isoformat()
3495 # Queue the suspension update
3496 queue_processor.queue_error_update(
3497 username=username,
3498 research_id=research_id,
3499 status=ResearchStatus.SUSPENDED,
3500 error_message="Research was terminated by user",
3501 metadata={"terminated_at": completed_at},
3502 completed_at=completed_at,
3503 report_path=None,
3504 )
3506 logger.info(f"Queued suspension update for research {research_id}")
3507 except Exception:
3508 logger.exception(
3509 f"Error queueing termination update for research {research_id}"
3510 )
3512 # Clean up resources (this already handles things properly).
3513 # Pass SUSPENDED so the final socket message reports the real terminal
3514 # status — not a spurious "completed" — to chat/progress subscribers.
3515 cleanup_kwargs = {"final_status": ResearchStatus.SUSPENDED}
3516 if preserve_termination_flag:
3517 cleanup_kwargs["preserve_termination_flag"] = True
3518 cleanup_research_resources(research_id, username, **cleanup_kwargs)
3521def cancel_research(research_id, username):
3522 """
3523 Cancel/terminate a research process using ORM.
3525 Args:
3526 research_id: The ID of the research to cancel
3527 username: The username of the user cancelling the research
3529 Returns:
3530 bool: True if the research was found and cancelled, False otherwise
3531 """
3532 try:
3533 from ..research_state import is_research_active, set_termination_flag
3535 # Ownership gate. research ids are per-user, but the global termination
3536 # registry (_termination_flags / _active_research, keyed by id alone) is
3537 # shared across all users. Without this check any signed-in user could
3538 # terminate ANOTHER user's active research by id: the active branch
3539 # below sets the flag and calls handle_termination before the only
3540 # pre-existing ownership query (which lives in the not-active branch).
3541 # Confirm the caller owns this research in their OWN database first.
3542 with get_user_db_session(username) as _owner_check_session:
3543 owns_research = (
3544 _owner_check_session.query(ResearchHistory.id)
3545 .filter_by(id=research_id)
3546 .first()
3547 is not None
3548 )
3549 if not owns_research:
3550 # (#6938: research_id is the raw path parameter on this
3551 # not-found branch.)
3552 from ..routers.notes import _log_value_preview
3554 logger.info(
3555 "Research {} not found for this user; refusing to cancel.",
3556 _log_value_preview(research_id),
3557 )
3558 return False
3560 # Set termination flag
3561 set_termination_flag(research_id)
3563 # Check if the research is active
3564 if is_research_active(research_id):
3565 # Update the database and release the active slot immediately, but
3566 # keep the signal set until the still-running worker sees it. The
3567 # worker's own handle_termination call performs final flag cleanup.
3568 handle_termination(
3569 research_id,
3570 username,
3571 preserve_termination_flag=True,
3572 )
3573 return True
3574 try:
3575 with get_user_db_session(username) as db_session:
3576 research = (
3577 db_session.query(ResearchHistory)
3578 .filter_by(id=research_id)
3579 .first()
3580 )
3581 if not research:
3582 logger.info(f"Research {research_id} not found in database")
3583 return False
3585 # Check if already in a terminal state
3586 if research.status in (
3587 ResearchStatus.COMPLETED,
3588 ResearchStatus.SUSPENDED,
3589 ResearchStatus.FAILED,
3590 ResearchStatus.ERROR,
3591 ):
3592 logger.info(
3593 f"Research {research_id} already in terminal state: {research.status}"
3594 )
3595 return True # Consider this a success since it's already stopped
3597 # If it exists but isn't in active_research, still update status
3598 if research.status == ResearchStatus.QUEUED:
3599 from ...database.models import QueuedResearch
3601 claimed_queue_row = (
3602 db_session.query(QueuedResearch.id)
3603 .filter(
3604 QueuedResearch.research_id == research_id,
3605 QueuedResearch.is_processing.is_(True),
3606 )
3607 .exists()
3608 )
3609 suspended = (
3610 db_session.query(ResearchHistory)
3611 .filter(
3612 ResearchHistory.id == research_id,
3613 ResearchHistory.status == ResearchStatus.QUEUED,
3614 ~claimed_queue_row,
3615 )
3616 .update(
3617 {ResearchHistory.status: ResearchStatus.SUSPENDED},
3618 synchronize_session=False,
3619 )
3620 )
3621 if suspended:
3622 from ..queue.lifecycle_cleanup import (
3623 cleanup_queued_research_state,
3624 )
3626 cleanup_queued_research_state(
3627 db_session, [research_id], acting_user=username
3628 )
3629 else:
3630 research.status = ResearchStatus.SUSPENDED
3631 db_session.commit()
3632 logger.info(f"Successfully suspended research {research_id}")
3633 except Exception:
3634 logger.exception(
3635 f"Error accessing database for research {research_id}"
3636 )
3637 return False
3639 return True
3640 except Exception:
3641 from ..routers.notes import _log_value_preview
3643 logger.exception(
3644 "Unexpected error in cancel_research for {}",
3645 _log_value_preview(research_id),
3646 )
3647 return False