From cc400ffc6392e8f5c5587700e2a1decfae044b78 Mon Sep 17 00:00:00 2001
From: mukunda katta <mukunda.vjcs6@gmail.com>
Date: Thu, 14 May 2026 19:39:19 -0700
Subject: [PATCH] fix: atomically write forked session transcripts

---
 .../_internal/session_mutations.py            | 38 ++++++++++++++++---
 tests/test_session_mutations.py               | 27 +++++++++++++
 2 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/src/claude_agent_sdk/_internal/session_mutations.py b/src/claude_agent_sdk/_internal/session_mutations.py
index 55a7f2132..749f65e73 100644
--- a/src/claude_agent_sdk/_internal/session_mutations.py
+++ b/src/claude_agent_sdk/_internal/session_mutations.py
@@ -26,6 +26,7 @@
 import unicodedata
 import uuid as uuid_mod
 from collections.abc import Callable
+from contextlib import suppress
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
@@ -336,15 +337,42 @@ def _derive_title() -> str | None:
     )
 
     fork_path = project_dir / f"{forked_session_id}.jsonl"
-    fd = os.open(fork_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
-    try:
-        os.write(fd, ("\n".join(lines) + "\n").encode("utf-8"))
-    finally:
-        os.close(fd)
+    _write_new_file_atomically(fork_path, ("\n".join(lines) + "\n").encode("utf-8"))
 
     return ForkSessionResult(session_id=forked_session_id)
 
 
+def _write_new_file_atomically(path: Path, data: bytes) -> None:
+    """Create ``path`` atomically with ``data``.
+
+    The final path is only published after the complete payload has been written
+    to a temporary sibling file. This keeps failed fork_session() writes from
+    leaving partial transcripts behind.
+    """
+
+    if path.exists():
+        raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), str(path))
+
+    tmp_path = path.with_name(f".{path.name}.{uuid_mod.uuid4().hex}.tmp")
+    fd = os.open(tmp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
+    try:
+        offset = 0
+        while offset < len(data):
+            written = os.write(fd, data[offset:])
+            if written == 0:
+                raise OSError("os.write returned 0 bytes")
+            offset += written
+        os.close(fd)
+        fd = -1
+        tmp_path.replace(path)
+    except BaseException:
+        if fd != -1:
+            os.close(fd)
+        with suppress(FileNotFoundError):
+            tmp_path.unlink()
+        raise
+
+
 def _build_fork_lines(
     transcript: list[dict[str, Any]],
     content_replacements: list[Any],
diff --git a/tests/test_session_mutations.py b/tests/test_session_mutations.py
index 18ad8aeb5..5db7fd824 100644
--- a/tests/test_session_mutations.py
+++ b/tests/test_session_mutations.py
@@ -633,6 +633,33 @@ def test_fork_creates_new_session(self, claude_config_dir: Path, tmp_path: Path)
         fork_path = project_dir / f"{result.session_id}.jsonl"
         assert fork_path.exists()
 
+    def test_fork_write_failure_leaves_no_partial_session(
+        self,
+        claude_config_dir: Path,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """A failed fork write does not publish a partial transcript."""
+        project_path = str(tmp_path / "proj")
+        Path(project_path).mkdir(parents=True)
+        project_dir = _make_project_dir(
+            claude_config_dir, os.path.realpath(project_path)
+        )
+        sid, source_path, _ = _make_transcript_session(project_dir)
+        real_write = os.write
+
+        def fail_after_partial_write(fd: int, data: bytes) -> int:
+            real_write(fd, data[:5])
+            raise OSError("simulated write failure")
+
+        monkeypatch.setattr(os, "write", fail_after_partial_write)
+
+        with pytest.raises(OSError, match="simulated write failure"):
+            fork_session(sid, directory=project_path)
+
+        assert sorted(project_dir.glob("*.jsonl")) == [source_path]
+        assert not [path for path in project_dir.iterdir() if path.suffix == ".tmp"]
+
     def test_fork_remaps_uuids(self, claude_config_dir: Path, tmp_path: Path):
         """uuid and parentUuid fields are remapped; originals only appear in forkedFrom."""
         project_path = str(tmp_path / "proj")