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

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

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

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

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(