diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index e9127ea7..4dcc2b37 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/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) 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/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index 0c8e2a47..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 @@ -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}" diff --git a/src/agent_env/env/legacy_protocol.py b/src/agent_env/env/legacy_protocol.py index 3d5ba02c..47f951ba 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 @@ -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,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, 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/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/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( 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"] 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..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,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 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.""" @@ -18,8 +19,11 @@ 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_provider import SANDBOX_MODE_CONTAINER, SANDBOX_MODE_VM @@ -62,6 +66,35 @@ 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_load_is_timed_by_the_size_its_store_holds_for_the_payload(sent, monkeypatch): + env, sandboxes, size = _env(), _containers(), 3_500 * 1024 * 1024 + 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()) + + 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 +@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(_artifact().get_file_artifact()) is None + + +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 async def test_a_vm_deploy_record_is_unchanged(): env, vm = _env(), _sandbox("vm-1", mode="vm", sandbox_type="modal_vm") 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