From f62032fb43de7934fa43c1ad8d9ef03ae9e9eb53 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 09:48:13 -0700 Subject: [PATCH 1/7] Check whether a store's grants reach a sandbox in one place Every caller that hands a transfer grant on checked both supports_transfer_grants and grants_reach(sandbox_type), and a caller that checked only one half would pass most tests. issues_grants_to answers both, and choose_transfer, _require_object_form, transfer_store, readable_url and _snapshot_upload use it. _require_object_form keeps its two error messages. No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/a2a_agent/object_transfer.py | 20 ++++++++----------- src/agent_env/a2a_agent/staging.py | 3 ++- .../store/object_store/object_store.py | 8 +++++++- .../task_step/task_steps/snapshot_env.py | 3 ++- 4 files changed, 19 insertions(+), 15 deletions(-) diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index 63f76abe..26e46a46 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/src/agent_env/a2a_agent/object_transfer.py @@ -42,7 +42,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 logger = logging.getLogger(__name__) @@ -117,12 +117,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) @@ -148,13 +143,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) diff --git a/src/agent_env/a2a_agent/staging.py b/src/agent_env/a2a_agent/staging.py index cb392206..c7700fb9 100644 --- a/src/agent_env/a2a_agent/staging.py +++ b/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) diff --git a/src/agent_env/store/object_store/object_store.py b/src/agent_env/store/object_store/object_store.py index 071992ba..edab2be5 100644 --- a/src/agent_env/store/object_store/object_store.py +++ b/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 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 9708cfa7..d732af65 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/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( From 5af191520c5ca316733fed03bf87e0dd057cc404 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 09:48:13 -0700 Subject: [PATCH 2/7] Size a load's timeout from the staged payload on container sandboxes too _staged_artifact_size ran `docker exec ... stat` through exec_script, which only VM sandboxes have. On a container sandbox (Modal) the stat raised, so every load there fell back to the floor timeout however large its payload. A container sandbox is the env's own container, so stat the file in it directly; the VM path is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/envs/mcp_server.py | 14 +++++--- .../envs/test_mcp_server_container_mode.py | 34 +++++++++++++++++++ 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index c80a7aa9..1d037d90 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -195,10 +195,16 @@ async def _staged_artifact_size(self, container_path: str) -> int | None: if self._sandbox is None: return None 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)}" - ) + # A container-mode sandbox is the env's container, staged into directly; a VM runs it under docker. + if self._sandbox.mode == SANDBOX_MODE_CONTAINER: + exit_code, out, err = await self._sandbox.exec_with_output("stat", "-c", "%s", container_path) + if exit_code != 0: + raise RuntimeError(f"stat exited {exit_code}: {err.strip()[-200:]}") + else: + 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()) except Exception as e: # noqa: BLE001 -- a measurement, not a dependency logger.warning( diff --git a/tst/unit/env/envs/test_mcp_server_container_mode.py b/tst/unit/env/envs/test_mcp_server_container_mode.py index 1308cdd0..3d0ae5b7 100644 --- a/tst/unit/env/envs/test_mcp_server_container_mode.py +++ b/tst/unit/env/envs/test_mcp_server_container_mode.py @@ -1,6 +1,7 @@ """A standalone MCPServerEnv on Modal containers stages into its server's own container and records every container it created, so a later process can reattach them: loads stage into the server again, and close() and the reapers reach them all. A dead servicedb or sidecar is skipped on restore; a dead server fails it. A VM deploy's record is unchanged. +A load is timed by the size of the payload it staged, read in the server's container, or in the env's docker container on a VM. Deployed without a gateway, the server is the env: its own card and container are the record's, and loads go straight to it. Only the sandboxes, the HTTP boundary and each topology's deploy path are faked.""" @@ -20,6 +21,7 @@ from agent_env.env.env import DeployedEnv, DeployedSandboxEnv from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.env.gateway import GatewayMode +from agent_env.env.gateway.constants import data_plane_load_timeout_s from agent_env.providers.env_providers.env_gateway_provider import DeployedGateway from agent_env.providers.env_providers.env_provider import _builtin_env_providers, build_env_provider from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SANDBOX_MODE_VM @@ -62,6 +64,38 @@ async def test_a_container_deploy_stages_into_the_server_and_records_every_conta sandboxes["gw"].write_file_from_s3.assert_not_awaited() +@pytest.mark.asyncio +async def test_a_container_load_is_timed_by_the_payload_staged_in_the_server(sent): + env, sandboxes, size = _env(), _containers(), 3_500 * 1024 * 1024 + sandboxes["srv"].exec_with_output = AsyncMock(return_value=(0, f"{size}\n", "")) + + await _deploy(env, sandboxes) + await env.load_environment_artifact(_artifact()) + + sandboxes["srv"].exec_with_output.assert_awaited_once_with("stat", "-c", "%s", "/data/email.json") + assert sent and {r.extensions["timeout"]["read"] for r in sent} == {data_plane_load_timeout_s(size)} + assert data_plane_load_timeout_s(size) > data_plane_load_timeout_s(None) + + +@pytest.mark.asyncio +async def test_a_container_stat_that_fails_leaves_the_load_its_floor_timeout(): + env, server = _env(), _sandbox("srv") + server.exec_with_output = AsyncMock(return_value=(1, "", "stat: cannot statx '/data/email.json'")) + env._sandbox = server + + assert await env._staged_artifact_size("/data/email.json") is None + + +@pytest.mark.asyncio +async def test_a_vm_payload_is_measured_in_the_envs_docker_container(): + env, vm = _env(), _sandbox("vm-1", mode=SANDBOX_MODE_VM, sandbox_type="modal_vm") + vm.exec_script = AsyncMock(return_value="1234\n") + env._sandbox, env._env_provider._get_container_id = vm, AsyncMock(return_value="ctr-1") + + assert await env._staged_artifact_size("/data/email.json") == 1234 + vm.exec_script.assert_awaited_once_with("docker exec ctr-1 stat -c %s /data/email.json") + + @pytest.mark.asyncio async def test_a_vm_deploy_record_is_unchanged(): env, vm = _env(), _sandbox("vm-1", mode="vm", sandbox_type="modal_vm") From 2e2902a8cf30d620c81ae19034d4d21d32e9fc3d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 09:48:13 -0700 Subject: [PATCH 3/7] Read a service's JSON state from /export-state when data/get answers with a file snapshot_agent_state and verify_universe_roundtrip took a v1 service's state from data/get's first part as a DataPart. A service that exports a file bundle of its database answers with a FilePart instead, so snapshot_agent_state skipped it with a warning and verify_universe_roundtrip failed. legacy_protocol.service_state gives both the same answer: the data data/get returns when it is data, else the JSON the service serves at /export-state, at the address the env card gives it. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/legacy_protocol.py | 24 ++++++++++- .../verify_universe_roundtrip.py | 13 +----- .../task_steps/snapshot_agent_state.py | 10 +---- tst/unit/env/test_legacy_protocol.py | 33 +++++++++++++++ .../task_step/test_snapshot_agent_state.py | 40 +++++++++++++++++++ 5 files changed, 99 insertions(+), 21 deletions(-) diff --git a/src/agent_env/env/legacy_protocol.py b/src/agent_env/env/legacy_protocol.py index 3d5ba02c..1c600c0e 100644 --- a/src/agent_env/env/legacy_protocol.py +++ b/src/agent_env/env/legacy_protocol.py @@ -12,7 +12,7 @@ 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 @@ -82,13 +82,33 @@ async def add_via_rest(base_url: str, file_path: str, timeout: int = 120, verify async def export_state(gateway_url: str, environment_name: str, timeout: int = 60, 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=60, verify=True) + + async def reset_via_mcp_tool( mcp_url: str, environment_name: str, diff --git a/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py b/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py index fc69eb0c..c1ada510 100644 --- a/src/agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py +++ b/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.""" diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index 7bb50f60..25b1b7a3 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/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-/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, diff --git a/tst/unit/env/test_legacy_protocol.py b/tst/unit/env/test_legacy_protocol.py index 69f983af..6f5a60a1 100644 --- a/tst/unit/env/test_legacy_protocol.py +++ b/tst/unit/env/test_legacy_protocol.py @@ -4,7 +4,10 @@ from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from agentenv_protocol import RPC_PATH, WELL_KNOWN_PATH, DataPart, uploaded_file_part +from agentenv_protocol.client import GetDataResponse from agent_env.env import legacy_protocol @@ -123,3 +126,33 @@ async def test_a_gateway_path_without_a_gateway_fails_readably(): await legacy_protocol.child_env_card(record, None, "slack") with pytest.raises(EnvNeedsGateway, match=message): await legacy_protocol.export_state(None, "slack") + + +def _v1_service(monkeypatch, answer: list, export_state: dict) -> list[str]: + """A v1 service at http://gw/svc/mcp-slack that answers data/get with ``answer`` and serves ``export_state``; + returns each request it got, as ``METHOD path``.""" + asked, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + asked.append(f"{request.method} {request.url.path}") + if request.url.path.endswith(WELL_KNOWN_PATH): + return httpx.Response(200, json={"name": "slack", "url": RPC_PATH}) + if request.url.path.endswith("/export-state"): + return httpx.Response(200, json=export_state) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": GetDataResponse(parts=answer).model_dump(mode="json")}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return asked + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answer, state, reads_export_state", [ + ([DataPart(data={"messages": 2})], {"messages": 2}, False), + ([uploaded_file_part("slack.zip", name="slack.zip", mime_type="application/zip")], {"messages": 3}, True), + ([], {}, False), +], ids=["data", "file-bundle", "nothing"]) +async def test_a_v1_services_state_is_the_data_it_answers_with_else_its_export_state(monkeypatch, answer, state, reads_export_state): + asked = _v1_service(monkeypatch, answer, export_state={"messages": 3}) + + assert await legacy_protocol.service_state(None, "http://gw", "slack") == state + assert ("GET /svc/mcp-slack/export-state" in asked) is reads_export_state diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index c598ebe0..ca96aab2 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -179,6 +179,46 @@ async def _export_state(gw, name): assert set(json.loads(ctx.metadata["snapshot_json_url"])) == {"calendar", "contacts"} +@pytest.mark.asyncio +async def test_a_service_that_exports_a_file_bundle_is_captured_from_its_export_state(monkeypatch): + """calendar answers data/get with a bundle of its database, which is not JSON state: its state is read from + /export-state at the card's address rather than skipped. contacts answers with its state.""" + import httpx + from agentenv_protocol import DataPart, uploaded_file_part + from agentenv_protocol.client import GetDataResponse + + bodies: dict[str, dict] = {} + + class _BodyS3(_StubS3): + def put_object(self, Bucket, Key, Body, ContentType): # noqa: N803 + bodies[Key.rsplit("/", 1)[-1]] = json.loads(Body) + + store = S3ObjectStore(_BodyS3(), "artifact-bucket") + monkeypatch.setattr(artifact_mod, "EnvironmentUniverseArtifact", type("U", (), {"get": staticmethod(lambda _id: _StubUniverse())})) + monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store": lambda self: store})()) + + async def _get_data(base, timeout=30): + if base.endswith("/mcp-calendar"): + return GetDataResponse(parts=[uploaded_file_part("calendar.zip", name="calendar.zip", mime_type="application/zip")]) + return GetDataResponse(parts=[DataPart(data={"contacts": 4})]) + + real = httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + if str(request.url) == "https://sandbox.example/sb-1/svc/mcp-calendar/export-state": + return httpx.Response(200, json={"events": 7}) + return httpx.Response(404) + + monkeypatch.setattr(protocol_v1, "get_data", _get_data) + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + + ctx = _carded_ctx(await _card("mcp-calendar", "mcp-contacts")) + await _step()._capture_universe_state(ctx, _PREFIX) + + assert bodies == {"calendar.json": {"events": 7}, "contacts.json": {"contacts": 4}} + assert set(json.loads(ctx.metadata["snapshot_json_url"])) == {"calendar", "contacts"} + + async def _card(*keys: str) -> dict: """The real gateway's composed card over child envs at the given gateway keys (`mcp-{name}` or `{name}`).""" from agent_env.env.gateway.gateway import Gateway From 1d559a89fc17d333cd3074d6efab350a08d915a8 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 09:57:41 -0700 Subject: [PATCH 4/7] Stat a container-mode payload through docker exec on a VM-backed sandbox A VM-backed sandbox in container mode (the local one) runs exec on the container's Docker host, not inside the container, so stat there missed the staged file and the load fell back to the floor. Reach the container the way write_file_from_s3 does, by name. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/envs/mcp_server.py | 8 ++++++-- tst/unit/env/envs/test_mcp_server_container_mode.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index 1d037d90..5064ce16 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -20,6 +20,7 @@ load_by_signed_url, plugin_provider_like_a_builtin, provider_or_class, ) from agent_env.env.gateway import GatewayMode +from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER from agent_env.attribution import Attribution @@ -195,9 +196,12 @@ async def _staged_artifact_size(self, container_path: str) -> int | None: if self._sandbox is None: return None try: - # A container-mode sandbox is the env's container, staged into directly; a VM runs it under docker. + # A container-mode sandbox is the env's container; a VM runs the env under docker. if self._sandbox.mode == SANDBOX_MODE_CONTAINER: - exit_code, out, err = await self._sandbox.exec_with_output("stat", "-c", "%s", container_path) + command = ("stat", "-c", "%s", container_path) + if isinstance(self._sandbox, VmSandbox): # its exec runs on the container's Docker host + command = ("sudo", "docker", "exec", self._sandbox.container_name, *command) + exit_code, out, err = await self._sandbox.exec_with_output(*command) if exit_code != 0: raise RuntimeError(f"stat exited {exit_code}: {err.strip()[-200:]}") else: diff --git a/tst/unit/env/envs/test_mcp_server_container_mode.py b/tst/unit/env/envs/test_mcp_server_container_mode.py index 3d0ae5b7..93d93259 100644 --- a/tst/unit/env/envs/test_mcp_server_container_mode.py +++ b/tst/unit/env/envs/test_mcp_server_container_mode.py @@ -24,6 +24,7 @@ from agent_env.env.gateway.constants import data_plane_load_timeout_s from agent_env.providers.env_providers.env_gateway_provider import DeployedGateway from agent_env.providers.env_providers.env_provider import _builtin_env_providers, build_env_provider +from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SANDBOX_MODE_VM from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.teardown_sandboxes import _env_sandbox_ids @@ -86,6 +87,18 @@ async def test_a_container_stat_that_fails_leaves_the_load_its_floor_timeout(): assert await env._staged_artifact_size("/data/email.json") is None +@pytest.mark.asyncio +async def test_a_container_on_a_docker_host_is_measured_through_docker_exec(): + """A VM-backed sandbox in container mode (the local one) execs on the container's host, not in it.""" + env = _env() + env._sandbox = MagicMock(spec=VmSandbox, mode=SANDBOX_MODE_CONTAINER, container_name="agent-env-srv") + env._sandbox.exec_with_output = AsyncMock(return_value=(0, "1234\n", "")) + + assert await env._staged_artifact_size("/data/email.json") == 1234 + env._sandbox.exec_with_output.assert_awaited_once_with( + "sudo", "docker", "exec", "agent-env-srv", "stat", "-c", "%s", "/data/email.json") + + @pytest.mark.asyncio async def test_a_vm_payload_is_measured_in_the_envs_docker_container(): env, vm = _env(), _sandbox("vm-1", mode=SANDBOX_MODE_VM, sandbox_type="modal_vm") From 3046e9a2fd9115d53c926bbceba0714fa1793d88 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 20:31:10 -0700 Subject: [PATCH 5/7] test(env): a load is timed by its staged payload, and a server's state reads as JSON from a file export The integration tiers ran the staged-payload stat on every local server-env load but never checked it, and no test read a service whose data/get answers with a file. The items server now serves its state at /export-state and, with urn:agentenv:export-as-file/v1 enabled, answers data/get with a file bundle of it. On the local backend (CI) and on Modal (when the resolved config reaches it), a server env's load measures the payload it staged through the deploy's own handle and a restored one, and service_state reads the state as JSON in both modes. Co-Authored-By: Claude Opus 5.5 (1M context) --- tst/data/agentenv_mcp/server.py | 22 +++++- .../env/envs/server_env_local_test.py | 30 ++++++- .../env/envs/server_env_modal_test.py | 79 +++++++++++++++++++ tst/integration/env/gateway/gateway_test.py | 1 + 4 files changed, 128 insertions(+), 4 deletions(-) create mode 100644 tst/integration/env/envs/server_env_modal_test.py diff --git a/tst/data/agentenv_mcp/server.py b/tst/data/agentenv_mcp/server.py index 58afcf6a..c52eedf0 100644 --- a/tst/data/agentenv_mcp/server.py +++ b/tst/data/agentenv_mcp/server.py @@ -1,4 +1,8 @@ -"""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 @@ -6,8 +10,9 @@ 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 diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index 217a5e27..8a25a903 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -1,8 +1,9 @@ """A server env on the local backend: its one MCP server in a LocalSandbox container, with no gateway and no Modal. The env deploys, a later process restores its record and loads data into the server, and the teardown step, which -rebuilds each sandbox from disk, removes the container. Requires a docker daemon; spins up a throwaway ``registry:2`` -and skips if it can't start. +rebuilds each sandbox from disk, removes the container. A load is timed by the size of the payload it staged, through +the deploy's own handle and a restored one; and the server's state reads as JSON whether ``data/get`` answers with +data or with a file bundle. Requires a docker daemon; spins up a throwaway ``registry:2`` and skips if it can't start. """ import json @@ -24,6 +25,7 @@ from agent_env.env import legacy_protocol from agent_env.env.env import DeployedSandboxEnv, Env from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.env.gateway import constants from agent_env.store.image_store import LocalRegistryImageStore from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.teardown_sandboxes import TORN_DOWN_KEY, TeardownSandboxesTaskStep @@ -110,6 +112,30 @@ async def test_a_server_env_deploys_restores_loads_and_tears_down_on_local_conta assert _docker("ps", "-aq", "--filter", f"name=^/{container}$").stdout.strip() == "" +@pytest.mark.asyncio +async def test_a_load_is_timed_by_its_payload_and_the_servers_state_reads_as_json_from_a_file_export(local_stack, monkeypatch): + staged_sizes, timeout_for = [], constants.data_plane_load_timeout_s + monkeypatch.setattr(constants, "data_plane_load_timeout_s", lambda size: staged_sizes.append(size) or timeout_for(size)) + uid = uuid.uuid4().hex[:8] + env = MCPServerEnv.put(id=f"server-items-{uid}", docker_image_artifact=_put_items_image(f"server-items-{uid}"), + environment_name="items", env_provider_type="server") + deployed = await env.deploy(sandbox_type="local", ttl_seconds=900) + local_stack.append(f"agent-{deployed.sandbox_id}") + artifact = _items_artifact(uid) + + await env.load_environment_artifact(artifact) # the deploy's container-mode handle + await (await Env.from_instance_id(deployed.instance_id)).load_environment_artifact(artifact) # a restored one + + assert staged_sizes == [len(artifact.get_file_artifact().load())] * 2 + items = {"items": ["snap-x", "snap-y"]} + assert await legacy_protocol.service_state(deployed, None, "items") == items + base_url = await legacy_protocol.v1_base_url(deployed, None, "items") + card = await protocol_v1.get_card(base_url) + await protocol_v1.invoke_extension(base_url, card, "urn:agentenv:export-as-file/v1", {"enabled": True}) + assert [part.kind for part in (await protocol_v1.get_data(base_url)).parts] == ["file"] + assert await legacy_protocol.service_state(deployed, None, "items") == items + + def _put_items_image(artifact_id: str) -> DockerImageArtifact: """The in-memory items server (card 'items'), built with the vendored agentenv_protocol and pushed to the registry.""" data = _REPO / "tst" / "data" / "agentenv_mcp" diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py new file mode 100644 index 00000000..9895b746 --- /dev/null +++ b/tst/integration/env/envs/server_env_modal_test.py @@ -0,0 +1,79 @@ +"""A server env on Modal: its one MCP server in a Modal container, with no gateway. + +A load is timed by the size of the payload it staged, measured in the server's container through the deploy's own +handle and a restored one; and the server's state reads as JSON whether ``data/get`` answers with data or with a file +bundle. Needs the resolved config to reach Modal and an image store Modal can pull from; skipped otherwise. +""" + +import json +import shutil +import tempfile +import uuid +from pathlib import Path + +import pytest +from agentenv_protocol import client as protocol_v1 + +from agent_env.artifact import DockerImageArtifact, EnvironmentArtifact, FileArtifact +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.gateway import constants +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.teardown_sandboxes import TeardownSandboxesTaskStep +from tst.util.capabilities import skip_without_remote_sandbox +from tst.util.image_cache import build_or_reuse + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_remote_sandbox("modal")] + +_REPO = Path(__file__).resolve().parents[4] + + +@pytest.mark.asyncio +async def test_a_modal_load_is_timed_by_its_payload_and_the_servers_state_reads_as_json_from_a_file_export(monkeypatch): + staged_sizes, timeout_for = [], constants.data_plane_load_timeout_s + monkeypatch.setattr(constants, "data_plane_load_timeout_s", lambda size: staged_sizes.append(size) or timeout_for(size)) + uid = uuid.uuid4().hex[:8] + env = MCPServerEnv.put(id=f"server-items-modal-{uid}", docker_image_artifact=_items_image(), + environment_name="items", env_provider_type="server") + deployed = await env.deploy(sandbox_type="modal", ttl_seconds=900) + try: + artifact = _items_artifact(uid) + await env.load_environment_artifact(artifact) + await (await Env.from_instance_id(deployed.instance_id)).load_environment_artifact(artifact) + + assert staged_sizes == [len(artifact.get_file_artifact().load())] * 2 + items = {"items": ["snap-x", "snap-y"]} + assert await legacy_protocol.service_state(deployed, None, "items") == items + base_url = await legacy_protocol.v1_base_url(deployed, None, "items") + card = await protocol_v1.get_card(base_url) + await protocol_v1.invoke_extension(base_url, card, "urn:agentenv:export-as-file/v1", {"enabled": True}) + assert [part.kind for part in (await protocol_v1.get_data(base_url)).parts] == ["file"] + assert await legacy_protocol.service_state(deployed, None, "items") == items + finally: + await TeardownSandboxesTaskStep(id="teardown", version=None, env_ids=[env.id]).execute( + TaskStepContext(deployed_envs=[deployed]) + ) + + +def _items_image() -> DockerImageArtifact: + """The in-memory items server (card 'items'), built for linux/amd64 with the vendored agentenv_protocol.""" + data = _REPO / "tst" / "data" / "agentenv_mcp" + with tempfile.TemporaryDirectory() as build_dir: + context = Path(build_dir) + shutil.copytree(_REPO / "packages" / "agentenv-protocol" / "src" / "agentenv_protocol", context / "agentenv_protocol", + ignore=shutil.ignore_patterns("__pycache__")) + for name in ("server.py", "Dockerfile", "seed.json"): + shutil.copy(data / name, context / name) + return build_or_reuse(artifact_id="server-items-modal", description="server env on Modal", + dockerfile=context / "Dockerfile", context=context, tag="server-items-modal") + + +def _items_artifact(uid: str) -> EnvironmentArtifact: + return EnvironmentArtifact.put( + id=f"server-items-modal-data-{uid}", + environment_name="items", + file_artifact=FileArtifact.put_bytes(id=f"server-items-modal-file-{uid}", description="server env on Modal", + filename="items.json", content=json.dumps({"items": ["snap-x", "snap-y"]}).encode(), + content_type="application/json"), + ) diff --git a/tst/integration/env/gateway/gateway_test.py b/tst/integration/env/gateway/gateway_test.py index c299eae0..8d9cedbc 100644 --- a/tst/integration/env/gateway/gateway_test.py +++ b/tst/integration/env/gateway/gateway_test.py @@ -1966,6 +1966,7 @@ async def test_gateway_with_agentenv_environment(agentenv_items_env): assert card["additionalInterfaces"] == [{"url": "/mcp", "transport": "mcp"}] assert {e["uri"] for e in card["capabilities"]["extensions"]} == { "urn:agentenv:disable-tool/v1", "urn:agentenv:set-errors/v1", "urn:agentenv:clock/v1", + "urn:agentenv:export-as-file/v1", } # Decorator tools ride capabilities.tools ({environment_name} resolved); imperative tools stay handshake-only. assert [t["name"] for t in card["capabilities"]["tools"]] == ["items_add_item"] From f2bf992876354ed0e26848ca96a3642490ebbbc3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 20:45:09 -0700 Subject: [PATCH 6/7] refactor(env): name the /export-state timeout both readers use Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/legacy_protocol.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/agent_env/env/legacy_protocol.py b/src/agent_env/env/legacy_protocol.py index 1c600c0e..47f951ba 100644 --- a/src/agent_env/env/legacy_protocol.py +++ b/src/agent_env/env/legacy_protocol.py @@ -20,6 +20,7 @@ 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,7 +81,9 @@ 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.""" return await _export_state_at(environment_base_url(gateway_url, environment_name, mcp=True), timeout, verify) @@ -106,7 +109,7 @@ async def service_state(deployed: Optional[DeployedEnv], gateway_url: Optional[s return {} if isinstance(response.parts[0], DataPart): return response.parts[0].data - return await _export_state_at(base_url, timeout=60, verify=True) + return await _export_state_at(base_url, timeout=EXPORT_STATE_TIMEOUT_SECONDS, verify=True) async def reset_via_mcp_tool( From 7e40368c782304f93b7cdbd51d0fa57896aa4951 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 6 Oct 2026 21:36:06 -0700 Subject: [PATCH 7/7] refactor(env): size a load's timeout from the store's metadata for its payload, not a stat in the sandbox The staged file is a byte-for-byte copy of the artifact's object, so the store already knows its size. Asking it, as load_by_signed_url does for plugin-deployed envs, replaces the stat that had to know how to reach the env's container on each kind of sandbox (in it on Modal, through docker on a VM-backed sandbox in container mode, through the gateway's container on a VM) and that a new sandbox backend could break again. A store that can't say still leaves the floor timeout. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/envs/mcp_server.py | 44 +++++++----------- .../envs/test_mcp_server_container_mode.py | 46 +++++++------------ 2 files changed, 32 insertions(+), 58 deletions(-) diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index d961dd6b..5dff7545 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -2,9 +2,9 @@ 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 @@ -12,6 +12,7 @@ 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 @@ -21,7 +22,6 @@ load_by_signed_url, plugin_provider_like_a_builtin, provider_or_class, ) from agent_env.env.gateway import GatewayMode -from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER from agent_env.attribution import Attribution @@ -156,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. @@ -205,35 +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: - # A container-mode sandbox is the env's container; a VM runs the env under docker. - if self._sandbox.mode == SANDBOX_MODE_CONTAINER: - command = ("stat", "-c", "%s", container_path) - if isinstance(self._sandbox, VmSandbox): # its exec runs on the container's Docker host - command = ("sudo", "docker", "exec", self._sandbox.container_name, *command) - exit_code, out, err = await self._sandbox.exec_with_output(*command) - if exit_code != 0: - raise RuntimeError(f"stat exited {exit_code}: {err.strip()[-200:]}") - else: - 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}" diff --git a/tst/unit/env/envs/test_mcp_server_container_mode.py b/tst/unit/env/envs/test_mcp_server_container_mode.py index 93d93259..b7449d37 100644 --- a/tst/unit/env/envs/test_mcp_server_container_mode.py +++ b/tst/unit/env/envs/test_mcp_server_container_mode.py @@ -1,7 +1,7 @@ """A standalone MCPServerEnv on Modal containers stages into its server's own container and records every container it created, so a later process can reattach them: loads stage into the server again, and close() and the reapers reach them all. A dead servicedb or sidecar is skipped on restore; a dead server fails it. A VM deploy's record is unchanged. -A load is timed by the size of the payload it staged, read in the server's container, or in the env's docker container on a VM. +A load is timed by the size its object store holds for the payload it stages. Deployed without a gateway, the server is the env: its own card and container are the record's, and loads go straight to it. Only the sandboxes, the HTTP boundary and each topology's deploy path are faked.""" @@ -19,12 +19,13 @@ from agent_env.artifact import Artifact from agent_env.env.env import DeployedEnv, DeployedSandboxEnv +from agent_env.env.envs import mcp_server from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.env.gateway import GatewayMode from agent_env.env.gateway.constants import data_plane_load_timeout_s +from agent_env.store.object_store import ObjectMetadata from agent_env.providers.env_providers.env_gateway_provider import DeployedGateway from agent_env.providers.env_providers.env_provider import _builtin_env_providers, build_env_provider -from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SANDBOX_MODE_VM from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.teardown_sandboxes import _env_sandbox_ids @@ -66,47 +67,32 @@ async def test_a_container_deploy_stages_into_the_server_and_records_every_conta @pytest.mark.asyncio -async def test_a_container_load_is_timed_by_the_payload_staged_in_the_server(sent): +async def test_a_load_is_timed_by_the_size_its_store_holds_for_the_payload(sent, monkeypatch): env, sandboxes, size = _env(), _containers(), 3_500 * 1024 * 1024 - sandboxes["srv"].exec_with_output = AsyncMock(return_value=(0, f"{size}\n", "")) + store = _store_reporting(ObjectMetadata(size=size)) + monkeypatch.setattr(mcp_server, "get_config", lambda: MagicMock(get_object_store_at=MagicMock(return_value=store))) await _deploy(env, sandboxes) await env.load_environment_artifact(_artifact()) - sandboxes["srv"].exec_with_output.assert_awaited_once_with("stat", "-c", "%s", "/data/email.json") + store.get_object_metadata_at.assert_called_once_with("s3://bucket/email.json") assert sent and {r.extensions["timeout"]["read"] for r in sent} == {data_plane_load_timeout_s(size)} assert data_plane_load_timeout_s(size) > data_plane_load_timeout_s(None) @pytest.mark.asyncio -async def test_a_container_stat_that_fails_leaves_the_load_its_floor_timeout(): - env, server = _env(), _sandbox("srv") - server.exec_with_output = AsyncMock(return_value=(1, "", "stat: cannot statx '/data/email.json'")) - env._sandbox = server +@pytest.mark.parametrize("lookup", [None, RuntimeError("store unreachable")], ids=["no-metadata", "lookup-fails"]) +async def test_a_store_that_cannot_size_the_payload_leaves_the_load_its_floor_timeout(monkeypatch, lookup): + store = _store_reporting(lookup) + monkeypatch.setattr(mcp_server, "get_config", lambda: MagicMock(get_object_store_at=MagicMock(return_value=store))) - assert await env._staged_artifact_size("/data/email.json") is None + assert await _env()._staged_artifact_size(_artifact().get_file_artifact()) is None -@pytest.mark.asyncio -async def test_a_container_on_a_docker_host_is_measured_through_docker_exec(): - """A VM-backed sandbox in container mode (the local one) execs on the container's host, not in it.""" - env = _env() - env._sandbox = MagicMock(spec=VmSandbox, mode=SANDBOX_MODE_CONTAINER, container_name="agent-env-srv") - env._sandbox.exec_with_output = AsyncMock(return_value=(0, "1234\n", "")) - - assert await env._staged_artifact_size("/data/email.json") == 1234 - env._sandbox.exec_with_output.assert_awaited_once_with( - "sudo", "docker", "exec", "agent-env-srv", "stat", "-c", "%s", "/data/email.json") - - -@pytest.mark.asyncio -async def test_a_vm_payload_is_measured_in_the_envs_docker_container(): - env, vm = _env(), _sandbox("vm-1", mode=SANDBOX_MODE_VM, sandbox_type="modal_vm") - vm.exec_script = AsyncMock(return_value="1234\n") - env._sandbox, env._env_provider._get_container_id = vm, AsyncMock(return_value="ctr-1") - - assert await env._staged_artifact_size("/data/email.json") == 1234 - vm.exec_script.assert_awaited_once_with("docker exec ctr-1 stat -c %s /data/email.json") +def _store_reporting(metadata): + """A store whose metadata lookup returns ``metadata``, or raises it.""" + lookup = MagicMock(side_effect=metadata) if isinstance(metadata, Exception) else MagicMock(return_value=metadata) + return MagicMock(get_object_metadata_at=lookup) @pytest.mark.asyncio