diff --git a/gateway/session.py b/gateway/session.py index c760520dce27e..38616c78e1ef0 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -3416,14 +3416,44 @@ class SessionStore: pending = self._dirty_transcripts.setdefault(session_id, []) pending.append(dict(message)) # Cap pending messages per session to avoid unbounded memory - # growth when the DB is persistently broken. Drop the oldest. + # growth when the DB is persistently broken. Spool the evicted + # oldest message to the on-disk pending spool (same machinery + # flush_pending_to_file uses at shutdown) so a runtime cap + # rotation does not silently discard it (#78182); it is + # replayed on the next successful transcript flush. if len(pending) > self._MAX_PENDING_PER_SESSION: - pending.pop(0) - logger.warning( - "Session DB transcript pending queue full for %s " - "(cap=%d); dropping oldest message to make room", - session_id, self._MAX_PENDING_PER_SESSION, - ) + dropped = pending.pop(0) + spool_path = None + try: + from gateway.shutdown_flush import ( + spool_dropped_transcript_message, + ) + spool_path = spool_dropped_transcript_message( + session_id, dropped + ) + except Exception: + spool_path = None + if spool_path is not None: + spooled_sessions = getattr( + self, "_spooled_drop_sessions", None + ) + if spooled_sessions is None: + spooled_sessions = set() + self._spooled_drop_sessions = spooled_sessions + spooled_sessions.add(session_id) + logger.warning( + "Session DB transcript pending queue full for %s " + "(cap=%d); spooled oldest message to %s for replay " + "after DB recovery", + session_id, self._MAX_PENDING_PER_SESSION, spool_path, + ) + else: + logger.warning( + "Session DB transcript pending queue full for %s " + "(cap=%d); dropping oldest message to make room " + "(on-disk spool unavailable)", + session_id, self._MAX_PENDING_PER_SESSION, + ) # Snapshot the first pending message, then release the lock # before the DB write so other sessions are not blocked. msg = pending[0] @@ -3526,10 +3556,43 @@ class SessionStore: if not pending: self._dirty_transcripts.pop(queue_session_id, None) self._transcript_append_failures.pop(session_id, None) - return - msg = pending[0] + queue_empty = True + else: + queue_empty = False + msg = pending[0] + if queue_empty: + # DB write just succeeded and the in-memory backlog is + # clear: replay any cap-dropped messages spooled to disk + # for this session (#78182). + self._drain_spooled_drops(session_id) + return continue + def _drain_spooled_drops(self, session_id: str) -> None: + """Replay cap-dropped spooled transcript messages after DB recovery. + + Best-effort: replay failures keep the spool files for the next + successful flush; nothing here may raise into the caller. + """ + spooled_sessions = getattr(self, "_spooled_drop_sessions", None) + if not spooled_sessions or session_id not in spooled_sessions: + return + try: + from gateway.shutdown_flush import drain_transcript_spool + + _replayed, remaining = drain_transcript_spool( + session_id, + lambda message: self._append_transcript_message( + session_id, message + ), + ) + if not remaining: + spooled_sessions.discard(session_id) + except Exception as exc: + logger.warning( + "Failed to drain transcript spool for %s: %s", session_id, exc + ) + def _append_transcript_message(self, session_id: str, message: Dict[str, Any]) -> None: """Write one transcript row. Caller handles retry queuing.""" self._db.append_message( diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 7b5c57be9e773..14b04cd722c7f 100644 --- a/gateway/shutdown_flush.py +++ b/gateway/shutdown_flush.py @@ -25,6 +25,7 @@ See issue #72680 for the full incident report. from __future__ import annotations +import itertools import json import logging import os @@ -58,8 +59,11 @@ def _fsync_directory(path: Path) -> None: os.close(directory_fd) -def _write_payload(flush_dir: Path, payload: Dict[str, Any]) -> None: - """Atomically write one private, uniquely named recovery payload.""" +def _write_payload(flush_dir: Path, payload: Dict[str, Any]) -> Path: + """Atomically write one private, uniquely named recovery payload. + + Returns the path of the published payload file. + """ from utils import atomic_json_write file_id = uuid.uuid4().hex @@ -77,6 +81,7 @@ def _write_payload(flush_dir: Path, payload: Dict[str, Any]) -> None: # The atomically published file is still the only recovery copy. # Keep it even if this filesystem cannot persist directory entries. logger.debug("Failed to fsync pending-message directory: %s", exc) + return final_path def flush_pending_to_file( @@ -137,6 +142,118 @@ def flush_pending_to_file( return flushed +# Reason tag for transcript messages dropped by the in-memory pending cap +# during live operation (#78182). These payloads carry the full transcript +# message dict so they can be replayed verbatim once the DB recovers. +TRANSCRIPT_CAP_DROP_REASON = "transcript_cap_drop" + + +def spool_dropped_transcript_message( + session_id: str, + message: Dict[str, Any], +) -> Optional[Path]: + """Spool a transcript message evicted by the runtime pending cap. + + Uses the same on-disk pending spool as :func:`flush_pending_to_file` + (one atomic JSON payload per message under + ``/pending_messages/``), so a runtime cap rotation no + longer silently discards user data while the process stays up + (#78182). + + Returns the written spool path, or ``None`` when spooling failed — + callers must degrade to the previous drop-and-log behaviour. + """ + try: + flush_dir = _get_flush_dir() + return _write_payload( + flush_dir, + { + "session_key": session_id, + "reason": TRANSCRIPT_CAP_DROP_REASON, + "ts": int(time.time()), + "seq": next(_TRANSCRIPT_SPOOL_SEQ), + "data": { + "session_id": session_id, + "message": message, + }, + }, + ) + except Exception as exc: + logger.debug( + "Failed to spool cap-dropped transcript message for %s: %s", + session_id, exc, + ) + return None + + +# Monotonic tiebreaker so same-second spool files replay in drop order. +_TRANSCRIPT_SPOOL_SEQ = itertools.count() + + +def drain_transcript_spool(session_id: str, replay) -> tuple[int, int]: + """Replay cap-dropped transcript messages spooled for *session_id*. + + ``replay(message_dict)`` is invoked for each spooled message in drop + order; the spool file is deleted only after a successful replay. On + the first replay failure the drain stops and remaining files are kept + for the next attempt (the DB is likely still unhealthy). + + Returns ``(replayed, remaining)`` — messages replayed and spool files + left behind for a later retry. + """ + try: + flush_dir = _get_flush_dir() + candidates = list(flush_dir.glob("pending-*.json")) + except Exception as exc: + logger.debug("Cannot scan transcript spool: %s", exc) + return 0, 0 + + entries = [] + for path in candidates: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except Exception: + continue + if payload.get("reason") != TRANSCRIPT_CAP_DROP_REASON: + continue + if payload.get("session_key") != session_id: + continue + message = (payload.get("data") or {}).get("message") + if not isinstance(message, dict): + logger.warning( + "Removing structurally invalid transcript spool file %s", path, + ) + path.unlink(missing_ok=True) + continue + entries.append( + (payload.get("ts", 0), payload.get("seq", 0), path.name, path, message) + ) + + replayed = 0 + ordered = sorted(entries, key=lambda e: e[:3]) + remaining = 0 + for idx, (_ts, _seq, _name, path, message) in enumerate(ordered): + try: + replay(message) + except Exception as exc: + logger.warning( + "Replay of spooled transcript message %s for %s failed; " + "keeping spool file for retry: %s", + path, session_id, exc, + ) + remaining = len(ordered) - idx + break + path.unlink(missing_ok=True) + replayed += 1 + + if replayed: + logger.info( + "Replayed %d spooled transcript message(s) for %s after DB recovery", + replayed, session_id, + ) + return replayed, remaining + + def _serialise_value(value: Any) -> Optional[dict]: """Convert a pending message value to a JSON-serialisable dict.""" # MessageEvent objects have a .text attribute and other fields @@ -208,6 +325,29 @@ def recover_pending_to_db( # not automatic DB insertion. Skip them silently. if payload.get("reason") == "shutdown-with-unpersisted-agent-history": continue + # Cap-dropped transcript payloads carry the full message dict + # keyed by session_id — replay directly (#78182). This handles + # spool files that were never drained before a restart. + if payload.get("reason") == TRANSCRIPT_CAP_DROP_REASON: + data = payload.get("data", {}) or {} + spooled_sid = data.get("session_id", "") + message = data.get("message") + if not spooled_sid or not isinstance(message, dict): + logger.warning( + "Cannot recover structurally invalid transcript spool " + "file %s; preserved for manual inspection", + path, + ) + continue + session_db.append_message( + session_id=spooled_sid, + role=message.get("role", "unknown"), + content=message.get("content") or "", + timestamp=message.get("timestamp") or payload.get("ts"), + ) + recovered += 1 + path.unlink(missing_ok=True) + continue session_key = payload.get("session_key", "") data = payload.get("data", {}) text = data.get("text", "") diff --git a/tests/gateway/test_pending_queue_spool.py b/tests/gateway/test_pending_queue_spool.py new file mode 100644 index 0000000000000..cec51cb2780fe --- /dev/null +++ b/tests/gateway/test_pending_queue_spool.py @@ -0,0 +1,223 @@ +"""Regression tests for runtime spool-on-drop of the pending transcript queue. + +When the per-session pending cap (``_MAX_PENDING_PER_SESSION``) forces the +gateway to evict the oldest queued transcript message during live operation, +the message must be spooled to the on-disk pending spool (the same machinery +``flush_pending_to_file`` uses at shutdown) and replayed on the next +successful transcript flush — not silently discarded (#78182, #82616). +""" +import json +import logging +import threading + +import pytest + +from gateway import shutdown_flush +from gateway.session import SessionStore + + +def _make_store(db): + store = object.__new__(SessionStore) + store._db = db + store._transcript_retry_lock = threading.Lock() + store._dirty_transcripts = {} + store._transcript_append_failures = {} + store._fts_rebuild_attempted = True + return store + + +class BrokenThenHealedDb: + """append_message fails while ``broken`` is True, then records rows.""" + + def __init__(self): + self.broken = True + self.rows = [] + + def append_message(self, **kwargs): + if self.broken: + raise RuntimeError("db unavailable") + self.rows.append(kwargs) + + +@pytest.fixture() +def spool_home(tmp_path, monkeypatch): + """Point the pending spool at an isolated HERMES_HOME.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + import hermes_constants + monkeypatch.setattr( + hermes_constants, "get_hermes_home", lambda: tmp_path, raising=True + ) + return tmp_path + + +def _spool_files(home): + d = home / "pending_messages" + return sorted(d.glob("pending-*.json")) if d.exists() else [] + + +class TestSpoolOnDrop: + def test_drop_spool_drain_roundtrip(self, spool_home, caplog, monkeypatch): + # Small cap so the test stays fast. + monkeypatch.setattr(SessionStore, "_MAX_PENDING_PER_SESSION", 5) + db = BrokenThenHealedDb() + store = _make_store(db) + + n_extra = 3 + with caplog.at_level(logging.WARNING, logger="gateway.session"): + for i in range(SessionStore._MAX_PENDING_PER_SESSION + n_extra): + store.append_to_transcript( + "sess-1", {"role": "user", "content": f"msg{i}"} + ) + + # The oldest n_extra messages were evicted — and spooled, not lost. + files = _spool_files(spool_home) + assert len(files) == n_extra + payloads = [json.loads(p.read_text()) for p in files] + assert all( + p["reason"] == shutdown_flush.TRANSCRIPT_CAP_DROP_REASON + for p in payloads + ) + spooled_contents = sorted( + p["data"]["message"]["content"] for p in payloads + ) + assert spooled_contents == ["msg0", "msg1", "msg2"] + + # Drop log escalated to WARNING and includes the spool path. + drop_warnings = [ + r for r in caplog.records + if r.levelno == logging.WARNING and "spooled oldest message" in r.getMessage() + ] + assert len(drop_warnings) == n_extra + assert str(spool_home / "pending_messages") in drop_warnings[0].getMessage() + + # DB heals; the next successful flush drains the backlog AND + # replays the spooled messages in drop order. + db.broken = False + store.append_to_transcript( + "sess-1", {"role": "assistant", "content": "recovered"} + ) + + contents = [r["content"] for r in db.rows] + # All surviving in-memory messages plus the recovery trigger... + for i in range(n_extra, SessionStore._MAX_PENDING_PER_SESSION + n_extra): + assert f"msg{i}" in contents + assert "recovered" in contents + # ...and the previously dropped messages, replayed in drop order. + replayed = [c for c in contents if c in ("msg0", "msg1", "msg2")] + assert replayed == ["msg0", "msg1", "msg2"] + # Spool files consumed after successful replay. + assert _spool_files(spool_home) == [] + # Nothing pending in memory. + assert "sess-1" not in store._dirty_transcripts + + def test_drain_only_touches_own_session(self, spool_home, monkeypatch): + monkeypatch.setattr(SessionStore, "_MAX_PENDING_PER_SESSION", 3) + db = BrokenThenHealedDb() + store = _make_store(db) + + for i in range(SessionStore._MAX_PENDING_PER_SESSION + 1): + store.append_to_transcript("sess-a", {"role": "user", "content": f"a{i}"}) + store.append_to_transcript("sess-b", {"role": "user", "content": f"b{i}"}) + + assert len(_spool_files(spool_home)) == 2 # one drop per session + + db.broken = False + store.append_to_transcript("sess-a", {"role": "user", "content": "go-a"}) + + # Only sess-a's spooled drop was replayed; sess-b's remains on disk. + remaining = [ + json.loads(p.read_text()) for p in _spool_files(spool_home) + ] + assert len(remaining) == 1 + assert remaining[0]["session_key"] == "sess-b" + a_rows = [r["content"] for r in db.rows if r["session_id"] == "sess-a"] + assert "a0" in a_rows + + def test_spool_failure_degrades_to_plain_drop( + self, spool_home, caplog, monkeypatch + ): + """If the spool cannot be written, behave exactly like the old + drop-oldest path: cap enforced, WARNING logged, no crash.""" + monkeypatch.setattr(SessionStore, "_MAX_PENDING_PER_SESSION", 4) + + def _boom(): + raise OSError("disk full") + + monkeypatch.setattr(shutdown_flush, "_get_flush_dir", _boom) + + db = BrokenThenHealedDb() + store = _make_store(db) + + with caplog.at_level(logging.WARNING, logger="gateway.session"): + for i in range(SessionStore._MAX_PENDING_PER_SESSION + 5): + store.append_to_transcript( + "sess-x", {"role": "user", "content": f"msg{i}"} + ) + + pending = store._dirty_transcripts.get("sess-x", []) + assert len(pending) <= SessionStore._MAX_PENDING_PER_SESSION + assert _spool_files(spool_home) == [] + degraded = [ + r for r in caplog.records + if r.levelno == logging.WARNING + and "on-disk spool unavailable" in r.getMessage() + ] + assert len(degraded) == 5 + # No spool bookkeeping means recovery must not attempt a drain. + db.broken = False + store.append_to_transcript("sess-x", {"role": "user", "content": "fin"}) + assert [r["content"] for r in db.rows][-1] == "fin" + + def test_replay_failure_keeps_spool_files(self, spool_home, monkeypatch): + """A failed replay must preserve the spool files for a later retry.""" + monkeypatch.setattr(SessionStore, "_MAX_PENDING_PER_SESSION", 3) + db = BrokenThenHealedDb() + store = _make_store(db) + + for i in range(SessionStore._MAX_PENDING_PER_SESSION + 2): + store.append_to_transcript("sess-r", {"role": "user", "content": f"m{i}"}) + assert len(_spool_files(spool_home)) == 2 + + # DB heals only for live writes; replayed (spooled) rows still fail. + class FlakyDb(BrokenThenHealedDb): + def append_message(self, **kwargs): + if kwargs["content"] in ("m0", "m1", "m2"): + raise RuntimeError("still broken for replays") + self.rows.append(kwargs) + + flaky = FlakyDb() + flaky.broken = False + store._db = flaky + # This append pushes pending over the cap again (dropping/spooling + # m2) before the successful flush triggers the drain. + store.append_to_transcript("sess-r", {"role": "user", "content": "go"}) + + # Spool files survive the failed replay for the next attempt. + assert len(_spool_files(spool_home)) == 3 + assert "sess-r" in getattr(store, "_spooled_drop_sessions", set()) + + +class TestSpoolPrimitives: + def test_drain_skips_other_reasons(self, spool_home): + # A shutdown-format flush file must not be consumed by the drain. + shutdown_flush.flush_pending_to_file({"key1": "hello"}, reason="shutdown") + assert len(_spool_files(spool_home)) == 1 + replayed, remaining = shutdown_flush.drain_transcript_spool( + "key1", lambda m: None + ) + assert replayed == 0 + assert len(_spool_files(spool_home)) == 1 + + def test_roundtrip_order(self, spool_home): + for i in range(3): + shutdown_flush.spool_dropped_transcript_message( + "s", {"role": "user", "content": f"c{i}"} + ) + seen = [] + replayed, remaining = shutdown_flush.drain_transcript_spool( + "s", lambda m: seen.append(m["content"]) + ) + assert replayed == 3 + assert remaining == 0 + assert seen == ["c0", "c1", "c2"] + assert _spool_files(spool_home) == []