feat(snapshot_env): hand a service that takes parts grants an object of its own to upload to

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)}")
 

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")

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(

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")

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)