🪲 BUG-#19: Persist only new messages instead of rewriting the whole session

pycodeloop/core/session.py

@@ -15,6 +15,7 @@ class Session:
     system_prompt: str = ""
     messages: list[Message] = field(default_factory=list)
     cwd: str = "."
+    dirty: bool = field(default=False, repr=False, compare=False)
     _lock: threading.Lock = field(
         default_factory=threading.Lock, repr=False, compare=False
     )
@@ -77,15 +78,21 @@ def _repair_dangling_tool_calls(self) -> None:
                         tool_call_id=call["id"],
                     ),
                 )
+            if missing:
+                self.dirty = True
             i = j + len(missing)
 
     def replace_messages(self, messages: list[Message]) -> None:
         """Swap in a whole new message list — e.g. compaction replacing
         older history with a condensed summary. Holding the same lock
         `add_*` uses closes the race where a tool-result thread appends
-        to the list being discarded right as it's replaced."""
+        to the list being discarded right as it's replaced. Marks the
+        session `dirty` so storage backends that append incrementally
+        (tracking a message-count watermark) know to do a full rewrite
+        on the next save instead of trusting the watermark."""
         with self._lock:
             self.messages = messages
+            self.dirty = True
 
     def trim(self, max_turns: int) -> None:
         """Keep only the most recent `max_turns` user-initiated turns,
@@ -102,3 +109,4 @@ def trim(self, max_turns: int) -> None:
 
             cutoff = turn_starts[-max_turns]
             self.messages = self.messages[cutoff:]
+            self.dirty = True

pycodeloop/store/sqlite_sessions.py

@@ -37,6 +37,7 @@ def __init__(self, path: str | Path | None = None) -> None:
         )
         self._migrate_legacy_messages()
         self._migrate_images_column()
+        self._session_identity: dict[str, int] = {}
 
     def _db(self) -> OrmSession:
         return self._session_factory()
@@ -119,22 +120,63 @@ def close(self) -> None:
         self._engine.dispose()
 
     def post(self, key: str, session: Session) -> None:
+        """Persists `session`, writing only the messages appended since
+        the last save when possible — the common case, since `on_message`
+        fires after every message including every tool result. Falls
+        back to a full delete-and-reinsert when the session is new, a
+        *different* `Session` object is posted for a key we've already
+        saved (an intentional overwrite, not the same object growing —
+        checked via identity, since a message-count watermark alone
+        can't tell the two apart), `session.dirty` (compaction/trim/
+        repair replaced or reordered existing messages, not just
+        appended), or the message count shrank.
+
+        `dirty` and `messages` are snapshotted together under
+        `session._lock` before touching the DB — `_repair_dangling_tool_calls`
+        sets `dirty` while holding that same lock, so reading the two
+        fields separately and unlocked could catch `dirty` before a
+        concurrent repair flips it, silently dropping the repair from
+        this save."""
+        with session._lock:
+            is_dirty = session.dirty
+            messages_snapshot = list(session.messages)
+
         with self._db() as db:
             record = db.get(SessionRecord, key)
+            is_new = record is None
 
-            if record is None:
+            if is_new:
                 record = SessionRecord(key=key)
                 db.add(record)
 
             record.system_prompt = session.system_prompt
             record.cwd = session.cwd
             record.updated_at = time.time()
-            record.message_count = len(session.messages)
 
-            db.query(MessageRecord).filter(
-                MessageRecord.session_key == key
-            ).delete()
-            for position, message in enumerate(session.messages):
+            previous_count = 0 if is_new else record.message_count
+            same_object = self._session_identity.get(key) == id(session)
+            full_rewrite = (
+                is_new
+                or not same_object
+                or is_dirty
+                or len(messages_snapshot) < previous_count
+            )
+            record.message_count = len(messages_snapshot)
+
+            if full_rewrite:
+                db.query(MessageRecord).filter(
+                    MessageRecord.session_key == key
+                ).delete()
+                new_messages = list(enumerate(messages_snapshot))
+            else:
+                new_messages = list(
+                    enumerate(
+                        messages_snapshot[previous_count:],
+                        start=previous_count,
+                    )
+                )
+
+            for position, message in new_messages:
                 db.add(
                     MessageRecord(
                         session_key=key,
@@ -156,6 +198,8 @@ def post(self, key: str, session: Session) -> None:
                 )
 
             db.commit()
+            session.dirty = False
+            self._session_identity[key] = id(session)
 
     def get(self, key: str) -> Session | None:
         with self._db() as db:
@@ -191,6 +235,7 @@ def get(self, key: str) -> Session | None:
             )
 
     def delete(self, key: str) -> None:
+        self._session_identity.pop(key, None)
         with self._db() as db:
             record = db.get(SessionRecord, key)
 

tests/core/test_session.py

@@ -119,6 +119,62 @@ def test_message_images_defaults_to_none(self):
         self.assertIsNone(message.images)
 
 
