Expose filesystem limits in Python SDK (#214)

Signed-off-by: James Sturtevant <jsturtevant@gmail.com>

Author: jsturtevant

src/sdk/python/core/hyperlight_sandbox/__init__.py

@@ -10,7 +10,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from importlib import metadata
-from typing import Any
+from typing import Any, Literal
 
 from ._module_resolver import DEFAULT_MODULE_REF, resolve_module_path
 
@@ -83,7 +83,12 @@ def success(self) -> bool:
 
 @dataclass
 class SandboxEnvironment:
-    """Configuration for creating a sandbox."""
+    """Configuration for creating a sandbox.
+
+    Filesystem quotas use human-readable binary size strings such as ``"64Mi"``.
+    Leave all quota fields unset to use balanced backend defaults, or set
+    ``filesystem_limits="unlimited"`` to explicitly disable filesystem quotas.
+    """
 
     input_dir: str | None = None
     output_dir: str | None = None
@@ -93,10 +98,20 @@ class SandboxEnvironment:
     module_path: str | None = None
     heap_size: str = field(default_factory=lambda: _DEFAULT_HEAP_SIZE)
     stack_size: str = field(default_factory=lambda: _DEFAULT_STACK_SIZE)
+    filesystem_limits: Literal["unlimited"] | None = None
+    max_file_size: str | None = None
+    max_total_size: str | None = None
+    max_file_count: int | None = None
 
 
 class Sandbox:
-    """Stable Python API over swappable Hyperlight backends."""
+    """Stable Python API over swappable Hyperlight backends.
+
+    ``max_file_size`` and ``max_total_size`` accept binary size strings such as
+    ``"64Mi"``. Partial numeric overrides retain balanced defaults for omitted
+    values. Use ``filesystem_limits="unlimited"`` to explicitly disable all
+    filesystem quotas; it cannot be combined with numeric overrides.
+    """
 
     def __init__(
         self,
@@ -109,7 +124,30 @@ def __init__(
         module_path: str | None = None,
         heap_size: str | None = None,
         stack_size: str | None = None,
+        filesystem_limits: Literal["unlimited"] | None = None,
+        max_file_size: str | None = None,
+        max_total_size: str | None = None,
+        max_file_count: int | None = None,
     ) -> None:
+        numeric_limits = (max_file_size, max_total_size, max_file_count)
+        if filesystem_limits is not None and filesystem_limits != "unlimited":
+            raise ValueError("filesystem_limits must be 'unlimited' when provided")
+        if filesystem_limits == "unlimited" and any(value is not None for value in numeric_limits):
+            raise ValueError(
+                "filesystem_limits='unlimited' cannot be combined with max_file_size, max_total_size, or max_file_count"
+            )
+        for name, value in (
+            ("max_file_size", max_file_size),
+            ("max_total_size", max_total_size),
+        ):
+            if value is not None and not isinstance(value, str):
+                raise TypeError(f"{name} must be a size string")
+        if max_file_count is not None:
+            if isinstance(max_file_count, bool) or not isinstance(max_file_count, int):
+                raise TypeError("max_file_count must be an integer")
+            if max_file_count < 0:
+                raise ValueError("max_file_count must be non-negative")
+
         if heap_size is None:
             heap_size = _DEFAULT_HEAP_SIZE
         if stack_size is None:
@@ -129,6 +167,14 @@ def __init__(
             kwargs["output_dir"] = output_dir
         if temp_output:
             kwargs["temp_output"] = True
+        if filesystem_limits is not None:
+            kwargs["filesystem_limits"] = filesystem_limits
+        if max_file_size is not None:
+            kwargs["max_file_size"] = max_file_size
+        if max_total_size is not None:
+            kwargs["max_total_size"] = max_total_size
+        if max_file_count is not None:
+            kwargs["max_file_count"] = max_file_count
 
         if normalized_backend == "wasm":
             resolved_module_path = resolve_module_path(module=effective_module, module_path=module_path)
@@ -218,6 +264,10 @@ def _get_sandbox(self) -> Sandbox:
                 module_path=self.environment.module_path,
                 heap_size=self.environment.heap_size,
                 stack_size=self.environment.stack_size,
+                filesystem_limits=self.environment.filesystem_limits,
+                max_file_size=self.environment.max_file_size,
+                max_total_size=self.environment.max_total_size,
+                max_file_count=self.environment.max_file_count,
             )
             for tool_fn in self.tools:
                 self._sandbox.register_tool(tool_fn)

src/sdk/python/core/tests/test_filesystem_limits.py

@@ -0,0 +1,126 @@
+"""Tests for stable filesystem quota validation and forwarding."""
+
+import unittest
+from unittest.mock import patch
+
+from hyperlight_sandbox import CodeExecutionTool, Sandbox, SandboxEnvironment
+
+
+class _FakeNativeSandbox:
+    calls: list[dict] = []
+
+    def __init__(self, **kwargs):
+        self.calls.append(kwargs)
+
+    def register_tool(self, *args):
+        pass
+
+
+class TestFilesystemLimitForwarding(unittest.TestCase):
+    def setUp(self):
+        _FakeNativeSandbox.calls.clear()
+
+    def test_omitted_limits_leave_backend_defaults_unchanged(self):
+        with patch(
+            "hyperlight_sandbox._load_backend",
+            return_value=("hyperlight-js", _FakeNativeSandbox),
+        ):
+            Sandbox(backend="hyperlight-js")
+
+        kwargs = _FakeNativeSandbox.calls[-1]
+        self.assertNotIn("filesystem_limits", kwargs)
+        self.assertNotIn("max_file_size", kwargs)
+        self.assertNotIn("max_total_size", kwargs)
+        self.assertNotIn("max_file_count", kwargs)
+
+    def test_custom_limits_and_zero_are_forwarded(self):
+        with patch(
+            "hyperlight_sandbox._load_backend",
+            return_value=("hyperlight-js", _FakeNativeSandbox),
+        ):
+            Sandbox(
+                backend="hyperlight-js",
+                max_file_size="8Mi",
+                max_total_size="20Mi",
+                max_file_count=0,
+            )
+
+        self.assertEqual(_FakeNativeSandbox.calls[-1]["max_file_size"], "8Mi")
+        self.assertEqual(_FakeNativeSandbox.calls[-1]["max_total_size"], "20Mi")
+        self.assertEqual(_FakeNativeSandbox.calls[-1]["max_file_count"], 0)
+
+    def test_unlimited_is_forwarded(self):
+        with patch(
+            "hyperlight_sandbox._load_backend",
+            return_value=("hyperlight-js", _FakeNativeSandbox),
+        ):
+            Sandbox(backend="hyperlight-js", filesystem_limits="unlimited")
+
+        self.assertEqual(_FakeNativeSandbox.calls[-1]["filesystem_limits"], "unlimited")
+
+    def test_wasm_backend_forwards_limits(self):
+        with (
+            patch(
+                "hyperlight_sandbox._load_backend",
+                return_value=("wasm", _FakeNativeSandbox),
+            ),
+            patch("hyperlight_sandbox.resolve_module_path", return_value="/guest.wasm"),
+        ):
+            Sandbox(backend="wasm", max_total_size="4Mi")
+
+        kwargs = _FakeNativeSandbox.calls[-1]
+        self.assertEqual(kwargs["module_path"], "/guest.wasm")
+        self.assertEqual(kwargs["max_total_size"], "4Mi")
+
+    def test_environment_forwards_limits(self):
+        environment = SandboxEnvironment(max_file_size="1Ki", max_file_count=2)
+        tool = CodeExecutionTool(environment=environment)
+        with patch(
+            "hyperlight_sandbox._load_backend",
+            return_value=("hyperlight-js", _FakeNativeSandbox),
+        ):
+            tool._get_sandbox()
+
+        kwargs = _FakeNativeSandbox.calls[-1]
+        self.assertEqual(kwargs["max_file_size"], "1Ki")
+        self.assertEqual(kwargs["max_file_count"], 2)
+
+
+class TestFilesystemLimitValidation(unittest.TestCase):
+    def test_unlimited_rejects_numeric_overrides(self):
+        with self.assertRaisesRegex(ValueError, "cannot be combined"):
+            Sandbox(filesystem_limits="unlimited", max_file_count=0)
+
+    def test_unknown_mode_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, "must be 'unlimited'"):
+            Sandbox(filesystem_limits="default")
+
+    def test_negative_file_count_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, "non-negative"):
+            Sandbox(max_file_count=-1)
+
+    def test_file_count_rejects_bool(self):
+        with self.assertRaisesRegex(TypeError, "must be an integer"):
+            Sandbox(max_file_count=True)
+
+    def test_size_override_must_be_a_string(self):
+        with self.assertRaisesRegex(TypeError, "must be a size string"):
+            Sandbox(max_file_size=1024)  # type: ignore[arg-type]
+
+    def test_new_kwargs_are_not_silently_dropped_for_old_backends(self):
+        class _OldBackend:
+            def __init__(self, *, heap_size, stack_size):
+                pass
+
+        with (
+            patch(
+                "hyperlight_sandbox._load_backend",
+                return_value=("hyperlight-js", _OldBackend),
+            ),
+            self.assertRaisesRegex(TypeError, "max_file_count"),
+        ):
+            Sandbox(backend="hyperlight-js", max_file_count=1)
+
+
+if __name__ == "__main__":
+    unittest.main()

