mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-19 05:50:41 +08:00
feat(python-sdk): add provider/policy/upload helpers and demo notebook
Extend python/openshell/ with the verbs needed to express the multi-agent-notepad flow entirely in Python: - ProviderClient (create/get/list/delete); credentials_from_env resolves env var names to mirror the CLI's --credential flag. - policy_from_yaml() loads the CLI's YAML format into a SandboxPolicy proto. Lazy-imports PyYAML, added to runtime deps. - SandboxClient.upload() gzip-tars the source and pipes the tarball over ExecSandbox(stdin=...) to 'tar xzf -'. No paramiko dependency. - SandboxClient.create() gains friendly kwargs (name, image, providers, policy, environment, gpu, labels) alongside the existing spec= path. architecture/python-sdk.md documents the design principle (mirror CLI verbs, not raw protos), the public surface, and what's out of scope (image build, gateway bootstrap, draft policy, port-forward, RelayStream). examples/multi-agent-notepad/notebook.ipynb is a Python equivalent of demo.sh: env validation, policy render, provider creation, parallel worker sandboxes via ThreadPoolExecutor, synthesis sandbox, cleanup. 13 new unit tests; all 23 sandbox tests pass; ruff/ty/markdownlint clean.
This commit is contained in:
Max Dubrinsky
@@ -10,12 +10,15 @@ from .sandbox import (
|
||||
ExecChunk,
|
||||
ExecResult,
|
||||
InferenceRouteClient,
|
||||
ProviderClient,
|
||||
ProviderRef,
|
||||
Sandbox,
|
||||
SandboxClient,
|
||||
SandboxError,
|
||||
SandboxRef,
|
||||
SandboxSession,
|
||||
TlsConfig,
|
||||
policy_from_yaml,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -30,6 +33,8 @@ __all__ = [
|
||||
"ExecChunk",
|
||||
"ExecResult",
|
||||
"InferenceRouteClient",
|
||||
"ProviderClient",
|
||||
"ProviderRef",
|
||||
"Sandbox",
|
||||
"SandboxClient",
|
||||
"SandboxError",
|
||||
@@ -37,4 +42,5 @@ __all__ = [
|
||||
"SandboxSession",
|
||||
"TlsConfig",
|
||||
"__version__",
|
||||
"policy_from_yaml",
|
||||
]
|
||||
|
||||
+204
-4
@@ -4,22 +4,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
import tarfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import grpc
|
||||
from google.protobuf import json_format
|
||||
|
||||
from ._proto import (
|
||||
datamodel_pb2,
|
||||
inference_pb2,
|
||||
inference_pb2_grpc,
|
||||
openshell_pb2,
|
||||
openshell_pb2_grpc,
|
||||
sandbox_pb2,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -41,6 +46,13 @@ class SandboxRef:
|
||||
phase: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProviderRef:
|
||||
id: str
|
||||
name: str
|
||||
type: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExecChunk:
|
||||
stream: str
|
||||
@@ -137,6 +149,13 @@ class SandboxClient:
|
||||
)
|
||||
self._channel = grpc.secure_channel(endpoint, credentials)
|
||||
self._stub = openshell_pb2_grpc.OpenShellStub(self._channel)
|
||||
self._providers: ProviderClient | None = None
|
||||
|
||||
@property
|
||||
def providers(self) -> ProviderClient:
|
||||
if self._providers is None:
|
||||
self._providers = ProviderClient(self._stub, timeout=self._timeout)
|
||||
return self._providers
|
||||
|
||||
@classmethod
|
||||
def from_active_cluster(
|
||||
@@ -186,12 +205,48 @@ class SandboxClient:
|
||||
self,
|
||||
*,
|
||||
spec: openshell_pb2.SandboxSpec | None = None,
|
||||
name: str | None = None,
|
||||
image: str | None = None,
|
||||
providers: Sequence[str] | None = None,
|
||||
policy: sandbox_pb2.SandboxPolicy | None = None,
|
||||
environment: Mapping[str, str] | None = None,
|
||||
gpu: bool = False,
|
||||
labels: Mapping[str, str] | None = None,
|
||||
) -> SandboxRef:
|
||||
request_spec = spec if spec is not None else _default_spec()
|
||||
response = self._stub.CreateSandbox(
|
||||
openshell_pb2.CreateSandboxRequest(spec=request_spec),
|
||||
timeout=self._timeout,
|
||||
kwargs_used = (
|
||||
image is not None
|
||||
or providers is not None
|
||||
or policy is not None
|
||||
or environment is not None
|
||||
or gpu
|
||||
)
|
||||
if spec is not None and kwargs_used:
|
||||
raise SandboxError(
|
||||
"create() accepts either spec= or the convenience kwargs "
|
||||
"(image, providers, policy, environment, gpu), not both"
|
||||
)
|
||||
|
||||
if spec is not None:
|
||||
request_spec = spec
|
||||
elif kwargs_used:
|
||||
request_spec = openshell_pb2.SandboxSpec(
|
||||
environment=dict(environment or {}),
|
||||
providers=list(providers or []),
|
||||
gpu=gpu,
|
||||
)
|
||||
if image is not None:
|
||||
request_spec.template.image = image
|
||||
if policy is not None:
|
||||
request_spec.policy.CopyFrom(policy)
|
||||
else:
|
||||
request_spec = _default_spec()
|
||||
|
||||
request = openshell_pb2.CreateSandboxRequest(
|
||||
spec=request_spec,
|
||||
name=name or "",
|
||||
labels=dict(labels or {}),
|
||||
)
|
||||
response = self._stub.CreateSandbox(request, timeout=self._timeout)
|
||||
sandbox_ref = _sandbox_ref(response.sandbox)
|
||||
if sandbox_ref.id == "":
|
||||
raise SandboxError("CreateSandbox returned empty sandbox id")
|
||||
@@ -259,6 +314,39 @@ class SandboxClient:
|
||||
time.sleep(1)
|
||||
raise SandboxError(f"sandbox {sandbox_name} was not ready within timeout")
|
||||
|
||||
def upload(
|
||||
self,
|
||||
sandbox_id: str,
|
||||
src: str | pathlib.Path,
|
||||
dst: str,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> None:
|
||||
src_path = pathlib.Path(src)
|
||||
if not src_path.exists():
|
||||
raise SandboxError(f"upload source does not exist: {src_path}")
|
||||
|
||||
buffer = io.BytesIO()
|
||||
with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
|
||||
tar.add(str(src_path), arcname=src_path.name)
|
||||
tarball = buffer.getvalue()
|
||||
|
||||
# Stream the tarball over ExecSandbox stdin instead of using the
|
||||
# gateway SSH tunnel; sufficient for notebook-sized payloads and
|
||||
# avoids an SSH client dependency.
|
||||
quoted_dst = "'" + dst.replace("'", "'\\''") + "'"
|
||||
result = self.exec(
|
||||
sandbox_id,
|
||||
["sh", "-c", f"mkdir -p {quoted_dst} && tar xzf - -C {quoted_dst}"],
|
||||
stdin=tarball,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise SandboxError(
|
||||
f"upload to {dst} failed (exit {result.exit_code}): "
|
||||
f"{result.stderr.strip() or result.stdout.strip()}"
|
||||
)
|
||||
|
||||
def exec_stream(
|
||||
self,
|
||||
sandbox_id: str,
|
||||
@@ -374,6 +462,80 @@ class SandboxClient:
|
||||
)
|
||||
|
||||
|
||||
class ProviderClient:
|
||||
"""gRPC client for provider CRUD."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stub: openshell_pb2_grpc.OpenShellStub,
|
||||
*,
|
||||
timeout: float = 30.0,
|
||||
) -> None:
|
||||
self._stub = stub
|
||||
self._timeout = timeout
|
||||
|
||||
def create(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
provider_type: str,
|
||||
credentials: Mapping[str, str] | None = None,
|
||||
credentials_from_env: Sequence[str] | None = None,
|
||||
config: Mapping[str, str] | None = None,
|
||||
labels: Mapping[str, str] | None = None,
|
||||
) -> ProviderRef:
|
||||
creds: dict[str, str] = dict(credentials or {})
|
||||
for env_name in credentials_from_env or ():
|
||||
value = os.environ.get(env_name)
|
||||
if value is None:
|
||||
raise SandboxError(
|
||||
f"environment variable {env_name!r} is not set; "
|
||||
f"cannot resolve credential for provider {name!r}"
|
||||
)
|
||||
creds[env_name] = value
|
||||
|
||||
provider = datamodel_pb2.Provider(
|
||||
metadata=datamodel_pb2.ObjectMeta(
|
||||
name=name,
|
||||
labels=dict(labels or {}),
|
||||
),
|
||||
type=provider_type,
|
||||
credentials=creds,
|
||||
config=dict(config or {}),
|
||||
)
|
||||
response = self._stub.CreateProvider(
|
||||
openshell_pb2.CreateProviderRequest(provider=provider),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return _provider_ref(response.provider)
|
||||
|
||||
def get(self, name: str) -> ProviderRef:
|
||||
response = self._stub.GetProvider(
|
||||
openshell_pb2.GetProviderRequest(name=name),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return _provider_ref(response.provider)
|
||||
|
||||
def list(
|
||||
self,
|
||||
*,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> builtins.list[ProviderRef]:
|
||||
response = self._stub.ListProviders(
|
||||
openshell_pb2.ListProvidersRequest(limit=limit, offset=offset),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return [_provider_ref(item) for item in response.providers]
|
||||
|
||||
def delete(self, name: str) -> bool:
|
||||
response = self._stub.DeleteProvider(
|
||||
openshell_pb2.DeleteProviderRequest(name=name),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return bool(response.deleted)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClusterInferenceConfig:
|
||||
provider_name: str
|
||||
@@ -549,6 +711,36 @@ class Sandbox:
|
||||
)
|
||||
|
||||
|
||||
def policy_from_yaml(
|
||||
source: str | pathlib.Path,
|
||||
) -> sandbox_pb2.SandboxPolicy:
|
||||
try:
|
||||
import yaml
|
||||
except ImportError as exc:
|
||||
raise SandboxError(
|
||||
"policy_from_yaml requires PyYAML; install with `pip install pyyaml`"
|
||||
) from exc
|
||||
|
||||
if isinstance(source, pathlib.Path):
|
||||
text = source.read_text()
|
||||
elif "\n" not in source and pathlib.Path(source).exists():
|
||||
text = pathlib.Path(source).read_text()
|
||||
else:
|
||||
text = source
|
||||
|
||||
parsed = yaml.safe_load(text)
|
||||
if not isinstance(parsed, dict):
|
||||
raise SandboxError("policy YAML must parse to a mapping at the top level")
|
||||
|
||||
# CLI YAML uses `filesystem_policy`; the proto field is `filesystem`.
|
||||
if "filesystem_policy" in parsed and "filesystem" not in parsed:
|
||||
parsed["filesystem"] = parsed.pop("filesystem_policy")
|
||||
|
||||
policy = sandbox_pb2.SandboxPolicy()
|
||||
json_format.ParseDict(parsed, policy, ignore_unknown_fields=False)
|
||||
return policy
|
||||
|
||||
|
||||
_PYTHON_CLOUDPICKLE_BOOTSTRAP = (
|
||||
"import base64,cloudpickle,os;"
|
||||
"payload=base64.b64decode(os.environ['OPENSHELL_PYFUNC_B64']);"
|
||||
@@ -583,6 +775,14 @@ def _sandbox_ref(sandbox: openshell_pb2.Sandbox) -> SandboxRef:
|
||||
)
|
||||
|
||||
|
||||
def _provider_ref(provider: datamodel_pb2.Provider) -> ProviderRef:
|
||||
return ProviderRef(
|
||||
id=provider.metadata.id if provider.metadata else "",
|
||||
name=provider.metadata.name if provider.metadata else "",
|
||||
type=provider.type,
|
||||
)
|
||||
|
||||
|
||||
def _default_spec() -> openshell_pb2.SandboxSpec:
|
||||
# Omit the policy field so the sandbox container discovers its policy
|
||||
# from /etc/openshell/policy.yaml (baked into the image at build time).
|
||||
|
||||
@@ -3,15 +3,22 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import tarfile
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
from openshell._proto import openshell_pb2
|
||||
import pytest
|
||||
|
||||
from openshell._proto import datamodel_pb2, openshell_pb2, sandbox_pb2
|
||||
from openshell.sandbox import (
|
||||
_PYTHON_CLOUDPICKLE_BOOTSTRAP,
|
||||
_SANDBOX_PYTHON_BIN,
|
||||
InferenceRouteClient,
|
||||
ProviderClient,
|
||||
SandboxClient,
|
||||
SandboxError,
|
||||
policy_from_yaml,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -21,6 +28,10 @@ if TYPE_CHECKING:
|
||||
class _FakeStub:
|
||||
def __init__(self) -> None:
|
||||
self.request: openshell_pb2.ExecSandboxRequest | None = None
|
||||
self.create_sandbox_request: openshell_pb2.CreateSandboxRequest | None = None
|
||||
self.exec_exit_code: int = 0
|
||||
self.exec_stdout: bytes = b""
|
||||
self.exec_stderr: bytes = b""
|
||||
|
||||
def ExecSandbox(
|
||||
self,
|
||||
@@ -29,10 +40,89 @@ class _FakeStub:
|
||||
):
|
||||
self.request = request
|
||||
_ = timeout
|
||||
if self.exec_stdout:
|
||||
yield openshell_pb2.ExecSandboxEvent(
|
||||
stdout=openshell_pb2.ExecSandboxStdout(data=self.exec_stdout)
|
||||
)
|
||||
if self.exec_stderr:
|
||||
yield openshell_pb2.ExecSandboxEvent(
|
||||
stderr=openshell_pb2.ExecSandboxStderr(data=self.exec_stderr)
|
||||
)
|
||||
yield openshell_pb2.ExecSandboxEvent(
|
||||
exit=openshell_pb2.ExecSandboxExit(exit_code=0)
|
||||
exit=openshell_pb2.ExecSandboxExit(exit_code=self.exec_exit_code)
|
||||
)
|
||||
|
||||
def CreateSandbox(
|
||||
self,
|
||||
request: openshell_pb2.CreateSandboxRequest,
|
||||
timeout: float | None = None,
|
||||
):
|
||||
self.create_sandbox_request = request
|
||||
_ = timeout
|
||||
sandbox = openshell_pb2.Sandbox(
|
||||
metadata=datamodel_pb2.ObjectMeta(
|
||||
id="sandbox-id-1",
|
||||
name=request.name or "generated-name",
|
||||
),
|
||||
phase=openshell_pb2.SANDBOX_PHASE_PROVISIONING,
|
||||
)
|
||||
return openshell_pb2.SandboxResponse(sandbox=sandbox)
|
||||
|
||||
|
||||
class _FakeProviderStub:
|
||||
def __init__(self) -> None:
|
||||
self.create_request: openshell_pb2.CreateProviderRequest | None = None
|
||||
self.get_request: openshell_pb2.GetProviderRequest | None = None
|
||||
self.list_request: openshell_pb2.ListProvidersRequest | None = None
|
||||
self.delete_request: openshell_pb2.DeleteProviderRequest | None = None
|
||||
self.delete_response_flag: bool = True
|
||||
self.providers_to_return: list[datamodel_pb2.Provider] = []
|
||||
|
||||
def CreateProvider(
|
||||
self,
|
||||
request: openshell_pb2.CreateProviderRequest,
|
||||
timeout: float | None = None,
|
||||
):
|
||||
self.create_request = request
|
||||
_ = timeout
|
||||
echo = datamodel_pb2.Provider()
|
||||
echo.CopyFrom(request.provider)
|
||||
echo.metadata.id = "provider-id-1"
|
||||
return openshell_pb2.ProviderResponse(provider=echo)
|
||||
|
||||
def GetProvider(
|
||||
self,
|
||||
request: openshell_pb2.GetProviderRequest,
|
||||
timeout: float | None = None,
|
||||
):
|
||||
self.get_request = request
|
||||
_ = timeout
|
||||
provider = datamodel_pb2.Provider(
|
||||
metadata=datamodel_pb2.ObjectMeta(id="get-id", name=request.name),
|
||||
type="generic",
|
||||
)
|
||||
return openshell_pb2.ProviderResponse(provider=provider)
|
||||
|
||||
def ListProviders(
|
||||
self,
|
||||
request: openshell_pb2.ListProvidersRequest,
|
||||
timeout: float | None = None,
|
||||
):
|
||||
self.list_request = request
|
||||
_ = timeout
|
||||
return openshell_pb2.ListProvidersResponse(
|
||||
providers=self.providers_to_return,
|
||||
)
|
||||
|
||||
def DeleteProvider(
|
||||
self,
|
||||
request: openshell_pb2.DeleteProviderRequest,
|
||||
timeout: float | None = None,
|
||||
):
|
||||
self.delete_request = request
|
||||
_ = timeout
|
||||
return openshell_pb2.DeleteProviderResponse(deleted=self.delete_response_flag)
|
||||
|
||||
|
||||
class _FakeInferenceStub:
|
||||
def __init__(self) -> None:
|
||||
@@ -54,6 +144,14 @@ def _client_with_fake_stub(stub: _FakeStub) -> SandboxClient:
|
||||
client = cast("SandboxClient", object.__new__(SandboxClient))
|
||||
client._timeout = 30.0
|
||||
client._stub = cast("Any", stub)
|
||||
client._providers = None
|
||||
return client
|
||||
|
||||
|
||||
def _provider_client_with_fake_stub(stub: _FakeProviderStub) -> ProviderClient:
|
||||
client = cast("ProviderClient", object.__new__(ProviderClient))
|
||||
client._timeout = 30.0
|
||||
client._stub = cast("Any", stub)
|
||||
return client
|
||||
|
||||
|
||||
@@ -153,3 +251,298 @@ def test_inference_set_cluster_forwards_no_verify_flag() -> None:
|
||||
|
||||
assert stub.request is not None
|
||||
assert stub.request.no_verify is True
|
||||
|
||||
|
||||
def test_provider_create_resolves_credentials_from_env(monkeypatch: Any) -> None:
|
||||
monkeypatch.setenv("CODEX_AUTH_ACCESS_TOKEN", "access-secret")
|
||||
monkeypatch.setenv("CODEX_AUTH_REFRESH_TOKEN", "refresh-secret")
|
||||
|
||||
stub = _FakeProviderStub()
|
||||
client = _provider_client_with_fake_stub(stub)
|
||||
|
||||
ref = client.create(
|
||||
name="codex-oauth",
|
||||
provider_type="generic",
|
||||
credentials_from_env=[
|
||||
"CODEX_AUTH_ACCESS_TOKEN",
|
||||
"CODEX_AUTH_REFRESH_TOKEN",
|
||||
],
|
||||
)
|
||||
|
||||
assert stub.create_request is not None
|
||||
provider = stub.create_request.provider
|
||||
assert provider.metadata.name == "codex-oauth"
|
||||
assert provider.type == "generic"
|
||||
assert dict(provider.credentials) == {
|
||||
"CODEX_AUTH_ACCESS_TOKEN": "access-secret",
|
||||
"CODEX_AUTH_REFRESH_TOKEN": "refresh-secret",
|
||||
}
|
||||
assert ref.id == "provider-id-1"
|
||||
assert ref.name == "codex-oauth"
|
||||
assert ref.type == "generic"
|
||||
|
||||
|
||||
def test_provider_create_passes_explicit_credentials(monkeypatch: Any) -> None:
|
||||
monkeypatch.setenv("DEMO_GITHUB_TOKEN", "from-env")
|
||||
|
||||
stub = _FakeProviderStub()
|
||||
client = _provider_client_with_fake_stub(stub)
|
||||
|
||||
client.create(
|
||||
name="github-memory",
|
||||
provider_type="generic",
|
||||
credentials={"DEMO_GITHUB_TOKEN": "explicit-override"},
|
||||
config={"endpoint": "api.github.com"},
|
||||
)
|
||||
|
||||
assert stub.create_request is not None
|
||||
provider = stub.create_request.provider
|
||||
assert dict(provider.credentials) == {"DEMO_GITHUB_TOKEN": "explicit-override"}
|
||||
assert dict(provider.config) == {"endpoint": "api.github.com"}
|
||||
|
||||
|
||||
def test_provider_create_missing_env_raises(monkeypatch: Any) -> None:
|
||||
monkeypatch.delenv("MISSING_TOKEN", raising=False)
|
||||
|
||||
stub = _FakeProviderStub()
|
||||
client = _provider_client_with_fake_stub(stub)
|
||||
|
||||
with pytest.raises(SandboxError, match="MISSING_TOKEN"):
|
||||
client.create(
|
||||
name="x",
|
||||
provider_type="generic",
|
||||
credentials_from_env=["MISSING_TOKEN"],
|
||||
)
|
||||
|
||||
assert stub.create_request is None
|
||||
|
||||
|
||||
def test_provider_delete_returns_response_flag() -> None:
|
||||
stub = _FakeProviderStub()
|
||||
stub.delete_response_flag = False
|
||||
client = _provider_client_with_fake_stub(stub)
|
||||
|
||||
deleted = client.delete("github-memory")
|
||||
|
||||
assert deleted is False
|
||||
assert stub.delete_request is not None
|
||||
assert stub.delete_request.name == "github-memory"
|
||||
|
||||
|
||||
def test_provider_list_returns_refs() -> None:
|
||||
stub = _FakeProviderStub()
|
||||
stub.providers_to_return = [
|
||||
datamodel_pb2.Provider(
|
||||
metadata=datamodel_pb2.ObjectMeta(id="p1", name="alpha"),
|
||||
type="generic",
|
||||
),
|
||||
datamodel_pb2.Provider(
|
||||
metadata=datamodel_pb2.ObjectMeta(id="p2", name="beta"),
|
||||
type="claude",
|
||||
),
|
||||
]
|
||||
client = _provider_client_with_fake_stub(stub)
|
||||
|
||||
refs = client.list(limit=50)
|
||||
|
||||
assert stub.list_request is not None
|
||||
assert stub.list_request.limit == 50
|
||||
assert [(r.id, r.name, r.type) for r in refs] == [
|
||||
("p1", "alpha", "generic"),
|
||||
("p2", "beta", "claude"),
|
||||
]
|
||||
|
||||
|
||||
def test_policy_from_yaml_loads_demo_template() -> None:
|
||||
yaml_text = """
|
||||
version: 1
|
||||
filesystem_policy:
|
||||
include_workdir: true
|
||||
read_only: [/usr, /lib]
|
||||
read_write: [/sandbox, /tmp]
|
||||
landlock:
|
||||
compatibility: best_effort
|
||||
process:
|
||||
run_as_user: sandbox
|
||||
run_as_group: sandbox
|
||||
network_policies:
|
||||
github_memory:
|
||||
name: github-memory
|
||||
endpoints:
|
||||
- host: api.github.com
|
||||
port: 443
|
||||
protocol: rest
|
||||
enforcement: enforce
|
||||
rules:
|
||||
- allow:
|
||||
method: GET
|
||||
path: "/repos/owner/repo/contents/runs/run-id"
|
||||
"""
|
||||
|
||||
policy = policy_from_yaml(yaml_text)
|
||||
|
||||
assert policy.version == 1
|
||||
assert policy.filesystem.include_workdir is True
|
||||
assert list(policy.filesystem.read_only) == ["/usr", "/lib"]
|
||||
assert list(policy.filesystem.read_write) == ["/sandbox", "/tmp"]
|
||||
assert policy.landlock.compatibility == "best_effort"
|
||||
assert policy.process.run_as_user == "sandbox"
|
||||
|
||||
rule = policy.network_policies["github_memory"]
|
||||
assert rule.name == "github-memory"
|
||||
assert len(rule.endpoints) == 1
|
||||
endpoint = rule.endpoints[0]
|
||||
assert endpoint.host == "api.github.com"
|
||||
assert endpoint.port == 443
|
||||
assert endpoint.protocol == "rest"
|
||||
assert len(endpoint.rules) == 1
|
||||
assert endpoint.rules[0].allow.method == "GET"
|
||||
assert endpoint.rules[0].allow.path == "/repos/owner/repo/contents/runs/run-id"
|
||||
|
||||
|
||||
def test_policy_from_yaml_reads_path(tmp_path: Path) -> None:
|
||||
policy_file = tmp_path / "policy.yaml"
|
||||
policy_file.write_text("version: 2\nfilesystem_policy:\n include_workdir: false\n")
|
||||
|
||||
policy = policy_from_yaml(policy_file)
|
||||
|
||||
assert policy.version == 2
|
||||
assert policy.filesystem.include_workdir is False
|
||||
|
||||
|
||||
def test_policy_from_yaml_rejects_non_mapping() -> None:
|
||||
with pytest.raises(SandboxError, match="mapping"):
|
||||
policy_from_yaml("- just-a-list\n- item-2\n")
|
||||
|
||||
|
||||
def test_upload_packages_tarball_and_pipes_to_exec(tmp_path: Path) -> None:
|
||||
src = tmp_path / "payload"
|
||||
src.mkdir()
|
||||
(src / "runner.sh").write_text("#!/bin/sh\necho hello\n")
|
||||
(src / "prompts").mkdir()
|
||||
(src / "prompts" / "worker.md").write_text("# worker\n")
|
||||
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
client.upload("sandbox-1", src, "/sandbox")
|
||||
|
||||
assert stub.request is not None
|
||||
assert stub.request.sandbox_id == "sandbox-1"
|
||||
assert stub.request.command == [
|
||||
"sh",
|
||||
"-c",
|
||||
"mkdir -p '/sandbox' && tar xzf - -C '/sandbox'",
|
||||
]
|
||||
assert stub.request.stdin
|
||||
with tarfile.open(fileobj=io.BytesIO(stub.request.stdin), mode="r:gz") as tar:
|
||||
names = sorted(tar.getnames())
|
||||
assert "payload" in names
|
||||
assert "payload/runner.sh" in names
|
||||
assert "payload/prompts/worker.md" in names
|
||||
|
||||
|
||||
def test_upload_missing_source_raises(tmp_path: Path) -> None:
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
with pytest.raises(SandboxError, match="upload source does not exist"):
|
||||
client.upload("sandbox-1", tmp_path / "nope", "/sandbox")
|
||||
|
||||
assert stub.request is None
|
||||
|
||||
|
||||
def test_upload_failure_surfaces_stderr(tmp_path: Path) -> None:
|
||||
src = tmp_path / "payload"
|
||||
src.mkdir()
|
||||
(src / "f.txt").write_text("ok")
|
||||
|
||||
stub = _FakeStub()
|
||||
stub.exec_exit_code = 2
|
||||
stub.exec_stderr = b"tar: cannot write: read-only filesystem"
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
with pytest.raises(SandboxError, match="read-only filesystem"):
|
||||
client.upload("sandbox-1", src, "/sandbox")
|
||||
|
||||
|
||||
def test_upload_quotes_destination_with_single_quote(tmp_path: Path) -> None:
|
||||
src = tmp_path / "f.txt"
|
||||
src.write_text("ok")
|
||||
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
client.upload("sandbox-1", src, "/sand'box")
|
||||
|
||||
assert stub.request is not None
|
||||
assert stub.request.command[2] == (
|
||||
"mkdir -p '/sand'\\''box' && tar xzf - -C '/sand'\\''box'"
|
||||
)
|
||||
|
||||
|
||||
def test_create_with_kwargs_constructs_spec() -> None:
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
policy = sandbox_pb2.SandboxPolicy(version=3)
|
||||
policy.filesystem.include_workdir = True
|
||||
|
||||
ref = client.create(
|
||||
name="agent-1",
|
||||
image="base",
|
||||
providers=["codex-oauth", "github-memory"],
|
||||
policy=policy,
|
||||
environment={"AGENT_INDEX": "1"},
|
||||
labels={"run-id": "20260506-001"},
|
||||
)
|
||||
|
||||
assert stub.create_sandbox_request is not None
|
||||
request = stub.create_sandbox_request
|
||||
assert request.name == "agent-1"
|
||||
assert dict(request.labels) == {"run-id": "20260506-001"}
|
||||
assert request.spec.template.image == "base"
|
||||
assert list(request.spec.providers) == ["codex-oauth", "github-memory"]
|
||||
assert dict(request.spec.environment) == {"AGENT_INDEX": "1"}
|
||||
assert request.spec.policy.version == 3
|
||||
assert request.spec.policy.filesystem.include_workdir is True
|
||||
assert ref.id == "sandbox-id-1"
|
||||
assert ref.name == "agent-1"
|
||||
|
||||
|
||||
def test_create_default_uses_empty_spec_with_name() -> None:
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
client.create(name="just-a-name")
|
||||
|
||||
assert stub.create_sandbox_request is not None
|
||||
request = stub.create_sandbox_request
|
||||
assert request.name == "just-a-name"
|
||||
assert request.spec.template.image == ""
|
||||
assert list(request.spec.providers) == []
|
||||
|
||||
|
||||
def test_create_rejects_spec_and_kwargs_together() -> None:
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
spec = openshell_pb2.SandboxSpec(
|
||||
template=openshell_pb2.SandboxTemplate(image="explicit"),
|
||||
)
|
||||
|
||||
with pytest.raises(SandboxError, match="spec="):
|
||||
client.create(spec=spec, image="conflicting")
|
||||
|
||||
assert stub.create_sandbox_request is None
|
||||
|
||||
|
||||
def test_sandbox_client_providers_property_caches() -> None:
|
||||
stub = _FakeStub()
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
first = client.providers
|
||||
second = client.providers
|
||||
|
||||
assert first is second
|
||||
assert isinstance(first, ProviderClient)
|
||||
|
||||
Reference in New Issue
Block a user