fix(session): fsync sessions on graceful shutdown to prevent data loss

On filesystems with write-back caching (rclone VFS, NFS, FUSE mounts)
the OS page cache may buffer recent session writes. If the process is
killed before the cache flushes, the most recent conversation turns are
silently lost — causing the agent to "forget" recent context and
respond to stale history on the next startup.

Changes:

- session/manager.py: add fsync=True option to save() that flushes the
  file and its parent directory to durable storage. Add flush_all() that
  re-saves every cached session with fsync. Default save() behavior is
  unchanged (no fsync) to avoid performance regression in normal
  operation.

- cli/commands.py: call agent.sessions.flush_all() in the gateway
  shutdown finally block, after stopping heartbeat/cron/channels.

- tests/session/test_session_fsync.py: 8 tests covering fsync flag
  behavior, flush_all with empty/multiple/errored sessions, and
  data survival across simulated process restart.

- tests/cli/test_commands.py: add sessions attribute to _FakeAgentLoop
  so the gateway health endpoint test passes with the new shutdown
  flush.
This commit is contained in:
hussein1362
2026-04-22 13:19:53 +08:00
committed by Xubin Ren
parent ef8bbab7b3
commit 512bf59b3c
5 changed files with 173 additions and 2 deletions
+6
View File
@@ -946,6 +946,12 @@ def _run_gateway(
cron.stop()
agent.stop()
await channels.stop_all()
# Flush all cached sessions to durable storage before exit.
# This prevents data loss on filesystems with write-back
# caching (rclone VFS, NFS, FUSE mounts, etc.).
flushed = agent.sessions.flush_all()
if flushed:
logger.info("Shutdown: flushed {} session(s) to disk", flushed)
asyncio.run(run())
+37 -2
View File
@@ -262,8 +262,16 @@ class SessionManager:
"messages": session.messages,
}
def save(self, session: Session) -> None:
"""Save a session to disk atomically."""
def save(self, session: Session, *, fsync: bool = False) -> None:
"""Save a session to disk atomically.
When *fsync* is ``True`` the final file and its parent directory are
explicitly flushed to durable storage. This is intentionally off by
default (the OS page-cache is sufficient for normal operation) but
should be enabled during graceful shutdown so that filesystems with
write-back caching (e.g. rclone VFS, NFS, FUSE mounts) do not lose
the most recent writes.
"""
path = self._get_session_path(session.key)
tmp_path = path.with_suffix(".jsonl.tmp")
@@ -280,14 +288,41 @@ class SessionManager:
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
for msg in session.messages:
f.write(json.dumps(msg, ensure_ascii=False) + "\n")
if fsync:
f.flush()
os.fsync(f.fileno())
os.replace(tmp_path, path)
if fsync:
# fsync the directory so the rename is durable.
fd = os.open(str(path.parent), os.O_RDONLY)
try:
os.fsync(fd)
finally:
os.close(fd)
except BaseException:
tmp_path.unlink(missing_ok=True)
raise
self._cache[session.key] = session
def flush_all(self) -> int:
"""Re-save every cached session with fsync for durable shutdown.
Returns the number of sessions flushed. Errors on individual
sessions are logged but do not prevent other sessions from being
flushed.
"""
flushed = 0
for key, session in list(self._cache.items()):
try:
self.save(session, fsync=True)
flushed += 1
except Exception:
logger.warning("Failed to flush session {}", key, exc_info=True)
return flushed
def invalidate(self, key: str) -> None:
"""Remove a session from the in-memory cache."""
self._cache.pop(key, None)
+5
View File
@@ -1288,10 +1288,15 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
async def run(self) -> None:
return None
class _FakeSessionManager:
def flush_all(self) -> int:
return 0
class _FakeAgentLoop:
def __init__(self, **_kwargs) -> None:
self.model = "test-model"
self.dream = _FakeDream()
self.sessions = _FakeSessionManager()
async def run(self) -> None:
await asyncio.Event().wait()
View File
+125
View File
@@ -0,0 +1,125 @@
"""Tests for session fsync and flush_all on graceful shutdown."""
from __future__ import annotations
from pathlib import Path
from unittest.mock import patch
import pytest
from nanobot.session.manager import SessionManager
@pytest.fixture
def sessions_dir(tmp_path: Path) -> Path:
d = tmp_path / "sessions"
d.mkdir()
return tmp_path
@pytest.fixture
def manager(sessions_dir: Path) -> SessionManager:
return SessionManager(workspace=sessions_dir)
class TestSaveFsync:
"""Verify that save(fsync=True) calls os.fsync."""
def test_save_without_fsync_does_not_call_fsync(self, manager: SessionManager):
session = manager.get_or_create("test:no-fsync")
session.add_message("user", "hello")
with patch("os.fsync") as mock_fsync:
manager.save(session, fsync=False)
mock_fsync.assert_not_called()
def test_save_with_fsync_calls_fsync(self, manager: SessionManager):
session = manager.get_or_create("test:with-fsync")
session.add_message("user", "hello")
with patch("os.fsync") as mock_fsync:
manager.save(session, fsync=True)
# Should be called twice: once for the file, once for the directory
assert mock_fsync.call_count == 2
def test_save_default_no_fsync(self, manager: SessionManager):
"""Default save() should not fsync (backward compat)."""
session = manager.get_or_create("test:default")
session.add_message("user", "hello")
with patch("os.fsync") as mock_fsync:
manager.save(session)
mock_fsync.assert_not_called()
class TestFlushAll:
"""Verify flush_all re-saves all cached sessions with fsync."""
def test_flush_all_empty_cache(self, manager: SessionManager):
assert manager.flush_all() == 0
def test_flush_all_saves_cached_sessions(self, manager: SessionManager):
s1 = manager.get_or_create("test:session-1")
s1.add_message("user", "msg 1")
manager.save(s1)
s2 = manager.get_or_create("test:session-2")
s2.add_message("user", "msg 2")
manager.save(s2)
flushed = manager.flush_all()
assert flushed == 2
def test_flush_all_uses_fsync(self, manager: SessionManager):
session = manager.get_or_create("test:fsync-check")
session.add_message("user", "important")
manager.save(session)
with patch("os.fsync") as mock_fsync:
manager.flush_all()
# file fsync + directory fsync
assert mock_fsync.call_count == 2
def test_flush_all_continues_on_error(self, manager: SessionManager):
"""One broken session should not prevent others from flushing."""
s1 = manager.get_or_create("test:good")
s1.add_message("user", "ok")
manager.save(s1)
s2 = manager.get_or_create("test:bad")
s2.add_message("user", "ok")
manager.save(s2)
original_save = manager.save
call_count = {"n": 0}
def patched_save(session, *, fsync=False):
call_count["n"] += 1
if session.key == "test:bad":
raise OSError("disk on fire")
original_save(session, fsync=fsync)
manager.save = patched_save
flushed = manager.flush_all()
# One succeeded, one failed — flush_all returns successful count
assert flushed == 1
assert call_count["n"] == 2
def test_flush_all_data_survives_reload(self, sessions_dir: Path):
"""Data flushed by flush_all should survive a fresh SessionManager load."""
mgr1 = SessionManager(workspace=sessions_dir)
session = mgr1.get_or_create("test:persist")
session.add_message("user", "remember this")
session.add_message("assistant", "noted")
mgr1.save(session)
mgr1.flush_all()
# Simulate process restart — new manager, cold cache
mgr2 = SessionManager(workspace=sessions_dir)
reloaded = mgr2.get_or_create("test:persist")
history = reloaded.get_history(max_messages=100)
assert len(history) == 2
assert history[0]["content"] == "remember this"
assert history[1]["content"] == "noted"