feat(hf): add dedicated sandbox runtime mode

Author: zkwentz

docs/source/guides/runtime-providers.md

@@ -14,6 +14,7 @@ is a one-line change.
 | `UVProvider` | Local process via `uv` (no container) | core | ✅ |
 | `DaytonaProvider` | Daytona cloud sandboxes | `pip install openenv[daytona]` | ✅ |
 | `ACASandboxProvider` | Azure Container Apps Sandboxes | `pip install openenv[aca]` | ✅ |
+| `HFSandboxProvider` | Hugging Face Sandboxes | `pip install openenv[hf-sandbox]` | ✅ |
 | `ModalProvider` | Modal sandboxes | `pip install openenv[modal]` | ✅ |
 | `KubernetesProvider` | Kubernetes cluster | core | 🚧 planned |
 
@@ -159,6 +160,26 @@ from openenv.core.containers.runtime import DockerSwarmProvider
 provider = DockerSwarmProvider()
 ```
 
+### HFSandboxProvider
+
+Runs a registry image in a Hugging Face Sandbox. Pooled mode is the default and
+is appropriate for trusted same-user rollout fan-out. Dedicated mode allocates
+one VM per sandbox and is suited to untrusted or GPU workloads.
+
+```python
+from openenv.core.containers.runtime.hf_sandbox_provider import HFSandboxProvider
+
+provider = HFSandboxProvider(
+    image="hf.co/spaces/openenv/coding_env",
+    mode="dedicated",
+)
+base_url = provider.start_container()
+provider.wait_for_ready(base_url)
+```
+
+Plain `env_vars` are visible as Job environment metadata. Use the dedicated
+mode's `secrets=` argument for encrypted runtime secrets.
+
 ### KubernetesProvider
 
 🚧 Not yet implemented. The class exists as a placeholder for the planned

docs/source/reference/core.md

@@ -240,4 +240,6 @@ For a high-level explanation of how MCP-backed environments move through `step()
 
 [[autodoc]] openenv.core.containers.runtime.aca_provider.ACASandboxProvider
 
+[[autodoc]] openenv.core.containers.runtime.hf_sandbox_provider.HFSandboxProvider
+
 [[autodoc]] openenv.core.containers.runtime.modal_provider.ModalProvider

pyproject.toml

