diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 0e2e3758e..134d04964 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -632,6 +632,7 @@ async def _load_from_snapshot(self, snapshot) -> None: host_ips=self._sandbox.host_ips, extra_hosts=self._sandbox.extra_hosts, mcp_server_name=self._mcp_server_name or self.name, + trust_files=await gw._stage_trust_files(self._sandbox), ) compose_content = compose_content.replace( " volumes:\n - ./init-schemas.sql:/docker-entrypoint-initdb.d/init-schemas.sql\n", "", diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 90b1fbe1e..dac3bc46a 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -49,7 +49,12 @@ _size_kwargs, _tool_names, ) -from agent_env.providers.sandbox_providers.local_docker.provider import remove_local_work_dir +from agent_env.providers.sandbox_providers.local_docker.provider import ( + LOCAL_TRUST_DIR, + LOCAL_TRUST_ENV, + local_grant_trust, + remove_local_work_dir, +) from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_CONTAINER, @@ -75,6 +80,8 @@ _LOG_LINE_CHARS = 300 # Gathering it may stall with the connection; past this, the deploy raises the error it had. _DOWN_SERVICES_TIMEOUT_S = 60.0 +# Where beside the compose file its services' copy of the local transfer CA's trust files is kept. +_TRUST_FILES_DIR = "trust" # MCP server health-check budget. Some servers wait on an upstream before they bind; a server # that is up passes its first probe, so the long start period only delays failing a broken one. @@ -270,8 +277,12 @@ def create_docker_compose( host_ips: tuple[str, ...] = (), extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, + trust_files: str | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. + With ``trust_files``, where the compose finds the local transfer CA's trust files, each MCP server mounts + them where ``LOCAL_TRUST_ENV`` points, unless its own variables say otherwise, and carries ``extra_hosts``, so + it reaches and trusts the local grant server. Supports MCP servers, websites, or both. Website containers (backend + frontend) are added to the same docker network. The gateway receives WEBSITE_URLS env var for discovery. @@ -418,9 +429,12 @@ def create_docker_compose( f" - ENVIRONMENT_NAME={server.environment_name}", f" - DATABASE_URL={db_url}", ]) - if server.extra_env_vars: - for key, value in server.extra_env_vars.items(): - lines.append(f" - {key}={value}") + server_env = {**(LOCAL_TRUST_ENV if trust_files is not None else {}), **(server.extra_env_vars or {})} + lines.extend(f" - {key}={value}" for key, value in server_env.items()) + if trust_files is not None: + lines.extend([" volumes:", f' - "{trust_files}:{LOCAL_TRUST_DIR}:ro"']) + if extra_hosts: + lines.extend([" extra_hosts:", *(f' - "{entry}"' for entry in extra_hosts)]) lines.extend([ " healthcheck:", f' test: ["CMD-SHELL", "python3 -c \\"import socket; s=socket.create_connection((\'localhost\',{gateway_port}),2); s.close()\\""]', @@ -799,6 +813,7 @@ async def _stand_up(_spec): self._all_environments(mcp_servers, website_configs), instance=state_instance ) + trust_files = await self._stage_trust_files(sandbox) compose_content = self.create_docker_compose( mcp_servers=mcp_servers, gateway_image=gateway_env.docker_image_artifact.image_name, @@ -812,6 +827,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, + trust_files=trust_files, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") @@ -1121,6 +1137,20 @@ async def read_trajectory(self, sandbox: VmSandbox) -> list[dict]: events.append(json.loads(line)) return events + @staticmethod + async def _stage_trust_files(sandbox: VmSandbox) -> str | None: + """Write the local transfer CA's trust files beside the compose file and return where it finds them, on a sandbox + on this machine when the store hands out local grants, which its services are given; None otherwise. They are + made readable by any user a service runs as, whatever the umask they were written under.""" + if not sandbox.ON_THIS_MACHINE or (trust_dir := await asyncio.to_thread(local_grant_trust)) is None: + return None + directory = shlex.quote(f"{GATEWAY_APP_DIR}/{_TRUST_FILES_DIR}") + await sandbox.exec_script(f"mkdir -p {directory}", max_retries=2) + for path in sorted(trust_dir.iterdir()): + await sandbox.write_host_file(await asyncio.to_thread(path.read_bytes), f"{GATEWAY_APP_DIR}/{_TRUST_FILES_DIR}/{path.name}") + await sandbox.exec_script(f"chmod 755 {directory} && chmod 644 {directory}/*", max_retries=2) + return f"./{_TRUST_FILES_DIR}" + @staticmethod async def _write_file(sandbox: VmSandbox, path: str, content: str) -> None: """Write ``content`` and a newline to ``path`` on the sandbox host, the content one quoted argument so no byte diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 86c740656..711032afd 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -597,7 +597,8 @@ async def _write_bytes_to_vm_path(self, data: bytes, vm_path: str) -> None: await self.exec_script(f": > {shlex.quote(vm_b64)}") for i in range(0, len(encoded), self._WFT_CHUNK_BYTES): await self.exec_script(f"printf '%s' {shlex.quote(encoded[i:i + self._WFT_CHUNK_BYTES])} >> {shlex.quote(vm_b64)}") - await self.exec_script(f"base64 -d {shlex.quote(vm_b64)} > {shlex.quote(vm_path)} && rm -f {shlex.quote(vm_b64)}") + # From stdin: macOS's base64, a local sandbox's, takes no file argument. + await self.exec_script(f"base64 -d < {shlex.quote(vm_b64)} > {shlex.quote(vm_path)} && rm -f {shlex.quote(vm_b64)}") async def write_host_file(self, data: bytes, vm_path: str) -> None: """Write bytes to vm_path on the VM host itself, not into the agent container.""" diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index 16c325d22..dc696a29a 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -8,6 +8,7 @@ import mimetypes import os import tempfile +import time import uuid from dataclasses import dataclass from datetime import datetime, timezone @@ -16,16 +17,16 @@ import httpx -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object_path -from agentenv_protocol.client import find_extension, find_extension_method -from agentenv_protocol.transfers import WriteNamespaceGrant +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object, uploaded_object_path +from agentenv_protocol.client import find_extension, find_extension_method, write_kinds +from agentenv_protocol.transfers import HttpPartsPutGrant, WriteNamespaceGrant, part_ranges from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant from agent_env.env.gateway.constants import CHANGES_PAGE_MAX, EXT_CHANGES_URI, EXT_TRAJECTORY_URI from agent_env.store import get_config from agent_env.store.base import GrantUnavailableError from agent_env.store.ids import derive_id, is_local_id, validate_local_id -from agent_env.store.object_store import MIN_GRANT_LIFETIME_SECONDS, ObjectStore +from agent_env.store.object_store import MIN_GRANT_LIFETIME_SECONDS, ObjectStore, PendingWrite from agent_env.store.object_store.object_store import issues_grants_to from agent_env.store.routing import in_local_run from agent_env.task_step.context import TaskStepContext @@ -48,6 +49,9 @@ # What a service may upload through its snapshot grant: one bundle, as large as one upload to the # store can be up to this. ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=1, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) +# The most a service may upload as its bundle through an object write: room for the largest bundles seen, +# about twice over, in a grant of a few hundred part URLs. +ENV_SNAPSHOT_MAX_BYTES = 16 * 1024**3 _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" @@ -89,6 +93,48 @@ def _snapshot_upload(card: dict, timeout_seconds: float, sandbox_type: Optional[ return _SnapshotUpload(store, grant) +def _snapshot_write( + card: dict, environment_name: str, timeout_seconds: float, sandbox_type: Optional[str] +) -> Optional[PendingWrite]: + """A write of a service's bundle to an object of its own, for a service whose card takes parts grants; + None when it takes none, or the object store issues no grant that reaches the service's sandbox, of + ``sandbox_type`` (None: unknown), and lasts the export.""" + kinds = write_kinds(card) + if "http-put-parts" not in kinds: + return None + config = get_config() + store = config.get_object_store() + if not issues_grants_to(store, sandbox_type): + return None + object_url = store.object_url( + f"{config.get_artifact_key_prefix()}{_SNAPSHOT_KEY_PREFIX}/{uuid.uuid4().hex}/{environment_name}.zip" + ) + try: + return store.begin_write( + object_url, + media_type="application/zip", + max_bytes=ENV_SNAPSHOT_MAX_BYTES, + kinds=kinds, + expires_in=max(int(timeout_seconds), MIN_GRANT_LIFETIME_SECONDS), + ) + except GrantUnavailableError as exc: + logger.info("snapshot_env: the object store issues no snapshot write grant (%s)", exc) + return None + + +async def _finish_in_thread(fn, *args): + """``await asyncio.to_thread(fn, *args)``, except that a caller cancelled meanwhile first waits for ``fn`` to + return, and aborts a write it began, so neither a write nor its completion outlives the caller's cleanup.""" + task = asyncio.ensure_future(asyncio.to_thread(fn, *args)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + await asyncio.wait([task]) + if task.exception() is None and isinstance(task.result(), PendingWrite): + await asyncio.to_thread(task.result().abort) + raise + + async def _push_s3_credentials(base_url: str, card: dict, timeout_seconds: float) -> None: """Best-effort: push the object store's shared AWS creds + bucket to a service advertising add_s3_credentials, if the store shares any. Never raises (get_data falls back). @@ -467,28 +513,55 @@ async def _export_environment_to_file( ``tmp_path``, or the object url a FilePart names when the service uploaded the bundle itself (nothing written to ``tmp_path`` then — export_one registers it). - v1 ``get_data``, handed a namespace grant when the service takes the data-objects - form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` → - object url (no download) / base64 ``bytes`` / streamed relative-or-http ``uri`` → - ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON - guard. Errors are caught by the caller (fails just this service).""" + v1 ``get_data``, handed an object write when the service takes parts grants, else a + namespace grant when it takes the data-objects form: an upload through either → its + object url; ``DataPart`` → json; ``FilePart`` → object url (no download) / base64 + ``bytes`` / streamed relative-or-http ``uri`` → ".zip"; else legacy + ``GET /export-state`` streamed to disk with a first-byte JSON guard. Logs which way the + export came and how long it took. Errors are caught by the caller (fails just this + service).""" from agent_env.env import legacy_protocol from agentenv_protocol import client as protocol_v1 + started = time.monotonic() + + def exported(tier: str, result: str, detail: str = "") -> str: + logger.info( + "snapshot_env: %s exported through %s%s in %.1fs", + environment_name, tier, detail, time.monotonic() - started, + ) + return result + base_url = await legacy_protocol.v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is not None: card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) - upload = await asyncio.to_thread( - _snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None) - ) - # Push S3 creds (no-op unless the service advertises the extension), with a grant too: a - # bundle the grant cannot hold can still be uploaded with them. - await _push_s3_credentials(card_base, card or {}, timeout_seconds) - grant = {} if upload is None else {"write_namespace": upload.grant} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) - part = resp.parts[0] if resp.parts else None + sandbox_type = getattr(deployed, "sandbox_type", None) + write = None + try: + write = await _finish_in_thread(_snapshot_write, card or {}, environment_name, timeout_seconds, sandbox_type) + upload = None if write is not None else await asyncio.to_thread( + _snapshot_upload, card or {}, timeout_seconds, sandbox_type + ) + # Push S3 creds (no-op unless the service advertises the extension) beside either grant: a + # bundle the grant cannot hold, or an upload through it that fails, can still use them. + await _push_s3_credentials(card_base, card or {}, timeout_seconds) + if write is not None: + grant = {"write_object": write.grant} + else: + grant = {} if upload is None else {"write_namespace": upload.grant} + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) + part = resp.parts[0] if resp.parts else None + if write is not None and (uploaded := uploaded_object(part)) is not None: + await _finish_in_thread(write.complete, uploaded) + if not isinstance(grant := write.grant.write, HttpPartsPutGrant): + return exported("object", write.object_url, f" ({uploaded.size_bytes} bytes)") + parts = len(part_ranges(uploaded.size_bytes, grant.part_bytes)) + return exported("parts", write.object_url, f" ({uploaded.size_bytes} bytes in {parts} parts)") + finally: + if write is not None: + await asyncio.to_thread(write.abort) if upload is not None and (path := uploaded_object_path(part)) is not None: - return upload.object_url(path) + return exported("namespace", upload.object_url(path)) if isinstance(part, FilePart): raw = getattr(part.file, "bytes", None) if raw is not None: @@ -507,7 +580,7 @@ async def _export_environment_to_file( # last segment carries the shape (e.g. gdrive.zip). scheme = urlparse(uri).scheme if scheme and scheme not in ("http", "https"): - return uri + return exported("service upload", uri) # A relative uri (e.g. "export-snapshot") means "stream it from # my service endpoint" — resolve against the service base_url so # multi-GB bundles never ride inline in the JSON-RPC response. @@ -525,11 +598,12 @@ async def _export_environment_to_file( with open(tmp_path, "wb") as f: async for chunk in resp_file.aiter_bytes(): f.write(chunk) - return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + suffix = os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + return exported("inline" if raw is not None else "stream", suffix) state = part.data if part is not None else {} with open(tmp_path, "w") as f: json.dump(state, f, default=str) - return ".json" + return exported("data", ".json") legacy_url = f"{legacy_protocol.environment_base_url(gateway_url, environment_name, mcp=True)}/export-state" async with httpx.AsyncClient() as client: @@ -550,7 +624,7 @@ async def _export_environment_to_file( ) seen_first_bytes = True f.write(chunk) - return ".json" + return exported("export-state", ".json") @classmethod async def snapshot_env_state( diff --git a/tst/integration/env/envs/multi_env_snapshot_local_test.py b/tst/integration/env/envs/multi_env_snapshot_local_test.py new file mode 100644 index 000000000..4b2ef71f6 --- /dev/null +++ b/tst/integration/env/envs/multi_env_snapshot_local_test.py @@ -0,0 +1,63 @@ +"""snapshot_env on the local backend, end to end, through a parts grant. + +A MultiEnv's items server takes ``write_object``, so the snapshot hands it an object of its own on the local store; the +server uploads its bundle from its container, the snapshot registers that object as the server's export, and loading the +snapshot back restores the state it held. Requires a docker daemon; spins up a throwaway ``registry:2`` and skips if it +can't start. +""" + +import io +import json +import logging +import uuid +import zipfile + +import pytest +from agentenv_protocol import DataPart +from agentenv_protocol import client as protocol_v1 + +from agent_env.artifact import EnvironmentUniverseArtifact +from agent_env.env import legacy_protocol +from agent_env.env.env import Env +from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.env.envs.multi_env import MultiEnv +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps import snapshot_env as snapshot_mod +from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep +from tst.integration.env.envs.server_env_local_test import _put_items_image +from tst.integration.store.local_parallel_deploy_test import _docker, _seed_bootstrap_envs, local_stack as local_stack + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + + +@pytest.mark.asyncio +async def test_a_snapshot_of_a_server_that_takes_parts_registers_its_upload_and_loads_back(local_stack, caplog): + _seed_bootstrap_envs() + uid = uuid.uuid4().hex[:8] + items = MCPServerEnv.put(id=f"snap-items-{uid}", environment_name="items", + docker_image_artifact=_put_items_image(f"snap-items-{uid}")) + env = MultiEnv.put(id=f"snap-{uid}", mcp_server_envs=[items], name="suite") + + deployed = await env.deploy(sandbox_type="local", ttl_seconds=900) + work_dir = env._sandbox.work_dir + try: + base_url = await legacy_protocol.v1_base_url(deployed, deployed.gateway_url, "items") + await protocol_v1.add_data(base_url, [DataPart(data={"items": ["snap-x", "snap-y"]})]) + + step = SnapshotEnvTaskStep(id="snap", version=None, snapshot_id=f"snap-universe-{uid}") + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + context = await step.execute(TaskStepContext(deployed_envs=[deployed])) + + assert "items exported through parts" in caplog.text + universe = EnvironmentUniverseArtifact.get(context.metadata["env_snapshotted_universes"]["snap"]["id"]) + (bundle,) = universe.get_file_artifacts().values() + assert "/agentenv-snapshots/" in bundle.object_url and bundle.object_url.endswith("/items.zip") + assert (bundle.filename, bundle.content_type) == ("items.zip", "application/zip") + with zipfile.ZipFile(io.BytesIO(bundle.load())) as zf: + assert json.loads(zf.read("items.json")) == {"items": ["snap-x", "snap-y"]} + + await protocol_v1.add_data(base_url, [DataPart(data={"items": ["after-the-snapshot"]})]) + await (await Env.from_instance_id(deployed.instance_id)).load_environment_universe_artifact(universe) + assert (await protocol_v1.get_data(base_url)).model_dump()["parts"][0]["data"] == {"items": ["snap-x", "snap-y"]} + finally: + _docker("compose", "--project-directory", str(work_dir), "down", "-v", "--remove-orphans") diff --git a/tst/integration/store/local_parallel_deploy_test.py b/tst/integration/store/local_parallel_deploy_test.py index 4964e80c4..acada97df 100644 --- a/tst/integration/store/local_parallel_deploy_test.py +++ b/tst/integration/store/local_parallel_deploy_test.py @@ -24,6 +24,7 @@ from agent_env.config import configure, reset_config, set_image_store from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.store.image_store import LocalRegistryImageStore +from agent_env.store.object_store.local.grant_server import grant_server from agent_env.task import Task from agent_env.task.store import get_task_instance_store from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep @@ -90,6 +91,9 @@ def local_stack_at(monkeypatch: pytest.MonkeyPatch, path: Path): finally: for cid in _docker("ps", "-aq", "--filter", "name=agent-env-local-").stdout.split(): _docker("rm", "-f", cid) + # Grants handed out here came from this process's grant server, whose certificate is from this test's state + # root; a later test's containers trust another root's CA, so it must start afresh. + grant_server(None, None).close() reset_artifact_store() reset_config() _docker("rm", "-f", reg_name) diff --git a/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py b/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py index 3ee8f8da2..0b11024a9 100644 --- a/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py +++ b/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py @@ -41,6 +41,26 @@ async def test_snapshot_reload_keeps_the_env_name(declared, deployed_as, expecte assert compose.call_args.kwargs["host_port"] == env._sandbox.host_port +@pytest.mark.asyncio +async def test_snapshot_reload_keeps_its_services_trusting_the_local_transfer_ca(): + env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") + env._sandbox = MagicMock(mode="vm", load_docker_images=AsyncMock(), host_ips=("127.0.0.1",)) + db_image = MagicMock(spec=DockerImageArtifact, image_name="snap:1") + snapshot = MagicMock(db_image_artifact_id="snap", db_image_artifact_version=1) + + with patch("agent_env.artifact.Artifact.get", return_value=db_image), \ + patch("agent_env.env.env.Env.get", return_value=MagicMock()), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id", default_service_db_env_id="db-id")), \ + patch("agent_env.providers.env_state.LocalPostgresStateProvider"), \ + patch.object(EnvironmentGatewayProvider, "_stage_trust_files", AsyncMock(return_value="./trust")) as stage, \ + patch.object(EnvironmentGatewayProvider, "create_docker_compose", side_effect=_Stop) as compose: + with pytest.raises(_Stop): + await env._load_from_snapshot(snapshot) + + stage.assert_awaited_once_with(env._sandbox) + assert compose.call_args.kwargs["trust_files"] == "./trust" + + @pytest.mark.asyncio async def test_snapshot_reload_waits_for_the_gateway_where_it_is_published(): env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 1f882a41c..466a0a809 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -7,9 +7,11 @@ import dataclasses import json import logging +import os import re import shlex import socket +import stat import subprocess from datetime import datetime from types import SimpleNamespace @@ -41,6 +43,7 @@ ) from agent_env.providers.env_providers.env_provider import DeployRefused from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider +from agent_env.providers.sandbox_providers.local_docker.provider import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV, LocalSandbox from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.sandbox_providers.sandbox_provider import Accepts @@ -1098,6 +1101,111 @@ def test_compose_reads_a_servers_card_on_its_card_port_and_its_mcp_on_the_mcp_po assert "18765" in services["website-browser"]["healthcheck"]["test"][1] +@pytest.mark.parametrize("trust_files", [None, "./agentenv-trust"], ids=["no-local-grants", "local-grants"]) +def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(trust_files): + """Services on a local compose stack reach the local grant server over TLS, so they mount its CA's trust files + and carry the sandbox's host mappings, which on Linux are how host.docker.internal resolves.""" + from agent_env.env.envs.service_db import ServiceDBConfig + + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack"), + MCPServerConfig(image="mcp-items", environment_name="items")], + gateway_image="agent-gateway", + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + extra_hosts=("host.docker.internal:host-gateway",), + trust_files=trust_files, + ) + + services = yaml.safe_load(compose)["services"] + for name in ("slack", "items"): + env = dict(entry.split("=", 1) for entry in services[name]["environment"]) + trusted = {key: env.get(key) for key in LOCAL_TRUST_ENV} + assert trusted == (LOCAL_TRUST_ENV if trust_files else dict.fromkeys(LOCAL_TRUST_ENV)) + assert services[name].get("volumes") == ([f"{trust_files}:{LOCAL_TRUST_DIR}:ro"] if trust_files else None) + assert services[name].get("extra_hosts") == (["host.docker.internal:host-gateway"] if trust_files else None) + + +def test_a_servers_own_variables_win_over_the_local_transfer_ca(): + from agent_env.env.envs.service_db import ServiceDBConfig + + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-items", environment_name="items", + extra_env_vars={"SSL_CERT_FILE": "/srv/own-roots.pem"})], + gateway_image="agent-gateway", + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + trust_files="./trust", + ) + + keys = [entry.split("=", 1)[0] for entry in yaml.safe_load(compose)["services"]["items"]["environment"]] + env = dict(entry.split("=", 1) for entry in yaml.safe_load(compose)["services"]["items"]["environment"]) + assert len(keys) == len(set(keys)) + assert env["SSL_CERT_FILE"] == "/srv/own-roots.pem" + assert env["NODE_EXTRA_CA_CERTS"] == LOCAL_TRUST_ENV["NODE_EXTRA_CA_CERTS"] + + +class _HostSandbox: + """Records what a deploy writes on its host.""" + + def __init__(self, on_this_machine: bool) -> None: + self.ON_THIS_MACHINE = on_this_machine + self.scripts: list[str] = [] + self.files: dict[str, bytes] = {} + + async def exec_script(self, script: str, **_) -> str: + self.scripts.append(script) + return "" + + async def write_host_file(self, data: bytes, vm_path: str) -> None: + self.files[vm_path] = data + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("on_this_machine", "local_grants", "staged"), [ + (True, True, "./trust"), + (True, False, None), + (False, True, None), +], ids=["on-this-machine", "no-local-grants", "elsewhere"]) +async def test_a_deploy_on_this_machine_writes_the_trust_files_beside_its_compose( + tmp_path, monkeypatch, on_this_machine, local_grants, staged +): + """Any sandbox on this machine, a local VM's as well as local Docker's, gets them through its own host, which its + compose mounts from; a sandbox elsewhere is handed no local grants, so it gets none.""" + (tmp_path / "ca.pem").write_bytes(b"ca") + (tmp_path / "ca-bundle.pem").write_bytes(b"roots+ca") + monkeypatch.setattr(env_gateway_provider, "local_grant_trust", lambda: tmp_path if local_grants else None) + sandbox = _HostSandbox(on_this_machine) + + assert await EnvironmentGatewayProvider._stage_trust_files(sandbox) == staged + + expected = {"/app/trust/ca-bundle.pem": b"roots+ca", "/app/trust/ca.pem": b"ca"} if staged else {} + assert sandbox.files == expected + assert sandbox.scripts == (["mkdir -p /app/trust", "chmod 755 /app/trust && chmod 644 /app/trust/*"] if staged else []) + + +@pytest.mark.asyncio +async def test_a_local_sandbox_keeps_the_trust_files_in_its_work_dir(tmp_path, monkeypatch): + """The CA's own folder may sit outside what Docker shares; the work dir, where the compose runs, is shared. The bundle + of public roots and the CA is a few hundred KB, so it is written in chunks, and a service running as any user reads + the files, whatever the deploy's umask.""" + (tmp_path / "ca").mkdir() + trust = {"ca.pem": b"ca", "ca-bundle.pem": b"r" * (300 * 1024)} + for name, content in trust.items(): + (tmp_path / "ca" / name).write_bytes(content) + monkeypatch.setattr(env_gateway_provider, "local_grant_trust", lambda: tmp_path / "ca") + + umask = os.umask(0o077) + try: + assert await EnvironmentGatewayProvider._stage_trust_files(LocalSandbox(work_dir=tmp_path / "work")) == "./trust" + finally: + os.umask(umask) + staged = tmp_path / "work" / "trust" + assert {p.name: p.read_bytes() for p in staged.iterdir()} == trust + assert {p.name: stat.S_IMODE(p.stat().st_mode) for p in [staged, *staged.iterdir()]} == { + "trust": 0o755, "ca.pem": 0o644, "ca-bundle.pem": 0o644} + + # --- sidecar rendering -------------------------------------------------------- class _FakeImageArtifact: @@ -1355,6 +1463,41 @@ def _fake_env_get(env_id, *a, **k): assert browser.card_port == 18766 # its card is the recorder's +@pytest.mark.asyncio +async def test_deploy_via_vm_mounts_the_trust_files_it_staged_on_its_sandbox(): + from agent_env.env.gateway import GatewayMode + from agent_env.providers.env_state import DatabaseStateProvider + + gp = EnvironmentGatewayProvider() + gp._state_provider = MagicMock(spec=DatabaseStateProvider, prepare=AsyncMock()) + gp._state_instance = MagicMock() + sandbox = AsyncMock() + sandbox_provider = MagicMock(create_vm=AsyncMock(return_value=sandbox)) + + class _Stop(Exception): + """Abort once the compose is rendered.""" + + with patch("agent_env.env.env.Env.get", side_effect=lambda env_id, *a, **k: MagicMock(docker_image_artifact=_loadable(env_id))), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id")), \ + patch.object(EnvironmentGatewayProvider, "_stage_trust_files", AsyncMock(return_value="./trust")) as stage, \ + patch.object(gp, "create_docker_compose", side_effect=_Stop) as compose: + with pytest.raises(_Stop): + await gp._deploy_via_vm( + sandbox_provider=sandbox_provider, + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], + mcp_server_images=[_loadable("mcp-slack")], + gateway_port=18765, + website_configs=None, + website_images=None, + gateway_mode=GatewayMode.PERFORMANCE, + ttl_seconds=60, + disk_size_gb=10, + ) + + stage.assert_awaited_once_with(sandbox) + assert compose.call_args.kwargs["trust_files"] == "./trust" + + @pytest.mark.asyncio async def test_deploy_via_vm_side_loads_sidecar_image(): """A sidecar's image_artifact is side-loaded onto the VM via load_docker_images alongside diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py index 2e77e3540..632b79953 100644 --- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py @@ -469,7 +469,7 @@ async def test_write_file_from_text_large_chunks_reconstruct_exactly(): # The reconstructed base64 must equal the original, byte-for-byte. assert _b64_from_chunk_scripts(sandbox.scripts, b64_path) == expected_b64 # And it's decoded into the real path before the docker cp. - assert any(f"base64 -d {b64_path} > " in s for s in sandbox.scripts) + assert any(f"base64 -d < {b64_path} > " in s for s in sandbox.scripts) class _CleanupFailingSandbox(VmSandbox): @@ -715,7 +715,7 @@ async def test_write_host_file_writes_on_the_host_in_bounded_chunks(): assert not any("docker" in s for s in sandbox.scripts) b64_path = "/tmp/agentenv_run_code/input.json.b64" assert _b64_from_chunk_scripts(sandbox.scripts, b64_path) == base64.b64encode(data).decode() - assert sandbox.scripts[-1] == f"base64 -d {b64_path} > /tmp/agentenv_run_code/input.json && rm -f {b64_path}" + assert sandbox.scripts[-1] == f"base64 -d < {b64_path} > /tmp/agentenv_run_code/input.json && rm -f {b64_path}" class _ArgsRecorder(VmSandbox): diff --git a/tst/unit/task_step/snapshot_env_object_write_test.py b/tst/unit/task_step/snapshot_env_object_write_test.py new file mode 100644 index 000000000..7b0ca9d65 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_object_write_test.py @@ -0,0 +1,272 @@ +"""snapshot_env hands a service that takes parts grants an object of its own to upload its bundle to, makes +the object from what it reports, and registers it there; every other reply aborts the write.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import logging +import threading +from urllib.parse import unquote, urlsplit + +import httpx +import pytest +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, TRANSFERS_EXTENSION_URI, uploaded_object_part +from agentenv_protocol.transfers import Uploaded + +from agent_env.artifact import EnvironmentUniverseArtifact +from agent_env.config import set_object_store +from agent_env.store.base import GrantUnavailableError +from agent_env.task_step.task_steps import snapshot_env as snapshot_mod +from agent_env.task_step.task_steps.snapshot_env import ENV_SNAPSHOT_MAX_BYTES, SnapshotEnvTaskStep +from tst.unit.task_step.test_snapshot_env import _CARD, _FakeMultiEnv, _card, _carded +from tst.util.granting_object_store import GrantingObjectStore + +_PARTS = {"uri": TRANSFERS_EXTENSION_URI, "params": {"write": ["http-post-policy", "http-put", "http-put-parts"]}} +_DATA_OBJECTS = {"uri": DATA_OBJECTS_EXTENSION_URI} +_S3_CREDENTIALS = {"uri": "urn:agentenv:add-s3-credentials/v1", + "params": {"endpoint": "/agentenv/ext/add_s3_credentials", "methods": {"add_s3_credentials": {"method": "POST"}}}} +_BUNDLE = b"PK\x03\x04bundle" + + +class _UngrantingStore(GrantingObjectStore): + def begin_write(self, object_url, **kwargs): + raise GrantUnavailableError("this store issues no object writes") + + +class _WatchedStore(GrantingObjectStore): + """Records how each write it begins ends.""" + + def __init__(self, root): + super().__init__(root) + self.ended: list[str] = [] + + def begin_write(self, object_url, **kwargs): + write = super().begin_write(object_url, **kwargs) + complete, abort, ended = write.complete, write.abort, self.ended + write.complete = lambda uploaded: (complete(uploaded), ended.append("completed"))[0] + write.abort = lambda: (ended.append("aborted") if not write._settled else None, abort())[1] + return write + + +class _PutOnlyStore(_WatchedStore): + """Answers a service that takes parts with one PUT, as a store may.""" + + def begin_write(self, object_url, *, kinds, **kwargs): + return super().begin_write(object_url, kinds=kinds - {"http-put-parts"}, **kwargs) + + +class _GatedStore(_WatchedStore): + """Holds the export at ``at`` (beginning the write, pushing creds, completing it) until released.""" + + def __init__(self, root, at): + super().__init__(root) + self.at, self.reached, self.release = at, threading.Event(), threading.Event() + + def hold(self, at): + if at == self.at: + self.reached.set() + self.release.wait(5) + + def begin_write(self, object_url, **kwargs): + self.hold("begin") + write = super().begin_write(object_url, **kwargs) + complete = write.complete + write.complete = lambda uploaded: (self.hold("complete"), complete(uploaded))[1] + return write + + +@pytest.fixture +def store(local_stores, tmp_path): + granting = _WatchedStore(str(tmp_path / "objects")) + set_object_store(granting) + return granting + + +@pytest.fixture +def pushes(monkeypatch): + pushed: list[str] = [] + + async def push(base_url, card, timeout_seconds): + pushed.append(base_url) + + monkeypatch.setattr(snapshot_mod, "_push_s3_credentials", push) + return pushed + + +async def _record(children: dict, *, sandbox_type: str | None = "local"): + return dataclasses.replace(_carded(await _card(children)), sandbox_type=sandbox_type) + + +def _serve(monkeypatch, answers: dict): + """Answer each service's data/get with ``answers[service](params)``, recording the params sent.""" + sent, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + service = next(seg for seg in request.url.path.split("/") if seg.startswith("mcp-")).removeprefix("mcp-") + params = json.loads(request.content)["params"] + sent.append((service, params)) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": answers[service](params)}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return sent + + +def _key(params) -> str: + return unquote(urlsplit(params["write_object"]["write"]["urls"][0]).path.lstrip("/")) + + +def _uploading(store, *, reported: int | None = None): + """A service that PUTs its bundle to the object it is handed and says what it sent.""" + def answer(params): + store.put(_key(params), _BUNDLE, content_type="application/zip") + uploaded = Uploaded(size_bytes=len(_BUNDLE) if reported is None else reported) + part = uploaded_object_part(uploaded, name="gdrive.zip", mime_type="application/zip") + return {"parts": [part.model_dump(mode="json", exclude_none=True)]} + return answer + + +def _data(params): + return {"parts": [{"kind": "data", "data": {"rows": 1}}]} + + +@pytest.mark.asyncio +async def test_a_service_taking_parts_uploads_to_an_object_of_its_own_and_is_registered_there( + store, pushes, monkeypatch, caplog +): + record = await _record({"mcp-gdrive": [_PARTS, _DATA_OBJECTS, _S3_CREDENTIALS]}) + sent = _serve(monkeypatch, {"gdrive": _uploading(store)}) + + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + result = await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive"]), gateway_url="https://gw-1", snapshot_id="snap-o", deployed=record, + ) + + ((_, params),) = sent + assert list(params) == ["write_object"] + grant = params["write_object"] + assert (grant["media_type"], grant["max_bytes"], grant["write"]["kind"]) == ("application/zip", ENV_SNAPSHOT_MAX_BYTES, "http-put-parts") + key = _key(params) + assert key.startswith("agentenv-snapshots/") and key.endswith("/gdrive.zip") + assert pushes == [_CARD] + assert store.ended == ["completed"] + (bundle,) = EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id).get_file_artifacts().values() + assert bundle.object_url == store.object_url(key) + assert (bundle.filename, bundle.content_type, bundle.load()) == ("gdrive.zip", "application/zip", _BUNDLE) + assert f"gdrive exported through parts ({len(_BUNDLE)} bytes in 1 parts)" in caplog.text + + +@pytest.mark.asyncio +async def test_a_service_taking_parts_that_is_handed_one_put_is_registered_there_too( + local_stores, tmp_path, pushes, monkeypatch, caplog +): + store = _PutOnlyStore(str(tmp_path / "objects")) + set_object_store(store) + record = await _record({"mcp-gdrive": [_PARTS]}) + + def answer(params): + grant = params["write_object"]["write"] + assert grant["kind"] == "http-put" + store.put(unquote(urlsplit(grant["url"]).path.lstrip("/")), _BUNDLE, content_type="application/zip") + part = uploaded_object_part(Uploaded(size_bytes=len(_BUNDLE)), name="gdrive.zip", mime_type="application/zip") + return {"parts": [part.model_dump(mode="json", exclude_none=True)]} + + _serve(monkeypatch, {"gdrive": answer}) + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + result = await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive"]), gateway_url="https://gw-1", snapshot_id="snap-put", deployed=record, + ) + + assert store.ended == ["completed"] + (bundle,) = EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id).get_file_artifacts().values() + assert bundle.load() == _BUNDLE + assert f"gdrive exported through object ({len(_BUNDLE)} bytes)" in caplog.text + + +@pytest.mark.asyncio +async def test_each_object_write_gets_an_object_of_its_own(store, pushes, monkeypatch): + record = await _record({"mcp-gdrive": [_PARTS]}) + sent = _serve(monkeypatch, {"gdrive": _uploading(store)}) + + for _ in range(2): + await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive"]), gateway_url="https://gw-1", snapshot_id="snap-1", deployed=record, + ) + + first, second = (_key(params) for _, params in sent) + assert first != second + + +@pytest.mark.asyncio +async def test_a_reply_that_names_no_object_upload_aborts_the_write_and_exports_as_before( + store, pushes, monkeypatch, tmp_path +): + record = await _record({"mcp-gdrive": [_PARTS]}) + _serve(monkeypatch, {"gdrive": _data}) + + out = tmp_path / "export-tmp" + suffix = await SnapshotEnvTaskStep._export_environment_to_file("https://gw-1", "gdrive", str(out), 30, record) + + assert (suffix, json.loads(out.read_text())) == (".json", {"rows": 1}) + assert store.ended == ["aborted"] + + +@pytest.mark.asyncio +async def test_an_upload_the_store_cannot_make_into_the_object_fails_only_that_service(store, pushes, monkeypatch): + record = await _record({"mcp-gdrive": [_PARTS], "mcp-slack": []}) + _serve(monkeypatch, {"gdrive": _uploading(store, reported=len(_BUNDLE) + 1), "slack": _data}) + + with pytest.raises(RuntimeError, match=r"1/2 services were not exported.*gdrive"): + await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive", "slack"]), gateway_url="https://gw-1", snapshot_id="snap-f", deployed=record, + ) + assert store.ended == ["aborted"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("extensions", "store_cls", "sandbox_type", "params"), [ + ([_DATA_OBJECTS], _WatchedStore, "local", "write_namespace"), + ([_PARTS, _DATA_OBJECTS], _WatchedStore, "modal", None), + ([_PARTS, _DATA_OBJECTS], _UngrantingStore, "local", "write_namespace"), +], ids=["parts-not-declared", "grants-do-not-reach-the-sandbox", "store-issues-no-object-write"]) +async def test_otherwise_the_export_is_offered_as_before( + local_stores, tmp_path, pushes, monkeypatch, caplog, extensions, store_cls, sandbox_type, params +): + set_object_store(store_cls(str(tmp_path / "objects"))) + record = await _record({"mcp-gdrive": extensions}, sandbox_type=sandbox_type) + sent = _serve(monkeypatch, {"gdrive": _data}) + + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + suffix = await SnapshotEnvTaskStep._export_environment_to_file( + "https://gw-1", "gdrive", str(tmp_path / "export-tmp"), 30, record + ) + + assert suffix == ".json" + ((_, sent_params),) = sent + assert list(sent_params) == ([] if params is None else [params]) + assert pushes == [_CARD] + assert "gdrive exported through data in" in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("at", "ended"), [("begin", ["aborted"]), ("push", ["aborted"]), ("complete", ["completed"])]) +async def test_an_export_cancelled_midway_leaves_no_write_open_or_aborted_mid_completion( + local_stores, tmp_path, monkeypatch, at, ended +): + store = _GatedStore(str(tmp_path / "objects"), at) + set_object_store(store) + monkeypatch.setattr(snapshot_mod, "_push_s3_credentials", lambda *_: asyncio.to_thread(store.hold, "push")) + record = await _record({"mcp-gdrive": [_PARTS]}) + _serve(monkeypatch, {"gdrive": _uploading(store)}) + + export = asyncio.ensure_future( + SnapshotEnvTaskStep._export_environment_to_file("https://gw-1", "gdrive", str(tmp_path / "out"), 30, record) + ) + assert await asyncio.to_thread(store.reached.wait, 5) + export.cancel() + await asyncio.sleep(0.05) + store.release.set() + with pytest.raises(asyncio.CancelledError): + await export + assert store.ended == ended