Python: Improve atomic writes for Memory and Todos

python/packages/core/AGENTS.md

@@ -16,7 +16,7 @@ agent_framework/
 ├── _vectors.py          # Vector store models, CRUD/search abstractions, and protocols
 ├── _middleware.py       # Middleware system for request/response interception
 ├── _sessions.py         # AgentSession and context provider abstractions
-├── _filesystem.py       # Private filesystem safety helpers (link detection, storage-key derivation)
+├── _filesystem.py       # Private filesystem safety helpers (atomic writes, link detection, storage-key derivation)
 ├── _skills.py           # Agent Skills system (models, executors, provider)
 ├── _mcp.py              # Model Context Protocol support
 ├── _telemetry.py        # User-Agent identity and internal feature-usage mask

python/packages/core/agent_framework/_filesystem.py

@@ -5,7 +5,9 @@
 from __future__ import annotations
 
 import hashlib
+import os
 import stat
+import uuid
 from base64 import b32encode
 from pathlib import Path
 
@@ -44,6 +46,33 @@
 _DIGEST_SEGMENT_MARKER = "sha256-"
 
 
+def _atomic_write_text(path: Path, text: str, *, encoding: str = "utf-8") -> None:  # pyright: ignore[reportUnusedFunction]
+    """Atomically write text through a securely created sibling temporary file.
+
+    Exclusive creation prevents an existing file or symbolic link from redirecting the write.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temp_path = path.with_name(f".maf-state-{uuid.uuid4().hex}.tmp")
+    cleanup_path: Path | None = None
+    file_descriptor: int | None = None
+    try:
+        file_descriptor = os.open(temp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o666)
+        cleanup_path = temp_path
+        file_handle = os.fdopen(file_descriptor, "w", encoding=encoding, closefd=False)
+        with file_handle:
+            file_handle.write(text)
+        descriptor_to_close = file_descriptor
+        file_descriptor = None
+        os.close(descriptor_to_close)
+        os.replace(temp_path, path)
+        cleanup_path = None
+    finally:
+        if file_descriptor is not None:
+            os.close(file_descriptor)
+        if cleanup_path is not None:
+            cleanup_path.unlink(missing_ok=True)
+
+
 def _is_link_or_reparse_point(path: Path) -> bool:  # pyright: ignore[reportUnusedFunction]
     """Return whether ``path`` is a symbolic link, junction, or other reparse point."""
     path_stat = path.lstat()

python/packages/core/agent_framework/_harness/_memory.py

@@ -5,7 +5,6 @@
 import asyncio
 import json
 import logging
-import os
 import re
 import threading
 import weakref
@@ -20,7 +19,10 @@
 from .._clients import SupportsChatGetResponse
 from .._compaction import group_messages
 from .._feature_stage import ExperimentalFeature, experimental
-from .._filesystem import _storage_key_segment  # pyright: ignore[reportPrivateUsage]
+from .._filesystem import (
+    _atomic_write_text,  # pyright: ignore[reportPrivateUsage]
+    _storage_key_segment,  # pyright: ignore[reportPrivateUsage]
+)
 from .._sessions import AgentSession, FileHistoryProvider, HistoryProvider, JsonDumps, JsonLoads, SessionContext
 from .._telemetry import FeatureIndex, mark_feature_used
 from .._tools import tool
@@ -122,22 +124,6 @@ def _unescape_markdown_line(line: str) -> str:
     return line
 
 
-def _atomic_write_text(path: Path, text: str, *, encoding: str = "utf-8") -> None:
-    """Write text to ``path`` atomically via a sibling temp file plus :func:`os.replace`.
-
-    A crash, OOM, or disk-full mid-write leaves the previous file (if any) intact instead of
-    producing a truncated file that breaks every subsequent read.
-    """
-    path.parent.mkdir(parents=True, exist_ok=True)
-    temp_path = path.with_name(f"{path.name}.tmp.{os.getpid()}")
-    try:
-        temp_path.write_text(text, encoding=encoding)
-        os.replace(temp_path, path)
-    finally:
-        if temp_path.exists():
-            temp_path.unlink(missing_ok=True)
-
-
 def _slugify_topic(topic: str) -> str:
     slug = re.sub(r"[^a-z0-9]+", "-", _normalize_topic(topic).lower()).strip("-")
     return slug or "memory-topic"
@@ -661,7 +647,11 @@ def search_transcripts(
 
 @experimental(feature_id=ExperimentalFeature.HARNESS)
 class MemoryFileStore(MemoryStore):
-    """Store memory files under ``MEMORY.md``, ``topics/``, and ``transcripts/``."""
+    """Store memory files under ``MEMORY.md``, ``topics/``, and ``transcripts/``.
+
+    The storage directory must not be writable by untrusted users, because they could
+    replace the final memory files independently of this store's secure temporary writes.
+    """
 
     def __init__(
         self,

python/packages/core/agent_framework/_harness/_todo.py

@@ -4,7 +4,6 @@
 
 import asyncio
 import json
-import os
 import weakref
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, MutableMapping
@@ -15,6 +14,7 @@
 
 from .._feature_stage import ExperimentalFeature, experimental
 from .._filesystem import (
+    _atomic_write_text,  # pyright: ignore[reportPrivateUsage]
     _is_literal_storage_key_segment_safe,  # pyright: ignore[reportPrivateUsage]
     _storage_key_segment,  # pyright: ignore[reportPrivateUsage]
 )
@@ -289,7 +289,11 @@ async def save_state(self, session: AgentSession, items: list[TodoItem], *, next
 
 @experimental(feature_id=ExperimentalFeature.HARNESS)
 class TodoFileStore(TodoStore):
-    """Store todo state in one JSON file per session and source ID."""
+    """Store todo state in one JSON file per session and source ID.
+
+    The storage directory must not be writable by untrusted users, because they could
+    replace the final todo files independently of this store's secure temporary writes.
+    """
 
     def __init__(
         self,
@@ -404,16 +408,7 @@ async def save_state(self, session: AgentSession, items: list[TodoItem], *, next
     @staticmethod
     def _save_state_sync(state_path: Path, payload: str) -> None:
         """Synchronous helper that atomically writes the JSON state file."""
-        state_path.parent.mkdir(parents=True, exist_ok=True)
-        # Write to a sibling temp file then atomically replace, so a crash mid-write cannot leave
-        # a truncated state file that breaks every subsequent tool call.
-        temp_path = state_path.with_name(f"{state_path.name}.tmp.{os.getpid()}")
-        try:
-            temp_path.write_text(payload, encoding="utf-8")
-            os.replace(temp_path, state_path)
-        finally:
-            if temp_path.exists():
-                temp_path.unlink(missing_ok=True)
+        _atomic_write_text(state_path, payload)
 
 
 class TodoProvider(ContextProvider):

python/packages/core/tests/core/test_filesystem.py

@@ -19,8 +19,10 @@
 
 import pytest
 
+from agent_framework import _filesystem as filesystem_module
 from agent_framework._filesystem import (  # pyright: ignore[reportPrivateUsage]
     _MAX_ENCODED_STORAGE_KEY_SEGMENT_LENGTH,
+    _atomic_write_text,
     _is_literal_storage_key_segment_safe,
     _storage_key_segment,
 )
@@ -80,6 +82,43 @@
     "0198F0C5-1A2B-7C3D-8E4F-5A6B7C8D9E0F",
 )
 
+
+def test_atomic_write_text_rejects_preexisting_symbolic_link(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """A planted temporary-file link must not redirect a state write."""
+
+    class _FixedUuid:
+        hex = "0" * 32
+
+    monkeypatch.setattr(filesystem_module.uuid, "uuid4", _FixedUuid)
+    target_path = tmp_path / "target.txt"
+    target_path.write_text("original", encoding="utf-8")
+    temp_path = tmp_path / f".maf-state-{'0' * 32}.tmp"
+    try:
+        temp_path.symlink_to(target_path)
+    except (NotImplementedError, OSError):
+        pytest.skip("Symbolic links are not supported in this environment.")
+
+    destination_path = tmp_path / "state.json"
+    with pytest.raises(FileExistsError):
+        _atomic_write_text(destination_path, "replacement")
+
+    assert target_path.read_text(encoding="utf-8") == "original"
+    assert not destination_path.exists()
+    assert temp_path.is_symlink()
+
+
+def test_atomic_write_text_preserves_stream_initialization_error(tmp_path) -> None:
+    """A stream setup failure must retain its error and remove the temporary file."""
+    destination_path = tmp_path / "state.json"
+    destination_path.write_text("original", encoding="utf-8")
+
+    with pytest.raises(LookupError, match="unknown encoding"):
+        _atomic_write_text(destination_path, "replacement", encoding="not-a-real-encoding")
+
+    assert destination_path.read_text(encoding="utf-8") == "original"
+    assert list(tmp_path.glob(".maf-state-*.tmp")) == []
+
+
 SAFE_IDENTIFIERS: tuple[str, ...] = (
     "customer-42",
     "customer_42",

python/packages/core/tests/core/test_harness_memory.py

@@ -27,6 +27,7 @@
     MemoryTopicRecord,
     Message,
 )
+from agent_framework import _filesystem as filesystem_module
 
 from .test_filesystem import COLLIDING_IDENTIFIERS
 
@@ -820,8 +821,6 @@ def test_memory_topic_record_round_trips_when_text_contains_section_markers() ->
 
 async def test_memory_file_store_atomic_write_preserves_prior_topic_on_failure(tmp_path, monkeypatch) -> None:
     """If ``os.replace`` fails mid-write, the previous topic file must remain intact."""
-    from agent_framework._harness import _memory as memory_module
-
     store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
     session = AgentSession(session_id="session-1")
     session.state["owner_id"] = "alice"
@@ -834,12 +833,12 @@ async def test_memory_file_store_atomic_write_preserves_prior_topic_on_failure(t
     )
     store.write_topic(session, original, source_id=DEFAULT_MEMORY_SOURCE_ID)
 
-    real_replace = memory_module.os.replace
+    real_replace = filesystem_module.os.replace
 
     def _boom(*args: object, **kwargs: object) -> None:
         raise OSError("simulated disk-full")
 
-    monkeypatch.setattr(memory_module.os, "replace", _boom)
+    monkeypatch.setattr(filesystem_module.os, "replace", _boom)
     with pytest.raises(OSError, match="simulated disk-full"):
         store.write_topic(
             session,
@@ -853,12 +852,12 @@ def _boom(*args: object, **kwargs: object) -> None:
             source_id=DEFAULT_MEMORY_SOURCE_ID,
         )
 
-    monkeypatch.setattr(memory_module.os, "replace", real_replace)
+    monkeypatch.setattr(filesystem_module.os, "replace", real_replace)
     surviving = store.get_topic(session, source_id=DEFAULT_MEMORY_SOURCE_ID, topic="preferences")
     assert surviving.summary == "Prefers concise answers."
     # Temp file should not be left behind.
     topics_dir = surviving_dir = tmp_path
-    leftover = [path for path in topics_dir.rglob("*.tmp.*")]
+    leftover = [path for path in topics_dir.rglob(".maf-state-*.tmp")]
     assert leftover == []
     del surviving_dir