fix(memory): ensure atomic write for history.jsonl

Use temp file + os.replace + fsync to prevent partial writes on crash.
Add tests for atomic write behavior and tmp file cleanup on exception.
This commit is contained in:
yorkhellen 2026-04-28 19:31:13 +08:00 committed by Xubin Ren
parent 74270bb8a8
commit 2af45945e2
2 changed files with 61 additions and 4 deletions

View File

@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio
import json
import os
import re
import weakref
import tiktoken
@ -359,10 +360,18 @@ class MemoryStore:
return None
def _write_entries(self, entries: list[dict[str, Any]]) -> None:
"""Overwrite history.jsonl with the given entries."""
with open(self.history_file, "w", encoding="utf-8") as f:
for entry in entries:
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
"""Overwrite history.jsonl with the given entries (atomic write)."""
tmp_path = self.history_file.with_suffix(self.history_file.suffix + ".tmp")
try:
with open(tmp_path, "w", encoding="utf-8") as f:
for entry in entries:
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
f.flush()
os.fsync(f.fileno())
os.replace(tmp_path, self.history_file)
except BaseException:
tmp_path.unlink(missing_ok=True)
raise
# -- dream cursor --------------------------------------------------------

View File

@ -141,6 +141,54 @@ class TestHistoryWithCursor:
assert len(entries) == 2
assert entries[0]["cursor"] in {4, 5}
def test_write_entries_uses_atomic_write(self, tmp_path):
"""_write_entries uses temp file + os.replace for atomicity."""
store = MemoryStore(tmp_path)
store.append_history("event 1")
store.append_history("event 2")
store.append_history("event 3")
entries = store.read_unprocessed_history(since_cursor=0)
# Monitor temp file existence
tmp_path_obj = store.history_file.with_suffix(".jsonl.tmp")
assert not tmp_path_obj.exists() # Should not exist initially
# Call _write_entries
store._write_entries(entries)
# Temp file should be cleaned up
assert not tmp_path_obj.exists()
# Original file should exist
assert store.history_file.exists()
def test_write_entries_cleans_up_tmp_on_exception(self, tmp_path, monkeypatch):
"""Exception during _write_entries cleans up the temp file."""
store = MemoryStore(tmp_path)
store.append_history("event 1")
entries = store.read_unprocessed_history(since_cursor=0)
tmp_path_obj = store.history_file.with_suffix(".jsonl.tmp")
# Mock os.replace to raise an exception
original_replace = __import__('os').replace
def failing_replace(*args, **kwargs):
raise RuntimeError("Simulated failure")
monkeypatch.setattr('os.replace', failing_replace)
try:
store._write_entries(entries)
assert False, "Should have raised"
except RuntimeError:
pass
# Temp file should be cleaned up
assert not tmp_path_obj.exists()
# Original file should still exist (because replace failed)
assert store.history_file.exists()
class TestAppendHistoryHardCap:
"""append_history has a defensive cap that catches new callers who forgot