From 69fd846ef86c034c74a855e98733f7edd3ce433e Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Tue, 14 Jul 2026 00:09:41 +0530
Subject: [PATCH] fix(session): serialize direct persistence flushes

---
 run_agent.py                      | 12 +++++++
 tests/run_agent/test_run_agent.py | 52 +++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/run_agent.py b/run_agent.py
index 82e93a5dc04a..e815971d35f5 100644
--- a/run_agent.py
+++ b/run_agent.py
@@ -1765,6 +1765,18 @@ class AIAgent:
         return repair_message_sequence(self, messages)
 
     def _flush_messages_to_session_db(self, messages: List[Dict], conversation_history: List[Dict] = None):
+        """Serialize direct and turn-boundary session flushes per agent."""
+        persist_lock = getattr(self, "_session_persist_lock", None)
+        if persist_lock is None:
+            return self._flush_messages_to_session_db_unlocked(messages, conversation_history)
+        with persist_lock:
+            return self._flush_messages_to_session_db_unlocked(messages, conversation_history)
+
+    def _flush_messages_to_session_db_unlocked(
+        self,
+        messages: List[Dict],
+        conversation_history: List[Dict] = None,
+    ):
         """Persist any un-flushed messages to the SQLite session store.
 
         Deduplicates via an intrinsic ``_DB_PERSISTED_MARKER`` stamped on each
diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py
index 07e54a50ded4..b0b9a2b93162 100644
--- a/tests/run_agent/test_run_agent.py
+++ b/tests/run_agent/test_run_agent.py
@@ -11,6 +11,7 @@ import io
 import json
 import logging
 import re
+import threading
 import uuid
 from logging.handlers import RotatingFileHandler
 from pathlib import Path
@@ -168,6 +169,57 @@ def test_flush_persist_override_replaces_api_local_multimodal_note(agent):
     assert api_content[0]["text"] == "[MODEL SWITCH NOTE]\n\nDescribe this screenshot"
 
 
+def test_direct_session_db_flushes_share_marker_claim(agent):
+    """A direct flush cannot interleave its marker check with `_persist_session`."""
+    class _BarrierDB:
+        def __init__(self):
+            self.rows = []
+            self.entered = threading.Event()
+            self.release = threading.Event()
+            self.calls = 0
+            self._lock = threading.Lock()
+
+        def append_message(self, **kwargs):
+            with self._lock:
+                self.calls += 1
+                first = self.calls == 1
+            if first:
+                self.entered.set()
+                assert self.release.wait(timeout=5)
+            self.rows.append(kwargs["content"])
+
+    db = _BarrierDB()
+    agent._session_db = db
+    agent._session_db_created = True
+    agent.session_id = "session-123"
+    agent._last_flushed_db_idx = 0
+    agent._flushed_db_message_ids = set()
+    agent._flushed_db_message_session_id = None
+    agent._persist_user_message_idx = None
+    agent._persist_user_message_override = None
+    agent._persist_user_message_timestamp = None
+    agent._persist_disabled = False
+    agent._session_persist_lock = threading.RLock()
+    agent._session_json_enabled = False
+
+    message = {"role": "user", "content": "exactly once"}
+    normal = threading.Thread(target=lambda: agent._persist_session([message], []))
+    direct = threading.Thread(target=lambda: agent._flush_messages_to_session_db([message], []))
+    normal.start()
+    assert db.entered.wait(timeout=5)
+    direct.start()
+    # Direct flush is blocked by the agent-wide persistence lock until the
+    # normal writer stamps the message's durable marker.
+    assert db.calls == 1
+    db.release.set()
+    normal.join(timeout=5)
+    direct.join(timeout=5)
+
+    assert not normal.is_alive()
+    assert not direct.is_alive()
+    assert db.rows == ["exactly once"]
+
+
 @pytest.fixture()
 def agent_with_memory_tool():
     """Agent whose valid_tool_names includes 'memory'."""