From efb04a1712055eecce939f0d9275f664048fdf08 Mon Sep 17 00:00:00 2001 From: aiguozhi123456 <126325311+aiguozhi123456@users.noreply.github.com> Date: Sun, 19 Apr 2026 22:36:16 +0800 Subject: [PATCH] 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) --- nanobot/session/manager.py | 96 +++++++++++-- tests/agent/test_session_atomic.py | 220 +++++++++++++++++++++++++++++ 2 files changed, 303 insertions(+), 13 deletions(-) create mode 100644 tests/agent/test_session_atomic.py diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index c91eabcb..c264649f 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -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 diff --git a/tests/agent/test_session_atomic.py b/tests/agent/test_session_atomic.py new file mode 100644 index 00000000..4b84d8b6 --- /dev/null +++ b/tests/agent/test_session_atomic.py @@ -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"