feat: atomic temp+swap file writes (crash-safe)

2026-05-31 10:13:04 +04:00
parent a063c89718
commit 7ad29c3e01
2 changed files with 60 additions and 1 deletions
@@ -31,6 +31,7 @@ import java.nio.file.Files
 import java.nio.file.InvalidPathException
 import java.nio.file.Path
 import java.nio.file.Paths
+import java.nio.file.StandardCopyOption

 class FileWriteTool(
    allowedPaths: Set<Path> = emptySet(),
@@ -172,7 +173,7 @@ class FileWriteTool(
        "write" -> {
            val content = request.parameters["content"] as String
            withContext(Dispatchers.IO) {
-                Files.writeString(path, content)
+                atomicWrite(path, content)
            }
            storeArtifact(pathString, content, request.parameters["mode"] as? String ?: "0644")
            ToolResult.Success(
@@ -206,6 +207,27 @@ class FileWriteTool(
        )
    }

+    /**
+     * Crash-safe write: stage into a temp file in the target's OWN directory (same
+     * filesystem, so the rename is atomic) then atomically swap it into place. A
+     * partial write can never be observed at [target]; on failure the temp is
+     * cleaned and the original is left untouched. Falls back to a non-atomic
+     * replace only if the platform rejects ATOMIC_MOVE.
+     */
+    private fun atomicWrite(target: Path, content: String) {
+        target.parent?.let { Files.createDirectories(it) }
+        val dir = target.parent ?: Paths.get(".")
+        val tmp = Files.createTempFile(dir, ".${target.fileName}", ".tmp")
+        runCatching {
+            Files.writeString(tmp, content)
+            runCatching { Files.move(tmp, target, StandardCopyOption.ATOMIC_MOVE) }
+                .getOrElse { Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING) }
+        }.onFailure {
+            Files.deleteIfExists(tmp)
+            throw it
+        }
+    }
+
    @Suppress("ReturnCount")
    private suspend fun storeArtifact(pathString: String, content: String, mode: String) {
        val writer = materializingWriter ?: return
@@ -229,4 +229,41 @@ class FileWriteToolTest {
        )
        assertEquals(ValidationResult.Valid, tool.validateRequest(request))
    }
+
+    @Test
+    fun `write creates missing parent directories and writes content`(): Unit = runBlocking {
+        val tempDir = Files.createTempDirectory("file_write_atomic")
+        val target = tempDir.resolve("nested/deep/file.txt")
+        val tool = FileWriteTool(allowedPaths = setOf(tempDir))
+        val request = createRequest(
+            mapOf("operation" to "write", "path" to target.toString(), "content" to "hello"),
+        )
+
+        val result = tool.execute(request)
+
+        assertTrue(result is ToolResult.Success)
+        assertTrue(Files.exists(target))
+        assertEquals("hello", Files.readString(target))
+    }
+
+    @Test
+    fun `write atomically overwrites an existing file and leaves no temp files`(): Unit = runBlocking {
+        val tempDir = Files.createTempDirectory("file_write_overwrite")
+        val target = tempDir.resolve("file.txt")
+        Files.writeString(target, "old")
+        val tool = FileWriteTool(allowedPaths = setOf(tempDir))
+        val request = createRequest(
+            mapOf("operation" to "write", "path" to target.toString(), "content" to "new"),
+        )
+
+        val result = tool.execute(request)
+
+        assertTrue(result is ToolResult.Success)
+        assertEquals("new", Files.readString(target))
+        // the staged temp file must not survive the swap
+        val leftovers = Files.list(tempDir).use { stream ->
+            stream.filter { it.fileName.toString().endsWith(".tmp") }.count()
+        }
+        assertEquals(0L, leftovers)
+    }
 }