fix(#4767): review fixes - more edge cases

This commit is contained in:
AlexeyDsov 2026-06-24 16:08:28 +03:00 committed by nesquena-hermes
parent 195d53ea90
commit 037253e4fc
6 changed files with 407 additions and 41 deletions

View file

@ -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),
)

View file

@ -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",

View file

@ -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:

View file

@ -52,6 +52,7 @@ _METADATA_FIELDS = frozenset(
"threshold_tokens",
"last_prompt_tokens",
"truncation_watermark",
"truncation_boundary",
"gateway_routing",
"gateway_routing_history",
"llm_title_generated",

View file

@ -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"
)

View file

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