From ea592997e25472ee6ca6564fe4024f328dcd7e49 Mon Sep 17 00:00:00 2001 From: Franco Garcia Date: Tue, 13 Oct 2026 20:41:52 -0300 Subject: [PATCH] feat(sandbox): add a Tenki Sandbox VM sandbox provider Adds `tenki` as a built-in VM sandbox provider behind a `tenki` extra. Each sandbox is one Tenki VM running agent-env's docker-in-VM flows: commands run as root via privileged exec, objects go over exec stdin, and exposed ports get public HTTPS URLs. Allowlist egress is enforced outside the VM and set at create, so bytes the policy doesn't cover go over stdin. Config: api_key (required), base_url, snapshot_id or image, min_disk_gb. --- AGENTS.md | 6 +- README.md | 2 +- THIRD_PARTY_NOTICES.md | 9 +- pyproject.toml | 6 +- src/agent_env/cli/a2a_agent/deploy.py | 2 +- src/agent_env/cli/env/deploy.py | 2 +- .../sandbox_providers/sandbox_provider.py | 1 + .../sandbox_providers/tenki/README.md | 156 +++++++ .../sandbox_providers/tenki/__init__.py | 1 + .../sandbox_providers/tenki/provider.py | 415 ++++++++++++++++++ .../sandbox_providers/tenki/sandbox.py | 331 ++++++++++++++ tst/integration/env/gateway/gateway_test.py | 8 + .../env/task_step/task_steps_test.py | 9 + .../tenki_sandbox_smoke_test.py | 104 +++++ .../sandbox_providers/egress_hosts_test.py | 4 +- .../local_sandbox_host_ports_test.py | 3 +- .../network_policy_stamping_test.py | 12 + .../provider_capabilities_test.py | 10 +- .../sandbox_providers/tenki_fakes.py | 161 +++++++ .../tenki_sandbox_provider_test.py | 385 ++++++++++++++++ .../sandbox_providers/tenki_sandbox_test.py | 366 +++++++++++++++ .../test_config_sandbox_providers.py | 30 ++ tst/unit/test_capabilities.py | 2 +- .../test_optional_extras_stay_optional.py | 8 +- tst/util/capabilities.py | 14 +- uv.lock | 28 +- 26 files changed, 2051 insertions(+), 24 deletions(-) create mode 100644 src/agent_env/providers/sandbox_providers/tenki/README.md create mode 100644 src/agent_env/providers/sandbox_providers/tenki/__init__.py create mode 100644 src/agent_env/providers/sandbox_providers/tenki/provider.py create mode 100644 src/agent_env/providers/sandbox_providers/tenki/sandbox.py create mode 100644 tst/integration/providers/sandbox_providers/tenki_sandbox_smoke_test.py create mode 100644 tst/unit/providers/sandbox_providers/tenki_fakes.py create mode 100644 tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py create mode 100644 tst/unit/providers/sandbox_providers/tenki_sandbox_test.py diff --git a/AGENTS.md b/AGENTS.md index c204c595..42d3574e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -60,8 +60,8 @@ or one of them only where it is used. - Defaults are local and carry no external coordinates: SQLite document store, filesystem object store, an OCI registry at `localhost:5000`, env-var secret store, `LocalRunner`, `local` Docker sandboxes for envs and agents. MongoDB, S3, Cloud Storage, ECR, AWS Secrets Manager, Google Cloud - Secret Manager, Modal, E2B and Tensorlake exist as implementations and are selected by config - (Tensorlake through the `tensorlake` extra). + Secret Manager, Modal, E2B, Tensorlake and Tenki exist as implementations and are selected by config + (Tensorlake through the `tensorlake` extra, Tenki through the `tenki` extra). - The model endpoint is unset until `[model] base_url` / `api_key` (or `LITELLM_BASE_URL` / `LITELLM_API_KEY`) is configured. A loopback URL handed into a `local` container (the model endpoint, an agent, peer or gateway on this machine) becomes `host.docker.internal`. On Linux that @@ -96,7 +96,7 @@ and deserialized through a registry. | `task/`, `task_step/` | A `Task` holds its `TaskStep`s inline; `Task.run()` executes them as a DAG. `depends_on` (None means all prior steps) gates a step, independent steps run concurrently, `fail_task_on_error` makes a failure fatal or tolerated, `retry_config` rolls a failed span back through the step journal and re-dispatches it. Built-in steps live in `task_step/task_steps/` (`deploy_env`, `deploy_agent`, `prompt_agent`, the verifiers under `verifiers/`, and more). | | `store/` | Four store ABCs with local and cloud implementations: `DocumentStore` (SQLite, MongoDB), `ObjectStore` (filesystem, S3, Cloud Storage), `ImageStore` (local OCI registry, ECR), `SecretStore` (env vars or file, AWS Secrets Manager, Google Cloud Secret Manager). `VersionedEntityStore` implements the shared versioned get/put logic, `QueryBuilder` is the immutable chained query API, `store/base.py` holds the error types. A new backend must pass the conformance kits in `tst/store/`. | | `config/` | The `Config` singleton (`get_config`, `configure`, `reset_config`) in `config/runtime.py`, file discovery in `config/loader.py`, and `load_impl`, which resolves `module:Class` pointers. `agent_env.store` re-exports the config names for compatibility. | -| `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `tensorlake` (the `tensorlake` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store) and `EnvironmentGatewayProvider`, which renders a docker-compose for the gateway and its MCP servers inside the sandbox; `providers/env_state/` holds env-state providers (`local_postgres` built in). | +| `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store) and `EnvironmentGatewayProvider`, which renders a docker-compose for the gateway and its MCP servers inside the sandbox; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns `(run_id, instance_id)`; `LocalRunner` is built in. | | `explorer/` | Optional local web UI: `agent-env up`, needs the `explorer` extra, binds loopback `:8234`. | diff --git a/README.md b/README.md index 1132669a..51df7650 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ Or install it with pip, into a virtualenv: pip install agentenv-framework ``` -To use the SDK in your own project, add it as a dependency with `uv add agentenv-framework`. The explorer, `agent-env up`, needs the `explorer` extra (`uv tool install 'agentenv-framework[explorer]'`, or the same extra with uvx or pip) and an `.agentenv/config.toml` in the current folder or above it; an empty one keeps every local default. The `tensorlake` sandbox backend needs the `tensorlake` extra (`pip install 'agentenv-framework[tensorlake]'`) and an `api_key` in `[sandbox.providers.tensorlake.config]`. Its sandboxes boot from the public image `agentenv-dind-host-v1`; set `image` there to use another registered Docker-capable image. +To use the SDK in your own project, add it as a dependency with `uv add agentenv-framework`. The explorer, `agent-env up`, needs the `explorer` extra (`uv tool install 'agentenv-framework[explorer]'`, or the same extra with uvx or pip) and an `.agentenv/config.toml` in the current folder or above it; an empty one keeps every local default. The `tensorlake` sandbox backend needs the `tensorlake` extra (`pip install 'agentenv-framework[tensorlake]'`) and an `api_key` in `[sandbox.providers.tensorlake.config]`. Its sandboxes boot from the public image `agentenv-dind-host-v1`; set `image` there to use another registered Docker-capable image. The `tenki` sandbox backend needs the `tenki` extra (`pip install 'agentenv-framework[tenki]'`) and an `api_key` in `[sandbox.providers.tenki.config]`; see [its README](src/agent_env/providers/sandbox_providers/tenki/README.md). A plain install runs on local stores and needs no cloud SDK. The cloud store backends are extras: `aws` for S3, Secrets Manager, DynamoDB and ECR (`pip install 'agentenv-framework[aws]'`), and `gcp` for Cloud Storage, Secret Manager, Firestore and Artifact Registry (`pip install 'agentenv-framework[gcp]'`). A config that names a backend without its extra fails, naming the extra to install. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index e0f4b0ff..5a103bc8 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -14,7 +14,7 @@ marks, or contributor names to endorse or promote Scale AI, or related products. ## Python dependencies -The runtime dependencies of `agentenv-framework` and its `aws`, `explorer`, `gcp`, `sail` and `tensorlake` extras, at the versions resolved in `uv.lock`. Development-only dependencies are not listed. +The runtime dependencies of `agentenv-framework` and its `aws`, `explorer`, `gcp`, `sail`, `tensorlake` and `tenki` extras, at the versions resolved in `uv.lock`. Development-only dependencies are not listed. ### a2a-sdk 0.3.26 @@ -716,6 +716,13 @@ The runtime dependencies of `agentenv-framework` and its `aws`, `explorer`, `gcp - Source: - [License text 1](#license-text-1) +### tenki 1.4.1 + +- License: MIT +- Author: Tenki Cloud +- Source: +- License text: not shipped with the package; see its source + ### tensorlake 0.5.148 - License: Apache-2.0 diff --git a/pyproject.toml b/pyproject.toml index 86d72f42..f67ebef1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,8 +76,12 @@ sail = [ tensorlake = [ "tensorlake~=0.5.140", # AsyncSandbox APIs validated against 0.5.140 to 0.5.149 ] +# The Tenki Sandbox VM provider, selected as the built-in `tenki` backend. +tenki = [ + "tenki[async]~=1.4.1", # AsyncClient/AsyncSandbox exec, stdin, ports and egress APIs validated against the 1.4 line +] dev = [ - "agentenv-framework[aws,explorer,gcp,sail,tensorlake]", # so the AWS, explorer, GCP backend, Sail and Tensorlake provider tests run in CI + "agentenv-framework[aws,explorer,gcp,sail,tensorlake,tenki]", # so the AWS, explorer, GCP backend, Sail, Tensorlake and Tenki provider tests run in CI "fastmcp>=3.4,<4", # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x "griffelib==2.3.0", # the plugin API check; exact, since what it reports as a break changes between releases "moto>=5.0.0", diff --git a/src/agent_env/cli/a2a_agent/deploy.py b/src/agent_env/cli/a2a_agent/deploy.py index 70603161..831d7d5d 100644 --- a/src/agent_env/cli/a2a_agent/deploy.py +++ b/src/agent_env/cli/a2a_agent/deploy.py @@ -18,7 +18,7 @@ @click.option("--ttl-seconds", type=click.IntRange(min=MIN_TTL_SECONDS, max=MAX_TTL_SECONDS), default=DEFAULT_TTL_SECONDS, help=f"VM lifetime in seconds (default {DEFAULT_TTL_SECONDS})") @click.option("--sandbox", default=None, - help="Sandbox backend(s): a built-in (modal, modal_vm, e2b, tensorlake, local) " + help="Sandbox backend(s): a built-in (modal, modal_vm, e2b, tensorlake, tenki, local) " "or a name from [sandbox.providers] in .agentenv/config.toml; comma-separated for a " "fallback chain. Defaults to [sandbox].agent_default (else local) when omitted.") def deploy(agent_id: str, agent_version: int | None, env_var_pairs: tuple[str, ...], ttl_seconds: int, sandbox: str): diff --git a/src/agent_env/cli/env/deploy.py b/src/agent_env/cli/env/deploy.py index 470bbf1b..7d84f754 100644 --- a/src/agent_env/cli/env/deploy.py +++ b/src/agent_env/cli/env/deploy.py @@ -19,7 +19,7 @@ @click.option("--gateway-mode", type=click.Choice([m.value for m in GatewayMode], case_sensitive=False), default=GatewayMode.PERFORMANCE.value, help="Gateway mode (performance or consistent)") @click.option("--sandbox", default=None, - help="Sandbox backend(s): a built-in (modal, modal_vm, e2b, tensorlake, local) " + help="Sandbox backend(s): a built-in (modal, modal_vm, e2b, tensorlake, tenki, local) " "or a name from [sandbox.providers] in .agentenv/config.toml; comma-separated for a " "fallback chain. Defaults to [sandbox].default (else local) when omitted.") @click.option("--service-db", "service_db_env_id", default=None, diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 7850da48..26539ea0 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -296,6 +296,7 @@ def effective_network_policy(cls, policy: Optional[NetworkPolicy]) -> NetworkPol "e2b": "agent_env.providers.sandbox_providers.e2b:E2BSandboxProvider", "sail_vm": "agent_env.providers.sandbox_providers.sail_vm.provider:SailVmSandboxProvider", "tensorlake": "agent_env.providers.sandbox_providers.tensorlake.provider:TensorlakeSandboxProvider", + "tenki": "agent_env.providers.sandbox_providers.tenki.provider:TenkiSandboxProvider", "local": "agent_env.providers.sandbox_providers.local_sandbox:LocalSandboxProvider", } diff --git a/src/agent_env/providers/sandbox_providers/tenki/README.md b/src/agent_env/providers/sandbox_providers/tenki/README.md new file mode 100644 index 00000000..79406c2d --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/tenki/README.md @@ -0,0 +1,156 @@ +# `tenki` sandbox provider + +Runs agent-env sandboxes on [Tenki Sandbox](https://tenki.cloud/docs/sandbox) VMs. Each environment and each agent +gets its own Linux VM, with Docker inside and public HTTPS URLs for its ports. It is a VM provider: agent-env's +docker-in-VM flows run on it unchanged. That covers the gateway's docker-compose, an agent's `docker run`, image loading +and artifact collection. + +## Set up + +### 1. Install the extra + +The Tenki SDK is an optional extra: + +```bash +pip install 'agentenv-framework[tenki]' # or: uv add 'agentenv-framework[tenki]' +``` + +Without it, selecting `tenki` fails with a `ConfigError` that names this extra. Nothing else needs it. + +### 2. Get a Tenki API key + +Create a key as described in [API keys](https://tenki.cloud/docs/account/api-keys). For local work, export it: + +```bash +export TENKI_API_KEY=tk_... +``` + +For a shared deployment, put it in your secret store instead, for example as `tenki_api_key`, and reference it with +`secret:tenki_api_key`. Never put the key itself in `config.toml`. The provider reads the key only from its config, +never from the process environment. + +### 3. Configure + +A complete `.agentenv/config.toml` that runs every environment and agent on Tenki, with the local stores: + +```toml +[sandbox] +default = "tenki" # environments and sandboxes +agent_default = "tenki" # agents + +[sandbox.providers.tenki.config] +api_key = "env:TENKI_API_KEY" # or "secret:tenki_api_key" + +[model] +base_url = "https://litellm.example.com" +api_key = "env:LITELLM_API_KEY" +``` + +To keep the local default and use Tenki per run instead, add only the `[sandbox.providers.tenki.config]` table and +pass `--sandbox tenki`, or `--env-sandbox tenki` / `--agent-sandbox tenki` on `agent-env task run`. +`agent-env config show` prints the file in effect and masks the key. + +### 4. Run something + +The bundled `hello` task deploys a sandbox, loads a file into it and checks it, with no model needed: + +```bash +agent-env run hello --sandbox tenki +``` + +`agent-env -v run …` also logs each VM as it starts +(`Tenki VM started: session_id=… cpu=… memory=…MB disk=…GB ports=… attribution=…`). + +## Configuration reference + +All keys go under `[sandbox.providers.tenki.config]`. Unknown keys are refused. + +| Key | Default | Meaning | +| ------------- | ------------------------- | ----------------------------------------------------------------------------------------------------------- | +| `api_key` | required | The Tenki API key, as an `env:` or `secret:` reference. | +| `base_url` | `https://api.tenki.cloud` | The Tenki API endpoint. | +| `snapshot_id` | none | Boot every VM from this snapshot, for example one with Docker already installed (see below). | +| `image` | the Tenki base image | Boot every VM from this Tenki registry image instead. At most one of `image` and `snapshot_id`. | +| `min_disk_gb` | `20` | The smallest VM disk. Docker keeps compressed and extracted layers, so agent-env's 10 GB default is raised. | + +## Faster startup + +A VM booted from the Tenki base image installs Docker and Compose v2 the first time agent-env uses it. Snapshot a VM +that has them once, and every VM booted from the snapshot skips that step: + +```bash +tenki sandbox create --name agentenv-docker --cpu 2 --memory-mb 4096 --disk-size-gb 20 --max-duration 30m +tenki sandbox exec --session agentenv-docker -c 'sudo apt-get update -qq && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq docker.io docker-compose-v2 docker-buildx curl' +tenki sandbox exec --session agentenv-docker -c 'sudo systemctl start docker && sudo docker info --format "{{.ServerVersion}}"' +tenki sandbox snapshot create agentenv-docker --name agentenv-docker --wait-durable +tenki sandbox terminate agentenv-docker +``` + +Then set `snapshot_id` to the snapshot's id. A VM can be restored larger than its snapshot but not smaller, so keep +the snapshot's VM at 2 vCPU and 4 GiB, the smallest the provider starts. Images your environments and agents pull +from a registry can be pulled into the same VM before the snapshot; a VM booted from it then finds them in place. + +## Resources and lifetime + +- **CPU and memory:** whole vCPUs, memory rounded up to 2 MiB. A request below 2 vCPU or 4 GiB is raised to it, the + smallest VM Tenki starts from a warm snapshot. +- **Disk:** at least `min_disk_gb`. +- **Refused requests:** a request beyond Tenki's limits (128 vCPU, 512 GiB memory, 100 GB disk, or the workspace's own + lower limit) is refused before anything is created, not shrunk. +- **Ignored requests:** `boot_mode` has no Tenki equivalent, so a `boot_mode` that is set logs a warning and has no + effect. +- **Lifetime:** the sandbox's `timeout` is the VM's hard maximum lifetime, so a crashed run can't leave a VM running + past it. It is also every command's execution budget. A workspace with a lower maximum session duration grants that + instead; the provider logs a warning and budgets commands by the lifetime granted. + +## Networking + +- **Ports:** each exposed port gets a public HTTPS URL under `sb.tenki.sh`, which fills `tunnel_urls`. It serves HTTP + and WebSocket with no platform authentication. +- **Egress:** allow-all, or an allowlist of up to 256 hostnames and 64 IPv4 CIDRs. An IP address listed as a host + becomes a /32. A policy with IPv6 or malformed entries is refused with `NetworkPolicyUnsupportedError` before + anything is created, and a fallback chain skips to its next provider. Tenki applies the policy outside the VM, so + one policy covers `dockerd` and every container. +- **Wildcards:** `*.example.com` covers exactly one label: `files.example.com`, not `a.b.example.com`. List deeper names + explicitly. +- **Tenki's own hosts:** `*.sb.tenki.sh` and `*.us.sb.tenki.sh` are added to every allowlist, so a restricted agent can + reach the environments it talks to. A workspace with wildcard preview domains gets URLs one label deeper + (`..us.sb.tenki.sh`); the provider adds `*..us.sb.tenki.sh` once it has exposed a port + there. A restricted VM created before any environment in the run needs that entry in its own `allow_hosts`. +- **Image and object downloads:** Tenki sets a VM's egress when it is created. Objects and image tarballs whose host the + allowlist doesn't cover go over the exec's stdin instead, which needs no egress. Image pulls and builds inside the + VM still need their registries allowed, and so does installing Docker on the base image, so boot allowlisted VMs + from a snapshot that has Docker. +- **Reconnect:** `get_sandbox` restores the ports and the applied egress policy. + +## Attribution and cost + +Every VM is tagged `agentenv` and carries agent-env's attribution, including `[sandbox.attribution]` defaults, in its +session metadata. `tenki sandbox list --tags agentenv` lists them, and `tenki sandbox get --json` shows a +VM's metadata. + +## Cleanup + +- A VM whose setup fails after it is created (readiness, ports, Docker) is terminated. +- If the caller is cancelled while the create is in flight, the provider terminates the VM that create makes once it + answers. +- If a create fails after Tenki may have made the VM, the provider finds that VM by its unique name and terminates it. +- `terminate` is idempotent: a VM that is already gone or terminating is not an error. + +## Things to know + +- **Commands run as root** through Tenki's privileged exec, so the `sudo` agent-env prefixes is dropped. +- **Lost connections:** when a command's stream ends without an exit status, `exec` returns exit code -1. + `exec_script` retries idempotent scripts on -1 only. +- **Images:** `create_vm(image=...)` is refused. VMs boot the configured `image` or `snapshot_id`, and OCI images run + under Docker inside them. +- **GPUs** are not supported. +- **The model key enters the sandbox:** an agent's `LITELLM_API_KEY` is passed in, as on other providers. + +## Testing + +`tst/integration/providers/sandbox_providers/tenki_sandbox_smoke_test.py` runs against a real Tenki workspace when +`[sandbox.providers.tenki.config]` resolves, for example with `make int-test-slow`. Otherwise it skips with +`agentenv-capability-missing: remote_sandbox`. The unit tests are +`tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py` and `tenki_sandbox_test.py`, and run offline +against in-memory stand-ins for the SDK. diff --git a/src/agent_env/providers/sandbox_providers/tenki/__init__.py b/src/agent_env/providers/sandbox_providers/tenki/__init__.py new file mode 100644 index 00000000..688e8917 --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/tenki/__init__.py @@ -0,0 +1 @@ +"""Tenki Sandbox VM sandbox provider (the ``tenki`` extra); see README.md.""" diff --git a/src/agent_env/providers/sandbox_providers/tenki/provider.py b/src/agent_env/providers/sandbox_providers/tenki/provider.py new file mode 100644 index 00000000..bf00a79d --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/tenki/provider.py @@ -0,0 +1,415 @@ +"""Tenki Sandbox VM sandbox provider. + +Each sandbox is one Tenki session, a full Linux VM. The gateway and agent containers run under Docker inside it +through the inherited ``VmSandbox`` helpers; a VM booted from an image without Docker installs it first. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import logging +import math +import re +import uuid +import weakref +from datetime import UTC, datetime +from typing import Any, ClassVar, Self +from urllib.parse import urlparse + +from tenki import AsyncClient + +from agent_env.attribution import Attribution +from agent_env.config.errors import ConfigError +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, NetworkPolicyUnsupportedError +from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider, apply_default_attribution +from agent_env.providers.sandbox_providers.tenki.sandbox import ENDED_STATES, TenkiSandbox, policy_from_info + +logger = logging.getLogger(__name__) + +DEFAULT_BASE_URL = "https://api.tenki.cloud" +SESSION_TAG = "agentenv" +_SECTION = "[sandbox.providers.tenki.config]" +_STRING_KEYS = ("api_key", "base_url", "image", "snapshot_id") +_CONFIG_KEYS = frozenset({*_STRING_KEYS, "min_disk_gb"}) + +# Tenki's API bounds; a workspace's own limit may be lower and is enforced at create. +# Tenki restores VMs of at least this size from a warm snapshot; a smaller one cold boots, which is slower. +MIN_CPUS = 2 +MIN_MEMORY_MB = 4096 +MAX_CPUS = 128 +MAX_MEMORY_MB = 512 * 1024 +MIN_DISK_GB = 5 +MAX_DISK_GB = 100 +# Docker keeps both the compressed and the extracted layers, so 10 GB runs out once a gateway and a few MCP server +# images are on the VM. +DEFAULT_MIN_DISK_GB = 20 +MAX_ALLOW_DOMAINS = 256 +MAX_ALLOW_CIDRS = 64 + +READY_TIMEOUT_SEC = 300 +# Clock skew and rounding between this machine and Tenki, below which a shorter lifetime isn't a cap. +_LIFETIME_SLACK_SEC = 60 +_REAP_ATTEMPTS = 3 +_reapers: set[asyncio.Task] = set() + +# A workspace can serve its ports under its own domain, two labels under the regional suffix +# (..us.sb.tenki.sh). The parents seen on exposed ports are added to later allowlists, so a +# restricted agent reaches environments served there. +_preview_parents: set[str] = set() +# A process sees a handful of workspaces; the cap keeps the learned entries well inside MAX_ALLOW_DOMAINS. +_MAX_PREVIEW_PARENTS = 16 + + +def vm_shape(cpu: float, memory_mb: int, disk_size_gb: float, *, min_disk_gb: int) -> tuple[int, int, int]: + """``(cpu_cores, memory_mb, disk_size_gb)``: whole cores and memory rounded up to Tenki's 2 MiB alignment, raised to + the warm-start floor, and a disk of at least ``min_disk_gb``. A request beyond Tenki's bounds is refused, not + shrunk.""" + cores = max(MIN_CPUS, math.ceil(cpu)) + memory = max(MIN_MEMORY_MB, math.ceil(memory_mb / 2) * 2) + disk = max(MIN_DISK_GB, min_disk_gb, math.ceil(disk_size_gb)) + if cores > MAX_CPUS or memory > MAX_MEMORY_MB or disk > MAX_DISK_GB: + raise ValueError( + f"Tenki cannot size a VM for cpu={cpu}, memory={memory_mb}MB, disk={disk_size_gb}GB " + f"(at most {MAX_CPUS} vCPU, {MAX_MEMORY_MB} MB memory, {MAX_DISK_GB} GB disk)" + ) + return cores, memory, disk + + +def _split_entries(policy: NetworkPolicy) -> tuple[list[str], list[str]]: + """Tenki's ``allow_domains`` and ``allow_cidrs``, CIDRs normalized; an IP address listed as a host becomes a /32. + A malformed CIDR raises ValueError.""" + domains: list[str] = [] + cidrs = [str(ipaddress.ip_network(cidr, strict=False)) for cidr in policy.allow_cidrs] + for host in policy.allow_hosts: + try: + cidrs.append(str(ipaddress.ip_network(host, strict=False))) + except ValueError: + domains.append(host) + return domains, cidrs + + +def egress_kwargs(policy: NetworkPolicy) -> dict[str, Any]: + """``AsyncClient.create`` egress arguments for ``policy``. Tenki reads empty allowlists as unrestricted, so an + allowlist with no entries turns outbound off instead.""" + if not policy.restricts_egress: + return {} + domains, cidrs = _split_entries(policy) + if not domains and not cidrs: + return {"allow_outbound": False} + return {"allow_domains": domains or None, "allow_cidrs": cidrs or None} + + +def session_name(attribution: Attribution) -> str: + """``agentenv-`` and the attribution values, slugged into Tenki's 63-character name.""" + slugs = [re.sub(r"[^a-z0-9]+", "-", str(attribution[key]).lower()).strip("-") for key in sorted(attribution)] + return "-".join(["agentenv", uuid.uuid4().hex[:8], *filter(None, slugs)])[:63].rstrip("-") + + +def _tunnel_url(url: str) -> str: + url = url if url.startswith(("http://", "https://")) else f"https://{url}" + parent = (urlparse(url).hostname or "").partition(".")[2] + entry = f"*.{parent}" + if (parent.endswith(".sb.tenki.sh") and entry not in TenkiSandboxProvider.EGRESS_HOSTS + and len(_preview_parents) < _MAX_PREVIEW_PARENTS): + _preview_parents.add(entry) + return url + + +def _granted_lifetime(timeout_at: datetime | None, requested_at: datetime) -> float: + """The lifetime Tenki gave a VM requested at ``requested_at``, boot included; infinite when it reports none.""" + if timeout_at is None: + return math.inf + if timeout_at.tzinfo is None: + timeout_at = timeout_at.replace(tzinfo=UTC) + return (timeout_at - requested_at).total_seconds() + + +def _remaining_lifetime(timeout_at: datetime | None) -> int: + if timeout_at is None: + return 2 * 3600 + if timeout_at.tzinfo is None: + timeout_at = timeout_at.replace(tzinfo=UTC) + return max(60, math.ceil((timeout_at - datetime.now(UTC)).total_seconds())) + + +async def _reap(sandbox: TenkiSandbox) -> None: + """Terminate ``sandbox``, retrying. A cancellation, as ``asyncio.run`` sends every task left when it ends, is + honored only after the terminate settles, so Ctrl-C during a create doesn't leave its VM running.""" + cancelled = False + for attempt in range(_REAP_ATTEMPTS): + try: + await sandbox.terminate() + logger.info("Terminated Tenki VM %s after its create failed or was cancelled", sandbox.sandbox_id) + break + except asyncio.CancelledError: + cancelled = True + except Exception as exc: # noqa: BLE001 - retried, then logged + logger.warning("Terminating Tenki VM %s failed (attempt %s): %s", sandbox.sandbox_id, attempt + 1, exc) + if attempt + 1 < _REAP_ATTEMPTS and not cancelled: + try: + await asyncio.sleep(2**attempt) + except asyncio.CancelledError: + cancelled = True + else: + logger.error("Tenki VM %s from a failed create is still running; it stops at its max lifetime", + sandbox.sandbox_id) + with contextlib.suppress(Exception): + await sandbox.release() + if cancelled: + raise asyncio.CancelledError + + +async def _reclaim_by_name(client: Any, name: str, policy: NetworkPolicy, timeout: int) -> None: + """Terminate the VM a failed create may have made anyway, its response lost: the active session tagged + ``agentenv`` with the create's unique name.""" + try: + sessions = await client.list(tags=[SESSION_TAG]) + except Exception as exc: # noqa: BLE001 - best effort; the VM still stops at its max lifetime + logger.warning("Could not look for a Tenki VM %s left by a failed create: %s", name, exc) + return + for session in sessions: + if session.info.name == name and session.info.state not in ENDED_STATES: + await _reap(TenkiSandbox(session, tunnel_urls={}, network_policy=policy, command_timeout=timeout)) + + +def _spawn(cleanup: Any) -> asyncio.Task: + """Run ``cleanup`` as a task no caller's cancellation reaches, kept referenced until it ends.""" + task = asyncio.ensure_future(cleanup) + _reapers.add(task) + task.add_done_callback(_reapers.discard) + return task + + +class TenkiSandboxProvider(SandboxProvider): + """Provision Docker-capable Tenki VMs. + + ``api_key`` comes from resolved provider config, never from the process environment: configure it with an + ``env:`` or ``secret:`` reference in ``[sandbox.providers.tenki.config]``. Workloads never see it. + """ + + # Exposed ports get public URLs under these domains. Tenki wildcards cover one label, hence the regional entry. + EGRESS_HOSTS: ClassVar[tuple[str, ...]] = ("*.sb.tenki.sh", "*.us.sb.tenki.sh") + + def __init__( + self, + *, + api_key: str, + base_url: str = DEFAULT_BASE_URL, + image: str | None = None, + snapshot_id: str | None = None, + min_disk_gb: int = DEFAULT_MIN_DISK_GB, + client_cls: Any = AsyncClient, + ): + if not isinstance(api_key, str) or not api_key.strip(): + raise ValueError("Tenki api_key must be a non-empty resolved secret") + if image and snapshot_id: + raise ValueError("Pass at most one of image and snapshot_id") + self._api_key = api_key + self._base_url = base_url + self._image = image + self._snapshot_id = snapshot_id + self._min_disk_gb = min_disk_gb + # Injectable so unit tests never build an SDK client. + self._client_cls = client_cls + # One client per event loop, shared by every sandbox this provider creates or reconnects to: an SDK client + # holds a gRPC channel bound to the loop it was made in, and the CLI runs several loops in one process. + self._clients: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, Any] = weakref.WeakKeyDictionary() + + @classmethod + def from_config(cls, **config: Any) -> Self: + unknown = set(config) - _CONFIG_KEYS + if unknown: + raise ConfigError(f"{_SECTION} has unknown key(s): {sorted(unknown)}") + if not config.get("api_key"): + raise ConfigError(f"{_SECTION} requires a non-empty 'api_key'") + for key in _STRING_KEYS: + value = config.get(key) + if value is not None and (not isinstance(value, str) or not value.strip()): + raise ConfigError(f"{_SECTION} '{key}' must be a non-empty string") + if config.get("image") and config.get("snapshot_id"): + raise ConfigError(f"{_SECTION} takes at most one of 'image' and 'snapshot_id'") + min_disk_gb = config.get("min_disk_gb", DEFAULT_MIN_DISK_GB) + if isinstance(min_disk_gb, bool) or not isinstance(min_disk_gb, int) or not ( + MIN_DISK_GB <= min_disk_gb <= MAX_DISK_GB + ): + raise ConfigError(f"{_SECTION} 'min_disk_gb' must be an integer from {MIN_DISK_GB} to {MAX_DISK_GB}") + return cls(**config) + + @classmethod + def effective_network_policy(cls, policy: NetworkPolicy | None) -> NetworkPolicy: + """Every platform's hosts, as on the base, and the workspace preview domains seen so far.""" + return super().effective_network_policy(policy).with_hosts(sorted(_preview_parents)) + + @classmethod + def supports_network_policy(cls, policy: NetworkPolicy) -> bool: + """Allow-all, or an allowlist of hostnames and IPv4 addresses or CIDRs within Tenki's limits.""" + if not policy.restricts_egress: + return True + try: + domains, cidrs = _split_entries(policy) + except ValueError: + return False + return ( + len(domains) <= MAX_ALLOW_DOMAINS + and len(cidrs) <= MAX_ALLOW_CIDRS + and all(ipaddress.ip_network(cidr, strict=False).version == 4 for cidr in cidrs) + ) + + def _client(self) -> Any: + loop = asyncio.get_running_loop() + client = self._clients.get(loop) + if client is None: + client = self._clients[loop] = self._client_cls(api_key=self._api_key, base_url=self._base_url) + return client + + async def create_vm( + self, + *, + image: str | None = None, + boot_mode: str | None = None, + cpu: float = 1.0, + memory: int = 8192, + disk_size_gb: float = 10, + timeout: int = 3600 * 2, + exposed_ports: list[int] | None = None, + setup_for_gateway: bool = True, + attribution: Attribution | None = None, + network_policy: NetworkPolicy | None = None, + ) -> TenkiSandbox: + """Create a Tenki VM and expose ``exposed_ports`` publicly. ``timeout`` is the VM's hard maximum lifetime. The + VM boots the configured image or snapshot, so a per-call ``image`` is refused rather than ignored; + ``boot_mode`` has no Tenki equivalent and is ignored.""" + if image is not None: + raise ValueError( + "Tenki VMs boot the image or snapshot_id in [sandbox.providers.tenki.config]; " + "a per-call image is unsupported" + ) + if boot_mode is not None: + logger.warning("Ignoring boot_mode=%s: Tenki VMs have no boot modes", boot_mode) + cores, memory_mb, disk_gb = vm_shape(cpu, memory, disk_size_gb, min_disk_gb=self._min_disk_gb) + effective_policy = self.effective_network_policy(network_policy) + if not self.supports_network_policy(effective_policy): + raise NetworkPolicyUnsupportedError( + f"Tenki enforces allow-all or an allowlist of up to {MAX_ALLOW_DOMAINS} hostnames and " + f"{MAX_ALLOW_CIDRS} IPv4 CIDRs, not {effective_policy.to_dict()}" + ) + metadata = { + key: str(value) + for key, value in apply_default_attribution(dict(attribution or {})).items() + if value is not None + } + ports = list(dict.fromkeys(exposed_ports or [])) + create_kwargs: dict[str, Any] = { + **egress_kwargs(effective_policy), + "name": session_name(metadata), + "cpu_cores": cores, + "memory_mb": memory_mb, + "disk_size_gb": disk_gb, + "max_duration": timeout, + "metadata": metadata, + "tags": [SESSION_TAG], + "wait": False, + } + if self._image: + create_kwargs["image"] = self._image + if self._snapshot_id: + create_kwargs["snapshot_id"] = self._snapshot_id + + client = self._client() + requested_at = datetime.now(UTC) + session = await self._create_or_reclaim(client, create_kwargs, effective_policy, timeout) + sandbox = TenkiSandbox(session, tunnel_urls={}, network_policy=effective_policy, command_timeout=timeout) + try: + await session.wait_ready(timeout=READY_TIMEOUT_SEC) + if (granted := _granted_lifetime(session.info.timeout_at, requested_at)) < timeout - _LIFETIME_SLACK_SEC: + logger.warning( + "Tenki capped VM %s's lifetime to %ss of the %ss requested: the workspace's maximum session " + "duration is lower", session.id, granted, timeout, + ) + sandbox.command_timeout = _remaining_lifetime(session.info.timeout_at) + sandbox.tunnel_urls = await self._expose(session, ports) + logger.info( + "Tenki VM started: session_id=%s cpu=%s memory=%sMB disk=%sGB ports=%s attribution=%s", + session.id, cores, memory_mb, disk_gb, ports, metadata, + ) + if setup_for_gateway: + await sandbox.setup_vm_for_gateway(ports) + return sandbox + except BaseException as error: + # The VM is terminated, with retries, whatever happens to this task; a failure waits for it and a + # cancellation leaves it running. The create's own error is the one raised. + reaper = _spawn(_reap(sandbox)) + if not isinstance(error, asyncio.CancelledError): + # A cancellation while waiting propagates; the reaper keeps running on its own. + await asyncio.shield(reaper) + raise + + @staticmethod + async def _create_or_reclaim( + client: Any, create_kwargs: dict[str, Any], policy: NetworkPolicy, timeout: int + ) -> Any: + """Await the create. A VM that would otherwise run with no handle until its max lifetime is terminated: the + one a create yields after its caller was cancelled, and the one a create that failed may have made anyway, + its response lost, which is found by its unique name.""" + task = asyncio.ensure_future(client.create(**create_kwargs)) + + def terminate_orphan(done: asyncio.Future) -> None: + # A create cancelled itself, as the loop's shutdown cancels it, may have made its VM all the same. + if done.cancelled() or done.exception() is not None: + _spawn(_reclaim_by_name(client, create_kwargs["name"], policy, timeout)) + else: + _spawn(_reap(TenkiSandbox(done.result(), tunnel_urls={}, network_policy=policy, + command_timeout=timeout))) + + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + task.add_done_callback(terminate_orphan) + raise + except Exception: + await asyncio.shield(_spawn(_reclaim_by_name(client, create_kwargs["name"], policy, timeout))) + raise + + @staticmethod + async def _expose(session: Any, ports: list[int]) -> dict[int, str]: + exposed = await asyncio.gather(*(session.expose_port(port) for port in ports)) + return {port: _tunnel_url(item.url) for port, item in zip(ports, exposed, strict=True)} + + async def create_sandbox( + self, + *, + image_name: str, + port: int, + env: dict[str, str], + cpu: float = 1.0, + memory: int = 8192, + disk_size_gb: float = 10, + timeout: int = 3600 * 2, + attribution: Attribution | None = None, + network_policy: NetworkPolicy | None = None, + ) -> TenkiSandbox: + # A VM backend: the gateway and agent paths start image_name inside the VM afterwards. + del image_name, env + return await self.create_vm( + cpu=cpu, + memory=memory, + disk_size_gb=disk_size_gb, + timeout=timeout, + exposed_ports=[port], + attribution=attribution, + network_policy=network_policy, + ) + + async def get_sandbox(self, sandbox_id: str) -> TenkiSandbox: + """Reconnect, restoring the exposed ports and the applied egress policy.""" + session = await self._client().get(sandbox_id) + exposed = await session.list_exposed_ports() + info = session.info + return TenkiSandbox( + session, + tunnel_urls={item.port: _tunnel_url(item.url) for item in exposed}, + network_policy=policy_from_info(info), + command_timeout=_remaining_lifetime(info.timeout_at), + ) diff --git a/src/agent_env/providers/sandbox_providers/tenki/sandbox.py b/src/agent_env/providers/sandbox_providers/tenki/sandbox.py new file mode 100644 index 00000000..841997b0 --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/tenki/sandbox.py @@ -0,0 +1,331 @@ +"""A Tenki Sandbox VM presented through the :class:`VmSandbox` contract. + +Tenki VMs boot Tenki images, not OCI images, so the gateway's docker-compose, an agent's ``docker run`` and image +loading run under Docker inside the VM. Every command runs as root through Tenki's privileged exec. Tenki applies the +egress policy outside the VM, so one policy covers ``dockerd`` and every container it runs. +""" + +from __future__ import annotations + +import asyncio +import ipaddress +import logging +import shlex +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse + +from tenki import InvalidStateError, SandboxError, SessionNotFoundError, SessionTerminatedError + +from agent_env.config import get_config +from agent_env.providers.sandbox_providers.sandbox import ( + CURL_RETRY_FLAGS, + NetworkMode, + NetworkPolicy, + VmSandbox, + push_object_over_stdin, +) +from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_VM + +if TYPE_CHECKING: + from agent_env.store.object_store import ObjectStore + +logger = logging.getLogger(__name__) + +# Installs Docker and Compose v2 only when the image lacks them, then starts the daemon. Every fd of the fallback +# dockerd is redirected, since a child holding the exec's output stream keeps the exec from returning. +DOCKER_BOOTSTRAP = r""" +set -e +sysctl -qw vm.overcommit_memory=1 || true +if ! command -v docker >/dev/null 2>&1 || ! docker compose version >/dev/null 2>&1; then + export DEBIAN_FRONTEND=noninteractive + apt-get -o DPkg::Lock::Timeout=180 update -qq + apt-get -o DPkg::Lock::Timeout=180 install -y -qq --no-install-recommends \ + docker.io docker-compose-v2 docker-buildx curl ca-certificates >/dev/null +fi +if ! iptables -t nat -L >/dev/null 2>&1; then + echo "this VM's kernel has no netfilter, which Docker networking needs; boot from a snapshot taken on a VM" \ + "where Docker runs" >&2 + exit 1 +fi +if ! docker info >/dev/null 2>&1 && ! systemctl start docker >/dev/null 2>&1; then + setsid nohup dockerd >/var/log/dockerd.log 2>&1 /dev/null 2>&1 && exit 0 + sleep 1 +done +echo "docker daemon did not become ready" >&2 +tail -n 40 /var/log/dockerd.log >&2 2>/dev/null || true +exit 1 +""" + +BOOTSTRAP_TIMEOUT_SEC = 900 +ENDED_STATES = frozenset({"TERMINATING", "TERMINATED", "GONE"}) +# What exec returns when the stream drops without an exit status; exec_script retries only this. +_UNKNOWN_EXIT_CODE = -1 +_TIMED_OUT_EXIT_CODE = 124 +_STDIN_QUEUE_LIMIT = 16 +# How the SDK reports a run stream that ended before the command did. +_STREAM_CLOSED = "stream closed before exit" +# Tenki reports a command killed by a signal, or one that never started, as exit -1; agent-env reads -1 as a lost +# stream, so these get the shell's codes instead: 128 + the signal's number, or 127. +_SIGNAL_NUMBERS = {"hangup": 1, "interrupt": 2, "quit": 3, "aborted": 6, "killed": 9, "segmentation fault": 11, + "broken pipe": 13, "terminated": 15} +_NOT_STARTED_EXIT_CODE = 127 +_REAP_TIMEOUT_SEC = 10 + + +def allows_host(policy: NetworkPolicy, host: str) -> bool: + """Whether ``policy`` lets the VM reach ``host`` by Tenki's rules: an exact name, or a ``*.domain`` entry, which + covers exactly one label under ``domain``.""" + if not policy.restricts_egress: + return True + host = host.lower().rstrip(".") + for entry in policy.allow_hosts: + entry = entry.lower().rstrip(".") + if entry == host: + return True + if entry.startswith("*."): + suffix = entry[1:] + label = host.removesuffix(suffix) + if host.endswith(suffix) and label and "." not in label: + return True + try: + address = ipaddress.ip_address(host) + except ValueError: + return False + return any(address in ipaddress.ip_network(cidr, strict=False) for cidr in policy.allow_cidrs) + + +def policy_from_info(info: Any) -> NetworkPolicy: + """The policy a Tenki session applies. Outbound off is an allowlist with no entries; empty allowlists with outbound + on leave egress open.""" + if not info.outbound_enabled: + return NetworkPolicy(mode=NetworkMode.ALLOWLIST) + egress = info.egress + if not egress.allow_domains and not egress.allow_cidrs: + return NetworkPolicy() + return NetworkPolicy( + mode=NetworkMode.ALLOWLIST, allow_hosts=tuple(egress.allow_domains), allow_cidrs=tuple(egress.allow_cidrs), + ) + + +def _is_transient(exc: BaseException) -> bool: + """A dropped stream or an unavailable data plane, which a rerun can get past; a terminated session is not.""" + if isinstance(exc, SessionTerminatedError): + return False + return bool(getattr(exc, "retryable", False)) or _STREAM_CLOSED in str(exc) + + +def _exit_status(result: Any) -> tuple[int, bytes]: + """``result``'s exit code and a note for its stderr, with Tenki's -1 for a signal or a failed start mapped.""" + if result.exit_code != _UNKNOWN_EXIT_CODE: + return result.exit_code, b"" + if result.signal: + return 128 + _SIGNAL_NUMBERS.get(str(result.signal).lower(), 0), f"\nkilled by signal: {result.signal}".encode() + return _NOT_STARTED_EXIT_CODE, f"\ncommand did not run: {result.reason or 'unknown reason'}".encode() + + +class _BytesReader: + """The small ``StreamReader`` subset consumed by ``Sandbox.exec_with_output``.""" + + def __init__(self, value: bytes): + self._value = value + + async def read(self) -> bytes: + return self._value + + +class _CompletedProcess: + """A finished Tenki exec in the common process protocol.""" + + def __init__(self, stdout: bytes, stderr: bytes, exit_code: int): + self.stdout = _BytesReader(stdout) + self.stderr = _BytesReader(stderr) + self._exit_code = exit_code + + async def wait(self) -> int: + return self._exit_code + + +class TenkiSandbox(VmSandbox): + """A Tenki VM with Docker; its client is its provider's, shared with the provider's other sandboxes. + ``command_timeout`` is every command's execution budget, normally the VM's remaining lifetime: Tenki bounds a + command that declares none to one hour.""" + + type = "tenki" + _DOCKER_PROBE_TIMEOUT = 15 + + def __init__( + self, + session: Any, + *, + tunnel_urls: dict[int, str], + network_policy: NetworkPolicy | None, + command_timeout: int, + ): + self._session = session + self.command_timeout = command_timeout + self.sandbox_id = session.id + self.mode = SANDBOX_MODE_VM + self.vnc_url = None + self.tunnel_urls = tunnel_urls + self.network_policy = network_policy + + async def terminate(self) -> None: + """Terminate the VM. One another handle already terminated, while this handle's cached state still says + running, is not an error.""" + try: + await self._session.close_if_open() + except SessionNotFoundError: + logger.info("Tenki session %s was already gone at terminate", self.sandbox_id) + except (InvalidStateError, SessionTerminatedError): + if await self._current_state() not in ENDED_STATES: + raise + logger.info("Tenki session %s was already terminating at terminate", self.sandbox_id) + await self.release() + + async def release(self) -> None: + """Stop the session's data-plane credential renewal, which close_if_open leaves to the caller for a session + that is already ending.""" + reset = getattr(self._session, "_reset_data_plane", None) + if reset is not None: + await reset() + + async def _current_state(self) -> str: + try: + return (await self._session.refresh()).state + except SessionNotFoundError: + return "GONE" + + async def exec(self, *command: str) -> _CompletedProcess: + """Run argv to completion as root, so a leading ``sudo`` is dropped. A lost stream is exit -1, which + ``exec_script`` retries.""" + return await self._run(*command, timeout=self.command_timeout) + + async def exec_with_output(self, *args: str) -> tuple[int, str, str]: + """An exec whose stream was cut (exit -1) returns no stdout, so a caller that skips the exit code can't mistake + partial output for the whole.""" + exit_code, stdout, stderr = await self._exec_with_output_bounded(*args, timeout=self.command_timeout) + return exit_code, "" if exit_code == _UNKNOWN_EXIT_CODE else stdout, stderr + + async def _run(self, *command: str, timeout: int) -> _CompletedProcess: + argv = list(command[1:] if command[:1] == ("sudo",) else command) + try: + result = await self._session.exec(*argv, timeout=timeout, privileged=True) + except SandboxError as exc: + if not _is_transient(exc): + raise + return _CompletedProcess(b"", f"{type(exc).__name__}: {exc}".encode(), _UNKNOWN_EXIT_CODE) + if result.timed_out: + note = f"\ncommand timed out after {timeout}s".encode() + return _CompletedProcess(result.stdout, result.stderr + note, _TIMED_OUT_EXIT_CODE) + exit_code, note = _exit_status(result) + return _CompletedProcess(result.stdout, result.stderr + note, exit_code) + + async def _exec_with_output_bounded(self, *command: str, timeout: int) -> tuple[int, str, str]: + process = await self._run(*command, timeout=timeout) + stdout, stderr = await asyncio.gather(process.stdout.read(), process.stderr.read()) + return await process.wait(), stdout.decode(errors="replace"), stderr.decode(errors="replace") + + async def _exec_with_stdin(self, script: str, stdin: AsyncIterator[bytes]) -> tuple[int, str, str]: + """Run ``script`` with bash as root, its standard input the bytes ``stdin`` yields.""" + process = await self._session.start("bash", "-c", script, timeout=self.command_timeout, privileged=True) + try: + async for piece in stdin: + await process.write_stdin(piece) + if await self._wait_for_stdin_drain(process): + break + await process.close_stdin() + result = await process.wait() + except BaseException: + await self._reap(process) + raise + exit_code, note = _exit_status(result) + return exit_code, result.stdout_text, result.stderr_text + note.decode() + + @staticmethod + async def _reap(process: Any) -> None: + """Close stdin, kill and wait for a process whose input failed, so the remote command and the SDK stream end + now rather than at the command's timeout. Its own failures are logged, never raised over the original one; a + cancellation is honored once every step has run.""" + cancelled = False + for step in (process.close_stdin, process.kill, lambda: process.wait(timeout=_REAP_TIMEOUT_SEC)): + try: + await step() + except asyncio.CancelledError: + cancelled = True + except Exception as exc: # noqa: BLE001 - the input's failure is the one to report + logger.debug("Reaping a failed stdin exec: %s: %s", type(exc).__name__, exc) + if cancelled: + raise asyncio.CancelledError + + @staticmethod + async def _wait_for_stdin_drain(process: Any) -> bool: + """Stdin is queued as it is written, so a fast reader would buffer a whole object in memory; wait for the + stream to catch up. True once the command has ended, so the caller stops reading input nobody will consume; + a stream that ended in error raises it.""" + queue = getattr(process, "_outbound", None) + done = getattr(process, "_done", None) + if not isinstance(queue, asyncio.Queue) or not isinstance(done, asyncio.Event): + return False + while not done.is_set() and queue.qsize() > _STDIN_QUEUE_LIMIT: + await asyncio.sleep(0.005) + if done.is_set() and (error := getattr(process, "_error", None)) is not None: + raise error + return done.is_set() + + async def wait_for_vm(self) -> None: + """Install Docker and Compose v2 if the image lacks them, start the daemon and wait for it.""" + exit_code, stdout, stderr = await self._exec_with_output_bounded( + "bash", "-c", DOCKER_BOOTSTRAP, timeout=BOOTSTRAP_TIMEOUT_SEC, + ) + if exit_code != 0: + raise RuntimeError(f"Docker not ready in Tenki VM {self.sandbox_id}: {(stderr or stdout).strip()[-1500:]}") + logger.info("Docker ready in Tenki VM %s", self.sandbox_id) + + async def setup_vm_for_gateway(self, exposed_ports: list[int] | None = None) -> None: + """Ready Docker and check for Compose v2, which ``deploy_env`` needs; Tenki routes exposed ports, so no + firewall.""" + await self.wait_for_vm() + exit_code, stdout, stderr = await self._exec_with_output_bounded( + "docker", "compose", "version", timeout=self._DOCKER_PROBE_TIMEOUT, + ) + if exit_code != 0: + raise RuntimeError(f"Tenki VM {self.sandbox_id} has no Docker Compose v2: {(stderr or stdout).strip()}") + + # Bytes go over exec stdin as root, which reaches any path on the VM and needs no egress, so they arrive whatever + # the VM's egress policy, which is set when the VM is created. + + async def _write_bytes_to_vm_path(self, data: bytes, vm_path: str) -> None: + async def pieces() -> AsyncIterator[bytes]: + yield data + + exit_code, stdout, stderr = await self._exec_with_stdin(f"cat > {shlex.quote(vm_path)}", pieces()) + if exit_code != 0: + raise RuntimeError(f"Writing {vm_path} failed (exit {exit_code}): {(stderr or stdout)[-1500:]}") + + async def _write_unsigned_object(self, object_store: ObjectStore, object_url: str, vm_path: str) -> None: + await push_object_over_stdin(self, object_store, object_url, vm_path) + + def _can_download(self, url: str | None) -> bool: + """Whether the VM can fetch ``url`` itself; an unknown policy counts as no.""" + if url is None or self.network_policy is None: + return False + host = urlparse(url).hostname + return host is not None and allows_host(self.network_policy, host) + + async def _download_object_to_vm(self, object_url: str, vm_path: str) -> None: + object_store = get_config().get_object_store_at(object_url) + signed = await asyncio.to_thread(object_store.signed_get_url, object_url) + if self._can_download(signed): + await self.exec_script(f"curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(signed)} -o {shlex.quote(vm_path)}") + return + await push_object_over_stdin(self, object_store, object_url, vm_path) + + async def _load_tarballs(self, artifacts: list) -> None: + """Load image tarballs, pushing over stdin those the VM's egress policy can't download.""" + signed_urls = await self._signed_image_urls(artifacts) + reachable = [url if self._can_download(url) else None for url in signed_urls] + await self._load_docker_images(artifacts, reachable) diff --git a/tst/integration/env/gateway/gateway_test.py b/tst/integration/env/gateway/gateway_test.py index a1763333..4b326a12 100644 --- a/tst/integration/env/gateway/gateway_test.py +++ b/tst/integration/env/gateway/gateway_test.py @@ -454,6 +454,11 @@ def agentenv_website_env() -> WebsiteEnv: id="tensorlake", marks=skip_without_remote_sandbox("tensorlake"), ), + pytest.param( + "tenki", + id="tenki", + marks=skip_without_remote_sandbox("tenki"), + ), ]) def sandbox_provider(request): from agent_env.providers import ( @@ -478,6 +483,9 @@ def sandbox_provider(request): elif request.param == "tensorlake": set_sandbox_provider(build_sandbox_provider("tensorlake")) set_env_sandbox_provider(build_sandbox_provider("tensorlake")) + elif request.param == "tenki": + set_sandbox_provider(build_sandbox_provider("tenki")) + set_env_sandbox_provider(build_sandbox_provider("tenki")) try: yield request.param finally: diff --git a/tst/integration/env/task_step/task_steps_test.py b/tst/integration/env/task_step/task_steps_test.py index 03fd2174..e586538c 100644 --- a/tst/integration/env/task_step/task_steps_test.py +++ b/tst/integration/env/task_step/task_steps_test.py @@ -284,6 +284,11 @@ def echo_agent() -> A2AAgent: id="tensorlake", marks=skip_without_remote_sandbox("tensorlake"), ), + pytest.param( + "tenki", + id="tenki", + marks=skip_without_remote_sandbox("tenki"), + ), ]) def sandbox_provider(request): from agent_env.providers import ( @@ -314,6 +319,10 @@ def sandbox_provider(request): set_sandbox_provider(build_sandbox_provider("tensorlake")) set_env_sandbox_provider(build_sandbox_provider("tensorlake")) set_agent_sandbox_provider(build_sandbox_provider("tensorlake")) + elif request.param == "tenki": + set_sandbox_provider(build_sandbox_provider("tenki")) + set_env_sandbox_provider(build_sandbox_provider("tenki")) + set_agent_sandbox_provider(build_sandbox_provider("tenki")) try: yield request.param finally: diff --git a/tst/integration/providers/sandbox_providers/tenki_sandbox_smoke_test.py b/tst/integration/providers/sandbox_providers/tenki_sandbox_smoke_test.py new file mode 100644 index 00000000..731578e1 --- /dev/null +++ b/tst/integration/providers/sandbox_providers/tenki_sandbox_smoke_test.py @@ -0,0 +1,104 @@ +"""Tenki VM smoke test. + +Runs when the resolved config can build the ``tenki`` sandbox provider (the ``tenki`` extra and +``[sandbox.providers.tenki.config]`` with ``api_key = "secret:tenki_api_key"``), the ``remote_sandbox`` +capability, and skips with the declared reason otherwise. The test never reads or prints the API key. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import os + +import httpx +import pytest +import pytest_asyncio + +from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy +from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider +from tst.util.capabilities import skip_without_remote_sandbox + +_PORT = 8080 + +pytestmark = [ + pytest.mark.integration, + pytest.mark.int_test_slow, + pytest.mark.asyncio, + skip_without_remote_sandbox("tenki"), +] + + +async def _wait_for_200(url: str) -> None: + async with httpx.AsyncClient(timeout=10) as client: + for _ in range(60): + try: + if (await client.get(url)).status_code == 200: + return + except httpx.HTTPError: + pass + await asyncio.sleep(2) + pytest.fail(f"{url} never answered 200") + + +@pytest.fixture(scope="module") +def tenki_provider(): + return build_sandbox_provider("tenki") + + +@pytest_asyncio.fixture(scope="module") +async def tenki_sandbox(tenki_provider): + """One Tenki VM serving on ``_PORT`` from a container, always terminated after the module.""" + sandbox = await tenki_provider.create_vm(cpu=1, memory=2048, exposed_ports=[_PORT], timeout=900) + try: + await sandbox.exec_script( + f"docker run -d --name smoke -p {_PORT}:{_PORT} public.ecr.aws/docker/library/python:3.12-slim " + f"python -m http.server {_PORT}" + ) + await _wait_for_200(sandbox.tunnel_urls[_PORT]) + yield sandbox + finally: + await sandbox.terminate() + + +async def test_tenki_vm_docker_ports_upload_and_reconnect(tenki_provider, tenki_sandbox): + """Nested Docker, a public port end to end, an upload over stdin to a root-owned path, and reconnect.""" + assert (tenki_sandbox.type, tenki_sandbox.mode) == ("tenki", "vm") + assert tenki_sandbox.network_policy == NetworkPolicy() + + exit_code, version, stderr = await tenki_sandbox.exec_with_output( + "sudo", "docker", "info", "--format", "{{.ServerVersion}}" + ) + assert exit_code == 0, stderr + assert version.strip() + + payload = os.urandom(1024 * 1024 + 17) + await tenki_sandbox.write_host_file(payload, "/opt/smoke/payload.bin") + digest = await tenki_sandbox.exec_script("sha256sum /opt/smoke/payload.bin") + assert digest.split()[0] == hashlib.sha256(payload).hexdigest() + + reconnected = await build_sandbox_provider("tenki").get_sandbox(tenki_sandbox.sandbox_id) + assert reconnected.tunnel_urls == tenki_sandbox.tunnel_urls + assert reconnected.network_policy == NetworkPolicy() + assert (await reconnected.exec_with_output("true"))[0] == 0 + + +async def test_tenki_allowlisted_vm_reaches_its_environment_and_nothing_else(tenki_provider, tenki_sandbox): + """An allowlisted VM reaches an allowed host and another Tenki VM's public port, and nothing else.""" + policy = NetworkPolicy(mode=NetworkMode.ALLOWLIST, allow_hosts=("example.com",)) + agent = await tenki_provider.create_vm(cpu=1, memory=1024, timeout=600, setup_for_gateway=False, + network_policy=policy) + try: + async def fetch(url: str) -> int: + return (await agent.exec_with_output("curl", "-fsS", "-o", "/dev/null", "--max-time", "15", url))[0] + + assert await fetch("https://example.com") == 0 + assert await fetch(tenki_sandbox.tunnel_urls[_PORT]) == 0 + assert await fetch("https://www.wikipedia.org") != 0 + + reconnected = await build_sandbox_provider("tenki").get_sandbox(agent.sandbox_id) + assert reconnected.network_policy.restricts_egress + assert set(reconnected.network_policy.allow_hosts) == set(agent.network_policy.allow_hosts) + finally: + await agent.terminate() + await reconnected.terminate() diff --git a/tst/unit/providers/sandbox_providers/egress_hosts_test.py b/tst/unit/providers/sandbox_providers/egress_hosts_test.py index 81278575..477d0c58 100644 --- a/tst/unit/providers/sandbox_providers/egress_hosts_test.py +++ b/tst/unit/providers/sandbox_providers/egress_hosts_test.py @@ -47,6 +47,7 @@ def test_union_covers_every_registered_platforms_tunnel_hosts(platform_registry) assert "*.modal.host" in hosts # a Modal EGRESS_HOSTS entry assert "*.e2b.app" in hosts # E2B's - ingress URL assert "*.sandbox.tensorlake.ai" in hosts + assert "*.sb.tenki.sh" in hosts def test_union_holds_only_platform_hosts(): @@ -97,5 +98,6 @@ def test_the_builtin_floor_is_pinned(builtin_registry): Pinned rather than counted: a new entry should be argued for in review. Platforms registered from config add theirs on top.""" assert set(all_sandbox_egress_hosts()) == { - "*.modal.host", "*.w.modal.host", "*.e2b.app", "*.sail.box", "*.sandbox.tensorlake.ai", + "*.modal.host", "*.w.modal.host", "*.e2b.app", "*.sail.box", "*.sandbox.tensorlake.ai", "*.sb.tenki.sh", + "*.us.sb.tenki.sh", } diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py index 8e406696..09fdf07f 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py @@ -28,6 +28,7 @@ from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandbox from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandbox from agent_env.providers.sandbox_providers.sandbox import port_bindings +from agent_env.providers.sandbox_providers.tenki.sandbox import TenkiSandbox from agent_env.providers.sandbox_providers.tensorlake.sandbox import TensorlakeSandbox RESERVED = [AGENT_ENV_GATEWAY_MCP_PORT, DB_WEB_PORT, DB_MCP_PORT] @@ -51,7 +52,7 @@ def test_base_sandbox_host_port_is_identity(): assert Sandbox.host_port(None, 18765) == 18765 -@pytest.mark.parametrize("sandbox_class", [ModalVmSandbox, ModalSandbox, E2BSandbox, TensorlakeSandbox]) +@pytest.mark.parametrize("sandbox_class", [ModalVmSandbox, ModalSandbox, E2BSandbox, TensorlakeSandbox, TenkiSandbox]) def test_remote_sandboxes_publish_on_every_interface(sandbox_class): """Their tunnels reach a published port from outside the VM, so it stays on every interface.""" assert port_bindings(sandbox_class.host_ips, 18765, 18765) == ["18765:18765"] diff --git a/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py b/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py index 467b21f5..963ac2f9 100644 --- a/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py +++ b/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py @@ -21,6 +21,7 @@ from agent_env.providers.sandbox_providers.sail_vm.provider import SailVmSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.sandbox_providers.sandbox_provider import _BUILTIN_SANDBOX_PROVIDERS +from agent_env.providers.sandbox_providers.tenki.provider import TenkiSandboxProvider from agent_env.providers.sandbox_providers.tensorlake.provider import TensorlakeSandboxProvider @@ -98,12 +99,23 @@ async def _tensorlake() -> Sandbox: return await provider.create_vm(exposed_ports=[], setup_for_gateway=False) +async def _tenki() -> Sandbox: + session = MagicMock(id="tenki-test") + session.wait_ready = AsyncMock() + session.info.timeout_at = None + client = MagicMock() + client.create = AsyncMock(return_value=session) + provider = TenkiSandboxProvider(api_key="test-key", client_cls=MagicMock(return_value=client)) + return await provider.create_vm(exposed_ports=[], setup_for_gateway=False) + + FACTORIES = { "modal": _modal, "modal_vm": _modal_vm, "e2b": _e2b, "sail_vm": _sail, "tensorlake": _tensorlake, + "tenki": _tenki, "local": _local, } diff --git a/tst/unit/providers/sandbox_providers/provider_capabilities_test.py b/tst/unit/providers/sandbox_providers/provider_capabilities_test.py index 8faeffdc..524734a1 100644 --- a/tst/unit/providers/sandbox_providers/provider_capabilities_test.py +++ b/tst/unit/providers/sandbox_providers/provider_capabilities_test.py @@ -27,6 +27,8 @@ SandboxProvider, on_this_machine_provider, ) +from agent_env.providers.sandbox_providers.tenki.provider import TenkiSandboxProvider +from agent_env.providers.sandbox_providers.tenki.sandbox import TenkiSandbox from agent_env.providers.sandbox_providers.tensorlake.provider import TensorlakeSandboxProvider from agent_env.providers.sandbox_providers.tensorlake.sandbox import TensorlakeSandbox @@ -35,7 +37,7 @@ LOADABLE, NAME, NAME_OR_CONTEXT = Accepts.LOADABLE, Accepts.NAME, Accepts.NAME_OR_CONTEXT PROVIDER_CLASSES = [LocalSandboxProvider, ModalSandboxProvider, ModalVmSandboxProvider, E2BSandboxProvider, - SailVmSandboxProvider, TensorlakeSandboxProvider, ChainedSandboxProvider] + SailVmSandboxProvider, TensorlakeSandboxProvider, TenkiSandboxProvider, ChainedSandboxProvider] @pytest.mark.parametrize("cls, on_this_machine, creates_vms, sandbox, container, private_network", [ @@ -45,7 +47,8 @@ (E2BSandboxProvider, False, True, LOADABLE, NAME, False), (SailVmSandboxProvider, False, True, LOADABLE, NAME, False), (TensorlakeSandboxProvider, False, True, LOADABLE, NAME, False), -], ids=["local", "modal", "modal_vm", "e2b", "sail_vm", "tensorlake"]) + (TenkiSandboxProvider, False, True, LOADABLE, NAME, False), +], ids=["local", "modal", "modal_vm", "e2b", "sail_vm", "tensorlake", "tenki"]) def test_each_provider_declares_where_it_runs_and_the_images_it_runs(cls, on_this_machine, creates_vms, sandbox, container, private_network): assert ((cls.ON_THIS_MACHINE, SANDBOX_MODE_VM in cls.SUPPORTED_SANDBOX_MODES, cls.PRIVATE_NETWORK) @@ -266,7 +269,8 @@ def test_an_image_that_records_no_platform_isnt_checked(): assert image.platform is None and image.load_problem(on_this_machine=False, platform="linux/arm64") is None -@pytest.mark.parametrize("sandbox", [Sandbox, ModalSandbox, ModalVmSandbox, E2BSandbox, SailVmSandbox, TensorlakeSandbox]) +@pytest.mark.parametrize("sandbox", [Sandbox, ModalSandbox, ModalVmSandbox, E2BSandbox, SailVmSandbox, TensorlakeSandbox, + TenkiSandbox]) def test_a_sandbox_runs_linux_amd64_unless_its_provider_sets_another(sandbox): assert sandbox.platform == "linux/amd64" diff --git a/tst/unit/providers/sandbox_providers/tenki_fakes.py b/tst/unit/providers/sandbox_providers/tenki_fakes.py new file mode 100644 index 00000000..6c96af37 --- /dev/null +++ b/tst/unit/providers/sandbox_providers/tenki_fakes.py @@ -0,0 +1,161 @@ +"""In-memory stand-ins for the Tenki SDK's AsyncClient, AsyncSandbox and AsyncProcess.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + +from tenki_sandbox.models import CommandResult, ExposedPort, SessionEgressPolicy + + +@dataclass +class FakeInfo: + id: str + name: str = "" + state: str = "CREATING" + outbound_enabled: bool = True + egress: SessionEgressPolicy = field(default_factory=SessionEgressPolicy) + timeout_at: datetime | None = None + + +class FakeProcess: + def __init__(self, result: CommandResult): + self._result = result + self.stdin: list[bytes] = [] + self.stdin_closed = False + self.killed = False + self.calls: list[str] = [] + self.wait_error: BaseException | None = None + + async def write_stdin(self, data: bytes) -> None: + self.stdin.append(data) + + async def close_stdin(self) -> None: + self.calls.append("close_stdin") + self.stdin_closed = True + + async def kill(self) -> None: + self.calls.append("kill") + self.killed = True + + async def wait(self, timeout: float | None = None) -> CommandResult: + self.calls.append("wait") + if self.wait_error is not None and len([c for c in self.calls if c == "wait"]) == 1: + raise self.wait_error + return self._result + + +class FakeSession: + """Records every call; ``results`` answers exec and start in order, raising an exception found there.""" + + def __init__(self, session_id: str = "sess-1", **info: Any): + self.id = session_id + self.info = FakeInfo(id=session_id, **info) + self.exec_calls: list[dict[str, Any]] = [] + self.results: list[CommandResult | BaseException] = [] + self.processes: list[FakeProcess] = [] + self.exposed: dict[int, str] = {} + self.closed = 0 + self.close_error: BaseException | None = None + self.ready_delay = 0.0 + self.close_delay = 0.0 + self.preview_parent = "sb.tenki.sh" + + @property + def state(self) -> str: + return self.info.state + + async def wait_ready(self, timeout: float = 180) -> None: + if self.ready_delay: + await asyncio.sleep(self.ready_delay) + self.info.state = "RUNNING" + + def _next(self, argv: tuple[str, ...], timeout: Any, privileged: bool) -> CommandResult: + self.exec_calls.append({"argv": list(argv), "timeout": timeout, "privileged": privileged}) + outcome = self.results.pop(0) if self.results else CommandResult(argv=list(argv), exit_code=0) + if isinstance(outcome, BaseException): + raise outcome + return outcome + + async def exec(self, *argv: str, timeout: Any = None, privileged: bool = False) -> CommandResult: + return self._next(argv, timeout, privileged) + + async def start(self, *argv: str, timeout: Any = None, privileged: bool = False) -> FakeProcess: + process = FakeProcess(self._next(argv, timeout, privileged)) + self.processes.append(process) + return process + + async def expose_port(self, port: int) -> ExposedPort: + url = f"p{port}-{self.id}.{self.preview_parent}" + self.exposed[port] = url + return ExposedPort(port=port, url=url) + + async def list_exposed_ports(self) -> list[ExposedPort]: + return [ExposedPort(port=port, url=f"https://{url}") for port, url in self.exposed.items()] + + async def refresh(self) -> FakeInfo: + return self.info + + async def close_if_open(self) -> None: + if self.state in {"TERMINATING", "TERMINATED"}: + return + if self.close_delay: + await asyncio.sleep(self.close_delay) + if self.close_error is not None: + raise self.close_error + self.closed += 1 + self.info.state = "TERMINATED" + + +class FakeClient: + """Each instance records itself on its class's ``instances``; ``on_init`` stages what a new one returns.""" + + instances: list[FakeClient] + on_init: Any = None + + def __init__(self, **kwargs: Any): + self.kwargs = kwargs + self.create_calls: list[dict[str, Any]] = [] + self.sessions: dict[str, FakeSession] = {} + self.closed = False + self.create_delay = 0.0 + self.create_error: BaseException | None = None + # The server made the VM but the response was lost: create records the session, then raises create_error. + self.create_commits = False + self.next_session: FakeSession | None = None + type(self).instances.append(self) + if type(self).on_init is not None: + type(self).on_init(self) + + async def create(self, **kwargs: Any) -> FakeSession: + self.create_calls.append(kwargs) + if self.create_delay: + await asyncio.sleep(self.create_delay) + session = self.next_session or FakeSession() + session.info.name = kwargs.get("name", "") + if self.create_error is not None and not self.create_commits: + raise self.create_error + self.sessions[session.id] = session + if self.create_error is not None: + raise self.create_error + return session + + async def list(self, tags: list[str] | None = None) -> list[FakeSession]: + return list(self.sessions.values()) + + async def get(self, sandbox_id: str) -> FakeSession: + return self.sessions[sandbox_id] + + async def close(self) -> None: + self.closed = True + + +def fake_client_cls() -> type[FakeClient]: + """A fresh ``FakeClient`` class, so no test sees another's clients.""" + return type("FakeClient", (FakeClient,), {"instances": [], "on_init": None}) + + +def result(stdout: bytes = b"", stderr: bytes = b"", exit_code: int = 0, **extra: Any) -> CommandResult: + return CommandResult(argv=[], exit_code=exit_code, stdout=stdout, stderr=stderr, **extra) diff --git a/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py new file mode 100644 index 00000000..9830f0b5 --- /dev/null +++ b/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py @@ -0,0 +1,385 @@ +"""Offline unit tests for the Tenki sandbox provider.""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta + +import pytest +from tenki_sandbox.models import SessionEgressPolicy + +from agent_env.config.errors import ConfigError +from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy, NetworkPolicyUnsupportedError +from agent_env.providers.sandbox_providers.tenki import provider as provider_module +from agent_env.providers.sandbox_providers.tenki.provider import ( + DEFAULT_BASE_URL, + DEFAULT_MIN_DISK_GB, + SESSION_TAG, + TenkiSandboxProvider, + egress_kwargs, + session_name, + vm_shape, +) +from agent_env.providers.sandbox_providers.tenki.sandbox import DOCKER_BOOTSTRAP, TenkiSandbox +from tenki import SandboxError + +from tst.unit.providers.sandbox_providers.tenki_fakes import FakeSession, fake_client_cls, result + +ALLOWLIST = NetworkMode.ALLOWLIST + + +@pytest.fixture(autouse=True) +def _no_preview_parents(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(provider_module, "_preview_parents", set()) + + +@pytest.fixture +def client_cls(): + return fake_client_cls() + + +def _provider(client_cls, **kwargs) -> TenkiSandboxProvider: + return TenkiSandboxProvider(api_key="tk_test", client_cls=client_cls, **kwargs) + + +def test_vm_shape_rounds_up_to_tenki_units(): + assert vm_shape(2.1, 8193, 42.2, min_disk_gb=DEFAULT_MIN_DISK_GB) == (3, 8194, 43) + assert vm_shape(4, 8192, 10, min_disk_gb=50) == (4, 8192, 50) + + +def test_vm_shape_raises_a_small_request_to_the_warm_start_floor(): + """A VM smaller than 2 vCPU / 4 GiB cold boots instead of restoring Tenki's warm snapshot.""" + assert vm_shape(0.5, 1001, 3, min_disk_gb=DEFAULT_MIN_DISK_GB) == (2, 4096, DEFAULT_MIN_DISK_GB) + assert vm_shape(1, 8192, 10, min_disk_gb=DEFAULT_MIN_DISK_GB) == (2, 8192, DEFAULT_MIN_DISK_GB) + + +@pytest.mark.parametrize(("cpu", "memory", "disk"), [(129, 1024, 10), (1, 524290, 10), (1, 1024, 101)]) +def test_vm_shape_refuses_rather_than_shrinks(cpu, memory, disk): + with pytest.raises(ValueError, match="Tenki cannot size a VM"): + vm_shape(cpu, memory, disk, min_disk_gb=DEFAULT_MIN_DISK_GB) + + +def test_egress_allow_all_sends_nothing(): + assert egress_kwargs(NetworkPolicy()) == {} + + +def test_egress_allowlist_splits_hosts_and_cidrs(): + policy = NetworkPolicy(mode=ALLOWLIST, allow_hosts=("pypi.org", "10.1.2.3"), allow_cidrs=("192.168.0.0/16",)) + assert egress_kwargs(policy) == {"allow_domains": ["pypi.org"], "allow_cidrs": ["192.168.0.0/16", "10.1.2.3/32"]} + + +def test_an_empty_allowlist_turns_outbound_off(): + """Tenki reads empty allowlists as unrestricted, which would open egress.""" + assert egress_kwargs(NetworkPolicy(mode=ALLOWLIST)) == {"allow_outbound": False} + + +def test_supports_allow_all_and_ipv4_allowlists_within_tenki_limits(): + assert TenkiSandboxProvider.supports_network_policy(NetworkPolicy()) + assert TenkiSandboxProvider.supports_network_policy(NetworkPolicy(mode=ALLOWLIST, allow_hosts=("a.com",))) + assert not TenkiSandboxProvider.supports_network_policy( + NetworkPolicy(mode=ALLOWLIST, allow_cidrs=("2001:db8::/32",)) + ) + too_many = tuple(f"h{i}.example.com" for i in range(257)) + assert not TenkiSandboxProvider.supports_network_policy(NetworkPolicy(mode=ALLOWLIST, allow_hosts=too_many)) + + +def test_session_name_fits_tenki_limits(): + name = session_name({"run_id": "Run 42/" + "x" * 80, "team": "Env Pod"}) + assert name.startswith("agentenv-") + assert len(name) <= 63 + assert name == name.lower() and " " not in name + + +def test_from_config_takes_the_known_keys(): + provider = TenkiSandboxProvider.from_config(api_key="tk_x", snapshot_id="snap", min_disk_gb=40) + assert (provider._api_key, provider._snapshot_id, provider._min_disk_gb) == ("tk_x", "snap", 40) + assert provider._base_url == DEFAULT_BASE_URL + + +@pytest.mark.parametrize(("config", "message"), [ + ({"api_key": "k", "region": "us"}, "unknown key"), + ({}, "requires a non-empty 'api_key'"), + ({"api_key": "k", "base_url": " "}, "non-empty string"), + ({"api_key": "k", "image": "a", "snapshot_id": "b"}, "at most one"), + ({"api_key": "k", "min_disk_gb": 2}, "min_disk_gb"), + ({"api_key": "k", "min_disk_gb": True}, "min_disk_gb"), +]) +def test_from_config_refuses_bad_config(config, message): + with pytest.raises(ConfigError, match=message): + TenkiSandboxProvider.from_config(**config) + + +@pytest.mark.asyncio +async def test_create_vm_sizes_tags_and_exposes_ports(client_cls): + sandbox = await _provider(client_cls, snapshot_id="snap").create_vm( + cpu=1.5, memory=4096, disk_size_gb=10, timeout=900, exposed_ports=[8080, 8080, 9000], + setup_for_gateway=False, attribution={"run_id": "r1"}, + ) + client = client_cls.instances[0] + assert client.kwargs == {"api_key": "tk_test", "base_url": DEFAULT_BASE_URL} + call = client.create_calls[0] + assert (call["cpu_cores"], call["memory_mb"], call["disk_size_gb"]) == (2, 4096, DEFAULT_MIN_DISK_GB) + assert call["max_duration"] == 900 + assert call["tags"] == [SESSION_TAG] + assert call["snapshot_id"] == "snap" + assert call["metadata"]["run_id"] == "r1" + assert call["wait"] is False + assert "allow_domains" not in call and "allow_outbound" not in call + assert isinstance(sandbox, TenkiSandbox) + assert (sandbox.type, sandbox.mode) == ("tenki", "vm") + assert sandbox.tunnel_urls == {8080: "https://p8080-sess-1.sb.tenki.sh", 9000: "https://p9000-sess-1.sb.tenki.sh"} + assert sandbox.command_timeout == 900 + + +@pytest.mark.asyncio +async def test_create_vm_applies_an_allowlist_with_the_tenki_hosts(client_cls): + sandbox = await _provider(client_cls).create_vm( + setup_for_gateway=False, network_policy=NetworkPolicy(mode=ALLOWLIST, allow_hosts=("pypi.org",)), + ) + domains = client_cls.instances[0].create_calls[0]["allow_domains"] + assert {"pypi.org", "*.sb.tenki.sh", "*.us.sb.tenki.sh"} <= set(domains) + assert sandbox.network_policy.restricts_egress + + +@pytest.mark.asyncio +async def test_workspace_preview_domains_join_later_allowlists(client_cls): + def wildcard_workspace(client): + client.next_session = FakeSession("env-1") + client.next_session.preview_parent = "acme.us.sb.tenki.sh" + + provider = _provider(client_cls) + client_cls.on_init = wildcard_workspace + env = await provider.create_vm(setup_for_gateway=False, exposed_ports=[8080]) + assert env.tunnel_urls[8080] == "https://p8080-env-1.acme.us.sb.tenki.sh" + + client_cls.instances[0].next_session = None + await provider.create_vm( + setup_for_gateway=False, network_policy=NetworkPolicy(mode=ALLOWLIST, allow_hosts=("a.io",)), + ) + assert "*.acme.us.sb.tenki.sh" in client_cls.instances[0].create_calls[-1]["allow_domains"] + + +@pytest.mark.parametrize(("url", "learned"), [ + ("https://tok.acme.us.sb.tenki.sh", {"*.acme.us.sb.tenki.sh"}), + ("https://tok.sb.tenki.sh", set()), + ("https://tok.us.sb.tenki.sh", set()), + ("https://tok.evil-sb.tenki.sh", set()), + ("https://tok.example.com", set()), +]) +def test_only_tenki_preview_parents_are_learned(url, learned): + provider_module._tunnel_url(url) + assert provider_module._preview_parents == learned + + +def test_learned_preview_parents_are_capped(): + for i in range(provider_module._MAX_PREVIEW_PARENTS + 5): + provider_module._tunnel_url(f"https://tok.ws{i}.us.sb.tenki.sh") + assert len(provider_module._preview_parents) == provider_module._MAX_PREVIEW_PARENTS + + +def test_a_malformed_cidr_is_unsupported_rather_than_an_error(): + policy = NetworkPolicy(mode=ALLOWLIST, allow_cidrs=("10.0.0.300/8",)) + assert not TenkiSandboxProvider.supports_network_policy(policy) + + +def test_cidrs_are_normalized(): + assert egress_kwargs(NetworkPolicy(mode=ALLOWLIST, allow_cidrs=("10.0.0.1/8",)))["allow_cidrs"] == ["10.0.0.0/8"] + + +@pytest.mark.asyncio +async def test_a_lifetime_tenki_capped_becomes_the_command_budget(client_cls, caplog): + session = FakeSession(timeout_at=datetime.now(UTC) + timedelta(hours=2)) + client_cls.on_init = lambda client: setattr(client, "next_session", session) + sandbox = await _provider(client_cls).create_vm(timeout=3 * 3600, setup_for_gateway=False) + assert 7100 <= sandbox.command_timeout <= 7200 + assert "capped" in caplog.text + + +@pytest.mark.asyncio +async def test_a_lifetime_shortened_only_by_the_boot_is_not_a_cap(client_cls, caplog): + session = FakeSession(timeout_at=datetime.now(UTC) + timedelta(seconds=3 * 3600 - 20)) + client_cls.on_init = lambda client: setattr(client, "next_session", session) + sandbox = await _provider(client_cls).create_vm(timeout=3 * 3600, setup_for_gateway=False) + assert sandbox.command_timeout == 3 * 3600 + assert "capped" not in caplog.text + + +@pytest.mark.asyncio +async def test_a_create_cancelled_by_itself_still_reclaims_the_vm_it_made(client_cls): + def cancelled_after_commit(client): + client.create_error = asyncio.CancelledError() + client.create_commits = True + + client_cls.on_init = cancelled_after_commit + with pytest.raises(asyncio.CancelledError): + await _provider(client_cls).create_vm(setup_for_gateway=False) + await asyncio.sleep(0.05) + assert client_cls.instances[0].sessions["sess-1"].closed == 1 + + +@pytest.mark.asyncio +async def test_cancelling_a_create_while_it_cleans_up_raises_the_cancellation(client_cls): + session = FakeSession() + session.results = [result(stderr=b"no space left", exit_code=1)] + session.close_delay = 0.1 + client_cls.on_init = lambda client: setattr(client, "next_session", session) + task = asyncio.ensure_future(_provider(client_cls).create_vm()) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0.15) + assert session.closed == 1 + + +@pytest.mark.asyncio +async def test_a_reaper_that_gives_up_stops_the_sessions_renewal_without_a_last_sleep(monkeypatch): + slept, released = [], [] + + async def no_sleep(seconds): + slept.append(seconds) + + async def reset(): + released.append(True) + + monkeypatch.setattr(provider_module.asyncio, "sleep", no_sleep) + session = FakeSession(state="RUNNING") + session.close_error = SandboxError("unavailable") + session._reset_data_plane = reset + sandbox = TenkiSandbox(session, tunnel_urls={}, network_policy=NetworkPolicy(), command_timeout=60) + await provider_module._reap(sandbox) + assert slept == [1, 2] + assert released == [True] + + +@pytest.mark.asyncio +async def test_every_sandbox_of_a_provider_shares_one_client_per_event_loop(client_cls): + provider = _provider(client_cls) + created = await provider.create_vm(setup_for_gateway=False) + await provider.get_sandbox(created.sandbox_id) + await created.terminate() + assert len(client_cls.instances) == 1 + assert not client_cls.instances[0].closed + + +@pytest.mark.asyncio +async def test_create_vm_refuses_an_unenforceable_policy_before_creating(client_cls): + with pytest.raises(NetworkPolicyUnsupportedError): + await _provider(client_cls).create_vm( + network_policy=NetworkPolicy(mode=ALLOWLIST, allow_cidrs=("2001:db8::/32",)), + ) + assert client_cls.instances == [] + + +@pytest.mark.asyncio +async def test_create_vm_refuses_a_per_call_image(client_cls): + with pytest.raises(ValueError, match="per-call image"): + await _provider(client_cls).create_vm(image="ubuntu:24.04") + assert client_cls.instances == [] + + +@pytest.mark.asyncio +async def test_create_vm_readies_docker_for_the_gateway_as_root(client_cls): + await _provider(client_cls).create_vm(exposed_ports=[8080]) + session = client_cls.instances[0].sessions["sess-1"] + argvs = [call["argv"] for call in session.exec_calls] + assert ["bash", "-c", DOCKER_BOOTSTRAP] in argvs + assert ["docker", "compose", "version"] in argvs + assert all(call["privileged"] for call in session.exec_calls) + + +@pytest.mark.asyncio +async def test_a_setup_failure_terminates_the_vm(client_cls): + session = FakeSession() + session.results = [result(stderr=b"no space left", exit_code=1)] + client_cls.on_init = lambda client: setattr(client, "next_session", session) + with pytest.raises(RuntimeError, match="no space left"): + await _provider(client_cls).create_vm() + assert session.closed == 1 + + +@pytest.mark.asyncio +async def test_a_create_cancelled_while_the_vm_boots_still_terminates_it(client_cls): + session = FakeSession() + session.ready_delay = 0.05 + client_cls.on_init = lambda client: setattr(client, "next_session", session) + task = asyncio.ensure_future(_provider(client_cls).create_vm(setup_for_gateway=False)) + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0.05) + assert session.closed == 1 + + +@pytest.mark.asyncio +async def test_a_cancelled_create_reaps_the_vm(client_cls): + client_cls.on_init = lambda client: setattr(client, "create_delay", 0.05) + task = asyncio.ensure_future(_provider(client_cls).create_vm(setup_for_gateway=False)) + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0.1) + assert client_cls.instances[0].sessions["sess-1"].closed == 1 + + +def _lost_response(client): + client.create_error = SandboxError("connection reset") + client.create_commits = True + + +@pytest.mark.asyncio +async def test_a_create_whose_response_was_lost_terminates_the_vm_it_made(client_cls): + client_cls.on_init = _lost_response + with pytest.raises(SandboxError, match="connection reset"): + await _provider(client_cls).create_vm(setup_for_gateway=False) + assert client_cls.instances[0].sessions["sess-1"].closed == 1 + + +@pytest.mark.asyncio +async def test_a_cancelled_create_whose_response_was_lost_terminates_the_vm_it_made(client_cls): + def slow_lost_response(client): + _lost_response(client) + client.create_delay = 0.05 + + client_cls.on_init = slow_lost_response + task = asyncio.ensure_future(_provider(client_cls).create_vm(setup_for_gateway=False)) + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0.1) + assert client_cls.instances[0].sessions["sess-1"].closed == 1 + + +@pytest.mark.asyncio +async def test_a_create_rejected_outright_leaves_nothing_to_reclaim(client_cls): + client_cls.on_init = lambda client: setattr(client, "create_error", SandboxError("quota exceeded")) + with pytest.raises(SandboxError, match="quota exceeded"): + await _provider(client_cls).create_vm(setup_for_gateway=False) + assert client_cls.instances[0].sessions == {} + + +@pytest.mark.asyncio +async def test_create_sandbox_is_a_vm_with_its_port(client_cls): + sandbox = await _provider(client_cls).create_sandbox(image_name="agent:latest", port=8000, env={"K": "V"}) + assert list(sandbox.tunnel_urls) == [8000] + assert "env" not in client_cls.instances[0].create_calls[0] + + +@pytest.mark.asyncio +async def test_get_sandbox_restores_ports_policy_and_lifetime(client_cls): + session = FakeSession( + "sess-9", state="RUNNING", + egress=SessionEgressPolicy(allow_domains=("pypi.org",), allow_cidrs=("10.0.0.0/8",)), + timeout_at=datetime.now(UTC) + timedelta(minutes=30), + ) + session.exposed = {8080: "p8080.sb.tenki.sh"} + client_cls.on_init = lambda client: client.sessions.update({"sess-9": session}) + sandbox = await _provider(client_cls).get_sandbox("sess-9") + assert sandbox.tunnel_urls == {8080: "https://p8080.sb.tenki.sh"} + assert sandbox.network_policy == NetworkPolicy( + mode=ALLOWLIST, allow_hosts=("pypi.org",), allow_cidrs=("10.0.0.0/8",), + ) + assert 1700 <= sandbox.command_timeout <= 1800 diff --git a/tst/unit/providers/sandbox_providers/tenki_sandbox_test.py b/tst/unit/providers/sandbox_providers/tenki_sandbox_test.py new file mode 100644 index 00000000..67671f1c --- /dev/null +++ b/tst/unit/providers/sandbox_providers/tenki_sandbox_test.py @@ -0,0 +1,366 @@ +"""Offline unit tests for the Tenki sandbox: exec, stdin transfer, lifecycle and egress.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest +from tenki import InvalidStateError, SandboxError, SessionNotFoundError, SessionTerminatedError, StreamClosedError +from tenki_sandbox.models import SessionEgressPolicy + +from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy +from agent_env.providers.sandbox_providers.tenki import sandbox as sandbox_module +from agent_env.providers.sandbox_providers.tenki.sandbox import TenkiSandbox, allows_host, policy_from_info +from tst.unit.providers.sandbox_providers.tenki_fakes import FakeSession, result + +ALLOWLIST = NetworkMode.ALLOWLIST + + +def _sandbox(session: FakeSession | None = None, *, policy=NetworkPolicy()) -> TenkiSandbox: + return TenkiSandbox( + session or FakeSession(state="RUNNING"), tunnel_urls={}, network_policy=policy, + command_timeout=600, + ) + + +def _transient() -> SandboxError: + error = SandboxError("unavailable") + error.retryable = True + return error + + +async def _no_sleep(_seconds): + return None + + +@pytest.mark.asyncio +async def test_exec_runs_as_root_with_the_lifetime_budget_and_drops_sudo(): + session = FakeSession() + session.results = [result(stdout=b"hi\n")] + assert (await _sandbox(session).exec_with_output("sudo", "echo", "hi"))[:2] == (0, "hi\n") + assert session.exec_calls == [{"argv": ["echo", "hi"], "timeout": 600, "privileged": True}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_transient(), SandboxError("sandbox run stream closed before exit")]) +async def test_a_lost_stream_is_exit_minus_one_with_no_stdout(error): + session = FakeSession() + session.results = [error] + exit_code, stdout, stderr = await _sandbox(session).exec_with_output("true") + assert (exit_code, stdout) == (-1, "") + assert type(error).__name__ in stderr + + +@pytest.mark.asyncio +async def test_a_terminated_session_raises_rather_than_reads_as_a_lost_stream(): + session = FakeSession() + session.results = [SessionTerminatedError("session terminated")] + with pytest.raises(SessionTerminatedError): + await _sandbox(session).exec_with_output("true") + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("fields", "exit_code", "note"), [ + ({"signal": "killed", "reason": "signaled"}, 137, "killed by signal: killed"), + ({"reason": 'exec: "nope": executable file not found in $PATH'}, 127, "executable file not found"), +]) +async def test_tenkis_minus_one_for_a_signal_or_a_failed_start_gets_the_shells_code(fields, exit_code, note): + session = FakeSession() + session.results = [result(exit_code=-1, **fields)] + code, _, stderr = await _sandbox(session).exec_with_output("nope") + assert code == exit_code + assert note in stderr + + +@pytest.mark.asyncio +async def test_exec_script_retries_a_lost_stream(monkeypatch): + monkeypatch.setattr("asyncio.sleep", _no_sleep) + session = FakeSession() + session.results = [_transient(), result(stdout=b"done")] + assert await _sandbox(session).exec_script("echo done", max_retries=1) == "done" + assert len(session.exec_calls) == 2 + + +@pytest.mark.asyncio +async def test_other_errors_propagate(): + session = FakeSession() + session.results = [SessionNotFoundError("gone")] + with pytest.raises(SessionNotFoundError): + await _sandbox(session).exec_with_output("true") + + +@pytest.mark.asyncio +async def test_a_timed_out_command_is_exit_124(): + session = FakeSession() + session.results = [result(stdout=b"partial", timed_out=True)] + exit_code, stdout, stderr = await _sandbox(session).exec_with_output("sleep", "999") + assert (exit_code, stdout) == (124, "partial") + assert "timed out after 600s" in stderr + + +@pytest.mark.asyncio +async def test_exec_with_stdin_streams_every_piece_as_root(): + session = FakeSession() + session.results = [result(stdout=b"ok")] + + async def pieces(): + yield b"ab" + yield b"cd" + + assert await _sandbox(session)._exec_with_stdin("cat > /x", pieces()) == (0, "ok", "") + assert (session.processes[0].stdin, session.processes[0].stdin_closed) == ([b"ab", b"cd"], True) + assert session.exec_calls[0] == {"argv": ["bash", "-c", "cat > /x"], "timeout": 600, "privileged": True} + + +@pytest.mark.asyncio +async def test_exec_with_stdin_kills_the_process_when_the_reader_fails(): + session = FakeSession() + + async def pieces(): + yield b"ab" + raise OSError("disk read failed") + + with pytest.raises(OSError): + await _sandbox(session)._exec_with_stdin("cat > /x", pieces()) + assert session.processes[0].calls == ["close_stdin", "kill", "wait"] + + +@pytest.mark.asyncio +async def test_a_wait_that_fails_reaps_the_process(): + session = FakeSession() + process_start = session.start + + async def start(*argv, **kwargs): + process = await process_start(*argv, **kwargs) + process.wait_error = StreamClosedError("stream closed") + return process + + session.start = start + + async def pieces(): + yield b"ab" + + with pytest.raises(StreamClosedError): + await _sandbox(session)._exec_with_stdin("cat > /x", pieces()) + assert session.processes[0].calls == ["close_stdin", "wait", "close_stdin", "kill", "wait"] + + +@pytest.mark.asyncio +async def test_a_reap_that_is_cancelled_still_kills_and_waits(): + calls = [] + + async def cancelled(): + calls.append("close_stdin") + raise asyncio.CancelledError + + async def record(name): + calls.append(name) + + process = SimpleNamespace(close_stdin=cancelled, kill=lambda: record("kill"), + wait=lambda timeout=None: record("wait")) + with pytest.raises(asyncio.CancelledError): + await TenkiSandbox._reap(process) + assert calls == ["close_stdin", "kill", "wait"] + + +@pytest.mark.asyncio +async def test_a_failing_reap_does_not_mask_the_input_error(): + session = FakeSession() + + async def pieces(): + yield b"ab" + raise OSError("disk read failed") + + async def stream_gone(*_args, **_kwargs): + raise StreamClosedError("stream closed") + + process_start = session.start + + async def start(*argv, **kwargs): + process = await process_start(*argv, **kwargs) + process.close_stdin = process.kill = process.wait = stream_gone + return process + + session.start = start + with pytest.raises(OSError, match="disk read failed"): + await _sandbox(session)._exec_with_stdin("cat > /x", pieces()) + + +@pytest.mark.asyncio +async def test_stdin_stops_reading_once_the_command_ended(monkeypatch): + session = FakeSession() + read = [] + + async def pieces(): + for i in range(5): + read.append(i) + yield b"x" + + async def ended(_process): + return True + + monkeypatch.setattr(TenkiSandbox, "_wait_for_stdin_drain", staticmethod(ended)) + await _sandbox(session)._exec_with_stdin("cat > /x", pieces()) + assert read == [0] + + +@pytest.mark.asyncio +async def test_terminate_stops_the_data_plane_renewal_of_a_session_already_ending(): + calls = [] + session = FakeSession(state="TERMINATING") + + async def reset(): + calls.append("reset") + + session._reset_data_plane = reset + await _sandbox(session).terminate() + assert calls == ["reset"] + + +def _stalled_process(*, error: BaseException | None) -> SimpleNamespace: + queue: asyncio.Queue = asyncio.Queue() + for _ in range(sandbox_module._STDIN_QUEUE_LIMIT + 4): + queue.put_nowait(object()) + done = asyncio.Event() + done.set() + return SimpleNamespace(_outbound=queue, _done=done, _error=error) + + +@pytest.mark.asyncio +async def test_the_stdin_drain_raises_the_error_of_a_stream_that_died(): + with pytest.raises(StreamClosedError): + await asyncio.wait_for( + TenkiSandbox._wait_for_stdin_drain(_stalled_process(error=StreamClosedError("stream closed"))), 1, + ) + + +@pytest.mark.asyncio +async def test_the_stdin_drain_returns_once_the_command_ended(): + assert await asyncio.wait_for(TenkiSandbox._wait_for_stdin_drain(_stalled_process(error=None)), 1) + + +@pytest.mark.asyncio +async def test_write_host_file_goes_over_stdin_to_any_path(): + session = FakeSession() + await _sandbox(session).write_host_file(b"bytes", "/etc/agent/config.json") + assert session.exec_calls[-1]["argv"] == ["bash", "-c", "cat > /etc/agent/config.json"] + assert session.processes[-1].stdin == [b"bytes"] + + +@pytest.mark.asyncio +async def test_terminate_closes_the_session_once(): + session = FakeSession(state="RUNNING") + sandbox = _sandbox(session) + await sandbox.terminate() + await sandbox.terminate() + assert session.closed == 1 + + +@pytest.mark.asyncio +async def test_a_failed_terminate_can_be_retried(): + session = FakeSession(state="RUNNING") + session.close_error = SandboxError("unavailable") + sandbox = _sandbox(session) + with pytest.raises(SandboxError): + await sandbox.terminate() + session.close_error = None + await sandbox.terminate() + assert session.closed == 1 + + +@pytest.mark.asyncio +async def test_terminate_tolerates_a_session_that_is_already_gone(): + session = FakeSession(state="RUNNING") + session.close_error = SessionNotFoundError("gone") + await _sandbox(session).terminate() + + +class _TerminatedElsewhere(FakeSession): + """Another handle terminated it: this one's cached state is stale and close is refused.""" + + def __init__(self, server_state: str): + super().__init__(state="RUNNING") + self.close_error = InvalidStateError("cannot terminate session in current state") + self.server_state = server_state + + async def refresh(self): + self.info.state = self.server_state + return self.info + + +@pytest.mark.asyncio +@pytest.mark.parametrize("server_state", ["TERMINATING", "TERMINATED"]) +async def test_terminate_tolerates_a_session_another_handle_terminated(server_state): + await _sandbox(_TerminatedElsewhere(server_state)).terminate() + + +@pytest.mark.asyncio +async def test_terminate_still_raises_when_the_session_is_live(): + with pytest.raises(InvalidStateError): + await _sandbox(_TerminatedElsewhere("PAUSING")).terminate() + + +@pytest.mark.asyncio +async def test_wait_for_vm_reports_the_bootstrap_failure(): + session = FakeSession() + session.results = [result(stderr=b"E: Unable to locate package docker.io", exit_code=100)] + with pytest.raises(RuntimeError, match="Unable to locate package"): + await _sandbox(session).wait_for_vm() + + +@pytest.mark.asyncio +async def test_setup_requires_compose_v2(): + session = FakeSession() + session.results = [result(), result(stderr=b"docker: 'compose' is not a docker command.", exit_code=1)] + with pytest.raises(RuntimeError, match="no Docker Compose v2"): + await _sandbox(session).setup_vm_for_gateway([8080]) + + +@pytest.mark.parametrize(("host", "allowed"), [ + ("pypi.org", True), + ("files.pythonhosted.org", True), + ("a.b.pythonhosted.org", False), + ("pythonhosted.org", False), + ("10.2.3.4", True), + ("11.0.0.1", False), + ("evil.com", False), +]) +def test_allows_host_follows_tenki_one_label_wildcards(host, allowed): + policy = NetworkPolicy(mode=ALLOWLIST, allow_hosts=("pypi.org", "*.pythonhosted.org"), allow_cidrs=("10.0.0.0/8",)) + assert allows_host(policy, host) is allowed + assert allows_host(NetworkPolicy(), host) + + +def test_policy_from_info(): + def info(**fields): + return SimpleNamespace(**{"outbound_enabled": True, "egress": SessionEgressPolicy(), **fields}) + + assert policy_from_info(info()) == NetworkPolicy() + assert policy_from_info(info(outbound_enabled=False)) == NetworkPolicy(mode=ALLOWLIST) + assert policy_from_info(info(egress=SessionEgressPolicy(allow_domains=("a.com",)))) == NetworkPolicy( + mode=ALLOWLIST, allow_hosts=("a.com",), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("policy", "signed", "curled"), [ + (NetworkPolicy(), "https://bucket.s3.amazonaws.com/obj?sig", True), + (NetworkPolicy(mode=ALLOWLIST, allow_hosts=("pypi.org",)), "https://bucket.s3.amazonaws.com/obj?sig", False), + (NetworkPolicy(mode=ALLOWLIST, allow_hosts=("*.s3.amazonaws.com",)), "https://bucket.s3.amazonaws.com/o", True), + (None, "https://bucket.s3.amazonaws.com/obj?sig", False), + (NetworkPolicy(), None, False), +]) +async def test_a_download_the_vm_cannot_fetch_goes_over_stdin(monkeypatch, policy, signed, curled): + store = SimpleNamespace(signed_get_url=lambda _url: signed) + monkeypatch.setattr(sandbox_module, "get_config", lambda: SimpleNamespace(get_object_store_at=lambda _url: store)) + pushed = [] + + async def push(sandbox, object_store, object_url, vm_path): + pushed.append(vm_path) + + monkeypatch.setattr(sandbox_module, "push_object_over_stdin", push) + session = FakeSession() + await _sandbox(session, policy=policy)._download_object_to_vm("s3://b/obj", "/tmp/obj") + assert any("curl" in " ".join(call["argv"]) for call in session.exec_calls) is curled + assert bool(pushed) is not curled diff --git a/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py b/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py index 18aafecb..c929420e 100644 --- a/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py +++ b/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py @@ -18,6 +18,7 @@ SandboxProviderTypeError, build_sandbox_provider, ) +from agent_env.providers.sandbox_providers.tenki.provider import TenkiSandboxProvider from agent_env.providers.sandbox_providers.tensorlake.provider import TensorlakeSandboxProvider @@ -285,6 +286,35 @@ def test_tensorlake_missing_api_key_is_a_config_error(monkeypatch, tmp_path): build_sandbox_provider("tensorlake") +def test_tenki_builtin_receives_interpolated_key_and_snapshot(monkeypatch, tmp_path): + cfg = _write_config(tmp_path, """ + [sandbox.providers.tenki.config] + api_key = "env:TENKI_TEST_API_KEY" + snapshot_id = "snap-docker" + min_disk_gb = 40 + """) + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + monkeypatch.setenv("TENKI_TEST_API_KEY", "resolved-test-key") + + provider = build_sandbox_provider("tenki") + + assert isinstance(provider, TenkiSandboxProvider) + assert provider._api_key == "resolved-test-key" + assert provider._snapshot_id == "snap-docker" + assert provider._min_disk_gb == 40 + + +def test_tenki_missing_api_key_is_a_config_error(monkeypatch, tmp_path): + cfg = _write_config(tmp_path, """ + [sandbox.providers.tenki.config] + snapshot_id = "snap-docker" + """) + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + + with pytest.raises(ConfigError, match="requires a non-empty 'api_key'"): + build_sandbox_provider("tenki") + + def test_builtin_config_reaches_chain_members(monkeypatch, tmp_path): cfg = _write_config(tmp_path, """ [sandbox.providers.modal.config] diff --git a/tst/unit/test_capabilities.py b/tst/unit/test_capabilities.py index 6710db6f..08f65acc 100644 --- a/tst/unit/test_capabilities.py +++ b/tst/unit/test_capabilities.py @@ -151,7 +151,7 @@ def build(spec): monkeypatch.setattr("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", build) -_CONFIGURED_PROVIDERS = pytest.mark.parametrize("provider", ["e2b", "sail_vm", "tensorlake"]) +_CONFIGURED_PROVIDERS = pytest.mark.parametrize("provider", ["e2b", "sail_vm", "tensorlake", "tenki"]) @_CONFIGURED_PROVIDERS diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index b3b1a8d3..0994e983 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -185,8 +185,14 @@ def test_core_and_every_store_module_import_without_any_extra(): "agent_env.providers.sandbox_providers.tensorlake.provider:TensorlakeSandboxProvider", "agent_env.providers.sandbox_providers.sandbox_provider:SandboxProvider", ), + ( + "tenki", + "agent_env.providers.sandbox_providers.tenki.provider:TenkiSandboxProvider", + "agent_env.providers.sandbox_providers.sandbox_provider:SandboxProvider", + ), ], - ids=["gcs", "secret-manager", "s3", "s3-reexport", "secrets-manager-reexport", "dynamodb", "dynamodb-reexport", "tensorlake"], + ids=["gcs", "secret-manager", "s3", "s3-reexport", "secrets-manager-reexport", "dynamodb", "dynamodb-reexport", "tensorlake", + "tenki"], ) def test_an_impl_without_its_extra_names_the_extra_to_install(extra, impl, base): message = _run(_LOAD_IMPL, _only_in(extra), impl, base) diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index 3222a1b2..cdcf3d19 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -19,8 +19,8 @@ #: Only the agent-driving paths need a model endpoint. MODEL_ENDPOINT = "model_endpoint_configured" -#: The remote sandbox providers (``modal``, ``modal_vm``, ``e2b``, ``sail_vm``, ``tensorlake``) need credentials the resolved -#: config may not carry; the local default never does. +#: The remote sandbox providers (``modal``, ``modal_vm``, ``e2b``, ``sail_vm``, ``tensorlake``, ``tenki``) need +#: credentials the resolved config may not carry; the local default never does. REMOTE_SANDBOX = "remote_sandbox" #: The A2A agent a ``deploy_agent`` step without an id resolves to (``default_a2a_agent_id``); the @@ -76,16 +76,16 @@ def skip_without_model_endpoint() -> pytest.MarkDecorator: def remote_sandbox_is_available(provider: str) -> bool: """Whether the resolved config can build the ``modal`` / ``modal_vm`` / ``e2b`` / ``sail_vm`` / - ``tensorlake`` sandbox provider, credentials included. + ``tensorlake`` / ``tenki`` sandbox provider, credentials included. - Modal: any failure to resolve the credentials answers False. E2B, Sail and Tensorlake: building - the provider resolves ``[sandbox.providers..config]`` and its ``secret:`` references; only a + Modal: any failure to resolve the credentials answers False. E2B, Sail, Tensorlake and Tenki: + building the provider resolves ``[sandbox.providers..config]`` and its ``secret:`` references; only a ``ConfigError`` (absent or incomplete config, or the provider's extra not installed) answers False, any other failure answers True and lets the test fail on the real problem, as ``model_endpoint_is_configured`` does.""" - if provider not in ("modal", "modal_vm", "e2b", "sail_vm", "tensorlake"): + if provider not in ("modal", "modal_vm", "e2b", "sail_vm", "tensorlake", "tenki"): raise ValueError(f"unknown remote sandbox provider {provider!r}") - if provider in ("e2b", "sail_vm", "tensorlake"): + if provider in ("e2b", "sail_vm", "tensorlake", "tenki"): from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider diff --git a/uv.lock b/uv.lock index 7a13c552..1e2b75fc 100644 --- a/uv.lock +++ b/uv.lock @@ -88,6 +88,7 @@ dev = [ { name = "pytest-xdist" }, { name = "requests" }, { name = "sail" }, + { name = "tenki", extra = ["async"] }, { name = "tensorlake" }, { name = "uvicorn" }, ] @@ -105,6 +106,9 @@ gcp = [ sail = [ { name = "sail" }, ] +tenki = [ + { name = "tenki", extra = ["async"] }, +] tensorlake = [ { name = "tensorlake" }, ] @@ -112,7 +116,7 @@ tensorlake = [ [package.metadata] requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, - { name = "agentenv-framework", extras = ["aws", "explorer", "gcp", "sail", "tensorlake"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "explorer", "gcp", "sail", "tensorlake", "tenki"], marker = "extra == 'dev'" }, { name = "agentenv-framework-protocol", editable = "packages/agentenv-protocol" }, { name = "boto3", marker = "extra == 'aws'", specifier = ">=1.35.0" }, { name = "botocore", marker = "extra == 'aws'", specifier = ">=1.35.0" }, @@ -147,11 +151,12 @@ requires-dist = [ { name = "requests", marker = "extra == 'gcp'", specifier = ">=2.31" }, { name = "sail", marker = "extra == 'sail'", specifier = "~=0.12.8" }, { name = "starlette", specifier = ">=0.40" }, + { name = "tenki", extras = ["async"], marker = "extra == 'tenki'", specifier = "~=1.4.1" }, { name = "tensorlake", marker = "extra == 'tensorlake'", specifier = "~=0.5.140" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "uvicorn", marker = "extra == 'explorer'", specifier = ">=0.29" }, ] -provides-extras = ["aws", "explorer", "gcp", "sail", "tensorlake", "dev"] +provides-extras = ["aws", "explorer", "gcp", "sail", "tensorlake", "tenki", "dev"] [[package]] name = "agentenv-framework-protocol" @@ -3429,6 +3434,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/9c/8408129632f2cbd70018cff4a7eb78cc105b141de233b699feeda15d44c9/synchronicity-0.12.6-py3-none-any.whl", hash = "sha256:bc2bab6dde31f6bd9389912cef573b7fa7a5cb4a120ec786275e3284f06473ae", size = 41971, upload-time = "2026-10-02T15:41:42.672Z" }, ] +[[package]] +name = "tenki" +version = "1.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "grpcio" }, + { name = "protobuf" }, + { name = "websocket-client" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3c/59/93036224f9c175af8d7653e618f385ded93f873da5b727755f519ca5b75c/tenki-1.4.1.tar.gz", hash = "sha256:ff361ce1ef7b1676ce8fea519aa63cbce8cb1c1d98fc989d1fde1ff6ef8a79ae", size = 290340, upload-time = "2026-10-01T15:26:19.115Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4c/00/f5c5d6b8f75acf6cc2e25128e9418eddadf63571497b6b04e524d5ac09c8/tenki-1.4.1-py3-none-any.whl", hash = "sha256:c48352be3d1c514c7e47aa5264e55723675510ce4a20a025a2dedb1d83dcfada", size = 199117, upload-time = "2026-10-01T15:26:17.517Z" }, +] + +[package.optional-dependencies] +async = [ + { name = "websockets" }, +] + [[package]] name = "tensorlake" version = "0.5.149"