Handle short writes when creating memory files

src/anthropic/lib/tools/_beta_builtin_memory_tool.py

@@ -278,6 +278,16 @@ async def clear_all_memory(self) -> BetaFunctionToolResultType:
         raise NotImplementedError("clear_all_memory not implemented")
 
 
+def _write_all(fd: int, data: bytes) -> None:
+    """Write the complete buffer to ``fd``, retrying legal short writes."""
+    offset = 0
+    while offset < len(data):
+        written = os.write(fd, data[offset:])
+        if written == 0:
+            raise OSError("os.write returned 0")
+        offset += written
+
+
 def _atomic_write_file(target_path: Path, content: str) -> None:
     dir_path = target_path.parent
     temp_path = dir_path / f".tmp-{os.getpid()}-{uuid.uuid4()}"
@@ -286,13 +296,7 @@ def _atomic_write_file(target_path: Path, content: str) -> None:
     try:
         fd = os.open(temp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, _FILE_CREATE_MODE)
         try:
-            offset = 0
-            while offset < len(data):
-                written = os.write(fd, data[offset:])
-                if written == 0:
-                    raise OSError("os.write returned 0")
-                offset += written
-
+            _write_all(fd, data)
             os.fsync(fd)
         finally:
             os.close(fd)
@@ -479,7 +483,7 @@ def create(self, command: BetaMemoryTool20250818CreateCommand) -> str:
         try:
             fd = os.open(full_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, _FILE_CREATE_MODE)
             try:
-                os.write(fd, command.file_text.encode("utf-8"))
+                _write_all(fd, command.file_text.encode("utf-8"))
                 os.fsync(fd)
             finally:
                 os.close(fd)
@@ -618,13 +622,7 @@ async def _async_atomic_write_file(target_path: AsyncPath, content: str) -> None
         def write_replace_and_sync() -> None:
             fd = os.open(sync_temp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, _FILE_CREATE_MODE)
             try:
-                offset = 0
-                while offset < len(data):
-                    written = os.write(fd, data[offset:])
-                    if written == 0:
-                        raise OSError("os.write returned 0")
-                    offset += written
-
+                _write_all(fd, data)
                 os.fsync(fd)
             finally:
                 os.close(fd)
@@ -776,7 +774,7 @@ async def create(self, command: BetaMemoryTool20250818CreateCommand) -> str:
             def create_exclusive() -> None:
                 fd = os.open(sync_full_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, _FILE_CREATE_MODE)
                 try:
-                    os.write(fd, command.file_text.encode("utf-8"))
+                    _write_all(fd, command.file_text.encode("utf-8"))
                     os.fsync(fd)
                 finally:
                     os.close(fd)

tests/lib/tools/memory_tools/test_create_short_writes.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+from anthropic.lib.tools._beta_builtin_memory_tool import (
+    BetaAsyncLocalFilesystemMemoryTool,
+    BetaLocalFilesystemMemoryTool,
+)
+from anthropic.types.beta import BetaMemoryTool20250818CreateCommand
+
+
+def _force_short_writes(monkeypatch: pytest.MonkeyPatch) -> list[int]:
+    real_write = os.write
+    write_sizes: list[int] = []
+
+    def short_write(fd: int, data: bytes) -> int:
+        chunk_size = min(3, len(data))
+        written = real_write(fd, data[:chunk_size])
+        write_sizes.append(written)
+        return written
+
+    monkeypatch.setattr(os, "write", short_write)
+    return write_sizes
+
+
+def test_create_retries_short_writes(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    write_sizes = _force_short_writes(monkeypatch)
+    tool = BetaLocalFilesystemMemoryTool(base_path=str(tmp_path))
+    text = "prefix-☃-suffix"
+
+    result = tool.create(
+        BetaMemoryTool20250818CreateCommand(
+            command="create",
+            file_text=text,
+            path="/memories/short-write.txt",
+        )
+    )
+
+    assert result == "File created successfully at: /memories/short-write.txt"
+    assert (tmp_path / "memories" / "short-write.txt").read_text(encoding="utf-8") == text
+    assert len(write_sizes) > 1
+
+
+@pytest.mark.asyncio
+async def test_async_create_retries_short_writes(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    write_sizes = _force_short_writes(monkeypatch)
+    tool = BetaAsyncLocalFilesystemMemoryTool(base_path=str(tmp_path))
+    text = "prefix-☃-suffix"
+
+    result = await tool.create(
+        BetaMemoryTool20250818CreateCommand(
+            command="create",
+            file_text=text,
+            path="/memories/short-write.txt",
+        )
+    )
+
+    assert result == "File created successfully at: /memories/short-write.txt"
+    assert (tmp_path / "memories" / "short-write.txt").read_text(encoding="utf-8") == text
+    assert len(write_sizes) > 1