diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 6997f342..f427a8a3 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -29,7 +29,6 @@ start_trusting, ) from agent_env.providers.sandbox_providers.sandbox_provider import all_sandbox_container_env, all_sandbox_url_rewrites -from agent_env.utils.paths import validate_relative_filename from agent_env.attribution import Attribution if TYPE_CHECKING: @@ -263,7 +262,7 @@ async def load_file_artifact_universe( destination_path: str, ) -> dict[str, str]: """Stage every FileArtifact in `universe` into the agent sandbox under destination_path.""" - from agent_env.providers.sandbox_providers.sandbox import VmSandbox + from agent_env.providers.sandbox_providers.sandbox import stage_files_into_container from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider destination = destination_path.rstrip("/") or "/" @@ -277,41 +276,11 @@ async def load_file_artifact_universe( ) return {} - dirs_to_make: set[str] = {destination} - loaded: dict[str, str] = {} - for filename in file_artifacts: - validate_relative_filename(filename) - dest_path = posixpath.join(destination, filename) - parent = posixpath.dirname(dest_path) - if parent: - dirs_to_make.add(parent) - loaded[filename] = dest_path - - for d in sorted(dirs_to_make): - if isinstance(sandbox, VmSandbox): - await sandbox.exec_script(f"docker exec -u 0 {shlex.quote(sandbox.container_name)} mkdir -p {shlex.quote(d)}") - else: - await sandbox.exec("mkdir", "-p", d) - - total = len(file_artifacts) logger.info( f"Loading FileArtifactUniverse '{universe.id}' v{universe.version} " - f"({total} file(s)) into agent at {destination}" + f"({len(file_artifacts)} file(s)) into agent at {destination}" ) - # Stage files in parallel with a bounded semaphore so a universe with - # many files doesn't accumulate per-file presign + docker cp latency. - # 8 matches the default concurrency used by the FAU `get-many` CLI. - sem = asyncio.Semaphore(8) - - async def _stage(filename: str, file_artifact: "Any") -> None: - dest_path = loaded[filename] - async with sem: - logger.info(f" {file_artifact.object_url} -> {dest_path}") - await sandbox.write_file_from_s3(file_artifact.object_url, dest_path) - - await asyncio.gather(*[_stage(fn, fa) for fn, fa in file_artifacts.items()]) - - return loaded + return await stage_files_into_container(sandbox, file_artifacts, destination) def __init__( self, diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index df726708..d89f8b09 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -5,6 +5,7 @@ import asyncio import logging import os +import posixpath import shlex import time import uuid @@ -15,6 +16,7 @@ from typing import TYPE_CHECKING, Any, Iterable, Optional from agent_env.config import get_config +from agent_env.utils.paths import validate_relative_filename if TYPE_CHECKING: from agent_env.store.object_store import ObjectStore @@ -406,3 +408,32 @@ async def write_file_from_text(self, content: str, destination_path: str) -> Non def port_bindings(host_ips: Iterable[str], host_port: int, container_port: int) -> list[str]: """Docker publish specs for one port: one per host IP, or a bare one (every interface) when there are none.""" return [f"{ip}:{host_port}:{container_port}" for ip in host_ips] or [f"{host_port}:{container_port}"] + + +async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, Any], destination: str) -> dict[str, str]: + """Write each object-store file into the sandbox's container under ``destination``; returns ``{name: path}``. + A VM-backed sandbox stages through its host into ``container_name``; a container sandbox is written directly.""" + loaded: dict[str, str] = {} + dirs_to_make: set[str] = {destination} + for filename in file_artifacts: + validate_relative_filename(filename) + dest_path = posixpath.join(destination, filename) + dirs_to_make.add(posixpath.dirname(dest_path)) + loaded[filename] = dest_path + + for d in sorted(dirs_to_make): + if isinstance(sandbox, VmSandbox): + await sandbox.exec_script(f"docker exec -u 0 {shlex.quote(sandbox.container_name)} mkdir -p {shlex.quote(d)}") + else: + await sandbox.exec("mkdir", "-p", d) + + # Bounded so a universe with many files doesn't serialize per-file presign + copy latency. + sem = asyncio.Semaphore(8) + + async def _stage(filename: str, file_artifact: Any) -> None: + async with sem: + logger.info(f" {file_artifact.object_url} -> {loaded[filename]}") + await sandbox.write_file_from_s3(file_artifact.object_url, loaded[filename]) + + await asyncio.gather(*(_stage(fn, fa) for fn, fa in file_artifacts.items())) + return loaded diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index 15f4e250..9fefa90b 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -152,7 +152,8 @@ def _exec_args(sandbox, container: Optional[str], cmd: tuple) -> tuple: if sandbox.mode != SANDBOX_MODE_VM: return cmd # container-mode sandboxes are the runtime and already root if container is not None: - return ("sudo", "docker", "exec", container, *cmd) + # As root, like the steps that write into the container: a non-root image user can't read their files. + return ("sudo", "docker", "exec", "-u", "0", container, *cmd) return ("sudo", *cmd) # host-mode agents write as uid 0, so read them back as root @@ -786,8 +787,13 @@ async def _collect_via_agent_container(self, context, store, artifact_id, versio ) async def _collect_via_vm_host(self, context, store, artifact_id, version): - """Collect off the VM's own filesystem — host-mode agents leave no container to exec into.""" - from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_sandbox_provider + """Collect off a deploy_sandbox sandbox: a VM's own filesystem, or a container-mode sandbox's container.""" + from agent_env.providers.sandbox_providers.sandbox import VmSandbox + from agent_env.providers.sandbox_providers.sandbox_provider import ( + SANDBOX_MODE_VM, + build_sandbox_provider, + get_sandbox_provider, + ) ds = next( (sb for sb in context.deployed_sandboxes if sb.sandbox_name == self.sandbox_name), None @@ -806,10 +812,18 @@ async def _collect_via_vm_host(self, context, store, artifact_id, version): provider = build_sandbox_provider(ds.sandbox_type) if ds.sandbox_type else get_sandbox_provider() sandbox = await self._resolve_live_sandbox(provider, ds.sandbox_id) - logger.info(f"Collecting from the VM host of sandbox '{self.sandbox_name}' ({ds.sandbox_id})") + # A container-mode sandbox on a VM-backed provider runs its image beside the host: read from the container. + container = ( + sandbox.container_name + if ds.sandbox_mode != SANDBOX_MODE_VM and isinstance(sandbox, VmSandbox) + else None + ) + logger.info( + f"Collecting from sandbox '{self.sandbox_name}' ({ds.sandbox_id}, {container or ds.sandbox_mode})" + ) return await self._collect_items( - provider, sandbox, None, items, context, store, artifact_id, version, + provider, sandbox, container, items, context, store, artifact_id, version, ) async def _collect_via_sandbox_container(self, context, store, artifact_id, version): diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index cb3cfc1c..6bc26a1c 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -288,6 +288,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: from agent_env.env.env import Env, require_gateway_url from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import ( + SANDBOX_MODE_VM, build_sandbox_provider, get_agent_sandbox_provider, get_sandbox_provider, @@ -329,6 +330,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: build_sandbox_provider(ds.sandbox_type) if ds.sandbox_type else get_sandbox_provider() ) target_sandbox = await provider.get_sandbox(ds.sandbox_id) + # A VM-mode sandbox with no container_name is loaded onto its host; a container-mode one, into its container. + onto_vm_host = ( + self.sandbox_name is not None and self.container_name is None and ds.sandbox_mode == SANDBOX_MODE_VM + ) env = None for ref in resolved_artifacts: @@ -352,8 +357,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: files = await _load_universe_into_container( target_sandbox, self.container_name, artifact, destination, ) - else: + elif onto_vm_host: files = await _load_universe_onto_vm(target_sandbox, artifact, destination) + else: + files = await _load_universe_into_sandbox_container(target_sandbox, artifact, destination) context.metadata.setdefault("loaded_file_artifact_universes", []).append({ "id": artifact.id, "version": artifact.version, @@ -418,8 +425,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # involved. if self.sandbox_name is not None: sandbox = target_sandbox - # None targets the VM host itself. + # None targets the VM host itself, or a container sandbox directly. container = self.container_name + if container is None and not onto_vm_host and isinstance(sandbox, VmSandbox): + container = sandbox.container_name else: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if agent is None or not agent.sandbox_id: @@ -523,15 +532,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: async def _load_one(url: str, filename: str) -> None: async with semaphore: dest = f"{destination}/{filename}" - if self.sandbox_name is not None and self.container_name is None: + if onto_vm_host: await self._load_url_onto_vm(sandbox, url, dest) else: await sandbox.write_file_from_url(url, dest) - target_desc = ( - f"VM sandbox '{self.sandbox_name}'" - if self.sandbox_name is not None and self.container_name is None - else "agent" - ) + target_desc = f"VM sandbox '{self.sandbox_name}'" if onto_vm_host else "agent" logger.info(f"Loaded URL into {target_desc}: {url} -> {dest}") await asyncio.gather(*(_load_one(u, f) for u, f in zip(urls, files))) @@ -734,6 +739,23 @@ async def _load_universe_onto_vm(sandbox, universe, destination: str) -> list[st return loaded +async def _load_universe_into_sandbox_container(sandbox, universe, destination: str) -> list[str]: + """Stage each file in `universe` into a container-mode sandbox, as an agent's container is loaded.""" + from agent_env.providers.sandbox_providers.sandbox import stage_files_into_container + + file_artifacts = universe.get_file_artifacts() + if not file_artifacts: + logger.warning( + f"FileArtifactUniverse '{universe.id}' v{universe.version} has no files; nothing to load" + ) + return [] + logger.info( + f"Loading FileArtifactUniverse '{universe.id}' v{universe.version} ({len(file_artifacts)} file(s)) " + f"into the container sandbox at {destination}" + ) + return list(await stage_files_into_container(sandbox, file_artifacts, destination)) + + async def _load_universe_into_container(sandbox, container_name: str, universe, destination: str) -> list[str]: """For each file in `universe`: pull from S3 to a VM temp path, then `docker cp` into `container_name` at `destination/`.""" file_artifacts = universe.get_file_artifacts() diff --git a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py index aeaa7d38..7290dbfa 100644 --- a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py +++ b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py @@ -28,6 +28,7 @@ import uuid from typing import ClassVar, Optional +from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_VM, build_sandbox_provider, @@ -118,13 +119,21 @@ def from_dict(cls, data: dict) -> "VerifySandboxTaskStep": async def execute(self, context: TaskStepContext) -> TaskStepContext: on_agent = self.sandbox_name is None - sandbox = await (self._agent_sandbox(context) if on_agent else self._deployed_sandbox(context)) + if on_agent: + sandbox = await self._agent_sandbox(context) + else: + deployed, sandbox = await self._deployed_sandbox(context) logger.info(f"Connected to sandbox {sandbox.sandbox_id} (mode={sandbox.mode})") - container = ( - await find_agent_container(sandbox) - if on_agent and sandbox.mode == SANDBOX_MODE_VM - else None - ) + if on_agent: + container = await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM else None + else: + # A container-mode sandbox on a VM-backed provider runs its image beside the host, and reattaches as a VM: + # its files are reached in the container, decided by the mode deploy_sandbox recorded. + container = ( + sandbox.container_name + if deployed.sandbox_mode != SANDBOX_MODE_VM and isinstance(sandbox, VmSandbox) + else None + ) if container: logger.info(f"Using agent container: {container}") @@ -204,7 +213,7 @@ async def _deployed_sandbox(self, context: TaskStepContext): if deployed.sandbox_type else get_sandbox_provider() ) - return await provider.get_sandbox(deployed.sandbox_id) + return deployed, await provider.get_sandbox(deployed.sandbox_id) async def _eval_criterion(self, sandbox, container: Optional[str], criterion: dict) -> dict: rtype = criterion["type"] @@ -225,7 +234,8 @@ async def _eval_criterion(self, sandbox, container: Optional[str], criterion: di @staticmethod def _probe_args(sandbox, container: Optional[str], cmd: tuple[str, ...]) -> tuple[str, ...]: if container is not None: - return ("sudo", "docker", "exec", container, *cmd) + # As root, like the steps that write into the container: a non-root image user can't read their files. + return ("sudo", "docker", "exec", "-u", "0", container, *cmd) if sandbox.mode == SANDBOX_MODE_VM: return ("sudo", *cmd) return cmd diff --git a/tst/integration/task_step/test_container_sandbox_files_local.py b/tst/integration/task_step/test_container_sandbox_files_local.py new file mode 100644 index 00000000..e3b141a5 --- /dev/null +++ b/tst/integration/task_step/test_container_sandbox_files_local.py @@ -0,0 +1,153 @@ +"""Files loaded into a container sandbox are found there by verify_sandbox and collect_artifacts, on the local provider. + +A local container sandbox, a deploy_sandbox one or an agent's, runs beside the host machine, so each step has to reach +into the container rather than act on the host. Every path here is one only the container has. Needs Docker and a +throwaway local registry. +""" + +import re +import shutil +import socket +import subprocess +import tempfile +import time +import uuid +from pathlib import Path + +import httpx +import pytest + +from agent_env.artifact import DockerImageArtifact, FileArtifact, FileArtifactUniverse +from agent_env.artifact.store import reset_artifact_store +from agent_env.config import configure, reset_config, set_image_store +from agent_env.store.image_store import LocalRegistryImageStore +from agent_env.task import Task +from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep +from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep +from agent_env.task_step.task_steps.verifiers.verify_sandbox import VerifySandboxTaskStep +from tst.util.a2a_test_agent import put_test_agent + +pytestmark = pytest.mark.integration + + +def _docker(*args): + return subprocess.run(["docker", *args], capture_output=True, text=True) + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@pytest.fixture +def local_registry(monkeypatch, tmp_path): + sandboxes = tmp_path / "sandboxes" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_IMAGE_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + port = _free_port() + name = f"files-local-registry-{uuid.uuid4().hex[:8]}" + started = _docker("run", "-d", "--rm", "--name", name, "-p", f"127.0.0.1:{port}:5000", "registry:2") + assert started.returncode == 0, started.stderr + host = f"localhost:{port}" + deadline = time.time() + 45 + while time.time() < deadline: + try: + if httpx.get(f"http://{host}/v2/", timeout=2).status_code in (200, 401): + break + except httpx.HTTPError: + pass + time.sleep(0.5) + configure() + set_image_store(LocalRegistryImageStore(host)) + reset_artifact_store() + try: + yield tmp_path + finally: + # Only this test's containers: every local sandbox it deployed has a work dir in its own sandbox root, and its + # container is named after that sandbox. Images are shared, so matching by image would reach other runs. + for work_dir in sandboxes.glob("agent-env-*"): + if m := re.match(r"agent-env-(local-[0-9a-f]+)-", work_dir.name): + _docker("rm", "-f", f"agent-{m.group(1)}") + _docker("rm", "-f", name) + shutil.rmtree(sandboxes, ignore_errors=True) + reset_artifact_store() + reset_config() + + +def _probe_universe(tmp_path, suffix): + src = tmp_path / "probe.txt" + src.write_text("hello from the probe\n") + return FileArtifactUniverse.put( + id=f"probe-files-{suffix}", + file_artifacts={"probe.txt": FileArtifact.put(id=f"probe-file-{suffix}", description="probe", file_path=str(src))}, + ) + + +def _load_verify_collect(target: dict, universe, probe_dir: str): + return [ + LoadArtifactTaskStep(id="load", version=None, artifact_id=universe.id, destination_path=probe_dir, **target), + VerifySandboxTaskStep( + id="verify", version=None, verifier_id="probe", base_dir=probe_dir, **target, + criteria=[{"type": "probe_file_contains", "criterion": "loaded", "paths": ["probe.txt"], "expected": "hello"}], + ), + CollectArtifactsTaskStep(id="collect", version=None, base_path=probe_dir, artifact_paths=["probe.txt"], **target), + ] + + +def _assert_found_in_the_container(ctx, probe_dir): + assert not Path(probe_dir).exists() + assert ctx.metadata["verifications"]["probe"]["score"] == 1.0 + universe = FileArtifactUniverse.get(ctx.metadata["collected_artifacts"]["collect"]["file_artifact_universe"]["id"]) + assert universe.get_file_artifacts()["probe.txt"].load() == b"hello from the probe\n" + assert Task.get_instance(ctx.instance_id).status == "completed" + + +@pytest.mark.asyncio +async def test_a_container_mode_deploy_sandbox(local_registry): + tmp_path = local_registry + suffix = uuid.uuid4().hex[:8] + port = _free_port() + probe_dir = f"/probe-{suffix}" + with tempfile.TemporaryDirectory() as d: + # A non-root image user and a root-only probe dir: the steps only reach the loaded files as root. + (Path(d) / "Dockerfile").write_text( + f"FROM python:3.12-slim\nRUN mkdir -m 700 {probe_dir}\nUSER nobody\n" + f'CMD ["python", "-m", "http.server", "{port}"]\n' + ) + built = _docker("build", "-t", f"files-local-box-{suffix}", d) + assert built.returncode == 0, built.stderr + image = DockerImageArtifact.put(id=f"files-local-box-{suffix}", description="box", image_name=f"files-local-box-{suffix}") + task = Task.put(id=f"files-box-{suffix}", steps=[ + DeploySandboxTaskStep( + id="deploy", version=None, sandbox_name="box", sandbox_mode="container", sandbox_type="local", + image=image.image_name, port=port, + ), + *_load_verify_collect({"sandbox_name": "box"}, _probe_universe(tmp_path, suffix), probe_dir), + ]) + ctx = await task.run() + + _assert_found_in_the_container(ctx, probe_dir) + + +@pytest.mark.asyncio +async def test_a_local_agent(local_registry): + tmp_path = local_registry + suffix = uuid.uuid4().hex[:8] + agent = put_test_agent(f"files-local-agent-{suffix}") + probe_dir = f"/probe-{suffix}" + task = Task.put(id=f"files-agent-{suffix}", steps=[ + DeployAgentTaskStep( + id="deploy", version=None, env_ids=[], a2a_agent_id=agent.id, a2a_agent_version=agent.version, + agent_name="solver", sandbox_type="local", + ), + *_load_verify_collect({"agent_name": "solver"}, _probe_universe(tmp_path, suffix), probe_dir), + ]) + ctx = await task.run() + + _assert_found_in_the_container(ctx, probe_dir) diff --git a/tst/unit/task_step/test_collect_artifacts.py b/tst/unit/task_step/test_collect_artifacts.py index ff8861fc..f458e6f5 100644 --- a/tst/unit/task_step/test_collect_artifacts.py +++ b/tst/unit/task_step/test_collect_artifacts.py @@ -553,10 +553,10 @@ def test_a_reattached_local_agent_is_read_inside_its_own_container(self, tmp_pat assert calls == [ ("sudo", "docker", "ps", "--format", "{{.Names}}"), - ("sudo", "docker", "exec", "agent-local-agent1", "find", "/app/artifact", "-type", "f", "-printf", "%P\n"), + ("sudo", "docker", "exec", "-u", "0", "agent-local-agent1", "find", "/app/artifact", "-type", "f", "-printf", "%P\n"), ] assert _exec_args(sandbox, container, ("bash", "-c", "base64 < /app/artifact/a.txt")) == ( - "sudo", "docker", "exec", "agent-local-agent1", "bash", "-c", "base64 < /app/artifact/a.txt") + "sudo", "docker", "exec", "-u", "0", "agent-local-agent1", "bash", "-c", "base64 < /app/artifact/a.txt") def test_a_reattached_local_agent_never_borrows_another_runs_container(self, tmp_path, monkeypatch): sandbox, calls = _reattached_local_agent(tmp_path, monkeypatch, running="a2a-agent-other\n") diff --git a/tst/unit/task_step/test_collect_artifacts_vm_host.py b/tst/unit/task_step/test_collect_artifacts_vm_host.py index ef6a6425..945b06c1 100644 --- a/tst/unit/task_step/test_collect_artifacts_vm_host.py +++ b/tst/unit/task_step/test_collect_artifacts_vm_host.py @@ -224,7 +224,7 @@ def _args(self, mode, container): return _exec_args(type("S", (), {"mode": mode})(), container, ("find", "/app")) def test_vm_with_container_execs_into_it(self): - assert self._args("vm", "agent-api") == ("sudo", "docker", "exec", "agent-api", "find", "/app") + assert self._args("vm", "agent-api") == ("sudo", "docker", "exec", "-u", "0", "agent-api", "find", "/app") def test_vm_without_container_still_sudoes(self): assert self._args("vm", None) == ("sudo", "find", "/app") diff --git a/tst/unit/task_step/test_container_sandbox_routing.py b/tst/unit/task_step/test_container_sandbox_routing.py new file mode 100644 index 00000000..0c6bb92d --- /dev/null +++ b/tst/unit/task_step/test_container_sandbox_routing.py @@ -0,0 +1,122 @@ +"""Which container collect_artifacts and verify_sandbox reach into, per kind of sandbox. + +A container sandbox on a container provider (Modal) is the runtime itself; one on a VM-backed provider (local, modal_vm, +e2b) runs beside its host, so its files are only reachable through `docker exec `. A VM-backed provider +reattaches every sandbox as a VM, so a deploy_sandbox sandbox's mode is the one deploy_sandbox recorded. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agent_env.providers.sandbox_providers import sandbox_provider as sp_mod +from agent_env.providers.sandbox_providers.sandbox import VmSandbox +from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext +from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep +from agent_env.task_step.task_steps.verifiers import verify_sandbox +from agent_env.task_step.task_steps.verifiers.verify_sandbox import VerifySandboxTaskStep + +CONTAINER = "agent-local-1" + + +def _sandbox(kind: str): + """`vm`: a VM with an agent container; `vm-backed-container`: local/modal_vm-style; `container`: Modal-style.""" + sandbox = AsyncMock(spec=VmSandbox) if kind != "container" else AsyncMock() + sandbox.sandbox_id = "sb-1" + sandbox.mode = "vm" if kind == "vm" else "container" + sandbox.container_name = CONTAINER + sandbox.exec_with_output = AsyncMock(return_value=(0, f"{CONTAINER}\n", "")) + return sandbox + + +def _provide(monkeypatch, sandbox): + provider = MagicMock() + provider.get_sandbox = AsyncMock(return_value=sandbox) + for module in (sp_mod, verify_sandbox): # verify_sandbox binds the getters at import + monkeypatch.setattr(module, "get_agent_sandbox_provider", lambda: provider) + monkeypatch.setattr(module, "get_sandbox_provider", lambda: provider) + + +def _agent_context(): + ctx = TaskStepContext() + ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1")] + return ctx + + +def _sandbox_context(recorded_mode: str): + ctx = TaskStepContext() + ctx.deployed_sandboxes = [DeployedSandbox(sandbox_name="box", sandbox_id="sb-1", sandbox_mode=recorded_mode)] + return ctx + + +async def _collect_container(monkeypatch, step, ctx): + seen = {} + + async def _items(self, provider, sandbox, container, *rest): + seen["container"] = container + return {}, {}, [] + + monkeypatch.setattr(CollectArtifactsTaskStep, "_collect_items", _items) + # Nothing is collected from the stub, so the step then fails; the container it chose is already recorded. + with pytest.raises(RuntimeError, match="produced no file_artifact_universe"): + await step.execute(ctx) + return seen["container"] + + +async def _verify_container(monkeypatch, step, ctx): + seen = {} + + async def _eval(self, sandbox, container, criterion): + seen["container"] = container + return {"passed": True, "score": 1.0, "justification": ""} + + monkeypatch.setattr(VerifySandboxTaskStep, "_eval_criterion", _eval) + await step.execute(ctx) + return seen["container"] + + +_CRITERIA = [{"type": "probe_file_exists", "criterion": "c", "paths": ["out.txt"]}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind,expected", [("vm", CONTAINER), ("container", None)]) +async def test_collect_reaches_an_agents_container(monkeypatch, kind, expected): + _provide(monkeypatch, _sandbox(kind)) + step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver", artifact_paths=["out.txt"]) + assert await _collect_container(monkeypatch, step, _agent_context()) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind,recorded,expected", [ + ("vm", "vm", None), + ("vm", "container", CONTAINER), # a VM-backed provider reattaches a container sandbox as a VM + ("vm-backed-container", "container", CONTAINER), + ("container", "container", None), +]) +async def test_collect_reaches_a_named_sandboxs_container(monkeypatch, kind, recorded, expected): + _provide(monkeypatch, _sandbox(kind)) + step = CollectArtifactsTaskStep(id="c", version=None, sandbox_name="box", artifact_paths=["out.txt"]) + assert await _collect_container(monkeypatch, step, _sandbox_context(recorded)) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind,expected", [("vm", CONTAINER), ("container", None)]) +async def test_verify_sandbox_probes_an_agents_container(monkeypatch, kind, expected): + _provide(monkeypatch, _sandbox(kind)) + step = VerifySandboxTaskStep(id="v", version=None, agent_name="solver", criteria=_CRITERIA) + assert await _verify_container(monkeypatch, step, _agent_context()) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind,recorded,expected", [ + ("vm", "vm", None), + ("vm", "container", CONTAINER), + ("vm-backed-container", "container", CONTAINER), + ("container", "container", None), +]) +async def test_verify_sandbox_probes_a_named_sandboxs_container(monkeypatch, kind, recorded, expected): + _provide(monkeypatch, _sandbox(kind)) + step = VerifySandboxTaskStep(id="v", version=None, sandbox_name="box", criteria=_CRITERIA) + assert await _verify_container(monkeypatch, step, _sandbox_context(recorded)) == expected diff --git a/tst/unit/task_step/test_load_artifact_container_sandbox.py b/tst/unit/task_step/test_load_artifact_container_sandbox.py new file mode 100644 index 00000000..5d80e64f --- /dev/null +++ b/tst/unit/task_step/test_load_artifact_container_sandbox.py @@ -0,0 +1,109 @@ +"""LoadArtifactTaskStep with `sandbox_name` naming a container-mode deploy_sandbox sandbox: files land in its container +the way an agent's container is loaded, not on a VM host the sandbox does not have.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from agent_env.artifact.artifact import Artifact +from agent_env.artifact.artifacts.file import FileArtifact +from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +from agent_env.providers.sandbox_providers import sandbox_provider as sp_mod +from agent_env.providers.sandbox_providers.sandbox import VmSandbox +from agent_env.task_step.context import DeployedSandbox, TaskStepContext +from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep + +UNIVERSE_ID = "bundler-inputs" + + +@pytest.fixture +def universe(monkeypatch): + uni = FileArtifactUniverse(id=UNIVERSE_ID, version=2, file_artifact_refs={}) + staged = { + "task.json": FileArtifact( + id="fa-1", version=1, description="task", filename="task.json", + content_type="application/json", s3_url="s3://bucket/task.json", + ), + "kb/INDEX.md": FileArtifact( + id="fa-2", version=1, description="index", filename="INDEX.md", + content_type="text/markdown", s3_url="s3://bucket/INDEX.md", + ), + } + monkeypatch.setattr(type(uni), "get_file_artifacts", lambda self: staged) + monkeypatch.setattr(Artifact, "get", classmethod(lambda cls, id, version=None: uni)) + return uni + + +def _context(sandbox, monkeypatch): + async def _get_sandbox(sandbox_id): + return sandbox + + monkeypatch.setattr( + sp_mod, "get_sandbox_provider", + lambda: type("P", (), {"get_sandbox": staticmethod(_get_sandbox)})(), + ) + ctx = TaskStepContext() + ctx.deployed_sandboxes = [DeployedSandbox(sandbox_name="bundler", sandbox_id="sb-1", sandbox_mode="container")] + return ctx + + +def _container_sandbox(*, vm_backed): + sandbox = AsyncMock(spec=VmSandbox) if vm_backed else AsyncMock() + sandbox.mode = "container" + sandbox.container_name = "agent-local-1" + return sandbox + + +@pytest.mark.asyncio +async def test_a_container_sandbox_is_written_directly(universe, monkeypatch): + sandbox = _container_sandbox(vm_backed=False) + step = LoadArtifactTaskStep( + id="load", version=None, sandbox_name="bundler", artifact_id=UNIVERSE_ID, destination_path="/work/inputs", + ) + + ctx = await step.execute(_context(sandbox, monkeypatch)) + + assert sorted(c.args for c in sandbox.exec.await_args_list) == [ + ("mkdir", "-p", "/work/inputs"), ("mkdir", "-p", "/work/inputs/kb"), + ] + assert sorted(c.args for c in sandbox.write_file_from_s3.await_args_list) == [ + ("s3://bucket/INDEX.md", "/work/inputs/kb/INDEX.md"), + ("s3://bucket/task.json", "/work/inputs/task.json"), + ] + sandbox.exec_script.assert_not_awaited() + sandbox.load_s3_file.assert_not_awaited() + (loaded,) = ctx.metadata["loaded_file_artifact_universes"] + assert loaded["sandbox_name"] == "bundler" + assert loaded["destination_path"] == "/work/inputs" + assert sorted(loaded["files"]) == ["kb/INDEX.md", "task.json"] + + +@pytest.mark.asyncio +async def test_a_container_on_a_vm_backed_provider_is_loaded_through_its_container(universe, monkeypatch): + sandbox = _container_sandbox(vm_backed=True) + step = LoadArtifactTaskStep( + id="load", version=None, sandbox_name="bundler", artifact_id=UNIVERSE_ID, destination_path="/work/inputs", + ) + + await step.execute(_context(sandbox, monkeypatch)) + + scripts = [c.args[0] for c in sandbox.exec_script.await_args_list] + assert scripts and all(s.startswith("docker exec -u 0 agent-local-1 mkdir -p ") for s in scripts) + assert len(sandbox.write_file_from_s3.await_args_list) == 2 + sandbox.load_s3_file.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_urls_go_into_a_container_sandbox(monkeypatch): + sandbox = _container_sandbox(vm_backed=False) + step = LoadArtifactTaskStep( + id="load", version=None, sandbox_name="bundler", urls=["https://example.com/data.zip"], + destination_path="/work/inputs", + ) + + await step.execute(_context(sandbox, monkeypatch)) + + sandbox.write_file_from_url.assert_awaited_once_with("https://example.com/data.zip", "/work/inputs/data.zip") + sandbox.exec_script.assert_not_awaited() diff --git a/tst/unit/task_step/verifiers/test_verify_sandbox.py b/tst/unit/task_step/verifiers/test_verify_sandbox.py index 71878eac..1f703d4e 100644 --- a/tst/unit/task_step/verifiers/test_verify_sandbox.py +++ b/tst/unit/task_step/verifiers/test_verify_sandbox.py @@ -157,9 +157,9 @@ async def test_agent_source_probes_inside_the_agent_container(provide): assert sandbox.calls == [ ("sudo", "docker", "ps", "--format", "{{.Names}}"), - ("sudo", "docker", "exec", "agent-api", "test", "-e", "/app/greeting/hello.txt"), - ("sudo", "docker", "exec", "agent-api", "test", "-e", "/etc/hosts"), - ("sudo", "docker", "exec", "agent-api", "cat", "/app/greeting/hello.txt"), + ("sudo", "docker", "exec", "-u", "0", "agent-api", "test", "-e", "/app/greeting/hello.txt"), + ("sudo", "docker", "exec", "-u", "0", "agent-api", "test", "-e", "/etc/hosts"), + ("sudo", "docker", "exec", "-u", "0", "agent-api", "cat", "/app/greeting/hello.txt"), ("sudo", "docker", "exec", "-w", "/app/greeting", "agent-api", "bash", "-c", "python3 check.py"), ] assert ctx.metadata["verifications"]["hello"]["score"] == 1.0 @@ -174,9 +174,9 @@ async def test_a_reattached_local_agent_is_probed_inside_its_own_container(provi assert calls == [ ("sudo", "docker", "ps", "--format", "{{.Names}}"), - ("sudo", "docker", "exec", "agent-local-agent1", "test", "-e", "/app/greeting/hello.txt"), - ("sudo", "docker", "exec", "agent-local-agent1", "test", "-e", "/etc/hosts"), - ("sudo", "docker", "exec", "agent-local-agent1", "cat", "/app/greeting/hello.txt"), + ("sudo", "docker", "exec", "-u", "0", "agent-local-agent1", "test", "-e", "/app/greeting/hello.txt"), + ("sudo", "docker", "exec", "-u", "0", "agent-local-agent1", "test", "-e", "/etc/hosts"), + ("sudo", "docker", "exec", "-u", "0", "agent-local-agent1", "cat", "/app/greeting/hello.txt"), ("sudo", "docker", "exec", "-w", "/app/greeting", "agent-local-agent1", "bash", "-c", "python3 check.py"), ]