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 3e96fd46..825a8e49 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -33,6 +33,7 @@ GATEWAY_SERVICE_NAME, ) from agent_env.providers.env_providers.env_provider import _SandboxEnvironmentProvider, _size_kwargs, _tool_names +from agent_env.providers.sandbox_providers.local_sandbox 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, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SandboxProvider @@ -215,8 +216,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 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. @@ -361,6 +366,11 @@ def create_docker_compose( if server.extra_env_vars: for key, value in server.extra_env_vars.items(): lines.append(f" - {key}={value}") + if trust_files is not None: + lines.extend(f" - {key}={value}" for key, value in LOCAL_TRUST_ENV.items()) + 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()\\""]', @@ -737,6 +747,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, + trust_files=await asyncio.to_thread(sandbox.stage_local_trust) if isinstance(sandbox, LocalSandbox) else None, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 7289aa16..e9870a32 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -74,11 +74,13 @@ def _runs_in_container(cmd: list[str]) -> bool: # Where a container finds the local transfer CA's trust files, and the variables that point TLS clients at them: # SSL_CERT_FILE replaces a client's roots, so it gets the public roots plus the CA; NODE_EXTRA_CA_CERTS adds. -_TRUST_DIR = "/etc/agentenv" +LOCAL_TRUST_DIR = "/etc/agentenv" +# Where in a sandbox's work dir a compose there finds a copy of those trust files to mount. +_STAGED_TRUST_DIR = "agentenv-trust" LOCAL_TRUST_ENV = { - "SSL_CERT_FILE": f"{_TRUST_DIR}/ca-bundle.pem", - "REQUESTS_CA_BUNDLE": f"{_TRUST_DIR}/ca-bundle.pem", - "NODE_EXTRA_CA_CERTS": f"{_TRUST_DIR}/ca.pem", + "SSL_CERT_FILE": f"{LOCAL_TRUST_DIR}/ca-bundle.pem", + "REQUESTS_CA_BUNDLE": f"{LOCAL_TRUST_DIR}/ca-bundle.pem", + "NODE_EXTRA_CA_CERTS": f"{LOCAL_TRUST_DIR}/ca.pem", } # Marker dropped in the work dir when this sandbox runs a container, so a later get_sandbox() @@ -204,6 +206,15 @@ def owns_container(self) -> bool: def work_dir(self) -> Path: return self._work_dir + def stage_local_trust(self) -> str | None: + """Copy the local transfer CA's trust files into the work dir, which Docker shares, for a compose there to + mount; their path relative to the work dir, or None where the store hands out no local grants.""" + trust_dir = local_grant_trust() + if trust_dir is None: + return None + shutil.copytree(trust_dir, self._work_dir / _STAGED_TRUST_DIR, dirs_exist_ok=True) + return f"./{_STAGED_TRUST_DIR}" + def _rewrite_app_arg(self, arg: str) -> str: if arg == "/app" or arg.startswith(("/app/", "/app:")): return str(self._work_dir) + arg[4:] @@ -460,7 +471,7 @@ async def start_trusting(sandbox: VmSandbox, container: str, trust_dir: Path) -> """Copy ``trust_dir`` into the created ``container`` where ``LOCAL_TRUST_ENV`` points, then start it. A container that cannot be given the files or started is removed, so its name is free for the next attempt.""" try: - await asyncio.to_thread(_copy_into_container, trust_dir, container, _TRUST_DIR) + await asyncio.to_thread(_copy_into_container, trust_dir, container, LOCAL_TRUST_DIR) await sandbox.exec_script(f"docker start {shlex.quote(container)} > /dev/null") except Exception: await sandbox.exec_script(f"docker rm -f {shlex.quote(container)} >/dev/null 2>&1 || true") 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 c78c5042..a84605ad 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 -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, write_kinds +from agentenv_protocol.transfers import WriteNamespaceGrant, part_ranges from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant from agent_env.env.gateway.constants import 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 @@ -46,6 +47,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" @@ -87,6 +91,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). @@ -412,28 +458,53 @@ 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) + parts = len(part_ranges(uploaded.size_bytes, write.grant.write.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: @@ -452,7 +523,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. @@ -470,11 +541,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: @@ -495,7 +567,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 00000000..f5517b21 --- /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 # noqa: F401 (fixture) + +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 e31738d3..ea01a82e 100644 --- a/tst/integration/store/local_parallel_deploy_test.py +++ b/tst/integration/store/local_parallel_deploy_test.py @@ -23,6 +23,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 @@ -81,6 +82,9 @@ def local_stack(monkeypatch, tmp_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/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 9600c5f1..15a30b8f 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -37,6 +37,7 @@ _redact_compose_secrets, ) from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider +from agent_env.providers.sandbox_providers.local_sandbox import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.env_state import LocalPostgresStateProvider @@ -817,6 +818,31 @@ def test_compose_maps_extra_hosts_into_the_gateway(extra_hosts): assert mapped == ({GATEWAY_SERVICE_NAME: list(extra_hosts)} if extra_hosts else {}) +@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) + + # --- sidecar rendering -------------------------------------------------------- class _FakeImageArtifact: diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_test.py index f98ff7c8..372a7810 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_test.py @@ -1165,3 +1165,14 @@ def fail(*args): _HOST_IPS = ls._host_ips + + +def test_a_compose_on_a_local_sandbox_mounts_a_copy_of_the_trust_files_its_work_dir_holds(tmp_path, monkeypatch): + """The CA's own folder may sit outside what Docker shares; the work dir is kept on a shared path.""" + sandbox = LocalSandbox(work_dir=tmp_path / "work") + assert sandbox.stage_local_trust() == "./agentenv-trust" + staged = tmp_path / "work" / "agentenv-trust" + assert {p.name: p.read_bytes() for p in staged.iterdir()} == {p.name: p.read_bytes() for p in local_ca().trust_dir.iterdir()} + + monkeypatch.setattr(ls, "local_grant_trust", lambda: None) + assert LocalSandbox(work_dir=tmp_path / "other").stage_local_trust() is None 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 00000000..adbf93b0 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_object_write_test.py @@ -0,0 +1,238 @@ +"""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 _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_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