Harden FileAccess search and atomic save in store API

- Add wall-clock timeout (10s) around regex scans so a pathological pattern (e.g. `(a+)+`) below the length cap cannot stall the event loop.
- Offload the InMemoryAgentFileStore regex scan to a worker thread, matching the filesystem store.
- Fail closed when `Path.is_symlink` raises during the safe-path probe so a permission error cannot silently bypass the symlink/reparse-point rejection.
- Add `overwrite: bool = True` to `AgentFileStore.write_file`; the in-memory store performs the check under the existing lock and the filesystem store uses `open(mode='x')` so concurrent callers cannot race past `overwrite=False`.
- `file_access_save_file` now relies on the atomic store call instead of a separate `file_exists` round-trip.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: alliscode

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

@@ -25,7 +25,7 @@
 import os
 import re
 from abc import ABC, abstractmethod
-from collections.abc import Mapping, MutableMapping
+from collections.abc import Callable, Mapping, MutableMapping
 from pathlib import Path
 from typing import Any, cast
 
@@ -53,11 +53,14 @@
 
 # Hard cap on the length of a user-supplied search regex. Python's ``re`` module
 # has no built-in timeout, so a catastrophic-backtracking pattern (such as
-# ``(a+)+$``) submitted by the model could block the event loop indefinitely.
-# Capping pattern length is a simple, predictable mitigation; pathological
-# ReDoS patterns are typically far shorter than this limit, so the cap mostly
-# rejects obviously-malformed input while still allowing realistic queries.
+# ``(a+)+$``) submitted by the model could spin the CPU indefinitely. The cap
+# alone does not stop short pathological patterns, so :meth:`search_files`
+# additionally executes the regex scan in a worker thread and bounds the wall
+# clock with :data:`_SEARCH_TIMEOUT_SECONDS`. The thread itself cannot be
+# safely interrupted from Python, so a runaway scan continues until the
+# regex engine returns, but the caller and event loop stay responsive.
 _MAX_SEARCH_PATTERN_LENGTH = 256
+_SEARCH_TIMEOUT_SECONDS = 10.0
 
 
 def _compile_search_regex(pattern: str) -> re.Pattern[str]:
@@ -75,6 +78,24 @@ def _compile_search_regex(pattern: str) -> re.Pattern[str]:
     return re.compile(pattern, flags=re.IGNORECASE)
 
 
+async def _run_search_with_timeout(
+    fn: Callable[[], list[FileSearchResult]],
+) -> list[FileSearchResult]:
+    """Run ``fn`` in a worker thread with a bounded wall-clock timeout.
+
+    Raises:
+        ValueError: When the search does not complete within
+            :data:`_SEARCH_TIMEOUT_SECONDS` seconds.
+    """
+    try:
+        return await asyncio.wait_for(asyncio.to_thread(fn), timeout=_SEARCH_TIMEOUT_SECONDS)
+    except TimeoutError as exc:
+        raise ValueError(
+            f"Regex search did not complete within {_SEARCH_TIMEOUT_SECONDS:g} seconds. "
+            "Use a more specific pattern (avoid nested quantifiers such as '(a+)+')."
+        ) from exc
+
+
 def _normalize_relative_path(path: str, *, is_directory: bool = False) -> str:
     """Normalize and validate a relative store path.
 
@@ -316,12 +337,22 @@ class AgentFileStore(ABC):
     """
 
     @abstractmethod
-    async def write_file(self, path: str, content: str) -> None:
-        """Write ``content`` to the file at ``path``, creating or overwriting it.
+    async def write_file(self, path: str, content: str, *, overwrite: bool = True) -> None:
+        """Write ``content`` to the file at ``path``.
 
         Args:
             path: The relative path of the file to write.
             content: The content to write to the file.
