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.
This commit is contained in:
kshitijk4poor
2026-09-08 20:12:36 +05:30
committed by kshitij
parent 4126b144bb
commit 73e3547ffd
3 changed files with 44 additions and 348 deletions

View File

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

View File

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

View File

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