src/sdk/python/hyperlight_js_backend/src/lib.rs

@@ -2,11 +2,12 @@ use std::collections::HashMap;
 
 use hyperlight_javascript_sandbox::HyperlightJs;
 use hyperlight_sandbox::{
-    DEFAULT_HEAP_SIZE, DEFAULT_STACK_SIZE, DirPerms, FilePerms, HttpMethod, Sandbox,
-    SandboxBuilder, SandboxConfig,
+    DEFAULT_HEAP_SIZE, DEFAULT_STACK_SIZE, DirPerms, FilePerms, FilesystemLimits, HttpMethod,
+    Sandbox, SandboxBuilder, SandboxConfig,
 };
 use hyperlight_sandbox_pyo3_common::{
-    PyExecutionResult, build_tool_registry, parse_size, parse_tool_registration,
+    PyExecutionResult, build_tool_registry, parse_filesystem_limits, parse_size,
+    parse_tool_registration,
 };
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -30,19 +31,24 @@ pub struct JSSandbox {
     input_dir: Option<String>,
     output_dir: Option<String>,
     temp_output: bool,
+    filesystem_limits: Option<FilesystemLimits>,
 }
 
 #[pymethods]
 impl JSSandbox {
     #[new]
-    #[pyo3(signature = (input_dir=None, output_dir=None, temp_output=false, module_path="", heap_size=None, stack_size=None))]
+    #[pyo3(signature = (input_dir=None, output_dir=None, temp_output=false, module_path="", heap_size=None, stack_size=None, filesystem_limits=None, max_file_size=None, max_total_size=None, max_file_count=None))]
     fn new(
         input_dir: Option<&str>,
         output_dir: Option<&str>,
         temp_output: bool,
         module_path: &str,
         heap_size: Option<&str>,
         stack_size: Option<&str>,
+        filesystem_limits: Option<&str>,
+        max_file_size: Option<&str>,
+        max_total_size: Option<&str>,
+        max_file_count: Option<usize>,
     ) -> PyResult<Self> {
         if !module_path.is_empty() {
             return Err(PyRuntimeError::new_err(
@@ -68,6 +74,12 @@ impl JSSandbox {
             input_dir: input_dir.map(|s| s.to_string()),
             output_dir: output_dir.map(|s| s.to_string()),
             temp_output,
+            filesystem_limits: parse_filesystem_limits(
+                filesystem_limits,
+                max_file_size,
+                max_total_size,
+                max_file_count,
+            )?,
         })
     }
 
@@ -99,6 +111,9 @@ impl JSSandbox {
                 .stack_size(self.config.stack_size)
                 .with_tools(registry)
                 .guest(HyperlightJs);
+            if let Some(limits) = self.filesystem_limits {
+                builder = builder.filesystem_limits(limits);
+            }
             if let Some(ref dir) = self.input_dir {
                 builder = builder.input_dir(dir);
             }

src/sdk/python/pyo3_common/src/lib.rs

@@ -1,10 +1,14 @@
-use hyperlight_sandbox::{ArgType, ToolRegistry, ToolSchema};
-use pyo3::exceptions::{PyRuntimeError, PyTypeError};
+use hyperlight_sandbox::{ArgType, FilesystemLimits, ToolRegistry, ToolSchema};
+use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyModule};
 
 /// Convert a human-readable size string (e.g. `"200Mi"`) to bytes.
 pub fn parse_size(size: &str) -> PyResult<u64> {
+    parse_size_value(size).map_err(PyRuntimeError::new_err)
+}
+
+fn parse_size_value(size: &str) -> Result<u64, String> {
     let size = size.trim();
     let (value, multiplier) = if let Some(value) = size.strip_suffix("Gi") {
         (value, 1024u64.pow(3))
@@ -15,12 +19,65 @@ pub fn parse_size(size: &str) -> PyResult<u64> {
     } else {
         (size, 1)
     };
-    let parsed: u64 = value
-        .parse()
-        .map_err(|e| PyRuntimeError::new_err(format!("invalid size: {e}")))?;
+    let parsed: u64 = value.parse().map_err(|e| format!("invalid size: {e}"))?;
     parsed
         .checked_mul(multiplier)
-        .ok_or_else(|| PyRuntimeError::new_err("invalid size: value is too large"))
+        .ok_or_else(|| "invalid size: value is too large".to_string())
+}
+
+/// Parse the Python filesystem quota arguments into an optional builder override.
+///
+/// Returning `None` preserves the builder's balanced default limits.
+pub fn parse_filesystem_limits(
+    filesystem_limits: Option<&str>,
+    max_file_size: Option<&str>,
+    max_total_size: Option<&str>,
+    max_file_count: Option<usize>,
+) -> PyResult<Option<FilesystemLimits>> {
+    let has_numeric_override =
+        max_file_size.is_some() || max_total_size.is_some() || max_file_count.is_some();
+
+    if let Some(mode) = filesystem_limits {
+        if mode != "unlimited" {
+            return Err(PyValueError::new_err(
+                "filesystem_limits must be 'unlimited' when provided",
+            ));
+        }
+        if has_numeric_override {
+            return Err(PyValueError::new_err(
+                "filesystem_limits='unlimited' cannot be combined with max_file_size, \
+                 max_total_size, or max_file_count",
+            ));
+        }
+        return Ok(Some(FilesystemLimits::unlimited()));
+    }
+
+    if !has_numeric_override {
+        return Ok(None);
+    }
+
+    let defaults = FilesystemLimits::default();
+    let max_file_size = match max_file_size {
+        Some(size) => parse_size_value(size).map_err(PyValueError::new_err)?,
+        None => defaults
+            .max_file_size()
+            .expect("default file size is limited"),
+    };
+    let max_total_size = match max_total_size {
+        Some(size) => parse_size_value(size).map_err(PyValueError::new_err)?,
+        None => defaults
+            .max_total_size()
+            .expect("default total size is limited"),
+    };
+    let max_file_count = max_file_count.unwrap_or_else(|| {
+        defaults
+            .max_file_count()
+            .expect("default file count is limited")
+    });
+
+    FilesystemLimits::new(max_file_size, max_total_size, max_file_count)
+        .map(Some)
+        .map_err(|e| PyValueError::new_err(format!("invalid filesystem limits: {e}")))
 }
 
 /// Wrap an SDK-style tool object (with `.handler` and `.name`) into a plain `**kwargs` callable.
@@ -456,3 +513,18 @@ pub fn parse_tool_registration(
         Ok((name, cb))
     }
 }
+
+#[cfg(test)]
+mod filesystem_limit_tests {
+    use super::*;
+
+    #[test]
+    fn malformed_size_is_rejected() {
+        assert!(parse_size_value("invalid").is_err());
+    }
+
+    #[test]
+    fn overflowing_size_is_rejected() {
+        assert!(parse_size_value("18446744073709551615Gi").is_err());
+    }
+}