Coverage for src/local_deep_research/chat/service.py: 92%
420 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Service layer for chat functionality.
4This service handles the business logic for chat sessions and messages,
5including session management, message handling, and context building.
6"""
8from typing import Dict, Any, List, Optional
9from concurrent.futures import (
10 ThreadPoolExecutor,
11 TimeoutError as FuturesTimeoutError,
12)
13from datetime import datetime, UTC
14from math import isfinite
15import uuid
16from loguru import logger
17from sqlalchemy import String, and_, cast, func, literal, or_, tuple_, update
18from sqlalchemy.exc import SQLAlchemyError
19from sqlalchemy.sql import ColumnElement
21from ..database.models import (
22 ChatMessage,
23 ChatMessageType,
24 ChatProgressStep,
25 ChatRole,
26 ChatSession,
27 ChatSessionStatus,
28 ResearchHistory,
29 UserActiveResearch,
30)
31from ..database.session_context import get_user_db_session
32from ..security.log_sanitizer import redact_and_bound_for_log
33from ..web.research_state import user_research_start_gate
34from ..web.routes.globals import (
35 cleanup_research,
36 is_research_thread_alive,
37 set_termination_flag,
38)
39from ..database.research_cascade import (
40 collect_research_document_rag_targets,
41 delete_research_with_dependents,
42)
43from ..constants import ResearchStatus
44from ..utilities.formatting import short_id
46# Standard exception tuple for service-layer DB operations
47DB_EXCEPTIONS = (ValueError, RuntimeError, SQLAlchemyError)
50class ChatCursorError(ValueError):
51 """Raised for an invalid chat pagination cursor; maps to HTTP 400."""
54class ArchiveBlockedError(RuntimeError):
55 """Raised when archive_session is called while a research is in_progress.
57 Archive flips a session to read-only; allowing it while research is
58 still running would leave an orphaned research tied to a session the
59 user thinks is frozen. The route layer translates this to HTTP 409.
60 """
63class ChatSessionNotFound(LookupError):
64 """Raised by get_session when no row matches the supplied session_id.
66 The route layer translates this to HTTP 404. Distinct from
67 ChatRepositoryError so a transient DB failure cannot masquerade as
68 "session does not exist".
69 """
72class ChatRepositoryError(RuntimeError):
73 """Raised by get_session when the underlying DB query fails.
75 The route layer translates this to HTTP 500. Keeping this separate
76 from ChatSessionNotFound prevents false 404s on infrastructure
77 errors (locked DB file, encryption key failure, etc.).
78 """
81class AttemptNotFound(LookupError):
82 """Raised when no ResearchHistory row matches the supplied research_id
83 inside the session scoped by ``session_id``.
85 The route layer translates this to HTTP 404. Distinct from
86 ChatSessionNotFound (also 404) so the diagnostic log distinguishes
87 "session missing" from "research missing".
88 """
91class AttemptInProgress(RuntimeError):
92 """Raised by delete_attempt when the target research is IN_PROGRESS
93 AND its worker thread is alive.
95 Deleting a research out from under a live worker would orphan the
96 thread (it would keep burning LLM cycles against a ResearchHistory
97 row that no longer exists). The route layer translates this to HTTP
98 409 and returns ``research_id`` so the client can offer a Stop +
99 retry flow.
101 A stale IN_PROGRESS row whose thread is dead does NOT raise: the
102 sweep logic inside ``delete_attempt`` reclaims it (mirrors
103 send_message's stale-row reclaim at chat/routes.py:903-923).
104 """
107# Title generation should not block the request thread on a slow LLM.
108# Wrap the synchronous llm.invoke() call in a worker future with a hard
109# wall-clock timeout. The default matches the rest of the codebase's
110# 30s LLM-timeout convention; tune via chat.title_llm_timeout_seconds.
111_DEFAULT_TITLE_LLM_TIMEOUT_SECONDS = 30.0
114def _serialize_dt(value):
115 """Return an ISO-8601 string for a datetime, or None."""
116 return value.isoformat() if value is not None else None
119def _timestamp_order_key(column):
120 """Compare SQLite UTC timestamps at their full stored precision.
122 ``utcnow()`` stores three fractional digits; bound datetimes store six.
123 Padding the fraction makes equal instants compare equal without rounding
124 away microseconds (as SQLite's date/time functions would).
125 """
126 stored = cast(column, String)
127 return (
128 func.substr(stored, 1, 19, type_=String)
129 + "."
130 + func.substr(
131 func.substr(stored, 21, type_=String) + "000000",
132 1,
133 6,
134 type_=String,
135 )
136 )
139class ChatService:
140 """Service for managing chat conversations and messages."""
142 def __init__(self, username: str):
143 """
144 Initialize the chat service.
146 Args:
147 username: Username for database access
148 """
149 self.username = username
151 @staticmethod
152 def _atomic_increment(db, counter_col, where_clause):
153 """Atomically increment an integer counter column by one and return
154 its new value.
156 Emits a single ``UPDATE <table> SET col = col + 1 WHERE <clause>
157 RETURNING col`` so the next sequence number is allocated without a
158 read-modify-write race between concurrent writers. ``counter_col`` is
159 a mapped column attribute (e.g. ``ChatSession.message_count``); the
160 target table and column name are derived from it.
162 Returns the post-increment value, or ``None`` when no row matched
163 ``where_clause`` — the caller decides whether a miss is a 404/ValueError.
164 Extracted from the previously-duplicated counter logic in
165 ``insert_message_in_db`` (message_count) and ``add_progress_step``
166 (step_count).
167 """
168 # Key the values() dict by the column object (not its string name)
169 # so the table/column are both derived straight from counter_col.
170 stmt = (
171 update(counter_col.class_)
172 .where(where_clause)
173 .values({counter_col: counter_col + 1})
174 .returning(counter_col)
175 )
176 return db.execute(stmt).scalar_one_or_none()
178 def create_session(
179 self,
180 initial_query: Optional[str] = None,
181 title: Optional[str] = None,
182 settings_snapshot: Optional[Dict[str, Any]] = None,
183 ) -> str:
184 """
185 Create a new chat session.
187 Args:
188 initial_query: Optional initial query (used for title generation)
189 title: Optional custom title for the session
190 settings_snapshot: Optional settings for LLM title generation
192 Returns:
193 Session ID (UUID string)
194 """
195 try:
196 session_id = str(uuid.uuid4())
198 # Use fast, non-LLM fallback title synchronously. If the caller
199 # wants an LLM-generated title they trigger it asynchronously via
200 # POST /api/chat/sessions/<id>/generate-title so the creation
201 # request isn't blocked on an LLM round-trip.
202 resolved_title = title or self._fallback_title(initial_query)
204 with get_user_db_session(self.username) as db:
205 # created_at is populated by the utcnow() default on the
206 # column — no need to pass it explicitly.
207 session = ChatSession(
208 id=session_id,
209 title=resolved_title,
210 status=ChatSessionStatus.ACTIVE.value,
211 accumulated_context={
212 "key_entities": [],
213 "topics": [],
214 "summary": "",
215 },
216 message_count=0,
217 )
218 db.add(session)
219 db.commit()
221 logger.info(
222 f"Created chat {session_id[:8]}... for user {self.username}"
223 )
224 return session_id
226 except DB_EXCEPTIONS:
227 logger.exception("Error creating chat session")
228 raise
230 def regenerate_title_with_llm(
231 self,
232 session_id: str,
233 query: Optional[str],
234 settings_snapshot: Optional[Dict[str, Any]],
235 ) -> Optional[str]:
236 """
237 Regenerate a session's title using the LLM.
239 Intended to be called from a dedicated endpoint the frontend fires
240 after session creation, so the create request doesn't block on the
241 LLM round-trip.
243 Idempotency: if the current title no longer matches the non-LLM
244 fallback (i.e. the user manually edited it, or a sibling tab's
245 LLM-gen already ran), skip the LLM call so we don't spend credits
246 only to overwrite the user's deliberate edit on the way back.
248 Returns the new title on success, or None on failure / no-op.
249 """
250 if not query:
251 return None
252 # Check whether the session still has the fallback title. If the
253 # user (or a concurrent generate-title request) has already moved
254 # past the fallback, don't burn an LLM call to overwrite their work.
255 try:
256 current = self.get_session(session_id)
257 except ChatSessionNotFound:
258 return None
259 current_title = (current or {}).get("title") or ""
260 fallback = self._fallback_title(query)
261 if current_title and current_title != fallback:
262 logger.info(
263 "Skipping LLM title gen for {}...: title already set ('{}')",
264 session_id[:8],
265 # (#6938: redacted before it is cut.)
266 redact_and_bound_for_log(str(current_title), 30),
267 )
268 return None
269 new_title = self._generate_title(query, settings_snapshot)
270 if not new_title: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 return None
272 updated = self.update_session_title(session_id, new_title)
273 return new_title if updated else None
275 def add_message(
276 self,
277 session_id: str,
278 role: str,
279 content: str,
280 message_type: str,
281 research_id: Optional[str] = None,
282 allow_archived: bool = False,
283 ) -> str:
284 """
285 Add a durable message (query/followup/response) to a chat session.
287 Content is required and stored inline. Step rows live in
288 chat_progress_steps and are written via add_progress_step().
290 Args:
291 session_id: ID of the session to add message to
292 role: Message role (user or assistant)
293 content: Message content (required, non-empty)
294 message_type: Type of message (query, followup, response)
295 research_id: Optional ID of associated research
296 allow_archived: When True, the atomic-update WHERE clause omits
297 the ``status='active'`` filter so a system-written assistant
298 response can land even if the session was archived between
299 research start and completion. Use ONLY for system writes
300 (final assistant response, terminate-partial) — the user
301 send-message path MUST keep the default False so that
302 archiving a session in one browser tab still blocks a
303 concurrent user reply from another tab mid-flight.
305 Returns:
306 Message ID (UUID string)
308 Raises:
309 ValueError: if content is None. Validated before opening the DB
310 session so callers (e.g. the route layer) can return HTTP 400
311 without paying SQLCipher cold-open cost on a doomed request.
312 """
313 if content is None:
314 raise ValueError("content is required for chat messages")
315 try:
316 with get_user_db_session(self.username) as db:
317 message_id = self.insert_message_in_db(
318 db,
319 session_id=session_id,
320 role=role,
321 content=content,
322 message_type=message_type,
323 research_id=research_id,
324 allow_archived=allow_archived,
325 )
326 db.commit()
327 return message_id
329 except DB_EXCEPTIONS:
330 logger.exception("Error adding message to chat session")
331 raise
333 def insert_message_in_db(
334 self,
335 db,
336 session_id: str,
337 role: str,
338 content: str,
339 message_type: str,
340 research_id: Optional[str] = None,
341 allow_archived: bool = False,
342 ) -> str:
343 """
344 Insert a durable chat message in an active SQLAlchemy session WITHOUT
345 committing. The caller owns the transaction lifecycle and is
346 responsible for commit/rollback.
348 This exists so the route layer can atomically commit the user message
349 together with the research-history row in a single transaction —
350 avoiding the orphan-message bug that occurs if the user-message
351 commit succeeds and the research insert later raises.
353 Validation and the atomic message_count increment are identical to
354 ``add_message``; only the commit responsibility differs.
356 Raises:
357 ValueError: if role/message_type are invalid, content is None,
358 or the session row does not exist.
359 """
360 # Content is required (NOT NULL on the column).
361 # Empty string is permitted by the column (NOT NULL only rejects
362 # SQL NULL); reject Python None.
363 if content is None: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 raise ValueError("content is required for chat messages")
366 # Authoritative validation via the enum constructors — raises
367 # ValueError for unknown values, which the route layer maps to HTTP
368 # 400 via ROUTE_EXCEPTIONS. Keeps failure fast (before DB hit) and
369 # avoids the HTTP-500 regression we'd get from letting SQLAlchemy's
370 # StatementError surface at commit time.
371 try:
372 ChatRole(role)
373 except ValueError as exc:
374 raise ValueError(f"Invalid role: {role!r}") from exc
375 try:
376 ChatMessageType(message_type)
377 except ValueError as exc:
378 raise ValueError(f"Invalid message_type: {message_type!r}") from exc
380 message_id = str(uuid.uuid4())
381 # Atomic increment-and-return on ChatSession.message_count.
382 # By default the WHERE clause re-checks `status='active'` so an
383 # archive PATCH racing with a user-message send cannot land on
384 # a now-archived session. When ``allow_archived=True`` (system-
385 # written assistant responses), the filter is relaxed so a
386 # final-report save can complete even if the session flipped to
387 # archived mid-research — losing the answer is worse than the
388 # "archive means stop" semantic for system writes.
389 where_clause: ColumnElement[bool]
390 if allow_archived:
391 where_clause = ChatSession.id == session_id
392 not_found_msg = f"Chat session {session_id} not found"
393 else:
394 where_clause = (ChatSession.id == session_id) & (
395 ChatSession.status == ChatSessionStatus.ACTIVE.value
396 )
397 not_found_msg = f"Chat session {session_id} not found or not active"
398 message_count = self._atomic_increment(
399 db, ChatSession.message_count, where_clause
400 )
401 if message_count is None:
402 raise ValueError(not_found_msg)
404 # message_count tracks surviving rows, so deleting an older attempt
405 # makes it smaller than the largest surviving sequence. The UPDATE
406 # above holds SQLite's write lock until this transaction commits;
407 # concurrent senders cannot allocate the same next sequence here.
408 last_sequence = (
409 db.query(func.max(ChatMessage.sequence_number))
410 .filter(ChatMessage.session_id == session_id)
411 .scalar()
412 )
413 # A caller may disable autoflush while staging several messages in
414 # this transaction. Include those reservations without forcing its
415 # pending work to flush or committing on its behalf.
416 pending_sequence = max(
417 (
418 message.sequence_number
419 for message in db.new
420 if isinstance(message, ChatMessage)
421 and message.session_id == session_id
422 and message.sequence_number is not None
423 ),
424 default=0,
425 )
426 sequence = max(last_sequence or 0, pending_sequence) + 1
428 # created_at populated by column default (utcnow()).
429 message = ChatMessage(
430 id=message_id,
431 session_id=session_id,
432 research_id=research_id,
433 role=role,
434 message_type=message_type,
435 content=content,
436 sequence_number=sequence,
437 )
438 db.add(message)
439 logger.debug(
440 f"Staged message {sequence} for chat {session_id[:8]}... (uncommitted)"
441 )
442 return message_id
444 def add_progress_step(
445 self,
446 session_id: str,
447 research_id: str,
448 content: str,
449 phase: Optional[str] = None,
450 ) -> str:
451 """
452 Add a transient research-progress step for a chat session.
454 Step rows live in chat_progress_steps and have their
455 own per-research sequence (allocated atomically against
456 ResearchHistory.step_count). They do NOT increment the chat
457 session's message_count.
459 Args:
460 session_id: ID of the parent chat session
461 research_id: ID of the research producing the step
462 content: Rendered step text (e.g. "Searching for ...")
463 phase: Optional phase tag from research_service._STEP_PHASES
465 Returns:
466 Step ID (UUID string)
467 """
468 if content is None:
469 raise ValueError("content is required for progress steps")
471 try:
472 step_id = str(uuid.uuid4())
474 with get_user_db_session(self.username) as db:
475 # Atomic increment-and-return on research_history.step_count.
476 sequence = self._atomic_increment(
477 db,
478 ResearchHistory.step_count,
479 ResearchHistory.id == research_id,
480 )
481 if sequence is None:
482 raise ValueError( # noqa: TRY301
483 f"Research {research_id} not found"
484 )
486 step = ChatProgressStep(
487 id=step_id,
488 research_id=research_id,
489 session_id=session_id,
490 phase=phase,
491 content=content,
492 sequence_number=sequence,
493 )
494 db.add(step)
495 db.commit()
497 logger.debug(
498 f"Added progress step {sequence} for research "
499 f"{research_id[:8]}... in chat {session_id[:8]}..."
500 )
501 return step_id
503 except DB_EXCEPTIONS:
504 logger.exception("Error adding progress step")
505 raise
507 def get_session(self, session_id: str) -> Dict[str, Any]:
508 """
509 Get a chat session by ID.
511 Args:
512 session_id: ID of the session
514 Returns:
515 Session data dictionary.
517 Raises:
518 ChatSessionNotFound: if no row matches ``session_id``.
519 Route layer maps to HTTP 404.
520 ChatRepositoryError: if the DB query itself fails. Route
521 layer maps to HTTP 500. Keeping these separate avoids
522 masking transient DB errors as "not found".
523 """
524 try:
525 with get_user_db_session(self.username) as db:
526 session = db.query(ChatSession).filter_by(id=session_id).first()
528 if not session:
529 logger.warning(f"Chat not found: {session_id[:8]}...")
530 # noqa: TRY301 — re-raised by the outer except
531 # ChatSessionNotFound below to propagate as 404.
532 raise ChatSessionNotFound(session_id) # noqa: TRY301
534 return {
535 "id": session.id,
536 "title": session.title,
537 "status": session.status,
538 "message_count": session.message_count,
539 "created_at": _serialize_dt(session.created_at),
540 "accumulated_context": session.accumulated_context,
541 }
543 except ChatSessionNotFound:
544 # Propagate as-is; this is the genuine 404 signal.
545 raise
546 except DB_EXCEPTIONS as exc:
547 logger.exception("Error getting chat session")
548 raise ChatRepositoryError(
549 f"DB error reading session {session_id[:8]}..."
550 ) from exc
552 def get_session_messages(
553 self,
554 session_id: str,
555 limit: int = 50,
556 offset: int = 0,
557 before_created_at: Optional[str] = None,
558 before_id: Optional[str] = None,
559 include_steps: bool = True,
560 ) -> List[Dict[str, Any]]:
561 """
562 Get messages for a session, server-side merged with progress steps.
564 chat_messages.content is always inline; step rows live in
565 chat_progress_steps. This method merges both into a single ordered
566 stream by created_at so the client renderer (chat.js) sees a
567 unified message list with `message_type='step'` rows interleaved.
569 Pagination is SQL-level via per-table LIMIT + Python merge: each
570 table fetches at most ``limit`` rows ordered by created_at DESC,
571 the two streams are merged on (timestamp, sequence, kind, id), and
572 the latest ``limit`` are returned in ASC order so the client
573 renders oldest→newest as before.
575 Pass ``before_created_at`` to fetch the page IMMEDIATELY older
576 than the given ISO timestamp (use the oldest currently-displayed
577 ``created_at`` to implement a "load older messages" trigger).
578 Without the cursor, ``offset`` selects which DESC slice to return
579 (offset=0 → newest, offset=limit → next older window, …).
581 Args:
582 session_id: ID of the session
583 limit: Maximum number of (merged) entries to return
584 offset: Number of entries to skip (DESC-ordered slice index)
585 before_created_at: Optional timezone-aware ISO timestamp — return only
586 entries strictly older than this. Useful for cursor-based
587 "load older" pagination instead of offset arithmetic.
588 before_id: Optional ID of an entry at the cursor timestamp in this
589 session. Resolves its sequence and kind to break timestamp ties;
590 requires before_created_at. Missing/deleted entries are invalid.
591 include_steps: When False, ``chat_progress_steps`` rows are
592 excluded entirely (not merged in) so ``limit`` counts only
593 genuine ``ChatMessage`` conversation turns. Callers that
594 build LLM context (rather than render the client timeline)
595 want this — otherwise a single retried/verbose attempt's
596 progress-step rows can fill the whole merged window and
597 push real conversation history out of it, even though
598 downstream context building discards step rows anyway.
600 Returns:
601 List of message + step data dictionaries, ordered by
602 created_at ascending.
604 Raises:
605 ChatCursorError: The cursor is malformed, lacks a timezone, or
606 identifies no matching boundary entry in this session.
607 """
608 if before_id and not before_created_at:
609 raise ChatCursorError("before_id requires before_created_at")
610 cutoff = None
611 if before_created_at:
612 if not isinstance(before_created_at, str): 612 ↛ 613line 612 didn't jump to line 613 because the condition on line 612 was never true
613 raise ChatCursorError(
614 "before_created_at must be an ISO timestamp with a timezone"
615 )
616 try:
617 cutoff = datetime.fromisoformat(
618 before_created_at.replace("Z", "+00:00")
619 )
620 cutoff = (
621 cutoff.astimezone(UTC)
622 if cutoff.utcoffset() is not None
623 else None
624 )
625 except (ValueError, OverflowError) as exc:
626 raise ChatCursorError(
627 "before_created_at must be an ISO timestamp with a timezone"
628 ) from exc
629 if cutoff is None:
630 raise ChatCursorError(
631 "before_created_at must be an ISO timestamp with a timezone"
632 )
634 try:
635 with get_user_db_session(self.username) as db:
636 msg_q = db.query(ChatMessage).filter_by(session_id=session_id)
637 step_q = db.query(ChatProgressStep).filter_by(
638 session_id=session_id
639 )
640 msg_timestamp = _timestamp_order_key(ChatMessage.created_at)
641 step_timestamp = _timestamp_order_key(
642 ChatProgressStep.created_at
643 )
645 if cutoff is not None:
646 cutoff_key = cutoff.replace(tzinfo=None).isoformat(
647 sep=" ", timespec="microseconds"
648 )
649 if before_id:
650 # UUID order is unrelated to the sequence ordering
651 # used by the page. Resolve the boundary in this session
652 # and compare the same complete key as ORDER BY/merge.
653 # Steps have their own ID namespace in the response.
654 is_step = before_id.startswith("step-")
655 cursor_model = (
656 ChatProgressStep if is_step else ChatMessage
657 )
658 cursor_id = before_id[5:] if is_step else before_id
659 cursor_entry = (
660 db.query(cursor_model)
661 .filter_by(session_id=session_id, id=cursor_id)
662 .first()
663 )
664 if (
665 cursor_entry is None
666 or cursor_entry.created_at != cutoff
667 ):
668 raise ChatCursorError( # noqa: TRY301 - bypasses DB error logging below
669 "before_id must identify an entry in this session "
670 "at before_created_at"
671 )
672 cursor_key = (
673 cutoff_key,
674 cursor_entry.sequence_number,
675 0 if is_step else 1,
676 cursor_id,
677 )
678 msg_q = msg_q.filter(
679 tuple_(
680 msg_timestamp,
681 ChatMessage.sequence_number,
682 literal(1),
683 ChatMessage.id,
684 )
685 < cursor_key
686 )
687 step_q = step_q.filter(
688 tuple_(
689 step_timestamp,
690 ChatProgressStep.sequence_number,
691 literal(0),
692 ChatProgressStep.id,
693 )
694 < cursor_key
695 )
696 else:
697 msg_q = msg_q.filter(msg_timestamp < cutoff_key)
698 step_q = step_q.filter(step_timestamp < cutoff_key)
700 # Pull at most ``limit`` rows from EACH table in DESC
701 # order. The merged window is at most 2 * limit rows
702 # (one extreme: all from one table), which we trim to
703 # ``limit`` after the Python merge. This bounds rows loaded
704 # into Python; SQLite still sorts this session's matching rows
705 # to normalize timestamps with different stored precision.
706 fetch_n = limit + offset
707 # Secondary ORDER BY on sequence_number stabilises rows
708 # whose created_at collide at SQLite's millisecond
709 # precision (sqlalchemy_utc stores `%Y-%m-%d %H:%M:%S.fff`).
710 # Without it, rapid-fire inserts (paste-and-submit,
711 # auto-retries) can be returned in arbitrary order even
712 # though sequence_number is monotonic.
713 messages = (
714 msg_q.order_by(
715 msg_timestamp.desc(),
716 ChatMessage.sequence_number.desc(),
717 ChatMessage.id.desc(),
718 )
719 .limit(fetch_n)
720 .all()
721 )
722 # Skip the step query entirely when the caller only wants
723 # conversation turns (include_steps=False) — an empty steps
724 # list here means `limit` is spent solely on ChatMessage
725 # rows below, instead of step rows crowding them out.
726 steps = (
727 step_q.order_by(
728 step_timestamp.desc(),
729 ChatProgressStep.sequence_number.desc(),
730 ChatProgressStep.id.desc(),
731 )
732 .limit(fetch_n)
733 .all()
734 if include_steps
735 else []
736 )
738 merged: List[Dict[str, Any]] = []
739 for msg in messages:
740 merged.append(
741 {
742 "id": msg.id,
743 "session_id": msg.session_id,
744 "role": msg.role,
745 "message_type": msg.message_type,
746 "content": msg.content,
747 "sequence_number": msg.sequence_number,
748 "research_id": msg.research_id,
749 "created_at": _serialize_dt(msg.created_at),
750 }
751 )
752 for step in steps:
753 merged.append(
754 {
755 "id": f"step-{step.id}",
756 "session_id": step.session_id,
757 "role": "assistant",
758 "message_type": "step",
759 "content": step.content,
760 "phase": step.phase,
761 "sequence_number": step.sequence_number,
762 "research_id": step.research_id,
763 "created_at": _serialize_dt(step.created_at),
764 }
765 )
767 # Sort DESC by (created_at, sequence_number,
768 # step-before-message on tie, id), take the newest
769 # [offset:offset+limit] slice, then flip to ASC so the
770 # client still renders oldest→newest. Including
771 # sequence_number in the Python tie-break mirrors the
772 # SQL ORDER BY above and prevents same-timestamp messages
773 # from rendering out of insertion order. The final ID breaks
774 # ties between steps from different research runs, which can
775 # share both a timestamp and their per-research sequence.
776 merged.sort(
777 key=lambda m: (
778 m["created_at"] or "",
779 m.get("sequence_number") or 0,
780 0 if m["message_type"] == "step" else 1,
781 m["id"],
782 ),
783 reverse=True,
784 )
785 window = merged[offset : offset + limit]
786 window.reverse()
787 return window
789 except ChatCursorError:
790 raise
791 except DB_EXCEPTIONS:
792 # Re-raise so the route returns HTTP 500 instead of a
793 # misleading 200 + []. An empty list here would be
794 # indistinguishable from a session that genuinely has no
795 # messages, hiding infrastructure failures from the client.
796 logger.exception("Error getting chat messages")
797 raise
799 def get_in_progress_research_id(self, session_id: str) -> Optional[str]:
800 """Return the id of the in-progress research for this chat session,
801 or ``None`` if no research is currently running.
803 Used by the GET messages endpoint so the client can restore the
804 live "thinking" indicator on reload without inferring it from
805 message metadata (which fails during the wrapper-strategy
806 preprocessing window before any progress step has persisted).
808 The partial-unique index
809 ``ux_research_history_chat_session_in_progress`` (migration 0010)
810 guarantees at most one matching row exists and turns this into
811 an O(1) index lookup.
812 """
813 try:
814 with get_user_db_session(self.username) as db:
815 row = (
816 db.query(ResearchHistory.id)
817 .filter(
818 ResearchHistory.chat_session_id == session_id,
819 ResearchHistory.status == ResearchStatus.IN_PROGRESS,
820 )
821 .first()
822 )
823 return row[0] if row else None
824 except DB_EXCEPTIONS:
825 # Re-raise rather than swallow → the route handler can
826 # surface a 500 so the client shows an error banner. Returning
827 # None here is indistinguishable from "no research running",
828 # which leaves the send button enabled and lets the user
829 # double-submit into the unique-index guard.
830 logger.exception(
831 "Error fetching in-progress research id for chat session"
832 )
833 raise
835 def list_sessions(
836 self,
837 status: str = ChatSessionStatus.ACTIVE.value,
838 limit: int = 20,
839 offset: int = 0,
840 ) -> List[Dict[str, Any]]:
841 """
842 List chat sessions for the user.
844 Args:
845 status: Filter by status (active, archived, deleted, or all)
846 limit: Maximum number of sessions to return
847 offset: Number of sessions to skip
849 Returns:
850 List of session data dictionaries
851 """
852 try:
853 with get_user_db_session(self.username) as db:
854 query = db.query(ChatSession)
856 if status != "all":
857 query = query.filter_by(status=status)
859 sessions = (
860 query.order_by(ChatSession.created_at.desc())
861 .offset(offset)
862 .limit(limit)
863 .all()
864 )
866 return [
867 {
868 "id": s.id,
869 "title": s.title,
870 "status": s.status,
871 "message_count": s.message_count,
872 "created_at": _serialize_dt(s.created_at),
873 }
874 for s in sessions
875 ]
877 except DB_EXCEPTIONS:
878 # Re-raise so the route returns HTTP 500. Silently returning
879 # [] would make a real DB failure look like a brand-new user
880 # with no sessions, hiding the problem from operators and
881 # confusing the UI.
882 logger.exception("Error listing chat sessions")
883 raise
885 def update_session_title(self, session_id: str, title: str) -> bool:
886 """
887 Update the title of a chat session.
889 Args:
890 session_id: ID of the session
891 title: New title
893 Returns:
894 True if updated successfully
895 """
896 try:
897 with get_user_db_session(self.username) as db:
898 session = db.query(ChatSession).filter_by(id=session_id).first()
899 if session:
900 session.title = title
902 db.commit()
903 return True
904 return False
906 except DB_EXCEPTIONS:
907 logger.exception("Error updating chat session title")
908 return False
910 def reactivate_session(self, session_id: str) -> bool:
911 """
912 Reactivate an archived or deleted chat session.
914 Args:
915 session_id: ID of the session to reactivate
917 Returns:
918 True if reactivated successfully
919 """
920 try:
921 with get_user_db_session(self.username) as db:
922 session = db.query(ChatSession).filter_by(id=session_id).first()
923 if session:
924 session.status = ChatSessionStatus.ACTIVE.value
926 db.commit()
927 logger.info(f"Reactivated chat: {session_id[:8]}...")
928 return True
929 return False
931 except DB_EXCEPTIONS:
932 logger.exception("Error reactivating chat session")
933 return False
935 def archive_session(self, session_id: str) -> bool:
936 """
937 Archive a chat session.
939 Refuses to archive while a research is still in_progress for the
940 session: archive flips the session read-only, and an in-flight
941 research would otherwise survive as an orphaned process writing
942 back into a session the user believes is frozen. The caller (route layer)
943 must stop the research first (or use delete, which terminates
944 in-flight research as a side effect).
946 Args:
947 session_id: ID of the session to archive
949 Returns:
950 True if archived successfully, False if the session does not
951 exist or a DB error occurred.
953 Raises:
954 ArchiveBlockedError: if the session has an in_progress
955 research tied to it. The route layer maps this to HTTP
956 409, mirroring the existing send-to-archived 409 rule.
957 """
958 try:
959 with get_user_db_session(self.username) as db:
960 session = db.query(ChatSession).filter_by(id=session_id).first()
961 if not session:
962 return False
964 in_flight = (
965 db.query(ResearchHistory.id)
966 .filter(
967 ResearchHistory.chat_session_id == session_id,
968 ResearchHistory.status == ResearchStatus.IN_PROGRESS,
969 )
970 .first()
971 )
972 if in_flight is not None:
973 # Bubble up to route layer for 409 mapping. Caught
974 # and re-raised by the inner ``except
975 # ArchiveBlockedError`` below — the broad
976 # ``except DB_EXCEPTIONS`` must not swallow it.
977 raise ArchiveBlockedError( # noqa: TRY301 — re-raised by inner except ArchiveBlockedError
978 "Cannot archive: research in_progress. Stop it first."
979 )
981 session.status = ChatSessionStatus.ARCHIVED.value
982 db.commit()
983 logger.info(f"Archived chat: {session_id[:8]}...")
984 return True
986 except ArchiveBlockedError:
987 # Bubble up so the route layer can produce a 409 response.
988 raise
989 except DB_EXCEPTIONS:
990 logger.exception("Error archiving chat session")
991 return False
993 def delete_session(self, session_id: str) -> bool:
994 """
995 Permanently delete a chat session.
997 Cascades: ChatMessages deleted (CASCADE), ResearchHistory.chat_session_id set NULL.
999 Args:
1000 session_id: ID of the session to delete
1002 Returns:
1003 True if deleted successfully
1004 """
1005 try:
1006 # Terminate any in-progress research tied to this session, so the
1007 # FK's ON DELETE SET NULL doesn't leave it alive with a null
1008 # chat_session_id — an orphan that keeps burning LLM cycles for a
1009 # conversation the user already discarded.
1010 #
1011 # Order matters: collect the in-flight ids inside the transaction,
1012 # but set the (in-memory, non-transactional) termination flags only
1013 # AFTER the delete commits. Flagging before the commit would, on a
1014 # commit failure, kill the research of a session that still exists.
1015 # Share the admission lock with worker registration. A thread
1016 # cannot start between the liveness check and committed cleanup.
1017 with (
1018 user_research_start_gate(self.username),
1019 get_user_db_session(self.username) as db,
1020 ):
1021 session = db.query(ChatSession).filter_by(id=session_id).first()
1022 if not session:
1023 return False
1024 in_flight = (
1025 db.query(ResearchHistory.id)
1026 .filter(
1027 ResearchHistory.chat_session_id == session_id,
1028 ResearchHistory.status == ResearchStatus.IN_PROGRESS,
1029 )
1030 .all()
1031 )
1032 crashed = {
1033 rid
1034 for (rid,) in in_flight
1035 if not is_research_thread_alive(rid)
1036 }
1037 if crashed:
1038 # No worker remains to run its finally block. Keep the
1039 # research history, but release its active slot in the
1040 # same transaction as the session deletion. The status
1041 # predicate preserves a concurrently completed result.
1042 db.query(ResearchHistory).filter(
1043 ResearchHistory.id.in_(crashed),
1044 ResearchHistory.status == ResearchStatus.IN_PROGRESS,
1045 ).update(
1046 {
1047 "status": ResearchStatus.SUSPENDED,
1048 "completed_at": datetime.now(UTC).isoformat(),
1049 },
1050 synchronize_session=False,
1051 )
1052 db.query(UserActiveResearch).filter(
1053 UserActiveResearch.username == self.username,
1054 UserActiveResearch.research_id.in_(crashed),
1055 ).delete(synchronize_session=False)
1056 db.delete(session)
1057 db.commit()
1058 for (rid,) in in_flight:
1059 set_termination_flag(rid)
1060 if rid in crashed:
1061 # An admitted request may still be waiting to register
1062 # its worker. Preserve the stop request so that late
1063 # worker exits at its first cancellation checkpoint.
1064 cleanup_research(rid, preserve_termination_flag=True)
1065 # Include the (truncated) username so a stolen-token bulk
1066 # delete leaves a forensic trail tying each deletion to an
1067 # account, not just an opaque session id (L_SEC1).
1068 logger.info(
1069 f"Deleted chat: user={short_id(self.username)} "
1070 f"session={short_id(session_id)}"
1071 )
1072 return True
1074 except DB_EXCEPTIONS:
1075 logger.exception("Error deleting chat session")
1076 return False
1078 def delete_attempt(self, session_id: str, research_id: str) -> bool:
1079 """Delete one chat attempt, then purge its documents' RAG data.
1081 Thin wrapper over ``delete_attempt_rows`` followed by
1082 ``purge_attempt_rag``; see ``delete_attempt_rows`` for the row
1083 semantics and raised exceptions. Callers holding the per-user
1084 admission gate must call the two halves separately and run the
1085 purge after releasing the gate: the purge does per-collection
1086 vector-store IO under FAISS locks, and the queue tick acquires
1087 every user's gate in turn.
1089 Returns:
1090 True if the attempt existed and was deleted.
1091 """
1092 rag_targets = self.delete_attempt_rows(session_id, research_id)
1093 self.purge_attempt_rag(rag_targets)
1094 return True
1096 def purge_attempt_rag(
1097 self, rag_targets: list[tuple[str, list[str]]]
1098 ) -> None:
1099 """Purge vectors/chunks for documents a deleted attempt owned.
1101 Best-effort: the rows are already committed as deleted, so a
1102 failure is logged rather than raised.
1103 """
1104 if not rag_targets:
1105 return
1106 try:
1107 from ..research_library.deletion.services.document_deletion import (
1108 DocumentDeletionService,
1109 )
1111 DocumentDeletionService(
1112 self.username
1113 ).purge_deleted_research_documents(rag_targets)
1114 except Exception:
1115 logger.exception("Failed to purge RAG data for a deleted attempt")
1117 def delete_attempt_rows(
1118 self, session_id: str, research_id: str
1119 ) -> list[tuple[str, list[str]]]:
1120 """Permanently delete one chat attempt (research + its messages + steps).
1122 An "attempt" is the user message that triggered a research, the
1123 research_history row itself, any assistant response message(s)
1124 tagged with that research, and all chat_progress_steps. The
1125 assistant message(s) and progress steps carry ``research_id``
1126 directly; the user message is inserted with ``research_id=NULL``
1127 (see ``_spawn_chat_research``) and is reachable only via
1128 ``research_meta.submission.message_id`` — both linkages are
1129 resolved below so the user bubble is removed too (issue #4659).
1131 Refuses to delete while the target research is IN_PROGRESS and
1132 its worker thread is alive (raises ``AttemptInProgress`` → 409).
1133 A stale IN_PROGRESS row whose thread is dead is reclaimed: the
1134 status flips to FAILED inside the same transaction that deletes
1135 the rows (mirrors the stale-reclaim sweep in
1136 chat/routes.py:903-923).
1138 Unlike ``delete_session``, this MUST decrement
1139 ``ChatSession.message_count`` — the session still exists, so the
1140 counter would otherwise drift permanently upward. Mirrors the
1141 per-message decrement in ``_cleanup_chat_send_rows``
1142 (chat/routes.py:175-179).
1144 Args:
1145 session_id: ID of the parent chat session (scoped lookup).
1146 research_id: ID of the research attempt to delete.
1148 Does NOT purge RAG data for library documents the attempt owned;
1149 pass the returned targets to ``purge_attempt_rag`` once any
1150 admission gate has been released.
1152 Returns:
1153 The ``(document_id, collection_ids)`` targets whose RAG data
1154 should be purged (empty when there is nothing to purge).
1156 Raises:
1157 AttemptNotFound: research_id does not belong to session_id.
1158 Route layer maps to HTTP 404.
1159 AttemptInProgress: research is IN_PROGRESS and its worker
1160 thread is alive. Route layer maps to HTTP 409 with
1161 ``active_research_id``.
1162 """
1163 try:
1164 # Phase 1 — load + liveness check, OUTSIDE the delete tx so
1165 # the (rare) set_termination_flag call below doesn't have to
1166 # roll back if the worker is still mid-flight. Also lets us
1167 # return AttemptNotFound before touching any rows.
1168 with get_user_db_session(self.username) as db:
1169 research = (
1170 db.query(ResearchHistory)
1171 .filter(
1172 ResearchHistory.id == research_id,
1173 ResearchHistory.chat_session_id == session_id,
1174 )
1175 .first()
1176 )
1177 if research is None:
1178 raise AttemptNotFound(research_id) # noqa: TRY301 — re-raised by outer except
1179 if (
1180 research.status == ResearchStatus.IN_PROGRESS
1181 and is_research_thread_alive(research_id)
1182 ):
1183 # Signal the worker to drain; the route layer tells
1184 # the client to Stop+retry. Don't hard-delete while a
1185 # live thread owns the row — the worker's finally
1186 # block would otherwise write back to a deleted id.
1187 set_termination_flag(research_id)
1188 raise AttemptInProgress(research_id) # noqa: TRY301 — re-raised by outer except
1190 # Phase 2 — atomic delete. The row may be FAILED, COMPLETED,
1191 # SUSPENDED, or stale-IN_PROGRESS (thread dead); all of those
1192 # are safe to delete. Count the ChatMessage rows first so the
1193 # message_count decrement matches the rows removed.
1194 with get_user_db_session(self.username) as db:
1195 # Re-load inside this tx (the row may have changed
1196 # status between Phase 1 and Phase 2 — e.g. the worker
1197 # finished). If the session itself was deleted by a
1198 # concurrent request, the rows we're about to delete
1199 # CASCADE away anyway, so a missing row here is a clean
1200 # 404.
1201 research = (
1202 db.query(ResearchHistory)
1203 .filter(
1204 ResearchHistory.id == research_id,
1205 ResearchHistory.chat_session_id == session_id,
1206 )
1207 .first()
1208 )
1209 if research is None: 1209 ↛ 1210line 1209 didn't jump to line 1210 because the condition on line 1209 was never true
1210 raise AttemptNotFound(research_id) # noqa: TRY301 — re-raised by outer except
1212 # Belt-and-braces: re-check liveness. A thread that was
1213 # dead in Phase 1 can't come back, but a thread that was
1214 # alive (and tripped AttemptInProgress above) would have
1215 # returned already, so this branch only fires for the
1216 # narrow race where the worker went from dead→alive
1217 # between the two phases — not actually possible, kept
1218 # as a defensive guard.
1219 if ( 1219 ↛ 1223line 1219 didn't jump to line 1223 because the condition on line 1219 was never true
1220 research.status == ResearchStatus.IN_PROGRESS
1221 and is_research_thread_alive(research_id)
1222 ):
1223 set_termination_flag(research_id)
1224 raise AttemptInProgress(research_id) # noqa: TRY301 — re-raised by outer except
1226 # Resolve the user (query) message id. In current
1227 # production the user message is inserted by
1228 # _spawn_chat_research with research_id=NULL and linked to
1229 # the attempt only via research_meta.submission.message_id
1230 # (mirrors get_original_attempt_query). Assistant
1231 # response(s) — and legacy pre-research_meta user rows —
1232 # carry research_id directly. We must delete BOTH, otherwise
1233 # the user bubble lingers orphaned after the attempt is
1234 # removed (issue #4659).
1235 user_message_id = None
1236 meta = research.research_meta or {}
1237 submission = meta.get("submission") or {}
1238 if isinstance(submission, dict): 1238 ↛ 1247line 1238 didn't jump to line 1247 because the condition on line 1238 was always true
1239 candidate = submission.get("message_id")
1240 if isinstance(candidate, str) and candidate:
1241 user_message_id = candidate
1243 # research_id matches assistant rows (+ legacy user rows);
1244 # the id branch matches the NULL-research_id user message.
1245 # session_id scopes the id branch so a forged/corrupt
1246 # message_id can't reach another session's row.
1247 msg_filter: ColumnElement[bool] = (
1248 ChatMessage.research_id == research_id
1249 )
1250 if user_message_id:
1251 msg_filter = or_(
1252 msg_filter,
1253 and_(
1254 ChatMessage.id == user_message_id,
1255 ChatMessage.session_id == session_id,
1256 ),
1257 )
1259 # Count messages BEFORE deleting them so the
1260 # message_count decrement is exact.
1261 removed_messages = (
1262 db.query(ChatMessage).filter(msg_filter).count()
1263 )
1265 # ChatMessage rows that reference the research must
1266 # be removed (not SET-NULL-preserved — the chat UI
1267 # treats an attempt's bubbles as part of the attempt,
1268 # so we want them gone with the row). The helper
1269 # below SET-NULLs any ``chat_messages.research_id``
1270 # rows that survived the user_message_id branch; by
1271 # then the assistant rows are gone so the SET NULL
1272 # is a no-op against ``research_id``.
1273 db.query(ChatMessage).filter(msg_filter).delete(
1274 synchronize_session=False
1275 )
1277 # user_active_research row (per-user cap counter) is
1278 # not a child of ``research_history.id`` via FK so
1279 # the helper does not touch it. Delete it here so
1280 # the cap counter stays accurate.
1281 db.query(UserActiveResearch).filter(
1282 UserActiveResearch.research_id == research_id
1283 ).delete(synchronize_session=False)
1285 # ``_delete_research_with_dependents`` performs the
1286 # explicit, application-level cascade (see its
1287 # docstring for why we don't trust the schema's
1288 # ``ON DELETE CASCADE`` rules alone). The two
1289 # ``require_*`` guards are off: this path has
1290 # already established that (a) no live worker
1291 # thread owns the row (``AttemptInProgress`` was
1292 # raised above otherwise), and (b) the row may be
1293 # ``IN_PROGRESS`` *stale* (thread dead) — exactly
1294 # the shape ``#5074`` was protecting against, but
1295 # applied through the chat service's own liveness
1296 # gate rather than the helper's IN_PROGRESS check.
1297 rag_targets = collect_research_document_rag_targets(
1298 db, research_id
1299 )
1300 deleted = delete_research_with_dependents(
1301 db,
1302 research_id,
1303 require_status_not_in_progress=False,
1304 require_no_claimed_queue_row=False,
1305 )
1307 # Decrement message_count. Skipped when no messages were
1308 # removed (e.g. an attempt that crashed before any
1309 # assistant response landed) — avoids a pointless UPDATE.
1310 if removed_messages > 0:
1311 db.query(ChatSession).filter(
1312 ChatSession.id == session_id
1313 ).update(
1314 {
1315 ChatSession.message_count: (
1316 ChatSession.message_count - removed_messages
1317 )
1318 },
1319 synchronize_session=False,
1320 )
1322 db.commit()
1324 # Phase 3 — post-commit in-memory cleanup. The worker's own
1325 # finally block calls cleanup_research too, so this is a
1326 # no-op for the in-progress path; for the stale-IN_PROGRESS
1327 # path it frees the slot immediately.
1328 cleanup_research(research_id)
1330 logger.info(
1331 f"Deleted chat attempt: user={short_id(self.username)} "
1332 f"session={short_id(session_id)} "
1333 f"research={short_id(research_id)} "
1334 f"({removed_messages} messages)"
1335 )
1336 return list(rag_targets) if deleted else []
1338 except (AttemptNotFound, AttemptInProgress):
1339 raise
1340 except DB_EXCEPTIONS:
1341 logger.exception("Error deleting chat attempt")
1342 raise
1344 def get_original_attempt_query(
1345 self, session_id: str, research_id: str
1346 ) -> str:
1347 """Return the original user message content for a chat research.
1349 Used by the retry route to re-submit the same query without the
1350 client echoing it back. Looks up ``research_meta.submission.\
1351 message_id`` first (set at send time by chat/routes.py:1058);
1352 falls back to a query on ``ChatMessage.research_id == X AND
1353 role='user'`` for older rows that predate the meta field.
1355 Args:
1356 session_id: ID of the parent chat session (scoped lookup).
1357 research_id: ID of the research attempt.
1359 Returns:
1360 The original user message content as a string.
1362 Raises:
1363 AttemptNotFound: research_id does not belong to session_id,
1364 or no user message is reachable from it. Route layer
1365 maps to HTTP 404.
1366 """
1367 try:
1368 with get_user_db_session(self.username) as db:
1369 # Scope-by-session first: a research_id from another
1370 # session (e.g. user-supplied path param) should 404,
1371 # not silently return that other session's content.
1372 research = (
1373 db.query(ResearchHistory)
1374 .filter(
1375 ResearchHistory.id == research_id,
1376 ResearchHistory.chat_session_id == session_id,
1377 )
1378 .first()
1379 )
1380 if research is None:
1381 raise AttemptNotFound(research_id) # noqa: TRY301 — re-raised by outer except
1383 # Fast path: research_meta carries the original
1384 # message_id set at send time.
1385 user_message_id = None
1386 meta = research.research_meta or {}
1387 submission = meta.get("submission") or {}
1388 if isinstance(submission, dict): 1388 ↛ 1393line 1388 didn't jump to line 1393 because the condition on line 1388 was always true
1389 candidate = submission.get("message_id")
1390 if isinstance(candidate, str) and candidate:
1391 user_message_id = candidate
1393 if user_message_id:
1394 msg = (
1395 db.query(ChatMessage)
1396 .filter(
1397 ChatMessage.id == user_message_id,
1398 ChatMessage.session_id == session_id,
1399 ChatMessage.role == ChatRole.USER.value,
1400 )
1401 .first()
1402 )
1403 if msg is not None and msg.content: 1403 ↛ 1408line 1403 didn't jump to line 1408 because the condition on line 1403 was always true
1404 return str(msg.content)
1406 # Fallback: pre-research_meta rows. Look up the user
1407 # message by research_id + role.
1408 msg = (
1409 db.query(ChatMessage)
1410 .filter(
1411 ChatMessage.research_id == research_id,
1412 ChatMessage.session_id == session_id,
1413 ChatMessage.role == ChatRole.USER.value,
1414 )
1415 .order_by(ChatMessage.created_at.asc())
1416 .first()
1417 )
1418 if msg is None or not msg.content: 1418 ↛ 1419line 1418 didn't jump to line 1419 because the condition on line 1418 was never true
1419 raise AttemptNotFound(research_id) # noqa: TRY301 — re-raised by outer except
1420 return str(msg.content)
1422 except AttemptNotFound:
1423 raise
1424 except DB_EXCEPTIONS:
1425 logger.exception("Error fetching original attempt query")
1426 raise
1428 def update_accumulated_context(
1429 self,
1430 session_id: str,
1431 new_entities: Optional[List[str]] = None,
1432 new_topics: Optional[List[str]] = None,
1433 summary_addition: Optional[str] = None,
1434 ) -> bool:
1435 """
1436 Update the accumulated context for a session.
1438 Args:
1439 session_id: ID of the session
1440 new_entities: New entities to add
1441 new_topics: New topics to add
1442 summary_addition: Text to append to summary
1444 Returns:
1445 True if updated successfully
1446 """
1447 try:
1448 with get_user_db_session(self.username) as db:
1449 # with_for_update() is a no-op on SQLite but provides
1450 # row locking on PostgreSQL/MySQL if ever used
1451 session = (
1452 db.query(ChatSession)
1453 .filter_by(id=session_id)
1454 .with_for_update()
1455 .first()
1456 )
1457 if not session:
1458 return False
1460 # Build a NEW dict and reassign so SQLAlchemy's plain JSON
1461 # column marks the row dirty. In-place mutation of the existing
1462 # dict (or reassigning the same object identity) is not
1463 # detected without MutableDict.as_mutable() — at flush time
1464 # the loaded snapshot equals the current value and no UPDATE
1465 # is emitted. Same convention as research_sources_service.py.
1466 existing_ctx = session.accumulated_context
1467 if not isinstance(existing_ctx, dict):
1468 # Corrupt/legacy non-dict value: start fresh so this
1469 # write repairs the row instead of failing on dict().
1470 existing_ctx = {}
1471 ctx = dict(existing_ctx)
1473 # A dict row can still hold corrupt values (key_entities: 7,
1474 # topics: "abc", [1]). Only keep stored items that are
1475 # strings in a list, so the merge repairs the row instead of
1476 # raising (set(7), unhashable items) or exploding a string
1477 # into characters (set("abc")).
1479 # Merge entities (deduplicate)
1480 if new_entities:
1481 stored = ctx.get("key_entities", [])
1482 existing = (
1483 {x for x in stored if isinstance(x, str)}
1484 if isinstance(stored, list)
1485 else set()
1486 )
1487 existing.update(new_entities)
1488 ctx["key_entities"] = list(existing)[:50]
1490 # Merge topics
1491 if new_topics:
1492 stored = ctx.get("topics", [])
1493 existing = (
1494 {x for x in stored if isinstance(x, str)}
1495 if isinstance(stored, list)
1496 else set()
1497 )
1498 existing.update(new_topics)
1499 ctx["topics"] = list(existing)[:20]
1501 # Append to summary (with size limit)
1502 if summary_addition:
1503 current = ctx.get("summary", "")
1504 if not isinstance(current, str):
1505 current = ""
1506 new_summary = (
1507 f"{current}\n\n{summary_addition}"
1508 if current
1509 else summary_addition
1510 )
1511 ctx["summary"] = new_summary[-8000:] # Keep last 8000 chars
1513 session.accumulated_context = ctx
1514 db.commit()
1515 return True
1517 except DB_EXCEPTIONS:
1518 logger.exception("Error updating accumulated context")
1519 return False
1521 def _fallback_title(self, query: Optional[str]) -> str:
1522 """Non-LLM title used at creation time (never blocks on I/O)."""
1523 if not query:
1524 return f"Chat {datetime.now(UTC).strftime('%Y-%m-%d %H:%M')}"
1525 if len(query) > 100:
1526 return query[:97].strip() + "..."
1527 return query.strip()
1529 def _generate_title(
1530 self,
1531 query: Optional[str],
1532 settings_snapshot: Optional[Dict[str, Any]] = None,
1533 ) -> str:
1534 """
1535 Generate a title from the initial query.
1537 When chat.llm_title_generation is enabled and settings_snapshot is
1538 provided, uses an LLM for concise titles. Otherwise returns the
1539 non-LLM fallback title.
1540 """
1541 if not query:
1542 return self._fallback_title(query)
1544 if settings_snapshot:
1545 from ..config.llm_config import get_llm
1546 from ..config.thread_settings import get_setting_from_snapshot
1547 from ..security.log_sanitizer import strip_control_chars
1549 if get_setting_from_snapshot(
1550 "chat.llm_title_generation",
1551 False,
1552 settings_snapshot=settings_snapshot,
1553 ):
1554 try:
1555 timeout = float(
1556 get_setting_from_snapshot(
1557 "chat.title_llm_timeout_seconds",
1558 _DEFAULT_TITLE_LLM_TIMEOUT_SECONDS,
1559 settings_snapshot=settings_snapshot,
1560 )
1561 )
1562 except (TypeError, ValueError, OverflowError):
1563 timeout = 0
1564 if not isfinite(timeout) or timeout <= 0:
1565 logger.warning(
1566 "Invalid title-generation timeout; falling back to truncation"
1567 )
1568 return self._fallback_title(query)
1569 # Run the blocking invoke in a worker thread so the request
1570 # thread isn't parked past `timeout` by an unresponsive LLM.
1571 # `with ThreadPoolExecutor(...) as pool:` would call
1572 # shutdown(wait=True) on __exit__, defeating the timeout —
1573 # use wait=False + cancel_futures so the timeout actually fires.
1574 pool = ThreadPoolExecutor(
1575 max_workers=1,
1576 thread_name_prefix="chat-title",
1577 )
1578 try:
1579 llm = get_llm(settings_snapshot=settings_snapshot)
1580 prompt = (
1581 "Generate a concise 3-7 word title for this research "
1582 "query. Return ONLY the title, no quotes or "
1583 f"explanation.\n\nQuery: {query[:200]}"
1584 )
1585 future = pool.submit(llm.invoke, prompt)
1586 try:
1587 response = future.result(timeout=timeout)
1588 except FuturesTimeoutError:
1589 logger.warning(
1590 "LLM title generation exceeded {}s timeout; "
1591 "falling back to truncation",
1592 timeout,
1593 )
1594 return self._fallback_title(query)
1595 # Store generated titles on one line and remove control
1596 # characters while preserving visible Unicode.
1597 title = (
1598 strip_control_chars(
1599 " ".join(str(response.content).splitlines())
1600 )
1601 .strip()
1602 .strip("\"'")[:100]
1603 )
1604 if title:
1605 return title
1606 except Exception:
1607 # User opted into LLM title generation via the
1608 # `chat.llm_title_generation` setting; a silent
1609 # debug-level swallow would hide provider misconfig,
1610 # auth failures, or response-shape regressions in
1611 # production where stderr level is INFO. Log with
1612 # traceback so operators can diagnose, then fall back
1613 # to truncation for UX continuity.
1614 logger.exception(
1615 "LLM title generation failed, falling back to truncation"
1616 )
1617 finally:
1618 pool.shutdown(wait=False, cancel_futures=True)
1620 return self._fallback_title(query)