fix(#2914): prevent state.db from replaying edited/undone messages to agent context

Author: AlexeyDsov

api/models.py

@@ -605,16 +605,20 @@ def path(self):
         return SESSION_DIR / f'{self.session_id}.json'
 
     def _maybe_clear_truncation_watermark(self) -> None:
-        watermark = _message_timestamp_as_float({"timestamp": self.truncation_watermark})
-        if watermark is None:
-            return
-        max_message_timestamp = None
-        for msg in self.messages or []:
-            timestamp = _message_timestamp_as_float(msg)
-            if timestamp is not None:
-                max_message_timestamp = timestamp if max_message_timestamp is None else max(max_message_timestamp, timestamp)
-        if max_message_timestamp is not None and max_message_timestamp > watermark:
-            self.truncation_watermark = None
+        # Do NOT auto-clear truncation_watermark here.  Once the user has
+        # truncated this session (Edit / Regenerate / undo / retry), the
+        # watermark must stay active to prevent state.db from replaying the
+        # rows the user deliberately removed.  Auto-clearing on every save()
+        # after the agent appends a newer message defeated the watermark
+        # entirely — the next reload would merge the full state.db tail
+        # back in (#2914).
+        #
+        # The watermark is cleared explicitly when the sidecar itself grows
+        # beyond it (new messages appended with timestamp > watermark), which
+        # means the user has continued the conversation past the truncation
+        # point.  That check is done in _append_message / save paths that
+        # know a genuine new turn was added, not in this blanket save() hook.
+        pass
 
     def save(self, touch_updated_at: bool = True, skip_index: bool = False) -> None:
         if not is_safe_session_id(self.session_id):
@@ -3798,6 +3802,20 @@ def merge_session_messages_append_only(
             and key not in seen_message_keys
         ):
             continue
+        # When a truncation watermark is active, state.db may contain original
+        # messages that were replaced by Edit (old content with old timestamp).
+        # The timestamp-based filter above catches messages AFTER the watermark,
+        # but messages BEFORE it (like the original pre-edit content) slip through.
+        # If a state.db message's content is not present in the sidecar and its
+        # timestamp is before the watermark, it's a replaced/stale row — skip it.
+        if (
+            watermark_timestamp is not None
+            and timestamp is not None
+            and timestamp < watermark_timestamp
+            and key not in seen_message_keys
+            and _session_message_content_key(msg) not in seen_content_keys
+        ):
+            continue
         if max_sidecar_timestamp is not None and timestamp is not None and timestamp <= max_sidecar_timestamp:
             if key in seen_message_keys:
                 continue

api/routes.py

@@ -5730,13 +5730,28 @@ def handle_post(handler, parsed) -> bool:
             return bad(handler, "Session not found", 404)
         keep = int(body["keep_count"])
         with _get_session_agent_lock(body["session_id"]):
+            old_msg_count = len(s.messages or [])
+            old_ctx_count = len(getattr(s, 'context_messages', None) or [])
             s.messages = s.messages[:keep]
+            # Truncate context_messages in sync with messages so the agent's
+            # model-facing context doesn't retain rows the user removed via
+            # Edit / Regenerate.  Without this, context_messages still contains
+            # the full pre-truncation history and the agent sees "deleted"
+            # turns on the next turn (#2914).
+            if isinstance(getattr(s, 'context_messages', None), list):
+                s.context_messages = s.context_messages[:keep]
             try:
                 from api.session_ops import _truncation_watermark_for
                 s.truncation_watermark = _truncation_watermark_for(s.messages)
             except Exception:
                 s.truncation_watermark = 0.0
             s.save()
+            logger.info(
+                "truncate %s: messages %d→%d, context_messages %d→%d, watermark=%.2f",
+                body["session_id"], old_msg_count, len(s.messages or []),
+                old_ctx_count, len(getattr(s, 'context_messages', None) or []),
+                s.truncation_watermark or 0,
+            )
         return j(
             handler, {"ok": True, "session": s.compact() | {"messages": s.messages}}
         )

api/streaming.py

@@ -45,6 +45,7 @@
 from api.usage import prompt_cache_hit_percent
 from api.models import (
     _is_empty_partial_activity_message,
+    _message_timestamp_as_float,
     get_state_db_session_messages,
     reconciled_state_db_messages_for_session,
 )
@@ -2618,6 +2619,32 @@ def _restore_display_reasoning_metadata(previous_messages, updated_messages):
     return updated_messages
 
 
+def _clamp_context_to_watermark(session, messages: list) -> list:
+    """Filter context_messages to the truncation watermark boundary (#2914).
+
+    When a user edits, regenerates, or undoes messages, the agent's result
+    may contain the full state.db history including turns the user deliberately
+    removed.  This helper drops messages whose timestamp exceeds the active
+    watermark so the next turn doesn't feed the agent "deleted" rows again.
+    """
+    _tw = getattr(session, 'truncation_watermark', None)
+    if _tw is None:
+        return messages
+    _tw_ts = _message_timestamp_as_float({'timestamp': _tw})
+    if _tw_ts is None:
+        return messages
+    _clamped = [
+        m for m in messages
+        if (m_ts := _message_timestamp_as_float(m)) is None or m_ts <= _tw_ts
+    ]
+    if len(_clamped) != len(messages):
+        logger.info(
+            "clamping context_messages: %d → %d (watermark=%.2f, session=%s)",
+            len(messages), len(_clamped), _tw_ts, session.session_id,
+        )
+    return _clamped
+
+
 def _session_context_messages(session):
     """Return model-facing history without assuming it matches the UI transcript."""
     context_messages = getattr(session, 'context_messages', None)
@@ -5148,6 +5175,7 @@ def _periodic_checkpoint():
                     _previous_context_messages,
                     _next_context_messages,
                 )
+                _next_context_messages = _clamp_context_to_watermark(s, _next_context_messages)
                 s.context_messages = _deduplicate_context_messages(_next_context_messages)
                 s.messages = _merge_display_messages_after_agent_result(
                     _previous_messages,
@@ -5295,6 +5323,7 @@ def _periodic_checkpoint():
                                     _previous_context_messages,
                                     _next_context_messages,
                                 )
+                                _next_context_messages = _clamp_context_to_watermark(s, _next_context_messages)
                                 s.context_messages = _deduplicate_context_messages(_next_context_messages)
                                 s.messages = _merge_display_messages_after_agent_result(
                                     _previous_messages,
@@ -6172,6 +6201,7 @@ def _periodic_checkpoint():
                                     _previous_context_messages,
                                     _next_context_messages,
                                 )
+                                _next_context_messages = _clamp_context_to_watermark(s, _next_context_messages)
                                 s.context_messages = _deduplicate_context_messages(_next_context_messages)
                                 s.messages = _merge_display_messages_after_agent_result(
                                     _previous_messages,