fix(session): use atomic writes and add corrupt-file repair
SessionManager.save() previously used bare open("w") which could
truncate the JSONL file if the process crashed mid-write. Now writes
to a .tmp file and atomically replaces via os.replace(), matching the
pattern already used in qq.py.
_load() now attempts _repair() before returning None, recovering
valid lines from partially-written files. 12 new tests cover atomic
save correctness, temp-file cleanup on failure, and repair of
truncated/corrupt JSONL.
cowork-with:opencode(glm-5.1)
This commit is contained in:
aiguozhi123456
committed by
Xubin Ren
Xubin Ren
parent
5d976d79ff
commit
efb04a1712
+83
-13
@@ -1,6 +1,7 @@
|
||||
"""Session management for conversation history."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
@@ -187,24 +188,93 @@ class SessionManager:
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load session {}: {}", key, e)
|
||||
repaired = self._repair(key)
|
||||
if repaired is not None:
|
||||
logger.info("Recovered session {} from corrupt file ({} messages)", key, len(repaired.messages))
|
||||
return repaired
|
||||
|
||||
def _repair(self, key: str) -> Session | None:
|
||||
"""Attempt to recover a session from a corrupt JSONL file."""
|
||||
path = self._get_session_path(key)
|
||||
if not path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
messages: list[dict[str, Any]] = []
|
||||
metadata: dict[str, Any] = {}
|
||||
created_at: datetime | None = None
|
||||
updated_at: datetime | None = None
|
||||
last_consolidated = 0
|
||||
skipped = 0
|
||||
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
if data.get("_type") == "metadata":
|
||||
metadata = data.get("metadata", {})
|
||||
if data.get("created_at"):
|
||||
try:
|
||||
created_at = datetime.fromisoformat(data["created_at"])
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
if data.get("updated_at"):
|
||||
try:
|
||||
updated_at = datetime.fromisoformat(data["updated_at"])
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
last_consolidated = data.get("last_consolidated", 0)
|
||||
else:
|
||||
messages.append(data)
|
||||
|
||||
if skipped:
|
||||
logger.warning("Skipped {} corrupt lines in session {}", skipped, key)
|
||||
|
||||
if not messages and not metadata:
|
||||
return None
|
||||
|
||||
return Session(
|
||||
key=key,
|
||||
messages=messages,
|
||||
created_at=created_at or datetime.now(),
|
||||
updated_at=updated_at or datetime.now(),
|
||||
metadata=metadata,
|
||||
last_consolidated=last_consolidated
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Repair failed for session {}: {}", key, e)
|
||||
return None
|
||||
|
||||
def save(self, session: Session) -> None:
|
||||
"""Save a session to disk."""
|
||||
"""Save a session to disk atomically."""
|
||||
path = self._get_session_path(session.key)
|
||||
tmp_path = path.with_suffix(".jsonl.tmp")
|
||||
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
metadata_line = {
|
||||
"_type": "metadata",
|
||||
"key": session.key,
|
||||
"created_at": session.created_at.isoformat(),
|
||||
"updated_at": session.updated_at.isoformat(),
|
||||
"metadata": session.metadata,
|
||||
"last_consolidated": session.last_consolidated
|
||||
}
|
||||
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
|
||||
for msg in session.messages:
|
||||
f.write(json.dumps(msg, ensure_ascii=False) + "\n")
|
||||
try:
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
metadata_line = {
|
||||
"_type": "metadata",
|
||||
"key": session.key,
|
||||
"created_at": session.created_at.isoformat(),
|
||||
"updated_at": session.updated_at.isoformat(),
|
||||
"metadata": session.metadata,
|
||||
"last_consolidated": session.last_consolidated
|
||||
}
|
||||
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
|
||||
for msg in session.messages:
|
||||
f.write(json.dumps(msg, ensure_ascii=False) + "\n")
|
||||
|
||||
os.replace(tmp_path, path)
|
||||
except BaseException:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
self._cache[session.key] = session
|
||||
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
"""Tests for atomic session save and corrupt-file repair."""
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
|
||||
|
||||
class TestAtomicSave:
|
||||
def test_save_creates_valid_jsonl(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
session = Session(key="test:1")
|
||||
session.add_message("user", "hello")
|
||||
session.add_message("assistant", "hi")
|
||||
|
||||
mgr.save(session)
|
||||
|
||||
path = mgr._get_session_path("test:1")
|
||||
lines = path.read_text(encoding="utf-8").strip().split("\n")
|
||||
assert len(lines) == 3
|
||||
|
||||
meta = json.loads(lines[0])
|
||||
assert meta["_type"] == "metadata"
|
||||
assert meta["key"] == "test:1"
|
||||
|
||||
msg1 = json.loads(lines[1])
|
||||
assert msg1["role"] == "user"
|
||||
assert msg1["content"] == "hello"
|
||||
|
||||
def test_no_tmp_file_left_after_successful_save(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
session = Session(key="test:clean")
|
||||
mgr.save(session)
|
||||
|
||||
tmp_files = list(mgr.sessions_dir.glob("*.tmp"))
|
||||
assert tmp_files == []
|
||||
|
||||
def test_tmp_file_cleaned_up_on_write_failure(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
session = Session(key="test:fail")
|
||||
path = mgr._get_session_path("test:fail")
|
||||
tmp_path_file = path.with_suffix(".jsonl.tmp")
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path_file.write_text("stale")
|
||||
|
||||
class BadMessage:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
original_dumps = json.dumps
|
||||
|
||||
def failing_dumps(obj, **kwargs):
|
||||
if isinstance(obj, dict) and obj.get("role") == "assistant":
|
||||
raise OSError("simulated disk full")
|
||||
return original_dumps(obj, **kwargs)
|
||||
|
||||
session = Session(key="test:fail")
|
||||
session.messages = [
|
||||
{"role": "user", "content": "ok"},
|
||||
{"role": "assistant", "content": "will fail"},
|
||||
]
|
||||
|
||||
import unittest.mock
|
||||
with unittest.mock.patch("nanobot.session.manager.json.dumps", side_effect=failing_dumps):
|
||||
try:
|
||||
mgr.save(session)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
assert not tmp_path_file.exists()
|
||||
|
||||
def test_overwrite_preserves_latest_data(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
session = Session(key="test:overwrite")
|
||||
|
||||
session.add_message("user", "first")
|
||||
mgr.save(session)
|
||||
|
||||
session.add_message("user", "second")
|
||||
mgr.save(session)
|
||||
|
||||
mgr.invalidate("test:overwrite")
|
||||
loaded = mgr.get_or_create("test:overwrite")
|
||||
assert len(loaded.messages) == 2
|
||||
assert loaded.messages[0]["content"] == "first"
|
||||
assert loaded.messages[1]["content"] == "second"
|
||||
|
||||
def test_consecutive_saves_are_consistent(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
session = Session(key="test:consistency")
|
||||
|
||||
for i in range(5):
|
||||
session.add_message("user", f"msg{i}")
|
||||
mgr.save(session)
|
||||
|
||||
mgr.invalidate("test:consistency")
|
||||
loaded = mgr.get_or_create("test:consistency")
|
||||
assert len(loaded.messages) == 5
|
||||
for i in range(5):
|
||||
assert loaded.messages[i]["content"] == f"msg{i}"
|
||||
|
||||
|
||||
class TestRepairCorruptFile:
|
||||
def _write_corrupt_jsonl(self, path: Path, lines: list[str]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
def test_truncated_last_line_recovered(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:trunc")
|
||||
|
||||
valid_meta = json.dumps({
|
||||
"_type": "metadata",
|
||||
"key": "test:trunc",
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"updated_at": datetime.now().isoformat(),
|
||||
"metadata": {},
|
||||
"last_consolidated": 0,
|
||||
})
|
||||
valid_msg = json.dumps({"role": "user", "content": "hello"})
|
||||
|
||||
self._write_corrupt_jsonl(path, [
|
||||
valid_meta,
|
||||
valid_msg,
|
||||
'{"role": "assistant", "content": "partial...',
|
||||
])
|
||||
|
||||
session = mgr._load("test:trunc")
|
||||
assert session is not None
|
||||
assert len(session.messages) == 1
|
||||
assert session.messages[0]["content"] == "hello"
|
||||
|
||||
def test_corrupt_metadata_line_skipped(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:badmeta")
|
||||
|
||||
self._write_corrupt_jsonl(path, [
|
||||
"NOT VALID JSON!!!",
|
||||
'{"role": "user", "content": "survived"}',
|
||||
])
|
||||
|
||||
session = mgr._load("test:badmeta")
|
||||
assert session is not None
|
||||
assert len(session.messages) == 1
|
||||
assert session.messages[0]["content"] == "survived"
|
||||
|
||||
def test_all_corrupt_lines_returns_none(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:allbad")
|
||||
|
||||
self._write_corrupt_jsonl(path, [
|
||||
"garbage line 1",
|
||||
"garbage line 2",
|
||||
"{{invalid json",
|
||||
])
|
||||
|
||||
session = mgr._load("test:allbad")
|
||||
assert session is None
|
||||
|
||||
def test_empty_file_returns_empty_session(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:empty")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text("", encoding="utf-8")
|
||||
|
||||
session = mgr._load("test:empty")
|
||||
assert session is not None
|
||||
assert session.messages == []
|
||||
assert session.key == "test:empty"
|
||||
|
||||
def test_repair_preserves_valid_messages_amid_corruption(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:mixed")
|
||||
|
||||
self._write_corrupt_jsonl(path, [
|
||||
json.dumps({"_type": "metadata", "key": "test:mixed",
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"updated_at": datetime.now().isoformat(),
|
||||
"metadata": {}, "last_consolidated": 0}),
|
||||
"BROKEN",
|
||||
json.dumps({"role": "user", "content": "msg1"}),
|
||||
'{"role": "assistant", "content": "broken',
|
||||
json.dumps({"role": "user", "content": "msg2"}),
|
||||
])
|
||||
|
||||
session = mgr._load("test:mixed")
|
||||
assert session is not None
|
||||
assert len(session.messages) == 2
|
||||
assert session.messages[0]["content"] == "msg1"
|
||||
assert session.messages[1]["content"] == "msg2"
|
||||
|
||||
def test_repair_with_bad_timestamp_uses_fallback(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:badts")
|
||||
|
||||
self._write_corrupt_jsonl(path, [
|
||||
json.dumps({"_type": "metadata", "key": "test:badts",
|
||||
"created_at": "not-a-date",
|
||||
"updated_at": "also-bad",
|
||||
"metadata": {}, "last_consolidated": 5}),
|
||||
json.dumps({"role": "user", "content": "hi"}),
|
||||
])
|
||||
|
||||
session = mgr._load("test:badts")
|
||||
assert session is not None
|
||||
assert session.last_consolidated == 5
|
||||
assert isinstance(session.created_at, datetime)
|
||||
|
||||
def test_get_or_create_returns_new_session_for_corrupt_file(self, tmp_path: Path):
|
||||
mgr = SessionManager(tmp_path)
|
||||
path = mgr._get_session_path("test:fallback")
|
||||
|
||||
self._write_corrupt_jsonl(path, ["{{{{"])
|
||||
|
||||
session = mgr.get_or_create("test:fallback")
|
||||
assert session is not None
|
||||
assert session.messages == []
|
||||
assert session.key == "test:fallback"
|
||||
Reference in New Issue
Block a user