From 7e4d52eabfe809b4973c8350ce1b43d0f7b5efc1 Mon Sep 17 00:00:00 2001 From: sshekhar563 Date: Tue, 7 Jul 2026 19:39:06 +0530 Subject: [PATCH] fix(storage): make save and export crash-safe (fixes #282) Replace delete-before-move with backup-then-swap pattern. Adds: - _replace_with_backup(): atomic swap with rollback on failure - write_bytes_atomic(): temp file + fsync + atomic replace for exports - _fsync_directory(): platform-aware dir fsync (no-op on Windows) - Regression tests for failed replacements and temp cleanup --- src/soul_protocol/runtime/soul.py | 3 +- src/soul_protocol/runtime/storage/file.py | 125 ++++++++++++++++++---- tests/test_soul.py | 13 +++ tests/test_storage/test_storage.py | 113 +++++++++++++++++++ 4 files changed, 235 insertions(+), 19 deletions(-) diff --git a/src/soul_protocol/runtime/soul.py b/src/soul_protocol/runtime/soul.py index 564b233..36b5c29 100644 --- a/src/soul_protocol/runtime/soul.py +++ b/src/soul_protocol/runtime/soul.py @@ -3285,6 +3285,7 @@ async def export( devices). See docs/trust-chain.md for the threat model. """ from .exceptions import SoulExportError + from .storage.file import write_bytes_atomic try: memory_data = self._memory.to_dict() @@ -3297,7 +3298,7 @@ async def export( trust_chain_data=trust_chain_data, key_files=key_files, ) - Path(path).write_bytes(data) + write_bytes_atomic(Path(path), data, keep_backup=True) logger.info( "Soul exported: name=%s, path=%s, size=%d bytes", self.name, diff --git a/src/soul_protocol/runtime/storage/file.py b/src/soul_protocol/runtime/storage/file.py index 922dd4a..152f5e1 100644 --- a/src/soul_protocol/runtime/storage/file.py +++ b/src/soul_protocol/runtime/storage/file.py @@ -19,6 +19,7 @@ import json import logging +import os import shutil import tempfile import warnings @@ -32,6 +33,102 @@ DEFAULT_SOUL_DIR: Path = Path.home() / ".soul" +def _backup_path(path: Path) -> Path: + """Return the sibling backup path used for crash recovery.""" + return path.with_name(f"{path.name}.bak") + + +def _remove_path(path: Path) -> None: + if not path.exists(): + return + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + else: + path.unlink() + + +def _fsync_directory(path: Path) -> None: + """Best-effort fsync for directory metadata. + + Windows does not support opening directories with os.open(), so this + is a no-op on win32. On POSIX systems it ensures the directory entry + changes (renames, new files) have been flushed to stable storage. + """ + if os.name == "nt": + return + try: + fd = os.open(str(path), os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _fsync_tree(path: Path) -> None: + """Best-effort fsync of all files and directory metadata under ``path``.""" + for child in path.rglob("*"): + if child.is_file(): + try: + with child.open("rb") as f: + os.fsync(f.fileno()) + except OSError: + logger.debug("Could not fsync file during soul save: %s", child) + for child in sorted((p for p in path.rglob("*") if p.is_dir()), reverse=True): + _fsync_directory(child) + _fsync_directory(path) + + +def _replace_with_backup(source: Path, target: Path, *, keep_backup: bool) -> None: + """Replace ``target`` with ``source`` while preserving the old target first.""" + backup = _backup_path(target) + target_was_backed_up = False + + if backup.exists(): + _remove_path(backup) + + if target.exists(): + os.replace(target, backup) + target_was_backed_up = True + _fsync_directory(target.parent) + + try: + os.replace(source, target) + except Exception: + if target_was_backed_up and not target.exists() and backup.exists(): + os.replace(backup, target) + raise + + _fsync_directory(target.parent) + + if target_was_backed_up and not keep_backup: + _remove_path(backup) + _fsync_directory(target.parent) + + +def write_bytes_atomic(path: Path, data: bytes, *, keep_backup: bool = True) -> None: + """Write bytes via temp file + fsync + atomic replace.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + fd, tmp_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + tmp_path = Path(tmp_name) + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + _replace_with_backup(tmp_path, path, keep_backup=keep_backup) + except Exception: + _remove_path(tmp_path) + raise + + def _safe_soul_id(config: SoulConfig) -> str: """Extract soul_id from config and sanitize for filesystem use. @@ -320,15 +417,14 @@ async def save_soul_full( base.mkdir(parents=True, exist_ok=True) soul_dir = base / soul_id - # Atomic write: build in temp dir, then move into place - with tempfile.TemporaryDirectory(dir=base) as tmp: + # Atomic-ish directory write: build a sibling temp dir, fsync it, then + # replace the live directory through a backup instead of deleting first. + with tempfile.TemporaryDirectory(prefix=f".{soul_id}.tmp-", dir=base) as tmp: tmp_dir = Path(tmp) / soul_id tmp_dir.mkdir() _write_soul_files(tmp_dir, config, memory_data) - - if soul_dir.exists(): - shutil.rmtree(soul_dir) - shutil.move(str(tmp_dir), str(soul_dir)) + _fsync_tree(tmp_dir) + _replace_with_backup(tmp_dir, soul_dir, keep_backup=False) logger.debug("Soul saved (full): path=%s", soul_dir) @@ -464,18 +560,11 @@ async def save_soul_flat( path = Path(path) path.mkdir(parents=True, exist_ok=True) - # Atomic write: build in temp dir, then move individual files into place - with tempfile.TemporaryDirectory(dir=path.parent) as tmp: + # Atomic-ish directory write: build a complete replacement directory, then + # swap the old directory aside before installing the new one. + with tempfile.TemporaryDirectory(prefix=f".{path.name}.tmp-", dir=path.parent) as tmp: tmp_dir = Path(tmp) / path.name tmp_dir.mkdir() _write_soul_files(tmp_dir, config, memory_data) - - # Move files into target (preserve existing dir structure) - for item in tmp_dir.iterdir(): - dest = path / item.name - if dest.exists(): - if dest.is_dir(): - shutil.rmtree(dest) - else: - dest.unlink() - shutil.move(str(item), str(dest)) + _fsync_tree(tmp_dir) + _replace_with_backup(tmp_dir, path, keep_backup=False) diff --git a/tests/test_soul.py b/tests/test_soul.py index 36c35f3..bdcb4c8 100644 --- a/tests/test_soul.py +++ b/tests/test_soul.py @@ -106,6 +106,19 @@ async def test_export_and_awaken(tmp_path): assert restored.lifecycle == LifecycleState.ACTIVE +async def test_export_overwrite_keeps_previous_file_backup(tmp_path): + """Overwriting an export keeps the previous archive as a .bak fallback.""" + soul = await Soul.birth("Aria") + soul_path = tmp_path / "aria.soul" + old_data = b"previous archive" + soul_path.write_bytes(old_data) + + await soul.export(soul_path) + + assert soul_path.read_bytes() != old_data + assert (tmp_path / "aria.soul.bak").read_bytes() == old_data + + async def test_system_prompt(): """to_system_prompt() returns a non-empty string containing the soul's name.""" soul = await Soul.birth("Aria", archetype="The Compassionate Creator") diff --git a/tests/test_storage/test_storage.py b/tests/test_storage/test_storage.py index 3d092e4..63134ea 100644 --- a/tests/test_storage/test_storage.py +++ b/tests/test_storage/test_storage.py @@ -4,8 +4,11 @@ from __future__ import annotations +from pathlib import Path + import pytest +from soul_protocol.runtime.storage import file as storage_file from soul_protocol.runtime.storage.file import FileStorage from soul_protocol.runtime.storage.memory_store import InMemoryStorage from soul_protocol.runtime.types import Identity, SoulConfig @@ -118,3 +121,113 @@ async def test_file_storage_list(config: SoulConfig, tmp_path): # Loading from non-existent returns None missing = await store.load("missing-soul") assert missing is None + + +async def test_save_soul_full_restores_previous_directory_on_replace_failure( + config: SoulConfig, + tmp_path, + monkeypatch, +): + """A failed final replacement leaves the previous full save readable.""" + await storage_file.save_soul_full(config, {"core": {"persona": "old"}}, path=tmp_path) + + soul_dir = tmp_path / "did_soul_aria-abc123" + sentinel = soul_dir / "sentinel.txt" + sentinel.write_text("old copy", encoding="utf-8") + + real_replace = storage_file.os.replace + + def fail_final_replace(src, dst): + if Path(dst) == soul_dir and Path(src).name == soul_dir.name: + raise OSError("simulated replace failure") + return real_replace(src, dst) + + monkeypatch.setattr(storage_file.os, "replace", fail_final_replace) + + with pytest.raises(OSError, match="simulated replace failure"): + await storage_file.save_soul_full(config, {"core": {"persona": "new"}}, path=tmp_path) + + assert sentinel.read_text(encoding="utf-8") == "old copy" + assert soul_dir.exists() + assert not (tmp_path / "did_soul_aria-abc123.bak").exists() + + +async def test_save_soul_flat_restores_previous_directory_on_replace_failure( + config: SoulConfig, + tmp_path, + monkeypatch, +): + """A failed final replacement leaves the previous flat save readable.""" + soul_dir = tmp_path / ".soul" + await storage_file.save_soul_flat(config, {"core": {"persona": "old"}}, soul_dir) + + sentinel = soul_dir / "sentinel.txt" + sentinel.write_text("old copy", encoding="utf-8") + + real_replace = storage_file.os.replace + + def fail_final_replace(src, dst): + if Path(dst) == soul_dir and Path(src).name == soul_dir.name: + raise OSError("simulated replace failure") + return real_replace(src, dst) + + monkeypatch.setattr(storage_file.os, "replace", fail_final_replace) + + with pytest.raises(OSError, match="simulated replace failure"): + await storage_file.save_soul_flat(config, {"core": {"persona": "new"}}, soul_dir) + + assert sentinel.read_text(encoding="utf-8") == "old copy" + assert soul_dir.exists() + assert not (tmp_path / ".soul.bak").exists() + + +def test_write_bytes_atomic_cleans_temp_on_write_failure(tmp_path, monkeypatch): + """If the write itself fails, no temp file is left behind.""" + target = tmp_path / "test.soul" + + def exploding_fdopen(fd, mode): + # Close the fd so it doesn't leak, then raise + storage_file.os.close(fd) + raise OSError("simulated disk full") + + monkeypatch.setattr(storage_file.os, "fdopen", exploding_fdopen) + + with pytest.raises(OSError, match="simulated disk full"): + storage_file.write_bytes_atomic(target, b"data") + + # No temp files should remain + leftover = [p for p in tmp_path.iterdir() if ".tmp" in p.name] + assert leftover == [], f"Temp files left behind: {leftover}" + assert not target.exists() + + +async def test_save_soul_full_first_time_creates_no_backup( + config: SoulConfig, + tmp_path, +): + """First-time save (no pre-existing target) creates the directory with no .bak.""" + await storage_file.save_soul_full(config, {"core": {"persona": "v1"}}, path=tmp_path) + + soul_dir = tmp_path / "did_soul_aria-abc123" + assert soul_dir.exists() + assert (soul_dir / "soul.json").exists() + + # No backup should exist on first save + bak_dir = tmp_path / "did_soul_aria-abc123.bak" + assert not bak_dir.exists() + + +def test_write_bytes_atomic_overwrites_stale_backup(tmp_path): + """A stale .bak from a previous run is replaced on the next write.""" + target = tmp_path / "test.dat" + bak = tmp_path / "test.dat.bak" + + # Simulate a stale backup from a previous crash + bak.write_bytes(b"stale backup") + target.write_bytes(b"current") + + storage_file.write_bytes_atomic(target, b"new data") + + assert target.read_bytes() == b"new data" + # .bak should now contain the previous "current", not "stale backup" + assert bak.read_bytes() == b"current"