kaxil commented on code in PR #71676: URL: https://github.com/apache/airflow/pull/71676#discussion_r4042195887
########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), Review Comment: The fail-closed guarantee here is the server's rather than ours. The SDK only forwards `networkPolicy` on the create request and documents nothing about refusing a policy it cannot enforce, so on a deployment without the egress sidecar this is a successful create of a sandbox with open egress, which is exactly the case `SandboxSpec` says a backend must raise on rather than provision something weaker ([base.py L85-88](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L85-L88)). The recording shows the matrix holding on a server that has the sidecar; did you get to try one that doesn't? `SandboxSync.get_egress_policy()` reads the enforced policy back off the sidecar, so a read-back after create (deny default, rules present) would make this ours to guarantee instead of a property of whichever server version someone deploys. Same for the docs sentence at `toolsets.rst` L842, which states the rejection as fact. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), + connection_config=self._get_connection_config(), + ) + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + def _get_sandbox(self, sandbox_id: str) -> SandboxSync: + if sandbox := self._sandboxes.get(sandbox_id): + return sandbox + with _translate_opensandbox_errors("connect to a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.connect( + sandbox_id, + connection_config=self._get_connection_config(), + connect_timeout=timedelta(seconds=self._ready_timeout), + ) + self._sandboxes[sandbox_id] = sandbox + return sandbox + + def run_command( + self, sandbox: str, command: str, *, timeout: float, max_output_bytes: int + ) -> SandboxExecResult: + _validate_positive_finite(timeout, "timeout") + _validate_positive_finite(max_output_bytes, "max_output_bytes") + stdout = _BoundedTail(max_output_bytes) + stderr = _BoundedTail(max_output_bytes) + started = time.monotonic() + with _translate_opensandbox_errors("run a sandbox command"): + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.execd_sync import ExecutionHandlersSync + + execution = self._get_sandbox(sandbox).commands.run( Review Comment: Nothing bounds this call on our side. `RunCommandOpts(timeout=...)` is a server-side instruction, and the SDK streams over an SSE client built with `read=None` and a comment saying so ("SSE client (read timeout disabled)", `opensandbox/sync/adapters/command_adapter.py` L159-176), so `request_timeout` only covers connect and write. If execd stops sending events the call never returns, and the toolset awaits it in a thread it can't cancel, so a model-requested 60s command can hold a worker slot until the sandbox's own TTL (`sandbox_timeout`, default 3600) or indefinitely on a blackholed connection. `SbxSandboxBackend` guards exactly this: it caps the subprocess at `timeout + _EXEC_GRACE` and returns `timed_out=True, sandbox_terminated=True` ([sbx.py L374-394](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py#L374-L394)). `sandbox_terminated` is never set anywhere in this backend, so the toolset's "provision a fresh one" recovery is dead here. Separately, `max_output_bytes` can't actually bound worker memory the way [base.py L180-183](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L180-L183) says it should. `_BoundedTail` holds to the cap for what it's handed, but the SSE layer materialises a whole output line before handing it over (the SDK's frame normaliser accumulates into an unbounded `bytearray` until a `\n\n` boundary, `opensandbox/adapters/sse.py` L36-60). One newline-free line, which model-written code can produce with `head -c 2G /dev/zero`, is fully resident first. Not fixable here without patching the SDK, but worth a docstring note since `sbx` is immune (64 KiB pipe reads, trimmed as it goes). ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( Review Comment: Nothing tags the sandbox as Airflow's. `SandboxSync.create` takes `metadata: dict[str, str]`, `SandboxInfo` returns it and `SandboxFilter` can filter on it, so something like `{"created-by": "airflow"}` would let an operator on a shared server find these. `SbxSandboxBackend` already does this through `_new_sandbox_name()`'s `airflow-sandbox-` prefix, which [base.py L75-77](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L75-L77) describes as being there for correlation and cleanup. The server TTL covers reclamation, but not attribution. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), + connection_config=self._get_connection_config(), + ) + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + def _get_sandbox(self, sandbox_id: str) -> SandboxSync: + if sandbox := self._sandboxes.get(sandbox_id): + return sandbox + with _translate_opensandbox_errors("connect to a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.connect( + sandbox_id, + connection_config=self._get_connection_config(), + connect_timeout=timedelta(seconds=self._ready_timeout), + ) + self._sandboxes[sandbox_id] = sandbox + return sandbox + + def run_command( + self, sandbox: str, command: str, *, timeout: float, max_output_bytes: int + ) -> SandboxExecResult: + _validate_positive_finite(timeout, "timeout") + _validate_positive_finite(max_output_bytes, "max_output_bytes") + stdout = _BoundedTail(max_output_bytes) + stderr = _BoundedTail(max_output_bytes) + started = time.monotonic() + with _translate_opensandbox_errors("run a sandbox command"): + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.execd_sync import ExecutionHandlersSync + + execution = self._get_sandbox(sandbox).commands.run( + command, + opts=RunCommandOpts(timeout=timedelta(seconds=timeout)), + handlers=ExecutionHandlersSync( + on_stdout=stdout.add_message, + on_stderr=stderr.add_message, + skip_accumulation=True, + ), + ) + + if execution.exit_code is None: + raise SandboxTerminalError("OpenSandbox returned no terminal status for the command.") + if execution.error is not None and not stderr.get_text(): + details = "\n".join(execution.error.traceback) or execution.error.value + stderr.add_text(details) + timed_out = execution.exit_code < 0 and time.monotonic() - started >= timeout + return SandboxExecResult( + exit_code=execution.exit_code, + stdout=stdout.get_text(), + stderr=stderr.get_text(), + timed_out=timed_out, + stdout_truncated=stdout.truncated, + stderr_truncated=stderr.truncated, + ) + + @staticmethod + def _confirm_sandbox_exists(sandbox: SandboxSync) -> None: + with _translate_opensandbox_errors("confirm that a sandbox still exists"): + sandbox.get_info() + + def read_file(self, sandbox: str, path: str, *, max_bytes: int) -> bytes: + _validate_positive_finite(max_bytes, "max_bytes") + sandbox_client = self._get_sandbox(sandbox) + chunks = None + data = bytearray() + try: + chunks = sandbox_client.files.read_bytes_stream( + path, + chunk_size=min(65536, max_bytes + 1), + range_header=f"bytes=0-{max_bytes}", + ) + for chunk in chunks: + data.extend(chunk[: max_bytes + 1 - len(data)]) + if len(data) > max_bytes: + raise SandboxFileTooLargeError(path, len(data), max_bytes) + except SandboxFileTooLargeError: + raise + except Exception as e: + if _get_status_code(e) == 404: + self._confirm_sandbox_exists(sandbox_client) + raise SandboxError(f"{path!r} does not exist in the sandbox, or is not readable.") from e + with _translate_opensandbox_errors("read a sandbox file", recoverable_statuses=frozenset({400})): + raise + finally: + close = getattr(chunks, "close", None) + if close is not None: + with suppress(Exception): + close() + return bytes(data) + + def write_file(self, sandbox: str, path: str, content: bytes) -> None: + sandbox_client = self._get_sandbox(sandbox) + try: + from opensandbox.models.filesystem import WriteEntry + + parent = posixpath.dirname(path) + if parent and parent != "/": + sandbox_client.files.create_directories([WriteEntry(path=parent, mode=755)]) + sandbox_client.files.write_file(path, content, mode=644) + except Exception as e: + if _get_status_code(e) == 404: + self._confirm_sandbox_exists(sandbox_client) + raise SandboxError(f"Could not write {path!r} in the sandbox.") from e + with _translate_opensandbox_errors("write a sandbox file", recoverable_statuses=frozenset({400})): + raise + + def list_directory(self, sandbox: str, path: str) -> list[tuple[str, bool]]: + sandbox_client = self._get_sandbox(sandbox) + try: + from opensandbox.models.filesystem import DirectoryListEntry + + entries = sandbox_client.files.list_directory(DirectoryListEntry(path=path, depth=1)) Review Comment: This override drops the cap the base class puts on guest-controlled output. The base runs `find` with `max_output_bytes=_FILE_OP_OUTPUT_CAP`, and [base.py L65-67](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L65-L67) gives the reason: "a small cap bounds what a hostile guest can push into worker memory". The native call has no equivalent and the SDK parses the whole JSON array into `FileInfo` objects eagerly, so a directory the model just filled with a few hundred thousand entries is fully resident before the toolset's `truncate_output` ever runs. `DirectoryListEntry` has no limit field, so this may need a cap on the returned list plus a note to the model that it was truncated. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), + connection_config=self._get_connection_config(), + ) + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + def _get_sandbox(self, sandbox_id: str) -> SandboxSync: + if sandbox := self._sandboxes.get(sandbox_id): + return sandbox + with _translate_opensandbox_errors("connect to a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.connect( + sandbox_id, + connection_config=self._get_connection_config(), + connect_timeout=timedelta(seconds=self._ready_timeout), + ) + self._sandboxes[sandbox_id] = sandbox + return sandbox + + def run_command( + self, sandbox: str, command: str, *, timeout: float, max_output_bytes: int + ) -> SandboxExecResult: + _validate_positive_finite(timeout, "timeout") + _validate_positive_finite(max_output_bytes, "max_output_bytes") + stdout = _BoundedTail(max_output_bytes) + stderr = _BoundedTail(max_output_bytes) + started = time.monotonic() + with _translate_opensandbox_errors("run a sandbox command"): + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.execd_sync import ExecutionHandlersSync + + execution = self._get_sandbox(sandbox).commands.run( + command, + opts=RunCommandOpts(timeout=timedelta(seconds=timeout)), + handlers=ExecutionHandlersSync( + on_stdout=stdout.add_message, + on_stderr=stderr.add_message, + skip_accumulation=True, + ), + ) + + if execution.exit_code is None: Review Comment: `exit_code` isn't reported by the server on this path. For every foreground run the SDK synthesises it with `_infer_foreground_exit_code`, which returns `int(execution.error.value)`, `0`, or `None` (`opensandbox/sync/adapters/command_adapter.py` L68-76 and L215-216, reached unconditionally here because `background` is never set). `error.value` is free text, documented as "Error message explaining what went wrong". So `None` means "the SDK couldn't parse an integer out of a free-text field", not "the sandbox is unusable", which is what [base.py L43-50](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L43-L50) reserves terminal for. I ran the inference directly: `'1'` gives 1, `'-9'` gives -9, and `'signal: killed'`, `'exit status 1'` and no-events-at-all all give `None`. Any execd build that words `error.value` as prose fails the task on an ordinary non-zero exit. The timeout branch on L286 has the same root: it needs `error.value` to be a *negative integer string*, so `'124'` leaves `timed_out` False and the model sees `[exit code: 124]`, and prose lands on the raise above instead of the `[timed out after Ns]` the toolset documents as normal tool output. The fixtures can't catch it because they mock `commands.run` past the inference: `_execution(exit_code=-9)` leaves `error=None`, a pair the SDK never returns, and `test_execution_error_is_returned_on_stderr` sets `exit_code=1` with `error.value="failed"`, where the real `exit_code` would be `None` and this raise would fire before the stderr fallback that test is covering. Did a real timeout and a signal kill get exercised against the server? ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host Review Comment: Worth rejecting an empty host? `request_timeout` and `use_server_proxy` just below both get strict validation, but a connection with no host leaves `domain` as `None`, and `ConnectionConfigSync` then falls back to `OPEN_SANDBOX_DOMAIN` and finally to `localhost:8080` (`opensandbox/config/connection_sync.py` L107), so a half-filled connection quietly points the worker at itself. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), + connection_config=self._get_connection_config(), + ) + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + def _get_sandbox(self, sandbox_id: str) -> SandboxSync: + if sandbox := self._sandboxes.get(sandbox_id): + return sandbox + with _translate_opensandbox_errors("connect to a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.connect( + sandbox_id, + connection_config=self._get_connection_config(), + connect_timeout=timedelta(seconds=self._ready_timeout), + ) + self._sandboxes[sandbox_id] = sandbox + return sandbox + + def run_command( + self, sandbox: str, command: str, *, timeout: float, max_output_bytes: int + ) -> SandboxExecResult: + _validate_positive_finite(timeout, "timeout") + _validate_positive_finite(max_output_bytes, "max_output_bytes") + stdout = _BoundedTail(max_output_bytes) + stderr = _BoundedTail(max_output_bytes) + started = time.monotonic() + with _translate_opensandbox_errors("run a sandbox command"): + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.execd_sync import ExecutionHandlersSync + + execution = self._get_sandbox(sandbox).commands.run( + command, + opts=RunCommandOpts(timeout=timedelta(seconds=timeout)), + handlers=ExecutionHandlersSync( + on_stdout=stdout.add_message, + on_stderr=stderr.add_message, + skip_accumulation=True, + ), + ) + + if execution.exit_code is None: + raise SandboxTerminalError("OpenSandbox returned no terminal status for the command.") + if execution.error is not None and not stderr.get_text(): + details = "\n".join(execution.error.traceback) or execution.error.value + stderr.add_text(details) + timed_out = execution.exit_code < 0 and time.monotonic() - started >= timeout + return SandboxExecResult( + exit_code=execution.exit_code, + stdout=stdout.get_text(), + stderr=stderr.get_text(), + timed_out=timed_out, + stdout_truncated=stdout.truncated, + stderr_truncated=stderr.truncated, + ) + + @staticmethod + def _confirm_sandbox_exists(sandbox: SandboxSync) -> None: + with _translate_opensandbox_errors("confirm that a sandbox still exists"): + sandbox.get_info() + + def read_file(self, sandbox: str, path: str, *, max_bytes: int) -> bytes: + _validate_positive_finite(max_bytes, "max_bytes") + sandbox_client = self._get_sandbox(sandbox) + chunks = None + data = bytearray() + try: + chunks = sandbox_client.files.read_bytes_stream( + path, + chunk_size=min(65536, max_bytes + 1), + range_header=f"bytes=0-{max_bytes}", + ) + for chunk in chunks: + data.extend(chunk[: max_bytes + 1 - len(data)]) + if len(data) > max_bytes: + raise SandboxFileTooLargeError(path, len(data), max_bytes) + except SandboxFileTooLargeError: + raise + except Exception as e: + if _get_status_code(e) == 404: + self._confirm_sandbox_exists(sandbox_client) + raise SandboxError(f"{path!r} does not exist in the sandbox, or is not readable.") from e + with _translate_opensandbox_errors("read a sandbox file", recoverable_statuses=frozenset({400})): + raise + finally: + close = getattr(chunks, "close", None) + if close is not None: + with suppress(Exception): + close() + return bytes(data) + + def write_file(self, sandbox: str, path: str, content: bytes) -> None: + sandbox_client = self._get_sandbox(sandbox) + try: + from opensandbox.models.filesystem import WriteEntry + + parent = posixpath.dirname(path) + if parent and parent != "/": + sandbox_client.files.create_directories([WriteEntry(path=parent, mode=755)]) + sandbox_client.files.write_file(path, content, mode=644) + except Exception as e: + if _get_status_code(e) == 404: + self._confirm_sandbox_exists(sandbox_client) + raise SandboxError(f"Could not write {path!r} in the sandbox.") from e + with _translate_opensandbox_errors("write a sandbox file", recoverable_statuses=frozenset({400})): + raise + + def list_directory(self, sandbox: str, path: str) -> list[tuple[str, bool]]: + sandbox_client = self._get_sandbox(sandbox) + try: + from opensandbox.models.filesystem import DirectoryListEntry + + entries = sandbox_client.files.list_directory(DirectoryListEntry(path=path, depth=1)) + except Exception as e: + if _get_status_code(e) == 404: + self._confirm_sandbox_exists(sandbox_client) + raise SandboxError(f"{path!r} does not exist in the sandbox, or is not readable.") from e + with _translate_opensandbox_errors( + "list a sandbox directory", recoverable_statuses=frozenset({400}) + ): + raise + return [ + (posixpath.basename(entry.path.rstrip("/")), entry.entry_type == "directory") for entry in entries + ] + + def destroy(self, sandbox: str) -> None: + sandbox_client = self._sandboxes.pop(sandbox, None) + try: + if sandbox_client is None: + from opensandbox import SandboxSync + + sandbox_client = SandboxSync.connect( Review Comment: `connect` waits for readiness unless you pass `skip_health_check=True`, so this spends up to `ready_timeout` (120s) polling before it even tries to kill. For a sandbox that exists but isn't healthy, paused or with execd not answering, the poll then raises `SandboxReadyTimeoutException`, which carries no `status_code`, so the 404 check below can't fire and it becomes terminal. That's 120s spent to fail a method [base.py L287-289](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L287-L289) requires to be idempotent, and you don't need a healthy sandbox to destroy one. `skip_health_check=True` on this path looks right. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), + connection_config=self._get_connection_config(), + ) + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + def _get_sandbox(self, sandbox_id: str) -> SandboxSync: + if sandbox := self._sandboxes.get(sandbox_id): + return sandbox + with _translate_opensandbox_errors("connect to a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.connect( + sandbox_id, + connection_config=self._get_connection_config(), + connect_timeout=timedelta(seconds=self._ready_timeout), + ) + self._sandboxes[sandbox_id] = sandbox + return sandbox + + def run_command( + self, sandbox: str, command: str, *, timeout: float, max_output_bytes: int + ) -> SandboxExecResult: + _validate_positive_finite(timeout, "timeout") + _validate_positive_finite(max_output_bytes, "max_output_bytes") + stdout = _BoundedTail(max_output_bytes) + stderr = _BoundedTail(max_output_bytes) + started = time.monotonic() + with _translate_opensandbox_errors("run a sandbox command"): + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.execd_sync import ExecutionHandlersSync + + execution = self._get_sandbox(sandbox).commands.run( + command, + opts=RunCommandOpts(timeout=timedelta(seconds=timeout)), + handlers=ExecutionHandlersSync( + on_stdout=stdout.add_message, + on_stderr=stderr.add_message, + skip_accumulation=True, + ), + ) + + if execution.exit_code is None: + raise SandboxTerminalError("OpenSandbox returned no terminal status for the command.") + if execution.error is not None and not stderr.get_text(): + details = "\n".join(execution.error.traceback) or execution.error.value + stderr.add_text(details) + timed_out = execution.exit_code < 0 and time.monotonic() - started >= timeout + return SandboxExecResult( + exit_code=execution.exit_code, + stdout=stdout.get_text(), + stderr=stderr.get_text(), + timed_out=timed_out, + stdout_truncated=stdout.truncated, + stderr_truncated=stderr.truncated, + ) + + @staticmethod + def _confirm_sandbox_exists(sandbox: SandboxSync) -> None: + with _translate_opensandbox_errors("confirm that a sandbox still exists"): + sandbox.get_info() + + def read_file(self, sandbox: str, path: str, *, max_bytes: int) -> bytes: + _validate_positive_finite(max_bytes, "max_bytes") + sandbox_client = self._get_sandbox(sandbox) + chunks = None + data = bytearray() + try: + chunks = sandbox_client.files.read_bytes_stream( + path, + chunk_size=min(65536, max_bytes + 1), + range_header=f"bytes=0-{max_bytes}", + ) + for chunk in chunks: + data.extend(chunk[: max_bytes + 1 - len(data)]) + if len(data) > max_bytes: + raise SandboxFileTooLargeError(path, len(data), max_bytes) Review Comment: `len(data)` is always `max_bytes + 1` here, so every oversized file is reported as one byte over the limit, and the toolset renders `size_bytes` straight into the message the model sees. At the default 5 MiB budget a 1 GB file comes back as "is 5.0MB, over the 5.0MB read limit", which isn't much to plan a slicing strategy from. The base implementation reports the `stat` size ([base.py L232-240](https://github.com/apache/airflow/blob/4eba473d6c0518e019b22adac4c913b3a723cb75/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py#L232-L240)), and `files.get_file_info([path])` returns `EntryInfo.size` if you want the same here. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""OpenSandbox backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import posixpath +import time +from contextlib import contextmanager, suppress +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from opensandbox import SandboxSync + from opensandbox.config import ConnectionConfigSync + from opensandbox.models.sandboxes import NetworkPolicy + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + + +def _get_status_code(error: Exception) -> int | None: + status_code = getattr(error, "status_code", None) + return status_code if isinstance(status_code, int) else None + + +@contextmanager +def _translate_opensandbox_errors( + operation: str, *, recoverable_statuses: frozenset[int] = frozenset() +) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from opensandbox.exceptions import SandboxApiException + except ImportError: + raise SandboxTerminalError( + "The OpenSandbox SDK is not installed. Install " + '"apache-airflow-providers-common-ai[sandbox-opensandbox]".' + ) from e + status_code = _get_status_code(e) if isinstance(e, SandboxApiException) else None + status = f" (HTTP {status_code})" if status_code is not None else "" + message = f"OpenSandbox could not {operation}{status}." + if status_code in recoverable_statuses: + raise SandboxError(message) from e + raise SandboxTerminalError(message) from e + + +class _BoundedTail: + def __init__(self, max_bytes: int) -> None: + self._max_bytes = max_bytes + self._data = bytearray() + self.truncated = False + + def add_text(self, text: str) -> None: + self._data.extend(text.encode("utf-8")) + if len(self._data) > self._max_bytes: + del self._data[: len(self._data) - self._max_bytes] + self.truncated = True + + def add_message(self, message: Any) -> None: + # execd streams one message per output line with the delimiter stripped, + # so the newline has to be put back or every line runs together. A blank + # line already arrives as "\n", hence the guard. + text = message.text + self.add_text(text if text.endswith("\n") else text + "\n") + + def get_text(self) -> str: + return bytes(self._data).decode("utf-8", errors="ignore") + + +def _parse_bool(value: Any, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise SandboxTerminalError(f"The OpenSandbox connection extra {name} must be a boolean.") + + +class OpenSandboxBackend(SandboxBackend): + """ + Run sandbox tools through an OpenSandbox server. + + OpenSandbox supports Docker and Kubernetes runtimes behind the same API. + Airflow workers need only network access to that API; the OpenSandbox + deployment owns container provisioning and isolation. + + A generic Airflow connection supplies the server configuration. ``host`` + and ``port`` identify the lifecycle API, ``schema`` selects ``http`` or + ``https``, and ``password`` carries the optional API key. Connection extras + may set ``request_timeout`` and ``use_server_proxy``. + + Strict network policy requires the OpenSandbox egress sidecar. The server + rejects a requested policy when that component or runtime support is + unavailable, preserving :class:`~airflow.providers.common.ai.sandbox.SandboxSpec`'s + fail-closed contract. + + :param opensandbox_conn_id: Generic Airflow connection ID. ``None`` lets the + SDK resolve ``OPEN_SANDBOX_DOMAIN`` and ``OPEN_SANDBOX_API_KEY``. + :param image: Container image used for each sandbox. + :param cpu: OpenSandbox CPU resource limit. + :param memory: OpenSandbox memory resource limit. + :param sandbox_timeout: Server-side sandbox lifetime in seconds. + :param ready_timeout: Seconds to wait for a newly created sandbox to become healthy. + :param use_server_proxy: Route sandbox service calls through the lifecycle + server. ``None`` reads the connection extra and otherwise defaults to ``True``. + """ + + name = "opensandbox" + + def __init__( + self, + opensandbox_conn_id: str | None = "opensandbox_default", + *, + image: str = "python:3.12-slim", + cpu: str = "1", + memory: str = "2Gi", + sandbox_timeout: float = 3600.0, + ready_timeout: float = 120.0, + use_server_proxy: bool | None = None, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not cpu: + raise ValueError("cpu must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + _validate_positive_finite(sandbox_timeout, "sandbox_timeout") + _validate_positive_finite(ready_timeout, "ready_timeout") + self._opensandbox_conn_id = opensandbox_conn_id + self._image = image + self._resource = {"cpu": cpu, "memory": memory} + self._sandbox_timeout = sandbox_timeout + self._ready_timeout = ready_timeout + self._use_server_proxy = use_server_proxy + self._connection_config: ConnectionConfigSync | None = None + self._sandboxes: dict[str, SandboxSync] = {} + + def _get_connection_config(self) -> ConnectionConfigSync: + if self._connection_config is not None: + return self._connection_config + with _translate_opensandbox_errors("initialize its client"): + from opensandbox.config import ConnectionConfigSync + + if self._opensandbox_conn_id is None: + self._connection_config = ConnectionConfigSync( + use_server_proxy=True if self._use_server_proxy is None else self._use_server_proxy + ) + return self._connection_config + + conn = BaseHook.get_connection(self._opensandbox_conn_id) + extra = conn.extra_dejson + request_timeout = extra.get("request_timeout", 30) + try: + request_timeout = float(request_timeout) + _validate_positive_finite(request_timeout, "connection extra request_timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The OpenSandbox connection extra request_timeout must be a positive finite number." + ) from e + + use_server_proxy = self._use_server_proxy + if use_server_proxy is None: + value = extra.get("use_server_proxy", True) + use_server_proxy = _parse_bool(value, "use_server_proxy") + + domain = conn.host + if domain and conn.port: + domain = f"{domain}:{conn.port}" + self._connection_config = ConnectionConfigSync( + api_key=conn.password or None, + domain=domain, + protocol=conn.schema or "http", + request_timeout=timedelta(seconds=request_timeout), + use_server_proxy=use_server_proxy, + ) + return self._connection_config + + @staticmethod + def _get_network_policy(spec: SandboxSpec | None) -> NetworkPolicy | None: + if spec is None: + return None + if not spec.block_network and spec.allow_egress_to: + raise SandboxTerminalError( + "SandboxSpec.allow_egress_to only narrows a deny-by-default policy; " + "set block_network=True or remove the allowlist." + ) + from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule + + rules = [NetworkRule(action="allow", target=target) for target in spec.allow_egress_to or ()] + # default_action is declared under its wire alias. populate_by_name means both + # spellings work at runtime, but only the alias is in the typed signature. + return NetworkPolicy( + defaultAction="deny" if spec.block_network else "allow", + egress=rules or None, + ) + + def create(self, *, spec: SandboxSpec | None = None) -> str: + with _translate_opensandbox_errors("create a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.create( + self._image, + timeout=timedelta(seconds=self._sandbox_timeout), + ready_timeout=timedelta(seconds=self._ready_timeout), + env=dict(spec.env) if spec is not None and spec.env else None, + resource=dict(self._resource), + network_policy=self._get_network_policy(spec), + connection_config=self._get_connection_config(), + ) + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + def _get_sandbox(self, sandbox_id: str) -> SandboxSync: + if sandbox := self._sandboxes.get(sandbox_id): + return sandbox + with _translate_opensandbox_errors("connect to a sandbox"): + from opensandbox import SandboxSync + + sandbox = SandboxSync.connect( + sandbox_id, + connection_config=self._get_connection_config(), + connect_timeout=timedelta(seconds=self._ready_timeout), + ) + self._sandboxes[sandbox_id] = sandbox + return sandbox + + def run_command( + self, sandbox: str, command: str, *, timeout: float, max_output_bytes: int + ) -> SandboxExecResult: + _validate_positive_finite(timeout, "timeout") + _validate_positive_finite(max_output_bytes, "max_output_bytes") + stdout = _BoundedTail(max_output_bytes) + stderr = _BoundedTail(max_output_bytes) + started = time.monotonic() Review Comment: `started` is taken before the sandbox lookup on L271, and for an uncached handle `_get_sandbox` calls `SandboxSync.connect(connect_timeout=ready_timeout)`, which polls for readiness up to 120s by default. That reconnect time lands in the elapsed value feeding `timed_out` on L286, so a command that was signal-killed early after a slow reconnect gets reported to the model as having hit its budget. Moving this below the lookup fixes it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
