mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-16 14:32:34 +00:00
383 lines
13 KiB
Python
383 lines
13 KiB
Python
"""Regression: blocking I/O must not run while session_store._lock is held.
|
|
|
|
``get_or_create_session`` previously held the store lock during SQLite
|
|
SELECTs (``_is_session_ended_in_db``), a full routing-index rewrite +
|
|
``os.fsync`` (``_save``), and a recovery DB query
|
|
(``_recover_session_from_db``) -- all on every inbound message.
|
|
|
|
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
|
|
|
|
import pytest
|
|
|
|
from gateway.config import GatewayConfig, Platform, SessionResetPolicy
|
|
from gateway.session import SessionEntry, SessionSource, SessionStore
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class _TrackedLock:
|
|
"""Drop-in replacement for ``threading.Lock`` that tracks hold state.
|
|
|
|
Used to assert that blocking I/O runs only when the lock is released.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self._held = False
|
|
|
|
def acquire(self, *a, **kw):
|
|
r = self._lock.acquire(*a, **kw)
|
|
if r:
|
|
self._held = True
|
|
return r
|
|
|
|
def release(self):
|
|
self._held = False
|
|
self._lock.release()
|
|
|
|
def __enter__(self):
|
|
self.acquire()
|
|
return self
|
|
|
|
def __exit__(self, *a):
|
|
self.release()
|
|
|
|
@property
|
|
def held(self) -> bool:
|
|
return self._held
|
|
|
|
|
|
def _db_with_rows(rows: dict) -> MagicMock:
|
|
"""Mock SessionDB where ``get_session`` maps session_id -> row dict."""
|
|
db = MagicMock()
|
|
db.get_session.side_effect = lambda sid: rows.get(sid)
|
|
db.find_latest_gateway_session_for_peer.return_value = None
|
|
db.reopen_session.return_value = None
|
|
db.create_session.return_value = None
|
|
# Identity compression tip (no child session).
|
|
db.get_compression_tip.side_effect = lambda sid: sid
|
|
return db
|
|
|
|
|
|
def _make_store(tmp_path, db_mock=None) -> SessionStore:
|
|
"""Build a SessionStore with a ``_TrackedLock``, bypassing disk load."""
|
|
config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none"))
|
|
with patch("gateway.session.SessionStore._ensure_loaded"):
|
|
store = SessionStore(sessions_dir=tmp_path, config=config)
|
|
if db_mock is not None:
|
|
store._db = db_mock
|
|
store._loaded = True
|
|
store._lock = _TrackedLock()
|
|
return store
|
|
|
|
|
|
def _source() -> SessionSource:
|
|
return SessionSource(
|
|
platform=Platform.TELEGRAM,
|
|
chat_id="12345",
|
|
chat_type="dm",
|
|
user_id="12345",
|
|
)
|
|
|
|
|
|
def _seed_entry(store, key, session_id) -> SessionEntry:
|
|
now = datetime.now()
|
|
entry = SessionEntry(
|
|
session_key=key,
|
|
session_id=session_id,
|
|
created_at=now - timedelta(hours=2),
|
|
updated_at=now - timedelta(hours=1),
|
|
platform=Platform.TELEGRAM,
|
|
chat_type="dm",
|
|
)
|
|
store._entries[key] = entry
|
|
return entry
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestStaleCheckOutsideLock:
|
|
def test_is_session_ended_not_holding_lock(self, tmp_path):
|
|
"""``_is_session_ended_in_db`` must run with the lock released."""
|
|
source = _source()
|
|
db = _db_with_rows({
|
|
"sid_alive": {"end_reason": None, "id": "sid_alive"},
|
|
})
|
|
store = _make_store(tmp_path, db)
|
|
key = store._generate_session_key(source)
|
|
_seed_entry(store, key, "sid_alive")
|
|
|
|
lock = store._lock
|
|
calls_under_lock = []
|
|
|
|
orig = store._is_session_ended_in_db
|
|
|
|
def tracking(sid):
|
|
if lock.held:
|
|
calls_under_lock.append(sid)
|
|
return orig(sid)
|
|
|
|
store._is_session_ended_in_db = tracking # type: ignore[method-assign]
|
|
|
|
store.get_or_create_session(source)
|
|
|
|
assert not calls_under_lock, (
|
|
f"_is_session_ended_in_db called {len(calls_under_lock)} "
|
|
f"time(s) while lock was held"
|
|
)
|
|
|
|
|
|
class TestSaveOutsideLock:
|
|
def test_save_not_holding_lock(self, tmp_path):
|
|
"""``_save`` must run with the lock released."""
|
|
source = _source()
|
|
db = _db_with_rows({})
|
|
store = _make_store(tmp_path, db)
|
|
|
|
lock = store._lock
|
|
save_calls_under_lock = []
|
|
|
|
orig_save = store._save_entries
|
|
|
|
def tracking_save():
|
|
if lock.held:
|
|
save_calls_under_lock.append(True)
|
|
orig_save()
|
|
|
|
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)
|
|
|
|
assert not save_calls_under_lock, (
|
|
f"_save called {len(save_calls_under_lock)} time(s) "
|
|
f"while lock was held"
|
|
)
|
|
|
|
|
|
class TestRecoverOutsideLock:
|
|
def test_recover_not_holding_lock(self, tmp_path):
|
|
"""``_recover_session_from_db`` must run with the lock released."""
|
|
source = _source()
|
|
db = _db_with_rows({})
|
|
db.find_latest_gateway_session_for_peer.return_value = {
|
|
"id": "sid_recovered",
|
|
"started_at": datetime.now().timestamp(),
|
|
}
|
|
store = _make_store(tmp_path, db)
|
|
# No entry seeded -- forces the recovery path.
|
|
|
|
lock = store._lock
|
|
recover_calls_under_lock = []
|
|
|
|
orig = store._query_recoverable_session
|
|
|
|
def tracking(**kw):
|
|
if getattr(lock, "held", False):
|
|
recover_calls_under_lock.append(True)
|
|
return orig(**kw)
|
|
|
|
store._query_recoverable_session = tracking # type: ignore[method-assign]
|
|
|
|
store.get_or_create_session(source)
|
|
|
|
assert not recover_calls_under_lock, (
|
|
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)
|
|
owner_started = threading.Event()
|
|
release_owner = threading.Event()
|
|
original_query = store._query_recoverable_session
|
|
|
|
def synchronized_query(**kwargs):
|
|
owner_started.set()
|
|
assert release_owner.wait(timeout=2)
|
|
return original_query(**kwargs)
|
|
|
|
store._query_recoverable_session = synchronized_query # type: ignore[method-assign]
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
owner = pool.submit(store.get_or_create_session, source)
|
|
assert owner_started.wait(timeout=2)
|
|
follower = pool.submit(store.get_or_create_session, source)
|
|
release_owner.set()
|
|
entries = [owner.result(timeout=2), follower.result(timeout=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_concurrent_force_new_returns_one_published_session(tmp_path):
|
|
"""Concurrent /new delivery must not create orphan SQLite sessions."""
|
|
source = _source()
|
|
db = _db_with_rows({})
|
|
store = _make_store(tmp_path, db)
|
|
owner_started = threading.Event()
|
|
release_owner = threading.Event()
|
|
original_impl = store._get_or_create_session_impl
|
|
|
|
def synchronized_impl(*args, **kwargs):
|
|
owner_started.set()
|
|
assert release_owner.wait(timeout=2)
|
|
return original_impl(*args, **kwargs)
|
|
|
|
store._get_or_create_session_impl = synchronized_impl # type: ignore[method-assign]
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
owner = pool.submit(store.get_or_create_session, source, True)
|
|
assert owner_started.wait(timeout=2)
|
|
follower = pool.submit(store.get_or_create_session, source, True)
|
|
release_owner.set()
|
|
entries = [owner.result(timeout=2), follower.result(timeout=2)]
|
|
|
|
assert entries[0] is entries[1]
|
|
created_ids = {call.kwargs["session_id"] for call in db.create_session.call_args_list}
|
|
assert created_ids == {entries[0].session_id}
|
|
|
|
|
|
def test_auto_reset_does_not_recover_session_being_ended(tmp_path):
|
|
source = _source()
|
|
db = _db_with_rows({})
|
|
store = _make_store(tmp_path, db)
|
|
key = store._generate_session_key(source)
|
|
old = _seed_entry(store, key, "old-session")
|
|
old.suspended = True
|
|
db.find_latest_gateway_session_for_peer.return_value = {
|
|
"id": old.session_id,
|
|
"session_key": key,
|
|
"started_at": old.created_at.timestamp(),
|
|
}
|
|
|
|
entry = store.get_or_create_session(source)
|
|
|
|
assert entry.session_id != old.session_id
|
|
assert entry.was_auto_reset is True
|
|
db.reopen_session.assert_not_called()
|
|
db.end_session.assert_called_once_with(old.session_id, "session_reset")
|
|
|
|
|
|
def test_legacy_and_off_lock_saves_share_one_serialization_lock(tmp_path):
|
|
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)
|
|
_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)
|
|
_seed_entry(store, key_b, "sid-b")
|
|
future_b = pool.submit(store._save)
|
|
release_first_write.set()
|
|
future_a.result(timeout=2)
|
|
future_b.result(timeout=2)
|
|
|
|
assert set(persisted) == {key_a, key_b}
|
|
|
|
|
|
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()
|