diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 2a8b7c9459151..323c564299840 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -318,6 +318,7 @@ contrib convertToIri cooldown copyable +coreutils CoreV coroutine cosmosdb @@ -879,6 +880,7 @@ irreproducible IRSA isfile ish +islo isn isort istio @@ -1069,6 +1071,7 @@ mgmt microservice microservices microsoft +microVM middleware middlewares midnights @@ -1481,6 +1484,7 @@ sanitization sas Sasl sasl +sbx scala scalability scalable @@ -1753,6 +1757,7 @@ tis TLS tls tmp +tmpfs tnsnames todo tokenization diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index d1bff2ce31c6a..8998e86934fbc 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -239,29 +239,30 @@ Install them when installing from PyPI. For example: pip install apache-airflow-providers-common-ai[anthropic] -============== ======================================================================================================================================= -Extra Dependencies -============== ======================================================================================================================================= -``anthropic`` ``pydantic-ai-slim[anthropic]`` -``bedrock`` ``pydantic-ai-slim[bedrock]`` -``google`` ``pydantic-ai-slim[google]`` -``openai`` ``pydantic-ai-slim[openai]`` -``mcp`` ``pydantic-ai-slim[mcp]`` -``code-mode`` ``pydantic-ai-harness[codemode]>=0.3.0`` -``shields`` ``pydantic-ai-shields>=0.3.4`` -``skills`` ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0`` -``avro`` ``fastavro>=1.10.0; python_version < "3.14"``, ``fastavro>=1.12.1; python_version >= "3.14"`` -``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``, ``pyarrow>=22.0.0; python_version >= '3.14'`` -``sql`` ``apache-airflow-providers-common-sql>=1.33.0``, ``sqlglot>=30.0.0`` -``aws`` ``apache-airflow-providers-amazon>=9.0.0`` -``common.sql`` ``apache-airflow-providers-common-sql>=1.33.0`` -``langchain`` ``langchain>=1.0.0`` -``llamaindex`` ``dataclasses-json>=0.6.7``, ``llama-index-core>=0.13.0``, ``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.0`` -``pdf`` ``pypdf>=4.0.0`` -``docx`` ``python-docx>=1.0.0`` -``git`` ``apache-airflow-providers-git`` -``amazon`` ``apache-airflow-providers-amazon`` -============== ======================================================================================================================================= +================ ======================================================================================================================================= +Extra Dependencies +================ ======================================================================================================================================= +``anthropic`` ``pydantic-ai-slim[anthropic]`` +``bedrock`` ``pydantic-ai-slim[bedrock]`` +``google`` ``pydantic-ai-slim[google]`` +``openai`` ``pydantic-ai-slim[openai]`` +``mcp`` ``pydantic-ai-slim[mcp]`` +``code-mode`` ``pydantic-ai-harness[codemode]>=0.3.0`` +``shields`` ``pydantic-ai-shields>=0.3.4`` +``skills`` ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0`` +``sandbox-islo`` ``islo>=0.3.9`` +``avro`` ``fastavro>=1.10.0; python_version < "3.14"``, ``fastavro>=1.12.1; python_version >= "3.14"`` +``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``, ``pyarrow>=22.0.0; python_version >= '3.14'`` +``sql`` ``apache-airflow-providers-common-sql>=1.33.0``, ``sqlglot>=30.0.0`` +``aws`` ``apache-airflow-providers-amazon>=9.0.0`` +``common.sql`` ``apache-airflow-providers-common-sql>=1.33.0`` +``langchain`` ``langchain>=1.0.0`` +``llamaindex`` ``dataclasses-json>=0.6.7``, ``llama-index-core>=0.13.0``, ``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.0`` +``pdf`` ``pypdf>=4.0.0`` +``docx`` ``python-docx>=1.0.0`` +``git`` ``apache-airflow-providers-git`` +``amazon`` ``apache-airflow-providers-amazon`` +================ ======================================================================================================================================= Downloading official packages ----------------------------- diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index d14e61772b5a3..f326c3d5296e1 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -24,7 +24,7 @@ Airflow's 350+ provider hooks already have typed methods, rich docstrings, and managed credentials. Toolsets expose them as pydantic-ai tools so that LLM agents can call them during multi-turn reasoning. -Four toolsets are included: +The toolsets include: - :class:`~airflow.providers.common.ai.toolsets.aws.AWSToolset` — configured AWS services toolset for agent access to AWS APIs through Airflow-managed @@ -36,8 +36,10 @@ Four toolsets are included: connections. - :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated 4-tool database toolset. +- :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` — + execute agent-written Python in an isolated sandbox, off the Airflow worker. -All four implement pydantic-ai's +All of them implement pydantic-ai's `AbstractToolset `__ interface and can be passed to any pydantic-ai ``Agent``, including via :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`. @@ -543,6 +545,166 @@ resolves sources to local ``SKILL.md`` directories that any loader accepts: ``resolve_skills`` needs the Git provider (for ``GitSkills``) but not pydantic-ai, and removes any cloned directories when the ``with`` block exits. + +``SandboxToolset`` +------------------ + +:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` gives the +agent one tool, ``run_python_in_sandbox``, that executes agent-written Python in a +disposable sandbox provisioned by a +:class:`~airflow.providers.common.ai.sandbox.SandboxBackend` — off the Airflow +worker process. Airflow does not inject its context, connections, credentials, +or worker environment into the sandbox. Custom images can still contain +credentials, and hosted backends can provide their own identity, so image and +backend configuration remain the Deployment Manager's responsibility. + +It is the counterpart to :ref:`code mode `. Code mode's ``run_code`` +executes in-process on the worker via the Monty interpreter: it is glue between +tool calls, restricted to a subset of Python with no third-party imports. +``SandboxToolset`` instead ships the code to a real Python interpreter in an +isolated environment, so the agent can crunch data with the full language and +whatever the sandbox image provides — at the cost of provisioning a microVM +per run. + +The sandbox is created lazily on the first ``run_python_in_sandbox`` call, shared by every +call within the agent run, and destroyed when the run ends. Each call runs a +fresh interpreter (``python3 -c``), so files written by earlier calls persist +in the sandbox while in-memory variables do not. A nonzero exit code or a +timeout is normal tool output — the model reads ``stderr`` and fixes its own +code. If a backend cannot confirm that timed-out code stopped, it destroys the +sandbox and returns ``sandbox_terminated: true``; the next call receives a +fresh sandbox and files from earlier calls are no longer available. Only +infrastructure errors (the ``sbx`` CLI missing, the islo API unreachable) fail +the task. + +sbx backend (Docker Sandboxes) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +:class:`~airflow.providers.common.ai.sandbox.SbxSandboxBackend` runs each +sandbox in a `Docker Sandboxes `__ +microVM by driving the ``sbx`` CLI (``create`` / ``exec`` / ``rm``). Each +sandbox is a genuine microVM with its own kernel, so agent code is isolated by +a hardware boundary rather than a shared kernel. Command output is capped at +the first 64 KiB of each stream, and the microVM is torn down when the run +ends. + +It is a Deployment Manager prerequisite to install the ``sbx`` binary +(``brew install docker/tap/sbx`` / ``winget install Docker.sbx``) and run +``sbx policy init`` once on the worker host; the backend needs no Python +dependency. Outbound network egress is governed by that host ``sbx policy`` +rather than by the backend — use ``sbx policy init deny-all`` for a no-egress +default. The template image must provide the GNU ``timeout`` utility, which +enforces the per-command timeout (any Debian/Ubuntu-based image, including +``python:*-slim``, does). + +.. code-block:: python + + from airflow.providers.common.ai.operators.agent import AgentOperator + from airflow.providers.common.ai.sandbox import SbxSandboxBackend + from airflow.providers.common.ai.toolsets import SandboxToolset + + AgentOperator( + task_id="sandboxed_analyst", + prompt="Estimate pi with a Monte Carlo simulation of one million points.", + llm_conn_id="pydanticai_default", + toolsets=[SandboxToolset(SbxSandboxBackend())], + ) + +Constructor parameters: + +- ``image``: Container image for the sandbox (``sbx --template``). + Default ``"python:3.12-slim"``. +- ``memory``: Memory limit in binary units (e.g. ``"2g"``). ``sbx`` enforces a + 1 GiB minimum. Default ``"2g"``. +- ``cpus``: Number of CPUs to allocate. ``None`` (default) uses the ``sbx`` + default. +- ``sbx_path``: Path to the ``sbx`` binary. Default ``"sbx"``. +- ``create_timeout``: Seconds allowed for provisioning (first-run microVM boot + plus image pull can be slow). Default ``600``. + +islo backend +^^^^^^^^^^^^ + +:class:`~airflow.providers.common.ai.sandbox.IsloSandboxBackend` runs each +sandbox in an `islo.dev `__ microVM — hardware-level +isolation with no local Docker daemon required. + +Requires the ``sandbox-islo`` extra: +``pip install "apache-airflow-providers-common-ai[sandbox-islo]"`` + +.. code-block:: python + + from airflow.providers.common.ai.sandbox import IsloSandboxBackend + + SandboxToolset(IsloSandboxBackend(islo_conn_id="islo_default")) + +Credentials come from a generic Airflow connection (no custom connection type +is needed): + +- ``password``: the islo API key. Required. +- ``host``: the compute URL. Optional. +- Extra: optional ``base_url`` and ``timeout`` (request timeout in seconds) + keys. + +Constructor parameters: + +- ``islo_conn_id``: Airflow connection ID for islo. Default ``"islo_default"``. + ``None`` lets the SDK resolve credentials from its environment variables + (``ISLO_API_KEY`` etc.). +- ``image``, ``vcpus``, ``memory_mb``: sandbox image and sizing. ``None`` + (default) uses the server default for each. +- ``internet_enabled``: outbound internet access. Default ``False`` (no egress); + ``None`` defers to the server default. +- ``delete_after``: Server-side TTL in seconds after which the sandbox is + deleted even if the worker never got to destroy it (killed mid-run). + Default ``3600``. + +The command timeout is enforced by the backend, not the islo API (the API's +``timeout_secs`` is only a compatibility hint). If no terminal command state +arrives within the timeout plus a short grace period, the backend deletes the +microVM, marks the sandbox terminated, and returns a timeout. + +Parameters +^^^^^^^^^^ + +- ``backend``: The :class:`~airflow.providers.common.ai.sandbox.SandboxBackend` + that provisions and runs the sandbox. +- ``timeout``: Timeout in seconds for a single ``run_python_in_sandbox`` call. + Default ``300``. +- ``python_command``: Python executable inside the sandbox. + Default ``"python3"``. + +Bringing your own backend +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Any vendor that can create a sandbox, run a command in it, and destroy it can +plug in: subclass +:class:`~airflow.providers.common.ai.sandbox.SandboxBackend` in your own +package and pass an instance to ``SandboxToolset``: + +.. code-block:: python + + from airflow.providers.common.ai.sandbox import SandboxBackend, SandboxResult + + + class AcmeSandboxBackend(SandboxBackend): + name = "acme" + + def create(self) -> str: + return acme_sdk.create_sandbox().id + + def run(self, sandbox: str, command: list[str], *, timeout: float) -> SandboxResult: + r = acme_sdk.exec(sandbox, command, timeout=timeout) + return SandboxResult(exit_code=r.exit_code, stdout=r.stdout, stderr=r.stderr) + + def destroy(self, sandbox: str) -> None: + acme_sdk.delete_sandbox(sandbox) + +Constructors run at Dag-parse time, so resolve credentials lazily, on first +use, and make ``destroy`` idempotent (destroying an already-gone sandbox must +not raise). + + Working with LangChain ---------------------- @@ -692,6 +854,16 @@ No single layer is sufficient — they work together. - Does **not** constrain what those tools do. An MCP server can expose shell, filesystem, or network access. Run only trusted servers and audit the tools they expose. + * - **SandboxToolset: off-worker execution** + - Executes agent-written code in a disposable microVM, never in the + worker process. Airflow does not inject its context, connections, + credentials, or worker environment. The backends enforce command + timeouts, cap each output stream at 64 KiB, and tear down the sandbox + after the run; the islo backend also sets a server-side cleanup TTL. + - Does not sanitize what the code computes or returns. Custom images can + contain secrets, hosted backends can expose their own identity, and + Docker Sandboxes network access follows the host ``sbx policy``. Its + CPU allocation defaults to all host CPUs unless explicitly limited. * - **pydantic-ai: tool call budget** - pydantic-ai's ``max_result_retries`` and ``model_settings`` control how many tool-call rounds the agent can make before stopping. diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index abb5ca41d3d7d..37b86411a61ae 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -62,6 +62,12 @@ integrations: - /docs/apache-airflow-providers-common-ai/operators/llamaindex_embedding.rst - /docs/apache-airflow-providers-common-ai/operators/llamaindex_retrieval.rst tags: [ai] + - integration-name: Docker Sandboxes + external-doc-url: https://docs.docker.com/ai/sandboxes/ + tags: [software] + - integration-name: Islo + external-doc-url: https://docs.islo.dev/ + tags: [service] hooks: - integration-name: Pydantic AI @@ -460,6 +466,7 @@ toolsets: - airflow.providers.common.ai.toolsets.datafusion - airflow.providers.common.ai.toolsets.logging - airflow.providers.common.ai.toolsets.mcp + - airflow.providers.common.ai.toolsets.sandbox - airflow.providers.common.ai.toolsets.skills - airflow.providers.common.ai.toolsets.langchain_bridge diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index cbaa9e82ee4d4..3b512f1e9280a 100644 --- a/providers/common/ai/pyproject.toml +++ b/providers/common/ai/pyproject.toml @@ -96,6 +96,10 @@ dependencies = [ "apache-airflow-providers-git>=0.4.0", "pydantic-ai-skills>=1.2.0", ] +# Isolated sandbox code execution for agents (SandboxToolset). The islo backend +# needs its SDK; the sbx backend only shells out to the `sbx` CLI, so it needs +# no Python dependency and has no extra. +"sandbox-islo" = ["islo>=0.3.9"] "avro" = [ 'fastavro>=1.10.0; python_version < "3.14"', 'fastavro>=1.12.1; python_version >= "3.14"', @@ -153,6 +157,7 @@ dev = [ "llama-index-core>=0.13.0", "llama-index-embeddings-openai>=0.6.0", "llama-index-llms-openai>=0.6.0", + "islo>=0.3.9", ] # To build docs: diff --git a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py index 89d96737d94fa..f2fd9000e7090 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py +++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py @@ -65,6 +65,12 @@ def get_provider_info(): ], "tags": ["ai"], }, + { + "integration-name": "Docker Sandboxes", + "external-doc-url": "https://docs.docker.com/ai/sandboxes/", + "tags": ["software"], + }, + {"integration-name": "Islo", "external-doc-url": "https://docs.islo.dev/", "tags": ["service"]}, ], "hooks": [ { diff --git a/providers/common/ai/src/airflow/providers/common/ai/sandbox/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/sandbox/__init__.py new file mode 100644 index 0000000000000..23d13c6565902 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/__init__.py @@ -0,0 +1,39 @@ +# 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. +"""Isolated sandbox backends for the SandboxToolset.""" + +from __future__ import annotations + +from airflow.providers.common.ai.sandbox.base import SandboxBackend, SandboxResult + +# SbxSandboxBackend only shells out to the `sbx` CLI (stdlib imports), so it is +# always importable; IsloSandboxBackend needs the optional `islo` SDK. +from airflow.providers.common.ai.sandbox.sbx import SbxSandboxBackend + +__all__ = ["IsloSandboxBackend", "SandboxBackend", "SandboxResult", "SbxSandboxBackend"] + + +def __getattr__(name: str): + if name == "IsloSandboxBackend": + try: + from airflow.providers.common.ai.sandbox.islo import IsloSandboxBackend + except ImportError as e: + from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + return IsloSandboxBackend + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py b/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py new file mode 100644 index 0000000000000..2dd88f613bae2 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py @@ -0,0 +1,94 @@ +# 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. +"""Vendor-neutral contract for running agent-written code in an isolated sandbox.""" + +from __future__ import annotations + +import math +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import ClassVar + +# Cap on each of stdout/stderr a backend returns from a command, protecting the +# LLM context window from unbounded output. Shared by every backend. +_MAX_OUTPUT_CHARS = 64 * 1024 + + +def _validate_positive_finite(value: float, name: str) -> None: + if not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be a positive finite number, got {value!r}.") + + +def _new_sandbox_name() -> str: + """Generate a unique sandbox name, ``airflow-sandbox-`` prefixed for correlation and cleanup.""" + return f"airflow-sandbox-{uuid.uuid4().hex[:12]}" + + +def _cap_output(text: str) -> tuple[str, bool]: + """Return ``text`` capped at ``_MAX_OUTPUT_CHARS`` and whether it was truncated.""" + if len(text) <= _MAX_OUTPUT_CHARS: + return text, False + return text[:_MAX_OUTPUT_CHARS], True + + +@dataclass(frozen=True) +class SandboxResult: + """ + Outcome of one command execution inside a sandbox. + + ``timed_out`` means the command hit the timeout (``exit_code`` is then + meaningless). ``truncated`` means the backend capped ``stdout``/``stderr``. + ``sandbox_terminated`` means the backend destroyed the sandbox while + stopping the command, so the toolset must provision a fresh one next time. + """ + + exit_code: int + stdout: str + stderr: str + timed_out: bool = False + truncated: bool = False + sandbox_terminated: bool = False + + +class SandboxBackend(ABC): + """ + Contract for running commands in an isolated sandbox. + + The lifecycle is create → run (any number of times) → destroy, driven by + :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`. + Implementations must be cheap to construct — constructors run at Dag-parse + time — so credentials and connections are resolved lazily, on first use. + ``destroy`` must be idempotent: destroying an already-gone sandbox is not + an error. All methods are synchronous; the toolset offloads them to a + thread. + """ + + name: ClassVar[str] + """Short backend identifier (e.g. ``"docker"``), used in the toolset id.""" + + @abstractmethod + def create(self) -> str: + """Provision one sandbox and return its handle (name or id).""" + + @abstractmethod + def run(self, sandbox: str, command: list[str], *, timeout: float) -> SandboxResult: + """Run ``command`` in the sandbox, bounded by ``timeout`` seconds.""" + + @abstractmethod + def destroy(self, sandbox: str) -> None: + """Tear down the sandbox. Must be idempotent.""" diff --git a/providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py b/providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py new file mode 100644 index 0000000000000..33cd7ab1d3919 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py @@ -0,0 +1,183 @@ +# 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. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import time +from contextlib import suppress +from typing import Any + +# The missing-``islo`` case is turned into AirflowOptionalProviderFeatureException +# by the package __init__'s lazy import, so import the SDK plainly here. +from islo import Islo +from islo.errors import NotFoundError +from islo.types import LifecyclePolicy + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxResult, + _cap_output, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_EXEC_POLL_INTERVAL = 0.2 +_EXEC_POLL_GRACE = 5.0 + + +class IsloSandboxBackend(SandboxBackend): + """ + Sandbox backend that runs agent code in an `islo.dev `__ microVM. + + Credentials come from an Airflow connection: ``password`` is the islo API + key (required), ``host`` the compute URL (optional), and the extra field + may set ``base_url`` and ``timeout`` (request timeout in seconds). + + :param islo_conn_id: Airflow connection ID for islo. ``None`` lets the SDK + resolve credentials from its environment variables (``ISLO_API_KEY`` etc.). + :param image: Sandbox image. ``None`` (default) uses the server default. + :param vcpus: Number of virtual CPUs. ``None`` uses the server default. + :param memory_mb: Memory in MB. ``None`` uses the server default. + :param internet_enabled: Allow outbound internet access from the sandbox. + Default ``False`` (no egress) — a safe default matching an isolated + sandbox; pass ``True`` to allow it, or ``None`` to defer to the server. + :param delete_after: Server-side TTL in seconds after which the sandbox is + deleted even if the worker never got to destroy it (killed mid-run). + Default ``3600``. + """ + + name = "islo" + + def __init__( + self, + islo_conn_id: str | None = "islo_default", + *, + image: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + internet_enabled: bool | None = False, + delete_after: int = 3600, + ) -> None: + _validate_positive_finite(delete_after, "delete_after") + if vcpus is not None: + _validate_positive_finite(vcpus, "vcpus") + if memory_mb is not None: + _validate_positive_finite(memory_mb, "memory_mb") + if image == "": + raise ValueError("image must not be empty.") + self._islo_conn_id = islo_conn_id + self._image = image + self._vcpus = vcpus + self._memory_mb = memory_mb + self._internet_enabled = internet_enabled + self._delete_after = delete_after + self._client: Islo | None = None + + def _get_client(self) -> Islo: + if self._client is not None: + return self._client + if self._islo_conn_id is None: + self._client = Islo() + return self._client + conn = BaseHook.get_connection(self._islo_conn_id) + api_key = (conn.password or "").strip() + if not api_key: + raise ValueError( + f"Connection {self._islo_conn_id!r} has no password; set it to the islo API key." + ) + kwargs: dict[str, Any] = {"api_key": api_key} + if conn.host: + kwargs["compute_url"] = conn.host + extra = conn.extra_dejson + if extra.get("base_url"): + kwargs["base_url"] = extra["base_url"] + if extra.get("timeout") is not None: + request_timeout = float(extra["timeout"]) + _validate_positive_finite(request_timeout, "connection extra timeout") + kwargs["timeout"] = request_timeout + self._client = Islo(**kwargs) + return self._client + + def create(self) -> str: + # Pass only the kwargs that were set; the SDK treats absent kwargs as + # "omit" and the server fills in its defaults. + kwargs: dict[str, Any] = {} + if self._image is not None: + kwargs["image"] = self._image + if self._vcpus is not None: + kwargs["vcpus"] = self._vcpus + if self._memory_mb is not None: + kwargs["memory_mb"] = self._memory_mb + if self._internet_enabled is not None: + kwargs["internet_enabled"] = self._internet_enabled + sandbox = self._get_client().sandboxes.create_sandbox( + name=_new_sandbox_name(), + lifecycle=LifecyclePolicy(delete_after=self._delete_after), + **kwargs, + ) + # The server may normalize the name; the response is authoritative. + return sandbox.name + + def run(self, sandbox: str, command: list[str], *, timeout: float) -> SandboxResult: + _validate_positive_finite(timeout, "timeout") + client = self._get_client() + # ``timeout_secs`` is only a client-side hint to the islo API (not + # server-enforced), so the real bound is our poll deadline below: if no + # terminal result arrives in time we delete the microVM ourselves. + response = client.sandboxes.exec_in_sandbox( + sandbox, + command=list(command), + timeout_secs=max(1, math.ceil(timeout)), + ) + deadline = time.monotonic() + timeout + _EXEC_POLL_GRACE + while time.monotonic() < deadline: + result = client.sandboxes.get_exec_result(sandbox, response.exec_id) + if result.status in _TERMINAL_EXEC_STATUSES: + stdout, stdout_truncated = _cap_output(result.stdout or "") + stderr, stderr_truncated = _cap_output(result.stderr or "") + return SandboxResult( + exit_code=result.exit_code if result.exit_code is not None else -1, + stdout=stdout, + stderr=stderr, + timed_out=result.status == "timeout", + truncated=result.truncated or stdout_truncated or stderr_truncated, + ) + time.sleep(_EXEC_POLL_INTERVAL) + + # No terminal result within our poll deadline. Delete the microVM so the + # command cannot continue in the shared sandbox, then tell the toolset + # to provision a fresh one. A + # transient delete failure must not fail the task — the server-side TTL + # (delete_after) is the backstop — so this teardown is best-effort. + with suppress(Exception): + self.destroy(sandbox) + return SandboxResult( + exit_code=-1, + stdout="", + stderr="", + timed_out=True, + sandbox_terminated=True, + ) + + def destroy(self, sandbox: str) -> None: + # Already-gone is fine — destroy is idempotent. + with suppress(NotFoundError): + self._get_client().sandboxes.delete_sandbox(sandbox_name=sandbox) diff --git a/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py b/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py new file mode 100644 index 0000000000000..1116c3a80e4c9 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py @@ -0,0 +1,248 @@ +# 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. +"""Docker Sandboxes (``sbx``) microVM backend for the SandboxToolset.""" + +from __future__ import annotations + +import math +import shutil +import subprocess +import tempfile +import threading +import time +from contextlib import suppress + +from airflow.providers.common.ai.sandbox.base import ( + _MAX_OUTPUT_CHARS, + SandboxBackend, + SandboxResult, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +# Extra wall-clock beyond the per-command budget to absorb CLI + microVM +# round-trip overhead before we treat the sbx call itself as hung. +_EXEC_GRACE = 30.0 +# Grace after the timeout fires before the command is force-killed (SIGKILL) if +# it ignored SIGTERM. Kept well under _EXEC_GRACE so the outer call still returns. +_KILL_AFTER = 10 + + +class SbxSandboxBackend(SandboxBackend): + """ + Sandbox backend that runs agent code in a Docker Sandboxes (``sbx``) microVM. + + Drives the ``sbx`` CLI: ``create`` provisions a per-session microVM, ``exec`` + runs commands in it (``docker exec`` semantics), and ``rm`` tears it down. + Each sandbox is a genuine microVM with its own kernel, so agent code is far + better isolated than a shared-kernel container — the same isolation grade as + :class:`~airflow.providers.common.ai.sandbox.IsloSandboxBackend`, but local. + + Requires the ``sbx`` binary on ``PATH`` (``brew install docker/tap/sbx`` / + ``winget install Docker.sbx``) and a one-time ``sbx policy init`` on the host; + it is a Deployment Manager prerequisite, not something Airflow configures. The + template image must provide the GNU coreutils ``timeout`` binary, which + enforces the per-command timeout (any Debian/Ubuntu-based image, including + ``python:*-slim``, does). + + Airflow injects none of its context, connections, or worker environment into + the sandbox; a custom template image may still bundle its own tools. Outbound + network egress is governed by the host ``sbx policy`` (a Deployment Manager + setting), not by this backend — run ``sbx policy init deny-all`` for a + no-egress default. + + :param image: Container image for the sandbox (``sbx --template``). + Default ``"python:3.12-slim"``. + :param memory: Memory limit in binary units (e.g. ``"2g"``). ``sbx`` enforces + a 1 GiB minimum. Default ``"2g"``. + :param cpus: Number of CPUs to allocate. ``None`` (default) uses the ``sbx`` + default (all host CPUs). + :param sbx_path: Path to the ``sbx`` binary. Default ``"sbx"``. + :param create_timeout: Seconds to allow for provisioning (first-run microVM + boot plus image pull can be slow). Default ``600``. + """ + + name = "sbx" + + def __init__( + self, + *, + image: str = "python:3.12-slim", + memory: str = "2g", + cpus: int | None = None, + sbx_path: str = "sbx", + create_timeout: float = 600.0, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + if cpus is not None: + _validate_positive_finite(cpus, "cpus") + if not sbx_path: + raise ValueError("sbx_path must not be empty.") + _validate_positive_finite(create_timeout, "create_timeout") + self._image = image + self._memory = memory + self._cpus = cpus + self._sbx_path = sbx_path + self._create_timeout = create_timeout + # Each sandbox mounts a throwaway host workspace; remember it so destroy + # can remove it. ``sbx create`` requires a workspace path but the agent + # never needs host files, so an empty temp dir keeps the host untouched. + self._workspaces: dict[str, str] = {} + + def _run_cli(self, args: list[str], *, timeout: float) -> subprocess.CompletedProcess[str]: + # Only for small, bounded output (create/rm). Command execution uses + # _exec_capped, which bounds memory against unbounded agent output. + return subprocess.run( + [self._sbx_path, *args], + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + + def _exec_capped(self, args: list[str], *, timeout: float) -> tuple[int, str, str, bool]: + """ + Run the CLI with the output retained bounded to ``_MAX_OUTPUT_CHARS`` per stream. + + Agent code can print unbounded output; ``subprocess.run`` would buffer + all of it and OOM the worker, so drain incrementally and keep only the + cap. Returns ``(returncode, stdout, stderr, truncated)``. + """ + proc = subprocess.Popen( + [self._sbx_path, *args], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + def drain(stream, buf: bytearray, flag: list[bool]) -> None: + for chunk in iter(lambda: stream.read(65536), b""): + room = _MAX_OUTPUT_CHARS - len(buf) + if room > 0: + buf.extend(chunk[:room]) + if len(chunk) > room: + flag[0] = True + + out, err = bytearray(), bytearray() + out_trunc, err_trunc = [False], [False] + threads = [ + threading.Thread(target=drain, args=(proc.stdout, out, out_trunc)), + threading.Thread(target=drain, args=(proc.stderr, err, err_trunc)), + ] + for t in threads: + t.start() + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + for t in threads: + t.join() + raise + for t in threads: + t.join() + return ( + proc.returncode, + out.decode(errors="replace"), + err.decode(errors="replace"), + out_trunc[0] or err_trunc[0], + ) + + def create(self) -> str: + if shutil.which(self._sbx_path) is None: + raise AirflowOptionalProviderFeatureException( + f"The {self._sbx_path!r} binary was not found on PATH. Install Docker Sandboxes " + "(https://docs.docker.com/ai/sandboxes/) and run 'sbx policy init' once." + ) + name = _new_sandbox_name() + workspace = tempfile.mkdtemp(prefix="airflow-sandbox-ws-") + args = ["create", "--name", name, "--memory", self._memory, "--template", self._image] + if self._cpus is not None: + args += ["--cpus", str(int(self._cpus))] + args += ["shell", workspace] + try: + result = self._run_cli(args, timeout=self._create_timeout) + if result.returncode != 0: + raise RuntimeError(f"'sbx create' failed ({result.returncode}): {result.stderr.strip()}") + except BaseException: + # A timeout, a nonzero exit, or a partial provision all leave a + # possibly-orphaned microVM and a workspace tempdir. Best-effort + # clean both so nothing leaks on the host or in sbx. + with suppress(Exception): + self._run_cli(["rm", "-f", name], timeout=120.0) + shutil.rmtree(workspace, ignore_errors=True) + raise + self._workspaces[name] = workspace + return name + + def run(self, sandbox: str, command: list[str], *, timeout: float) -> SandboxResult: + _validate_positive_finite(timeout, "timeout") + # Round up: GNU timeout treats 0 as "no timeout", so a sub-second value + # must not truncate to it. + seconds = max(1, math.ceil(timeout)) + # GNU ``timeout`` exits 124 when the budget is hit and the command dies to + # the SIGTERM; if the command ignores it, ``--kill-after`` escalates to + # SIGKILL and ``timeout`` exits 137 instead — ambiguous with an OOM kill, + # so 137 counts as a timeout only when the call also outlived the budget. + exec_args = [ + "exec", + sandbox, + "timeout", + f"--kill-after={_KILL_AFTER}", + str(seconds), + *command, + ] + start = time.monotonic() + try: + returncode, stdout, stderr, truncated = self._exec_capped( + exec_args, timeout=timeout + _EXEC_GRACE + ) + except subprocess.TimeoutExpired: + # The CLI never returned: the command may still be running in the + # shared microVM. Destroy it so it cannot continue, and tell the + # toolset to provision a fresh one. + self.destroy(sandbox) + return SandboxResult( + exit_code=-1, + stdout="", + stderr="", + timed_out=True, + sandbox_terminated=True, + ) + elapsed = time.monotonic() - start + return SandboxResult( + exit_code=returncode, + stdout=stdout, + stderr=stderr, + timed_out=returncode == 124 or (returncode == 137 and elapsed >= seconds), + truncated=truncated, + ) + + def destroy(self, sandbox: str) -> None: + # Already-gone is fine — 'sbx rm -f' exits nonzero for a missing sandbox, + # which we ignore so destroy stays idempotent. + try: + self._run_cli(["rm", "-f", sandbox], timeout=120.0) + except subprocess.TimeoutExpired: + pass + finally: + workspace = self._workspaces.pop(sandbox, None) + if workspace is not None: + shutil.rmtree(workspace, ignore_errors=True) diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py index e36690b036df0..2c5845e83ac97 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py @@ -19,8 +19,16 @@ from __future__ import annotations from airflow.providers.common.ai.toolsets.hook import HookToolset +from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset -__all__ = ["AWSToolset", "HookToolset", "MCPToolset", "SQLToolset", "airflow_toolset_to_langchain_tools"] +__all__ = [ + "AWSToolset", + "HookToolset", + "MCPToolset", + "SQLToolset", + "SandboxToolset", + "airflow_toolset_to_langchain_tools", +] def __getattr__(name: str): diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py new file mode 100644 index 0000000000000..640edfe69e3f2 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py @@ -0,0 +1,220 @@ +# 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. +"""Toolset that executes agent-written Python in an isolated sandbox, off the worker.""" + +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING, Any + +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool +from pydantic_core import SchemaValidator, core_schema +from typing_extensions import Self + +from airflow.providers.common.ai.sandbox.base import _validate_positive_finite +from airflow.providers.common.ai.utils.tool_definition import return_schema_kwargs + +if TYPE_CHECKING: + from pydantic_ai._run_context import RunContext + + from airflow.providers.common.ai.sandbox.base import SandboxBackend + +# Deliberately not "run_code": that name is reserved by the Monty code_mode +# meta-tool, and the pinned pydantic-ai-harness rejects a second tool claiming +# it. A distinct name lets this toolset and code_mode coexist. +_TOOL_NAME = "run_python_in_sandbox" + +_RUN_CODE_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "code": {"type": "string", "description": "Python code to execute."}, + }, + "required": ["code"], +} + +# Must stay in sync with _RUN_CODE_SCHEMA. A ValidationError raised here becomes +# a retry prompt (honoring max_retries) instead of a KeyError in call_tool that +# would fail the whole task. +_RUN_CODE_VALIDATOR = SchemaValidator( + core_schema.typed_dict_schema( + {"code": core_schema.typed_dict_field(core_schema.str_schema())}, + extra_behavior="ignore", + ) +) + +_RUN_CODE_DESCRIPTION = ( + "Execute Python code in an isolated sandbox and return a JSON object with " + "exit_code, stdout, stderr, timed_out, and optionally truncated or " + "sandbox_terminated. Airflow does not inject its context, connections, or " + "worker environment. Each call runs a fresh interpreter, but files written " + "by earlier calls persist unless sandbox_terminated is true. Print anything " + "you need to see." +) + + +class SandboxToolset(AbstractToolset[Any]): + """ + Toolset that executes agent-written Python in an isolated sandbox. + + Provides one tool, ``run_python_in_sandbox``. Unlike ``code_mode`` — whose + ``run_code`` executes in-process on the worker via the Monty interpreter — + this ships the code to a disposable sandbox provisioned by the given + :class:`~airflow.providers.common.ai.sandbox.SandboxBackend`. Airflow does + not inject its context, connections, credentials, or worker environment; + custom images and hosted backends can still provide their own credentials. + + The sandbox is created lazily on the first ``run_python_in_sandbox`` call, shared by + every call within the agent run, and destroyed when the run ends. A run + that never calls the tool never provisions one. + + A nonzero exit code or a timeout is normal tool output — the model reads + ``stderr`` and fixes its own code. A backend can terminate and transparently + replace the sandbox when that is required to stop timed-out code. Only + infrastructure exceptions from the backend (daemon unreachable, API + failure) propagate and fail the task. + + :param backend: Sandbox backend that provisions and runs the sandbox, e.g. + :class:`~airflow.providers.common.ai.sandbox.SbxSandboxBackend` or + :class:`~airflow.providers.common.ai.sandbox.IsloSandboxBackend`. + :param timeout: Timeout in seconds for a single ``run_python_in_sandbox`` call. + Default ``300``. + :param python_command: Python executable inside the sandbox. + Default ``"python3"``. + """ + + def __init__( + self, + backend: SandboxBackend, + *, + timeout: float = 300.0, + python_command: str = "python3", + ) -> None: + _validate_positive_finite(timeout, "timeout") + if not python_command: + raise ValueError("python_command must not be empty.") + self._backend = backend + self._timeout = timeout + self._python_command = python_command + self._sandbox: str | None = None + self._create_task: asyncio.Task[str] | None = None + + @property + def id(self) -> str: + return f"sandbox-{self._backend.name}" + + async def for_run(self, ctx: RunContext[Any]) -> AbstractToolset[Any]: + # Per-run isolation: pydantic-ai shares one toolset instance across runs, + # but each run holds its own sandbox in ``_sandbox``/``_create_task``. + # Hand every run a fresh instance so concurrent runs never share a + # sandbox or destroy each other's. The backend keys all state by unique + # sandbox name, so sharing it across runs is safe. + return SandboxToolset(self._backend, timeout=self._timeout, python_command=self._python_command) + + async def __aenter__(self) -> Self: + # The sandbox is provisioned lazily in call_tool, not here: a durable + # replay that only serves cached tool results must not provision one, + # and nothing leaks if the run fails before any tool executes. + return self + + async def __aexit__(self, *args: Any) -> bool | None: + if self._sandbox is not None: + try: + await asyncio.to_thread(self._backend.destroy, self._sandbox) + finally: + # Re-entering the same instance (HITL regenerate_with_feedback + # does) must never reuse a sandbox whose cleanup was attempted. + self._sandbox = None + return None + + async def _ensure_sandbox(self) -> str: + if self._sandbox is not None: + return self._sandbox + if self._create_task is None: + self._create_task = asyncio.create_task(asyncio.to_thread(self._backend.create)) + create_task = self._create_task + try: + sandbox = await asyncio.shield(create_task) + except asyncio.CancelledError: + # A thread cannot be cancelled. Wait until it publishes the handle + # so __aexit__ can destroy a sandbox created during cancellation. + self._sandbox = await create_task + raise + else: + self._sandbox = sandbox + return sandbox + finally: + if self._create_task is create_task: + self._create_task = None + + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: + # sequential=True: every call shares one sandbox and may depend on + # state left by earlier calls — they must not run concurrently. + # return_schema is "string": the tool returns a JSON-encoded string + # (json.dumps), so code mode renders `-> str` instead of `-> Any`. + tool_def = ToolDefinition( + name=_TOOL_NAME, + description=_RUN_CODE_DESCRIPTION, + parameters_json_schema=_RUN_CODE_SCHEMA, + sequential=True, + **return_schema_kwargs({"type": "string"}), + ) + return { + _TOOL_NAME: ToolsetTool( + toolset=self, + tool_def=tool_def, + max_retries=2, + args_validator=_RUN_CODE_VALIDATOR, + ) + } + + async def call_tool( + self, + name: str, + tool_args: dict[str, Any], + ctx: RunContext[Any], + tool: ToolsetTool[Any], + ) -> Any: + if name != _TOOL_NAME: + raise ValueError(f"Unknown tool: {name!r}") + # Backend calls are synchronous and a sandbox run can take minutes, so + # offload them to a thread instead of blocking the event loop (unlike + # SQLToolset, whose queries are expected to be short). + sandbox = await self._ensure_sandbox() + result = await asyncio.to_thread( + self._backend.run, + sandbox, + [self._python_command, "-c", tool_args["code"]], + timeout=self._timeout, + ) + if result.sandbox_terminated: + self._sandbox = None + payload: dict[str, Any] = { + "exit_code": result.exit_code, + "stdout": result.stdout, + "stderr": result.stderr, + "timed_out": result.timed_out, + } + if result.truncated: + payload["truncated"] = True + if result.sandbox_terminated: + payload["sandbox_terminated"] = True + # A nonzero exit or a timeout is still a normal return: a failed *user + # code* run is valid tool output, not a tool failure — the model reads + # stderr and fixes its own code. + return json.dumps(payload) diff --git a/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py b/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py index c48f9ac877af1..7ebaa370e2203 100644 --- a/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py +++ b/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py @@ -52,8 +52,8 @@ AgentSessionData, SessionStatus, ) +from airflow.providers.common.compat.sdk import timezone from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.types import DagRunType diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/__init__.py b/providers/common/ai/tests/unit/common/ai/sandbox/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/sandbox/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py b/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py new file mode 100644 index 0000000000000..84c276d931f85 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py @@ -0,0 +1,74 @@ +# 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. +from __future__ import annotations + +import dataclasses + +import pytest + +from airflow.providers.common.ai.sandbox.base import ( + _MAX_OUTPUT_CHARS, + SandboxBackend, + SandboxResult, + _cap_output, + _new_sandbox_name, + _validate_positive_finite, +) + + +class TestSandboxResult: + def test_defaults(self): + result = SandboxResult(exit_code=0, stdout="out", stderr="err") + assert result.timed_out is False + assert result.truncated is False + assert result.sandbox_terminated is False + + def test_frozen(self): + result = SandboxResult(exit_code=0, stdout="", stderr="") + with pytest.raises(dataclasses.FrozenInstanceError): + result.exit_code = 1 # type: ignore[misc] + + +class TestNewSandboxName: + def test_prefix(self): + assert _new_sandbox_name().startswith("airflow-sandbox-") + + def test_unique(self): + names = {_new_sandbox_name() for _ in range(100)} + assert len(names) == 100 + + +class TestCapOutput: + def test_short_text_is_untouched(self): + assert _cap_output("hello") == ("hello", False) + + def test_long_text_is_capped_and_flagged(self): + text, truncated = _cap_output("x" * (_MAX_OUTPUT_CHARS + 1)) + assert len(text) == _MAX_OUTPUT_CHARS + assert truncated is True + + +class TestSandboxBackend: + def test_abstract_not_instantiable(self): + with pytest.raises(TypeError): + SandboxBackend() # type: ignore[abstract] + + +@pytest.mark.parametrize("value", [0, -1, float("inf"), float("-inf"), float("nan")]) +def test_validate_positive_finite_rejects_invalid_values(value): + with pytest.raises(ValueError, match="timeout must be a positive finite number"): + _validate_positive_finite(value, "timeout") diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_islo.py b/providers/common/ai/tests/unit/common/ai/sandbox/test_islo.py new file mode 100644 index 0000000000000..63ea080c718c5 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_islo.py @@ -0,0 +1,285 @@ +# 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. +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +pytest.importorskip("islo") + +from islo.errors import NotFoundError +from islo.types import LifecyclePolicy + +from airflow.providers.common.ai.sandbox.islo import IsloSandboxBackend + +_BASE_HOOK_PATH = "airflow.providers.common.ai.sandbox.islo.BaseHook" +_ISLO_PATH = "airflow.providers.common.ai.sandbox.islo.Islo" + + +def _make_mock_connection(password="secret-key", host=None, extra=None): + conn = MagicMock() + conn.password = password + conn.host = host + conn.extra_dejson = extra or {} + return conn + + +def _backend_with_mock_client(**kwargs) -> tuple[IsloSandboxBackend, MagicMock]: + backend = IsloSandboxBackend(**kwargs) + mock_client = MagicMock() + mock_client.sandboxes.exec_in_sandbox.return_value.exec_id = "exec-1" + backend._client = mock_client + return backend, mock_client + + +def _make_exec_result( + *, + status="completed", + stdout="out", + stderr="err", + exit_code=0, + truncated=False, +): + result = MagicMock() + result.status = status + result.stdout = stdout + result.stderr = stderr + result.exit_code = exit_code + result.truncated = truncated + return result + + +class TestIsloSandboxBackendClient: + @patch(_ISLO_PATH, autospec=True) + @patch(_BASE_HOOK_PATH, autospec=True) + def test_init_resolves_nothing(self, mock_base_hook, mock_islo_cls): + IsloSandboxBackend() + mock_base_hook.get_connection.assert_not_called() + mock_islo_cls.assert_not_called() + + @patch(_ISLO_PATH, autospec=True) + @patch(_BASE_HOOK_PATH, autospec=True) + def test_connection_maps_to_client_kwargs(self, mock_base_hook, mock_islo_cls): + mock_base_hook.get_connection.return_value = _make_mock_connection( + password=" secret-key ", + host="https://compute.islo.dev", + extra={"base_url": "https://api.islo.dev", "timeout": 30}, + ) + + backend = IsloSandboxBackend("islo_conn") + client = backend._get_client() + + mock_base_hook.get_connection.assert_called_once_with("islo_conn") + mock_islo_cls.assert_called_once_with( + api_key="secret-key", + compute_url="https://compute.islo.dev", + base_url="https://api.islo.dev", + timeout=30.0, + ) + assert client is mock_islo_cls.return_value + + @patch(_ISLO_PATH, autospec=True) + @patch(_BASE_HOOK_PATH, autospec=True) + def test_connection_omits_unset_optional_fields(self, mock_base_hook, mock_islo_cls): + mock_base_hook.get_connection.return_value = _make_mock_connection() + + IsloSandboxBackend()._get_client() + + mock_islo_cls.assert_called_once_with(api_key="secret-key") + + @patch(_ISLO_PATH, autospec=True) + @patch(_BASE_HOOK_PATH, autospec=True) + def test_none_conn_id_uses_sdk_environment_resolution(self, mock_base_hook, mock_islo_cls): + IsloSandboxBackend(islo_conn_id=None)._get_client() + + mock_base_hook.get_connection.assert_not_called() + mock_islo_cls.assert_called_once_with() + + @pytest.mark.parametrize("password", [None, " "]) + @patch(_BASE_HOOK_PATH, autospec=True) + def test_missing_password_raises_clear_error(self, mock_base_hook, password): + mock_base_hook.get_connection.return_value = _make_mock_connection(password=password) + + with pytest.raises(ValueError, match="no password.*islo API key"): + IsloSandboxBackend("islo_conn")._get_client() + + @patch(_ISLO_PATH, autospec=True) + @patch(_BASE_HOOK_PATH, autospec=True) + def test_caches_client(self, mock_base_hook, mock_islo_cls): + mock_base_hook.get_connection.return_value = _make_mock_connection() + + backend = IsloSandboxBackend() + assert backend._get_client() is backend._get_client() + + mock_base_hook.get_connection.assert_called_once() + mock_islo_cls.assert_called_once() + + @patch(_BASE_HOOK_PATH, autospec=True) + def test_rejects_invalid_connection_timeout(self, mock_base_hook): + mock_base_hook.get_connection.return_value = _make_mock_connection(extra={"timeout": 0}) + + with pytest.raises(ValueError, match="connection extra timeout"): + IsloSandboxBackend()._get_client() + + +class TestIsloSandboxBackendCreate: + def test_create_defaults_disable_internet(self): + backend, mock_client = _backend_with_mock_client() + + backend.create() + + kwargs = mock_client.sandboxes.create_sandbox.call_args.kwargs + assert set(kwargs) == {"name", "lifecycle", "internet_enabled"} + assert kwargs["name"].startswith("airflow-sandbox-") + assert kwargs["internet_enabled"] is False + assert kwargs["lifecycle"] == LifecyclePolicy(delete_after=3600) + + def test_create_defers_internet_to_server_when_none(self): + backend, mock_client = _backend_with_mock_client(internet_enabled=None) + + backend.create() + + kwargs = mock_client.sandboxes.create_sandbox.call_args.kwargs + assert set(kwargs) == {"name", "lifecycle"} + + def test_create_passes_set_kwargs(self): + backend, mock_client = _backend_with_mock_client( + image="python:3.12", vcpus=2, memory_mb=1024, internet_enabled=False, delete_after=60 + ) + + backend.create() + + kwargs = mock_client.sandboxes.create_sandbox.call_args.kwargs + assert kwargs["image"] == "python:3.12" + assert kwargs["vcpus"] == 2 + assert kwargs["memory_mb"] == 1024 + assert kwargs["internet_enabled"] is False + assert kwargs["lifecycle"] == LifecyclePolicy(delete_after=60) + + def test_create_returns_server_normalized_name(self): + backend, mock_client = _backend_with_mock_client() + mock_client.sandboxes.create_sandbox.return_value.name = "airflow-sbx-normalized" + + assert backend.create() == "airflow-sbx-normalized" + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"image": ""}, "image"), + ({"vcpus": 0}, "vcpus"), + ({"memory_mb": -1}, "memory_mb"), + ({"delete_after": float("nan")}, "delete_after"), + ], + ) + def test_init_rejects_invalid_resources(self, kwargs, match): + with pytest.raises(ValueError, match=match): + IsloSandboxBackend(**kwargs) + + +class TestIsloSandboxBackendRun: + def test_run_maps_exec_result_and_preserves_truncation(self): + backend, mock_client = _backend_with_mock_client() + mock_client.sandboxes.get_exec_result.return_value = _make_exec_result(exit_code=3, truncated=True) + + result = backend.run("sbx", ["python3", "-c", "x"], timeout=10.0) + + mock_client.sandboxes.exec_in_sandbox.assert_called_once_with( + "sbx", + command=["python3", "-c", "x"], + timeout_secs=10, + ) + mock_client.sandboxes.get_exec_result.assert_called_once_with("sbx", "exec-1") + assert result.exit_code == 3 + assert result.stdout == "out" + assert result.stderr == "err" + assert result.timed_out is False + assert result.truncated is True + + def test_run_maps_server_enforced_timeout(self): + backend, mock_client = _backend_with_mock_client() + mock_client.sandboxes.get_exec_result.return_value = _make_exec_result( + status="timeout", stdout="", stderr="", exit_code=None + ) + + result = backend.run("sbx", ["sleep", "60"], timeout=0.5) + + assert result.timed_out is True + assert result.exit_code == -1 + mock_client.sandboxes.exec_in_sandbox.assert_called_once_with( + "sbx", + command=["sleep", "60"], + timeout_secs=1, + ) + + @patch("airflow.providers.common.ai.sandbox.islo.time.monotonic", side_effect=[0.0, 10.0]) + def test_run_deletes_sandbox_when_server_does_not_stop_command(self, mock_monotonic): + backend, mock_client = _backend_with_mock_client() + + result = backend.run("sbx", ["sleep", "60"], timeout=1.0) + + assert result.timed_out is True + assert result.sandbox_terminated is True + mock_client.sandboxes.get_exec_result.assert_not_called() + mock_client.sandboxes.delete_sandbox.assert_called_once_with(sandbox_name="sbx") + + @patch("airflow.providers.common.ai.sandbox.islo.time.monotonic", side_effect=[0.0, 10.0]) + def test_run_termination_cleanup_failure_does_not_fail_task(self, mock_monotonic): + # A transient delete failure during the timeout-termination path must not + # propagate — the server-side TTL is the backstop. + backend, mock_client = _backend_with_mock_client() + mock_client.sandboxes.delete_sandbox.side_effect = RuntimeError("transient 503") + + result = backend.run("sbx", ["sleep", "60"], timeout=1.0) + + assert result.timed_out is True + assert result.sandbox_terminated is True + + def test_run_caps_oversized_output(self): + from airflow.providers.common.ai.sandbox.base import _MAX_OUTPUT_CHARS + + backend, mock_client = _backend_with_mock_client() + mock_client.sandboxes.get_exec_result.return_value = _make_exec_result( + stdout="x" * (_MAX_OUTPUT_CHARS + 50), stderr="y", truncated=False + ) + + result = backend.run("sbx", ["python3", "-c", "x"], timeout=10.0) + + assert len(result.stdout) == _MAX_OUTPUT_CHARS + assert result.truncated is True + + @pytest.mark.parametrize("timeout", [0, -1, float("nan")]) + def test_run_rejects_invalid_timeout(self, timeout): + backend, _ = _backend_with_mock_client() + + with pytest.raises(ValueError, match="timeout"): + backend.run("sbx", ["true"], timeout=timeout) + + +class TestIsloSandboxBackendDestroy: + def test_destroy_deletes_sandbox(self): + backend, mock_client = _backend_with_mock_client() + + backend.destroy("airflow-sbx-abc") + + mock_client.sandboxes.delete_sandbox.assert_called_once_with(sandbox_name="airflow-sbx-abc") + + def test_destroy_swallows_missing_sandbox(self): + backend, mock_client = _backend_with_mock_client() + mock_client.sandboxes.delete_sandbox.side_effect = NotFoundError(body=None) + + backend.destroy("airflow-sbx-abc") # Does not raise. diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py b/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py new file mode 100644 index 0000000000000..856115991b154 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py @@ -0,0 +1,243 @@ +# 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. +from __future__ import annotations + +import subprocess +import sys +from unittest.mock import patch + +import pytest + +from airflow.providers.common.ai.sandbox.base import _MAX_OUTPUT_CHARS +from airflow.providers.common.ai.sandbox.sbx import SbxSandboxBackend +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +_MOD = "airflow.providers.common.ai.sandbox.sbx" + + +def _completed(returncode: int, stdout: str = "", stderr: str = "") -> subprocess.CompletedProcess: + return subprocess.CompletedProcess(args=["sbx"], returncode=returncode, stdout=stdout, stderr=stderr) + + +class TestSbxSandboxBackendInit: + def test_init_makes_no_cli_calls(self): + with patch(f"{_MOD}.subprocess.run") as run, patch(f"{_MOD}.shutil.which") as which: + SbxSandboxBackend() + run.assert_not_called() + which.assert_not_called() + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"image": ""}, "image"), + ({"memory": ""}, "memory"), + ({"cpus": 0}, "cpus"), + ({"sbx_path": ""}, "sbx_path"), + ({"create_timeout": float("nan")}, "create_timeout"), + ], + ) + def test_init_rejects_invalid(self, kwargs, match): + with pytest.raises(ValueError, match=match): + SbxSandboxBackend(**kwargs) + + +class TestSbxSandboxBackendCreate: + @patch(f"{_MOD}.tempfile.mkdtemp", return_value="/tmp/ws") + @patch(f"{_MOD}.shutil.which", return_value="/usr/local/bin/sbx") + @patch(f"{_MOD}.subprocess.run", return_value=_completed(0)) + def test_create_builds_command_and_returns_name(self, run, which, mkdtemp): + backend = SbxSandboxBackend() + + name = backend.create() + + assert name.startswith("airflow-sandbox-") + args = run.call_args.args[0] + assert args == [ + "sbx", "create", "--name", name, "--memory", "2g", + "--template", "python:3.12-slim", "shell", "/tmp/ws", + ] # fmt: skip + assert backend._workspaces[name] == "/tmp/ws" + + @patch(f"{_MOD}.tempfile.mkdtemp", return_value="/tmp/ws") + @patch(f"{_MOD}.shutil.which", return_value="/usr/local/bin/sbx") + @patch(f"{_MOD}.subprocess.run", return_value=_completed(0)) + def test_create_passes_cpus_image_and_memory(self, run, which, mkdtemp): + backend = SbxSandboxBackend(image="python:3.13-slim", memory="4g", cpus=2) + + backend.create() + + args = run.call_args.args[0] + assert args[:8] == [ + "sbx", "create", "--name", args[3], "--memory", "4g", "--template", "python:3.13-slim", + ] # fmt: skip + assert "--cpus" in args + assert args[args.index("--cpus") + 1] == "2" + + @patch(f"{_MOD}.shutil.which", return_value=None) + def test_create_raises_when_binary_missing(self, which): + with pytest.raises(AirflowOptionalProviderFeatureException, match="sbx"): + SbxSandboxBackend().create() + + @patch(f"{_MOD}.shutil.rmtree") + @patch(f"{_MOD}.tempfile.mkdtemp", return_value="/tmp/ws") + @patch(f"{_MOD}.shutil.which", return_value="/usr/local/bin/sbx") + @patch(f"{_MOD}.subprocess.run", return_value=_completed(1, stderr="boom")) + def test_create_cleans_up_on_nonzero_exit(self, run, which, mkdtemp, rmtree): + backend = SbxSandboxBackend() + + with pytest.raises(RuntimeError, match="sbx create.*boom"): + backend.create() + + rmtree.assert_called_once_with("/tmp/ws", ignore_errors=True) + # A best-effort 'sbx rm' of the possibly-orphaned sandbox was attempted. + assert any(call.args[0][:2] == ["sbx", "rm"] for call in run.call_args_list) + assert backend._workspaces == {} + + @patch(f"{_MOD}.shutil.rmtree") + @patch(f"{_MOD}.tempfile.mkdtemp", return_value="/tmp/ws") + @patch(f"{_MOD}.shutil.which", return_value="/usr/local/bin/sbx") + @patch(f"{_MOD}.subprocess.run", side_effect=subprocess.TimeoutExpired(cmd="sbx", timeout=600)) + def test_create_cleans_up_on_timeout(self, run, which, mkdtemp, rmtree): + backend = SbxSandboxBackend() + + with pytest.raises(subprocess.TimeoutExpired): + backend.create() + + rmtree.assert_called_once_with("/tmp/ws", ignore_errors=True) + assert backend._workspaces == {} + + +class TestSbxSandboxBackendRun: + def test_run_wraps_in_timeout_and_maps_output(self): + backend = SbxSandboxBackend() + with patch.object(backend, "_exec_capped", return_value=(0, "42\n", "", False)) as ec: + result = backend.run("sbx-1", ["python3", "-c", "print(42)"], timeout=10.0) + + args = ec.call_args.args[0] + assert args == [ + "exec", "sbx-1", "timeout", "--kill-after=10", "10", "python3", "-c", "print(42)", + ] # fmt: skip + assert ec.call_args.kwargs["timeout"] == pytest.approx(40.0) # 10 + _EXEC_GRACE + assert result.exit_code == 0 + assert result.stdout == "42\n" + assert result.timed_out is False + assert result.truncated is False + + def test_run_exit_124_is_timeout(self): + # GNU timeout exits exactly 124 when the budget is hit. + backend = SbxSandboxBackend() + with patch.object(backend, "_exec_capped", return_value=(124, "", "", False)): + result = backend.run("sbx-1", ["sleep", "60"], timeout=10.0) + assert result.timed_out is True + + @pytest.mark.parametrize("exit_code", [0, 1]) + def test_run_normal_exit_is_not_timeout(self, exit_code): + backend = SbxSandboxBackend() + with patch.object(backend, "_exec_capped", return_value=(exit_code, "", "x", False)): + result = backend.run("sbx-1", ["python3", "-c", "..."], timeout=10.0) + assert result.timed_out is False + assert result.exit_code == exit_code + + @pytest.mark.parametrize( + ("elapsed", "expected_timed_out"), + [ + # SIGKILL escalation after the budget elapsed (command ignored SIGTERM). + (15.0, True), + # Fast SIGKILL well before the budget — an OOM kill, not a timeout. + (2.0, False), + ], + ) + def test_run_exit_137_is_timeout_only_after_budget(self, elapsed, expected_timed_out): + backend = SbxSandboxBackend() + with ( + patch(f"{_MOD}.time.monotonic", side_effect=[100.0, 100.0 + elapsed]), + patch.object(backend, "_exec_capped", return_value=(137, "", "", False)), + ): + result = backend.run("sbx-1", ["python3", "-c", "..."], timeout=10.0) + assert result.timed_out is expected_timed_out + assert result.exit_code == 137 + + def test_run_propagates_truncation_flag(self): + backend = SbxSandboxBackend() + with patch.object(backend, "_exec_capped", return_value=(0, "x", "", True)): + result = backend.run("sbx-1", ["yes"], timeout=10.0) + assert result.truncated is True + + def test_run_cli_hang_destroys_and_replaces(self): + backend = SbxSandboxBackend() + with ( + patch.object(backend, "_exec_capped", side_effect=subprocess.TimeoutExpired("sbx", 40)), + patch(f"{_MOD}.subprocess.run", return_value=_completed(0)) as run, + ): + result = backend.run("sbx-1", ["sleep", "999"], timeout=10.0) + + assert result.timed_out is True + assert result.sandbox_terminated is True + assert any(call.args[0][:2] == ["sbx", "rm"] for call in run.call_args_list) + + @pytest.mark.parametrize("timeout", [0, -1, float("nan")]) + def test_run_rejects_invalid_timeout(self, timeout): + with pytest.raises(ValueError, match="timeout"): + SbxSandboxBackend().run("sbx-1", ["true"], timeout=timeout) + + +class TestSbxSandboxBackendExecCapped: + """Exercises the real Popen + capped-drain path (no sbx needed: run Python directly).""" + + def test_bounds_stdout_and_stderr_and_flags_truncation(self): + backend = SbxSandboxBackend(sbx_path=sys.executable) + code = "import sys; sys.stdout.write('x'*200000); sys.stderr.write('y'*200000)" + rc, out, err, truncated = backend._exec_capped(["-c", code], timeout=30.0) + + assert rc == 0 + assert len(out) == _MAX_OUTPUT_CHARS + assert len(err) == _MAX_OUTPUT_CHARS + assert truncated is True + + def test_small_output_is_untruncated(self): + backend = SbxSandboxBackend(sbx_path=sys.executable) + rc, out, err, truncated = backend._exec_capped(["-c", "print('hi')"], timeout=30.0) + assert rc == 0 + assert out.strip() == "hi" + assert truncated is False + + def test_raises_on_timeout(self): + backend = SbxSandboxBackend(sbx_path=sys.executable) + with pytest.raises(subprocess.TimeoutExpired): + backend._exec_capped(["-c", "import time; time.sleep(30)"], timeout=1.0) + + +class TestSbxSandboxBackendDestroy: + @patch(f"{_MOD}.shutil.rmtree") + @patch(f"{_MOD}.subprocess.run", return_value=_completed(0)) + def test_destroy_removes_sandbox_and_workspace(self, run, rmtree): + backend = SbxSandboxBackend() + backend._workspaces["sbx-1"] = "/tmp/ws" + + backend.destroy("sbx-1") + + assert run.call_args.args[0] == ["sbx", "rm", "-f", "sbx-1"] + rmtree.assert_called_once_with("/tmp/ws", ignore_errors=True) + assert backend._workspaces == {} + + @patch(f"{_MOD}.subprocess.run", return_value=_completed(1, stderr="not found")) + def test_destroy_ignores_missing_sandbox(self, run): + SbxSandboxBackend().destroy("gone") # Does not raise. + + @patch(f"{_MOD}.subprocess.run", side_effect=subprocess.TimeoutExpired(cmd="sbx", timeout=120)) + def test_destroy_swallows_cli_timeout(self, run): + SbxSandboxBackend().destroy("sbx-1") # Does not raise. diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py new file mode 100644 index 0000000000000..c445da887d8e4 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py @@ -0,0 +1,348 @@ +# 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. +from __future__ import annotations + +import asyncio +import json +import threading +from unittest.mock import MagicMock + +import pytest +from pydantic_core import ValidationError + +from airflow.providers.common.ai.sandbox.base import SandboxBackend, SandboxResult +from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset +from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA + + +class _RecordingBackend(SandboxBackend): + """In-test backend recording every lifecycle call; each create hands out a fresh handle.""" + + name = "fake" + + def __init__( + self, + result: SandboxResult | None = None, + run_error: Exception | None = None, + destroy_error: Exception | None = None, + ) -> None: + self.result = result or SandboxResult(exit_code=0, stdout="hello", stderr="") + self.run_error = run_error + self.destroy_error = destroy_error + self.created: list[str] = [] + self.runs: list[tuple[str, list[str], float]] = [] + self.destroyed: list[str] = [] + + def create(self) -> str: + handle = f"sbx-{len(self.created) + 1}" + self.created.append(handle) + return handle + + def run(self, sandbox: str, command: list[str], *, timeout: float) -> SandboxResult: + if self.run_error is not None: + raise self.run_error + self.runs.append((sandbox, list(command), timeout)) + return self.result + + def destroy(self, sandbox: str) -> None: + self.destroyed.append(sandbox) + if self.destroy_error is not None: + raise self.destroy_error + + +def _call_run_code(ts: SandboxToolset, code: str = "print('hi')"): + return asyncio.run( + ts.call_tool("run_python_in_sandbox", {"code": code}, ctx=MagicMock(), tool=MagicMock()) + ) + + +class TestSandboxToolsetInit: + def test_id_includes_backend_name(self): + ts = SandboxToolset(_RecordingBackend()) + assert ts.id == "sandbox-fake" + + @pytest.mark.parametrize("timeout", [0, -1, float("nan")]) + def test_rejects_invalid_timeout(self, timeout): + with pytest.raises(ValueError, match="timeout"): + SandboxToolset(_RecordingBackend(), timeout=timeout) + + def test_rejects_empty_python_command(self): + with pytest.raises(ValueError, match="python_command"): + SandboxToolset(_RecordingBackend(), python_command="") + + +class TestSandboxToolsetForRun: + def test_for_run_hands_each_run_a_fresh_instance(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend, timeout=42.0, python_command="py") + + a = asyncio.run(ts.for_run(MagicMock())) + b = asyncio.run(ts.for_run(MagicMock())) + + # Distinct instances so concurrent runs never share sandbox state... + assert a is not ts + assert a is not b + # ...but the same backend and configuration are carried over. + assert a._backend is backend + assert a._timeout == 42.0 + assert a._python_command == "py" + assert a._sandbox is None + + +class TestSandboxToolsetGetTools: + def test_returns_run_code_tool(self): + ts = SandboxToolset(_RecordingBackend()) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + assert set(tools.keys()) == {"run_python_in_sandbox"} + + def test_tool_definition_shape(self): + ts = SandboxToolset(_RecordingBackend()) + tool = asyncio.run(ts.get_tools(ctx=MagicMock()))["run_python_in_sandbox"] + assert tool.tool_def.description + assert tool.tool_def.sequential is True + assert tool.tool_def.parameters_json_schema["required"] == ["code"] + assert tool.tool_def.parameters_json_schema["properties"]["code"]["type"] == "string" + assert tool.max_retries == 2 + + @pytest.mark.parametrize("bad_args", [{}, {"code": 42}, {"script": "print(1)"}]) + def test_args_validator_rejects_malformed_args(self, bad_args): + """A malformed call must raise ValidationError so pydantic-ai retries instead of failing the task.""" + ts = SandboxToolset(_RecordingBackend()) + tool = asyncio.run(ts.get_tools(ctx=MagicMock()))["run_python_in_sandbox"] + + with pytest.raises(ValidationError): + tool.args_validator.validate_python(bad_args) + with pytest.raises(ValidationError): + tool.args_validator.validate_json(json.dumps(bad_args)) + + def test_args_validator_accepts_valid_args_and_ignores_extras(self): + ts = SandboxToolset(_RecordingBackend()) + tool = asyncio.run(ts.get_tools(ctx=MagicMock()))["run_python_in_sandbox"] + + args = {"code": "print(1)", "unexpected": True} + assert tool.args_validator.validate_python(args) == {"code": "print(1)"} + assert tool.args_validator.validate_json(json.dumps(args)) == {"code": "print(1)"} + + @pytest.mark.skipif( + not _SUPPORTS_RETURN_SCHEMA, reason="pydantic-ai too old for ToolDefinition.return_schema" + ) + def test_tool_declares_string_return_schema(self): + # run_code returns a JSON-encoded string, so code mode should see `-> str`. + ts = SandboxToolset(_RecordingBackend()) + tool = asyncio.run(ts.get_tools(ctx=MagicMock()))["run_python_in_sandbox"] + assert tool.tool_def.return_schema == {"type": "string"} + + +class TestSandboxToolsetCallTool: + def test_creates_sandbox_lazily_on_first_call(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend) + assert backend.created == [] + + _call_run_code(ts) + + assert backend.created == ["sbx-1"] + + def test_reuses_one_sandbox_across_calls(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend) + + _call_run_code(ts, "a = 1") + _call_run_code(ts, "print(a)") + + assert backend.created == ["sbx-1"] + assert [run[0] for run in backend.runs] == ["sbx-1", "sbx-1"] + + def test_runs_code_with_python_command_and_timeout(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend, timeout=10.0, python_command="python3.12") + + _call_run_code(ts, "print('hi')") + + assert backend.runs == [("sbx-1", ["python3.12", "-c", "print('hi')"], 10.0)] + + def test_returns_json_payload(self): + backend = _RecordingBackend(result=SandboxResult(exit_code=0, stdout="hello", stderr="")) + ts = SandboxToolset(backend) + + payload = json.loads(_call_run_code(ts)) + + assert payload == {"exit_code": 0, "stdout": "hello", "stderr": "", "timed_out": False} + + def test_truncated_key_only_present_when_set(self): + backend = _RecordingBackend(result=SandboxResult(exit_code=0, stdout="x", stderr="", truncated=True)) + ts = SandboxToolset(backend) + + payload = json.loads(_call_run_code(ts)) + + assert payload["truncated"] is True + + def test_nonzero_exit_is_returned_not_raised(self): + """A failed user-code run is valid tool output — the model reads stderr and fixes its code.""" + backend = _RecordingBackend( + result=SandboxResult(exit_code=1, stdout="", stderr="NameError: name 'x' is not defined") + ) + ts = SandboxToolset(backend) + + payload = json.loads(_call_run_code(ts)) + + assert payload["exit_code"] == 1 + assert "NameError" in payload["stderr"] + + def test_timeout_is_returned_not_raised(self): + backend = _RecordingBackend(result=SandboxResult(exit_code=-1, stdout="", stderr="", timed_out=True)) + ts = SandboxToolset(backend) + + payload = json.loads(_call_run_code(ts)) + + assert payload["timed_out"] is True + + def test_terminated_sandbox_is_reported_and_recreated(self): + backend = _RecordingBackend( + result=SandboxResult( + exit_code=-1, + stdout="", + stderr="", + timed_out=True, + sandbox_terminated=True, + ) + ) + ts = SandboxToolset(backend) + + first_payload = json.loads(_call_run_code(ts)) + _call_run_code(ts) + + assert first_payload["sandbox_terminated"] is True + assert backend.created == ["sbx-1", "sbx-2"] + + def test_backend_exception_propagates(self): + """Infrastructure failures are not tool output; they fail the task.""" + backend = _RecordingBackend(run_error=RuntimeError("daemon unreachable")) + ts = SandboxToolset(backend) + + with pytest.raises(RuntimeError, match="daemon unreachable"): + _call_run_code(ts) + + def test_unknown_tool_raises(self): + ts = SandboxToolset(_RecordingBackend()) + + with pytest.raises(ValueError, match="Unknown tool"): + asyncio.run(ts.call_tool("other_tool", {}, ctx=MagicMock(), tool=MagicMock())) + + +class TestSandboxToolsetLifecycle: + def test_destroys_sandbox_on_exit(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend) + + async def scenario(): + async with ts: + await ts.call_tool("run_python_in_sandbox", {"code": "1"}, ctx=MagicMock(), tool=MagicMock()) + + asyncio.run(scenario()) + + assert backend.destroyed == ["sbx-1"] + + def test_exit_without_use_destroys_nothing(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend) + + async def scenario(): + async with ts: + pass + + asyncio.run(scenario()) + + assert backend.created == [] + assert backend.destroyed == [] + + def test_destroys_sandbox_even_when_run_raises(self): + backend = _RecordingBackend() + ts = SandboxToolset(backend) + + async def scenario(): + async with ts: + backend.run_error = RuntimeError("boom") + await ts.call_tool("run_python_in_sandbox", {"code": "1"}, ctx=MagicMock(), tool=MagicMock()) + + with pytest.raises(RuntimeError, match="boom"): + asyncio.run(scenario()) + + assert backend.destroyed == ["sbx-1"] + + def test_reentry_creates_fresh_sandbox(self): + """HITL regenerate_with_feedback re-enters the same instance; it must not reuse a destroyed sandbox.""" + backend = _RecordingBackend() + ts = SandboxToolset(backend) + + async def one_run(): + async with ts: + await ts.call_tool("run_python_in_sandbox", {"code": "1"}, ctx=MagicMock(), tool=MagicMock()) + + asyncio.run(one_run()) + asyncio.run(one_run()) + + assert backend.created == ["sbx-1", "sbx-2"] + assert backend.destroyed == ["sbx-1", "sbx-2"] + + def test_reentry_creates_fresh_sandbox_after_destroy_error(self): + backend = _RecordingBackend(destroy_error=RuntimeError("delete failed")) + ts = SandboxToolset(backend) + + async def one_run(): + async with ts: + await ts.call_tool("run_python_in_sandbox", {"code": "1"}, ctx=MagicMock(), tool=MagicMock()) + + with pytest.raises(RuntimeError, match="delete failed"): + asyncio.run(one_run()) + + backend.destroy_error = None + asyncio.run(one_run()) + + assert backend.created == ["sbx-1", "sbx-2"] + assert backend.destroyed == ["sbx-1", "sbx-2"] + + def test_cancellation_during_create_still_destroys_sandbox(self): + create_started = threading.Event() + allow_create = threading.Event() + + class BlockingCreateBackend(_RecordingBackend): + def create(self) -> str: + create_started.set() + if not allow_create.wait(timeout=5): + raise RuntimeError("test timed out waiting to release create") + return super().create() + + backend = BlockingCreateBackend() + ts = SandboxToolset(backend) + + async def scenario(): + async with ts: + call = asyncio.create_task( + ts.call_tool("run_python_in_sandbox", {"code": "1"}, ctx=MagicMock(), tool=MagicMock()) + ) + assert await asyncio.to_thread(create_started.wait, 1) + call.cancel() + allow_create.set() + with pytest.raises(asyncio.CancelledError): + await call + + asyncio.run(scenario()) + + assert backend.created == ["sbx-1"] + assert backend.runs == [] + assert backend.destroyed == ["sbx-1"] diff --git a/uv.lock b/uv.lock index 8b5cea2e6a4a0..f8c5659552c72 100644 --- a/uv.lock +++ b/uv.lock @@ -64,9 +64,9 @@ apache-airflow-providers-apache-cassandra = false apache-airflow-providers-asana = false apache-airflow-providers-oracle = false apache-airflow-providers-mysql = false +apache-airflow-providers-teradata = false apache-airflow-providers-alibaba = false apache-airflow-providers-microsoft-mssql = false -apache-airflow-providers-teradata = false apache-airflow-providers-jdbc = false apache-airflow-helm-chart = false apache-airflow-providers-anthropic = false @@ -4447,6 +4447,9 @@ parquet = [ pdf = [ { name = "pypdf" }, ] +sandbox-islo = [ + { name = "islo" }, +] shields = [ { name = "pydantic-ai-shields" }, ] @@ -4469,6 +4472,7 @@ dev = [ { name = "apache-airflow-providers-git" }, { name = "apache-airflow-providers-standard" }, { name = "apache-airflow-task-sdk" }, + { name = "islo" }, { name = "langchain" }, { name = "llama-index-core" }, { name = "llama-index-embeddings-openai" }, @@ -4495,6 +4499,7 @@ requires-dist = [ { name = "dataclasses-json", marker = "extra == 'llamaindex'", specifier = ">=0.6.7" }, { name = "fastavro", marker = "python_full_version >= '3.14' and extra == 'avro'", specifier = ">=1.12.1" }, { name = "fastavro", marker = "python_full_version < '3.14' and extra == 'avro'", specifier = ">=1.10.0" }, + { name = "islo", marker = "extra == 'sandbox-islo'", specifier = ">=0.3.9" }, { name = "langchain", marker = "extra == 'langchain'", specifier = ">=1.0.0" }, { name = "llama-index-core", marker = "extra == 'llamaindex'", specifier = ">=0.13.0" }, { name = "llama-index-embeddings-openai", marker = "extra == 'llamaindex'", specifier = ">=0.6.0" }, @@ -4514,7 +4519,7 @@ requires-dist = [ { name = "python-docx", marker = "extra == 'docx'", specifier = ">=1.0.0" }, { name = "sqlglot", marker = "extra == 'sql'", specifier = ">=30.0.0" }, ] -provides-extras = ["anthropic", "bedrock", "google", "openai", "mcp", "code-mode", "shields", "skills", "avro", "parquet", "sql", "aws", "common-sql", "langchain", "llamaindex", "pdf", "docx", "git", "amazon"] +provides-extras = ["anthropic", "bedrock", "google", "openai", "mcp", "code-mode", "shields", "skills", "sandbox-islo", "avro", "parquet", "sql", "aws", "common-sql", "langchain", "llamaindex", "pdf", "docx", "git", "amazon"] [package.metadata.requires-dev] dev = [ @@ -4527,6 +4532,7 @@ dev = [ { name = "apache-airflow-providers-git", editable = "providers/git" }, { name = "apache-airflow-providers-standard", editable = "providers/standard" }, { name = "apache-airflow-task-sdk", editable = "task-sdk" }, + { name = "islo", specifier = ">=0.3.9" }, { name = "langchain", specifier = ">=1.0.0" }, { name = "llama-index-core", specifier = ">=0.13.0" }, { name = "llama-index-embeddings-openai", specifier = ">=0.6.0" }, @@ -14769,6 +14775,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" }, ] +[[package]] +name = "islo" +version = "0.3.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8d/2a/8a1fc0d5a056deda8a54a47ddb62a0f75fff7fd341f7f1542f7587b2d4cc/islo-0.3.11.tar.gz", hash = "sha256:7eda29022686429f6308c1ffa71da6d3747d3594935bb8941184be3d1587119c", size = 111539, upload-time = "2026-07-12T14:39:02.85Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/f9/b162c490115383fe83409c2cecb86b6a731a82dc079b9aace85505b3f759/islo-0.3.11-py3-none-any.whl", hash = "sha256:fbfd3ea7048d0cf820ed33897d6556edff38eed8598ef0b89d61ee803ed5d3a3", size = 258479, upload-time = "2026-07-12T14:39:01.348Z" }, +] + [[package]] name = "isodate" version = "0.7.2"