diff --git a/src/anthropic/lib/tools/_beta_builtin_memory_tool.py b/src/anthropic/lib/tools/_beta_builtin_memory_tool.py index b995fbb3d..16f27ccf9 100644 --- a/src/anthropic/lib/tools/_beta_builtin_memory_tool.py +++ b/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) diff --git a/tests/lib/tools/memory_tools/test_create_short_writes.py b/tests/lib/tools/memory_tools/test_create_short_writes.py new file mode 100644 index 000000000..857dc5568 --- /dev/null +++ b/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