Merge pull request #2652 from arc53/fix/stream-durability-triad

feat: improve tool call durability on long calls
This commit is contained in:
Alex authored and GitHub committed 2026-08-10 15:38:52 +01:00
commit bac62b04b9
16 files changed
+1840 -116

No files matched your search

+80 -5
View File
@@ -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(
+174 -46
View File
@@ -1,6 +1,7 @@
import datetime
import json
import logging
import threading
import time
import uuid
from typing import Any, Dict, Generator, List, Optional
@@ -23,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
@@ -37,6 +41,25 @@ 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
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="/")
@@ -318,12 +341,13 @@ 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
# 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.
@@ -335,7 +359,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 +372,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 +383,92 @@ 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
cancelled = stream_cancelled
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:
state = service.heartbeat_message_state(message_id)
except Exception:
# Swallowed deliberately: a transient DB blip must not
# 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,
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 +590,18 @@ 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()
# 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:
@@ -882,7 +959,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 +979,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,
@@ -1196,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:
@@ -1227,6 +1347,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,
@@ -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,
@@ -402,11 +419,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(
+7
View File
@@ -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),
+42 -2
View File
@@ -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)
+69 -10
View File
@@ -3,12 +3,71 @@
Contains parsers for tabular data files.
"""
import datetime
import math
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 ""
# 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
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():
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 +173,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 +259,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 +274,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)
@@ -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
@@ -873,6 +886,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 +896,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 +940,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 +959,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 +983,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 +1028,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
@@ -993,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."""
@@ -1041,11 +1153,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 "
@@ -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(
+21 -6
View File
@@ -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),
)
@@ -0,0 +1,283 @@
"""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 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
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",
"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
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, 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()
resource.conversation_service = service or _service()
agent = MagicMock()
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 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))
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_state.call_count
time.sleep(0.3)
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
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
service = _service()
resource.conversation_service = service
agent = MagicMock()
agent.gen.side_effect = _chatty()
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_state.call_count
time.sleep(0.3)
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."""
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()
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_state.assert_not_called()
def test_ticker_stops_when_row_goes_terminal(self, mock_mongo_db, flask_app):
"""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 = _service(HeartbeatState.MISSING)
resource.conversation_service = service
agent = MagicMock()
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):
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()
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()
+103
View File
@@ -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)
+14 -14
View File
@@ -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"))
@@ -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)
@@ -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
@@ -0,0 +1,120 @@
"""``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
+49
View File
@@ -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.