diff --git a/gateway/session.py b/gateway/session.py index c4e571c1a76..58b8368d7c6 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -991,6 +991,10 @@ class SessionStore: self._entries: Dict[str, SessionEntry] = {} self._loaded = False self._lock = threading.Lock() + # Serialize whole-index persistence without holding ``_lock`` across + # SQLite / fsync. Each writer snapshots the latest state only after + # acquiring this lock, preventing stale delayed writes. + self._save_lock = threading.Lock() self._has_active_processes_fn = has_active_processes_fn # Whether to keep writing the legacy sessions.json mirror alongside # the primary gateway_routing table in state.db. Default True for @@ -1272,34 +1276,36 @@ class SessionStore: logger.debug("Could not remove temp file %s: %s", tmp_path, e) raise - def _save_entries(self, entries_snapshot: Dict[str, "SessionEntry"]) -> None: - """Persist a caller-captured snapshot of entries (callable off-lock). + def _save_entries(self) -> None: + """Persist the latest routing index without holding the state lock for I/O.""" + save_lock = getattr(self, "_save_lock", None) + if save_lock is None: + save_lock = threading.Lock() + self._save_lock = save_lock + with save_lock: + # Capture immutable dictionaries under the state lock, then release + # it before SQLite replacement and JSON fsync. + with self._lock: + data = {key: entry.to_dict() for key, entry in self._entries.items()} - Same write logic as :meth:`_save` but operates on a snapshot dict - instead of reading ``self._entries`` live. This avoids - ``RuntimeError: dictionary changed size during iteration`` when - another thread mutates ``_entries`` concurrently. - """ - data = {key: entry.to_dict() for key, entry in entries_snapshot.items()} + db_saved = False + _db = getattr(self, "_db", None) + if _db: + replacer = getattr(_db, "replace_gateway_routing_entries", None) + if callable(replacer): + try: + replacer( + {k: json.dumps(v) for k, v in data.items()}, + scope=self._routing_scope(), + ) + db_saved = True + except Exception as exc: + logger.warning( + "gateway.session: state.db routing save failed: %s", exc + ) - db_saved = False - _db = getattr(self, "_db", None) - if _db: - replacer = getattr(_db, "replace_gateway_routing_entries", None) - if callable(replacer): - try: - replacer( - {k: json.dumps(v) for k, v in data.items()}, - scope=self._routing_scope(), - ) - db_saved = True - except Exception as exc: - logger.warning( - "gateway.session: state.db routing save failed: %s", exc - ) - - if getattr(self, "_write_sessions_json", True) or not db_saved: - self._save_sessions_json(data) + if getattr(self, "_write_sessions_json", True) or not db_saved: + self._save_sessions_json(data) def _resolve_profile_for_key(self, source: Optional[SessionSource] = None) -> Optional[str]: """Return the profile namespace for session keys, or None when off. @@ -1464,7 +1470,19 @@ class SessionStore: logger.debug("Gateway session DB recovery failed for %s: %s", session_key, exc) return None - if not recovered: + if not isinstance(recovered, dict): + return None + if not self._recovered_row_allowed_for_active_profile( + requested_session_key=session_key, + recovered=recovered, + ): + logger.warning( + "Gateway session DB recovery ignored %s for %s because " + "multiplex_profiles is disabled and the row belongs to a " + "different profile", + recovered.get("session_key"), + session_key, + ) return None try: self._db.reopen_session(str(recovered["id"])) @@ -1831,7 +1849,6 @@ class SessionStore: # ---- Phase 2: lock write -- apply decisions to _entries ---- _needs_save = False _needs_recover = False - _entries_snapshot: Optional[Dict[str, SessionEntry]] = None entry: Optional[SessionEntry] = None was_auto_reset = False auto_reset_reason = None @@ -1885,8 +1902,6 @@ class SessionStore: if not force_new: _needs_recover = True - _entries_snapshot = dict(self._entries) - # ---- Phase 3: no-lock I/O -- recovery + create + save + DB ops ---- if _needs_recover: recovered = self._query_recoverable_session( @@ -1894,15 +1909,18 @@ class SessionStore: ) if recovered is not None: with self._lock: - self._entries[session_key] = recovered - _entries_snapshot = dict(self._entries) - entry = recovered + published = self._entries.get(session_key) + if published is None: + self._entries[session_key] = recovered + published = recovered + entry = published _needs_save = True if entry is None: - # Create new session. + # Create a candidate outside the lock, then publish only if another + # worker has not already populated this routing key. session_id = f"{now.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}" - entry = SessionEntry( + candidate = SessionEntry( session_key=session_key, session_id=session_id, created_at=now, @@ -1916,22 +1934,25 @@ class SessionStore: reset_had_activity=reset_had_activity, ) with self._lock: - self._entries[session_key] = entry - _entries_snapshot = dict(self._entries) + published = self._entries.get(session_key) if not force_new else None + if published is None: + self._entries[session_key] = candidate + published = candidate + entry = published _needs_save = True - db_create_kwargs = { - "session_id": session_id, - "source": source.platform.value, - "user_id": source.user_id, - "session_key": session_key, - "chat_id": source.chat_id, - "chat_type": source.chat_type, - "thread_id": source.thread_id, - } + if entry is candidate: + db_create_kwargs = { + "session_id": session_id, + "source": source.platform.value, + "user_id": source.user_id, + "session_key": session_key, + "chat_id": source.chat_id, + "chat_type": source.chat_type, + "thread_id": source.thread_id, + } - # Persist routing index from snapshot (was _save under lock). - if _needs_save and _entries_snapshot is not None: - self._save_entries(_entries_snapshot) + if _needs_save: + self._save_entries() # SQLite operations outside the lock (unchanged). if self._db and db_end_session_id: diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index df72a5e0367..8d4eee356f9 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -4044,7 +4044,9 @@ class GatewaySlashCommandsMixin: # Same engine the desktop popover uses (PR #54907). The system # prompt / tools / skills / memory slices read off the live agent; # the conversation slice is estimated from the session transcript. - breakdown_lines = self._context_breakdown_lines(agent, source) + breakdown_lines = await asyncio.to_thread( + self._context_breakdown_lines, agent, source + ) if breakdown_lines: lines.append("") lines.extend(breakdown_lines) diff --git a/tests/gateway/test_session_store_lock_io.py b/tests/gateway/test_session_store_lock_io.py index ef8cbdbe3ff..64d6291c974 100644 --- a/tests/gateway/test_session_store_lock_io.py +++ b/tests/gateway/test_session_store_lock_io.py @@ -8,7 +8,9 @@ SELECTs (``_is_session_ended_in_db``), a full routing-index rewrite + These tests assert those three I/O calls are invoked *outside* the lock. They follow the mock-DB idiom from ``test_session_store_runtime_stale_guard``. """ +import json import threading +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta from unittest.mock import MagicMock, patch @@ -146,14 +148,14 @@ class TestSaveOutsideLock: lock = store._lock save_calls_under_lock = [] - orig_save = store._save + orig_save = store._save_entries def tracking_save(): if lock.held: save_calls_under_lock.append(True) orig_save() - store._save = tracking_save # type: ignore[method-assign] + store._save_entries = tracking_save # type: ignore[method-assign] # force_new bypasses the existing-entry path, goes straight to create. store.get_or_create_session(source, force_new=True) @@ -179,14 +181,14 @@ class TestRecoverOutsideLock: lock = store._lock recover_calls_under_lock = [] - orig = store._recover_session_from_db + orig = store._query_recoverable_session def tracking(**kw): if getattr(lock, "held", False): recover_calls_under_lock.append(True) return orig(**kw) - store._recover_session_from_db = tracking # type: ignore[method-assign] + store._query_recoverable_session = tracking # type: ignore[method-assign] store.get_or_create_session(source) @@ -194,3 +196,91 @@ class TestRecoverOutsideLock: f"_recover_session_from_db called " f"{len(recover_calls_under_lock)} time(s) while lock was held" ) + + +def test_concurrent_same_key_returns_one_published_session(tmp_path): + """Concurrent first messages for one routing key must converge on one ID.""" + source = _source() + db = _db_with_rows({}) + store = _make_store(tmp_path, db) + barrier = threading.Barrier(2) + original_query = store._query_recoverable_session + + def synchronized_query(**kwargs): + barrier.wait(timeout=2) + return original_query(**kwargs) + + store._query_recoverable_session = synchronized_query # type: ignore[method-assign] + with ThreadPoolExecutor(max_workers=2) as pool: + entries = list(pool.map(lambda _: store.get_or_create_session(source), range(2))) + + key = store._generate_session_key(source) + assert entries[0] is entries[1] + assert entries[0].session_id == store._entries[key].session_id + created_ids = {call.kwargs["session_id"] for call in db.create_session.call_args_list} + assert created_ids == {entries[0].session_id} + + +def test_save_serialization_snapshots_latest_routing_index(tmp_path): + """A delayed earlier writer must snapshot the state visible when it writes.""" + db = _db_with_rows({}) + persisted: dict[str, str] = {} + first_write_started = threading.Event() + release_first_write = threading.Event() + write_count = 0 + count_lock = threading.Lock() + + def replace(entries, *, scope): + nonlocal write_count, persisted + with count_lock: + write_count += 1 + call_number = write_count + if call_number == 1: + first_write_started.set() + assert release_first_write.wait(timeout=2) + persisted = dict(entries) + + db.replace_gateway_routing_entries.side_effect = replace + store = _make_store(tmp_path, db) + source_a = _source() + source_b = SessionSource( + platform=Platform.TELEGRAM, + chat_id="67890", + chat_type="dm", + user_id="67890", + ) + key_a = store._generate_session_key(source_a) + key_b = store._generate_session_key(source_b) + entry_a = _seed_entry(store, key_a, "sid-a") + + with ThreadPoolExecutor(max_workers=2) as pool: + future_a = pool.submit(store._save_entries) + assert first_write_started.wait(timeout=2) + entry_b = _seed_entry(store, key_b, "sid-b") + future_b = pool.submit(store._save_entries) + release_first_write.set() + future_a.result(timeout=2) + future_b.result(timeout=2) + + assert set(store._entries) == {key_a, key_b} + assert set(persisted) == {key_a, key_b} + assert json.loads(persisted[key_a])["session_id"] == entry_a.session_id + assert json.loads(persisted[key_b])["session_id"] == entry_b.session_id + + +def test_recovery_rejects_other_profile_row(tmp_path, monkeypatch): + """The lock-free recovery path must retain the canonical profile guard.""" + source = _source() + db = _db_with_rows({}) + db.find_latest_gateway_session_for_peer.return_value = { + "id": "foreign-session", + "session_key": "agent:other:telegram:dm:12345", + "started_at": datetime.now().timestamp(), + } + store = _make_store(tmp_path, db) + monkeypatch.setattr(store, "_active_profile_name", lambda: "default") + + entry = store.get_or_create_session(source) + + assert entry.session_id != "foreign-session" + db.reopen_session.assert_not_called()