fix(serve): bounded flush-on-SIGTERM + periodic incremental session flush
A hermes serve killed mid-update lost every un-flushed in-memory session (#94724 item 2, reported by @ruangraung): the next RPC failed with 'session-scoped RPC rejected: not in memory (detached/reaped runtime)' and no store held the transcript. #95576 made serves survive future updates; this closes the kill path itself: - install chaining SIGTERM/SIGINT handlers (hermes serve / dashboard startup, before uvicorn's capture_signals) that first persist in-memory session transcripts to state.db — bounded by HERMES_TUI_EXIT_FLUSH_BUDGET_S (default 5s, daemon worker + join) so a hung SQLite write can never block exit - _shutdown_sessions (atexit) runs the same bounded flush FIRST, before the slow per-session teardown a supervisor may SIGKILL mid-way - the idle-reaper scan piggybacks a periodic incremental flush (marker-deduped agent._persist_session, running sessions skipped) so even a SIGKILL loses at most one flush interval — no new timer subsystem Refs #94724
This commit is contained in:
Teknium
@@ -19770,6 +19770,19 @@ def start_server(
|
||||
)
|
||||
server = uvicorn.Server(config)
|
||||
|
||||
# Flush-on-kill guard (#94724 item 2): install chaining SIGTERM/SIGINT
|
||||
# handlers that first persist in-memory session transcripts to state.db
|
||||
# (bounded, best-effort) before the normal shutdown story runs. Installed
|
||||
# on the main thread BEFORE uvicorn's capture_signals() so uvicorn saves
|
||||
# these as the "original" handlers and re-raises into them after its own
|
||||
# graceful shutdown — kills outside the serve window are covered too.
|
||||
try:
|
||||
from tui_gateway.server import install_exit_flush_signal_handlers
|
||||
|
||||
install_exit_flush_signal_handlers()
|
||||
except Exception as exc:
|
||||
_log.debug("exit-flush signal handlers not installed: %s", exc)
|
||||
|
||||
# ── #93608: machine-readable port-conflict detection ──────────────
|
||||
# uvicorn's own bind_socket() would catch the EADDRINUSE and exit 1
|
||||
# with a bare ERROR line — indistinguishable from "backend broken".
|
||||
|
||||
200
tests/tui_gateway/test_serve_exit_flush.py
Normal file
200
tests/tui_gateway/test_serve_exit_flush.py
Normal file
@@ -0,0 +1,200 @@
|
||||
"""A killed ``hermes serve`` must not lose in-memory session transcripts.
|
||||
|
||||
Regression for #94724 (item 2, @ruangraung): a serve terminated mid-update
|
||||
lost every un-flushed in-memory session — the next RPC failed with
|
||||
"session-scoped RPC rejected: not in memory (detached/reaped runtime)" and no
|
||||
store held the transcript. #95576 made serves survive *future* updates; this
|
||||
covers the kill path itself:
|
||||
|
||||
* SIGTERM/SIGINT first flush in-memory sessions to state.db (bounded,
|
||||
best-effort, chained to the previously installed handler so uvicorn's
|
||||
graceful shutdown still runs).
|
||||
* The idle-reaper tick piggybacks a periodic incremental flush so even a
|
||||
SIGKILL loses at most one flush interval.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import signal
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from tui_gateway import server
|
||||
|
||||
|
||||
class _FlushAgent:
|
||||
"""Minimal agent exposing the real ``_persist_session`` flush contract."""
|
||||
|
||||
def __init__(self, messages=None):
|
||||
self.session_id = "flush-agent"
|
||||
self.flush_calls: list[list] = []
|
||||
self._session_messages = (
|
||||
messages
|
||||
if messages is not None
|
||||
else [{"role": "user", "content": "unflushed turn"}]
|
||||
)
|
||||
|
||||
def _persist_session(self, messages, conversation_history=None):
|
||||
self.flush_calls.append(list(messages))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registered_session():
|
||||
"""Register a fake in-memory session; always deregister on exit."""
|
||||
registered: list[str] = []
|
||||
|
||||
def _register(sid: str, agent, **extra):
|
||||
session = {"agent": agent, "session_key": sid, "running": False}
|
||||
session.update(extra)
|
||||
with server._sessions_lock:
|
||||
server._sessions[sid] = session
|
||||
registered.append(sid)
|
||||
return session
|
||||
|
||||
yield _register
|
||||
|
||||
with server._sessions_lock:
|
||||
for sid in registered:
|
||||
server._sessions.pop(sid, None)
|
||||
|
||||
|
||||
def _restore_signal_state(prev_handlers):
|
||||
for signum, handler in prev_handlers.items():
|
||||
signal.signal(signum, handler)
|
||||
server._exit_flush_prev_handlers.clear()
|
||||
server._exit_flush_handlers_installed = False
|
||||
|
||||
|
||||
def test_sigterm_flushes_populated_session_into_state_db(
|
||||
registered_session, tmp_path, monkeypatch
|
||||
):
|
||||
"""A populated in-memory session survives a SIGTERM into state.db."""
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
||||
from hermes_state import SessionDB
|
||||
|
||||
db = SessionDB(db_path=tmp_path / "state.db")
|
||||
sid = "sess-sigterm-flush"
|
||||
db.create_session(sid, source="tui")
|
||||
|
||||
class _DbAgent(_FlushAgent):
|
||||
def _persist_session(self, messages, conversation_history=None):
|
||||
super()._persist_session(messages, conversation_history)
|
||||
for msg in messages:
|
||||
if msg.get("_db_persisted"):
|
||||
continue
|
||||
db.append_message(sid, msg["role"], msg["content"])
|
||||
msg["_db_persisted"] = True
|
||||
|
||||
agent = _DbAgent(messages=[{"role": "user", "content": "survive the kill"}])
|
||||
registered_session(sid, agent)
|
||||
|
||||
chained = {"called": False}
|
||||
|
||||
def _prev_handler(signum, frame):
|
||||
chained["called"] = True
|
||||
|
||||
prev = {signal.SIGTERM: signal.signal(signal.SIGTERM, _prev_handler)}
|
||||
try:
|
||||
assert server.install_exit_flush_signal_handlers() is True
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
# The handler runs synchronously on the main thread at the next
|
||||
# bytecode boundary; poll briefly for robustness.
|
||||
deadline = time.monotonic() + 5.0
|
||||
while not chained["called"] and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
finally:
|
||||
_restore_signal_state(prev)
|
||||
|
||||
assert chained["called"], "previous SIGTERM handler must still be chained"
|
||||
assert agent.flush_calls, "SIGTERM must flush in-memory sessions"
|
||||
rows = db.get_messages(sid)
|
||||
assert any("survive the kill" in str(r.get("content", "")) for r in rows)
|
||||
|
||||
|
||||
def test_exit_flush_is_bounded(registered_session):
|
||||
"""A hung persist must never block exit longer than the budget."""
|
||||
|
||||
class _HangingAgent(_FlushAgent):
|
||||
def _persist_session(self, messages, conversation_history=None):
|
||||
time.sleep(5.0)
|
||||
|
||||
registered_session("sess-hang", _HangingAgent())
|
||||
|
||||
start = time.monotonic()
|
||||
server._flush_sessions_before_exit(budget_s=0.3)
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 2.0, f"exit flush blocked {elapsed:.1f}s past its budget"
|
||||
|
||||
|
||||
def test_shutdown_sessions_flushes_before_teardown(monkeypatch):
|
||||
"""The atexit path persists transcripts BEFORE slow per-session teardown."""
|
||||
order: list[str] = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
server, "_release_gateway_wake_owner", lambda: None, raising=False
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_flush_sessions_before_exit",
|
||||
lambda budget_s=None: order.append("flush") or 0,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_close_session_by_id",
|
||||
lambda sid, **kw: order.append(f"close:{sid}"),
|
||||
)
|
||||
with server._sessions_lock:
|
||||
server._sessions["sess-order"] = {"agent": None, "session_key": "sess-order"}
|
||||
try:
|
||||
server._shutdown_sessions()
|
||||
finally:
|
||||
with server._sessions_lock:
|
||||
server._sessions.pop("sess-order", None)
|
||||
|
||||
assert order and order[0] == "flush"
|
||||
assert "close:sess-order" in order
|
||||
|
||||
|
||||
def test_periodic_flush_respects_interval_with_fake_clock(
|
||||
registered_session, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(server, "_INCREMENTAL_FLUSH_INTERVAL_S", 300.0)
|
||||
agent = _FlushAgent()
|
||||
registered_session("sess-interval", agent)
|
||||
|
||||
assert server._flush_dirty_sessions(now=1_000.0) == 1
|
||||
assert len(agent.flush_calls) == 1
|
||||
|
||||
# Within the interval: no re-flush.
|
||||
assert server._flush_dirty_sessions(now=1_000.0 + 299.0) == 0
|
||||
assert len(agent.flush_calls) == 1
|
||||
|
||||
# Past the interval: flushes again — SIGKILL loses at most one interval.
|
||||
assert server._flush_dirty_sessions(now=1_000.0 + 301.0) == 1
|
||||
assert len(agent.flush_calls) == 2
|
||||
|
||||
|
||||
def test_periodic_flush_skips_running_sessions(registered_session, monkeypatch):
|
||||
"""Mid-turn sessions are the turn thread's to persist — never race them."""
|
||||
monkeypatch.setattr(server, "_INCREMENTAL_FLUSH_INTERVAL_S", 300.0)
|
||||
agent = _FlushAgent()
|
||||
registered_session("sess-running", agent, running=True)
|
||||
|
||||
assert server._flush_dirty_sessions(now=1_000.0) == 0
|
||||
assert agent.flush_calls == []
|
||||
|
||||
|
||||
def test_idle_reaper_scan_piggybacks_incremental_flush(monkeypatch):
|
||||
"""The existing reaper tick drives the flush — no new timer subsystem."""
|
||||
called = {"flush": 0}
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_flush_dirty_sessions",
|
||||
lambda now=None: called.__setitem__("flush", called["flush"] + 1) or 0,
|
||||
)
|
||||
monkeypatch.setattr(server, "_enforce_session_cap", lambda: None)
|
||||
monkeypatch.setattr(server, "_reclaim_orphaned_leases", lambda: None)
|
||||
server._reap_idle_sessions()
|
||||
assert called["flush"] == 1
|
||||
@@ -1441,6 +1441,15 @@ def _close_sessions_for_transport(
|
||||
|
||||
|
||||
def _shutdown_sessions() -> None:
|
||||
# Durable-first (#94724 item 2): persist every session's un-flushed
|
||||
# transcript within a bounded budget BEFORE the slow per-session
|
||||
# teardown below (plugin hooks, memory commit, delegation interrupts,
|
||||
# agent.close). A supervisor that SIGKILLs a slow shutdown mid-way can
|
||||
# then no longer lose the transcripts — the flush already landed.
|
||||
try:
|
||||
_flush_sessions_before_exit()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
_release_gateway_wake_owner()
|
||||
except Exception:
|
||||
@@ -1462,6 +1471,176 @@ _SESSION_TTL_S = max(0.0, _SESSION_TTL_S)
|
||||
_REAPER_SCAN_S = 300.0
|
||||
|
||||
|
||||
# ── Flush-on-kill + periodic incremental flush (#94724 item 2) ───────────
|
||||
# A `hermes serve` killed mid-update used to lose every un-flushed in-memory
|
||||
# session: the next RPC failed with "session-scoped RPC rejected: not in
|
||||
# memory (detached/reaped runtime)" and NO store held the transcript. #95576
|
||||
# made serves survive *future* updates; this closes the kill path itself:
|
||||
# (a) SIGTERM/SIGINT run a bounded, best-effort flush of in-memory session
|
||||
# transcripts to state.db BEFORE the normal shutdown path, chained to
|
||||
# whatever handler was installed before (uvicorn's included);
|
||||
# (b) the idle-reaper scan piggybacks a periodic incremental flush so even
|
||||
# a SIGKILL loses at most one flush interval.
|
||||
try:
|
||||
_EXIT_FLUSH_BUDGET_S = float(
|
||||
os.environ.get("HERMES_TUI_EXIT_FLUSH_BUDGET_S") or 5.0
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
_EXIT_FLUSH_BUDGET_S = 5.0
|
||||
_EXIT_FLUSH_BUDGET_S = max(0.0, _EXIT_FLUSH_BUDGET_S)
|
||||
|
||||
try:
|
||||
_INCREMENTAL_FLUSH_INTERVAL_S = float(
|
||||
os.environ.get("HERMES_TUI_SESSION_FLUSH_INTERVAL_S") or _REAPER_SCAN_S
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
_INCREMENTAL_FLUSH_INTERVAL_S = _REAPER_SCAN_S
|
||||
_INCREMENTAL_FLUSH_INTERVAL_S = max(0.0, _INCREMENTAL_FLUSH_INTERVAL_S)
|
||||
|
||||
|
||||
def _flush_session_messages(session: dict | None) -> bool:
|
||||
"""Best-effort durable flush of one session's in-memory transcript.
|
||||
|
||||
Rides ``agent._persist_session`` — the same marker-deduped persist
|
||||
contract ``_finalize_session`` uses (#13121) — so repeated calls only
|
||||
write genuinely-unflushed messages and never duplicate durable rows.
|
||||
"""
|
||||
if not session:
|
||||
return False
|
||||
agent = session.get("agent")
|
||||
if agent is None or not hasattr(agent, "_persist_session"):
|
||||
return False
|
||||
snapshot = getattr(agent, "_session_messages", None)
|
||||
if not snapshot:
|
||||
return False
|
||||
try:
|
||||
agent._persist_session(snapshot)
|
||||
return True
|
||||
except Exception:
|
||||
logger.debug("incremental session flush failed", exc_info=True)
|
||||
return False
|
||||
|
||||
|
||||
def _flush_dirty_sessions(now: float | None = None) -> int:
|
||||
"""Periodic incremental flush, driven by the idle-reaper scan.
|
||||
|
||||
Skips ``running`` sessions: the turn thread owns mid-turn persistence
|
||||
(it already flushes at every persist point) and
|
||||
``_drop_trailing_empty_response_scaffolding`` mutates the live message
|
||||
list, so racing an in-flight turn from the reaper thread is never safe.
|
||||
Idle/detached sessions — precisely the ones a kill strands — are flushed
|
||||
at most once per ``_INCREMENTAL_FLUSH_INTERVAL_S``. ``now`` is injectable
|
||||
for tests (monotonic clock).
|
||||
"""
|
||||
if _INCREMENTAL_FLUSH_INTERVAL_S <= 0:
|
||||
return 0
|
||||
if now is None:
|
||||
now = time.monotonic()
|
||||
with _sessions_lock:
|
||||
sessions = list(_sessions.values())
|
||||
flushed = 0
|
||||
for session in sessions:
|
||||
if not isinstance(session, dict) or session.get("running"):
|
||||
continue
|
||||
last = float(session.get("_last_incremental_flush") or 0.0)
|
||||
if last and (now - last) < _INCREMENTAL_FLUSH_INTERVAL_S:
|
||||
continue
|
||||
if _flush_session_messages(session):
|
||||
flushed += 1
|
||||
session["_last_incremental_flush"] = now
|
||||
return flushed
|
||||
|
||||
|
||||
def _flush_sessions_before_exit(budget_s: float | None = None) -> int:
|
||||
"""Bounded flush of ALL in-memory sessions on the way out.
|
||||
|
||||
Runs on a daemon worker joined with the budget so a hung SQLite write
|
||||
can never block exit longer than ``HERMES_TUI_EXIT_FLUSH_BUDGET_S``
|
||||
(default 5s). Running sessions are included — the process is dying, so
|
||||
a best-effort partial transcript beats guaranteed loss.
|
||||
"""
|
||||
budget = _EXIT_FLUSH_BUDGET_S if budget_s is None else max(0.0, budget_s)
|
||||
if budget <= 0:
|
||||
return 0
|
||||
result = {"flushed": 0}
|
||||
|
||||
def _run() -> None:
|
||||
deadline = time.monotonic() + budget
|
||||
with _sessions_lock:
|
||||
sessions = list(_sessions.values())
|
||||
for session in sessions:
|
||||
if time.monotonic() >= deadline:
|
||||
break
|
||||
if _flush_session_messages(session):
|
||||
result["flushed"] += 1
|
||||
|
||||
worker = threading.Thread(target=_run, daemon=True, name="hermes-exit-flush")
|
||||
worker.start()
|
||||
worker.join(budget)
|
||||
return result["flushed"]
|
||||
|
||||
|
||||
_exit_flush_prev_handlers: dict[int, Any] = {}
|
||||
_exit_flush_handlers_installed = False
|
||||
|
||||
|
||||
def _handle_exit_flush_signal(signum, frame) -> None:
|
||||
"""Flush in-memory sessions, then hand off to the prior handler.
|
||||
|
||||
Chaining preserves the pre-existing signal story (uvicorn's graceful
|
||||
shutdown, a supervisor's handler, or the default terminate disposition)
|
||||
— this handler only *prepends* a bounded durable flush.
|
||||
"""
|
||||
try:
|
||||
_flush_sessions_before_exit()
|
||||
except Exception:
|
||||
pass
|
||||
import signal as _signal
|
||||
|
||||
prev = _exit_flush_prev_handlers.get(signum)
|
||||
if callable(prev):
|
||||
prev(signum, frame)
|
||||
return
|
||||
if prev is _signal.SIG_IGN:
|
||||
return
|
||||
# Default disposition: restore it and re-raise so the process still dies
|
||||
# with the correct signal (exit status visible to supervisors).
|
||||
try:
|
||||
_signal.signal(signum, _signal.SIG_DFL)
|
||||
os.kill(os.getpid(), signum)
|
||||
except Exception:
|
||||
raise SystemExit(128 + int(signum)) from None
|
||||
|
||||
|
||||
def install_exit_flush_signal_handlers() -> bool:
|
||||
"""Install chaining SIGTERM/SIGINT flush handlers (main thread only).
|
||||
|
||||
Called by ``hermes serve`` / dashboard startup before uvicorn takes over
|
||||
signals: uvicorn's ``capture_signals()`` saves these as the "original"
|
||||
handlers and restores + re-raises into them after its graceful shutdown,
|
||||
so the flush also covers terminations outside uvicorn's serve window.
|
||||
Idempotent; returns False off the main thread or when installation fails.
|
||||
"""
|
||||
global _exit_flush_handlers_installed
|
||||
if _exit_flush_handlers_installed:
|
||||
return True
|
||||
if threading.current_thread() is not threading.main_thread():
|
||||
return False
|
||||
import signal as _signal
|
||||
|
||||
installed = False
|
||||
for signum in (_signal.SIGTERM, _signal.SIGINT):
|
||||
try:
|
||||
prev = _signal.getsignal(signum)
|
||||
_signal.signal(signum, _handle_exit_flush_signal)
|
||||
_exit_flush_prev_handlers[signum] = prev
|
||||
installed = True
|
||||
except (ValueError, OSError, RuntimeError):
|
||||
continue
|
||||
_exit_flush_handlers_installed = installed
|
||||
return installed
|
||||
|
||||
|
||||
def _transport_is_dead(transport) -> bool:
|
||||
# _detached_ws_transport is the post-WS-disconnect drop sentinel; a session
|
||||
# parked on it has no live client. _stdio_transport is the REAL transport
|
||||
@@ -1491,6 +1670,13 @@ def _session_is_evictable(sid: str, session: dict, now: float) -> bool:
|
||||
|
||||
def _reap_idle_sessions() -> None:
|
||||
now = time.time()
|
||||
# Piggyback the periodic incremental flush on the existing reaper tick
|
||||
# (#94724 item 2) — no new timer subsystem. Even a SIGKILL then loses at
|
||||
# most one flush interval of un-persisted messages.
|
||||
try:
|
||||
_flush_dirty_sessions()
|
||||
except Exception:
|
||||
logger.debug("periodic incremental session flush failed", exc_info=True)
|
||||
with _sessions_lock:
|
||||
victims = [sid for sid, s in _sessions.items() if _session_is_evictable(sid, s, now)]
|
||||
for sid in victims:
|
||||
|
||||
Reference in New Issue
Block a user