From 9a502c19d39a423c94b178acb179747c1c260182 Mon Sep 17 00:00:00 2001 From: Ioodu Date: Sun, 12 Apr 2026 11:57:15 +0800 Subject: [PATCH 1/3] feat(gateway): add crash checkpoint for precise session recovery Persist in-flight agent runs to agent_checkpoints.json so that on restart the gateway can precisely identify interrupted sessions instead of relying solely on the suspend_recently_active time-window heuristic. Sessions found in the checkpoint are suspended and the checkpoint is cleared; the time-window heuristic remains as fallback. --- gateway/run.py | 25 ++++ gateway/session.py | 75 +++++++++++ tests/gateway/test_session_crash_recovery.py | 134 +++++++++++++++++++ 3 files changed, 234 insertions(+) create mode 100644 tests/gateway/test_session_crash_recovery.py diff --git a/gateway/run.py b/gateway/run.py index ac8f763b7fc7f..665ada721a8f4 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -695,6 +695,13 @@ def __init__(self, config: Optional[GatewayConfig] = None): self._busy_ack_ts: Dict[str, float] = {} # last busy-ack timestamp per session (debounce) self._session_run_generation: Dict[str, int] = {} + # Crash-recovery checkpoint — records in-flight agent runs to disk + # so that a gateway restart can detect interrupted sessions. + from gateway.session import SessionCrashCheckpoint + from hermes_constants import HERMES_HOME + _checkpoint_path = os.path.join(HERMES_HOME, "agent_checkpoints.json") + self._crash_checkpoint = SessionCrashCheckpoint(path=_checkpoint_path) + # Cache AIAgent instances per session to preserve prompt caching. # Without this, a new AIAgent is created per message, rebuilding the # system prompt (including memory) every turn — breaking prefix cache @@ -2249,6 +2256,22 @@ async def start(self) -> bool: except Exception: pass else: + # Primary: use the crash checkpoint for precise detection of sessions + # that were in-flight when the gateway crashed. + try: + interrupted = self._crash_checkpoint.get_active_sessions() + if interrupted: + for session_key in interrupted: + self.session_store.suspend_session(session_key) + logger.info( + "Suspended %d interrupted session(s) from crash checkpoint", + len(interrupted), + ) + self._crash_checkpoint.clear() + except Exception as e: + logger.warning("Crash checkpoint recovery failed: %s", e) + + # Fallback: time-window heuristic for sessions not tracked by checkpoint. try: suspended = self.session_store.suspend_recently_active() if suspended: @@ -10691,6 +10714,7 @@ async def track_agent(): ) return self._running_agents[session_key] = agent_holder[0] + self._crash_checkpoint.mark_running(session_key, session_id=session_id) if self._draining: self._update_runtime_status("draining") @@ -11213,6 +11237,7 @@ async def _notify_long_running(): self._release_running_agent_state( session_key, run_generation=run_generation ) + self._crash_checkpoint.mark_completed(session_key) if self._draining: self._update_runtime_status("draining") diff --git a/gateway/session.py b/gateway/session.py index 02d4eb3ed01f8..a85e3e1d30553 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -12,6 +12,7 @@ import logging import os import json +import time import threading import uuid from pathlib import Path @@ -1370,3 +1371,77 @@ def build_session_context( context.updated_at = session_entry.updated_at return context + + +# --------------------------------------------------------------------------- +# Crash-recovery checkpoint +# --------------------------------------------------------------------------- + + +class SessionCrashCheckpoint: + """Persist a record of in-flight agent runs so that a gateway restart + can detect sessions that were interrupted by a crash. + + The checkpoint file (``agent_checkpoints.json``) maps session keys to + a dict with ``session_id`` and ``started_at`` (epoch). Entries are + added when an agent starts and removed on clean completion. If the + gateway crashes, the file retains the entries for runs that never + completed, allowing precise identification of interrupted sessions + on the next startup. + """ + + def __init__(self, path: str): + self.path = path + self._lock = threading.Lock() + + # ── Read / Write helpers ────────────────────────────────────────── + + def _read(self) -> Dict[str, Any]: + """Load checkpoint data from disk. Returns {} on any error.""" + try: + with open(self.path, "r", encoding="utf-8") as f: + return json.load(f) + except (OSError, json.JSONDecodeError): + return {} + + def _write(self, data: Dict[str, Any]) -> None: + """Atomically write checkpoint data to disk.""" + tmp_path = self.path + ".tmp" + with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + os.replace(tmp_path, self.path) + + # ── Public API ──────────────────────────────────────────────────── + + def mark_running(self, session_key: str, *, session_id: str) -> None: + """Record that an agent run has started for *session_key*.""" + with self._lock: + data = self._read() + data[session_key] = { + "session_id": session_id, + "started_at": time.time(), + } + self._write(data) + + def mark_completed(self, session_key: str) -> None: + """Remove the checkpoint entry for *session_key*.""" + with self._lock: + data = self._read() + if session_key in data: + del data[session_key] + self._write(data) + + def get_active_sessions(self) -> Dict[str, Dict[str, Any]]: + """Return all sessions that started but never completed. + + On a clean shutdown this dict is empty. After a crash it + contains exactly the sessions that were in-flight. + """ + with self._lock: + return dict(self._read()) + + def clear(self) -> None: + """Remove all entries — called on clean shutdown.""" + with self._lock: + self._write({}) + diff --git a/tests/gateway/test_session_crash_recovery.py b/tests/gateway/test_session_crash_recovery.py new file mode 100644 index 0000000000000..db3e9fa168d19 --- /dev/null +++ b/tests/gateway/test_session_crash_recovery.py @@ -0,0 +1,134 @@ +"""Tests for session state recovery on gateway crash. + +When the gateway crashes while agent runs are in-flight, a checkpoint +file records which sessions were active. On restart, the gateway reads +the checkpoint to identify interrupted sessions with precision, rather +than relying solely on the blunt suspend_recently_active() time window. +""" + +import json +import os +import tempfile +import time +from pathlib import Path +from unittest.mock import patch, MagicMock + +import pytest + +from gateway.session import SessionCrashCheckpoint + + +# ── Helpers ──────────────────────────────────────────────────────────────── + + +def _tmp_checkpoint_path(tmp_path): + """Return a checkpoint file path inside a temp directory.""" + return str(tmp_path / "agent_checkpoints.json") + + +# ── SessionCrashCheckpoint tests ────────────────────────────────────────── + + +class TestSessionCrashCheckpointWriteRead: + """Checkpoint file is written on agent start and read on restart.""" + + def test_write_creates_file_with_session_entry(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + cp.mark_running("telegram:123:456", session_id="sess_abc") + data = json.loads(Path(cp.path).read_text()) + assert "telegram:123:456" in data + assert data["telegram:123:456"]["session_id"] == "sess_abc" + + def test_mark_completed_removes_entry(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + cp.mark_running("telegram:123:456", session_id="sess_abc") + cp.mark_completed("telegram:123:456") + data = json.loads(Path(cp.path).read_text()) + assert "telegram:123:456" not in data + + def test_read_returns_active_sessions(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + cp.mark_running("telegram:123:456", session_id="sess_abc") + cp.mark_running("discord:789", session_id="sess_def") + cp.mark_completed("discord:789") + active = cp.get_active_sessions() + assert "telegram:123:456" in active + assert "discord:789" not in active + + def test_checkpoint_includes_timestamp(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + before = time.time() + cp.mark_running("telegram:123:456", session_id="sess_abc") + data = json.loads(Path(cp.path).read_text()) + assert data["telegram:123:456"]["started_at"] >= before + + +class TestSessionCrashCheckpointEdgeCases: + """Edge cases for crash checkpoint persistence.""" + + def test_mark_completed_nonexistent_is_noop(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + cp.mark_completed("nonexistent:key") # Should not raise + # No file created since there was nothing to remove + assert not os.path.exists(cp.path) or cp.get_active_sessions() == {} + + def test_file_does_not_exist_on_init(self, tmp_path): + path = str(tmp_path / "nonexistent.json") + cp = SessionCrashCheckpoint(path=path) + assert cp.get_active_sessions() == {} + + def test_corrupted_file_returns_empty(self, tmp_path): + path = _tmp_checkpoint_path(tmp_path) + Path(path).write_text("not valid json{{{") + cp = SessionCrashCheckpoint(path=path) + assert cp.get_active_sessions() == {} + + def test_multiple_running_sessions(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + for i in range(5): + cp.mark_running(f"platform:{i}", session_id=f"sess_{i}") + active = cp.get_active_sessions() + assert len(active) == 5 + + def test_clear_removes_all_entries(self, tmp_path): + cp = SessionCrashCheckpoint(path=_tmp_checkpoint_path(tmp_path)) + cp.mark_running("a:1", session_id="s1") + cp.mark_running("b:2", session_id="s2") + cp.clear() + assert cp.get_active_sessions() == {} + + +class TestCrashRecoveryIntegration: + """Integration: restart detects interrupted sessions from checkpoint.""" + + def test_interrupted_sessions_detected_after_crash(self, tmp_path): + """Simulate: gateway crashes with active agents, then restarts.""" + path = _tmp_checkpoint_path(tmp_path) + + # Phase 1: Gateway running, agents active + cp1 = SessionCrashCheckpoint(path=path) + cp1.mark_running("telegram:100:200", session_id="sess_active1") + cp1.mark_running("discord:300", session_id="sess_active2") + cp1.mark_completed("telegram:100:200") # This one finished + # Simulate crash — checkpoint file remains with discord:300 + + # Phase 2: Gateway restarts, reads checkpoint + cp2 = SessionCrashCheckpoint(path=path) + interrupted = cp2.get_active_sessions() + assert "discord:300" in interrupted + assert "telegram:100:200" not in interrupted + assert interrupted["discord:300"]["session_id"] == "sess_active2" + + def test_no_interrupted_sessions_on_clean_shutdown(self, tmp_path): + """After a clean shutdown, checkpoint should be empty.""" + path = _tmp_checkpoint_path(tmp_path) + + cp = SessionCrashCheckpoint(path=path) + cp.mark_running("telegram:100:200", session_id="sess_1") + cp.mark_completed("telegram:100:200") # Clean completion + # On clean shutdown, clear is called + cp.clear() + + # Restart reads empty checkpoint + cp2 = SessionCrashCheckpoint(path=path) + assert cp2.get_active_sessions() == {} From 60f01f8d028c5e1ce2d5c3c8feb98dbae5ec3630 Mon Sep 17 00:00:00 2001 From: Ioodu Date: Mon, 27 Apr 2026 15:14:24 +0800 Subject: [PATCH 2/3] fix(gateway): harden crash checkpoint after code review - Fix startup ImportError: replace non-existent HERMES_HOME import with module-level _hermes_home variable (gateway was failing to start) - Gate mark_completed on generation ownership: stale runs no longer clear the checkpoint entry when a newer generation owns the slot - Add fsync + unique mkstemp to _write: checkpoint is now durable across power failures and safe under concurrent gateway instances - Add defensive mark_completed on /stop sentinel fast-path for future-proofing if mark_running timing ever shifts --- gateway/run.py | 9 +++++---- gateway/session.py | 22 +++++++++++++++++----- 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index 665ada721a8f4..17fb9c150f476 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -698,8 +698,7 @@ def __init__(self, config: Optional[GatewayConfig] = None): # Crash-recovery checkpoint — records in-flight agent runs to disk # so that a gateway restart can detect interrupted sessions. from gateway.session import SessionCrashCheckpoint - from hermes_constants import HERMES_HOME - _checkpoint_path = os.path.join(HERMES_HOME, "agent_checkpoints.json") + _checkpoint_path = os.path.join(str(_hermes_home), "agent_checkpoints.json") self._crash_checkpoint = SessionCrashCheckpoint(path=_checkpoint_path) # Cache AIAgent instances per session to preserve prompt caching. @@ -3861,6 +3860,7 @@ async def _handle_message(self, event: MessageEvent) -> Optional[str]: if event.get_command() == "stop": # Force-clean the sentinel so the session is unlocked. self._release_running_agent_state(_quick_key) + self._crash_checkpoint.mark_completed(_quick_key) logger.info("HARD STOP (pending) for session %s — sentinel cleared", _quick_key) return "⚡ Force-stopped. The agent was still starting — session unlocked." # Queue the message so it will be picked up after the @@ -11234,10 +11234,11 @@ async def _notify_long_running(): # were unwinding has already installed its own state; this # guard prevents an old run from clobbering it on the way # out. - self._release_running_agent_state( + released = self._release_running_agent_state( session_key, run_generation=run_generation ) - self._crash_checkpoint.mark_completed(session_key) + if released: + self._crash_checkpoint.mark_completed(session_key) if self._draining: self._update_runtime_status("draining") diff --git a/gateway/session.py b/gateway/session.py index a85e3e1d30553..219df23819907 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -1405,11 +1405,23 @@ def _read(self) -> Dict[str, Any]: return {} def _write(self, data: Dict[str, Any]) -> None: - """Atomically write checkpoint data to disk.""" - tmp_path = self.path + ".tmp" - with open(tmp_path, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=2) - os.replace(tmp_path, self.path) + """Atomically write checkpoint data to disk with fsync for crash safety.""" + import tempfile + dir_path = os.path.dirname(self.path) or "." + os.makedirs(dir_path, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(dir=dir_path, suffix=".tmp", prefix=".cp_") + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, self.path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise # ── Public API ──────────────────────────────────────────────────── From cfb07641c9a6c4bfb86fed9b83fe3e2a6331cca5 Mon Sep 17 00:00:00 2001 From: Ioodu Date: Mon, 27 Apr 2026 16:19:08 +0800 Subject: [PATCH 3/3] fix(gateway): address review findings in crash checkpoint - Clear checkpoint on clean shutdown to prevent stale entries accumulating - Add mark_completed after stale-eviction and interrupt-clear paths - Move tempfile import to module level --- gateway/run.py | 18 ++++++++++ gateway/session.py | 2 +- tests/gateway/test_session_crash_recovery.py | 36 ++++++++++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/gateway/run.py b/gateway/run.py index 17fb9c150f476..2b443da6fb817 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -3006,6 +3006,13 @@ def _kill_tool_subprocesses(phase: str) -> None: (_hermes_home / ".clean_shutdown").touch() except Exception: pass + # Clear the crash checkpoint on clean shutdown so that stale + # entries left by /stop or stale-eviction paths do not cause + # false-positive session suspensions if a crash occurs later. + try: + self._crash_checkpoint.clear() + except Exception as _e: + logger.debug("Failed to clear crash checkpoint on clean shutdown: %s", _e) else: logger.info( "Skipping .clean_shutdown marker — drain timed out with " @@ -3637,6 +3644,12 @@ async def _handle_message(self, event: MessageEvent) -> Optional[str]: reason="stale_running_agent_eviction", ) self._release_running_agent_state(_quick_key) + # The generation was just invalidated, so the agent's finally + # block will see released=False and skip mark_completed. Do + # it here so the checkpoint entry is not left as a stale + # false-positive that would cause an unwarranted suspension on + # the next crash restart. + self._crash_checkpoint.mark_completed(_quick_key) if _quick_key in self._running_agents: if event.get_command() == "status": @@ -9099,6 +9112,11 @@ async def _interrupt_and_clear_session( self._pending_messages.pop(session_key, None) if release_running_state: self._release_running_agent_state(session_key) + # The generation was invalidated above, so the interrupted agent's + # finally block will see released=False and skip mark_completed. + # Remove the checkpoint entry here to avoid a stale false-positive + # that would cause an unwarranted suspension on the next crash restart. + self._crash_checkpoint.mark_completed(session_key) def _evict_cached_agent(self, session_key: str) -> None: """Remove a cached agent for a session (called on /new, /model, etc).""" diff --git a/gateway/session.py b/gateway/session.py index 219df23819907..ee41ccabe986f 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -12,6 +12,7 @@ import logging import os import json +import tempfile import time import threading import uuid @@ -1406,7 +1407,6 @@ def _read(self) -> Dict[str, Any]: def _write(self, data: Dict[str, Any]) -> None: """Atomically write checkpoint data to disk with fsync for crash safety.""" - import tempfile dir_path = os.path.dirname(self.path) or "." os.makedirs(dir_path, exist_ok=True) fd, tmp_path = tempfile.mkstemp(dir=dir_path, suffix=".tmp", prefix=".cp_") diff --git a/tests/gateway/test_session_crash_recovery.py b/tests/gateway/test_session_crash_recovery.py index db3e9fa168d19..876800b8b10cf 100644 --- a/tests/gateway/test_session_crash_recovery.py +++ b/tests/gateway/test_session_crash_recovery.py @@ -132,3 +132,39 @@ def test_no_interrupted_sessions_on_clean_shutdown(self, tmp_path): # Restart reads empty checkpoint cp2 = SessionCrashCheckpoint(path=path) assert cp2.get_active_sessions() == {} + + def test_stale_stop_entry_cleared_by_clean_shutdown(self, tmp_path): + """Entries left by /stop (generation-invalidated paths) are purged by + the clean-shutdown clear() so they do not cause false-positive + suspensions after a subsequent crash.""" + path = _tmp_checkpoint_path(tmp_path) + + cp = SessionCrashCheckpoint(path=path) + cp.mark_running("telegram:100:200", session_id="sess_1") + # Simulate /stop: mark_completed is called directly (not via finally block) + cp.mark_completed("telegram:100:200") + + # A second session that was stopped but whose entry was cleaned up + # by _interrupt_and_clear_session + cp.mark_running("discord:300", session_id="sess_2") + cp.mark_completed("discord:300") + + # Clean shutdown calls clear() — any residual entries are removed + cp.clear() + + # After crash on next run, no false-positive suspensions + cp2 = SessionCrashCheckpoint(path=path) + assert cp2.get_active_sessions() == {} + + def test_stale_eviction_entry_does_not_persist(self, tmp_path): + """Entries for stale-evicted agents must be removed so they do not + cause false-positive suspensions on the next crash restart.""" + path = _tmp_checkpoint_path(tmp_path) + + cp = SessionCrashCheckpoint(path=path) + cp.mark_running("telegram:100:200", session_id="sess_1") + # Stale eviction calls mark_completed directly (generation invalidated) + cp.mark_completed("telegram:100:200") + + active = cp.get_active_sessions() + assert "telegram:100:200" not in active