fix(env): one grant-reach check, size-scaled load timeouts on container sandboxes, and JSON state for services that export a file

src/agent_env/a2a_agent/object_transfer.py

@@ -43,7 +43,7 @@
 from agent_env.config import get_config
 from agent_env.store.base import GrantUnavailableError
 from agent_env.store.object_store import DEFAULT_CONTENT_TYPE, ObjectStore
-from agent_env.store.object_store.object_store import readable_url
+from agent_env.store.object_store.object_store import issues_grants_to, readable_url
 from agent_env.store.object_store.local.grant_server import unreachable_hint
 
 if TYPE_CHECKING:
@@ -121,12 +121,7 @@ def choose_transfer(
     (None: unknown). A method without a declared request predates variant negotiation and takes
     the legacy form. None: the agent takes neither form.
     """
-    if (
-        objects is not None
-        and _accepts(method, objects)
-        and store.supports_transfer_grants
-        and store.grants_reach(sandbox_type)
-    ):
+    if objects is not None and _accepts(method, objects) and issues_grants_to(store, sandbox_type):
         return "objects"
     if legacy is not None and (
         method is None or "request" not in method or _accepts(method, legacy)
@@ -152,13 +147,14 @@ def _require_object_form(
     that reach the agent's sandbox, of ``sandbox_type``: agent-env moves objects no other way."""
     if not _accepts(method, _fields(model)):
         raise RuntimeError(f"{operation}: the agent does not advertise the object form")
+    if issues_grants_to(store, sandbox_type):
+        return
     if not store.supports_transfer_grants:
         raise RuntimeError(f"{operation}: the object store does not issue transfer grants")
-    if not store.grants_reach(sandbox_type):
-        raise RuntimeError(
-            f"{operation}: the object store's grants do not reach agents on the "
-            f"{sandbox_type or 'unknown'!r} sandbox provider"
-        )
+    raise RuntimeError(
+        f"{operation}: the object store's grants do not reach agents on the "
+        f"{sandbox_type or 'unknown'!r} sandbox provider"
+    )
 
 
 @dataclass(frozen=True)

src/agent_env/a2a_agent/staging.py

@@ -32,6 +32,7 @@
 from agent_env.config import get_config
 from agent_env.store.base import ObjectAlreadyExistsError
 from agent_env.store.object_store import DEFAULT_CONTENT_TYPE, ObjectMetadata, ObjectStore, UploadPolicy
+from agent_env.store.object_store.object_store import issues_grants_to
 
 logger = logging.getLogger(__name__)
 
@@ -76,7 +77,7 @@ def transfer_store(
     store that stages them through the agent when the agent serves staging, else ``store`` itself, so the
     call falls back to the forms that carry no grants."""
     endpoint = staging_endpoint(a2a_url, card)
-    if endpoint is None or (store.supports_transfer_grants and store.grants_reach(sandbox_type)):
+    if endpoint is None or issues_grants_to(store, sandbox_type):
         return store
     return _staged(store, endpoint, card)
 

src/agent_env/env/envs/mcp_server.py

@@ -2,16 +2,17 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
-import shlex
 import uuid
 from importlib.metadata import version as pkg_version
 from typing import TYPE_CHECKING, Callable, ClassVar, Optional
 
 from agentenv_protocol import FilePart, client as protocol_v1
 from agent_env.artifact import Artifact, DockerImageArtifact
 from agent_env.artifact.artifacts.docker_image import GitHubBuildResult, ProgressCallback, refuse_local_github_build
+from agent_env.config import get_config
 from agent_env.entity_refs import EntityRef
 from agent_env.env.env import Env, gateway_url_of
 from agent_env.store.ids import derive_id
@@ -155,11 +156,10 @@ async def load_environment_artifact(self, environment_artifact: EnvironmentArtif
         if base_url is not None:
             from agent_env.env.gateway.constants import data_plane_load_timeout_s
             container_path = await self._copy_artifact_into_container(file_artifact)
-            # Sized from the payload actually staged on the VM rather than a flat constant:
-            # a 3.5GB service and a 30KB one were previously given the same 600s, which
-            # killed healthy large loads. Measured after staging so the number reflects real
-            # bytes rather than whatever the artifact document claims.
-            timeout = data_plane_load_timeout_s(await self._staged_artifact_size(container_path))
+            # Sized from the payload rather than a flat constant: a 3.5GB service and a 30KB one
+            # were previously given the same 600s, which killed healthy large loads. The size is
+            # the store's for the object staging copied, not whatever the artifact document claims.
+            timeout = data_plane_load_timeout_s(await self._staged_artifact_size(file_artifact))
             # reset_data previously took the client default (30s), tighter than the load it
             # precedes -- dropping and recreating a large service's schemas can plausibly
             # exceed that. Give it the same budget as the load it belongs to.
@@ -204,26 +204,24 @@ async def load_artifact(self, artifact):
             return await self.load_environment_artifact(artifact)
         raise ValueError(f"{type(self).__name__} '{self.id}' cannot load artifact '{getattr(artifact, 'id', '?')}' of type '{getattr(artifact, 'type', '?')}' — expected a EnvironmentArtifact or EnvironmentUniverseArtifact")
 
-    async def _staged_artifact_size(self, container_path: str) -> int | None:
-        """Bytes of the staged payload, or None if it can't be measured.
+    async def _staged_artifact_size(self, file_artifact) -> int | None:
+        """Bytes of the payload staged for ``file_artifact``: its object's size as the store holds it, which
+        staging copies unchanged. None if the store can't say.
 
-        Best-effort on purpose: the size only picks a timeout, so a failed stat should fall
+        Best-effort on purpose: the size only picks a timeout, so a failed lookup should fall
         back to the floor (today's behaviour) rather than fail a load that is fine.
         """
-        if self._sandbox is None:
-            return None
+        object_url = file_artifact.object_url
         try:
-            container_id = await self._env_provider._get_container_id(self._sandbox, self.environment_name)
-            out = await self._sandbox.exec_script(
-                f"docker exec {container_id} stat -c %s {shlex.quote(container_path)}"
-            )
-            return int(out.strip())
+            store = get_config().get_object_store_at(object_url)
+            metadata = await asyncio.to_thread(store.get_object_metadata_at, object_url)
         except Exception as e:  # noqa: BLE001 -- a measurement, not a dependency
             logger.warning(
-                "[%s] could not stat staged payload %s (%s); using the floor timeout",
-                self.environment_name, container_path, e,
+                "[%s] could not size staged payload %s (%s); using the floor timeout",
+                self.environment_name, file_artifact.filename, e,
             )
             return None
+        return metadata.size if metadata is not None else None
 
     async def _copy_artifact_into_container(self, file_artifact) -> str:
         container_path = f"/data/{file_artifact.filename}"

src/agent_env/env/legacy_protocol.py

@@ -12,14 +12,15 @@
 from typing import TYPE_CHECKING, Optional
 
 import httpx
-from agentenv_protocol import RPC_PATH, client as protocol_v1
+from agentenv_protocol import RPC_PATH, DataPart, client as protocol_v1
 
 if TYPE_CHECKING:
     from agent_env.env.env import DeployedEnv
 
 logger = logging.getLogger(__name__)
 
 DEFAULT_MCP_MAX_RETRIES = 5
+EXPORT_STATE_TIMEOUT_SECONDS = 60
 
 
 def environment_base_url(gateway_url: Optional[str], environment_name: str, mcp: bool = True) -> str:
@@ -80,15 +81,37 @@ async def add_via_rest(base_url: str, file_path: str, timeout: int = 120, verify
         return result
 
 
-async def export_state(gateway_url: str, environment_name: str, timeout: int = 60, verify: bool = True) -> dict:
+async def export_state(
+    gateway_url: str, environment_name: str, timeout: int = EXPORT_STATE_TIMEOUT_SECONDS, verify: bool = True
+) -> dict:
     """GET ``/svc/mcp-{name}/export-state`` — the legacy state snapshot for one service."""
-    base_url = environment_base_url(gateway_url, environment_name, mcp=True)
+    return await _export_state_at(environment_base_url(gateway_url, environment_name, mcp=True), timeout, verify)
+
+
+async def _export_state_at(base_url: str, timeout: int, verify: bool) -> dict:
     async with httpx.AsyncClient(verify=verify) as client:
         response = await client.get(f"{base_url}/export-state", timeout=timeout)
         response.raise_for_status()
         return response.json()
 
 
+async def service_state(deployed: Optional[DeployedEnv], gateway_url: Optional[str], environment_name: str) -> dict:
+    """An MCP service's state as JSON: its v1 ``data/get`` answer when that is data, else ``GET /export-state``.
+
+    A service that exports a file from ``data/get`` (a bundle of its database, to reload from) still serves its state
+    as JSON at ``/export-state``.
+    """
+    base_url = await v1_base_url(deployed, gateway_url, environment_name, mcp=True)
+    if base_url is None:
+        return await export_state(gateway_url, environment_name)
+    response = await protocol_v1.get_data(base_url)
+    if not response.parts:
+        return {}
+    if isinstance(response.parts[0], DataPart):
+        return response.parts[0].data
+    return await _export_state_at(base_url, timeout=EXPORT_STATE_TIMEOUT_SECONDS, verify=True)
+
+
 async def reset_via_mcp_tool(
     mcp_url: str,
     environment_name: str,

src/agent_env/store/object_store/object_store.py

@@ -227,13 +227,19 @@ def issue_upload_policy(
         )
 
 
+def issues_grants_to(store: ObjectStore, sandbox_type: str | None) -> bool:
+    """Whether ``store`` can hand a transfer grant to a remote party on the ``sandbox_type`` sandbox
+    provider (None: unknown): it issues grants, and they reach that provider."""
+    return store.supports_transfer_grants and store.grants_reach(sandbox_type)
+
+
 def readable_url(store: ObjectStore, object_url: str, *, sandbox_type: str | None, expires_in: int) -> str | None:
     """An HTTPS URL that a remote party on the ``sandbox_type`` sandbox provider (None: unknown) can GET the
     object at ``object_url`` from: a read grant when the store's grants reach it, else a URL the store
     signs; None when the store offers neither. It lasts at least ``expires_in`` seconds and at least the
     store's grant lifetime, unless the store's signing credentials or limits end it sooner."""
     expires_in = max(expires_in, store.grant_lifetime_seconds)
-    if store.supports_transfer_grants and store.grants_reach(sandbox_type):
+    if issues_grants_to(store, sandbox_type):
         try:
             grant = store.issue_read_grant(object_url, expires_in=expires_in)
         except GrantUnavailableError:  # it cannot last that long; a signed URL may

src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py

@@ -215,19 +215,10 @@ def _create_file_artifact_universe(self, original_environment_artifacts: list, e
 
     @staticmethod
     async def _export_all(deployed: DeployedEnv, environment_names: list[str]) -> dict[str, dict]:
-        """Export each service: v1 ``data/get`` at the base the env card gives it, else legacy ``GET /export-state``."""
+        """Export each service's state as JSON (see ``legacy_protocol.service_state``)."""
         from agent_env.env import legacy_protocol
         from agent_env.env.env import gateway_url_of
-        from agentenv_protocol import client as protocol_v1
-        result = {}
-        for name in environment_names:
-            base_url = await legacy_protocol.v1_base_url(deployed, gateway_url_of(deployed), name)
-            if base_url is not None:
-                resp = await protocol_v1.get_data(base_url)
-                result[name] = resp.parts[0].data if resp.parts else {}
-            else:
-                result[name] = await legacy_protocol.export_state(gateway_url_of(deployed), name)
-        return result
+        return {name: await legacy_protocol.service_state(deployed, gateway_url_of(deployed), name) for name in environment_names}
 
     def _create_universe_artifact(self, original_environment_artifacts: list, export_data: dict[str, dict], env_version: int, universe_version: int) -> Any:
         """Create FileArtifact + EnvironmentArtifact per service, bundle into EnvironmentUniverseArtifact."""

src/agent_env/task_step/task_steps/snapshot_agent_state.py

@@ -144,7 +144,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
     async def _capture_universe_state(
         self, context: TaskStepContext, capture_prefix: str
     ) -> None:
-        """GET /svc/mcp-<name>/export-state for each service, upload, publish `s3://` refs."""
+        """Read each service's state as JSON, upload it, publish `s3://` refs."""
         from agent_env.artifact import EnvironmentUniverseArtifact
         from agent_env.config import get_config
 
@@ -177,16 +177,10 @@ async def _capture_universe_state(
 
         from agent_env.env import legacy_protocol
         from agent_env.env.env import gateway_url_of
-        from agentenv_protocol import client as protocol_v1
         urls: dict[str, str] = {}
         for name in environment_names:
             try:
-                base_url = await legacy_protocol.v1_base_url(deployed_env, gateway_url_of(deployed_env), name, mcp=True)
-                if base_url is not None:
-                    resp = await protocol_v1.get_data(base_url)
-                    state = resp.parts[0].data if resp.parts else {}
-                else:
-                    state = await legacy_protocol.export_state(gateway_url_of(deployed_env), name)
+                state = await legacy_protocol.service_state(deployed_env, gateway_url_of(deployed_env), name)
                 urls[name] = await asyncio.to_thread(
                     store.put, f"{key_prefix}services/{name}.json", json.dumps(state).encode(),
                     content_type="application/json", allow_overwrite=True,

src/agent_env/task_step/task_steps/snapshot_env.py

@@ -26,6 +26,7 @@
 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, S3ObjectStore
+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
 from agent_env.entity_refs import EntityRef, RefRole
@@ -67,7 +68,7 @@ def _snapshot_upload(card: dict, timeout_seconds: float, sandbox_type: Optional[
         return None
     config = get_config()
     store = config.get_object_store()
-    if not (store.supports_transfer_grants and store.grants_reach(sandbox_type)):
+    if not issues_grants_to(store, sandbox_type):
         return None
     cap = store.max_single_upload_bytes or ENV_SNAPSHOT_LIMITS.max_object_bytes
     limits = ObjectLimits(

tst/data/agentenv_mcp/server.py

@@ -1,13 +1,18 @@
-"""Minimal in-memory MCP server built on AgentEnvEnvironment, for integration testing."""
+"""Minimal in-memory MCP server built on AgentEnvEnvironment, for integration testing.
+
+It also serves its state as JSON at ``GET /export-state``, and with ``urn:agentenv:export-as-file/v1`` enabled it
+answers ``data/get`` with a file bundle of that state, as a service that exports its database does."""
+import base64
 import json
 import random
 from typing import Annotated
 from urllib.parse import urlparse
 
 import httpx
 from pydantic import Field
+from starlette.responses import JSONResponse
 
-from agentenv_protocol import AgentEnvEnvironment, DataPart, EnvironmentCapabilities, EnvironmentExtension, add_data, environment_card, extension, get_data, reset_data, tool
+from agentenv_protocol import AgentEnvEnvironment, DataPart, EnvironmentCapabilities, EnvironmentExtension, FilePart, add_data, environment_card, extension, get_data, reset_data, tool
 
 
 @environment_card(
@@ -39,7 +44,9 @@ def __init__(self) -> None:
         self.store: list = []
         self.errors: dict = {}
         self.env_get_time_url: str | None = None
+        self.export_as_file = False
         self.create_app()
+        self.mcp.custom_route("/export-state", methods=["GET"])(self._export_state)
         # list_items stays imperatively registered — @tool is additive; both styles coexist.
         self.mcp.tool(name="list_items")(self.list_items)
 
@@ -67,8 +74,19 @@ async def _add(self, parts: list) -> None:
 
     @get_data
     async def _state(self) -> list:
+        if self.export_as_file:
+            bundle = base64.b64encode(json.dumps({"items": self.store}).encode()).decode()
+            return [FilePart(file={"bytes": bundle, "name": "items.json", "mimeType": "application/json"})]
         return [DataPart(data={"items": self.store})]
 
+    async def _export_state(self, request) -> JSONResponse:
+        return JSONResponse({"items": self.store})
+
+    @extension(uri="urn:agentenv:export-as-file/v1", description="Answer data/get with a file bundle of the state.")
+    async def set_export_as_file(self, enabled: bool) -> dict:
+        self.export_as_file = enabled
+        return {"export_as_file": enabled}
+
     @extension(uri="urn:agentenv:set-errors/v1", description="Make a tool start raising at a given error rate.")
     async def set_errors(self, tool_name: str, error_rate: float) -> dict:
         self.errors[tool_name] = error_rate