fix(session): serialize direct persistence flushes

fix/verification-admin-route-recovery
kshitijk4poor 2026-07-14 00:09:41 +05:30
parent 0b422559f3
commit 69fd846ef8
2 changed files with 64 additions and 0 deletions

View File

@ -1765,6 +1765,18 @@ class AIAgent:
return repair_message_sequence(self, messages) return repair_message_sequence(self, messages)
def _flush_messages_to_session_db(self, messages: List[Dict], conversation_history: List[Dict] = None): 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. """Persist any un-flushed messages to the SQLite session store.
Deduplicates via an intrinsic ``_DB_PERSISTED_MARKER`` stamped on each Deduplicates via an intrinsic ``_DB_PERSISTED_MARKER`` stamped on each

View File

@ -11,6 +11,7 @@ import io
import json import json
import logging import logging
import re import re
import threading
import uuid import uuid
from logging.handlers import RotatingFileHandler from logging.handlers import RotatingFileHandler
from pathlib import Path 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" 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() @pytest.fixture()
def agent_with_memory_tool(): def agent_with_memory_tool():
"""Agent whose valid_tool_names includes 'memory'.""" """Agent whose valid_tool_names includes 'memory'."""