feat: auto cancel on retry or edit during reconciliation

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:

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,

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."""

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()