fix(gateway): spool cap-dropped pending transcript messages instead of discarding
When the per-session pending transcript queue hits _MAX_PENDING_PER_SESSION (200) while the session DB is broken, the gateway previously popped the oldest message and discarded it permanently — silent user data loss during live operation (#78182). The on-disk pending spool only ran at shutdown via flush_pending_to_file. Extend that existing spool machinery for runtime drops: - gateway/shutdown_flush.py: add spool_dropped_transcript_message() and drain_transcript_spool(), reusing _get_flush_dir/_write_payload (same atomic-JSON pending_messages/ spool format). recover_pending_to_db() now also replays transcript_cap_drop payloads left over across restarts. - gateway/session.py: on cap eviction, spool the dropped message and log a WARNING that includes the spool path; if spooling fails, degrade to the previous drop-and-warn behavior. On the next fully successful transcript flush for that session, drain and replay spooled messages in drop order; replay failures keep the spool files for the next attempt. - tests/gateway/test_pending_queue_spool.py: drop→spool→drain roundtrip, per-session drain isolation, spool-failure degradation, replay-failure retention, and spool primitive ordering/reason filtering. No new config; extends existing flush_pending_to_file infrastructure per AGENTS.md guidance. Refs #82616, #78182
This commit is contained in:
Teknium
parent
c790ed2a5d
commit
de0f20ff05
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
``<hermes_home>/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", "")
|
||||
|
|
|
|||
|
|
@ -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) == []
|
||||
Loading…
Reference in New Issue