@@ -55,6 +55,9 @@ modal = [
 inspect = [
     "inspect-ai>=0.3.0",
 ]
+hf-sandbox = [
+    "huggingface_hub>=1.22.0",
+]
 
 [project.scripts]
 openenv = "openenv.cli.__main__:main"

src/openenv/core/containers/runtime/hf_sandbox_provider.py

@@ -10,11 +10,12 @@
 
 import asyncio
 import hashlib
+import re
 import socket
 import threading
 import time
 from contextlib import suppress
-from typing import Any
+from typing import Any, Literal
 
 import requests
 import uvicorn
@@ -27,12 +28,21 @@
 
 
 _DEFAULT_PORT = 8000
-_SERVER_COMMAND = (
+_POOLED_SERVER_COMMAND = (
     'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; '
     'mkdir -p "$SBX_PROXY_DIR"; '
-    'nohup server >"$HOME/openenv-server.log" 2>&1 &'
+    'exec server >"$HOME/openenv-server.log" 2>&1'
 )
+_DEDICATED_SERVER_COMMAND = (
+    'unset SBX_PROXY_DIR; exec server >"$HOME/openenv-server.log" 2>&1'
+)
+# Compatibility alias for callers that imported the old private constant.
+_SERVER_COMMAND = _POOLED_SERVER_COMMAND
 _MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024
+_MAX_HTTP_RESPONSE_SIZE = 100 * 1024 * 1024
+_CREDENTIAL_ENV_NAME = re.compile(
+    r"(?i)(authorization|credential|password|private.?key|secret|token|api.?key)"
+)
 _HOP_BY_HOP_HEADERS = {
     "connection",
     "content-encoding",
@@ -50,6 +60,11 @@
 _POOL_LOCK = threading.Lock()
 
 
+class _NoRedirectWebSocketConnect(ws_connect):
+    def process_redirect(self, exc: Exception) -> Exception | str:
+        return exc
+
+
 def _find_available_port() -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
         sock.bind(("127.0.0.1", 0))
@@ -81,6 +96,22 @@ def _get_sandbox_pool_cls() -> Any:
     return SandboxPool
 
 
+def _get_sandbox_cls() -> Any:
+    try:
+        from huggingface_hub import Sandbox
+    except ImportError as exc:
+        raise RuntimeError(
+            "Dedicated HFSandboxProvider execution requires huggingface_hub>=1.22.0. "
+            "Install it with `pip install 'openenv[hf-sandbox]'`."
+        ) from exc
+    if not hasattr(Sandbox, "create"):
+        raise RuntimeError(
+            "Dedicated HFSandboxProvider execution requires huggingface_hub>=1.22.0. "
+            "Install it with `pip install 'openenv[hf-sandbox]'`."
+        )
+    return Sandbox
+
+
 def _get_pool(image: str, flavor: str) -> Any:
     key = (image, flavor)
     with _POOL_LOCK:
@@ -133,23 +164,40 @@ async def proxy_http(path: str, request: Request) -> Response:
                     data=body,
                     headers=headers,
                     timeout=60.0,
-                    allow_redirects=True,
+                    allow_redirects=False,
+                    stream=True,
                 )
             except requests.RequestException:
                 return Response(
                     content=b"upstream HF job unreachable",
                     status_code=502,
                 )
-            response_headers = {
-                key: value
-                for key, value in upstream.headers.items()
-                if key.lower() not in _HOP_BY_HOP_HEADERS
-            }
-            return Response(
-                content=upstream.content,
-                status_code=upstream.status_code,
-                headers=response_headers,
-            )
+            try:
+                if 300 <= upstream.status_code < 400:
+                    return Response(
+                        content=b"upstream HF job redirects are not allowed",
+                        status_code=502,
+                    )
+                content = upstream.raw.read(
+                    _MAX_HTTP_RESPONSE_SIZE + 1, decode_content=True
+                )
+                if len(content) > _MAX_HTTP_RESPONSE_SIZE:
+                    return Response(
+                        content=b"upstream HF job response exceeded proxy limit",
+                        status_code=502,
+                    )
+                response_headers = {
+                    key: value
+                    for key, value in upstream.headers.items()
+                    if key.lower() not in _HOP_BY_HOP_HEADERS
+                }
+                return Response(
+                    content=content,
+                    status_code=upstream.status_code,
+                    headers=response_headers,
+                )
+            finally:
+                upstream.close()
 
         @app.websocket("/{path:path}")
         async def proxy_websocket(path: str, websocket: WebSocket) -> None:
@@ -202,7 +250,7 @@ async def _connect_upstream_websocket(
         last_error: Exception | None = None
         for _ in range(5):
             try:
-                return await ws_connect(
+                return await _NoRedirectWebSocketConnect(
                     target,
                     additional_headers=headers,
                     max_size=_MAX_WS_MESSAGE_SIZE,
@@ -243,14 +291,87 @@ def __init__(
         *,
         image: str,
         env_vars: dict[str, str] | None = None,
+        secrets: dict[str, str] | None = None,
         flavor: str = "cpu-basic",
+        mode: Literal["pooled", "dedicated"] = "pooled",
     ):
+        if mode not in {"pooled", "dedicated"}:
+            raise ValueError("HFSandboxProvider mode must be 'pooled' or 'dedicated'")
+        if mode == "pooled" and secrets:
+            raise ValueError("Encrypted secrets require dedicated HF sandbox mode")
         self.image = image
-        self.env_vars = env_vars
+        self._env_vars = dict(env_vars) if env_vars is not None else None
+        self._secrets = dict(secrets) if secrets is not None else None
         self.flavor = flavor
+        self._mode = mode
+        self._official_certification = False
+        self._official_image: str | None = None
+        self._official_flavor: str | None = None
         self._sandbox: Any = None
+        self._server_process: Any = None
         self._proxy: _LocalAuthProxy | None = None
 
+    @property
+    def mode(self) -> Literal["pooled", "dedicated"]:
+        return self._mode
+
+    @property
+    def env_vars(self) -> dict[str, str] | None:
+        return dict(self._env_vars) if self._env_vars is not None else None
+
+    @property
+    def secrets(self) -> dict[str, str] | None:
+        return dict(self._secrets) if self._secrets is not None else None
+
+    @property
+    def isolation_mode(self) -> Literal["pooled", "dedicated", "unknown"]:
+        if self._sandbox is None:
+            return self._mode
+        try:
+            host_id = self._sandbox.host_id
+        except (AttributeError, RuntimeError):
+            return "unknown"
+        return "dedicated" if host_id is None else "pooled"
+
+    def _lock_for_official_certification(self) -> None:
+        """Freeze the settings that cross the official runner trust boundary."""
+        if self._sandbox is not None:
+            raise RuntimeError(
+                "Official certification must validate and lock the provider "
+                "before a Hugging Face sandbox is started."
+            )
+        self._official_certification = True
+        self._official_image = self.image
+        self._official_flavor = self.flavor
+
+    def _verify_official_runtime_isolation(self) -> None:
+        try:
+            host_id = self._sandbox.host_id
+        except (AttributeError, RuntimeError) as exc:
+            raise RuntimeError(
+                "Official certification could not verify dedicated isolation "
+                "for the created Hugging Face sandbox."
+            ) from exc
+        if host_id is not None:
+            raise RuntimeError(
+                "Official certification requires a dedicated Hugging Face sandbox; "
+                "the created sandbox is attached to a shared host."
+            )
+
+    def _create_sandbox(self, *, image: str, env_vars: dict[str, str] | None) -> Any:
+        if self.mode == "dedicated":
+            Sandbox = _get_sandbox_cls()
+            create_kwargs: dict[str, Any] = {
+                "image": image,
+                "flavor": self.flavor,
+                "env": env_vars,
+            }
+            if self._secrets is not None:
+                create_kwargs["secrets"] = dict(self._secrets)
+            return Sandbox.create(**create_kwargs)
+        pool = _get_pool(image, self.flavor)
+        return pool.create(env=env_vars)
+
     def start_container(
         self,
         image: str | None = None,
@@ -271,47 +392,93 @@ def start_container(
                 f"HFSandboxProvider only supports port {_DEFAULT_PORT} (got {port})."
             )
 
+        if self._official_certification:
+            if (
+                self.image != self._official_image
+                or self.flavor != self._official_flavor
+                or (image is not None and image != self._official_image)
+            ):
+                raise RuntimeError(
+                    "Official certification does not allow image or flavor overrides after "
+                    "the provider trust check."
+                )
+            if env_vars is not None:
+                raise RuntimeError(
+                    "Official certification does not allow environment overrides "
+                    "after the provider trust check."
+                )
+            credential_names = sorted(
+                name
+                for name in self._env_vars or {}
+                if _CREDENTIAL_ENV_NAME.search(name)
+            )
+            if credential_names or self._secrets:
+                raise RuntimeError(
+                    "Official subject sandboxes cannot receive coordinator credentials."
+                )
+
         effective_image = self.image if image is None else image
-        effective_env = self.env_vars if env_vars is None else env_vars
-        pool = _get_pool(effective_image, self.flavor)
-        self._sandbox = pool.create(env=effective_env)
+        effective_env = self._env_vars if env_vars is None else env_vars
+        self._sandbox = self._create_sandbox(
+            image=effective_image, env_vars=effective_env
+        )
         try:
+            if self._official_certification:
+                self._verify_official_runtime_isolation()
             if not hasattr(self._sandbox, "proxy_url_for") or not hasattr(
                 self._sandbox, "proxy_headers"
             ):
                 raise RuntimeError(
                     "HFSandboxProvider requires a huggingface_hub version with "
                     "Sandbox.proxy_url_for and Sandbox.proxy_headers support."
                 )
-            self._sandbox.run(
-                _SERVER_COMMAND,
+            self._server_process = self._sandbox.run(
+                (
+                    _DEDICATED_SERVER_COMMAND
+                    if self.mode == "dedicated"
+                    else _POOLED_SERVER_COMMAND
+                ),
                 shell=True,
+                background=True,
             )
             self._proxy = _LocalAuthProxy(
                 target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"),
                 headers=self._sandbox.proxy_headers,
             )
             return self._proxy.start()
         except Exception:
-            self.stop_container()
+            with suppress(Exception):
+                self.stop_container()
             raise
 
     def stop_container(self) -> None:
+        proxy_error: Exception | None = None
         if self._proxy is not None:
-            self._proxy.stop()
-            self._proxy = None
+            try:
+                self._proxy.stop()
+            except Exception as exc:
+                proxy_error = exc
+            finally:
+                self._proxy = None
         if self._sandbox is not None:
             sandbox = self._sandbox
+            sandbox.kill()
+            if getattr(sandbox, "_killed", None) is not True:
+                raise RuntimeError(
+                    "Hugging Face did not confirm termination of the sandbox; "
+                    "the provider retained its handle so cleanup can be retried."
+                )
             self._sandbox = None
-            with suppress(Exception):
-                sandbox.kill()
+            self._server_process = None
+        if proxy_error is not None:
+            raise proxy_error
 
     def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
         deadline = time.time() + timeout_s
         health_url = f"{base_url}/health"
         while time.time() < deadline:
             try:
-                response = requests.get(health_url, timeout=5.0)
+                response = requests.get(health_url, timeout=5.0, allow_redirects=False)
                 if response.status_code == 200:
                     return
             except requests.exceptions.RequestException: