mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 20:15:26 +00:00
feat: improve tool call durability on long calls
This commit is contained in:
1 parent
7137196212
commit
ca345e3bf8
15 files changed
+1563
-96
No files matched your search
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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 "
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,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
|
||||
@@ -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)
|
||||
@@ -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,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
|
||||
Reference in new issue
Block a user