diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 48909e4d..63a44cd2 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -9,8 +9,10 @@ the entry, defaulting to the sole ``.py`` member or to ``run.py``. Target: with ``env_id`` set, the script runs on that env's VM host (no agent, no -LLM key); otherwise it runs inside the agent container (``agent_name``), where it -can see the agent's workspace. +LLM key); with ``sandbox_name`` set, on a ``deploy_sandbox`` sandbox (its VM host, or its +container for a container-mode sandbox); +otherwise it runs inside the agent container (``agent_name``), where it can see the +agent's workspace. """ from __future__ import annotations @@ -82,6 +84,7 @@ class RunCodeTaskStep(TaskStep): "timeout_seconds", "env_id", "agent_name", + "sandbox_name", ) def __init__( @@ -98,10 +101,20 @@ def __init__( timeout_seconds: Optional[int] = None, env_id: Optional[str] = None, agent_name: Optional[str] = None, + sandbox_name: Optional[str] = None, # --- base TaskStep params --- depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, ): + # A blank field (as a form submits it) is unset, not a second target. + env_id, sandbox_name, agent_name = env_id or None, sandbox_name or None, agent_name or None + targets = [ + name + for name, value in (("env_id", env_id), ("sandbox_name", sandbox_name), ("agent_name", agent_name)) + if value is not None + ] + if len(targets) > 1: + raise ValueError(f"run_code: set at most one of `env_id`, `sandbox_name` and `agent_name`; got {targets}") super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) self.script_artifact_id = script_artifact_id self.script_artifact_version = script_artifact_version @@ -111,9 +124,9 @@ def __init__( self.args = args or {} self.result_id = result_id or id self.timeout_seconds = timeout_seconds or self.DEFAULT_TIMEOUT_SECONDS - # `env_id` set → host mode (no agent); otherwise agent mode via `agent_name`. self.env_id = env_id - self.agent_name = agent_name or self.DEFAULT_AGENT_NAME + self.sandbox_name = sandbox_name + self.agent_name = None if env_id or sandbox_name else agent_name or self.DEFAULT_AGENT_NAME def to_dict(self) -> dict: base = super().to_dict() @@ -122,10 +135,14 @@ def to_dict(self) -> dict: @classmethod def from_dict(cls, data: dict) -> "RunCodeTaskStep": - return cls( - **cls._base_from_dict(data), - **{f: data[f] for f in cls._SERIALIZED_FIELDS if f in data}, - ) + fields = {f: data[f] for f in cls._SERIALIZED_FIELDS if f in data} + # Steps stored before the targets were exclusive can carry an `agent_name` beside `env_id`, which `env_id` + # overrode: drop it so they still load. A sandbox_name step only ever stored the default. + if fields.get("env_id") or ( + fields.get("sandbox_name") and fields.get("agent_name") == cls.DEFAULT_AGENT_NAME + ): + fields.pop("agent_name", None) + return cls(**cls._base_from_dict(data), **fields) def preflight(self) -> list[str]: """Config problems detectable without a sandbox; resolves documents, never content.""" @@ -155,7 +172,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: context.metadata.setdefault("script_results", {})[self.result_id] = result logger.info( "run_code '%s' completed (result_id=%s, target=%s)", - self.id, self.result_id, "host" if self.env_id else "agent", + self.id, self.result_id, self._target_kind(), ) return context @@ -278,10 +295,19 @@ def _build_input(self, context: TaskStepContext) -> dict: "results": context.metadata.get("script_results", {}), } + def _target_kind(self) -> str: + if self.env_id is not None: + return "host" + if self.sandbox_name is not None: + return "sandbox" + return "agent" + async def _resolve_target(self, context: TaskStepContext) -> _ExecTarget: - """Pick where the script runs: the env's VM host or the agent container.""" + """Pick where the script runs: the env's VM host, a deployed sandbox's host, or the agent container.""" if self.env_id is not None: return await self._host_target(context) + if self.sandbox_name is not None: + return await self._sandbox_target(context) return await self._agent_target(context) async def _stage( @@ -362,6 +388,35 @@ async def _host_target(self, context: TaskStepContext) -> _ExecTarget: sandbox = await provider.get_sandbox(deployed.sandbox_id) return _HostTarget(sandbox) + async def _sandbox_target(self, context: TaskStepContext) -> _ExecTarget: + 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, + ) + + deployed = next( + (s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name), None + ) + if deployed is None: + raise RuntimeError(f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes") + provider = ( + build_sandbox_provider(deployed.sandbox_type) + if deployed.sandbox_type + else get_sandbox_provider() + ) + sandbox = await provider.get_sandbox(deployed.sandbox_id) + if deployed.sandbox_mode == SANDBOX_MODE_VM: + return _HostTarget(sandbox) + # A container sandbox on a VM-backed provider runs its image beside the host: exec into it. + prefix = ( + ("sudo", "docker", "exec", "-u", "0", sandbox.container_name) + if isinstance(sandbox, VmSandbox) + else () + ) + return _AgentTarget(sandbox, prefix) + async def _agent_target(self, context: TaskStepContext) -> _ExecTarget: # Runs inside the agent container. VM mode needs `sudo docker exec -u 0`: # docker on the host needs sudo, and -u 0 reads the root-owned copied files. @@ -425,7 +480,7 @@ async def run(self, *command: str) -> tuple[int, str, str]: ... @dataclass class _HostTarget: - """Runs the script on the env's VM host (no agent).""" + """Runs the script on a VM host (an env's or a deployed sandbox's; no agent).""" sandbox: Any diff --git a/tst/integration/task_step/test_run_code_named_sandbox_local.py b/tst/integration/task_step/test_run_code_named_sandbox_local.py new file mode 100644 index 00000000..7c8b98f8 --- /dev/null +++ b/tst/integration/task_step/test_run_code_named_sandbox_local.py @@ -0,0 +1,212 @@ +"""run_code on a deploy_sandbox sandbox: files loaded onto it, a script that builds from them, and the result +collected off it and loaded back by step id, on a local VM-mode sandbox and on a local container-mode sandbox. + +A local VM-mode sandbox maps only `/app` paths onto its work dir, and a script's own file reads are not rewritten, +so its paths are absolute host paths under the test's tmp dir. The container-mode run needs Docker and a throwaway +local registry; its paths are inside the container. +""" + +import io +import json +import re +import shutil +import socket +import subprocess +import tempfile +import time +import uuid +import zipfile +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_sandbox import DeploySandboxTaskStep +from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep +from agent_env.task_step.task_steps.run_code import RunCodeTaskStep + +pytestmark = pytest.mark.integration + +_BUILD_PY = """\ +import json, pathlib, zipfile + +def run(inp): + args = inp["args"] + src, out = pathlib.Path(args["src"]), pathlib.Path(args["out"]) + out.parent.mkdir(parents=True, exist_ok=True) + names = sorted(p.name for p in src.iterdir()) + with zipfile.ZipFile(out, "w") as zf: + for name in names: + zf.write(src / name, f"bundle/{name}") + out.chmod(0o600) + return {"members": names} +""" + +_READ_PY = """\ +import pathlib, zipfile + +def run(inp): + (path,) = pathlib.Path(inp["args"]["dir"]).iterdir() + return {"name": path.name, "members": sorted(zipfile.ZipFile(path).namelist())} +""" + + +@pytest.fixture +def local_backends(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_LOCAL_SANDBOX_DIR", str(sandboxes)) + configure() + reset_artifact_store() + try: + yield tmp_path + finally: + shutil.rmtree(sandboxes, ignore_errors=True) + reset_artifact_store() + reset_config() + + +def _universe(src_dir, uid, files): + src_dir.mkdir() + for name, body in files.items(): + (src_dir / name).write_text(body) + return FileArtifactUniverse.put( + id=uid, + file_artifacts={ + name: FileArtifact.put(id=f"{uid}-{name}", description=name, file_path=str(src_dir / name)) + for name in files + }, + ) + + +def _chain(tmp_path, suffix, deploy, work): + inputs = _universe(tmp_path / "inputs", f"inputs-{suffix}", { + "task.json": json.dumps({"initial_prompt": "hi"}), + "rubrics.json": "[]", + }) + build = _universe(tmp_path / "build", f"build-{suffix}", {"build.py": _BUILD_PY}) + read = _universe(tmp_path / "read", f"read-{suffix}", {"read.py": _READ_PY}) + return Task.put(id=f"named-sandbox-{suffix}", steps=[ + deploy, + LoadArtifactTaskStep( + id="load-inputs", version=None, artifact_id=inputs.id, sandbox_name="box", + destination_path=f"{work}/inputs", + ), + RunCodeTaskStep( + id="build", version=None, script_artifact_id=build.id, sandbox_name="box", + args={"src": f"{work}/inputs", "out": f"{work}/out/bundle.zip"}, + ), + CollectArtifactsTaskStep( + id="collect", version=None, sandbox_name="box", base_path=f"{work}/out", artifact_paths=["bundle.zip"], + ), + LoadArtifactTaskStep( + id="load-back", version=None, artifact_from_step_id="collect", sandbox_name="box", + destination_path=f"{work}/loaded", + ), + RunCodeTaskStep( + id="read", version=None, script_artifact_id=read.id, sandbox_name="box", + args={"dir": f"{work}/loaded"}, + ), + ]) + + +def _assert_round_trip(ctx): + results = ctx.metadata["script_results"] + assert results["build"] == {"members": ["rubrics.json", "task.json"]} + assert results["read"] == {"name": "bundle.zip", "members": ["bundle/rubrics.json", "bundle/task.json"]} + universe = FileArtifactUniverse.get(ctx.metadata["collected_artifacts"]["collect"]["file_artifact_universe"]["id"]) + stored = universe.get_file_artifacts()["bundle.zip"].load() + assert sorted(zipfile.ZipFile(io.BytesIO(stored)).namelist()) == ["bundle/rubrics.json", "bundle/task.json"] + assert Task.get_instance(ctx.instance_id).status == "completed" + + +@pytest.mark.asyncio +async def test_a_script_builds_on_a_named_vm_sandbox_and_its_output_round_trips(local_backends): + tmp_path = local_backends + deploy = DeploySandboxTaskStep(id="deploy", version=None, sandbox_name="box", sandbox_mode="vm", sandbox_type="local") + ctx = await _chain(tmp_path, uuid.uuid4().hex[:8], deploy, str(tmp_path / "work")).run() + _assert_round_trip(ctx) + + +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(local_backends): + port = _free_port() + name = f"run-code-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) + set_image_store(LocalRegistryImageStore(host)) + try: + yield local_backends + finally: + _remove_sandbox_containers(local_backends / "sandboxes") + _docker("rm", "-f", name) + + +def _remove_sandbox_containers(sandboxes: Path): + # 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. Matching by image could reach another run's containers. + 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)}") + + +@pytest.mark.asyncio +async def test_a_script_builds_in_a_named_container_sandbox_and_its_output_round_trips(local_registry): + tmp_path = local_registry + suffix = uuid.uuid4().hex[:8] + port = _free_port() + with tempfile.TemporaryDirectory() as d: + # A non-root image user: the steps write as root, so collect has to read as root too. + (Path(d) / "Dockerfile").write_text( + f'FROM python:3.12-slim\nUSER nobody\nCMD ["python", "-m", "http.server", "{port}"]\n' + ) + built = _docker("build", "-t", f"run-code-box-{suffix}", d) + assert built.returncode == 0, built.stderr + image = DockerImageArtifact.put(id=f"run-code-box-{suffix}", description="box", image_name=f"run-code-box-{suffix}") + deploy = DeploySandboxTaskStep( + id="deploy", version=None, sandbox_name="box", sandbox_mode="container", sandbox_type="local", + image=image.image_name, port=port, + ) + try: + ctx = await _chain(tmp_path, suffix, deploy, "/work").run() + finally: + # The tag is unique to this run, as is its registry copy: drop them so repeated runs don't pile up images. + # This test's container still holds them, so remove it first. + _remove_sandbox_containers(tmp_path / "sandboxes") + refs = [r for ref in (f"run-code-box-{suffix}", f"*/run-code-box-{suffix}") for r in _docker( + "image", "ls", "--format", "{{.Repository}}:{{.Tag}}", "--filter", f"reference={ref}", + ).stdout.split()] + if refs: + _docker("rmi", "-f", *refs) + + (deployed,) = ctx.deployed_sandboxes + assert deployed.sandbox_mode == "container" + _assert_round_trip(ctx) diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index 9ea293d8..99ffdf5c 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -17,7 +17,8 @@ from agent_env.artifact import FileArtifact from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse from agent_env.env.env import DeployedEnv, DeployedGatewayEnv, EnvNeedsSandbox -from agent_env.task_step.context import DeployedAgent, TaskStepContext +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 import run_code as _run_code_module from tst.unit.event_loop_probe import on_event_loop from agent_env.task_step.task_steps.run_code import ( @@ -277,6 +278,165 @@ async def test_host_target_honors_custom_env_sandbox_type(): assert ctx.metadata["script_results"]["filter"] == {"keep": True} +def _sandbox_context(*, sandbox_type=None): + ctx = _make_context(with_agent=False) + ctx.deployed_sandboxes = [ + DeployedSandbox(sandbox_name="bundler", sandbox_id="sb-bundler", sandbox_mode="vm", sandbox_type=sandbox_type) + ] + return ctx + + +def _host_sandbox(): + sandbox = AsyncMock() + sandbox.mode = "vm" + sandbox.exec_with_output = AsyncMock( + side_effect=[(0, "", ""), (0, "13", ""), (0, '{"keep": true}', "")] + ) + provider = MagicMock() + provider.get_sandbox = AsyncMock(return_value=sandbox) + return sandbox, provider + + +@pytest.mark.asyncio +async def test_sandbox_target_runs_on_the_named_sandboxs_host(): + step = _step(sandbox_name="bundler", result_id="filter") + sandbox, provider = _host_sandbox() + with patch( + "agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", + return_value=provider, + ), patch("agent_env.artifact.Artifact.get", return_value=_file_artifact()), patch.object( + FileArtifact, "load", _load_body + ): + ctx = await step.execute(_sandbox_context()) + + provider.get_sandbox.assert_awaited_once_with("sb-bundler") + commands = [" ".join(map(str, c.args)) for c in sandbox.exec_with_output.await_args_list] + assert not any("docker" in c for c in commands) + sandbox.write_file_from_text.assert_not_awaited() + written = [c.args[1].rsplit("/", 1)[-1] for c in sandbox.write_host_file.await_args_list] + assert {"input.json", "runner.py"} <= set(written) + assert ctx.metadata["script_results"]["filter"] == {"keep": True} + + +@pytest.mark.asyncio +async def test_sandbox_target_honors_the_sandboxs_type(): + step = _step(sandbox_name="bundler", result_id="filter") + _, provider = _host_sandbox() + with patch( + "agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", + return_value=provider, + ) as build, patch( + "agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", + ) as default, patch("agent_env.artifact.Artifact.get", return_value=_file_artifact()), patch.object( + FileArtifact, "load", _load_body + ): + ctx = await step.execute(_sandbox_context(sandbox_type="modal_vm")) + + build.assert_called_once_with("modal_vm") + default.assert_not_called() + assert ctx.metadata["script_results"]["filter"] == {"keep": True} + + +@pytest.mark.asyncio +async def test_sandbox_target_names_a_sandbox_the_run_never_deployed(): + with pytest.raises(RuntimeError, match="^Sandbox 'ghost' not found in context.deployed_sandboxes$"): + await _step(sandbox_name="ghost")._sandbox_target(_sandbox_context()) + + +def _container_sandbox(*, vm_backed): + sandbox = AsyncMock(spec=VmSandbox) if vm_backed else AsyncMock() + sandbox.mode = "container" + sandbox.container_name = "agent-local-1" + # reset_dir's `rm -rf` + `mkdir -p`, then the script run, `wc -c`, then `cat`. + sandbox.exec_with_output = AsyncMock( + side_effect=[(0, "", ""), (0, "", ""), (0, "", ""), (0, "13", ""), (0, '{"keep": true}', "")] + ) + provider = MagicMock() + provider.get_sandbox = AsyncMock(return_value=sandbox) + return sandbox, provider + + +@pytest.mark.asyncio +@pytest.mark.parametrize("vm_backed,prefix", [ + (False, ()), + (True, ("sudo", "docker", "exec", "-u", "0", "agent-local-1")), +]) +async def test_sandbox_target_runs_inside_a_container_mode_sandbox(vm_backed, prefix): + step = _step(sandbox_name="bundler", result_id="filter") + ctx = _sandbox_context() + ctx.deployed_sandboxes[0].sandbox_mode = "container" + sandbox, provider = _container_sandbox(vm_backed=vm_backed) + with patch( + "agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", + return_value=provider, + ), patch("agent_env.artifact.Artifact.get", return_value=_file_artifact()), patch.object( + FileArtifact, "load", _load_body + ): + ctx = await step.execute(ctx) + + calls = [c.args for c in sandbox.exec_with_output.await_args_list] + assert all(c[: len(prefix)] == prefix for c in calls) + assert any("runner.py" in " ".join(c) for c in calls) + sandbox.write_host_file.assert_not_awaited() + written = [c.args[1].rsplit("/", 1)[-1] for c in sandbox.write_file_from_text.await_args_list] + assert {"input.json", "runner.py"} <= set(written) + assert ctx.metadata["script_results"]["filter"] == {"keep": True} + + +@pytest.mark.parametrize("targets", [ + {"env_id": "env-x", "sandbox_name": "bundler"}, + {"env_id": "env-x", "agent_name": "solver"}, + {"sandbox_name": "bundler", "agent_name": "solver"}, + {"sandbox_name": "bundler", "agent_name": RunCodeTaskStep.DEFAULT_AGENT_NAME}, + {"env_id": "env-x", "sandbox_name": "bundler", "agent_name": "solver"}, +]) +def test_the_three_targets_are_mutually_exclusive(targets): + with pytest.raises(ValueError, match="set at most one of `env_id`, `sandbox_name` and `agent_name`"): + _step(**targets) + + +@pytest.mark.parametrize("target", [{"env_id": "env-x"}, {"sandbox_name": "bundler"}]) +@pytest.mark.parametrize("blank", ["", None]) +def test_a_blank_agent_name_beside_a_host_target_is_unset(target, blank): + step = _step(**target, agent_name=blank) + assert step.agent_name is None + assert RunCodeTaskStep.from_dict({**step.to_dict(), "agent_name": blank}).to_dict() == step.to_dict() + + +def test_blank_targets_fall_back_to_the_default_agent(): + step = _step(env_id="", sandbox_name="", agent_name="") + assert (step.env_id, step.sandbox_name, step.agent_name) == (None, None, RunCodeTaskStep.DEFAULT_AGENT_NAME) + + +@pytest.mark.parametrize("target", [{"env_id": "env-x"}, {"sandbox_name": "bundler"}]) +def test_a_host_target_stores_no_agent_name(target): + assert _step(**target).to_dict()["agent_name"] is None + + +def test_no_target_runs_in_the_default_agent(): + assert _step().agent_name == RunCodeTaskStep.DEFAULT_AGENT_NAME + + +@pytest.mark.parametrize("target", [{"env_id": "env-x"}, {"sandbox_name": "bundler"}]) +def test_a_stored_step_with_the_default_agent_name_beside_its_target_still_loads(target): + stored = {**_step(**target).to_dict(), "agent_name": RunCodeTaskStep.DEFAULT_AGENT_NAME} + step = RunCodeTaskStep.from_dict(stored) + assert step.agent_name is None + assert {k: getattr(step, k) for k in target} == target + + +def test_a_stored_step_with_a_named_agent_beside_its_env_id_still_loads_on_the_host(): + stored = {**_step(env_id="env-x").to_dict(), "agent_name": "solver"} + step = RunCodeTaskStep.from_dict(stored) + assert (step.env_id, step.agent_name) == ("env-x", None) + + +def test_a_stored_step_with_a_named_agent_beside_its_sandbox_name_is_rejected(): + stored = {**_step(sandbox_name="bundler").to_dict(), "agent_name": "solver"} + with pytest.raises(ValueError, match="set at most one of"): + RunCodeTaskStep.from_dict(stored) + + @pytest.mark.asyncio async def test_agent_target_honors_the_agents_sandbox_type(): # an agent deployed on another provider is reached there, not through the configured default @@ -483,6 +643,12 @@ def test_config_round_trips_through_to_dict_from_dict(): assert RunCodeTaskStep.from_dict(d).to_dict() == d +def test_sandbox_name_round_trips_through_to_dict_from_dict(): + d = _step(sandbox_name="bundler").to_dict() + assert d["sandbox_name"] == "bundler" + assert RunCodeTaskStep.from_dict(d).to_dict() == d + + # --- multi-file script artifacts (FileArtifactUniverse) ---