+
+        Keyword Args:
+            overwrite: When ``True`` (default) any existing file is replaced.
+                When ``False`` the implementation must perform an atomic
+                exclusive create and raise :class:`FileExistsError` if a file
+                already exists at ``path``.
+
+        Raises:
+            FileExistsError: When ``overwrite`` is ``False`` and a file already
+                exists at ``path``.
         """
 
     @abstractmethod
@@ -413,10 +444,17 @@ def __init__(self) -> None:
     def _key(path: str) -> str:
         return _normalize_relative_path(path).lower()
 
-    async def write_file(self, path: str, content: str) -> None:
-        """Write ``content`` to the file at ``path``."""
+    async def write_file(self, path: str, content: str, *, overwrite: bool = True) -> None:
+        """Write ``content`` to the file at ``path``.
+
+        When ``overwrite`` is ``False`` the check-and-write happens under the
+        store lock so concurrent callers cannot both observe a missing file
+        and race to create it.
+        """
         key = self._key(path)
         async with self._lock:
+            if not overwrite and key in self._files:
+                raise FileExistsError(f"File already exists: {path!r}")
             self._files[key] = content
 
     async def read_file(self, path: str) -> str | None:
@@ -452,7 +490,12 @@ async def search_files(
         regex_pattern: str,
         file_pattern: str | None = None,
     ) -> list[FileSearchResult]:
-        """Search file contents for ``regex_pattern`` matches."""
+        """Search file contents for ``regex_pattern`` matches.
+
+        Snapshots the entries under the store lock and offloads the regex scan
+        to a worker thread with a bounded timeout so a pathological pattern
+        cannot stall the event loop.
+        """
         prefix = _normalize_relative_path(directory, is_directory=True).lower()
         if prefix and not prefix.endswith("/"):
             prefix += "/"
@@ -461,19 +504,22 @@ async def search_files(
         async with self._lock:
             entries = list(self._files.items())
 
-        results: list[FileSearchResult] = []
-        for key, file_content in entries:
-            if not key.startswith(prefix):
-                continue
-            relative_name = key[len(prefix) :]
-            if "/" in relative_name:
-                continue
-            if not _matches_glob(relative_name, file_pattern):
-                continue
-            result = _search_file_content(relative_name, file_content, regex)
-            if result is not None:
-                results.append(result)
-        return results
+        def scan() -> list[FileSearchResult]:
+            results: list[FileSearchResult] = []
+            for key, file_content in entries:
+                if not key.startswith(prefix):
+                    continue
+                relative_name = key[len(prefix) :]
+                if "/" in relative_name:
+                    continue
+                if not _matches_glob(relative_name, file_pattern):
+                    continue
+                result = _search_file_content(relative_name, file_content, regex)
+                if result is not None:
+                    results.append(result)
+            return results
+
+        return await _run_search_with_timeout(scan)
 
     async def create_directory(self, path: str) -> None:
         """No-op: directories are implicit from file paths in the in-memory store."""
@@ -567,26 +613,38 @@ def _throw_if_contains_symlink(self, candidate: Path) -> None:
         for segment in relative_parts:
             current = current / segment
             try:
-                if current.is_symlink():
-                    raise ValueError("Invalid path: the resolved path contains a symbolic link or reparse point.")
-            except OSError:
-                # Permission errors and similar transient OS errors during the
-                # symlink probe should not silently allow the access; treat as
-                # missing and stop checking so the underlying I/O surfaces the
-                # real error.
-                break
+                is_link = current.is_symlink()
+            except OSError as exc:
+                # Fail closed: if we cannot verify whether a segment is a
+                # symlink/reparse point we refuse the operation rather than
+                # silently allow access that may escape the root.
+                raise ValueError(
+                    f"Invalid path: unable to verify whether '{segment}' is a symbolic link or reparse point."
+                ) from exc
+            if is_link:
+                raise ValueError("Invalid path: the resolved path contains a symbolic link or reparse point.")
             if not current.exists():
                 break
 
-    async def write_file(self, path: str, content: str) -> None:
-        """Write ``content`` to the file at ``path``."""
+    async def write_file(self, path: str, content: str, *, overwrite: bool = True) -> None:
+        """Write ``content`` to the file at ``path``.
+
+        When ``overwrite`` is ``False`` the file is created using ``mode="x"``
+        so the underlying ``open`` call performs an atomic exclusive create
+        (``O_EXCL`` on POSIX, ``CREATE_NEW`` on Windows) and raises
+        :class:`FileExistsError` if a file already exists.
+        """
         full_path = self._resolve_safe_path(path)
-        await asyncio.to_thread(self._write_file_sync, full_path, content)
+        await asyncio.to_thread(self._write_file_sync, full_path, content, overwrite)
 
     @staticmethod
-    def _write_file_sync(full_path: Path, content: str) -> None:
+    def _write_file_sync(full_path: Path, content: str, overwrite: bool) -> None:
         full_path.parent.mkdir(parents=True, exist_ok=True)
-        full_path.write_text(content, encoding="utf-8")
+        if overwrite:
+            full_path.write_text(content, encoding="utf-8")
+            return
+        with full_path.open("x", encoding="utf-8") as handle:
+            handle.write(content)
 
     async def read_file(self, path: str) -> str | None:
         """Return the file content, or ``None`` if the file does not exist."""
@@ -646,7 +704,7 @@ async def search_files(
         """Search file contents for ``regex_pattern`` matches."""
         full_dir = self._resolve_safe_directory_path(directory)
         regex = _compile_search_regex(regex_pattern)
-        return await asyncio.to_thread(self._search_files_sync, full_dir, regex, file_pattern)
+        return await _run_search_with_timeout(lambda: self._search_files_sync(full_dir, regex, file_pattern))
 
     @staticmethod
     def _search_files_sync(full_dir: Path, regex: re.Pattern[str], file_pattern: str | None) -> list[FileSearchResult]:
@@ -736,9 +794,10 @@ async def before_run(
         async def file_access_save_file(file_name: str, content: str, overwrite: bool = False) -> str:
             """Save a file with the given name and content. By default, does not overwrite an existing file unless overwrite is set to true."""  # noqa: E501
             normalized = _normalize_relative_path(file_name)
-            if not overwrite and await self.store.file_exists(normalized):
+            try:
+                await self.store.write_file(normalized, content, overwrite=overwrite)
+            except FileExistsError:
                 return f"File '{file_name}' already exists. To replace it, save again with overwrite set to true."
-            await self.store.write_file(normalized, content)
             return f"File '{file_name}' saved."
 
         @tool(name="file_access_read_file", approval_mode="never_require")

python/packages/core/tests/core/test_harness_file_access.py

@@ -4,6 +4,7 @@
 
 import json
 import re
+import time
 from pathlib import Path
 
 import pytest
@@ -21,11 +22,13 @@
     Message,
     SupportsChatGetResponse,
 )
+from agent_framework._harness import _file_access as _file_access_module
 from agent_framework._harness._file_access import (
     DEFAULT_FILE_ACCESS_INSTRUCTIONS,
     DEFAULT_FILE_ACCESS_SOURCE_ID,
     _matches_glob,
     _normalize_relative_path,
+    _run_search_with_timeout,
 )
 
 
@@ -393,6 +396,64 @@ async def test_file_access_provider_accepts_custom_instructions() -> None:
     assert provider.source_id == DEFAULT_FILE_ACCESS_SOURCE_ID
 
 
+async def test_in_memory_store_write_file_raises_when_exists_and_no_overwrite() -> None:
+    """The atomic exclusive-create path should raise ``FileExistsError`` under the lock."""
+    store = InMemoryAgentFileStore()
+    await store.write_file("plan.md", "v1")
+
+    with pytest.raises(FileExistsError):
+        await store.write_file("plan.md", "v2", overwrite=False)
+
+    # The original content is preserved.
+    assert await store.read_file("plan.md") == "v1"
+
+    # Default ``overwrite=True`` still replaces.
+    await store.write_file("plan.md", "v3")
+    assert await store.read_file("plan.md") == "v3"
+
+
+async def test_filesystem_store_write_file_raises_when_exists_and_no_overwrite(tmp_path: Path) -> None:
+    """The filesystem store should use exclusive-create semantics when ``overwrite=False``."""
+    store = FileSystemAgentFileStore(tmp_path)
+    await store.write_file("plan.md", "v1")
+
+    with pytest.raises(FileExistsError):
+        await store.write_file("plan.md", "v2", overwrite=False)
+
+    assert (tmp_path / "plan.md").read_text(encoding="utf-8") == "v1"
+
+    await store.write_file("plan.md", "v3", overwrite=True)
+    assert (tmp_path / "plan.md").read_text(encoding="utf-8") == "v3"
+
+
+async def test_run_search_with_timeout_raises_value_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A scan that exceeds the timeout should surface a clean ``ValueError``."""
+    monkeypatch.setattr(_file_access_module, "_SEARCH_TIMEOUT_SECONDS", 0.01)
+
+    def slow() -> list[FileSearchResult]:
+        time.sleep(0.5)
+        return []
+
+    with pytest.raises(ValueError, match="did not complete"):
+        await _run_search_with_timeout(slow)
+
+
+async def test_filesystem_store_symlink_probe_fails_closed_on_oserror(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """If ``Path.is_symlink`` raises during the probe, the operation must be refused."""
+    store = FileSystemAgentFileStore(tmp_path)
+    await store.write_file("ok.txt", "content")
+
+    def boom(self: Path) -> bool:
+        raise PermissionError("access denied")
+
+    monkeypatch.setattr(Path, "is_symlink", boom)
+
+    with pytest.raises(ValueError, match="symbolic link or reparse point"):
+        await store.read_file("ok.txt")
+
+
 def test_file_access_harness_classes_are_marked_experimental() -> None:
     """File-access harness public classes should expose HARNESS experimental metadata."""
     assert AgentFileStore.__feature_id__ == ExperimentalFeature.HARNESS.value