fix(checkpoint): preserve state on failed saves

plexe/checkpointing.py

@@ -14,6 +14,7 @@
 import base64
 import json
 import logging
+import tempfile
 from datetime import datetime, timezone
 from pathlib import Path
 
@@ -93,6 +94,7 @@ def save_checkpoint(
     """
     logger.info(f"Saving checkpoint for phase: {phase_name}")
 
+    temporary_path = None
     try:
         # Extract user feedback from context scratch space (if provided)
         user_feedback = context.scratch.get("_user_feedback") if hasattr(context, "scratch") else None
@@ -115,8 +117,14 @@ def save_checkpoint(
         checkpoint_dir.mkdir(parents=True, exist_ok=True)
 
         checkpoint_path = checkpoint_dir / f"{phase_name}.json"
-        with open(checkpoint_path, "w") as f:
+        # Publish only a complete checkpoint. A failed write must not truncate
+        # the previous phase state, which may be the only resumable copy.
+        with tempfile.NamedTemporaryFile(
+            mode="w", encoding="utf-8", dir=checkpoint_dir, prefix=f".{phase_name}.", suffix=".tmp", delete=False
+        ) as f:
+            temporary_path = Path(f.name)
             json.dump(checkpoint_data, f, indent=2, default=str)
+        temporary_path.replace(checkpoint_path)
 
         logger.info(f"✓ Checkpoint saved locally: {checkpoint_path}")
         if user_feedback:
@@ -126,6 +134,12 @@ def save_checkpoint(
     except Exception as e:
         logger.error(f"Failed to save checkpoint locally: {e}", exc_info=True)
         return None
+    finally:
+        if temporary_path is not None:
+            try:
+                temporary_path.unlink(missing_ok=True)
+            except OSError:
+                logger.warning("Could not remove temporary checkpoint %s", temporary_path, exc_info=True)
 
 
 def load_checkpoint(

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "plexe"
-version = "1.4.4"
+version = "1.4.5"
 description = "An agentic framework for building ML models from natural language"
 authors = [
     "Marcello De Bernardi <mdebernardi@plexe.ai>",

tests/unit/test_checkpoint_atomicity.py

@@ -0,0 +1,73 @@
+"""Checkpoint publication must preserve the last resumable state on failure."""
+
+import json
+from pathlib import Path
+
+import pytest
+
+from plexe import checkpointing
+from plexe.models import BuildContext
+
+
+@pytest.fixture
+def context(tmp_path):
+    """Build real workflow state without running a model or a provider."""
+    return BuildContext("user", "experiment", "data.parquet", tmp_path, "Predict a value")
+
+
+@pytest.mark.parametrize("existing", [False, True])
+@pytest.mark.parametrize("failure", [OSError("disk full"), ValueError("serialization failed"), KeyboardInterrupt()])
+def test_failed_save_preserves_previous_checkpoint(tmp_path, context, monkeypatch, existing, failure):
+    """A partial write must never replace a valid checkpoint or appear as one."""
+    checkpoint_dir = tmp_path / "checkpoints"
+    checkpoint_dir.mkdir()
+    destination = checkpoint_dir / "prepare_data.json"
+    previous = b'{"phase": "prepare_data", "context": {"intent": "previous"}}'
+    if existing:
+        destination.write_bytes(previous)
+
+    def failing_dump(data, stream, **kwargs):
+        stream.write('{"context": ')
+        raise failure
+
+    monkeypatch.setattr(checkpointing.json, "dump", failing_dump)
+    if isinstance(failure, KeyboardInterrupt):
+        with pytest.raises(KeyboardInterrupt):
+            checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path)
+    else:
+        assert checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) is None
+
+    assert list(checkpoint_dir.iterdir()) == ([destination] if existing else [])
+    if existing:
+        assert destination.read_bytes() == previous
+        assert checkpointing.load_checkpoint("prepare_data", tmp_path)["context"]["intent"] == "previous"
+
+
+def test_successful_save_replaces_previous_checkpoint(tmp_path, context):
+    """The public path and persisted feedback remain compatible across saves."""
+    first = checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path)
+    assert first == tmp_path / "checkpoints" / "prepare_data.json"
+    context.intent = "Updated intent"
+    context.scratch["_user_feedback"] = {"message": "保留 • feedback"}
+    second = checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path)
+    assert second == first
+    result = checkpointing.load_checkpoint("prepare_data", tmp_path)
+    assert result["context"]["intent"] == "Updated intent"
+    assert result["user_feedback"] == context.scratch["_user_feedback"]
+    assert json.loads(second.read_text())["version"] == "plexe-v1"
+    assert list(second.parent.iterdir()) == [second]
+
+
+def test_failed_replace_preserves_previous_checkpoint(tmp_path, context, monkeypatch):
+    """An OS replacement error must clean the temporary file and preserve state."""
+    destination = checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path)
+    previous = destination.read_bytes()
+
+    def fail_replace(self, target):
+        raise PermissionError("destination temporarily unavailable")
+
+    monkeypatch.setattr(Path, "replace", fail_replace)
+    context.intent = "Not published"
+    assert checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) is None
+    assert destination.read_bytes() == previous
+    assert list(destination.parent.iterdir()) == [destination]