mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-15 14:22:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

fix(gateway): close SessionStore concurrency gaps

Browse source
This commit is contained in:
kshitijk4poor 2026-07-10 11:55:12 +05:30 • committed by kshitij
parent 9d38a2309e
commit b3f77f5c82
3 changed files with 167 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

119
gateway/session.py
View file

@ -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:

4
gateway/slash_commands.py
View file

@ -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)

98
tests/gateway/test_session_store_lock_io.py
View file

@ -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()