mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-01 21:37:33 +08:00
dcadd8d837
### Summary Adds a `tenki` sandbox provider that runs each agent code execution in a disposable Tenki (https://tenki.cloud) microVM (create → exec → destroy, no volumes or snapshots). Registration mirrors PR #15039, configure `api_key` and `project_id` in Admin > Sandbox Settings. Both runtimes are covered: - Python: `agent/sandbox/providers/tenki.py` (structured results + artifact collection). - Go: `internal/agent/sandbox/tenki.go`, mirroring the e2b provider and wired into the provider manager. `tenki-sandbox` is an optional dependency (it requires `protobuf>=6.31`, which differs from RAGFlow's pinned gRPC stack), lazily imported with a clear error when missing; installation is documented in the sandbox quickstart. Unit tests cover execution, structured results, artifacts (symlink/size/extension limits), non-zero exit, timeout, error mapping, and idempotent destroy. --------- Co-authored-by: yiming.wang <yiming.wang@luxor.com>
429 lines
15 KiB
Python
429 lines
15 KiB
Python
import base64
|
|
import posixpath
|
|
import stat
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from agent.sandbox.providers.base import SandboxProviderConfigError
|
|
from agent.sandbox.providers.tenki import TenkiProvider
|
|
from agent.sandbox.result_protocol import RESULT_MARKER_PREFIX
|
|
|
|
pytestmark = pytest.mark.p3
|
|
|
|
|
|
class _FakeErrors:
|
|
"""Stand-in for tenki_sandbox.errors so tests need not install the SDK."""
|
|
|
|
class QuotaExceededError(Exception):
|
|
pass
|
|
|
|
class RateLimitedError(Exception):
|
|
pass
|
|
|
|
class UnauthorizedError(Exception):
|
|
pass
|
|
|
|
class CommandTimeoutError(Exception):
|
|
pass
|
|
|
|
class PrimitiveTimeoutError(Exception):
|
|
pass
|
|
|
|
class SessionTerminatedError(Exception):
|
|
pass
|
|
|
|
class SessionNotFoundError(Exception):
|
|
pass
|
|
|
|
class FileNotFoundError(Exception):
|
|
pass
|
|
|
|
|
|
class _FakeFS:
|
|
def __init__(self):
|
|
self.files: dict[str, bytes] = {}
|
|
|
|
def write_text(self, path: str, content: str):
|
|
self.files[path] = content.encode("utf-8")
|
|
|
|
def read_bytes(self, path: str) -> bytes:
|
|
return self.files[path]
|
|
|
|
def list(self, path: str, *, include_hidden: bool = False):
|
|
# Mirror the real SDK: entry.path is the basename within `path`,
|
|
# not an absolute path.
|
|
prefix = path.rstrip("/") + "/"
|
|
entries = []
|
|
seen_dirs: set[str] = set()
|
|
for file_path, payload in self.files.items():
|
|
if not file_path.startswith(prefix):
|
|
continue
|
|
relative = file_path[len(prefix) :]
|
|
head, _, tail = relative.partition("/")
|
|
if tail: # nested -> surface the intermediate directory once
|
|
if head not in seen_dirs:
|
|
seen_dirs.add(head)
|
|
entries.append(SimpleNamespace(path=head, size=4096, mode=0o40755, is_dir=True, is_symlink=False, symlink_target=""))
|
|
continue
|
|
entries.append(
|
|
SimpleNamespace(
|
|
path=head,
|
|
size=len(payload),
|
|
mode=0o100644,
|
|
is_dir=False,
|
|
is_symlink=False,
|
|
symlink_target="",
|
|
)
|
|
)
|
|
return entries
|
|
|
|
|
|
class _FakeCommandResult:
|
|
def __init__(self, exit_code: int, stdout: str = "", stderr: str = ""):
|
|
self.exit_code = exit_code
|
|
self.stdout_text = stdout
|
|
self.stderr_text = stderr
|
|
self.stdout = stdout.encode("utf-8")
|
|
self.stderr = stderr.encode("utf-8")
|
|
|
|
|
|
class _FakeSandbox:
|
|
def __init__(self, run_handler=None):
|
|
self.id = "sbx-fake-1"
|
|
self.state = "RUNNING"
|
|
self.fs = _FakeFS()
|
|
self.terminated = False
|
|
self._run_handler = run_handler
|
|
self.exec_calls: list[tuple] = []
|
|
|
|
def exec(self, *argv, cwd=None, timeout=None, env=None, input=None, check=False, privileged=False):
|
|
self.exec_calls.append((argv, cwd, timeout))
|
|
if argv and argv[0] == "mkdir":
|
|
return _FakeCommandResult(0)
|
|
if self._run_handler is not None:
|
|
return self._run_handler(self, argv, cwd)
|
|
return _FakeCommandResult(0)
|
|
|
|
def terminate(self):
|
|
self.terminated = True
|
|
|
|
|
|
class _FakeClient:
|
|
def __init__(self, sandbox: _FakeSandbox):
|
|
self._sandbox = sandbox
|
|
self.create_kwargs = None
|
|
|
|
def who_am_i(self):
|
|
return SimpleNamespace(owner_type="WORKSPACE")
|
|
|
|
def create(self, **kwargs):
|
|
self.create_kwargs = kwargs
|
|
return self._sandbox
|
|
|
|
|
|
def _build_provider(sandbox: _FakeSandbox, monkeypatch) -> tuple[TenkiProvider, _FakeClient]:
|
|
provider = TenkiProvider()
|
|
provider.api_key = "tk_test"
|
|
provider.project_id = "proj-1"
|
|
provider.timeout = 30
|
|
provider.max_output_bytes = 1024 * 1024
|
|
provider.max_artifacts = 20
|
|
provider.max_artifact_bytes = 1024 * 1024
|
|
provider._initialized = True
|
|
client = _FakeClient(sandbox)
|
|
provider._client = client
|
|
monkeypatch.setattr(provider, "_tenki_errors", lambda: _FakeErrors)
|
|
return provider, client
|
|
|
|
|
|
def test_tenki_provider_executes_python_main_and_collects_artifacts(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
# Simulate the script run: emit a structured result + write an artifact.
|
|
sandbox.fs.files[posixpath.join(cwd, "artifacts", "chart.png")] = b"PNGDATA"
|
|
payload = base64.b64encode(b'{"present":true,"value":{"message":"hello tenki"},"type":"json"}').decode("ascii")
|
|
return _FakeCommandResult(0, stdout=f"debug line\n{RESULT_MARKER_PREFIX}{payload}\n")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
|
|
instance = provider.create_instance("python")
|
|
result = provider.execute_code(
|
|
instance.instance_id,
|
|
'def main() -> dict:\n return {"message": "hello tenki"}\n',
|
|
"python",
|
|
timeout=5,
|
|
)
|
|
provider.destroy_instance(instance.instance_id)
|
|
|
|
assert result.exit_code == 0
|
|
assert result.stdout == "debug line\n"
|
|
assert result.metadata["result_present"] is True
|
|
assert result.metadata["result_value"] == {"message": "hello tenki"}
|
|
assert result.metadata["artifacts"] == [
|
|
{
|
|
"name": "chart.png",
|
|
"content_b64": base64.b64encode(b"PNGDATA").decode("ascii"),
|
|
"mime_type": "image/png",
|
|
"size": 7,
|
|
}
|
|
]
|
|
assert sandbox.terminated is True
|
|
|
|
|
|
def test_tenki_provider_reports_nonzero_exit(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
return _FakeCommandResult(7, stdout="", stderr="boom\n")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
|
|
instance = provider.create_instance("python")
|
|
result = provider.execute_code(instance.instance_id, "def main():\n raise SystemExit(7)\n", "python", timeout=5)
|
|
|
|
assert result.exit_code == 7
|
|
assert result.stderr == "boom\n"
|
|
assert result.metadata["status"] == "error"
|
|
assert result.metadata["result_present"] is False
|
|
|
|
|
|
def test_tenki_provider_propagates_timeout(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
raise _FakeErrors.CommandTimeoutError("timed out")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
|
|
instance = provider.create_instance("python")
|
|
with pytest.raises(TimeoutError, match="timed out after"):
|
|
provider.execute_code(instance.instance_id, "def main():\n return 1\n", "python", timeout=5)
|
|
|
|
|
|
def test_tenki_provider_create_maps_quota_error(monkeypatch):
|
|
sandbox = _FakeSandbox()
|
|
provider, client = _build_provider(sandbox, monkeypatch)
|
|
|
|
def _raise_quota(**kwargs):
|
|
raise _FakeErrors.QuotaExceededError("no credits")
|
|
|
|
client.create = _raise_quota
|
|
with pytest.raises(RuntimeError, match="quota exceeded"):
|
|
provider.create_instance("python")
|
|
|
|
|
|
def test_tenki_provider_destroy_is_idempotent(monkeypatch):
|
|
sandbox = _FakeSandbox()
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
# Destroying an unknown instance returns True without error.
|
|
assert provider.destroy_instance("nonexistent") is True
|
|
|
|
|
|
def test_tenki_provider_create_passes_project_and_outbound(monkeypatch):
|
|
sandbox = _FakeSandbox()
|
|
provider, client = _build_provider(sandbox, monkeypatch)
|
|
provider.image = "my-image"
|
|
provider.cpu_cores = 4
|
|
provider.allow_outbound = True
|
|
|
|
provider.create_instance("python")
|
|
|
|
assert client.create_kwargs["project_id"] == "proj-1"
|
|
assert client.create_kwargs["allow_outbound"] is True
|
|
assert client.create_kwargs["image"] == "my-image"
|
|
assert client.create_kwargs["cpu_cores"] == 4
|
|
|
|
|
|
def test_tenki_provider_config_schema_and_validation():
|
|
schema = TenkiProvider.get_config_schema()
|
|
assert schema["api_key"]["required"] is True
|
|
assert schema["api_key"]["secret"] is True
|
|
assert schema["project_id"]["required"] is True
|
|
|
|
provider = TenkiProvider()
|
|
ok, _ = provider.validate_config({"api_key": "tk", "project_id": "p", "timeout": 30, "max_lifetime": 3600, "max_output_bytes": 1024, "max_artifacts": 5, "max_artifact_bytes": 1024})
|
|
assert ok is True
|
|
bad, msg = provider.validate_config({"api_key": "", "project_id": "p"})
|
|
assert bad is False
|
|
assert "API key" in msg
|
|
|
|
|
|
def test_tenki_provider_supported_languages():
|
|
assert TenkiProvider().get_supported_languages() == ["python", "javascript"]
|
|
|
|
|
|
def test_tenki_provider_initialize_maps_auth_error(monkeypatch):
|
|
provider = TenkiProvider()
|
|
monkeypatch.setattr(provider, "_tenki_errors", lambda: _FakeErrors)
|
|
|
|
class _UnauthorizedClient:
|
|
def who_am_i(self):
|
|
raise _FakeErrors.UnauthorizedError("bad token")
|
|
|
|
monkeypatch.setattr(provider, "_create_client", lambda: _UnauthorizedClient())
|
|
|
|
with pytest.raises(SandboxProviderConfigError, match="authentication failed"):
|
|
provider.initialize({"api_key": "tk_bad", "project_id": "proj-1"})
|
|
|
|
|
|
def test_tenki_provider_instances_are_independent(monkeypatch):
|
|
sandboxes: list[_FakeSandbox] = []
|
|
|
|
class _MultiClient(_FakeClient):
|
|
def create(self, **kwargs):
|
|
sb = _FakeSandbox()
|
|
sb.id = f"sbx-{len(sandboxes)}"
|
|
sandboxes.append(sb)
|
|
return sb
|
|
|
|
provider = TenkiProvider()
|
|
provider.api_key = "tk_test"
|
|
provider.project_id = "proj-1"
|
|
provider.timeout = 30
|
|
provider._initialized = True
|
|
provider._client = _MultiClient(None)
|
|
monkeypatch.setattr(provider, "_tenki_errors", lambda: _FakeErrors)
|
|
|
|
a = provider.create_instance("python")
|
|
b = provider.create_instance("python")
|
|
|
|
assert a.instance_id != b.instance_id
|
|
assert a.metadata["sandbox_id"] != b.metadata["sandbox_id"]
|
|
|
|
provider.destroy_instance(a.instance_id)
|
|
assert sandboxes[0].terminated is True
|
|
assert sandboxes[1].terminated is False # destroying one leaves the other running
|
|
|
|
provider.destroy_instance(b.instance_id)
|
|
assert sandboxes[1].terminated is True
|
|
|
|
|
|
def test_tenki_provider_executes_javascript(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
payload = base64.b64encode(b'{"present":true,"value":42,"type":"json"}').decode("ascii")
|
|
return _FakeCommandResult(0, stdout=f"{RESULT_MARKER_PREFIX}{payload}\n")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
|
|
inst = provider.create_instance("javascript")
|
|
assert inst.metadata["language"] == "nodejs"
|
|
|
|
result = provider.execute_code(inst.instance_id, "function main(args){return 42;}", "javascript", timeout=5)
|
|
|
|
# The script is written as main.js and run with node.
|
|
assert any(k.endswith("main.js") for k in sandbox.fs.files)
|
|
run_argv = [c[0] for c in sandbox.exec_calls if c[0] and c[0][0] == "node"]
|
|
assert run_argv, "expected a node invocation"
|
|
assert result.metadata["result_value"] == 42
|
|
|
|
|
|
def test_tenki_provider_rejects_symlink_artifacts(monkeypatch):
|
|
provider, _ = _build_provider(_FakeSandbox(), monkeypatch)
|
|
|
|
class _SymlinkFS:
|
|
def list(self, path, *, include_hidden=False):
|
|
# is_symlink False but mode marks a symlink -> must still be rejected.
|
|
return [SimpleNamespace(path="evil.json", size=10, mode=stat.S_IFLNK | 0o777, is_dir=False, is_symlink=False, symlink_target="/etc/passwd")]
|
|
|
|
fake_sb = SimpleNamespace(fs=_SymlinkFS())
|
|
with pytest.raises(RuntimeError, match="symlinks are not allowed"):
|
|
provider._collect_artifacts(fake_sb, "/home/tenki/wk/artifacts")
|
|
|
|
|
|
def test_tenki_provider_enforces_output_size_limit(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
return _FakeCommandResult(0, stdout="x" * 5000)
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
provider.max_output_bytes = 100
|
|
|
|
inst = provider.create_instance("python")
|
|
with pytest.raises(RuntimeError, match="output exceeded"):
|
|
provider.execute_code(inst.instance_id, "def main():\n return 1\n", "python", timeout=5)
|
|
|
|
|
|
def test_tenki_provider_rejects_disallowed_artifact_extension(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
sandbox.fs.files[posixpath.join(cwd, "artifacts", "malware.exe")] = b"MZ"
|
|
return _FakeCommandResult(0, stdout="")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
|
|
inst = provider.create_instance("python")
|
|
with pytest.raises(RuntimeError, match="Unsupported artifact type"):
|
|
provider.execute_code(inst.instance_id, "def main():\n return 1\n", "python", timeout=5)
|
|
|
|
|
|
def test_tenki_provider_allow_outbound_configurable(monkeypatch):
|
|
sandbox = _FakeSandbox()
|
|
provider, client = _build_provider(sandbox, monkeypatch)
|
|
provider.allow_outbound = False
|
|
provider.create_instance("python")
|
|
assert client.create_kwargs["allow_outbound"] is False
|
|
|
|
|
|
def test_tenki_provider_maps_session_terminated(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
raise _FakeErrors.SessionTerminatedError("reclaimed")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
inst = provider.create_instance("python")
|
|
with pytest.raises(RuntimeError, match="no longer available"):
|
|
provider.execute_code(inst.instance_id, "def main():\n return 1\n", "python", timeout=5)
|
|
|
|
|
|
def test_tenki_provider_create_maps_rate_limit(monkeypatch):
|
|
sandbox = _FakeSandbox()
|
|
provider, client = _build_provider(sandbox, monkeypatch)
|
|
|
|
def _raise_rate(**kwargs):
|
|
raise _FakeErrors.RateLimitedError("slow down")
|
|
|
|
client.create = _raise_rate
|
|
with pytest.raises(RuntimeError, match="rate limited"):
|
|
provider.create_instance("python")
|
|
|
|
|
|
def test_tenki_provider_create_wraps_unexpected_sdk_error(monkeypatch):
|
|
sandbox = _FakeSandbox()
|
|
provider, client = _build_provider(sandbox, monkeypatch)
|
|
|
|
def _raise_other(**kwargs):
|
|
raise ValueError("weird sdk failure")
|
|
|
|
client.create = _raise_other
|
|
with pytest.raises(RuntimeError, match="Failed to create Tenki sandbox"):
|
|
provider.create_instance("python")
|
|
|
|
|
|
def test_tenki_provider_enforces_max_artifacts(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
for i in range(3):
|
|
sandbox.fs.files[posixpath.join(cwd, "artifacts", f"f{i}.json")] = b"{}"
|
|
return _FakeCommandResult(0, stdout="")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
provider.max_artifacts = 2
|
|
|
|
inst = provider.create_instance("python")
|
|
with pytest.raises(RuntimeError, match="more than 2 artifacts"):
|
|
provider.execute_code(inst.instance_id, "def main():\n return 1\n", "python", timeout=5)
|
|
|
|
|
|
def test_tenki_provider_enforces_max_artifact_bytes(monkeypatch):
|
|
def run_handler(sandbox, argv, cwd):
|
|
sandbox.fs.files[posixpath.join(cwd, "artifacts", "big.json")] = b"x" * 5000
|
|
return _FakeCommandResult(0, stdout="")
|
|
|
|
sandbox = _FakeSandbox(run_handler)
|
|
provider, _ = _build_provider(sandbox, monkeypatch)
|
|
provider.max_artifact_bytes = 100
|
|
|
|
inst = provider.create_instance("python")
|
|
with pytest.raises(RuntimeError, match="exceeds 100 bytes"):
|
|
provider.execute_code(inst.instance_id, "def main():\n return 1\n", "python", timeout=5)
|