From ca345e3bf80c44eb09ef877393348fc1964236e6 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 14:32:30 +0100 Subject: [PATCH 1/4] feat: improve tool call durability on long calls --- application/agents/tool_executor.py | 85 ++++- application/api/answer/routes/base.py | 153 +++++--- .../answer/services/conversation_service.py | 9 + application/api/user/reconciliation.py | 7 + application/parser/file/docling_parser.py | 44 ++- application/parser/file/tabular_parser.py | 71 +++- .../storage/db/repositories/conversations.py | 123 +++++- .../storage/db/repositories/reconciliation.py | 130 ++++++- application/streaming/message_journal.py | 27 +- .../answer/test_stream_heartbeat_ticker.py | 176 +++++++++ tests/parser/file/test_docling_parser.py | 103 +++++ tests/parser/file/test_tabular_parser.py | 28 +- .../file/test_tabular_parser_real_fixtures.py | 189 +++++++++ .../test_reconciliation_liveness.py | 359 ++++++++++++++++++ .../repositories/test_truncate_supersede.py | 155 ++++++++ 15 files changed, 1563 insertions(+), 96 deletions(-) create mode 100644 tests/api/answer/test_stream_heartbeat_ticker.py create mode 100644 tests/parser/file/test_tabular_parser_real_fixtures.py create mode 100644 tests/storage/db/repositories/test_reconciliation_liveness.py create mode 100644 tests/storage/db/repositories/test_truncate_supersede.py diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index 09c3377a..98ee481b 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -1,9 +1,11 @@ import logging import re import uuid -from collections import Counter +from collections import Counter, OrderedDict from typing import Any, Dict, List, Optional, Tuple +from sqlalchemy.exc import IntegrityError + from application.agents.default_tools import ( is_headless_excluded_tool, resolve_tool_by_id, @@ -24,6 +26,16 @@ from application.storage.db.session import db_readonly, db_session logger = logging.getLogger(__name__) +def _is_foreign_key_violation(exc: BaseException) -> bool: + """Whether ``exc`` is a Postgres FK violation (SQLSTATE 23503).""" + if not isinstance(exc, IntegrityError): + return False + pgcode = getattr(getattr(exc, "orig", None), "sqlstate", None) or getattr( + getattr(exc, "orig", None), "pgcode", None + ) + return pgcode == "23503" + + # Tightest provider limit on function-call names (OpenAI: ^[a-zA-Z0-9_-]{1,64}$). _MAX_LLM_NAME_LEN = 64 @@ -122,6 +134,41 @@ def _redact_args_for_log(args: Any) -> Any: return redacted +# ``message_id``s whose ``conversation_messages`` parent has been found +# missing. Every tool-call journal write for such a message FK-fails +# identically, and a long tool loop produces one ERROR *pair* per call — 60 +# ERRORs from a single stream in production, the whole api-tier error volume +# for that day. Latch the first failure and log the rest at debug. +# Bounded so a long-lived worker cannot grow this without limit. +_MISSING_PARENTS: "OrderedDict[str, None]" = OrderedDict() +_MISSING_PARENTS_MAX = 256 + + +def _is_missing_parent(message_id: Optional[str]) -> bool: + return bool(message_id) and message_id in _MISSING_PARENTS + + +def _note_missing_parent(message_id: Optional[str], exc: BaseException) -> bool: + """Record a FK failure against ``message_id``. True if newly latched. + + Args: + message_id: The message the write was scoped to, if any. + exc: The exception raised by the write. + + Returns: + True when this is the first sighting for the message (so the caller + should log loudly), False when already latched or not a FK error. + """ + if not message_id or not _is_foreign_key_violation(exc): + return False + if message_id in _MISSING_PARENTS: + return False + _MISSING_PARENTS[message_id] = None + while len(_MISSING_PARENTS) > _MISSING_PARENTS_MAX: + _MISSING_PARENTS.popitem(last=False) + return True + + def _journal_key(call_id: str, message_id: Optional[str]) -> str: """Namespace the durability-journal key by the per-turn ``message_id``. @@ -177,8 +224,26 @@ def _record_proposed( extra={"alert": "tool_call_id_collision", "call_id": call_id}, ) return inserted - except Exception: - logger.exception("tool_call_attempts proposed write failed for %s", call_id) + except Exception as exc: + if _note_missing_parent(message_id, exc): + logger.warning( + "tool_call_attempts: parent message row %s is gone " + "(deleted mid-stream); suppressing further journal errors " + "for this message. First failing call: %s", + message_id, + call_id, + ) + elif _is_missing_parent(message_id): + logger.debug( + "tool_call_attempts proposed write skipped for %s " + "(parent %s already known missing)", + call_id, + message_id, + ) + else: + logger.exception( + "tool_call_attempts proposed write failed for %s", call_id + ) return False @@ -229,8 +294,18 @@ def _mark_executed( user_id=user_id, agent_id=(str(agent_id) if agent_id and looks_like_uuid(str(agent_id)) else None), ) - except Exception: - logger.exception("tool_call_attempts executed write failed for %s", call_id) + except Exception as exc: + if _note_missing_parent(message_id, exc) or _is_missing_parent(message_id): + logger.debug( + "tool_call_attempts executed write skipped for %s " + "(parent %s gone)", + call_id, + message_id, + ) + else: + logger.exception( + "tool_call_attempts executed write failed for %s", call_id + ) def _mark_failed( diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 1a2301f8..341dee9f 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -1,6 +1,7 @@ import datetime import json import logging +import threading import time import uuid from typing import Any, Dict, Generator, List, Optional @@ -37,6 +38,16 @@ from application.utils import check_required_fields logger = logging.getLogger(__name__) +# Seconds between liveness stamps on an in-flight message row. 30 s keeps +# three stamps inside the reconciler's 5-minute staleness window and inside +# the replay watchdog's 90 s producer-idle window, so a single missed tick is +# never enough to trip either. +STREAM_HEARTBEAT_INTERVAL = 30 +# Ceiling on how long the ticker will keep a row alive. Above the realistic +# worst case for a 25-round tool loop, but finite, so a wedged-but-alive +# stream still gets swept eventually. +STREAM_HEARTBEAT_MAX_SECONDS = 3600 + answer_ns = Namespace("answer", description="Answer related operations", path="/") @@ -318,12 +329,10 @@ class BaseAnswerResource: # ``streaming``, yet must still count as live). streaming_marked = False # Heartbeat goes into ``metadata.last_heartbeat_at`` (not - # ``updated_at``, which reconciler-side writes share) and uses - # ``time.monotonic`` so a blocked event loop can't fake fresh. + # ``updated_at``, which reconciler-side writes share). # ``heartbeat_message`` only touches non-terminal rows, so stamping a # still-``pending`` row is safe and does NOT change its status. - STREAM_HEARTBEAT_INTERVAL = 60 - last_heartbeat_at = time.monotonic() + heartbeat_stop: Optional[threading.Event] = None def _mark_streaming_once() -> None: """Flip the reserved row ``pending → streaming`` exactly once. @@ -335,7 +344,7 @@ class BaseAnswerResource: ``_heartbeat_streaming``), so a thought-only reasoning phase that never reaches this point still stays live. """ - nonlocal streaming_marked, last_heartbeat_at + nonlocal streaming_marked if streaming_marked or not reserved_message_id: return try: @@ -348,7 +357,7 @@ class BaseAnswerResource: reserved_message_id, ) # Re-stamp last_heartbeat_at on the transition too; harmless given - # the seed at generation start and the per-interval pump below. + # the seed at generation start and the background ticker below. try: self.conversation_service.heartbeat_message( reserved_message_id, @@ -359,43 +368,68 @@ class BaseAnswerResource: reserved_message_id, ) streaming_marked = True - last_heartbeat_at = time.monotonic() - def _heartbeat_streaming() -> None: - """Pump the liveness heartbeat once per ``STREAM_HEARTBEAT_INTERVAL``. + def _start_heartbeat_ticker() -> "Optional[threading.Event]": + """Stamp the liveness heartbeat on a timer for the stream's life. - Deliberately gated on ``reserved_message_id`` only — NOT on - ``streaming_marked``. The loop calls this for *every* chunk - (including ``thought``/``metadata``), so a reasoning model that - streams only ``thought`` chunks while it "thinks" keeps a still- - ``pending`` row's ``last_heartbeat_at`` fresh and stays out of the - reconciler's staleness sweep. ``heartbeat_message`` only updates - non-terminal rows, so this never resurrects or restatuses a - terminal row. + This replaces a per-chunk pump that could only stamp when a chunk + flowed, which made liveness a function of *output* rather than of + the stream actually being alive. Four windows are routinely silent + — a provider round that emits only tool-call deltas, the body of a + tool call, a same-primary retry or cross-provider fallback, and + mid-execution compression — and any of them longer than the + reconciler's threshold got a healthy stream swept and its answer + discarded. The old pump was also interval-gated against the *last + stamp*, so bursty output could leave the row 60 s staler than the + last chunk suggested. - Residual: a model that emits NO chunks at all (not even - ``thought``) for longer than the reconciler threshold still goes - stale, because this pump only ticks when a chunk flows. Covering a - fully-silent stream would require a background-thread heartbeat or - a higher staleness threshold; both are out of scope here. The - realistic reasoning case (``thought`` chunks streaming) is covered. + Liveness stays honest because the ticker is an in-process daemon + thread: every real death mode we see in production (gunicorn + SIGKILL, worker OOM, ``max_requests`` recycle, host freeze) takes + the thread with it, so the row goes stale on schedule and the + reconciler still does its job. ``STREAM_HEARTBEAT_MAX_SECONDS`` + bounds the other direction — a hung-but-alive stream cannot keep a + row alive forever. + + Returns: + The stop event for the ticker, or None when there is no + reserved row to stamp (headless/``/v1`` continuation rounds). """ - nonlocal last_heartbeat_at if not reserved_message_id: - return - now_mono = time.monotonic() - if now_mono - last_heartbeat_at < STREAM_HEARTBEAT_INTERVAL: - return - try: - self.conversation_service.heartbeat_message( - reserved_message_id, - ) - except Exception: - logger.exception( - "stream heartbeat update failed for %s", - reserved_message_id, - ) - last_heartbeat_at = now_mono + return None + stop = threading.Event() + message_id = reserved_message_id + service = self.conversation_service + + def _tick() -> None: + deadline = time.monotonic() + STREAM_HEARTBEAT_MAX_SECONDS + while not stop.wait(STREAM_HEARTBEAT_INTERVAL): + if time.monotonic() > deadline: + logger.warning( + "stream heartbeat ticker hit its %ss ceiling for " + "message_id=%s; stopping", + STREAM_HEARTBEAT_MAX_SECONDS, + message_id, + ) + return + try: + # Returns False once the row is terminal, which is the + # natural stop signal for a reclaimed or finished row. + if not service.heartbeat_message(message_id): + return + except Exception: + # Swallowed deliberately: a transient DB blip must not + # kill the stream. The reconciler is the backstop. + logger.exception( + "stream heartbeat update failed for %s", message_id, + ) + + threading.Thread( + target=_tick, + daemon=True, + name=f"stream-heartbeat-{message_id[:8]}", + ).start() + return stop # Correlates tool_call_attempts rows with this message. if reserved_message_id and getattr(agent, "tool_executor", None): @@ -517,14 +551,12 @@ class BaseAnswerResource: "generation-start heartbeat seed failed for %s", reserved_message_id, ) - last_heartbeat_at = time.monotonic() + + # The seed above covers t=0; the ticker takes over from the first + # interval onwards, independently of whether anything is flowing. + heartbeat_stop = _start_heartbeat_ticker() for line in gen_iter: - # Cheap closure check that only hits the DB when the heartbeat - # interval has elapsed. Runs for *every* chunk (incl. ``thought`` - # / ``metadata``), so a still-``pending`` reasoning stream stays - # live without waiting for the ``streaming`` status flip. - _heartbeat_streaming() if "metadata" in line: query_metadata.update(line["metadata"]) elif "answer" in line: @@ -882,7 +914,7 @@ class BaseAnswerResource: if should_persist: if reserved_message_id is not None: - self.conversation_service.finalize_message( + finalize_outcome = self.conversation_service.finalize_message( reserved_message_id, response_full, thought=thought, @@ -902,6 +934,29 @@ class BaseAnswerResource: } if llm is not None else None, async_title_generation=llm is not None, ) + # The outcome used to be discarded here, which is how a + # finished answer could vanish silently: if the row was + # deleted mid-stream (retry/edit truncation) the write + # lands nowhere and `activity_finished` still reports + # `ok`, because activity logging never observes the DB. + # Emit a distinct, countable signal instead. + if finalize_outcome is MessageUpdateOutcome.NOT_FOUND: + logger.error( + "answer_persist_failed: message row %s no longer " + "exists; %d chars of answer were produced and " + "could not be saved (conversation=%s)", + reserved_message_id, + len(response_full or ""), + conversation_id, + extra={ + "alert": "answer_persist_failed", + "message_id": reserved_message_id, + "conversation_id": ( + str(conversation_id) if conversation_id else None + ), + "answer_length": len(response_full or ""), + }, + ) else: conversation_id = self.conversation_service.save_conversation( conversation_id, @@ -1227,6 +1282,14 @@ class BaseAnswerResource: if journal_writer is not None: journal_writer.close() return + finally: + # Every exit path — normal, client abort, error — must stop the + # ticker, or a leaked thread keeps stamping a row nobody owns. + # Harmless if it ever does leak (``heartbeat_message`` no-ops on + # terminal rows) but the thread would live until the worker + # recycles. + if heartbeat_stop is not None: + heartbeat_stop.set() def _finalize_stateless_tool_pause( self, diff --git a/application/api/answer/services/conversation_service.py b/application/api/answer/services/conversation_service.py index d89b503a..0cb94ab4 100644 --- a/application/api/answer/services/conversation_service.py +++ b/application/api/answer/services/conversation_service.py @@ -402,11 +402,20 @@ class ConversationService: # Atomic message update + tool_call_attempts confirm; the # ``only_if_non_terminal`` guard prevents a late stream from # retracting a row the reconciler already escalated. + # + # Exception: a *successful* finalize is allowed to reclaim a row the + # reconciler failed for staleness. A stream that reaches this point + # with an answer was self-evidently not stuck, so refusing it throws + # away real output and leaves the user staring at "Response was + # terminated prior to completion". Only on success — a late failure + # must never overwrite the reconciler's failure with a different one. + reclaim = status == "complete" with db_session() as conn: repo = ConversationsRepository(conn) outcome = repo.update_message_by_id( message_id, update_fields, only_if_non_terminal=True, + reclaim_reconciler_failed=reclaim, ) if outcome is not MessageUpdateOutcome.UPDATED: logger.warning( diff --git a/application/api/user/reconciliation.py b/application/api/user/reconciliation.py index d9656b10..83cf3aa9 100644 --- a/application/api/user/reconciliation.py +++ b/application/api/user/reconciliation.py @@ -86,6 +86,13 @@ def run_reconciliation() -> Dict[str, Any]: and conversation_id and pt_repo.delete_state(str(conversation_id), str(user_id)) ): + # Mark the row as having had its approval revoked. A late + # finalize is allowed to reclaim a reconciler-failed row + # (see ``finalize_message``), but must not reclaim this + # one: the resumable state is gone and the UI has already + # been told the approval was cleared, so landing an answer + # here would contradict what the user was shown. + repo.mark_message_approval_cleared(msg["id"]) events.append( ( str(user_id), diff --git a/application/parser/file/docling_parser.py b/application/parser/file/docling_parser.py index 52eaf170..74c2018e 100644 --- a/application/parser/file/docling_parser.py +++ b/application/parser/file/docling_parser.py @@ -94,6 +94,46 @@ def _exceeds_tabular_gate(file: Path) -> bool: return _tabular_content_size(file) > max_bytes +def _delegate( + parser: BaseParser, file: Path, errors: str +) -> Union[str, List[str]]: + """Run a lightweight fallback parser, normalizing its failures. + + Anything the fallback raises is converted to ``DocumentParseError``, + which matters twice over. The Celery tasks that own these paths list + ``DocumentParseError`` in ``dont_autoretry_for``, so a deterministic + content error now fails once instead of retrying four times (each retry + re-downloading the file from S3); and ``SimpleDirectoryReader.load_data`` + catches only ``DocumentParseError``, so one poison spreadsheet is skipped + into ``failed_files`` instead of aborting an entire multi-file ingest. + + This restores the symmetry with ``DoclingParser.parse_file``, whose own + failures are already wrapped — the delegation branch bypassed that. + + Args: + parser: The lightweight parser to delegate to. + file: Path to the file being parsed. + errors: Decoding error policy, forwarded to the parser. + + Returns: + The parsed text, or list of row strings. + + Raises: + DocumentParseError: If the fallback parser fails for any reason. + """ + try: + return parser.parse_file(file, errors) + except DocumentParseError: + raise + except Exception as e: + logger.error( + f"Lightweight parse of {file.name} failed: {e}", exc_info=True + ) + raise DocumentParseError( + f"Failed to parse {file.name}: the file could not be read." + ) from e + + def _capped_markup_copy(file: Path) -> Optional[str]: """Head-truncate an oversized markup file to a temp copy for docling. @@ -413,7 +453,7 @@ class DoclingXLSXParser(DoclingParser): ) from application.parser.file.tabular_parser import ExcelParser - return ExcelParser().parse_file(file, errors) + return _delegate(ExcelParser(), file, errors) return super().parse_file(file, errors) @@ -466,7 +506,7 @@ class DoclingCSVParser(DoclingParser): ) from application.parser.file.tabular_parser import CSVParser - return CSVParser().parse_file(file, errors) + return _delegate(CSVParser(), file, errors) return super().parse_file(file, errors) diff --git a/application/parser/file/tabular_parser.py b/application/parser/file/tabular_parser.py index cc9c821e..d47e69c6 100644 --- a/application/parser/file/tabular_parser.py +++ b/application/parser/file/tabular_parser.py @@ -3,12 +3,63 @@ Contains parsers for tabular data files. """ +import datetime from pathlib import Path from typing import Any, Dict, List, Union from application.parser.file.base_parser import BaseParser +def cell_to_text(value: Any) -> str: + """Render one spreadsheet/CSV cell as text that is always safe to join. + + ``Series.astype(str)`` cannot be trusted for this. Up to pandas 2.x it + rendered missing values as the string ``"nan"``; pandas 3.0 *preserves* + them, so ``.tolist()`` hands back Python ``float`` NaNs and ``str.join`` + raises ``TypeError: sequence item N: expected str instance, float found``. + A single blank cell anywhere in a sheet was therefore enough to destroy + an entire upload. + + Missing values (blank cells, ``NaT``, Excel error values, and formulas + with no cached result) become the empty string. Dates and times are + rendered ISO-8601 rather than pandas' repr. Integral floats drop the + ``.0`` that pandas adds when a blank upcasts an integer column — an ID + column with one gap otherwise turns every id into ``1001.0``. NUL bytes + are stripped here because they crash the downstream Postgres text write. + + Args: + value: A single cell value, of any type pandas or openpyxl may yield. + + Returns: + str: The cell rendered as text; ``""`` for missing values. + """ + if value is None: + return "" + # NaN is the only float that is not equal to itself. Checked before the + # pandas call because it is by far the common case and needs no import. + if isinstance(value, float) and value != value: + return "" + try: + import pandas as pd + + # Raises on array-like cells, which are legitimate values, not NA. + if pd.isna(value): + return "" + except (ImportError, TypeError, ValueError): + pass + if isinstance(value, (datetime.datetime, datetime.date, datetime.time)): + return value.isoformat() + if isinstance(value, float) and value.is_integer(): + return str(int(value)) + text = str(value) + return text.replace("\x00", "") if "\x00" in text else text + + +def _row_to_texts(row: Any) -> List[str]: + """Render a pandas row as a list of join-safe strings.""" + return [cell_to_text(v) for v in row.tolist()] + + class CSVParser(BaseParser): """CSV parser. @@ -114,22 +165,22 @@ class PandasCSVParser(BaseParser): raise ValueError("pandas module is required to read CSV files.") df = pd.read_csv(file, **self._pandas_config) - headers = df.columns.tolist() + headers = [cell_to_text(h) for h in df.columns.tolist()] header_row = f"{self._header_prefix}{self._col_joiner.join(headers)}" if not self._concat_rows: return df.apply( - lambda row: (self._col_joiner).join(row.astype(str).tolist()), axis=1 + lambda row: (self._col_joiner).join(_row_to_texts(row)), axis=1 ).tolist() - + text_list = [] if self._header_period != 1: text_list.append(header_row) - + for i, row in df.iterrows(): if (self._header_period > 1 and i > 0 and i % self._header_period == 0): text_list.append(header_row) - text_list.append(self._col_joiner.join(row.astype(str).tolist())) + text_list.append(self._col_joiner.join(_row_to_texts(row))) if self._header_period == 1 and i < len(df) - 1: text_list.append(header_row) @@ -200,14 +251,14 @@ class ExcelParser(BaseParser): raise ValueError("pandas module is required to read Excel files.") df = pd.read_excel(file, **self._pandas_config) - headers = [str(h) for h in df.columns.tolist()] + headers = [cell_to_text(h) for h in df.columns.tolist()] header_row = f"{self._header_prefix}{self._col_joiner.join(headers)}" - + if not self._concat_rows: return df.apply( - lambda row: (self._col_joiner).join(row.astype(str).tolist()), axis=1 + lambda row: (self._col_joiner).join(_row_to_texts(row)), axis=1 ).tolist() - + text_list = [] if self._header_period != 1: text_list.append(header_row) @@ -215,7 +266,7 @@ class ExcelParser(BaseParser): for i, row in df.iterrows(): if (self._header_period > 1 and i > 0 and i % self._header_period == 0): text_list.append(header_row) - text_list.append(self._col_joiner.join(row.astype(str).tolist())) + text_list.append(self._col_joiner.join(_row_to_texts(row))) if self._header_period == 1 and i < len(df) - 1: text_list.append(header_row) return self._row_joiner.join(text_list) \ No newline at end of file diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index a9e46825..5b504564 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -873,6 +873,7 @@ class ConversationsRepository: def update_message_by_id( self, message_id: str, fields: dict, *, only_if_non_terminal: bool = False, + reclaim_reconciler_failed: bool = False, ) -> MessageUpdateOutcome: """Update specific fields on a message identified by its UUID. @@ -882,9 +883,36 @@ class ConversationsRepository: True, the update is gated so a late finalize cannot retract a reconciler-set ``failed`` (or a prior ``complete``). + ``reclaim_reconciler_failed`` punches one hole in that gate: a row + the **reconciler** failed for staleness may still be finalized by the + stream that owns it, because such a stream demonstrably was not dead. + Production showed a 20-minute tool loop swept at minute 6 and then + completing normally at minute 20, its answer refused and replaced by + "Response was terminated prior to completion". The hole is deliberately + narrow — it matches only the reconciler's own error string, so genuine + terminal failures (client disconnect, provider errors, the pending-tool + cleanup task) stay terminal, and it skips rows whose awaiting-approval + prompt the reconciler already revoked. + + Reclaiming also strips the reconciler's bookkeeping keys, which + otherwise survive the metadata merge and leave a ``complete`` row + carrying a stale ``error``. + The return value discriminates "I updated the row" from "the row was already at a terminal state" so the abort handler can - journal ``end`` when the normal-path finalize already ran. + journal ``end`` when the normal-path finalize already ran. A reclaim + returns plain ``UPDATED``: callers that branch on the outcome should + treat it as an ordinary successful finalize. + + Args: + message_id: UUID of the message row. + fields: Field/value pairs to write. + only_if_non_terminal: Gate the write on a non-terminal status. + reclaim_reconciler_failed: Also allow rows the reconciler failed + for staleness. + + Returns: + MessageUpdateOutcome: What happened to the row. """ if not looks_like_uuid(message_id): return MessageUpdateOutcome.INVALID @@ -899,6 +927,16 @@ class ConversationsRepository: api_to_col = {"metadata": "message_metadata"} + # On reclaim the reconciler's bookkeeping must not survive into the + # finished row: ``error`` would show a stale failure in the API + # response, and leaving the counters behind corrupts any "how often + # does this fire" query. + reclaim_strip = ( + " - 'error' - 'reconcile_attempts' - 'last_reconcile_seen_heartbeat'" + if reclaim_reconciler_failed + else "" + ) + set_parts = [] params: dict = {"id": message_id} for key, val in filtered.items(): @@ -908,7 +946,7 @@ class ConversationsRepository: set_parts.append(f"{col} = NULL") else: set_parts.append( - f"{col} = COALESCE({col}, '{{}}'::jsonb) " + f"{col} = (COALESCE({col}, '{{}}'::jsonb){reclaim_strip}) " f"|| CAST(:{col} AS jsonb)" ) params[col] = ( @@ -932,9 +970,28 @@ class ConversationsRepository: params[col] = val set_parts.append("updated_at = now()") + # Metadata may not be among the written fields (a finalize with no + # query metadata passes none), but the strip must still happen. + if reclaim_reconciler_failed and "metadata" not in filtered: + set_parts.append( + f"message_metadata = COALESCE(message_metadata, '{{}}'::jsonb)" + f"{reclaim_strip}" + ) update_where = ["id = CAST(:id AS uuid)"] if only_if_non_terminal: - update_where.append("status NOT IN ('complete', 'failed')") + terminal_gate = "status NOT IN ('complete', 'failed')" + if reclaim_reconciler_failed: + terminal_gate = ( + "(" + terminal_gate + " OR (" + "status = 'failed' " + "AND message_metadata->>'error' LIKE " + "'reconciler: stuck in pending/streaming%' " + "AND COALESCE(" + "(message_metadata->>'reconciler_cleared_approval')::boolean," + " false) = false" + "))" + ) + update_where.append(terminal_gate) # Single-statement attempt + prior-status probe. Both CTEs see # the same MVCC snapshot, so ``prior.status`` reflects the row # state before the UPDATE — exactly what we need to tell @@ -958,6 +1015,15 @@ class ConversationsRepository: return MessageUpdateOutcome.NOT_FOUND updated, prior_status = row[0], row[1] if updated: + if reclaim_reconciler_failed and prior_status == "failed": + # Distinct, greppable signal: each of these is one user-visible + # answer that would previously have been discarded. Should + # trend to zero once false sweeps stop. + logger.info( + "finalize_message: reclaimed reconciler-failed row " + "message_id=%s", + message_id, + ) return MessageUpdateOutcome.UPDATED if prior_status is None: return MessageUpdateOutcome.NOT_FOUND @@ -1041,11 +1107,62 @@ class ConversationsRepository: Mirrors Mongo's ``$push`` + ``$slice`` that trims queries after an index-based update. + + Rows being deleted may still be **owned by a running stream**: this + runs on the retry/edit path, and a stream survives client disconnect + (the SSE keepalive pump drains the generator in a daemon thread), while + a stream paused on a client-side tool can resume many minutes later. + Deleting such a row cascades its ``message_events`` away and makes + every subsequent journal and ``tool_call_attempts`` write FK-fail — in + production one paused stream, deleted 9 minutes before it resumed, + produced 342 WARNs plus 60 ERRORs and finalized into nothing. + + Deleting is still correct — the user explicitly asked to replace these + turns — so the rows go, but any non-terminal one is first stamped + terminal. That makes the owning stream's late ``finalize_message`` + report ``ALREADY_FAILED`` (a known, quiet outcome) rather than + ``NOT_FOUND``, and leaves an audit trail of what was superseded. + + Args: + conversation_id: Conversation whose tail is being trimmed. + keep_up_to: Highest position to keep. + + Returns: + Number of rows deleted. """ # Shape-gate: see ``rename`` — prevents transaction poisoning when # a non-UUID id reaches this code path. if not looks_like_uuid(conversation_id): return 0 + superseded = self._conn.execute( + text( + """ + UPDATE conversation_messages + SET status = 'failed', + message_metadata = jsonb_set( + COALESCE(message_metadata, '{}'::jsonb), + '{error}', + to_jsonb( + CAST('superseded by a newer turn at this position' + AS text) + ) + ) + WHERE conversation_id = CAST(:conv_id AS uuid) + AND position > :pos + AND status NOT IN ('complete', 'failed') + RETURNING id + """ + ), + {"conv_id": conversation_id, "pos": keep_up_to}, + ).fetchall() + if superseded: + logger.info( + "truncate_after: superseding %d in-flight message(s) in " + "conversation %s: %s", + len(superseded), + conversation_id, + ", ".join(str(r[0]) for r in superseded), + ) result = self._conn.execute( text( "DELETE FROM conversation_messages " diff --git a/application/storage/db/repositories/reconciliation.py b/application/storage/db/repositories/reconciliation.py index b4be9c14..019ce910 100644 --- a/application/storage/db/repositories/reconciliation.py +++ b/application/storage/db/repositories/reconciliation.py @@ -14,7 +14,8 @@ class ReconciliationRepository: self._conn = conn def find_and_lock_stuck_messages( - self, *, age_minutes: int = 5, limit: int = 100, + self, *, age_minutes: int = 5, tool_grace_minutes: int = 15, + limit: int = 100, ) -> list[dict]: """Lock stuck pending/streaming messages skipping live resumes. @@ -27,6 +28,33 @@ class ReconciliationRepository: waiting for resume) and ``resuming`` (actively executing) ``pending_tool_state`` rows so a paused message survives until the PT row's own TTL retires it. + + A second exemption covers **server-executed** tools. + ``pending_tool_state`` is only written on the pause path (client-side + tools, approval gates), so an agent grinding through + ``google_search``/``read_webpage``/``code_executor`` rounds has no PT + row and was swept as "stuck" while perfectly healthy. A message with a + ``proposed``/``executed`` ``tool_call_attempts`` row that transitioned + within ``tool_grace_minutes`` is therefore also exempt. This is + deliberately redundant with the stream's timed heartbeat: that write + shares the app's connection pool, so pool pressure is exactly the + moment liveness should not depend on it, whereas these rows were + committed earlier by a different code path. + + Note the coupling — the ``proposed`` (5 min) and ``executed`` (15 min) + sweeps below both flip rows to ``failed``, which drops them out of this + exemption. That is intentional (a genuinely dead stream's exemption + expires rather than lasting forever), but it means shortening either + sweep silently shortens this grace too. + + Args: + age_minutes: Staleness threshold for the message row. + tool_grace_minutes: How recently a tool call must have transitioned + for its message to count as live. + limit: Maximum rows to lock per tick. + + Returns: + The locked rows as dicts. """ result = self._conn.execute( text( @@ -54,12 +82,24 @@ class ReconciliationRepository: > now() - interval '10 minutes') ) ) + AND NOT EXISTS ( + SELECT 1 + FROM tool_call_attempts tca + WHERE tca.message_id = cm.id + AND tca.status IN ('proposed', 'executed') + AND tca.updated_at + > now() - make_interval(mins => :tool_grace) + ) ORDER BY cm.timestamp ASC LIMIT :limit FOR UPDATE OF cm SKIP LOCKED """ ), - {"age": age_minutes, "limit": limit}, + { + "age": age_minutes, + "tool_grace": tool_grace_minutes, + "limit": limit, + }, ) return [row_to_dict(r) for r in result.fetchall()] @@ -149,21 +189,60 @@ class ReconciliationRepository: return result.rowcount > 0 def increment_message_reconcile_attempts(self, message_id: str) -> int: - """Bump ``message_metadata.reconcile_attempts`` and return the new count.""" + """Bump ``reconcile_attempts``, resetting it if the row proved alive. + + The counter means *consecutive* stale ticks, not a lifetime total. It + used to be a pure accumulator, so a long tool loop that stalled, then + recovered, then stalled again died on its second stall — with a + counter whose name lied about what it counted. Each tick records the + ``last_heartbeat_at`` it observed; if the heartbeat has advanced since + the previous tick, the stream demonstrably produced life in between + and the count restarts at 1. + + This also absorbs a race the timed heartbeat introduces: the + reconciler can lock a row microseconds before the ticker's stamp + commits, and without the reset that phantom attempt would be + permanent. + + Safe against masking a dead stream: ``last_heartbeat_at`` is only ever + written by ``heartbeat_message``, which runs in the process producing + the stream. A dead producer cannot advance it, so the reset condition + can never be met and escalation proceeds unchanged. + + Args: + message_id: The message row to bump. + + Returns: + The new consecutive-stale-tick count. + """ result = self._conn.execute( text( """ UPDATE conversation_messages - SET message_metadata = jsonb_set( - COALESCE(message_metadata, '{}'::jsonb), - '{reconcile_attempts}', - to_jsonb( + SET message_metadata = + jsonb_set( + jsonb_set( + COALESCE(message_metadata, '{}'::jsonb), + '{reconcile_attempts}', + to_jsonb( + CASE + WHEN message_metadata->>'last_heartbeat_at' + IS DISTINCT FROM + message_metadata->>'last_reconcile_seen_heartbeat' + THEN 1 + ELSE COALESCE( + (message_metadata->>'reconcile_attempts')::int, + 0 + ) + 1 + END + ) + ), + '{last_reconcile_seen_heartbeat}', COALESCE( - (message_metadata->>'reconcile_attempts')::int, - 0 - ) + 1 + message_metadata->'last_heartbeat_at', + 'null'::jsonb + ) ) - ) WHERE id = CAST(:message_id AS uuid) RETURNING (message_metadata->>'reconcile_attempts')::int AS new_count @@ -193,6 +272,35 @@ class ReconciliationRepository: ) return result.rowcount > 0 + def mark_message_approval_cleared(self, message_id: str) -> bool: + """Record that this sweep also revoked an awaiting-approval prompt. + + Read by ``finalize_message``'s reclaim path, which refuses to revive a + row whose ``pending_tool_state`` the reconciler already deleted and + whose ``tool.approval.cleared`` event has already been published. + + Args: + message_id: The message row to stamp. + + Returns: + True when the row was stamped. + """ + result = self._conn.execute( + text( + """ + UPDATE conversation_messages + SET message_metadata = jsonb_set( + COALESCE(message_metadata, '{}'::jsonb), + '{reconciler_cleared_approval}', + 'true'::jsonb + ) + WHERE id = CAST(:message_id AS uuid) + """ + ), + {"message_id": message_id}, + ) + return result.rowcount > 0 + def mark_tool_call_failed(self, call_id: str, *, error: str) -> bool: """Flip a tool_call_attempts row to ``failed`` with ``error``.""" result = self._conn.execute( diff --git a/application/streaming/message_journal.py b/application/streaming/message_journal.py index ef0e94b4..9e3e07c7 100644 --- a/application/streaming/message_journal.py +++ b/application/streaming/message_journal.py @@ -233,6 +233,9 @@ class BatchedJournalWriter: self._buffer: list[tuple[int, str, dict[str, Any]]] = [] self._last_flush_mono_ms = time.monotonic() * 1000.0 self._closed = False + # Latched once the parent row is found to be gone. A deleted parent + # never returns, so every subsequent write would FK-fail identically. + self._parent_missing = False def record( self, @@ -250,6 +253,10 @@ class BatchedJournalWriter: event_type, ) return False + if self._parent_missing: + # Parent row deleted mid-stream; nothing can be replayed and no + # future write can succeed. Already logged once by ``flush``. + return False if not event_type: logger.warning( "BatchedJournalWriter.record without event_type: " @@ -317,15 +324,23 @@ class BatchedJournalWriter: ) except IntegrityError as exc: if _is_foreign_key_violation(exc): - # The whole batch references a conversation_messages - # parent that doesn't exist — every per-row retry would - # FK-fail too, so drop the batch once instead of N - # pointless re-attempts. The buffer is already cleared. + # The parent ``conversation_messages`` row is gone. In + # production this means the row was DELETED out from under a + # still-running stream — ``truncate_after`` on a retry or an + # edited question removes the row a live (or paused) stream + # still owns, and ``message_events`` cascades with it. + # + # It never comes back, so latch: one WARN for the stream, then + # stop re-attempting and stop re-logging. Un-latched, a single + # 3.5-minute research stream produced 342 identical WARNs and + # was the entire WARN volume for the day. + self._parent_missing = True logger.warning( "BatchedJournalWriter: no conversation_messages row " "for message_id=%s (foreign-key violation); dropping " - "%d event(s). The parent message row must be committed " - "before its events are journaled.", + "%d event(s) and all further events for this stream. " + "The parent row was deleted while the stream was still " + "running (retry/edit truncation) or never committed.", self._message_id, len(pending), ) diff --git a/tests/api/answer/test_stream_heartbeat_ticker.py b/tests/api/answer/test_stream_heartbeat_ticker.py new file mode 100644 index 00000000..908d0bf5 --- /dev/null +++ b/tests/api/answer/test_stream_heartbeat_ticker.py @@ -0,0 +1,176 @@ +"""The stream liveness heartbeat must be time-based, not output-based. + +Regression cover for the production failure where a 20-minute agent tool +loop was force-failed by the reconciler at minute 6 — because the old +heartbeat only stamped when a chunk flowed, and a tool call emits nothing +while it runs. +""" + +import time +from unittest.mock import MagicMock, patch + +import pytest + + +def _reservation(message_id="11111111-1111-1111-1111-111111111111"): + return {"conversation_id": "22222222-2222-2222-2222-222222222222", + "message_id": message_id} + + +def _silent_then_answer(silence_seconds): + """A generator that emits nothing for a while, then answers. + + Models the real silent windows: a provider round emitting only tool-call + deltas, or the body of a ``read_webpage``/``code_executor`` call. + """ + + def _gen(*args, **kwargs): + time.sleep(silence_seconds) + yield {"answer": "done"} + + return _gen + + +@pytest.mark.unit +class TestHeartbeatTicker: + def _run(self, flask_app, gen, interval=0.05): + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = MagicMock() + service.save_user_question.return_value = _reservation() + service.heartbeat_message.return_value = True + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = gen + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", interval): + list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + ) + return service + + def test_heartbeats_during_a_fully_silent_stream( + self, mock_mongo_db, flask_app + ): + """No chunks flow for 0.4 s; the ticker must still stamp.""" + service = self._run(flask_app, _silent_then_answer(0.4)) + + # 1 seed at generation start + several ticker stamps. + assert service.heartbeat_message.call_count > 2 + + def test_ticker_stops_when_the_stream_ends(self, mock_mongo_db, flask_app): + service = self._run(flask_app, _silent_then_answer(0.1)) + settled = service.heartbeat_message.call_count + + time.sleep(0.3) + + assert service.heartbeat_message.call_count == settled + + def test_ticker_stops_on_client_abort(self, mock_mongo_db, flask_app): + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = MagicMock() + service.save_user_question.return_value = _reservation() + service.heartbeat_message.return_value = True + resource.conversation_service = service + + def _endless(*args, **kwargs): + while True: + time.sleep(0.02) + yield {"answer": "chunk"} + + agent = MagicMock() + agent.gen.side_effect = _endless + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): + gen = resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + next(gen) + next(gen) + time.sleep(0.15) + gen.close() + + settled = service.heartbeat_message.call_count + time.sleep(0.3) + + assert service.heartbeat_message.call_count == settled + + def test_no_ticker_without_a_reserved_row(self, mock_mongo_db, flask_app): + """Headless/continuation rounds have no row to stamp.""" + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = MagicMock() + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = _silent_then_answer(0.2) + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): + list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=False, + ) + ) + + service.heartbeat_message.assert_not_called() + + def test_ticker_stops_when_row_goes_terminal(self, mock_mongo_db, flask_app): + """A False return (row already terminal) is the natural stop signal.""" + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = MagicMock() + service.save_user_question.return_value = _reservation() + # The row is already terminal: every stamp reports "no row updated". + service.heartbeat_message.return_value = False + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = _silent_then_answer(0.6) + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): + list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + ) + + # The seed and the status-flip stamp are unconditional; the ticker + # itself must bail on its first False rather than stamping ~12 + # times over the 0.6 s of silence. + assert service.heartbeat_message.call_count <= 3 diff --git a/tests/parser/file/test_docling_parser.py b/tests/parser/file/test_docling_parser.py index e4e8beee..83612e3a 100644 --- a/tests/parser/file/test_docling_parser.py +++ b/tests/parser/file/test_docling_parser.py @@ -788,3 +788,106 @@ class TestDoclingMarkupGate: dp.DoclingVTTParser().parse_file(path) assert seen["path"] == str(path), "disabled gate must parse the original" + + +# ===================================================================== +# Gate seam: the oversized path through the REAL lightweight parser +# ===================================================================== + + +@pytest.mark.unit +class TestTabularGateSeam: + """Exercise the gate and the lightweight parser together. + + Every test in ``TestDoclingTabularSizeGate`` monkeypatches + ``ExcelParser.parse_file`` away, so the seam between the size gate and + the real fallback had zero coverage — which is how a crash on blank + cells reached production and silently destroyed an upload. + """ + + def _write_xlsx_with_hole(self, path: Path) -> Path: + openpyxl = pytest.importorskip("openpyxl") + wb = openpyxl.Workbook() + ws = wb.active + ws.append([f"col{i}" for i in range(14)]) + ws.append(list(range(14))) + holed = list(range(14)) + holed[11] = None + ws.append(holed) + wb.save(path) + return path + + def test_oversized_xlsx_with_blank_cells_parses_end_to_end( + self, tmp_path, monkeypatch + ): + """The real incident, end to end: gate trips, blanks do not crash.""" + from application.core.settings import settings + from application.parser.file.docling_parser import ( + DoclingParser, + DoclingXLSXParser, + ) + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 64) + docling_parse = MagicMock(name="docling_parse") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + + path = self._write_xlsx_with_hole(tmp_path / "focus.xlsx") + + out = DoclingXLSXParser().parse_file(path) + + docling_parse.assert_not_called() + assert isinstance(out, str) and out + assert "nan" not in out + + def test_lightweight_failure_becomes_document_parse_error( + self, tmp_path, monkeypatch + ): + """A fallback crash must be non-retryable and batch-skippable. + + ``DocumentParseError`` is listed in ``dont_autoretry_for`` on both + ``ingest`` and ``store_attachment``, and is the only exception + ``SimpleDirectoryReader.load_data`` skips rather than propagating. + """ + from application.core.settings import settings + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import ( + DoclingParser, + DoclingXLSXParser, + ) + from application.parser.file.tabular_parser import ExcelParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 64) + monkeypatch.setattr(DoclingParser, "parse_file", MagicMock()) + monkeypatch.setattr( + ExcelParser, + "parse_file", + MagicMock(side_effect=TypeError("sequence item 11")), + ) + + path = self._write_xlsx_with_hole(tmp_path / "boom.xlsx") + + with pytest.raises(DocumentParseError): + DoclingXLSXParser().parse_file(path) + + def test_oversized_csv_failure_becomes_document_parse_error( + self, tmp_path, monkeypatch + ): + from application.core.settings import settings + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import ( + DoclingCSVParser, + DoclingParser, + ) + from application.parser.file.tabular_parser import CSVParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 8) + monkeypatch.setattr(DoclingParser, "parse_file", MagicMock()) + monkeypatch.setattr( + CSVParser, "parse_file", MagicMock(side_effect=ValueError("bad")) + ) + + path = tmp_path / "data.csv" + path.write_text("a,b\n1,2\n" * 20) + + with pytest.raises(DocumentParseError): + DoclingCSVParser().parse_file(path) diff --git a/tests/parser/file/test_tabular_parser.py b/tests/parser/file/test_tabular_parser.py index 30adc505..8cc0cd1d 100644 --- a/tests/parser/file/test_tabular_parser.py +++ b/tests/parser/file/test_tabular_parser.py @@ -66,8 +66,8 @@ def test_pandas_csv_parser_concat_rows(pandas_csv_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["col1", "col2"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value1", "value2"]))), - (1, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value3", "value4"]))) + (0, MagicMock(tolist=lambda: ["value1", "value2"])), + (1, MagicMock(tolist=lambda: ["value3", "value4"])) ] with patch("pandas.read_csv", return_value=mock_df): @@ -92,9 +92,9 @@ def test_pandas_csv_parser_header_period(pandas_csv_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["col1", "col2"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value1", "value2"]))), - (1, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value3", "value4"]))), - (2, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value5", "value6"]))) + (0, MagicMock(tolist=lambda: ["value1", "value2"])), + (1, MagicMock(tolist=lambda: ["value3", "value4"])), + (2, MagicMock(tolist=lambda: ["value5", "value6"])) ] mock_df.__len__.return_value = 3 @@ -108,8 +108,8 @@ def test_excel_parser_concat_rows(excel_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["col1", "col2"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value1", "value2"]))), - (1, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value3", "value4"]))) + (0, MagicMock(tolist=lambda: ["value1", "value2"])), + (1, MagicMock(tolist=lambda: ["value3", "value4"])) ] with patch("pandas.read_excel", return_value=mock_df): @@ -134,8 +134,8 @@ def test_excel_parser_header_period(excel_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["col1", "col2"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value1", "value2"]))), - (1, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value3", "value4"]))) + (0, MagicMock(tolist=lambda: ["value1", "value2"])), + (1, MagicMock(tolist=lambda: ["value3", "value4"])) ] mock_df.__len__.return_value = 2 @@ -163,8 +163,8 @@ def test_pandas_csv_parser_header_period_zero(pandas_csv_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["c1", "c2"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["v1", "v2"]))), - (1, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["v3", "v4"]))), + (0, MagicMock(tolist=lambda: ["v1", "v2"])), + (1, MagicMock(tolist=lambda: ["v3", "v4"])), ] with patch("pandas.read_csv", return_value=mock_df): result = pandas_csv_parser.parse_file(Path("f.csv")) @@ -176,8 +176,8 @@ def test_pandas_csv_parser_header_period_one(pandas_csv_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["a", "b"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["x", "y"]))), - (1, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["m", "n"]))), + (0, MagicMock(tolist=lambda: ["x", "y"])), + (1, MagicMock(tolist=lambda: ["m", "n"])), ] mock_df.__len__.return_value = 2 with patch("pandas.read_csv", return_value=mock_df): @@ -202,7 +202,7 @@ def test_excel_parser_custom_joiners_and_prefix(excel_parser): mock_df = MagicMock() mock_df.columns.tolist.return_value = ["A", "B"] mock_df.iterrows.return_value = [ - (0, MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["x", "y"]))), + (0, MagicMock(tolist=lambda: ["x", "y"])), ] with patch("pandas.read_excel", return_value=mock_df): result = excel_parser.parse_file(Path("t.xlsx")) diff --git a/tests/parser/file/test_tabular_parser_real_fixtures.py b/tests/parser/file/test_tabular_parser_real_fixtures.py new file mode 100644 index 00000000..bfe23893 --- /dev/null +++ b/tests/parser/file/test_tabular_parser_real_fixtures.py @@ -0,0 +1,189 @@ +"""Real-fixture tests for the tabular parsers. + +The mock-based tests in ``test_tabular_parser.py`` stub pandas out entirely +(``MagicMock(astype=lambda _: MagicMock(tolist=lambda: ["value1", "value2"]))``), +so real pandas semantics were never exercised. That is why the pandas 3.0 +change to ``Series.astype(str)`` — which now *preserves* missing values instead +of rendering them as the string ``"nan"`` — reached production as a +``TypeError: sequence item N: expected str instance, float found`` that +destroyed an entire spreadsheet upload. + +These tests drive the parsers with real files. +""" + +import datetime +from pathlib import Path + +import pytest + +from application.parser.file.tabular_parser import ( + ExcelParser, + PandasCSVParser, + cell_to_text, +) + +openpyxl = pytest.importorskip("openpyxl") + + +def _write_xlsx(path: Path, rows: list[list], sheets: dict[str, list[list]] | None = None) -> Path: + """Write ``rows`` to ``path`` as the first sheet, plus any extra ``sheets``.""" + wb = openpyxl.Workbook() + ws = wb.active + for row in rows: + ws.append(row) + for name, extra_rows in (sheets or {}).items(): + extra = wb.create_sheet(title=name) + for row in extra_rows: + extra.append(row) + wb.save(path) + return path + + +def test_excel_blank_cell_in_row_does_not_crash(tmp_path): + """A single blank cell must not destroy the whole document. + + Reproduces the production failure verbatim: a 14-column sheet whose + second data row is blank at column index 11 raises + ``sequence item 11: expected str instance, float found``. + """ + headers = [f"col{i}" for i in range(14)] + full_row = list(range(14)) + holed_row = list(range(14)) + holed_row[11] = None + + path = _write_xlsx(tmp_path / "focus.xlsx", [headers, full_row, holed_row]) + + text = ExcelParser().parse_file(path) + + assert isinstance(text, str) + assert "nan" not in text + # The blank renders as an empty field between two joiners, not as a crash. + assert ", , " in text + + +def test_excel_all_blank_row_does_not_crash(tmp_path): + headers = ["a", "b", "c"] + path = _write_xlsx( + tmp_path / "blanks.xlsx", [headers, [1, 2, 3], [None, None, None]] + ) + + text = ExcelParser().parse_file(path) + + assert isinstance(text, str) + assert "nan" not in text + + +def test_excel_mixed_types_render_readably(tmp_path): + """Ints stay ints, dates are ISO, blanks are empty, nothing is ``nan``.""" + headers = ["id", "name", "ratio", "when", "flag", "blank"] + path = _write_xlsx( + tmp_path / "mixed.xlsx", + [ + headers, + [1001, "alpha", 1.5, datetime.datetime(2024, 1, 2, 3, 4), True, "x"], + [1002, "beta", 2.25, datetime.datetime(2024, 5, 6, 7, 8), False, None], + ], + ) + + text = ExcelParser().parse_file(path) + + assert "nan" not in text + assert "NaT" not in text + # A blank in the last column upcasts nothing: ids must not become 1001.0. + assert "1001" in text and "1001.0" not in text + assert "2024-01-02T03:04:00" in text + + +def test_excel_uncalculated_formula_cells_do_not_crash(tmp_path): + """openpyxl never calculates formulas, so cached values are absent. + + Any workbook written by a non-Excel tool has all-NaN formula columns, + which makes this a guaranteed crash on the pre-fix code. + """ + path = _write_xlsx( + tmp_path / "formulas.xlsx", + [["a", "b", "total"], [1, 2, "=A2+B2"], [3, 4, "=A3+B3"]], + ) + + text = ExcelParser().parse_file(path) + + assert isinstance(text, str) + assert "nan" not in text + + +def test_excel_error_values_do_not_crash(tmp_path): + path = _write_xlsx( + tmp_path / "errors.xlsx", + [["a", "b"], [1, "#REF!"], [2, "#N/A"]], + ) + + text = ExcelParser().parse_file(path) + + assert isinstance(text, str) + assert "nan" not in text + + +def test_excel_non_concat_rows_blank_cell_does_not_crash(tmp_path): + """Covers the ``concat_rows=False`` branch, a separate join site.""" + path = _write_xlsx( + tmp_path / "rows.xlsx", [["a", "b", "c"], [1, None, 3]] + ) + + rows = ExcelParser(concat_rows=False).parse_file(path) + + assert isinstance(rows, list) + assert all(isinstance(r, str) for r in rows) + assert not any("nan" in r for r in rows) + + +def test_nul_byte_in_cell_is_stripped(tmp_path): + """NUL bytes crash the Postgres text write further downstream. + + openpyxl refuses to *write* a NUL into a worksheet + (``IllegalCharacterError``), so the reachable route is CSV — plus the + normalizer itself, which is what a hand-crafted or corrupt workbook + would hit. + """ + assert cell_to_text("di\x00rty") == "dirty" + + path = tmp_path / "nul.csv" + path.write_bytes(b"a,b\nclean,di\x00rty\n") + + text = PandasCSVParser().parse_file(path) + + # pandas' C parser truncates the field at the NUL, so the tail is lost + # before the normalizer sees it. What matters is that no NUL survives + # into the text we hand to Postgres. + assert "\x00" not in text + assert "clean" in text + + +def test_pandas_csv_blank_cell_does_not_crash(tmp_path): + path = tmp_path / "holed.csv" + path.write_text("a,b,c\n1,,3\n4,5,6\n") + + text = PandasCSVParser().parse_file(path) + + assert isinstance(text, str) + assert "nan" not in text + + +def test_pandas_csv_non_string_headers_do_not_crash(tmp_path): + """``header=None`` yields integer column labels; the header join was never hardened.""" + path = tmp_path / "noheader.csv" + path.write_text("1,2,3\n4,5,6\n") + + text = PandasCSVParser(pandas_config={"header": None}).parse_file(path) + + assert isinstance(text, str) + assert text.startswith("HEADERS: ") + + +def test_pandas_csv_non_concat_rows_blank_cell_does_not_crash(tmp_path): + path = tmp_path / "holed2.csv" + path.write_text("a,b,c\n1,,3\n") + + rows = PandasCSVParser(concat_rows=False).parse_file(path) + + assert isinstance(rows, list) + assert not any("nan" in r for r in rows) diff --git a/tests/storage/db/repositories/test_reconciliation_liveness.py b/tests/storage/db/repositories/test_reconciliation_liveness.py new file mode 100644 index 00000000..cf579e05 --- /dev/null +++ b/tests/storage/db/repositories/test_reconciliation_liveness.py @@ -0,0 +1,359 @@ +"""Liveness fixes for the reconciler's stuck-message sweep. + +Production evidence behind these: two healthy streams were force-failed +mid-flight and shown to users as "Response was terminated prior to +completion", one of which went on to finish successfully 20 minutes later +with 1,866 completion tokens. Both were heavy tool loops, whose silent +windows the chunk-gated heartbeat could not cover. +""" + +from __future__ import annotations + +import json + +from sqlalchemy import text + +from application.storage.db.repositories.conversations import ( + ConversationsRepository, + MessageUpdateOutcome, +) +from application.storage.db.repositories.reconciliation import ( + ReconciliationRepository, +) + +RECONCILER_ERROR = ( + "reconciler: stuck in pending/streaming for >5 min after 3 attempts" +) + + +def _seed_message( + conn, + *, + status: str = "pending", + age_minutes: int = 6, + user_id: str = "u", + metadata: dict | None = None, +) -> dict: + conv = ConversationsRepository(conn).create(user_id, "liveness test") + row = conn.execute( + text( + """ + INSERT INTO conversation_messages ( + conversation_id, position, prompt, response, status, user_id, + timestamp, message_metadata + ) + VALUES ( + CAST(:cid AS uuid), 0, 'p', '', :status, :uid, + clock_timestamp() - make_interval(mins => :age), + CAST(:meta AS jsonb) + ) + RETURNING id + """ + ), + { + "cid": conv["id"], + "status": status, + "uid": user_id, + "age": age_minutes, + "meta": json.dumps(metadata or {}), + }, + ).fetchone() + return {"id": str(row[0]), "conversation_id": conv["id"], "user_id": user_id} + + +def _seed_tool_call_for( + conn, message_id: str, *, status: str, age_minutes: int, call_id: str, +) -> None: + conn.execute( + text( + """ + INSERT INTO tool_call_attempts ( + call_id, tool_name, action_name, arguments, status, message_id + ) + VALUES (:cid, 'googlesearch', 'search', '{}'::jsonb, :st, + CAST(:mid AS uuid)) + """ + ), + {"cid": call_id, "st": status, "mid": message_id}, + ) + # The BEFORE-UPDATE trigger resets updated_at; disable it so the + # backdate lands as written. + conn.execute(text("ALTER TABLE tool_call_attempts DISABLE TRIGGER USER")) + try: + conn.execute( + text( + """ + UPDATE tool_call_attempts + SET attempted_at = clock_timestamp() - make_interval(mins => :age), + updated_at = clock_timestamp() - make_interval(mins => :age) + WHERE call_id = :cid + """ + ), + {"cid": call_id, "age": age_minutes}, + ) + finally: + conn.execute(text("ALTER TABLE tool_call_attempts ENABLE TRIGGER USER")) + + +def _meta(conn, message_id: str) -> dict: + row = conn.execute( + text( + "SELECT message_metadata FROM conversation_messages " + "WHERE id = CAST(:id AS uuid)" + ), + {"id": message_id}, + ).fetchone() + return row[0] or {} + + +def _status(conn, message_id: str) -> str: + return conn.execute( + text( + "SELECT status FROM conversation_messages WHERE id = CAST(:id AS uuid)" + ), + {"id": message_id}, + ).fetchone()[0] + + +class TestServerToolLivenessExemption: + """Fix D — server-executed tools leave no ``pending_tool_state`` row.""" + + def test_recent_executed_tool_call_exempts_message(self, pg_conn): + msg = _seed_message(pg_conn, age_minutes=10) + _seed_tool_call_for( + pg_conn, msg["id"], status="executed", age_minutes=1, call_id="c1", + ) + + rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages() + + assert not any(str(r["id"]) == msg["id"] for r in rows) + + def test_recent_proposed_tool_call_exempts_message(self, pg_conn): + msg = _seed_message(pg_conn, age_minutes=10) + _seed_tool_call_for( + pg_conn, msg["id"], status="proposed", age_minutes=2, call_id="c2", + ) + + rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages() + + assert not any(str(r["id"]) == msg["id"] for r in rows) + + def test_stale_tool_call_does_not_exempt(self, pg_conn): + """The exemption expires, so a genuinely dead stream is still swept.""" + msg = _seed_message(pg_conn, age_minutes=30) + _seed_tool_call_for( + pg_conn, msg["id"], status="executed", age_minutes=20, call_id="c3", + ) + + rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages() + + assert any(str(r["id"]) == msg["id"] for r in rows) + + def test_terminal_tool_call_does_not_exempt(self, pg_conn): + """Confirmed/failed rows are not in-flight.""" + msg = _seed_message(pg_conn, age_minutes=10) + _seed_tool_call_for( + pg_conn, msg["id"], status="confirmed", age_minutes=1, call_id="c4", + ) + + rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages() + + assert any(str(r["id"]) == msg["id"] for r in rows) + + +class TestConsecutiveAttemptCounter: + """Fix F — the counter means consecutive stale ticks, not a lifetime total.""" + + def test_counter_accumulates_while_heartbeat_is_unchanged(self, pg_conn): + msg = _seed_message( + pg_conn, metadata={"last_heartbeat_at": "2026-08-10T10:00:00+00:00"}, + ) + repo = ReconciliationRepository(pg_conn) + + assert repo.increment_message_reconcile_attempts(msg["id"]) == 1 + assert repo.increment_message_reconcile_attempts(msg["id"]) == 2 + assert repo.increment_message_reconcile_attempts(msg["id"]) == 3 + + def test_counter_resets_when_heartbeat_advanced(self, pg_conn): + """A stream that proved life between ticks starts over.""" + msg = _seed_message( + pg_conn, metadata={"last_heartbeat_at": "2026-08-10T10:00:00+00:00"}, + ) + repo = ReconciliationRepository(pg_conn) + + assert repo.increment_message_reconcile_attempts(msg["id"]) == 1 + assert repo.increment_message_reconcile_attempts(msg["id"]) == 2 + + pg_conn.execute( + text( + """ + UPDATE conversation_messages + SET message_metadata = jsonb_set( + message_metadata, '{last_heartbeat_at}', + to_jsonb('2026-08-10T10:05:00+00:00'::text) + ) + WHERE id = CAST(:id AS uuid) + """ + ), + {"id": msg["id"]}, + ) + + assert repo.increment_message_reconcile_attempts(msg["id"]) == 1 + + def test_missing_heartbeat_still_accumulates(self, pg_conn): + """A row that never heartbeats must still escalate.""" + msg = _seed_message(pg_conn, metadata={}) + repo = ReconciliationRepository(pg_conn) + + assert repo.increment_message_reconcile_attempts(msg["id"]) == 1 + assert repo.increment_message_reconcile_attempts(msg["id"]) == 2 + assert repo.increment_message_reconcile_attempts(msg["id"]) == 3 + + +class TestFinalizeReclaim: + """Fix C — a late, successful finalize may reclaim a stale-swept row.""" + + def test_reclaims_reconciler_failed_row(self, pg_conn): + msg = _seed_message( + pg_conn, + status="failed", + metadata={"error": RECONCILER_ERROR, "reconcile_attempts": 3}, + ) + repo = ConversationsRepository(pg_conn) + + outcome = repo.update_message_by_id( + msg["id"], + {"response": "the real answer", "status": "complete"}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + assert outcome is MessageUpdateOutcome.UPDATED + assert _status(pg_conn, msg["id"]) == "complete" + + def test_reclaim_strips_reconciler_bookkeeping(self, pg_conn): + msg = _seed_message( + pg_conn, + status="failed", + metadata={ + "error": RECONCILER_ERROR, + "reconcile_attempts": 3, + "last_reconcile_seen_heartbeat": "2026-08-10T10:00:00+00:00", + "keep_me": "yes", + }, + ) + repo = ConversationsRepository(pg_conn) + + repo.update_message_by_id( + msg["id"], + {"response": "answer", "status": "complete", "metadata": {"m": 1}}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + meta = _meta(pg_conn, msg["id"]) + assert "error" not in meta + assert "reconcile_attempts" not in meta + assert "last_reconcile_seen_heartbeat" not in meta + assert meta["keep_me"] == "yes" + assert meta["m"] == 1 + + def test_reclaim_strips_even_without_metadata_field(self, pg_conn): + """A finalize with no query metadata must still clear the keys.""" + msg = _seed_message( + pg_conn, + status="failed", + metadata={"error": RECONCILER_ERROR, "reconcile_attempts": 3}, + ) + repo = ConversationsRepository(pg_conn) + + repo.update_message_by_id( + msg["id"], + {"response": "answer", "status": "complete"}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + meta = _meta(pg_conn, msg["id"]) + assert "error" not in meta + assert "reconcile_attempts" not in meta + + def test_does_not_reclaim_genuine_failure(self, pg_conn): + """Client disconnect, provider errors etc. stay terminal.""" + msg = _seed_message( + pg_conn, + status="failed", + metadata={"error": "ConnectionError: client disconnected"}, + ) + repo = ConversationsRepository(pg_conn) + + outcome = repo.update_message_by_id( + msg["id"], + {"response": "late answer", "status": "complete"}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + assert outcome is MessageUpdateOutcome.ALREADY_FAILED + assert _status(pg_conn, msg["id"]) == "failed" + + def test_does_not_reclaim_when_approval_was_revoked(self, pg_conn): + """The UI was already told the approval was cleared; don't contradict it.""" + msg = _seed_message( + pg_conn, + status="failed", + metadata={ + "error": RECONCILER_ERROR, + "reconciler_cleared_approval": True, + }, + ) + repo = ConversationsRepository(pg_conn) + + outcome = repo.update_message_by_id( + msg["id"], + {"response": "late answer", "status": "complete"}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + assert outcome is MessageUpdateOutcome.ALREADY_FAILED + assert _status(pg_conn, msg["id"]) == "failed" + + def test_does_not_reclaim_a_completed_row(self, pg_conn): + msg = _seed_message(pg_conn, status="complete") + repo = ConversationsRepository(pg_conn) + + outcome = repo.update_message_by_id( + msg["id"], + {"response": "second answer", "status": "complete"}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + assert outcome is MessageUpdateOutcome.ALREADY_COMPLETE + + def test_without_reclaim_flag_reconciler_row_stays_failed(self, pg_conn): + """The old behaviour is unchanged when the flag is off.""" + msg = _seed_message( + pg_conn, + status="failed", + metadata={"error": RECONCILER_ERROR}, + ) + repo = ConversationsRepository(pg_conn) + + outcome = repo.update_message_by_id( + msg["id"], + {"response": "answer", "status": "complete"}, + only_if_non_terminal=True, + ) + + assert outcome is MessageUpdateOutcome.ALREADY_FAILED + + +class TestApprovalClearedStamp: + def test_marks_metadata_flag(self, pg_conn): + msg = _seed_message(pg_conn, status="failed") + repo = ReconciliationRepository(pg_conn) + + assert repo.mark_message_approval_cleared(msg["id"]) is True + assert _meta(pg_conn, msg["id"])["reconciler_cleared_approval"] is True diff --git a/tests/storage/db/repositories/test_truncate_supersede.py b/tests/storage/db/repositories/test_truncate_supersede.py new file mode 100644 index 00000000..81717e5e --- /dev/null +++ b/tests/storage/db/repositories/test_truncate_supersede.py @@ -0,0 +1,155 @@ +"""``truncate_after`` must not silently orphan a still-running stream. + +Reconstructed from production: on 2026-08-09 a paused message row was +deleted by the user's edit at the same position, and 9 minutes later the +original stream resumed and wrote 342 FK-violating journal events before +finalizing into nothing. A second victim the same day was deleted 235 ms +after a concurrent retry started, mid-emit. +""" + +from __future__ import annotations + +from sqlalchemy import text + +from application.storage.db.repositories.conversations import ( + ConversationsRepository, + MessageUpdateOutcome, +) + + +def _seed(conn, conv_id: str, position: int, status: str, user_id="u") -> str: + row = conn.execute( + text( + """ + INSERT INTO conversation_messages ( + conversation_id, position, prompt, response, status, user_id, + timestamp, message_metadata + ) + VALUES (CAST(:cid AS uuid), :pos, 'p', '', :st, :uid, + clock_timestamp(), '{}'::jsonb) + RETURNING id + """ + ), + {"cid": conv_id, "pos": position, "st": status, "uid": user_id}, + ).fetchone() + return str(row[0]) + + +def _row(conn, message_id: str): + return conn.execute( + text( + "SELECT status, message_metadata FROM conversation_messages " + "WHERE id = CAST(:id AS uuid)" + ), + {"id": message_id}, + ).fetchone() + + +class TestTruncateAfterSupersede: + def test_in_flight_row_is_superseded_before_deletion(self, pg_conn): + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "supersede test") + _seed(pg_conn, conv["id"], 0, "complete") + live = _seed(pg_conn, conv["id"], 1, "streaming") + + # The user edits the turn at position 1. + deleted = repo.truncate_after(conv["id"], keep_up_to=0) + + assert deleted == 1 + # The row is gone, as the user asked. + assert _row(pg_conn, live) is None + + def test_late_finalize_reports_not_found_and_is_countable(self, pg_conn): + """The owning stream's late write must be distinguishable. + + ``NOT_FOUND`` is what the route now turns into an + ``answer_persist_failed`` alert, rather than discarding silently. + """ + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "late finalize") + live = _seed(pg_conn, conv["id"], 0, "streaming") + + repo.truncate_after(conv["id"], keep_up_to=-1) + + outcome = repo.update_message_by_id( + live, + {"response": "3115 chars of answer", "status": "complete"}, + only_if_non_terminal=True, + reclaim_reconciler_failed=True, + ) + + assert outcome is MessageUpdateOutcome.NOT_FOUND + + def test_terminal_rows_are_not_restamped(self, pg_conn): + """Only in-flight rows get the supersede stamp.""" + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "terminal") + done = _seed(pg_conn, conv["id"], 0, "complete") + + # keep_up_to=-1 targets everything, but nothing is in flight. + repo.truncate_after(conv["id"], keep_up_to=-1) + + assert _row(pg_conn, done) is None + + def test_pending_row_is_superseded(self, pg_conn): + """A paused row (client-side tool) counts as in flight. + + This is the exact shape of the 2026-08-09 victim: reserved, paused + awaiting a client tool result, deleted by an edit at the same + position, resumed minutes later. + """ + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "paused") + _seed(pg_conn, conv["id"], 0, "complete") + paused = _seed(pg_conn, conv["id"], 1, "pending") + + repo.truncate_after(conv["id"], keep_up_to=0) + + assert _row(pg_conn, paused) is None + + def test_nothing_to_truncate_is_a_noop(self, pg_conn): + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "noop") + kept = _seed(pg_conn, conv["id"], 0, "complete") + + assert repo.truncate_after(conv["id"], keep_up_to=5) == 0 + assert _row(pg_conn, kept) is not None + + def test_invalid_conversation_id_is_rejected(self, pg_conn): + repo = ConversationsRepository(pg_conn) + assert repo.truncate_after("not-a-uuid", keep_up_to=0) == 0 + + +class TestJournalParentMissingLatch: + """One WARN per stream, not one per flush.""" + + def test_latches_after_first_foreign_key_failure(self, caplog): + import logging + + from application.streaming.message_journal import BatchedJournalWriter + + writer = BatchedJournalWriter( + "99999999-9999-9999-9999-999999999999", batch_size=1, + ) + + with caplog.at_level(logging.WARNING): + for seq in range(25): + writer.record(seq, "answer", {"answer": "x"}) + + warnings = [ + r for r in caplog.records + if "no conversation_messages row" in r.getMessage() + ] + assert len(warnings) == 1, ( + f"expected a single latched warning, got {len(warnings)}" + ) + + def test_record_returns_false_once_latched(self): + from application.streaming.message_journal import BatchedJournalWriter + + writer = BatchedJournalWriter( + "88888888-8888-8888-8888-888888888888", batch_size=1, + ) + writer.record(0, "answer", {"answer": "x"}) + + assert writer.record(1, "answer", {"answer": "y"}) is False From 924864f9646288bb0abff1be23ab8d7fa2a00697 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 14:51:33 +0100 Subject: [PATCH 2/4] fix: tests --- .../repositories/test_truncate_supersede.py | 35 ------------- tests/test_message_journal.py | 49 +++++++++++++++++++ 2 files changed, 49 insertions(+), 35 deletions(-) diff --git a/tests/storage/db/repositories/test_truncate_supersede.py b/tests/storage/db/repositories/test_truncate_supersede.py index 81717e5e..8eb800f2 100644 --- a/tests/storage/db/repositories/test_truncate_supersede.py +++ b/tests/storage/db/repositories/test_truncate_supersede.py @@ -118,38 +118,3 @@ class TestTruncateAfterSupersede: def test_invalid_conversation_id_is_rejected(self, pg_conn): repo = ConversationsRepository(pg_conn) assert repo.truncate_after("not-a-uuid", keep_up_to=0) == 0 - - -class TestJournalParentMissingLatch: - """One WARN per stream, not one per flush.""" - - def test_latches_after_first_foreign_key_failure(self, caplog): - import logging - - from application.streaming.message_journal import BatchedJournalWriter - - writer = BatchedJournalWriter( - "99999999-9999-9999-9999-999999999999", batch_size=1, - ) - - with caplog.at_level(logging.WARNING): - for seq in range(25): - writer.record(seq, "answer", {"answer": "x"}) - - warnings = [ - r for r in caplog.records - if "no conversation_messages row" in r.getMessage() - ] - assert len(warnings) == 1, ( - f"expected a single latched warning, got {len(warnings)}" - ) - - def test_record_returns_false_once_latched(self): - from application.streaming.message_journal import BatchedJournalWriter - - writer = BatchedJournalWriter( - "88888888-8888-8888-8888-888888888888", batch_size=1, - ) - writer.record(0, "answer", {"answer": "x"}) - - assert writer.record(1, "answer", {"answer": "y"}) is False diff --git a/tests/test_message_journal.py b/tests/test_message_journal.py index 7461b096..147e18a9 100644 --- a/tests/test_message_journal.py +++ b/tests/test_message_journal.py @@ -648,6 +648,55 @@ class TestBatchedJournalWriter: # Buffer cleared so close() is a no-op and memory stays bounded. assert writer._buffer == [] + def test_missing_parent_is_latched_after_the_first_failure(self, caplog): + """One WARN per stream, not one per flush. + + A parent row deleted mid-stream (retry/edit truncation) never comes + back, so every later flush FK-fails identically. Un-latched, a single + 3.5-minute research stream logged 342 identical WARNs in production + and was that day's entire WARN volume. + """ + import logging + + session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io() + with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls: + mock_session.return_value.__enter__.return_value = MagicMock() + repo = MagicMock(name="repo") + repo.bulk_record.side_effect = _integrity_error("23503") + mock_repo_cls.return_value = repo + mock_topic_cls.return_value = MagicMock() + + writer = BatchedJournalWriter("msg-1", batch_size=1) + with caplog.at_level( + logging.WARNING, logger="application.streaming.message_journal" + ): + for seq in range(25): + writer.record(seq, "answer", {"text": "x"}) + + warnings = [ + r for r in caplog.records + if "no conversation_messages row" in r.getMessage() + ] + assert len(warnings) == 1, ( + f"expected a single latched warning, got {len(warnings)}" + ) + # And it stopped hitting the DB after the first failure. + assert repo.bulk_record.call_count == 1 + + def test_record_returns_false_once_parent_is_known_missing(self): + session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io() + with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls: + mock_session.return_value.__enter__.return_value = MagicMock() + repo = MagicMock(name="repo") + repo.bulk_record.side_effect = _integrity_error("23503") + mock_repo_cls.return_value = repo + mock_topic_cls.return_value = MagicMock() + + writer = BatchedJournalWriter("msg-1", batch_size=1) + + assert writer.record(0, "answer", {"text": "a"}) is True + assert writer.record(1, "answer", {"text": "b"}) is False + def test_bulk_collision_23505_still_falls_back_to_per_row(self): """An explicit 23505 (unique/PK collision) keeps the per-row fallback — the parent exists, only a ``sequence_no`` collided. From b6cc07ad0496784f21141a89405eed86c04fd991 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 15:07:31 +0100 Subject: [PATCH 3/4] fix: mini quality issue --- application/parser/file/tabular_parser.py | 26 +++++++++++++++-------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/application/parser/file/tabular_parser.py b/application/parser/file/tabular_parser.py index d47e69c6..de00f38a 100644 --- a/application/parser/file/tabular_parser.py +++ b/application/parser/file/tabular_parser.py @@ -4,6 +4,7 @@ Contains parsers for tabular data files. """ import datetime +import math from pathlib import Path from typing import Any, Dict, List, Union @@ -35,18 +36,25 @@ def cell_to_text(value: Any) -> str: """ if value is None: return "" - # NaN is the only float that is not equal to itself. Checked before the - # pandas call because it is by far the common case and needs no import. - if isinstance(value, float) and value != value: + # Handled before the pandas lookup below: it is by far the common case + # (every blank cell) and needs no import. + if isinstance(value, float) and math.isnan(value): return "" try: import pandas as pd - - # Raises on array-like cells, which are legitimate values, not NA. - if pd.isna(value): - return "" - except (ImportError, TypeError, ValueError): - pass + except ImportError: + # pandas is optional here — plain float NaN is already covered above, + # and pd.NA/NaT can only reach us via pandas in the first place. + pd = None + if pd is not None: + try: + if pd.isna(value): + return "" + except (TypeError, ValueError): + # pd.isna raises on array-like cells (lists, ndarrays). Those are + # real values rather than missing data, so fall through and + # stringify them below. + pass if isinstance(value, (datetime.datetime, datetime.date, datetime.time)): return value.isoformat() if isinstance(value, float) and value.is_integer(): From 26038da49942262f4f9767e1b49304c53e093af3 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 15:26:03 +0100 Subject: [PATCH 4/4] feat: auto cancel on retry or edit during reconciliation --- application/api/answer/routes/base.py | 77 ++++++- .../answer/services/conversation_service.py | 17 ++ .../storage/db/repositories/conversations.py | 84 ++++++-- .../answer/test_stream_heartbeat_ticker.py | 191 ++++++++++++++---- 4 files changed, 302 insertions(+), 67 deletions(-) diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 341dee9f..64c1b615 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -24,7 +24,10 @@ from application.core.settings import settings from application.error import sanitize_api_error from application.llm.llm_creator import LLMCreator from application.storage.db.repositories.agents import AgentsRepository -from application.storage.db.repositories.conversations import MessageUpdateOutcome +from application.storage.db.repositories.conversations import ( + HeartbeatState, + MessageUpdateOutcome, +) from application.storage.db.repositories.token_usage import TokenUsageRepository from application.storage.db.repositories.user_logs import UserLogsRepository from application.storage.db.session import db_readonly, db_session @@ -49,6 +52,15 @@ STREAM_HEARTBEAT_INTERVAL = 30 STREAM_HEARTBEAT_MAX_SECONDS = 3600 +class StreamSuperseded(Exception): + """Raised to unwind a stream whose message row was deleted mid-flight. + + Not an error condition: the user replaced this turn (retry or edited + question) and `truncate_after` removed the row. Carries the message id + purely for logging. + """ + + answer_ns = Namespace("answer", description="Answer related operations", path="/") @@ -333,6 +345,9 @@ class BaseAnswerResource: # ``heartbeat_message`` only touches non-terminal rows, so stamping a # still-``pending`` row is safe and does NOT change its status. heartbeat_stop: Optional[threading.Event] = None + # Set by the heartbeat ticker when it finds the row gone. Checked by + # the emit loop, which is where the stream regains control. + stream_cancelled = threading.Event() def _mark_streaming_once() -> None: """Flip the reserved row ``pending → streaming`` exactly once. @@ -400,6 +415,7 @@ class BaseAnswerResource: stop = threading.Event() message_id = reserved_message_id service = self.conversation_service + cancelled = stream_cancelled def _tick() -> None: deadline = time.monotonic() + STREAM_HEARTBEAT_MAX_SECONDS @@ -413,16 +429,39 @@ class BaseAnswerResource: ) return try: - # Returns False once the row is terminal, which is the - # natural stop signal for a reclaimed or finished row. - if not service.heartbeat_message(message_id): - return + state = service.heartbeat_message_state(message_id) except Exception: # Swallowed deliberately: a transient DB blip must not - # kill the stream. The reconciler is the backstop. + # kill the stream, and must never be mistaken for the + # row being gone. The reconciler is the backstop. logger.exception( "stream heartbeat update failed for %s", message_id, ) + continue + if state is HeartbeatState.MISSING: + # The row was deleted mid-stream — a retry or an + # edited question truncated this position away. Nothing + # this stream produces can ever be read, so stop the + # work instead of burning tool calls and LLM rounds + # into a void. Production saw a superseded stream run + # 4 further minutes and 12 further rounds. + logger.info( + "stream superseded: message row %s was deleted " + "mid-stream; cancelling", + message_id, + extra={ + "alert": "stream_superseded", + "message_id": message_id, + }, + ) + cancelled.set() + return + if state is HeartbeatState.TERMINAL: + # Row exists but is complete/failed — usually the + # reconciler having swept it. Deliberately NOT + # cancelled: if this stream finishes, finalize is + # allowed to reclaim the row and land the real answer. + return threading.Thread( target=_tick, @@ -557,6 +596,12 @@ class BaseAnswerResource: heartbeat_stop = _start_heartbeat_ticker() for line in gen_iter: + # The emit loop is where the stream regains control between + # rounds, so this is the cheapest place to honour a cancel. + # Bound: a stream sitting inside one long tool call emits + # nothing and is only cancelled when that call returns. + if stream_cancelled.is_set(): + raise StreamSuperseded(reserved_message_id or "") if "metadata" in line: query_metadata.update(line["metadata"]) elif "answer" in line: @@ -1251,6 +1296,26 @@ class BaseAnswerResource: exc_info=True, ) raise + except StreamSuperseded as e: + # Deliberately ahead of the generic handler below: this is not a + # failure and must not be finalized as one. The row is gone, so + # there is nothing to write and nothing to journal (the writer has + # already latched on the same FK violation). The client that + # replaced this turn is watching a different stream. + logger.info( + "stream superseded mid-flight for message_id=%s after " + "%d chars; abandoning without persisting", + str(e), + len(response_full or ""), + extra={ + "alert": "stream_superseded", + "message_id": str(e), + "answer_length": len(response_full or ""), + }, + ) + if journal_writer is not None: + journal_writer.close() + return except Exception as e: logger.error(f"Error in stream: {str(e)}", exc_info=True) if reserved_message_id is not None: diff --git a/application/api/answer/services/conversation_service.py b/application/api/answer/services/conversation_service.py index 0cb94ab4..56920e6f 100644 --- a/application/api/answer/services/conversation_service.py +++ b/application/api/answer/services/conversation_service.py @@ -18,6 +18,7 @@ from application.storage.db.base_repository import looks_like_uuid from application.storage.db.repositories.agents import AgentsRepository from application.storage.db.repositories.conversations import ( ConversationsRepository, + HeartbeatState, MessageUpdateOutcome, ) from application.storage.db.session import db_readonly, db_session @@ -348,6 +349,22 @@ class ConversationService: with db_session() as conn: return ConversationsRepository(conn).heartbeat_message(message_id) + def heartbeat_message_state(self, message_id: str) -> HeartbeatState: + """Heartbeat, reporting whether the row is live, terminal, or gone. + + Args: + message_id: UUID of the message row. + + Returns: + HeartbeatState: What the row was at stamp time. + """ + if not message_id: + return HeartbeatState.MISSING + with db_session() as conn: + return ConversationsRepository(conn).heartbeat_message_state( + message_id + ) + def finalize_message( self, message_id: str, diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index 5b504564..ca10dfd8 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -45,6 +45,19 @@ class MessageUpdateOutcome(str, Enum): INVALID = "invalid" +class HeartbeatState(str, Enum): + """What a liveness heartbeat found when it tried to stamp a row. + + ``TERMINAL`` and ``MISSING`` both mean "not stamped", but they want + opposite handling: a terminal row may still be reclaimed by a finishing + stream, whereas a missing one can never be written again. + """ + + STAMPED = "stamped" + TERMINAL = "terminal" + MISSING = "missing" + + def _message_row_to_dict(row) -> dict: """Like ``row_to_dict`` but renames the DB column ``message_metadata`` back to the public API key ``metadata`` so callers keep the Mongo-era @@ -1059,34 +1072,67 @@ class ConversationsRepository: ) return result.rowcount > 0 - def heartbeat_message(self, message_id: str) -> bool: - """Stamp ``message_metadata.last_heartbeat_at`` with ``clock_timestamp()``. + def heartbeat_message_state(self, message_id: str) -> HeartbeatState: + """Stamp the liveness heartbeat and report what the row looks like. - The reconciler's staleness check uses ``GREATEST(timestamp, - last_heartbeat_at)``, so this call extends a long-running - stream's effective freshness without touching ``timestamp`` (the - creation time, used for history sort) or ``status`` (the WAL - marker). Skips terminal rows so a late heartbeat can't silently - retract a reconciler-set ``failed``. + One round trip, three outcomes. The distinction matters because the + two "not stamped" cases want opposite handling: + + - ``MISSING`` — the row was deleted mid-stream (``truncate_after`` on + a retry/edit). Nothing this stream produces can ever be read, so the + caller should stop working. + - ``TERMINAL`` — the row exists but is ``complete``/``failed``, + typically because the reconciler swept it. The stream must NOT be + cancelled: if it finishes, ``finalize_message`` is allowed to reclaim + a reconciler-failed row and land the real answer. + - ``STAMPED`` — normal in-flight row, freshness extended. + + The stamp deliberately avoids ``timestamp`` (creation time, used for + history sort) and ``status`` (the WAL marker), and skips terminal rows + so a late heartbeat can't retract a reconciler-set ``failed``. + + Args: + message_id: UUID of the message row. + + Returns: + HeartbeatState: What the row was at stamp time. """ if not looks_like_uuid(message_id): - return False - result = self._conn.execute( + return HeartbeatState.MISSING + # Both CTEs read the same MVCC snapshot, so ``present`` describes the + # row as it was when the UPDATE ran — no follow-up SELECT needed. + row = self._conn.execute( text( """ - UPDATE conversation_messages - SET message_metadata = jsonb_set( - COALESCE(message_metadata, '{}'::jsonb), - '{last_heartbeat_at}', - to_jsonb(clock_timestamp()) + WITH stamped AS ( + UPDATE conversation_messages + SET message_metadata = jsonb_set( + COALESCE(message_metadata, '{}'::jsonb), + '{last_heartbeat_at}', + to_jsonb(clock_timestamp()) + ) + WHERE id = CAST(:id AS uuid) + AND status NOT IN ('complete', 'failed') + RETURNING 1 AS ok ) - WHERE id = CAST(:id AS uuid) - AND status NOT IN ('complete', 'failed') + SELECT (SELECT ok FROM stamped) AS stamped, + EXISTS( + SELECT 1 FROM conversation_messages + WHERE id = CAST(:id AS uuid) + ) AS present """ ), {"id": message_id}, - ) - return result.rowcount > 0 + ).fetchone() + if row is not None and row[0]: + return HeartbeatState.STAMPED + if row is not None and row[1]: + return HeartbeatState.TERMINAL + return HeartbeatState.MISSING + + def heartbeat_message(self, message_id: str) -> bool: + """Stamp the heartbeat; True when an in-flight row was updated.""" + return self.heartbeat_message_state(message_id) is HeartbeatState.STAMPED def confirm_executed_tool_calls(self, message_id: str) -> int: """Flip ``tool_call_attempts.status='executed' → 'confirmed'`` for the message.""" diff --git a/tests/api/answer/test_stream_heartbeat_ticker.py b/tests/api/answer/test_stream_heartbeat_ticker.py index 908d0bf5..cf3b0446 100644 --- a/tests/api/answer/test_stream_heartbeat_ticker.py +++ b/tests/api/answer/test_stream_heartbeat_ticker.py @@ -1,9 +1,13 @@ -"""The stream liveness heartbeat must be time-based, not output-based. +"""The stream liveness heartbeat must be time-based, not output-based — +and it is also how a superseded stream learns to stop. -Regression cover for the production failure where a 20-minute agent tool -loop was force-failed by the reconciler at minute 6 — because the old -heartbeat only stamped when a chunk flowed, and a tool call emits nothing -while it runs. +Regression cover for two production failures: + +- a 20-minute agent tool loop force-failed by the reconciler at minute 6, + because the old heartbeat only stamped when a chunk flowed and a tool call + emits nothing while it runs; +- a stream whose row was deleted by the user's retry running four further + minutes and twelve further LLM rounds into a void. """ import time @@ -11,6 +15,8 @@ from unittest.mock import MagicMock, patch import pytest +from application.storage.db.repositories.conversations import HeartbeatState + def _reservation(message_id="11111111-1111-1111-1111-111111111111"): return {"conversation_id": "22222222-2222-2222-2222-222222222222", @@ -31,24 +37,40 @@ def _silent_then_answer(silence_seconds): return _gen +def _chatty(chunks=200, gap=0.02): + """A long multi-round stream that keeps emitting.""" + + def _gen(*args, **kwargs): + for i in range(chunks): + time.sleep(gap) + yield {"answer": f"chunk{i}"} + + return _gen + + +def _service(state=HeartbeatState.STAMPED): + service = MagicMock() + service.save_user_question.return_value = _reservation() + service.heartbeat_message.return_value = True + service.heartbeat_message_state.return_value = state + return service + + @pytest.mark.unit class TestHeartbeatTicker: - def _run(self, flask_app, gen, interval=0.05): + def _run(self, flask_app, gen, service=None, interval=0.05): from application.api.answer.routes import base as base_mod from application.api.answer.routes.base import BaseAnswerResource with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() - service.save_user_question.return_value = _reservation() - service.heartbeat_message.return_value = True - resource.conversation_service = service + resource.conversation_service = service or _service() agent = MagicMock() agent.gen.side_effect = gen with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", interval): - list( + out = list( resource.complete_stream( question="q", agent=agent, @@ -58,24 +80,23 @@ class TestHeartbeatTicker: should_persist=True, ) ) - return service + return resource.conversation_service, out def test_heartbeats_during_a_fully_silent_stream( self, mock_mongo_db, flask_app ): """No chunks flow for 0.4 s; the ticker must still stamp.""" - service = self._run(flask_app, _silent_then_answer(0.4)) + service, _ = self._run(flask_app, _silent_then_answer(0.4)) - # 1 seed at generation start + several ticker stamps. - assert service.heartbeat_message.call_count > 2 + assert service.heartbeat_message_state.call_count > 2 def test_ticker_stops_when_the_stream_ends(self, mock_mongo_db, flask_app): - service = self._run(flask_app, _silent_then_answer(0.1)) - settled = service.heartbeat_message.call_count + service, _ = self._run(flask_app, _silent_then_answer(0.1)) + settled = service.heartbeat_message_state.call_count time.sleep(0.3) - assert service.heartbeat_message.call_count == settled + assert service.heartbeat_message_state.call_count == settled def test_ticker_stops_on_client_abort(self, mock_mongo_db, flask_app): from application.api.answer.routes import base as base_mod @@ -83,18 +104,11 @@ class TestHeartbeatTicker: with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() - service.save_user_question.return_value = _reservation() - service.heartbeat_message.return_value = True + service = _service() resource.conversation_service = service - def _endless(*args, **kwargs): - while True: - time.sleep(0.02) - yield {"answer": "chunk"} - agent = MagicMock() - agent.gen.side_effect = _endless + agent.gen.side_effect = _chatty() with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): gen = resource.complete_stream( @@ -110,10 +124,10 @@ class TestHeartbeatTicker: time.sleep(0.15) gen.close() - settled = service.heartbeat_message.call_count + settled = service.heartbeat_message_state.call_count time.sleep(0.3) - assert service.heartbeat_message.call_count == settled + assert service.heartbeat_message_state.call_count == settled def test_no_ticker_without_a_reserved_row(self, mock_mongo_db, flask_app): """Headless/continuation rounds have no row to stamp.""" @@ -122,7 +136,7 @@ class TestHeartbeatTicker: with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() + service = _service() resource.conversation_service = service agent = MagicMock() @@ -140,26 +154,92 @@ class TestHeartbeatTicker: ) ) - service.heartbeat_message.assert_not_called() + service.heartbeat_message_state.assert_not_called() def test_ticker_stops_when_row_goes_terminal(self, mock_mongo_db, flask_app): - """A False return (row already terminal) is the natural stop signal.""" + """A terminal row stops the ticker but must NOT cancel the stream.""" + service = _service(HeartbeatState.TERMINAL) + service, out = self._run( + flask_app, _silent_then_answer(0.6), service=service, + ) + + # The ticker bailed on its first TERMINAL rather than stamping ~12 + # times over 0.6 s of silence... + assert service.heartbeat_message_state.call_count == 1 + # ...and the stream still finished and finalized, so a + # reconciler-swept row can still be reclaimed. + assert any('"type": "end"' in chunk for chunk in out) + service.finalize_message.assert_called_once() + + +@pytest.mark.unit +class TestSupersededStreamCancellation: + """A deleted row must stop the work, not just quiet the logs.""" + + def _run_with_missing_row(self, flask_app, gen, interval=0.05): from application.api.answer.routes import base as base_mod from application.api.answer.routes.base import BaseAnswerResource with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() - service.save_user_question.return_value = _reservation() - # The row is already terminal: every stamp reports "no row updated". - service.heartbeat_message.return_value = False + service = _service(HeartbeatState.MISSING) resource.conversation_service = service agent = MagicMock() - agent.gen.side_effect = _silent_then_answer(0.6) + agent.gen.side_effect = gen + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", interval): + out = list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + ) + return service, out + + def test_stream_stops_early_when_its_row_is_deleted( + self, mock_mongo_db, flask_app + ): + """The generator must not be drained to completion.""" + service, out = self._run_with_missing_row(flask_app, _chatty(chunks=400)) + + # 400 chunks × 0.02 s ≈ 8 s if drained; cancellation lands far sooner. + assert len(out) < 200, f"stream was not cancelled early ({len(out)})" + + def test_superseded_stream_does_not_persist(self, mock_mongo_db, flask_app): + """Nothing to write — the row is gone. No finalize, no save.""" + service, _ = self._run_with_missing_row(flask_app, _chatty(chunks=400)) + + service.finalize_message.assert_not_called() + service.save_conversation.assert_not_called() + + def test_superseded_stream_emits_no_error_to_the_client( + self, mock_mongo_db, flask_app + ): + """It is not a failure: the user replaced this turn deliberately.""" + _, out = self._run_with_missing_row(flask_app, _chatty(chunks=400)) + + assert not any('"type": "error"' in chunk for chunk in out) + + def test_live_row_is_never_cancelled(self, mock_mongo_db, flask_app): + """The common case must be untouched.""" + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = _service(HeartbeatState.STAMPED) + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = _chatty(chunks=30) with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): - list( + out = list( resource.complete_stream( question="q", agent=agent, @@ -170,7 +250,34 @@ class TestHeartbeatTicker: ) ) - # The seed and the status-flip stamp are unconditional; the ticker - # itself must bail on its first False rather than stamping ~12 - # times over the 0.6 s of silence. - assert service.heartbeat_message.call_count <= 3 + assert any('"type": "end"' in chunk for chunk in out) + service.finalize_message.assert_called_once() + + def test_db_error_does_not_cancel_the_stream(self, mock_mongo_db, flask_app): + """A transient blip must never be mistaken for a deleted row.""" + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = _service() + service.heartbeat_message_state.side_effect = RuntimeError("pool") + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = _chatty(chunks=30) + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): + out = list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + ) + + assert any('"type": "end"' in chunk for chunk in out) + service.finalize_message.assert_called_once()