feat: auto cancel on retry or edit during reconciliation

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