Python: Clean Hyperlight output files after execution

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,14 +156,15 @@ 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)
         # _sandbox/_snapshot are accessed/mutated ONLY from worker-side closures.
         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()