Python: bound Monty file capture

python/packages/monty/README.md

@@ -195,12 +195,22 @@ codeact = MontyCodeActProvider(
   after each `execute_code` call and returned as `Content.from_data(...)`
   attachments (with a `path` annotation in `additional_properties`),
   mirroring Hyperlight's `/output` flow.
+- Capture is limited to **20 files and 20 MiB of raw file data per call across
+  all read-write mounts**, with a **5 MiB per-file limit**. Empty files count
+  toward the file limit. Files that do not fit the remaining byte budget are
+  skipped; later smaller files can still be captured. Exact fits are accepted.
+  Returned base64 data and Python objects add memory overhead beyond the raw
+  byte budget.
+- Files omitted by capture limits produce one text summary, not one warning
+  per file. Omission does not undo filesystem writes. These limits bound
+  attachment capture, not filesystem traversal, disk use, or all host memory.
 - `overlay` mounts buffer writes in memory (nothing leaks to the host and
   nothing is captured). `read-only` mounts reject writes.
 - **`resource_limits`** is forwarded straight to Monty's
   [`ResourceLimits`](https://github.com/pydantic/monty) TypedDict
   (`max_duration_secs`, `max_memory`, `gc_interval`, `max_recursion_depth`,
-  `max_suspensions`).
+  `max_suspensions`). These interpreter limits do not independently bound
+  host-side post-execution attachment capture.
 
 ## DSL inside `execute_code`
 

python/packages/monty/agent_framework_monty/_execute_code_tool.py

@@ -6,8 +6,8 @@
 applies to a pure-Python interpreter (no backends to choose from). By default
 the Monty sandbox rejects OS / filesystem / network calls with
 ``PermissionError``; pass ``workspace_root`` or ``file_mounts`` to expose
-scoped host directories, and the tool will capture any files written under
-``read-write`` mounts as ``Content`` items in the response.
+scoped host directories, and the tool will capture files written under
+``read-write`` mounts as ``Content`` items within per-execution capture limits.
 """
 
 from __future__ import annotations
@@ -56,10 +56,18 @@
 #: matching the Hyperlight default. Use ``file_mounts`` for any other path.
 WORKSPACE_MOUNT_PATH = "/input"
 
-#: Maximum bytes per captured output file. Files larger than this are skipped
-#: and a ``Content.from_text`` warning is appended in their place.
+#: Maximum bytes per captured output file. Oversized files are skipped and
+#: reported in one omission summary per execution.
 MAX_CAPTURED_FILE_BYTES = 5 * 1024 * 1024  # 5 MiB
 
+#: Maximum captured files across all read-write mounts in one execution.
+MAX_CAPTURED_FILES = 20
+
+#: Maximum raw captured bytes across all read-write mounts in one execution.
+MAX_CAPTURED_TOTAL_BYTES = 20 * 1024 * 1024  # 20 MiB
+
+_CAPTURE_OMISSION_MESSAGE = "[files omitted: capture file-count or byte limits exceeded]"
+
 EXECUTE_CODE_INPUT_SCHEMA: dict[str, Any] = {
     "type": "object",
     "title": "_ExecuteCodeInput",
@@ -552,14 +560,25 @@ def _safe_capture_file_stat(*, root: Path, entry: Path) -> os.stat_result | None
     return entry_stat if stat.S_ISREG(entry_stat.st_mode) else None
 
 
-def _read_validated_file(*, root: Path, entry: Path, expected: os.stat_result) -> bytes | None:
+class _CaptureSizeLimitError(ValueError):
+    """A file exceeded its remaining capture byte allowance."""
+
+
+def _read_validated_file(
+    *,
+    root: Path,
+    entry: Path,
+    expected: os.stat_result,
+    max_bytes: int,
+) -> bytes | None:
     """Read ``entry`` through one descriptor that is checked against the validated file.
 
     ``Path.read_bytes()`` reopens the path, so a component swapped for a link or junction after
     validation would be followed. Opening first, comparing the descriptor with the validated ``lstat``
     identity, and revalidating the path afterwards rejects such a swap. Windows has no no-follow open,
     so the check is not atomic there, but an undetected swap would have to be applied, reverted, and
-    applied again around the checks. Returns ``None`` when the file changed or cannot be read safely.
+    applied again around the checks. Returns ``None`` when the file identity changed or cannot be read
+    safely. Raises ``_CaptureSizeLimitError`` if the file exceeds ``max_bytes``.
     """
     flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0)
     try:
@@ -574,8 +593,13 @@ def _read_validated_file(*, root: Path, entry: Path, expected: os.stat_result) -
         rechecked = _safe_capture_file_stat(root=root, entry=entry)
         if rechecked is None or (rechecked.st_dev, rechecked.st_ino) != identity:
             return None
+        if opened.st_size > max_bytes:
+            raise _CaptureSizeLimitError("File exceeds the capture byte limit.")
         with os.fdopen(fd, "rb", closefd=False) as handle:
-            return handle.read()
+            data = handle.read(max_bytes + 1)
+        if len(data) > max_bytes:
+            raise _CaptureSizeLimitError("File grew beyond the capture byte limit.")
+        return data
     except OSError:
         return None
     finally:
@@ -614,15 +638,16 @@ def _capture_written_files(
 ) -> list[Content]:
     """Return :class:`Content` items for files the sandbox wrote during the run.
 
-    Mirrors Hyperlight's ``/output`` capture flow: any new or modified real
-    regular file under a read-write mount is read back as binary and
-    surfaced as ``Content.from_data`` with a ``path`` annotation in
-    ``additional_properties``. Links and reparse points are skipped at every
-    directory level, and each file is read through a descriptor verified
-    against the validated path, so a malicious workspace cannot trick us into
-    capturing host files outside the configured mount root.
+    New or modified real regular files under read-write mounts are surfaced as
+    ``Content.from_data`` with a ``path`` annotation in ``additional_properties``.
+    File-count and raw-byte budgets apply across all mounts, with one summary
+    for files omitted by capture limits. Links and reparse points are skipped at
+    every directory level, and bounded reads use descriptors verified against
+    the validated paths.
     """
     captured: list[Content] = []
+    total_bytes = 0
+    omitted = False
     for mount in mounts:
         if mount.mode != "read-write":
             continue
@@ -636,16 +661,19 @@ def _capture_written_files(
             current = (int(entry_stat.st_size), int(entry_stat.st_mtime_ns))
             if before.get(relative) == current:
                 continue  # Unchanged.
+            if len(captured) >= MAX_CAPTURED_FILES:
+                captured.append(Content.from_text(_CAPTURE_OMISSION_MESSAGE))
+                return captured
             sandbox_path = f"{mount.mount_path.rstrip('/')}/{relative}"
-            if entry_stat.st_size > MAX_CAPTURED_FILE_BYTES:
-                captured.append(
-                    Content.from_text(
-                        f"[file {sandbox_path} omitted: {entry_stat.st_size} bytes "
-                        f"exceeds MAX_CAPTURED_FILE_BYTES={MAX_CAPTURED_FILE_BYTES}]"
-                    )
-                )
+            max_bytes = min(MAX_CAPTURED_FILE_BYTES, MAX_CAPTURED_TOTAL_BYTES - total_bytes)
+            if entry_stat.st_size > max_bytes:
+                omitted = True
+                continue
+            try:
+                data = _read_validated_file(root=host_root, entry=entry, expected=entry_stat, max_bytes=max_bytes)
+            except _CaptureSizeLimitError:
+                omitted = True
                 continue
-            data = _read_validated_file(root=host_root, entry=entry, expected=entry_stat)
             if data is None:
                 continue
             media_type = mimetypes.guess_type(entry.name)[0] or "application/octet-stream"
@@ -656,4 +684,7 @@ def _capture_written_files(
                     additional_properties={"path": sandbox_path},
                 )
             )
+            total_bytes += len(data)
+    if omitted:
+        captured.append(Content.from_text(_CAPTURE_OMISSION_MESSAGE))
     return captured

python/packages/monty/agent_framework_monty/_instructions.py

@@ -76,7 +76,8 @@ def _format_filesystem_capabilities(mounts: Sequence[FileMount]) -> str:
         writable_paths = ", ".join(f"`{m.mount_path}`" for m in writable)
         lines.append(
             f"Files written to {writable_paths} are returned to the caller as attached files; "
-            "use these paths for any output artifacts."
+            "use these paths for any output artifacts. Capture file-count and byte limits apply across "
+            "all mounts per execution; omitted files are reported in one summary."
         )
 
     return "\n".join(lines)