refactor(gateway/session): split SessionStore into persistence/recovery/lifecycle/transcript mixins by call-graph cohesion; compact wire helpers
This commit is contained in:
Teknium
@@ -0,0 +1,656 @@
|
||||
"""SessionStore storage plumbing: per-profile SessionDB handle resolution and the
|
||||
routing-index load/save paths (state.db gateway_routing primary, sessions.json
|
||||
legacy mirror).
|
||||
|
||||
Mixin split out of ``gateway/session.py``; bound onto ``SessionStore`` via the MRO.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
from utils import atomic_replace
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from gateway.session import SessionEntry
|
||||
|
||||
# Log-record parity with the origin module.
|
||||
logger = logging.getLogger("gateway.session")
|
||||
|
||||
|
||||
class SessionPersistenceMixin:
|
||||
"""SessionStore storage plumbing: per-profile SessionDB handle resolution and
|
||||
the routing-index load/save paths (state.db gateway_routing primary,
|
||||
sessions.json legacy mirror).
|
||||
"""
|
||||
|
||||
def _open_session_db_for_active_scope(self, db_path: Optional[Path] = None):
|
||||
"""SessionDB for the profile scope active on this task.
|
||||
|
||||
``db_path`` pins the store; otherwise ``_default_db_path()`` follows the
|
||||
context-local HERMES_HOME from ``_profile_runtime_scope`` (resolved per
|
||||
call so multiplexed profiles reach their own store). Handles are cached
|
||||
per path; failed opens enter a bounded backoff during which callers keep
|
||||
using the JSONL fallback.
|
||||
"""
|
||||
from hermes_state import _default_db_path, get_shared_session_db
|
||||
|
||||
path = Path(db_path) if db_path is not None else Path(_default_db_path())
|
||||
def _open():
|
||||
try:
|
||||
# Process-wide shared registry: one writer connection per path.
|
||||
return get_shared_session_db(path)
|
||||
except Exception as e:
|
||||
if isinstance(e, RuntimeError) and "live-system guard" in str(e):
|
||||
# Test-isolation guard: must stay a loud failure and is
|
||||
# deliberately not cached so it fires again next attempt.
|
||||
raise
|
||||
print(f"[gateway] Warning: SQLite session store unavailable, falling back to JSONL: {e}")
|
||||
raise
|
||||
|
||||
return self._db_handle_cache.get(
|
||||
path,
|
||||
_open,
|
||||
non_cacheable=lambda exc: (
|
||||
isinstance(exc, RuntimeError) and "live-system guard" in str(exc)
|
||||
),
|
||||
)
|
||||
|
||||
def _pinned_db(self):
|
||||
"""Return the explicitly pinned DB (``store._db = x``), else ``_DB_UNPINNED``."""
|
||||
from gateway.session import _DB_UNPINNED
|
||||
return getattr(self, "_db_pinned", _DB_UNPINNED)
|
||||
|
||||
@property
|
||||
def _db(self):
|
||||
"""The SessionDB for the active profile scope, or a pinned override.
|
||||
|
||||
Assigning ``store._db`` pins that value for every subsequent read
|
||||
(tests install a fake or disable the DB with ``store._db = None``).
|
||||
Unpinned, each read resolves the scope so a multiplexed profile's
|
||||
writes reach its own store.
|
||||
"""
|
||||
from gateway.session import _DB_UNPINNED
|
||||
pinned = self._pinned_db()
|
||||
if pinned is not _DB_UNPINNED:
|
||||
return pinned
|
||||
return self._open_session_db_for_active_scope()
|
||||
|
||||
@_db.setter
|
||||
def _db(self, value) -> None:
|
||||
self._db_pinned = value
|
||||
|
||||
@property
|
||||
def _routing_db(self):
|
||||
"""The one store that owns the routing index, whatever scope is active.
|
||||
|
||||
``_entries`` is one flat dict holding every profile's keys, so it must
|
||||
persist to ONE file (``_routing_home``), not whichever profile is
|
||||
scoped — otherwise a mid-turn rewrite and the unscoped startup load see
|
||||
different copies and crash markers under a secondary profile go
|
||||
unrecovered. A pinned handle still wins; bare test instances lacking
|
||||
the handle cache report no DB.
|
||||
"""
|
||||
from gateway.session import _DB_UNPINNED
|
||||
pinned = self._pinned_db()
|
||||
if pinned is not _DB_UNPINNED:
|
||||
return pinned
|
||||
home = getattr(self, "_routing_home", None)
|
||||
try:
|
||||
if home is None:
|
||||
return self._db
|
||||
return self._open_session_db_for_active_scope(db_path=home / "state.db")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _named_profile_for_key(self, session_key: Optional[str]) -> Optional[str]:
|
||||
"""The non-default profile that owns *session_key*, or None.
|
||||
|
||||
None means the ambient store is authoritative (multiplexing off, or
|
||||
legacy ``agent:main`` namespace). It deliberately does NOT cover "that
|
||||
profile has no directory" — ownership and resolvability are separate
|
||||
questions that ``_db_for_key`` answers separately.
|
||||
"""
|
||||
if not getattr(self.config, "multiplex_profiles", False):
|
||||
return None
|
||||
profile = self._profile_from_session_key(session_key)
|
||||
if not profile or profile == "default":
|
||||
return None
|
||||
return profile
|
||||
|
||||
def _profile_home_for_key(self, session_key: Optional[str]) -> Optional[Path]:
|
||||
"""HERMES_HOME of the profile that owns *session_key*, or None.
|
||||
|
||||
None means only "no live home to point at" — no named owner, or the
|
||||
owner's directory could not be resolved.
|
||||
"""
|
||||
profile = self._named_profile_for_key(session_key)
|
||||
if profile is None:
|
||||
return None
|
||||
cache = self._profile_home_cache
|
||||
if profile in cache:
|
||||
return cache[profile]
|
||||
home: Optional[Path] = None
|
||||
try:
|
||||
from hermes_cli.profiles import get_profile_dir, profile_exists
|
||||
|
||||
if profile_exists(profile):
|
||||
home = Path(get_profile_dir(profile))
|
||||
except Exception as exc:
|
||||
logger.debug("Could not resolve profile home for %r: %s", session_key, exc)
|
||||
home = None
|
||||
# Only hits are memoized: a profile directory can be provisioned
|
||||
# *after* startup (enrollment bridge), and a cached miss would pin
|
||||
# that profile's rows to the ambient store for the process lifetime.
|
||||
if home is not None:
|
||||
cache[profile] = home
|
||||
return home
|
||||
|
||||
def _db_for_key(self, session_key: Optional[str]):
|
||||
"""The SessionDB holding *session_key*'s rows, whatever scope is active.
|
||||
|
||||
``_db`` follows the ambient HERMES_HOME that only the inbound message
|
||||
path installs; background work (expiry watcher) runs unscoped and would
|
||||
write profile rows into the ROOT store until the stale-route self-heal
|
||||
drops a live conversation. The owning profile is encoded in the key.
|
||||
"""
|
||||
from gateway.session import _DB_UNPINNED
|
||||
pinned = self._pinned_db()
|
||||
if pinned is not _DB_UNPINNED:
|
||||
return pinned
|
||||
profile = self._named_profile_for_key(session_key)
|
||||
if profile is None:
|
||||
return self._db
|
||||
home = self._profile_home_for_key(session_key)
|
||||
if home is None:
|
||||
# Named owner we cannot resolve (not provisioned yet, or lookup
|
||||
# failed). Falling back to the ambient store would split ONE
|
||||
# session identity across two physical stores — fail closed;
|
||||
# callers already handle a missing DB.
|
||||
logger.warning(
|
||||
"gateway.session: profile %r has no resolvable home (key %r); "
|
||||
"refusing to fall back to the ambient store",
|
||||
profile, session_key,
|
||||
)
|
||||
return None
|
||||
try:
|
||||
return self._open_session_db_for_active_scope(db_path=home / "state.db")
|
||||
except Exception:
|
||||
# Same contract as ``_db``: a failed open degrades to JSONL fallback.
|
||||
return None
|
||||
|
||||
def _owner_key_for_session_id(self, session_id: Optional[str]) -> Optional[str]:
|
||||
"""The routing key that owns *session_id*, or None.
|
||||
|
||||
The published index is authoritative; ``_session_owner_hints`` covers
|
||||
the window where ownership is proven but routing not yet published.
|
||||
Deliberately lock-free: several callers already hold ``_lock``.
|
||||
"""
|
||||
if not session_id:
|
||||
return None
|
||||
try:
|
||||
for entry in list(self._entries.values()):
|
||||
if entry.session_id == session_id:
|
||||
return entry.session_key
|
||||
except Exception:
|
||||
pass # bare stores / foreign entry objects in suites
|
||||
return (getattr(self, "_session_owner_hints", None) or {}).get(session_id)
|
||||
|
||||
def _db_for_session_id(self, session_id: Optional[str]):
|
||||
"""The SessionDB holding *session_id*'s row (owner recovered from the
|
||||
index or a pre-published hint; unknown ids fall back to the ambient store)."""
|
||||
if not session_id:
|
||||
return self._db
|
||||
return self._db_for_key(self._owner_key_for_session_id(session_id))
|
||||
|
||||
def close_all_db_handles(self) -> None:
|
||||
"""Close every SessionDB handle this store opened (one per path).
|
||||
|
||||
Closing only ``store._db`` would strand secondary profiles' handles with
|
||||
their WAL lock held ('database is locked' on restart). Drained under the
|
||||
lock, closed outside it; a pinned handle is the pinner's to close.
|
||||
"""
|
||||
def _close(db) -> None:
|
||||
# Shared instances no-op on close(); release the refcount instead.
|
||||
from hermes_state import release_or_close
|
||||
try:
|
||||
release_or_close(db)
|
||||
except Exception as exc:
|
||||
logger.debug("SessionDB close error during handle sweep: %s", exc)
|
||||
|
||||
self._db_handle_cache.close_all(_close)
|
||||
|
||||
def _ensure_loaded(self) -> None:
|
||||
"""Load sessions index from disk if not already loaded."""
|
||||
with self._lock:
|
||||
self._ensure_loaded_locked()
|
||||
|
||||
def _entry_locked(self, session_key: str) -> Optional[SessionEntry]:
|
||||
"""Load the index and return the entry for *session_key*. Lock held."""
|
||||
self._ensure_loaded_locked()
|
||||
return self._entries.get(session_key)
|
||||
|
||||
def _routing_scope(self) -> str:
|
||||
"""Namespace for this store's gateway_routing rows: the resolved
|
||||
sessions_dir, so stores with different dirs never share entries."""
|
||||
try:
|
||||
return str(Path(self.sessions_dir).resolve())
|
||||
except Exception:
|
||||
return str(self.sessions_dir)
|
||||
|
||||
def _routing_db_method(self, name: str):
|
||||
"""Bound ``_routing_db.<name>`` if the handle exists and has it, else None."""
|
||||
db = self._routing_db
|
||||
method = getattr(db, name, None) if db else None
|
||||
return method if callable(method) else None
|
||||
|
||||
@staticmethod
|
||||
def _routing_entry_from_json(key: str, entry_json: str) -> Optional[SessionEntry]:
|
||||
"""Parse one gateway_routing row; None (with a warning) when invalid."""
|
||||
from gateway.session import SessionEntry
|
||||
try:
|
||||
entry_data = json.loads(entry_json)
|
||||
if isinstance(entry_data, dict):
|
||||
return SessionEntry.from_dict(entry_data)
|
||||
except (ValueError, KeyError, TypeError) as e:
|
||||
logger.warning("Skipping invalid routing entry %r: %s", key, e)
|
||||
return None
|
||||
|
||||
def _ensure_loaded_locked(self) -> None:
|
||||
"""Load the routing index. Must be called with self._lock held.
|
||||
|
||||
state.db ``gateway_routing`` is primary; sessions.json is the legacy
|
||||
import for keys the DB lacks (persisted to the DB on the next _save).
|
||||
"""
|
||||
if self._loaded:
|
||||
self._reconcile_recovered_routing_locked()
|
||||
return
|
||||
|
||||
self.sessions_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
db_had_entries = False
|
||||
db_load_succeeded = False
|
||||
loader = self._routing_db_method("load_gateway_routing_entries")
|
||||
if loader is not None:
|
||||
try:
|
||||
for key, entry_json in loader(scope=self._routing_scope()).items():
|
||||
entry = self._routing_entry_from_json(key, entry_json)
|
||||
if entry is not None:
|
||||
self._entries[key] = entry
|
||||
db_had_entries = bool(self._entries)
|
||||
db_load_succeeded = True
|
||||
except Exception as e:
|
||||
logger.warning("gateway.session: state.db routing load failed: %s", e)
|
||||
|
||||
self._import_legacy_sessions_json(db_had_entries)
|
||||
|
||||
self._loaded = True
|
||||
self._routing_db_loaded = db_load_succeeded
|
||||
self._routing_fallback_baseline = (
|
||||
None
|
||||
if db_load_succeeded
|
||||
else {key: entry.to_dict() for key, entry in self._entries.items()}
|
||||
)
|
||||
|
||||
# A hard crash skips graceful shutdown and leaves sessions.json
|
||||
# pointing at ended sessions; self-heal before the first message.
|
||||
self._prune_stale_sessions_locked()
|
||||
|
||||
def _import_legacy_sessions_json(self, db_had_entries: bool) -> None:
|
||||
"""Legacy import: sessions.json fills only keys the DB lacks. Lock held."""
|
||||
from gateway.session import SessionEntry
|
||||
sessions_file = self.sessions_dir / "sessions.json"
|
||||
if not sessions_file.exists():
|
||||
return
|
||||
try:
|
||||
with open(sessions_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
imported = 0
|
||||
for key, entry_data in data.items():
|
||||
# "_"-prefixed keys are sentinels (e.g. "_README"), not entries.
|
||||
if key.startswith("_") or key in self._entries:
|
||||
continue
|
||||
# A non-dict entry (corrupt file) must not abort the whole load.
|
||||
if not isinstance(entry_data, dict):
|
||||
logger.warning(
|
||||
"Skipping invalid session entry %r: "
|
||||
"expected dict, got %s",
|
||||
key, type(entry_data).__name__,
|
||||
)
|
||||
continue
|
||||
try:
|
||||
self._entries[key] = SessionEntry.from_dict(entry_data)
|
||||
imported += 1
|
||||
except (ValueError, KeyError, TypeError) as e:
|
||||
logger.warning("Skipping invalid session entry %r: %s", key, e)
|
||||
if imported and db_had_entries:
|
||||
logger.info(
|
||||
"gateway.session: imported %d legacy sessions.json "
|
||||
"entr%s missing from state.db routing table",
|
||||
imported, "y" if imported == 1 else "ies",
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[gateway] Warning: Failed to load sessions: {e}")
|
||||
|
||||
def _prune_stale_sessions_locked(self) -> None:
|
||||
"""Remove routing entries whose session has ended in state.db (startup, lock held).
|
||||
|
||||
Stale == ``end_reason IS NOT NULL``. Rows absent from the DB are kept;
|
||||
a ``None`` DB handle is a no-op; DB errors are non-fatal.
|
||||
"""
|
||||
if not self._entries:
|
||||
return
|
||||
|
||||
stale_keys: list = []
|
||||
recovered_keys = 0
|
||||
try:
|
||||
for key, entry in self._entries.items():
|
||||
# Ask the store that owns the key, not the ambient handle, or a
|
||||
# live secondary-profile session gets pruned on the root copy.
|
||||
db = self._db_for_key(key)
|
||||
if db is None:
|
||||
continue
|
||||
row = db.get_session(entry.session_id)
|
||||
if row is None or row.get("end_reason") is None:
|
||||
continue
|
||||
verdict = self._stale_entry_verdict(key, entry, row)
|
||||
if verdict == "prune":
|
||||
stale_keys.append(key)
|
||||
elif verdict is not None:
|
||||
self._entries[key] = verdict
|
||||
recovered_keys += 1
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"gateway.session: stale-entry pruning skipped due to DB error: %s",
|
||||
exc,
|
||||
)
|
||||
return
|
||||
|
||||
for key in stale_keys:
|
||||
del self._entries[key]
|
||||
|
||||
if stale_keys or recovered_keys:
|
||||
self._save()
|
||||
|
||||
def _stale_entry_verdict(self, key: str, entry, row):
|
||||
"""For a routing entry whose row has ended: ``"prune"``, a replacement
|
||||
entry (repoint), or None (keep as-is)."""
|
||||
from gateway.session import _now
|
||||
recovered_entry = None
|
||||
if entry.origin is not None:
|
||||
try:
|
||||
recovered_entry = self._recover_session_from_db(
|
||||
session_key=key,
|
||||
source=entry.origin,
|
||||
now=_now(),
|
||||
raise_on_lookup_error=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
# Indeterminate: keep the only routing handle.
|
||||
logger.debug(
|
||||
"gateway.session: recovery lookup failed for stale "
|
||||
"sessions.json entry %r -> %s: %s",
|
||||
key,
|
||||
entry.session_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
# Compression-ended parent with a newer live child for the same peer:
|
||||
# repoint instead of dropping, or queued/resume-pending work vanishes
|
||||
# until the next message.
|
||||
if recovered_entry is not None and recovered_entry.session_id != entry.session_id:
|
||||
logger.warning(
|
||||
"gateway.session: repointing stale sessions.json entry "
|
||||
"%r from ended %s (end_reason=%r) to recovered %s",
|
||||
key,
|
||||
entry.session_id,
|
||||
row["end_reason"],
|
||||
recovered_entry.session_id,
|
||||
)
|
||||
return recovered_entry
|
||||
|
||||
# Same-id recovery == successful resume: keep the ORIGINAL entry object
|
||||
# (the recovered one is rebuilt minimal and would drop counters,
|
||||
# model_override, resume markers, metadata). Nothing changes, no save.
|
||||
if recovered_entry is not None:
|
||||
logger.info(
|
||||
"gateway.session: reopened ended session %s for "
|
||||
"sessions.json entry %r (end_reason=%r); keeping route",
|
||||
entry.session_id, key, row["end_reason"],
|
||||
)
|
||||
return None
|
||||
|
||||
logger.warning(
|
||||
"gateway.session: pruning stale sessions.json entry "
|
||||
"%r -> %s (end_reason=%r); left by a crashed gateway",
|
||||
key, entry.session_id, row["end_reason"],
|
||||
)
|
||||
return "prune"
|
||||
|
||||
def _save(self) -> None:
|
||||
"""Persist the routing index while the caller holds ``_lock``."""
|
||||
data, generation = self._snapshot_routing_locked()
|
||||
self._persist_routing_data(data, generation)
|
||||
|
||||
def _next_routing_generation_locked(self) -> int:
|
||||
"""Bump and return the shared routing counter. Caller holds ``_lock``.
|
||||
|
||||
Full snapshots AND single-entry fast saves MUST allocate from this one
|
||||
counter: the stale-write protection is a total order over
|
||||
serialization times and silently breaks otherwise.
|
||||
"""
|
||||
self._routing_generation = getattr(self, "_routing_generation", 0) + 1
|
||||
return self._routing_generation
|
||||
|
||||
def _reconcile_recovered_routing_locked(self) -> None:
|
||||
"""Merge authoritative rows after a fallback-only startup load."""
|
||||
baseline = getattr(self, "_routing_fallback_baseline", None)
|
||||
if getattr(self, "_routing_db_loaded", False) or baseline is None:
|
||||
return
|
||||
|
||||
loader = self._routing_db_method("load_gateway_routing_entries")
|
||||
if loader is None:
|
||||
return
|
||||
try:
|
||||
durable = loader(scope=self._routing_scope())
|
||||
except Exception as exc:
|
||||
logger.warning("gateway.session: recovered state.db routing load failed: %s", exc)
|
||||
return
|
||||
|
||||
current = {key: entry.to_dict() for key, entry in self._entries.items()}
|
||||
for key, entry_json in durable.items():
|
||||
durable_entry = self._routing_entry_from_json(key, entry_json)
|
||||
if durable_entry is None:
|
||||
continue
|
||||
|
||||
if key not in baseline:
|
||||
# A key created while on fallback wins over a DB-only key;
|
||||
# otherwise restore the authoritative row that fallback never saw.
|
||||
self._entries.setdefault(key, durable_entry)
|
||||
elif key not in current:
|
||||
# The key was loaded from fallback and deliberately removed.
|
||||
continue
|
||||
elif current[key] == baseline[key]:
|
||||
# Unchanged fallback data yields to the authoritative DB copy.
|
||||
self._entries[key] = durable_entry
|
||||
|
||||
self._routing_db_loaded = True
|
||||
self._routing_fallback_baseline = None
|
||||
|
||||
def _snapshot_routing_locked(self) -> tuple[Dict[str, Any], int]:
|
||||
"""Capture immutable routing data and a monotonic generation."""
|
||||
self._reconcile_recovered_routing_locked()
|
||||
return (
|
||||
{key: entry.to_dict() for key, entry in self._entries.items()},
|
||||
self._next_routing_generation_locked(),
|
||||
)
|
||||
|
||||
def _persist_routing_data(self, data: Dict[str, Any], generation: int) -> None:
|
||||
"""Serialize all whole-index writers through one durable write lock."""
|
||||
with self._lazy("_save_lock", threading.Lock):
|
||||
if generation <= getattr(self, "_persisted_routing_generation", 0):
|
||||
return
|
||||
# Fold in fast upserts numbered above this snapshot: they were
|
||||
# serialized after us and a delayed full rewrite must not regress them.
|
||||
fast_persisted = getattr(self, "_fast_persisted_entries", None)
|
||||
if fast_persisted:
|
||||
for key, (revision, entry_json) in fast_persisted.items():
|
||||
if revision > generation:
|
||||
data[key] = json.loads(entry_json)
|
||||
db_saved = False
|
||||
replacer = self._routing_db_method("replace_gateway_routing_entries")
|
||||
if replacer is not None:
|
||||
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:
|
||||
try:
|
||||
self._save_sessions_json(data)
|
||||
except Exception as exc:
|
||||
if not db_saved:
|
||||
raise
|
||||
# state.db is authoritative. A failed legacy mirror must not
|
||||
# report the already-committed primary write as failed.
|
||||
logger.warning(
|
||||
"gateway.session: sessions.json mirror save failed "
|
||||
"after state.db commit: %s",
|
||||
exc,
|
||||
)
|
||||
self._persisted_routing_generation = generation
|
||||
# This rewrite supersedes fast records at or below its
|
||||
# generation; newer ones stay for the next delayed full writer.
|
||||
if fast_persisted:
|
||||
for key in [
|
||||
k for k, (rev, _) in fast_persisted.items()
|
||||
if rev <= generation
|
||||
]:
|
||||
del fast_persisted[key]
|
||||
|
||||
def _save_sessions_json(self, data: Dict[str, Any]) -> None:
|
||||
"""Write the legacy sessions.json mirror of the routing index."""
|
||||
import tempfile
|
||||
self.sessions_dir.mkdir(parents=True, exist_ok=True)
|
||||
sessions_file = self.sessions_dir / "sessions.json"
|
||||
|
||||
# Self-documenting sentinel; "_" keys are skipped on load. Ordered
|
||||
# first so it renders at the top of the file.
|
||||
data = {
|
||||
"_README": (
|
||||
"LEGACY MIRROR of the gateway routing index (the primary copy "
|
||||
"lives in the gateway_routing table in ~/.hermes/state.db). "
|
||||
"Maps messaging session keys (agent:main:<platform>:...) to "
|
||||
"active session IDs. This is NOT the session list. ALL "
|
||||
"sessions (CLI, TUI, and gateway) live in ~/.hermes/state.db "
|
||||
"and are shown by `hermes sessions list` and `/sessions`. "
|
||||
"Disable this file with `gateway.write_sessions_json: false` "
|
||||
"in config.yaml."
|
||||
),
|
||||
**data,
|
||||
}
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
dir=str(self.sessions_dir), suffix=".tmp", prefix=".sessions_"
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
atomic_replace(tmp_path, sessions_file)
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError as e:
|
||||
logger.debug("Could not remove temp file %s: %s", tmp_path, e)
|
||||
raise
|
||||
|
||||
def _save_entries(self) -> None:
|
||||
"""Snapshot latest state under ``_lock`` and persist after releasing it."""
|
||||
with self._lock:
|
||||
data, generation = self._snapshot_routing_locked()
|
||||
self._persist_routing_data(data, generation)
|
||||
|
||||
def _save_entry(
|
||||
self,
|
||||
session_key: str,
|
||||
*,
|
||||
entry_data: Optional[Dict[str, Any]] = None,
|
||||
lock_held: bool = False,
|
||||
) -> None:
|
||||
"""Persist ONE routing entry via UPSERT — the per-turn fast path
|
||||
(a full rewrite fsyncs a multi-MB sessions.json, ~50ms at ~1100 keys).
|
||||
|
||||
Invariants: the key -> session_id mapping never changes here —
|
||||
structural transitions (create/recover/reset/switch/prune/heal) use the
|
||||
full rewrite, which also refreshes the sessions.json mirror (it may lag
|
||||
in metadata only; state.db stays primary). The entry is serialized under
|
||||
``_lock`` with a revision from the shared routing generation counter
|
||||
(higher == same-or-newer); under ``_save_lock`` the upsert is skipped if
|
||||
a full snapshot or a newer fast save of this key already persisted (the
|
||||
reverse case lives in ``_persist_routing_data``). No DB or a failed
|
||||
upsert falls back to the full rewrite so DB-less installs stay durable.
|
||||
``entry_data`` persists a candidate BEFORE it is published to the live
|
||||
entry (failure-atomic transitions); the fallback carries the same candidate.
|
||||
"""
|
||||
def _capture() -> Optional[tuple[str, int, Optional[Dict[str, Any]]]]:
|
||||
entry = self._entries.get(session_key)
|
||||
if entry is None:
|
||||
return None
|
||||
serialized_entry = (
|
||||
dict(entry_data) if entry_data is not None else entry.to_dict()
|
||||
)
|
||||
entry_json = json.dumps(serialized_entry)
|
||||
revision = self._next_routing_generation_locked()
|
||||
# The O(n) full snapshot is deferred to the fallback branch.
|
||||
return entry_json, revision, serialized_entry if entry_data is not None else None
|
||||
|
||||
if lock_held:
|
||||
captured = _capture()
|
||||
else:
|
||||
with self._lock:
|
||||
captured = _capture()
|
||||
if captured is None:
|
||||
return
|
||||
entry_json, revision, candidate_entry = captured
|
||||
saver = self._routing_db_method("save_gateway_routing_entry")
|
||||
if saver is not None:
|
||||
try:
|
||||
with self._lazy("_save_lock", threading.Lock):
|
||||
if getattr(self, "_persisted_routing_generation", 0) >= revision:
|
||||
return
|
||||
fast_persisted = self._lazy("_fast_persisted_entries", dict)
|
||||
persisted = fast_persisted.get(session_key)
|
||||
if persisted is not None and persisted[0] >= revision:
|
||||
return
|
||||
saver(session_key, entry_json, scope=self._routing_scope())
|
||||
fast_persisted[session_key] = (revision, entry_json)
|
||||
return
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"gateway.session: single-entry routing save failed for %r "
|
||||
"(%s); falling back to full index rewrite",
|
||||
session_key, exc,
|
||||
)
|
||||
if candidate_entry is not None:
|
||||
# Full-snapshot fallback carrying the candidate transition.
|
||||
def _snapshot() -> Dict[str, Any]:
|
||||
return {key: current.to_dict() for key, current in self._entries.items()}
|
||||
if lock_held:
|
||||
fallback_data = _snapshot()
|
||||
else:
|
||||
with self._lock:
|
||||
fallback_data = _snapshot()
|
||||
fallback_data[session_key] = candidate_entry
|
||||
self._persist_routing_data(fallback_data, revision)
|
||||
else:
|
||||
self._save_entries()
|
||||
Reference in New Issue
Block a user