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