diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 0bc85e26f19..368ab5acc73 100644 --- a/python/packages/core/AGENTS.md +++ b/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 diff --git a/python/packages/core/agent_framework/_filesystem.py b/python/packages/core/agent_framework/_filesystem.py index 512ca17ff7f..c4bf4c89244 100644 --- a/python/packages/core/agent_framework/_filesystem.py +++ b/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() diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index 66957756261..c3f02f883b5 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/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, diff --git a/python/packages/core/agent_framework/_harness/_todo.py b/python/packages/core/agent_framework/_harness/_todo.py index daa554a585e..7aaeeecb2a3 100644 --- a/python/packages/core/agent_framework/_harness/_todo.py +++ b/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): diff --git a/python/packages/core/tests/core/test_filesystem.py b/python/packages/core/tests/core/test_filesystem.py index 764f8307aac..0278e6c3a78 100644 --- a/python/packages/core/tests/core/test_filesystem.py +++ b/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", diff --git a/python/packages/core/tests/core/test_harness_memory.py b/python/packages/core/tests/core/test_harness_memory.py index 400c9fc6149..464255d0b43 100644 --- a/python/packages/core/tests/core/test_harness_memory.py +++ b/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