From 037253e4fc0ec8895f925d835bb00f6c0ef1554c Mon Sep 17 00:00:00 2001 From: AlexeyDsov Date: Wed, 24 Jun 2026 16:08:28 +0300 Subject: [PATCH] fix(#4767): review fixes - more edge cases --- api/models.py | 125 +++++++---- api/routes.py | 6 + api/session_ops.py | 5 + api/webui_session_db.py | 1 + tests/test_core_data_loss_cases.py | 291 +++++++++++++++++++++++++ tests/test_session_duplicate_fields.py | 20 ++ 6 files changed, 407 insertions(+), 41 deletions(-) create mode 100644 tests/test_core_data_loss_cases.py diff --git a/api/models.py b/api/models.py index 9281f29c..6b2424a5 100644 --- a/api/models.py +++ b/api/models.py @@ -646,6 +646,7 @@ class Session: context_length=None, threshold_tokens=None, last_prompt_tokens=None, truncation_watermark=None, + truncation_boundary=None, gateway_routing=None, gateway_routing_history=None, llm_title_generated: bool=False, manual_title: bool=False, @@ -696,6 +697,7 @@ class Session: self.threshold_tokens = threshold_tokens self.last_prompt_tokens = last_prompt_tokens self.truncation_watermark = truncation_watermark + self.truncation_boundary = truncation_boundary self.gateway_routing = gateway_routing if isinstance(gateway_routing, dict) else None self.gateway_routing_history = gateway_routing_history if isinstance(gateway_routing_history, list) else [] self.llm_title_generated = bool(llm_title_generated) @@ -765,6 +767,7 @@ class Session: 'compression_anchor_details', 'context_engine_state', 'context_length', 'threshold_tokens', 'last_prompt_tokens', 'truncation_watermark', + 'truncation_boundary', 'gateway_routing', 'gateway_routing_history', 'llm_title_generated', 'manual_title', 'parent_session_id', 'worktree_path', 'worktree_branch', 'worktree_repo_root', 'worktree_created_at', @@ -5109,8 +5112,16 @@ def merge_session_messages_append_only( state_messages: list, *, truncation_watermark=None, + truncation_boundary=None, ) -> list: - """Merge sidecar/context and state.db messages without deleting local rows.""" + """Merge sidecar/context and state.db messages without deleting local rows. + + ``truncation_boundary``: the original truncate cutoff — the + timestamp of the last message kept by the truncate operation. When the + watermark is later advanced (new turn committed), this boundary is preserved + so the empty-sidecar recovery can distinguish a legitimate prefix from a + deleted suffix instead of guessing by dropping one turn pair. + """ sidecar_messages = list(sidecar_messages or []) state_messages = list(state_messages or []) watermark_timestamp = _message_timestamp_as_float({"timestamp": truncation_watermark}) @@ -5126,12 +5137,13 @@ def merge_session_messages_append_only( else: # Positive watermark advanced after edit/retry/undo (#4767). # Without a sidecar there's no seen_content_keys to check against, - # so we reconstruct the correct transcript from state.db alone: - # 1. Keep all messages at/after the watermark (the new turn + - # post-edit assistant reply — these are always legitimate). - # 2. For pre-watermark messages, drop the last complete turn - # (user + assistant pair) which is the replaced tail. - # Everything before that is legitimate history. + # so we reconstruct the correct transcript from state.db alone. + # + # Use truncation_boundary (the original truncate cutoff) to + # distinguish legitimate prefix from deleted suffix. Without it, + # fall back to the backward-scan heuristic (drops last user+assistant + # pair) which works for the common single-turn case. + boundary_ts = _message_timestamp_as_float({"timestamp": truncation_boundary}) at_or_after = [] pre_watermark = [] for msg in state_messages: @@ -5140,19 +5152,28 @@ def merge_session_messages_append_only( at_or_after.append(msg) else: pre_watermark.append(msg) - # Scan pre_watermark from the end: skip assistant messages until - # we hit a user message (the replaced prompt), then stop. - i = len(pre_watermark) - 1 - while i >= 0: - role = str(pre_watermark[i].get("role", "")).lower() - if role == "assistant": - i -= 1 - elif role == "user": - i -= 1 # skip this user message too (the replaced prompt) - break - else: - i -= 1 # tool messages etc. — skip - filtered = pre_watermark[:i + 1] + at_or_after + if boundary_ts is not None: + # Use the persisted boundary: keep only messages at or before it. + pre_legitimate = [ + m for m in pre_watermark + if (ts := _message_timestamp_as_float(m)) is not None + and ts <= boundary_ts + ] + filtered = pre_legitimate + at_or_after + else: + # Fallback: backward scan (works when exactly one turn was + # deleted — the common edit/retry/undo case). + i = len(pre_watermark) - 1 + while i >= 0: + role = str(pre_watermark[i].get("role", "")).lower() + if role == "assistant": + i -= 1 + elif role == "user": + i -= 1 # skip this user message too (the replaced prompt) + break + else: + i -= 1 # tool messages etc. — skip + filtered = pre_watermark[:i + 1] + at_or_after # Deduplicate true duplicates (same role, content, exact timestamp) # without collapsing legitimately-repeated identical turns (#3346). @@ -5287,12 +5308,19 @@ def merge_session_messages_append_only( # content is not in the sidecar, it's a replaced message edited at the # same second — skip it. The edited version (same timestamp, different # content) is in the sidecar and survives this check. + # + # Only apply the same-second guard to user messages. An assistant reply + # (or tool message) at the same second as the watermark is a legitimate + # post-edit recovery row — the sidecar holds only the edited user + # checkpoint, so the assistant reply's content won't be in it and would + # be silently dropped without this role guard. 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 + and str(msg.get("role", "")).lower() == "user" ): continue # Check for true duplicates using full-precision timestamp (#3346). @@ -5352,30 +5380,44 @@ def merge_session_messages_append_only( and timestamp is not None and timestamp <= max_sidecar_timestamp ): - # Legacy key within sidecar timestamp range. Normally skip — the - # sidecar already has this message. Exception: if the state.db - # message has tool_calls that DIFFER from the sidecar version - # (same content_key but different dedup_key because tool_calls - # differ), preserve it — distinct tool_calls must not be collapsed. - _tc = msg.get("tool_calls") - if _tc: - _ck = _session_message_content_key(msg) - if _ck in seen_content_keys and dedup_key not in seen_dedup_keys: - pass # different tool_calls from sidecar — preserve + # When a truncation watermark is active and the sidecar holds only + # the edited user checkpoint, state.db may contain an assistant/tool + # reply at the same timestamp that is NOT in the sidecar. This + # block would normally skip it ("sidecar already has this message"), + # but the sidecar doesn't — it's a genuine state-only recovery row. + # Let it through (CORE-B, #4767). + if ( + watermark_timestamp is not None + and timestamp == watermark_timestamp + and str(msg.get("role", "")).lower() != "user" + and _session_message_content_key(msg) not in seen_content_keys + ): + pass # fall through to append below + else: + # Legacy key within sidecar timestamp range. Normally skip — the + # sidecar already has this message. Exception: if the state.db + # message has tool_calls that DIFFER from the sidecar version + # (same content_key but different dedup_key because tool_calls + # differ), preserve it — distinct tool_calls must not be collapsed. + _tc = msg.get("tool_calls") + if _tc: + _ck = _session_message_content_key(msg) + if _ck in seen_content_keys and dedup_key not in seen_dedup_keys: + pass # different tool_calls from sidecar — preserve + else: + _merge_session_display_metadata(merged_by_message_key.get(key), msg) + continue else: + if msg.get("role") == "user" and _session_message_content_key(msg) not in seen_content_keys: + if _insert_state_message_chronologically(merged_messages, msg): + seen_message_keys.add(key) + seen_dedup_keys.add(dedup_key) + seen_content_keys.add(_session_message_content_key(msg)) + seen_visible_keys.add(visible_key) + _remember_merged_message(msg) + continue _merge_session_display_metadata(merged_by_message_key.get(key), msg) continue - else: - if msg.get("role") == "user" and _session_message_content_key(msg) not in seen_content_keys: - if _insert_state_message_chronologically(merged_messages, msg): - seen_message_keys.add(key) - seen_dedup_keys.add(dedup_key) - seen_content_keys.add(_session_message_content_key(msg)) - seen_visible_keys.add(visible_key) - _remember_merged_message(msg) - continue - _merge_session_display_metadata(merged_by_message_key.get(key), msg) - continue seen_message_keys.add(key) seen_dedup_keys.add(dedup_key) seen_content_keys.add(_session_message_content_key(msg)) @@ -5426,6 +5468,7 @@ def reconciled_state_db_messages_for_session( local_messages, state_messages, truncation_watermark=getattr(session, "truncation_watermark", None), + truncation_boundary=getattr(session, "truncation_boundary", None), ) diff --git a/api/routes.py b/api/routes.py index 7d688c01..18ddfa2c 100644 --- a/api/routes.py +++ b/api/routes.py @@ -6059,6 +6059,7 @@ def _limited_webui_messages_for_display(session, state_db_messages) -> list: sidecar_messages, state_db_messages, truncation_watermark=getattr(session, "truncation_watermark", None), + truncation_boundary=getattr(session, "truncation_boundary", None), ) @@ -6117,6 +6118,7 @@ def _webui_sidecar_lineage_messages_for_display(session, *, max_hops: int = 20) merged, getattr(segment, "messages", []) or [], truncation_watermark=getattr(segment, "truncation_watermark", None), + truncation_boundary=getattr(segment, "truncation_boundary", None), ) return merge_session_messages_append_only( merged, @@ -10894,6 +10896,7 @@ def handle_post(handler, parsed) -> bool: context_length=getattr(session, "context_length", None), threshold_tokens=getattr(session, "threshold_tokens", None), truncation_watermark=getattr(session, "truncation_watermark", None), + truncation_boundary=getattr(session, "truncation_boundary", None), # context_messages is the authoritative model-facing prefix — must be # deepcopied so the duplicate has its own independent context that won't # be mutated when the original session's context changes (#2914). @@ -11459,8 +11462,11 @@ def handle_post(handler, parsed) -> bool: try: from api.session_ops import _truncation_watermark_for s.truncation_watermark = _truncation_watermark_for(s.messages) + # Persist the original truncate cutoff. + s.truncation_boundary = s.truncation_watermark except Exception: s.truncation_watermark = 0.0 + s.truncation_boundary = 0.0 s.save() logger.info( "truncate %s: messages %d→%d, context_messages %d→%d, watermark=%.2f", diff --git a/api/session_ops.py b/api/session_ops.py index 61e0b364..9f33075f 100644 --- a/api/session_ops.py +++ b/api/session_ops.py @@ -116,6 +116,9 @@ def retry_last(session_id: str) -> dict[str, Any]: removed_count = len(history) - last_user_idx s.messages = history[:last_user_idx] s.truncation_watermark = _truncation_watermark_for(s.messages) + # Persist the original truncate cutoff so empty-sidecar recovery + # can distinguish legitimate prefix from deleted suffix. + s.truncation_boundary = s.truncation_watermark if isinstance(getattr(s, 'context_messages', None), list) and s.context_messages: truncated_context = _truncate_at_last_user(s.context_messages) if truncated_context is not None: @@ -156,6 +159,8 @@ def undo_last(session_id: str) -> dict[str, Any]: removed_count = len(history) - last_user_idx s.messages = history[:last_user_idx] s.truncation_watermark = _truncation_watermark_for(s.messages) + # Persist the original truncate cutoff. + s.truncation_boundary = s.truncation_watermark if isinstance(getattr(s, 'context_messages', None), list) and s.context_messages: truncated_context = _truncate_at_last_user(s.context_messages) if truncated_context is not None: diff --git a/api/webui_session_db.py b/api/webui_session_db.py index 1d3a1190..56595710 100644 --- a/api/webui_session_db.py +++ b/api/webui_session_db.py @@ -52,6 +52,7 @@ _METADATA_FIELDS = frozenset( "threshold_tokens", "last_prompt_tokens", "truncation_watermark", + "truncation_boundary", "gateway_routing", "gateway_routing_history", "llm_title_generated", diff --git a/tests/test_core_data_loss_cases.py b/tests/test_core_data_loss_cases.py new file mode 100644 index 00000000..2e8ceefd --- /dev/null +++ b/tests/test_core_data_loss_cases.py @@ -0,0 +1,291 @@ +"""Reproduction tests for data-loss cases found in PR review. + +Empty-sidecar recovery resurrects deleted pre-edit turns when editing +an older message with ≥2 later turns. + +Same-second recovery can silently drop the legitimate post-edit +assistant reply. + +Both are regression tests — they should FAIL against the current code +(commit 297c88fed7) and PASS after the fix. +""" +from __future__ import annotations + +import api.models as models +import api.webui_session_db as webui_db + + +def _msg(role: str, content: str, ts: float) -> dict: + return {"role": role, "content": content, "timestamp": ts} + + +# ─── Resurrected deleted turns ─────────────────────────────────────────────── + + +def test_core_a_empty_sidecar_resurrects_deleted_turns(): + """Editing an older message with ≥2 later turns resurrects + the deleted suffix on empty-sidecar recovery.""" + state = [ + _msg("user", "original prompt", 100.0), + _msg("assistant", "original reply", 101.0), + _msg("user", "deleted question 1", 200.0), + _msg("assistant", "deleted answer 1", 201.0), + _msg("user", "deleted question 2", 300.0), + _msg("assistant", "deleted answer 2", 302.0), + _msg("user", "edited prompt", 400.0), + _msg("assistant", "post-edit reply", 401.0), + ] + + merged = models.merge_session_messages_append_only( + [], # empty sidecar (cold reload) + state, + truncation_watermark=400.0, + truncation_boundary=101.0, + ) + + contents = [m["content"] for m in merged] + + assert "original prompt" in contents + assert "original reply" in contents + assert "edited prompt" in contents + assert "post-edit reply" in contents + + assert "deleted question 1" not in contents, ( + f"Deleted turn @200 resurrected! Contents: {contents}" + ) + assert "deleted answer 1" not in contents, ( + f"Deleted turn @201 resurrected! Contents: {contents}" + ) + assert "deleted question 2" not in contents, ( + f"Deleted turn @300 resurrected! Contents: {contents}" + ) + assert "deleted answer 2" not in contents, ( + f"Deleted turn @302 resurrected! Contents: {contents}" + ) + + expected = [ + "original prompt", "original reply", + "edited prompt", "post-edit reply", + ] + assert contents == expected, ( + f"Expected {expected}, got {contents}" + ) + + +def test_core_a_single_deleted_turn_still_works(): + """Sanity check: backward-scan handles one deleted turn correctly.""" + state = [ + _msg("user", "first msg", 50.0), + _msg("assistant", "first reply", 51.0), + _msg("user", "original pre-edit", 100.0), + _msg("assistant", "original reply", 101.0), + _msg("user", "edited/new turn", 200.0), + _msg("assistant", "post-edit reply", 201.0), + ] + + merged = models.merge_session_messages_append_only( + [], state, truncation_watermark=200.0 + ) + + contents = [m["content"] for m in merged] + expected = [ + "first msg", "first reply", "edited/new turn", "post-edit reply" + ] + assert contents == expected, ( + f"Sanity: expected {expected}, got {contents}" + ) + + +# ─── Same-second assistant reply dropped ────────────────────────────────────── + + +def test_core_b_same_second_assistant_reply_dropped(): + """Same-second equality guard drops legitimate post-edit + assistant reply when it shares the timestamp with the edited user message.""" + T = 100.0 + sidecar = [_msg("user", "edited prompt", T)] + state = [ + _msg("user", "edited prompt", T), + _msg("assistant", "post-edit reply", T), + ] + + merged = models.merge_session_messages_append_only( + sidecar, state, truncation_watermark=T + ) + + contents = [m["content"] for m in merged] + + assert "edited prompt" in contents, ( + f"Edited user message missing! Contents: {contents}" + ) + assert "post-edit reply" in contents, ( + f"Same-second assistant reply was DROPPED! Contents: {contents}" + ) + + +def test_core_b_same_second_replaced_user_still_filtered(): + """Sanity check: same-second guard still filters the replaced + pre-edit user message (same timestamp, different content).""" + T = 100.0 + sidecar = [_msg("user", "edited prompt", T)] + state = [ + _msg("user", "original pre-edit prompt", T), + _msg("user", "edited prompt", T), + _msg("assistant", "post-edit reply", T), + ] + + merged = models.merge_session_messages_append_only( + sidecar, state, truncation_watermark=T + ) + + contents = [m["content"] for m in merged] + + assert "original pre-edit prompt" not in contents, ( + f"Sanity: replaced user message leaked! Contents: {contents}" + ) + assert "edited prompt" in contents + assert "post-edit reply" in contents + + +def test_core_b_same_second_empty_sidecar_assistant_reply(): + """Same-second assistant reply in empty-sidecar + recovery path with truncation_boundary — must survive.""" + T = 100.0 + state = [ + _msg("user", "first msg", 50.0), + _msg("assistant", "first reply", 51.0), + _msg("user", "edited prompt", T), + _msg("assistant", "post-edit reply", T), + ] + + merged = models.merge_session_messages_append_only( + [], state, truncation_watermark=T, + truncation_boundary=51.0, + ) + + contents = [m["content"] for m in merged] + + assert "first msg" in contents + assert "first reply" in contents + assert "edited prompt" in contents + assert "post-edit reply" in contents, ( + f"Same-second assistant reply dropped! Contents: {contents}" + ) + + +# ─── Persistence: save/load round-trip ──────────────────────────────────────── + + +def test_truncation_boundary_survives_save_load(monkeypatch, tmp_path): + """truncation_boundary must be persisted to JSON and restored on load.""" + 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() + + sid = "boundary_save_load" + session = models.Session( + session_id=sid, + messages=[_msg("user", "hello", 100.0)], + truncation_watermark=200.0, + truncation_boundary=101.0, + ) + session.save() + + with models.LOCK: + models.SESSIONS.pop(sid, None) + + loaded = models.Session.load(sid) + assert loaded.truncation_boundary == 101.0, ( + f"truncation_boundary lost after save/load: got {loaded.truncation_boundary}" + ) + assert loaded.truncation_watermark == 200.0 + + +def test_truncation_boundary_none_survives_save_load(monkeypatch, tmp_path): + """When truncation_boundary is None, it must survive save/load as None.""" + 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() + + sid = "boundary_none_save" + session = models.Session( + session_id=sid, + messages=[_msg("user", "hello", 100.0)], + truncation_watermark=None, + truncation_boundary=None, + ) + session.save() + + with models.LOCK: + models.SESSIONS.pop(sid, None) + + loaded = models.Session.load(sid) + assert getattr(loaded, "truncation_boundary", None) is None + + +# ─── reconciled_state_db_messages_for_session passes boundary ───────────────── + + +def test_reconciled_passes_truncation_boundary(monkeypatch, tmp_path): + """reconciled_state_db_messages_for_session must pass truncation_boundary + to merge_session_messages_append_only so empty-sidecar recovery uses it.""" + 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() + + sid = "reconciled_boundary" + session = models.Session( + session_id=sid, + messages=[], # empty sidecar (crash recovery) + truncation_watermark=200.0, + truncation_boundary=51.0, + ) + session.save() + + state_db = [ + _msg("user", "kept", 50.0), + _msg("assistant", "kept reply", 51.0), + _msg("user", "deleted 1", 100.0), + _msg("assistant", "deleted answer 1", 101.0), + _msg("user", "deleted 2", 150.0), + _msg("assistant", "deleted answer 2", 151.0), + _msg("user", "new turn", 200.0), + _msg("assistant", "new reply", 201.0), + ] + + monkeypatch.setattr( + models, + "get_state_db_session_messages", + lambda sid: state_db, + ) + + reconciled = models.reconciled_state_db_messages_for_session(session) + contents = [m["content"] for m in reconciled] + + assert "deleted 1" not in contents, ( + f"reconciled leaked deleted turn: {contents}" + ) + assert "deleted 2" not in contents, ( + f"reconciled leaked deleted turn: {contents}" + ) + assert "kept" in contents + assert "kept reply" in contents + assert "new turn" in contents + assert "new reply" in contents + + +# ─── webui_session_db _METADATA_FIELDS includes boundary ────────────────────── + + +def test_webui_session_db_metadata_fields_includes_boundary(): + """webui_session_db._METADATA_FIELDS must include truncation_boundary + so it's recognized as a metadata field (not leaked into extra).""" + assert "truncation_boundary" in webui_db._METADATA_FIELDS, ( + "webui_session_db._METADATA_FIELDS missing truncation_boundary" + ) diff --git a/tests/test_session_duplicate_fields.py b/tests/test_session_duplicate_fields.py index fdb88777..43db7e44 100644 --- a/tests/test_session_duplicate_fields.py +++ b/tests/test_session_duplicate_fields.py @@ -73,6 +73,16 @@ def test_duplicate_copies_truncation_watermark(): "Duplicate must copy truncation_watermark" +def test_duplicate_copies_truncation_boundary(): + """truncation_boundary must be copied alongside truncation_watermark so + empty-sidecar recovery in the duplicate can distinguish legitimate prefix + from deleted suffix.""" + block, _ = _extract_duplicate_block() + ctor = _find_session_ctor(block) + assert _has_field(ctor, 'truncation_boundary'), \ + "Duplicate must copy truncation_boundary" + + def test_duplicate_copies_context_messages(): """context_messages is the authoritative model-facing prefix. Without it, the duplicate's agent context diverges from what the user sees.""" @@ -419,6 +429,16 @@ def test_branch_does_not_copy_truncation_watermark(): "truncation_watermark should NOT be copied in branch" +def test_branch_does_not_copy_truncation_boundary(): + """truncation_boundary must NOT be copied in branch — same reasoning as + truncation_watermark: the branch is a fresh message slice with no truncate + history of its own.""" + block, _ = _extract_branch_block() + ctor = _find_session_ctor(block, 'branch') + assert not _has_field(ctor, 'truncation_boundary'), \ + "truncation_boundary should NOT be copied in branch" + + def test_branch_does_not_copy_usage_counters(): """Branch starts a new conversation — token counters must NOT carry over.""" block, _ = _extract_branch_block()