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): move all I/O out of session_store._lock in get_or_create_session

Browse source 
The second lock block in get_or_create_session held self._lock during six
blocking operations on every inbound message: _is_session_ended_in_db
(SQLite SELECT), _should_reset (callback), _save (SQLite write + JSON write
+ os.fsync), and _recover_session_from_db (SQLite SELECT + UPDATE).

A code comment at line 1607 claimed 'SQLite calls are made outside the
lock' -- true only for _compression_tip_for_session_id, which was moved
out in a prior fix. The remaining I/O was never addressed.

Restructure into a four-phase lock/no-lock split that mirrors the pattern
already established at the bottom of the function:

  Phase 1  (lock)    -- read entry + session_id
  Phase 1b (no lock) -- stale check + reset policy
  Phase 2  (lock)    -- apply decisions to _entries, capture snapshot + flags
  Phase 3  (no lock) -- recovery DB query, _save from snapshot, end/create

_save_entries(snapshot) replaces _save() to avoid dict-mutation races when
called outside the lock. _query_recoverable_session splits the DB I/O out
of _recover_session_from_db so only the _entries assignment needs the lock.

Three early returns inside the lock block are eliminated in favour of a
unified save + return path.
This commit is contained in:
kenyonxu 2026-07-10 11:34:30 +08:00 • committed by kshitij
parent 94c2a4016b
commit 08e9dcf182
2 changed files with 351 additions and 95 deletions
Download patch file Download diff file Expand all files Collapse all files

250
gateway/session.py
View file

