fix(chat): honor edit watermarks for distinct image rows

Apply the existing truncation and checkpoint-order eligibility to marked
state.db image rows before insertion. Keep removed turns out of full,
limited, and model-context reads without changing no-watermark replay.
Preserve conflicting provider payloads for a uniquely proven retained
image row across Undo without restoring distinct removed rows.

Assisted-by: Hermes Agent:gpt-6-sol
Assisted-by: Codex:gpt-6-luna
This commit is contained in:
starship-s committed 2026-09-24 03:04:50 -06:00
1 parent 5d42c2acd2
commit 089fb57d25
3 files changed
+340 -80

No files matched your search

+120 -78
View File
@@ -11581,6 +11581,95 @@ def _merge_session_messages_append_only_impl(
and boundary_ts is not None and boundary_ts is not None
and boundary_ts < watermark_timestamp and boundary_ts < watermark_timestamp
) )
def _state_row_is_truncated(
msg, key, content_key, timestamp, checkpoint_consumed,
*, retained_native_image_row: bool | None = None,
):
# Skip rows ABOVE the watermark only while the sidecar has NOT advanced
# past the watermark. Because Session.save() no longer auto-clears the
# watermark, an unconditional `timestamp > watermark` skip would become
# permanent and silently drop legitimate future state.db-only recovery
# rows once the session moves forward past the edit boundary. Once the
# sidecar's own max timestamp is beyond the watermark, allow state rows
# newer than the sidecar tail to merge.
#
# The sidecar's max timestamp can also EQUAL the watermark when the new
# post-edit USER turn has been checkpointed into the sidecar (its
# timestamp == the advanced watermark) but its ASSISTANT reply exists
# only in state.db (recovery before the sidecar tail advances). In that
# state truncation_boundary < watermark proves the session is genuinely
# advanced, so the post-watermark state-only reply is legitimate
# post-edit content and must merge through (not be dropped as a replaced
# tail). The conservative skip still applies for boundary is None and
# boundary == watermark (not-advanced / legacy).
#
# CRITICAL: the boundary-advanced signal may only bypass the skip AFTER
# state replay has consumed the sidecar's visible checkpoint
# (state_replay_idx >= len(sidecar_visible_sequence)). A deleted suffix
# row with ts > watermark that appears in state.db BEFORE the edited
# checkpoint must still be skipped — otherwise the advanced signal would
# resurrect it. The sidecar-max-timestamp signal needs no such gate (a
# sidecar tail beyond the watermark is itself proof the checkpoint has
# advanced).
sidecar_advanced_past_watermark = (
watermark_timestamp is not None
and (
(max_sidecar_timestamp is not None
and max_sidecar_timestamp > watermark_timestamp)
or (watermark_advanced_by_boundary and checkpoint_consumed)
)
)
message_key_seen = key in seen_message_keys
content_key_seen = content_key in seen_content_keys
if retained_native_image_row is not None:
# A marked image mirror needs durable row proof; an equal scalar
# projection from another row cannot exempt it from the watermark.
message_key_seen = content_key_seen = retained_native_image_row
if (
watermark_timestamp is not None
and timestamp is not None
and timestamp > watermark_timestamp
and not message_key_seen
and (
not sidecar_advanced_past_watermark
or (max_sidecar_timestamp is not None and timestamp <= max_sidecar_timestamp)
)
):
return True
# 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 not message_key_seen
and not content_key_seen
):
return True
# Same-second edit: if timestamp equals the watermark and the message
# 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.
return (
watermark_timestamp is not None
and timestamp is not None
and timestamp == watermark_timestamp
and not message_key_seen
and not content_key_seen
and str(msg.get("role", "")).lower() == "user"
)
for source_message in state_messages: for source_message in state_messages:
preserve_native_image_row = ( preserve_native_image_row = (
isinstance(source_message, dict) isinstance(source_message, dict)
@@ -11607,11 +11696,37 @@ def _merge_session_messages_append_only_impl(
if row_id_valid and row_id is not None if row_id_valid and row_id is not None
else None else None
) )
if ( row_id_fast_path_allowed = (
existing is not None existing is not None
and row_id not in ambiguous_row_ids and row_id not in ambiguous_row_ids
and _row_id_fast_path_allowed(existing, msg) and _row_id_fast_path_allowed(existing, msg)
)
existing_timestamp, existing_timestamp_valid = (
_message_exact_timestamp_details(existing)
)
incoming_timestamp, incoming_timestamp_valid = (
_message_exact_timestamp_details(msg)
)
retained_native_image_row = (
row_id_fast_path_allowed
and sidecar_row_id_counts.get(row_id, 0) == 1
and state_row_id_counts.get(row_id, 0) == 1
and existing_timestamp_valid
and incoming_timestamp_valid
and existing_timestamp is not None
and existing_timestamp == incoming_timestamp
and str(msg.get("role") or "").lower() == "user"
)
if _state_row_is_truncated(
msg,
key,
content_key,
timestamp,
state_replay_idx >= len(sidecar_visible_sequence),
retained_native_image_row=retained_native_image_row,
): ):
continue
if row_id_fast_path_allowed:
existing_api_content = _session_message_api_content_key(existing) existing_api_content = _session_message_api_content_key(existing)
incoming_api_content = _session_message_api_content_key(msg) incoming_api_content = _session_message_api_content_key(msg)
if ( if (
@@ -11620,7 +11735,8 @@ def _merge_session_messages_append_only_impl(
and existing_api_content != incoming_api_content and existing_api_content != incoming_api_content
): ):
# Row identity does not establish which provider payload is # Row identity does not establish which provider payload is
# newer. Fall through so both conflicting versions survive. # newer. A unique exact row match proves this row survived
# truncation, so preserve both versions for model context.
pass pass
else: else:
if existing_api_content is None and incoming_api_content is not None: if existing_api_content is None and incoming_api_content is not None:
@@ -11743,83 +11859,9 @@ def _merge_session_messages_append_only_impl(
_copy_api_content_sidecar(existing, msg) _copy_api_content_sidecar(existing, msg)
_merge_session_display_metadata(existing, msg) _merge_session_display_metadata(existing, msg)
continue continue
# Skip rows ABOVE the watermark only while the sidecar has NOT advanced
# past the watermark. Because Session.save() no longer auto-clears the
# watermark, an unconditional `timestamp > watermark` skip would become
# permanent and silently drop legitimate future state.db-only recovery
# rows once the session moves forward past the edit boundary. Once the
# sidecar's own max timestamp is beyond the watermark (the session has
# advanced), allow state rows newer than the sidecar tail to merge.
#
# The sidecar's max timestamp can also EQUAL the watermark when the new
# post-edit USER turn has been checkpointed into the sidecar (its
# timestamp == the advanced watermark) but its ASSISTANT reply exists
# only in state.db (recovery before the sidecar tail advances). In that
# state truncation_boundary < watermark proves the session is genuinely
# advanced, so the post-watermark state-only reply is legitimate
# post-edit content and must merge through (not be dropped as a replaced
# tail). The conservative skip still applies for boundary is None and
# boundary == watermark (not-advanced / legacy).
#
# CRITICAL: the boundary-advanced signal may only bypass the skip AFTER
# state replay has consumed the sidecar's visible checkpoint
# (state_replay_idx >= len(sidecar_visible_sequence)). A deleted suffix
# row with ts > watermark that appears in state.db BEFORE the edited
# checkpoint must still be skipped — otherwise the advanced signal would
# resurrect it. The sidecar-max-timestamp signal needs no such gate (a
# sidecar tail beyond the watermark is itself proof the checkpoint has
# advanced).
checkpoint_consumed = state_replay_idx >= len(sidecar_visible_sequence) checkpoint_consumed = state_replay_idx >= len(sidecar_visible_sequence)
sidecar_advanced_past_watermark = ( if _state_row_is_truncated(
watermark_timestamp is not None msg, key, content_key, timestamp, checkpoint_consumed,
and (
(max_sidecar_timestamp is not None
and max_sidecar_timestamp > watermark_timestamp)
or (watermark_advanced_by_boundary and checkpoint_consumed)
)
)
if (
watermark_timestamp is not None
and timestamp is not None
and timestamp > watermark_timestamp
and key not in seen_message_keys
and (
not sidecar_advanced_past_watermark
or (max_sidecar_timestamp is not None and timestamp <= max_sidecar_timestamp)
)
):
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 content_key not in seen_content_keys
):
continue
# Same-second edit: if timestamp equals the watermark and the message
# 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 content_key not in seen_content_keys
and str(msg.get("role", "")).lower() == "user"
): ):
continue continue
# Check for true duplicates using full-precision timestamp (#3346). # Check for true duplicates using full-precision timestamp (#3346).
@@ -199,8 +199,11 @@ and 5; it does not mark every run-state boundary implemented.
share the existing sidecar bubble only when its valid durable row ID, exact share the existing sidecar bubble only when its valid durable row ID, exact
timestamp, and exact visible user content match; keep the sidecar-owned row timestamp, and exact visible user content match; keep the sidecar-owned row
and its display metadata. Distinct row IDs, ambiguous or invalid identities, and its display metadata. Distinct row IDs, ambiguous or invalid identities,
and different visible user text remain separate. Fill a missing payload from and different visible user text remain separate only while eligible under
the other copy; repeated reconciliation must remain bounded and idempotent. the existing edit/undo truncation watermark and checkpoint-order rules;
removed rows must not reappear in display or model replay. Fill a missing
payload from the other copy; repeated reconciliation must remain bounded
and idempotent.
Agent state.db alone cannot restore the original attachment if the WebUI Agent state.db alone cannot restore the original attachment if the WebUI
sidecar is lost. sidecar is lost.
Visible interim assistant progress must remain visible timeline content; a Visible interim assistant progress must remain visible timeline content; a
@@ -338,6 +338,221 @@ def test_get_session_projects_marked_payload_conflict_in_full_and_limited_paths(
} }
@pytest.mark.parametrize("msg_limit", [None, 5])
@pytest.mark.parametrize("removed_payload_matches_retained", [False, True])
def test_truncation_watermark_keeps_proven_retained_image_row_only(
monkeypatch, tmp_path, msg_limit, removed_payload_matches_retained,
):
import api.config
import api.models as models
import api.session_ops
from api import routes
session_id = f"native-image-truncated-route-{msg_limit}"
timestamp = 860.0
retained_tail_timestamp = 880.0
sidecar_payload = "SIDECAR-IMAGE-PAYLOAD-A"
state_payload = "STATE-DB-IMAGE-PAYLOAD-B"
removed_payload = (
sidecar_payload if removed_payload_matches_retained
else "STATE-DB-REMOVED-IMAGE-PAYLOAD"
)
seed, identity, _ = _settle_image_turn(
session_id=session_id,
timestamp=timestamp,
agent_row_id=41,
)
context_user = next(
message for message in seed.context_messages
if message.get("_active_turn_token") == identity["token"]
)
mirror = _durable_agent_content(context_user["content"])
older_rows = [
{"role": "user", "content": "Earlier prompt", "timestamp": 100.0},
{"role": "assistant", "content": "Earlier response", "timestamp": 101.0},
]
seed.messages[:0] = [dict(row) for row in older_rows]
seed.context_messages[:0] = [dict(row) for row in older_rows]
sidecar_attachments = [{
"name": "sidecar-owned.png",
"mime": "image/png",
"is_image": True,
}]
retained_sidecar_row = {
"role": "user",
"content": mirror,
"timestamp": timestamp,
"_state_db_row_id": 42,
"api_content": sidecar_payload,
"attachments": sidecar_attachments,
}
retained_tail_row = {
"role": "assistant",
"content": "Retained after Undo",
"timestamp": retained_tail_timestamp,
}
seed.messages.extend([retained_sidecar_row, retained_tail_row])
seed.context_messages.extend([retained_sidecar_row, retained_tail_row])
session_dir = tmp_path / "sessions"
session_dir.mkdir()
db_path = tmp_path / "state.db"
with sqlite3.connect(db_path) as conn:
conn.execute(
"CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id TEXT, role TEXT, "
"content TEXT, timestamp REAL, active INTEGER DEFAULT 1, api_content TEXT)"
)
conn.executemany(
"INSERT INTO messages (id, session_id, role, content, timestamp, api_content) "
"VALUES (?, ?, ?, ?, ?, ?)",
[
(42, session_id, "user", mirror, timestamp, state_payload),
(43, session_id, "user", mirror, timestamp, removed_payload),
(44, session_id, "user", mirror, retained_tail_timestamp, "SAME-SECOND-REMOVED-PAYLOAD"),
],
)
monkeypatch.setattr(models, "SESSION_DIR", session_dir)
monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
monkeypatch.setattr(routes, "SESSION_DIR", session_dir)
monkeypatch.setattr(models, "_active_state_db_path", lambda: db_path)
session = models.Session(
session_id=session_id,
workspace="/fixture",
model="fixture-model",
context_length=128_000,
messages=seed.messages,
context_messages=seed.context_messages,
truncation_watermark=retained_tail_timestamp,
truncation_boundary=retained_tail_timestamp,
source_tag="webui",
session_source="webui",
)
session.save(skip_index=True)
with models.LOCK:
models.SESSIONS.pop(session_id, None)
state_rows = models.get_state_db_session_messages(session_id)
marked_state_rows = models._suppress_native_image_display_mirrors(
session,
state_rows,
)
assert all(
message["_webui_unmatched_native_image_mirror"] is True
for message in marked_state_rows[:2]
)
monkeypatch.setattr(routes, "_session_visible_to_active_profile", lambda *_: True)
monkeypatch.setattr(routes, "_active_stream_ids", lambda: set())
monkeypatch.setattr(routes, "_lookup_cli_session_metadata", lambda *_: {})
monkeypatch.setattr(routes, "find_run_summary", lambda *_: None)
monkeypatch.setattr(api.config, "load_settings", lambda: {"api_redact_enabled": False})
monkeypatch.setattr(api.session_ops, "regeneration_state", lambda _session: ([], []))
monkeypatch.setattr(api.session_ops, "regeneration_authority", lambda *_args, **_kwargs: None)
response = {}
monkeypatch.setattr(
routes,
"j",
lambda _handler, payload, status=200, **_kwargs: (
response.update(payload=payload, status=status) or payload
),
)
query = f"session_id={session_id}&resolve_model=0"
if msg_limit is not None:
query += f"&msg_limit={msg_limit}"
routes._handle_session_get(None, SimpleNamespace(path="/api/session", query=query))
assert response["status"] == 200
public_session = response["payload"]["session"]
public_messages = public_session["messages"]
if msg_limit is not None:
assert len(public_messages) == msg_limit
assert public_session["message_count"] > len(public_messages)
public_mirrors = [
message for message in public_messages
if message.get("content") == mirror
]
assert len(public_mirrors) == 1
assert public_mirrors[0]["attachments"] == sidecar_attachments
assert any(
message.get("content") == "Retained after Undo"
for message in public_messages
)
assert next(
message for message in public_messages
if message.get("content") == "Describe this image"
)["attachments"][0]["name"] == "sample.png"
first_public = None
first_context = None
for _ in range(3):
display = models.reconciled_state_db_messages_for_session(
session,
state_messages=state_rows,
)
context = models.reconciled_state_db_messages_for_session(
session,
prefer_context=True,
state_messages=state_rows,
)
display_row_42 = [
message for message in display
if message.get("_state_db_row_id") == 42
]
context_row_42 = [
message for message in context
if message.get("_state_db_row_id") == 42
]
assert len(display_row_42) == 1
assert display_row_42[0]["api_content"] == sidecar_payload
assert display_row_42[0]["attachments"] == sidecar_attachments
assert len(context_row_42) == 2
assert {message["api_content"] for message in context_row_42} == {
sidecar_payload, state_payload,
}
assert not any(
message.get("_state_db_row_id") in (43, 44)
for message in (*display, *context)
)
assert any(
message.get("_active_turn_token") == identity["token"]
for message in context
)
public = public_session_projection({"messages": display})["messages"]
replay = _sanitize_messages_for_agent(context)
assert sum(message.get("content") == mirror for message in public) == 1
replay_row_42 = [
message for message in replay
if message.get("content") == mirror
]
assert len(replay_row_42) == 2
assert {message["api_content"] for message in replay_row_42} == {
sidecar_payload, state_payload,
}
if first_public is None:
first_public = public
first_context = replay
else:
assert public == first_public
assert replay == first_context
session.truncation_boundary = timestamp - 10
advanced_state_copy = {
"role": "user",
"content": mirror,
"timestamp": retained_tail_timestamp + 5,
"_state_db_row_id": 42,
"api_content": state_payload,
}
advanced_context = models.reconciled_state_db_messages_for_session(
session,
prefer_context=True,
state_messages=[advanced_state_copy],
)
assert not any(
message.get("timestamp") == retained_tail_timestamp + 5
for message in advanced_context
)
def test_get_session_projects_parent_only_payload_conflict_without_losing_parent_rows( def test_get_session_projects_parent_only_payload_conflict_without_losing_parent_rows(
monkeypatch, tmp_path, monkeypatch, tmp_path,
): ):