mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 08:28:36 +00:00
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:
parent
74270bb8a8
commit
2af45945e2
@ -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 --------------------------------------------------------
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user