zozo123 commented on code in PR #73961: URL: https://github.com/apache/airflow/pull/73961#discussion_r4262873904
########## providers/common/ai/tests/unit/common/ai/sandbox/test_openshell.py: ########## @@ -0,0 +1,1243 @@ +# 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 errno +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import time +from collections.abc import Callable +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, + _NO_SETSID, + _NO_SETSID_STATUS, + _RUN_WRAPPER, + _STAGING_FAILED, + _SYSTEM_PATH, + OpenShellSandboxBackend, + _wrapper_line, +) + +_MODULE = "airflow.providers.common.ai.sandbox.openshell" +_MONOTONIC_PATH = f"{_MODULE}.time.monotonic" +_GATEWAY_DIR = "/home/airflow/.config/openshell/gateways/prod" +# The temp file the SDK writes a refreshed OIDC token through, next to the registration. +_TOKEN_TEMP_FILE = f"{_GATEWAY_DIR}/.oidc_token.k2x9w1qz.tmp" + + +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] + + +def _nonce(request) -> str: + return request.stdin.partition(b"\n")[0].decode() + + +def _reply(code: int, stderr: Callable[[str], bytes]): + """An ExecSandbox side effect whose stderr is built from the nonce the call sent.""" + + def reply(request, timeout): + return _result(code, err=stderr(_nonce(request))) + + return reply + + [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() + + @pytest.mark.parametrize( + "error", + [ + pytest.param( + PermissionError(errno.EACCES, "Permission denied", _TOKEN_TEMP_FILE), id="not-writable" + ), + pytest.param( + OSError(errno.EROFS, "Read-only file system", _TOKEN_TEMP_FILE), id="read-only-mount" + ), + ], + ) + def test_an_oidc_token_that_cannot_be_written_back_is_terminal(self, error): + backend, client = _backend() + client._stub.GetSandboxConfig.side_effect = error + + with pytest.raises(SandboxTerminalError, match="registration is read-only.*mTLS is the supported"): + backend.run_command("box", "true", timeout=5, max_output_bytes=100) + + client._stub.ExecSandbox.assert_not_called() + + @pytest.mark.parametrize( + "error", + [ + pytest.param( + PermissionError(errno.EACCES, "Permission denied", f"{_GATEWAY_DIR}/mtls/tls.key"), + id="an-unreadable-mtls-key", + ), + pytest.param( + PermissionError(errno.EACCES, "Permission denied", f"{_GATEWAY_DIR}/oidc_token.json"), + id="the-token-file-itself", + ), + pytest.param(OSError(errno.ENOSPC, "No space left on device"), id="no-file-named"), + pytest.param(OSError(errno.ENOSPC, "No space left on device", _TOKEN_TEMP_FILE), id="disk-full"), + pytest.param(OSError(errno.EDQUOT, "Disk quota exceeded", _TOKEN_TEMP_FILE), id="quota-used-up"), + ], + ) + def test_other_os_errors_are_not_reported_as_a_read_only_registration(self, error): + backend, client = _backend() + client._stub.GetSandboxConfig.side_effect = error + + with pytest.raises(SandboxTerminalError) as raised: + backend.run_command("box", "true", timeout=5, max_output_bytes=100) + + assert str(raised.value) == ( + f"OpenShell could not read the network policy of sandbox box ({type(error).__name__}: {error})." + ) + + +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 spec.policy.filesystem.include_workdir is True + assert list(spec.policy.filesystem.read_only) == [ + "/bin", + "/usr", + "/lib", + "/proc", + "/dev/urandom", + "/etc", + "/var/log", + ] + assert list(spec.policy.filesystem.read_write) == ["/tmp", "/dev/null", "/dev/shm"] + 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"] + nonce, _, command = request.stdin.partition(b"\n") + assert re.fullmatch(rb"[0-9a-f]{32}", nonce) + assert command == 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 + + def test_each_command_carries_a_fresh_nonce(self): + backend, client = _backend() + client._stub.ExecSandbox.side_effect = _reply(0, lambda nonce: b"") + + backend.run_command("box", "true", timeout=5, max_output_bytes=100) + backend.run_command("box", "true", timeout=5, max_output_bytes=100) + + assert _nonce(_exec_request(client, 0)) != _nonce(_exec_request(client, 1)) + + @pytest.mark.parametrize( + ("exit_code", "stderr"), + [ + pytest.param( + 125, lambda nonce: b"docker: invalid reference format\n", id="a-command-exiting-125" + ), + pytest.param( + 1, + lambda nonce: _wrapper_line(_STAGING_FAILED, nonce), + id="the-staging-line-with-another-status", + ), + pytest.param( + 125, + lambda nonce: _wrapper_line(_STAGING_FAILED, nonce) + b"more\n", + id="the-staging-line-not-last", + ), + pytest.param( + 125, + lambda nonce: f"{_STAGING_FAILED}\n".encode(), + id="the-staging-text-printed-by-the-command", + ), + pytest.param( + 125, + lambda nonce: _wrapper_line(_STAGING_FAILED, "0" * 32), + id="the-staging-line-with-another-nonce", + ), + pytest.param(127, lambda nonce: b"sh: 1: setsid: not found\n", id="a-command-exiting-127"), + pytest.param( + 1, lambda nonce: _wrapper_line(_NO_SETSID, nonce), id="the-setsid-line-with-another-status" + ), + pytest.param( + 127, lambda nonce: _wrapper_line(_NO_SETSID, nonce) + b"more\n", id="the-setsid-line-not-last" + ), + pytest.param( + 127, lambda nonce: f"{_NO_SETSID}\n".encode(), id="the-setsid-text-printed-by-the-command" + ), + pytest.param( + 127, + lambda nonce: _wrapper_line(_NO_SETSID, "0" * 32), + id="the-setsid-line-with-another-nonce", + ), + ], + ) + def test_a_result_unlike_a_wrapper_failure_is_the_commands_own(self, exit_code, stderr): + backend, client = _backend() + client._stub.ExecSandbox.side_effect = _reply(exit_code, stderr) + + result = backend.run_command("box", "cmd", timeout=5, max_output_bytes=100) + + expected = stderr(_nonce(_exec_request(client))).decode() + assert (result.exit_code, result.stderr) == (exit_code, expected) + + @pytest.mark.parametrize( + ("max_output_bytes", "message"), + [ + pytest.param( + 200, + "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.side_effect = _reply( + 125, + lambda nonce: ( + b"cat: write error: No space left on device\n" + _wrapper_line(_STAGING_FAILED, nonce) + ), + ) + + 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_an_image_without_setsid_is_terminal(self): + backend, client = _backend() + client._stub.ExecSandbox.side_effect = _reply( + _NO_SETSID_STATUS, lambda nonce: _wrapper_line(_NO_SETSID, nonce) + ) + + # A cap shorter than the line, so the line is still recognised past it. + with pytest.raises(SandboxTerminalError, match="no setsid.*util-linux"): + backend.run_command("box", "true", timeout=5, max_output_bytes=10) + + 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 Review Comment: Done. The test now records `len(tail._data)` after every `add`, so it fails with the trim removed. --- Drafted-by: Claude Code (Opus 5.5) (no human review before posting) ########## providers/common/ai/docs/sandbox/index.rst: ########## @@ -117,9 +118,10 @@ The four tools: - Lists a directory. Directories are shown with a trailing ``/``. A :class:`~airflow.providers.common.ai.sandbox.SandboxBackend` provisions the sandbox -on the model's first tool call and tears it down when the agent run ends. Three +on the model's first tool call and tears it down when the agent run ends. Four backends ship: `Modal <https://modal.com/docs/guide/sandbox>`__ (hosted) and -`OpenSandbox <https://open-sandbox.ai/>`__ (self-hosted) for production and +`OpenSandbox <https://open-sandbox.ai/>`__ and +`NVIDIA OpenShell <https://github.com/NVIDIA/OpenShell>`__ (self-hosted) for production and Review Comment: Done. Both places now leave Kubernetes out for OpenShell and say it has been verified on its Docker driver only. --- Drafted-by: Claude Code (Opus 5.5) (no human review before posting) ########## providers/common/ai/tests/unit/common/ai/sandbox/test_openshell.py: ########## @@ -0,0 +1,1243 @@ +# 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 errno +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import time +from collections.abc import Callable +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, + _NO_SETSID, + _NO_SETSID_STATUS, + _RUN_WRAPPER, + _STAGING_FAILED, + _SYSTEM_PATH, + OpenShellSandboxBackend, + _wrapper_line, +) + +_MODULE = "airflow.providers.common.ai.sandbox.openshell" +_MONOTONIC_PATH = f"{_MODULE}.time.monotonic" +_GATEWAY_DIR = "/home/airflow/.config/openshell/gateways/prod" +# The temp file the SDK writes a refreshed OIDC token through, next to the registration. +_TOKEN_TEMP_FILE = f"{_GATEWAY_DIR}/.oidc_token.k2x9w1qz.tmp" + + +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] + + +def _nonce(request) -> str: + return request.stdin.partition(b"\n")[0].decode() + + +def _reply(code: int, stderr: Callable[[str], bytes]): + """An ExecSandbox side effect whose stderr is built from the nonce the call sent.""" + + def reply(request, timeout): + return _result(code, err=stderr(_nonce(request))) + + return reply + + [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 Review Comment: Done. --- Drafted-by: Claude Code (Opus 5.5) (no human review before posting) -- 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]