@ -1254,6 +1254,34 @@ 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).
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
)
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.
@ -1395,6 +1423,39 @@ class SessionStore:
now=now,
)
def _query_recoverable_session(self, *, session_key, source, now):
"""DB-only half of _recover_session_from_db (no lock needed).
Returns a SessionEntry or None. Caller assigns _entries[key] under lock.
"""
if not self._db:
return None
finder = getattr(self._db, "find_latest_gateway_session_for_peer", None)
if not callable(finder):
return None
try:
recovered = finder(
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,
)
except Exception as exc:
logger.debug("Gateway session DB recovery failed for %s: %s",
session_key, exc)
return None
if not recovered:
return None
try:
self._db.reopen_session(str(recovered["id"]))
except Exception as exc:
logger.debug("Gateway session DB reopen failed for %s: %s",
session_key, exc)
return self._create_entry_from_recovered_row(
row=recovered, session_key=session_key, source=source, now=now,
)
def _record_gateway_session_peer(
self,
session_id: str,
@ -1689,21 +1750,23 @@ class SessionStore:
source: SessionSource,
force_new: bool = False
) -> SessionEntry:
"""
Get an existing session or create a new one.
"""Get an existing session or create a new one.
Evaluates reset policy to determine if the existing session is stale.
Creates a session record in SQLite when a new session starts.
All blocking I/O (SQLite SELECTs, routing-index rewrite + ``os.fsync``,
recovery DB queries) is performed *outside* ``self._lock``. The lock
protects only ``_entries`` / ``_loaded`` mutations.
"""
session_key = self._generate_session_key(source)
now = _now()
# SQLite calls are made outside the lock to avoid holding it during I/O.
# All _entries / _loaded mutations are protected by self._lock.
db_end_session_id = None
db_create_kwargs = None
existing_session_id = None
# ---- Phase 0: lock read -- existing session_id for compression tip ----
if not force_new:
with self._lock:
self._ensure_loaded_locked()
@ -1711,13 +1774,51 @@ class SessionStore:
if entry is not None:
existing_session_id = entry.session_id
# Look up the compression continuation outside the lock (DB I/O).
# Compression tip lookup outside the lock (DB I/O).
canonical_existing_session_id = (
self._compression_tip_for_session_id(existing_session_id)
if existing_session_id
else None
)
# ---- Phase 1: lock read -- get entry snapshot for stale/reset checks ----
_stale_session_id = None
_entry_for_checks = None
with self._lock:
self._ensure_loaded_locked()
if session_key in self._entries and not force_new:
_entry_for_checks = self._entries[session_key]
_stale_session_id = _entry_for_checks.session_id
# ---- Phase 1b: no-lock I/O -- stale check + reset policy ----
_is_stale = False
_reset_reason = None
if _entry_for_checks is not None:
_is_stale = self._is_session_ended_in_db(_stale_session_id)
if _entry_for_checks.suspended:
_reset_reason = "suspended"
elif _entry_for_checks.resume_pending:
_reset_reason = self._should_reset(_entry_for_checks, source)
if not _reset_reason:
_fw = auto_continue_freshness_window()
_ref_time = (
_entry_for_checks.last_resume_marked_at
or _entry_for_checks.updated_at
)
if _fw > 0 and (now - _ref_time).total_seconds() > _fw:
_reset_reason = "resume_pending_expired"
else:
_reset_reason = self._should_reset(_entry_for_checks, source)
# ---- 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
reset_had_activity = False
with self._lock:
self._ensure_loaded_locked()
@ -1727,25 +1828,13 @@ class SessionStore:
entry, existing_session_id, canonical_existing_session_id
)
# Self-heal stale routing: if this session_key still points at
# a session that has ALREADY been ended in state.db (end_reason
# set), the in-memory sessions.json entry is stale. Reusing it
# would route every incoming message into a closed session and
# silently drop it — with no log, no error, no response — until
# the gateway restarts and _prune_stale_sessions_locked() clears
# it (#54878 — the live-gateway variant of #52804/FM9, which
# only the startup prune previously caught).
#
# Drop the stale entry and fall through to the recovery path
# below. Leaving db_end_session_id None routes us into
# _recover_session_from_db, whose finder
# (hermes_state.find_latest_gateway_session_for_peer) selects
# rows WHERE `ended_at IS NULL OR end_reason = 'agent_close'`
# — so it REOPENS gateway-cleanup-ended ('agent_close') rows and
# resumes the SAME session_id (transcript preserved), but returns
# None for any other end_reason (e.g. /new), which then correctly
# starts a fresh session.
if self._is_session_ended_in_db(entry.session_id):
if _is_stale and entry.session_id == _stale_session_id:
# Stale routing self-heal (#54878): the in-memory entry
# points at a session that has ALREADY been ended in
# state.db. Drop it and fall through to recovery/create.
# Recovery finder reopens ``agent_close`` rows (preserving
# the transcript) but returns None for other end_reasons
# (e.g. /new), starting a fresh session.
logger.warning(
"gateway.session: routing key %r -> %s is ended in "
"state.db but still live in sessions.json; dropping "
@ -1754,81 +1843,47 @@ class SessionStore:
session_key, entry.session_id,
)
self._entries.pop(session_key, None)
was_auto_reset = False
auto_reset_reason = None
reset_had_activity = False
# Fall through to the recovery/create path below; the
# stale entry is gone so we must NOT consult its
# suspended/resume/reset state.
entry = None
_needs_recover = True
elif entry.session_id != _stale_session_id:
# Another thread handled this entry during our lock-free
# window. Treat as healthy -- bump updated_at and save.
entry.updated_at = now
_needs_save = True
else:
# Auto-reset sessions marked as suspended (e.g. after /stop
# broke a stuck loop — #7536). ``suspended`` is the hard
# forced-wipe signal and always wins over ``resume_pending``,
# so repeated interrupted restarts that escalate via the
# existing ``.restart_failure_counts`` stuck-loop counter
# still converge to a clean slate.
if entry.suspended:
reset_reason = "suspended"
elif entry.resume_pending:
# Restart-interrupted session: preserve the session_id
# and return the existing entry so the transcript reloads
# intact, but still honour normal daily/idle reset policy.
#
# Freshness gate (#46934): the idle/daily policy checks
# ``updated_at``, which is bumped to ``now`` on every
# message — so a zombie session that keeps receiving
# messages never trips it and would resume stale context
# forever. ``last_resume_marked_at`` is set once when
# resume was marked and never bumped per-message, so it
# correctly measures how long resume has been pending.
# If that exceeds the auto-continue freshness window, the
# recovery turn either never ran or failed — treat the
# session as a zombie and fall through to auto-reset.
reset_reason = self._should_reset(entry, source)
if not reset_reason:
_fw = auto_continue_freshness_window()
_ref_time = entry.last_resume_marked_at or entry.updated_at
if _fw > 0 and (now - _ref_time).total_seconds() > _fw:
reset_reason = "resume_pending_expired"
else:
entry.updated_at = now
self._save()
return entry
else:
reset_reason = self._should_reset(entry, source)
if not reset_reason:
entry.updated_at = now
self._save()
return entry
else:
# Session is being auto-reset.
# Stale check clean. Apply reset decision.
if _reset_reason:
was_auto_reset = True
auto_reset_reason = reset_reason
# Track whether the expired session had any real
# conversation. total_tokens is never written (token
# counts migrated to agent-direct persistence) so it is
# always 0 — use last_prompt_tokens, updated every turn.
auto_reset_reason = _reset_reason
reset_had_activity = entry.last_prompt_tokens > 0
db_end_session_id = entry.session_id
self._entries.pop(session_key, None)
entry = None
_needs_recover = True
else:
entry.updated_at = now
_needs_save = True
else:
was_auto_reset = False
auto_reset_reason = None
reset_had_activity = False
if not force_new:
_needs_recover = True
if not force_new and not db_end_session_id:
recovered_entry = self._recover_session_from_db(
session_key=session_key,
source=source,
now=now,
)
if recovered_entry is not None:
self._entries[session_key] = recovered_entry
self._save()
return recovered_entry
_entries_snapshot = dict(self._entries)
# Create new session
# ---- Phase 3: no-lock I/O -- recovery + create + save + DB ops ----
if _needs_recover:
recovered = self._query_recoverable_session(
session_key=session_key, source=source, now=now,
)
if recovered is not None:
with self._lock:
self._entries[session_key] = recovered
_entries_snapshot = dict(self._entries)
entry = recovered
_needs_save = True
if entry is None:
# Create new session.
session_id = f"{now.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
entry = SessionEntry(
session_key=session_key,
session_id=session_id,
@ -1842,9 +1897,10 @@ class SessionStore:
auto_reset_reason=auto_reset_reason,
reset_had_activity=reset_had_activity,
)
self._entries[session_key] = entry
self._save()
with self._lock:
self._entries[session_key] = entry
_entries_snapshot = dict(self._entries)
_needs_save = True
db_create_kwargs = {
"session_id": session_id,
"source": source.platform.value,
@ -1855,7 +1911,11 @@ class SessionStore:
"thread_id": source.thread_id,
}
# SQLite operations outside the lock
# Persist routing index from snapshot (was _save under lock).
if _needs_save and _entries_snapshot is not None:
self._save_entries(_entries_snapshot)
# SQLite operations outside the lock (unchanged).
if self._db and db_end_session_id:
try:
self._db.end_session(db_end_session_id, "session_reset")

196
tests/gateway/test_session_store_lock_io.py Normal file
View file

@ -0,0 +1,196 @@
"""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 threading
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
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
def tracking_save():
if lock.held:
save_calls_under_lock.append(True)
orig_save()
store._save = tracking_save
# 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._recover_session_from_db
def tracking(**kw):
if lock.held:
recover_calls_under_lock.append(True)
return orig(**kw)
store._recover_session_from_db = tracking
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"
)