From 4385ebdba135231dd57c344b903f81b85d5898df Mon Sep 17 00:00:00 2001 From: quifox <289420841+quifox@users.noreply.github.com> Date: Fri, 2 Oct 2026 11:02:21 +0000 Subject: [PATCH 1/2] Python: Bound harness file memory with quotas and retention Add opt-in quotas, shared retention management, and server-owned cleanup with request lifetime protection. Preserve SDK defaults and cover limits, TTL transitions, concurrent operations, interruptions, and hosted streams. --- python/packages/core/AGENTS.md | 2 + .../packages/core/agent_framework/__init__.py | 2 + .../core/agent_framework/__init__.pyi | 2 + .../agent_framework/_harness/_file_access.py | 168 ++- .../agent_framework/_harness/_file_memory.py | 170 +++- .../_harness/_file_memory_retention.py | 955 ++++++++++++++++++ .../_harness/_file_store_limits.py | 63 ++ python/packages/core/pyproject.toml | 1 + .../tests/core/test_harness_file_memory.py | 24 + .../test_harness_file_memory_retention.py | 803 +++++++++++++++ .../core/test_harness_file_store_limits.py | 195 ++++ .../test_harness_hosted_memory_lifetime.py | 175 ++++ .../claw_step04_production_ready/README.md | 83 ++ .../claw_step04_production_ready/agent.py | 9 + .../claw_step04_production_ready/hosted.py | 45 +- python/uv.lock | 11 + 16 files changed, 2662 insertions(+), 46 deletions(-) create mode 100644 python/packages/core/agent_framework/_harness/_file_memory_retention.py create mode 100644 python/packages/core/agent_framework/_harness/_file_store_limits.py create mode 100644 python/packages/core/tests/core/test_harness_file_memory_retention.py create mode 100644 python/packages/core/tests/core/test_harness_file_store_limits.py create mode 100644 python/packages/core/tests/core/test_harness_hosted_memory_lifetime.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 13361bbeaa9..9dbf104945b 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -223,6 +223,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **`AgentFileStore`** - Abstract async store backing the file-access harness. Implementations expose `write`, `read`, `delete`, `list_children`, `file_exists`, `search`, and `create_directory` over forward-slash relative paths. `list_children` returns the direct children (files and subdirectories, subdirectories first) as `FileStoreEntry` instances; `search` accepts a keyword-only `recursive` flag (default `False`) and, when `recursive=True`, walks all descendants and returns `file_name` values relative to the search directory. The line-numbering contract lives on the base class: `split_lines` publishes the `\n`-only keepends split that every `line_number` addresses, `scan_content` is the numbering primitive both shipped stores report through, and `search` is now **concrete** — it asks the overridable `find_matching_files` hook which files to consider (superset semantics; a backend with a native index overrides it and prunes server-side) and then reads and numbers them itself. A store may still override `search` outright, but then it owns numbering: it must report `line_number` as a 1-based coordinate into `split_lines` of the content `read` returns. Nothing checks that at run time, so a store that numbers differently makes a later edit land on the wrong line silently. The same applies to ReDoS: the grep pattern is model-supplied, so `search` compiles it through the `regex` module against a single monotonic deadline that bounds the whole scan (CPython's `re` holds the GIL for an entire match, so `asyncio.wait_for` around it bounds nothing). A store overriding `search` with bare `re` silently opts itself back out of that guarantee. - **`InMemoryAgentFileStore`** - Dict-backed store suitable for tests and lightweight scenarios. - **`FileSystemAgentFileStore`** - Disk-backed store rooted under a configurable directory. Enforces relative-path normalization, root containment, and rejects symlink/reparse-point segments to prevent escape. +- **Opt-in file-store capacity** - The two built-in stores accept keyword-only `max_file_bytes`, `max_files`, and `max_total_bytes`; each defaults to `None`. Limits count stored content bytes and files (including file-memory sidecars/indexes) under the store root, not filesystem overhead. The filesystem implementation measures existing usage under a local inter-process lock; cooperating writers sharing a root must use the same policy. The fixed root lock file is reserved and excluded from tool listings and quota usage. Whole-file reads enforce `max_file_bytes` before loading oversized data. Third-party stores retain the existing interface and own their backend's capacity guarantees. + FileMemoryRetentionManager is an opt-in local-filesystem owner of shared capacity, durable TTL policy, request protection, periodic GC, and explicit repair; FileMemoryProvider keeps its existing default behavior. - **`FileSearchResult`** / **`FileSearchMatch`** - `SerializationMixin` DTOs returned by `search`, carrying the matching file name, a context snippet, and the matching lines with 1-based line numbers. Implementers should report each matching line verbatim, including its own terminator, so it can be reused as a `file_access_replace_lines` `new_line`; the pattern itself is matched against the line with its whole terminator removed, so `^`/`$` anchor to the line's text on a CRLF file as they already did on an LF one. A custom store populates these DTOs from its own `search`; the verbatim text is a recommendation, but the line number is not — it must address `split_lines`. - **`FileStoreEntry`** - `SerializationMixin` DTO returned by `list_children`, carrying an entry `name` and `type` (`"file"` or `"directory"`). - **`FileAccessProvider`** - `ContextProvider` that adds shared file-access tools (`file_access_write`, `file_access_read`, `file_access_read_lines`, `file_access_delete`, `file_access_ls`, `file_access_grep`, `file_access_replace`, `file_access_replace_lines`) plus default usage instructions to each invocation. `file_access_ls` enumerates direct children (both files and subdirectories) as `{name, type}` entries with an optional `glob_pattern`, so the agent can walk the tree level by level; `file_access_grep` searches recursively from an optional base `directory` and returns relative `file_name` paths, scoped via an `fnmatch` `glob_pattern` (where `*` crosses `/`, e.g. `*.md`, `reports/*`). `file_access_replace` substitutes `old_string` with `new_string` (failing if not found, or if multiple matches and `replace_all` is false); `file_access_replace_lines` replaces whole 1-based lines with literal text (each `new_line` includes its own trailing newline; an empty `new_line` deletes the line, including its line break). `file_access_read_lines` returns a 1-based inclusive line range, one line per row as `\t`; `end_line` may be omitted to read to the end of the file, and an `end_line` past the last line clamps to it. Everything after the tab is verbatim, including the line's own terminator (which therefore doubles as the row separator), so a row's text can be fed straight back as a `file_access_replace_lines` `new_line` without losing a `\r\n`. Its line numbering comes from the same `_split_lines_keepends` split as `file_access_replace_lines` and as the stores in this package, so with one of those a number reported by grep addresses the same line in all three tools, including the trailing empty line of a newline-terminated file; grep itself runs through `AgentFileStore.search`, which must number by the same split but does not inherit it, so a store overriding `search` owns its numbering and nothing verifies it at run time. All tools are registered with `approval_mode="always_require"` by default, so every file operation needs host approval. Pass `disable_write_tools=True` to advertise only the read-only tools. To run unattended you can disable approval at the source with `disable_readonly_tool_approval=True` (read, read_lines, ls, grep) and/or `disable_write_tool_approval=True` (write, delete, replace, replace_lines), which register the affected tools with `approval_mode="never_require"`; alternatively, keep approval on and pass one of the static auto-approval rules to `ToolApprovalMiddleware` (via `auto_approval_rules`): `FileAccessProvider.read_only_tools_auto_approval_rule` approves only the read-only tools (read, read_lines, ls, grep), while `FileAccessProvider.all_tools_auto_approval_rule` approves every file-access tool including the write tools. Both rules reject any call carrying a `server_label` so they stay scoped to this provider's local tools and never auto-approve a same-named hosted tool. The tool names are also exposed as class constants (`WRITE_TOOL_NAME`, `READ_TOOL_NAME`, `READ_LINES_TOOL_NAME`, `DELETE_TOOL_NAME`, `LS_TOOL_NAME`, `GREP_TOOL_NAME`, `REPLACE_TOOL_NAME`, `REPLACE_LINES_TOOL_NAME`). Unlike `MemoryContextProvider`, the store is intentionally shared across sessions and agents; pass `session_scoped=True` (with an optional explicit `scope`) to confine tool operations to a working folder derived from the session id or scope via the shared `_storage_key_segment` derivation (the provider fails closed when neither is available). diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 3953187578e..867bed804fb 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -126,6 +126,7 @@ "FileSystemAgentFileStore", "InMemoryAgentFileStore", ), + "._harness._file_memory_retention": ("FileMemoryRetentionManager",), "._harness._file_memory": ( "DEFAULT_FILE_MEMORY_INSTRUCTIONS", "DEFAULT_FILE_MEMORY_SOURCE_ID", @@ -511,6 +512,7 @@ "FileCheckpointStorage", "FileHistoryProvider", "FileMemoryProvider", + "FileMemoryRetentionManager", "FileSearchMatch", "FileSearchResult", "FileSessionStore", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index e8d1bc6e411..f3bd560e742 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -89,6 +89,7 @@ from ._harness._file_access import ( InMemoryAgentFileStore, ) from ._harness._file_memory import DEFAULT_FILE_MEMORY_INSTRUCTIONS, DEFAULT_FILE_MEMORY_SOURCE_ID, FileMemoryProvider +from ._harness._file_memory_retention import FileMemoryRetentionManager from ._harness._loop import ( AgentLoopMiddleware, JudgeVerdict, @@ -466,6 +467,7 @@ __all__ = [ "FileCheckpointStorage", "FileHistoryProvider", "FileMemoryProvider", + "FileMemoryRetentionManager", "FileSearchMatch", "FileSearchResult", "FileSessionStore", diff --git a/python/packages/core/agent_framework/_harness/_file_access.py b/python/packages/core/agent_framework/_harness/_file_access.py index 244ee5b983a..278255a115f 100644 --- a/python/packages/core/agent_framework/_harness/_file_access.py +++ b/python/packages/core/agent_framework/_harness/_file_access.py @@ -29,7 +29,8 @@ import threading import time from abc import ABC, abstractmethod -from collections.abc import Awaitable, Mapping, MutableMapping +from collections.abc import Awaitable, Generator, Mapping, MutableMapping +from contextlib import contextmanager from pathlib import Path from typing import Annotated, Any, ClassVar, Final, Protocol, cast @@ -43,10 +44,12 @@ from .._telemetry import FeatureIndex, mark_feature_used from .._tools import ApprovalMode, tool from .._types import Content +from ._file_store_limits import _FileStoreLimits, _FileStoreQuotaError # pyright: ignore[reportPrivateUsage] logger = logging.getLogger(__name__) DEFAULT_FILE_ACCESS_SOURCE_ID = "file_access" +_FILE_STORE_QUOTA_LOCK: Final[str] = ".agent-framework-quota.lock" DEFAULT_FILE_ACCESS_INSTRUCTIONS = ( "## File Access\n" "You have access to a shared file storage area via the `file_access_*` tools " @@ -1017,6 +1020,8 @@ def scan_batch(pending: list[tuple[str, str]]) -> list[FileSearchResult]: continue try: content = await self.read(_combine_search_path(directory, name)) + except _FileStoreQuotaError: + raise except (OSError, ValueError): # ``read`` raises ValueError for non-UTF-8 bytes and for symlinked paths, and # OSError for a file deleted or made unreadable under us. Skip either so one @@ -1056,8 +1061,25 @@ class InMemoryAgentFileStore(AgentFileStore): directory structure is maintained. """ - def __init__(self) -> None: - """Initialize an empty in-memory file store.""" + def __init__( + self, + *, + max_file_bytes: int | None = None, + max_files: int | None = None, + max_total_bytes: int | None = None, + ) -> None: + """Initialize an empty in-memory file store. + + Keyword Args: + max_file_bytes: Maximum UTF-8 bytes per file and whole-file read, or ``None`` for no limit. + max_files: Maximum physical file count under this store, or ``None`` for no limit. + max_total_bytes: Maximum combined UTF-8 bytes under this store, or ``None`` for no limit. + + Raises: + ValueError: If an explicit limit is not a positive integer. + """ + self._limits = _FileStoreLimits.create(max_file_bytes, max_files, max_total_bytes) + self._file_sizes: dict[str, int] = {} # Keys are case-insensitive (normalized + lowercased) so the store # behaves consistently on case-insensitive deployments. Each entry # also records the *original* normalized path so ``list_children`` and @@ -1083,6 +1105,15 @@ async def write(self, path: str, content: str, *, overwrite: bool = True) -> Non async with self._lock: if not overwrite and key in self._files: raise FileExistsError(f"File already exists: {path!r}") + if self._limits is not None: + size = self._limits.content_size(content) + self._limits.check_usage( + len(self._files) + int(key not in self._files), + sum(self._file_sizes.values()) - self._file_sizes.get(key, 0) + size + if self._limits.max_total_bytes is not None + else 0, + ) + self._file_sizes[key] = size self._files[key] = (display, content) async def read(self, path: str) -> str | None: @@ -1090,12 +1121,15 @@ async def read(self, path: str) -> str | None: key = self._key(path) async with self._lock: entry = self._files.get(key) + if entry is not None and self._limits is not None: + self._limits.check_read(self._file_sizes[key]) 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.""" key = self._key(path) async with self._lock: + self._file_sizes.pop(key, None) return self._files.pop(key, None) is not None async def list_children(self, directory: str = "") -> list[FileStoreEntry]: @@ -1222,7 +1256,14 @@ class FileSystemAgentFileStore(AgentFileStore): # Case aliases can identify the same file while having different Path hashes. _DELETE_LOCK: ClassVar[threading.Lock] = threading.Lock() - def __init__(self, root_directory: str | os.PathLike[str]) -> None: + def __init__( + self, + root_directory: str | os.PathLike[str], + *, + max_file_bytes: int | None = None, + max_files: int | None = None, + max_total_bytes: int | None = None, + ) -> None: """Initialize the file-system store. The root directory is **not** created here; construction performs no @@ -1234,7 +1275,22 @@ def __init__(self, root_directory: str | os.PathLike[str]) -> None: Args: root_directory: The directory under which all files are stored. Created lazily on first write if it does not exist. + + Keyword Args: + max_file_bytes: Maximum bytes per file and whole-file read, or ``None`` for no limit. + max_files: Maximum regular-file count recursively under the root, or ``None`` for no limit. + max_total_bytes: Maximum regular-file content bytes under the root, or ``None`` for no limit. + + Limits include existing files, memory descriptions, and indexes. Configured stores sharing + a root must use the same limits; all writers must cooperate through these stores. Accounting + is serialized across local processes, and the reserved root lock file is hidden from tools. + Limits do not cover filesystem overhead, other roots, or direct external filesystem writes. + No quota metadata is created when all limits are ``None``. + + Raises: + ValueError: If an explicit limit is not a positive integer. """ + self._limits = _FileStoreLimits.create(max_file_bytes, max_files, max_total_bytes) raw_root = os.fspath(root_directory) if not raw_root or not raw_root.strip(): raise ValueError("root_directory must not be empty or whitespace-only.") @@ -1347,7 +1403,77 @@ async def write(self, path: str, content: str, *, overwrite: bool = True) -> Non the obvious probe-then-open race for the file itself. """ full_path = self._resolve_safe_path(path) - await asyncio.to_thread(self._write_file_sync, full_path, content, overwrite) + self._reject_quota_metadata(full_path) + if self._limits is None: + await asyncio.to_thread(self._write_file_sync, full_path, content, overwrite) + else: + await asyncio.to_thread(self._write_limited_sync, full_path, content, overwrite) + + def _is_quota_metadata(self, path: Path) -> bool: + lock_path = self._root_path / _FILE_STORE_QUOTA_LOCK + return path == lock_path or (path.exists() and lock_path.exists() and path.samefile(lock_path)) + + def _reject_quota_metadata(self, path: Path) -> None: + if self._limits is not None and self._is_quota_metadata(path): + raise ValueError("This path is reserved for file-store quota accounting.") + + @contextmanager + def _quota_lock(self) -> Generator[None]: + """Coordinate cooperating local processes without allocating quota state per request. + + FileLock uses OS-backed local locks; do not remove the lock file after release. + See https://py-filelock.readthedocs.io/en/latest/ for supported filesystem guarantees. + """ + from filelock import FileLock, Timeout + + root = self._resolve_safe_directory_path("") + root.mkdir(parents=True, exist_ok=True) + lock_path = self._resolve_safe_path(_FILE_STORE_QUOTA_LOCK) + if lock_path.exists() and lock_path.stat().st_size: + raise ValueError("An existing file conflicts with file-store quota metadata.") + try: + with FileLock(lock_path, timeout=10, preserve_lock_file=True, fallback_to_soft=False): + yield + except Timeout as exc: + raise TimeoutError("Timed out acquiring the file-store quota lock.") from exc + + def _quota_usage(self) -> tuple[int, int]: + """Read current sizes under the root lock, including data left by earlier instances.""" + file_count = total_bytes = 0 + directories = [self._root_path] + while directories: + for entry in directories.pop().iterdir(): + if self._is_quota_metadata(entry) or _is_link_or_reparse_point(entry): + continue + if entry.is_dir(): + directories.append(entry) + elif entry.is_file(): + file_count += 1 + total_bytes += entry.stat().st_size + return file_count, total_bytes + + def _write_limited_sync(self, path: Path, content: str, overwrite: bool) -> None: + limits = cast(_FileStoreLimits, self._limits) + with self._quota_lock(): + path = self._resolve_safe_path(path.relative_to(self._root_path).as_posix()) + if path.is_dir(): + raise IsADirectoryError(f"Path is a directory: {path}") + exists = path.is_file() + if not overwrite and exists: + raise FileExistsError(f"File already exists: {path.name!r}") + size = limits.content_size(content) + if limits.max_files is not None or limits.max_total_bytes is not None: + file_count, total_bytes = self._quota_usage() + previous_size = path.stat().st_size if exists else 0 + limits.check_usage(file_count + int(not exists), total_bytes - previous_size + size) + self._write_file_sync(path, content, overwrite) + + def _delete_limited_sync(self, path: Path) -> bool: + if not self._root_path.exists(): + return False + with self._quota_lock(): + path = self._resolve_safe_path(path.relative_to(self._root_path).as_posix()) + return self._delete_file_sync(path) @staticmethod def _write_file_sync(full_path: Path, content: str, overwrite: bool) -> None: @@ -1406,10 +1532,13 @@ async def read(self, path: str) -> str | None: recoverable string response rather than a stack trace. """ full_path = self._resolve_safe_path(path) - return await asyncio.to_thread(self._read_file_sync, full_path) + self._reject_quota_metadata(full_path) + if self._limits is None or self._limits.max_file_bytes is None: + return await asyncio.to_thread(self._read_file_sync, full_path) + return await asyncio.to_thread(self._read_file_sync, full_path, self._limits.max_file_bytes) @staticmethod - def _read_file_sync(full_path: Path) -> str | None: + def _read_file_sync(full_path: Path, max_file_bytes: int | None = None) -> str | None: if not full_path.is_file(): return None nofollow = getattr(os, "O_NOFOLLOW", 0) @@ -1420,7 +1549,15 @@ def _read_file_sync(full_path: Path) -> str | None: raise ValueError("Invalid path: the resolved path contains a symbolic link or reparse point.") from exc raise with os.fdopen(fd, "rb") as handle: - raw = handle.read() + if max_file_bytes is None: + raw = handle.read() + else: + size = os.fstat(handle.fileno()).st_size + if size > max_file_bytes: + raise _FileStoreQuotaError("max_file_bytes", max_file_bytes, size) + raw = handle.read(max_file_bytes + 1) + if len(raw) > max_file_bytes: + raise _FileStoreQuotaError("max_file_bytes", max_file_bytes, len(raw)) try: return raw.decode("utf-8") except UnicodeDecodeError as exc: @@ -1429,7 +1566,10 @@ def _read_file_sync(full_path: Path) -> str | None: async def delete(self, path: str) -> bool: """Delete the file and return whether anything was removed.""" full_path = self._resolve_safe_path(path) - return await asyncio.to_thread(self._delete_file_sync, full_path) + self._reject_quota_metadata(full_path) + if self._limits is None: + return await asyncio.to_thread(self._delete_file_sync, full_path) + return await asyncio.to_thread(self._delete_limited_sync, full_path) @classmethod def _delete_file_sync(cls, full_path: Path) -> bool: @@ -1450,7 +1590,10 @@ async def list_children(self, directory: str = "") -> list[FileStoreEntry]: escapes the root. An empty list is returned for a non-existent directory. """ full_dir = self._resolve_safe_directory_path(directory) - return await asyncio.to_thread(self._list_sync, full_dir) + entries = await asyncio.to_thread(self._list_sync, full_dir) + if self._limits is not None and full_dir == self._root_path: + entries = [entry for entry in entries if not self._is_quota_metadata(full_dir / entry.name)] + return entries @staticmethod def _list_sync(full_dir: Path) -> list[FileStoreEntry]: @@ -1475,6 +1618,8 @@ def _list_sync(full_dir: Path) -> list[FileStoreEntry]: async def file_exists(self, path: str) -> bool: """Return whether the file exists.""" full_path = self._resolve_safe_path(path) + if self._limits is not None and self._is_quota_metadata(full_path): + return False return await asyncio.to_thread(self._file_exists_sync, full_path) @staticmethod @@ -1500,6 +1645,8 @@ async def search( ``True`` all descendant files are searched, otherwise only the direct children. """ + if self._limits is not None: + return await super().search(directory, regex_pattern, glob_pattern, recursive=recursive) full_dir = self._resolve_safe_directory_path(directory) search_pattern = _compile_search_regex(regex_pattern) return await _run_search_with_timeout( @@ -1584,6 +1731,7 @@ def _search_files_sync( async def create_directory(self, path: str) -> None: """Ensure the directory at ``path`` exists, creating it if necessary.""" full_path = self._resolve_safe_directory_path(path) + self._reject_quota_metadata(full_path) await asyncio.to_thread(lambda: full_path.mkdir(parents=True, exist_ok=True)) diff --git a/python/packages/core/agent_framework/_harness/_file_memory.py b/python/packages/core/agent_framework/_harness/_file_memory.py index bed59710b74..3524996e32e 100644 --- a/python/packages/core/agent_framework/_harness/_file_memory.py +++ b/python/packages/core/agent_framework/_harness/_file_memory.py @@ -44,7 +44,11 @@ import asyncio import logging -from typing import Annotated, Any +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager +from contextvars import ContextVar +from functools import wraps +from typing import Annotated, Any, ParamSpec, TypeVar, cast from pydantic import BaseModel, Field @@ -54,15 +58,24 @@ from .._types import Message from ._file_access import ( AgentFileStore, + FileAccessProvider, FileStoreEntry, + FileSystemAgentFileStore, _apply_replace, # pyright: ignore[reportPrivateUsage] _apply_replace_lines, # pyright: ignore[reportPrivateUsage] _line_edits, # pyright: ignore[reportPrivateUsage] _matches_glob, # pyright: ignore[reportPrivateUsage] _normalize_relative_path, # pyright: ignore[reportPrivateUsage] ) +from ._file_memory_retention import ( + FileMemoryRetentionManager, + _FileMemoryCommitError, # pyright: ignore[reportPrivateUsage] +) +from ._file_store_limits import _FileStoreQuotaError # pyright: ignore[reportPrivateUsage] logger = logging.getLogger(__name__) +_Params = ParamSpec("_Params") +_ResultT = TypeVar("_ResultT") DEFAULT_FILE_MEMORY_SOURCE_ID = "file_memory" @@ -274,6 +287,7 @@ def __init__( source_id: str = DEFAULT_FILE_MEMORY_SOURCE_ID, scope: str | None = None, instructions: str | None = None, + retention: FileMemoryRetentionManager | None = None, ) -> None: """Initialize the file memory provider. @@ -292,9 +306,15 @@ def __init__( instead of a nested directory. instructions: Optional instruction override. When ``None`` the default file-memory instructions are used. + retention: Optional local-filesystem lifecycle manager. Omitted by default. """ super().__init__(source_id) self.store = store + self.retention = retention + if retention is not None: + retention._validate_store(store) # pyright: ignore[reportPrivateUsage] + self._operation_store: ContextVar[AgentFileStore | None] = ContextVar("file_memory_operation", default=None) + self._retention_started_at: float | None = None self.scope = scope self.instructions = instructions or DEFAULT_FILE_MEMORY_INSTRUCTIONS # Serializes write/delete operations (and their index rebuilds) so the @@ -302,6 +322,62 @@ def __init__( # sufficient for v1; concurrent writes across scopes are rare in practice. self._write_lock = asyncio.Lock() + def _active_store(self) -> AgentFileStore: + return self._operation_store.get() or self.store + + @asynccontextmanager + async def protect(self) -> AsyncGenerator[None]: + """Protect a request-owned provider throughout its run and stream consumption. + + The host must close this context on success, failure, cancellation and stream + close. Sharing this provider between concurrent protected requests is unsupported. + """ + if self.retention is None: + yield + return + if self._retention_started_at is not None: + raise RuntimeError("This file-memory provider already protects a request.") + async with self.retention._protect_store(self.store) as started_at: # pyright: ignore[reportPrivateUsage] + self._retention_started_at = started_at + try: + yield + finally: + self._retention_started_at = None + + def _managed( + self, purpose: str + ) -> Callable[[Callable[_Params, Awaitable[_ResultT]]], Callable[_Params, Awaitable[_ResultT]]]: + def decorate(function: Callable[_Params, Awaitable[_ResultT]]) -> Callable[_Params, Awaitable[_ResultT]]: + @wraps(function) + async def invoke(*args: _Params.args, **kwargs: _Params.kwargs) -> _ResultT: + if self.retention is None: + return await function(*args, **kwargs) + + async def action(store: AgentFileStore) -> _ResultT: + token = self._operation_store.set(store) + try: + return await function(*args, **kwargs) + finally: + self._operation_store.reset(token) + + try: + return await self.retention._run( # pyright: ignore[reportPrivateUsage] + self.store, action, started_at=self._retention_started_at, purpose=purpose + ) + except (ValueError, OSError) as exc: + return cast(_ResultT, f"Could not complete memory operation: {self._error_text(exc)}") + + return invoke + + return decorate + + def _error_text(self, exc: ValueError | OSError) -> str: + if self.retention is None: + return str(exc.strerror or exc) if isinstance(exc, OSError) else str(exc) + if isinstance(exc, (_FileStoreQuotaError, _FileMemoryCommitError)): + return str(exc) + return "Storage or lifecycle state is unavailable; check host diagnostics." + def _resolve_working_folder(self, context: SessionContext) -> str: """Resolve the working folder for the current invocation. @@ -333,20 +409,22 @@ async def _rebuild_index(self, working_folder: str) -> None: Lists the non-internal files, sorts them deterministically, reads any companion descriptions, and writes a capped markdown summary. """ - entries = await self.store.list_children(working_folder) + entries = await self._active_store().list_children(working_folder) file_names = [entry.name for entry in entries if entry.type == FileStoreEntry.FILE] sorted_files = sorted((name for name in file_names if not _is_internal_file(name)), key=str.lower) lines = ["# Memory Index", ""] for file_name in sorted_files[:_MAX_INDEX_ENTRIES]: - description = await self.store.read(_combine_paths(working_folder, _description_file_name(file_name))) + description = await self._active_store().read( + _combine_paths(working_folder, _description_file_name(file_name)) + ) if description and description.strip(): lines.append(f"- **{file_name}**: {description.strip()}") else: lines.append(f"- **{file_name}**") index_path = _combine_paths(working_folder, _MEMORY_INDEX_FILE_NAME) - await self.store.write(index_path, "\n".join(lines) + "\n") + await self._active_store().write(index_path, "\n".join(lines) + "\n") async def before_run( self, @@ -357,12 +435,24 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject file-memory tools, instructions, and the memory index.""" + if self.retention is not None and agent is not None: + for provider in agent.context_providers: + if ( + isinstance(provider, FileAccessProvider) + and isinstance(provider.store, FileSystemAgentFileStore) + and ( + self.retention.directory.is_relative_to(provider.store.root_path) + or provider.store.root_path.is_relative_to(self.retention.directory) + ) + ): + raise ValueError("File access must not expose the managed file-memory storage tree.") working_folder = self._resolve_working_folder(context) - if working_folder: + if working_folder and self.retention is None: await self.store.create_directory(working_folder) @tool(name="file_memory_write", schema=_WriteFileInput, approval_mode="never_require") + @self._managed("write") async def file_memory_write(file_name: str, content: str, description: str | None = None) -> str: """Write a memory file with the given name and content. Overwrites the file if it already exists. Include a description for large files to provide a summary that helps with future discovery.""" # ruff:ignore[line-too-long] try: @@ -383,22 +473,32 @@ async def file_memory_write(file_name: str, content: str, description: str | Non path = _combine_paths(working_folder, normalized) desc_path = _combine_paths(working_folder, _description_file_name(normalized)) async with self._write_lock: + file_written = False try: - await self.store.write(path, content) + await self._active_store().write(path, content) + file_written = True if description and description.strip(): - await self.store.write(desc_path, description) + await self._active_store().write(desc_path, description) else: - await self.store.delete(desc_path) + await self._active_store().delete(desc_path) await self._rebuild_index(working_folder) + except _FileStoreQuotaError as exc: + if file_written: + return ( + f"File '{file_name}' was written, but its description or index could not be updated: {exc} " + "Delete unused memories to free storage before retrying." + ) + return f"Could not write file '{file_name}': {exc}" except ValueError as exc: return f"Could not write file '{file_name}': {exc}" except OSError as exc: - return f"Could not write file '{file_name}': {exc.strerror or exc}" + return f"Could not write file '{file_name}': {self._error_text(exc)}" if description and description.strip(): return f"File '{file_name}' written with description." return f"File '{file_name}' written." @tool(name="file_memory_read", schema=_ReadFileInput, approval_mode="never_require") + @self._managed("read") async def file_memory_read(file_name: str) -> str: r"""Read the content of a memory file by name. Returns the file content or a message indicating the file was not found. Line numbers count lines split on \n only: a lone \r never starts a new line, each line keeps its own terminator, and content ending in a newline has a final empty line.""" # ruff:ignore[line-too-long] try: @@ -408,14 +508,15 @@ async def file_memory_read(file_name: str) -> str: if _is_nested_path(normalized): return f"File '{file_name}' not found." try: - content = await self.store.read(_combine_paths(working_folder, normalized)) + content = await self._active_store().read(_combine_paths(working_folder, normalized)) except ValueError as exc: return f"Could not read file '{file_name}': {exc}" except OSError as exc: - return f"Could not read file '{file_name}': {exc.strerror or exc}" + return f"Could not read file '{file_name}': {self._error_text(exc)}" return content if content is not None else f"File '{file_name}' not found." @tool(name="file_memory_delete", schema=_DeleteFileInput, approval_mode="never_require") + @self._managed("delete") async def file_memory_delete(file_name: str) -> str: """Delete a memory file by name. Also removes its companion description file if one exists.""" try: @@ -429,22 +530,25 @@ async def file_memory_delete(file_name: str) -> str: desc_path = _combine_paths(working_folder, _description_file_name(normalized)) async with self._write_lock: try: - deleted = await self.store.delete(path) - await self.store.delete(desc_path) + deleted = await self._active_store().delete(path) + await self._active_store().delete(desc_path) await self._rebuild_index(working_folder) + except _FileStoreQuotaError as exc: + return f"File '{file_name}' is absent, but the memory index could not be updated: {exc}" except ValueError as exc: return f"Could not delete file '{file_name}': {exc}" except OSError as exc: - return f"Could not delete file '{file_name}': {exc.strerror or exc}" + return f"Could not delete file '{file_name}': {self._error_text(exc)}" return f"File '{file_name}' deleted." if deleted else f"File '{file_name}' not found." @tool(name="file_memory_ls", schema=_ListInput, approval_mode="never_require") + @self._managed("list") async def file_memory_ls(glob_pattern: str | None = None) -> list[dict[str, Any]] | str: """List all memory files with their descriptions (if available). Optionally filter file names with a glob_pattern (e.g. "*.md"). Internal files (description sidecars and the memory index) are not shown. Each entry is {"name": , "type": "file", "description": }.""" # ruff:ignore[line-too-long] try: - entries = await self.store.list_children(working_folder) + entries = await self._active_store().list_children(working_folder) except OSError as exc: - return f"Could not list memory files: {exc.strerror or exc}" + return f"Could not list memory files: {self._error_text(exc)}" file_names = [entry.name for entry in entries if entry.type == FileStoreEntry.FILE] available = set(file_names) @@ -457,11 +561,12 @@ async def file_memory_ls(glob_pattern: str | None = None) -> list[dict[str, Any] description: str | None = None desc_file_name = _description_file_name(file_name) if desc_file_name in available: - description = await self.store.read(_combine_paths(working_folder, desc_file_name)) + description = await self._active_store().read(_combine_paths(working_folder, desc_file_name)) results.append({"name": file_name, "type": "file", "description": description}) return results @tool(name="file_memory_replace", schema=_ReplaceInput, approval_mode="never_require") + @self._managed("write") async def file_memory_replace( file_name: str, old_string: str, new_string: str, replace_all: bool = False ) -> str: @@ -480,18 +585,19 @@ async def file_memory_replace( path = _combine_paths(working_folder, normalized) async with self._write_lock: try: - content = await self.store.read(path) + content = await self._active_store().read(path) if content is None: return f"File '{file_name}' not found." new_content, count = _apply_replace(content, old_string, new_string, replace_all) - await self.store.write(path, new_content) + await self._active_store().write(path, new_content) except ValueError as exc: return f"Could not replace in file '{file_name}': {exc}" except OSError as exc: - return f"Could not replace in file '{file_name}': {exc.strerror or exc}" + return f"Could not replace in file '{file_name}': {self._error_text(exc)}" return f"Replaced {count} occurrence(s) in '{file_name}'." @tool(name="file_memory_replace_lines", schema=_ReplaceLinesInput, approval_mode="never_require") + @self._managed("write") async def file_memory_replace_lines(file_name: str, edits: list[_LineEdit]) -> str: r"""Replace lines in a memory file. Provide a list of edits, each with a 1-based line_number and a literal new_line (include your own trailing newline); an empty new_line deletes the line, including its line break. Fails on out-of-range or duplicate line numbers. Line numbers count lines split on \n only: a lone \r never starts a new line, each line keeps its own terminator, and content ending in a newline has a final empty line.""" # ruff:ignore[line-too-long] try: @@ -508,18 +614,19 @@ async def file_memory_replace_lines(file_name: str, edits: list[_LineEdit]) -> s path = _combine_paths(working_folder, normalized) async with self._write_lock: try: - content = await self.store.read(path) + content = await self._active_store().read(path) if content is None: return f"File '{file_name}' not found." new_content = _apply_replace_lines(content, _line_edits(edits)) - await self.store.write(path, new_content) + await self._active_store().write(path, new_content) except ValueError as exc: return f"Could not edit file '{file_name}': {exc}" except OSError as exc: - return f"Could not edit file '{file_name}': {exc.strerror or exc}" + return f"Could not edit file '{file_name}': {self._error_text(exc)}" return f"Replaced {len(edits)} line(s) in '{file_name}'." @tool(name="file_memory_grep", schema=_SearchFilesInput, approval_mode="never_require") + @self._managed("search") async def file_memory_grep( regex_pattern: str, glob_pattern: str | None = None, @@ -527,7 +634,7 @@ async def file_memory_grep( """Search memory file contents using a case-insensitive regular expression. Optionally filter which files to search using a glob pattern (e.g., "*.md", "research*"). Returns matching file names, content snippets, and matching lines with line numbers. The regex_pattern must be 256 characters or fewer.""" # ruff:ignore[line-too-long] glob_filter = glob_pattern if glob_pattern and glob_pattern.strip() else None try: - results = await self.store.search(working_folder, regex_pattern, glob_filter, recursive=False) + results = await self._active_store().search(working_folder, regex_pattern, glob_filter, recursive=False) # The index and the description sidecars are the provider's own bookkeeping and # are never shown to the agent, so a match inside one is dropped rather than # reported. @@ -535,7 +642,7 @@ async def file_memory_grep( except ValueError as exc: return f"Could not search memory files: {exc}" except OSError as exc: - return f"Could not search memory files: {exc.strerror or exc}" + return f"Could not search memory files: {self._error_text(exc)}" return [result.to_dict() for result in visible] context.extend_instructions(self.source_id, [self.instructions]) @@ -553,12 +660,21 @@ async def file_memory_grep( ) try: - index_content = await self.store.read(_combine_paths(working_folder, _MEMORY_INDEX_FILE_NAME)) + if self.retention is None: + index_content = await self.store.read(_combine_paths(working_folder, _MEMORY_INDEX_FILE_NAME)) + else: + + async def read_index(store: AgentFileStore) -> str | None: + return await store.read(_combine_paths(working_folder, _MEMORY_INDEX_FILE_NAME)) + + index_content = await self.retention._run( # pyright: ignore[reportPrivateUsage] + self.store, read_index, started_at=self._retention_started_at, purpose="index" + ) except (OSError, ValueError) as exc: # A corrupt/unavailable index (e.g. non-UTF8 bytes on disk or a store # error) must not block the run. Skip index injection for this run; it # self-heals on the next successful write/delete that rebuilds the index. - logger.warning("Could not read memory index; skipping index injection: %s", exc) + logger.warning("Could not read memory index; skipping index injection (%s).", type(exc).__name__) index_content = None if index_content and index_content.strip(): context.extend_messages( diff --git a/python/packages/core/agent_framework/_harness/_file_memory_retention.py b/python/packages/core/agent_framework/_harness/_file_memory_retention.py new file mode 100644 index 00000000000..1a9c0f936f3 --- /dev/null +++ b/python/packages/core/agent_framework/_harness/_file_memory_retention.py @@ -0,0 +1,955 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Opt-in local-filesystem coordination for hosted file-memory lifecycle management.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import logging +import math +import os +import time +import uuid +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager, suppress +from pathlib import Path +from typing import Any, TypeVar, cast + +from filelock import AsyncFileLock, Timeout + +from .._feature_stage import ExperimentalFeature, experimental +from .._filesystem import _is_link_or_reparse_point # pyright: ignore[reportPrivateUsage] +from ._file_access import AgentFileStore, FileSearchResult, FileStoreEntry, FileSystemAgentFileStore +from ._file_store_limits import _FileStoreLimits, _FileStoreQuotaError # pyright: ignore[reportPrivateUsage] + +logger = logging.getLogger(__name__) +_ResultT = TypeVar("_ResultT") + + +class _FileMemoryCommitError(ValueError): + """An operation changed content but did not commit its lifecycle metadata.""" + + +async def _complete(work: Awaitable[_ResultT]) -> _ResultT: + """Finish filesystem work before releasing its lock, including on cancellation.""" + task = asyncio.ensure_future(work) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + continue + except BaseException: + break + if not task.cancelled(): + task.exception() + raise + + +@experimental(feature_id=ExperimentalFeature.HARNESS) +class FileMemoryRetentionManager: + """Coordinate quotas and sliding expiry for cooperating local file-memory stores. + + Store roots must be direct children of directory. Control files live outside + those roots. Existing unregistered files remain readable and count toward capacity. + All participating writers must use this manager on the same host/local filesystem. + Construction performs no writes. Use the async context for periodic GC and + FileMemoryProvider.protect() throughout request/stream lifetime. + """ + + def __init__( + self, + directory: str | os.PathLike[str], + *, + retention_seconds: int | None = None, + max_total_bytes: int | None = None, + max_files: int | None = None, + sweep_interval_seconds: int = 300, + ) -> None: + """Configure a host-owned memory area. + + Args: + directory: Common parent of participating store roots. + + Keyword Args: + retention_seconds: Initial idle lifetime for a new directory. Existing policy is preserved; + use set_retention() for explicit configuration changes. None initially disables expiry. + max_total_bytes: Shared file bytes, including control metadata and an atomic-copy reserve. + max_files: Shared regular-file count, including control files and one staging slot. + sweep_interval_seconds: Interval between server-driven sweeps. + + Raises: + ValueError: If a limit is invalid or the managed directory is empty. + """ + _FileStoreLimits.create(retention_seconds, sweep_interval_seconds, None) + if not os.fspath(directory).strip(): + raise ValueError("The managed memory directory must not be empty.") + self.directory = Path(directory).absolute() + self._control = self.directory / ".retention" + self.retention_seconds = retention_seconds + self._initial_retention = retention_seconds + self._policy: dict[str, Any] = {} + self._capacity = _FileStoreLimits.create(None, max_files, max_total_bytes) + self.sweep_interval_seconds = sweep_interval_seconds + self._now: Callable[[], float] = time.time + self._sweeper: asyncio.Task[None] | None = None + + def _validate_store(self, store: AgentFileStore) -> FileSystemAgentFileStore: + """Require an explicitly supported store without creating directories.""" + if not isinstance(store, FileSystemAgentFileStore) or type(store) is not FileSystemAgentFileStore: + raise ValueError("File-memory retention requires FileSystemAgentFileStore without backend overrides.") + if store.root_path.parent != self.directory.resolve() or store.root_path.name == ".retention": + raise ValueError("The memory store must be a direct child of the managed directory.") + return store + + @staticmethod + def _digest(value: str) -> str: + return hashlib.sha256(os.path.normcase(value).replace("\\", "/").encode("utf-8")).hexdigest() + + def _scope(self, store: FileSystemAgentFileStore) -> str: + return self._digest(store.root_path.name) + + def _key(self, store: FileSystemAgentFileStore, path: Path) -> str: + return self._digest(path.relative_to(store.root_path).as_posix()) + + def _manifest_path(self, store: FileSystemAgentFileStore) -> Path: + return self._control / f"{self._scope(store)}.json" + + @staticmethod + def _safe(path: Path) -> None: + if os.path.lexists(path) and _is_link_or_reparse_point(path): + raise ValueError("Memory lifecycle paths must not contain links or reparse points.") + + def _prepare_control(self) -> None: + self._safe(self.directory) + if self.directory.resolve() != self.directory: + raise ValueError("The managed memory directory must not resolve through a link.") + self.directory.mkdir(parents=True, exist_ok=True) + self._safe(self._control) + self._control.mkdir(exist_ok=True) + lock_path = self._control / "manager.lock" + self._safe(lock_path) + if lock_path.exists() and lock_path.stat().st_size: + raise ValueError("Existing data conflicts with the lifecycle control lock.") + staging = self._control / "staging" + self._safe(staging) + staging.mkdir(exist_ok=True) + + @asynccontextmanager + async def _locked(self, *, load_policy: bool = True) -> AsyncGenerator[None]: + self._prepare_control() + # Keep the coordination pathname stable, including on Windows. Native OS + # locks release after process death; a soft-file fallback cannot promise that. + # https://py-filelock.readthedocs.io/en/latest/changelog.html + lock = AsyncFileLock( + self._control / "manager.lock", timeout=10, preserve_lock_file=True, fallback_to_soft=False + ) + try: + async with lock: + for path in (self._control / "staging").iterdir(): + self._safe(path) + if path.is_file() and path.name.endswith(".tmp"): + path.unlink() + if load_policy: + self._load_policy() + yield + except Timeout as exc: + raise TimeoutError("Timed out acquiring the file-memory management lock.") from exc + + def _load_policy(self) -> None: + path = self._control / "policy.json" + self._safe(path) + if not path.exists(): + self._policy = { + "schema_version": 1, + "retention_seconds": self._initial_retention, + "generation": 0, + "resumed_expires_at": None, + } + self._save_policy() + else: + if path.stat().st_size > 16384: + raise ValueError("The file-memory policy exceeds its supported size.") + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (UnicodeError, json.JSONDecodeError) as exc: + raise ValueError("The file-memory policy is invalid; management is paused.") from exc + if ( + not isinstance(raw, dict) + or cast(dict[str, Any], raw).get("schema_version") != 1 + or isinstance(cast(dict[str, Any], raw).get("schema_version"), bool) + ): + raise ValueError("Unsupported file-memory policy version; management is paused.") + raw = cast(dict[str, Any], raw) + if not {"retention_seconds", "generation", "resumed_expires_at"}.issubset(raw): + raise ValueError("The file-memory policy is missing required fields.") + _FileStoreLimits.create(raw["retention_seconds"], None, None) + generation, deadline = raw.get("generation"), raw.get("resumed_expires_at") + if ( + type(generation) is not int + or generation < 0 + or ( + deadline is not None + and ( + isinstance(deadline, bool) + or not isinstance(deadline, (int, float)) + or not math.isfinite(deadline) + ) + ) + or (generation > 0 and deadline is None) + ): + raise ValueError("The file-memory policy has invalid recovery metadata.") + self._policy = raw + self.retention_seconds = self._policy["retention_seconds"] + + def _save_policy(self) -> None: + path = self._control / "policy.json" + payload = json.dumps(self._policy, sort_keys=True, separators=(",", ":")).encode("utf-8") + if len(payload) > 16384: + raise ValueError("The file-memory policy exceeds its supported size.") + self._check_growth(path, len(payload)) + self._atomic(path, payload) + + async def set_retention(self, retention_seconds: int | None) -> None: + """Explicitly enable, disable or change the shared idle lifetime. + + Args: + retention_seconds: Positive lifetime in seconds, or None to disable TTL. + + Disabling retains deadlines and permits ordinary reads without renewal. + Reenabling grants every retained ready record a complete new lifetime. + Same-value calls are no-ops. Changing one positive value to another only + changes subsequent successful writes/uses. Construction/restart reads an + existing policy without changing it; the constructor value initializes a + new managed directory only. + """ + _FileStoreLimits.create(retention_seconds, None, None) + + async def work() -> None: + async with self._locked(): + previous = self.retention_seconds + if previous == retention_seconds: + return + policy = self._policy.copy() + policy["retention_seconds"] = retention_seconds + if previous is None and retention_seconds is not None: + policy["generation"] += 1 + policy["resumed_expires_at"] = self._now() + retention_seconds + self._policy = policy + self._save_policy() + self.retention_seconds = retention_seconds + + await _complete(work()) + + def _deadline(self, record: dict[str, Any]) -> float | None: + if record.get("generation", 0) < self._policy["generation"]: + return self._policy["resumed_expires_at"] + return record["expires_at"] + + def _load(self, store: FileSystemAgentFileStore) -> dict[str, Any]: + path = self._manifest_path(store) + self._safe(path) + if not path.exists(): + return {"schema_version": 1, "entries": {}} + if path.stat().st_size > 16 * 1024 * 1024: + raise ValueError("The file-memory manifest exceeds its supported size.") + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (UnicodeError, json.JSONDecodeError) as exc: + raise ValueError("The file-memory manifest is invalid; lifecycle operations are paused.") from exc + if ( + not isinstance(raw, dict) + or cast(dict[str, Any], raw).get("schema_version") != 1 + or isinstance(cast(dict[str, Any], raw).get("schema_version"), bool) + ): + raise ValueError("Unsupported file-memory manifest version; lifecycle operations are paused.") + state = cast(dict[str, Any], raw) + entries = state.get("entries") + if not isinstance(entries, dict): + raise ValueError("The file-memory manifest has invalid entries.") + for key, value in cast(dict[str, Any], entries).items(): + if len(key) != 64 or any(c not in "0123456789abcdef" for c in key) or not isinstance(value, dict): + raise ValueError("The file-memory manifest contains an unsupported record.") + record = cast(dict[str, Any], value) + if record.get("state") not in ("ready", "updating", "deleting"): + raise ValueError("The file-memory manifest contains an unsupported state.") + for name in ("folder", "description"): + digest = record.get(name) + if not isinstance(digest, str) or len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise ValueError("The file-memory manifest contains an unsupported reference.") + if key in (record["folder"], record["description"]) or record["folder"] == record["description"]: + raise ValueError("The file-memory record contains inconsistent references.") + generation = record.get("generation", 0) + if type(generation) is not int or generation < 0 or generation > self._policy["generation"]: + raise ValueError("The file-memory record has an unsupported recovery generation.") + if "expires_at" not in record: + raise ValueError("The file-memory record is missing its expiry field.") + deadline = record["expires_at"] + if ( + deadline is None + and record["state"] == "ready" + and self.retention_seconds is not None + and generation == self._policy["generation"] + ): + raise ValueError("An active ready record must have an expiry time.") + if deadline is not None and ( + isinstance(deadline, bool) or not isinstance(deadline, (float, int)) or not math.isfinite(deadline) + ): + raise ValueError("The file-memory manifest has an invalid expiry time.") + self._stored_limits(state) + return state + + def _usage(self) -> tuple[int, int, int]: + count = total = largest = 0 + directories = [self.directory] + while directories: + directory = directories.pop() + if not directory.exists(): + continue + for path in directory.iterdir(): + self._safe(path) + if path.is_dir(): + if path != self._control / "staging": + directories.append(path) + elif path.is_file(): + count += 1 + size = path.stat().st_size + total += size + largest = max(largest, size) + return count, total, largest + + def _check_growth(self, path: Path, size: int) -> None: + self._check_combined_growth([(path, size)]) + + def _check_combined_growth(self, changes: list[tuple[Path, int]]) -> None: + if self._capacity is None: + return + count, total, largest = self._usage() + resulting_count, resulting_bytes, reserve = count, total, largest + for path, size in changes: + exists = path.is_file() + resulting_count += int(not exists) + resulting_bytes += size - (path.stat().st_size if exists else 0) + reserve = max(reserve, size) + # Serialized mutations reserve room for one complete atomic-copy file. + if self._capacity.max_files is not None and resulting_count + 1 > max(count + 1, self._capacity.max_files): + raise _FileStoreQuotaError("shared_max_files", self._capacity.max_files, resulting_count + 1) + if self._capacity.max_total_bytes is not None and ( + resulting_bytes + reserve > max(total + largest, self._capacity.max_total_bytes) + ): + raise _FileStoreQuotaError( + "shared_max_total_bytes", self._capacity.max_total_bytes, resulting_bytes + largest + ) + + @staticmethod + def _stored_limits(state: dict[str, Any]) -> _FileStoreLimits | None: + raw = state.get("store_limits") + if raw is None: + return None + if not isinstance(raw, dict): + raise ValueError("Invalid stored file-memory capacity configuration.") + limits = cast(dict[str, Any], raw) + return _FileStoreLimits.create( + limits.get("max_file_bytes"), limits.get("max_files"), limits.get("max_total_bytes") + ) + + def _check_store_growth(self, store: FileSystemAgentFileStore, path: Path, size: int) -> None: + limits = store._limits # pyright: ignore[reportPrivateUsage] + if limits is not None: + limits.check_read(size) + count, total = store._quota_usage() if store.root_path.exists() else (0, 0) # pyright: ignore[reportPrivateUsage] + limits.check_usage( + count + int(not path.exists()), total - (path.stat().st_size if path.exists() else 0) + size + ) + + async def _write_index(self, store: FileSystemAgentFileStore, view: _RetainedFileStore, folder: Path) -> None: + path = folder / "memories.md" + payload = (await view.index_text(folder.relative_to(store.root_path).as_posix())).encode("utf-8") + self._check_store_growth(store, path, len(payload)) + self._check_growth(path, len(payload)) + self._atomic(path, payload) + + def _atomic(self, path: Path, payload: bytes, *, overwrite: bool = True) -> None: + temporary = self._control / "staging" / f"{uuid.uuid4().hex}.tmp" + try: + with temporary.open("xb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + self._safe(path) + path.parent.mkdir(parents=True, exist_ok=True) + if overwrite: + os.replace(temporary, path) + else: + os.link(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + def _save(self, store: FileSystemAgentFileStore, state: dict[str, Any], *, reclaim: bool = False) -> None: + path = self._manifest_path(store) + payload = json.dumps(state, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") + if len(payload) > 16 * 1024 * 1024: + raise ValueError("The file-memory manifest exceeds its supported size.") + if not reclaim: + self._check_growth(path, len(payload)) + self._atomic(path, payload) + + @staticmethod + def _files(root: Path) -> list[Path]: + files: list[Path] = [] + directories = [root] + while directories: + directory = directories.pop() + if not directory.exists(): + continue + FileMemoryRetentionManager._safe(directory) + for path in directory.iterdir(): + FileMemoryRetentionManager._safe(path) + if path.is_dir(): + directories.append(path) + elif path.is_file() and path != root / ".agent-framework-quota.lock": + files.append(path) + return files + + async def _run( + self, + store: AgentFileStore, + action: Callable[[AgentFileStore], Awaitable[_ResultT]], + *, + started_at: float | None = None, + purpose: str = "", + ) -> _ResultT: + """Run a provider operation under shared capacity and lifecycle coordination.""" + backing = self._validate_store(store) + + async def work() -> _ResultT: + async with self._locked(): + self._safe(backing.root_path) + state = self._load(backing) + view = _RetainedFileStore(self, backing, state, started_at, purpose) + result = await action(view) + if view.changed: + if not view.failed: + for key in view.deleted: + state["entries"].pop(key, None) + for key in view.pending: + record = state["entries"][key] + record["state"] = "ready" + record["generation"] = self._policy["generation"] + record["expires_at"] = ( + self._now() + self.retention_seconds if self.retention_seconds is not None else None + ) + try: + self._save(backing, state) + except (ValueError, OSError) as exc: + if view.pending or view.deleted: + raise _FileMemoryCommitError( + "Memory content may have changed, but lifecycle metadata was not committed. " + "The interrupted operation is protected from expiry; host repair is required." + ) from exc + raise + return result + + return await _complete(work()) + + @asynccontextmanager + async def _protect_store(self, store: AgentFileStore) -> AsyncGenerator[float]: + """Hold an OS-backed request token until the host closes the request/stream.""" + backing = self._validate_store(store) + token_lock: AsyncFileLock | None = None + token_path: Path | None = None + try: + async with self._locked(): + self._load(backing) + tokens = self._control / "active" / self._scope(backing) + self._safe(self._control / "active") + self._safe(tokens) + tokens.mkdir(parents=True, exist_ok=True) + token_path = tokens / f"{uuid.uuid4().hex}.lock" + self._check_growth(token_path, 0) + token_lock = AsyncFileLock(token_path, timeout=0, preserve_lock_file=True, fallback_to_soft=False) + await token_lock.acquire() + started_at = self._now() + yield started_at # ruff: ignore[unnecessary-assign-before-yield] - capture visibility while registration holds the lock + finally: + if token_lock is not None: + + async def release() -> None: + try: + async with self._locked(load_policy=False): + await token_lock.release() + if token_path is not None: + token_path.unlink(missing_ok=True) + if token_path.parent.exists() and not any(token_path.parent.iterdir()): + token_path.parent.rmdir() + finally: + # A broken policy/control directory must not pin an OS lease. + await token_lock.release() + + await _complete(release()) + + async def _active(self, store: FileSystemAgentFileStore) -> bool: + return await self._tokens_active(self._scope(store)) + + async def _tokens_active(self, scope: str) -> bool: + tokens = self._control / "active" / scope + self._safe(tokens) + if not tokens.exists(): + return False + active = False + for path in tokens.iterdir(): + self._safe(path) + if ( + not path.is_file() + or path.suffix != ".lock" + or len(path.stem) != 32 + or any(c not in "0123456789abcdef" for c in path.stem) + or path.stat().st_size + ): + raise ValueError("Unexpected file in the request protection directory.") + probe = AsyncFileLock(path, timeout=0, preserve_lock_file=True, fallback_to_soft=False) + try: + await probe.acquire() + except Timeout: + active = True + else: + await probe.release() + path.unlink(missing_ok=True) + if not active: + tokens.rmdir() + return active + + async def collect(self) -> int: + """Delete registered expired memories, retaining active or uncertain scopes.""" + + async def work() -> int: + async with self._locked(): + return await self._collect_locked() + + return await _complete(work()) + + async def _collect_locked(self) -> int: + removed = 0 + active_scopes: set[str] = set() + active = self._control / "active" + self._safe(active) + if active.exists(): + for tokens in active.iterdir(): + self._safe(tokens) + if ( + not tokens.is_dir() + or len(tokens.name) != 64 + or any(c not in "0123456789abcdef" for c in tokens.name) + ): + raise ValueError("Unexpected request protection metadata.") + if await self._tokens_active(tokens.name): + active_scopes.add(tokens.name) + if self.retention_seconds is None: + return 0 + for root in list(self.directory.iterdir()): + if root == self._control or not root.is_dir(): + continue + try: + self._safe(root) + backing = FileSystemAgentFileStore(root) + state = self._load(backing) + limits = self._stored_limits(state) + if limits is not None: + backing = FileSystemAgentFileStore( + root, + max_file_bytes=limits.max_file_bytes, + max_files=limits.max_files, + max_total_bytes=limits.max_total_bytes, + ) + if self._scope(backing) in active_scopes: + continue + view = _RetainedFileStore(self, backing, state, None, "gc") + files = self._files(root) + paths = {self._key(backing, p): p for p in files} + folders = {self._key(backing, p.parent): p.parent for p in files} + for key, record in list(state["entries"].items()): + deadline = self._deadline(record) + if record["state"] == "updating": + logger.warning("Retaining interrupted memory update; explicit repair is required.") + continue + if record["state"] != "deleting" and (deadline is None or deadline > self._now()): + continue + body, folder = paths.get(key), folders.get(record["folder"]) + if body is not None and ( + view.internal(body.name) + or record["folder"] != self._key(backing, body.parent) + or record["description"] != self._key(backing, body.with_name(view._description(body.name))) # pyright: ignore[reportPrivateUsage] + ): + raise ValueError("Inconsistent file-memory lifecycle references.") + record["state"] = "deleting" + self._save(backing, state, reclaim=True) + if body is not None: + body.unlink(missing_ok=True) + description = paths.get(record["description"]) + if description is not None: + from ._file_memory import _description_file_name # pyright: ignore[reportPrivateUsage] + + referenced = any( + p.exists() + and not view.internal(p.name) + and p.parent == description.parent + and os.path.normcase(_description_file_name(p.name)) == os.path.normcase(description.name) + for p in files + ) + if not referenced: + description.unlink(missing_ok=True) + del state["entries"][key] + if folder is not None and folder.exists(): + index = folder / "memories.md" + if any(not view.internal(p.name) for p in folder.iterdir() if p.is_file()): + try: + await self._write_index(backing, view, folder) + except _FileStoreQuotaError: + # The index is a projection; insufficient space must not + # leave an already-deleted body permanently tombstoned. + index.unlink(missing_ok=True) + logger.warning("Deferred memory index rebuild due to capacity; stale index removed.") + else: + index.unlink(missing_ok=True) + if not any(folder.iterdir()): + folder.rmdir() + self._save(backing, state, reclaim=True) + removed += 1 + directories: list[Path] = [] + pending = [root] + while pending: + directory = pending.pop() + self._safe(directory) + directories.append(directory) + pending.extend(path for path in directory.iterdir() if path.is_dir()) + for directory in reversed(directories[1:]): + if not any(directory.iterdir()): + directory.rmdir() + if not state["entries"] and not self._files(root): + self._manifest_path(backing).unlink(missing_ok=True) + if not any(root.iterdir()): + root.rmdir() + except (ValueError, OSError) as exc: + logger.warning("Skipped file-memory cleanup (%s).", type(exc).__name__) + return removed + + async def repair(self, store: AgentFileStore) -> int: + """Confirm retained interrupted updates after the operator inspects their content. + + Args: + store: Supported filesystem store belonging to this manager. + + Returns: + Number of interrupted updates confirmed or removed if the body is absent. + + This is an explicit confirmation, not proof that a multi-file write completed. + Present bodies are accepted as their current content, indexes are rebuilt, and + a fresh lifetime begins. Pending deletions remain deletions. Active scopes + cannot be repaired. + """ + backing = self._validate_store(store) + + async def work() -> int: + async with self._locked(): + if await self._active(backing): + raise ValueError("Cannot repair a memory scope while a request is active.") + state = self._load(backing) + files = self._files(backing.root_path) + paths = {self._key(backing, path): path for path in files} + folders = {self._key(backing, path.parent): path.parent for path in files} + view = _RetainedFileStore(self, backing, state, None, "repair") + affected: set[Path] = set() + repaired = 0 + for key, record in list(state["entries"].items()): + if record["state"] != "updating": + continue + body, folder = paths.get(key), folders.get(record["folder"]) + if body is None: + state["entries"].pop(key) + else: + record["state"] = "ready" + record["generation"] = self._policy["generation"] + record["expires_at"] = ( + self._now() + self.retention_seconds if self.retention_seconds is not None else None + ) + if folder is not None: + affected.add(folder) + repaired += 1 + for folder in affected: + await self._write_index(backing, view, folder) + if repaired: + self._save(backing, state) + return repaired + + return await _complete(work()) + + async def __aenter__(self) -> FileMemoryRetentionManager: + """Start one server-owned sweeper; each sweep uses current on-disk state.""" + if self._sweeper is not None: + raise RuntimeError("The file-memory retention manager is already running.") + + async with self._locked(): + pass + + async def sweep() -> None: + while True: + await asyncio.sleep(self.sweep_interval_seconds) + try: + await self.collect() + except (ValueError, OSError) as exc: + logger.warning("File-memory sweep failed (%s).", type(exc).__name__) + + self._sweeper = asyncio.create_task(sweep()) + return self + + async def __aexit__(self, *exc_info: object) -> None: + """Stop GC; task shutdown does not itself select a retention-resume policy.""" + if self._sweeper is not None: + self._sweeper.cancel() + with suppress(asyncio.CancelledError): + await self._sweeper + self._sweeper = None + + +class _RetainedFileStore(AgentFileStore): + """Request-local view; all operations hold the manager lock.""" + + def __init__( + self, + manager: FileMemoryRetentionManager, + backing: FileSystemAgentFileStore, + state: dict[str, Any], + started_at: float | None, + purpose: str, + ) -> None: + self.manager, self.backing, self.state = manager, backing, state + self.started_at, self.purpose = started_at, purpose + self.pending: set[str] = set() + self.deleted: set[str] = set() + self.changed = self.failed = False + + @staticmethod + def internal(name: str) -> bool: + return name.lower() == "memories.md" or name.lower().endswith("_description.md") + + def _path(self, path: str) -> Path: + return self.backing._resolve_safe_path(path) # pyright: ignore[reportPrivateUsage] + + def _visible(self, path: str) -> bool: + record = self.state["entries"].get(self.manager._key(self.backing, self._path(path))) # pyright: ignore[reportPrivateUsage] + if record is None: + return True + if record["state"] == "deleting": + return False + if record["state"] == "updating" or self.manager.retention_seconds is None: + return True + deadline = self.manager._deadline(record) # pyright: ignore[reportPrivateUsage] + cutoff = self.started_at if self.started_at is not None else self.manager._now() # pyright: ignore[reportPrivateUsage] + return deadline is None or deadline > cutoff + + def _touch(self, path: str) -> None: + record = self.state["entries"].get(self.manager._key(self.backing, self._path(path))) # pyright: ignore[reportPrivateUsage] + if record is not None and record["state"] == "ready" and self.manager.retention_seconds is not None: + record["expires_at"] = self.manager._now() + self.manager.retention_seconds # pyright: ignore[reportPrivateUsage] + record["generation"] = self.manager._policy["generation"] # pyright: ignore[reportPrivateUsage] + self.changed = True + + async def write(self, path: str, content: str, *, overwrite: bool = True) -> None: + target = self._path(path) + try: + limits = self.backing._limits # pyright: ignore[reportPrivateUsage] + sizing = limits or _FileStoreLimits( + None, + None, + self.manager._capacity.max_total_bytes if self.manager._capacity else None, # pyright: ignore[reportPrivateUsage] + ) + if not overwrite and target.exists(): + raise FileExistsError("The memory file already exists.") + size = sizing.content_size(content) + + def check() -> None: + self.manager._check_store_growth(self.backing, target, size) # pyright: ignore[reportPrivateUsage] + if self.internal(target.name): + self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] + return + key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] + previous = self.state["entries"].get(key) + record: dict[str, Any] = { + **(previous or {}), + "state": "ready", + "expires_at": self.manager._now() + self.manager.retention_seconds # pyright: ignore[reportPrivateUsage] + if self.manager.retention_seconds is not None + else None, + "generation": self.manager._policy["generation"], # pyright: ignore[reportPrivateUsage] + "folder": self.manager._key(self.backing, target.parent), # pyright: ignore[reportPrivateUsage] + "description": self.manager._key(self.backing, target.with_name(self._description(target.name))), # pyright: ignore[reportPrivateUsage] + } + self.state["entries"][key] = record + config: dict[str, Any] = { + **(self.state.get("store_limits") or {}), + "max_file_bytes": limits.max_file_bytes if limits else None, + "max_files": limits.max_files if limits else None, + "max_total_bytes": limits.max_total_bytes if limits else None, + } + self.state["store_limits"] = config + payload_size = len(json.dumps(self.state, sort_keys=True, separators=(",", ":")).encode("utf-8")) + # A finite float's JSON representation is shorter than 32 bytes; + # reserve its maximum field width before touching the body. + if record["expires_at"] is not None: + payload_size += max(0, 32 - len(json.dumps(record["expires_at"]))) + if previous is None: + self.state["entries"].pop(key) + else: + self.state["entries"][key] = previous + self.manager._check_combined_growth( # pyright: ignore[reportPrivateUsage] + [(target, size), (self.manager._manifest_path(self.backing), payload_size)] # pyright: ignore[reportPrivateUsage] + ) + + try: + check() + except _FileStoreQuotaError: + await self.manager._collect_locked() # pyright: ignore[reportPrivateUsage] + loaded = self.manager._load(self.backing) # pyright: ignore[reportPrivateUsage] + self.state.clear() + self.state.update(loaded) + check() + if not self.internal(target.name): + key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] + record = self.state["entries"].get(key) + if record is not None and record["state"] == "deleting": + raise ValueError("This memory has a pending deletion; host cleanup must finish first.") + self.state["entries"][key] = { + **(record or {}), + "state": "updating", + "expires_at": None, + "generation": self.manager._policy["generation"], # pyright: ignore[reportPrivateUsage] + "folder": self.manager._key(self.backing, target.parent), # pyright: ignore[reportPrivateUsage] + "description": self.manager._key(self.backing, target.with_name(self._description(target.name))), # pyright: ignore[reportPrivateUsage] + } + self.manager._save(self.backing, self.state) # pyright: ignore[reportPrivateUsage] + try: + self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] + except _FileStoreQuotaError: + if record is None: + self.state["entries"].pop(key, None) + else: + self.state["entries"][key] = record + self.manager._save(self.backing, self.state, reclaim=True) # pyright: ignore[reportPrivateUsage] + raise + self.pending.add(key) + self.changed = True + self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] + await asyncio.to_thread(self.manager._atomic, target, content.encode("utf-8"), overwrite=overwrite) # pyright: ignore[reportPrivateUsage] + except BaseException: + self.failed = True + raise + + @staticmethod + def _description(name: str) -> str: + from ._file_memory import _description_file_name # pyright: ignore[reportPrivateUsage] + + return _description_file_name(name) + + async def read(self, path: str) -> str | None: + if Path(path).name.lower() == "memories.md": + directory = Path(path).parent.as_posix() + content = await self.index_text("" if directory == "." else directory) + return content if content.strip() != "# Memory Index" else None + if not self.internal(Path(path).name) and not self._visible(path): + return None + target = self._path(path) + if target.name.lower().endswith("_description.md") and target.parent.exists(): + references = [ + body + for body in target.parent.iterdir() + if body.is_file() + and not self.internal(body.name) + and os.path.normcase(self._description(body.name)) == os.path.normcase(target.name) + ] + key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] + registered = any(record["description"] == key for record in self.state["entries"].values()) + if (references or registered) and not any( + self._visible(body.relative_to(self.backing.root_path).as_posix()) for body in references + ): + return None + try: + content = await self.backing.read(path) + except ValueError: + self.failed = self.failed or bool(self.pending or self.deleted) + raise + if content is not None and not self.internal(Path(path).name) and self.purpose == "read": + self._touch(path) + return content + + async def delete(self, path: str) -> bool: + target = self._path(path) + if not self.internal(target.name): + key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] + record = self.state["entries"].get(key) + if record is not None: + record["state"] = "deleting" + self.manager._save(self.backing, self.state, reclaim=True) # pyright: ignore[reportPrivateUsage] + self.deleted.add(key) + self.changed = True + elif target.name.lower().endswith("_description.md") and target.parent.exists(): + if any( + p.is_file() + and not self.internal(p.name) + and os.path.normcase(self._description(p.name)) == os.path.normcase(target.name) + for p in target.parent.iterdir() + ): + return False + try: + return await asyncio.to_thread(self.backing._delete_file_sync, target) # pyright: ignore[reportPrivateUsage] + except BaseException: + self.failed = True + raise + + async def list_children(self, directory: str = "") -> list[FileStoreEntry]: + entries = await self.backing.list_children(directory) + return [ + entry + for entry in entries + if entry.type != FileStoreEntry.FILE + or self.internal(entry.name) + or self._visible(f"{directory}/{entry.name}" if directory else entry.name) + ] + + async def file_exists(self, path: str) -> bool: + return self._visible(path) and await self.backing.file_exists(path) + + async def create_directory(self, path: str) -> None: + await self.backing.create_directory(path) + + async def search( + self, directory: str, regex_pattern: str, glob_pattern: str | None = None, *, recursive: bool = False + ) -> list[FileSearchResult]: + results = await super().search(directory, regex_pattern, glob_pattern, recursive=recursive) + for result in results: + if not self.internal(Path(result.file_name).name): + self._touch(f"{directory}/{result.file_name}" if directory else result.file_name) + return results + + async def index_text(self, directory: str) -> str: + from ._file_memory import _MAX_INDEX_ENTRIES # pyright: ignore[reportPrivateUsage] + + entries = await self.list_children(directory) + names = sorted( + (entry.name for entry in entries if entry.type == FileStoreEntry.FILE and not self.internal(entry.name)), + key=str.lower, + ) + lines = ["# Memory Index", ""] + for name in names[:_MAX_INDEX_ENTRIES]: + path = f"{directory}/{self._description(name)}" if directory else self._description(name) + description = await self.backing.read(path) + lines.append( + f"- **{name}**: {description.strip()}" if description and description.strip() else f"- **{name}**" + ) + content = "\n".join(lines) + "\n" + limits = self.backing._limits # pyright: ignore[reportPrivateUsage] + if limits is not None: + limits.content_size(content) + return content diff --git a/python/packages/core/agent_framework/_harness/_file_store_limits.py b/python/packages/core/agent_framework/_harness/_file_store_limits.py new file mode 100644 index 00000000000..9f3beb7910c --- /dev/null +++ b/python/packages/core/agent_framework/_harness/_file_store_limits.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Capacity accounting shared by the built-in file stores.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +class _FileStoreQuotaError(ValueError): + """A configured storage limit would be exceeded, without revealing paths or identities.""" + + def __init__(self, name: str, limit: int, requested: int) -> None: + super().__init__(f"File store quota exceeded: {name}={limit}, requested={requested}.") + + +@dataclass(frozen=True) +class _FileStoreLimits: + max_file_bytes: int | None + max_files: int | None + max_total_bytes: int | None + + @classmethod + def create( + cls, max_file_bytes: int | None, max_files: int | None, max_total_bytes: int | None + ) -> _FileStoreLimits | None: + """Validate explicit limits; an entirely unconfigured store keeps its original path.""" + values = (max_file_bytes, max_files, max_total_bytes) + for name, value in zip(("max_file_bytes", "max_files", "max_total_bytes"), values, strict=True): + if value is not None and (isinstance(value, bool) or not isinstance(value, int) or value <= 0): + raise ValueError(f"{name} must be a positive integer or None.") + return cls(*values) if any(value is not None for value in values) else None + + def content_size(self, content: str) -> int: + """Count UTF-8 bytes incrementally before allocating the complete encoded payload.""" + limits = [ + (name, limit) + for name, limit in (("max_file_bytes", self.max_file_bytes), ("max_total_bytes", self.max_total_bytes)) + if limit is not None + ] + ceiling = min((limit for _, limit in limits), default=None) + chunk_size = min(4096, ceiling + 1) if ceiling is not None else 4096 + size = 0 + for offset in range(0, len(content), chunk_size): + size += len(content[offset : offset + chunk_size].encode("utf-8")) + for name, limit in limits: + if size > limit: + raise _FileStoreQuotaError(name, limit, size) + return size + + def check_usage(self, file_count: int, total_bytes: int) -> None: + """Check the resulting root usage before modifying a file.""" + for name, limit, requested in ( + ("max_files", self.max_files, file_count), + ("max_total_bytes", self.max_total_bytes, total_bytes), + ): + if limit is not None and requested > limit: + raise _FileStoreQuotaError(name, limit, requested) + + def check_read(self, file_bytes: int) -> None: + """Bound whole-file reads, including files written before quotas were enabled.""" + if self.max_file_bytes is not None and file_bytes > self.max_file_bytes: + raise _FileStoreQuotaError("max_file_bytes", self.max_file_bytes, file_bytes) diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index 5406165ae9c..0464fbdc026 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ + "filelock>=3.30.3,<4", "msgspec>=0.20.0,<0.22", "typing-extensions>=4.15.0,<5", "pydantic>=2,<3", diff --git a/python/packages/core/tests/core/test_harness_file_memory.py b/python/packages/core/tests/core/test_harness_file_memory.py index 4f2e810c858..5f7043103f3 100644 --- a/python/packages/core/tests/core/test_harness_file_memory.py +++ b/python/packages/core/tests/core/test_harness_file_memory.py @@ -663,3 +663,27 @@ async def test_file_memory_replace_lines_honours_expected_line() -> None: # endregion + + +async def test_quota_failure_reports_committed_memory_and_unfinished_metadata() -> None: + """A later quota failure must not falsely claim the memory content was never written.""" + store = InMemoryAgentFileStore(max_files=1) + _, tools = await _prepare(FileMemoryProvider(store=store, scope="user-1")) + result = await tools["file_memory_write"].invoke( + arguments={"file_name": "plan.md", "content": "plan", "description": "summary"} + ) + text = _text(result) + assert "was written" in text and "description or index" in text and "max_files" in text + assert await store.read("user-1/plan.md") == "plan" + assert await store.read("user-1/plan_description.md") is None + assert await store.read("user-1/memories.md") is None + + +async def test_quota_rejection_preserves_existing_memory_content() -> None: + """An oversized replacement must leave the existing memory and index intact.""" + store = InMemoryAgentFileStore(max_total_bytes=60) + _, tools = await _prepare(FileMemoryProvider(store=store, scope="user-1")) + await tools["file_memory_write"].invoke(arguments={"file_name": "plan.md", "content": "old"}) + result = await tools["file_memory_write"].invoke(arguments={"file_name": "plan.md", "content": "x" * 61}) + assert "Could not write" in _text(result) and "max_total_bytes" in _text(result) + assert await store.read("user-1/plan.md") == "old" diff --git a/python/packages/core/tests/core/test_harness_file_memory_retention.py b/python/packages/core/tests/core/test_harness_file_memory_retention.py new file mode 100644 index 00000000000..113bf3bddad --- /dev/null +++ b/python/packages/core/tests/core/test_harness_file_memory_retention.py @@ -0,0 +1,803 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Persistence and failure contracts for opt-in file-memory lifecycle management.""" + +from __future__ import annotations + +import asyncio +import json +import sys +from pathlib import Path +from typing import Any + +import pytest + +from agent_framework import ( + AgentSession, + FileMemoryProvider, + FileMemoryRetentionManager, + FileSystemAgentFileStore, +) +from agent_framework._sessions import SessionContext + + +async def _prepare(provider: FileMemoryProvider) -> tuple[SessionContext, dict[str, Any]]: + context = SessionContext(session_id="conversation", input_messages=[]) + await provider.before_run(agent=None, session=AgentSession(session_id="conversation"), context=context, state={}) + return context, {tool.name: tool for tool in context.tools} + + +async def _call(tools: dict[str, Any], name: str, **arguments: object) -> str: + result = await tools[name].invoke(arguments=arguments) + return str(result[0].text or "") + + +async def test_retention_refreshes_read_and_hides_expiry_on_every_surface(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + provider = FileMemoryProvider(store, retention=manager) + _, tools = await _prepare(provider) + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + now[0] = 109 + assert await _call(tools, "file_memory_read", file_name="notes.md") == "hello" + now[0] = 118 + assert await manager.collect() == 0 + now[0] = 119 + context, tools = await _prepare(provider) + assert "not found" in await _call(tools, "file_memory_read", file_name="notes.md") + assert await _call(tools, "file_memory_ls") == "[]" + assert await _call(tools, "file_memory_grep", regex_pattern="hello") == "[]" + assert not context.context_messages + assert (store.root_path / "conversation/notes.md").exists() + assert await manager.collect() == 1 + assert not (store.root_path / "conversation/notes.md").exists() + + +async def test_old_memory_without_metadata_is_not_adopted_by_read_or_gc(tmp_path: Path) -> None: + store = FileSystemAgentFileStore(tmp_path / "owner") + await store.write("conversation/legacy.md", "preserve") + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=1) + manager._now = lambda: 1000.0 + provider = FileMemoryProvider(store, retention=manager) + _, tools = await _prepare(provider) + assert await _call(tools, "file_memory_read", file_name="legacy.md") == "preserve" + assert await manager.collect() == 0 + assert await store.read("conversation/legacy.md") == "preserve" + for state in (tmp_path / ".retention").glob("*.json"): + assert "legacy.md" not in state.read_text() + + +async def test_expiry_preserves_description_still_referenced_by_another_memory(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + provider = FileMemoryProvider(store, retention=manager) + _, tools = await _prepare(provider) + await _call(tools, "file_memory_write", file_name="notes.md", content="first", description="shared") + now[0] = 105 + await _call(tools, "file_memory_write", file_name="notes.txt", content="second", description="shared") + now[0] = 111 + assert await manager.collect() == 1 + assert await store.read("conversation/notes_description.md") == "shared" + assert "notes.txt" in (await store.read("conversation/memories.md") or "") + now[0] = 116 + assert await manager.collect() == 1 + assert await store.read("conversation/notes_description.md") is None + + +async def test_active_request_protects_storage_and_visibility_past_deadline(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + provider = FileMemoryProvider(store, retention=manager) + _, tools = await _prepare(provider) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + now[0] = 109 + async with provider.protect(): + now[0] = 120 + assert await manager.collect() == 0 + assert await _call(tools, "file_memory_read", file_name="notes.md") == "hello" + now[0] = 131 + assert await manager.collect() == 1 + + +async def test_disable_resume_and_restart_have_distinct_persistent_behavior(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + provider = FileMemoryProvider(store, retention=manager) + _, tools = await _prepare(provider) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + now[0] = 105 + await manager.set_retention(None) + manifests = list((tmp_path / ".retention").glob("*.json")) + manifest = next(path for path in manifests if path.name != "policy.json") + before = manifest.read_bytes() + now[0] = 200 + assert await _call(tools, "file_memory_read", file_name="notes.md") == "hello" + assert await manager.collect() == 0 + assert manifest.read_bytes() == before + + # A restarted worker reads the durable disabled policy even with an initial TTL. + restarted = FileMemoryRetentionManager(tmp_path, retention_seconds=999) + restarted._now = lambda: now[0] + assert await restarted.collect() == 0 + assert restarted.retention_seconds is None + + # An explicit resume grants ten complete seconds, visible to other workers. + await restarted.set_retention(10) + policy = json.loads((tmp_path / ".retention/policy.json").read_text()) + assert policy["resumed_expires_at"] == 210 + assert policy["generation"] == 1 + now[0] = 209 + assert await manager.collect() == 0 + await restarted.set_retention(10) + assert json.loads((tmp_path / ".retention/policy.json").read_text()) == policy + now[0] = 210 + assert await manager.collect() == 1 + + +async def test_multiple_resumes_and_successful_use_keep_individual_deadlines(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + await manager.set_retention(None) + now[0] = 200 + await manager.set_retention(10) + now[0] = 205 + assert await _call(tools, "file_memory_read", file_name="notes.md") == "hello" + now[0] = 211 + assert await manager.collect() == 0 + await manager.set_retention(None) + now[0] = 300 + await manager.set_retention(20) + assert json.loads((tmp_path / ".retention/policy.json").read_text())["generation"] == 2 + now[0] = 319 + assert await manager.collect() == 0 + now[0] = 320 + assert await manager.collect() == 1 + + +async def test_positive_policy_change_does_not_recalculate_existing_deadlines(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + _, tools = await _prepare(FileMemoryProvider(FileSystemAgentFileStore(tmp_path / "owner"), retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + await manager.set_retention(100) + now[0] = 110 + assert await manager.collect() == 1 + + +async def test_metadata_commit_failure_protects_actual_partial_write( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="old") + original = manager._save + + def fail_commit(backing: FileSystemAgentFileStore, state: dict[str, Any], *, reclaim: bool = False) -> None: + if any(record["state"] == "ready" for record in state["entries"].values()): + raise OSError("injected failure") + original(backing, state, reclaim=reclaim) + + monkeypatch.setattr(manager, "_save", fail_commit) + result = await _call(tools, "file_memory_write", file_name="notes.md", content="new") + assert "content may have changed" in result + assert await store.read("conversation/notes.md") == "new" + monkeypatch.setattr(manager, "_save", original) + now[0] = 1000 + assert await manager.collect() == 0 + assert await store.read("conversation/notes.md") == "new" + + +async def test_registration_failure_preserves_previous_body(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="old") + + def fail(*args: Any, **kwargs: Any) -> None: + raise OSError("failure containing a private host path") + + monkeypatch.setattr(manager, "_save", fail) + result = await _call(tools, "file_memory_write", file_name="notes.md", content="new") + assert "private host path" not in result + assert await store.read("conversation/notes.md") == "old" + + +@pytest.mark.parametrize( + "payload", + [ + "{", + '{"schema_version":99}', + '{"schema_version":true}', + '{"schema_version":1,"generation":0,"resumed_expires_at":null}', + ], +) +async def test_unknown_or_corrupt_policy_is_preserved_and_cannot_delete_files(tmp_path: Path, payload: str) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=1) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="preserve") + path = tmp_path / ".retention/policy.json" + path.write_text(payload) + with pytest.raises(ValueError): + await manager.collect() + assert path.read_text() == payload + assert await store.read("conversation/notes.md") == "preserve" + + +async def test_resume_does_not_adopt_unregistered_legacy_files(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=1) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + await store.write("conversation/legacy.md", "legacy") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="managed") + await manager.set_retention(None) + now[0] = 200 + await manager.set_retention(1) + now[0] = 201 + assert await manager.collect() == 1 + assert await store.read("conversation/legacy.md") == "legacy" + + +def _manifest(directory: Path) -> Path: + return next(path for path in (directory / ".retention").glob("*.json") if path.name != "policy.json") + + +async def test_listing_and_unmatched_search_do_not_renew_other_records(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="a.md", content="alpha") + await _call(tools, "file_memory_write", file_name="b.md", content="beta") + now[0] = 109 + assert "a.md" in await _call(tools, "file_memory_ls") + assert await _call(tools, "file_memory_grep", regex_pattern="missing") == "[]" + assert "a.md" in await _call(tools, "file_memory_grep", regex_pattern="alpha") + now[0] = 110 + assert await manager.collect() == 1 + assert await store.read("conversation/a.md") == "alpha" + assert await store.read("conversation/b.md") is None + now[0] = 119 + assert await manager.collect() == 1 + + +@pytest.mark.parametrize("corruption", ["version", "state", "missing_expiry", "null_expiry", "reference"]) +async def test_invalid_manifest_preserves_original_bytes_and_body(tmp_path: Path, corruption: str) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=1) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="preserve") + path = _manifest(tmp_path) + state = json.loads(path.read_text()) + record = next(iter(state["entries"].values())) + if corruption == "version": + state["schema_version"] = 99 + elif corruption == "state": + record["state"] = "unknown" + elif corruption == "missing_expiry": + del record["expires_at"] + elif corruption == "null_expiry": + record["expires_at"] = None + else: + record["description"] = next(iter(state["entries"])) + path.write_text(json.dumps(state)) + original = path.read_bytes() + assert await manager.collect() == 0 + assert "Could not" in await _call(tools, "file_memory_write", file_name="notes.md", content="replace") + assert await store.read("conversation/notes.md") == "preserve" + assert path.read_bytes() == original + + +async def test_compatible_extension_fields_survive_overwrite(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + _, tools = await _prepare(FileMemoryProvider(FileSystemAgentFileStore(tmp_path / "owner"), retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="old") + path = _manifest(tmp_path) + state = json.loads(path.read_text()) + state["future_top_level"] = "keep" + state["store_limits"]["future_limit"] = "keep" + next(iter(state["entries"].values()))["future_field"] = "keep" + path.write_text(json.dumps(state)) + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="new") + after = json.loads(path.read_text()) + assert after["future_top_level"] == "keep" + assert after["store_limits"]["future_limit"] == "keep" + assert next(iter(after["entries"].values()))["future_field"] == "keep" + + +async def test_quota_refusal_before_body_write_does_not_leave_updating(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10, max_total_bytes=2000) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + # A body alone fits, but registration plus its atomic-copy reserve does not. + result = await _call(tools, "file_memory_write", file_name="notes.md", content="x" * 900) + assert "shared_max_total_bytes" in result + assert await store.read("conversation/notes.md") is None + for path in (tmp_path / ".retention").glob("*.json"): + if path.name != "policy.json": + assert not json.loads(path.read_text())["entries"] + + +async def test_interrupted_update_requires_explicit_repair_before_expiry(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="inspect me") + path = _manifest(tmp_path) + state = json.loads(path.read_text()) + next(iter(state["entries"].values()))["state"] = "updating" + path.write_text(json.dumps(state)) + now[0] = 1000 + assert await manager.collect() == 0 + assert await manager.repair(store) == 1 + assert await manager.repair(store) == 0 + now[0] = 1010 + assert await manager.collect() == 1 + + +async def test_empty_unleased_scope_directories_are_reclaimed(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + store = FileSystemAgentFileStore(tmp_path / "owner") + await store.create_directory("conversation") + await _prepare(FileMemoryProvider(store, retention=manager)) + assert (store.root_path / "conversation").is_dir() + assert await manager.collect() == 0 + assert not store.root_path.exists() + + +async def test_nested_quota_lock_name_is_an_ordinary_expiring_memory(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=1) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name=".agent-framework-quota.lock", content="ordinary") + now[0] = 101 + assert await manager.collect() == 1 + assert await store.read("conversation/.agent-framework-quota.lock") is None + + +async def test_corrupt_policy_does_not_pin_request_protection(tmp_path: Path) -> None: + from filelock import FileLock + + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + store = FileSystemAgentFileStore(tmp_path / "owner") + async with manager._protect_store(store): + token = next((tmp_path / ".retention/active").rglob("*.lock")) + (tmp_path / ".retention/policy.json").write_text("{") + assert not token.exists() + with FileLock(tmp_path / ".retention/manager.lock", timeout=0, preserve_lock_file=True): + pass + + +_WORKER = """ +import asyncio, json, os, sys +from pathlib import Path +from agent_framework import AgentSession, FileMemoryProvider, FileMemoryRetentionManager, FileSystemAgentFileStore +from agent_framework._sessions import SessionContext +directory, mode, root, name, budget, phase = sys.argv[1:] +async def main(): + manager = FileMemoryRetentionManager(directory, retention_seconds=10, + max_total_bytes=int(budget) if budget else None) + manager._now = lambda: 100.0 + store = FileSystemAgentFileStore(Path(directory) / root) + if mode == "lease": + manager._now = lambda: 109.0 + async with manager._protect_store(store): + print("ready", flush=True) + await asyncio.to_thread(sys.stdin.readline) + return + provider = FileMemoryProvider(store, retention=manager) + if mode in ("race_read", "race_gc") and phase == "no_hold": + print("attempting", flush=True) + context = SessionContext(session_id="conversation", input_messages=[]) + await provider.before_run(agent=None, session=AgentSession(session_id="conversation"), context=context, state={}) + tools = {tool.name: tool for tool in context.tools} + if mode in ("race_read", "race_gc"): + manager._now = lambda: 109.0 if mode == "race_read" else 111.0 + original = manager._save + held = False + def hold_after_save(backing, state, *, reclaim=False): + nonlocal held + original(backing, state, reclaim=reclaim) + if not held and phase != "no_hold": + held = True + print("held", flush=True) + sys.stdin.readline() + manager._save = hold_after_save + if mode == "race_gc": + print(json.dumps(await manager.collect()), flush=True) + else: + response = await tools["file_memory_read"].invoke(arguments={"file_name": name}) + print(json.dumps(response[0].text), flush=True) + return + if mode == "write": + print("ready", flush=True) + await asyncio.to_thread(sys.stdin.readline) + response = await tools["file_memory_write"].invoke(arguments={"file_name": name, "content": "x" * 600}) + print(json.dumps(response[0].text), flush=True) + return + if mode == "crash_write": + original = manager._atomic + def crash(path, payload, *, overwrite=True): + if phase == "before_body" and path.name == name: + os._exit(17) + original(path, payload, overwrite=overwrite) + if phase == "after_body" and path.name == name: + os._exit(17) + manager._atomic = crash + await tools["file_memory_write"].invoke(arguments={"file_name": name, "content": "new"}) + if mode == "crash_gc": + manager._now = lambda: 111.0 + original = manager._save + def crash(backing, state, *, reclaim=False): + original(backing, state, reclaim=reclaim) + if any(record["state"] == "deleting" for record in state["entries"].values()): + os._exit(17) + manager._save = crash + await manager.collect() +asyncio.run(main()) +""" + + +async def _worker( + directory: Path, mode: str, *, root: str = "owner", name: str = "notes.md", budget: str = "", phase: str = "" +) -> asyncio.subprocess.Process: + return await asyncio.create_subprocess_exec( + sys.executable, + "-c", + _WORKER, + str(directory), + mode, + root, + name, + budget, + phase, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + +async def test_independent_processes_cannot_lose_manifest_entries(tmp_path: Path) -> None: + processes = [await _worker(tmp_path, "write", name=f"{n}.md") for n in range(2)] + try: + for process in processes: + assert process.stdout is not None and process.stdin is not None + assert await asyncio.wait_for(process.stdout.readline(), 15) == b"ready\n" + for process in processes: + assert process.stdin is not None + process.stdin.write(b"go\n") + await process.stdin.drain() + for process in processes: + output, error = await asyncio.wait_for(process.communicate(), 15) + assert process.returncode == 0, error.decode() + assert "written" in json.loads(output) + assert len(json.loads(_manifest(tmp_path).read_text())["entries"]) == 2 + assert (tmp_path / "owner/conversation/0.md").read_text() == "x" * 600 + assert (tmp_path / "owner/conversation/1.md").read_text() == "x" * 600 + assert "0.md" in (tmp_path / "owner/conversation/memories.md").read_text() + assert "1.md" in (tmp_path / "owner/conversation/memories.md").read_text() + finally: + for process in processes: + if process.returncode is None: + process.kill() + await process.wait() + + +async def test_independent_roots_cannot_oversubscribe_shared_capacity(tmp_path: Path) -> None: + processes = [await _worker(tmp_path, "write", root=f"owner-{n}", budget="1900") for n in range(2)] + try: + for process in processes: + assert process.stdout is not None + assert await asyncio.wait_for(process.stdout.readline(), 15) == b"ready\n" + for process in processes: + assert process.stdin is not None + process.stdin.write(b"go\n") + await process.stdin.drain() + outcomes = [] + for process in processes: + output, error = await asyncio.wait_for(process.communicate(), 15) + assert process.returncode == 0, error.decode() + outcomes.append(json.loads(output)) + assert sum("written" in outcome for outcome in outcomes) == 1 + assert ( + await asyncio.to_thread(lambda: sum(p.stat().st_size for p in tmp_path.rglob("*") if p.is_file())) <= 1900 + ) + assert await asyncio.to_thread(lambda: len(list(tmp_path.glob("*/conversation/notes.md")))) == 1 + finally: + for process in processes: + if process.returncode is None: + process.kill() + await process.wait() + + +@pytest.mark.parametrize("phase", ["before_body", "after_body"]) +async def test_process_death_during_write_preserves_atomic_body_and_protects_update(tmp_path: Path, phase: str) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + manager._now = lambda: 100.0 + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="old") + process = await _worker(tmp_path, "crash_write", phase=phase) + _, error = await asyncio.wait_for(process.communicate(), 15) + assert process.returncode == 17, error.decode() + assert await store.read("conversation/notes.md") == ("old" if phase == "before_body" else "new") + manager._now = lambda: 1000.0 + assert await manager.collect() == 0 + assert next(iter(json.loads(_manifest(tmp_path).read_text())["entries"].values()))["state"] == "updating" + + +async def test_process_death_releases_os_lease_for_later_gc(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + process = await _worker(tmp_path, "lease") + try: + assert process.stdout is not None + assert await asyncio.wait_for(process.stdout.readline(), 15) == b"ready\n" + now[0] = 120 + assert await manager.collect() == 0 + process.kill() + await asyncio.wait_for(process.wait(), 15) + assert await manager.collect() == 1 + finally: + if process.returncode is None: + process.kill() + await process.wait() + + +async def test_interrupted_deletion_is_completed_idempotently_and_not_revived(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello", description="summary") + process = await _worker(tmp_path, "crash_gc") + _, error = await asyncio.wait_for(process.communicate(), 15) + assert process.returncode == 17, error.decode() + await manager.set_retention(None) + now[0] = 200 + await manager.set_retention(10) + assert "not found" in await _call(tools, "file_memory_read", file_name="notes.md") + assert "Could not" in await _call(tools, "file_memory_write", file_name="notes.md", content="revive") + assert await manager.collect() == 1 + assert await manager.collect() == 0 + assert await store.read("conversation/notes.md") is None + assert await store.read("conversation/notes_description.md") is None + + +async def test_cancelled_operation_finishes_commit_before_releasing_manager_lock(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + store = FileSystemAgentFileStore(tmp_path / "owner") + entered, finish = asyncio.Event(), asyncio.Event() + + async def write(view: Any) -> None: + await view.write("conversation/notes.md", "complete") + entered.set() + await finish.wait() + + operation = asyncio.create_task(manager._run(store, write, purpose="write")) + await asyncio.wait_for(entered.wait(), 5) + operation.cancel() + finish.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(operation, 5) + assert await store.read("conversation/notes.md") == "complete" + assert next(iter(json.loads(_manifest(tmp_path).read_text())["entries"].values()))["state"] == "ready" + + +async def test_cancellation_during_registration_exit_releases_token( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from collections.abc import AsyncGenerator + from contextlib import asynccontextmanager + + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + store = FileSystemAgentFileStore(tmp_path / "owner") + original = manager._locked + + @asynccontextmanager + async def cancel_on_exit(*, load_policy: bool = True) -> AsyncGenerator[None]: + async with original(load_policy=load_policy): + yield + if load_policy: + raise asyncio.CancelledError + + monkeypatch.setattr(manager, "_locked", cancel_on_exit) + with pytest.raises(asyncio.CancelledError): + async with manager._protect_store(store): + pytest.fail("registration should have been cancelled") + assert not list((tmp_path / ".retention/active").rglob("*.lock")) + + +async def test_file_access_cannot_expose_manager_control_directory(tmp_path: Path) -> None: + from types import SimpleNamespace + + from agent_framework import FileAccessProvider + + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + provider = FileMemoryProvider(FileSystemAgentFileStore(tmp_path / "owner"), retention=manager) + agent = SimpleNamespace(context_providers=[FileAccessProvider(FileSystemAgentFileStore(tmp_path))]) + with pytest.raises(ValueError, match="must not expose"): + await provider.before_run( + agent=agent, + session=AgentSession(session_id="conversation"), + context=SessionContext(session_id="conversation", input_messages=[]), + state={}, + ) + assert not (tmp_path / ".retention").exists() + + +async def test_description_read_cannot_bypass_body_expiry(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + _, tools = await _prepare(FileMemoryProvider(FileSystemAgentFileStore(tmp_path / "owner"), retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="body", description="private summary") + now[0] = 110 + assert "not found" in await _call(tools, "file_memory_read", file_name="notes_description.md") + + +def test_retention_rejects_unsupported_store_before_any_io(tmp_path: Path) -> None: + from agent_framework import InMemoryAgentFileStore + + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + with pytest.raises(ValueError, match="requires FileSystemAgentFileStore"): + FileMemoryProvider(InMemoryAgentFileStore(), retention=manager) + assert not (tmp_path / ".retention").exists() + + +async def test_root_quota_pressure_reclaims_expired_records_before_refusing(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner", max_files=2) + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + assert "written" in await _call(tools, "file_memory_write", file_name="old.md", content="old") + now[0] = 110 + assert "written" in await _call(tools, "file_memory_write", file_name="new.md", content="new") + assert await store.read("conversation/old.md") is None + assert await store.read("conversation/new.md") == "new" + + +@pytest.mark.parametrize("root_limit", [True, False]) +async def test_gc_index_rebuild_cannot_bypass_capacity(tmp_path: Path, root_limit: bool) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10, max_total_bytes=None if root_limit else 2000) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner", max_file_bytes=200 if root_limit else None) + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="old.md", content="old") + # Legacy files count toward storage but are not automatically given TTL. + for n in range(10): + await store.write(f"conversation/{n}.md", "legacy") + await store.write(f"conversation/{n}_description.md", "d" * (150 if root_limit else 80)) + before = await store.read("conversation/memories.md") + now[0] = 110 + assert before is not None + assert await manager.collect() == 1 + assert await store.read("conversation/old.md") is None + assert await store.read("conversation/memories.md") is None + assert not json.loads(_manifest(tmp_path).read_text())["entries"] + assert await manager.collect() == 0 + + +@pytest.mark.parametrize("first", ["read", "gc"]) +async def test_renew_and_gc_recheck_serialized_state_in_both_process_orders(tmp_path: Path, first: str) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + manager._now = lambda: 100.0 + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + primary = await _worker(tmp_path, f"race_{first}") + secondary: asyncio.subprocess.Process | None = None + try: + assert primary.stdout is not None and primary.stdin is not None + assert await asyncio.wait_for(primary.stdout.readline(), 15) == b"held\n" + secondary = await _worker(tmp_path, "race_gc" if first == "read" else "race_read", phase="no_hold") + assert secondary.stdout is not None + assert await asyncio.wait_for(secondary.stdout.readline(), 15) == b"attempting\n" + primary.stdin.write(b"release\n") + await primary.stdin.drain() + primary_output, primary_error = await asyncio.wait_for(primary.communicate(), 15) + second_output, second_error = await asyncio.wait_for(secondary.communicate(), 15) + assert primary.returncode == secondary.returncode == 0, (primary_error, second_error) + if first == "read": + assert json.loads(primary_output) == "hello" + assert json.loads(second_output) == 0 + assert await store.read("conversation/notes.md") == "hello" + else: + assert json.loads(primary_output) == 1 + assert "not found" in json.loads(second_output) + assert await store.read("conversation/notes.md") is None + finally: + for process in (primary, secondary): + if process is not None and process.returncode is None: + process.kill() + await process.wait() + + +async def test_managed_limited_scope_does_not_leave_a_redundant_quota_lock(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=1) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner", max_file_bytes=1024, max_files=100, max_total_bytes=10000) + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + assert not (store.root_path / ".agent-framework-quota.lock").exists() + now[0] = 101 + assert await manager.collect() == 1 + assert not store.root_path.exists() + assert (tmp_path / ".retention/manager.lock").exists() + + +async def test_process_death_before_data_creation_reclaims_orphan_request_token(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + process = await _worker(tmp_path, "lease", root="no-data") + try: + assert process.stdout is not None + assert await asyncio.wait_for(process.stdout.readline(), 15) == b"ready\n" + assert not (tmp_path / "no-data").exists() + assert await manager.collect() == 0 + process.kill() + await asyncio.wait_for(process.wait(), 15) + assert await manager.collect() == 0 + assert not await asyncio.to_thread(lambda: list((tmp_path / ".retention/active").iterdir())) + finally: + if process.returncode is None: + process.kill() + await process.wait() + + +async def test_custom_filesystem_backend_keeps_default_behavior_and_fails_retention_early(tmp_path: Path) -> None: + class EncodedStore(FileSystemAgentFileStore): + async def write(self, path: str, content: str, *, overwrite: bool = True) -> None: + await super().write(path, "encoded:" + content, overwrite=overwrite) + + async def read(self, path: str) -> str | None: + content = await super().read(path) + return content.removeprefix("encoded:") if content is not None else None + + store = EncodedStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store)) + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + assert (store.root_path / "conversation/notes.md").read_text() == "encoded:hello" + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + with pytest.raises(ValueError, match="without backend overrides"): + FileMemoryProvider(store, retention=manager) + assert not (tmp_path / ".retention").exists() + + +async def test_read_only_managed_requests_do_not_allocate_empty_data_directories(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + assert await _call(tools, "file_memory_ls") == "[]" + assert "not found" in await _call(tools, "file_memory_read", file_name="notes.md") + assert not store.root_path.exists() diff --git a/python/packages/core/tests/core/test_harness_file_store_limits.py b/python/packages/core/tests/core/test_harness_file_store_limits.py new file mode 100644 index 00000000000..7eb2fe8a3f6 --- /dev/null +++ b/python/packages/core/tests/core/test_harness_file_store_limits.py @@ -0,0 +1,195 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Behavioral coverage for opt-in file-store capacity limits.""" + +import asyncio +import sys +from pathlib import Path +from typing import Any + +import pytest + +from agent_framework import AgentFileStore, FileSystemAgentFileStore, InMemoryAgentFileStore + + +def _store(kind: str, directory: Path, **limits: Any) -> AgentFileStore: + """Exercise the same capacity contract through either backing implementation.""" + return InMemoryAgentFileStore(**limits) if kind == "memory" else FileSystemAgentFileStore(directory, **limits) + + +@pytest.mark.parametrize("kind", ["memory", "filesystem"]) +async def test_file_size_limit_counts_utf8_and_preserves_previous_content(kind: str, tmp_path: Path) -> None: + store = _store(kind, tmp_path, max_file_bytes=4) + await store.write("notes.md", "old") + with pytest.raises(ValueError, match="max_file_bytes"): + await store.write("notes.md", "你好") + assert await store.read("notes.md") == "old" + + +@pytest.mark.parametrize("kind", ["memory", "filesystem"]) +async def test_store_limits_account_for_overwrite_delete_and_nested_files(kind: str, tmp_path: Path) -> None: + store = _store(kind, tmp_path, max_files=2, max_total_bytes=6) + await store.write("nested/first.txt", "abc") + await store.write("second.txt", "def") + with pytest.raises(ValueError, match="max_files"): + await store.write("third.txt", "") + with pytest.raises(ValueError, match="max_total_bytes"): + await store.write("second.txt", "defg") + assert await store.read("second.txt") == "def" + await store.write("nested/first.txt", "a") + await store.write("second.txt", "defgh") + await store.delete("nested/first.txt") + await store.write("third.txt", "x") + assert await store.read("third.txt") == "x" + + +async def test_filesystem_limits_include_existing_files_and_survive_new_instances(tmp_path: Path) -> None: + (tmp_path / "existing.txt").write_text("abc", encoding="utf-8") + store = FileSystemAgentFileStore(tmp_path, max_files=1, max_total_bytes=3) + with pytest.raises(ValueError, match="max_files"): + await store.write("new.txt", "x") + other = FileSystemAgentFileStore(tmp_path, max_total_bytes=3) + with pytest.raises(ValueError, match="max_total_bytes"): + await other.write("existing.txt", "abcd") + assert (tmp_path / "existing.txt").read_text(encoding="utf-8") == "abc" + + +@pytest.mark.parametrize("kind", ["memory", "filesystem"]) +@pytest.mark.parametrize("name", ["max_file_bytes", "max_files", "max_total_bytes"]) +@pytest.mark.parametrize("invalid", [0, -1, True, 1.5, "4"]) +def test_explicit_invalid_limits_fail_without_creating_storage( + kind: str, name: str, invalid: Any, tmp_path: Path +) -> None: + root = tmp_path / "uncreated" + with pytest.raises(ValueError, match=name): + _store(kind, root, **{name: invalid}) + assert not root.exists() + + +async def test_unconfigured_filesystem_store_keeps_original_paths_and_metadata_behavior(tmp_path: Path) -> None: + root = tmp_path / "uncreated" + store = FileSystemAgentFileStore(root) + assert not root.exists() + await store.write(".agent-framework-quota.lock", "existing application data") + assert await store.read(".agent-framework-quota.lock") == "existing application data" + assert [entry.name for entry in await store.list_children()] == [".agent-framework-quota.lock"] + + +@pytest.mark.parametrize("kind", ["memory", "filesystem"]) +async def test_exclusive_create_keeps_file_exists_error_with_limits(kind: str, tmp_path: Path) -> None: + store = _store(kind, tmp_path, max_file_bytes=3, max_files=1) + await store.write("first.txt", "abc") + with pytest.raises(FileExistsError): + await store.write("first.txt", "oversized", overwrite=False) + assert await store.read("first.txt") == "abc" + + +@pytest.mark.parametrize("kind", ["memory", "filesystem"]) +async def test_concurrent_writes_cannot_oversubscribe_root(kind: str, tmp_path: Path) -> None: + shared = _store(kind, tmp_path, max_files=1, max_total_bytes=3) + stores = ( + [shared] * 12 + if kind == "memory" + else [FileSystemAgentFileStore(tmp_path, max_files=1, max_total_bytes=3) for _ in range(12)] + ) + outcomes = await asyncio.gather( + *(store.write(f"file-{n}.txt", "abc") for n, store in enumerate(stores)), return_exceptions=True + ) + assert sum(result is None for result in outcomes) == 1 + assert all(result is None or isinstance(result, ValueError) for result in outcomes) + entries = await shared.list_children() + assert len(entries) == 1 + assert await shared.read(entries[0].name) == "abc" + + +async def test_separate_roots_have_independent_capacity(tmp_path: Path) -> None: + first = FileSystemAgentFileStore(tmp_path / "first", max_files=1, max_total_bytes=3) + second = FileSystemAgentFileStore(tmp_path / "second", max_files=1, max_total_bytes=3) + await asyncio.gather(first.write("note.txt", "abc"), second.write("note.txt", "def")) + assert await first.read("note.txt") == "abc" + assert await second.read("note.txt") == "def" + + +async def test_read_limit_preserves_oversized_existing_file_and_hides_identity(tmp_path: Path) -> None: + root = tmp_path / "private-user-identity" + root.mkdir() + (root / "note.txt").write_text("abcdef", encoding="utf-8") + store = FileSystemAgentFileStore(root, max_file_bytes=4) + with pytest.raises(ValueError, match="max_file_bytes") as error: + await store.read("note.txt") + assert "private-user-identity" not in str(error.value) + assert (root / "note.txt").read_text(encoding="utf-8") == "abcdef" + assert not (root / ".agent-framework-quota.lock").exists() + + +async def test_quota_metadata_is_hidden_and_cannot_replace_existing_data(tmp_path: Path) -> None: + store = FileSystemAgentFileStore(tmp_path, max_files=1) + await store.write("notes.txt", "hello") + assert [entry.name for entry in await store.list_children()] == ["notes.txt"] + assert not await store.file_exists(".agent-framework-quota.lock") + for operation in (store.read, store.delete): + with pytest.raises(ValueError, match="reserved"): + await operation(".agent-framework-quota.lock") + with pytest.raises(ValueError, match="reserved"): + await store.write(".agent-framework-quota.lock", "replacement") + assert [item.file_name for item in await store.search("", "hello")] == ["notes.txt"] + + other_root = tmp_path / "legacy" + other_root.mkdir() + lock_path = other_root / ".agent-framework-quota.lock" + lock_path.write_text("application data", encoding="utf-8") + other = FileSystemAgentFileStore(other_root, max_files=2) + with pytest.raises(ValueError, match="conflicts"): + await other.write("note.txt", "x") + assert lock_path.read_text(encoding="utf-8") == "application data" + + +async def test_filesystem_capacity_is_shared_between_processes(tmp_path: Path) -> None: + script = """import asyncio, sys +from agent_framework import FileSystemAgentFileStore +store = FileSystemAgentFileStore(sys.argv[1], max_files=1, max_total_bytes=3) +print('ready', flush=True) +sys.stdin.readline() +try: + asyncio.run(store.write(sys.argv[2], 'abc')) + print('written', flush=True) +except ValueError: + print('refused', flush=True) +""" + processes = [ + await asyncio.create_subprocess_exec( + sys.executable, + "-c", + script, + str(tmp_path), + f"child-{n}.txt", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + for n in range(2) + ] + try: + for process in processes: + assert process.stdout is not None + assert await asyncio.wait_for(process.stdout.readline(), 20) == b"ready\n" + for process in processes: + assert process.stdin is not None + process.stdin.write(b"go\n") + await process.stdin.drain() + process.stdin.close() + results = await asyncio.gather(*(asyncio.wait_for(p.communicate(), 20) for p in processes)) + assert sorted(out.decode().strip() for out, _ in results) == ["refused", "written"] + assert all(process.returncode == 0 for process in processes) + finally: + for process in processes: + if process.returncode is None: + process.kill() + await process.wait() + + +async def test_search_reports_configured_read_quota_instead_of_empty_matches(tmp_path: Path) -> None: + (tmp_path / "large.txt").write_text("matching", encoding="utf-8") + store = FileSystemAgentFileStore(tmp_path, max_file_bytes=3) + with pytest.raises(ValueError, match="max_file_bytes"): + await store.search("", "matching") diff --git a/python/packages/core/tests/core/test_harness_hosted_memory_lifetime.py b/python/packages/core/tests/core/test_harness_hosted_memory_lifetime.py new file mode 100644 index 00000000000..3de0da99bfc --- /dev/null +++ b/python/packages/core/tests/core/test_harness_hosted_memory_lifetime.py @@ -0,0 +1,175 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Offline verification of the hosted sample's actual request/stream protection.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +from collections.abc import AsyncIterator +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +import pytest + +from agent_framework import ( + AgentContext, + AgentResponse, + AgentResponseUpdate, + AgentSession, + FileMemoryProvider, + FileMemoryRetentionManager, + Message, + ResponseStream, +) +from agent_framework._sessions import SessionContext + + +def _load_hosted(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + pytest.importorskip("agent_framework_foundry_hosting") + pytest.importorskip("agent_framework_monty") + pytest.importorskip("azure.ai.agentserver.core") + sample = Path(__file__).parents[4] / "samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready" + monkeypatch.syspath_prepend(str(sample)) + + def load(name: str, path: Path) -> ModuleType: + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + load("agent", sample / "agent.py") + return load("_file_memory_hosted_sample", sample / "hosted.py") + + +async def _agent(monkeypatch: pytest.MonkeyPatch, directory: Path, manager: FileMemoryRetentionManager) -> Any: + hosted = _load_hosted(monkeypatch) + monkeypatch.setattr(hosted.Path, "home", lambda: directory) + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://example.invalid") + monkeypatch.setenv("AZURE_AI_MODEL_DEPLOYMENT_NAME", "offline-test") + monkeypatch.setattr(hosted, "AgentConfig", SimpleNamespace(from_env=lambda: SimpleNamespace(is_hosted=False))) + monkeypatch.setattr(hosted, "get_request_context", lambda: SimpleNamespace(platform_headers=lambda: {})) + monkeypatch.setattr( + hosted, + "FoundryRequestScope", + SimpleNamespace(from_context=lambda *args, **kwargs: SimpleNamespace(storage_key="owner")), + ) + monkeypatch.setattr(hosted, "AzureCliCredential", lambda: SimpleNamespace(close=lambda: None)) + monkeypatch.setattr(hosted, "_build_purview_middleware", lambda credential: []) + + class Client: + def __init__(self, **kwargs: Any) -> None: + pass + + monkeypatch.setattr(hosted, "FoundryChatClient", Client) + + async def build(**kwargs: Any) -> Any: + return SimpleNamespace(context_providers=[kwargs["file_memory_provider"]], middleware=[]) + + monkeypatch.setattr(hosted, "build_claw_agent", build) + return await hosted.create_agent(retention=manager) + + +@pytest.mark.parametrize("ending", ["normal", "close", "cancel", "error"]) +async def test_hosted_stream_keeps_lease_until_actual_consumption_ends( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ending: str +) -> None: + directory = tmp_path / ".claw/agent-file-memory" + manager = FileMemoryRetentionManager(directory, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + agent = await _agent(monkeypatch, tmp_path, manager) + provider: FileMemoryProvider = agent.context_providers[0] + context = SessionContext(session_id="conversation", input_messages=[]) + await provider.before_run(agent=None, session=AgentSession(session_id="conversation"), context=context, state={}) + tools = {tool.name: tool for tool in context.tools} + await tools["file_memory_write"].invoke(arguments={"file_name": "notes.md", "content": "hello"}) + now[0] = 109 + first, gate = asyncio.Event(), asyncio.Event() + + async def source() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant") + if ending == "cancel": + await gate.wait() + if ending == "error": + raise RuntimeError("upstream failure") + + inner: ResponseStream[AgentResponseUpdate, AgentResponse] = ResponseStream( + source(), finalizer=lambda _: AgentResponse(messages=[Message(role="assistant", contents=["done"])]) + ) + run_context = AgentContext(agent=agent, messages=[], session=AgentSession(session_id="conversation"), stream=True) + + async def next_call() -> None: + run_context.result = inner + + await agent.middleware[0](run_context, next_call) + stream = run_context.result + assert isinstance(stream, ResponseStream) + now[0] = 120 + assert await manager.collect() == 0 + + async def consume() -> None: + async for _ in stream: + first.set() + await stream.get_final_response() + + if ending == "close": + await stream.close() + await stream.close() + elif ending == "cancel": + consumer = asyncio.create_task(consume()) + await asyncio.wait_for(first.wait(), 5) + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(consumer, 5) + elif ending == "error": + with pytest.raises(RuntimeError, match="upstream failure"): + await consume() + else: + await consume() + + assert await manager.collect() == 1 + assert await provider.store.read("conversation/notes.md") is None + + +async def test_hosted_nonstream_failure_releases_request_lease(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + directory = tmp_path / ".claw/agent-file-memory" + manager = FileMemoryRetentionManager(directory, retention_seconds=10) + agent = await _agent(monkeypatch, tmp_path, manager) + context = AgentContext(agent=agent, messages=[], session=AgentSession(session_id="conversation")) + + async def fail() -> None: + raise RuntimeError("model failed") + + with pytest.raises(RuntimeError, match="model failed"): + await agent.middleware[0](context, fail) + assert not await asyncio.to_thread(lambda: list((directory / ".retention/active").rglob("*.lock"))) + + +async def test_hosted_server_runs_gc_without_receiving_requests( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + hosted = _load_hosted(monkeypatch) + monkeypatch.setattr(hosted.Path, "home", lambda: tmp_path) + monkeypatch.setenv("CLAW_FILE_MEMORY_RETENTION_SECONDS", "1") + monkeypatch.setenv("CLAW_FILE_MEMORY_SWEEP_INTERVAL_SECONDS", "1") + swept = asyncio.Event() + + async def collect(self: FileMemoryRetentionManager) -> int: + swept.set() + return 0 + + class Server: + def __init__(self, **kwargs: Any) -> None: + pass + + async def run_async(self) -> None: + await asyncio.wait_for(swept.wait(), 5) + + monkeypatch.setattr(hosted.FileMemoryRetentionManager, "collect", collect) + monkeypatch.setattr(hosted, "ResponsesHostServer", Server) + await hosted.main() + assert swept.is_set() diff --git a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md index 069cfa6492b..37e296f3d2a 100644 --- a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md +++ b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md @@ -134,6 +134,89 @@ resources require separate configuration. The unrelated Telegram sample also requires externally configured Telegram/Key Vault credentials. None of those credential-gated behaviors or deployments is proven by an offline smoke check. +### File-memory capacity and retention + +SDK limits are opt-in: omitted `max_file_bytes`, `max_files`, and `max_total_bytes` +remain unlimited, and `FileMemoryProvider` without a retention manager keeps its existing +behavior. Retention currently supports the builtin filesystem store; customized backends +continue using the existing interface and are rejected if retention is explicitly enabled. +The hosted sample supplies finite limits and one server-owned +`FileMemoryRetentionManager` shared by all request-owned stores. + +| Environment variable | Initial default | Meaning | +| --- | --- | --- | +| `CLAW_FILE_MEMORY_MAX_FILE_BYTES` | `1048576` (1 MiB) | Maximum stored file / whole-file read bytes | +| `CLAW_FILE_MEMORY_MAX_FILES` | `100` | Regular files under each trusted user/sandbox root | +| `CLAW_FILE_MEMORY_MAX_TOTAL_BYTES` | `10485760` (10 MiB) | Combined content bytes under each root | +| `CLAW_FILE_MEMORY_SHARED_MAX_BYTES` | `268435456` (256 MiB) | Shared contents, control metadata, and atomic-copy reserve | +| `CLAW_FILE_MEMORY_RETENTION_SECONDS` | `2592000` (30 days) | Initial idle lifetime for a new managed area; `none` initially disables TTL | +| `CLAW_FILE_MEMORY_SWEEP_INTERVAL_SECONDS` | `300` (5 minutes) | Server-driven cleanup interval, including when there are no requests | + +These are example starting points, not universal deployment recommendations. Numeric values +must be positive integers; invalid configuration fails instead of disabling protection. +Root limits count existing files, descriptions, and `memories.md`, so `max_files` does not mean +that many user-visible memories. Bytes count UTF-8 content, not filesystem allocation. +The shared budget covers all roots and `.retention` control files; it reserves one staging +file and the size of the largest file for an atomic replacement. Filesystem allocation, +external writers, and unrelated storage are outside these limits. Reducing a budget does +not delete unexpired data. New growth is refused; request registration still requires +space for its protection token, and index repair may require additional capacity. + +Retention applies to a logical memory: its body, optional description, and index entry. +Successful body reads, edits, and actual grep matches renew that memory. Listing, +unmatched search, index injection, and background scanning do not renew every memory. +Expired memories disappear from read/list/search/index before the next physical cleanup. +GC deletes eligible bodies and descriptions still unreferenced by other bodies; it only +prunes empty directories. Existing unregistered files stay readable and count toward +capacity, but are not enrolled or deleted by reading/scanning. A successful overwrite +starts their managed lifetime. + +`.retention/policy.json` records the shared policy; per-root versioned manifests use hashed +record identifiers rather than filenames or identity values. This directory is host-owned +and must not be exposed through file access, shell, or custom tools. The provider rejects +overlapping filesystem `FileAccessProvider` roots. Every writer sharing the managed area +must use the manager with identical shared limits and consistent per-root limits; +direct store writes, older workers without management, multi-host +deployment, and network filesystems do not provide these guarantees. File locks coordinate +local processes and release after process death. Fixed coordination lock files stay in place. + +The host keeps request protection through the full run and stream consumption/close. +GC skips scopes with an active request. That request can still use records that were valid +when it began; new requests cannot revive an expired record by reading it. Empty scopes +are reclaimed after their requests end. An existing root quota-lock marker is retained. + +A normal restart loads the stored TTL policy and preserves deadlines. The TTL environment +value initializes a new area only; change an existing policy explicitly: + +```python +# Use the same host-owned manager/directory as the server. +await retention.set_retention(None) # disable expiration; retain metadata +await retention.set_retention(30 * 86400) # explicitly resume with a complete new TTL +``` + +While disabled, reads remain available without renewal and TTL GC deletes no memories; +cleanup of dead request-protection markers continues. +Resuming grants still-present, registered `ready` records a complete lifetime from resume. +Repeated same-value calls do nothing. Changing one positive TTL to another affects new +writes and subsequent successful uses, not all existing deadlines. Deleted records cannot +be restored; interrupted updates/deletions are not revived by resuming. + +Content, description, index, and manifest are not one filesystem transaction. Managed +files use staged atomic replacement; an interrupted update remains protected from expiry. +Failures after content changes report partial completion. Unknown/corrupt metadata is +preserved and management for that scope stops. Operators inspect current content and then +explicitly confirm interrupted updates with `await retention.repair(store)`; this rebuilds +the index and starts a fresh lifetime, without claiming to recover the original operation. +Pending deletions remain deletions and GC retries interrupted I/O. If a rebuilt index would +exceed capacity, GC removes the stale index and completes confirmed deletion; the provider +reconstructs its index projection on reads when capacity permits. Explicit update repair may +still need additional capacity. Maintenance does not bypass file or shared quotas. + +The existing description naming replaces extensions, so `notes.md` and `notes.txt` share +`notes_description.md`. Descriptions retain this existing last-write-wins behavior; cleanup +preserves a sidecar while any other existing body still references it. Use distinct stems +when memories need independent descriptions. + ### Deploy to Foundry ```bash diff --git a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/agent.py b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/agent.py index 1bb6ababfc1..9cec4a8824b 100644 --- a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/agent.py +++ b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/agent.py @@ -48,6 +48,7 @@ AggregatingSkillsSource, DeduplicatingSkillsSource, FileAccessProvider, + FileMemoryProvider, FileSkillsSource, FileSystemAgentFileStore, HistoryProvider, @@ -300,6 +301,7 @@ async def build_claw_agent( enable_file_access: bool = True, file_access_store: Any = None, file_memory_store: Any = None, + file_memory_provider: FileMemoryProvider | None = None, enable_shell: bool = True, purview_credential: TokenCredential | None = None, auto_approve_skill_scripts: bool = False, @@ -324,6 +326,8 @@ async def build_claw_agent( None, the harness default is used, which writes to ``{cwd}/agent-file-memory``. Hosted deployments must supply a store rooted at a writable path, because the deployed code directory is mounted read-only. + file_memory_provider: Optional request-owned provider with lifecycle management. Mutually + exclusive with file_memory_store; the host owns its protection lifetime. enable_shell: When True (default), the agent can run shell commands. Disable it on shared/hosted deployments: arbitrary command execution inside the container is a serious security risk (data exfiltration, persistence, tampering) even behind a deny-list. @@ -338,6 +342,8 @@ async def build_claw_agent( Returns: A fully configured harness agent with Step 03 capabilities plus opt-in Purview middleware. """ + if file_memory_store is not None and file_memory_provider is not None: + raise ValueError("Supply either file_memory_store or file_memory_provider.") load_dotenv() # @@ -380,6 +386,8 @@ async def build_claw_agent( # context_providers: list[Any] = [MontyCodeActProvider(approval_mode="never_require")] + if file_memory_provider is not None: + context_providers.append(file_memory_provider) logger.info("CodeAct enabled (Monty).") # @@ -400,6 +408,7 @@ async def build_claw_agent( history_provider=history_provider or InMemoryHistoryProvider(), file_access_store=access_store, file_memory_store=file_memory_store, + disable_file_memory=file_memory_provider is not None, skills_provider=skills_provider, background_agents=[research_agent], shell_executor=shell, diff --git a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/hosted.py b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/hosted.py index 7192ed02b57..552eb882493 100644 --- a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/hosted.py +++ b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/hosted.py @@ -58,6 +58,8 @@ AgentContext, AgentResponseUpdate, BackgroundAgentsProvider, + FileMemoryProvider, + FileMemoryRetentionManager, FileSystemAgentFileStore, InMemoryHistoryProvider, ResponseStream, @@ -102,10 +104,18 @@ def _log_environment() -> None: ) -async def create_agent() -> Agent: +async def create_agent(*, retention: FileMemoryRetentionManager | None = None) -> Agent: """Build request-owned tools and scope file memory by trusted user and sandbox.""" config, context = AgentConfig.from_env(), get_request_context() scope = FoundryRequestScope.from_context(config, context, local_session_id="local-development") + memory_dir = Path.home() / ".claw" / "agent-file-memory" / scope.storage_key + memory_store = FileSystemAgentFileStore( + memory_dir, + max_file_bytes=int(os.environ.get("CLAW_FILE_MEMORY_MAX_FILE_BYTES", "1048576")), + max_files=int(os.environ.get("CLAW_FILE_MEMORY_MAX_FILES", "100")), + max_total_bytes=int(os.environ.get("CLAW_FILE_MEMORY_MAX_TOTAL_BYTES", "10485760")), + ) + memory_provider = FileMemoryProvider(memory_store, retention=retention) endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"] model = os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"] credential = ( @@ -113,7 +123,6 @@ async def create_agent() -> Agent: if config.is_hosted else AzureCliCredential() ) - memory_dir = Path.home() / ".claw" / "agent-file-memory" / scope.storage_key class RequestClient(FoundryChatClient): async def __aenter__(self) -> RequestClient: @@ -145,7 +154,7 @@ async def __aexit__( # external file_access_store (e.g. one backed by Azure Blob Storage) instead of the disk. enable_file_access=False, enable_shell=False, - file_memory_store=FileSystemAgentFileStore(memory_dir), + file_memory_provider=memory_provider, # Purview authenticates via the container's managed identity; InteractiveBrowserCredential # cannot run on a headless hosted container. purview_credential=credential, @@ -159,6 +168,8 @@ async def request_background_lifetime(context: AgentContext, call_next: Callable if context.session is None: raise RuntimeError("The hosted harness requires a request-owned AgentSession.") session = context.session + protection = memory_provider.protect() + await protection.__aenter__() @asynccontextmanager async def release_background() -> AsyncIterator[None]: @@ -170,8 +181,10 @@ async def release_background() -> AsyncIterator[None]: raise finally: try: - for provider in background_providers: - await provider.release_session(session) + async with AsyncExitStack() as release: + release.push_async_callback(protection.__aexit__, None, None, None) + for provider in background_providers: + await provider.release_session(session) except BaseException as cleanup_error: logger.error("Failed to release request-owned background tasks (%s).", type(cleanup_error).__name__) if not run_failed: @@ -184,10 +197,11 @@ async def release_background() -> AsyncIterator[None]: except BaseException: async with release_background(): raise - inner = context.result - if not isinstance(inner, ResponseStream): + result = context.result + if not isinstance(result, ResponseStream): async with release_background(): raise RuntimeError("The streaming hosted harness must return a ResponseStream.") + inner = result cleaned_up = False @@ -232,8 +246,21 @@ async def main() -> None: _configure_logging() load_dotenv() _log_environment() - server = ResponsesHostServer(agent=create_agent, history_source="agent_server") - await server.run_async() + value = os.environ.get("CLAW_FILE_MEMORY_RETENTION_SECONDS", "2592000") + retention_seconds = None if value.lower() == "none" else int(value) + retention = FileMemoryRetentionManager( + Path.home() / ".claw" / "agent-file-memory", + retention_seconds=retention_seconds, + max_total_bytes=int(os.environ.get("CLAW_FILE_MEMORY_SHARED_MAX_BYTES", "268435456")), + sweep_interval_seconds=int(os.environ.get("CLAW_FILE_MEMORY_SWEEP_INTERVAL_SECONDS", "300")), + ) + async with retention: + + async def request_agent() -> Agent: + return await create_agent(retention=retention) + + server = ResponsesHostServer(agent=request_agent, history_source="agent_server") + await server.run_async() if __name__ == "__main__": diff --git a/python/uv.lock b/python/uv.lock index 580a37e3565..bc361ebedf1 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -443,6 +443,7 @@ name = "agent-framework-core" version = "1.19.0" source = { editable = "packages/core" } dependencies = [ + { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -496,6 +497,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "filelock", specifier = ">=3.30.3,<4" }, { name = "agent-framework-a2a", marker = "extra == 'all'", editable = "packages/a2a" }, { name = "agent-framework-ag-ui", marker = "extra == 'all'", editable = "packages/ag-ui" }, { name = "agent-framework-anthropic", marker = "extra == 'all'", editable = "packages/anthropic" }, @@ -2221,6 +2223,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/03/10388a42375ee7e4ac9b94eb2c5c569c8b5795e377e701c9ac3ad63de890/fastapi-0.141.1-py3-none-any.whl", hash = "sha256:bfb91aa2d334c61cb35ba9a116fc123b3d3df31640b801cf57a7a78ec3f603b3", size = 131954, upload-time = "2026-07-29T17:18:04.364Z" }, ] +[[package]] +name = "filelock" +version = "3.32.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/59/e19834834cb01a32febfbb0f8a23a9088088f5d45991824ff2bc3b5e8acb/filelock-3.32.7.tar.gz", hash = "sha256:37b8a3d9811b0f9aef7e5ec5c71bb320de52df51e6ca9bcd6f5ad81187660da7", size = 225154, upload-time = "2026-09-16T00:24:20.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/df/31098c5aeb4d966b553641472bd55fcf5fdfac953549894b8a765ba44e91/filelock-3.32.7-py3-none-any.whl", hash = "sha256:65ff0d0190ea42038b32bda4b77834fb05be2cad4c5b9b01aa4dfb3614536e52", size = 100157, upload-time = "2026-09-16T00:24:19.543Z" }, +] + [[package]] name = "filetype" version = "1.2.0" From d67a8b3c9e2b12b0cbfd07c17a862610e721055f Mon Sep 17 00:00:00 2001 From: quifox <289420841+quifox@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:00:41 +0000 Subject: [PATCH 2/2] Python: Address file-memory retention review findings Complete confirmed deletions while TTL is disabled, clear omitted descriptions on overwrite, and report list/quota errors consistently. Move filesystem accounting and lifecycle I/O off the event loop while retaining coordination through cancellation. Cover legacy limits, shared descriptions, deletion recovery, quota diagnostics, and request cancellation with regression tests. --- .../agent_framework/_harness/_file_memory.py | 5 +- .../_harness/_file_memory_retention.py | 379 ++++++++++-------- .../tests/core/test_harness_file_memory.py | 17 + .../test_harness_file_memory_retention.py | 161 ++++++++ .../claw_step04_production_ready/README.md | 5 +- 5 files changed, 405 insertions(+), 162 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_file_memory.py b/python/packages/core/agent_framework/_harness/_file_memory.py index 3524996e32e..161c5a2bf6d 100644 --- a/python/packages/core/agent_framework/_harness/_file_memory.py +++ b/python/packages/core/agent_framework/_harness/_file_memory.py @@ -561,7 +561,10 @@ async def file_memory_ls(glob_pattern: str | None = None) -> list[dict[str, Any] description: str | None = None desc_file_name = _description_file_name(file_name) if desc_file_name in available: - description = await self._active_store().read(_combine_paths(working_folder, desc_file_name)) + try: + description = await self._active_store().read(_combine_paths(working_folder, desc_file_name)) + except _FileStoreQuotaError as exc: + return f"Could not list memory files: {exc}" results.append({"name": file_name, "type": "file", "description": description}) return results diff --git a/python/packages/core/agent_framework/_harness/_file_memory_retention.py b/python/packages/core/agent_framework/_harness/_file_memory_retention.py index 1a9c0f936f3..519c373a8f3 100644 --- a/python/packages/core/agent_framework/_harness/_file_memory_retention.py +++ b/python/packages/core/agent_framework/_harness/_file_memory_retention.py @@ -140,7 +140,7 @@ def _prepare_control(self) -> None: @asynccontextmanager async def _locked(self, *, load_policy: bool = True) -> AsyncGenerator[None]: - self._prepare_control() + await _complete(asyncio.to_thread(self._prepare_control)) # Keep the coordination pathname stable, including on Windows. Native OS # locks release after process death; a soft-file fallback cannot promise that. # https://py-filelock.readthedocs.io/en/latest/changelog.html @@ -149,65 +149,68 @@ async def _locked(self, *, load_policy: bool = True) -> AsyncGenerator[None]: ) try: async with lock: - for path in (self._control / "staging").iterdir(): - self._safe(path) - if path.is_file() and path.name.endswith(".tmp"): - path.unlink() + + def cleanup_staging() -> None: + for path in (self._control / "staging").iterdir(): + self._safe(path) + if path.is_file() and path.name.endswith(".tmp"): + path.unlink() + + await _complete(asyncio.to_thread(cleanup_staging)) if load_policy: - self._load_policy() + policy = await _complete(asyncio.to_thread(self._load_policy)) + self._policy = policy + self.retention_seconds = policy["retention_seconds"] yield except Timeout as exc: raise TimeoutError("Timed out acquiring the file-memory management lock.") from exc - def _load_policy(self) -> None: + def _load_policy(self) -> dict[str, Any]: path = self._control / "policy.json" self._safe(path) if not path.exists(): - self._policy = { + policy = { "schema_version": 1, "retention_seconds": self._initial_retention, "generation": 0, "resumed_expires_at": None, } - self._save_policy() - else: - if path.stat().st_size > 16384: - raise ValueError("The file-memory policy exceeds its supported size.") - try: - raw = json.loads(path.read_text(encoding="utf-8")) - except (UnicodeError, json.JSONDecodeError) as exc: - raise ValueError("The file-memory policy is invalid; management is paused.") from exc - if ( - not isinstance(raw, dict) - or cast(dict[str, Any], raw).get("schema_version") != 1 - or isinstance(cast(dict[str, Any], raw).get("schema_version"), bool) - ): - raise ValueError("Unsupported file-memory policy version; management is paused.") - raw = cast(dict[str, Any], raw) - if not {"retention_seconds", "generation", "resumed_expires_at"}.issubset(raw): - raise ValueError("The file-memory policy is missing required fields.") - _FileStoreLimits.create(raw["retention_seconds"], None, None) - generation, deadline = raw.get("generation"), raw.get("resumed_expires_at") - if ( - type(generation) is not int - or generation < 0 - or ( - deadline is not None - and ( - isinstance(deadline, bool) - or not isinstance(deadline, (int, float)) - or not math.isfinite(deadline) - ) + self._save_policy(policy) + return policy + if path.stat().st_size > 16384: + raise ValueError("The file-memory policy exceeds its supported size.") + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (UnicodeError, json.JSONDecodeError) as exc: + raise ValueError("The file-memory policy is invalid; management is paused.") from exc + if ( + not isinstance(raw, dict) + or cast(dict[str, Any], raw).get("schema_version") != 1 + or isinstance(cast(dict[str, Any], raw).get("schema_version"), bool) + ): + raise ValueError("Unsupported file-memory policy version; management is paused.") + raw = cast(dict[str, Any], raw) + if not {"retention_seconds", "generation", "resumed_expires_at"}.issubset(raw): + raise ValueError("The file-memory policy is missing required fields.") + _FileStoreLimits.create(raw["retention_seconds"], None, None) + generation, deadline = raw.get("generation"), raw.get("resumed_expires_at") + if ( + type(generation) is not int + or generation < 0 + or ( + deadline is not None + and ( + isinstance(deadline, bool) or not isinstance(deadline, (int, float)) or not math.isfinite(deadline) ) - or (generation > 0 and deadline is None) - ): - raise ValueError("The file-memory policy has invalid recovery metadata.") - self._policy = raw - self.retention_seconds = self._policy["retention_seconds"] + ) + or (generation > 0 and deadline is None) + ): + raise ValueError("The file-memory policy has invalid recovery metadata.") + return raw - def _save_policy(self) -> None: + def _save_policy(self, policy: dict[str, Any]) -> None: path = self._control / "policy.json" - payload = json.dumps(self._policy, sort_keys=True, separators=(",", ":")).encode("utf-8") + payload = json.dumps(policy, sort_keys=True, separators=(",", ":")).encode("utf-8") if len(payload) > 16384: raise ValueError("The file-memory policy exceeds its supported size.") self._check_growth(path, len(payload)) @@ -220,6 +223,7 @@ async def set_retention(self, retention_seconds: int | None) -> None: retention_seconds: Positive lifetime in seconds, or None to disable TTL. Disabling retains deadlines and permits ordinary reads without renewal. + Already-confirmed deletions continue to completion while TTL is disabled. Reenabling grants every retained ready record a complete new lifetime. Same-value calls are no-ops. Changing one positive value to another only changes subsequent successful writes/uses. Construction/restart reads an @@ -238,8 +242,8 @@ async def work() -> None: if previous is None and retention_seconds is not None: policy["generation"] += 1 policy["resumed_expires_at"] = self._now() + retention_seconds + await asyncio.to_thread(self._save_policy, policy) self._policy = policy - self._save_policy() self.retention_seconds = retention_seconds await _complete(work()) @@ -341,7 +345,7 @@ def _check_combined_growth(self, changes: list[tuple[Path, int]]) -> None: resulting_bytes + reserve > max(total + largest, self._capacity.max_total_bytes) ): raise _FileStoreQuotaError( - "shared_max_total_bytes", self._capacity.max_total_bytes, resulting_bytes + largest + "shared_max_total_bytes", self._capacity.max_total_bytes, resulting_bytes + reserve ) @staticmethod @@ -368,9 +372,13 @@ def _check_store_growth(self, store: FileSystemAgentFileStore, path: Path, size: async def _write_index(self, store: FileSystemAgentFileStore, view: _RetainedFileStore, folder: Path) -> None: path = folder / "memories.md" payload = (await view.index_text(folder.relative_to(store.root_path).as_posix())).encode("utf-8") - self._check_store_growth(store, path, len(payload)) - self._check_growth(path, len(payload)) - self._atomic(path, payload) + + def write_index() -> None: + self._check_store_growth(store, path, len(payload)) + self._check_growth(path, len(payload)) + self._atomic(path, payload) + + await asyncio.to_thread(write_index) def _atomic(self, path: Path, payload: bytes, *, overwrite: bool = True) -> None: temporary = self._control / "staging" / f"{uuid.uuid4().hex}.tmp" @@ -428,7 +436,7 @@ async def _run( async def work() -> _ResultT: async with self._locked(): self._safe(backing.root_path) - state = self._load(backing) + state = await asyncio.to_thread(self._load, backing) view = _RetainedFileStore(self, backing, state, started_at, purpose) result = await action(view) if view.changed: @@ -443,7 +451,7 @@ async def work() -> _ResultT: self._now() + self.retention_seconds if self.retention_seconds is not None else None ) try: - self._save(backing, state) + await asyncio.to_thread(self._save, backing, state) except (ValueError, OSError) as exc: if view.pending or view.deleted: raise _FileMemoryCommitError( @@ -461,33 +469,44 @@ async def _protect_store(self, store: AgentFileStore) -> AsyncGenerator[float]: backing = self._validate_store(store) token_lock: AsyncFileLock | None = None token_path: Path | None = None - try: + + async def register() -> float: + nonlocal token_lock, token_path async with self._locked(): - self._load(backing) + await asyncio.to_thread(self._load, backing) tokens = self._control / "active" / self._scope(backing) self._safe(self._control / "active") self._safe(tokens) - tokens.mkdir(parents=True, exist_ok=True) + await asyncio.to_thread(tokens.mkdir, parents=True, exist_ok=True) token_path = tokens / f"{uuid.uuid4().hex}.lock" - self._check_growth(token_path, 0) + await asyncio.to_thread(self._check_growth, token_path, 0) token_lock = AsyncFileLock(token_path, timeout=0, preserve_lock_file=True, fallback_to_soft=False) await token_lock.acquire() - started_at = self._now() + return self._now() + + try: + started_at = await _complete(register()) yield started_at # ruff: ignore[unnecessary-assign-before-yield] - capture visibility while registration holds the lock finally: if token_lock is not None: + held_lock = cast(AsyncFileLock, token_lock) async def release() -> None: try: async with self._locked(load_policy=False): - await token_lock.release() + await held_lock.release() if token_path is not None: - token_path.unlink(missing_ok=True) - if token_path.parent.exists() and not any(token_path.parent.iterdir()): - token_path.parent.rmdir() + released_path = token_path + + def remove_token() -> None: + released_path.unlink(missing_ok=True) + if released_path.parent.exists() and not any(released_path.parent.iterdir()): + released_path.parent.rmdir() + + await asyncio.to_thread(remove_token) finally: # A broken policy/control directory must not pin an OS lease. - await token_lock.release() + await held_lock.release() await _complete(release()) @@ -500,7 +519,7 @@ async def _tokens_active(self, scope: str) -> bool: if not tokens.exists(): return False active = False - for path in tokens.iterdir(): + for path in await asyncio.to_thread(lambda: list(tokens.iterdir())): self._safe(path) if ( not path.is_file() @@ -517,13 +536,17 @@ async def _tokens_active(self, scope: str) -> bool: active = True else: await probe.release() - path.unlink(missing_ok=True) + await asyncio.to_thread(path.unlink, missing_ok=True) if not active: - tokens.rmdir() + await asyncio.to_thread(tokens.rmdir) return active async def collect(self) -> int: - """Delete registered expired memories, retaining active or uncertain scopes.""" + """Delete registered expired memories and finish confirmed deletions. + + Disabling TTL stops new expiry decisions; already-confirmed deletions still finish. + Active and uncertain scopes remain protected. + """ async def work() -> int: async with self._locked(): @@ -537,7 +560,7 @@ async def _collect_locked(self) -> int: active = self._control / "active" self._safe(active) if active.exists(): - for tokens in active.iterdir(): + for tokens in await asyncio.to_thread(lambda: list(active.iterdir())): self._safe(tokens) if ( not tokens.is_dir() @@ -547,15 +570,17 @@ async def _collect_locked(self) -> int: raise ValueError("Unexpected request protection metadata.") if await self._tokens_active(tokens.name): active_scopes.add(tokens.name) - if self.retention_seconds is None: - return 0 - for root in list(self.directory.iterdir()): + for root in await asyncio.to_thread(lambda: list(self.directory.iterdir())): if root == self._control or not root.is_dir(): continue try: self._safe(root) backing = FileSystemAgentFileStore(root) - state = self._load(backing) + state = await asyncio.to_thread(self._load, backing) + if self.retention_seconds is None and not any( + record["state"] == "deleting" for record in state["entries"].values() + ): + continue limits = self._stored_limits(state) if limits is not None: backing = FileSystemAgentFileStore( @@ -567,7 +592,7 @@ async def _collect_locked(self) -> int: if self._scope(backing) in active_scopes: continue view = _RetainedFileStore(self, backing, state, None, "gc") - files = self._files(root) + files = await asyncio.to_thread(self._files, root) paths = {self._key(backing, p): p for p in files} folders = {self._key(backing, p.parent): p.parent for p in files} for key, record in list(state["entries"].items()): @@ -575,7 +600,9 @@ async def _collect_locked(self) -> int: if record["state"] == "updating": logger.warning("Retaining interrupted memory update; explicit repair is required.") continue - if record["state"] != "deleting" and (deadline is None or deadline > self._now()): + if record["state"] != "deleting" and ( + self.retention_seconds is None or deadline is None or deadline > self._now() + ): continue body, folder = paths.get(key), folders.get(record["folder"]) if body is not None and ( @@ -585,53 +612,70 @@ async def _collect_locked(self) -> int: ): raise ValueError("Inconsistent file-memory lifecycle references.") record["state"] = "deleting" - self._save(backing, state, reclaim=True) + await asyncio.to_thread(self._save, backing, state, reclaim=True) if body is not None: - body.unlink(missing_ok=True) + await asyncio.to_thread(body.unlink, missing_ok=True) description = paths.get(record["description"]) if description is not None: from ._file_memory import _description_file_name # pyright: ignore[reportPrivateUsage] - referenced = any( - p.exists() - and not view.internal(p.name) - and p.parent == description.parent - and os.path.normcase(_description_file_name(p.name)) == os.path.normcase(description.name) - for p in files - ) - if not referenced: - description.unlink(missing_ok=True) + def remove_description(path: Path, sources: list[Path]) -> None: + referenced = any( + p.exists() + and not _RetainedFileStore.internal(p.name) + and p.parent == path.parent + and os.path.normcase(_description_file_name(p.name)) == os.path.normcase(path.name) + for p in sources + ) + if not referenced: + path.unlink(missing_ok=True) + + await asyncio.to_thread(remove_description, description, files) del state["entries"][key] - if folder is not None and folder.exists(): + if folder is not None and await asyncio.to_thread(folder.exists): index = folder / "memories.md" - if any(not view.internal(p.name) for p in folder.iterdir() if p.is_file()): + + def has_bodies(path: Path) -> bool: + return any(not _RetainedFileStore.internal(p.name) for p in path.iterdir() if p.is_file()) + + if await asyncio.to_thread(has_bodies, folder): try: await self._write_index(backing, view, folder) except _FileStoreQuotaError: # The index is a projection; insufficient space must not # leave an already-deleted body permanently tombstoned. - index.unlink(missing_ok=True) + await asyncio.to_thread(index.unlink, missing_ok=True) logger.warning("Deferred memory index rebuild due to capacity; stale index removed.") else: - index.unlink(missing_ok=True) - if not any(folder.iterdir()): - folder.rmdir() - self._save(backing, state, reclaim=True) + + def remove_empty_index(path: Path) -> None: + (path / "memories.md").unlink(missing_ok=True) + if not any(path.iterdir()): + path.rmdir() + + await asyncio.to_thread(remove_empty_index, folder) + await asyncio.to_thread(self._save, backing, state, reclaim=True) removed += 1 - directories: list[Path] = [] - pending = [root] - while pending: - directory = pending.pop() - self._safe(directory) - directories.append(directory) - pending.extend(path for path in directory.iterdir() if path.is_dir()) - for directory in reversed(directories[1:]): - if not any(directory.iterdir()): - directory.rmdir() - if not state["entries"] and not self._files(root): - self._manifest_path(backing).unlink(missing_ok=True) - if not any(root.iterdir()): - root.rmdir() + + def prune_empty_directories(root_path: Path, manifest: Path, empty: bool) -> None: + directories: list[Path] = [] + pending = [root_path] + while pending: + directory = pending.pop() + self._safe(directory) + directories.append(directory) + pending.extend(path for path in directory.iterdir() if path.is_dir()) + for directory in reversed(directories[1:]): + if not any(directory.iterdir()): + directory.rmdir() + if empty and not self._files(root_path): + manifest.unlink(missing_ok=True) + if not any(root_path.iterdir()): + root_path.rmdir() + + await asyncio.to_thread( + prune_empty_directories, root, self._manifest_path(backing), not state["entries"] + ) except (ValueError, OSError) as exc: logger.warning("Skipped file-memory cleanup (%s).", type(exc).__name__) return removed @@ -656,8 +700,8 @@ async def work() -> int: async with self._locked(): if await self._active(backing): raise ValueError("Cannot repair a memory scope while a request is active.") - state = self._load(backing) - files = self._files(backing.root_path) + state = await asyncio.to_thread(self._load, backing) + files = await asyncio.to_thread(self._files, backing.root_path) paths = {self._key(backing, path): path for path in files} folders = {self._key(backing, path.parent): path.parent for path in files} view = _RetainedFileStore(self, backing, state, None, "repair") @@ -681,7 +725,7 @@ async def work() -> int: for folder in affected: await self._write_index(backing, view, folder) if repaired: - self._save(backing, state) + await asyncio.to_thread(self._save, backing, state) return repaired return await _complete(work()) @@ -758,7 +802,7 @@ def _touch(self, path: str) -> None: self.changed = True async def write(self, path: str, content: str, *, overwrite: bool = True) -> None: - target = self._path(path) + target = await asyncio.to_thread(self._path, path) try: limits = self.backing._limits # pyright: ignore[reportPrivateUsage] sizing = limits or _FileStoreLimits( @@ -766,56 +810,67 @@ async def write(self, path: str, content: str, *, overwrite: bool = True) -> Non None, self.manager._capacity.max_total_bytes if self.manager._capacity else None, # pyright: ignore[reportPrivateUsage] ) - if not overwrite and target.exists(): - raise FileExistsError("The memory file already exists.") - size = sizing.content_size(content) - - def check() -> None: - self.manager._check_store_growth(self.backing, target, size) # pyright: ignore[reportPrivateUsage] - if self.internal(target.name): - self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] - return - key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] - previous = self.state["entries"].get(key) - record: dict[str, Any] = { - **(previous or {}), - "state": "ready", - "expires_at": self.manager._now() + self.manager.retention_seconds # pyright: ignore[reportPrivateUsage] - if self.manager.retention_seconds is not None - else None, - "generation": self.manager._policy["generation"], # pyright: ignore[reportPrivateUsage] - "folder": self.manager._key(self.backing, target.parent), # pyright: ignore[reportPrivateUsage] - "description": self.manager._key(self.backing, target.with_name(self._description(target.name))), # pyright: ignore[reportPrivateUsage] - } - self.state["entries"][key] = record + size = await asyncio.to_thread(sizing.content_size, content) + + async def check() -> None: + candidate: dict[str, Any] | None = None config: dict[str, Any] = { **(self.state.get("store_limits") or {}), "max_file_bytes": limits.max_file_bytes if limits else None, "max_files": limits.max_files if limits else None, "max_total_bytes": limits.max_total_bytes if limits else None, } - self.state["store_limits"] = config - payload_size = len(json.dumps(self.state, sort_keys=True, separators=(",", ":")).encode("utf-8")) - # A finite float's JSON representation is shorter than 32 bytes; - # reserve its maximum field width before touching the body. - if record["expires_at"] is not None: - payload_size += max(0, 32 - len(json.dumps(record["expires_at"]))) - if previous is None: - self.state["entries"].pop(key) - else: - self.state["entries"][key] = previous - self.manager._check_combined_growth( # pyright: ignore[reportPrivateUsage] - [(target, size), (self.manager._manifest_path(self.backing), payload_size)] # pyright: ignore[reportPrivateUsage] - ) + deadline: float | None = None + if not self.internal(target.name): + deadline = ( + self.manager._now() + self.manager.retention_seconds # pyright: ignore[reportPrivateUsage] + if self.manager.retention_seconds is not None + else None + ) + key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] + previous = self.state["entries"].get(key) + description_path = target.with_name(self._description(target.name)) + record: dict[str, Any] = { + **(previous or {}), + "state": "ready", + "expires_at": deadline, + "generation": self.manager._policy["generation"], # pyright: ignore[reportPrivateUsage] + "folder": self.manager._key(self.backing, target.parent), # pyright: ignore[reportPrivateUsage] + "description": self.manager._key(self.backing, description_path), # pyright: ignore[reportPrivateUsage] + } + candidate = { + **self.state, + "entries": {**self.state["entries"], key: record}, + "store_limits": config, + } + + def check_files() -> None: + if not overwrite and target.exists(): + raise FileExistsError("The memory file already exists.") + self.manager._check_store_growth(self.backing, target, size) # pyright: ignore[reportPrivateUsage] + if candidate is None: + self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] + return + payload_size = len(json.dumps(candidate, sort_keys=True, separators=(",", ":")).encode("utf-8")) + # Reserve the maximum finite-float width before changing content. + if deadline is not None: + payload_size += max(0, 32 - len(json.dumps(deadline))) + self.manager._check_combined_growth( # pyright: ignore[reportPrivateUsage] + [(target, size), (self.manager._manifest_path(self.backing), payload_size)] # pyright: ignore[reportPrivateUsage] + ) + + await asyncio.to_thread(check_files) + if candidate is not None: + self.state["store_limits"] = config try: - check() + await check() except _FileStoreQuotaError: await self.manager._collect_locked() # pyright: ignore[reportPrivateUsage] - loaded = self.manager._load(self.backing) # pyright: ignore[reportPrivateUsage] + loaded = await asyncio.to_thread(self.manager._load, self.backing) # pyright: ignore[reportPrivateUsage] self.state.clear() self.state.update(loaded) - check() + await check() if not self.internal(target.name): key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] record = self.state["entries"].get(key) @@ -829,20 +884,21 @@ def check() -> None: "folder": self.manager._key(self.backing, target.parent), # pyright: ignore[reportPrivateUsage] "description": self.manager._key(self.backing, target.with_name(self._description(target.name))), # pyright: ignore[reportPrivateUsage] } - self.manager._save(self.backing, self.state) # pyright: ignore[reportPrivateUsage] + await asyncio.to_thread(self.manager._save, self.backing, self.state) # pyright: ignore[reportPrivateUsage] try: - self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] + await asyncio.to_thread(self.manager._check_growth, target, size) # pyright: ignore[reportPrivateUsage] except _FileStoreQuotaError: if record is None: self.state["entries"].pop(key, None) else: self.state["entries"][key] = record - self.manager._save(self.backing, self.state, reclaim=True) # pyright: ignore[reportPrivateUsage] + await asyncio.to_thread(self.manager._save, self.backing, self.state, reclaim=True) # pyright: ignore[reportPrivateUsage] raise self.pending.add(key) self.changed = True - self.manager._check_growth(target, size) # pyright: ignore[reportPrivateUsage] - await asyncio.to_thread(self.manager._atomic, target, content.encode("utf-8"), overwrite=overwrite) # pyright: ignore[reportPrivateUsage] + await asyncio.to_thread(self.manager._check_growth, target, size) # pyright: ignore[reportPrivateUsage] + payload = await asyncio.to_thread(content.encode, "utf-8") + await asyncio.to_thread(self.manager._atomic, target, payload, overwrite=overwrite) # pyright: ignore[reportPrivateUsage] except BaseException: self.failed = True raise @@ -885,22 +941,27 @@ async def read(self, path: str) -> str | None: return content async def delete(self, path: str) -> bool: - target = self._path(path) + target = await asyncio.to_thread(self._path, path) if not self.internal(target.name): key = self.manager._key(self.backing, target) # pyright: ignore[reportPrivateUsage] record = self.state["entries"].get(key) if record is not None: record["state"] = "deleting" - self.manager._save(self.backing, self.state, reclaim=True) # pyright: ignore[reportPrivateUsage] + await asyncio.to_thread(self.manager._save, self.backing, self.state, reclaim=True) # pyright: ignore[reportPrivateUsage] self.deleted.add(key) self.changed = True elif target.name.lower().endswith("_description.md") and target.parent.exists(): - if any( - p.is_file() - and not self.internal(p.name) - and os.path.normcase(self._description(p.name)) == os.path.normcase(target.name) - for p in target.parent.iterdir() - ): + + def referenced() -> bool: + return any( + p.is_file() + and not self.internal(p.name) + and self.manager._key(self.backing, p) not in self.pending # pyright: ignore[reportPrivateUsage] + and os.path.normcase(self._description(p.name)) == os.path.normcase(target.name) + for p in target.parent.iterdir() + ) + + if await asyncio.to_thread(referenced): return False try: return await asyncio.to_thread(self.backing._delete_file_sync, target) # pyright: ignore[reportPrivateUsage] diff --git a/python/packages/core/tests/core/test_harness_file_memory.py b/python/packages/core/tests/core/test_harness_file_memory.py index 5f7043103f3..f3debab2819 100644 --- a/python/packages/core/tests/core/test_harness_file_memory.py +++ b/python/packages/core/tests/core/test_harness_file_memory.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from pathlib import Path import pytest import regex @@ -12,6 +13,7 @@ AgentSession, Content, FileMemoryProvider, + FileSystemAgentFileStore, FunctionTool, InMemoryAgentFileStore, ) @@ -687,3 +689,18 @@ async def test_quota_rejection_preserves_existing_memory_content() -> None: result = await tools["file_memory_write"].invoke(arguments={"file_name": "plan.md", "content": "x" * 61}) assert "Could not write" in _text(result) and "max_total_bytes" in _text(result) assert await store.read("user-1/plan.md") == "old" + + +async def test_file_memory_list_reports_legacy_description_quota(tmp_path: Path) -> None: + """Listing refuses an oversized legacy description without discarding stored content.""" + unlimited = FileSystemAgentFileStore(tmp_path) + await unlimited.write("user-1/notes.md", "body") + await unlimited.write("user-1/notes_description.md", "x" * 128) + store = FileSystemAgentFileStore(tmp_path, max_file_bytes=64) + _, tools = await _prepare(FileMemoryProvider(store=store, scope="user-1")) + + result = _text(await tools["file_memory_ls"].invoke()) + + assert "Could not list memory files" in result and "max_file_bytes=64" in result + assert await unlimited.read("user-1/notes.md") == "body" + assert await unlimited.read("user-1/notes_description.md") == "x" * 128 diff --git a/python/packages/core/tests/core/test_harness_file_memory_retention.py b/python/packages/core/tests/core/test_harness_file_memory_retention.py index 113bf3bddad..92fed9898aa 100644 --- a/python/packages/core/tests/core/test_harness_file_memory_retention.py +++ b/python/packages/core/tests/core/test_harness_file_memory_retention.py @@ -6,11 +6,14 @@ import asyncio import json +import re import sys +import threading from pathlib import Path from typing import Any import pytest +from filelock import AsyncFileLock, Timeout from agent_framework import ( AgentSession, @@ -801,3 +804,161 @@ async def test_read_only_managed_requests_do_not_allocate_empty_data_directories assert await _call(tools, "file_memory_ls") == "[]" assert "not found" in await _call(tools, "file_memory_read", file_name="notes.md") assert not store.root_path.exists() + + +@pytest.mark.parametrize("other_body", ["absent", "managed", "legacy"]) +async def test_managed_overwrite_clears_omitted_description_unless_shared(tmp_path: Path, other_body: str) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="old", description="summary") + if other_body == "managed": + await _call(tools, "file_memory_write", file_name="notes.txt", content="other", description="summary") + elif other_body == "legacy": + await store.write("conversation/notes.txt", "other") + + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="new") + + expected = None if other_body == "absent" else "summary" + assert await store.read("conversation/notes_description.md") == expected + listed = json.loads(await _call(tools, "file_memory_ls")) + assert next(item for item in listed if item["name"] == "notes.md")["description"] == expected + index = await store.read("conversation/memories.md") or "" + assert ("summary" in index) == (other_body != "absent") + if other_body != "absent": + assert await store.read("conversation/notes.txt") == "other" + + +@pytest.mark.parametrize("body_present", [True, False]) +async def test_disabled_ttl_finishes_pending_deletion_and_preserves_ready(tmp_path: Path, body_present: bool) -> None: + manager = FileMemoryRetentionManager(tmp_path, retention_seconds=10) + now = [100.0] + manager._now = lambda: now[0] + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + await _call(tools, "file_memory_write", file_name="notes.md", content="delete", description="old summary") + now[0] = 105 + await _call(tools, "file_memory_write", file_name="keep.md", content="preserve", description="keep summary") + process = await _worker(tmp_path, "crash_gc") + _, error = await asyncio.wait_for(process.communicate(), 15) + assert process.returncode == 17, error.decode() + if not body_present: + (store.root_path / "conversation/notes.md").unlink() + + await manager.set_retention(None) + now[0] = 1000.0 + + assert await manager.collect() == 1 + assert await manager.collect() == 0 + assert await store.read("conversation/notes.md") is None + assert await store.read("conversation/notes_description.md") is None + assert await _call(tools, "file_memory_read", file_name="keep.md") == "preserve" + assert await store.read("conversation/keep_description.md") == "keep summary" + assert "notes.md" not in (await store.read("conversation/memories.md") or "") + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="replacement") + + +async def test_shared_quota_message_accounts_for_new_atomic_copy(tmp_path: Path) -> None: + manager = FileMemoryRetentionManager(tmp_path, max_total_bytes=4000) + store = FileSystemAgentFileStore(tmp_path / "owner") + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="old") + + message = await _call(tools, "file_memory_write", file_name="notes.md", content="x" * 3000) + + assert "shared_max_total_bytes=4000" in message + requested = re.search(r"requested=(\d+)", message) + assert requested is not None + # The replacement and its complete atomic copy alone require at least 6000 bytes. + assert int(requested.group(1)) >= 6000 + assert await store.read("conversation/notes.md") == "old" + + +@pytest.mark.parametrize("accounting", ["shared", "root"]) +async def test_capacity_accounting_allows_other_coroutines_to_progress( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, accounting: str +) -> None: + manager = FileMemoryRetentionManager(tmp_path, max_total_bytes=50000 if accounting == "shared" else None) + store = FileSystemAgentFileStore(tmp_path / "owner", max_total_bytes=30000 if accounting == "root" else None) + _, tools = await _prepare(FileMemoryProvider(store, retention=manager)) + started, release, timed_out = threading.Event(), threading.Event(), threading.Event() + target, name = (manager, "_usage") if accounting == "shared" else (store, "_quota_usage") + original = getattr(target, name) + first = True + + def slow_usage() -> Any: + nonlocal first + if first: + first = False + started.set() + if not release.wait(2): + timed_out.set() + return original() + + monkeypatch.setattr(target, name, slow_usage) + + async def foreground() -> None: + assert await asyncio.to_thread(started.wait, 5) + release.set() + + heartbeat = asyncio.create_task(foreground()) + try: + assert "written" in await _call(tools, "file_memory_write", file_name="notes.md", content="hello") + finally: + release.set() + await heartbeat + + assert not timed_out.is_set(), "Storage accounting blocked the foreground coroutine" + assert await store.read("conversation/notes.md") == "hello" + + +@pytest.mark.parametrize("operation", ["write", "registration"]) +async def test_cancelled_accounting_keeps_coordination_until_worker_finishes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str +) -> None: + manager = FileMemoryRetentionManager(tmp_path, max_total_bytes=50000) + store = FileSystemAgentFileStore(tmp_path / "owner") + provider = FileMemoryProvider(store, retention=manager) + _, tools = await _prepare(provider) + started, release = threading.Event(), threading.Event() + original = manager._usage + first = True + + def slow_usage() -> tuple[int, int, int]: + nonlocal first + if first: + first = False + started.set() + if not release.wait(5): + raise AssertionError("The test did not release the accounting worker") + return original() + + monkeypatch.setattr(manager, "_usage", slow_usage) + protection = provider.protect() + task: asyncio.Task[Any] + if operation == "write": + task = asyncio.create_task(_call(tools, "file_memory_write", file_name="notes.md", content="complete")) + else: + task = asyncio.create_task(protection.__aenter__()) + probe = AsyncFileLock( + tmp_path / ".retention/manager.lock", timeout=0, preserve_lock_file=True, fallback_to_soft=False + ) + try: + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + with pytest.raises(Timeout): + await probe.acquire() + finally: + await probe.release() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + await probe.acquire() + await probe.release() + if operation == "write": + assert await store.read("conversation/notes.md") == "complete" + else: + assert not list((tmp_path / ".retention/active").glob("*/*.lock")) diff --git a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md index 37e296f3d2a..3cc5da4d5ed 100644 --- a/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md +++ b/python/samples/02-agents/harness/build_your_own_claw/claw_step04_production_ready/README.md @@ -194,8 +194,9 @@ await retention.set_retention(None) # disable expiration; retain metadat await retention.set_retention(30 * 86400) # explicitly resume with a complete new TTL ``` -While disabled, reads remain available without renewal and TTL GC deletes no memories; -cleanup of dead request-protection markers continues. +While disabled, reads remain available without renewal and GC starts no new expiry +deletions. Already-confirmed deletions continue to completion, including after an +interruption; cleanup of dead request-protection markers also continues. Resuming grants still-present, registered `ready` records a complete lifetime from resume. Repeated same-value calls do nothing. Changing one positive TTL to another affects new writes and subsequent successful uses, not all existing deadlines. Deleted records cannot