Python: fix(core): prevent path collisions between files and directories in InMemoryAgentFileStore

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

@@ -90,6 +90,15 @@
 # regex match when building a result snippet.
 _SEARCH_SNIPPET_RADIUS = 50
 
+# Path complexity bounds shared by every store backend. Real filesystems cap
+# these too (PATH_MAX and practical directory depth); the in-memory store's
+# directory index materializes the full key of every ancestor directory, so an
+# unbounded path would let a single write retain quadratic space in the path
+# depth. Bounding here keeps that retention constant regardless of input, and
+# rejects the same extremes the filesystem backend's OS layer would refuse.
+_MAX_PATH_SEGMENTS: Final[int] = 255
+_MAX_PATH_LENGTH: Final[int] = 4096
+
 # Hard cap on the length of a user-supplied search regex. The cap is a coarse bound on
 # how much work a single pattern can describe; it does not stop a short pathological
 # pattern, which is what :class:`_BoundedSearchPattern` is for.
@@ -296,6 +305,10 @@ def _normalize_relative_path(path: str, *, is_directory: bool = False) -> str:
         ValueError: When the path is rooted, starts with a drive letter, contains
             ``.``/``..`` segments, is empty for a file path, or ends with a
             separator for a file path.
+        OSError: With :data:`errno.ENAMETOOLONG` when the normalized path
+            exceeds ``_MAX_PATH_SEGMENTS`` segments or ``_MAX_PATH_LENGTH``
+            characters, matching what a real filesystem raises for the same
+            input.
     """
     if not path or not path.strip():
         if not is_directory:
@@ -333,6 +346,9 @@ def _normalize_relative_path(path: str, *, is_directory: bool = False) -> str:
     result = "/".join(clean_segments)
     if not is_directory and not result:
         raise ValueError(f"Invalid path: {path!r}. A file path must not be empty.")
+    if len(clean_segments) > _MAX_PATH_SEGMENTS or len(result) > _MAX_PATH_LENGTH:
+        # Check before any caller materializes ancestor prefixes of this path.
+        raise OSError(errno.ENAMETOOLONG, "File name too long", path)
     return result
 
 
@@ -1048,6 +1064,23 @@ def scan_batch(pending: list[tuple[str, str]]) -> list[FileSearchResult]:
     async def create_directory(self, path: str) -> None:
         """Ensure ``path`` exists as a directory, creating it if necessary."""
 
+    async def delete_directory(self, path: str) -> bool:
+        """Delete the empty directory at ``path``.
+
+        Args:
+            path: The relative path of the directory to delete.
+
+        Returns:
+            ``True`` if an empty directory was removed; ``False`` if it did not
+            exist, was not a directory, or still held entries. Backends that
+            cannot remove directories return ``False`` rather than raising.
+
+        ``delete`` stays file-only so backends remain interchangeable; this is
+        the supported way to release an explicitly created directory (for
+        example a session-scoped working folder) once it is empty.
+        """
+        return False
+
 
 @experimental(feature_id=ExperimentalFeature.HARNESS)
 class InMemoryAgentFileStore(AgentFileStore):
@@ -1067,12 +1100,24 @@ def __init__(self) -> None:
         # wrote, mirroring how :class:`FileSystemAgentFileStore` preserves the
         # on-disk casing.
         self._files: dict[str, tuple[str, str]] = {}
+        self._dir_counts: dict[str, int] = {}
+        self._explicit_dirs: dict[str, str] = {}
         self._lock = asyncio.Lock()
 
     @staticmethod
     def _key(path: str) -> str:
         return _normalize_relative_path(path).lower()
 
+    @staticmethod
+    def _ancestor_keys(key: str) -> list[str]:
+        """Return all strict ancestor directory keys in top-down order in linear time."""
+        ancestors: list[str] = []
+        idx = key.find("/")
+        while idx != -1:
+            ancestors.append(key[:idx])
+            idx = key.find("/", idx + 1)
+        return ancestors
+
     async def write(self, path: str, content: str, *, overwrite: bool = True) -> None:
         """Write ``content`` to the file at ``path``.
 
@@ -1083,9 +1128,19 @@ async def write(self, path: str, content: str, *, overwrite: bool = True) -> Non
         display = _normalize_relative_path(path)
         key = display.lower()
         async with self._lock:
+            ancestors = self._ancestor_keys(key)
+            for ancestor in ancestors:
+                if ancestor in self._files:
+                    raise NotADirectoryError(f"Not a directory: {path!r}")
+            if self._dir_counts.get(key, 0) > 0 or key in self._explicit_dirs:
+                raise IsADirectoryError(f"Is a directory: {path!r}")
             if not overwrite and key in self._files:
                 raise FileExistsError(f"File already exists: {path!r}")
+            is_new_file = key not in self._files
             self._files[key] = (display, content)
+            if is_new_file:
+                for ancestor in ancestors:
+                    self._dir_counts[ancestor] = self._dir_counts.get(ancestor, 0) + 1
 
     async def read(self, path: str) -> str | None:
         """Return the file content, or ``None`` if the file does not exist."""
@@ -1095,10 +1150,48 @@ async def read(self, path: str) -> str | None:
         return entry[1] if entry is not None else None
 
     async def delete(self, path: str) -> bool:
