diff --git a/python/packages/hyperlight/README.md b/python/packages/hyperlight/README.md index 08b188ddb87..02332e100f4 100644 --- a/python/packages/hyperlight/README.md +++ b/python/packages/hyperlight/README.md @@ -118,6 +118,35 @@ codeact = HyperlightCodeActProvider( ) ``` +### Output attachment limits + +Files written under `/output` are returned as inline data attachments. Hyperlight +limits each invocation to 20 files, 5 MiB per file, and 20 MiB of cumulative raw +file data by default. Oversized output is returned as a structured execution error +without partial data attachments. +Output discovery also has finite internal entry and nesting-depth safeguards; +directory-heavy output that exceeds them is rejected as an execution error. + +Trusted applications can raise these limits with positive integers on either +`HyperlightExecuteCodeTool` or `HyperlightCodeActProvider`: + +```python +codeact = HyperlightCodeActProvider( + workspace_root="./workspace", + max_output_files=40, + max_output_file_bytes=10 * 1024 * 1024, + max_output_total_bytes=50 * 1024 * 1024, +) +``` + +Limits are always finite. Increasing them also increases host memory use because +file data is encoded as inline base64, and may increase model context cost when +attachments are included in subsequent requests. + +Nested output paths require secure directory-relative file opening. On platforms +without that capability, nested attachments fail closed; write attachment files +directly under `/output` for portable behavior. + ## Notes - This package is intentionally separate from `agent-framework-core` so CodeAct diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py b/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py index 022227e9a0c..3dfb20f0afe 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py @@ -31,6 +31,13 @@ EXECUTE_CODE_TOOL_DESCRIPTION = "Execute Python in an isolated Hyperlight sandbox." OUTPUT_FILE_RETRY_ATTEMPTS = 10 OUTPUT_FILE_RETRY_DELAY_SECONDS = 0.1 +DEFAULT_MAX_OUTPUT_FILES = 20 +DEFAULT_MAX_OUTPUT_FILE_BYTES = 5 * 1024 * 1024 +DEFAULT_MAX_OUTPUT_TOTAL_BYTES = 20 * 1024 * 1024 +OUTPUT_TRAVERSAL_MIN_ENTRIES = 100 +OUTPUT_TRAVERSAL_ENTRIES_PER_FILE = 10 +OUTPUT_TRAVERSAL_MAX_ENTRIES = 10_000 +OUTPUT_TRAVERSAL_MAX_DEPTH = 32 EXECUTE_CODE_INPUT_SCHEMA: dict[str, Any] = { "type": "object", @@ -46,6 +53,17 @@ } +class _OutputMaterializationError(RuntimeError): + """Raised when sandbox output cannot be safely materialized.""" + + +@dataclass(frozen=True, slots=True) +class _ValidatedOutputFile: + relative_path: str + media_type: str + data: bytes + + @dataclass(frozen=True, slots=True) class _NormalizedFileMount: host_path: Path @@ -64,6 +82,9 @@ class _RunConfig: workspace_signature: tuple[tuple[str, int, int], ...] file_mounts: tuple[_NormalizedFileMount, ...] allowed_domains: tuple[AllowedDomain, ...] + max_output_files: int = DEFAULT_MAX_OUTPUT_FILES + max_output_file_bytes: int = DEFAULT_MAX_OUTPUT_FILE_BYTES + max_output_total_bytes: int = DEFAULT_MAX_OUTPUT_TOTAL_BYTES @property def mounted_paths(self) -> tuple[str, ...]: @@ -74,6 +95,8 @@ def filesystem_enabled(self) -> bool: return self.workspace_root is not None or bool(self.file_mounts) def cache_key(self) -> tuple[Any, ...]: + # Output limits are invocation-scoped and do not change sandbox construction, + # so they intentionally do not participate in the shared runtime cache key. return ( self.backend, self.module, @@ -200,6 +223,9 @@ def execute( code: str, output_dir: TemporaryDirectory[str] | None, build_contents: Callable[..., list[Content]], + max_output_files: int, + max_output_file_bytes: int, + max_output_total_bytes: int, ) -> list[Content]: """Restore + run + build sendable contents — all on the worker thread. @@ -216,9 +242,11 @@ def _on_worker() -> list[Content]: try: return build_contents( result=result, - sandbox=sandbox, output_dir=output_dir, code=code, + max_output_files=max_output_files, + max_output_file_bytes=max_output_file_bytes, + max_output_total_bytes=max_output_total_bytes, ) finally: # ``result`` may carry a back-reference to the sandbox. Force its @@ -346,6 +374,14 @@ def _resolve_execute_code_approval_mode( return "never_require" +def _validate_positive_integer(*, name: str, value: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be a positive integer.") + if value <= 0: + raise ValueError(f"{name} must be a positive integer.") + return value + + def _resolve_existing_path(value: str | Path) -> Path: return Path(value).expanduser().resolve(strict=True) @@ -557,7 +593,7 @@ def _display_mount_path(mount_path: str) -> str: return f"/input/{mount_path}" -def _iter_real_entries(root: Path, *, reject_links: bool = False) -> Iterator[Path]: +def _iter_real_entries(root: Path, *, reject_links: bool = False, files_only: bool = False) -> Iterator[Path]: """Walk ``root`` recursively, yielding directories and regular files only. ``Path.rglob`` follows directory links by default, which combined with @@ -570,36 +606,62 @@ def _iter_real_entries(root: Path, *, reject_links: bool = False) -> Iterator[Pa Non-regular files (sockets, FIFOs, devices) are also filtered out so the signature mirrors exactly what ``_copy_path`` actually stages. """ - stack: list[Path] = [root] - while stack: - current = stack.pop() + scan_stack: list[tuple[Path, Any]] = [] + try: try: - children = list(current.iterdir()) + scan_stack.append((root, os.scandir(root))) except OSError as exc: if reject_links: - raise ValueError(f"Could not inspect Hyperlight sandbox input directory: {current}") from exc - continue - for child in children: + raise ValueError(f"Could not inspect Hyperlight sandbox input directory: {root}") from exc + return + + while scan_stack: + current, entries = scan_stack[-1] + try: + entry = next(entries) + except StopIteration: + entries.close() + scan_stack.pop() + continue + except OSError as exc: + entries.close() + scan_stack.pop() + if reject_links: + raise ValueError(f"Could not inspect Hyperlight sandbox input directory: {current}") from exc + continue + + child = Path(entry.path) try: child_stat = child.lstat() - if _is_link_or_reparse_point(child, child_stat): - if reject_links: - raise ValueError( - f"Refusing to stage linked or reparse-point path for Hyperlight sandbox input: {child}" - ) - continue - if stat.S_ISDIR(child_stat.st_mode): - stack.append(child) - yield child - elif stat.S_ISREG(child_stat.st_mode): - yield child - # Non-regular files (sockets/FIFOs/devices) are skipped to - # match ``_copy_path``'s staging behaviour. except OSError as exc: if reject_links: raise ValueError(f"Could not inspect Hyperlight sandbox input path: {child}") from exc continue + if _is_link_or_reparse_point(child, child_stat): + if reject_links: + raise ValueError( + f"Refusing to stage linked or reparse-point path for Hyperlight sandbox input: {child}" + ) + continue + + if stat.S_ISDIR(child_stat.st_mode): + if not files_only: + yield child + try: + scan_stack.append((child, os.scandir(child))) + except OSError as exc: + if reject_links: + raise ValueError(f"Could not inspect Hyperlight sandbox input directory: {child}") from exc + continue + elif stat.S_ISREG(child_stat.st_mode): + yield child + # Non-regular files (sockets/FIFOs/devices) are skipped to + # match ``_copy_path``'s staging behaviour. + finally: + for _, entries in scan_stack: + entries.close() + def _path_tree_signature(path: Path) -> tuple[tuple[str, int, int], ...]: """Return a stable signature of the real (non-symlink) file tree under ``path``. @@ -682,64 +744,166 @@ def _populate_input_dir(*, config: _RunConfig, input_root: Path) -> None: _copy_path(mount.host_path, input_root / mount.mount_path, source_root=mount_root) -def _read_output_file_bytes(file_path: Path) -> bytes: - """Read ``file_path`` without following a link, even under a TOCTOU swap. +def _supports_secure_output_dir_fd() -> bool: + return ( + os.open in os.supports_dir_fd + and os.stat in os.supports_dir_fd + and os.stat in os.supports_follow_symlinks + and hasattr(os, "O_DIRECTORY") + and hasattr(os, "O_NOFOLLOW") + ) - ``Path.read_bytes`` follows links, so a sandbox payload that replaces an - output file with ``/output/leak.txt -> /host/secret`` or a Windows reparse - point between validation and read could still exfiltrate a host file. Two - layers defend against this: - * ``os.O_NOFOLLOW`` makes the kernel reject a final-component symlink with - ``ELOOP``. The flag is absent on some platforms (notably Windows), where - it degrades to ``0``, so it cannot be the only defense. - * The file is ``lstat``-ed before opening and ``fstat``-ed after; if the - ``(st_dev, st_ino)`` identity changed, or the pre-open entry is a link or - reparse point, the read is refused. This closes the swap window on every - platform. - """ +def _open_direct_output_file(*, root: Path, file_name: str) -> tuple[int, os.stat_result]: + file_path = root / file_name pre_stat = file_path.lstat() - if _is_link_or_reparse_point(file_path, pre_stat): - raise OSError(f"refusing to read linked or reparse-point output file: {file_path}") + if _is_link_or_reparse_point(file_path, pre_stat) or not stat.S_ISREG(pre_stat.st_mode): + raise OSError(f"refusing to read linked, reparse-point, or non-regular output file: {file_path}") fd = os.open(file_path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + keep_fd = False try: opened_stat = os.fstat(fd) if (opened_stat.st_dev, opened_stat.st_ino) != (pre_stat.st_dev, pre_stat.st_ino): raise OSError(f"output file changed between validation and read: {file_path}") - except BaseException: - os.close(fd) - raise + if not stat.S_ISREG(opened_stat.st_mode): + raise OSError(f"refusing to read non-regular output file: {file_path}") + keep_fd = True + return fd, opened_stat + finally: + if not keep_fd: + os.close(fd) + + +def _open_output_file_with_dir_fd(*, root: Path, path_parts: tuple[str, ...]) -> tuple[int, os.stat_result]: + directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + file_flags = os.O_RDONLY | os.O_NOFOLLOW + directory_fds: list[int] = [] + file_fd = -1 + keep_file_fd = False - with os.fdopen(fd, "rb") as handle: - return handle.read() + try: + root_fd = os.open(root, directory_flags) + directory_fds.append(root_fd) + parent_fd = root_fd + + for part in path_parts[:-1]: + pre_stat = os.stat(part, dir_fd=parent_fd, follow_symlinks=False) + if stat.S_ISLNK(pre_stat.st_mode) or not stat.S_ISDIR(pre_stat.st_mode): + raise OSError(f"refusing to traverse linked or non-directory output component: {part}") + + next_fd = os.open(part, directory_flags, dir_fd=parent_fd) + directory_fds.append(next_fd) + opened_stat = os.fstat(next_fd) + if (opened_stat.st_dev, opened_stat.st_ino) != (pre_stat.st_dev, pre_stat.st_ino): + raise OSError(f"output directory changed while opening component: {part}") + if not stat.S_ISDIR(opened_stat.st_mode): + raise OSError(f"refusing to traverse non-directory output component: {part}") + + parent_fd = next_fd + + file_name = path_parts[-1] + pre_stat = os.stat(file_name, dir_fd=parent_fd, follow_symlinks=False) + if stat.S_ISLNK(pre_stat.st_mode) or not stat.S_ISREG(pre_stat.st_mode): + raise OSError(f"refusing to read linked or non-regular output file: {file_name}") + + file_fd = os.open(file_name, file_flags, dir_fd=parent_fd) + opened_stat = os.fstat(file_fd) + if (opened_stat.st_dev, opened_stat.st_ino) != (pre_stat.st_dev, pre_stat.st_ino): + raise OSError(f"output file changed while opening: {file_name}") + if not stat.S_ISREG(opened_stat.st_mode): + raise OSError(f"refusing to read non-regular output file: {file_name}") + + keep_file_fd = True + return file_fd, opened_stat + finally: + if file_fd >= 0 and not keep_file_fd: + os.close(file_fd) + for directory_fd in reversed(directory_fds): + os.close(directory_fd) + + +def _open_output_file(*, root: Path, relative_path: str) -> tuple[int, os.stat_result]: + path_parts = tuple(PurePosixPath(relative_path).parts) + if not path_parts: + raise OSError(f"invalid output path: {relative_path}") + if any(part in {"", ".", ".."} for part in path_parts): + raise OSError(f"invalid output path: {relative_path}") + + if _supports_secure_output_dir_fd(): + return _open_output_file_with_dir_fd(root=root, path_parts=path_parts) + + if len(path_parts) > 1: + raise _OutputMaterializationError( + "Nested output attachments cannot be opened safely on this platform; write files directly under /output." + ) + return _open_direct_output_file(root=root, file_name=next(iter(path_parts))) -def _create_file_content(file_path: Path, *, relative_path: str) -> Content: - media_type = mimetypes.guess_type(file_path.name)[0] or "application/octet-stream" - return Content.from_data( - data=_read_output_file_bytes(file_path), - media_type=media_type, - additional_properties={"path": f"/output/{relative_path}"}, - ) +def _read_output_file_bytes( + root: Path, + *, + relative_path: str, + max_file_bytes: int, + remaining_total_bytes: int, + max_total_bytes: int, +) -> bytes: + """Open and read an output file relative to a pinned output root. + + Platforms with ``dir_fd`` support walk every component relative to verified directory + descriptors. Other platforms accept only direct children of the trusted output root, + where lstat/open/fstat identity checks fail closed on final-component replacements. + """ + output_path = f"/output/{relative_path}" + fd, opened_stat = _open_output_file(root=root, relative_path=relative_path) + try: + if opened_stat.st_size > max_file_bytes: + raise _OutputMaterializationError( + f"Output file {output_path[:200]!r} exceeds the {max_file_bytes}-byte per-file output limit." + ) + if opened_stat.st_size > remaining_total_bytes: + raise _OutputMaterializationError( + f"Output files exceed the {max_total_bytes}-byte cumulative output limit." + ) + with os.fdopen(fd, "rb", closefd=False) as handle: + read_allowance = min(opened_stat.st_size, max_file_bytes, remaining_total_bytes) + try: + data = handle.read(read_allowance + 1) + except (MemoryError, OverflowError): + raise _OutputMaterializationError( + "Sandbox output could not be read within the configured byte limits." + ) from None + finally: + os.close(fd) -def _normalize_output_relative_path(*, output_file: object, root: Path) -> str | None: - candidate_path = Path(str(output_file)) - if candidate_path.is_absolute(): - try: - return candidate_path.relative_to(root).as_posix() - except ValueError: - return None + if len(data) > max_file_bytes: + raise _OutputMaterializationError( + f"Output file {output_path[:200]!r} exceeds the {max_file_bytes}-byte per-file output limit." + ) + if len(data) > remaining_total_bytes: + raise _OutputMaterializationError(f"Output files exceed the {max_total_bytes}-byte cumulative output limit.") + if len(data) > opened_stat.st_size: + raise _OutputMaterializationError(f"Output file {output_path[:200]!r} grew while it was being read.") + return data - raw_path = str(output_file).replace("\\", "/") - pure_path = PurePosixPath(raw_path) - parts = [part for part in pure_path.parts if part not in {"", "/", "."}] - if parts and parts[0] == "output": - parts = parts[1:] - if not parts or any(part == ".." for part in parts): - return None - return "/".join(parts) + +def _materialize_output_files(output_files: Sequence[_ValidatedOutputFile]) -> list[Content]: + contents: list[Content] = [] + try: + for output_file in output_files: + contents.append( + Content.from_data( + data=output_file.data, + media_type=output_file.media_type, + additional_properties={"path": f"/output/{output_file.relative_path}"}, + ) + ) + except MemoryError: + raise _OutputMaterializationError( + "Sandbox output could not be encoded because the host did not have enough memory." + ) from None + return contents def _is_safe_output_file(*, root: Path, host_path: Path) -> bool: @@ -790,35 +954,100 @@ def _is_safe_output_file(*, root: Path, host_path: Path) -> bool: return stat.S_ISREG(final_stat.st_mode) -def _collect_output_relative_paths(*, sandbox: Any, root: Path) -> set[str]: - relative_paths: set[str] = set() +def _output_traversal_entry_limit(max_output_files: int) -> int: + return min( + OUTPUT_TRAVERSAL_MAX_ENTRIES, + max(OUTPUT_TRAVERSAL_MIN_ENTRIES, max_output_files * OUTPUT_TRAVERSAL_ENTRIES_PER_FILE), + ) - if hasattr(sandbox, "get_output_files"): + +def _iter_bounded_output_files( + root: Path, + *, + max_entries: int, + max_depth: int, +) -> Iterator[Path]: + """Yield regular output files while bounding traversal work and open scanners.""" + scan_stack: list[tuple[Any, int]] = [] + entries_visited = 0 + + def _open_scanner(path: Path) -> Any: try: - output_files = cast(Sequence[object], sandbox.get_output_files()) - except Exception: - output_files = () + return os.scandir(path) + except OSError as exc: + raise _OutputMaterializationError("Could not enumerate output directory safely.") from exc - for output_file in output_files: - if (relative_path := _normalize_output_relative_path(output_file=output_file, root=root)) is not None: - relative_paths.add(relative_path) + try: + scan_stack.append((_open_scanner(root), 0)) + while scan_stack: + entries, depth = scan_stack[-1] + try: + entry = next(entries) + except StopIteration: + entries.close() + scan_stack.pop() + continue + except OSError as exc: + raise _OutputMaterializationError("Could not enumerate output directory safely.") from exc - # ``Path.rglob`` follows directory symlinks and ``Path.is_file`` follows - # symlinks, both of which would surface paths outside the sandbox-controlled - # output tree. ``_iter_real_entries`` skips symlinks and never descends - # through a symlinked directory, yielding only real entries under ``root``. - for host_path in _iter_real_entries(root): - if host_path.is_file(): - relative_paths.add(host_path.relative_to(root).as_posix()) + entries_visited += 1 + if entries_visited > max_entries: + raise _OutputMaterializationError( + f"Sandbox output exceeded the traversal entry limit of {max_entries}." + ) + + child = Path(entry.path) + try: + child_stat = child.lstat() + except OSError as exc: + raise _OutputMaterializationError("Could not inspect output entry safely.") from exc + + if _is_link_or_reparse_point(child, child_stat): + continue + if stat.S_ISREG(child_stat.st_mode): + yield child + continue + if not stat.S_ISDIR(child_stat.st_mode): + continue + + child_depth = depth + 1 + if child_depth > max_depth: + raise _OutputMaterializationError(f"Sandbox output exceeded the nesting depth limit of {max_depth}.") + scan_stack.append((_open_scanner(child), child_depth)) + finally: + for entries, _ in scan_stack: + with suppress(OSError): + entries.close() + + +def _collect_output_relative_paths( + *, + root: Path, + max_output_files: int = DEFAULT_MAX_OUTPUT_FILES, +) -> set[str]: + relative_paths: set[str] = set() + + for host_path in _iter_bounded_output_files( + root, + max_entries=_output_traversal_entry_limit(max_output_files), + max_depth=OUTPUT_TRAVERSAL_MAX_DEPTH, + ): + if len(relative_paths) >= max_output_files: + raise _OutputMaterializationError( + f"Sandbox exceeded the output file count limit of {max_output_files} while enumerating candidates." + ) + relative_paths.add(host_path.relative_to(root).as_posix()) return relative_paths def _parse_output_files( *, - sandbox: Any, output_dir: _NamedDirectory | None, expect_output_files: bool, + max_output_files: int = DEFAULT_MAX_OUTPUT_FILES, + max_output_file_bytes: int = DEFAULT_MAX_OUTPUT_FILE_BYTES, + max_output_total_bytes: int = DEFAULT_MAX_OUTPUT_TOTAL_BYTES, ) -> list[Content]: if output_dir is None: return [] @@ -826,22 +1055,61 @@ def _parse_output_files( root = Path(output_dir.name) for attempt in range(OUTPUT_FILE_RETRY_ATTEMPTS): - relative_paths = _collect_output_relative_paths(sandbox=sandbox, root=root) + try: + relative_paths = _collect_output_relative_paths( + root=root, + max_output_files=max_output_files, + ) + except MemoryError: + raise _OutputMaterializationError( + "Sandbox output could not be enumerated because the host did not have enough memory." + ) from None missing_files = expect_output_files and not relative_paths - contents: list[Content] = [] + safe_output_files: list[tuple[str, Path]] = [] for relative_path in sorted(relative_paths): host_path = root.joinpath(*PurePosixPath(relative_path).parts) if not _is_safe_output_file(root=root, host_path=host_path): missing_files = True continue + safe_output_files.append((relative_path, host_path)) + + if len(safe_output_files) > max_output_files: + raise _OutputMaterializationError( + f"Sandbox produced {len(safe_output_files)} safe output files, exceeding the " + f"output file count limit of {max_output_files}." + ) + + validated_output_files: list[_ValidatedOutputFile] = [] + total_bytes_read = 0 + for relative_path, host_path in safe_output_files: try: - contents.append(_create_file_content(host_path, relative_path=relative_path)) + data = _read_output_file_bytes( + root, + relative_path=relative_path, + max_file_bytes=max_output_file_bytes, + remaining_total_bytes=max_output_total_bytes - total_bytes_read, + max_total_bytes=max_output_total_bytes, + ) + except MemoryError: + raise _OutputMaterializationError( + "Sandbox output could not be read because the host did not have enough memory." + ) from None except (PermissionError, OSError): missing_files = True + continue + + total_bytes_read += len(data) + validated_output_files.append( + _ValidatedOutputFile( + relative_path=relative_path, + media_type=mimetypes.guess_type(host_path.name)[0] or "application/octet-stream", + data=data, + ) + ) if not missing_files or attempt == OUTPUT_FILE_RETRY_ATTEMPTS - 1: - return contents + return _materialize_output_files(validated_output_files) time.sleep(OUTPUT_FILE_RETRY_DELAY_SECONDS) @@ -868,9 +1136,11 @@ def _result_snapshot(result: Any) -> dict[str, Any]: def _build_execution_contents( *, result: Any, - sandbox: Any, output_dir: TemporaryDirectory[str] | None, code: str, + max_output_files: int, + max_output_file_bytes: int, + max_output_total_bytes: int, ) -> list[Content]: success = bool(getattr(result, "success", False)) stdout = str(getattr(result, "stdout", "") or "").replace("\r\n", "\n") or None @@ -881,13 +1151,25 @@ def _build_execution_contents( if stdout is not None: outputs.append(Content.from_text(stdout, raw_representation=snapshot)) - outputs.extend( - _parse_output_files( - sandbox=sandbox, + try: + output_files = _parse_output_files( output_dir=output_dir, expect_output_files="/output" in code, + max_output_files=max_output_files, + max_output_file_bytes=max_output_file_bytes, + max_output_total_bytes=max_output_total_bytes, ) - ) + except _OutputMaterializationError as exc: + outputs.append( + Content.from_error( + message="Execution error", + error_details=str(exc), + raw_representation=snapshot, + ) + ) + return outputs + + outputs.extend(output_files) if success: if stderr is not None: @@ -985,6 +1267,9 @@ def execute(self, *, config: _RunConfig, code: str) -> list[Content]: code=code, output_dir=entry.output_dir, build_contents=_build_execution_contents, + max_output_files=config.max_output_files, + max_output_file_bytes=config.max_output_file_bytes, + max_output_total_bytes=config.max_output_total_bytes, ) def _get_or_create_entry(self, config: _RunConfig) -> _SandboxEntry: @@ -1093,11 +1378,17 @@ def __init__( workspace_root: str | Path | None = None, file_mounts: FileMountInput | Sequence[FileMountInput] | None = None, allowed_domains: AllowedDomainInput | Sequence[AllowedDomainInput] | None = None, + max_output_files: int = DEFAULT_MAX_OUTPUT_FILES, + max_output_file_bytes: int = DEFAULT_MAX_OUTPUT_FILE_BYTES, + max_output_total_bytes: int = DEFAULT_MAX_OUTPUT_TOTAL_BYTES, backend: str = DEFAULT_HYPERLIGHT_BACKEND, module: str | None = DEFAULT_HYPERLIGHT_MODULE, module_path: str | None = None, _registry: SandboxRuntime | None = None, ) -> None: + max_output_files = _validate_positive_integer(name="max_output_files", value=max_output_files) + max_output_file_bytes = _validate_positive_integer(name="max_output_file_bytes", value=max_output_file_bytes) + max_output_total_bytes = _validate_positive_integer(name="max_output_total_bytes", value=max_output_total_bytes) super().__init__( name="execute_code", description=EXECUTE_CODE_TOOL_DESCRIPTION, @@ -1112,6 +1403,9 @@ def __init__( self._backend: str = backend self._module: str | None = module self._module_path: str | None = module_path + self._max_output_files = max_output_files + self._max_output_file_bytes = max_output_file_bytes + self._max_output_total_bytes = max_output_total_bytes self._managed_tools: list[FunctionTool] = [] self._file_mounts: dict[str, FileMount] = {} self._allowed_domains: dict[str, AllowedDomain] = {} @@ -1264,6 +1558,9 @@ def create_run_tool(self) -> HyperlightExecuteCodeTool: workspace_root=self._workspace_root, file_mounts=file_mounts or None, allowed_domains=allowed_domains or None, + max_output_files=self._max_output_files, + max_output_file_bytes=self._max_output_file_bytes, + max_output_total_bytes=self._max_output_total_bytes, backend=self._backend, module=self._module, module_path=self._module_path, @@ -1278,6 +1575,9 @@ def build_serializable_state(self) -> dict[str, Any]: "module": config.module, "module_path": config.module_path, "approval_mode": config.approval_mode, + "max_output_files": config.max_output_files, + "max_output_file_bytes": config.max_output_file_bytes, + "max_output_total_bytes": config.max_output_total_bytes, "tool_names": [tool_obj.name for tool_obj in config.tools], "filesystem_enabled": config.filesystem_enabled, "workspace_root": str(config.workspace_root) if config.workspace_root is not None else None, @@ -1314,6 +1614,9 @@ def _build_run_config(self) -> _RunConfig: workspace_root = self._workspace_root stored_mounts = tuple(self._file_mounts.values()) allowed_domains = tuple(sorted(self._allowed_domains.values(), key=lambda value: value.target)) + max_output_files = self._max_output_files + max_output_file_bytes = self._max_output_file_bytes + max_output_total_bytes = self._max_output_total_bytes approval_mode = _resolve_execute_code_approval_mode( base_approval_mode=self._default_approval_mode, tools=managed_tools, @@ -1339,6 +1642,9 @@ def _build_run_config(self) -> _RunConfig: workspace_signature=workspace_signature, file_mounts=normalized_mounts, allowed_domains=allowed_domains, + max_output_files=max_output_files, + max_output_file_bytes=max_output_file_bytes, + max_output_total_bytes=max_output_total_bytes, ) async def _run_code(self, *, code: str) -> list[Content]: diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_provider.py b/python/packages/hyperlight/agent_framework_hyperlight/_provider.py index a4fb3a30d6b..7d26f295ef8 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_provider.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_provider.py @@ -10,7 +10,13 @@ from agent_framework._telemetry import mark_feature_used from agent_framework._tools import ApprovalMode -from ._execute_code_tool import HyperlightExecuteCodeTool, SandboxRuntime +from ._execute_code_tool import ( + DEFAULT_MAX_OUTPUT_FILE_BYTES, + DEFAULT_MAX_OUTPUT_FILES, + DEFAULT_MAX_OUTPUT_TOTAL_BYTES, + HyperlightExecuteCodeTool, + SandboxRuntime, +) from ._feature_usage import FeatureIndex from ._types import AllowedDomain, AllowedDomainInput, FileMount, FileMountInput @@ -29,6 +35,9 @@ def __init__( workspace_root: str | Path | None = None, file_mounts: FileMountInput | Sequence[FileMountInput] | None = None, allowed_domains: AllowedDomainInput | Sequence[AllowedDomainInput] | None = None, + max_output_files: int = DEFAULT_MAX_OUTPUT_FILES, + max_output_file_bytes: int = DEFAULT_MAX_OUTPUT_FILE_BYTES, + max_output_total_bytes: int = DEFAULT_MAX_OUTPUT_TOTAL_BYTES, backend: str = "wasm", module: str | None = "python_guest.path", module_path: str | None = None, @@ -41,6 +50,9 @@ def __init__( workspace_root=workspace_root, file_mounts=file_mounts, allowed_domains=allowed_domains, + max_output_files=max_output_files, + max_output_file_bytes=max_output_file_bytes, + max_output_total_bytes=max_output_total_bytes, backend=backend, module=module, module_path=module_path, diff --git a/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py b/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py index 984228d81b2..636ede98b78 100644 --- a/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py +++ b/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py @@ -359,6 +359,61 @@ def _write_file() -> None: return super().run(code) +class _FakeSandboxWithBoundedOutputs(_FakeSandbox): + def run(self, code: str) -> _FakeResult: + if code == "None": + return _FakeResult(success=True) + if self.output_dir is None: + raise AssertionError("Expected output directory for bounded output test.") + + output_root = Path(self.output_dir) + if code == "create-sparse-output": + sparse_file = output_root / "sparse.bin" + with sparse_file.open("wb") as handle: + handle.seek((2 << 30) - 1) + handle.write(b"X") + self.output_files = ["sparse.bin"] + elif code == "create-count-output": + for index in range(3): + (output_root / f"report-{index}.txt").write_bytes(b"x") + self.output_files = [f"report-{index}.txt" for index in range(3)] + elif code == "create-cumulative-output": + (output_root / "first.bin").write_bytes(b"abc") + (output_root / "second.bin").write_bytes(b"def") + self.output_files = ["first.bin", "second.bin"] + elif code == "create-exact-output": + (output_root / "first.bin").write_bytes(b"abc") + (output_root / "second.bin").write_bytes(b"de") + self.output_files = ["first.bin", "second.bin"] + elif code == "create-growing-output": + (output_root / "growing.bin").write_bytes(b"data") + self.output_files = ["growing.bin"] + elif code == "create-memory-output": + (output_root / "memory.bin").write_bytes(b"data") + self.output_files = ["memory.bin"] + elif code == "create-nested-output": + (output_root / "nested").mkdir() + (output_root / "nested" / "report.bin").write_bytes(b"data") + self.output_files = ["nested/report.bin"] + else: + return super().run(code) + + return _FakeResult(success=True, stdout="guest-finished\n") + + +class _FakeSandboxWithBoundedOutputsWithoutListing(_FakeSandboxWithBoundedOutputs): + def get_output_files(self) -> list[str]: + return [] + + +class _FakeSandboxWithEagerOutputListing(_FakeSandboxWithBoundedOutputs): + listing_calls = 0 + + def get_output_files(self) -> list[str]: + type(self).listing_calls += 1 + return [f"report-{index}.txt" for index in range(10_000)] + + class _FakeSessionContext: def __init__(self, *, tools: list[Any] | None = None) -> None: self.options: dict[str, Any] = {} @@ -834,14 +889,6 @@ def __init__(self, path: Path) -> None: self.name = str(path) -class _SandboxWithListing: - def __init__(self, output_files: list[str]) -> None: - self._output_files = output_files - - def get_output_files(self) -> list[str]: - return self._output_files - - def _decode_content_bytes(item: Content) -> bytes: import base64 @@ -850,6 +897,16 @@ def _decode_content_bytes(item: Content) -> bytes: return base64.b64decode(encoded) +def _assert_bounded_output_error(contents: list[Content], match: str) -> None: + assert any(item.type == "text" and item.text == "guest-finished\n" for item in contents) + assert not any(item.type == "data" for item in contents) + errors = [item for item in contents if item.type == "error"] + assert len(errors) == 1 + assert errors[0].message == "Execution error" + assert errors[0].error_details is not None + assert match in errors[0].error_details + + def test_collect_output_relative_paths_skips_symlinked_file(tmp_path: Path) -> None: """A final-component symlink planted in /output must not be surfaced.""" if not _symlinks_supported(tmp_path): @@ -861,7 +918,7 @@ def test_collect_output_relative_paths_skips_symlinked_file(tmp_path: Path) -> N secret.write_text("HOST_SECRET", encoding="utf-8") (output_root / "leak.txt").symlink_to(secret) - relative_paths = execute_code_module._collect_output_relative_paths(sandbox=object(), root=output_root) + relative_paths = execute_code_module._collect_output_relative_paths(root=output_root) assert "report.txt" in relative_paths assert "leak.txt" not in relative_paths @@ -878,7 +935,7 @@ def test_collect_output_relative_paths_skips_symlinked_directory(tmp_path: Path) (outside_dir / "deep.txt").write_text("deep-secret", encoding="utf-8") (output_root / "linked_dir").symlink_to(outside_dir, target_is_directory=True) - relative_paths = execute_code_module._collect_output_relative_paths(sandbox=object(), root=output_root) + relative_paths = execute_code_module._collect_output_relative_paths(root=output_root) assert relative_paths == set() @@ -892,11 +949,71 @@ def test_collect_output_relative_paths_skips_junctioned_directory(tmp_path: Path (outside_dir / "deep.txt").write_text("deep-secret", encoding="utf-8") _create_junction_or_skip(link=output_root / "linked_dir", target=outside_dir) - relative_paths = execute_code_module._collect_output_relative_paths(sandbox=object(), root=output_root) + relative_paths = execute_code_module._collect_output_relative_paths(root=output_root) assert relative_paths == set() +def test_collect_output_relative_paths_bounds_directory_only_breadth(tmp_path: Path) -> None: + output_root = tmp_path / "output" + output_root.mkdir() + for index in range(101): + (output_root / f"directory-{index}").mkdir() + + with pytest.raises(execute_code_module._OutputMaterializationError, match="traversal entry limit of 100"): + execute_code_module._collect_output_relative_paths(root=output_root, max_output_files=1) + + +def test_collect_output_relative_paths_bounds_nesting_depth(tmp_path: Path) -> None: + output_root = tmp_path / "output" + output_root.mkdir() + current = output_root + for index in range(execute_code_module.OUTPUT_TRAVERSAL_MAX_DEPTH + 1): + current /= f"level-{index}" + current.mkdir() + + with pytest.raises(execute_code_module._OutputMaterializationError, match="nesting depth limit"): + execute_code_module._collect_output_relative_paths(root=output_root) + + +def test_collect_output_relative_paths_surfaces_scandir_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + output_root = tmp_path / "output" + output_root.mkdir() + original_scandir = os.scandir + + def fail_output_scan(path: str | os.PathLike[str]) -> Any: + if Path(path) == output_root: + raise PermissionError("simulated output scan failure") + return original_scandir(path) + + monkeypatch.setattr(execute_code_module.os, "scandir", fail_output_scan) + + with pytest.raises(execute_code_module._OutputMaterializationError, match="Could not enumerate output directory"): + execute_code_module._collect_output_relative_paths(root=output_root) + + +def test_parse_output_files_collects_legitimate_nested_file(tmp_path: Path) -> None: + if not execute_code_module._supports_secure_output_dir_fd(): + pytest.skip("Nested output attachments require secure dir_fd support") + output_root = tmp_path / "output" + nested_dir = output_root / "nested" + nested_dir.mkdir(parents=True) + (nested_dir / "report.txt").write_bytes(b"nested-report") + + contents = execute_code_module._parse_output_files( + output_dir=cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), + expect_output_files=True, + ) + + data_items = [item for item in contents if item.type == "data"] + assert len(data_items) == 1 + assert data_items[0].additional_properties["path"] == "/output/nested/report.txt" + assert _decode_content_bytes(data_items[0]) == b"nested-report" + + def test_parse_output_files_skips_symlink_to_host_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """End-to-end: a /output symlink to a host file is never returned as Content.""" if not _symlinks_supported(tmp_path): @@ -910,7 +1027,6 @@ def test_parse_output_files_skips_symlink_to_host_file(tmp_path: Path, monkeypat (output_root / "leak.txt").symlink_to(secret) contents = execute_code_module._parse_output_files( - sandbox=object(), output_dir=cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), expect_output_files=False, ) @@ -921,10 +1037,8 @@ def test_parse_output_files_skips_symlink_to_host_file(tmp_path: Path, monkeypat assert all(b"HOST_SECRET" not in _decode_content_bytes(item) for item in contents if item.type == "data") -def test_parse_output_files_rejects_intermediate_dir_symlink_from_listing( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """A backend-listed path traversing an intermediate dir symlink must be rejected.""" +def test_parse_output_files_rejects_intermediate_dir_symlink(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A path traversing an intermediate dir symlink must be rejected.""" if not _symlinks_supported(tmp_path): pytest.skip("Symlinks not supported on this platform/environment") monkeypatch.setattr(execute_code_module, "OUTPUT_FILE_RETRY_ATTEMPTS", 1) @@ -936,7 +1050,6 @@ def test_parse_output_files_rejects_intermediate_dir_symlink_from_listing( (output_root / "sub").symlink_to(outside_dir, target_is_directory=True) contents = execute_code_module._parse_output_files( - sandbox=_SandboxWithListing(["output/sub/leak.txt"]), output_dir=cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), expect_output_files=False, ) @@ -945,10 +1058,8 @@ def test_parse_output_files_rejects_intermediate_dir_symlink_from_listing( assert all(item.additional_properties.get("path") != "/output/sub/leak.txt" for item in contents) -def test_parse_output_files_rejects_intermediate_dir_junction_from_listing( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """A backend-listed path traversing an intermediate dir junction must be rejected.""" +def test_parse_output_files_rejects_intermediate_dir_junction(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A path traversing an intermediate dir junction must be rejected.""" monkeypatch.setattr(execute_code_module, "OUTPUT_FILE_RETRY_ATTEMPTS", 1) output_root = tmp_path / "output" output_root.mkdir() @@ -958,7 +1069,6 @@ def test_parse_output_files_rejects_intermediate_dir_junction_from_listing( _create_junction_or_skip(link=output_root / "sub", target=outside_dir) contents = execute_code_module._parse_output_files( - sandbox=_SandboxWithListing(["output/sub/leak.txt"]), output_dir=cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), expect_output_files=False, ) @@ -967,6 +1077,69 @@ def test_parse_output_files_rejects_intermediate_dir_junction_from_listing( assert all(item.additional_properties.get("path") != "/output/sub/leak.txt" for item in contents) +def test_parse_output_files_rejects_intermediate_dir_swap_after_validation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if not _symlinks_supported(tmp_path): + pytest.skip("Symlinks not supported on this platform/environment") + if not execute_code_module._supports_secure_output_dir_fd(): + pytest.skip("Atomic intermediate-directory swap test requires secure dir_fd support") + monkeypatch.setattr(execute_code_module, "OUTPUT_FILE_RETRY_ATTEMPTS", 1) + output_root = tmp_path / "output" + output_root.mkdir() + output_subdir = output_root / "sub" + output_subdir.mkdir() + output_file = output_subdir / "report.txt" + output_file.write_text("safe-report", encoding="utf-8") + outside_dir = tmp_path / "outside" + outside_dir.mkdir() + (outside_dir / "report.txt").write_text("HOST_SECRET", encoding="utf-8") + original_is_safe_output_file = execute_code_module._is_safe_output_file + swapped = False + + def swap_parent_after_validation(*, root: Path, host_path: Path) -> bool: + nonlocal swapped + is_safe = original_is_safe_output_file(root=root, host_path=host_path) + if is_safe and host_path == output_file and not swapped: + swapped = True + output_subdir.rename(output_root / "original-sub") + output_subdir.symlink_to(outside_dir, target_is_directory=True) + return is_safe + + monkeypatch.setattr(execute_code_module, "_is_safe_output_file", swap_parent_after_validation) + + contents = execute_code_module._parse_output_files( + output_dir=cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), + expect_output_files=False, + ) + + assert swapped + assert not any(item.type == "data" for item in contents) + assert all(b"HOST_SECRET" not in _decode_content_bytes(item) for item in contents if item.type == "data") + + +async def test_execute_code_tool_fails_closed_for_nested_output_without_secure_dir_fd( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + monkeypatch.setattr( + execute_code_module, + "_supports_secure_output_dir_fd", + lambda: False, + raising=False, + ) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path) + + try: + contents = await execute_code.invoke(arguments={"code": "create-nested-output"}) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error(contents, "Nested output attachments cannot be opened safely") + + def test_clear_directory_removes_junction_without_deleting_target(tmp_path: Path) -> None: output_root = tmp_path / "output" output_root.mkdir() @@ -1022,7 +1195,6 @@ def test_parse_output_files_collects_real_output_file(tmp_path: Path) -> None: (output_root / "report.txt").write_text("artifact", encoding="utf-8") contents = execute_code_module._parse_output_files( - sandbox=object(), output_dir=cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), expect_output_files=True, ) @@ -1033,6 +1205,337 @@ def test_parse_output_files_collects_real_output_file(tmp_path: Path) -> None: assert _decode_content_bytes(data_items[0]) == b"artifact" +async def test_execute_code_tool_rejects_sparse_output_without_unbounded_read( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + original_fdopen = os.fdopen + + class _BoundedReadGuard: + def __init__(self, handle: Any) -> None: + self._handle = handle + + def __enter__(self) -> _BoundedReadGuard: + return self + + def __exit__(self, *args: Any) -> None: + self._handle.close() + + def read(self, size: int = -1) -> bytes: + assert size >= 0, "output reads must always be bounded" + assert size <= 5 * 1024 * 1024 + 1 + return cast(bytes, self._handle.read(size)) + + def guarded_fdopen(fd: int, *args: Any, **kwargs: Any) -> _BoundedReadGuard: + return _BoundedReadGuard(original_fdopen(fd, *args, **kwargs)) + + monkeypatch.setattr(execute_code_module.os, "fdopen", guarded_fdopen) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path) + try: + contents = await execute_code.invoke(arguments={"code": "create-sparse-output"}) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error(contents, "per-file output limit") + + +async def test_execute_code_tool_checks_output_count_before_reading( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + + def fail_if_read(*args: Any, **kwargs: Any) -> bytes: + del args, kwargs + pytest.fail("output files were read before enforcing the count limit") + + monkeypatch.setattr(execute_code_module, "_read_output_file_bytes", fail_if_read) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path, max_output_files=2) + try: + contents = await execute_code.invoke(arguments={"code": "create-count-output"}) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error(contents, "output file count limit") + + +async def test_execute_code_tool_does_not_call_eager_backend_output_listing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _FakeSandboxWithEagerOutputListing.listing_calls = 0 + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithEagerOutputListing) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path, max_output_files=2) + + try: + contents = await execute_code.invoke(arguments={"code": "create-count-output"}) + finally: + _close_execute_code_registry(execute_code) + + assert _FakeSandboxWithEagerOutputListing.listing_calls == 0 + _assert_bounded_output_error(contents, "output file count limit") + + +async def test_execute_code_tool_streams_directory_enumeration_to_count_limit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + execute_code_module, + "_load_sandbox_class", + lambda: _FakeSandboxWithBoundedOutputsWithoutListing, + ) + original_scandir = os.scandir + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path, max_output_files=2) + config = execute_code._build_run_config() + output_root = Path(cast(Any, execute_code._registry)._get_or_create_entry(config).output_dir.name) + scanned_entries = 0 + + class _BoundedScandir: + def __init__(self, path: str | os.PathLike[str]) -> None: + self._entries = original_scandir(path) + self._track = Path(path) == output_root + + def __enter__(self) -> _BoundedScandir: + self._entries.__enter__() + return self + + def __exit__(self, *args: Any) -> None: + self._entries.__exit__(*args) + + def close(self) -> None: + self._entries.close() + + def __iter__(self) -> _BoundedScandir: + return self + + def __next__(self) -> os.DirEntry[str]: + nonlocal scanned_entries + entry = next(self._entries) + if self._track: + scanned_entries += 1 + if scanned_entries > 3: + pytest.fail("output directory enumeration continued past max_output_files + 1") + return entry + + monkeypatch.setattr(execute_code_module.os, "scandir", _BoundedScandir) + + try: + contents = await execute_code.invoke(arguments={"code": "create-count-output"}) + finally: + monkeypatch.setattr(execute_code_module.os, "scandir", original_scandir) + _close_execute_code_registry(execute_code) + + assert scanned_entries == 3 + _assert_bounded_output_error(contents, "output file count limit") + + +async def test_execute_code_tool_reads_only_observed_file_size_with_large_limits( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + original_fdopen = os.fdopen + read_sizes: list[int] = [] + + class _ReadSizeGuard: + def __init__(self, handle: Any) -> None: + self._handle = handle + + def __enter__(self) -> _ReadSizeGuard: + return self + + def __exit__(self, *args: Any) -> None: + self._handle.close() + + def read(self, size: int = -1) -> bytes: + read_sizes.append(size) + if size > 5: + pytest.fail("small output requested a quota-sized read") + return cast(bytes, self._handle.read(size)) + + monkeypatch.setattr( + execute_code_module.os, + "fdopen", + lambda fd, *args, **kwargs: _ReadSizeGuard(original_fdopen(fd, *args, **kwargs)), + ) + execute_code = HyperlightExecuteCodeTool( + workspace_root=tmp_path, + max_output_file_bytes=sys.maxsize, + max_output_total_bytes=sys.maxsize, + ) + + try: + contents = await execute_code.invoke(arguments={"code": "create-memory-output"}) + finally: + _close_execute_code_registry(execute_code) + + assert read_sizes == [5] + assert [_decode_content_bytes(item) for item in contents if item.type == "data"] == [b"data"] + + +async def test_execute_code_tool_rejects_cumulative_output_overflow( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + execute_code = HyperlightExecuteCodeTool( + workspace_root=tmp_path, + max_output_file_bytes=3, + max_output_total_bytes=5, + ) + try: + contents = await execute_code.invoke(arguments={"code": "create-cumulative-output"}) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error(contents, "cumulative output limit") + + +async def test_execute_code_tool_accepts_output_at_exact_limits( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + execute_code = HyperlightExecuteCodeTool( + workspace_root=tmp_path, + max_output_files=2, + max_output_file_bytes=3, + max_output_total_bytes=5, + ) + try: + contents = await execute_code.invoke(arguments={"code": "create-exact-output"}) + finally: + _close_execute_code_registry(execute_code) + + data_items = [item for item in contents if item.type == "data"] + assert [_decode_content_bytes(item) for item in data_items] == [b"abc", b"de"] + assert not any(item.type == "error" for item in contents) + + +async def test_execute_code_tool_bounds_file_growth_after_fstat( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + original_fdopen = os.fdopen + read_sizes: list[int] = [] + + class _GrowingHandle: + def __init__(self, handle: Any, file_path: Path) -> None: + self._handle = handle + self._file_path = file_path + + def __enter__(self) -> _GrowingHandle: + return self + + def __exit__(self, *args: Any) -> None: + self._handle.close() + + def read(self, size: int = -1) -> bytes: + read_sizes.append(size) + with self._file_path.open("ab") as growing_file: + growing_file.write(b"X") + return cast(bytes, self._handle.read(size)) + + execute_code = HyperlightExecuteCodeTool( + workspace_root=tmp_path, + max_output_file_bytes=8, + max_output_total_bytes=8, + ) + try: + config = execute_code._build_run_config() + output_root = Path(cast(Any, execute_code._registry)._get_or_create_entry(config).output_dir.name) + monkeypatch.setattr( + execute_code_module.os, + "fdopen", + lambda fd, *args, **kwargs: _GrowingHandle( + original_fdopen(fd, *args, **kwargs), output_root / "growing.bin" + ), + ) + contents = await execute_code.invoke(arguments={"code": "create-growing-output"}) + finally: + _close_execute_code_registry(execute_code) + + assert read_sizes == [5] + _assert_bounded_output_error(contents, "grew while it was being read") + + +@pytest.mark.parametrize("stage", ["read_memory", "read_overflow", "content_memory"]) +async def test_execute_code_tool_converts_output_allocation_error_to_content_error( + stage: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + + if stage != "content_memory": + original_fdopen = os.fdopen + + class _MemoryErrorHandle: + def __init__(self, handle: Any) -> None: + self._handle = handle + + def __enter__(self) -> _MemoryErrorHandle: + return self + + def __exit__(self, *args: Any) -> None: + self._handle.close() + + def read(self, size: int = -1) -> bytes: + del size + if stage == "read_memory": + raise MemoryError("simulated allocation failure") + raise OverflowError("simulated read size overflow") + + monkeypatch.setattr( + execute_code_module.os, + "fdopen", + lambda fd, *args, **kwargs: _MemoryErrorHandle(original_fdopen(fd, *args, **kwargs)), + ) + else: + + def raise_memory_error(*args: Any, **kwargs: Any) -> Any: + del args, kwargs + raise MemoryError("simulated allocation failure") + + monkeypatch.setattr(execute_code_module.Content, "from_data", raise_memory_error) + + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path) + try: + contents = await execute_code.invoke(arguments={"code": "create-memory-output"}) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error( + contents, + "configured byte limits" if stage != "content_memory" else "enough memory", + ) + + +@pytest.mark.parametrize( + ("option", "value", "error_type"), + [ + ("max_output_files", 0, ValueError), + ("max_output_file_bytes", -1, ValueError), + ("max_output_total_bytes", True, TypeError), + ("max_output_files", 1.5, TypeError), + ("max_output_file_bytes", "1024", TypeError), + ], +) +def test_hyperlight_output_limit_options_require_positive_integers( + option: str, + value: Any, + error_type: type[Exception], +) -> None: + kwargs = {option: value} + with pytest.raises(error_type, match=option): + HyperlightExecuteCodeTool(**kwargs) + with pytest.raises(error_type, match=option): + HyperlightCodeActProvider(**kwargs) + + def test_execute_code_tool_allowed_domains_use_structured_entries_and_replace_by_target() -> None: execute_code = HyperlightExecuteCodeTool(_registry=_FakeRuntime()) @@ -1304,6 +1807,75 @@ async def test_provider_injects_run_scoped_execute_code_tool() -> None: assert [tool_obj.name for tool_obj in run_tool.get_tools()] == ["compute"] +async def test_provider_forwards_output_limits_to_run_tool_and_serializable_state() -> None: + runtime = _FakeRuntime() + provider = HyperlightCodeActProvider( + max_output_files=7, + max_output_file_bytes=11, + max_output_total_bytes=13, + _registry=runtime, + ) + context = _FakeSessionContext() + state: dict[str, Any] = {} + + await provider.before_run(agent=object(), session=None, context=cast(Any, context), state=state) + run_tool = context.tools[0][1][0] + assert isinstance(run_tool, HyperlightExecuteCodeTool) + + result = await run_tool.invoke(arguments={"code": "None"}) + + assert result[0].text == "ok" + config = runtime.calls[0][0] + assert config.max_output_files == 7 + assert config.max_output_file_bytes == 11 + assert config.max_output_total_bytes == 13 + assert state[provider.source_id]["max_output_files"] == 7 + assert state[provider.source_id]["max_output_file_bytes"] == 11 + assert state[provider.source_id]["max_output_total_bytes"] == 13 + json.dumps(state) + + +async def test_output_limits_are_invocation_scoped_when_registry_is_shared( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _FakeSandbox.instances.clear() + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + registry = execute_code_module._SandboxRegistry() + restrictive_tool = HyperlightExecuteCodeTool( + workspace_root=tmp_path, + max_output_file_bytes=3, + max_output_total_bytes=10, + _registry=registry, + ) + permissive_tool = HyperlightExecuteCodeTool( + workspace_root=tmp_path, + max_output_file_bytes=4, + max_output_total_bytes=10, + _registry=registry, + ) + + try: + rejected = await restrictive_tool.invoke(arguments={"code": "create-growing-output"}) + accepted = await permissive_tool.invoke(arguments={"code": "create-growing-output"}) + finally: + registry.close() + + _assert_bounded_output_error(rejected, "per-file output limit") + assert [_decode_content_bytes(item) for item in accepted if item.type == "data"] == [b"data"] + assert len(_FakeSandbox.instances) == 1 + + +def test_execute_code_tool_uses_finite_default_output_limits() -> None: + execute_code = HyperlightExecuteCodeTool(_registry=_FakeRuntime()) + + state = execute_code.build_serializable_state() + + assert state["max_output_files"] == 20 + assert state["max_output_file_bytes"] == 5 * 1024 * 1024 + assert state["max_output_total_bytes"] == 20 * 1024 * 1024 + + def test_provider_delegates_file_mounts_and_allowed_domains_to_internal_tool(tmp_path: Path) -> None: provider = HyperlightCodeActProvider() @@ -1423,6 +1995,30 @@ async def test_provider_run_tool_writes_files_with_real_sandbox(tmp_path: Path) _close_execute_code_registry(run_tool) +@pytest.mark.integration +@skip_if_hyperlight_integration_tests_disabled +async def test_execute_code_tool_rejects_sparse_file_with_real_sandbox(tmp_path: Path) -> None: + _skip_if_hyperlight_integration_runtime_disabled() + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path) + + try: + contents = await execute_code.invoke( + arguments={ + "code": ( + 'with open("/output/sparse.bin", "wb") as output:\n' + " output.seek((2 << 30) - 1)\n" + ' output.write(b"X")\n' + 'print("guest-finished")\n' + ) + } + ) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error(contents, "per-file output limit") + + +@pytest.mark.flaky @pytest.mark.integration @skip_if_hyperlight_integration_tests_disabled @pytest.mark.skipif(sys.platform == "win32", reason="Hyperlight WASM sandbox lacks encodings.idna on Windows")