fix(chat): persist user message and partial reply when stream is aborted (eneo-ai#349) (eneo-ai#412)

* fix(chat): persist user message and partial reply when stream is aborted (eneo-ai#349)

Two-phase persistence so an aborted SSE stream no longer loses the user's
question or the partially-streamed assistant reply.

Phase 1: SessionService.create_question_placeholder inserts a Questions row
with the user's question and an empty answer before the LLM stream begins.
The router's existing setup transaction commits this row before any byte
streams out, so it is durable regardless of what happens during streaming.

Phase 2: complete_question_with_answer (normal completion) or the new
module-level persist_partial_question_answer (abort path) update the row.
The abort helper opens a fresh sessionmanager.session() instead of riding
the request-scoped AsyncSession, which FastAPI tears down concurrently
with the SSE close. The streaming generator's finally fires the partial
save via asyncio.create_task — no await across GeneratorExit, no reliance
on a dying session.

Pattern applied to both assistant_service._handle_response and the
group_chat_service selector-echo path. The main group-chat happy path
delegates to assistant_service.ask and is covered automatically.

Frontend: ChatService.askQuestion no longer early-returns on abort,
letting reloadHistory run so the sidebar reflects the persisted state
without a manual refresh.

Removed the now-orphan SessionService.add_question_to_session — zero
callers in src, tests, frontend, or docs after the swap. 2523 unit tests
remain green; 6 new unit tests + 4 new integration tests against real
Postgres cover the new paths.

Closes eneo-ai#349.

* style(sessions): apply ruff format to session_service.py

* fix(chat): harden abort persistence per review (eneo-ai#412)

Addresses four issues raised in PR review:

1. count_tokens safety. tiktoken can raise on unknown model names, which
   previously killed the asyncio.create_task call in the streaming
   generator's finally and silently collapsed the entire partial save.
   New safe_count_tokens helper catches and falls back to 0. Used both in
   create_question_placeholder (seeds num_tokens_question) and in the two
   _handle_response finally blocks.

2. Skip redundant UPDATE when no content streamed. If the LLM errors
   before yielding any chunks (or the user aborts that early), the
   placeholder already captures the question and an answer="" UPDATE is
   a no-op. Add `if not completed and response_string` guard. Document
   that placeholder rows can appear with answer="" after pre-stream LLM
   errors — that is intentional, the row reflects what the user asked.

3. Strong-ref background tasks. asyncio.create_task internally holds
   only a weak reference; without a strong ref a GC pass can cancel the
   task mid-flight (silent data loss on exactly the path this PR exists
   to protect). New schedule_background_save helper keeps a module-level
   set and removes via add_done_callback. Both assistant_service and
   group_chat_service now route through it.

4. Backfill num_tokens_question on placeholder insert. create_question_placeholder
   now seeds num_tokens_question via safe_count_tokens(question, model),
   so usage analytics no longer undercount aborted requests. The normal-
   completion path overwrites with provider-reported counts as before.

Adds 5 new unit tests:
  - placeholder seeds num_tokens_question from count_tokens
  - placeholder falls back to 0 when count_tokens raises
  - placeholder records 0 when no completion model
  - streaming generator skips partial save when response_string is empty
  - streaming generator schedules partial save even when count_tokens raises
  - schedule_background_save keeps strong ref until task completion

Author: Banald

backend/src/intric/assistants/assistant_service.py

@@ -337,62 +337,88 @@ async def _handle_response(
         completion_model: "CompletionModel",
         session: "SessionInDB",
         stream: bool,
+        question_id: "UUID",
         assistant_selector_tokens: int = 0,
     ):
         """Handle response for group chat selector, matching assistant_service"""
 
+        # Capture tenant_id outside the generator so the abort-path background save
+        # doesn't depend on self.user being safely accessible during teardown.
+        tenant_id = self.user.tenant_id
+
         if stream:
 
             async def response_stream():
-                chunk_response = response.split()
                 response_string = ""
-                for i, chunk in enumerate(chunk_response):
-                    if i < len(chunk_response):
-                        chunk_text = chunk + " "
-                    else:
-                        chunk_text = chunk
-
-                    response_string += chunk_text
-                    # yield empty references and chunk text, matching assistant_service format
-                    yield Completion(
-                        text=chunk_text,
-                        response_type=ResponseType.TEXT,
-                        reference_chunks=[],
+                completed = False
+
+                try:
+                    chunk_response = response.split()
+                    for i, chunk in enumerate(chunk_response):
+                        if i < len(chunk_response):
+                            chunk_text = chunk + " "
+                        else:
+                            chunk_text = chunk
+
+                        response_string += chunk_text
+                        # yield empty references and chunk text, matching assistant_service format
+                        yield Completion(
+                            text=chunk_text,
+                            response_type=ResponseType.TEXT,
+                            reference_chunks=[],
+                        )
+                        await asyncio.sleep(0.05)
+
+                    # NOTE: refactor question_token_count to include the whole contructed prompt.
+                    question_token_count = count_tokens(question, completion_model.name)
+                    token_count = count_tokens(response, completion_model.name)
+                    await self.session_service.complete_question_with_answer(
+                        question_id=question_id,
+                        answer=response,
+                        num_tokens_question=question_token_count
+                        + assistant_selector_tokens,
+                        num_tokens_answer=token_count,
+                        completion_model=completion_model,  # pyright: ignore[reportArgumentType]  # domain.CompletionModel vs ai_models.CompletionModel; structurally compatible at runtime
+                        info_blob_chunks=[],
+                        logging_details=None,
                     )
-                    await asyncio.sleep(0.05)
-
-                # NOTE: refactor question_token_count to include the whole contructed prompt.
-                question_token_count = count_tokens(question, completion_model.name)
-                token_count = count_tokens(response, completion_model.name)
-                await self.session_service.add_question_to_session(
-                    question=question,
-                    answer=response,
-                    num_tokens_question=question_token_count
-                    + assistant_selector_tokens,
-                    num_tokens_answer=token_count,
-                    session=session,
-                    completion_model=completion_model,  # pyright: ignore[reportArgumentType]  # domain.CompletionModel vs ai_models.CompletionModel; structurally compatible at runtime
-                    info_blob_chunks=[],
-                    files=[],
-                    logging_details=None,
-                )
+                    completed = True
+                finally:
+                    # Selector-echo stream did not reach normal completion. The
+                    # placeholder already captures the question; only schedule a
+                    # background UPDATE when there's actual content to persist.
+                    if not completed and response_string:
+                        from intric.sessions.session_service import (
+                            persist_partial_question_answer,
+                            safe_count_tokens,
+                            schedule_background_save,
+                        )
+
+                        partial_tokens_answer = safe_count_tokens(
+                            response_string, completion_model.name
+                        )
+                        schedule_background_save(
+                            persist_partial_question_answer(
+                                tenant_id=tenant_id,
+                                question_id=question_id,
+                                answer=response_string,
+                                num_tokens_answer=partial_tokens_answer,
+                            )
+                        )
 
             return response_stream()
         else:
             # NOTE: refactor question_token_count to include the whole contructed prompt.
             question_token_count = count_tokens(question, completion_model.name)
             token_count = count_tokens(response, completion_model.name)
-            await self.session_service.add_question_to_session(
-                question=question,
+            await self.session_service.complete_question_with_answer(
+                question_id=question_id,
                 answer=response,
                 num_tokens_question=question_token_count + assistant_selector_tokens,
                 num_tokens_answer=token_count,
-                session=session,
                 completion_model=completion_model,  # pyright: ignore[reportArgumentType]  # domain.CompletionModel vs ai_models.CompletionModel; structurally compatible at runtime
                 info_blob_chunks=[],
-                files=[],
                 logging_details=None,
-                assistant_id=None,
             )
             return response
 
@@ -466,12 +492,22 @@ async def ask_group_chat(
             assert (
                 first_completion_model is not None
             )  # assistant must have a model to be usable
+            # Persist a placeholder Question row before the selector echo streams out, so
+            # the user's question survives even if the stream is aborted.
+            question_id = await self.session_service.create_question_placeholder(
+                question=question,
+                session=session,
+                files=[],
+                assistant_id=None,
+                completion_model=first_completion_model,  # pyright: ignore[reportArgumentType]  # domain.CompletionModel vs ai_models.CompletionModel; structurally compatible at runtime
+            )
             final_response = await self._handle_response(
                 response=response_from_selector,
                 question=question,
                 completion_model=first_completion_model,  # pyright: ignore[reportArgumentType]  # domain.CompletionModel vs ai_models.CompletionModel; structurally compatible at runtime
                 session=session,
                 stream=stream,
+                question_id=question_id,
                 assistant_selector_tokens=selection_result.assistant_selector_tokens,
             )
             response = AssistantResponse(
@@ -484,6 +520,7 @@ async def ask_group_chat(
                 tools=UseTools(assistants=[]),
                 description=None,
                 web_search_results=[],
+                question_id=question_id,
             )
         else:
             response = await self.assistant_service.ask(

backend/src/intric/group_chat/application/group_chat_service.py

@@ -25,6 +25,8 @@
 
 if TYPE_CHECKING:
     from intric.completion_models.infrastructure.web_search import WebSearchResult
+    from intric.logging.logging import LoggingDetails
+    from intric.questions.question import ToolCallInfo
 
 
 class QuestionRepository:
@@ -129,6 +131,77 @@ async def _add_web_search_results(
     async def get(self, id: UUID):
         return await self.delegate.get(id)
 
+    async def update_with_answer(
+        self,
+        *,
+        question_id: UUID,
+        tenant_id: UUID,
+        answer: str,
+        num_tokens_question: int | None = None,
+        num_tokens_answer: int | None = None,
+        completion_model_id: UUID | None = None,
+        tool_calls: list["ToolCallInfo"] | None = None,
+        info_blob_chunks: list[InfoBlobChunkInDBWithScore] | None = None,
+        generated_files: list[File] | None = None,
+        web_search_results: list["WebSearchResult"] | None = None,
+        logging_details: "LoggingDetails | None" = None,
+    ) -> None:
+        """Update an existing placeholder Question row with the final or partial answer.
+
+        Used both for normal stream completion (full answer + token counts + late-bound rows)
+        and for partial persistence on abort (just the answer text + estimated token counts).
+
+        tenant_id is required in the WHERE clause to defend against cross-tenant writes if a
+        caller ever supplies a stale question_id.
+        """
+        logging_details_id: object = None
+        if logging_details is not None:
+            log_stmt = (  # pyright: ignore[reportUnknownVariableType]  # logging_table is imperatively mapped
+                sa.insert(logging_table)
+                .values(**logging_details.model_dump())
+                .returning(logging_table)
+            )
+            log_result = await self.session.execute(log_stmt)  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
+            logging_row = log_result.scalar_one()  # pyright: ignore[reportUnknownVariableType]
+            logging_details_id = logging_row.id  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+        update_values: dict[str, Any] = {"answer": answer}
+        if num_tokens_question is not None:
+            update_values["num_tokens_question"] = num_tokens_question
+        if num_tokens_answer is not None:
+            update_values["num_tokens_answer"] = num_tokens_answer
+        if completion_model_id is not None:
+            update_values["completion_model_id"] = completion_model_id
+        if tool_calls is not None:
+            update_values["tool_calls"] = [tc.model_dump() for tc in tool_calls]
+        if logging_details_id is not None:
+            update_values["logging_details_id"] = logging_details_id
+
+        update_stmt = (
+            sa.update(Questions)
+            .where(Questions.id == question_id)
+            .where(Questions.tenant_id == tenant_id)
+            .values(**update_values)
+        )
+        await self.session.execute(update_stmt)
+
+        if info_blob_chunks:
+            await self._add_references(
+                question_id=question_id,  # type: ignore[arg-type]  # helper annotated as int but ID is UUID
+                chunks=info_blob_chunks,
+            )
+        if generated_files:
+            await self._add_files(
+                question_id=question_id,  # type: ignore[arg-type]  # helper annotated as int but ID is UUID
+                files=list(generated_files),
+                file_type="assistant",
+            )
+        if web_search_results:
+            await self._add_web_search_results(
+                web_search_results=list(web_search_results),
+                question_id=question_id,
+            )
+
     async def add(
         self,
         question: QuestionAdd,