From 73e3547ffddd4d8bb9856953364865e9976af963 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:12:36 +0530 Subject: [PATCH] refactor(agent): one durable-row rule for the flush and the sidecar stamp; trim tests The turn-start stamp had grown its own copy of the "what does the current user row hold" rule (persist override = clean transcript, live content = wire bytes = sidecar when they differ) that _db_flush_row already implements. Two copies drift; extract durable_user_row_content() in session_persistence and call it from both. Also: reuse _persist_lock() instead of a third open-coded lock/nullcontext ladder; drop the hasattr guard on set_latest_user_api_content (it predates this fix and exists on every SessionDB); cut the comment to the WHY; trim the new test file from 18 cases to the 7 invariants (real close flush E2E, repeated-"ok" positional protection, API-only pre-flushed turn, normal path writes nothing, compaction keeps positional, store guards). Still 3 red / 4 green when agent/turn_context.py is swapped for main's copy. --- agent/session_persistence.py | 22 +- agent/turn_context.py | 109 ++------ ...test_api_content_row_addressed_backfill.py | 261 ------------------ 3 files changed, 44 insertions(+), 348 deletions(-) diff --git a/agent/session_persistence.py b/agent/session_persistence.py index 4817209b1a..7a5977acce 100644 --- a/agent/session_persistence.py +++ b/agent/session_persistence.py @@ -7,7 +7,7 @@ import logging import re from contextlib import nullcontext -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from agent.context_compressor import ( COMPRESSED_SUMMARY_METADATA_KEY, @@ -74,6 +74,19 @@ def _override_replaces_content(msg: Dict, content: Any, override: Any) -> bool: ) +def durable_user_row_content(agent, msg: Dict, content: Any, api_content: Any) -> Tuple[Any, Any]: + """``(content, api_content)`` as the current turn's user row is written: the persist override is the + clean transcript, the live content is what the wire sent — so when they differ and nothing else was + injected, the live bytes ARE the sidecar. Shared by the flush and the turn-start stamp so the stamp + matches the row the flush wrote.""" + override = getattr(agent, "_persist_user_message_override", None) + if _override_replaces_content(msg, content, override): + if api_content is None and isinstance(content, str) and content != override: + api_content = content + content = override + return content, api_content + + def _summary_display_kind(msg: Dict) -> Any: """Standalone handoffs are hidden so they never occupy the active user slot in retry/undo dispatch; merge-into-tail carriers keep their prior visibility.""" @@ -143,12 +156,7 @@ def _db_flush_row(agent, msg: Dict, is_current_turn_user: bool) -> Dict[str, Any api_content = msg.get("api_content") if isinstance(msg.get("api_content"), str) else None timestamp = msg.get("timestamp") if is_current_turn_user and role == "user": - override = getattr(agent, "_persist_user_message_override", None) - if _override_replaces_content(msg, content, override): - # Live content is what the wire sent, the override is the clean transcript; keep the sent bytes. - if api_content is None and isinstance(content, str) and content != override: - api_content = content - content = override + content, api_content = durable_user_row_content(agent, msg, content, api_content) ov_timestamp = getattr(agent, "_persist_user_message_timestamp", None) timestamp = timestamp if ov_timestamp is None else ov_timestamp if api_content == content: diff --git a/agent/turn_context.py b/agent/turn_context.py index 1a3cca15af..d947c27236 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -744,100 +744,49 @@ def _stamp_api_content_sidecar( API copy, so stamp the exact sent bytes on the live dict for replay.""" _turn_user_msg = messages[current_turn_user_idx] live_content = _turn_user_msg.get("content") - _api_content = compose_user_api_content( - live_content or "", ext_prefetch_cache, plugin_user_context + from agent.session_persistence import _persist_lock, durable_user_row_content + # Match the row the flush wrote (persist override = clean transcript), not the live bytes. + durable_content, _api_content = durable_user_row_content( + agent, _turn_user_msg, live_content, + compose_user_api_content(live_content or "", ext_prefetch_cache, plugin_user_context), ) - - durable_content = live_content - override = getattr(agent, "_persist_user_message_override", None) - from agent.session_persistence import _override_replaces_content - if _override_replaces_content(_turn_user_msg, durable_content, override): - # When an override replaces content in SQLite (e.g. voice prefix or - # model-switch note), live content is what the API sends while the - # override is the clean transcript stored in the DB row. - # If no memory/plugin context was injected, the API-only wire content - # itself is the sidecar. - if _api_content is None and isinstance(durable_content, str) and durable_content != override: - _api_content = durable_content - durable_content = override - if _api_content is None or _api_content == durable_content: return _turn_user_msg["api_content"] = _api_content - # When this turn's user row was ALREADY materialized before the - # sidecar could be composed, the crash persist below skips the - # message (marker/identity) and the stamp would never reach the - # DB — the next turn then replays clean content and the request - # prefix diverges at this message. Two writers get there first: - # in-place preflight compaction (archive_and_compact runs before - # prefetch/pre_llm_call) and a close/early flush that raced the - # prologue on the CLI path (#102194). Both stamp ``_row_id`` on - # the live dict when they write it (``_insert_message_rows`` - # directly, ``sync_flushed_message_markers`` after the batch - # commit), so that id is at once the proof a row exists and the - # address to update — no positional guess, and no extra write on - # the normal path where the row does not exist yet. + # When another writer materialized this turn's user row BEFORE the sidecar existed — in-place + # preflight compaction, or a close/early flush that raced the prologue (#102194) — the crash + # persist marker-skips the message and the stamp never reaches the DB, so the next turn replays + # clean content and the request prefix diverges here. Both writers stamp ``_row_id`` on the live + # dict, which is at once the proof a row exists and the address to update. # - # Do NOT widen this to an unconditional backfill: without a row - # id the store can only target the newest active user row, and a - # repeated user turn ("ok", "y", "continue") makes the PREVIOUS - # turn's row compare equal — this turn's bytes would overwrite - # its sidecar and be replayed as that turn forever. + # Do NOT widen this to an unconditional backfill: without a row id the store can only target the + # newest active user row, and a repeated user turn ("ok", "y", "continue") makes the PREVIOUS + # turn's row compare equal — this turn's bytes would overwrite its sidecar for good. # - # Rotation mode needs nothing here: its compacted copies flush to - # the child session after this stamp. - # - # ``_row_id`` is read under ``_session_persist_lock`` — the lock a close/early - # flush holds while it copies this row out, commits it, and writes ``_row_id`` - # back onto the dict (``sync_flushed_message_markers``). Read outside it, the - # stamp can land between the commit and the write-back: it sees no id and - # returns, the flush finishes with ``api_content = NULL`` and marks the message - # persisted, and turn-start persist skips it — nothing is left to correct the - # row. Under the lock either this runs first (no row yet, flush waits) or the - # flush already finished and ``_row_id`` is visible. - def _backfill_locked() -> None: + # ``_row_id`` is read under ``_session_persist_lock``: a close flush holds it while it commits + # the row and only then writes ``_row_id`` back (``sync_flushed_message_markers``). Read outside + # it, the stamp can land in between, see no id, return — and the flush then marks the message + # persisted with ``api_content = NULL``, leaving no writer to correct the row. + with _persist_lock(agent): _row_id = _turn_user_msg.get("_row_id") - _has_valid_row_id = ( - isinstance(_row_id, int) - and not isinstance(_row_id, bool) - and _row_id > 0 - ) - _in_place_compacted = preflight_compressed and bool( - getattr(agent, "_last_compaction_in_place", False) - ) - if not (_has_valid_row_id or _in_place_compacted): - return + _has_valid_row_id = isinstance(_row_id, int) and not isinstance(_row_id, bool) and _row_id > 0 + _in_place_compacted = preflight_compressed and bool(getattr(agent, "_last_compaction_in_place", False)) _db = getattr(agent, "_session_db", None) - if _db is None: + if _db is None or not (_has_valid_row_id or _in_place_compacted): return try: if _has_valid_row_id: + # Fail closed on a store wrapper without the row-addressed method: a positional + # fallback here is exactly the wrong-row write this function exists to prevent. if hasattr(_db, "set_message_api_content"): - _db.set_message_api_content( - agent.session_id, _row_id, durable_content, _api_content - ) - elif hasattr(_db, "set_latest_user_api_content"): - # Compacted copy that carries no row id: fall back - # to the positional backfill, which is safe here - # because archive_and_compact just made this - # message the newest active user row. - _db.set_latest_user_api_content( - agent.session_id, durable_content, _api_content - ) + _db.set_message_api_content(agent.session_id, _row_id, durable_content, _api_content) + else: + # Compacted copies carry no row id; positional is safe only because + # archive_and_compact just made this message the newest active user row. + _db.set_latest_user_api_content(agent.session_id, durable_content, _api_content) except Exception: - logger.warning( - "api_content backfill failed for session=%s", - agent.session_id or "none", - exc_info=True, - ) - - _lock = getattr(agent, "_session_persist_lock", None) - if _lock is None: - _backfill_locked() - else: - with _lock: - _backfill_locked() + logger.warning("api_content backfill failed for session=%s", agent.session_id or "none", exc_info=True) def _persist_turn_start( diff --git a/tests/agent/test_api_content_row_addressed_backfill.py b/tests/agent/test_api_content_row_addressed_backfill.py index 0e69e200d9..fff1558f98 100644 --- a/tests/agent/test_api_content_row_addressed_backfill.py +++ b/tests/agent/test_api_content_row_addressed_backfill.py @@ -38,16 +38,6 @@ class TestSetMessageApiContent: db.create_session("s1", source="cli") return db - def test_updates_the_addressed_row(self, tmp_path): - db = self._open(tmp_path) - try: - db.append_message("s1", "user", content="ok") - row_id = db.get_messages("s1")[0]["id"] - assert db.set_message_api_content("s1", row_id, "ok", "ok\n\nCTX") == 1 - assert db.get_messages("s1")[0]["api_content"] == "ok\n\nCTX" - finally: - db.close() - def test_older_identical_row_is_untouched(self, tmp_path): """Two user turns with the same text — the repeated-"ok" shape. @@ -94,68 +84,9 @@ class TestSetMessageApiContent: finally: db.close() - def test_survives_lone_surrogate(self, tmp_path): - db = self._open(tmp_path) - try: - db.append_message("s1", "user", content="turn text") - row_id = db.get_messages("s1")[0]["id"] - dirty = "text \ud83d\ude00 \ud83d more" - assert db.set_message_api_content("s1", row_id, "turn text", dirty) == 1 - stored = db.get_messages("s1")[0]["api_content"] - assert "\ud83d" not in stored or "\ud83d\ude00" in stored - finally: - db.close() - - def test_rejects_boolean_and_invalid_row_ids_and_empty_session(self, tmp_path): - db = self._open(tmp_path) - try: - db.append_message("s1", "user", content="turn text") - row_id = db.get_messages("s1")[0]["id"] - assert db.set_message_api_content("s1", True, "turn text", "sidecar") == 0 - assert db.set_message_api_content("s1", False, "turn text", "sidecar") == 0 - assert db.set_message_api_content("s1", 0, "turn text", "sidecar") == 0 - assert db.set_message_api_content("s1", -5, "turn text", "sidecar") == 0 - assert db.set_message_api_content("", row_id, "turn text", "sidecar") == 0 - assert db.set_message_api_content(None, row_id, "turn text", "sidecar") == 0 - assert db.get_messages("s1")[0]["api_content"] is None - finally: - db.close() - - class TestPrologueRowAddressedBackfill: """The prologue gate: backfill iff a durable row exists for this dict.""" - def test_preexisting_row_receives_the_sidecar(self, tmp_path): - """A close/early flush wrote the staged CLI input before the stamp and - synced ``_row_id`` back onto it. The crash persist then skips the - message, so the prologue must push the sidecar into that exact row.""" - db = SessionDB(db_path=tmp_path / "state.db") - db.create_session("s1", source="cli") - try: - db.append_message("s1", "user", content="hello") - row_id = db.get_messages("s1")[0]["id"] - - agent = _FakeAgent() - agent.session_id = "s1" - agent._session_db = db - agent._pending_cli_user_message = { - "role": "user", - "content": "hello", - "_db_persisted": True, - "_row_id": row_id, - } - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - ctx = _build(agent) - - expected = compose_user_api_content("hello", "", "PLUGIN-CTX") - assert ctx.messages[ctx.current_turn_user_idx]["api_content"] == expected - assert db.get_messages("s1")[0]["api_content"] == expected - finally: - db.close() - def test_no_row_id_and_no_compaction_writes_nothing(self): """The normal path: the row does not exist yet and the crash persist writes it WITH the sidecar. A backfill here has no row to address and @@ -175,60 +106,6 @@ class TestPrologueRowAddressedBackfill: agent._session_db.set_message_api_content.assert_not_called() agent._session_db.set_latest_user_api_content.assert_not_called() - def test_db_persisted_alone_does_not_arm_the_backfill(self): - """``_db_persisted`` is stamped on resumed history dicts whose row id - is unknown, so it cannot stand in for ``_row_id``: arming the - positional backfill from it re-opens the wrong-row write.""" - agent = _FakeAgent() - agent._session_db = MagicMock() - agent._pending_cli_user_message = { - "role": "user", - "content": "hello", - "_db_persisted": True, - } - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - _build(agent) - - agent._session_db.set_message_api_content.assert_not_called() - agent._session_db.set_latest_user_api_content.assert_not_called() - - def test_boolean_row_id_does_not_arm_the_backfill(self): - """In Python isinstance(True, int) is True; a boolean _row_id must not - be mistaken for a valid SQLite primary key.""" - agent = _FakeAgent() - agent._session_db = MagicMock() - agent._pending_cli_user_message = { - "role": "user", - "content": "hello", - "_db_persisted": True, - "_row_id": True, - } - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - _build(agent) - - agent._session_db.set_message_api_content.assert_not_called() - agent._session_db.set_latest_user_api_content.assert_not_called() - - def test_row_id_wins_over_the_compaction_fallback(self): - """A compacted copy that kept its fresh row id is addressed by id; the - positional fallback stays for a copy that carries none.""" - agent = _make_in_place_compaction_agent(row_id=41) - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - _build(agent) - agent._session_db.set_message_api_content.assert_called_once_with( - "sess-1", 41, "hello", "hello\n\nPLUGIN-CTX" - ) - agent._session_db.set_latest_user_api_content.assert_not_called() - def test_compaction_without_row_id_keeps_positional_fallback(self): agent = _make_in_place_compaction_agent(row_id=None) with patch( @@ -241,96 +118,6 @@ class TestPrologueRowAddressedBackfill: ) agent._session_db.set_message_api_content.assert_not_called() - def test_duck_typed_store_does_not_fall_back_to_positional_when_row_id_present(self): - """When a valid _row_id exists, a store lacking set_message_api_content - must NOT fall back to set_latest_user_api_content (fails closed to - prevent wrong-row corruption on repeated inputs).""" - agent = _FakeAgent() - # Mock defining ONLY set_latest_user_api_content (like older/external stores) - mock_db = MagicMock(spec=["set_latest_user_api_content"]) - agent._session_db = mock_db - agent._pending_cli_user_message = { - "role": "user", - "content": "hello", - "_db_persisted": True, - "_row_id": 42, - } - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - _build(agent) - - mock_db.set_latest_user_api_content.assert_not_called() - - def test_wrapper_lacking_set_message_api_content_fails_closed_without_corrupting_newer_row( - self, tmp_path - ): - """[ehz0ah blocking feedback]: A wrapper exposing only set_latest_user_api_content - and delegating to SessionDB must NOT be called when _row_id is present. - With repeated 'ok' rows and _row_id=1, falling back would update the newer row at id 3; - failing closed ensures row 3 is untouched and row 1 remains unchanged.""" - db = SessionDB(db_path=tmp_path / "state.db") - db.create_session("s1", source="cli") - try: - db.append_message("s1", "user", content="ok") # id=1 - db.append_message("s1", "assistant", content="reply") # id=2 - db.append_message("s1", "user", content="ok") # id=3 - - rows = db.get_messages("s1") - row_1_id, row_3_id = rows[0]["id"], rows[2]["id"] - - class _LegacyStoreWrapper: - def __init__(self, real_db): - self._real = real_db - - def set_latest_user_api_content(self, session_id, content, api_content): - return self._real.set_latest_user_api_content( - session_id, content, api_content - ) - - wrapper = _LegacyStoreWrapper(db) - agent = _FakeAgent() - agent.session_id = "s1" - agent._session_db = wrapper - agent._pending_cli_user_message = { - "role": "user", - "content": "ok", - "_db_persisted": True, - "_row_id": row_1_id, - } - - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "SIDE-1"}], - ): - _build(agent, user_message="ok") - - # Must fail closed: neither row 1 nor row 3 was updated - rows = {r["id"]: r for r in db.get_messages("s1")} - assert rows[row_1_id]["api_content"] is None - assert rows[row_3_id]["api_content"] is None - finally: - db.close() - - def test_duck_typed_store_safely_skips_when_neither_method_present(self): - """A store mock/wrapper defining neither method skips cleanly.""" - agent = _FakeAgent() - mock_db = MagicMock(spec=[]) - agent._session_db = mock_db - agent._pending_cli_user_message = { - "role": "user", - "content": "hello", - "_db_persisted": True, - "_row_id": 42, - } - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - # Must not raise AttributeError - _build(agent) - class _RealPersistenceAgent(SessionPersistenceMixin, _FakeAgent): """Stand-in agent with the real SessionPersistenceMixin flush implementation.""" @@ -433,54 +220,6 @@ class TestRealEarlyFlushAndOverrideLifecycle: finally: db.close() - def test_pre_flushed_api_only_turn_with_injections_guards_on_durable_content(self, tmp_path): - """[ehz0ah bug 2]: Pre-flushed clean input with API-only variant AND plugin context. - The row update must use the durable clean content ('hello') as the SQL guard, - not the restored API-facing content ('[voice] hello'), so the row is updated.""" - path = tmp_path / "state.db" - db = SessionDB(db_path=path) - sid = "sess-api-only-with-inj" - db.create_session(sid, source="cli") - try: - agent = _RealPersistenceAgent(db, sid) - - clean_text = "hello" - api_text = "[voice] hello" - - staged = {"role": "user", "content": clean_text} - agent._pending_cli_user_message = staged - agent._flush_messages_to_session_db([staged], None) - assert staged.get("_row_id") is not None - - with patch( - "hermes_cli.plugins.invoke_hook", - return_value=[{"context": "PLUGIN-CTX"}], - ): - ctx = _build( - agent, - user_message=api_text, - persist_user_message=clean_text, - ) - - expected_sidecar = compose_user_api_content(api_text, "", "PLUGIN-CTX") - turn_msg = ctx.messages[ctx.current_turn_user_idx] - assert turn_msg["api_content"] == expected_sidecar - - # Database row was updated successfully by row_id with durable content guard! - db_rows = db.get_messages(sid) - assert db_rows[0]["content"] == clean_text - assert db_rows[0]["api_content"] == expected_sidecar - - # Replay restores the clean content and the composed sidecar: - conv = db.get_messages_as_conversation(sid) - assert conv[0]["content"] == clean_text - assert conv[0]["api_content"] == expected_sidecar - from agent.turn_context import substitute_api_content - substitute_api_content(conv[0]) - assert conv[0]["content"] == expected_sidecar - finally: - db.close() - def test_repeated_prompt_protected_against_positional_overwrite(self, tmp_path): """Repeated prompts 'ok' across turns: Turn 1 has sidecar, Turn 2 is pre-flushed. Row-addressed backfill on Turn 2 never mutates Turn 1's stored sidecar."""