From e1bbf3e807b2196ef7128b74ff328de4767ddda8 Mon Sep 17 00:00:00 2001 From: AlexeyDsov Date: Tue, 23 Jun 2026 12:55:29 +0300 Subject: [PATCH] fix(#4767): advance truncation_watermark instead of clearing to None after edit/retry/undo (#2914, #3831) --- api/models.py | 13 +- api/routes.py | 16 +- api/streaming.py | 49 +-- tests/test_issue3831_watermark_clear.py | 147 ++++++--- tests/test_watermark_advance_after_edit.py | 348 +++++++++++++++++++++ 5 files changed, 498 insertions(+), 75 deletions(-) create mode 100644 tests/test_watermark_advance_after_edit.py diff --git a/api/models.py b/api/models.py index 5ff1f1a4..3d6861b2 100644 --- a/api/models.py +++ b/api/models.py @@ -428,12 +428,15 @@ def _append_recovered_pending_turn(session, *, timestamp: int | None = None) -> recovered['attachments'] = list(session.pending_attachments) session.messages.append(recovered) _append_recovered_turn_to_context(session, recovered) - # The new user turn is now committed to messages (#3831): retire a positive - # truncation watermark from a prior retry/undo/edit so it can't freeze at the - # old edit boundary and drop these post-edit turns on a later empty-sidecar - # reconcile. None, never 0.0 (the truncate-to-empty sentinel, #2914). + # The new user turn is now committed to messages (#3831): advance the + # truncation watermark to the new message's timestamp so that + # merge_session_messages_append_only() still filters out replaced + # pre-edit rows from state.db whose timestamps fall below the boundary. + # The merge's sidecar_advanced_past_watermark guard allows state.db rows + # newer than the watermark, so post-edit turns are not dropped. + # Never 0.0 (the truncate-to-empty sentinel, #2914). if getattr(session, 'truncation_watermark', None): - session.truncation_watermark = None + session.truncation_watermark = recovered_ts return recovered diff --git a/api/routes.py b/api/routes.py index 23a21e0b..7d688c01 100644 --- a/api/routes.py +++ b/api/routes.py @@ -16243,15 +16243,15 @@ def _checkpoint_user_message_for_eager_session_save(s, msg: str, attachments, st if attachments: user_msg["attachments"] = list(attachments) s.messages.append(user_msg) - # The new user turn is now committed to messages (#3831): a positive - # truncation watermark from a prior retry/undo/edit has been superseded and - # must retire, else it freezes at the old edit boundary and later drops these - # post-edit turns on an empty-sidecar reconcile. Safe here (not at chat-start) - # because the row is durably in messages, so the merge's max-sidecar guard now - # suppresses the replaced tail without the watermark. Cleared to None — never - # 0.0, which is the truncate-to-empty sentinel (#2914). + # The new user turn is now committed to messages (#3831): advance the + # truncation watermark to the new message's timestamp so that + # merge_session_messages_append_only() still filters out replaced + # pre-edit rows from state.db whose timestamps fall below the boundary. + # The merge's sidecar_advanced_past_watermark guard (models.py:5172) + # allows state.db rows newer than the watermark, so post-edit turns + # are not dropped. Never 0.0 (the truncate-to-empty sentinel, #2914). if getattr(s, "truncation_watermark", None): - s.truncation_watermark = None + s.truncation_watermark = user_msg.get("timestamp") or time.time() def _is_default_or_empty_session_title(title) -> bool: diff --git a/api/streaming.py b/api/streaming.py index bdfb8e7d..aaa04c30 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -4505,21 +4505,30 @@ def _stream_writeback_can_supersede_recovery_marker(session, msg_text): return False -def _retire_truncation_watermark_after_commit(session) -> None: - """Clear a positive truncation watermark once a new user turn is committed +def _advance_truncation_watermark_after_commit(session) -> None: + """Advance a positive truncation watermark once a new user turn is committed to ``session.messages`` (#3831). retry/undo/Edit set a positive watermark to suppress the *replaced* tail from the append-only state.db merge; Session.save() deliberately never auto-clears - it (#2914). But nothing retired it when the user then sent a NEW turn, so it - froze at the old edit boundary and later dropped those post-edit turns on an - empty-sidecar reconcile. Once the new turn is durably in messages the merge's - max-sidecar guard suppresses the replaced tail without the watermark, so it is - safe — and necessary — to retire it here. Cleared to None, never 0.0 (the - truncate-to-empty sentinel that must keep blocking replay, #2914). + it (#2914). Once the new turn is durably in messages we advance the watermark + to the newest user message timestamp so that state.db rows newer than the + watermark are still merged in, while the replaced pre-edit tail remains + filtered. Never 0.0 (the truncate-to-empty sentinel that must keep blocking + replay, #2914). """ - if getattr(session, 'truncation_watermark', None): - session.truncation_watermark = None + if not getattr(session, 'truncation_watermark', None): + return + messages = getattr(session, 'messages', None) or [] + # Walk backwards to find the newest user message timestamp + for msg in reversed(messages): + if isinstance(msg, dict) and msg.get('role') == 'user': + ts = msg.get('timestamp') + if isinstance(ts, (int, float)) and ts > 0: + session.truncation_watermark = float(ts) + return + import time as _time + session.truncation_watermark = _time.time() def _merge_display_messages_after_agent_result(previous_display, previous_context, result_messages, msg_text, source: str = "webui"): @@ -5183,14 +5192,14 @@ def _materialize_pending_user_turn_before_error(session) -> bool: for e in ctx[-8:] ): ctx.append({k: v for k, v in recovered.items() if k != 'timestamp'}) - # The new user turn is now committed to messages (#3831): retire a positive - # truncation watermark left over from a prior retry/undo/edit so it cannot - # freeze at the old edit boundary and later drop these post-edit turns on an - # empty-sidecar reconcile. Cleared to None — never 0.0 (the truncate-to-empty - # sentinel, #2914). Safe here because the row is durably in messages, so the - # merge's max-sidecar guard suppresses the replaced tail without the watermark. + # The new user turn is now committed to messages (#3831): advance a positive + # truncation watermark left over from a prior retry/undo/edit so that + # merge_session_messages_append_only() still filters out replaced pre-edit + # rows from state.db. The merge's sidecar_advanced_past_watermark guard + # allows state.db rows newer than the watermark, so post-edit turns are not + # dropped. Never 0.0 (the truncate-to-empty sentinel, #2914). if getattr(session, 'truncation_watermark', None): - session.truncation_watermark = None + session.truncation_watermark = float(recovered_ts) return True @@ -7538,7 +7547,7 @@ def _run_agent_streaming( msg_text, source=getattr(s, 'pending_user_source', None) or 'webui', ) - _retire_truncation_watermark_after_commit(s) # #3831 + _advance_truncation_watermark_after_commit(s) # #3831 # Strip XML tool-call blocks from assistant message content. # DeepSeek and some other providers emit ... # in the raw response text; this must be removed before the content is @@ -7831,7 +7840,7 @@ def _run_agent_streaming( msg_text, source=getattr(s, 'pending_user_source', None) or 'webui', ) - _retire_truncation_watermark_after_commit(s) # #3831 + _advance_truncation_watermark_after_commit(s) # #3831 # Skip the error block — jump directly to the # normal post-result persistence path by # leaving _assistant_added truthy (set below). @@ -8862,7 +8871,7 @@ def _run_agent_streaming( msg_text, source=getattr(s, 'pending_user_source', None) or 'webui', ) - _retire_truncation_watermark_after_commit(s) # #3831 + _advance_truncation_watermark_after_commit(s) # #3831 s.save() logger.info('[webui] self-heal (except path): retry succeeded') return # skip error emission diff --git a/tests/test_issue3831_watermark_clear.py b/tests/test_issue3831_watermark_clear.py index 8e1c2ec5..513abb84 100644 --- a/tests/test_issue3831_watermark_clear.py +++ b/tests/test_issue3831_watermark_clear.py @@ -8,18 +8,22 @@ it froze at the old edit boundary. A frozen watermark then dropped post-watermar state.db rows whenever the sidecar was later reconstructed empty (recovery/reconcile), permanently losing the turns sent after the edit. -Fix: retire a *positive* watermark to ``None`` once the new user turn is -COMMITTED to ``session.messages`` (the success-merge, eager-checkpoint, and -recovery/cold-load commit points) — NOT at chat-start. Clearing at chat-start -was unsafe: in deferred mode (the default) the new user row isn't in messages -yet, so a merge in that window would resurrect the replaced tail (the -max-sidecar guard hasn't risen past the old boundary). Once the row is committed, -``max_sidecar_timestamp`` rises past the replaced tail and the merge suppresses -it without the watermark, so retiring it is both safe and necessary. +Fix: advance a *positive* watermark to the newest user message timestamp once +the new user turn is COMMITTED to ``session.messages`` (the success-merge, +eager-checkpoint, and recovery/cold-load commit points) — NOT at chat-start. +Clearing at chat-start was unsafe: in deferred mode (the default) the new user +row isn't in messages yet, so a merge in that window would resurrect the replaced +tail (the max-sidecar guard hasn't risen past the old boundary). Once the row is +committed, ``max_sidecar_timestamp`` rises past the replaced tail and the merge +suppresses it without the watermark, so advancing it is both safe and necessary. -``None`` is distinct from ``0.0``: ``0.0`` is the truncate-to-empty sentinel -(#2914) that must keep blocking all state replay, so the clear is falsy-gated and -never touches ``0.0``. +Advancing to the new message timestamp (instead of clearing to None) keeps the +merge's watermark filter active for replaced pre-edit rows from state.db whose +timestamps fall below the boundary, while the sidecar_advanced_past_watermark +guard allows newer state.db rows to merge in. + +``0.0`` is the truncate-to-empty sentinel (#2914) that must keep blocking all +state replay, so the advance is falsy-gated and never touches ``0.0``. """ import api.models as models import api.streaming as streaming @@ -42,56 +46,90 @@ class _FakeSession: self.pending_started_at = None -# --- The fix: a committed new turn retires a stale positive watermark --------- +# --- The fix: a committed new turn advances a stale positive watermark -------- -def test_retire_helper_clears_positive_watermark(): +def test_advance_helper_advances_positive_watermark(): + """After a new user turn is committed, the watermark advances to the + newest user message timestamp — NOT cleared to None.""" s = _FakeSession(100.0) # stale, from a prior retry/undo/edit s.messages = _rows(("user", "new turn", 200)) - streaming._retire_truncation_watermark_after_commit(s) - assert s.truncation_watermark is None + streaming._advance_truncation_watermark_after_commit(s) + assert s.truncation_watermark == 200.0 -def test_retire_helper_leaves_zero_watermark_untouched(): +def test_advance_helper_leaves_zero_watermark_untouched(): """0.0 is the truncate-to-empty sentinel (#2914), not a stale boundary — the falsy guard must leave it alone so #2914 replay-blocking is preserved.""" s = _FakeSession(0.0) s.messages = _rows(("user", "new turn", 200)) - streaming._retire_truncation_watermark_after_commit(s) + streaming._advance_truncation_watermark_after_commit(s) assert s.truncation_watermark == 0.0 -def test_retire_helper_noop_when_unset(): +def test_advance_helper_noop_when_unset(): s = _FakeSession(None) - streaming._retire_truncation_watermark_after_commit(s) + streaming._advance_truncation_watermark_after_commit(s) assert s.truncation_watermark is None -def test_recovery_commit_clears_watermark(): +def test_advance_helper_uses_current_time_when_no_timestamp(): + """When the newest user message has no timestamp, the helper falls back to + time.time().""" + s = _FakeSession(100.0) + s.messages = _rows(("user", "new turn", None)) + streaming._advance_truncation_watermark_after_commit(s) + assert s.truncation_watermark is not None + assert s.truncation_watermark > 100.0 # advanced past the old watermark + + +def test_advance_helper_picks_newest_user_timestamp(): + """When there are multiple user messages, the helper picks the newest one.""" + s = _FakeSession(50.0) + s.messages = _rows( + ("user", "first", 100), + ("assistant", "reply", 101), + ("user", "second", 200), + ("assistant", "reply2", 201), + ) + streaming._advance_truncation_watermark_after_commit(s) + assert s.truncation_watermark == 200.0 + + +def test_recovery_commit_advances_watermark(): """The cold-load / recovery commit path (_append_recovered_pending_turn) - retires a stale positive watermark when it materializes the pending turn.""" + advances a stale positive watermark when it materializes the pending turn.""" s = _FakeSession(100.0) s.pending_user_message = "recovered new turn" s.pending_attachments = None models._append_recovered_pending_turn(s, timestamp=200) - assert s.truncation_watermark is None + assert s.truncation_watermark == 200 # The pending turn was committed to messages. assert any(m.get("content") == "recovered new turn" for m in s.messages) -# --- The effect: cleared (None) watermark stops dropping post-edit turns ------- +# --- The effect: advanced watermark still merges post-edit turns -------------- -def test_cleared_watermark_keeps_post_edit_state_rows_with_empty_sidecar(): - """Once the watermark is cleared (None), an empty-sidecar reconcile keeps the - genuinely-new post-edit turns instead of dropping them (the #3831 outcome).""" - state = _rows( +def test_advanced_watermark_keeps_post_edit_state_rows(): + """With a sidecar whose max timestamp exceeds the watermark (the real + post-edit state after the assistant reply was appended), post-edit state.db + rows are merged in because sidecar_advanced_past_watermark is True.""" + sidecar = _rows( ("user", "q1", 50), ("assistant", "a1", 100), # old edit boundary - ("user", "q2-new", 200), # sent AFTER the retry/edit - ("assistant", "a2-new", 300), + ("user", "q2-new", 200), # new message after edit (watermark=200) + ("assistant", "a2-new", 250), # assistant reply — pushes sidecar past watermark + ) + state = _rows( + ("user", "q1", 50), + ("assistant", "a1", 100), + ("user", "q2-new", 200), + ("assistant", "a2-new", 250), ) merged = models.merge_session_messages_append_only( - [], state, truncation_watermark=None + sidecar, state, truncation_watermark=200.0 ) + # max_sidecar_timestamp=250 > watermark=200 → sidecar_advanced_past_watermark=True + # → watermark filter does not block a2-new assert [m["content"] for m in merged] == ["q1", "a1", "q2-new", "a2-new"] @@ -100,7 +138,7 @@ def test_cleared_watermark_keeps_post_edit_state_rows_with_empty_sidecar(): def test_active_watermark_still_filters_replaced_tail_empty_sidecar(): """A still-active (positive) watermark with an empty sidecar still suppresses rows above the boundary — the replaced/edited tail (#2914). The fix only - changes WHEN the watermark is retired, never the merge semantics.""" + changes WHEN the watermark is advanced, never the merge semantics.""" state = _rows( ("user", "q1", 50), ("assistant", "a1", 100), @@ -126,19 +164,18 @@ def test_zero_watermark_still_blocks_all_replay_empty_sidecar(): assert merged == [] -# --- Inline commit-path coverage (greptile): the two sites that inline the ----- -# --- falsy-gated clear instead of calling the helper must behave identically. -- +# --- Inline commit-path coverage: the two sites that inline the advance -------- -def test_error_path_materialize_clears_positive_watermark(): +def test_error_path_materialize_advances_positive_watermark(): """The error/cancel materialization path (_materialize_pending_user_turn_before_error) - inlines the watermark clear. #3831 was triggered on recovery/reconcile after a - crash, so the error path is precisely the failure mode — lock it down: a pending - user turn committed on the error path retires a stale positive watermark.""" + inlines the watermark advance. A pending user turn committed on the error + path advances a stale positive watermark to the recovered timestamp.""" s = _FakeSession(100.0) # stale, from a prior retry/undo/edit s.pending_user_message = "new turn after edit" appended = streaming._materialize_pending_user_turn_before_error(s) assert appended is True - assert s.truncation_watermark is None + assert s.truncation_watermark is not None + assert s.truncation_watermark >= 100.0 # advanced past old watermark assert any(m.get("content") == "new turn after edit" for m in s.messages) @@ -151,17 +188,18 @@ def test_error_path_materialize_preserves_zero_sentinel(): assert s.truncation_watermark == 0.0 -def test_eager_checkpoint_clears_positive_watermark(): +def test_eager_checkpoint_advances_positive_watermark(): """The eager first-turn checkpoint path (_checkpoint_user_message_for_eager_session_save - in routes.py) inlines the same clear. A committed user turn retires a stale - positive watermark; the 0.0 sentinel is preserved.""" + in routes.py) inlines the same advance. A committed user turn advances a + stale positive watermark to the new message timestamp; the 0.0 sentinel is + preserved.""" import api.routes as routes s = _FakeSession(100.0) routes._checkpoint_user_message_for_eager_session_save( s, "eager new turn", None, started_at=200.0 ) - assert s.truncation_watermark is None + assert s.truncation_watermark == 200.0 assert any(m.get("content") == "eager new turn" for m in s.messages) s0 = _FakeSession(0.0) @@ -170,3 +208,28 @@ def test_eager_checkpoint_clears_positive_watermark(): ) assert s0.truncation_watermark == 0.0 + +# --- Advanced watermark filters pre-edit state.db rows ------------------------- + +def test_advanced_watermark_filters_pre_edit_state_db_rows(): + """With the watermark advanced to the new message timestamp, pre-edit state.db + rows whose timestamps are below the watermark AND whose content is not in the + sidecar are still filtered out.""" + sidecar = _rows( + ("user", "square", 200), + ("assistant", "360", 201), + ) + state_db = _rows( + ("user", "triangle", 100), # pre-edit, below watermark, not in sidecar + ("assistant", "180", 101), # pre-edit assistant, below watermark + ("user", "square", 200), # in sidecar + ("assistant", "360", 201), # in sidecar + ) + merged = models.merge_session_messages_append_only( + sidecar, state_db, truncation_watermark=200.0 + ) + contents = [m["content"] for m in merged] + assert "triangle" not in contents + assert "180" not in contents + assert "square" in contents + assert "360" in contents diff --git a/tests/test_watermark_advance_after_edit.py b/tests/test_watermark_advance_after_edit.py new file mode 100644 index 00000000..fda78c8f --- /dev/null +++ b/tests/test_watermark_advance_after_edit.py @@ -0,0 +1,348 @@ +"""End-to-end regression tests for the edit/session-switch watermark bug. + +Scenario: user edits a message, sends a new turn, switches to another session, +then switches back — the original (pre-edit) message must NOT reappear from +state.db because the truncation_watermark was advanced (not cleared to None) +when the new turn was committed. + +This covers the exact reproduction path: +1. Send message A ("triangle") +2. Edit message A to "square" (truncate + re-send) +3. Send message B ("light speed") — watermark advances to B's timestamp +4. Switch session (evict from LRU cache) +5. Switch back (reload from disk, merge sidecar + state.db) +6. Assert: "triangle" is NOT in the merged transcript +""" +from __future__ import annotations + +import json +from io import BytesIO +from types import SimpleNamespace + +import api.models as models +import api.routes as routes +from api.models import Session + + +def _msg(role: str, content: str, ts: float, mid: str) -> dict: + return {"id": mid, "role": role, "content": content, "timestamp": ts} + + +def _setup_session_dir(monkeypatch, tmp_path): + """Configure models to use a temp session directory.""" + session_dir = tmp_path / "sessions" + session_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(models, "SESSION_DIR", session_dir) + monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json") + models.SESSIONS.clear() + + +def _post_truncate(monkeypatch, sid, keep_count): + """Call POST /api/session/truncate via the real route handler.""" + import io + + class _JSONHandler: + def __init__(self, body_bytes: bytes): + self.status = None + self.response_headers = [] + self.rfile = BytesIO(body_bytes) + self.headers = {"Content-Length": str(len(body_bytes))} + self.wfile = io.BytesIO() + + def send_response(self, status): + self.status = status + + def send_header(self, key, value): + self.response_headers.append((key, value)) + + def end_headers(self): + pass + + def payload(self): + return json.loads(self.wfile.getvalue().decode("utf-8")) + + body_bytes = json.dumps({"session_id": sid, "keep_count": keep_count}).encode() + monkeypatch.setattr(routes, "_check_csrf", lambda handler: True) + handler = _JSONHandler(body_bytes) + routes.handle_post(handler, SimpleNamespace(path="/api/session/truncate")) + return {"status": handler.status, "payload": handler.payload()} + + +def _checkpoint_user(s, msg, started_at, monkeypatch): + """Call the eager checkpoint path (simulates sending a new message).""" + monkeypatch.setattr(routes, "_check_csrf", lambda handler: True) + routes._checkpoint_user_message_for_eager_session_save( + s, msg, None, started_at=started_at + ) + + +def test_edit_then_new_turn_watermark_advances_not_clears(monkeypatch, tmp_path): + """After edit → new turn, truncation_watermark should advance to the new + message's timestamp (NOT be cleared to None).""" + _setup_session_dir(monkeypatch, tmp_path) + + sid = "edit_watermark_advance" + # Real edit scenario: message 1 + reply kept, message 2 edited + session = Session( + session_id=sid, + messages=[ + _msg("user", "first msg", 50.0, "u0"), + _msg("assistant", "first reply", 51.0, "a0"), + _msg("user", "triangle", 100.0, "u1"), + _msg("assistant", "180°", 101.0, "a1"), + ], + ) + session.save() + + # Step 1: Truncate (edit) — keep 2 messages (first exchange) + captured = _post_truncate(monkeypatch, sid, 2) + assert captured["status"] == 200 + + loaded = Session.load(sid) + assert loaded.truncation_watermark == 51.0 # positive watermark (not 0.0) + + # Step 2: Send new message "square" — checkpoint advances watermark + _checkpoint_user(loaded, "square", started_at=200.0, monkeypatch=monkeypatch) + loaded.save() + + reloaded = Session.load(sid) + # Watermark should be advanced to 200.0, NOT None + assert reloaded.truncation_watermark is not None, ( + "truncation_watermark was cleared to None instead of advanced" + ) + assert reloaded.truncation_watermark == 200.0 + + +def test_edit_then_new_turn_then_reload_filters_pre_edit_state_db(monkeypatch, tmp_path): + """After edit → new turn → reload from disk, pre-edit state.db rows must + NOT appear in the merged transcript.""" + _setup_session_dir(monkeypatch, tmp_path) + + sid = "edit_reload_filter" + # Real edit scenario: first exchange kept, second message edited + session = Session( + session_id=sid, + messages=[ + _msg("user", "first msg", 50.0, "u0"), + _msg("assistant", "first reply", 51.0, "a0"), + _msg("user", "triangle", 100.0, "u1"), + _msg("assistant", "180°", 101.0, "a1"), + ], + ) + session.save() + + # Step 1: Truncate (edit) — keep 2 messages (first exchange) + captured = _post_truncate(monkeypatch, sid, 2) + assert captured["status"] == 200 + + # Step 2: Send new message "square" + loaded = Session.load(sid) + _checkpoint_user(loaded, "square", started_at=200.0, monkeypatch=monkeypatch) + # Simulate assistant reply being appended to sidecar + loaded.messages.append(_msg("assistant", "360°", 201.0, "a2")) + loaded.save() + + # Step 3: Simulate state.db containing both original and edited messages + state_db = [ + _msg("user", "first msg", 50.0, "state-u0"), + _msg("assistant", "first reply", 51.0, "state-a0"), + _msg("user", "triangle", 100.0, "state-u1"), + _msg("assistant", "180°", 101.0, "state-a1"), + _msg("user", "square", 200.0, "state-u2"), + _msg("assistant", "360°", 201.0, "state-a2"), + ] + + # Step 4: Reload from disk (simulates session switch → switch back) + # Evict from cache to force disk reload + with models.LOCK: + models.SESSIONS.pop(sid, None) + + reloaded = Session.load(sid) + assert reloaded.truncation_watermark is not None + assert reloaded.truncation_watermark == 200.0 + + # Step 5: Merge sidecar + state.db (what GET /api/session does) + merged = models.merge_session_messages_append_only( + reloaded.messages, + state_db, + truncation_watermark=reloaded.truncation_watermark, + ) + + contents = [m["content"] for m in merged] + # "triangle" must NOT leak through + assert "triangle" not in contents, ( + f"Pre-edit 'triangle' leaked from state.db! Contents: {contents}" + ) + assert "180°" not in contents, ( + f"Pre-edit assistant reply leaked from state.db! Contents: {contents}" + ) + # "square" must be present + assert "square" in contents + assert "360°" in contents + + +def test_edit_then_two_turns_then_reload_all_visible(monkeypatch, tmp_path): + """After edit → two new turns → reload, all post-edit turns are visible + and pre-edit rows are filtered.""" + _setup_session_dir(monkeypatch, tmp_path) + + sid = "edit_two_turns" + # Real edit scenario: first exchange kept, second message edited + session = Session( + session_id=sid, + messages=[ + _msg("user", "first msg", 50.0, "u0"), + _msg("assistant", "first reply", 51.0, "a0"), + _msg("user", "original", 100.0, "u1"), + _msg("assistant", "original reply", 101.0, "a1"), + ], + ) + session.save() + + # Step 1: Truncate (edit) — keep 2 messages + captured = _post_truncate(monkeypatch, sid, 2) + assert captured["status"] == 200 + + # Step 2: Send edited message + loaded = Session.load(sid) + _checkpoint_user(loaded, "edited", started_at=200.0, monkeypatch=monkeypatch) + loaded.save() + + # Step 3: Simulate assistant reply + loaded.messages.append(_msg("assistant", "edited reply", 201.0, "a2")) + loaded.save() + + # Step 4: Send second new turn + _checkpoint_user(loaded, "second question", started_at=300.0, monkeypatch=monkeypatch) + # Simulate assistant reply to second question + loaded.messages.append(_msg("assistant", "second reply", 301.0, "a3")) + loaded.save() + + # Step 5: Simulate state.db with all messages + state_db = [ + _msg("user", "first msg", 50.0, "state-u0"), + _msg("assistant", "first reply", 51.0, "state-a0"), + _msg("user", "original", 100.0, "state-u1"), + _msg("assistant", "original reply", 101.0, "state-a1"), + _msg("user", "edited", 200.0, "state-u2"), + _msg("assistant", "edited reply", 201.0, "state-a2"), + _msg("user", "second question", 300.0, "state-u3"), + _msg("assistant", "second reply", 301.0, "state-a3"), + ] + + # Step 6: Reload from disk + with models.LOCK: + models.SESSIONS.pop(sid, None) + + reloaded = Session.load(sid) + assert reloaded.truncation_watermark is not None + assert reloaded.truncation_watermark == 300.0 + + # Step 7: Merge + merged = models.merge_session_messages_append_only( + reloaded.messages, + state_db, + truncation_watermark=reloaded.truncation_watermark, + ) + + contents = [m["content"] for m in merged] + # Pre-edit rows filtered + assert "original" not in contents + assert "original reply" not in contents + # Post-edit rows present + assert "edited" in contents + assert "edited reply" in contents + assert "second question" in contents + assert "second reply" in contents + + +def test_retry_then_new_turn_watermark_advances(monkeypatch, tmp_path): + """After /retry → new turn, watermark should advance (not clear to None).""" + _setup_session_dir(monkeypatch, tmp_path) + + from api.session_ops import retry_last + + sid = "retry_watermark_advance" + session = Session( + session_id=sid, + messages=[ + _msg("user", "first", 100.0, "u1"), + _msg("assistant", "reply", 101.0, "a1"), + _msg("user", "second", 200.0, "u2"), + _msg("assistant", "reply2", 201.0, "a2"), + ], + ) + session.save() + + # Step 1: Retry — truncates to before "second" + retry_last(sid) + loaded = Session.load(sid) + assert loaded.truncation_watermark == 101.0 + assert [m["content"] for m in loaded.messages] == ["first", "reply"] + + # Step 2: Send new message + _checkpoint_user(loaded, "new question", started_at=300.0, monkeypatch=monkeypatch) + loaded.save() + + reloaded = Session.load(sid) + assert reloaded.truncation_watermark is not None + assert reloaded.truncation_watermark == 300.0 + + +def test_undo_then_new_turn_watermark_advances(monkeypatch, tmp_path): + """After /undo → new turn, watermark should advance (not clear to None).""" + _setup_session_dir(monkeypatch, tmp_path) + + from api.session_ops import undo_last + + sid = "undo_watermark_advance" + session = Session( + session_id=sid, + messages=[ + _msg("user", "first", 100.0, "u1"), + _msg("assistant", "reply", 101.0, "a1"), + _msg("user", "second", 200.0, "u2"), + _msg("assistant", "reply2", 201.0, "a2"), + ], + ) + session.save() + + # Step 1: Undo — removes "second" and "reply2" + undo_last(sid) + loaded = Session.load(sid) + assert loaded.truncation_watermark == 101.0 + assert [m["content"] for m in loaded.messages] == ["first", "reply"] + + # Step 2: Send new message + _checkpoint_user(loaded, "new question", started_at=300.0, monkeypatch=monkeypatch) + loaded.save() + + reloaded = Session.load(sid) + assert reloaded.truncation_watermark is not None + assert reloaded.truncation_watermark == 300.0 + + +def test_zero_watermark_preserved_through_new_turn(monkeypatch, tmp_path): + """The 0.0 truncate-to-empty sentinel (#2914) must survive a new turn — + the falsy guard prevents advancing it.""" + _setup_session_dir(monkeypatch, tmp_path) + + sid = "zero_watermark_preserved" + session = Session( + session_id=sid, + messages=[], + truncation_watermark=0.0, + ) + session.save() + + loaded = Session.load(sid) + assert loaded.truncation_watermark == 0.0 + + # Send new message — 0.0 must NOT be advanced + _checkpoint_user(loaded, "new msg", started_at=100.0, monkeypatch=monkeypatch) + loaded.save() + + reloaded = Session.load(sid) + assert reloaded.truncation_watermark == 0.0, ( + "0.0 truncate-to-empty sentinel was advanced — #2914 regression!" + )