diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 00c2527378f48..ed990a06a06be 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -17675,4 +17675,204 @@ def test_prompt_submit_row_id_real_sessiondb_unknown_refuses_despite_ordinal( server._sessions.pop(sid, None) +def test_prompt_submit_row_id_misaligned_memory_refuses_content_swap( + monkeypatch, tmp_path +): + """#82959 heal-path guard: equal-length but content-misaligned live memory + must refuse (4018), not zip-stamp durable ids positionally and cut the + wrong turn. Probe 4a from the PR #83202 review. + """ + from hermes_state import SessionDB + db = SessionDB(db_path=tmp_path / "rowid-misalign-content.db") + session_key = "real-db-row-misalign-content" + db.create_session(session_key, "cli") + msgs = [ + {"role": "user", "content": "A"}, + {"role": "assistant", "content": "ra"}, + {"role": "user", "content": "B"}, + {"role": "assistant", "content": "rb"}, + ] + with db._lock: + db._insert_message_rows(db._conn, session_key, msgs) + db._conn.commit() + rid_b = msgs[2]["_row_id"] + + # Same length + same role pattern, but content positions swapped: a + # positional stamp would mark live "B" with durable A's row id and the + # cut would keep the very turn the user rewound past. + live_history = [ + {"role": "user", "content": "B"}, + {"role": "assistant", "content": "rb"}, + {"role": "user", "content": "A"}, + {"role": "assistant", "content": "ra"}, + ] + sess = _session(history=list(live_history), session_key=session_key) + sid = "misalign-content-sid" + server._sessions[sid] = sess + monkeypatch.setattr(server, "_get_db", lambda: db) + monkeypatch.setattr( + server, "_start_agent_build", lambda *a, **k: pytest.fail("must not start a turn") + ) + monkeypatch.setattr( + server, "_start_inflight_turn", lambda *a, **k: pytest.fail("must not start a turn") + ) + + n_before = len(db.get_messages_as_conversation(session_key)) + try: + resp = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": { + "session_id": sid, + "text": "rewind B", + "truncate_before_row_id": rid_b, + "confirm_truncate": True, + }, + } + ) + assert resp.get("error") is not None + assert resp["error"]["code"] == 4018 + # Fail-closed: nothing stamped onto the misaligned dicts, nothing cut. + assert all("_row_id" not in m for m in sess["history"]) + assert len(sess["history"]) == 4 + assert len(db.get_messages_as_conversation(session_key)) == n_before + finally: + server._sessions.pop(sid, None) + + +def test_prompt_submit_row_id_misaligned_memory_role_shift_targets_real_turn( + monkeypatch, tmp_path +): + """#82959 heal-path guard: equal-length but role-misaligned live memory + must not be positionally stamped. The content-verified DB fallback still + resolves the REAL target turn in live order — the cut drops exactly the + addressed user turn, never a positionally mis-aimed one. Probe 4b from + the PR #83202 review. + """ + from hermes_state import SessionDB + + db = SessionDB(db_path=tmp_path / "rowid-misalign-role.db") + session_key = "real-db-row-misalign-role" + db.create_session(session_key, "cli") + msgs = [ + {"role": "user", "content": "A"}, + {"role": "assistant", "content": "ra"}, + {"role": "user", "content": "B"}, + {"role": "assistant", "content": "rb"}, + ] + with db._lock: + db._insert_message_rows(db._conn, session_key, msgs) + db._conn.commit() + rid_b = msgs[2]["_row_id"] + + live_history = [ + {"role": "user", "content": "A"}, + {"role": "assistant", "content": "ra"}, + {"role": "assistant", "content": "rb"}, + {"role": "user", "content": "B"}, + ] + sess = _session(history=list(live_history), session_key=session_key) + sid = "misalign-role-sid" + server._sessions[sid] = sess + monkeypatch.setattr(server, "_get_db", lambda: db) + monkeypatch.setattr(server, "_start_agent_build", lambda *a, **k: None) + monkeypatch.setattr(server, "_start_inflight_turn", lambda *a, **k: None) + + try: + resp = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": { + "session_id": sid, + "text": "rewind B", + "truncate_before_row_id": rid_b, + "confirm_truncate": True, + }, + } + ) + assert resp.get("error") is None, resp + # The addressed turn ("B") is dropped exactly; earlier live turns + # survive untouched. Before the guard, a positional zip-stamp put a + # user-row id on an assistant dict and the pre-guard fallback cut at + # a mis-aimed index. (Fresh _row_id stamps on survivors are expected — + # replace_messages re-inserts and re-stamps the surviving dicts.) + survivors = [(m["role"], m["content"]) for m in sess["history"]] + assert survivors == [ + ("user", "A"), + ("assistant", "ra"), + ("assistant", "rb"), + ] + finally: + server._sessions.pop(sid, None) + + +def test_prompt_submit_row_id_db_fallback_ordinal_mapping_verifies_content( + monkeypatch, +): + """#82959 db-fallback guard: when live/durable lengths differ, mapping the + durable user-ordinal onto live indices must verify the mapped turn shows + the durable target's content — a repaired user;user merge shifts ordinals + and would otherwise cut an extra turn silently. + """ + replaced = [] + + # Durable transcript (repaired): the merge collapsed two user turns, so + # durable user-ordinal 1 ("second") maps onto a DIFFERENT live user turn. + durable_history = [ + {"_row_id": 601, "role": "user", "content": "first\nfollow-up"}, + {"_row_id": 603, "role": "assistant", "content": "reply 1"}, + {"_row_id": 604, "role": "user", "content": "second"}, + {"_row_id": 605, "role": "assistant", "content": "reply 2"}, + ] + # Live memory (unrepaired, longer): user ordinal 1 here is "follow-up", + # NOT "second". + live_history = [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "follow-up"}, + {"role": "assistant", "content": "reply 1"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "reply 2"}, + ] + + class _FakeDB: + def replace_messages(self, key, messages, active_only=False, archive_dropped=False): + replaced.append((key, list(messages))) + + def get_messages_as_conversation(self, key, repair_alternation=False, include_row_ids=False): + return [dict(m) for m in durable_history] + + sess = _session(history=list(live_history), session_key="db-fallback-verify-key") + sid = "db-fallback-verify-sid" + server._sessions[sid] = sess + monkeypatch.setattr(server, "_get_db", lambda: _FakeDB()) + monkeypatch.setattr( + server, "_start_agent_build", lambda *a, **k: pytest.fail("must not start a turn") + ) + monkeypatch.setattr( + server, "_start_inflight_turn", lambda *a, **k: pytest.fail("must not start a turn") + ) + + try: + resp = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": { + "session_id": sid, + "text": "rewind to second", + "truncate_before_row_id": 604, + "confirm_truncate": True, + }, + } + ) + # Mapped live turn ("follow-up") does not match durable target + # ("second") — refuse instead of cutting the wrong turn. + assert resp.get("error") is not None + assert resp["error"]["code"] == 4018 + assert replaced == [] + assert len(sess["history"]) == 5 + finally: + server._sessions.pop(sid, None) diff --git a/tui_gateway/methods_prompt.py b/tui_gateway/methods_prompt.py index 9a3107c94547e..f35e58bfd0852 100644 --- a/tui_gateway/methods_prompt.py +++ b/tui_gateway/methods_prompt.py @@ -35,6 +35,35 @@ def _message_row_id(msg: dict): return None +def _mem_db_pair_agrees(mem, db_msg) -> bool: + """True when a live-memory entry plausibly corresponds to a durable row. + + Positional trust across the live and durable lists needs evidence, not + just equal lengths/ordinals: roles must match, display-marker status must + match (a marker living only on one side shifts every later position), and + an addressable user turn must show the same text. Non-string (multimodal) + content can't be compared cheaply — role/marker agreement suffices there. + Self-contained on builtins: register() rebinds callers onto server + globals, so any helper this calls must be in that namespace too. + """ + if not isinstance(mem, dict) or not isinstance(db_msg, dict): + return False + if mem.get("role") != db_msg.get("role"): + return False + if bool(mem.get("display_kind")) != bool(db_msg.get("display_kind")): + return False + if mem.get("role") == "user" and not mem.get("display_kind"): + mem_content = mem.get("content") + db_content = db_msg.get("content") + if ( + isinstance(mem_content, str) + and isinstance(db_content, str) + and mem_content.strip() != db_content.strip() + ): + return False + return True + + def _find_user_turn_by_row_id(history: list, target_row_id: int): """Return ``(user_ordinal, history_index)`` for ``target_row_id``, or None.""" for u_ord, h_idx in enumerate(_history_user_indices(history)): @@ -88,8 +117,18 @@ def _resolve_truncate_row_id(session: dict, history: list, target_row_id: int): return None # Heal missing in-memory stamps when the live list still lines up 1:1 with - # the durable transcript (common after turn-completion rewrites). - if len(db_history) == len(history): + # the durable transcript (common after turn-completion rewrites). Equal + # length alone is NOT proof of alignment: the durable copy above is loaded + # with repair_alternation=True (which can merge/drop rows) while the live + # list is unrepaired, and memory can carry optimistic/marker rows — so the + # two can coincide in length while position-shifted. A positional stamp on + # a misaligned pair is sticky and re-aims every later rewind at the wrong + # durable row. Stamp only when EVERY pair agrees (all-or-nothing): roles + # must match on every pair, and addressable user turns must match content. + if len(db_history) == len(history) and all( + _mem_db_pair_agrees(mem, db_msg) + for mem, db_msg in zip(history, db_history) + ): for mem, db_msg in zip(history, db_history): db_rid = _message_row_id(db_msg) if isinstance(db_msg, dict) else None if db_rid is not None and _message_row_id(mem) is None: @@ -101,11 +140,19 @@ def _resolve_truncate_row_id(session: dict, history: list, target_row_id: int): db_hit = _find_user_turn_by_row_id(db_history, target_row_id) if db_hit is None: return None - db_ord, _ = db_hit + db_ord, db_idx = db_hit mem_user_indices = _history_user_indices(history) if db_ord < 0 or db_ord >= len(mem_user_indices): return None - return db_ord, mem_user_indices[db_ord] + mem_idx = mem_user_indices[db_ord] + # Same-ordinal mapping across two lists that can diverge (the repaired + # durable copy may have merged a user;user pair, shifting every later + # user ordinal). Trust the mapping only when the mapped live turn shows + # the same content as the durable target — otherwise refuse (the caller + # returns fail-closed 4018) rather than cut the wrong turn (#82959). + if not _mem_db_pair_agrees(history[mem_idx], db_history[db_idx]): + return None + return db_ord, mem_idx def _pending_reaction_notes(session: dict) -> str: @@ -1286,6 +1333,7 @@ def register(server) -> None: for helper in ( _history_user_indices, _message_row_id, + _mem_db_pair_agrees, _find_user_turn_by_row_id, _resolve_truncate_row_id, _pending_reaction_notes,