Coverage for src/local_deep_research/web/routers/followup.py: 86%
187 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Flask routes for follow-up research functionality.
3"""
5from fastapi import APIRouter, Depends, Request
6from fastapi.responses import JSONResponse
7from ..dependencies.auth import require_auth
8from ..dependencies.threadpool import run_db_sync
10from loguru import logger
12from ...constants import ResearchStatus
13from ...exceptions import DuplicateResearchError, SystemAtCapacityError
14from ...llm.providers.base import normalize_provider
15from ...followup_research.service import FollowUpResearchService
16from ...followup_research.models import FollowUpRequest
17from ...utilities.url_utils import is_safe_custom_llm_endpoint
18from .notes import _log_value_preview
20from ..auth.password_utils import resolve_user_password
21from typing import Annotated
23# Create the router
24router = APIRouter(prefix="/api/followup", tags=["followup"])
26# NOTE: Routes use username (not .get()) intentionally.
27# Depends(require_auth) guarantees the key exists; direct access fails
28# fast if the dependency is ever removed.
31@router.post("/prepare")
32async def prepare_followup(
33 request: Request, username: Annotated[str, Depends(require_auth)]
34):
35 """
36 Prepare a follow-up research by loading parent context.
38 Request body:
39 {
40 "parent_research_id": "uuid",
41 "question": "follow-up question"
42 }
44 Returns:
45 {
46 "success": true,
47 "parent_summary": "...",
48 "available_sources": 10,
49 "suggested_strategy": "source-based"
50 }
51 """
52 try:
53 try:
54 data = await request.json()
55 except ValueError:
56 # Match start_followup: a malformed body is a client error, not
57 # a 500 with a logged stack trace.
58 return JSONResponse(
59 {"success": False, "error": "Request body must be valid JSON"},
60 status_code=400,
61 )
62 if not isinstance(data, dict):
63 # Flask's @require_json_body returned 400 for any non-dict body.
64 # Without this, a VALID but non-object body (e.g. `[1,2]` or a
65 # bare string) parses fine and then `data.get(...)` raises
66 # AttributeError, which the outer handler turns into a 500 plus a
67 # logged stack trace. The sibling routers kept the gate
68 # (chat._json_object_body, notes._notes_json_body); follow-up is
69 # the one that lost it.
70 return JSONResponse(
71 {
72 "success": False,
73 "error": "Request body must be a JSON object",
74 },
75 status_code=400,
76 )
77 parent_id = data.get("parent_research_id")
78 question = data.get("question")
80 if not parent_id or not question:
81 return JSONResponse(
82 {
83 "success": False,
84 "error": "Missing parent_research_id or question",
85 },
86 status_code=400,
87 )
89 from ...settings.manager import SettingsManager
90 from ...database.session_context import get_user_db_session
92 def _load_sync():
93 """Sync work: settings snapshot + parent research load."""
94 with get_user_db_session(username) as db_session:
95 settings_manager = SettingsManager(db_session=db_session)
96 settings_snapshot = settings_manager.get_all_settings()
98 strategy_from_settings = settings_snapshot.get(
99 "search.search_strategy", {}
100 ).get("value", "source-based")
102 service = FollowUpResearchService(username=username)
103 parent_data = service.load_parent_research(parent_id)
104 return strategy_from_settings, parent_data
106 strategy_from_settings, parent_data = await run_db_sync(_load_sync)
108 if not parent_data:
109 # Parent research doesn't exist (wrong ID, deleted, or belongs
110 # to another user whose DB we can't read). Return 404 so the
111 # caller doesn't submit a follow-up against a ghost parent.
112 # (#6938: parent_id is the raw, unchecked body value here.)
113 logger.warning(
114 "Parent research {} not found for user {}",
115 _log_value_preview(parent_id),
116 username,
117 )
118 return JSONResponse(
119 {"success": False, "error": "Parent research not found"},
120 status_code=404,
121 )
123 # Prepare response with parent context summary
124 return {
125 "success": True,
126 "parent_summary": parent_data.get("query", ""),
127 "available_sources": len(parent_data.get("resources", [])),
128 "suggested_strategy": strategy_from_settings, # Use strategy from settings
129 "parent_research": {
130 "id": parent_id,
131 "query": parent_data.get("query", ""),
132 "sources_count": len(parent_data.get("resources", [])),
133 },
134 }
136 except Exception:
137 logger.exception("Error preparing follow-up")
138 return JSONResponse(
139 {"success": False, "error": "An internal error has occurred."},
140 status_code=500,
141 )
144@router.post("/start")
145async def start_followup(
146 request: Request, username: Annotated[str, Depends(require_auth)]
147):
148 """
149 Start a follow-up research.
151 Request body:
152 {
153 "parent_research_id": "uuid",
154 "question": "follow-up question",
155 "strategy": "source-based", # optional
156 "max_iterations": 1, # optional
157 "questions_per_iteration": 3 # optional
158 }
160 Returns:
161 {
162 "success": true,
163 "research_id": "new-uuid",
164 "message": "Follow-up research started"
165 }
166 """
167 # Guard the body parse: the try/except lives in _start_followup_sync, so
168 # without this a malformed body escapes this route entirely instead of
169 # returning the 400 the Flask original (and prepare_followup) returns.
170 try:
171 data = await request.json()
172 except ValueError:
173 return JSONResponse(
174 {"success": False, "error": "Request body must be valid JSON"},
175 status_code=400,
176 )
177 if not isinstance(data, dict):
178 # See the identical guard in prepare_followup above.
179 return JSONResponse(
180 {"success": False, "error": "Request body must be a JSON object"},
181 status_code=400,
182 )
183 return await run_db_sync(_start_followup_sync, data, username)
186def _start_followup_sync(data, username):
187 admission_lock = None
188 try:
189 from ..services.research_service import (
190 start_research_process,
191 run_research_process,
192 clamp_user_max_concurrent,
193 )
194 from ..routes.globals import reclaim_stale_user_active_research
195 from ..research_state import get_user_research_start_lock
196 from ...database.models import UserActiveResearch
197 import threading
198 import uuid
200 # Get username from session
202 # Get settings snapshot first to use database values
203 from ...settings.manager import SettingsManager
204 from ...database.session_context import get_user_db_session
206 with get_user_db_session(username) as db_session:
207 settings_manager = SettingsManager(db_session=db_session)
208 settings_snapshot = settings_manager.get_all_settings()
210 # Get strategy from settings snapshot, fallback to source-based if not set
211 strategy_from_settings = settings_snapshot.get(
212 "search.search_strategy", {}
213 ).get("value", "source-based")
215 # Get iterations and questions from settings snapshot
216 iterations_from_settings = settings_snapshot.get(
217 "search.iterations", {}
218 ).get("value", 1)
219 questions_from_settings = settings_snapshot.get(
220 "search.questions_per_iteration", {}
221 ).get("value", 3)
223 # Create follow-up request using settings values
224 followup_request = FollowUpRequest(
225 parent_research_id=data.get("parent_research_id"),
226 question=data.get("question"),
227 strategy=strategy_from_settings, # Use strategy from settings
228 max_iterations=iterations_from_settings, # Use iterations from settings
229 questions_per_iteration=questions_from_settings, # Use questions from settings
230 )
232 # Initialize service
233 service = FollowUpResearchService(username=username)
235 # Resolve the user's password (needed for metrics DB access later) and
236 # decide authentication FIRST. An expired encrypted-DB session is a 401,
237 # and it must be settled BEFORE the parent-ownership check below so an
238 # unauthenticated caller sees "session expired" (401), not the
239 # authorization outcome "parent not found" (404). Auth precedes authz.
240 user_password, session_expired = resolve_user_password(username)
241 if session_expired:
242 # success/error keys match the followup API convention (the
243 # followup frontend checks data.success and data.error).
244 return JSONResponse(
245 {
246 "success": False,
247 "error": "Your session has expired. Please log out and log back in to start research.",
248 },
249 status_code=401,
250 )
252 # Reject a follow-up naming a parent research the caller does not own,
253 # mirroring /api/followup/prepare's 404 contract. Research ids are
254 # per-user (a different physical encrypted DB per user), so a parent id
255 # that isn't in the caller's DB is not theirs; without this a user could
256 # spawn a follow-up referencing another user's research_id (the parent
257 # context comes back empty, but a research thread would still start).
258 parent_id = data.get("parent_research_id")
259 if not parent_id or not service.load_parent_research(parent_id):
260 return JSONResponse(
261 {"success": False, "error": "Parent research not found"},
262 status_code=404,
263 )
265 # Prepare research parameters
266 research_params = service.perform_followup(followup_request)
268 # (#6938: the query is raw request input and a wide params dict has
269 # unbounded keys, so both go through the bounded preview instead of
270 # eager f-strings.)
271 params_keys = (
272 list(research_params.keys())
273 if isinstance(research_params, dict)
274 else "Not a dict"
275 )
276 query_value = (
277 research_params.get("query")
278 if isinstance(research_params, dict)
279 else "N/A"
280 )
281 logger.info("Research params type: {}", type(research_params))
282 logger.info("Research params keys: {}", _log_value_preview(params_keys))
283 logger.info("Query value: {}", _log_value_preview(query_value))
284 logger.info("Query type: {}", type(query_value))
286 # (user_password / session-expired 401 are resolved above, before the
287 # parent-ownership check, so auth precedes authz.)
289 # Pre-flight: refuse to spawn a research thread (and create an
290 # orphan ResearchHistory row) when llm.model is empty. Mirrors the
291 # empty-model check in web/routers/research.py's start_research —
292 # same contract: HTTP 400 with an actionable message before any DB
293 # writes or thread spawning. (This router uses success/error
294 # response keys rather than status/message, matching the followup
295 # API convention used by the other returns in this function.)
296 if not settings_snapshot.get("llm.model", {}).get("value"):
297 logger.error(
298 "Follow-up research blocked: llm.model is not configured"
299 )
300 return JSONResponse(
301 {
302 "success": False,
303 "error": "Model is required. Please configure a model in the settings.",
304 },
305 status_code=400,
306 )
308 # SSRF pre-flight on the LLM endpoint: reject metadata / link-local
309 # targets at the request boundary, before any DB row is written.
310 # This is fail-fast defense-in-depth — the OpenAI-compatible provider's
311 # assert_base_url_safe re-validates the same URL before the client is
312 # built. Private IPs and localhost pass because local LLMs
313 # (Ollama / LM Studio / vLLM) live there, including scheme-less
314 # endpoints (the helper normalizes exactly as the provider does).
315 custom_endpoint = settings_snapshot.get(
316 "llm.openai_endpoint.url", {}
317 ).get("value")
318 if not is_safe_custom_llm_endpoint(custom_endpoint):
319 return JSONResponse(
320 {
321 "success": False,
322 "error": "Invalid custom endpoint URL",
323 },
324 status_code=400,
325 )
327 # Per-user concurrency admission. Follow-ups previously enforced NO
328 # per-user cap: only the global research semaphore gated them, so a
329 # single authenticated user could fire many rapid
330 # /api/followup/start calls and monopolize the entire global
331 # research budget, starving other tenants with 429s. Route the
332 # follow-up through the SAME admission path
333 # research_routes.start_research uses -- reclaim dead-thread rows,
334 # count the user's live researches, and reject at the per-user cap.
335 # A UserActiveResearch row is created below (alongside the
336 # ResearchHistory) so both entry points keep consistent accounting.
337 max_concurrent_researches = clamp_user_max_concurrent(
338 settings_snapshot.get("app.max_concurrent_researches", {}).get(
339 "value", 3
340 )
341 )
343 # The helpers above reuse a thread-local session and may leave a read
344 # transaction open. End it before taking the process-wide lock order
345 # used by direct starts and queue replay: admission gate, then DB.
346 with get_user_db_session(username) as ordering_session:
347 ordering_session.rollback()
349 # Retain the same stable per-user gate from stale-row reconciliation
350 # through durable claim, worker registration, and failure cleanup.
351 # Otherwise this route can reclaim a direct start in its pre-register
352 # window, or have its own new claim reclaimed by queue replay.
353 pending_admission_lock = get_user_research_start_lock(username)
354 pending_admission_lock.acquire()
355 admission_lock = pending_admission_lock
357 # Reclaim stale rows + count active. Mirrors the try/except in
358 # research_routes.start_research: if the check itself fails we fall
359 # back to allowing the start (the global semaphore is still a
360 # backstop) rather than hard-failing the request.
361 try:
362 with get_user_db_session(username) as admission_session:
363 if reclaim_stale_user_active_research( 363 ↛ 366line 363 didn't jump to line 366 because the condition on line 363 was never true
364 admission_session, username, logger=logger
365 ):
366 admission_session.commit()
367 active_count = (
368 admission_session.query(UserActiveResearch)
369 .filter_by(
370 username=username,
371 status=ResearchStatus.IN_PROGRESS,
372 )
373 .count()
374 )
375 at_capacity = active_count >= max_concurrent_researches
376 except Exception:
377 logger.exception("Failed to check active follow-up researches")
378 at_capacity = False
380 if at_capacity:
381 logger.warning(
382 "Follow-up research rejected: user {} at per-user "
383 "concurrency cap ({}/{})",
384 username,
385 active_count,
386 max_concurrent_researches,
387 )
388 return JSONResponse(
389 {
390 "success": False,
391 "error": "Server is at research capacity. Please retry shortly.",
392 },
393 status_code=429,
394 )
396 # Generate new research ID
397 research_id = str(uuid.uuid4())
399 # Create database entry (settings_snapshot already captured above)
400 from ...database.models import ResearchHistory
401 from datetime import datetime, UTC
403 created_at = datetime.now(UTC).isoformat()
405 with get_user_db_session(username) as db_session:
406 # Create the database entry (required for tracking)
407 research_meta = {
408 "submission": {
409 "parent_research_id": data.get("parent_research_id"),
410 "question": data.get("question"),
411 "strategy": "contextual-followup",
412 },
413 # Persist the run's settings snapshot so view/export
414 # (report_assembly_service) renders the same
415 # report.uncited_sources_mode the generation-time render
416 # used. Mirrors the main-form creation path; completion
417 # merge back-fills when missing.
418 "settings_snapshot": settings_snapshot,
419 }
421 research = ResearchHistory(
422 id=research_id,
423 query=research_params["query"],
424 mode="quick", # Use 'quick' not 'quick_summary'
425 status=ResearchStatus.IN_PROGRESS,
426 created_at=created_at,
427 progress_log=[{"time": created_at, "progress": 0}],
428 research_meta=research_meta,
429 )
430 db_session.add(research)
432 # Record the active-research row in the SAME transaction so the
433 # per-user cap accounting matches research_routes.start_research.
434 # thread_id is the spawning (request) thread for now; the worker
435 # thread id isn't known until start_research_process runs.
436 active_record = UserActiveResearch(
437 username=username,
438 research_id=research_id,
439 status=ResearchStatus.IN_PROGRESS,
440 thread_id=str(threading.current_thread().ident),
441 settings_snapshot=settings_snapshot,
442 )
443 db_session.add(active_record)
444 db_session.commit()
445 logger.info(
446 f"Created follow-up research entry with ID: {research_id}"
447 )
449 # Post-commit race recheck. Two concurrent submissions can both
450 # pass the up-front admission check before either commits its
451 # row. If we are now over the cap, roll back BOTH rows and
452 # reject at 429 -- follow-ups have no queue fallback, so unlike
453 # research_routes.start_research (which re-queues) we simply ask
454 # the client to retry.
455 #
456 # The DETECTION query is wrapped so that a failure to *ask* never
457 # blocks a legitimate start (fail-open is right there: we do not
458 # know we are over capacity). The REMEDIATION below is
459 # deliberately NOT inside that handler. It used to be, and a
460 # failure in the rollback -- e.g. the commit raising under
461 # SQLCipher/WAL contention -- was swallowed by the same broad
462 # `except`, after which control fell through and started the
463 # research anyway, returning success: True for a run whose
464 # tracking rows had just been deleted. Once we KNOW we are over
465 # capacity, failing to clean up must not become a decision to
466 # proceed.
467 over_capacity = False
468 try:
469 final_count = (
470 db_session.query(UserActiveResearch)
471 .filter_by(
472 username=username,
473 status=ResearchStatus.IN_PROGRESS,
474 )
475 .count()
476 )
477 over_capacity = final_count > max_concurrent_researches
478 if over_capacity:
479 logger.warning(
480 "Race detected on follow-up start for {}: "
481 "{} > {}, rolling back and rejecting",
482 username,
483 final_count,
484 max_concurrent_researches,
485 )
486 except Exception:
487 logger.warning("Could not recheck active follow-up count")
489 if over_capacity:
490 # Best-effort rollback: if it fails we still reject, and the
491 # orphaned IN_PROGRESS row is reclaimed by
492 # reclaim_stale_user_active_research on this user's next
493 # start (its thread never spawned, so it reads as dead).
494 try:
495 db_session.delete(active_record)
496 db_session.query(ResearchHistory).filter_by(
497 id=research_id
498 ).delete()
499 db_session.commit()
500 except Exception:
501 from ...database.session_context import safe_rollback
503 logger.exception(
504 "Rollback of the over-capacity follow-up rows failed "
505 "for {}; still rejecting the request",
506 username,
507 )
508 safe_rollback(db_session, "followup over-capacity rollback")
509 return JSONResponse(
510 {
511 "success": False,
512 "error": "Server is at research capacity. Please retry shortly.",
513 },
514 status_code=429,
515 )
517 # Start the research process using the existing infrastructure
518 # Use quick_summary mode for follow-ups by default
519 logger.info(
520 f"Starting follow-up research for query of type: {type(research_params.get('query'))}"
521 )
523 # Get model and search settings from user's settings
524 model_provider = settings_snapshot.get("llm.provider", {}).get(
525 "value", "OLLAMA"
526 )
527 # Normalize provider to lowercase canonical form
528 model_provider = normalize_provider(model_provider)
529 model = settings_snapshot.get("llm.model", {}).get("value", "")
530 search_engine = settings_snapshot.get("search.tool", {}).get(
531 "value", "searxng"
532 )
534 # Spawn the research thread. If the spawn fails, the
535 # ResearchHistory row committed above would otherwise be
536 # permanently orphaned with status=IN_PROGRESS. Catch any
537 # exception, flip the status to FAILED, and return a clear
538 # error — same contract as the queue processor's terminal-
539 # failure branch (#3481) and the direct-UI spawn-failure path.
540 try:
541 start_research_process(
542 research_id,
543 research_params["query"],
544 "quick", # Use 'quick' for quick summary mode
545 run_research_process,
546 username=username,
547 user_password=user_password, # gitleaks:allow
548 model_provider=model_provider, # Pass model provider
549 model=model, # Pass model name
550 search_engine=search_engine, # Pass search engine
551 # llm.openai_endpoint.url is only the client's base URL for
552 # the openai_endpoint provider; don't label other providers
553 # with it.
554 custom_endpoint=(
555 custom_endpoint
556 if model_provider == "openai_endpoint"
557 else None
558 ),
559 strategy="enhanced-contextual-followup", # Use enhanced contextual follow-up strategy
560 iterations=research_params["max_iterations"],
561 questions_per_iteration=research_params[
562 "questions_per_iteration"
563 ],
564 delegate_strategy=research_params.get(
565 "delegate_strategy", "source-based"
566 ),
567 research_context=research_params["research_context"],
568 parent_research_id=research_params[
569 "parent_research_id"
570 ], # Pass parent research ID
571 settings_snapshot=settings_snapshot,
572 )
573 except DuplicateResearchError:
574 # A live thread already owns this research_id. Do NOT delete
575 # the row or mark it FAILED — the row belongs to the live
576 # thread and mutating it would terminate the running
577 # research from the user's perspective. Same contract as
578 # research_routes.start_research's duplicate-thread branch.
579 logger.warning(
580 f"Duplicate live thread detected for follow-up "
581 f"{research_id}; leaving state intact"
582 )
583 return JSONResponse(
584 {
585 "success": False,
586 "error": "Research is already running.",
587 },
588 status_code=409,
589 )
590 except SystemAtCapacityError:
591 # System at concurrent-research capacity. Roll back the rows
592 # committed above (UserActiveResearch + IN_PROGRESS history)
593 # and return 429.
594 logger.warning(
595 f"SystemAtCapacityError starting follow-up {research_id}"
596 )
597 try:
598 from ...database.session_context import get_user_db_session
599 from ...database.models import (
600 ResearchHistory,
601 UserActiveResearch,
602 )
604 with get_user_db_session(username) as cleanup_session:
605 stale_active = (
606 cleanup_session.query(UserActiveResearch)
607 .filter_by(username=username, research_id=research_id)
608 .first()
609 )
610 if stale_active:
611 cleanup_session.delete(stale_active)
612 cleanup_session.query(ResearchHistory).filter_by(
613 id=research_id
614 ).delete()
615 cleanup_session.commit()
616 except Exception:
617 logger.exception(
618 "Cleanup after follow-up capacity reject raised"
619 )
620 return JSONResponse(
621 {
622 "success": False,
623 "error": "Server is at research capacity. Please retry shortly.",
624 },
625 status_code=429,
626 )
627 except Exception:
628 logger.exception(
629 f"Failed to spawn follow-up research thread for {research_id}"
630 )
631 try:
632 from ...database.session_context import get_user_db_session
633 from ...database.models import (
634 ResearchHistory,
635 UserActiveResearch,
636 )
638 with get_user_db_session(username) as cleanup_session:
639 stale_active = (
640 cleanup_session.query(UserActiveResearch)
641 .filter_by(username=username, research_id=research_id)
642 .first()
643 )
644 if stale_active:
645 cleanup_session.delete(stale_active)
646 research_row = (
647 cleanup_session.query(ResearchHistory)
648 .filter_by(id=research_id)
649 .first()
650 )
651 if research_row:
652 research_row.status = ResearchStatus.FAILED
653 cleanup_session.commit()
654 except Exception:
655 logger.exception("Cleanup after follow-up spawn failure raised")
656 return JSONResponse(
657 content={
658 "success": False,
659 "error": "Failed to start follow-up research. Please try again.",
660 },
661 status_code=500,
662 )
664 return {
665 "success": True,
666 "research_id": research_id,
667 "message": "Follow-up research started",
668 }
670 except Exception:
671 logger.exception("Error starting follow-up")
672 return JSONResponse(
673 {"success": False, "error": "An internal error has occurred."},
674 status_code=500,
675 )
676 finally:
677 if admission_lock is not None:
678 admission_lock.release()