diff --git a/src/agent_env/env/envs/computer/env.py b/src/agent_env/env/envs/computer/env.py index 5881e030..9270f165 100644 --- a/src/agent_env/env/envs/computer/env.py +++ b/src/agent_env/env/envs/computer/env.py @@ -30,6 +30,7 @@ from agent_env.utils.paths import validate_relative_filename if TYPE_CHECKING: + from agent_env.artifact import EnvironmentArtifact, EnvironmentUniverseArtifact from agent_env.bundle.authoring import AuthoringContext from agent_env.env.env import DeployedEnv from agent_env.env.envs.computer.app_env import ComputerAppEnv @@ -206,6 +207,34 @@ async def load_file_artifact_universe(self, file_artifact_universe: Any, logger.info(f"Loaded {len(files)} file(s) of '{file_artifact_universe.id}' into {self.id}'s ~/{destination}") return LoadFileArtifactUniverseResult(destination_path=destination, files=paths) + async def load_environment_artifact(self, environment_artifact: EnvironmentArtifact) -> None: + """Write the files of ``environment_artifact``, the computer's own, at their paths relative to the session + user's home, through the driver's ``data/add``, as ``snapshot_computer_env`` kept them. Other files there + stay.""" + if environment_artifact.environment_name != self.environment_name: + raise ValueError(f"EnvironmentArtifact environment_name '{environment_artifact.environment_name}' does not " + f"match env environment_name '{self.environment_name}'") + files = environment_artifact.get_file_artifacts() + for path in files: + validate_relative_filename(path) + await inside.add_artifacts(self._deployed, files) + logger.info(f"Loaded {len(files)} file(s) of '{environment_artifact.id}' into {self.id}'s home") + + async def load_environment_universe_artifact(self, environment_universe_artifact: EnvironmentUniverseArtifact) -> None: + """Load the universe's environment named after the computer, ``load_environment_artifact``; any other is left + out, with a warning.""" + environments = environment_universe_artifact.get_environment_artifacts() + own = [environment for environment in environments if environment.environment_name == self.environment_name] + if not own: + raise RuntimeError(f"computer env '{self.id}' can't load '{environment_universe_artifact.id}': it has no " + f"environment named '{self.environment_name}' (it has " + f"{[environment.environment_name for environment in environments]})") + if others := [environment.environment_name for environment in environments + if environment.environment_name != self.environment_name]: + logger.warning(f"[{self.id}] loading only '{self.environment_name}' from '{environment_universe_artifact.id}'; " + f"leaving out {others}") + await self.load_environment_artifact(own[0]) + async def reset(self, deployed: DeployedEnv) -> None: """Return the computer to boot state: its recordings ended, its container replaced with a fresh one of the same image, on the same sandbox and port, so its address stays. An agent's MCP session ends with the old driver; one diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 33ad6fa3..86d84c92 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -150,7 +150,7 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s "aggregate_verifiers", "apply_server_config", "build_mcp_cli", "combine_universe_verdicts", "configure_scheduled_checks", "deploy_env", "deploy_sandbox", "end_computer_recording", "start_computer_recording", "env_outcome_verifier", "modify_env_tool_access", "reset_env", "review", - "run_container_unit_tests_verifier", "run_docker_container", "snapshot_env", "sync_env_clock", + "run_container_unit_tests_verifier", "run_docker_container", "snapshot_env", "snapshot_computer_env", "sync_env_clock", "validation_gate_aggregator", "verify_env_card", "verify_env_core_protocol", "verify_mcp_tool_schema", "verify_spec_conformance", "verify_universe_load_export_roundtrip", # They read only what an outside turn leaves: its reply. diff --git a/src/agent_env/task_step/registry.py b/src/agent_env/task_step/registry.py index 3c74f4fc..83821a63 100644 --- a/src/agent_env/task_step/registry.py +++ b/src/agent_env/task_step/registry.py @@ -76,6 +76,7 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: from agent_env.task_step.task_steps.verifiers.agent_prompt_response_verifier import AgentPromptResponseVerifierTaskStep from agent_env.task_step.task_steps.snapshot_agent_state import SnapshotAgentStateTaskStep from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep + from agent_env.task_step.task_steps.computer.snapshot_computer_env import SnapshotComputerEnvTaskStep from agent_env.task_step.task_steps.peer_agents import PeerAgentsTaskStep from agent_env.task_step.task_steps.verifiers.verify_sandbox import VerifySandboxTaskStep from agent_env.task_step.task_steps.verifiers.aggregate_verifiers import AggregateVerifiersTaskStep @@ -126,6 +127,7 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: _get_type(VerifyCoreA2AProtocolStep): VerifyCoreA2AProtocolStep, _get_type(SnapshotAgentStateTaskStep): SnapshotAgentStateTaskStep, _get_type(SnapshotEnvTaskStep): SnapshotEnvTaskStep, + _get_type(SnapshotComputerEnvTaskStep): SnapshotComputerEnvTaskStep, _get_type(RunCodeTaskStep): RunCodeTaskStep, _get_type(PeerAgentsTaskStep): PeerAgentsTaskStep, _get_type(VerifyA2ASnapshotStep): VerifyA2ASnapshotStep, diff --git a/src/agent_env/task_step/task_steps/computer/__init__.py b/src/agent_env/task_step/task_steps/computer/__init__.py new file mode 100644 index 00000000..16fef985 --- /dev/null +++ b/src/agent_env/task_step/task_steps/computer/__init__.py @@ -0,0 +1 @@ +"""Task steps for computer envs.""" diff --git a/src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py b/src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py new file mode 100644 index 00000000..f28c08bb --- /dev/null +++ b/src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py @@ -0,0 +1,107 @@ +"""snapshot_computer_env: a computer env's files, kept by their paths in its environment artifact, which load_artifact +loads back into a computer.""" + +from __future__ import annotations + +import asyncio +from typing import ClassVar, Optional + +from agentenv_protocol.transfers import RelativePath +from pydantic import TypeAdapter, ValidationError + +from agent_env.entity.refs import EntityRef, RefRole +from agent_env.env.env import Env +from agent_env.env.envs.computer.env import ComputerEnv +from agent_env.store.base import NotFoundError +from agent_env.store.ids import is_local_id, validate_local_id +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep, TaskStepDependency +from agent_env.task_step.task_steps.snapshot_env import DEFAULT_EXPORT_TIMEOUT_SECONDS, SnapshotEnvTaskStep + +_PATHS = TypeAdapter(list[RelativePath]) + + +class SnapshotComputerEnvTaskStep(SnapshotEnvTaskStep): + """Snapshot a deployed computer env's files into a new environment universe, as ``snapshot_env`` snapshots an env, + each file under its path relative to the session user's home in the computer's environment artifact: those under + ``file_paths``, by default the folders its card declares. ``load_artifact`` loads the snapshot back into a + computer, through ``data/add``. + + Output: ``context.metadata["env_snapshotted_universes"][self.id]`` = ``{"id", "version"}``.""" + + type: ClassVar[str] = "snapshot_computer_env" + entity_refs = ( + EntityRef.env("env_id"), + EntityRef.artifact("snapshot_id", role=RefRole.OUTPUT, artifact_type="environment_universe"), + ) + + def __init__( + self, + id: str, + version: Optional[int], + env_id: Optional[str] = None, + env_instance_id: Optional[str] = None, + env_step_id: Optional[str] = None, + snapshot_id: Optional[str] = None, + file_paths: Optional[list[str]] = None, + export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, + depends_on: Optional[list[TaskStepDependency]] = None, + fail_task_on_error: bool = True, + ): + super().__init__(id, version, env_id=env_id, env_instance_id=env_instance_id, env_step_id=env_step_id, + snapshot_id=snapshot_id, export_timeout_seconds=export_timeout_seconds, + depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.file_paths = file_paths + + def to_dict(self) -> dict: + return {**TaskStep.to_dict(self), "env_id": self.env_id, "env_instance_id": self.env_instance_id, + "env_step_id": self.env_step_id, "snapshot_id": self.snapshot_id, "file_paths": self.file_paths, + "export_timeout_seconds": self.export_timeout_seconds} + + @classmethod + def from_dict(cls, data: dict) -> SnapshotComputerEnvTaskStep: + return cls(**cls._base_from_dict(data), env_id=data.get("env_id"), env_instance_id=data.get("env_instance_id"), + env_step_id=data.get("env_step_id"), snapshot_id=data.get("snapshot_id"), + file_paths=data.get("file_paths"), + export_timeout_seconds=data.get("export_timeout_seconds", DEFAULT_EXPORT_TIMEOUT_SECONDS)) + + def config_problems(self) -> list[str]: + problems = super().config_problems() + if self.file_paths is not None: + try: + _PATHS.validate_python(self.file_paths) + except ValidationError: + problems.append(f"file_paths are normalized paths relative to the computer's session user's home, " + f"not {self.file_paths!r}") + return problems + + def preflight(self) -> list[str]: + """The env this step names when it isn't a computer env.""" + problems = super().preflight() + if self.env_id is None: + return problems + try: + env = Env.get(self.env_id) + except NotFoundError: # what the run reports + return problems + return problems if isinstance(env, ComputerEnv) else [*problems, _not_a_computer(env)] + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + snapshot_id = self._derive_snapshot_id(context) + if is_local_id(snapshot_id): + validate_local_id(snapshot_id) + deployed = self._resolve_deployed_env(context) + env = await asyncio.to_thread(Env.get, deployed.env_id, deployed.env_version) + if not isinstance(env, ComputerEnv): + raise RuntimeError(_not_a_computer(env)) + result = await self.snapshot_env_state( + env=env, gateway_url=None, snapshot_id=snapshot_id, deployed=deployed, + export_timeout_seconds=self.export_timeout_seconds, + data_get_params=None if self.file_paths is None else {"paths": self.file_paths}) + context.metadata.setdefault("env_snapshotted_universes", {})[self.id] = { + "id": result.environment_universe_artifact_id, "version": result.environment_universe_artifact_version} + return context + + +def _not_a_computer(env: Env) -> str: + return f"snapshot_computer_env snapshots a computer env, and {env.id} is a {env.type} env" 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 49a108d8..6bd8d205 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -1,4 +1,5 @@ -"""Snapshot a deployed env's live service state into a new EnvironmentUniverseArtifact.""" +"""Snapshot a deployed env's live service state into a new EnvironmentUniverseArtifact: through its gateway, or without +one through the card its deployment recorded.""" from __future__ import annotations import asyncio @@ -11,20 +12,23 @@ import uuid from dataclasses import dataclass from datetime import datetime, timezone +from pathlib import PurePosixPath from typing import ClassVar, Optional from urllib.parse import urlparse import httpx from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object_path -from agentenv_protocol.client import find_extension, find_extension_method +from agentenv_protocol import client as protocol_v1 +from agentenv_protocol.client import GetDataResponse, find_extension, find_extension_method, intake_declaration from agentenv_protocol.transfers import WriteNamespaceGrant +from agentenv_protocol.types import FileWithBytes from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant from agent_env.env.gateway.constants import CHANGES_PAGE_MAX, EXT_CHANGES_URI, EXT_TRAJECTORY_URI from agent_env.store import get_config from agent_env.store.base import GrantUnavailableError -from agent_env.store.ids import derive_id, is_local_id, validate_local_id +from agent_env.store.ids import derive_id, derive_path_id, is_local_id, validate_local_id from agent_env.store.object_store import MIN_GRANT_LIFETIME_SECONDS, ObjectStore from agent_env.store.object_store.object_store import issues_grants_to from agent_env.store.routing import in_local_run @@ -48,6 +52,11 @@ # What a service may upload through its snapshot grant: one bundle, as large as one upload to the # store can be up to this. ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=1, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) +# What an environment whose data/get returns a file tree, a computer, may upload through its grant: each of its files, +# each as large as one upload to the store can be, up to this in all. +FILE_TREE_LIMITS = ObjectLimits(max_objects=100_000, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) +# How many of a file tree's files are written to the store at once. +_WRITES_IN_FLIGHT = 8 _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" @@ -62,21 +71,22 @@ def object_url(self, path: str) -> str: return self.store.object_url(f"{self.grant.root_path}/{path}") -def _snapshot_upload(card: dict, timeout_seconds: float, sandbox_type: Optional[str]) -> Optional[_SnapshotUpload]: - """A grant for a service advertising the data-objects form to upload its bundle with, under a namespace - of its own; None when it does not advertise the form, or the object store issues no grant that reaches - the service's sandbox, of ``sandbox_type`` (None: unknown) and lasts the export.""" +def _snapshot_upload(card: dict, timeout_seconds: float, sandbox_type: Optional[str], + limits: ObjectLimits = ENV_SNAPSHOT_LIMITS) -> Optional[_SnapshotUpload]: + """A grant for a service advertising the data-objects form to upload its bundle with, or the files ``limits`` + allows, under a namespace of its own; None when it does not advertise the form, or the object store issues no + grant that reaches the service's sandbox, of ``sandbox_type`` (None: unknown) and lasts the export.""" if find_extension(card, DATA_OBJECTS_EXTENSION_URI) is None: return None config = get_config() store = config.get_object_store() if not issues_grants_to(store, sandbox_type): return None - cap = store.max_single_upload_bytes or ENV_SNAPSHOT_LIMITS.max_object_bytes + cap = store.max_single_upload_bytes or limits.max_object_bytes limits = ObjectLimits( - max_objects=ENV_SNAPSHOT_LIMITS.max_objects, - max_object_bytes=min(ENV_SNAPSHOT_LIMITS.max_object_bytes, cap), - max_total_bytes=min(ENV_SNAPSHOT_LIMITS.max_total_bytes, cap), + max_objects=limits.max_objects, + max_object_bytes=min(limits.max_object_bytes, cap), + max_total_bytes=min(limits.max_total_bytes, cap), ) namespace_url = store.object_url(f"{config.get_artifact_key_prefix()}{_SNAPSHOT_KEY_PREFIX}/{uuid.uuid4().hex}") try: @@ -157,6 +167,10 @@ class SnapshotEnvTaskStep(TaskStep): ``{"id", "version"}``. Fails the step if any service fails to export. ``include_env_changes`` also stores who changed what (urn:agentenv:changes/v1) under ``metadata["env_changes"]``. + + An env deployed without a gateway is reached through the card its deployment recorded, its environments the + card's children, or itself. An environment whose data/get returns a file tree, a computer's, is kept file by file, + each by its path, in its environment artifact. """ type: ClassVar[str] = "snapshot_env" @@ -238,9 +252,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if not env_id: raise RuntimeError("snapshot_env could not resolve an env_id") gateway_url = self.gateway_url or gateway_url_of(deployed) - if not gateway_url: + if not gateway_url and not (deployed is not None and deployed.environment_card): raise RuntimeError( - f"Deployed env '{env_id}' has no gateway_url and no override was provided" + f"Deployed env '{env_id}' has no gateway_url, no recorded card to reach it through, and no override " + "was provided" ) if self.include_env_trajectory: @@ -461,11 +476,14 @@ def _enumerate_environments(env) -> list[str]: @staticmethod async def _export_environment_to_file( - gateway_url: str, environment_name: str, tmp_path: str, timeout_seconds: float, deployed=None - ) -> str: + gateway_url: Optional[str], environment_name: str, tmp_path: str, timeout_seconds: float, deployed=None, + data_get_params: Optional[dict] = None, + ) -> str | tuple[GetDataResponse, Optional[_SnapshotUpload]]: """Export one service; return the suffix written (".json"/".zip") to ``tmp_path``, or the object url a FilePart names when the service uploaded the - bundle itself (nothing written to ``tmp_path`` then — export_one registers it). + bundle itself (nothing written to ``tmp_path`` then — export_one registers it), or, for + an environment whose card declares a file tree, its data/get answer and the upload it + answered under, which export_one keeps file by file. ``data_get_params`` go to its data/get. v1 ``get_data``, handed a namespace grant when the service takes the data-objects form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` → @@ -473,19 +491,21 @@ async def _export_environment_to_file( ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON guard. Errors are caught by the caller (fails just this service).""" from agent_env.env import legacy_protocol - from agentenv_protocol import client as protocol_v1 base_url = await legacy_protocol.v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is not None: card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) + if returns_file_tree(card or {}): + upload = await asyncio.to_thread(_snapshot_upload, card or {}, timeout_seconds, + getattr(deployed, "sandbox_type", None), FILE_TREE_LIMITS) + return await get_data(base_url, environment_name, timeout_seconds, upload, **(data_get_params or {})) upload = await asyncio.to_thread( _snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None) ) # Push S3 creds (no-op unless the service advertises the extension), with a grant too: a # bundle the grant cannot hold can still be uploaded with them. await _push_s3_credentials(card_base, card or {}, timeout_seconds) - grant = {} if upload is None else {"write_namespace": upload.grant} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) + resp, upload = await get_data(base_url, environment_name, timeout_seconds, upload, **(data_get_params or {})) part = resp.parts[0] if resp.parts else None if upload is not None and (path := uploaded_object_path(part)) is not None: return upload.object_url(path) @@ -562,27 +582,30 @@ async def snapshot_env_state( cls, *, env, - gateway_url: str, + gateway_url: Optional[str], snapshot_id: str, deployed=None, original_universe_artifact_id: Optional[str] = None, export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, concurrency: int = DEFAULT_EXPORT_CONCURRENCY, + data_get_params: Optional[dict] = None, ) -> EnvSnapshotResult: - """Export all services via the gateway into a new EnvironmentUniverseArtifact - ``snapshot_id``. All-or-nothing; the raised error carries public - per-service summaries (verbose details go to the log). + """Export all services via the gateway, or without one (``gateway_url`` None) through + ``deployed``'s recorded card, into a new EnvironmentUniverseArtifact ``snapshot_id``, + handing each service's data/get ``data_get_params``. + All-or-nothing; the raised error carries public per-service summaries (verbose details + go to the log). A classmethod so callers that have no step instance can reuse it.""" from agent_env.artifact import FileArtifact, EnvironmentArtifact, EnvironmentUniverseArtifact from agent_env.env.env import gateway_url_of - services = cls._enumerate_environments(env) - gateway = gateway_url.rstrip("/") + services = cls._enumerate_environments(env) if gateway_url else _card_environments(deployed) + gateway = gateway_url.rstrip("/") if gateway_url else None # `gateway` identifies which deployment was exported: env= alone doesn't when one env_id has several. logger.info( f"snapshot_env: env={env.id} v{env.version} snapshot_id={snapshot_id} " - f"gateway={gateway} services={services}" + f"gateway={gateway or 'none, through its recorded card'} services={services}" ) # Carry the source universe's metadata into the snapshot: explicit @@ -607,7 +630,7 @@ async def snapshot_env_state( logger.warning("snapshot_env: failed to carry metadata from %s", source_id, exc_info=True) # The stored card describes the deployment's own gateway; an override pointing elsewhere keeps the live probe. - card_record = deployed if deployed is not None and (gateway_url_of(deployed) or "").rstrip("/") == gateway else None + card_record = deployed if deployed is not None and (gateway is None or (gateway_url_of(deployed) or "").rstrip("/") == gateway) else None sem = asyncio.Semaphore(concurrency) async def export_one(environment_name: str): @@ -622,7 +645,7 @@ async def export_one(environment_name: str): try: result = await cls._export_environment_to_file( gateway, environment_name, tmp_path, export_timeout_seconds, - deployed=card_record, + deployed=card_record, data_get_params=data_get_params, ) except httpx.HTTPStatusError as e: body_snippet = "" @@ -638,31 +661,33 @@ async def export_one(environment_name: str): except ValueError as e: return str(e), str(e) + if isinstance(result, tuple): + try: + files = await _file_tree_files(*result, environment_name, f"{snapshot_id}-{environment_name}", + f"state snapshot of {environment_name} from env {env.id}") + return await asyncio.to_thread( + EnvironmentArtifact.put, id=f"{snapshot_id}-{environment_name}", + environment_name=environment_name, file_artifacts=files) + except ValueError as e: + return str(e), str(e) + except Exception as e: + public = f"failed to create artifacts ({type(e).__name__})" + return public, f"{public}: {e}" + # Service uploaded the bundle itself — register a FileArtifact # pointing at it directly (no download). Derive filename/content_type # from the url's last segment (the extension drives load-time parsing), # like FileArtifact.put. if urlparse(result).scheme: - from agent_env.artifact.store import get_artifact_store - object_url = result filename = object_url.rsplit("/", 1)[-1] or f"{environment_name}.zip" content_type = mimetypes.guess_type(filename)[0] or "application/zip" def _register_uploaded_bundle(): - if get_config().get_object_store().get_object_metadata_at(object_url) is None: - raise FileNotFoundError(f"snapshot bundle not found at {object_url}") - store = get_artifact_store() - fa_id = f"{snapshot_id}-{environment_name}-file" - get_config().check_object_url(fa_id, object_url) - fa = FileArtifact( - id=fa_id, - description=f"State snapshot of {environment_name} from env {env.id}", - filename=filename, - content_type=content_type, - object_url=object_url, + fa = registered_upload( + object_url, f"{snapshot_id}-{environment_name}-file", filename, content_type, + f"State snapshot of {environment_name} from env {env.id}", ) - fa = store.put_document(fa) return EnvironmentArtifact.put( id=f"{snapshot_id}-{environment_name}", environment_name=environment_name, @@ -748,6 +773,79 @@ def _register_uploaded_bundle(): ) +def returns_file_tree(card: dict) -> bool: + """Whether ``card`` declares its data/get returns a file tree: its files, each a part by its path.""" + return any(returned.get("format") == "file-tree" for returned in (intake_declaration(card) or {}).get("get") or []) + + +async def get_data(base_url: str, environment_name: str, timeout_seconds: float, upload: Optional[_SnapshotUpload], + **params) -> tuple[GetDataResponse, Optional[_SnapshotUpload]]: + """``data/get`` of the environment at ``base_url`` with ``params``, handed ``upload``'s grant when there is one, + and what it answered with the upload it answered under: asked again without the grant when it fails under it, as + an env that can't reach the store does, a computer whose network is closed say.""" + if upload is not None: + try: + return await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), write_namespace=upload.grant, + **params), upload + except RuntimeError as e: + logger.info("snapshot_env: %s's data/get failed under the upload grant (%s); asking again without one", + environment_name, e) + return await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **params), None + + +async def _file_tree_files(resp: GetDataResponse, upload: Optional[_SnapshotUpload], environment_name: str, + files_id: str, description: str) -> dict: + """Each file ``environment_name``'s file tree returned, by its path, written to the store as ``derive_path_id( + files_id, path)``: an upload under the grant in place, with no download, else from its inline bytes.""" + from agent_env.artifact import FileArtifact + + writing = asyncio.Semaphore(_WRITES_IN_FLIGHT) + + async def written(part) -> tuple[str, FileArtifact]: + path = uploaded_object_path(part) + if path is None: + raise ValueError(f"'{environment_name}' returned a {part.kind} part naming no path in its file tree") + file_id, filename = derive_path_id(files_id, path), PurePosixPath(path).name + content_type, about = part.file.mimeType or "application/octet-stream", f"{path} of the {description}" + async with writing: + if isinstance(part.file, FileWithBytes): + return path, await asyncio.to_thread(FileArtifact.put_bytes, file_id, description=about, filename=filename, + content=base64.b64decode(part.file.bytes), content_type=content_type) + if upload is None: + raise ValueError(f"'{environment_name}' returned {path} neither inline nor uploaded") + return path, await asyncio.to_thread(registered_upload, upload.object_url(path), file_id, filename, + content_type, about) + + files = dict(await asyncio.gather(*(written(part) for part in resp.parts))) + if not files: + raise ValueError(f"'{environment_name}' returned no files to snapshot") + return files + + +def registered_upload(object_url: str, file_id: str, filename: str, content_type: str, description: str): + """A file artifact ``file_id`` over ``object_url``, an object a service uploaded itself, written without + downloading it; refused when the store holds no object there.""" + from agent_env.artifact import FileArtifact + from agent_env.artifact.store import get_artifact_store + + if get_config().get_object_store().get_object_metadata_at(object_url) is None: + raise FileNotFoundError(f"snapshot bundle not found at {object_url}") + get_config().check_object_url(file_id, object_url) + return get_artifact_store().put_document(FileArtifact( + id=file_id, description=description, filename=filename, content_type=content_type, object_url=object_url)) + + +def _card_environments(deployed) -> list[str]: + """The environments a deployment without a gateway serves, as the card it recorded names them: its children, or + the env itself.""" + card = deployed.environment_card or {} + children = card.get("children_environments") + names = [child.get("name") for child in children] if isinstance(children, list) else [card.get("name")] + if not all(names): + raise RuntimeError(f"Deployed env '{deployed.env_id}' recorded a card naming no environment to snapshot") + return names + + def _changes_route(deployed, override_url: Optional[str]) -> tuple[str, str]: """The verb and URL of the env's changes: an override keeps `{override}/changes`, else the stored card names them.""" if override_url: diff --git a/tst/integration/computer/computer_env_test.py b/tst/integration/computer/computer_env_test.py index 5463a3ea..0f046784 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -3,7 +3,10 @@ computer, and gate 3 makes the image of this machine's arch ready; that image boots with everything running as the session user and passes the conformance probe at full three times in a row; another env that changes only what doesn't change the image reuses it; one with no network, deployed, is reached but reaches nothing, takes a file and runs -code that reaches its driver on loopback, comes back from a reset fresh, and leaves nothing once torn down; and a recording open when the computer resets or is torn down is ended by +code that reaches its driver on loopback, is snapshotted by snapshot_computer_env, its files under the paths it names +kept by their paths in its environment artifact after its upload under the grant fails, comes back from a reset fresh, takes the snapshot back +through load_artifact, and leaves nothing once torn down; an open one uploads each file under the snapshot's grant, kept where it was put; and a recording open when the +computer resets or is torn down is ended by the computer and registered whole. Needs Docker and the local registry on localhost:5000.""" import asyncio @@ -19,6 +22,7 @@ from agentenv_protocol.computer.conformance.client import Driver from agent_env.artifact import FileArtifact +from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact from agent_env.artifact.artifacts.computer import ComputerArtifact from agent_env.env import Env from agent_env.env.envs.computer.env import ComputerEnv @@ -35,6 +39,7 @@ from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep from agent_env.task_step.task_steps.reset_env import ResetEnvTaskStep from agent_env.task_step.task_steps.run_code import RunCodeTaskStep +from agent_env.task_step.task_steps.computer.snapshot_computer_env import SnapshotComputerEnvTaskStep from tst.integration.computer.conftest import docker pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -118,15 +123,23 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an driver = await Driver.connect(deployed.environment_url, 30) egress = await driver.run("python3", "-c", "import urllib.request; urllib.request.urlopen('http://1.1.1.1', timeout=5)") await driver.run("sh", "-c", "echo x > Desktop/left.txt") + context = await SnapshotComputerEnvTaskStep(id="snapshot", version=None, env_id=closed.id, + snapshot_id=f"{closed.id}-snapshot", + file_paths=["Desktop", "Downloads/note.txt"]).execute(context) + snapshot = {path: file.load() for path, file in _snapshotted(context, closed).items()} await ResetEnvTaskStep(id="reset", version=None, env_id=closed.id, fail_task_on_error=True).execute(context) left = await (await Driver.connect(deployed.environment_url, 30)).run("test", "-e", "Desktop/left.txt") + context = await LoadArtifactTaskStep(id="restore", version=None, env_id=closed.id, + artifact_from_step_id="snapshot").execute(context) + restored = await (await Driver.connect(deployed.environment_url, 30)).run("cat", "Desktop/left.txt") finally: report = await teardown_run(context) sandbox = deployed.sandbox_id assert context.metadata["script_results"]["read"] == {"note": "inside\n", "card": "closed"} + assert snapshot == {"Desktop/left.txt": b"x\n", "Downloads/note.txt": b"inside\n"} assert egress["exit_code"] != 0 and "Network is unreachable" in egress["stderr"], egress - assert left["exit_code"] != 0 + assert left["exit_code"] != 0 and (restored["exit_code"], restored["stdout"]) == (0, "x\n"), restored assert [s.sandbox_id for s in report.terminated] == [sandbox] and not report.failed for listing in (["ps", "-aq", "--filter", f"label=agentenv.sandbox={sandbox}"], ["ps", "-aq", "--filter", f"name=agent-{sandbox}"], @@ -135,6 +148,28 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an assert LocalSandbox.find_work_dir(sandbox) is None +@pytest.mark.asyncio +async def test_an_open_computer_uploads_each_file_under_the_snapshots_grant_and_it_is_kept_where_it_was_put(desktop): + with namespace_routing(): + deployed = await Env.get(desktop.id).deploy() + context = TaskStepContext(deployed_envs=[deployed]) + try: + driver = await Driver.connect(deployed.environment_url, 30) + await driver.run("sh", "-c", "mkdir -p Documents/reports && printf one > Documents/reports/a.txt " + "&& printf two > Downloads/b.txt") + context = await SnapshotComputerEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, + snapshot_id=f"{desktop.id}-snapshot").execute(context) + snapshot = {path: (file.object_url, file.load()) for path, file in _snapshotted(context, desktop).items()} + finally: + report = await teardown_run(context) + grant_server(None, None).close() + + assert not report.failed + assert {path: data for path, (_, data) in snapshot.items()} == {"Documents/reports/a.txt": b"one", + "Downloads/b.txt": b"two"} + assert all("/agentenv-snapshots/" in url and url.endswith(path) for path, (url, _) in snapshot.items()), snapshot + + @pytest.mark.asyncio async def test_a_recording_open_at_a_reset_or_at_teardown_is_ended_by_its_computer_and_registered_whole(desktop): with namespace_routing(): @@ -176,3 +211,11 @@ async def _ready(name: str) -> str: pass await asyncio.sleep(0.25) raise TimeoutError(f"{name}'s computer never became ready") + + +def _snapshotted(context: TaskStepContext, env) -> dict: + """The computer's files in the universe the step ``snapshot`` wrote, by path.""" + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + (computer,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() + assert computer.environment_name == env.environment_name + return computer.get_file_artifacts() diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index 87e14ad8..36c807f5 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -1,7 +1,7 @@ """Computer envs: gate 1 refuses a base, a driver or a declaration that can't compose, naming every problem, before anything is stored; the composer builds each form and arch, reuses an image of the same key whichever env built it, -records each result as it finishes, and writes the arches built as the env's computer; a deploy, a reset and a load; -and computer apps. Where and how a pair builds is test_computer_builder; a computer of several apps, test_computer_apps. +records each result as it finishes, and writes the arches built as the env's computer; a deploy, a reset and a load, +of a file universe or of a snapshot of the computer's own files; and computer apps. Where and how a pair builds is test_computer_builder; a computer of several apps, test_computer_apps. Docker and the registry are faked.""" import asyncio @@ -25,6 +25,8 @@ from agent_env.artifact.artifacts.computer import ComputerArtifact from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.artifacts.environment import EnvironmentArtifact +from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact from agent_env.artifact.artifacts.file import FileArtifact from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse from agent_env.artifact.store import get_artifact_store @@ -782,6 +784,45 @@ async def test_a_load_outside_the_home_is_refused_before_anything_is_sent(base, assert computer.calls == [] +def _snapshot(environment_name: str, contents: dict[str, bytes]) -> EnvironmentArtifact: + """An environment artifact of ``contents`` by path, as snapshot_computer_env keeps a computer's files.""" + files = {path: FileArtifact.put_bytes(f"snap-{environment_name}-{index}", description="snapshot", + filename=path.rsplit("/", 1)[-1], content=data, content_type="text/plain") + for index, (path, data) in enumerate(contents.items())} + return EnvironmentArtifact.put(f"snap-{environment_name}", environment_name=environment_name, file_artifacts=files) + + +@pytest.mark.asyncio +async def test_a_snapshot_of_its_files_loads_back_at_their_paths_leaving_out_other_environments( + base, driver, docker, computer, tmp_path, monkeypatch, caplog): + monkeypatch.setattr(inside, "issues_grants_to", lambda store, sandbox_type: False) + env = _deployed_env(base, driver, computer) + contents = {"Desktop/a.txt": b"a", "Documents/r/b.txt": b"b"} + universe = EnvironmentUniverseArtifact.put("snap", environment_artifacts=[ + _snapshot(env.environment_name, contents), _snapshot("mail", {"inbox/1.json": b"{}"})]) + + with caplog.at_level(logging.WARNING): + await env.load_environment_universe_artifact(universe) + + assert {path: (computer.home / path).read_bytes() for path in contents} == contents + assert not (computer.home / "inbox").exists() + assert "leaving out ['mail']" in caplog.text + + +@pytest.mark.asyncio +async def test_an_environment_not_named_after_the_computer_is_refused_before_anything_is_sent(base, driver, docker, + computer): + env = _deployed_env(base, driver, computer) + mail = _snapshot("mail", {"inbox/1.json": b"{}"}) + + with pytest.raises(ValueError, match="does not match env environment_name"): + await env.load_environment_artifact(mail) + with pytest.raises(RuntimeError, match=f"has no environment named '{env.environment_name}'"): + await env.load_environment_universe_artifact(EnvironmentUniverseArtifact.put("mail-only", + environment_artifacts=[mail])) + assert computer.calls == [] + + PROTOCOL_VERSION = importlib.metadata.version("agentenv-framework-protocol") diff --git a/tst/unit/task/local_run_write_fence_test.py b/tst/unit/task/local_run_write_fence_test.py index eb8896ff..d3f60293 100644 --- a/tst/unit/task/local_run_write_fence_test.py +++ b/tst/unit/task/local_run_write_fence_test.py @@ -264,7 +264,7 @@ def __init__(self, id, version=None): self.mcp_server_envs = [SimpleNamespace(environment_name="gmail")] -async def _export(gateway_url, name, tmp_path, timeout, deployed=None): +async def _export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): Path(tmp_path).write_text(json.dumps({"service": name})) return ".json" diff --git a/tst/unit/task_step/snapshot_computer_env_test.py b/tst/unit/task_step/snapshot_computer_env_test.py new file mode 100644 index 00000000..0fb423c7 --- /dev/null +++ b/tst/unit/task_step/snapshot_computer_env_test.py @@ -0,0 +1,142 @@ +"""snapshot_computer_env: a computer env's files kept by their paths in its environment artifact, in a new environment +universe, returned inline or uploaded under the grant and written where they were put, the computer asked again without +the grant when it fails under it, and file_paths sent to its data/get; refused for an env that isn't a computer, for +paths that aren't relative, and with no files to keep.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from agentenv_protocol import uploaded_file_part + +from agent_env.artifact import EnvironmentArtifact, EnvironmentUniverseArtifact +from agent_env.config import set_object_store +from agent_env.env.envs.computer.env import ComputerEnv +from agent_env.task_step.registry import get_task_step_registry +from agent_env.task_step.task_steps.computer.snapshot_computer_env import SnapshotComputerEnvTaskStep +from tst.unit.task_step.snapshot_env_without_gateway_test import (DATA_OBJECTS, UNREACHABLE, _computer, _inline, _record, + _serve) +from tst.unit.task_step.test_snapshot_env import _ctx +from tst.util.granting_object_store import GrantingObjectStore + + +class _Desk(ComputerEnv): + """A stored computer env, as far as a snapshot reads one.""" + + def __init__(self): + self.id, self.version, self.environment_name = "desk", 1, "desktop" + + +class _Mail: + id, version, type, environment_name = "mail", 1, "mcp_server", "mail" + + +async def _snapshot(monkeypatch, card: dict, env=None, **fields) -> EnvironmentArtifact: + """The computer's environment artifact in the universe the step snapshots it into.""" + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: env or _Desk())) + context = _ctx(_record(card)) + await SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk", **fields).execute(context) + snapshotted = context.metadata["env_snapshotted_universes"]["snap"] + (desktop,) = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]).get_environment_artifacts() + assert desktop.environment_name == "desktop" + return desktop + + +def _kept(desktop: EnvironmentArtifact) -> dict[str, bytes]: + return {path: file.load() for path, file in desktop.get_file_artifacts().items()} + + +def test_its_fields_round_trip_and_its_type_is_registered(): + step = SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk", file_paths=["Documents", "Desktop/a.pdf"], + export_timeout_seconds=60) + + data = step.to_dict() + restored = get_task_step_registry()["snapshot_computer_env"].from_dict(data) + + assert isinstance(restored, SnapshotComputerEnvTaskStep) and data["type"] == "snapshot_computer_env" + assert (restored.env_id, restored.file_paths, restored.export_timeout_seconds) == ( + "desk", ["Documents", "Desktop/a.pdf"], 60) + assert not {"gateway_url", "include_env_trajectory", "include_env_changes"} & set(data) + + +@pytest.mark.parametrize("paths", [["/etc/passwd"], ["../outside"], ["Documents/../.."]], + ids=["absolute", "parent", "escapes"]) +def test_file_paths_not_relative_to_the_home_are_refused_before_a_run(paths): + step = SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk", file_paths=paths) + + assert [problem for problem in step.config_problems() if "file_paths are normalized paths" in problem] + + +@pytest.mark.parametrize("stored, refused", [(_Mail(), True), (_Desk(), False)], ids=["mcp-server", "computer"]) +def test_an_env_that_isnt_a_computer_is_refused_before_a_run(monkeypatch, stored, refused): + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: stored)) + + problems = SnapshotComputerEnvTaskStep(id="snap", version=None, env_id=stored.id).preflight() + + assert problems == (["snapshot_computer_env snapshots a computer env, and mail is a mcp_server env"] + if refused else []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("file_paths, params", [(["Documents"], {"paths": ["Documents"]}), (None, {})], + ids=["file-paths", "the-cards-folders"]) +async def test_each_file_returned_inline_is_kept_by_its_path(local_stores, monkeypatch, file_paths, params): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello"), + _inline("Documents/b/c.bin", b"\x00\x01")]}) + + desktop = await _snapshot(monkeypatch, _computer(), file_paths=file_paths) + + assert _kept(desktop) == {"Documents/notes.txt": b"hello", "Documents/b/c.bin": b"\x00\x01"} + assert desktop.get_file_artifacts()["Documents/b/c.bin"].filename == "c.bin" + assert sent == [("/agentenv", params)] + + +@pytest.mark.asyncio +async def test_files_uploaded_under_the_grant_are_written_where_the_computer_put_them_without_a_download( + local_stores, tmp_path, monkeypatch): + store = GrantingObjectStore(str(tmp_path / "objects")) + set_object_store(store) + files = {"Documents/notes.txt": b"hello", "Downloads/c.bin": b"\x00"} + + def uploading(params): + root = params["write_namespace"]["root_path"] + for path, data in files.items(): + store.put(f"{root}/{path}", data, content_type="application/octet-stream") + return [uploaded_file_part(path, name=Path(path).name, mime_type="text/plain").model_dump( + mode="json", exclude_none=True) for path in files] + + sent = _serve(monkeypatch, {"/agentenv": uploading}) + + desktop = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) + + ((_, params),) = sent + root = params["write_namespace"]["root_path"] + assert params["write_namespace"]["max_objects"] == 100_000 + assert {path: (file.object_url, file.load()) for path, file in desktop.get_file_artifacts().items()} == { + path: (store.object_url(f"{root}/{path}"), data) for path, data in files.items()} + + +@pytest.mark.asyncio +async def test_a_computer_that_fails_under_the_grant_is_asked_again_without_one(local_stores, tmp_path, monkeypatch): + set_object_store(GrantingObjectStore(str(tmp_path / "objects"))) + sent = _serve(monkeypatch, {"/agentenv": lambda params: UNREACHABLE if "write_namespace" in params + else [_inline("Desktop/left.txt", b"x")]}) + + desktop = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) + + assert _kept(desktop) == {"Desktop/left.txt": b"x"} + assert [sorted(params) for _, params in sent] == [["write_namespace"], []] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("env, answers, refused", [ + (_Mail(), {}, "snapshots a computer env, and mail is a mcp_server env"), + (None, {"/agentenv": lambda params: []}, "'desktop' returned no files to snapshot"), +], ids=["not-a-computer", "no-files"]) +async def test_a_run_with_nothing_a_computer_snapshot_can_keep_is_refused(local_stores, monkeypatch, env, answers, + refused): + _serve(monkeypatch, answers) + + with pytest.raises(RuntimeError, match=refused): + await _snapshot(monkeypatch, _computer(), env=env, file_paths=["Documents"]) diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py new file mode 100644 index 00000000..d201c111 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -0,0 +1,154 @@ +"""snapshot_env on an env deployed without a gateway: reached through the card its deployment recorded, its +environments the card's children or the env itself, each kept as one file as through a gateway; an environment whose +data/get returns a file tree, a computer's, kept file by file by its paths under a grant sized for a file tree; and an +env that fails under the upload grant asked again without it.""" + +from __future__ import annotations + +import base64 +import json +from pathlib import Path + +import httpx +import pytest +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, IntakeDeclaration, IntakeFormat, intake_extension + +from agent_env.artifact import EnvironmentUniverseArtifact +from agent_env.config import set_object_store +from agent_env.env.env import DeployedSandboxEnv +from agent_env.env.gateway.constants import WELL_KNOWN_PATH +from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep +from tst.unit.task_step.test_snapshot_env import _ctx +from tst.util.granting_object_store import GrantingObjectStore + +COMPUTER = "http://127.0.0.1:18765" +FILE_TREE = intake_extension(IntakeDeclaration(get=[IntakeFormat(part="file", format="file-tree")])).model_dump( + mode="json", exclude_none=True) +DATA_OBJECTS = {"uri": DATA_OBJECTS_EXTENSION_URI} +UNREACHABLE = {"error": {"code": -32000, "message": "The object store is temporarily unavailable.", + "data": "get_failed"}} + + +class _Desk: + id, version, environment_name = "desk", 1, "desktop" + + +def _record(card: dict) -> DeployedSandboxEnv: + return DeployedSandboxEnv(env_id="desk", env_version=1, sandbox_id="sb-1", instance_id="desk-abc", + environment_card_url=f"{COMPUTER}{WELL_KNOWN_PATH}", environment_card=card, + sandbox_type="local") + + +def _computer(*extensions) -> dict: + return {"name": "desktop", "capabilities": {"extensions": [FILE_TREE, *extensions]}} + + +def _server(*extensions) -> dict: + return {"name": "desktop", "capabilities": {"extensions": list(extensions)}} + + +def _inline(path: str, data: bytes) -> dict: + return {"kind": "file", "file": {"bytes": base64.b64encode(data).decode(), "name": Path(path).name, + "mimeType": "text/plain"}, + "metadata": {DATA_OBJECTS_EXTENSION_URI: {"path": path}}} + + +def _serve(monkeypatch, answers: dict) -> list[tuple[str, dict]]: + """Answer each data/get with the parts ``answers`` gives for the path it's sent to, or the JSON-RPC error when it + gives one, recording each path and its params; any other request fails the test.""" + sent, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + params = json.loads(request.content)["params"] + sent.append((request.url.path, params)) + answer = answers[request.url.path](params) + body = {"error": answer["error"]} if isinstance(answer, dict) else {"result": {"parts": answer}} + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, **body}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return sent + + +async def _snapshot(card: dict) -> EnvironmentUniverseArtifact: + result = await SnapshotEnvTaskStep.snapshot_env_state(env=_Desk(), gateway_url=None, snapshot_id="snap-d", + deployed=_record(card)) + return EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id) + + +def _kept(universe: EnvironmentUniverseArtifact) -> dict[str, bytes]: + """Each environment's one file's bytes, by the environment's name.""" + return {name.split("/", 1)[0]: file.load() for name, file in universe.get_file_artifacts().items()} + + +@pytest.mark.asyncio +async def test_each_child_on_the_recorded_card_is_reached_and_kept_as_one_file(local_stores, monkeypatch): + card = {"name": "desk", "children_environments": [ + {"name": "mail", "url": "/svc/mail/agentenv", "capabilities": {"extensions": []}}, + {"name": "notes", "url": "/svc/notes/agentenv", "capabilities": {"extensions": []}}]} + sent = _serve(monkeypatch, {"/svc/mail/agentenv": lambda params: [{"kind": "data", "data": {"inbox": 2}}], + "/svc/notes/agentenv": lambda params: [_inline("notes.zip", b"PK")]}) + + universe = await _snapshot(card) + + assert _kept(universe) == {"mail": b'{"inbox": 2}', "notes": b"PK"} + assert sorted(path for path, _ in sent) == ["/svc/mail/agentenv", "/svc/notes/agentenv"] + + +@pytest.mark.asyncio +async def test_the_env_itself_is_reached_when_its_card_has_no_children(local_stores, monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [{"kind": "data", "data": {"rows": 1}}]}) + + universe = await _snapshot(_server()) + + assert _kept(universe) == {"desktop": b'{"rows": 1}'} + assert sent == [("/agentenv", {})] + + +@pytest.mark.asyncio +async def test_a_file_tree_is_kept_file_by_file_by_its_paths_under_a_grant_sized_for_it(local_stores, tmp_path, + monkeypatch): + set_object_store(GrantingObjectStore(str(tmp_path / "objects"))) + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Desktop/a.txt", b"a"), _inline("b.txt", b"b")]}) + + universe = await _snapshot(_computer(DATA_OBJECTS)) + + (desktop,) = universe.get_environment_artifacts() + assert {path: file.load() for path, file in desktop.get_file_artifacts().items()} == {"Desktop/a.txt": b"a", + "b.txt": b"b"} + ((_, params),) = sent + assert params["write_namespace"]["max_objects"] == 100_000 + + +@pytest.mark.asyncio +async def test_an_env_that_fails_under_the_grant_is_asked_again_without_one(local_stores, tmp_path, monkeypatch): + set_object_store(GrantingObjectStore(str(tmp_path / "objects"))) + sent = _serve(monkeypatch, {"/agentenv": lambda params: UNREACHABLE if "write_namespace" in params + else [{"kind": "data", "data": {"rows": 1}}]}) + + universe = await _snapshot(_server(DATA_OBJECTS)) + + assert _kept(universe) == {"desktop": b'{"rows": 1}'} + assert [(path, sorted(params)) for path, params in sent] == [("/agentenv", ["write_namespace"]), ("/agentenv", [])] + + +@pytest.mark.asyncio +async def test_an_env_that_fails_with_no_grant_is_asked_once(local_stores, monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: UNREACHABLE}) + + with pytest.raises(RuntimeError, match="data/get failed: The object store is temporarily unavailable"): + await _snapshot(_server()) + assert sent == [("/agentenv", {})] + + +@pytest.mark.asyncio +async def test_execute_snapshots_an_env_deployed_without_a_gateway_through_its_recorded_card(local_stores, + monkeypatch): + _serve(monkeypatch, {"/agentenv": lambda params: [{"kind": "data", "data": {"rows": 1}}]}) + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Desk())) + context = _ctx(_record(_server())) + + await SnapshotEnvTaskStep(id="snap", version=None, env_id="desk").execute(context) + + snapshotted = context.metadata["env_snapshotted_universes"]["snap"] + assert _kept(EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"])) == { + "desktop": b'{"rows": 1}'} diff --git a/tst/unit/task_step/test_field_scopes.py b/tst/unit/task_step/test_field_scopes.py index 86aa4948..54562a0e 100644 --- a/tst/unit/task_step/test_field_scopes.py +++ b/tst/unit/task_step/test_field_scopes.py @@ -23,6 +23,7 @@ "rubrics_verifier": {"prompt_id": "prompt", "criteria": [{"id": "c1", "description": "says hi"}]}, "run_docker_container": {"sandbox_name": "box"}, "snapshot_env": {}, + "snapshot_computer_env": {}, } # What a step without a user-scoped field that has no default says it needs, before a run starts. UNSOURCED = { diff --git a/tst/unit/task_step/test_load_artifact_service_universe.py b/tst/unit/task_step/test_load_artifact_service_universe.py index a4757eed..f102e4d7 100644 --- a/tst/unit/task_step/test_load_artifact_service_universe.py +++ b/tst/unit/task_step/test_load_artifact_service_universe.py @@ -25,6 +25,7 @@ from agent_env.env.env import DeployedEnv, DeployedGatewayEnv, Env from agent_env.env.envs.computer.env import ComputerEnv from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.env.envs.service_db import ServiceDBEnv from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps import load_artifact as load_artifact_mod from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep @@ -206,10 +207,11 @@ async def get_sandbox(self, sandbox_id): @pytest.mark.asyncio @pytest.mark.parametrize("env_class, problems", [ - (ComputerEnv, ["load_artifact 'load': computer env 'desk' loads no environment_universe artifacts, so not " - f"'{UNIVERSE_ID}'"]), + (ServiceDBEnv, ["load_artifact 'load': service_db env 'desk' loads no environment_universe artifacts, so not " + f"'{UNIVERSE_ID}'"]), + (ComputerEnv, []), (MCPServerEnv, []), -], ids=["computer", "mcp-server"]) +], ids=["service-db", "computer", "mcp-server"]) async def test_an_env_whose_type_loads_no_environment_artifacts_is_refused_at_save_and_at_run(universe, monkeypatch, env_class, problems): env = env_class.__new__(env_class) diff --git a/tst/unit/task_step/test_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index dd4087b8..b118bc39 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -145,9 +145,10 @@ def test_snapshot_id_falls_back_to_an_adhoc_run_without_instance_id(): @pytest.mark.asyncio -async def test_execute_refuses_an_env_deployed_without_a_gateway(): +async def test_execute_refuses_an_env_deployed_without_a_gateway_or_a_recorded_card(): bare = DeployedSandboxEnv(env_id="env-1", env_version=3, sandbox_id="srv", instance_id="env-1-abc123") - with pytest.raises(RuntimeError, match="Deployed env 'env-1' has no gateway_url and no override was provided"): + with pytest.raises(RuntimeError, match="Deployed env 'env-1' has no gateway_url, no recorded card to reach it " + "through, and no override was provided"): await _step(env_id="env-1").execute(_ctx(bare)) @pytest.mark.asyncio @@ -556,7 +557,7 @@ def __init__(self, names): @pytest.mark.asyncio async def test_core_all_or_nothing_on_partial_failure(monkeypatch): - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): if name == "slack": raise ValueError("export-state returned non-JSON body (starts with b'<')") with open(tmp_path, "w") as f: @@ -600,7 +601,7 @@ def put(id, *, environment_name, file_artifact): @pytest.mark.asyncio async def test_core_success_creates_universe_with_metadata(monkeypatch): - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): with open(tmp_path, "w") as f: json.dump({"ok": name}, f) return ".json" @@ -664,7 +665,7 @@ async def test_core_zip_export_creates_zip_filename_artifact(monkeypatch): """A FilePart(zip) export flows through as a .zip-suffixed FileArtifact path, so FileArtifact.filename (= basename) ends in .zip — the load side routes a .zip file:// artifact to reset_data() for a lossless round-trip.""" - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): # Simulate the lossless-zip case: write zip bytes and report ".zip". with open(tmp_path, "wb") as f: f.write(b"PK\x03\x04zip-bytes") @@ -902,7 +903,7 @@ async def test_core_registers_a_service_upload_in_place_and_it_loads_back(local_ bundle_file.write_bytes(b"PK\x03\x04bundle") uploaded = store.put_file_at(upload_url, str(bundle_file)) - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): return uploaded monkeypatch.setattr(SnapshotEnvTaskStep, "_export_environment_to_file", staticmethod(fake_export)) @@ -924,7 +925,7 @@ async def test_core_fails_a_service_whose_upload_the_store_cannot_find(local_sto set_object_store(FakeObjectStore()) reset_artifact_store() - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): return "fake://shared/snapshots/missing.zip" monkeypatch.setattr(SnapshotEnvTaskStep, "_export_environment_to_file", staticmethod(fake_export)) @@ -941,7 +942,7 @@ async def test_core_snapshots_under_an_id_longer_than_a_filename_can_be(local_st local store may be the whole id.""" snapshot_id = f"task-{'x' * 250}__snapshot-snap-1" - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): with open(tmp_path, "w") as f: json.dump({"service": name}, f) return ".json" @@ -1037,7 +1038,7 @@ async def test_export_of_a_child_env_missing_from_the_card_streams_legacy_export async def test_snapshot_uses_the_stored_card_only_for_the_deployments_own_gateway(monkeypatch): seen = [] - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): seen.append((gateway_url, deployed)) raise ValueError("stop after the export call") @@ -1130,7 +1131,7 @@ def get_environment_universe(self, instance_id): def _install_core_artifact_fakes(monkeypatch, universe_cls): """Wire the export + FileArtifact/EnvironmentArtifact fakes shared by the core carry-over tests; the caller supplies the EnvironmentUniverseArtifact fake.""" - async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None): + async def fake_export(gateway_url, name, tmp_path, timeout, deployed=None, data_get_params=None): with open(tmp_path, "w") as f: json.dump({"ok": name}, f) return ".json"