From 26038da49942262f4f9767e1b49304c53e093af3 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 15:26:03 +0100 Subject: [PATCH] feat: auto cancel on retry or edit during reconciliation --- application/api/answer/routes/base.py | 77 ++++++- .../answer/services/conversation_service.py | 17 ++ .../storage/db/repositories/conversations.py | 84 ++++++-- .../answer/test_stream_heartbeat_ticker.py | 191 ++++++++++++++---- 4 files changed, 302 insertions(+), 67 deletions(-) diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 341dee9f..64c1b615 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -24,7 +24,10 @@ from application.core.settings import settings from application.error import sanitize_api_error from application.llm.llm_creator import LLMCreator from application.storage.db.repositories.agents import AgentsRepository -from application.storage.db.repositories.conversations import MessageUpdateOutcome +from application.storage.db.repositories.conversations import ( + HeartbeatState, + MessageUpdateOutcome, +) from application.storage.db.repositories.token_usage import TokenUsageRepository from application.storage.db.repositories.user_logs import UserLogsRepository from application.storage.db.session import db_readonly, db_session @@ -49,6 +52,15 @@ STREAM_HEARTBEAT_INTERVAL = 30 STREAM_HEARTBEAT_MAX_SECONDS = 3600 +class StreamSuperseded(Exception): + """Raised to unwind a stream whose message row was deleted mid-flight. + + Not an error condition: the user replaced this turn (retry or edited + question) and `truncate_after` removed the row. Carries the message id + purely for logging. + """ + + answer_ns = Namespace("answer", description="Answer related operations", path="/") @@ -333,6 +345,9 @@ class BaseAnswerResource: # ``heartbeat_message`` only touches non-terminal rows, so stamping a # still-``pending`` row is safe and does NOT change its status. heartbeat_stop: Optional[threading.Event] = None + # Set by the heartbeat ticker when it finds the row gone. Checked by + # the emit loop, which is where the stream regains control. + stream_cancelled = threading.Event() def _mark_streaming_once() -> None: """Flip the reserved row ``pending → streaming`` exactly once. @@ -400,6 +415,7 @@ class BaseAnswerResource: stop = threading.Event() message_id = reserved_message_id service = self.conversation_service + cancelled = stream_cancelled def _tick() -> None: deadline = time.monotonic() + STREAM_HEARTBEAT_MAX_SECONDS @@ -413,16 +429,39 @@ class BaseAnswerResource: ) return try: - # Returns False once the row is terminal, which is the - # natural stop signal for a reclaimed or finished row. - if not service.heartbeat_message(message_id): - return + state = service.heartbeat_message_state(message_id) except Exception: # Swallowed deliberately: a transient DB blip must not - # kill the stream. The reconciler is the backstop. + # kill the stream, and must never be mistaken for the + # row being gone. The reconciler is the backstop. logger.exception( "stream heartbeat update failed for %s", message_id, ) + continue + if state is HeartbeatState.MISSING: + # The row was deleted mid-stream — a retry or an + # edited question truncated this position away. Nothing + # this stream produces can ever be read, so stop the + # work instead of burning tool calls and LLM rounds + # into a void. Production saw a superseded stream run + # 4 further minutes and 12 further rounds. + logger.info( + "stream superseded: message row %s was deleted " + "mid-stream; cancelling", + message_id, + extra={ + "alert": "stream_superseded", + "message_id": message_id, + }, + ) + cancelled.set() + return + if state is HeartbeatState.TERMINAL: + # Row exists but is complete/failed — usually the + # reconciler having swept it. Deliberately NOT + # cancelled: if this stream finishes, finalize is + # allowed to reclaim the row and land the real answer. + return threading.Thread( target=_tick, @@ -557,6 +596,12 @@ class BaseAnswerResource: heartbeat_stop = _start_heartbeat_ticker() for line in gen_iter: + # The emit loop is where the stream regains control between + # rounds, so this is the cheapest place to honour a cancel. + # Bound: a stream sitting inside one long tool call emits + # nothing and is only cancelled when that call returns. + if stream_cancelled.is_set(): + raise StreamSuperseded(reserved_message_id or "") if "metadata" in line: query_metadata.update(line["metadata"]) elif "answer" in line: @@ -1251,6 +1296,26 @@ class BaseAnswerResource: exc_info=True, ) raise + except StreamSuperseded as e: + # Deliberately ahead of the generic handler below: this is not a + # failure and must not be finalized as one. The row is gone, so + # there is nothing to write and nothing to journal (the writer has + # already latched on the same FK violation). The client that + # replaced this turn is watching a different stream. + logger.info( + "stream superseded mid-flight for message_id=%s after " + "%d chars; abandoning without persisting", + str(e), + len(response_full or ""), + extra={ + "alert": "stream_superseded", + "message_id": str(e), + "answer_length": len(response_full or ""), + }, + ) + if journal_writer is not None: + journal_writer.close() + return except Exception as e: logger.error(f"Error in stream: {str(e)}", exc_info=True) if reserved_message_id is not None: diff --git a/application/api/answer/services/conversation_service.py b/application/api/answer/services/conversation_service.py index 0cb94ab4..56920e6f 100644 --- a/application/api/answer/services/conversation_service.py +++ b/application/api/answer/services/conversation_service.py @@ -18,6 +18,7 @@ from application.storage.db.base_repository import looks_like_uuid from application.storage.db.repositories.agents import AgentsRepository from application.storage.db.repositories.conversations import ( ConversationsRepository, + HeartbeatState, MessageUpdateOutcome, ) from application.storage.db.session import db_readonly, db_session @@ -348,6 +349,22 @@ class ConversationService: with db_session() as conn: return ConversationsRepository(conn).heartbeat_message(message_id) + def heartbeat_message_state(self, message_id: str) -> HeartbeatState: + """Heartbeat, reporting whether the row is live, terminal, or gone. + + Args: + message_id: UUID of the message row. + + Returns: + HeartbeatState: What the row was at stamp time. + """ + if not message_id: + return HeartbeatState.MISSING + with db_session() as conn: + return ConversationsRepository(conn).heartbeat_message_state( + message_id + ) + def finalize_message( self, message_id: str, diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index 5b504564..ca10dfd8 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -45,6 +45,19 @@ class MessageUpdateOutcome(str, Enum): INVALID = "invalid" +class HeartbeatState(str, Enum): + """What a liveness heartbeat found when it tried to stamp a row. + + ``TERMINAL`` and ``MISSING`` both mean "not stamped", but they want + opposite handling: a terminal row may still be reclaimed by a finishing + stream, whereas a missing one can never be written again. + """ + + STAMPED = "stamped" + TERMINAL = "terminal" + MISSING = "missing" + + def _message_row_to_dict(row) -> dict: """Like ``row_to_dict`` but renames the DB column ``message_metadata`` back to the public API key ``metadata`` so callers keep the Mongo-era @@ -1059,34 +1072,67 @@ class ConversationsRepository: ) return result.rowcount > 0 - def heartbeat_message(self, message_id: str) -> bool: - """Stamp ``message_metadata.last_heartbeat_at`` with ``clock_timestamp()``. + def heartbeat_message_state(self, message_id: str) -> HeartbeatState: + """Stamp the liveness heartbeat and report what the row looks like. - The reconciler's staleness check uses ``GREATEST(timestamp, - last_heartbeat_at)``, so this call extends a long-running - stream's effective freshness without touching ``timestamp`` (the - creation time, used for history sort) or ``status`` (the WAL - marker). Skips terminal rows so a late heartbeat can't silently - retract a reconciler-set ``failed``. + One round trip, three outcomes. The distinction matters because the + two "not stamped" cases want opposite handling: + + - ``MISSING`` — the row was deleted mid-stream (``truncate_after`` on + a retry/edit). Nothing this stream produces can ever be read, so the + caller should stop working. + - ``TERMINAL`` — the row exists but is ``complete``/``failed``, + typically because the reconciler swept it. The stream must NOT be + cancelled: if it finishes, ``finalize_message`` is allowed to reclaim + a reconciler-failed row and land the real answer. + - ``STAMPED`` — normal in-flight row, freshness extended. + + The stamp deliberately avoids ``timestamp`` (creation time, used for + history sort) and ``status`` (the WAL marker), and skips terminal rows + so a late heartbeat can't retract a reconciler-set ``failed``. + + Args: + message_id: UUID of the message row. + + Returns: + HeartbeatState: What the row was at stamp time. """ if not looks_like_uuid(message_id): - return False - result = self._conn.execute( + return HeartbeatState.MISSING + # Both CTEs read the same MVCC snapshot, so ``present`` describes the + # row as it was when the UPDATE ran — no follow-up SELECT needed. + row = self._conn.execute( text( """ - UPDATE conversation_messages - SET message_metadata = jsonb_set( - COALESCE(message_metadata, '{}'::jsonb), - '{last_heartbeat_at}', - to_jsonb(clock_timestamp()) + WITH stamped AS ( + UPDATE conversation_messages + SET message_metadata = jsonb_set( + COALESCE(message_metadata, '{}'::jsonb), + '{last_heartbeat_at}', + to_jsonb(clock_timestamp()) + ) + WHERE id = CAST(:id AS uuid) + AND status NOT IN ('complete', 'failed') + RETURNING 1 AS ok ) - WHERE id = CAST(:id AS uuid) - AND status NOT IN ('complete', 'failed') + SELECT (SELECT ok FROM stamped) AS stamped, + EXISTS( + SELECT 1 FROM conversation_messages + WHERE id = CAST(:id AS uuid) + ) AS present """ ), {"id": message_id}, - ) - return result.rowcount > 0 + ).fetchone() + if row is not None and row[0]: + return HeartbeatState.STAMPED + if row is not None and row[1]: + return HeartbeatState.TERMINAL + return HeartbeatState.MISSING + + def heartbeat_message(self, message_id: str) -> bool: + """Stamp the heartbeat; True when an in-flight row was updated.""" + return self.heartbeat_message_state(message_id) is HeartbeatState.STAMPED def confirm_executed_tool_calls(self, message_id: str) -> int: """Flip ``tool_call_attempts.status='executed' → 'confirmed'`` for the message.""" diff --git a/tests/api/answer/test_stream_heartbeat_ticker.py b/tests/api/answer/test_stream_heartbeat_ticker.py index 908d0bf5..cf3b0446 100644 --- a/tests/api/answer/test_stream_heartbeat_ticker.py +++ b/tests/api/answer/test_stream_heartbeat_ticker.py @@ -1,9 +1,13 @@ -"""The stream liveness heartbeat must be time-based, not output-based. +"""The stream liveness heartbeat must be time-based, not output-based — +and it is also how a superseded stream learns to stop. -Regression cover for the production failure where a 20-minute agent tool -loop was force-failed by the reconciler at minute 6 — because the old -heartbeat only stamped when a chunk flowed, and a tool call emits nothing -while it runs. +Regression cover for two production failures: + +- a 20-minute agent tool loop force-failed by the reconciler at minute 6, + because the old heartbeat only stamped when a chunk flowed and a tool call + emits nothing while it runs; +- a stream whose row was deleted by the user's retry running four further + minutes and twelve further LLM rounds into a void. """ import time @@ -11,6 +15,8 @@ from unittest.mock import MagicMock, patch import pytest +from application.storage.db.repositories.conversations import HeartbeatState + def _reservation(message_id="11111111-1111-1111-1111-111111111111"): return {"conversation_id": "22222222-2222-2222-2222-222222222222", @@ -31,24 +37,40 @@ def _silent_then_answer(silence_seconds): return _gen +def _chatty(chunks=200, gap=0.02): + """A long multi-round stream that keeps emitting.""" + + def _gen(*args, **kwargs): + for i in range(chunks): + time.sleep(gap) + yield {"answer": f"chunk{i}"} + + return _gen + + +def _service(state=HeartbeatState.STAMPED): + service = MagicMock() + service.save_user_question.return_value = _reservation() + service.heartbeat_message.return_value = True + service.heartbeat_message_state.return_value = state + return service + + @pytest.mark.unit class TestHeartbeatTicker: - def _run(self, flask_app, gen, interval=0.05): + def _run(self, flask_app, gen, service=None, interval=0.05): from application.api.answer.routes import base as base_mod from application.api.answer.routes.base import BaseAnswerResource with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() - service.save_user_question.return_value = _reservation() - service.heartbeat_message.return_value = True - resource.conversation_service = service + resource.conversation_service = service or _service() agent = MagicMock() agent.gen.side_effect = gen with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", interval): - list( + out = list( resource.complete_stream( question="q", agent=agent, @@ -58,24 +80,23 @@ class TestHeartbeatTicker: should_persist=True, ) ) - return service + return resource.conversation_service, out def test_heartbeats_during_a_fully_silent_stream( self, mock_mongo_db, flask_app ): """No chunks flow for 0.4 s; the ticker must still stamp.""" - service = self._run(flask_app, _silent_then_answer(0.4)) + service, _ = self._run(flask_app, _silent_then_answer(0.4)) - # 1 seed at generation start + several ticker stamps. - assert service.heartbeat_message.call_count > 2 + assert service.heartbeat_message_state.call_count > 2 def test_ticker_stops_when_the_stream_ends(self, mock_mongo_db, flask_app): - service = self._run(flask_app, _silent_then_answer(0.1)) - settled = service.heartbeat_message.call_count + service, _ = self._run(flask_app, _silent_then_answer(0.1)) + settled = service.heartbeat_message_state.call_count time.sleep(0.3) - assert service.heartbeat_message.call_count == settled + assert service.heartbeat_message_state.call_count == settled def test_ticker_stops_on_client_abort(self, mock_mongo_db, flask_app): from application.api.answer.routes import base as base_mod @@ -83,18 +104,11 @@ class TestHeartbeatTicker: with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() - service.save_user_question.return_value = _reservation() - service.heartbeat_message.return_value = True + service = _service() resource.conversation_service = service - def _endless(*args, **kwargs): - while True: - time.sleep(0.02) - yield {"answer": "chunk"} - agent = MagicMock() - agent.gen.side_effect = _endless + agent.gen.side_effect = _chatty() with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): gen = resource.complete_stream( @@ -110,10 +124,10 @@ class TestHeartbeatTicker: time.sleep(0.15) gen.close() - settled = service.heartbeat_message.call_count + settled = service.heartbeat_message_state.call_count time.sleep(0.3) - assert service.heartbeat_message.call_count == settled + assert service.heartbeat_message_state.call_count == settled def test_no_ticker_without_a_reserved_row(self, mock_mongo_db, flask_app): """Headless/continuation rounds have no row to stamp.""" @@ -122,7 +136,7 @@ class TestHeartbeatTicker: with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() + service = _service() resource.conversation_service = service agent = MagicMock() @@ -140,26 +154,92 @@ class TestHeartbeatTicker: ) ) - service.heartbeat_message.assert_not_called() + service.heartbeat_message_state.assert_not_called() def test_ticker_stops_when_row_goes_terminal(self, mock_mongo_db, flask_app): - """A False return (row already terminal) is the natural stop signal.""" + """A terminal row stops the ticker but must NOT cancel the stream.""" + service = _service(HeartbeatState.TERMINAL) + service, out = self._run( + flask_app, _silent_then_answer(0.6), service=service, + ) + + # The ticker bailed on its first TERMINAL rather than stamping ~12 + # times over 0.6 s of silence... + assert service.heartbeat_message_state.call_count == 1 + # ...and the stream still finished and finalized, so a + # reconciler-swept row can still be reclaimed. + assert any('"type": "end"' in chunk for chunk in out) + service.finalize_message.assert_called_once() + + +@pytest.mark.unit +class TestSupersededStreamCancellation: + """A deleted row must stop the work, not just quiet the logs.""" + + def _run_with_missing_row(self, flask_app, gen, interval=0.05): from application.api.answer.routes import base as base_mod from application.api.answer.routes.base import BaseAnswerResource with flask_app.app_context(): resource = BaseAnswerResource() - service = MagicMock() - service.save_user_question.return_value = _reservation() - # The row is already terminal: every stamp reports "no row updated". - service.heartbeat_message.return_value = False + service = _service(HeartbeatState.MISSING) resource.conversation_service = service agent = MagicMock() - agent.gen.side_effect = _silent_then_answer(0.6) + agent.gen.side_effect = gen + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", interval): + out = list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + ) + return service, out + + def test_stream_stops_early_when_its_row_is_deleted( + self, mock_mongo_db, flask_app + ): + """The generator must not be drained to completion.""" + service, out = self._run_with_missing_row(flask_app, _chatty(chunks=400)) + + # 400 chunks × 0.02 s ≈ 8 s if drained; cancellation lands far sooner. + assert len(out) < 200, f"stream was not cancelled early ({len(out)})" + + def test_superseded_stream_does_not_persist(self, mock_mongo_db, flask_app): + """Nothing to write — the row is gone. No finalize, no save.""" + service, _ = self._run_with_missing_row(flask_app, _chatty(chunks=400)) + + service.finalize_message.assert_not_called() + service.save_conversation.assert_not_called() + + def test_superseded_stream_emits_no_error_to_the_client( + self, mock_mongo_db, flask_app + ): + """It is not a failure: the user replaced this turn deliberately.""" + _, out = self._run_with_missing_row(flask_app, _chatty(chunks=400)) + + assert not any('"type": "error"' in chunk for chunk in out) + + def test_live_row_is_never_cancelled(self, mock_mongo_db, flask_app): + """The common case must be untouched.""" + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = _service(HeartbeatState.STAMPED) + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = _chatty(chunks=30) with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): - list( + out = list( resource.complete_stream( question="q", agent=agent, @@ -170,7 +250,34 @@ class TestHeartbeatTicker: ) ) - # The seed and the status-flip stamp are unconditional; the ticker - # itself must bail on its first False rather than stamping ~12 - # times over the 0.6 s of silence. - assert service.heartbeat_message.call_count <= 3 + assert any('"type": "end"' in chunk for chunk in out) + service.finalize_message.assert_called_once() + + def test_db_error_does_not_cancel_the_stream(self, mock_mongo_db, flask_app): + """A transient blip must never be mistaken for a deleted row.""" + from application.api.answer.routes import base as base_mod + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + service = _service() + service.heartbeat_message_state.side_effect = RuntimeError("pool") + resource.conversation_service = service + + agent = MagicMock() + agent.gen.side_effect = _chatty(chunks=30) + + with patch.object(base_mod, "STREAM_HEARTBEAT_INTERVAL", 0.05): + out = list( + resource.complete_stream( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=True, + ) + ) + + assert any('"type": "end"' in chunk for chunk in out) + service.finalize_message.assert_called_once()