zozo123 commented on code in PR #73961: URL: https://github.com/apache/airflow/pull/73961#discussion_r4212119774
########## providers/common/ai/tests/unit/common/ai/sandbox/test_openshell.py: ########## @@ -0,0 +1,1002 @@ +# 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 builtins +import copy +import os +import shlex +import shutil +import signal +import subprocess +import sys +import time +from unittest import mock + +import pytest + +pytest.importorskip("openshell") + +import grpc +from openshell import SandboxClient +from openshell._proto import openshell_pb2, sandbox_pb2 + +from airflow.providers.common.ai.sandbox.base import ( + SandboxError, + SandboxFileTooLargeError, + SandboxSpec, + SandboxTerminalError, +) +from airflow.providers.common.ai.sandbox.openshell import ( + _EXEC_GRACE, + _READ_WRITE_PATHS, + _RUN_WRAPPER, + _STAGING_FAILED, + OpenShellSandboxBackend, +) + +_MODULE = "airflow.providers.common.ai.sandbox.openshell" +_MONOTONIC_PATH = f"{_MODULE}.time.monotonic" + + +class _RpcError(grpc.RpcError): + def __init__(self, code: grpc.StatusCode, details: str = "") -> None: + super().__init__(details) + self._code = code + self._details = details + + def code(self) -> grpc.StatusCode: + return self._code + + def details(self) -> str: + return self._details + + +class _Stream: + """An ExecSandbox response stream: events, optionally ending in an error.""" + + def __init__(self, events, error: Exception | None = None) -> None: + self._events = iter(events) + self._error = error + self.cancel = mock.Mock() + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._events) + except StopIteration: + if self._error is not None: + raise self._error from None + raise + + +def _stdout(data: bytes): + return openshell_pb2.ExecSandboxEvent(stdout=openshell_pb2.ExecSandboxStdout(data=data)) + + +def _stderr(data: bytes): + return openshell_pb2.ExecSandboxEvent(stderr=openshell_pb2.ExecSandboxStderr(data=data)) + + +def _exit(code: int): + return openshell_pb2.ExecSandboxEvent(exit=openshell_pb2.ExecSandboxExit(exit_code=code)) + + +def _result(code: int = 0, out: bytes = b"", err: bytes = b"") -> _Stream: + events = [] + if out: + events.append(_stdout(out)) + if err: + events.append(_stderr(err)) + events.append(_exit(code)) + return _Stream(events) + + +def _sandbox(phase: openshell_pb2.SandboxPhase = openshell_pb2.SANDBOX_PHASE_READY, *conditions): + response = openshell_pb2.SandboxResponse() + response.sandbox.status.phase = phase + for reason, message in conditions: + response.sandbox.status.conditions.add(type="Ready", status="False", reason=reason, message=message) + return response + + +def _config( + hosts=(), + *, + source: sandbox_pb2.PolicySource = sandbox_pb2.POLICY_SOURCE_SANDBOX, + admitted: bool = True, + landlock: str = "hard_requirement", + approval_mode: str | None = None, + proposals: bool = False, + extra_rule: tuple[str, int] | None = None, + allowed_ips=(), + binaries=("/**",), + middlewares=(), +): + config = sandbox_pb2.GetSandboxConfigResponse( + version=1, policy_source=source, configuration_admitted=admitted + ) + config.policy.landlock.compatibility = landlock + for name in middlewares: + # Reading a missing key of a message map adds it. + config.policy.network_middlewares[name] + if hosts or allowed_ips: + rule = config.policy.network_policies["airflow-egress"] + rule.name = "airflow-egress" + for host in hosts: + rule.endpoints.add(host=host, ports=[443], allowed_ips=list(allowed_ips)) + if not hosts: + rule.endpoints.add(ports=[443], allowed_ips=list(allowed_ips)) + rule.binaries.extend(sandbox_pb2.NetworkBinary(path=path) for path in binaries) + if extra_rule is not None: + rule = config.policy.network_policies["allow_extra"] + rule.endpoints.add(host=extra_rule[0], port=extra_rule[1]) + rule.binaries.add(path="/usr/local/bin/python3.12") + if approval_mode is not None: + config.settings["proposal_approval_mode"].value.string_value = approval_mode + config.settings["proposal_approval_mode"].scope = sandbox_pb2.SETTING_SCOPE_GLOBAL + if proposals: + config.settings["agent_policy_proposals_enabled"].value.bool_value = True + return config + + +def _backend(**kwargs) -> tuple[OpenShellSandboxBackend, mock.MagicMock]: + backend = OpenShellSandboxBackend(gateway="test-gateway", **kwargs) + client = mock.MagicMock(spec=SandboxClient) + client._stub = mock.MagicMock(spec=["GetSandbox", "GetSandboxConfig", "ExecSandbox"]) + client._stub.GetSandbox.return_value = _sandbox() + client._stub.GetSandboxConfig.return_value = _config() + client._stub.ExecSandbox.return_value = _result() + backend._client = client + return backend, client + + +def _exec_request(client, call: int = -1): + return client._stub.ExecSandbox.call_args_list[call].args[0] + + [email protected](autouse=True) +def _no_sleep(): + with mock.patch(f"{_MODULE}.time.sleep", autospec=True): + yield + + +def test_missing_sdk_error_is_actionable(): + real_import = builtins.__import__ + + def blocked_import(name, *args, **kwargs): + if name.startswith("openshell"): + raise ImportError("blocked for test") + return real_import(name, *args, **kwargs) + + backend = OpenShellSandboxBackend() + with mock.patch("builtins.__import__", side_effect=blocked_import): + with pytest.raises(SandboxTerminalError, match=r"\[openshell\]"): + backend.create() + + [email protected]( + ("kwargs", "message"), + [ + ({"gateway": ""}, "gateway"), + ({"workspace": ""}, "workspace"), + ({"image": ""}, "image"), + ({"cpu": ""}, "cpu"), + ({"memory": ""}, "memory"), + ({"ready_timeout": 0}, "ready_timeout"), + ({"request_timeout": float("inf")}, "request_timeout"), + ], +) +def test_constructor_rejects_invalid_values(kwargs, message): + with pytest.raises(ValueError, match=message): + OpenShellSandboxBackend(**kwargs) + + +class TestClient: + @mock.patch("openshell.SandboxClient.from_active_cluster", autospec=True) + def test_construction_reads_no_gateway_registration(self, from_active_cluster): + OpenShellSandboxBackend(gateway="prod") + + from_active_cluster.assert_not_called() + + def test_the_backend_can_be_deep_copied(self): + backend = OpenShellSandboxBackend(gateway="prod") + + assert copy.deepcopy(backend)._gateway == "prod" + + @mock.patch("openshell.SandboxClient.from_active_cluster", autospec=True) + def test_the_cli_gateway_registration_is_loaded_once(self, from_active_cluster): + backend = OpenShellSandboxBackend(gateway="prod", request_timeout=12) + + assert backend._get_client() is backend._get_client() + + from_active_cluster.assert_called_once_with(cluster="prod", timeout=12) + + @mock.patch("openshell.SandboxClient.from_active_cluster", autospec=True) + def test_a_missing_registration_is_terminal(self, from_active_cluster): + from openshell import SandboxError as SdkError + + from_active_cluster.side_effect = SdkError("gateway 'prod' not found") + + with pytest.raises(SandboxTerminalError, match="gateway 'prod' not found"): + OpenShellSandboxBackend(gateway="prod")._get_client() + + +class TestSpecRefusals: + @pytest.mark.parametrize( + ("spec", "message"), + [ + (SandboxSpec(owner="dag/run"), "owner"), + (SandboxSpec(allow_egress_to_cidrs=["203.0.113.0/24"]), "allow_egress_to_cidrs"), + (SandboxSpec(block_network=False), "open network"), + (SandboxSpec(block_network=False, allow_egress_to=["pypi.org"]), "open network"), + (SandboxSpec(allow_egress_to="pypi.org"), "not one string"), + (SandboxSpec(env={"FOO": 1}), "strings to strings"), # type: ignore[dict-item] + ], + ) + def test_unenforceable_specs_are_refused_before_the_gateway_is_asked(self, spec, message): + backend, client = _backend() + + with pytest.raises(SandboxTerminalError, match=message): + backend.create(spec=spec) + + client.create.assert_not_called() + + @pytest.mark.parametrize( + "host", + ["https://pypi.org";, "pypi.org:443", "203.0.113.7", "localhost", "*.com", "*", "pypi.org/simple", ""], + ) + def test_entries_that_are_not_hostnames_are_refused(self, host): + backend, client = _backend() + + with pytest.raises(SandboxTerminalError, match="bare hostnames"): + backend.create(spec=SandboxSpec(allow_egress_to=[host])) + + client.create.assert_not_called() + + @pytest.mark.parametrize( + "key", ["HTTPS_PROXY", "no_proxy", "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE", "OPENSHELL_SANDBOX"] + ) + def test_env_the_supervisor_would_silently_change_is_refused(self, key): + backend, client = _backend() + + with pytest.raises(SandboxTerminalError, match=f"sets {key}"): + backend.create(spec=SandboxSpec(env={key: "value"})) + + client.create.assert_not_called() + + +class TestCreate: + def test_default_spec_is_deny_all_under_hard_landlock(self): + backend, client = _backend(image="python:3.13-slim", cpu="2", memory="4Gi") + + with mock.patch(f"{_MODULE}.time.time", return_value=1_700_000_000.4): + name = backend.create(spec=SandboxSpec(env={"TOKEN": "value"})) + + kwargs = client.create.call_args.kwargs + assert kwargs["name"] == name + assert name.startswith("airflow-") + assert len(name) <= 19 + assert kwargs["workspace"] == "default" + assert kwargs["labels"] == {"created-by": "airflow", "airflow-created-at": "1700000000"} + spec = kwargs["spec"] + assert dict(spec.environment) == {"TOKEN": "value"} + assert spec.template.image == "python:3.13-slim" + assert dict(spec.template.resources)["limits"] == {"cpu": "2", "memory": "4Gi"} + assert len(spec.policy.network_policies) == 0 + assert spec.policy.landlock.compatibility == "hard_requirement" + assert tuple(spec.policy.filesystem.read_write) == _READ_WRITE_PATHS Review Comment: The policy test now spells out every path as a literal and asserts `read_only` and `include_workdir` (57aabdabf1). _🤖 Addressed by [Claude Code](https://claude.com/claude-code)_ ########## providers/common/ai/tests/unit/common/ai/sandbox/test_openshell.py: ########## @@ -0,0 +1,1002 @@ +# 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 builtins +import copy +import os +import shlex +import shutil +import signal +import subprocess +import sys +import time +from unittest import mock + +import pytest + +pytest.importorskip("openshell") + +import grpc +from openshell import SandboxClient +from openshell._proto import openshell_pb2, sandbox_pb2 + +from airflow.providers.common.ai.sandbox.base import ( + SandboxError, + SandboxFileTooLargeError, + SandboxSpec, + SandboxTerminalError, +) +from airflow.providers.common.ai.sandbox.openshell import ( + _EXEC_GRACE, + _READ_WRITE_PATHS, + _RUN_WRAPPER, + _STAGING_FAILED, + OpenShellSandboxBackend, +) + +_MODULE = "airflow.providers.common.ai.sandbox.openshell" +_MONOTONIC_PATH = f"{_MODULE}.time.monotonic" + + +class _RpcError(grpc.RpcError): + def __init__(self, code: grpc.StatusCode, details: str = "") -> None: + super().__init__(details) + self._code = code + self._details = details + + def code(self) -> grpc.StatusCode: + return self._code + + def details(self) -> str: + return self._details + + +class _Stream: + """An ExecSandbox response stream: events, optionally ending in an error.""" + + def __init__(self, events, error: Exception | None = None) -> None: + self._events = iter(events) + self._error = error + self.cancel = mock.Mock() + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._events) + except StopIteration: + if self._error is not None: + raise self._error from None + raise + + +def _stdout(data: bytes): + return openshell_pb2.ExecSandboxEvent(stdout=openshell_pb2.ExecSandboxStdout(data=data)) + + +def _stderr(data: bytes): + return openshell_pb2.ExecSandboxEvent(stderr=openshell_pb2.ExecSandboxStderr(data=data)) + + +def _exit(code: int): + return openshell_pb2.ExecSandboxEvent(exit=openshell_pb2.ExecSandboxExit(exit_code=code)) + + +def _result(code: int = 0, out: bytes = b"", err: bytes = b"") -> _Stream: + events = [] + if out: + events.append(_stdout(out)) + if err: + events.append(_stderr(err)) + events.append(_exit(code)) + return _Stream(events) + + +def _sandbox(phase: openshell_pb2.SandboxPhase = openshell_pb2.SANDBOX_PHASE_READY, *conditions): + response = openshell_pb2.SandboxResponse() + response.sandbox.status.phase = phase + for reason, message in conditions: + response.sandbox.status.conditions.add(type="Ready", status="False", reason=reason, message=message) + return response + + +def _config( + hosts=(), + *, + source: sandbox_pb2.PolicySource = sandbox_pb2.POLICY_SOURCE_SANDBOX, + admitted: bool = True, + landlock: str = "hard_requirement", + approval_mode: str | None = None, + proposals: bool = False, + extra_rule: tuple[str, int] | None = None, + allowed_ips=(), + binaries=("/**",), + middlewares=(), +): + config = sandbox_pb2.GetSandboxConfigResponse( + version=1, policy_source=source, configuration_admitted=admitted + ) + config.policy.landlock.compatibility = landlock + for name in middlewares: + # Reading a missing key of a message map adds it. + config.policy.network_middlewares[name] + if hosts or allowed_ips: + rule = config.policy.network_policies["airflow-egress"] + rule.name = "airflow-egress" + for host in hosts: + rule.endpoints.add(host=host, ports=[443], allowed_ips=list(allowed_ips)) + if not hosts: + rule.endpoints.add(ports=[443], allowed_ips=list(allowed_ips)) + rule.binaries.extend(sandbox_pb2.NetworkBinary(path=path) for path in binaries) + if extra_rule is not None: + rule = config.policy.network_policies["allow_extra"] + rule.endpoints.add(host=extra_rule[0], port=extra_rule[1]) + rule.binaries.add(path="/usr/local/bin/python3.12") + if approval_mode is not None: + config.settings["proposal_approval_mode"].value.string_value = approval_mode + config.settings["proposal_approval_mode"].scope = sandbox_pb2.SETTING_SCOPE_GLOBAL + if proposals: + config.settings["agent_policy_proposals_enabled"].value.bool_value = True + return config + + +def _backend(**kwargs) -> tuple[OpenShellSandboxBackend, mock.MagicMock]: + backend = OpenShellSandboxBackend(gateway="test-gateway", **kwargs) + client = mock.MagicMock(spec=SandboxClient) + client._stub = mock.MagicMock(spec=["GetSandbox", "GetSandboxConfig", "ExecSandbox"]) + client._stub.GetSandbox.return_value = _sandbox() + client._stub.GetSandboxConfig.return_value = _config() + client._stub.ExecSandbox.return_value = _result() + backend._client = client + return backend, client + + +def _exec_request(client, call: int = -1): + return client._stub.ExecSandbox.call_args_list[call].args[0] + + [email protected](autouse=True) +def _no_sleep(): + with mock.patch(f"{_MODULE}.time.sleep", autospec=True): + yield + + +def test_missing_sdk_error_is_actionable(): + real_import = builtins.__import__ + + def blocked_import(name, *args, **kwargs): + if name.startswith("openshell"): + raise ImportError("blocked for test") + return real_import(name, *args, **kwargs) + + backend = OpenShellSandboxBackend() + with mock.patch("builtins.__import__", side_effect=blocked_import): + with pytest.raises(SandboxTerminalError, match=r"\[openshell\]"): + backend.create() + + [email protected]( + ("kwargs", "message"), + [ + ({"gateway": ""}, "gateway"), + ({"workspace": ""}, "workspace"), + ({"image": ""}, "image"), + ({"cpu": ""}, "cpu"), + ({"memory": ""}, "memory"), + ({"ready_timeout": 0}, "ready_timeout"), + ({"request_timeout": float("inf")}, "request_timeout"), + ], +) +def test_constructor_rejects_invalid_values(kwargs, message): + with pytest.raises(ValueError, match=message): + OpenShellSandboxBackend(**kwargs) + + +class TestClient: + @mock.patch("openshell.SandboxClient.from_active_cluster", autospec=True) + def test_construction_reads_no_gateway_registration(self, from_active_cluster): + OpenShellSandboxBackend(gateway="prod") + + from_active_cluster.assert_not_called() + + def test_the_backend_can_be_deep_copied(self): + backend = OpenShellSandboxBackend(gateway="prod") + + assert copy.deepcopy(backend)._gateway == "prod" + + @mock.patch("openshell.SandboxClient.from_active_cluster", autospec=True) + def test_the_cli_gateway_registration_is_loaded_once(self, from_active_cluster): + backend = OpenShellSandboxBackend(gateway="prod", request_timeout=12) + + assert backend._get_client() is backend._get_client() + + from_active_cluster.assert_called_once_with(cluster="prod", timeout=12) + + @mock.patch("openshell.SandboxClient.from_active_cluster", autospec=True) + def test_a_missing_registration_is_terminal(self, from_active_cluster): + from openshell import SandboxError as SdkError + + from_active_cluster.side_effect = SdkError("gateway 'prod' not found") + + with pytest.raises(SandboxTerminalError, match="gateway 'prod' not found"): + OpenShellSandboxBackend(gateway="prod")._get_client() + + +class TestSpecRefusals: + @pytest.mark.parametrize( + ("spec", "message"), + [ + (SandboxSpec(owner="dag/run"), "owner"), + (SandboxSpec(allow_egress_to_cidrs=["203.0.113.0/24"]), "allow_egress_to_cidrs"), + (SandboxSpec(block_network=False), "open network"), + (SandboxSpec(block_network=False, allow_egress_to=["pypi.org"]), "open network"), + (SandboxSpec(allow_egress_to="pypi.org"), "not one string"), + (SandboxSpec(env={"FOO": 1}), "strings to strings"), # type: ignore[dict-item] + ], + ) + def test_unenforceable_specs_are_refused_before_the_gateway_is_asked(self, spec, message): + backend, client = _backend() + + with pytest.raises(SandboxTerminalError, match=message): + backend.create(spec=spec) + + client.create.assert_not_called() + + @pytest.mark.parametrize( + "host", + ["https://pypi.org";, "pypi.org:443", "203.0.113.7", "localhost", "*.com", "*", "pypi.org/simple", ""], + ) + def test_entries_that_are_not_hostnames_are_refused(self, host): + backend, client = _backend() + + with pytest.raises(SandboxTerminalError, match="bare hostnames"): + backend.create(spec=SandboxSpec(allow_egress_to=[host])) + + client.create.assert_not_called() + + @pytest.mark.parametrize( + "key", ["HTTPS_PROXY", "no_proxy", "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE", "OPENSHELL_SANDBOX"] + ) + def test_env_the_supervisor_would_silently_change_is_refused(self, key): + backend, client = _backend() + + with pytest.raises(SandboxTerminalError, match=f"sets {key}"): + backend.create(spec=SandboxSpec(env={key: "value"})) + + client.create.assert_not_called() + + +class TestCreate: + def test_default_spec_is_deny_all_under_hard_landlock(self): + backend, client = _backend(image="python:3.13-slim", cpu="2", memory="4Gi") + + with mock.patch(f"{_MODULE}.time.time", return_value=1_700_000_000.4): + name = backend.create(spec=SandboxSpec(env={"TOKEN": "value"})) + + kwargs = client.create.call_args.kwargs + assert kwargs["name"] == name + assert name.startswith("airflow-") + assert len(name) <= 19 + assert kwargs["workspace"] == "default" + assert kwargs["labels"] == {"created-by": "airflow", "airflow-created-at": "1700000000"} + spec = kwargs["spec"] + assert dict(spec.environment) == {"TOKEN": "value"} + assert spec.template.image == "python:3.13-slim" + assert dict(spec.template.resources)["limits"] == {"cpu": "2", "memory": "4Gi"} + assert len(spec.policy.network_policies) == 0 + assert spec.policy.landlock.compatibility == "hard_requirement" + assert tuple(spec.policy.filesystem.read_write) == _READ_WRITE_PATHS + assert "/dev/shm" in spec.policy.filesystem.read_write + client.delete.assert_not_called() + + def test_none_spec_still_gets_an_explicit_deny_all_policy(self): + # Without a policy OpenShell falls back to one baked into the image, which may be open. + backend, client = _backend() + + backend.create() + + spec = client.create.call_args.kwargs["spec"] + assert spec.HasField("policy") + assert len(spec.policy.network_policies) == 0 + + def test_hostnames_become_one_rule_on_port_443_for_any_binary(self): + backend, client = _backend() + client._stub.GetSandboxConfig.return_value = _config(["files.pythonhosted.org", "pypi.org"]) + + backend.create(spec=SandboxSpec(allow_egress_to=["PyPI.org", "files.pythonhosted.org", "pypi.org"])) + + rules = client.create.call_args.kwargs["spec"].policy.network_policies + assert list(rules) == ["airflow-egress"] + assert [(endpoint.host, list(endpoint.ports)) for endpoint in rules["airflow-egress"].endpoints] == [ + ("files.pythonhosted.org", [443]), + ("pypi.org", [443]), + ] + assert [binary.path for binary in rules["airflow-egress"].binaries] == ["/**"] + client.delete.assert_not_called() + + def test_waits_for_the_sandbox_to_become_ready(self): + backend, client = _backend() + client._stub.GetSandbox.side_effect = [ + _sandbox(openshell_pb2.SANDBOX_PHASE_PROVISIONING), + _RpcError(grpc.StatusCode.UNAVAILABLE, "restarting"), + _sandbox(), + ] + + backend.create(spec=SandboxSpec()) + + assert client._stub.GetSandbox.call_count == 3 + + def test_a_sandbox_that_fails_to_start_is_deleted_and_terminal(self): + backend, client = _backend() + client._stub.GetSandbox.return_value = _sandbox( + openshell_pb2.SANDBOX_PHASE_ERROR, ("IdentityResolutionFailed", "no passwd entry") + ) + + with pytest.raises(SandboxTerminalError, match="IdentityResolutionFailed: no passwd entry"): + backend.create(spec=SandboxSpec()) + + name = client.create.call_args.kwargs["name"] + client.delete.assert_called_once_with(name, workspace="default", allow_missing=True) + + def test_a_sandbox_not_ready_in_time_is_deleted_and_terminal(self): + backend, client = _backend(ready_timeout=5) + client._stub.GetSandbox.return_value = _sandbox(openshell_pb2.SANDBOX_PHASE_PROVISIONING) + + with mock.patch(_MONOTONIC_PATH, side_effect=[0.0, 1.0, 6.0]): + with pytest.raises(SandboxTerminalError, match="not ready: it is SANDBOX_PHASE_PROVISIONING"): + backend.create(spec=SandboxSpec()) + + client.delete.assert_called_once() + + def test_a_create_the_gateway_rejects_is_terminal_and_cleaned_up(self): + backend, client = _backend() + client.create.side_effect = _RpcError(grpc.StatusCode.INVALID_ARGUMENT, "name exceeds maximum length") + + with pytest.raises(SandboxTerminalError, match="INVALID_ARGUMENT: name exceeds maximum length"): + backend.create(spec=SandboxSpec()) + + client.delete.assert_called_once() + + @pytest.mark.parametrize( + ("config", "message"), + [ + (_config(source=sandbox_pb2.POLICY_SOURCE_GLOBAL), "gateway-wide policy"), + (_config(["pypi.org", "example.com"]), "admits"), + (_config(), "admits nothing where"), + (_config(extra_rule=("www.google.com", 443)), "www.google.com:443"), + (_config(["pypi.org"], binaries=("/usr/bin/curl",)), "admits"), + (_config(allowed_ips=["0.0.0.0/1"]), "admits addresses"), + (_config(["pypi.org"], approval_mode="auto"), "proposal_approval_mode is 'auto'"), + (_config(["pypi.org"], proposals=True), "agent_policy_proposals_enabled"), + (_config(["pypi.org"], landlock="best_effort"), "Landlock"), + (_config(["pypi.org"], admitted=False), "not admitted"), + (_config(["pypi.org"], middlewares=["inspect"]), "network middlewares"), + (_RpcError(grpc.StatusCode.PERMISSION_DENIED, "config:read"), "PERMISSION_DENIED"), + ], + ids=[ + "global-override", + "wider-allowlist", + "missing-rule", + "auto-approved-rule", + "narrower-binaries", + "address-rule", + "auto-approval", + "agent-proposals", + "landlock-best-effort", + "not-admitted", + "network-middlewares", + "unreadable", + ], + ) + def test_a_policy_other_than_the_requested_one_destroys_the_sandbox(self, config, message): + backend, client = _backend() + if isinstance(config, Exception): + client._stub.GetSandboxConfig.side_effect = config + else: + client._stub.GetSandboxConfig.return_value = config + + with pytest.raises(SandboxTerminalError, match=message): + backend.create(spec=SandboxSpec(allow_egress_to=["pypi.org"])) + + client.delete.assert_called_once() + assert backend._egress == {} + + @pytest.mark.parametrize("code", [grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED]) + def test_a_policy_read_rides_out_a_restarting_gateway(self, code): + backend, client = _backend() + client._stub.GetSandboxConfig.side_effect = [_RpcError(code, "restarting"), _config()] + + backend.create(spec=SandboxSpec()) + + assert client._stub.GetSandboxConfig.call_count == 2 + client.delete.assert_not_called() + + def test_manual_approval_mode_is_accepted(self): + backend, client = _backend() + client._stub.GetSandboxConfig.return_value = _config(approval_mode="manual") + + backend.create(spec=SandboxSpec()) + + client.delete.assert_not_called() + + +class TestRunCommand: + def test_the_command_travels_on_stdin_to_the_guest_wrapper(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _result(3, b"out\n", b"err\n") + + result = backend.run_command( + "box", "echo out; echo err >&2; exit 3", timeout=2.5, max_output_bytes=100 + ) + + request = _exec_request(client) + assert request.sandbox == "box" + assert list(request.command) == ["/bin/sh", "-c", _RUN_WRAPPER, "airflow-exec", "3", "100"] + assert request.stdin == b"echo out; echo err >&2; exit 3" + assert request.no_login_shell is True + assert request.request_id + assert not request.HasField("execution_timeout") + assert client._stub.ExecSandbox.call_args.kwargs["timeout"] == 3 + _EXEC_GRACE + assert (result.exit_code, result.stdout, result.stderr) == (3, "out\n", "err\n") + assert not result.timed_out + assert not result.sandbox_terminated + + @pytest.mark.parametrize( + ("exit_code", "elapsed", "timed_out"), + [(124, 3.1, True), (124, 0.2, False), (137, 3.1, False), (0, 3.1, False)], + ) + def test_only_the_wrappers_status_after_the_budget_is_a_timeout(self, exit_code, elapsed, timed_out): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _result(exit_code) + + # The policy read before and after the command each take a deadline reading too. + with mock.patch(_MONOTONIC_PATH, side_effect=[0.0, 100.0, 100.0 + elapsed, 0.0]): + result = backend.run_command("box", "sleep 300", timeout=3, max_output_bytes=100) + + assert result.timed_out is timed_out + + @pytest.mark.parametrize( + ("exit_code", "stderr"), + [ + pytest.param(125, b"docker: invalid reference format\n", id="a-command-exiting-125"), + pytest.param(1, f"{_STAGING_FAILED}\n".encode(), id="the-staging-line-with-another-status"), + pytest.param(125, f"{_STAGING_FAILED}\nmore\n".encode(), id="the-staging-line-not-last"), + ], + ) + def test_a_result_unlike_the_wrappers_staging_failure_is_the_commands_own(self, exit_code, stderr): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _result(exit_code, err=stderr) + + result = backend.run_command("box", "cmd", timeout=5, max_output_bytes=100) + + assert (result.exit_code, result.stderr) == (exit_code, stderr.decode()) + + @pytest.mark.parametrize( + ("max_output_bytes", "message"), + [ + pytest.param( + 100, + "The command was not run: the sandbox could not write it to /tmp " + "(cat: write error: No space left on device).", + id="with-the-cause", + ), + pytest.param( + 10, + "The command was not run: the sandbox could not write it to /tmp.", + id="under-a-cap-shorter-than-the-staging-line", + ), + ], + ) + def test_the_wrappers_own_staging_failure_is_a_recoverable_error(self, max_output_bytes, message): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _result( + 125, err=f"cat: write error: No space left on device\n{_STAGING_FAILED}\n".encode() + ) + + with pytest.raises(SandboxError) as error: + backend.run_command("box", "make", timeout=5, max_output_bytes=max_output_bytes) + + assert str(error.value) == message + assert not isinstance(error.value, SandboxTerminalError) + + def test_stderr_stays_capped_below_the_length_of_the_staging_line(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _result(1, err=b"0123456789" * 3) + + result = backend.run_command("box", "cmd", timeout=5, max_output_bytes=20) + + assert (result.stderr, result.stderr_truncated) == ("0123456789" * 2, True) + + def test_each_stream_is_capped_and_flagged_on_its_own(self): + backend, client = _backend() + # The wrapper sends one byte past the cap for a stream it had to cut. + client._stub.ExecSandbox.return_value = _Stream( + [_stdout(b"partial line\nkept 1\n"), _stdout(b"kept 2\n"), _stderr(b"short\n"), _exit(0)] + ) + + result = backend.run_command("box", "cmd", timeout=5, max_output_bytes=20) + + assert result.stdout == "kept 1\nkept 2\n" + assert result.stdout_truncated + assert result.stderr == "short\n" + assert not result.stderr_truncated + + def test_one_line_longer_than_the_budget_still_reaches_the_model(self): + """Dropping the leading partial line must not empty the window: the toolset prints "(no output)" for it.""" + backend, client = _backend() + client._stub.ExecSandbox.return_value = _Stream([_stdout(b"x" * 60000 + b"\n"), _exit(0)]) + + result = backend.run_command("box", "cat big.json", timeout=5, max_output_bytes=51200) + + assert result.stdout != "" + assert result.stdout_truncated + assert len(result.stdout.encode()) <= 51200 + + def test_a_long_line_followed_by_a_short_one_keeps_the_window(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _Stream( + [_stdout(b"L" * 200_000 + b"\nshort tail\n"), _exit(0)] + ) + + result = backend.run_command("box", "spew", timeout=5, max_output_bytes=51200) + + assert len(result.stdout.encode()) > 51200 // 2 + assert result.stdout.endswith("short tail\n") + + def test_the_whole_second_deadline_the_command_got_is_reported(self): + backend, client = _backend() + + result = backend.run_command("box", "true", timeout=2.2, max_output_bytes=100) + + assert result.applied_timeout == 3.0 + + def test_the_whole_second_deadline_is_reported_for_an_abandoned_command_too(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _Stream([], _RpcError(grpc.StatusCode.DEADLINE_EXCEEDED)) + + result = backend.run_command("box", "sleep 300", timeout=2.2, max_output_bytes=100) + + assert (result.sandbox_terminated, result.applied_timeout) == (True, 3.0) + + def test_output_injected_past_the_wrapper_stays_bounded_in_worker_memory(self): + backend, client = _backend() + chunk = b"y\n" * 32768 + client._stub.ExecSandbox.return_value = _Stream([_stdout(chunk) for _ in range(200)] + [_exit(0)]) + + result = backend.run_command("box", "yes > /proc/$PPID/fd/1", timeout=5, max_output_bytes=1024) + + assert len(result.stdout.encode()) <= 1024 + assert result.stdout_truncated + + def test_undecodable_output_is_replaced_not_raised(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _result(0, b"\xff\xfeok\n") + + assert backend.run_command("box", "cmd", timeout=5, max_output_bytes=100).stdout == "��ok\n" + + def test_a_command_over_the_request_limit_is_recoverable_and_not_sent(self): + backend, client = _backend() + + with pytest.raises(SandboxError, match="write_file") as error: + backend.run_command("box", "x" * 1_000_001, timeout=5, max_output_bytes=100) + + assert not isinstance(error.value, SandboxTerminalError) + client._stub.ExecSandbox.assert_not_called() + + def test_a_hung_exec_destroys_the_sandbox(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _Stream([], _RpcError(grpc.StatusCode.DEADLINE_EXCEEDED)) + + result = backend.run_command("box", "sleep 300", timeout=3, max_output_bytes=100) + + assert (result.exit_code, result.timed_out, result.sandbox_terminated) == (-1, True, True) + client.delete.assert_called_once_with("box", workspace="default", allow_missing=True) + + def test_a_sandbox_not_ready_after_a_gateway_restart_is_waited_for_and_retried(self): + backend, client = _backend() + client._stub.ExecSandbox.side_effect = [ + _RpcError(grpc.StatusCode.FAILED_PRECONDITION, "sandbox is not ready"), + _result(0, b"ok\n"), + ] + client._stub.GetSandbox.side_effect = [_sandbox(openshell_pb2.SANDBOX_PHASE_PROVISIONING), _sandbox()] + + result = backend.run_command("box", "echo ok", timeout=5, max_output_bytes=100) + + assert result.stdout == "ok\n" + assert client._stub.ExecSandbox.call_count == 2 + assert _exec_request(client, 0).request_id != _exec_request(client, 1).request_id + + def test_a_gateway_drop_during_the_not_ready_retry_is_reported_not_leaked(self): + backend, client = _backend() + client._stub.ExecSandbox.side_effect = [ + _RpcError(grpc.StatusCode.FAILED_PRECONDITION, "sandbox is not ready"), + _Stream([], _RpcError(grpc.StatusCode.UNAVAILABLE, "exec relay closed")), + ] + + with pytest.raises(SandboxError, match="may or may not have run") as error: + backend.run_command("box", "make install", timeout=5, max_output_bytes=100) + + assert not isinstance(error.value, SandboxTerminalError) + assert client._stub.ExecSandbox.call_count == 2 + + def test_an_exec_the_gateway_drops_is_reported_not_retried(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _Stream( + [], + _RpcError( + grpc.StatusCode.UNAVAILABLE, "exec relay closed before the command reported an exit status" + ), + ) + client._stub.GetSandbox.side_effect = [_RpcError(grpc.StatusCode.UNAVAILABLE), _sandbox()] + + with pytest.raises(SandboxError, match="may or may not have run") as error: + backend.run_command("box", "make install", timeout=5, max_output_bytes=100) + + assert not isinstance(error.value, SandboxTerminalError) + client._stub.ExecSandbox.assert_called_once() + + def test_a_dropped_exec_on_a_sandbox_that_does_not_come_back_is_terminal(self): + backend, client = _backend() + client._stub.ExecSandbox.return_value = _Stream([_stdout(b"x")]) + client._stub.GetSandbox.return_value = _sandbox(openshell_pb2.SANDBOX_PHASE_ERROR) + + with pytest.raises(SandboxTerminalError, match="SANDBOX_PHASE_ERROR"): + backend.run_command("box", "cmd", timeout=5, max_output_bytes=100) + + @pytest.mark.parametrize( + ("code", "terminal"), + [ + (grpc.StatusCode.NOT_FOUND, True), + (grpc.StatusCode.UNAUTHENTICATED, True), + (grpc.StatusCode.PERMISSION_DENIED, True), + (grpc.StatusCode.OUT_OF_RANGE, False), + (grpc.StatusCode.RESOURCE_EXHAUSTED, False), + ], + ) + def test_gateway_errors_are_classified(self, code, terminal): + backend, client = _backend() + client._stub.ExecSandbox.side_effect = _RpcError(code, "detail") Review Comment: Added `assert client._stub.ExecSandbox.call_count == 1` (57aabdabf1). _🤖 Addressed by [Claude Code](https://claude.com/claude-code)_ ########## providers/common/ai/tests/system/common/ai/example_sandbox_toolset_openshell.py: ########## @@ -0,0 +1,114 @@ +# 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. +"""End-to-end system test for SandboxToolset with NVIDIA OpenShell.""" + +from __future__ import annotations + +import os +from datetime import UTC, datetime + +from airflow.providers.common.compat.sdk import dag as airflow_dag, task + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") +DAG_ID = f"common_ai_sandbox_toolset_openshell_{ENV_ID}" if ENV_ID else "common_ai_sandbox_toolset_openshell" + +MARKER = "boundary-ok" +STATE_PATH = "/tmp/airflow_sandbox_e2e" +# Counts processes whose command line is exactly "sleep 300", without needing ps in the image. +COUNT_SLEEPERS = ( + "n=0; for f in /proc/[0-9]*/cmdline; do " + "[ \"$(tr '\\0' ' ' < \"$f\" 2>/dev/null)\" = 'sleep 300 ' ] && n=$((n + 1)); done; echo sleepers=$n" +) +CONNECT = ( + 'python3 -c "import socket\n' + "try:\n socket.create_connection(('1.1.1.1', 443), 5)\n print('egress=open')\n" + "except OSError as e:\n print('egress=denied', e.errno)\"" +) + + +@airflow_dag( + dag_id=DAG_ID, + schedule="@once", + start_date=datetime(2024, 1, 1, tzinfo=UTC), + catchup=False, + tags=["common.ai", "sandbox", "openshell", "system_test"], +) +def example_sandbox_toolset_openshell(): + @task + def run_sandbox_agent() -> str: + from pydantic_ai import Agent + from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart + from pydantic_ai.models.function import AgentInfo, FunctionModel + + from airflow.providers.common.ai.sandbox import OpenShellSandboxBackend + from airflow.providers.common.ai.toolsets import SandboxToolset + + # Each step is a tool call and a check of what the previous one returned. + steps: list[tuple[str, dict, str]] = [ + ("write_file", {"path": STATE_PATH, "content": MARKER}, "Wrote"), + ("run_command", {"command": f"cat {STATE_PATH} && python3 -c 'print(6 * 7)'"}, "42"), + ("read_file", {"path": STATE_PATH}, MARKER), + ("run_command", {"command": "echo to-stderr >&2; exit 3"}, "[exit code: 3]"), Review Comment: The system test now checks `to-stderr` and the marker with `42`, as the OpenSandbox example does (297cff7921). _🤖 Addressed by [Claude Code](https://claude.com/claude-code)_ -- 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]