+class TestSessionDirtyFlag(unittest.TestCase):
+    def test_starts_clean(self):
+        session = Session(system_prompt="sys")
+        session.add_user("hi")
+
+        self.assertFalse(session.dirty)
+
+    def test_replace_messages_marks_dirty(self):
+        session = Session(system_prompt="sys")
+        session.add_user("hi")
+
+        session.replace_messages([Message(role="user", content="summary")])
+
+        self.assertTrue(session.dirty)
+
+    def test_trim_that_actually_truncates_marks_dirty(self):
+        session = Session(system_prompt="sys")
+        for i in range(5):
+            _add_turn(session, f"user-{i}", f"assistant-{i}")
+
+        session.trim(max_turns=2)
+
+        self.assertTrue(session.dirty)
+
+    def test_trim_that_is_a_noop_does_not_mark_dirty(self):
+        session = Session(system_prompt="sys")
+        _add_turn(session, "user-0", "assistant-0")
+
+        session.trim(max_turns=5)
+
+        self.assertFalse(session.dirty)
+
+    def test_repairing_a_dangling_tool_call_marks_dirty(self):
+        session = Session(system_prompt="sys")
+        session.add_user("do the thing")
+        session.add_assistant(
+            "", tool_calls=[{"id": "c1", "name": "bash", "arguments": {}}]
+        )
+
+        session.history()
+
+        self.assertTrue(session.dirty)
+
+    def test_a_complete_history_call_does_not_mark_dirty(self):
+        session = Session(system_prompt="sys")
+        session.add_user("do the thing")
+        session.add_assistant(
+            "", tool_calls=[{"id": "c1", "name": "bash", "arguments": {}}]
+        )
+        session.add_tool_result("c1", "output")
+
+        session.history()
+
+        self.assertFalse(session.dirty)
+
+
 class TestSessionThreadSafety(unittest.TestCase):
     def test_concurrent_tool_results_are_never_lost(self):
         session = Session(system_prompt="sys")

tests/store/test_sqlite_sessions.py

@@ -3,6 +3,10 @@
 import tempfile
 import unittest
 from pathlib import Path
+from unittest import mock
+
+from sqlalchemy import text
+from sqlalchemy.orm import Query
 
 from pycodeloop.core.session import Message, Session
 from pycodeloop.store.sqlite_sessions import SqliteSessions
@@ -16,6 +20,27 @@ def setUp(self):
             path=Path(self._tmpdir.name) / "sessions.db"
         )
 
+    def _post_and_spy_on_delete(self, key: str, session: Session) -> bool:
+        """Posts `session` and returns whether the message table's
+        delete-and-reinsert path ran, by spying on `Query.delete` —
+        the only place `post()` deletes existing `MessageRecord` rows."""
+        with mock.patch.object(
+            Query, "delete", autospec=True, side_effect=Query.delete
+        ) as spy:
+            self.sessions.post(key, session)
+        return spy.called
+
+    def _row_ids(self, key: str) -> list[int]:
+        with self.sessions._engine.connect() as conn:
+            rows = conn.execute(
+                text(
+                    "SELECT id FROM messages WHERE session_key = :key "
+                    "ORDER BY position"
+                ),
+                {"key": key},
+            ).fetchall()
+        return [row[0] for row in rows]
+
     def test_post_then_get_roundtrips_session(self):
         session = Session(system_prompt="sys", cwd="/tmp")
         session.add_user("hi")
@@ -73,6 +98,52 @@ def test_list_sessions_returns_index(self):
         self.assertEqual(index["s1"]["message_count"], 2)
         self.assertEqual(index["s1"]["cwd"], "/tmp/proj")
 
+    def test_post_appends_incrementally_without_rewriting_existing_rows(
+        self,
+    ):
+        session = Session(system_prompt="sys")
+        session.add_user("hi")
+        session.add_assistant("hello")
+        self.sessions.post("s1", session)
+        original_ids = self._row_ids("s1")
+
+        session.add_tool_result("call-1", "ok")
+        did_rewrite = self._post_and_spy_on_delete("s1", session)
+
+        self.assertFalse(did_rewrite)
+        new_ids = self._row_ids("s1")
+        self.assertEqual(new_ids[: len(original_ids)], original_ids)
+        self.assertEqual(len(new_ids), len(original_ids) + 1)
+
+        restored = self.sessions.get("s1")
+        self.assertEqual(len(restored.messages), 3)
+        self.assertEqual(restored.messages[2].content, "ok")
+
+    def test_post_does_a_full_rewrite_when_the_session_is_marked_dirty(self):
+        session = Session(system_prompt="sys")
+        session.add_user("hi")
+        session.add_assistant("hello")
+        self.sessions.post("s1", session)
+
+        session.replace_messages([Message(role="user", content="summary")])
+        did_rewrite = self._post_and_spy_on_delete("s1", session)
+
+        self.assertTrue(did_rewrite)
+        restored = self.sessions.get("s1")
+        self.assertEqual(len(restored.messages), 1)
+        self.assertEqual(restored.messages[0].content, "summary")
+
+    def test_post_does_a_full_rewrite_for_a_different_session_object(self):
+        first = Session(system_prompt="sys")
+        first.add_user("first")
+        self.sessions.post("s1", first)
+
+        second = Session(system_prompt="sys")
+        second.add_user("second")
+        did_rewrite = self._post_and_spy_on_delete("s1", second)
+
+        self.assertTrue(did_rewrite)
+
     def test_creates_db_file_and_parent_dir(self):
         nested = Path(self._tmpdir.name) / "nested" / "sessions.db"
         SqliteSessions(path=nested)