-        """Delete the file and return whether anything was removed."""
+        """Delete the file and return whether anything was removed.
+
+        Deleting is file-only, matching :class:`FileSystemAgentFileStore`: a
+        path that names a directory (explicitly created or holding stored
+        descendants) is never removed and returns ``False``. Use
+        :meth:`delete_directory` to release an empty explicitly created
+        directory.
+        """
         key = self._key(path)
         async with self._lock:
-            return self._files.pop(key, None) is not None
+            if key in self._explicit_dirs or self._dir_counts.get(key, 0) > 0:
+                return False
+            removed = self._files.pop(key, None)
+            if removed is None:
+                return False
+            for ancestor in self._ancestor_keys(key):
+                new_count = self._dir_counts[ancestor] - 1
+                if new_count <= 0:
+                    del self._dir_counts[ancestor]
+                else:
+                    self._dir_counts[ancestor] = new_count
+            return True
+
+    async def delete_directory(self, path: str) -> bool:
+        """Delete the empty explicitly created directory at ``path``.
+
+        The directory is removed only when it is empty: nothing is stored under
+        it and no other explicit directory is nested below it. This keeps the
+        file-only ``delete`` contract while giving session-scoped directories a
+        lifecycle instead of staying in the shared map forever.
+        """
+        key = self._key(_normalize_relative_path(path, is_directory=True))
+        async with self._lock:
+            if key not in self._explicit_dirs:
+                return False
+            if self._dir_counts.get(key, 0) > 0:
+                return False
+            prefix = key + "/"
+            if any(other != key and other.startswith(prefix) for other in self._explicit_dirs):
+                return False
+            del self._explicit_dirs[key]
+            return True
 
     async def list_children(self, directory: str = "") -> list[FileStoreEntry]:
         """Return the direct child files and subdirectories of ``directory``.
@@ -1118,11 +1211,12 @@ async def list_children(self, directory: str = "") -> list[FileStoreEntry]:
             prefix += "/"
         prefix_depth = prefix.count("/")
         async with self._lock:
-            entries = [(key, display) for key, (display, _) in self._files.items()]
+            file_entries = [(key, display) for key, (display, _) in self._files.items()]
+            dir_entries = list(self._explicit_dirs.items())
         files: list[str] = []
         directories: list[str] = []
         seen_dirs: set[str] = set()
-        for key, display in entries:
+        for key, display in file_entries:
             if not key.startswith(prefix):
                 continue
             # Unicode lowercasing can change character counts, so key offsets
@@ -1137,6 +1231,17 @@ async def list_children(self, directory: str = "") -> list[FileStoreEntry]:
                     continue
                 seen_dirs.add(segment_key)
                 directories.append(segment)
+        for key, display in dir_entries:
+            if not key.startswith(prefix) or key == prefix.rstrip("/"):
+                continue
+            remainder = display.split("/", prefix_depth)[-1]
+            segment, _, _ = remainder.partition("/")
+            if segment:
+                segment_key = segment.lower()
+                if segment_key in seen_dirs:
+                    continue
+                seen_dirs.add(segment_key)
+                directories.append(segment)
         results: list[FileStoreEntry] = [FileStoreEntry(name, FileStoreEntry.DIRECTORY) for name in directories]
         results.extend(FileStoreEntry(name, FileStoreEntry.FILE) for name in files)
         return results
@@ -1194,8 +1299,30 @@ def scan() -> list[FileSearchResult]:
         return await _run_search_with_timeout(asyncio.to_thread(scan))
 
     async def create_directory(self, path: str) -> None:
-        """No-op: directories are implicit from file paths in the in-memory store."""
-        del path
+        """Ensure the directory at ``path`` exists."""
+        display = _normalize_relative_path(path, is_directory=True)
+        key = display.lower()
+        if not key:
+            return
+        async with self._lock:
+            ancestors = self._ancestor_keys(key)
+            for ancestor in ancestors:
+                if ancestor in self._files:
+                    raise FileExistsError(f"File already exists: {ancestor!r}")
+            if key in self._files:
+                raise FileExistsError(f"File already exists: {key!r}")
+            if key not in self._explicit_dirs:
+                idx_display = 0
+                for ancestor in ancestors:
+                    next_slash = display.find("/", idx_display)
+                    if next_slash != -1:
+                        ancestor_display = display[:next_slash]
+                        idx_display = next_slash + 1
+                    else:
+                        ancestor_display = display
+                    if ancestor not in self._explicit_dirs:
+                        self._explicit_dirs[ancestor] = ancestor_display
+                self._explicit_dirs[key] = display
 
 
 @experimental(feature_id=ExperimentalFeature.HARNESS)
@@ -1444,6 +1571,25 @@ def _delete_file_sync(cls, full_path: Path) -> bool:
                 return False
         return True
 
+    async def delete_directory(self, path: str) -> bool:
+        """Delete the empty directory at ``path``.
+
+        Returns ``True`` only when an empty directory was removed; a directory
+        that is missing, is a file, or still holds entries returns ``False``
+        instead of raising, matching :class:`InMemoryAgentFileStore`.
+        """
+        full_path = self._resolve_safe_directory_path(path)
+        return await asyncio.to_thread(self._delete_directory_sync, full_path)
+
+    @staticmethod
+    def _delete_directory_sync(full_path: Path) -> bool:
+        try:
+            full_path.rmdir()
+        except OSError:
+            # Missing path, a file, or a directory that still holds entries.
+            return False
+        return True
+
     async def list_children(self, directory: str = "") -> list[FileStoreEntry]:
         """Return the direct child files and subdirectories of ``directory``.