1947f74cc4
`Fixes #102374` was only true for the compacted transcript: the prune ran inside local `compress()` / `salvage_grown_transcript`, and under native compaction the server compacts LOCAL_TRIGGER_SAFETY_MARGIN below the local trigger, so local compaction rarely fires. Every assistant response still persisted its ~120 KB checkpoint to `messages.codex_reasoning_items` and the older rows' shadowed copies were never rewritten — the field path in the issue (1,967 checkpoint rows, 12.9 GiB in one lineage). `SessionDB._insert_message_rows` (the one INSERT path: append batch, replace, compact, import) now rewrites the older ACTIVE assistant rows of the same session to drop their `type: "compaction"` items once a newer carrier row lands, keeping every non-checkpoint item; wire-neutral because `prune_pre_checkpoint_items` never replays a shadowed checkpoint. The flush mirrors it on the live transcript (`drop_shadowed_checkpoints`) so forks and compaction built from memory carry one checkpoint too, and the marker contract holds: the dicts read exactly as their rows.
255 lines
9.8 KiB
Python
255 lines
9.8 KiB
Python
"""Tests for SessionDB.append_messages_batch (#23254 salvage).
|
|
|
|
The batch writer reuses _insert_message_rows (the same row-serialization
|
|
path as replace/compact/import), runs the same admission guards as
|
|
append_message, is atomic (all rows or none), and aggregates the session
|
|
counters in one UPDATE.
|
|
"""
|
|
|
|
import json
|
|
import sqlite3
|
|
|
|
import pytest
|
|
|
|
from hermes_state import SessionDB
|
|
from hermes_state_errors import CompressionSessionClosedError
|
|
|
|
|
|
@pytest.fixture()
|
|
def db(tmp_path):
|
|
d = SessionDB(db_path=tmp_path / "state.db")
|
|
d.create_session("sess-batch", source="cli")
|
|
yield d
|
|
d.close()
|
|
|
|
|
|
def _turn_messages():
|
|
return [
|
|
{"role": "user", "content": "question"},
|
|
{
|
|
"role": "assistant",
|
|
"content": "let me check",
|
|
"tool_calls": [{"name": "terminal", "arguments": "{}"}],
|
|
"reasoning_content": "thinking...",
|
|
"finish_reason": "tool_calls",
|
|
},
|
|
{
|
|
"role": "tool",
|
|
"content": "tool output",
|
|
"tool_name": "terminal",
|
|
"tool_call_id": "call_1",
|
|
},
|
|
{"role": "assistant", "content": "answer", "finish_reason": "stop"},
|
|
]
|
|
|
|
|
|
class TestAppendMessagesBatch:
|
|
def test_batch_rows_identical_to_single_appends(self, db, tmp_path):
|
|
"""The batch writer stores the same bytes append_message would."""
|
|
db2 = SessionDB(db_path=tmp_path / "state2.db")
|
|
db2.create_session("sess-batch", source="cli")
|
|
try:
|
|
msgs = _turn_messages()
|
|
db.append_messages_batch("sess-batch", msgs)
|
|
for m in msgs:
|
|
role = m["role"]
|
|
db2.append_message(
|
|
session_id="sess-batch",
|
|
role=role,
|
|
content=m.get("content"),
|
|
tool_name=m.get("tool_name"),
|
|
tool_calls=m.get("tool_calls"),
|
|
tool_call_id=m.get("tool_call_id"),
|
|
finish_reason=m.get("finish_reason"),
|
|
reasoning_content=(
|
|
m.get("reasoning_content") if role == "assistant" else None
|
|
),
|
|
)
|
|
cols = (
|
|
"role, content, tool_call_id, tool_calls, tool_name, "
|
|
"finish_reason, reasoning_content, observed, active"
|
|
)
|
|
rows_a = db._conn.execute(
|
|
f"SELECT {cols} FROM messages ORDER BY id"
|
|
).fetchall()
|
|
rows_b = db2._conn.execute(
|
|
f"SELECT {cols} FROM messages ORDER BY id"
|
|
).fetchall()
|
|
assert [tuple(r) for r in rows_a] == [tuple(r) for r in rows_b]
|
|
finally:
|
|
db2.close()
|
|
|
|
def test_reasoning_gated_to_assistant_rows(self, db):
|
|
"""_insert_message_rows role-gates reasoning fields; a tool row
|
|
carrying reasoning keys must not persist them."""
|
|
db.append_messages_batch(
|
|
"sess-batch",
|
|
[
|
|
{
|
|
"role": "tool",
|
|
"content": "out",
|
|
"tool_name": "t",
|
|
"tool_call_id": "c1",
|
|
"reasoning_content": "should not persist",
|
|
}
|
|
],
|
|
)
|
|
row = db._conn.execute(
|
|
"SELECT reasoning_content FROM messages"
|
|
).fetchone()
|
|
assert row[0] is None
|
|
|
|
def test_counters_aggregate_once(self, db):
|
|
db.append_messages_batch("sess-batch", _turn_messages())
|
|
row = db._conn.execute(
|
|
"SELECT message_count, tool_call_count FROM sessions WHERE id = ?",
|
|
("sess-batch",),
|
|
).fetchone()
|
|
assert row["message_count"] == 4
|
|
assert row["tool_call_count"] == 1
|
|
|
|
def test_returns_inserted_count(self, db):
|
|
assert db.append_messages_batch("sess-batch", _turn_messages()) == 4
|
|
|
|
def test_empty_batch_is_noop(self, db):
|
|
assert db.append_messages_batch("sess-batch", []) == 0
|
|
row = db._conn.execute(
|
|
"SELECT message_count FROM sessions WHERE id = ?", ("sess-batch",)
|
|
).fetchone()
|
|
assert row["message_count"] == 0
|
|
|
|
def test_atomicity_all_or_nothing(self, db, monkeypatch):
|
|
"""A failure mid-batch leaves ZERO rows and untouched counters."""
|
|
real_insert = SessionDB._insert_message_rows
|
|
|
|
def failing_insert(self_db, conn, session_id, messages):
|
|
real_conn_execute = conn.execute
|
|
calls = {"n": 0}
|
|
|
|
def exec_counting(sql, *args):
|
|
if sql.lstrip().startswith("INSERT INTO messages"):
|
|
calls["n"] += 1
|
|
if calls["n"] == 3:
|
|
raise sqlite3.OperationalError("boom mid-batch")
|
|
return real_conn_execute(sql, *args)
|
|
|
|
conn.execute = exec_counting
|
|
try:
|
|
return real_insert(self_db, conn, session_id, messages)
|
|
finally:
|
|
conn.execute = real_conn_execute
|
|
|
|
monkeypatch.setattr(SessionDB, "_insert_message_rows", failing_insert)
|
|
with pytest.raises(sqlite3.OperationalError):
|
|
db.append_messages_batch("sess-batch", _turn_messages())
|
|
monkeypatch.undo()
|
|
|
|
count = db._conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0]
|
|
assert count == 0
|
|
row = db._conn.execute(
|
|
"SELECT message_count, tool_call_count FROM sessions WHERE id = ?",
|
|
("sess-batch",),
|
|
).fetchone()
|
|
assert row["message_count"] == 0
|
|
assert row["tool_call_count"] == 0
|
|
|
|
def test_compression_closed_session_rejected(self, db):
|
|
db._conn.execute(
|
|
"UPDATE sessions SET ended_at = 1.0, end_reason = 'compression' "
|
|
"WHERE id = ?",
|
|
("sess-batch",),
|
|
)
|
|
db._conn.commit()
|
|
with pytest.raises(CompressionSessionClosedError):
|
|
db.append_messages_batch("sess-batch", _turn_messages())
|
|
|
|
def test_multimodal_content_encoded(self, db):
|
|
msgs = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "look"},
|
|
{"type": "image_url", "image_url": {"url": "data:x"}},
|
|
],
|
|
}
|
|
]
|
|
db.append_messages_batch("sess-batch", msgs)
|
|
raw = db._conn.execute("SELECT content FROM messages").fetchone()[0]
|
|
# encoded via _encode_content — same sentinel prefix as append_message
|
|
loaded = db.get_messages("sess-batch")
|
|
assert loaded, raw
|
|
|
|
def test_tool_calls_json_string_not_double_encoded(self, db):
|
|
msgs = [
|
|
{
|
|
"role": "assistant",
|
|
"content": "x",
|
|
"tool_calls": json.dumps([{"name": "t", "arguments": "{}"}]),
|
|
}
|
|
]
|
|
db.append_messages_batch("sess-batch", msgs)
|
|
raw = db._conn.execute("SELECT tool_calls FROM messages").fetchone()[0]
|
|
assert json.loads(raw) == [{"name": "t", "arguments": "{}"}]
|
|
|
|
|
|
class TestShadowedCheckpointRowsArePruned:
|
|
"""Under native compaction every assistant response persists a fresh ``type: "compaction"`` checkpoint
|
|
and local compaction (the only other prune site) rarely fires, so older rows kept ~120 KB of ciphertext
|
|
the wire builder never replays (#102374). Landing a newer carrier row rewrites the older active rows."""
|
|
|
|
@staticmethod
|
|
def _checkpoint(tag):
|
|
return {"type": "compaction", "encrypted_content": f"ckpt-{tag}"}
|
|
|
|
@staticmethod
|
|
def _reasoning(tag):
|
|
return {"type": "reasoning", "encrypted_content": f"rs-{tag}", "id": f"rs_{tag}"}
|
|
|
|
@staticmethod
|
|
def _agent(db):
|
|
from agent.session_persistence import SessionPersistenceMixin
|
|
|
|
class _Agent(SessionPersistenceMixin):
|
|
pass
|
|
|
|
agent = _Agent()
|
|
agent._session_db, agent._session_db_created, agent.session_id = db, True, "sess-batch"
|
|
agent._last_flushed_db_idx, agent._flushed_db_message_ids = 0, set()
|
|
agent._flushed_db_message_session_id, agent._persist_disabled = None, False
|
|
return agent
|
|
|
|
def _durable_items(self, db):
|
|
return [
|
|
(row["id"], json.loads(row["codex_reasoning_items"]) if row["codex_reasoning_items"] else None)
|
|
for row in db._conn.execute(
|
|
"SELECT id, codex_reasoning_items FROM messages WHERE session_id = ? AND role = 'assistant' "
|
|
"AND active = 1 ORDER BY id", ("sess-batch",)).fetchall()
|
|
]
|
|
|
|
def test_newer_carrier_row_prunes_the_older_rows_checkpoints(self, db):
|
|
"""The production flush path: turn 1 lands a carrier, turn 2 lands a newer one -> only the newest
|
|
row still holds a checkpoint, durably and in the live transcript; reasoning items are untouched."""
|
|
agent = self._agent(db)
|
|
messages = [
|
|
{"role": "user", "content": "u0"},
|
|
{"role": "assistant", "content": "a0", "codex_reasoning_items": [self._reasoning(0), self._checkpoint(0)]},
|
|
]
|
|
assert agent._flush_messages_to_session_db(messages) is True
|
|
assert self._durable_items(db) == [(2, [self._reasoning(0), self._checkpoint(0)])]
|
|
|
|
messages += [
|
|
{"role": "user", "content": "u1"},
|
|
{"role": "assistant", "content": "a1", "codex_reasoning_items": [self._reasoning(1), self._checkpoint(1)]},
|
|
]
|
|
assert agent._flush_messages_to_session_db(messages) is True
|
|
|
|
assert self._durable_items(db) == [
|
|
(2, [self._reasoning(0)]),
|
|
(4, [self._reasoning(1), self._checkpoint(1)]),
|
|
]
|
|
# The live dicts match the rows they were persisted as (the marker contract), so no re-write is queued.
|
|
assert messages[1]["codex_reasoning_items"] == [self._reasoning(0)]
|
|
assert messages[3]["codex_reasoning_items"] == [self._reasoning(1), self._checkpoint(1)]
|
|
assert agent._flush_messages_to_session_db(messages) is True
|
|
assert db.message_count("sess-batch") == 4
|