fix: write praisonai.code edit tools atomically so interrupts keep the original file (fixes #6204)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

src/praisonai/praisonai/code/tools/apply_diff.py

@@ -10,6 +10,7 @@ from typing import Optional, Dict, Any
 
 from ..diff.diff_strategy import apply_search_replace_diff
 from ..utils.file_utils import (
+    atomic_write,
     file_exists,
     is_path_within_directory,
 )
@@ -142,8 +143,7 @@ def apply_diff(
                 f.write(original_content)
         
         # Write the modified content
-        with open(abs_path, 'w', encoding=encoding) as f:
-            f.write(result.content)
+        atomic_write(abs_path, result.content, encoding=encoding)
         
         return {
             'success': True,

src/praisonai/praisonai/code/tools/search_replace.py

@@ -10,6 +10,7 @@ import re
 from typing import Optional, List, Dict, Any
 
 from ..utils.file_utils import (
+    atomic_write,
     file_exists,
     is_path_within_directory,
 )
@@ -191,8 +192,7 @@ def search_replace(
                 f.write(original_content)
         
         # Write the modified content
-        with open(abs_path, 'w', encoding=encoding) as f:
-            f.write(content)
+        atomic_write(abs_path, content, encoding=encoding)
         
         return {
             'success': True,

src/praisonai/praisonai/code/tools/write_file.py

@@ -8,6 +8,7 @@ import os
 from typing import Optional, Dict, Any
 
 from ..utils.file_utils import (
+    atomic_write,
     file_exists,
     create_directories_for_file,
     is_path_within_directory,
@@ -131,8 +132,7 @@ def write_file(
                 pass  # If we can't read, just use the content as-is
         
         # Write the file
-        with open(abs_path, 'w', encoding=encoding) as f:
-            f.write(processed_content)
+        atomic_write(abs_path, processed_content, encoding=encoding)
         
         return {
             'success': True,

src/praisonai/praisonai/code/utils/file_utils.py

@@ -10,6 +10,7 @@ Provides utilities for file operations including:
 
 import os
 import re
+import tempfile
 from typing import Optional, Tuple
 
 
@@ -305,3 +306,47 @@ def count_file_lines(file_path: str) -> int:
         for _ in f:
             count += 1
     return count
+
+
+def atomic_write(path: str, content: str, *, encoding: str = "utf-8") -> None:
+    """Write ``content`` to ``path`` atomically (temp file + fsync + os.replace).
+
+    ``open(path, 'w')`` truncates the target immediately, so any interrupt
+    (cancellation, SIGTERM, disk-full) between open and write leaves the file
+    zero-byte or half-written and the original content lost. The code editing
+    tools run on LLM-driven tool calls that regularly hit timeouts/cancellations,
+    so they need the same crash-safe write the wrapper already uses for configs.
+    """
+    target_dir = os.path.dirname(os.path.abspath(path)) or "."
+    fd, tmp_path = tempfile.mkstemp(prefix=".tmp_", dir=target_dir)
+    try:
+        # Atomic replace swaps in a fresh inode, so carry over the metadata that
+        # matters for a replaced file: mode (mkstemp creates 0600 by default) and
+        # ownership (uid/gid) so editing a group-/shared-owned file does not
+        # silently change its owner. Hard-link identity and xattrs/ACLs cannot be
+        # preserved through an inode swap and are out of scope for this tool.
+        try:
+            existing_stat = os.stat(path)
+        except OSError:
+            existing_stat = None
+        with os.fdopen(fd, "w", encoding=encoding) as f:
+            f.write(content)
+            f.flush()
+            os.fsync(f.fileno())  # durability before rename
+        if existing_stat is not None:
+            try:
+                os.chmod(tmp_path, existing_stat.st_mode)
+            except OSError:
+                pass  # best-effort
+            if hasattr(os, "chown"):
+                try:
+                    os.chown(tmp_path, existing_stat.st_uid, existing_stat.st_gid)
+                except (OSError, AttributeError):
+                    pass  # best-effort (needs privilege; no-op on Windows)
+        os.replace(tmp_path, path)  # atomic on POSIX + Windows
+    except BaseException:
+        try:
+            os.unlink(tmp_path)
+        except OSError:
+            pass
+        raise

src/praisonai/praisonai/tools/multiedit.py

@@ -8,55 +8,12 @@ without multiple file read/write cycles.
 import contextlib
 import difflib
 import os
-import tempfile
 from typing import Any, Dict, List, Optional
 
+from praisonai.code.utils.file_utils import atomic_write as _atomic_write
 from praisonai.code.utils.file_utils import is_path_within_directory
 
 
-def _atomic_write(path: str, content: str) -> None:
-    """Write ``content`` to ``path`` atomically (temp file + fsync + os.replace).
-
-    ``open(path, 'w')`` truncates the target immediately, so any interrupt
-    (cancellation, SIGTERM, disk-full) between open and write leaves the file
-    zero-byte or half-written and the original content lost. multiedit runs on
-    LLM-driven tool calls that regularly hit timeouts/cancellations, so it needs
-    the same crash-safe write the rest of the wrapper already uses for configs.
-    """
-    target_dir = os.path.dirname(os.path.abspath(path)) or "."
-    fd, tmp_path = tempfile.mkstemp(prefix=".tmp_", dir=target_dir)
-    try:
-        # Atomic replace swaps in a fresh inode, so carry over the metadata that
-        # matters for a replaced file: mode (mkstemp creates 0600 by default) and
-        # ownership (uid/gid) so editing a group-/shared-owned file does not
-        # silently change its owner. Hard-link identity and xattrs/ACLs cannot be
-        # preserved through an inode swap and are out of scope for this tool.
-        try:
-            existing_stat = os.stat(path)
-        except OSError:
-            existing_stat = None
-        with os.fdopen(fd, "w") as f:
-            f.write(content)
-            f.flush()
-            os.fsync(f.fileno())  # durability before rename
-        if existing_stat is not None:
-            try:
-                os.chmod(tmp_path, existing_stat.st_mode)
-            except OSError:
-                pass  # best-effort
-            if hasattr(os, "chown"):
-                try:
-                    os.chown(tmp_path, existing_stat.st_uid, existing_stat.st_gid)
-                except (OSError, AttributeError):
-                    pass  # best-effort (needs privilege; no-op on Windows)
-        os.replace(tmp_path, path)  # atomic on POSIX + Windows
-    except BaseException:
-        try:
-            os.unlink(tmp_path)
-        except OSError:
-            pass
-        raise
-
 
 @contextlib.contextmanager
 def _file_lock(path: str):

src/praisonai/tests/unit/code/test_tools.py

@@ -286,7 +286,31 @@ class TestSearchReplace:
             
             with open(file_path, 'r') as f:
                 assert f.read() == "Hello, new world!"
-    
+
+    def test_interrupted_write_keeps_original(self, monkeypatch):
+        """A failure mid-write must not truncate the original file."""
+        from praisonai.code.tools.search_replace import search_replace
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = os.path.join(temp_dir, "test.txt")
+            with open(file_path, 'w') as f:
+                f.write("Hello, old world!")
+
+            def boom(*args, **kwargs):
+                raise KeyboardInterrupt
+
+            monkeypatch.setattr(os, "fsync", boom)
+            try:
+                search_replace(file_path, [
+                    {'search': 'old', 'replace': 'new'}
+                ], workspace=temp_dir)
+            except KeyboardInterrupt:
+                pass
+
+            with open(file_path, 'r') as f:
+                assert f.read() == "Hello, old world!"
+            assert os.listdir(temp_dir) == ["test.txt"]
+
     def test_multiple_operations(self):
         """Test multiple search/replace operations."""
         from praisonai.code.tools.search_replace import search_replace