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 3dfb20f0af..05b63bf5cc 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py @@ -57,6 +57,14 @@ class _OutputMaterializationError(RuntimeError): """Raised when sandbox output cannot be safely materialized.""" +class _OutputCleanupError(RuntimeError): + """Raised when sandbox output cannot be safely cleaned within configured bounds.""" + + +def _output_cleanup_error_content(error: _OutputCleanupError) -> Content: + return Content.from_error(message="Execution error", error_details=str(error)) + + @dataclass(frozen=True, slots=True) class _ValidatedOutputFile: relative_path: str @@ -148,7 +156,7 @@ class _SandboxWorker: copy (preserving message and exception type) is re-raised on the caller. """ - __slots__ = ("_executor", "_initialized", "_sandbox", "_snapshot") + __slots__ = ("_executor", "_initialized", "_reusable", "_sandbox", "_snapshot") def __init__(self, *, name: str = "hl-sandbox") -> None: self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix=name) @@ -156,6 +164,7 @@ def __init__(self, *, name: str = "hl-sandbox") -> None: self._sandbox: Any = None self._snapshot: Any = None self._initialized = False + self._reusable = True def _run_on_worker(self, fn: Callable[[], _T]) -> _T: """Run ``fn`` on the worker thread; sanitize any exception's traceback there. @@ -232,30 +241,60 @@ def execute( Returns a plain ``list[Content]`` whose elements never carry strong references to the underlying sandbox or snapshot. """ + cleanup_max_entries = _output_traversal_entry_limit(max_output_files) def _on_worker() -> list[Content]: + if not self._reusable: + return [_output_cleanup_error_content(_OutputCleanupError("Could not clear sandbox output safely."))] + sandbox = self._sandbox snapshot = self._snapshot sandbox.restore(snapshot) - _clear_directory(output_dir) - result = sandbox.run(code=code) try: - return build_contents( - result=result, - 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, + _clear_directory( + output_dir, + max_entries=cleanup_max_entries, + max_depth=OUTPUT_TRAVERSAL_MAX_DEPTH, ) + except _OutputCleanupError as exc: + self._reusable = False + return [_output_cleanup_error_content(exc)] + + contents: list[Content] | None = None + try: + result = sandbox.run(code=code) + try: + contents = build_contents( + result=result, + 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, + ) + return contents + finally: + # ``result`` may carry a back-reference to the sandbox. Force its + # final dec_ref on this thread so Drop runs here, not on whatever + # thread later GCs the ``Content`` list. + del result finally: - # ``result`` may carry a back-reference to the sandbox. Force its - # final dec_ref on this thread so Drop runs here, not on whatever - # thread later GCs the ``Content`` list. - del result + try: + _clear_directory( + output_dir, + max_entries=cleanup_max_entries, + max_depth=OUTPUT_TRAVERSAL_MAX_DEPTH, + ) + except _OutputCleanupError as exc: + self._reusable = False + if contents is not None and not any(item.type == "error" for item in contents): + contents.append(_output_cleanup_error_content(exc)) return self._run_on_worker(_on_worker) + def is_reusable(self) -> bool: + return self._reusable + def is_alive(self) -> bool: """Return ``True`` while the worker thread can still accept new submissions. @@ -1224,33 +1263,63 @@ def _run() -> None: return _callback -def _clear_directory(output_dir: TemporaryDirectory[str] | None) -> None: - """Remove all contents of the output directory without deleting the directory itself.""" +def _clear_directory( + output_dir: TemporaryDirectory[str] | None, + *, + max_entries: int = OUTPUT_TRAVERSAL_MAX_ENTRIES, + max_depth: int = OUTPUT_TRAVERSAL_MAX_DEPTH, +) -> None: + """Remove output entries without following links or exceeding traversal bounds.""" if output_dir is None: return + root = Path(output_dir.name) - for child in root.iterdir(): + entries_visited = 0 + + def _remove_contents(path: Path, depth: int) -> None: + nonlocal entries_visited try: - child_stat = child.lstat() - if _is_link_or_reparse_point(child, child_stat): - if stat.S_ISDIR(child_stat.st_mode): - child.rmdir() - continue - try: - child.unlink() - except OSError: + with os.scandir(path) as entries: + for entry in entries: + entries_visited += 1 + if entries_visited > max_entries: + raise _OutputCleanupError(f"Sandbox output exceeded the cleanup entry limit of {max_entries}.") + + child = Path(entry.path) + child_stat = child.lstat() + if _is_link_or_reparse_point(child, child_stat): + if stat.S_ISDIR(child_stat.st_mode): + child.rmdir() + else: + try: + child.unlink() + except OSError: + child.rmdir() + continue + + if not stat.S_ISDIR(child_stat.st_mode): + child.unlink() + continue + + child_depth = depth + 1 + if child_depth > max_depth: + raise _OutputCleanupError( + f"Sandbox output exceeded the cleanup nesting depth limit of {max_depth}." + ) + _remove_contents(child, child_depth) child.rmdir() - elif stat.S_ISREG(child_stat.st_mode): - child.unlink() - elif stat.S_ISDIR(child_stat.st_mode): - shutil.rmtree(child, ignore_errors=True) - except OSError: - pass + except _OutputCleanupError: + raise + except OSError as exc: + raise _OutputCleanupError("Could not clear sandbox output safely.") from exc + + _remove_contents(root, 0) class _SandboxRegistry(SandboxRuntime): def __init__(self) -> None: self._entries: dict[tuple[Any, ...], _SandboxEntry] = {} + self._retired_entries: list[_SandboxEntry] = [] self._entries_lock = threading.RLock() def execute(self, *, config: _RunConfig, code: str) -> list[Content]: @@ -1261,16 +1330,30 @@ def execute(self, *, config: _RunConfig, code: str) -> list[Content]: both serializes concurrent callers and satisfies the PyO3 ``unsendable`` invariant that the sandbox can only be touched from the thread that created it. The unsendable objects never escape the worker; this method returns only sendable plain Python data. + Entries whose output cannot be cleaned safely are evicted before another invocation. """ + cache_key = config.cache_key() entry = self._get_or_create_entry(config) - return entry.worker.execute( - 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, - ) + try: + return entry.worker.execute( + 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, + ) + finally: + if not entry.worker.is_reusable(): + self._discard_entry(cache_key=cache_key, entry=entry) + + def _discard_entry(self, *, cache_key: tuple[Any, ...], entry: _SandboxEntry) -> None: + with self._entries_lock: + if self._entries.get(cache_key) is not entry: + return + del self._entries[cache_key] + self._retired_entries.append(entry) + entry.worker.dispose() def _get_or_create_entry(self, config: _RunConfig) -> _SandboxEntry: cache_key = config.cache_key() @@ -1288,8 +1371,9 @@ def close(self) -> None: worker thread that created it to honor the PyO3 ``unsendable`` invariant. """ with self._entries_lock: - entries = list(self._entries.values()) + entries = [*self._entries.values(), *self._retired_entries] self._entries.clear() + self._retired_entries.clear() try: for entry in entries: entry.dispose() diff --git a/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py b/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py index 636ede98b7..1cae1406a3 100644 --- a/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py +++ b/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py @@ -395,6 +395,14 @@ def run(self, code: str) -> _FakeResult: (output_root / "nested").mkdir() (output_root / "nested" / "report.bin").write_bytes(b"data") self.output_files = ["nested/report.bin"] + elif code == "create-wide-output": + for index in range(101): + (output_root / f"directory-{index}").mkdir() + elif code == "create-deep-output": + current = output_root + for index in range(execute_code_module.OUTPUT_TRAVERSAL_MAX_DEPTH + 1): + current /= f"level-{index}" + current.mkdir() else: return super().run(code) @@ -1155,6 +1163,35 @@ def test_clear_directory_removes_junction_without_deleting_target(tmp_path: Path assert not (output_root / "linked_dir").exists() +def test_clear_directory_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._OutputCleanupError, match="cleanup entry limit of 100"): + execute_code_module._clear_directory( + cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), + max_entries=100, + max_depth=execute_code_module.OUTPUT_TRAVERSAL_MAX_DEPTH, + ) + + assert sum(1 for _ in output_root.iterdir()) == 1 + + +def test_clear_directory_bounds_nesting_depth(tmp_path: Path) -> None: + output_root = tmp_path / "output" + current = output_root + for index in range(3): + current /= f"level-{index}" + current.mkdir(parents=True) + + with pytest.raises(execute_code_module._OutputCleanupError, match="cleanup nesting depth limit of 2"): + execute_code_module._clear_directory( + cast("TemporaryDirectory[str]", _OutputDirShim(output_root)), max_entries=100, max_depth=2 + ) + + def test_is_safe_output_file_rejects_parent_traversal(tmp_path: Path) -> None: """A lexical ``..`` component must be rejected even without any symlink.""" output_root = tmp_path / "output" @@ -1240,6 +1277,109 @@ def guarded_fdopen(fd: int, *args: Any, **kwargs: Any) -> _BoundedReadGuard: _assert_bounded_output_error(contents, "per-file output limit") +async def test_execute_code_tool_clears_output_after_rejection( + 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, + ) + config = execute_code._build_run_config() + output_root = Path(cast(Any, execute_code._registry)._get_or_create_entry(config).output_dir.name) + + try: + contents = await execute_code.invoke(arguments={"code": "create-memory-output"}) + + _assert_bounded_output_error(contents, "per-file output limit") + remaining_outputs = await asyncio.to_thread(lambda: list(output_root.iterdir())) + assert remaining_outputs == [] + finally: + _close_execute_code_registry(execute_code) + + +async def test_execute_code_tool_preserves_result_when_post_cleanup_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _FakeSandbox.instances.clear() + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandbox) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path) + config = execute_code._build_run_config() + registry = cast(Any, execute_code._registry) + output_root = Path(registry._get_or_create_entry(config).output_dir.name) + original_scandir = os.scandir + output_scan_calls = 0 + + def fail_post_cleanup(path: str | os.PathLike[str]) -> Any: + nonlocal output_scan_calls + if Path(path) == output_root: + output_scan_calls += 1 + if output_scan_calls == 3: + raise PermissionError("simulated cleanup failure") + return original_scandir(path) + + monkeypatch.setattr(execute_code_module.os, "scandir", fail_post_cleanup) + + try: + contents = await execute_code.invoke(arguments={"code": "create-output"}) + + assert any(item.type == "text" and item.text == "done\n" for item in contents) + assert any(item.type == "data" for item in contents) + cleanup_errors = [item for item in contents if item.type == "error"] + assert len(cleanup_errors) == 1 + assert cleanup_errors[0].error_details == "Could not clear sandbox output safely." + assert str(output_root) not in cleanup_errors[0].error_details + assert registry._entries == {} + + recovered = await execute_code.invoke(arguments={"code": "fail"}) + assert not any(item.type == "data" for item in recovered) + assert any(item.type == "error" and item.error_details == "sandbox boom" for item in recovered) + replacement_output_root = Path(registry._get_or_create_entry(config).output_dir.name) + remaining_outputs = await asyncio.to_thread(lambda: list(replacement_output_root.iterdir())) + assert len(_FakeSandbox.instances) == 2 + assert remaining_outputs == [] + finally: + _close_execute_code_registry(execute_code) + + +async def test_execute_code_tool_preserves_exception_when_post_cleanup_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _FakeSandbox.instances.clear() + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandbox) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path) + config = execute_code._build_run_config() + registry = cast(Any, execute_code._registry) + output_root = Path(registry._get_or_create_entry(config).output_dir.name) + original_scandir = os.scandir + output_scan_calls = 0 + + def fail_post_cleanup(path: str | os.PathLike[str]) -> Any: + nonlocal output_scan_calls + if Path(path) == output_root: + output_scan_calls += 1 + if output_scan_calls == 2: + raise PermissionError("simulated cleanup failure") + return original_scandir(path) + + def fail_build_contents(**kwargs: Any) -> list[Content]: + del kwargs + raise RuntimeError("primary build failure") + + monkeypatch.setattr(execute_code_module.os, "scandir", fail_post_cleanup) + monkeypatch.setattr(execute_code_module, "_build_execution_contents", fail_build_contents) + + try: + with pytest.raises(RuntimeError, match="primary build failure"): + await execute_code.invoke(arguments={"code": "create-output"}) + assert registry._entries == {} + finally: + _close_execute_code_registry(execute_code) + + async def test_execute_code_tool_checks_output_count_before_reading( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -1290,12 +1430,15 @@ async def test_execute_code_tool_streams_directory_enumeration_to_count_limit( 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 + scanned_entries: list[int] = [] class _BoundedScandir: def __init__(self, path: str | os.PathLike[str]) -> None: self._entries = original_scandir(path) - self._track = Path(path) == output_root + self._scan_index: int | None = None + if Path(path) == output_root: + self._scan_index = len(scanned_entries) + scanned_entries.append(0) def __enter__(self) -> _BoundedScandir: self._entries.__enter__() @@ -1311,12 +1454,11 @@ 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") + if self._scan_index is not None: + scanned_entries[self._scan_index] += 1 + if scanned_entries[self._scan_index] > 3: + pytest.fail("an output directory scan continued past max_output_files + 1") return entry monkeypatch.setattr(execute_code_module.os, "scandir", _BoundedScandir) @@ -1327,10 +1469,46 @@ def __next__(self) -> os.DirEntry[str]: monkeypatch.setattr(execute_code_module.os, "scandir", original_scandir) _close_execute_code_registry(execute_code) - assert scanned_entries == 3 + assert scanned_entries == [0, 3, 3] _assert_bounded_output_error(contents, "output file count limit") +@pytest.mark.parametrize( + ("code", "error_match"), + [ + ("create-wide-output", "traversal entry limit of 100"), + ("create-deep-output", "nesting depth limit of 32"), + ], +) +async def test_execute_code_tool_invalidates_entry_when_cleanup_exceeds_bounds( + code: str, + error_match: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _FakeSandbox.instances.clear() + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandboxWithBoundedOutputs) + execute_code = HyperlightExecuteCodeTool(workspace_root=tmp_path, max_output_files=1) + config = execute_code._build_run_config() + registry = cast(Any, execute_code._registry) + registry._get_or_create_entry(config) + + try: + rejected = await execute_code.invoke(arguments={"code": code}) + assert registry._entries == {} + + recovered = await execute_code.invoke(arguments={"code": "create-memory-output"}) + replacement_output_root = Path(registry._get_or_create_entry(config).output_dir.name) + remaining_after_recovery = await asyncio.to_thread(lambda: list(replacement_output_root.iterdir())) + finally: + _close_execute_code_registry(execute_code) + + _assert_bounded_output_error(rejected, error_match) + assert len(_FakeSandbox.instances) == 2 + assert [_decode_content_bytes(item) for item in recovered if item.type == "data"] == [b"data"] + assert remaining_after_recovery == [] + + async def test_execute_code_tool_reads_only_observed_file_size_with_large_limits( tmp_path: Path, monkeypatch: pytest.MonkeyPatch,