diff --git a/api/models.py b/api/models.py index 2488c9df..98cf40a5 100644 --- a/api/models.py +++ b/api/models.py @@ -11581,6 +11581,95 @@ def _merge_session_messages_append_only_impl( and boundary_ts is not None 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: preserve_native_image_row = ( 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 else None ) - if ( + row_id_fast_path_allowed = ( existing is not None and row_id not in ambiguous_row_ids 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) incoming_api_content = _session_message_api_content_key(msg) if ( @@ -11620,7 +11735,8 @@ def _merge_session_messages_append_only_impl( and existing_api_content != incoming_api_content ): # 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 else: 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) _merge_session_display_metadata(existing, msg) 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) - 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) - ) - ) - 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" + if _state_row_is_truncated( + msg, key, content_key, timestamp, checkpoint_consumed, ): continue # Check for true duplicates using full-precision timestamp (#3346). diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index b6301765..919ed855 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/docs/rfcs/webui-run-state-consistency-contract.md @@ -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 timestamp, and exact visible user content match; keep the sidecar-owned row and its display metadata. Distinct row IDs, ambiguous or invalid identities, - and different visible user text remain separate. Fill a missing payload from - the other copy; repeated reconciliation must remain bounded and idempotent. + and different visible user text remain separate only while eligible under + 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 sidecar is lost. Visible interim assistant progress must remain visible timeline content; a diff --git a/tests/test_native_image_turn_display_context.py b/tests/test_native_image_turn_display_context.py index cf3de3a0..ae9205b1 100644 --- a/tests/test_native_image_turn_display_context.py +++ b/tests/test_native_image_turn_display_context.py @@ -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( monkeypatch, tmp_path, ):