Expose cross-thread guest interruption in Python SDK

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

Author: jsturtevant

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

@@ -62,6 +62,7 @@ def _load_backend(backend: str):
 __all__ = [
     "CodeExecutionTool",
     "ExecutionResult",
+    "InterruptHandle",
     "Sandbox",
     "SandboxEnvironment",
     "__version__",
@@ -104,6 +105,24 @@ class SandboxEnvironment:
     max_file_count: int | None = None
 
 
+class InterruptHandle:
+    """An independently owned handle that can interrupt a guest from any thread."""
+
+    def __init__(self, inner: Any) -> None:
+        self._inner = inner
+
+    def kill(self) -> bool:
+        """Interrupt a running guest and wait for it to stop.
+
+        Returns whether the vCPU was running. False during a host callback can
+        still leave cancellation pending until that callback returns.
+        A request just before guest entry can be missed; callers can retry.
+        Host callbacks are not interrupted. Restore a snapshot before reusing
+        an interrupted sandbox. The handle can outlive its sandbox.
+        """
+        return self._inner.kill()
+
+
 class Sandbox:
     """Stable Python API over swappable Hyperlight backends.
 
@@ -199,6 +218,7 @@ def register_tool(self, name_or_tool: Any, callback: Any | None = None) -> None:
     MAX_CODE_SIZE: int = 10 * 1024 * 1024
 
     def run(self, code: str) -> ExecutionResult:
+        """Execute code, raising InterruptedError if a handle interrupts the guest."""
         if len(code) > self.MAX_CODE_SIZE:
             raise ValueError(f"code exceeds maximum size ({len(code)} > {self.MAX_CODE_SIZE} bytes)")
         native_result = self._inner.run(code)
@@ -208,6 +228,14 @@ def run(self, code: str) -> ExecutionResult:
             exit_code=native_result.exit_code,
         )
 
+    def interrupt_handle(self) -> InterruptHandle:
+        """Initialize the runtime if needed and return a thread-safe interrupt handle.
+
+        Register tools before obtaining the handle. Use the sandbox itself only
+        on its creating thread; the independent handle can be shared.
+        """
+        return InterruptHandle(self._inner.interrupt_handle())
+
     def get_output_files(self) -> list[str]:
         """List filenames written by the guest to the output directory.
 

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

@@ -0,0 +1,131 @@
+"""Cross-thread guest interruption for both Python SDK backends."""
+
+import threading
+import time
+import unittest
+
+from hyperlight_sandbox import Sandbox
+
+
+class TestInterruption(unittest.TestCase):
+    def test_idle_handle_and_handle_after_sandbox_drop(self):
+        for backend in ("wasm", "hyperlight-js"):
+            with self.subTest(backend=backend):
+                sandbox = Sandbox(backend=backend)
+                handle = sandbox.interrupt_handle()
+                self.assertFalse(handle.kill())
+                del sandbox
+                self.assertFalse(handle.kill())
+
+    def test_interrupt_from_another_thread_and_restore(self):
+        for backend in ("wasm", "hyperlight-js"):
+            with self.subTest(backend=backend):
+                self._interrupt_and_restore(backend)
+
+    def _interrupt_and_restore(self, backend):
+        sandbox = Sandbox(backend=backend)
+        started = threading.Event()
+        finished = threading.Event()
+        attempts = []
+
+        def ready() -> bool:
+            started.set()
+            return True
+
+        sandbox.register_tool("ready", ready)
+        handle = sandbox.interrupt_handle()
+        warmup = "print('ready')" if backend == "wasm" else "console.log('ready');"
+        self.assertTrue(sandbox.run(warmup).success)
+        snapshot = sandbox.snapshot()
+
+        def interrupt():
+            if not started.wait(5):
+                return
+            if finished.wait(0.05):
+                return
+            while not finished.is_set():
+                attempts.append(handle.kill())
+                if attempts[-1]:
+                    return
+                finished.wait(0.01)
+
+        worker = threading.Thread(target=interrupt, daemon=True)
+        worker.start()
+        code = (
+            "import time\ncall_tool('ready')\nend = time.monotonic() + 5\n"
+            "while time.monotonic() < end:\n"
+            "    for counter in range(100_000):\n        value = counter * counter"
+            if backend == "wasm"
+            else "call_tool('ready'); const end = Date.now() + 5000; let value = 0;"
+            "while (Date.now() < end) {"
+            "for (let counter = 0; counter < 100000; counter++) { value = counter * counter; }}"
+        )
+        start = time.monotonic()
+        try:
+            with self.assertRaises(InterruptedError):
+                sandbox.run(code)
+        finally:
+            finished.set()
+            worker.join(5)
+        self.assertFalse(worker.is_alive())
+        self.assertIn(True, attempts)
+        self.assertLess(time.monotonic() - start, 3)
+        with self.assertRaises(RuntimeError):
+            sandbox.run(warmup)
+        sandbox.restore(snapshot)
+        self.assertTrue(sandbox.run(warmup).success)
+
+    def test_interruption_does_not_abort_host_callback(self):
+        for backend in ("wasm", "hyperlight-js"):
+            with self.subTest(backend=backend):
+                self._interrupt_host_callback(backend)
+
+    def _interrupt_host_callback(self, backend):
+        sandbox = Sandbox(backend=backend)
+        entered = threading.Event()
+        release = threading.Event()
+        finished = threading.Event()
+        callback_completed = []
+        attempts = []
+
+        def blocked() -> bool:
+            entered.set()
+            released = release.wait(5)
+            callback_completed.append(released)
+            return released
+
+        sandbox.register_tool("blocked", blocked)
+        handle = sandbox.interrupt_handle()
+
+        def interrupt():
+            try:
+                if entered.wait(5):
+                    attempts.append(handle.kill())
+            finally:
+                release.set()
+            while not finished.is_set():
+                if handle.kill():
+                    return
+                finished.wait(0.01)
+
+        worker = threading.Thread(target=interrupt, daemon=True)
+        worker.start()
+        code = (
+            "import time\ncall_tool('blocked')\nend = time.monotonic() + 5\nwhile time.monotonic() < end:\n    pass"
+            if backend == "wasm"
+            else "call_tool('blocked'); const end = Date.now() + 5000; while (Date.now() < end) {}"
+        )
+        try:
+            with self.assertRaises(InterruptedError):
+                sandbox.run(code)
+        finally:
+            finished.set()
+            release.set()
+            worker.join(5)
+        self.assertFalse(worker.is_alive())
+        self.assertEqual(attempts, [False])
+        self.assertEqual(callback_completed, [True])
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -6,8 +6,8 @@ use hyperlight_sandbox::{
     MethodFilter, Sandbox, SandboxBuilder, SandboxConfig,
 };
 use hyperlight_sandbox_pyo3_common::{
-    PyExecutionResult, build_tool_registry, parse_filesystem_limits, parse_size,
-    parse_tool_registration,
+    PyExecutionResult, PyInterruptHandle, build_tool_registry, execution_error,
+    parse_filesystem_limits, parse_size, parse_tool_registration,
 };
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -34,6 +34,45 @@ pub struct JSSandbox {
     filesystem_limits: Option<FilesystemLimits>,
 }
 
+impl JSSandbox {
+    fn ensure_initialized(&mut self, py: Python<'_>) -> PyResult<()> {
+        if self.inner.is_none() {
+            let registry = build_tool_registry(py, &mut self.tools)?;
+            let mut builder = SandboxBuilder::new()
+                .module_path(&self.config.module_path)
+                .heap_size(self.config.heap_size)
+                .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);
+            }
+            if let Some(ref dir) = self.output_dir {
+                builder = builder.output_dir(
+                    dir,
+                    DirPerms::READ | DirPerms::MUTATE,
+                    FilePerms::READ | FilePerms::WRITE,
+                );
+            } else if self.temp_output {
+                builder = builder.temp_output();
+            }
+            let mut sandbox = builder
+                .build()
+                .map_err(|e| PyRuntimeError::new_err(format!("Failed to create sandbox: {e:#}")))?;
+            for (target, methods) in std::mem::take(&mut self.pending_networks) {
+                sandbox
+                    .allow_domain(&target, methods)
+                    .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
+            }
+            self.inner = Some(sandbox);
+        }
+        Ok(())
+    }
+}
+
 #[pymethods]
 impl JSSandbox {
     #[new]
@@ -101,45 +140,24 @@ impl JSSandbox {
         Ok(())
     }
 
+    fn interrupt_handle(&mut self, py: Python<'_>) -> PyResult<PyInterruptHandle> {
+        self.ensure_initialized(py)?;
+        let inner = self
+            .inner
+            .as_ref()
+            .unwrap()
+            .interrupt_handle()
+            .map_err(|error| {
+                PyRuntimeError::new_err(format!("Failed to get interrupt handle: {error:#}"))
+            })?;
+        Ok(PyInterruptHandle { inner })
+    }
+
     #[pyo3(signature = (code))]
-    fn run(&mut self, py: Python<'_>, code: &str) -> PyResult<PyExecutionResult> {
-        if self.inner.is_none() {
-            let registry = build_tool_registry(py, &mut self.tools)?;
-            let mut builder = SandboxBuilder::new()
-                .module_path(&self.config.module_path)
-                .heap_size(self.config.heap_size)
-                .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);
-            }
-            if let Some(ref dir) = self.output_dir {
-                builder = builder.output_dir(
-                    dir,
-                    DirPerms::READ | DirPerms::MUTATE,
-                    FilePerms::READ | FilePerms::WRITE,
-                );
-            } else if self.temp_output {
-                builder = builder.temp_output();
-            }
-            let mut sandbox = builder
-                .build()
-                .map_err(|e| PyRuntimeError::new_err(format!("Failed to create sandbox: {e:#}")))?;
-            for (target, methods) in std::mem::take(&mut self.pending_networks) {
-                sandbox
-                    .allow_domain(&target, methods)
-                    .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
-            }
-            self.inner = Some(sandbox);
-        }
+    fn run(&mut self, py: Python<'_>, code: String) -> PyResult<PyExecutionResult> {
+        self.ensure_initialized(py)?;
         let sandbox = self.inner.as_mut().unwrap();
-        let result = sandbox
-            .run(code)
-            .map_err(|e| PyRuntimeError::new_err(format!("Execution failed: {e}")))?;
+        let result = py.detach(|| sandbox.run(&code)).map_err(execution_error)?;
         Ok(PyExecutionResult {
             stdout: result.stdout,
             stderr: result.stderr,
@@ -210,6 +228,7 @@ fn _native_js(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<JSSandbox>()?;
     m.add_class::<PyExecutionResult>()?;
     m.add_class::<PySnapshot>()?;
+    m.add_class::<PyInterruptHandle>()?;
     m.add("__version__", "0.1.0")?;
     Ok(())
 }

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

@@ -1,8 +1,33 @@
-use hyperlight_sandbox::{ArgType, FilesystemLimits, ToolRegistry, ToolSchema};
-use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
+use hyperlight_sandbox::{
+    ArgType, FilesystemLimits, HyperlightError, InterruptHandle, ToolRegistry, ToolSchema,
+};
+use pyo3::exceptions::{PyInterruptedError, PyRuntimeError, PyTypeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyModule};
 
+#[pyclass(frozen, name = "InterruptHandle")]
+pub struct PyInterruptHandle {
+    pub inner: InterruptHandle,
+}
+
+#[pymethods]
+impl PyInterruptHandle {
+    fn kill(&self, py: Python<'_>) -> bool {
+        py.detach(|| self.inner.kill())
+    }
+}
+
+pub fn execution_error(error: anyhow::Error) -> PyErr {
+    if matches!(
+        error.downcast_ref::<HyperlightError>(),
+        Some(HyperlightError::ExecutionCanceledByHost())
+    ) {
+        PyInterruptedError::new_err(format!("Execution interrupted: {error:#}"))
+    } else {
+        PyRuntimeError::new_err(format!("Execution failed: {error:#}"))
+    }
+}
+
 /// 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)