From 87d462b6041d247e7091f6c6f2593b6164f7a25f Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 10:19:43 -0700 Subject: [PATCH 1/8] feat(snapshot): snapshot envs without a gateway, keeping each environment's export as its data/get returned it snapshot_env reaches an env deployed without a gateway, a computer env say, through the card its deployment recorded, and snapshots the environments the card lists as children, or the env itself. An export one file can't keep, several parts or the files of an env whose card declares its data/get returns a file tree, is stored as its parts (EnvironmentArtifact.parts: each part's kind, content, name, media type and metadata), so data/add can be handed them back; any other export is written exactly as before. computer_file_paths names what a computer env's own data/get returns. The snapshot grant allows a file tree's files, uploaded each where the computer put it, and a computer whose network keeps it from the store is handed none. Co-Authored-By: Claude Opus 5.5 --- .../artifact/artifacts/environment.py | 64 ++++- .../artifacts/environment_universe.py | 28 +- .../env/envs/computer/utils/inside.py | 12 +- .../task_step/task_steps/snapshot_env.py | 270 ++++++++++++------ tst/integration/computer/computer_env_test.py | 39 ++- tst/unit/artifact/environment_parts_test.py | 62 ++++ .../test_service_universe_file_view.py | 29 +- .../snapshot_env_upload_grant_test.py | 2 +- .../snapshot_env_without_gateway_test.py | 172 +++++++++++ tst/unit/task_step/test_snapshot_env.py | 9 +- 10 files changed, 573 insertions(+), 114 deletions(-) create mode 100644 tst/unit/artifact/environment_parts_test.py create mode 100644 tst/unit/task_step/snapshot_env_without_gateway_test.py diff --git a/src/agent_env/artifact/artifacts/environment.py b/src/agent_env/artifact/artifacts/environment.py index 957a27875..93303b0d7 100644 --- a/src/agent_env/artifact/artifacts/environment.py +++ b/src/agent_env/artifact/artifacts/environment.py @@ -1,4 +1,5 @@ -"""Environment artifact for wrapping a FileArtifact with environment identity.""" +"""Environment artifact for wrapping a FileArtifact with environment identity, or the parts an environment's data/get +returned.""" from __future__ import annotations @@ -6,16 +7,21 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional -from pydantic import ConfigDict, Field, model_serializer +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI +from pydantic import BaseModel, ConfigDict, Field, model_serializer from agent_env.artifact.artifact import Artifact from agent_env.artifact.artifacts.file import FileArtifact from agent_env.artifact.ref import ArtifactRef +from agent_env.entity.export import Unwritable from agent_env.entity.refs import EntityRef from agent_env.store.ids import derive_id logger = logging.getLogger(__name__) +#: Why an environment of parts has no artifact folder yet: a folder holds one file per environment. +PARTS_UNWRITABLE = "holds the parts its data/get returned, which an artifact folder can't hold yet" + if TYPE_CHECKING: from agent_env.bundle.authoring import AuthoringContext from agent_env.entity.export import ExportContext @@ -32,8 +38,28 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. path ``""`` for the folder itself. Raises BundleError when the folder can't be read.""" +class EnvironmentPart(BaseModel): + """One part of an environment's data as its data/get returned it, which data/add can be handed back: a file part + or a data part, its content stored as a file artifact, with the part's name, media type and metadata.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["file", "data"] + file_artifact_ref: ArtifactRef + name: Optional[str] = None + mime_type: Optional[str] = None + metadata: Optional[dict[str, Any]] = None + + @property + def path(self) -> Optional[str]: + """The path its metadata names it by under the data-objects marker, as a computer names each file.""" + marker = (self.metadata or {}).get(DATA_OBJECTS_EXTENSION_URI) + return marker.get("path") if isinstance(marker, dict) else None + + class EnvironmentArtifact(Artifact, SeedsEnvironments): - """A FileArtifact wrapped with the name of the environment it seeds.""" + """A FileArtifact wrapped with the name of the environment it seeds, or the ``parts`` its data/get returned when + one file would lose some of them: several parts, or the files of an environment that returns a file tree.""" model_config = ConfigDict(populate_by_name=True) @@ -56,6 +82,10 @@ class EnvironmentArtifact(Artifact, SeedsEnvironments): alias="file_artifact_id", description="Legacy unpinned wrapped-FileArtifact id", ) + parts: Optional[list[EnvironmentPart]] = Field( + default=None, + description="The parts the environment's data/get returned, in order, in place of one wrapped FileArtifact", + ) # No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields. @model_serializer(mode="wrap") @@ -65,8 +95,10 @@ def _serialize(self, handler: Any): data = handler(self) if "legacy_file_artifact_id" in data: data["file_artifact_id"] = data.pop("legacy_file_artifact_id") - if self.file_artifact_ref is not None: # new doc + if self.file_artifact_ref is not None or self.parts is not None: # new doc data.pop("file_artifact_id", None) + if self.parts is None: + data.pop("parts", None) # Dual-write. Sourced from the attribute, not from `data`, # because `alias=` emits exactly one spelling and which one depends on the # caller's `by_alias` — the store persists with by_alias=True @@ -81,16 +113,22 @@ def put( id: str, *, environment_name: Optional[str] = None, - file_artifact: FileArtifact, + file_artifact: Optional[FileArtifact] = None, + parts: Optional[list[EnvironmentPart]] = None, ) -> EnvironmentArtifact: + """Write the environment of ``environment_name`` over ``file_artifact``, or over ``parts``, the parts its + data/get returned.""" from agent_env.artifact.store import get_artifact_store if not environment_name: raise ValueError("environment_name cannot be empty") + if (file_artifact is None) == (parts is None): + raise ValueError("an environment artifact wraps a file_artifact or holds parts, one of the two") instance = cls( id=id, environment_name=environment_name, - file_artifact_ref=file_artifact.as_ref(), + file_artifact_ref=file_artifact.as_ref() if file_artifact is not None else None, + parts=parts, ) return get_artifact_store().put_document(instance) @@ -115,6 +153,8 @@ def from_toml(cls, data: dict, ctx: AuthoringContext) -> EnvironmentArtifact: def to_toml(self, out: ExportContext) -> dict: """The artifact.toml of a folder holding the environment's file, which ``from_toml`` reads back as this.""" + if self.parts is not None: + raise Unwritable(f"it {PARTS_UNWRITABLE}") file = self.get_file_artifact() out.file(file.filename, file.object_url) described = {} if file.description == file.filename else {"description": file.description} @@ -159,6 +199,9 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. return (("", name),) if isinstance(name, str) and name else () def get_file_artifact(self) -> FileArtifact: + if self.parts is not None: + raise ValueError(f"EnvironmentArtifact {self.id} v{self.version} holds the {len(self.parts)} parts its " + "environment's data/get returned, not one file") if self.file_artifact_ref is not None: return FileArtifact.get( self.file_artifact_ref.id, version=self.file_artifact_ref.version @@ -173,3 +216,12 @@ def get_file_artifact(self) -> FileArtifact: f"malformed EnvironmentArtifact {self.id} v{self.version}: " "neither file_artifact_ref nor file_artifact_id present" ) + + def get_parts(self) -> list[tuple[EnvironmentPart, FileArtifact]]: + """Each part and its content's file artifact, in order: the one file the environment wraps as a file part named + for it, or the parts its data/get returned.""" + if self.parts is None: + file = self.get_file_artifact() + return [(EnvironmentPart(kind="file", file_artifact_ref=file.as_ref(), name=file.filename), file)] + return [(part, FileArtifact.get(part.file_artifact_ref.id, version=part.file_artifact_ref.version)) + for part in self.parts] diff --git a/src/agent_env/artifact/artifacts/environment_universe.py b/src/agent_env/artifact/artifacts/environment_universe.py index 458677e0b..b0ad98c98 100644 --- a/src/agent_env/artifact/artifacts/environment_universe.py +++ b/src/agent_env/artifact/artifacts/environment_universe.py @@ -10,7 +10,7 @@ from pydantic import AliasChoices, ConfigDict, Field, model_serializer, model_validator from agent_env.artifact.artifact import FolderChild -from agent_env.artifact.artifacts.environment import EnvironmentArtifact, SeedsEnvironments +from agent_env.artifact.artifacts.environment import PARTS_UNWRITABLE, EnvironmentArtifact, SeedsEnvironments from agent_env.artifact.artifacts.environment_universe_folder import ( METADATA, UniverseFolder, @@ -184,6 +184,8 @@ def to_toml(self, out: ExportContext) -> dict: name = environment.environment_name if name.split("/", 1)[0] == METADATA: raise Unwritable(f"its environment {name!r} would read back as metadata, from its {METADATA}/ folder") + if environment.parts is not None: + raise Unwritable(f"its environment {name!r} {PARTS_UNWRITABLE}") file = environment.get_file_artifact() out.file(environment_path(Path(), name, file.filename).as_posix(), file.object_url) for key, file in (self.get_metadata() or {}).items(): @@ -260,18 +262,24 @@ def get_file_artifacts(self) -> dict[str, FileArtifact]: Matching ``FileArtifactUniverse.get_file_artifacts()`` is what lets the same loaders stage a service universe as a file tree (grading a frozen snapshot) rather than restoring it into an env. Returns the pinned artifacts as-is — - nothing is copied or re-registered. + nothing is copied or re-registered. An environment of parts lists each by the path + its metadata names, else its name, else its file's. """ out: dict[str, FileArtifact] = {} for sa in self.get_environment_artifacts(): - fa = sa.get_file_artifact() - name = f"{sa.environment_name}/{fa.filename}" - if name in out: - raise ValueError( - f"EnvironmentUniverseArtifact {self.id} v{self.version} has duplicate " - f"environment '{sa.environment_name}'; cannot stage it as a file tree" - ) - out[name] = fa + if sa.parts is None: + fa = sa.get_file_artifact() + named = [(fa.filename, fa)] + else: + named = [(part.path or part.name or fa.filename, fa) for part, fa in sa.get_parts()] + for relative, fa in named: + name = f"{sa.environment_name}/{relative}" + if name in out: + raise ValueError( + f"EnvironmentUniverseArtifact {self.id} v{self.version} has duplicate " + f"environment '{sa.environment_name}'; cannot stage it as a file tree" + ) + out[name] = fa return out def get_metadata(self) -> dict[str, FileArtifact]: diff --git a/src/agent_env/env/envs/computer/utils/inside.py b/src/agent_env/env/envs/computer/utils/inside.py index f1033e23c..7e3c92cd7 100644 --- a/src/agent_env/env/envs/computer/utils/inside.py +++ b/src/agent_env/env/envs/computer/utils/inside.py @@ -24,7 +24,7 @@ if TYPE_CHECKING: from agent_env.artifact import FileArtifact - from agent_env.env.env import DeployedSandboxEnv + from agent_env.env.env import DeployedEnv, DeployedSandboxEnv #: The most bytes one ``data/add`` carries: the driver reads a call's files into memory before it writes them. A file #: bigger than this goes alone. @@ -52,10 +52,16 @@ def loopback(deployed: DeployedSandboxEnv) -> dict: return {"id": deployed.env_id, "url": url, "mcp_url": url + deployed.mcp_url.removeprefix(deployed.environment_url)} +def egress_open(deployed: DeployedEnv) -> bool: + """Whether the network a computer's deploy recorded lets it reach the object store, to read or write through a + grant: one that allows everything, or none recorded, as for any env but a computer.""" + network = (deployed.metadata or {}).get("computer", {}).get("network") + return not (network and NetworkPolicy.from_dict(network).restricts_egress) + + async def add_artifacts(deployed: DeployedSandboxEnv, files: dict[str, FileArtifact]) -> None: """Write each file artifact at its path.""" - network = (deployed.metadata or {}).get("computer", {}).get("network") - granted = not (network and NetworkPolicy.from_dict(network).restricts_egress) + granted = egress_open(deployed) batch: list[FilePart] = [] size = 0 for path, artifact in files.items(): 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 49a108d81..f3966e368 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -1,8 +1,10 @@ -"""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 import base64 +import glob import json import logging import mimetypes @@ -11,16 +13,18 @@ import uuid from dataclasses import dataclass from datetime import datetime, timezone -from typing import ClassVar, Optional +from pathlib import Path +from typing import ClassVar, Optional, Union 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 DATA_OBJECTS_EXTENSION_URI, FilePart, Part, uploaded_object_path +from agentenv_protocol.client import find_extension, find_extension_method, intake_declaration from agentenv_protocol.transfers import WriteNamespaceGrant from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant +from agent_env.env.envs.computer.utils.inside import egress_open 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 @@ -45,9 +49,9 @@ # Must match the S3 credentials extension URI the server-side credentials mixin advertises. _S3_CREDENTIALS_EXTENSION_URI = "urn:agentenv:add-s3-credentials/v1" -# 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 a service may upload through its snapshot grant: one bundle, or the files of a file tree, each as large as one +# upload to the store can be, up to this in all. +ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=100_000, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" @@ -157,6 +161,12 @@ 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, a computer env say, is reached through the card its deployment recorded, its + environments the card's children, or itself. Each environment's export is kept as data/get returned it: one file + when it's one part, else its parts (``EnvironmentArtifact.parts``), as for a computer, whose data/get returns a + file tree. ``computer_file_paths``, relative to the session user's home, names what a computer env's own + data/get returns, by default the folders its card declares. """ type: ClassVar[str] = "snapshot_env" @@ -182,6 +192,7 @@ def __init__( export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, include_env_trajectory: bool = False, include_env_changes: bool = False, + computer_file_paths: Optional[list[str]] = None, depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, ): @@ -195,6 +206,7 @@ def __init__( self.export_timeout_seconds = export_timeout_seconds self.include_env_trajectory = include_env_trajectory self.include_env_changes = include_env_changes + self.computer_file_paths = computer_file_paths def to_dict(self) -> dict: base = super().to_dict() @@ -207,6 +219,8 @@ def to_dict(self) -> dict: base["export_timeout_seconds"] = self.export_timeout_seconds base["include_env_trajectory"] = self.include_env_trajectory base["include_env_changes"] = self.include_env_changes + if self.computer_file_paths is not None: + base["computer_file_paths"] = self.computer_file_paths return base @classmethod @@ -224,6 +238,7 @@ def from_dict(cls, data: dict) -> SnapshotEnvTaskStep: ), include_env_trajectory=data.get("include_env_trajectory", False), include_env_changes=data.get("include_env_changes", False), + computer_file_paths=data.get("computer_file_paths"), ) async def execute(self, context: TaskStepContext) -> TaskStepContext: @@ -238,9 +253,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: @@ -276,6 +292,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed=deployed, original_universe_artifact_id=self.original_universe_artifact_id, export_timeout_seconds=self.export_timeout_seconds, + computer_file_paths=self.computer_file_paths, ) logger.info( @@ -461,11 +478,15 @@ 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, + paths: Optional[list[str]] = None, + ) -> Union[str, list[tuple[Part, str]]]: """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). + When its data/get returns several parts, or a file tree (``_returns_file_tree``), + each part with the same for it, its file at ``-``; ``paths`` + goes to a data/get that takes it, a computer's. 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`` → @@ -478,63 +499,21 @@ async def _export_environment_to_file( base_url = await legacy_protocol.v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is not None: card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) - upload = await asyncio.to_thread( + # A grant is no use to an env whose network keeps it from the store, a closed computer say. + upload = None if deployed is not None and not egress_open(deployed) else await asyncio.to_thread( _snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None) ) # Push S3 creds (no-op unless the service advertises the extension), with a grant too: a # bundle the grant cannot hold can still be uploaded with them. await _push_s3_credentials(card_base, card or {}, timeout_seconds) grant = {} if upload is None else {"write_namespace": upload.grant} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) - part = resp.parts[0] if resp.parts else None - if upload is not None and (path := uploaded_object_path(part)) is not None: - return upload.object_url(path) - if isinstance(part, FilePart): - raw = getattr(part.file, "bytes", None) - if raw is not None: - with open(tmp_path, "wb") as f: - f.write(base64.b64decode(raw)) - else: - uri = getattr(part.file, "uri", None) - if not uri: - # ValueError, not TypeError, so export_one catches it and - # fails only this service — not the whole snapshot. - raise ValueError( - f"FilePart for '{environment_name}' has neither bytes nor uri" - ) - # Any url but http(s) is an object the service uploaded itself — hand - # it back as-is; export_one registers it directly (no download). Its - # last segment carries the shape (e.g. gdrive.zip). - scheme = urlparse(uri).scheme - if scheme and scheme not in ("http", "https"): - return uri - # A relative uri (e.g. "export-snapshot") means "stream it from - # my service endpoint" — resolve against the service base_url so - # multi-GB bundles never ride inline in the JSON-RPC response. An - # absolute one must be on that same origin: this process fetches it. - if not scheme: - uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" - elif urlparse(uri)[:2] != urlparse(base_url)[:2]: - raise ValueError( - f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" - ) - async with httpx.AsyncClient() as client: - async with client.stream( - "GET", - uri, - timeout=httpx.Timeout( - timeout_seconds, connect=_CONNECT_TIMEOUT_SECONDS - ), - ) as resp_file: - resp_file.raise_for_status() - with open(tmp_path, "wb") as f: - async for chunk in resp_file.aiter_bytes(): - f.write(chunk) - return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" - state = part.data if part is not None else {} - with open(tmp_path, "w") as f: - json.dump(state, f, default=str) - return ".json" + named = {} if paths is None else {"paths": paths} + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **named, **grant) + if len(resp.parts) > 1 or _returns_file_tree(card or {}): + return [(part, await _materialize(part, upload, base_url, f"{tmp_path}-{index}", environment_name, + timeout_seconds)) for index, part in enumerate(resp.parts)] + return await _materialize(resp.parts[0] if resp.parts else None, upload, base_url, tmp_path, + environment_name, timeout_seconds) legacy_url = f"{legacy_protocol.environment_base_url(gateway_url, environment_name, mcp=True)}/export-state" async with httpx.AsyncClient() as client: @@ -562,27 +541,29 @@ 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, + computer_file_paths: Optional[list[str]] = 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``. + All-or-nothing; the raised error carries public per-service summaries (verbose details + go to the log). ``computer_file_paths`` goes to the data/get of the env's own environment. 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 +588,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): @@ -620,9 +601,11 @@ async def export_one(environment_name: str): artifact_path = tmp_path try: try: + named = ({"paths": computer_file_paths} if computer_file_paths is not None + and environment_name == getattr(env, "environment_name", None) else {}) result = await cls._export_environment_to_file( gateway, environment_name, tmp_path, export_timeout_seconds, - deployed=card_record, + deployed=card_record, **named, ) except httpx.HTTPStatusError as e: body_snippet = "" @@ -638,31 +621,34 @@ async def export_one(environment_name: str): except ValueError as e: return str(e), str(e) + if isinstance(result, list): + try: + parts = await asyncio.to_thread( + _put_parts, result, tmp_path, 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, parts=parts, + ) + except Exception as e: + public = f"failed to create artifacts ({type(e).__name__})" + return public, f"{public} for {environment_name}: {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, @@ -700,7 +686,7 @@ def _register_uploaded_bundle(): public = f"failed to create artifacts ({type(e).__name__})" return public, f"{public}: {e}" finally: - for path in cleanup_paths: + for path in [*cleanup_paths, *glob.glob(f"{glob.escape(tmp_path)}-*")]: try: os.unlink(path) except OSError: @@ -748,6 +734,118 @@ def _register_uploaded_bundle(): ) +async def _materialize(part: Optional[Part], upload: Optional[_SnapshotUpload], base_url: str, tmp_path: str, + environment_name: str, timeout_seconds: float) -> str: + """One part of a service's export: the object url it was uploaded to under ``upload``, else the suffix written to + ``tmp_path`` (".json" for a data part or none, else its file's).""" + if upload is not None and (path := uploaded_object_path(part)) is not None: + return upload.object_url(path) + if isinstance(part, FilePart): + raw = getattr(part.file, "bytes", None) + if raw is not None: + with open(tmp_path, "wb") as f: + f.write(base64.b64decode(raw)) + else: + uri = getattr(part.file, "uri", None) + if not uri: + # ValueError, not TypeError, so export_one catches it and + # fails only this service — not the whole snapshot. + raise ValueError( + f"FilePart for '{environment_name}' has neither bytes nor uri" + ) + # Any url but http(s) is an object the service uploaded itself — hand + # it back as-is; export_one registers it directly (no download). Its + # last segment carries the shape (e.g. gdrive.zip). + scheme = urlparse(uri).scheme + if scheme and scheme not in ("http", "https"): + return uri + # A relative uri (e.g. "export-snapshot") means "stream it from + # my service endpoint" — resolve against the service base_url so + # multi-GB bundles never ride inline in the JSON-RPC response. An + # absolute one must be on that same origin: this process fetches it. + if not scheme: + uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" + elif urlparse(uri)[:2] != urlparse(base_url)[:2]: + raise ValueError( + f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" + ) + async with httpx.AsyncClient() as client: + async with client.stream( + "GET", + uri, + timeout=httpx.Timeout( + timeout_seconds, connect=_CONNECT_TIMEOUT_SECONDS + ), + ) as resp_file: + resp_file.raise_for_status() + with open(tmp_path, "wb") as f: + async for chunk in resp_file.aiter_bytes(): + f.write(chunk) + return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + state = part.data if part is not None else {} + with open(tmp_path, "w") as f: + json.dump(state, f, default=str) + return ".json" + + +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 _returns_file_tree(card: dict) -> bool: + """Whether ``card`` declares its data/get returns a file tree, each file a part named by its path, which one file + can't keep.""" + return any(returned.get("format") == "file-tree" for returned in (intake_declaration(card) or {}).get("get") or []) + + +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 _put_parts(exported: list[tuple[Part, str]], tmp_path: str, id: str, description: str) -> list: + """An ``EnvironmentPart`` of each part a service exported, in order, with what ``_materialize`` made of it: an + object it uploaded, written in place, or the file at ``-``, written as ``-part-`` + under the part's own name.""" + from agent_env.artifact import FileArtifact + from agent_env.artifact.artifacts.environment import EnvironmentPart + + parts = [] + with tempfile.TemporaryDirectory(prefix="agentenv-snapshot-parts-") as named: + for index, (part, result) in enumerate(exported): + file = part.file if isinstance(part, FilePart) else None + file_id = f"{id}-part-{index}" + if urlparse(result).scheme: + filename = getattr(file, "name", None) or result.rsplit("/", 1)[-1] or f"part-{index}" + stored = _registered_upload(result, file_id, filename, + mimetypes.guess_type(filename)[0] or "application/octet-stream", + description) + else: + path = Path(named, str(index), Path(getattr(file, "name", None) or f"part-{index}{result}").name) + path.parent.mkdir() + os.replace(f"{tmp_path}-{index}", path) + stored = FileArtifact.put(file_id, description=description, file_path=str(path)) + parts.append(EnvironmentPart(kind="file" if file is not None else "data", file_artifact_ref=stored.as_ref(), + name=getattr(file, "name", None), mime_type=getattr(file, "mimeType", None), + metadata=part.metadata or None)) + return parts + + 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 5463a3eac..b40cf6099 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -3,7 +3,9 @@ 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 without a gateway, each file under the paths named a part by its +path, comes back from a reset fresh, 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 @@ -18,7 +20,7 @@ from agentenv_protocol.computer.conformance import run_probe from agentenv_protocol.computer.conformance.client import Driver -from agent_env.artifact import FileArtifact +from agent_env.artifact import EnvironmentUniverseArtifact, FileArtifact 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 +37,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.snapshot_env import SnapshotEnvTaskStep from tst.integration.computer.conftest import docker pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -118,6 +121,12 @@ 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 SnapshotEnvTaskStep(id="snapshot", version=None, env_id=closed.id, + snapshot_id=f"{closed.id}-snapshot", + computer_file_paths=["Desktop", "Downloads/note.txt"]).execute(context) + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() + snapshot = {part.path: file.load() for part, file in environment.get_parts()} 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") finally: @@ -125,6 +134,8 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an sandbox = deployed.sandbox_id assert context.metadata["script_results"]["read"] == {"note": "inside\n", "card": "closed"} + assert (environment.environment_name, snapshot) == ("closed", {"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 [s.sandbox_id for s in report.terminated] == [sandbox] and not report.failed @@ -135,6 +146,30 @@ 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 SnapshotEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, + snapshot_id=f"{desktop.id}-snapshot").execute(context) + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() + snapshot = {part.path: (file.object_url, file.load()) for part, file in environment.get_parts()} + 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(): diff --git a/tst/unit/artifact/environment_parts_test.py b/tst/unit/artifact/environment_parts_test.py new file mode 100644 index 000000000..4c76a64ad --- /dev/null +++ b/tst/unit/artifact/environment_parts_test.py @@ -0,0 +1,62 @@ +"""An environment artifact of the parts an environment's data/get returned: written and read back in order, stored with +no wrapped file and nothing of the one-file shape, which stores as before; refused where one file is taken, and by an +artifact folder.""" + +import pytest +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI + +from agent_env.artifact.artifacts.environment import EnvironmentArtifact, EnvironmentPart +from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact +from agent_env.artifact.artifacts.file import FileArtifact +from agent_env.entity.export import Unwritable + + +def _file(tmp_path, name, data): + (tmp_path / name).write_bytes(data) + return FileArtifact.put(f"file-{name}", description=name, file_path=str(tmp_path / name)) + + +def test_an_environment_of_parts_is_written_and_read_back_in_order_with_nothing_of_one_file(local_stores, tmp_path): + notes, state = _file(tmp_path, "notes.txt", b"hello"), _file(tmp_path, "state.json", b'{"rows": 1}') + parts = [EnvironmentPart(kind="file", file_artifact_ref=notes.as_ref(), name="notes.txt", mime_type="text/plain", + metadata={DATA_OBJECTS_EXTENSION_URI: {"path": "Documents/notes.txt"}}), + EnvironmentPart(kind="data", file_artifact_ref=state.as_ref())] + + written = EnvironmentArtifact.put("desk", environment_name="desktop", parts=parts) + + read = EnvironmentArtifact.get("desk", version=written.version) + assert [(part.kind, part.path, file.load()) for part, file in read.get_parts()] == [ + ("file", "Documents/notes.txt", b"hello"), ("data", None, b'{"rows": 1}')] + assert read.parts == parts and read.file_artifact_ref is None + assert "file_artifact_id" not in read.model_dump(by_alias=True) + with pytest.raises(ValueError, match="holds the 2 parts its environment's data/get returned, not one file"): + read.get_file_artifact() + + +def test_one_file_is_stored_as_before_and_read_as_one_file_part(local_stores, tmp_path): + state = _file(tmp_path, "state.json", b"{}") + + written = EnvironmentArtifact.put("slack", environment_name="slack", file_artifact=state) + + assert "parts" not in written.model_dump(by_alias=True) and written.file_artifact_ref == state.as_ref() + assert [(part.kind, part.name, file.id) for part, file in EnvironmentArtifact.get("slack").get_parts()] == [ + ("file", "state.json", state.id)] + + +@pytest.mark.parametrize("given", [{}, {"file_artifact": "both", "parts": []}], ids=["neither", "both"]) +def test_an_environment_wraps_one_file_or_holds_parts_but_not_both_or_neither(local_stores, tmp_path, given): + if given: + given = {"file_artifact": _file(tmp_path, "state.json", b"{}"), "parts": []} + + with pytest.raises(ValueError, match="wraps a file_artifact or holds parts, one of the two"): + EnvironmentArtifact.put("x", environment_name="x", **given) + + +def test_an_artifact_folder_holds_no_environment_of_parts(local_stores, tmp_path): + desk = EnvironmentArtifact.put("desk", environment_name="desktop", parts=[]) + universe = EnvironmentUniverseArtifact.put("universe", environment_artifacts=[desk]) + + with pytest.raises(Unwritable, match="it holds the parts its data/get returned"): + desk.to_toml(None) + with pytest.raises(Unwritable, match="its environment 'desktop' holds the parts its data/get returned"): + universe.to_toml(None) diff --git a/tst/unit/artifact/test_service_universe_file_view.py b/tst/unit/artifact/test_service_universe_file_view.py index 6883da8a9..d5aa8197a 100644 --- a/tst/unit/artifact/test_service_universe_file_view.py +++ b/tst/unit/artifact/test_service_universe_file_view.py @@ -2,8 +2,9 @@ service universe be staged as a file tree (grading) by the same loaders that handle a FileArtifactUniverse. -Contract: keys are `/`, values are the -universe's EXISTING FileArtifacts (nothing copied or re-registered). +Contract: keys are `/`, or for an environment of the +parts its data/get returned, `/`; values +are the universe's EXISTING FileArtifacts (nothing copied or re-registered). """ from __future__ import annotations @@ -11,7 +12,9 @@ import pytest from agent_env.artifact.artifacts.file import FileArtifact -from agent_env.artifact.artifacts.environment import EnvironmentArtifact +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI + +from agent_env.artifact.artifacts.environment import EnvironmentArtifact, EnvironmentPart from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact from agent_env.artifact.ref import ArtifactRef @@ -91,3 +94,23 @@ def test_duplicate_service_names_raise(registry): universe = _universe(registry, [("slack", "a.json"), ("slack", "a.json")]) with pytest.raises(ValueError, match="duplicate environment 'slack'"): universe.get_file_artifacts() + + +def test_an_environment_of_parts_lists_each_part_by_its_path_else_its_name(registry): + files, services = registry + universe = _universe(registry, [("slack", "state.json")]) + parts = [] + for idx, (filename, name, path) in enumerate([("notes.txt", "notes.txt", "Documents/notes.txt"), + ("c.bin", "c.bin", "Downloads/b/c.bin"), + ("state.json", None, None)]): + files[f"part-{idx}"] = _file_artifact(f"part-{idx}", filename) + parts.append(EnvironmentPart(kind="file" if name else "data", file_artifact_ref=files[f"part-{idx}"].as_ref(), + name=name, metadata={DATA_OBJECTS_EXTENSION_URI: {"path": path}} if path else None)) + services["desk"] = EnvironmentArtifact(id="desk", version=1, environment_name="desktop", parts=parts) + universe.environment_artifact_refs.append(ArtifactRef(id="desk", version=1)) + + viewed = universe.get_file_artifacts() + + assert sorted(viewed) == ["desktop/Documents/notes.txt", "desktop/Downloads/b/c.bin", "desktop/state.json", + "slack/state.json"] + assert viewed["desktop/Downloads/b/c.bin"] is files["part-1"] diff --git a/tst/unit/task_step/snapshot_env_upload_grant_test.py b/tst/unit/task_step/snapshot_env_upload_grant_test.py index 9554f333c..5644b8904 100644 --- a/tst/unit/task_step/snapshot_env_upload_grant_test.py +++ b/tst/unit/task_step/snapshot_env_upload_grant_test.py @@ -94,7 +94,7 @@ async def test_a_service_taking_the_form_uploads_through_a_grant_and_is_register (params,) = _rpc_params(sent) grant = params["write_namespace"] assert grant["root_path"].startswith("agentenv-snapshots/") - assert (grant["max_objects"], grant["max_object_bytes"]) == (1, ENV_SNAPSHOT_LIMITS.max_object_bytes) + assert (grant["max_objects"], grant["max_object_bytes"]) == (100_000, ENV_SNAPSHOT_LIMITS.max_object_bytes) assert datetime.fromisoformat(grant["expires_at"]) >= datetime.now(UTC) + timedelta(seconds=MIN_GRANT_LIFETIME_SECONDS - 60) assert pushes == [_CARD] # still pushed, for a bundle too large for the grant assert [str(r.url) for r in sent] == [f"{_CARD}/svc/mcp-gdrive/agentenv"] 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 000000000..4d76f2175 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -0,0 +1,172 @@ +"""snapshot_env on an env deployed without a gateway, a computer env say: reached through the card its deployment +recorded, its environments the card's children or the env itself, and each environment's export kept as data/get +returned it, a file tree's files as parts by their paths, inline or uploaded under the grant without a download, and +what the env's own data/get returns named by computer_file_paths.""" + +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, + uploaded_file_part) + +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} + + +class _Desk: + id, version, environment_name = "desk", 1, "desktop" + + +def _record(card: dict, network: dict | None = None) -> 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", metadata={"computer": {"network": network}} if network else {}) + + +def _computer(*extensions) -> dict: + return {"name": "desktop", "capabilities": {"extensions": [FILE_TREE, *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, 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)) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, + "result": {"parts": answers[request.url.path](params)}}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return sent + + +async def _snapshot(card: dict, network: dict | None = None, **options) -> EnvironmentUniverseArtifact: + result = await SnapshotEnvTaskStep.snapshot_env_state(env=_Desk(), gateway_url=None, snapshot_id="snap-d", + deployed=_record(card, network), **options) + return EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id) + + +@pytest.mark.asyncio +async def test_a_computer_keeps_each_file_its_data_get_returns_as_a_part_by_its_path(local_stores, monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello"), + _inline("Downloads/b/c.bin", b"\x00\x01")]}) + + universe = await _snapshot(_computer()) + + (desk,) = universe.get_environment_artifacts() + assert (desk.environment_name, desk.file_artifact_ref) == ("desktop", None) + assert [(part.kind, part.path, part.name, part.mime_type, file.filename, file.load()) + for part, file in desk.get_parts()] == [ + ("file", "Documents/notes.txt", "notes.txt", "text/plain", "notes.txt", b"hello"), + ("file", "Downloads/b/c.bin", "c.bin", "text/plain", "c.bin", b"\x00\x01")] + assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt", "desktop/Downloads/b/c.bin"] + assert sent == [("/agentenv", {})] + + +@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}) + + (desk,) = (await _snapshot(_computer(DATA_OBJECTS))).get_environment_artifacts() + + ((path, params),) = sent + root = params["write_namespace"]["root_path"] + assert params["write_namespace"]["max_objects"] == 100_000 + assert [(part.path, file.object_url, file.load()) for part, file in desk.get_parts()] == [ + (path, store.object_url(f"{root}/{path}"), data) for path, data in files.items()] + + +@pytest.mark.asyncio +async def test_a_computer_whose_network_keeps_it_from_the_store_is_handed_no_grant_and_returns_its_files_inline( + local_stores, tmp_path, monkeypatch): + set_object_store(GrantingObjectStore(str(tmp_path / "objects"))) + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Desktop/left.txt", b"x")]}) + + (desk,) = (await _snapshot(_computer(DATA_OBJECTS), network={"mode": "allowlist"})).get_environment_artifacts() + + assert sent == [("/agentenv", {})] + assert [(part.path, file.load()) for part, file in desk.get_parts()] == [("Desktop/left.txt", b"x")] + + +@pytest.mark.asyncio +async def test_its_own_data_get_takes_computer_file_paths_and_each_child_on_its_card_is_snapshotted(local_stores, + monkeypatch): + card = {"name": "desktop", "children_environments": [ + {"name": "desktop", "url": "/agentenv", "capabilities": {"extensions": [FILE_TREE]}}, + {"name": "mail", "url": "/svc/mail/agentenv", "capabilities": {"extensions": []}}]} + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")], + "/svc/mail/agentenv": lambda params: [{"kind": "data", "data": {"inbox": 2}}]}) + + universe = await _snapshot(card, computer_file_paths=["Documents"]) + + desk, mail = sorted(universe.get_environment_artifacts(), key=lambda environment: environment.environment_name) + assert sorted(sent) == [("/agentenv", {"paths": ["Documents"]}), ("/svc/mail/agentenv", {})] + assert [part.path for part, _ in desk.get_parts()] == ["Documents/notes.txt"] + assert mail.parts is None and json.loads(mail.get_file_artifact().load()) == {"inbox": 2} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("card, answer, kept", [ + (_computer(), [], []), + ({"name": "desktop", "capabilities": {"extensions": []}}, + [{"kind": "data", "data": {"a": 1}}, {"kind": "data", "data": {"b": 2}}], + [("data", "part-0.json", b'{"a": 1}'), ("data", "part-1.json", b'{"b": 2}')]), +], ids=["a-file-tree-of-no-files", "several-parts-of-no-file-tree"]) +async def test_an_export_one_file_cant_keep_is_kept_as_its_parts(local_stores, monkeypatch, card, answer, kept): + _serve(monkeypatch, {"/agentenv": lambda params: answer}) + + (desk,) = (await _snapshot(card)).get_environment_artifacts() + + assert [(part.kind, file.filename, file.load()) for part, file in desk.get_parts()] == kept and desk.parts == [ + part for part, _ in desk.get_parts()] + + +@pytest.mark.asyncio +async def test_execute_snapshots_an_env_deployed_without_a_gateway_through_its_recorded_card(local_stores, + monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")]}) + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Desk())) + context = _ctx(_record(_computer())) + + await SnapshotEnvTaskStep(id="snap", version=None, env_id="desk", computer_file_paths=["Documents"]).execute( + context) + + snapshotted = context.metadata["env_snapshotted_universes"]["snap"] + universe = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]) + assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt"] + assert sent == [("/agentenv", {"paths": ["Documents"]})] diff --git a/tst/unit/task_step/test_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index dd4087b88..be7284ce1 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -62,6 +62,7 @@ def test_dict_roundtrip_and_registry(): export_timeout_seconds=120, include_env_trajectory=True, include_env_changes=True, + computer_file_paths=["Documents", "Desktop/report.pdf"], ) data = step.to_dict() assert data["type"] == "snapshot_env" @@ -71,7 +72,7 @@ def test_dict_roundtrip_and_registry(): for attr in ( "env_id", "env_instance_id", "gateway_url", "snapshot_id", "original_universe_artifact_id", "export_timeout_seconds", - "include_env_trajectory", "include_env_changes", + "include_env_trajectory", "include_env_changes", "computer_file_paths", ): assert getattr(restored, attr) == getattr(step, attr) @@ -82,6 +83,7 @@ def test_from_dict_defaults(): assert step.export_timeout_seconds == SnapshotEnvTaskStep.DEFAULT_EXPORT_TIMEOUT_SECONDS assert step.include_env_trajectory is False assert step.include_env_changes is False + assert step.computer_file_paths is None and "computer_file_paths" not in step.to_dict() # ── env resolution ─────────────────────────────────────────────────────────── @@ -145,9 +147,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 From bcd0ce56d9023db362d4a237e8b8498198067803 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 10:19:43 -0700 Subject: [PATCH 2/8] feat(snapshot): snapshot envs without a gateway, keeping each environment's export as its data/get returned it snapshot_env reaches an env deployed without a gateway, a computer env say, through the card its deployment recorded, and snapshots the environments the card lists as children, or the env itself. An export one file can't keep, several parts or the files of an env whose card declares its data/get returns a file tree, is stored as its parts (EnvironmentArtifact.parts: each part's kind, content, name, media type and metadata), so data/add can be handed them back; any other export is written exactly as before. computer_file_paths names what a computer env's own data/get returns. The snapshot grant allows a file tree's files, uploaded each where the computer put it, and a computer whose network keeps it from the store is handed none. Co-Authored-By: Claude Opus 5.5 --- .../artifact/artifacts/environment.py | 64 ++++- .../artifacts/environment_universe.py | 28 +- .../env/envs/computer/utils/inside.py | 12 +- .../task_step/task_steps/snapshot_env.py | 270 ++++++++++++------ tst/integration/computer/computer_env_test.py | 39 ++- tst/unit/artifact/environment_parts_test.py | 62 ++++ .../test_service_universe_file_view.py | 29 +- .../snapshot_env_upload_grant_test.py | 2 +- .../snapshot_env_without_gateway_test.py | 172 +++++++++++ tst/unit/task_step/test_snapshot_env.py | 9 +- 10 files changed, 573 insertions(+), 114 deletions(-) create mode 100644 tst/unit/artifact/environment_parts_test.py create mode 100644 tst/unit/task_step/snapshot_env_without_gateway_test.py diff --git a/src/agent_env/artifact/artifacts/environment.py b/src/agent_env/artifact/artifacts/environment.py index 957a27875..93303b0d7 100644 --- a/src/agent_env/artifact/artifacts/environment.py +++ b/src/agent_env/artifact/artifacts/environment.py @@ -1,4 +1,5 @@ -"""Environment artifact for wrapping a FileArtifact with environment identity.""" +"""Environment artifact for wrapping a FileArtifact with environment identity, or the parts an environment's data/get +returned.""" from __future__ import annotations @@ -6,16 +7,21 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional -from pydantic import ConfigDict, Field, model_serializer +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI +from pydantic import BaseModel, ConfigDict, Field, model_serializer from agent_env.artifact.artifact import Artifact from agent_env.artifact.artifacts.file import FileArtifact from agent_env.artifact.ref import ArtifactRef +from agent_env.entity.export import Unwritable from agent_env.entity.refs import EntityRef from agent_env.store.ids import derive_id logger = logging.getLogger(__name__) +#: Why an environment of parts has no artifact folder yet: a folder holds one file per environment. +PARTS_UNWRITABLE = "holds the parts its data/get returned, which an artifact folder can't hold yet" + if TYPE_CHECKING: from agent_env.bundle.authoring import AuthoringContext from agent_env.entity.export import ExportContext @@ -32,8 +38,28 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. path ``""`` for the folder itself. Raises BundleError when the folder can't be read.""" +class EnvironmentPart(BaseModel): + """One part of an environment's data as its data/get returned it, which data/add can be handed back: a file part + or a data part, its content stored as a file artifact, with the part's name, media type and metadata.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["file", "data"] + file_artifact_ref: ArtifactRef + name: Optional[str] = None + mime_type: Optional[str] = None + metadata: Optional[dict[str, Any]] = None + + @property + def path(self) -> Optional[str]: + """The path its metadata names it by under the data-objects marker, as a computer names each file.""" + marker = (self.metadata or {}).get(DATA_OBJECTS_EXTENSION_URI) + return marker.get("path") if isinstance(marker, dict) else None + + class EnvironmentArtifact(Artifact, SeedsEnvironments): - """A FileArtifact wrapped with the name of the environment it seeds.""" + """A FileArtifact wrapped with the name of the environment it seeds, or the ``parts`` its data/get returned when + one file would lose some of them: several parts, or the files of an environment that returns a file tree.""" model_config = ConfigDict(populate_by_name=True) @@ -56,6 +82,10 @@ class EnvironmentArtifact(Artifact, SeedsEnvironments): alias="file_artifact_id", description="Legacy unpinned wrapped-FileArtifact id", ) + parts: Optional[list[EnvironmentPart]] = Field( + default=None, + description="The parts the environment's data/get returned, in order, in place of one wrapped FileArtifact", + ) # No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields. @model_serializer(mode="wrap") @@ -65,8 +95,10 @@ def _serialize(self, handler: Any): data = handler(self) if "legacy_file_artifact_id" in data: data["file_artifact_id"] = data.pop("legacy_file_artifact_id") - if self.file_artifact_ref is not None: # new doc + if self.file_artifact_ref is not None or self.parts is not None: # new doc data.pop("file_artifact_id", None) + if self.parts is None: + data.pop("parts", None) # Dual-write. Sourced from the attribute, not from `data`, # because `alias=` emits exactly one spelling and which one depends on the # caller's `by_alias` — the store persists with by_alias=True @@ -81,16 +113,22 @@ def put( id: str, *, environment_name: Optional[str] = None, - file_artifact: FileArtifact, + file_artifact: Optional[FileArtifact] = None, + parts: Optional[list[EnvironmentPart]] = None, ) -> EnvironmentArtifact: + """Write the environment of ``environment_name`` over ``file_artifact``, or over ``parts``, the parts its + data/get returned.""" from agent_env.artifact.store import get_artifact_store if not environment_name: raise ValueError("environment_name cannot be empty") + if (file_artifact is None) == (parts is None): + raise ValueError("an environment artifact wraps a file_artifact or holds parts, one of the two") instance = cls( id=id, environment_name=environment_name, - file_artifact_ref=file_artifact.as_ref(), + file_artifact_ref=file_artifact.as_ref() if file_artifact is not None else None, + parts=parts, ) return get_artifact_store().put_document(instance) @@ -115,6 +153,8 @@ def from_toml(cls, data: dict, ctx: AuthoringContext) -> EnvironmentArtifact: def to_toml(self, out: ExportContext) -> dict: """The artifact.toml of a folder holding the environment's file, which ``from_toml`` reads back as this.""" + if self.parts is not None: + raise Unwritable(f"it {PARTS_UNWRITABLE}") file = self.get_file_artifact() out.file(file.filename, file.object_url) described = {} if file.description == file.filename else {"description": file.description} @@ -159,6 +199,9 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. return (("", name),) if isinstance(name, str) and name else () def get_file_artifact(self) -> FileArtifact: + if self.parts is not None: + raise ValueError(f"EnvironmentArtifact {self.id} v{self.version} holds the {len(self.parts)} parts its " + "environment's data/get returned, not one file") if self.file_artifact_ref is not None: return FileArtifact.get( self.file_artifact_ref.id, version=self.file_artifact_ref.version @@ -173,3 +216,12 @@ def get_file_artifact(self) -> FileArtifact: f"malformed EnvironmentArtifact {self.id} v{self.version}: " "neither file_artifact_ref nor file_artifact_id present" ) + + def get_parts(self) -> list[tuple[EnvironmentPart, FileArtifact]]: + """Each part and its content's file artifact, in order: the one file the environment wraps as a file part named + for it, or the parts its data/get returned.""" + if self.parts is None: + file = self.get_file_artifact() + return [(EnvironmentPart(kind="file", file_artifact_ref=file.as_ref(), name=file.filename), file)] + return [(part, FileArtifact.get(part.file_artifact_ref.id, version=part.file_artifact_ref.version)) + for part in self.parts] diff --git a/src/agent_env/artifact/artifacts/environment_universe.py b/src/agent_env/artifact/artifacts/environment_universe.py index 458677e0b..b0ad98c98 100644 --- a/src/agent_env/artifact/artifacts/environment_universe.py +++ b/src/agent_env/artifact/artifacts/environment_universe.py @@ -10,7 +10,7 @@ from pydantic import AliasChoices, ConfigDict, Field, model_serializer, model_validator from agent_env.artifact.artifact import FolderChild -from agent_env.artifact.artifacts.environment import EnvironmentArtifact, SeedsEnvironments +from agent_env.artifact.artifacts.environment import PARTS_UNWRITABLE, EnvironmentArtifact, SeedsEnvironments from agent_env.artifact.artifacts.environment_universe_folder import ( METADATA, UniverseFolder, @@ -184,6 +184,8 @@ def to_toml(self, out: ExportContext) -> dict: name = environment.environment_name if name.split("/", 1)[0] == METADATA: raise Unwritable(f"its environment {name!r} would read back as metadata, from its {METADATA}/ folder") + if environment.parts is not None: + raise Unwritable(f"its environment {name!r} {PARTS_UNWRITABLE}") file = environment.get_file_artifact() out.file(environment_path(Path(), name, file.filename).as_posix(), file.object_url) for key, file in (self.get_metadata() or {}).items(): @@ -260,18 +262,24 @@ def get_file_artifacts(self) -> dict[str, FileArtifact]: Matching ``FileArtifactUniverse.get_file_artifacts()`` is what lets the same loaders stage a service universe as a file tree (grading a frozen snapshot) rather than restoring it into an env. Returns the pinned artifacts as-is — - nothing is copied or re-registered. + nothing is copied or re-registered. An environment of parts lists each by the path + its metadata names, else its name, else its file's. """ out: dict[str, FileArtifact] = {} for sa in self.get_environment_artifacts(): - fa = sa.get_file_artifact() - name = f"{sa.environment_name}/{fa.filename}" - if name in out: - raise ValueError( - f"EnvironmentUniverseArtifact {self.id} v{self.version} has duplicate " - f"environment '{sa.environment_name}'; cannot stage it as a file tree" - ) - out[name] = fa + if sa.parts is None: + fa = sa.get_file_artifact() + named = [(fa.filename, fa)] + else: + named = [(part.path or part.name or fa.filename, fa) for part, fa in sa.get_parts()] + for relative, fa in named: + name = f"{sa.environment_name}/{relative}" + if name in out: + raise ValueError( + f"EnvironmentUniverseArtifact {self.id} v{self.version} has duplicate " + f"environment '{sa.environment_name}'; cannot stage it as a file tree" + ) + out[name] = fa return out def get_metadata(self) -> dict[str, FileArtifact]: diff --git a/src/agent_env/env/envs/computer/utils/inside.py b/src/agent_env/env/envs/computer/utils/inside.py index f1033e23c..7e3c92cd7 100644 --- a/src/agent_env/env/envs/computer/utils/inside.py +++ b/src/agent_env/env/envs/computer/utils/inside.py @@ -24,7 +24,7 @@ if TYPE_CHECKING: from agent_env.artifact import FileArtifact - from agent_env.env.env import DeployedSandboxEnv + from agent_env.env.env import DeployedEnv, DeployedSandboxEnv #: The most bytes one ``data/add`` carries: the driver reads a call's files into memory before it writes them. A file #: bigger than this goes alone. @@ -52,10 +52,16 @@ def loopback(deployed: DeployedSandboxEnv) -> dict: return {"id": deployed.env_id, "url": url, "mcp_url": url + deployed.mcp_url.removeprefix(deployed.environment_url)} +def egress_open(deployed: DeployedEnv) -> bool: + """Whether the network a computer's deploy recorded lets it reach the object store, to read or write through a + grant: one that allows everything, or none recorded, as for any env but a computer.""" + network = (deployed.metadata or {}).get("computer", {}).get("network") + return not (network and NetworkPolicy.from_dict(network).restricts_egress) + + async def add_artifacts(deployed: DeployedSandboxEnv, files: dict[str, FileArtifact]) -> None: """Write each file artifact at its path.""" - network = (deployed.metadata or {}).get("computer", {}).get("network") - granted = not (network and NetworkPolicy.from_dict(network).restricts_egress) + granted = egress_open(deployed) batch: list[FilePart] = [] size = 0 for path, artifact in files.items(): 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 49a108d81..f3966e368 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -1,8 +1,10 @@ -"""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 import base64 +import glob import json import logging import mimetypes @@ -11,16 +13,18 @@ import uuid from dataclasses import dataclass from datetime import datetime, timezone -from typing import ClassVar, Optional +from pathlib import Path +from typing import ClassVar, Optional, Union 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 DATA_OBJECTS_EXTENSION_URI, FilePart, Part, uploaded_object_path +from agentenv_protocol.client import find_extension, find_extension_method, intake_declaration from agentenv_protocol.transfers import WriteNamespaceGrant from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant +from agent_env.env.envs.computer.utils.inside import egress_open 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 @@ -45,9 +49,9 @@ # Must match the S3 credentials extension URI the server-side credentials mixin advertises. _S3_CREDENTIALS_EXTENSION_URI = "urn:agentenv:add-s3-credentials/v1" -# 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 a service may upload through its snapshot grant: one bundle, or the files of a file tree, each as large as one +# upload to the store can be, up to this in all. +ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=100_000, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" @@ -157,6 +161,12 @@ 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, a computer env say, is reached through the card its deployment recorded, its + environments the card's children, or itself. Each environment's export is kept as data/get returned it: one file + when it's one part, else its parts (``EnvironmentArtifact.parts``), as for a computer, whose data/get returns a + file tree. ``computer_file_paths``, relative to the session user's home, names what a computer env's own + data/get returns, by default the folders its card declares. """ type: ClassVar[str] = "snapshot_env" @@ -182,6 +192,7 @@ def __init__( export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, include_env_trajectory: bool = False, include_env_changes: bool = False, + computer_file_paths: Optional[list[str]] = None, depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, ): @@ -195,6 +206,7 @@ def __init__( self.export_timeout_seconds = export_timeout_seconds self.include_env_trajectory = include_env_trajectory self.include_env_changes = include_env_changes + self.computer_file_paths = computer_file_paths def to_dict(self) -> dict: base = super().to_dict() @@ -207,6 +219,8 @@ def to_dict(self) -> dict: base["export_timeout_seconds"] = self.export_timeout_seconds base["include_env_trajectory"] = self.include_env_trajectory base["include_env_changes"] = self.include_env_changes + if self.computer_file_paths is not None: + base["computer_file_paths"] = self.computer_file_paths return base @classmethod @@ -224,6 +238,7 @@ def from_dict(cls, data: dict) -> SnapshotEnvTaskStep: ), include_env_trajectory=data.get("include_env_trajectory", False), include_env_changes=data.get("include_env_changes", False), + computer_file_paths=data.get("computer_file_paths"), ) async def execute(self, context: TaskStepContext) -> TaskStepContext: @@ -238,9 +253,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: @@ -276,6 +292,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed=deployed, original_universe_artifact_id=self.original_universe_artifact_id, export_timeout_seconds=self.export_timeout_seconds, + computer_file_paths=self.computer_file_paths, ) logger.info( @@ -461,11 +478,15 @@ 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, + paths: Optional[list[str]] = None, + ) -> Union[str, list[tuple[Part, str]]]: """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). + When its data/get returns several parts, or a file tree (``_returns_file_tree``), + each part with the same for it, its file at ``-``; ``paths`` + goes to a data/get that takes it, a computer's. 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`` → @@ -478,63 +499,21 @@ async def _export_environment_to_file( base_url = await legacy_protocol.v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is not None: card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) - upload = await asyncio.to_thread( + # A grant is no use to an env whose network keeps it from the store, a closed computer say. + upload = None if deployed is not None and not egress_open(deployed) else await asyncio.to_thread( _snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None) ) # Push S3 creds (no-op unless the service advertises the extension), with a grant too: a # bundle the grant cannot hold can still be uploaded with them. await _push_s3_credentials(card_base, card or {}, timeout_seconds) grant = {} if upload is None else {"write_namespace": upload.grant} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) - part = resp.parts[0] if resp.parts else None - if upload is not None and (path := uploaded_object_path(part)) is not None: - return upload.object_url(path) - if isinstance(part, FilePart): - raw = getattr(part.file, "bytes", None) - if raw is not None: - with open(tmp_path, "wb") as f: - f.write(base64.b64decode(raw)) - else: - uri = getattr(part.file, "uri", None) - if not uri: - # ValueError, not TypeError, so export_one catches it and - # fails only this service — not the whole snapshot. - raise ValueError( - f"FilePart for '{environment_name}' has neither bytes nor uri" - ) - # Any url but http(s) is an object the service uploaded itself — hand - # it back as-is; export_one registers it directly (no download). Its - # last segment carries the shape (e.g. gdrive.zip). - scheme = urlparse(uri).scheme - if scheme and scheme not in ("http", "https"): - return uri - # A relative uri (e.g. "export-snapshot") means "stream it from - # my service endpoint" — resolve against the service base_url so - # multi-GB bundles never ride inline in the JSON-RPC response. An - # absolute one must be on that same origin: this process fetches it. - if not scheme: - uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" - elif urlparse(uri)[:2] != urlparse(base_url)[:2]: - raise ValueError( - f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" - ) - async with httpx.AsyncClient() as client: - async with client.stream( - "GET", - uri, - timeout=httpx.Timeout( - timeout_seconds, connect=_CONNECT_TIMEOUT_SECONDS - ), - ) as resp_file: - resp_file.raise_for_status() - with open(tmp_path, "wb") as f: - async for chunk in resp_file.aiter_bytes(): - f.write(chunk) - return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" - state = part.data if part is not None else {} - with open(tmp_path, "w") as f: - json.dump(state, f, default=str) - return ".json" + named = {} if paths is None else {"paths": paths} + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **named, **grant) + if len(resp.parts) > 1 or _returns_file_tree(card or {}): + return [(part, await _materialize(part, upload, base_url, f"{tmp_path}-{index}", environment_name, + timeout_seconds)) for index, part in enumerate(resp.parts)] + return await _materialize(resp.parts[0] if resp.parts else None, upload, base_url, tmp_path, + environment_name, timeout_seconds) legacy_url = f"{legacy_protocol.environment_base_url(gateway_url, environment_name, mcp=True)}/export-state" async with httpx.AsyncClient() as client: @@ -562,27 +541,29 @@ 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, + computer_file_paths: Optional[list[str]] = 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``. + All-or-nothing; the raised error carries public per-service summaries (verbose details + go to the log). ``computer_file_paths`` goes to the data/get of the env's own environment. 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 +588,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): @@ -620,9 +601,11 @@ async def export_one(environment_name: str): artifact_path = tmp_path try: try: + named = ({"paths": computer_file_paths} if computer_file_paths is not None + and environment_name == getattr(env, "environment_name", None) else {}) result = await cls._export_environment_to_file( gateway, environment_name, tmp_path, export_timeout_seconds, - deployed=card_record, + deployed=card_record, **named, ) except httpx.HTTPStatusError as e: body_snippet = "" @@ -638,31 +621,34 @@ async def export_one(environment_name: str): except ValueError as e: return str(e), str(e) + if isinstance(result, list): + try: + parts = await asyncio.to_thread( + _put_parts, result, tmp_path, 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, parts=parts, + ) + except Exception as e: + public = f"failed to create artifacts ({type(e).__name__})" + return public, f"{public} for {environment_name}: {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, @@ -700,7 +686,7 @@ def _register_uploaded_bundle(): public = f"failed to create artifacts ({type(e).__name__})" return public, f"{public}: {e}" finally: - for path in cleanup_paths: + for path in [*cleanup_paths, *glob.glob(f"{glob.escape(tmp_path)}-*")]: try: os.unlink(path) except OSError: @@ -748,6 +734,118 @@ def _register_uploaded_bundle(): ) +async def _materialize(part: Optional[Part], upload: Optional[_SnapshotUpload], base_url: str, tmp_path: str, + environment_name: str, timeout_seconds: float) -> str: + """One part of a service's export: the object url it was uploaded to under ``upload``, else the suffix written to + ``tmp_path`` (".json" for a data part or none, else its file's).""" + if upload is not None and (path := uploaded_object_path(part)) is not None: + return upload.object_url(path) + if isinstance(part, FilePart): + raw = getattr(part.file, "bytes", None) + if raw is not None: + with open(tmp_path, "wb") as f: + f.write(base64.b64decode(raw)) + else: + uri = getattr(part.file, "uri", None) + if not uri: + # ValueError, not TypeError, so export_one catches it and + # fails only this service — not the whole snapshot. + raise ValueError( + f"FilePart for '{environment_name}' has neither bytes nor uri" + ) + # Any url but http(s) is an object the service uploaded itself — hand + # it back as-is; export_one registers it directly (no download). Its + # last segment carries the shape (e.g. gdrive.zip). + scheme = urlparse(uri).scheme + if scheme and scheme not in ("http", "https"): + return uri + # A relative uri (e.g. "export-snapshot") means "stream it from + # my service endpoint" — resolve against the service base_url so + # multi-GB bundles never ride inline in the JSON-RPC response. An + # absolute one must be on that same origin: this process fetches it. + if not scheme: + uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" + elif urlparse(uri)[:2] != urlparse(base_url)[:2]: + raise ValueError( + f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" + ) + async with httpx.AsyncClient() as client: + async with client.stream( + "GET", + uri, + timeout=httpx.Timeout( + timeout_seconds, connect=_CONNECT_TIMEOUT_SECONDS + ), + ) as resp_file: + resp_file.raise_for_status() + with open(tmp_path, "wb") as f: + async for chunk in resp_file.aiter_bytes(): + f.write(chunk) + return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + state = part.data if part is not None else {} + with open(tmp_path, "w") as f: + json.dump(state, f, default=str) + return ".json" + + +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 _returns_file_tree(card: dict) -> bool: + """Whether ``card`` declares its data/get returns a file tree, each file a part named by its path, which one file + can't keep.""" + return any(returned.get("format") == "file-tree" for returned in (intake_declaration(card) or {}).get("get") or []) + + +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 _put_parts(exported: list[tuple[Part, str]], tmp_path: str, id: str, description: str) -> list: + """An ``EnvironmentPart`` of each part a service exported, in order, with what ``_materialize`` made of it: an + object it uploaded, written in place, or the file at ``-``, written as ``-part-`` + under the part's own name.""" + from agent_env.artifact import FileArtifact + from agent_env.artifact.artifacts.environment import EnvironmentPart + + parts = [] + with tempfile.TemporaryDirectory(prefix="agentenv-snapshot-parts-") as named: + for index, (part, result) in enumerate(exported): + file = part.file if isinstance(part, FilePart) else None + file_id = f"{id}-part-{index}" + if urlparse(result).scheme: + filename = getattr(file, "name", None) or result.rsplit("/", 1)[-1] or f"part-{index}" + stored = _registered_upload(result, file_id, filename, + mimetypes.guess_type(filename)[0] or "application/octet-stream", + description) + else: + path = Path(named, str(index), Path(getattr(file, "name", None) or f"part-{index}{result}").name) + path.parent.mkdir() + os.replace(f"{tmp_path}-{index}", path) + stored = FileArtifact.put(file_id, description=description, file_path=str(path)) + parts.append(EnvironmentPart(kind="file" if file is not None else "data", file_artifact_ref=stored.as_ref(), + name=getattr(file, "name", None), mime_type=getattr(file, "mimeType", None), + metadata=part.metadata or None)) + return parts + + 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 5463a3eac..b40cf6099 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -3,7 +3,9 @@ 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 without a gateway, each file under the paths named a part by its +path, comes back from a reset fresh, 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 @@ -18,7 +20,7 @@ from agentenv_protocol.computer.conformance import run_probe from agentenv_protocol.computer.conformance.client import Driver -from agent_env.artifact import FileArtifact +from agent_env.artifact import EnvironmentUniverseArtifact, FileArtifact 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 +37,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.snapshot_env import SnapshotEnvTaskStep from tst.integration.computer.conftest import docker pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -118,6 +121,12 @@ 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 SnapshotEnvTaskStep(id="snapshot", version=None, env_id=closed.id, + snapshot_id=f"{closed.id}-snapshot", + computer_file_paths=["Desktop", "Downloads/note.txt"]).execute(context) + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() + snapshot = {part.path: file.load() for part, file in environment.get_parts()} 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") finally: @@ -125,6 +134,8 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an sandbox = deployed.sandbox_id assert context.metadata["script_results"]["read"] == {"note": "inside\n", "card": "closed"} + assert (environment.environment_name, snapshot) == ("closed", {"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 [s.sandbox_id for s in report.terminated] == [sandbox] and not report.failed @@ -135,6 +146,30 @@ 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 SnapshotEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, + snapshot_id=f"{desktop.id}-snapshot").execute(context) + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() + snapshot = {part.path: (file.object_url, file.load()) for part, file in environment.get_parts()} + 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(): diff --git a/tst/unit/artifact/environment_parts_test.py b/tst/unit/artifact/environment_parts_test.py new file mode 100644 index 000000000..4c76a64ad --- /dev/null +++ b/tst/unit/artifact/environment_parts_test.py @@ -0,0 +1,62 @@ +"""An environment artifact of the parts an environment's data/get returned: written and read back in order, stored with +no wrapped file and nothing of the one-file shape, which stores as before; refused where one file is taken, and by an +artifact folder.""" + +import pytest +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI + +from agent_env.artifact.artifacts.environment import EnvironmentArtifact, EnvironmentPart +from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact +from agent_env.artifact.artifacts.file import FileArtifact +from agent_env.entity.export import Unwritable + + +def _file(tmp_path, name, data): + (tmp_path / name).write_bytes(data) + return FileArtifact.put(f"file-{name}", description=name, file_path=str(tmp_path / name)) + + +def test_an_environment_of_parts_is_written_and_read_back_in_order_with_nothing_of_one_file(local_stores, tmp_path): + notes, state = _file(tmp_path, "notes.txt", b"hello"), _file(tmp_path, "state.json", b'{"rows": 1}') + parts = [EnvironmentPart(kind="file", file_artifact_ref=notes.as_ref(), name="notes.txt", mime_type="text/plain", + metadata={DATA_OBJECTS_EXTENSION_URI: {"path": "Documents/notes.txt"}}), + EnvironmentPart(kind="data", file_artifact_ref=state.as_ref())] + + written = EnvironmentArtifact.put("desk", environment_name="desktop", parts=parts) + + read = EnvironmentArtifact.get("desk", version=written.version) + assert [(part.kind, part.path, file.load()) for part, file in read.get_parts()] == [ + ("file", "Documents/notes.txt", b"hello"), ("data", None, b'{"rows": 1}')] + assert read.parts == parts and read.file_artifact_ref is None + assert "file_artifact_id" not in read.model_dump(by_alias=True) + with pytest.raises(ValueError, match="holds the 2 parts its environment's data/get returned, not one file"): + read.get_file_artifact() + + +def test_one_file_is_stored_as_before_and_read_as_one_file_part(local_stores, tmp_path): + state = _file(tmp_path, "state.json", b"{}") + + written = EnvironmentArtifact.put("slack", environment_name="slack", file_artifact=state) + + assert "parts" not in written.model_dump(by_alias=True) and written.file_artifact_ref == state.as_ref() + assert [(part.kind, part.name, file.id) for part, file in EnvironmentArtifact.get("slack").get_parts()] == [ + ("file", "state.json", state.id)] + + +@pytest.mark.parametrize("given", [{}, {"file_artifact": "both", "parts": []}], ids=["neither", "both"]) +def test_an_environment_wraps_one_file_or_holds_parts_but_not_both_or_neither(local_stores, tmp_path, given): + if given: + given = {"file_artifact": _file(tmp_path, "state.json", b"{}"), "parts": []} + + with pytest.raises(ValueError, match="wraps a file_artifact or holds parts, one of the two"): + EnvironmentArtifact.put("x", environment_name="x", **given) + + +def test_an_artifact_folder_holds_no_environment_of_parts(local_stores, tmp_path): + desk = EnvironmentArtifact.put("desk", environment_name="desktop", parts=[]) + universe = EnvironmentUniverseArtifact.put("universe", environment_artifacts=[desk]) + + with pytest.raises(Unwritable, match="it holds the parts its data/get returned"): + desk.to_toml(None) + with pytest.raises(Unwritable, match="its environment 'desktop' holds the parts its data/get returned"): + universe.to_toml(None) diff --git a/tst/unit/artifact/test_service_universe_file_view.py b/tst/unit/artifact/test_service_universe_file_view.py index 6883da8a9..d5aa8197a 100644 --- a/tst/unit/artifact/test_service_universe_file_view.py +++ b/tst/unit/artifact/test_service_universe_file_view.py @@ -2,8 +2,9 @@ service universe be staged as a file tree (grading) by the same loaders that handle a FileArtifactUniverse. -Contract: keys are `/`, values are the -universe's EXISTING FileArtifacts (nothing copied or re-registered). +Contract: keys are `/`, or for an environment of the +parts its data/get returned, `/`; values +are the universe's EXISTING FileArtifacts (nothing copied or re-registered). """ from __future__ import annotations @@ -11,7 +12,9 @@ import pytest from agent_env.artifact.artifacts.file import FileArtifact -from agent_env.artifact.artifacts.environment import EnvironmentArtifact +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI + +from agent_env.artifact.artifacts.environment import EnvironmentArtifact, EnvironmentPart from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact from agent_env.artifact.ref import ArtifactRef @@ -91,3 +94,23 @@ def test_duplicate_service_names_raise(registry): universe = _universe(registry, [("slack", "a.json"), ("slack", "a.json")]) with pytest.raises(ValueError, match="duplicate environment 'slack'"): universe.get_file_artifacts() + + +def test_an_environment_of_parts_lists_each_part_by_its_path_else_its_name(registry): + files, services = registry + universe = _universe(registry, [("slack", "state.json")]) + parts = [] + for idx, (filename, name, path) in enumerate([("notes.txt", "notes.txt", "Documents/notes.txt"), + ("c.bin", "c.bin", "Downloads/b/c.bin"), + ("state.json", None, None)]): + files[f"part-{idx}"] = _file_artifact(f"part-{idx}", filename) + parts.append(EnvironmentPart(kind="file" if name else "data", file_artifact_ref=files[f"part-{idx}"].as_ref(), + name=name, metadata={DATA_OBJECTS_EXTENSION_URI: {"path": path}} if path else None)) + services["desk"] = EnvironmentArtifact(id="desk", version=1, environment_name="desktop", parts=parts) + universe.environment_artifact_refs.append(ArtifactRef(id="desk", version=1)) + + viewed = universe.get_file_artifacts() + + assert sorted(viewed) == ["desktop/Documents/notes.txt", "desktop/Downloads/b/c.bin", "desktop/state.json", + "slack/state.json"] + assert viewed["desktop/Downloads/b/c.bin"] is files["part-1"] diff --git a/tst/unit/task_step/snapshot_env_upload_grant_test.py b/tst/unit/task_step/snapshot_env_upload_grant_test.py index 9554f333c..5644b8904 100644 --- a/tst/unit/task_step/snapshot_env_upload_grant_test.py +++ b/tst/unit/task_step/snapshot_env_upload_grant_test.py @@ -94,7 +94,7 @@ async def test_a_service_taking_the_form_uploads_through_a_grant_and_is_register (params,) = _rpc_params(sent) grant = params["write_namespace"] assert grant["root_path"].startswith("agentenv-snapshots/") - assert (grant["max_objects"], grant["max_object_bytes"]) == (1, ENV_SNAPSHOT_LIMITS.max_object_bytes) + assert (grant["max_objects"], grant["max_object_bytes"]) == (100_000, ENV_SNAPSHOT_LIMITS.max_object_bytes) assert datetime.fromisoformat(grant["expires_at"]) >= datetime.now(UTC) + timedelta(seconds=MIN_GRANT_LIFETIME_SECONDS - 60) assert pushes == [_CARD] # still pushed, for a bundle too large for the grant assert [str(r.url) for r in sent] == [f"{_CARD}/svc/mcp-gdrive/agentenv"] 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 000000000..4d76f2175 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -0,0 +1,172 @@ +"""snapshot_env on an env deployed without a gateway, a computer env say: reached through the card its deployment +recorded, its environments the card's children or the env itself, and each environment's export kept as data/get +returned it, a file tree's files as parts by their paths, inline or uploaded under the grant without a download, and +what the env's own data/get returns named by computer_file_paths.""" + +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, + uploaded_file_part) + +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} + + +class _Desk: + id, version, environment_name = "desk", 1, "desktop" + + +def _record(card: dict, network: dict | None = None) -> 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", metadata={"computer": {"network": network}} if network else {}) + + +def _computer(*extensions) -> dict: + return {"name": "desktop", "capabilities": {"extensions": [FILE_TREE, *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, 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)) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, + "result": {"parts": answers[request.url.path](params)}}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return sent + + +async def _snapshot(card: dict, network: dict | None = None, **options) -> EnvironmentUniverseArtifact: + result = await SnapshotEnvTaskStep.snapshot_env_state(env=_Desk(), gateway_url=None, snapshot_id="snap-d", + deployed=_record(card, network), **options) + return EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id) + + +@pytest.mark.asyncio +async def test_a_computer_keeps_each_file_its_data_get_returns_as_a_part_by_its_path(local_stores, monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello"), + _inline("Downloads/b/c.bin", b"\x00\x01")]}) + + universe = await _snapshot(_computer()) + + (desk,) = universe.get_environment_artifacts() + assert (desk.environment_name, desk.file_artifact_ref) == ("desktop", None) + assert [(part.kind, part.path, part.name, part.mime_type, file.filename, file.load()) + for part, file in desk.get_parts()] == [ + ("file", "Documents/notes.txt", "notes.txt", "text/plain", "notes.txt", b"hello"), + ("file", "Downloads/b/c.bin", "c.bin", "text/plain", "c.bin", b"\x00\x01")] + assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt", "desktop/Downloads/b/c.bin"] + assert sent == [("/agentenv", {})] + + +@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}) + + (desk,) = (await _snapshot(_computer(DATA_OBJECTS))).get_environment_artifacts() + + ((path, params),) = sent + root = params["write_namespace"]["root_path"] + assert params["write_namespace"]["max_objects"] == 100_000 + assert [(part.path, file.object_url, file.load()) for part, file in desk.get_parts()] == [ + (path, store.object_url(f"{root}/{path}"), data) for path, data in files.items()] + + +@pytest.mark.asyncio +async def test_a_computer_whose_network_keeps_it_from_the_store_is_handed_no_grant_and_returns_its_files_inline( + local_stores, tmp_path, monkeypatch): + set_object_store(GrantingObjectStore(str(tmp_path / "objects"))) + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Desktop/left.txt", b"x")]}) + + (desk,) = (await _snapshot(_computer(DATA_OBJECTS), network={"mode": "allowlist"})).get_environment_artifacts() + + assert sent == [("/agentenv", {})] + assert [(part.path, file.load()) for part, file in desk.get_parts()] == [("Desktop/left.txt", b"x")] + + +@pytest.mark.asyncio +async def test_its_own_data_get_takes_computer_file_paths_and_each_child_on_its_card_is_snapshotted(local_stores, + monkeypatch): + card = {"name": "desktop", "children_environments": [ + {"name": "desktop", "url": "/agentenv", "capabilities": {"extensions": [FILE_TREE]}}, + {"name": "mail", "url": "/svc/mail/agentenv", "capabilities": {"extensions": []}}]} + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")], + "/svc/mail/agentenv": lambda params: [{"kind": "data", "data": {"inbox": 2}}]}) + + universe = await _snapshot(card, computer_file_paths=["Documents"]) + + desk, mail = sorted(universe.get_environment_artifacts(), key=lambda environment: environment.environment_name) + assert sorted(sent) == [("/agentenv", {"paths": ["Documents"]}), ("/svc/mail/agentenv", {})] + assert [part.path for part, _ in desk.get_parts()] == ["Documents/notes.txt"] + assert mail.parts is None and json.loads(mail.get_file_artifact().load()) == {"inbox": 2} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("card, answer, kept", [ + (_computer(), [], []), + ({"name": "desktop", "capabilities": {"extensions": []}}, + [{"kind": "data", "data": {"a": 1}}, {"kind": "data", "data": {"b": 2}}], + [("data", "part-0.json", b'{"a": 1}'), ("data", "part-1.json", b'{"b": 2}')]), +], ids=["a-file-tree-of-no-files", "several-parts-of-no-file-tree"]) +async def test_an_export_one_file_cant_keep_is_kept_as_its_parts(local_stores, monkeypatch, card, answer, kept): + _serve(monkeypatch, {"/agentenv": lambda params: answer}) + + (desk,) = (await _snapshot(card)).get_environment_artifacts() + + assert [(part.kind, file.filename, file.load()) for part, file in desk.get_parts()] == kept and desk.parts == [ + part for part, _ in desk.get_parts()] + + +@pytest.mark.asyncio +async def test_execute_snapshots_an_env_deployed_without_a_gateway_through_its_recorded_card(local_stores, + monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")]}) + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Desk())) + context = _ctx(_record(_computer())) + + await SnapshotEnvTaskStep(id="snap", version=None, env_id="desk", computer_file_paths=["Documents"]).execute( + context) + + snapshotted = context.metadata["env_snapshotted_universes"]["snap"] + universe = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]) + assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt"] + assert sent == [("/agentenv", {"paths": ["Documents"]})] diff --git a/tst/unit/task_step/test_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index dd4087b88..be7284ce1 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -62,6 +62,7 @@ def test_dict_roundtrip_and_registry(): export_timeout_seconds=120, include_env_trajectory=True, include_env_changes=True, + computer_file_paths=["Documents", "Desktop/report.pdf"], ) data = step.to_dict() assert data["type"] == "snapshot_env" @@ -71,7 +72,7 @@ def test_dict_roundtrip_and_registry(): for attr in ( "env_id", "env_instance_id", "gateway_url", "snapshot_id", "original_universe_artifact_id", "export_timeout_seconds", - "include_env_trajectory", "include_env_changes", + "include_env_trajectory", "include_env_changes", "computer_file_paths", ): assert getattr(restored, attr) == getattr(step, attr) @@ -82,6 +83,7 @@ def test_from_dict_defaults(): assert step.export_timeout_seconds == SnapshotEnvTaskStep.DEFAULT_EXPORT_TIMEOUT_SECONDS assert step.include_env_trajectory is False assert step.include_env_changes is False + assert step.computer_file_paths is None and "computer_file_paths" not in step.to_dict() # ── env resolution ─────────────────────────────────────────────────────────── @@ -145,9 +147,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 From 58beade88cedd30e8ff50a4b8fa7a2089b93aaee Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 12:14:41 -0700 Subject: [PATCH 3/8] fix(snapshot): keep a text part data/get returns, as a part of kind text A data/get answer may hold text parts too; one with a data part then a text part read the text part's data and failed the snapshot, and a lone text part did before this change. A text part is kept as a part of kind text, its text stored as a .txt file, so data/add can be handed it back. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/artifact/artifacts/environment.py | 6 +++--- src/agent_env/task_step/task_steps/snapshot_env.py | 13 +++++++++---- .../task_step/snapshot_env_without_gateway_test.py | 8 +++++++- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/src/agent_env/artifact/artifacts/environment.py b/src/agent_env/artifact/artifacts/environment.py index 93303b0d7..74add62dd 100644 --- a/src/agent_env/artifact/artifacts/environment.py +++ b/src/agent_env/artifact/artifacts/environment.py @@ -39,12 +39,12 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. class EnvironmentPart(BaseModel): - """One part of an environment's data as its data/get returned it, which data/add can be handed back: a file part - or a data part, its content stored as a file artifact, with the part's name, media type and metadata.""" + """One part of an environment's data as its data/get returned it, which data/add can be handed back: a file part, + a data part or a text part, its content stored as a file artifact, with the part's name, media type and metadata.""" model_config = ConfigDict(extra="forbid") - kind: Literal["file", "data"] + kind: Literal["file", "data", "text"] file_artifact_ref: ArtifactRef name: Optional[str] = None mime_type: Optional[str] = None 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 f3966e368..8c584f088 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -19,7 +19,7 @@ import httpx -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, Part, uploaded_object_path +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, Part, TextPart, uploaded_object_path from agentenv_protocol.client import find_extension, find_extension_method, intake_declaration from agentenv_protocol.transfers import WriteNamespaceGrant @@ -509,7 +509,8 @@ async def _export_environment_to_file( grant = {} if upload is None else {"write_namespace": upload.grant} named = {} if paths is None else {"paths": paths} resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **named, **grant) - if len(resp.parts) > 1 or _returns_file_tree(card or {}): + if len(resp.parts) > 1 or _returns_file_tree(card or {}) or any(isinstance(part, TextPart) + for part in resp.parts): return [(part, await _materialize(part, upload, base_url, f"{tmp_path}-{index}", environment_name, timeout_seconds)) for index, part in enumerate(resp.parts)] return await _materialize(resp.parts[0] if resp.parts else None, upload, base_url, tmp_path, @@ -737,9 +738,12 @@ def _register_uploaded_bundle(): async def _materialize(part: Optional[Part], upload: Optional[_SnapshotUpload], base_url: str, tmp_path: str, environment_name: str, timeout_seconds: float) -> str: """One part of a service's export: the object url it was uploaded to under ``upload``, else the suffix written to - ``tmp_path`` (".json" for a data part or none, else its file's).""" + ``tmp_path`` (".json" for a data part or none, ".txt" for a text part, else its file's).""" if upload is not None and (path := uploaded_object_path(part)) is not None: return upload.object_url(path) + if isinstance(part, TextPart): + Path(tmp_path).write_text(part.text) + return ".txt" if isinstance(part, FilePart): raw = getattr(part.file, "bytes", None) if raw is not None: @@ -840,7 +844,8 @@ def _put_parts(exported: list[tuple[Part, str]], tmp_path: str, id: str, descrip path.parent.mkdir() os.replace(f"{tmp_path}-{index}", path) stored = FileArtifact.put(file_id, description=description, file_path=str(path)) - parts.append(EnvironmentPart(kind="file" if file is not None else "data", file_artifact_ref=stored.as_ref(), + kind = "file" if file is not None else "text" if isinstance(part, TextPart) else "data" + parts.append(EnvironmentPart(kind=kind, file_artifact_ref=stored.as_ref(), name=getattr(file, "name", None), mime_type=getattr(file, "mimeType", None), metadata=part.metadata or None)) return parts diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py index 4d76f2175..2beb4ac2d 100644 --- a/tst/unit/task_step/snapshot_env_without_gateway_test.py +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -146,7 +146,13 @@ async def test_its_own_data_get_takes_computer_file_paths_and_each_child_on_its_ ({"name": "desktop", "capabilities": {"extensions": []}}, [{"kind": "data", "data": {"a": 1}}, {"kind": "data", "data": {"b": 2}}], [("data", "part-0.json", b'{"a": 1}'), ("data", "part-1.json", b'{"b": 2}')]), -], ids=["a-file-tree-of-no-files", "several-parts-of-no-file-tree"]) + ({"name": "desktop", "capabilities": {"extensions": []}}, + [{"kind": "data", "data": {"a": 1}}, {"kind": "text", "text": "hello"}], + [("data", "part-0.json", b'{"a": 1}'), ("text", "part-1.txt", b"hello")]), + ({"name": "desktop", "capabilities": {"extensions": []}}, [{"kind": "text", "text": "hello"}], + [("text", "part-0.txt", b"hello")]), +], ids=["a-file-tree-of-no-files", "several-parts-of-no-file-tree", "a-data-part-then-a-text-part", + "one-text-part"]) async def test_an_export_one_file_cant_keep_is_kept_as_its_parts(local_stores, monkeypatch, card, answer, kept): _serve(monkeypatch, {"/agentenv": lambda params: answer}) From e10c60164bd28f60e4d914cd4dc8cf653fdcbe1d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 12:34:58 -0700 Subject: [PATCH 4/8] refactor(snapshot): snapshot_env stays env-agnostic; snapshot_computer_env names a computer's files snapshot_env drops computer_file_paths. Its core takes more data/get params by environment name, passed as they are (never write_namespace, which it sets itself), and the step fills them from _data_get_params, empty for any env. The new snapshot_computer_env (task_steps/computer/) subclasses it with a typed file_paths, relative to the session user's home, sent as paths to the computer's own data/get only; an env that isn't a computer, and paths that aren't relative, are refused before a run. The protocol client's get_data takes extra params the same way. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/client.py | 9 +- .../tests/test_data_objects.py | 2 + src/agent_env/task/episode/rules.py | 2 +- src/agent_env/task_step/registry.py | 2 + .../task_step/task_steps/computer/__init__.py | 1 + .../computer/snapshot_computer_env.py | 89 +++++++++++++++++++ .../task_step/task_steps/snapshot_env.py | 40 +++++---- tst/integration/computer/computer_env_test.py | 11 +-- .../task_step/snapshot_computer_env_test.py | 80 +++++++++++++++++ .../snapshot_env_without_gateway_test.py | 22 +++-- tst/unit/task_step/test_field_scopes.py | 1 + tst/unit/task_step/test_snapshot_env.py | 4 +- 12 files changed, 226 insertions(+), 37 deletions(-) create mode 100644 src/agent_env/task_step/task_steps/computer/__init__.py create mode 100644 src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py create mode 100644 tst/unit/task_step/snapshot_computer_env_test.py diff --git a/packages/agentenv-protocol/src/agentenv_protocol/client.py b/packages/agentenv-protocol/src/agentenv_protocol/client.py index 0d0bfe648..9af27c042 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/client.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/client.py @@ -58,12 +58,15 @@ async def add_data(base_url: str, parts: list[Part], timeout: int = 120, verify: async def get_data( base_url: str, timeout: int = 30, verify: bool = True, *, write_namespace: WriteNamespaceGrant | None = None, - paths: list[str] | None = None, + paths: list[str] | None = None, extra: dict | None = None, ) -> GetDataResponse: """`data/get`. A server advertising `DATA_OBJECTS_EXTENSION_URI` may be handed `write_namespace` to upload its export under; `uploaded_object_path` reads from the answer where it did. `paths` names what a computer - returns, relative to its user's home; an env that returns all of its data refuses them.""" - params = {} if write_namespace is None else {"write_namespace": write_namespace.model_dump(mode="json", exclude_none=True)} + returns, relative to its user's home; an env that returns all of its data refuses them. `extra` holds any + other params the server's data/get takes, sent as they are; the named ones above win over its.""" + params = dict(extra or {}) + if write_namespace is not None: + params["write_namespace"] = write_namespace.model_dump(mode="json", exclude_none=True) if paths is not None: params["paths"] = paths return GetDataResponse.model_validate(await _rpc(base_url, METHOD_GET, params, timeout, verify)) diff --git a/packages/agentenv-protocol/tests/test_data_objects.py b/packages/agentenv-protocol/tests/test_data_objects.py index d19337bcb..6b6191e48 100644 --- a/packages/agentenv-protocol/tests/test_data_objects.py +++ b/packages/agentenv-protocol/tests/test_data_objects.py @@ -200,10 +200,12 @@ def handle(request: httpx.Request) -> httpx.Response: await protocol_v1.get_data("http://env") await protocol_v1.get_data("http://env", write_namespace=grant) await protocol_v1.get_data("http://env", paths=["Documents"]) + await protocol_v1.get_data("http://env", paths=["Documents"], extra={"paths": ["Desktop"], "tables": ["mail"]}) assert sent[0]["params"] == {} assert WriteNamespaceGrant.model_validate(sent[1]["params"]["write_namespace"]) == grant assert sent[2]["params"] == {"paths": ["Documents"]} + assert sent[3]["params"] == {"paths": ["Documents"], "tables": ["mail"]} @pytest.mark.asyncio diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 33ad6fa38..86d84c921 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 3c74f4fcd..83821a632 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 000000000..16fef9855 --- /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 000000000..a39d0f419 --- /dev/null +++ b/src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py @@ -0,0 +1,89 @@ +"""snapshot_computer_env: snapshot_env for a computer env, the files its snapshot keeps named by their paths.""" + +from __future__ import annotations + +from typing import ClassVar, Optional + +from agentenv_protocol.transfers import RelativePath +from pydantic import TypeAdapter, ValidationError + +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.task_step.task_step import TaskStepDependency +from agent_env.task_step.task_steps.snapshot_env import DEFAULT_EXPORT_TIMEOUT_SECONDS, SnapshotEnvTaskStep + +_PATHS = TypeAdapter(list[RelativePath]) + + +class SnapshotComputerEnvTaskStep(SnapshotEnvTaskStep): + """``snapshot_env`` for a computer env: ``file_paths``, each relative to its session user's home, names the files + and folders its own data/get returns, by default the folders its card declares. Any other environment it serves, + an app say, is snapshotted as ``snapshot_env`` snapshots it.""" + + type: ClassVar[str] = "snapshot_computer_env" + entity_refs = SnapshotEnvTaskStep.entity_refs + + 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, + gateway_url: Optional[str] = None, + snapshot_id: Optional[str] = None, + original_universe_artifact_id: Optional[str] = None, + export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, + include_env_trajectory: bool = False, + include_env_changes: bool = False, + file_paths: Optional[list[str]] = None, + 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, + gateway_url=gateway_url, snapshot_id=snapshot_id, + original_universe_artifact_id=original_universe_artifact_id, + export_timeout_seconds=export_timeout_seconds, include_env_trajectory=include_env_trajectory, + include_env_changes=include_env_changes, depends_on=depends_on, + fail_task_on_error=fail_task_on_error) + self.file_paths = file_paths + + def to_dict(self) -> dict: + return {**super().to_dict(), "file_paths": self.file_paths} + + @classmethod + def from_dict(cls, data: dict) -> SnapshotComputerEnvTaskStep: + step = super().from_dict(data) + step.file_paths = data.get("file_paths") + return step + + 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)] + + def _data_get_params(self, env) -> dict[str, dict]: + if not isinstance(env, ComputerEnv): + raise RuntimeError(_not_a_computer(env)) + return {env.environment_name: {"paths": self.file_paths}} if self.file_paths is not None else {} + + +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 8c584f088..8cb341fb7 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -53,6 +53,8 @@ # upload to the store can be, up to this in all. ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=100_000, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" +# The data/get params the snapshot sets itself, which data_get_params can't. +_SETS_ITSELF = {"write_namespace"} @dataclass(frozen=True) @@ -165,8 +167,8 @@ class SnapshotEnvTaskStep(TaskStep): An env deployed without a gateway, a computer env say, is reached through the card its deployment recorded, its environments the card's children, or itself. Each environment's export is kept as data/get returned it: one file when it's one part, else its parts (``EnvironmentArtifact.parts``), as for a computer, whose data/get returns a - file tree. ``computer_file_paths``, relative to the session user's home, names what a computer env's own - data/get returns, by default the folders its card declares. + file tree. A step for one kind of env can hand an environment's data/get more params (``_data_get_params``), as + ``snapshot_computer_env`` names a computer's files. """ type: ClassVar[str] = "snapshot_env" @@ -192,7 +194,6 @@ def __init__( export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, include_env_trajectory: bool = False, include_env_changes: bool = False, - computer_file_paths: Optional[list[str]] = None, depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, ): @@ -206,7 +207,6 @@ def __init__( self.export_timeout_seconds = export_timeout_seconds self.include_env_trajectory = include_env_trajectory self.include_env_changes = include_env_changes - self.computer_file_paths = computer_file_paths def to_dict(self) -> dict: base = super().to_dict() @@ -219,8 +219,6 @@ def to_dict(self) -> dict: base["export_timeout_seconds"] = self.export_timeout_seconds base["include_env_trajectory"] = self.include_env_trajectory base["include_env_changes"] = self.include_env_changes - if self.computer_file_paths is not None: - base["computer_file_paths"] = self.computer_file_paths return base @classmethod @@ -238,7 +236,6 @@ def from_dict(cls, data: dict) -> SnapshotEnvTaskStep: ), include_env_trajectory=data.get("include_env_trajectory", False), include_env_changes=data.get("include_env_changes", False), - computer_file_paths=data.get("computer_file_paths"), ) async def execute(self, context: TaskStepContext) -> TaskStepContext: @@ -292,7 +289,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed=deployed, original_universe_artifact_id=self.original_universe_artifact_id, export_timeout_seconds=self.export_timeout_seconds, - computer_file_paths=self.computer_file_paths, + data_get_params=self._data_get_params(env), ) logger.info( @@ -307,6 +304,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: } return context + def _data_get_params(self, env) -> dict[str, dict]: + """More params for the data/get of each of ``env``'s environments, by name, sent as they are: none here, as + for any env; a step for one kind of env names its own.""" + return {} + # ── internals ──────────────────────────────────────────────────────────── @staticmethod @@ -479,14 +481,14 @@ def _enumerate_environments(env) -> list[str]: @staticmethod async def _export_environment_to_file( gateway_url: Optional[str], environment_name: str, tmp_path: str, timeout_seconds: float, deployed=None, - paths: Optional[list[str]] = None, + extra: Optional[dict] = None, ) -> Union[str, list[tuple[Part, str]]]: """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). When its data/get returns several parts, or a file tree (``_returns_file_tree``), - each part with the same for it, its file at ``-``; ``paths`` - goes to a data/get that takes it, a computer's. + each part with the same for it, its file at ``-``; ``extra`` + holds more params for its data/get, sent as they are. 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`` → @@ -507,8 +509,8 @@ async def _export_environment_to_file( # 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} - named = {} if paths is None else {"paths": paths} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **named, **grant) + more = {} if extra is None else {"extra": extra} + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **more, **grant) if len(resp.parts) > 1 or _returns_file_tree(card or {}) or any(isinstance(part, TextPart) for part in resp.parts): return [(part, await _materialize(part, upload, base_url, f"{tmp_path}-{index}", environment_name, @@ -548,18 +550,21 @@ async def snapshot_env_state( original_universe_artifact_id: Optional[str] = None, export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, concurrency: int = DEFAULT_EXPORT_CONCURRENCY, - computer_file_paths: Optional[list[str]] = None, + data_get_params: Optional[dict[str, dict]] = None, ) -> EnvSnapshotResult: """Export all services via the gateway, or without one (``gateway_url`` None) through ``deployed``'s recorded card, into a new EnvironmentUniverseArtifact ``snapshot_id``. All-or-nothing; the raised error carries public per-service summaries (verbose details - go to the log). ``computer_file_paths`` goes to the data/get of the env's own environment. + go to the log). ``data_get_params`` holds more params for the data/get of each environment, by name, sent + as they are, never the ones the snapshot sets itself. 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) if gateway_url else _card_environments(deployed) + if reserved := sorted({key for params in (data_get_params or {}).values() for key in params} & _SETS_ITSELF): + raise ValueError(f"data_get_params sets {', '.join(reserved)}, which the snapshot sets itself") 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( @@ -602,11 +607,10 @@ async def export_one(environment_name: str): artifact_path = tmp_path try: try: - named = ({"paths": computer_file_paths} if computer_file_paths is not None - and environment_name == getattr(env, "environment_name", None) else {}) + extra = (data_get_params or {}).get(environment_name) result = await cls._export_environment_to_file( gateway, environment_name, tmp_path, export_timeout_seconds, - deployed=card_record, **named, + deployed=card_record, **({"extra": extra} if extra else {}), ) except httpx.HTTPStatusError as e: body_snippet = "" diff --git a/tst/integration/computer/computer_env_test.py b/tst/integration/computer/computer_env_test.py index b40cf6099..7f039f2c5 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -3,8 +3,8 @@ 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, is snapshotted without a gateway, each file under the paths named a part by its -path, comes back from a reset fresh, and leaves nothing once torn down; an open one uploads each file under the +code that reaches its driver on loopback, is snapshotted without a gateway by snapshot_computer_env, each file under the +paths it names a part by its path, comes back from a reset fresh, 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.""" @@ -37,6 +37,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 agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep from tst.integration.computer.conftest import docker @@ -121,9 +122,9 @@ 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 SnapshotEnvTaskStep(id="snapshot", version=None, env_id=closed.id, - snapshot_id=f"{closed.id}-snapshot", - computer_file_paths=["Desktop", "Downloads/note.txt"]).execute(context) + 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) taken = context.metadata["env_snapshotted_universes"]["snapshot"] (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() snapshot = {part.path: file.load() for part, file in environment.get_parts()} 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 000000000..9681a9468 --- /dev/null +++ b/tst/unit/task_step/snapshot_computer_env_test.py @@ -0,0 +1,80 @@ +"""snapshot_computer_env: snapshot_env for a computer env, its file_paths, relative to the session user's home, the +data/get params of the computer's own environment and of nothing else; refused for an env that isn't a computer, and +for paths that aren't relative, before a run.""" + +from __future__ import annotations + +import pytest + +from agent_env.artifact import EnvironmentUniverseArtifact +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 _computer, _inline, _record, _serve +from tst.unit.task_step.test_snapshot_env import _ctx + + +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" + + +def test_its_file_paths_round_trip_and_its_type_is_registered(): + step = SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk", file_paths=["Documents", "Desktop/a.pdf"], + include_env_trajectory=True) + + restored = get_task_step_registry()["snapshot_computer_env"].from_dict(step.to_dict()) + + assert isinstance(restored, SnapshotComputerEnvTaskStep) and step.to_dict()["type"] == "snapshot_computer_env" + assert (restored.env_id, restored.file_paths, restored.include_env_trajectory) == ( + "desk", ["Documents", "Desktop/a.pdf"], True) + + +@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_its_file_paths_are_the_computers_own_data_get_params(local_stores, monkeypatch, file_paths, params): + sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")]}) + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Desk())) + context = _ctx(_record(_computer())) + + await SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk", file_paths=file_paths).execute(context) + + snapshotted = context.metadata["env_snapshotted_universes"]["snap"] + universe = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]) + assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt"] + assert sent == [("/agentenv", params)] + + +@pytest.mark.asyncio +async def test_a_run_on_an_env_that_isnt_a_computer_is_refused_before_any_export(local_stores, monkeypatch): + sent = _serve(monkeypatch, {}) + monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Mail())) + + with pytest.raises(RuntimeError, match="snapshots a computer env, and mail is a mcp_server env"): + await SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk").execute(_ctx(_record(_computer()))) + assert sent == [] diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py index 2beb4ac2d..88c14f1e0 100644 --- a/tst/unit/task_step/snapshot_env_without_gateway_test.py +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -1,7 +1,7 @@ """snapshot_env on an env deployed without a gateway, a computer env say: reached through the card its deployment recorded, its environments the card's children or the env itself, and each environment's export kept as data/get returned it, a file tree's files as parts by their paths, inline or uploaded under the grant without a download, and -what the env's own data/get returns named by computer_file_paths.""" +more data/get params handed to the environment they name (snapshot_computer_env names a computer's files so).""" from __future__ import annotations @@ -124,15 +124,15 @@ async def test_a_computer_whose_network_keeps_it_from_the_store_is_handed_no_gra @pytest.mark.asyncio -async def test_its_own_data_get_takes_computer_file_paths_and_each_child_on_its_card_is_snapshotted(local_stores, - monkeypatch): +async def test_data_get_params_go_to_the_environment_they_name_and_each_child_on_its_card_is_snapshotted( + local_stores, monkeypatch): card = {"name": "desktop", "children_environments": [ {"name": "desktop", "url": "/agentenv", "capabilities": {"extensions": [FILE_TREE]}}, {"name": "mail", "url": "/svc/mail/agentenv", "capabilities": {"extensions": []}}]} sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")], "/svc/mail/agentenv": lambda params: [{"kind": "data", "data": {"inbox": 2}}]}) - universe = await _snapshot(card, computer_file_paths=["Documents"]) + universe = await _snapshot(card, data_get_params={"desktop": {"paths": ["Documents"]}}) desk, mail = sorted(universe.get_environment_artifacts(), key=lambda environment: environment.environment_name) assert sorted(sent) == [("/agentenv", {"paths": ["Documents"]}), ("/svc/mail/agentenv", {})] @@ -162,6 +162,15 @@ async def test_an_export_one_file_cant_keep_is_kept_as_its_parts(local_stores, m part for part, _ in desk.get_parts()] +@pytest.mark.asyncio +async def test_data_get_params_never_set_what_the_snapshot_sets_itself(local_stores, monkeypatch): + sent = _serve(monkeypatch, {"/agentenv": lambda params: []}) + + with pytest.raises(ValueError, match="data_get_params sets write_namespace, which the snapshot sets itself"): + await _snapshot(_computer(), data_get_params={"desktop": {"write_namespace": {}}}) + assert sent == [] + + @pytest.mark.asyncio async def test_execute_snapshots_an_env_deployed_without_a_gateway_through_its_recorded_card(local_stores, monkeypatch): @@ -169,10 +178,9 @@ async def test_execute_snapshots_an_env_deployed_without_a_gateway_through_its_r monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Desk())) context = _ctx(_record(_computer())) - await SnapshotEnvTaskStep(id="snap", version=None, env_id="desk", computer_file_paths=["Documents"]).execute( - context) + await SnapshotEnvTaskStep(id="snap", version=None, env_id="desk").execute(context) snapshotted = context.metadata["env_snapshotted_universes"]["snap"] universe = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]) assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt"] - assert sent == [("/agentenv", {"paths": ["Documents"]})] + assert sent == [("/agentenv", {})] diff --git a/tst/unit/task_step/test_field_scopes.py b/tst/unit/task_step/test_field_scopes.py index 86aa4948d..54562a0eb 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_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index be7284ce1..6e70635b1 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -62,7 +62,6 @@ def test_dict_roundtrip_and_registry(): export_timeout_seconds=120, include_env_trajectory=True, include_env_changes=True, - computer_file_paths=["Documents", "Desktop/report.pdf"], ) data = step.to_dict() assert data["type"] == "snapshot_env" @@ -72,7 +71,7 @@ def test_dict_roundtrip_and_registry(): for attr in ( "env_id", "env_instance_id", "gateway_url", "snapshot_id", "original_universe_artifact_id", "export_timeout_seconds", - "include_env_trajectory", "include_env_changes", "computer_file_paths", + "include_env_trajectory", "include_env_changes", ): assert getattr(restored, attr) == getattr(step, attr) @@ -83,7 +82,6 @@ def test_from_dict_defaults(): assert step.export_timeout_seconds == SnapshotEnvTaskStep.DEFAULT_EXPORT_TIMEOUT_SECONDS assert step.include_env_trajectory is False assert step.include_env_changes is False - assert step.computer_file_paths is None and "computer_file_paths" not in step.to_dict() # ── env resolution ─────────────────────────────────────────────────────────── From 076580d77fff82a72537fc6698a387fe6ef76c7e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 15:24:18 -0700 Subject: [PATCH 5/8] refactor(snapshot): ask again without the upload grant when data/get fails under it, instead of reading a computer's network The generic snapshot no longer reads a computer's recorded network to withhold its grant. An env whose data/get fails under the grant, one that can't reach the store such as a computer whose network is closed, is asked once more without it and returns its export inline; a failure without a grant is reported as before. inside.py is as on main. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/computer/utils/inside.py | 12 ++---- .../task_step/task_steps/snapshot_env.py | 16 ++++++-- .../snapshot_env_without_gateway_test.py | 39 +++++++++++++------ 3 files changed, 42 insertions(+), 25 deletions(-) diff --git a/src/agent_env/env/envs/computer/utils/inside.py b/src/agent_env/env/envs/computer/utils/inside.py index 7e3c92cd7..f1033e23c 100644 --- a/src/agent_env/env/envs/computer/utils/inside.py +++ b/src/agent_env/env/envs/computer/utils/inside.py @@ -24,7 +24,7 @@ if TYPE_CHECKING: from agent_env.artifact import FileArtifact - from agent_env.env.env import DeployedEnv, DeployedSandboxEnv + from agent_env.env.env import DeployedSandboxEnv #: The most bytes one ``data/add`` carries: the driver reads a call's files into memory before it writes them. A file #: bigger than this goes alone. @@ -52,16 +52,10 @@ def loopback(deployed: DeployedSandboxEnv) -> dict: return {"id": deployed.env_id, "url": url, "mcp_url": url + deployed.mcp_url.removeprefix(deployed.environment_url)} -def egress_open(deployed: DeployedEnv) -> bool: - """Whether the network a computer's deploy recorded lets it reach the object store, to read or write through a - grant: one that allows everything, or none recorded, as for any env but a computer.""" - network = (deployed.metadata or {}).get("computer", {}).get("network") - return not (network and NetworkPolicy.from_dict(network).restricts_egress) - - async def add_artifacts(deployed: DeployedSandboxEnv, files: dict[str, FileArtifact]) -> None: """Write each file artifact at its path.""" - granted = egress_open(deployed) + network = (deployed.metadata or {}).get("computer", {}).get("network") + granted = not (network and NetworkPolicy.from_dict(network).restricts_egress) batch: list[FilePart] = [] size = 0 for path, artifact in files.items(): 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 8cb341fb7..e150d1953 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -24,7 +24,6 @@ from agentenv_protocol.transfers import WriteNamespaceGrant from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant -from agent_env.env.envs.computer.utils.inside import egress_open 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 @@ -501,8 +500,7 @@ async def _export_environment_to_file( 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) - # A grant is no use to an env whose network keeps it from the store, a closed computer say. - upload = None if deployed is not None and not egress_open(deployed) else await asyncio.to_thread( + 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 @@ -510,7 +508,17 @@ async def _export_environment_to_file( await _push_s3_credentials(card_base, card or {}, timeout_seconds) grant = {} if upload is None else {"write_namespace": upload.grant} more = {} if extra is None else {"extra": extra} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **more, **grant) + try: + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **more, **grant) + except RuntimeError as e: + if upload is None: + raise + # An env that can't reach the store under the grant, a computer whose network is closed say, fails + # there; asked again without one, it returns its export inline. + logger.info("snapshot_env: %s's data/get failed under the upload grant (%s); asking again without one", + environment_name, e) + upload = None + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **more) if len(resp.parts) > 1 or _returns_file_tree(card or {}) or any(isinstance(part, TextPart) for part in resp.parts): return [(part, await _materialize(part, upload, base_url, f"{tmp_path}-{index}", environment_name, diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py index 88c14f1e0..5965043a1 100644 --- a/tst/unit/task_step/snapshot_env_without_gateway_test.py +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -32,10 +32,10 @@ class _Desk: id, version, environment_name = "desk", 1, "desktop" -def _record(card: dict, network: dict | None = None) -> DeployedSandboxEnv: +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", metadata={"computer": {"network": network}} if network else {}) + sandbox_type="local") def _computer(*extensions) -> dict: @@ -49,23 +49,24 @@ def _inline(path: str, data: bytes) -> dict: 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, recording each path and its - params; any other request fails the test.""" + """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)) - return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, - "result": {"parts": answers[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, network: dict | None = None, **options) -> EnvironmentUniverseArtifact: +async def _snapshot(card: dict, **options) -> EnvironmentUniverseArtifact: result = await SnapshotEnvTaskStep.snapshot_env_state(env=_Desk(), gateway_url=None, snapshot_id="snap-d", - deployed=_record(card, network), **options) + deployed=_record(card), **options) return EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id) @@ -111,18 +112,32 @@ def uploading(params): (path, store.object_url(f"{root}/{path}"), data) for path, data in files.items()] +UNREACHABLE = {"error": {"code": -32000, "message": "The object store is temporarily unavailable.", + "data": "get_failed"}} + + @pytest.mark.asyncio -async def test_a_computer_whose_network_keeps_it_from_the_store_is_handed_no_grant_and_returns_its_files_inline( +async def test_an_env_that_fails_under_the_grant_is_asked_again_without_one_and_returns_its_files_inline( local_stores, tmp_path, monkeypatch): set_object_store(GrantingObjectStore(str(tmp_path / "objects"))) - sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Desktop/left.txt", b"x")]}) + sent = _serve(monkeypatch, {"/agentenv": lambda params: UNREACHABLE if "write_namespace" in params + else [_inline("Desktop/left.txt", b"x")]}) - (desk,) = (await _snapshot(_computer(DATA_OBJECTS), network={"mode": "allowlist"})).get_environment_artifacts() + (desk,) = (await _snapshot(_computer(DATA_OBJECTS))).get_environment_artifacts() - assert sent == [("/agentenv", {})] + assert [(path, sorted(params)) for path, params in sent] == [("/agentenv", ["write_namespace"]), ("/agentenv", [])] assert [(part.path, file.load()) for part, file in desk.get_parts()] == [("Desktop/left.txt", b"x")] +@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(_computer()) + assert sent == [("/agentenv", {})] + + @pytest.mark.asyncio async def test_data_get_params_go_to_the_environment_they_name_and_each_child_on_its_card_is_snapshotted( local_stores, monkeypatch): From 8a57e749b63d28ff47587572bca5477d6cbd633a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 15:45:14 -0700 Subject: [PATCH 6/8] refactor(snapshot): keep EnvironmentArtifact as it is; snapshot_computer_env keeps a computer's files as a FileArtifactUniverse EnvironmentArtifact, EnvironmentUniverseArtifact and the protocol client are as on main. snapshot_env keeps one file per environment, as through a gateway, and refuses an environment whose data/get returns a file tree, pointing to snapshot_computer_env. That step writes a computer env's files into a FileArtifactUniverse by their paths, each upload under its grant written in place and any other file from its inline bytes, the computer asked again without the grant when it fails under it; load_artifact loads the universe back into a computer. Its fields are its own: no gateway, trajectory or changes. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/client.py | 9 +- .../tests/test_data_objects.py | 2 - .../artifact/artifacts/environment.py | 64 +--- .../artifacts/environment_universe.py | 28 +- .../computer/snapshot_computer_env.py | 123 ++++++-- .../task_step/task_steps/snapshot_env.py | 283 +++++++----------- tst/integration/computer/computer_env_test.py | 31 +- tst/unit/artifact/environment_parts_test.py | 62 ---- .../test_service_universe_file_view.py | 29 +- .../task_step/snapshot_computer_env_test.py | 109 +++++-- .../snapshot_env_upload_grant_test.py | 2 +- .../snapshot_env_without_gateway_test.py | 152 ++++------ 12 files changed, 380 insertions(+), 514 deletions(-) delete mode 100644 tst/unit/artifact/environment_parts_test.py diff --git a/packages/agentenv-protocol/src/agentenv_protocol/client.py b/packages/agentenv-protocol/src/agentenv_protocol/client.py index 9af27c042..0d0bfe648 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/client.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/client.py @@ -58,15 +58,12 @@ async def add_data(base_url: str, parts: list[Part], timeout: int = 120, verify: async def get_data( base_url: str, timeout: int = 30, verify: bool = True, *, write_namespace: WriteNamespaceGrant | None = None, - paths: list[str] | None = None, extra: dict | None = None, + paths: list[str] | None = None, ) -> GetDataResponse: """`data/get`. A server advertising `DATA_OBJECTS_EXTENSION_URI` may be handed `write_namespace` to upload its export under; `uploaded_object_path` reads from the answer where it did. `paths` names what a computer - returns, relative to its user's home; an env that returns all of its data refuses them. `extra` holds any - other params the server's data/get takes, sent as they are; the named ones above win over its.""" - params = dict(extra or {}) - if write_namespace is not None: - params["write_namespace"] = write_namespace.model_dump(mode="json", exclude_none=True) + returns, relative to its user's home; an env that returns all of its data refuses them.""" + params = {} if write_namespace is None else {"write_namespace": write_namespace.model_dump(mode="json", exclude_none=True)} if paths is not None: params["paths"] = paths return GetDataResponse.model_validate(await _rpc(base_url, METHOD_GET, params, timeout, verify)) diff --git a/packages/agentenv-protocol/tests/test_data_objects.py b/packages/agentenv-protocol/tests/test_data_objects.py index 6b6191e48..d19337bcb 100644 --- a/packages/agentenv-protocol/tests/test_data_objects.py +++ b/packages/agentenv-protocol/tests/test_data_objects.py @@ -200,12 +200,10 @@ def handle(request: httpx.Request) -> httpx.Response: await protocol_v1.get_data("http://env") await protocol_v1.get_data("http://env", write_namespace=grant) await protocol_v1.get_data("http://env", paths=["Documents"]) - await protocol_v1.get_data("http://env", paths=["Documents"], extra={"paths": ["Desktop"], "tables": ["mail"]}) assert sent[0]["params"] == {} assert WriteNamespaceGrant.model_validate(sent[1]["params"]["write_namespace"]) == grant assert sent[2]["params"] == {"paths": ["Documents"]} - assert sent[3]["params"] == {"paths": ["Documents"], "tables": ["mail"]} @pytest.mark.asyncio diff --git a/src/agent_env/artifact/artifacts/environment.py b/src/agent_env/artifact/artifacts/environment.py index 74add62dd..957a27875 100644 --- a/src/agent_env/artifact/artifacts/environment.py +++ b/src/agent_env/artifact/artifacts/environment.py @@ -1,5 +1,4 @@ -"""Environment artifact for wrapping a FileArtifact with environment identity, or the parts an environment's data/get -returned.""" +"""Environment artifact for wrapping a FileArtifact with environment identity.""" from __future__ import annotations @@ -7,21 +6,16 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI -from pydantic import BaseModel, ConfigDict, Field, model_serializer +from pydantic import ConfigDict, Field, model_serializer from agent_env.artifact.artifact import Artifact from agent_env.artifact.artifacts.file import FileArtifact from agent_env.artifact.ref import ArtifactRef -from agent_env.entity.export import Unwritable from agent_env.entity.refs import EntityRef from agent_env.store.ids import derive_id logger = logging.getLogger(__name__) -#: Why an environment of parts has no artifact folder yet: a folder holds one file per environment. -PARTS_UNWRITABLE = "holds the parts its data/get returned, which an artifact folder can't hold yet" - if TYPE_CHECKING: from agent_env.bundle.authoring import AuthoringContext from agent_env.entity.export import ExportContext @@ -38,28 +32,8 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. path ``""`` for the folder itself. Raises BundleError when the folder can't be read.""" -class EnvironmentPart(BaseModel): - """One part of an environment's data as its data/get returned it, which data/add can be handed back: a file part, - a data part or a text part, its content stored as a file artifact, with the part's name, media type and metadata.""" - - model_config = ConfigDict(extra="forbid") - - kind: Literal["file", "data", "text"] - file_artifact_ref: ArtifactRef - name: Optional[str] = None - mime_type: Optional[str] = None - metadata: Optional[dict[str, Any]] = None - - @property - def path(self) -> Optional[str]: - """The path its metadata names it by under the data-objects marker, as a computer names each file.""" - marker = (self.metadata or {}).get(DATA_OBJECTS_EXTENSION_URI) - return marker.get("path") if isinstance(marker, dict) else None - - class EnvironmentArtifact(Artifact, SeedsEnvironments): - """A FileArtifact wrapped with the name of the environment it seeds, or the ``parts`` its data/get returned when - one file would lose some of them: several parts, or the files of an environment that returns a file tree.""" + """A FileArtifact wrapped with the name of the environment it seeds.""" model_config = ConfigDict(populate_by_name=True) @@ -82,10 +56,6 @@ class EnvironmentArtifact(Artifact, SeedsEnvironments): alias="file_artifact_id", description="Legacy unpinned wrapped-FileArtifact id", ) - parts: Optional[list[EnvironmentPart]] = Field( - default=None, - description="The parts the environment's data/get returned, in order, in place of one wrapped FileArtifact", - ) # No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields. @model_serializer(mode="wrap") @@ -95,10 +65,8 @@ def _serialize(self, handler: Any): data = handler(self) if "legacy_file_artifact_id" in data: data["file_artifact_id"] = data.pop("legacy_file_artifact_id") - if self.file_artifact_ref is not None or self.parts is not None: # new doc + if self.file_artifact_ref is not None: # new doc data.pop("file_artifact_id", None) - if self.parts is None: - data.pop("parts", None) # Dual-write. Sourced from the attribute, not from `data`, # because `alias=` emits exactly one spelling and which one depends on the # caller's `by_alias` — the store persists with by_alias=True @@ -113,22 +81,16 @@ def put( id: str, *, environment_name: Optional[str] = None, - file_artifact: Optional[FileArtifact] = None, - parts: Optional[list[EnvironmentPart]] = None, + file_artifact: FileArtifact, ) -> EnvironmentArtifact: - """Write the environment of ``environment_name`` over ``file_artifact``, or over ``parts``, the parts its - data/get returned.""" from agent_env.artifact.store import get_artifact_store if not environment_name: raise ValueError("environment_name cannot be empty") - if (file_artifact is None) == (parts is None): - raise ValueError("an environment artifact wraps a file_artifact or holds parts, one of the two") instance = cls( id=id, environment_name=environment_name, - file_artifact_ref=file_artifact.as_ref() if file_artifact is not None else None, - parts=parts, + file_artifact_ref=file_artifact.as_ref(), ) return get_artifact_store().put_document(instance) @@ -153,8 +115,6 @@ def from_toml(cls, data: dict, ctx: AuthoringContext) -> EnvironmentArtifact: def to_toml(self, out: ExportContext) -> dict: """The artifact.toml of a folder holding the environment's file, which ``from_toml`` reads back as this.""" - if self.parts is not None: - raise Unwritable(f"it {PARTS_UNWRITABLE}") file = self.get_file_artifact() out.file(file.filename, file.object_url) described = {} if file.description == file.filename else {"description": file.description} @@ -199,9 +159,6 @@ def seeded_environments(cls, ctx: AuthoringContext) -> tuple[tuple[str, str], .. return (("", name),) if isinstance(name, str) and name else () def get_file_artifact(self) -> FileArtifact: - if self.parts is not None: - raise ValueError(f"EnvironmentArtifact {self.id} v{self.version} holds the {len(self.parts)} parts its " - "environment's data/get returned, not one file") if self.file_artifact_ref is not None: return FileArtifact.get( self.file_artifact_ref.id, version=self.file_artifact_ref.version @@ -216,12 +173,3 @@ def get_file_artifact(self) -> FileArtifact: f"malformed EnvironmentArtifact {self.id} v{self.version}: " "neither file_artifact_ref nor file_artifact_id present" ) - - def get_parts(self) -> list[tuple[EnvironmentPart, FileArtifact]]: - """Each part and its content's file artifact, in order: the one file the environment wraps as a file part named - for it, or the parts its data/get returned.""" - if self.parts is None: - file = self.get_file_artifact() - return [(EnvironmentPart(kind="file", file_artifact_ref=file.as_ref(), name=file.filename), file)] - return [(part, FileArtifact.get(part.file_artifact_ref.id, version=part.file_artifact_ref.version)) - for part in self.parts] diff --git a/src/agent_env/artifact/artifacts/environment_universe.py b/src/agent_env/artifact/artifacts/environment_universe.py index b0ad98c98..458677e0b 100644 --- a/src/agent_env/artifact/artifacts/environment_universe.py +++ b/src/agent_env/artifact/artifacts/environment_universe.py @@ -10,7 +10,7 @@ from pydantic import AliasChoices, ConfigDict, Field, model_serializer, model_validator from agent_env.artifact.artifact import FolderChild -from agent_env.artifact.artifacts.environment import PARTS_UNWRITABLE, EnvironmentArtifact, SeedsEnvironments +from agent_env.artifact.artifacts.environment import EnvironmentArtifact, SeedsEnvironments from agent_env.artifact.artifacts.environment_universe_folder import ( METADATA, UniverseFolder, @@ -184,8 +184,6 @@ def to_toml(self, out: ExportContext) -> dict: name = environment.environment_name if name.split("/", 1)[0] == METADATA: raise Unwritable(f"its environment {name!r} would read back as metadata, from its {METADATA}/ folder") - if environment.parts is not None: - raise Unwritable(f"its environment {name!r} {PARTS_UNWRITABLE}") file = environment.get_file_artifact() out.file(environment_path(Path(), name, file.filename).as_posix(), file.object_url) for key, file in (self.get_metadata() or {}).items(): @@ -262,24 +260,18 @@ def get_file_artifacts(self) -> dict[str, FileArtifact]: Matching ``FileArtifactUniverse.get_file_artifacts()`` is what lets the same loaders stage a service universe as a file tree (grading a frozen snapshot) rather than restoring it into an env. Returns the pinned artifacts as-is — - nothing is copied or re-registered. An environment of parts lists each by the path - its metadata names, else its name, else its file's. + nothing is copied or re-registered. """ out: dict[str, FileArtifact] = {} for sa in self.get_environment_artifacts(): - if sa.parts is None: - fa = sa.get_file_artifact() - named = [(fa.filename, fa)] - else: - named = [(part.path or part.name or fa.filename, fa) for part, fa in sa.get_parts()] - for relative, fa in named: - name = f"{sa.environment_name}/{relative}" - if name in out: - raise ValueError( - f"EnvironmentUniverseArtifact {self.id} v{self.version} has duplicate " - f"environment '{sa.environment_name}'; cannot stage it as a file tree" - ) - out[name] = fa + fa = sa.get_file_artifact() + name = f"{sa.environment_name}/{fa.filename}" + if name in out: + raise ValueError( + f"EnvironmentUniverseArtifact {self.id} v{self.version} has duplicate " + f"environment '{sa.environment_name}'; cannot stage it as a file tree" + ) + out[name] = fa return out def get_metadata(self) -> dict[str, FileArtifact]: 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 index a39d0f419..21f8049b2 100644 --- 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 @@ -1,28 +1,62 @@ -"""snapshot_computer_env: snapshot_env for a computer env, the files its snapshot keeps named by their paths.""" +"""snapshot_computer_env: a computer env's files, kept as a FileArtifactUniverse by their paths, which load_artifact +loads back into a computer.""" from __future__ import annotations +import asyncio +import base64 +import logging +from pathlib import PurePosixPath from typing import ClassVar, Optional +from agentenv_protocol import uploaded_object_path +from agentenv_protocol.types import FileWithBytes from agentenv_protocol.transfers import RelativePath from pydantic import TypeAdapter, ValidationError +from agent_env.a2a_agent.object_transfer import ObjectLimits +from agent_env.artifact import FileArtifact +from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +from agent_env.entity.refs import EntityRef, RefRole +from agent_env.env import legacy_protocol 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.task_step.task_step import TaskStepDependency -from agent_env.task_step.task_steps.snapshot_env import DEFAULT_EXPORT_TIMEOUT_SECONDS, SnapshotEnvTaskStep +from agent_env.store.ids import derive_id, 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, + _snapshot_upload, + get_data, + registered_upload, +) + +logger = logging.getLogger(__name__) _PATHS = TypeAdapter(list[RelativePath]) +#: What a computer may upload through its snapshot 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 its files are written to the store at once. +_WRITES_IN_FLIGHT = 8 class SnapshotComputerEnvTaskStep(SnapshotEnvTaskStep): - """``snapshot_env`` for a computer env: ``file_paths``, each relative to its session user's home, names the files - and folders its own data/get returns, by default the folders its card declares. Any other environment it serves, - an app say, is snapshotted as ``snapshot_env`` snapshots it.""" + """Snapshot a deployed computer env's files into a new ``FileArtifactUniverse``, each under its path relative to + the session user's home: those under ``file_paths``, by default the folders its card declares. The computer + uploads each file through a grant when the store's grants reach it, and returns it inline otherwise. + ``load_artifact`` loads the universe back into a computer, through ``data/add``. + + Picks the deployment and names the snapshot as ``snapshot_env`` does. Output: + ``context.metadata["env_snapshotted_universes"][self.id]`` = ``{"id", "version"}``.""" type: ClassVar[str] = "snapshot_computer_env" - entity_refs = SnapshotEnvTaskStep.entity_refs + entity_refs = ( + EntityRef.env("env_id"), + EntityRef.artifact("snapshot_id", role=RefRole.OUTPUT, artifact_type="file_artifact_universe"), + ) def __init__( self, @@ -31,32 +65,28 @@ def __init__( env_id: Optional[str] = None, env_instance_id: Optional[str] = None, env_step_id: Optional[str] = None, - gateway_url: Optional[str] = None, snapshot_id: Optional[str] = None, - original_universe_artifact_id: Optional[str] = None, - export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS, - include_env_trajectory: bool = False, - include_env_changes: bool = False, 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, - gateway_url=gateway_url, snapshot_id=snapshot_id, - original_universe_artifact_id=original_universe_artifact_id, - export_timeout_seconds=export_timeout_seconds, include_env_trajectory=include_env_trajectory, - include_env_changes=include_env_changes, depends_on=depends_on, - fail_task_on_error=fail_task_on_error) + 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 {**super().to_dict(), "file_paths": self.file_paths} + 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: - step = super().from_dict(data) - step.file_paths = data.get("file_paths") - return step + 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() @@ -79,10 +109,57 @@ def preflight(self) -> list[str]: return problems return problems if isinstance(env, ComputerEnv) else [*problems, _not_a_computer(env)] - def _data_get_params(self, env) -> dict[str, dict]: + 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)) - return {env.environment_name: {"paths": self.file_paths}} if self.file_paths is not None else {} + files = await _files(deployed, env, snapshot_id, self.file_paths, self.export_timeout_seconds) + if not files: + under = ", ".join(self.file_paths) if self.file_paths is not None else "the folders its card declares" + raise RuntimeError(f"computer env {env.id} has no files under {under} to snapshot") + universe = await asyncio.to_thread(FileArtifactUniverse.put, snapshot_id, file_artifacts=files) + logger.info(f"snapshot_computer_env: {env.id} -> {universe.id} v{universe.version} ({len(files)} files)") + context.metadata.setdefault("env_snapshotted_universes", {})[self.id] = {"id": universe.id, + "version": universe.version} + return context + + +async def _files(deployed, env: ComputerEnv, snapshot_id: str, paths: Optional[list[str]], + timeout_seconds: float) -> dict[str, FileArtifact]: + """Each file ``env``'s own data/get returns under ``paths``, by its path, written to the store as ``__file-``: an upload under the grant in place, without a download, else its inline bytes.""" + name = env.environment_name + base_url = await legacy_protocol.v1_base_url(deployed, None, name) + _, card = await legacy_protocol.child_env_card(deployed, None, name) + if base_url is None: + raise RuntimeError(f"computer env {env.id}'s recorded card names no environment {name!r}") + upload = await asyncio.to_thread(_snapshot_upload, card or {}, timeout_seconds, deployed.sandbox_type, + FILE_TREE_LIMITS) + resp, upload = await get_data(base_url, name, timeout_seconds, upload, + **({} if paths is None else {"paths": paths})) + writing = asyncio.Semaphore(_WRITES_IN_FLIGHT) + + async def written(index: int, part) -> tuple[str, FileArtifact]: + path = uploaded_object_path(part) + if path is None: + raise RuntimeError(f"computer env {env.id}'s data/get returned a {part.kind} part naming no path") + file_id, filename = derive_id(snapshot_id, f"file-{index}"), PurePosixPath(path).name + content_type, description = part.file.mimeType or "application/octet-stream", f"{path} of computer env {env.id}" + async with writing: + if isinstance(part.file, FileWithBytes): + return path, await asyncio.to_thread(FileArtifact.put_bytes, file_id, description=description, + filename=filename, content=base64.b64decode(part.file.bytes), + content_type=content_type) + if upload is None: + raise RuntimeError(f"computer env {env.id}'s data/get returned {path} neither inline nor uploaded") + return path, await asyncio.to_thread(registered_upload, upload.object_url(path), file_id, filename, + content_type, description) + + return dict(await asyncio.gather(*(written(index, part) for index, part in enumerate(resp.parts)))) def _not_a_computer(env: Env) -> str: 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 e150d1953..44822df77 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -4,7 +4,6 @@ import asyncio import base64 -import glob import json import logging import mimetypes @@ -13,14 +12,14 @@ import uuid from dataclasses import dataclass from datetime import datetime, timezone -from pathlib import Path -from typing import ClassVar, Optional, Union +from typing import ClassVar, Optional from urllib.parse import urlparse import httpx -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, Part, TextPart, uploaded_object_path -from agentenv_protocol.client import find_extension, find_extension_method, intake_declaration +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object_path +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 agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant @@ -48,12 +47,10 @@ # Must match the S3 credentials extension URI the server-side credentials mixin advertises. _S3_CREDENTIALS_EXTENSION_URI = "urn:agentenv:add-s3-credentials/v1" -# What a service may upload through its snapshot grant: one bundle, or the files of a file tree, each as large as one -# upload to the store can be, up to this in all. -ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=100_000, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) +# 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) _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" -# The data/get params the snapshot sets itself, which data_get_params can't. -_SETS_ITSELF = {"write_namespace"} @dataclass(frozen=True) @@ -67,21 +64,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: @@ -163,11 +161,9 @@ class SnapshotEnvTaskStep(TaskStep): ``include_env_changes`` also stores who changed what (urn:agentenv:changes/v1) under ``metadata["env_changes"]``. - An env deployed without a gateway, a computer env say, is reached through the card its deployment recorded, its - environments the card's children, or itself. Each environment's export is kept as data/get returned it: one file - when it's one part, else its parts (``EnvironmentArtifact.parts``), as for a computer, whose data/get returns a - file tree. A step for one kind of env can hand an environment's data/get more params (``_data_get_params``), as - ``snapshot_computer_env`` names a computer's files. + 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 refused, since + one file can't keep it: ``snapshot_computer_env`` snapshots a computer env's files. """ type: ClassVar[str] = "snapshot_env" @@ -288,7 +284,6 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed=deployed, original_universe_artifact_id=self.original_universe_artifact_id, export_timeout_seconds=self.export_timeout_seconds, - data_get_params=self._data_get_params(env), ) logger.info( @@ -303,11 +298,6 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: } return context - def _data_get_params(self, env) -> dict[str, dict]: - """More params for the data/get of each of ``env``'s environments, by name, sent as they are: none here, as - for any env; a step for one kind of env names its own.""" - return {} - # ── internals ──────────────────────────────────────────────────────────── @staticmethod @@ -479,52 +469,82 @@ def _enumerate_environments(env) -> list[str]: @staticmethod async def _export_environment_to_file( - gateway_url: Optional[str], environment_name: str, tmp_path: str, timeout_seconds: float, deployed=None, - extra: Optional[dict] = None, - ) -> Union[str, list[tuple[Part, str]]]: + gateway_url: Optional[str], environment_name: str, tmp_path: str, timeout_seconds: float, deployed=None + ) -> str: """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). - When its data/get returns several parts, or a file tree (``_returns_file_tree``), - each part with the same for it, its file at ``-``; ``extra`` - holds more params for its data/get, sent as they are. v1 ``get_data``, handed a namespace grant when the service takes the data-objects form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` → object url (no download) / base64 ``bytes`` / streamed relative-or-same-origin ``uri`` → ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON - guard. Errors are caught by the caller (fails just this service).""" + guard. Errors are caught by the caller (fails just this service): a file tree, which one file can't + keep, among them.""" 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 {}): + raise ValueError(f"'{environment_name}' returns a file tree from data/get, which one file can't keep; " + "snapshot_computer_env snapshots a computer env's files") 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} - more = {} if extra is None else {"extra": extra} - try: - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **more, **grant) - except RuntimeError as e: - if upload is None: - raise - # An env that can't reach the store under the grant, a computer whose network is closed say, fails - # there; asked again without one, it returns its export inline. - logger.info("snapshot_env: %s's data/get failed under the upload grant (%s); asking again without one", - environment_name, e) - upload = None - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **more) - if len(resp.parts) > 1 or _returns_file_tree(card or {}) or any(isinstance(part, TextPart) - for part in resp.parts): - return [(part, await _materialize(part, upload, base_url, f"{tmp_path}-{index}", environment_name, - timeout_seconds)) for index, part in enumerate(resp.parts)] - return await _materialize(resp.parts[0] if resp.parts else None, upload, base_url, tmp_path, - environment_name, timeout_seconds) + resp, upload = await get_data(base_url, environment_name, timeout_seconds, upload) + 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) + if isinstance(part, FilePart): + raw = getattr(part.file, "bytes", None) + if raw is not None: + with open(tmp_path, "wb") as f: + f.write(base64.b64decode(raw)) + else: + uri = getattr(part.file, "uri", None) + if not uri: + # ValueError, not TypeError, so export_one catches it and + # fails only this service — not the whole snapshot. + raise ValueError( + f"FilePart for '{environment_name}' has neither bytes nor uri" + ) + # Any url but http(s) is an object the service uploaded itself — hand + # it back as-is; export_one registers it directly (no download). Its + # last segment carries the shape (e.g. gdrive.zip). + scheme = urlparse(uri).scheme + if scheme and scheme not in ("http", "https"): + return uri + # A relative uri (e.g. "export-snapshot") means "stream it from + # my service endpoint" — resolve against the service base_url so + # multi-GB bundles never ride inline in the JSON-RPC response. An + # absolute one must be on that same origin: this process fetches it. + if not scheme: + uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" + elif urlparse(uri)[:2] != urlparse(base_url)[:2]: + raise ValueError( + f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" + ) + async with httpx.AsyncClient() as client: + async with client.stream( + "GET", + uri, + timeout=httpx.Timeout( + timeout_seconds, connect=_CONNECT_TIMEOUT_SECONDS + ), + ) as resp_file: + resp_file.raise_for_status() + with open(tmp_path, "wb") as f: + async for chunk in resp_file.aiter_bytes(): + f.write(chunk) + return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + state = part.data if part is not None else {} + with open(tmp_path, "w") as f: + json.dump(state, f, default=str) + return ".json" legacy_url = f"{legacy_protocol.environment_base_url(gateway_url, environment_name, mcp=True)}/export-state" async with httpx.AsyncClient() as client: @@ -558,21 +578,17 @@ async def snapshot_env_state( 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[str, dict]] = None, ) -> EnvSnapshotResult: """Export all services via the gateway, or without one (``gateway_url`` None) through ``deployed``'s recorded card, into a new EnvironmentUniverseArtifact ``snapshot_id``. All-or-nothing; the raised error carries public per-service summaries (verbose details - go to the log). ``data_get_params`` holds more params for the data/get of each environment, by name, sent - as they are, never the ones the snapshot sets itself. + 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) if gateway_url else _card_environments(deployed) - if reserved := sorted({key for params in (data_get_params or {}).values() for key in params} & _SETS_ITSELF): - raise ValueError(f"data_get_params sets {', '.join(reserved)}, which the snapshot sets itself") 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( @@ -615,10 +631,9 @@ async def export_one(environment_name: str): artifact_path = tmp_path try: try: - extra = (data_get_params or {}).get(environment_name) result = await cls._export_environment_to_file( gateway, environment_name, tmp_path, export_timeout_seconds, - deployed=card_record, **({"extra": extra} if extra else {}), + deployed=card_record, ) except httpx.HTTPStatusError as e: body_snippet = "" @@ -634,20 +649,6 @@ async def export_one(environment_name: str): except ValueError as e: return str(e), str(e) - if isinstance(result, list): - try: - parts = await asyncio.to_thread( - _put_parts, result, tmp_path, 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, parts=parts, - ) - except Exception as e: - public = f"failed to create artifacts ({type(e).__name__})" - return public, f"{public} for {environment_name}: {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), @@ -658,7 +659,7 @@ async def export_one(environment_name: str): content_type = mimetypes.guess_type(filename)[0] or "application/zip" def _register_uploaded_bundle(): - fa = _registered_upload( + fa = registered_upload( object_url, f"{snapshot_id}-{environment_name}-file", filename, content_type, f"State snapshot of {environment_name} from env {env.id}", ) @@ -699,7 +700,7 @@ def _register_uploaded_bundle(): public = f"failed to create artifacts ({type(e).__name__})" return public, f"{public}: {e}" finally: - for path in [*cleanup_paths, *glob.glob(f"{glob.escape(tmp_path)}-*")]: + for path in cleanup_paths: try: os.unlink(path) except OSError: @@ -747,81 +748,27 @@ def _register_uploaded_bundle(): ) -async def _materialize(part: Optional[Part], upload: Optional[_SnapshotUpload], base_url: str, tmp_path: str, - environment_name: str, timeout_seconds: float) -> str: - """One part of a service's export: the object url it was uploaded to under ``upload``, else the suffix written to - ``tmp_path`` (".json" for a data part or none, ".txt" for a text part, else its file's).""" - if upload is not None and (path := uploaded_object_path(part)) is not None: - return upload.object_url(path) - if isinstance(part, TextPart): - Path(tmp_path).write_text(part.text) - return ".txt" - if isinstance(part, FilePart): - raw = getattr(part.file, "bytes", None) - if raw is not None: - with open(tmp_path, "wb") as f: - f.write(base64.b64decode(raw)) - else: - uri = getattr(part.file, "uri", None) - if not uri: - # ValueError, not TypeError, so export_one catches it and - # fails only this service — not the whole snapshot. - raise ValueError( - f"FilePart for '{environment_name}' has neither bytes nor uri" - ) - # Any url but http(s) is an object the service uploaded itself — hand - # it back as-is; export_one registers it directly (no download). Its - # last segment carries the shape (e.g. gdrive.zip). - scheme = urlparse(uri).scheme - if scheme and scheme not in ("http", "https"): - return uri - # A relative uri (e.g. "export-snapshot") means "stream it from - # my service endpoint" — resolve against the service base_url so - # multi-GB bundles never ride inline in the JSON-RPC response. An - # absolute one must be on that same origin: this process fetches it. - if not scheme: - uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}" - elif urlparse(uri)[:2] != urlparse(base_url)[:2]: - raise ValueError( - f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}" - ) - async with httpx.AsyncClient() as client: - async with client.stream( - "GET", - uri, - timeout=httpx.Timeout( - timeout_seconds, connect=_CONNECT_TIMEOUT_SECONDS - ), - ) as resp_file: - resp_file.raise_for_status() - with open(tmp_path, "wb") as f: - async for chunk in resp_file.aiter_bytes(): - f.write(chunk) - return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" - state = part.data if part is not None else {} - with open(tmp_path, "w") as f: - json.dump(state, f, default=str) - return ".json" - - -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 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 []) -def _returns_file_tree(card: dict) -> bool: - """Whether ``card`` declares its data/get returns a file tree, each file a part named by its path, which one file - can't keep.""" - 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 -def _registered_upload(object_url: str, file_id: str, filename: str, content_type: str, description: str): +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 @@ -834,33 +781,15 @@ def _registered_upload(object_url: str, file_id: str, filename: str, content_typ id=file_id, description=description, filename=filename, content_type=content_type, object_url=object_url)) -def _put_parts(exported: list[tuple[Part, str]], tmp_path: str, id: str, description: str) -> list: - """An ``EnvironmentPart`` of each part a service exported, in order, with what ``_materialize`` made of it: an - object it uploaded, written in place, or the file at ``-``, written as ``-part-`` - under the part's own name.""" - from agent_env.artifact import FileArtifact - from agent_env.artifact.artifacts.environment import EnvironmentPart - - parts = [] - with tempfile.TemporaryDirectory(prefix="agentenv-snapshot-parts-") as named: - for index, (part, result) in enumerate(exported): - file = part.file if isinstance(part, FilePart) else None - file_id = f"{id}-part-{index}" - if urlparse(result).scheme: - filename = getattr(file, "name", None) or result.rsplit("/", 1)[-1] or f"part-{index}" - stored = _registered_upload(result, file_id, filename, - mimetypes.guess_type(filename)[0] or "application/octet-stream", - description) - else: - path = Path(named, str(index), Path(getattr(file, "name", None) or f"part-{index}{result}").name) - path.parent.mkdir() - os.replace(f"{tmp_path}-{index}", path) - stored = FileArtifact.put(file_id, description=description, file_path=str(path)) - kind = "file" if file is not None else "text" if isinstance(part, TextPart) else "data" - parts.append(EnvironmentPart(kind=kind, file_artifact_ref=stored.as_ref(), - name=getattr(file, "name", None), mime_type=getattr(file, "mimeType", None), - metadata=part.metadata or None)) - return parts +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]: diff --git a/tst/integration/computer/computer_env_test.py b/tst/integration/computer/computer_env_test.py index 7f039f2c5..3a9bb115e 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -3,9 +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, is snapshotted without a gateway by snapshot_computer_env, each file under the -paths it names a part by its path, comes back from a reset fresh, 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 +code that reaches its driver on loopback, is snapshotted by snapshot_computer_env, its files under the paths it names +kept by their paths 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 @@ -20,7 +21,8 @@ from agentenv_protocol.computer.conformance import run_probe from agentenv_protocol.computer.conformance.client import Driver -from agent_env.artifact import EnvironmentUniverseArtifact, FileArtifact +from agent_env.artifact import FileArtifact +from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse from agent_env.artifact.artifacts.computer import ComputerArtifact from agent_env.env import Env from agent_env.env.envs.computer.env import ComputerEnv @@ -38,7 +40,6 @@ 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 agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep from tst.integration.computer.conftest import docker pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -126,19 +127,21 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an snapshot_id=f"{closed.id}-snapshot", file_paths=["Desktop", "Downloads/note.txt"]).execute(context) taken = context.metadata["env_snapshotted_universes"]["snapshot"] - (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() - snapshot = {part.path: file.load() for part, file in environment.get_parts()} + snapshot = {path: file.load() for path, file in FileArtifactUniverse.get( + taken["id"], version=taken["version"]).get_file_artifacts().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", destination_path=".").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 (environment.environment_name, snapshot) == ("closed", {"Desktop/left.txt": b"x\n", - "Downloads/note.txt": b"inside\n"}) + 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}"], @@ -156,11 +159,11 @@ async def test_an_open_computer_uploads_each_file_under_the_snapshots_grant_and_ 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 SnapshotEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, - snapshot_id=f"{desktop.id}-snapshot").execute(context) + context = await SnapshotComputerEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, + snapshot_id=f"{desktop.id}-snapshot").execute(context) taken = context.metadata["env_snapshotted_universes"]["snapshot"] - (environment,) = EnvironmentUniverseArtifact.get(taken["id"], version=taken["version"]).get_environment_artifacts() - snapshot = {part.path: (file.object_url, file.load()) for part, file in environment.get_parts()} + snapshot = {path: (file.object_url, file.load()) for path, file in FileArtifactUniverse.get( + taken["id"], version=taken["version"]).get_file_artifacts().items()} finally: report = await teardown_run(context) grant_server(None, None).close() diff --git a/tst/unit/artifact/environment_parts_test.py b/tst/unit/artifact/environment_parts_test.py deleted file mode 100644 index 4c76a64ad..000000000 --- a/tst/unit/artifact/environment_parts_test.py +++ /dev/null @@ -1,62 +0,0 @@ -"""An environment artifact of the parts an environment's data/get returned: written and read back in order, stored with -no wrapped file and nothing of the one-file shape, which stores as before; refused where one file is taken, and by an -artifact folder.""" - -import pytest -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI - -from agent_env.artifact.artifacts.environment import EnvironmentArtifact, EnvironmentPart -from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact -from agent_env.artifact.artifacts.file import FileArtifact -from agent_env.entity.export import Unwritable - - -def _file(tmp_path, name, data): - (tmp_path / name).write_bytes(data) - return FileArtifact.put(f"file-{name}", description=name, file_path=str(tmp_path / name)) - - -def test_an_environment_of_parts_is_written_and_read_back_in_order_with_nothing_of_one_file(local_stores, tmp_path): - notes, state = _file(tmp_path, "notes.txt", b"hello"), _file(tmp_path, "state.json", b'{"rows": 1}') - parts = [EnvironmentPart(kind="file", file_artifact_ref=notes.as_ref(), name="notes.txt", mime_type="text/plain", - metadata={DATA_OBJECTS_EXTENSION_URI: {"path": "Documents/notes.txt"}}), - EnvironmentPart(kind="data", file_artifact_ref=state.as_ref())] - - written = EnvironmentArtifact.put("desk", environment_name="desktop", parts=parts) - - read = EnvironmentArtifact.get("desk", version=written.version) - assert [(part.kind, part.path, file.load()) for part, file in read.get_parts()] == [ - ("file", "Documents/notes.txt", b"hello"), ("data", None, b'{"rows": 1}')] - assert read.parts == parts and read.file_artifact_ref is None - assert "file_artifact_id" not in read.model_dump(by_alias=True) - with pytest.raises(ValueError, match="holds the 2 parts its environment's data/get returned, not one file"): - read.get_file_artifact() - - -def test_one_file_is_stored_as_before_and_read_as_one_file_part(local_stores, tmp_path): - state = _file(tmp_path, "state.json", b"{}") - - written = EnvironmentArtifact.put("slack", environment_name="slack", file_artifact=state) - - assert "parts" not in written.model_dump(by_alias=True) and written.file_artifact_ref == state.as_ref() - assert [(part.kind, part.name, file.id) for part, file in EnvironmentArtifact.get("slack").get_parts()] == [ - ("file", "state.json", state.id)] - - -@pytest.mark.parametrize("given", [{}, {"file_artifact": "both", "parts": []}], ids=["neither", "both"]) -def test_an_environment_wraps_one_file_or_holds_parts_but_not_both_or_neither(local_stores, tmp_path, given): - if given: - given = {"file_artifact": _file(tmp_path, "state.json", b"{}"), "parts": []} - - with pytest.raises(ValueError, match="wraps a file_artifact or holds parts, one of the two"): - EnvironmentArtifact.put("x", environment_name="x", **given) - - -def test_an_artifact_folder_holds_no_environment_of_parts(local_stores, tmp_path): - desk = EnvironmentArtifact.put("desk", environment_name="desktop", parts=[]) - universe = EnvironmentUniverseArtifact.put("universe", environment_artifacts=[desk]) - - with pytest.raises(Unwritable, match="it holds the parts its data/get returned"): - desk.to_toml(None) - with pytest.raises(Unwritable, match="its environment 'desktop' holds the parts its data/get returned"): - universe.to_toml(None) diff --git a/tst/unit/artifact/test_service_universe_file_view.py b/tst/unit/artifact/test_service_universe_file_view.py index d5aa8197a..6883da8a9 100644 --- a/tst/unit/artifact/test_service_universe_file_view.py +++ b/tst/unit/artifact/test_service_universe_file_view.py @@ -2,9 +2,8 @@ service universe be staged as a file tree (grading) by the same loaders that handle a FileArtifactUniverse. -Contract: keys are `/`, or for an environment of the -parts its data/get returned, `/`; values -are the universe's EXISTING FileArtifacts (nothing copied or re-registered). +Contract: keys are `/`, values are the +universe's EXISTING FileArtifacts (nothing copied or re-registered). """ from __future__ import annotations @@ -12,9 +11,7 @@ import pytest from agent_env.artifact.artifacts.file import FileArtifact -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI - -from agent_env.artifact.artifacts.environment import EnvironmentArtifact, EnvironmentPart +from agent_env.artifact.artifacts.environment import EnvironmentArtifact from agent_env.artifact.artifacts.environment_universe import EnvironmentUniverseArtifact from agent_env.artifact.ref import ArtifactRef @@ -94,23 +91,3 @@ def test_duplicate_service_names_raise(registry): universe = _universe(registry, [("slack", "a.json"), ("slack", "a.json")]) with pytest.raises(ValueError, match="duplicate environment 'slack'"): universe.get_file_artifacts() - - -def test_an_environment_of_parts_lists_each_part_by_its_path_else_its_name(registry): - files, services = registry - universe = _universe(registry, [("slack", "state.json")]) - parts = [] - for idx, (filename, name, path) in enumerate([("notes.txt", "notes.txt", "Documents/notes.txt"), - ("c.bin", "c.bin", "Downloads/b/c.bin"), - ("state.json", None, None)]): - files[f"part-{idx}"] = _file_artifact(f"part-{idx}", filename) - parts.append(EnvironmentPart(kind="file" if name else "data", file_artifact_ref=files[f"part-{idx}"].as_ref(), - name=name, metadata={DATA_OBJECTS_EXTENSION_URI: {"path": path}} if path else None)) - services["desk"] = EnvironmentArtifact(id="desk", version=1, environment_name="desktop", parts=parts) - universe.environment_artifact_refs.append(ArtifactRef(id="desk", version=1)) - - viewed = universe.get_file_artifacts() - - assert sorted(viewed) == ["desktop/Documents/notes.txt", "desktop/Downloads/b/c.bin", "desktop/state.json", - "slack/state.json"] - assert viewed["desktop/Downloads/b/c.bin"] is files["part-1"] diff --git a/tst/unit/task_step/snapshot_computer_env_test.py b/tst/unit/task_step/snapshot_computer_env_test.py index 9681a9468..000e6b35a 100644 --- a/tst/unit/task_step/snapshot_computer_env_test.py +++ b/tst/unit/task_step/snapshot_computer_env_test.py @@ -1,17 +1,24 @@ -"""snapshot_computer_env: snapshot_env for a computer env, its file_paths, relative to the session user's home, the -data/get params of the computer's own environment and of nothing else; refused for an env that isn't a computer, and -for paths that aren't relative, before a run.""" +"""snapshot_computer_env: a computer env's files kept as a FileArtifactUniverse by their paths, 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 EnvironmentUniverseArtifact +from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +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 _computer, _inline, _record, _serve +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): @@ -25,15 +32,29 @@ class _Mail: id, version, type, environment_name = "mail", 1, "mcp_server", "mail" -def test_its_file_paths_round_trip_and_its_type_is_registered(): +async def _snapshot(monkeypatch, card: dict, env=None, **fields) -> FileArtifactUniverse: + 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"] + return FileArtifactUniverse.get(snapshotted["id"], version=snapshotted["version"]) + + +def _kept(universe: FileArtifactUniverse) -> dict[str, bytes]: + return {path: file.load() for path, file in universe.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"], - include_env_trajectory=True) + export_timeout_seconds=60) - restored = get_task_step_registry()["snapshot_computer_env"].from_dict(step.to_dict()) + data = step.to_dict() + restored = get_task_step_registry()["snapshot_computer_env"].from_dict(data) - assert isinstance(restored, SnapshotComputerEnvTaskStep) and step.to_dict()["type"] == "snapshot_computer_env" - assert (restored.env_id, restored.file_paths, restored.include_env_trajectory) == ( - "desk", ["Documents", "Desktop/a.pdf"], True) + 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/../.."]], @@ -57,24 +78,62 @@ def test_an_env_that_isnt_a_computer_is_refused_before_a_run(monkeypatch, stored @pytest.mark.asyncio @pytest.mark.parametrize("file_paths, params", [(["Documents"], {"paths": ["Documents"]}), (None, {})], ids=["file-paths", "the-cards-folders"]) -async def test_its_file_paths_are_the_computers_own_data_get_params(local_stores, monkeypatch, file_paths, params): - sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")]}) - monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Desk())) - context = _ctx(_record(_computer())) +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")]}) - await SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk", file_paths=file_paths).execute(context) + universe = await _snapshot(monkeypatch, _computer(), file_paths=file_paths) - snapshotted = context.metadata["env_snapshotted_universes"]["snap"] - universe = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]) - assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt"] + assert _kept(universe) == {"Documents/notes.txt": b"hello", "Documents/b/c.bin": b"\x00\x01"} + assert universe.get_file_artifacts()["Documents/b/c.bin"].filename == "c.bin" assert sent == [("/agentenv", params)] @pytest.mark.asyncio -async def test_a_run_on_an_env_that_isnt_a_computer_is_refused_before_any_export(local_stores, monkeypatch): - sent = _serve(monkeypatch, {}) - monkeypatch.setattr("agent_env.env.env.Env.get", classmethod(lambda cls, id, version=None: _Mail())) +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}) + + universe = 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 universe.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")]}) - with pytest.raises(RuntimeError, match="snapshots a computer env, and mail is a mcp_server env"): - await SnapshotComputerEnvTaskStep(id="snap", version=None, env_id="desk").execute(_ctx(_record(_computer()))) - assert sent == [] + universe = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) + + assert _kept(universe) == {"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: []}, "computer env desk has no files under Documents 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_upload_grant_test.py b/tst/unit/task_step/snapshot_env_upload_grant_test.py index 5644b8904..9554f333c 100644 --- a/tst/unit/task_step/snapshot_env_upload_grant_test.py +++ b/tst/unit/task_step/snapshot_env_upload_grant_test.py @@ -94,7 +94,7 @@ async def test_a_service_taking_the_form_uploads_through_a_grant_and_is_register (params,) = _rpc_params(sent) grant = params["write_namespace"] assert grant["root_path"].startswith("agentenv-snapshots/") - assert (grant["max_objects"], grant["max_object_bytes"]) == (100_000, ENV_SNAPSHOT_LIMITS.max_object_bytes) + assert (grant["max_objects"], grant["max_object_bytes"]) == (1, ENV_SNAPSHOT_LIMITS.max_object_bytes) assert datetime.fromisoformat(grant["expires_at"]) >= datetime.now(UTC) + timedelta(seconds=MIN_GRANT_LIFETIME_SECONDS - 60) assert pushes == [_CARD] # still pushed, for a bundle too large for the grant assert [str(r.url) for r in sent] == [f"{_CARD}/svc/mcp-gdrive/agentenv"] diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py index 5965043a1..f964221d4 100644 --- a/tst/unit/task_step/snapshot_env_without_gateway_test.py +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -1,7 +1,7 @@ -"""snapshot_env on an env deployed without a gateway, a computer env say: reached through the card its deployment -recorded, its environments the card's children or the env itself, and each environment's export kept as data/get -returned it, a file tree's files as parts by their paths, inline or uploaded under the grant without a download, and -more data/get params handed to the environment they name (snapshot_computer_env names a computer's files so).""" +"""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, refused for snapshot_computer_env; and an env that fails under the upload +grant asked again without it.""" from __future__ import annotations @@ -11,8 +11,7 @@ import httpx import pytest -from agentenv_protocol import (DATA_OBJECTS_EXTENSION_URI, IntakeDeclaration, IntakeFormat, intake_extension, - uploaded_file_part) +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 @@ -26,6 +25,8 @@ 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: @@ -42,6 +43,10 @@ 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"}, @@ -64,69 +69,61 @@ def handle(request: httpx.Request) -> httpx.Response: return sent -async def _snapshot(card: dict, **options) -> EnvironmentUniverseArtifact: +async def _snapshot(card: dict) -> EnvironmentUniverseArtifact: result = await SnapshotEnvTaskStep.snapshot_env_state(env=_Desk(), gateway_url=None, snapshot_id="snap-d", - deployed=_record(card), **options) + deployed=_record(card)) return EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id) -@pytest.mark.asyncio -async def test_a_computer_keeps_each_file_its_data_get_returns_as_a_part_by_its_path(local_stores, monkeypatch): - sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello"), - _inline("Downloads/b/c.bin", b"\x00\x01")]}) - - universe = await _snapshot(_computer()) - - (desk,) = universe.get_environment_artifacts() - assert (desk.environment_name, desk.file_artifact_ref) == ("desktop", None) - assert [(part.kind, part.path, part.name, part.mime_type, file.filename, file.load()) - for part, file in desk.get_parts()] == [ - ("file", "Documents/notes.txt", "notes.txt", "text/plain", "notes.txt", b"hello"), - ("file", "Downloads/b/c.bin", "c.bin", "text/plain", "c.bin", b"\x00\x01")] - assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt", "desktop/Downloads/b/c.bin"] - assert sent == [("/agentenv", {})] +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_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"} +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")]}) - 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] + universe = await _snapshot(card) - sent = _serve(monkeypatch, {"/agentenv": uploading}) + assert _kept(universe) == {"mail": b'{"inbox": 2}', "notes": b"PK"} + assert sorted(path for path, _ in sent) == ["/svc/mail/agentenv", "/svc/notes/agentenv"] - (desk,) = (await _snapshot(_computer(DATA_OBJECTS))).get_environment_artifacts() - ((path, params),) = sent - root = params["write_namespace"]["root_path"] - assert params["write_namespace"]["max_objects"] == 100_000 - assert [(part.path, file.object_url, file.load()) for part, file in desk.get_parts()] == [ - (path, store.object_url(f"{root}/{path}"), data) for path, data in files.items()] +@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()) -UNREACHABLE = {"error": {"code": -32000, "message": "The object store is temporarily unavailable.", - "data": "get_failed"}} + assert _kept(universe) == {"desktop": b'{"rows": 1}'} + assert sent == [("/agentenv", {})] @pytest.mark.asyncio -async def test_an_env_that_fails_under_the_grant_is_asked_again_without_one_and_returns_its_files_inline( - local_stores, tmp_path, monkeypatch): +async def test_a_file_tree_is_refused_for_snapshot_computer_env_before_its_data_get(local_stores, monkeypatch): + sent = _serve(monkeypatch, {}) + + with pytest.raises(RuntimeError, match="'desktop' returns a file tree from data/get, which one file can't keep; " + "snapshot_computer_env snapshots a computer env's files"): + await _snapshot(_computer()) + assert sent == [] + + +@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 [_inline("Desktop/left.txt", b"x")]}) + else [{"kind": "data", "data": {"rows": 1}}]}) - (desk,) = (await _snapshot(_computer(DATA_OBJECTS))).get_environment_artifacts() + 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", [])] - assert [(part.path, file.load()) for part, file in desk.get_parts()] == [("Desktop/left.txt", b"x")] @pytest.mark.asyncio @@ -134,68 +131,19 @@ async def test_an_env_that_fails_with_no_grant_is_asked_once(local_stores, monke sent = _serve(monkeypatch, {"/agentenv": lambda params: UNREACHABLE}) with pytest.raises(RuntimeError, match="data/get failed: The object store is temporarily unavailable"): - await _snapshot(_computer()) + await _snapshot(_server()) assert sent == [("/agentenv", {})] -@pytest.mark.asyncio -async def test_data_get_params_go_to_the_environment_they_name_and_each_child_on_its_card_is_snapshotted( - local_stores, monkeypatch): - card = {"name": "desktop", "children_environments": [ - {"name": "desktop", "url": "/agentenv", "capabilities": {"extensions": [FILE_TREE]}}, - {"name": "mail", "url": "/svc/mail/agentenv", "capabilities": {"extensions": []}}]} - sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")], - "/svc/mail/agentenv": lambda params: [{"kind": "data", "data": {"inbox": 2}}]}) - - universe = await _snapshot(card, data_get_params={"desktop": {"paths": ["Documents"]}}) - - desk, mail = sorted(universe.get_environment_artifacts(), key=lambda environment: environment.environment_name) - assert sorted(sent) == [("/agentenv", {"paths": ["Documents"]}), ("/svc/mail/agentenv", {})] - assert [part.path for part, _ in desk.get_parts()] == ["Documents/notes.txt"] - assert mail.parts is None and json.loads(mail.get_file_artifact().load()) == {"inbox": 2} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("card, answer, kept", [ - (_computer(), [], []), - ({"name": "desktop", "capabilities": {"extensions": []}}, - [{"kind": "data", "data": {"a": 1}}, {"kind": "data", "data": {"b": 2}}], - [("data", "part-0.json", b'{"a": 1}'), ("data", "part-1.json", b'{"b": 2}')]), - ({"name": "desktop", "capabilities": {"extensions": []}}, - [{"kind": "data", "data": {"a": 1}}, {"kind": "text", "text": "hello"}], - [("data", "part-0.json", b'{"a": 1}'), ("text", "part-1.txt", b"hello")]), - ({"name": "desktop", "capabilities": {"extensions": []}}, [{"kind": "text", "text": "hello"}], - [("text", "part-0.txt", b"hello")]), -], ids=["a-file-tree-of-no-files", "several-parts-of-no-file-tree", "a-data-part-then-a-text-part", - "one-text-part"]) -async def test_an_export_one_file_cant_keep_is_kept_as_its_parts(local_stores, monkeypatch, card, answer, kept): - _serve(monkeypatch, {"/agentenv": lambda params: answer}) - - (desk,) = (await _snapshot(card)).get_environment_artifacts() - - assert [(part.kind, file.filename, file.load()) for part, file in desk.get_parts()] == kept and desk.parts == [ - part for part, _ in desk.get_parts()] - - -@pytest.mark.asyncio -async def test_data_get_params_never_set_what_the_snapshot_sets_itself(local_stores, monkeypatch): - sent = _serve(monkeypatch, {"/agentenv": lambda params: []}) - - with pytest.raises(ValueError, match="data_get_params sets write_namespace, which the snapshot sets itself"): - await _snapshot(_computer(), data_get_params={"desktop": {"write_namespace": {}}}) - assert sent == [] - - @pytest.mark.asyncio async def test_execute_snapshots_an_env_deployed_without_a_gateway_through_its_recorded_card(local_stores, monkeypatch): - sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello")]}) + _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(_computer())) + context = _ctx(_record(_server())) await SnapshotEnvTaskStep(id="snap", version=None, env_id="desk").execute(context) snapshotted = context.metadata["env_snapshotted_universes"]["snap"] - universe = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]) - assert sorted(universe.get_file_artifacts()) == ["desktop/Documents/notes.txt"] - assert sent == [("/agentenv", {})] + assert _kept(EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"])) == { + "desktop": b'{"rows": 1}'} From b37083df3a6f39cbc23a297b2ad35108338505b7 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 18:01:28 -0700 Subject: [PATCH 7/8] feat(snapshot): snapshot envs without a gateway, and a computer env's files with snapshot_computer_env The branch's change rebased onto main as one commit: snapshot_env reaches an env deployed without a gateway through the card its deployment recorded; a new snapshot_computer_env step keeps a computer env's files; an env failing data/get under the upload grant is asked again without it. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode/rules.py | 2 +- src/agent_env/task_step/registry.py | 2 + .../task_step/task_steps/computer/__init__.py | 1 + .../computer/snapshot_computer_env.py | 166 ++++++++++++++++++ .../task_step/task_steps/snapshot_env.py | 122 +++++++++---- tst/integration/computer/computer_env_test.py | 43 ++++- .../task_step/snapshot_computer_env_test.py | 139 +++++++++++++++ .../snapshot_env_without_gateway_test.py | 149 ++++++++++++++++ tst/unit/task_step/test_field_scopes.py | 1 + tst/unit/task_step/test_snapshot_env.py | 5 +- 10 files changed, 586 insertions(+), 44 deletions(-) create mode 100644 src/agent_env/task_step/task_steps/computer/__init__.py create mode 100644 src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py create mode 100644 tst/unit/task_step/snapshot_computer_env_test.py create mode 100644 tst/unit/task_step/snapshot_env_without_gateway_test.py diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 33ad6fa38..86d84c921 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 3c74f4fcd..83821a632 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 000000000..16fef9855 --- /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 000000000..21f8049b2 --- /dev/null +++ b/src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py @@ -0,0 +1,166 @@ +"""snapshot_computer_env: a computer env's files, kept as a FileArtifactUniverse by their paths, which load_artifact +loads back into a computer.""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +from pathlib import PurePosixPath +from typing import ClassVar, Optional + +from agentenv_protocol import uploaded_object_path +from agentenv_protocol.types import FileWithBytes +from agentenv_protocol.transfers import RelativePath +from pydantic import TypeAdapter, ValidationError + +from agent_env.a2a_agent.object_transfer import ObjectLimits +from agent_env.artifact import FileArtifact +from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +from agent_env.entity.refs import EntityRef, RefRole +from agent_env.env import legacy_protocol +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 derive_id, 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, + _snapshot_upload, + get_data, + registered_upload, +) + +logger = logging.getLogger(__name__) + +_PATHS = TypeAdapter(list[RelativePath]) +#: What a computer may upload through its snapshot 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 its files are written to the store at once. +_WRITES_IN_FLIGHT = 8 + + +class SnapshotComputerEnvTaskStep(SnapshotEnvTaskStep): + """Snapshot a deployed computer env's files into a new ``FileArtifactUniverse``, each under its path relative to + the session user's home: those under ``file_paths``, by default the folders its card declares. The computer + uploads each file through a grant when the store's grants reach it, and returns it inline otherwise. + ``load_artifact`` loads the universe back into a computer, through ``data/add``. + + Picks the deployment and names the snapshot as ``snapshot_env`` does. 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="file_artifact_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)) + files = await _files(deployed, env, snapshot_id, self.file_paths, self.export_timeout_seconds) + if not files: + under = ", ".join(self.file_paths) if self.file_paths is not None else "the folders its card declares" + raise RuntimeError(f"computer env {env.id} has no files under {under} to snapshot") + universe = await asyncio.to_thread(FileArtifactUniverse.put, snapshot_id, file_artifacts=files) + logger.info(f"snapshot_computer_env: {env.id} -> {universe.id} v{universe.version} ({len(files)} files)") + context.metadata.setdefault("env_snapshotted_universes", {})[self.id] = {"id": universe.id, + "version": universe.version} + return context + + +async def _files(deployed, env: ComputerEnv, snapshot_id: str, paths: Optional[list[str]], + timeout_seconds: float) -> dict[str, FileArtifact]: + """Each file ``env``'s own data/get returns under ``paths``, by its path, written to the store as ``__file-``: an upload under the grant in place, without a download, else its inline bytes.""" + name = env.environment_name + base_url = await legacy_protocol.v1_base_url(deployed, None, name) + _, card = await legacy_protocol.child_env_card(deployed, None, name) + if base_url is None: + raise RuntimeError(f"computer env {env.id}'s recorded card names no environment {name!r}") + upload = await asyncio.to_thread(_snapshot_upload, card or {}, timeout_seconds, deployed.sandbox_type, + FILE_TREE_LIMITS) + resp, upload = await get_data(base_url, name, timeout_seconds, upload, + **({} if paths is None else {"paths": paths})) + writing = asyncio.Semaphore(_WRITES_IN_FLIGHT) + + async def written(index: int, part) -> tuple[str, FileArtifact]: + path = uploaded_object_path(part) + if path is None: + raise RuntimeError(f"computer env {env.id}'s data/get returned a {part.kind} part naming no path") + file_id, filename = derive_id(snapshot_id, f"file-{index}"), PurePosixPath(path).name + content_type, description = part.file.mimeType or "application/octet-stream", f"{path} of computer env {env.id}" + async with writing: + if isinstance(part.file, FileWithBytes): + return path, await asyncio.to_thread(FileArtifact.put_bytes, file_id, description=description, + filename=filename, content=base64.b64decode(part.file.bytes), + content_type=content_type) + if upload is None: + raise RuntimeError(f"computer env {env.id}'s data/get returned {path} neither inline nor uploaded") + return path, await asyncio.to_thread(registered_upload, upload.object_url(path), file_id, filename, + content_type, description) + + return dict(await asyncio.gather(*(written(index, part) for index, part in enumerate(resp.parts)))) + + +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 49a108d81..44822df77 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 @@ -17,7 +18,8 @@ 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 agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant @@ -62,21 +64,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 +160,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 refused, since + one file can't keep it: ``snapshot_computer_env`` snapshots a computer env's files. """ type: ClassVar[str] = "snapshot_env" @@ -238,9 +245,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,7 +469,7 @@ 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 + gateway_url: Optional[str], environment_name: str, tmp_path: str, timeout_seconds: float, deployed=None ) -> str: """Export one service; return the suffix written (".json"/".zip") to ``tmp_path``, or the object url a FilePart names when the service uploaded the @@ -471,21 +479,23 @@ async def _export_environment_to_file( form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` → object url (no download) / base64 ``bytes`` / streamed relative-or-same-origin ``uri`` → ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON - guard. Errors are caught by the caller (fails just this service).""" + guard. Errors are caught by the caller (fails just this service): a file tree, which one file can't + keep, among them.""" 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 {}): + raise ValueError(f"'{environment_name}' returns a file tree from data/get, which one file can't keep; " + "snapshot_computer_env snapshots a computer env's files") 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) 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 +572,28 @@ 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, ) -> 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``. + 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 +618,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): @@ -643,26 +654,15 @@ async def export_one(environment_name: str): # 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 +748,50 @@ 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 + + +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 5463a3eac..3a9bb115e 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 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.file_artifact_universe import FileArtifactUniverse 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,25 @@ 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) + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + snapshot = {path: file.load() for path, file in FileArtifactUniverse.get( + taken["id"], version=taken["version"]).get_file_artifacts().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", destination_path=".").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 +150,30 @@ 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) + taken = context.metadata["env_snapshotted_universes"]["snapshot"] + snapshot = {path: (file.object_url, file.load()) for path, file in FileArtifactUniverse.get( + taken["id"], version=taken["version"]).get_file_artifacts().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(): 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 000000000..000e6b35a --- /dev/null +++ b/tst/unit/task_step/snapshot_computer_env_test.py @@ -0,0 +1,139 @@ +"""snapshot_computer_env: a computer env's files kept as a FileArtifactUniverse by their paths, 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.artifacts.file_artifact_universe import FileArtifactUniverse +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) -> FileArtifactUniverse: + 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"] + return FileArtifactUniverse.get(snapshotted["id"], version=snapshotted["version"]) + + +def _kept(universe: FileArtifactUniverse) -> dict[str, bytes]: + return {path: file.load() for path, file in universe.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")]}) + + universe = await _snapshot(monkeypatch, _computer(), file_paths=file_paths) + + assert _kept(universe) == {"Documents/notes.txt": b"hello", "Documents/b/c.bin": b"\x00\x01"} + assert universe.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}) + + universe = 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 universe.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")]}) + + universe = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) + + assert _kept(universe) == {"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: []}, "computer env desk has no files under Documents 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 000000000..f964221d4 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -0,0 +1,149 @@ +"""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, refused for snapshot_computer_env; 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_refused_for_snapshot_computer_env_before_its_data_get(local_stores, monkeypatch): + sent = _serve(monkeypatch, {}) + + with pytest.raises(RuntimeError, match="'desktop' returns a file tree from data/get, which one file can't keep; " + "snapshot_computer_env snapshots a computer env's files"): + await _snapshot(_computer()) + assert sent == [] + + +@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 86aa4948d..54562a0eb 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_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index dd4087b88..6e70635b1 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 From 7a611b82398b736eec451594d1870512bcf4ad1e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 31 Oct 2026 18:12:25 -0700 Subject: [PATCH 8/8] feat(snapshot): keep a file tree's files in its environment artifact, and load a computer's snapshot back into it Built on environment artifacts of several files (#424): - snapshot_env keeps an environment whose card declares a file tree, a computer's, file by file by its paths in its environment artifact, under a grant sized for a file tree, instead of refusing it. Each file is written as derive_path_id(-, path). snapshot_env_state takes data_get_params for every environment's data/get. - snapshot_computer_env is snapshot_env with the computer's file_paths as data/get's paths, and writes the same environment universe snapshot_env does, in place of a FileArtifactUniverse. - A computer env loads an environment artifact, its own, by writing its files at their paths under the session user's home through the driver's data/add, and a universe by loading its own environment and leaving the others out with a warning. load_artifact of a computer's snapshot so restores it. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/envs/computer/env.py | 29 ++++++ .../computer/snapshot_computer_env.py | 89 ++++--------------- .../task_step/task_steps/snapshot_env.py | 80 ++++++++++++++--- tst/integration/computer/computer_env_test.py | 22 +++-- tst/unit/env/envs/test_computer_env.py | 45 +++++++++- tst/unit/task/local_run_write_fence_test.py | 2 +- .../task_step/snapshot_computer_env_test.py | 37 ++++---- .../snapshot_env_without_gateway_test.py | 21 +++-- .../test_load_artifact_service_universe.py | 8 +- tst/unit/task_step/test_snapshot_env.py | 16 ++-- 10 files changed, 214 insertions(+), 135 deletions(-) diff --git a/src/agent_env/env/envs/computer/env.py b/src/agent_env/env/envs/computer/env.py index 5881e0308..9270f1655 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_step/task_steps/computer/snapshot_computer_env.py b/src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py index 21f8049b2..f28c08bb3 100644 --- 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 @@ -1,61 +1,38 @@ -"""snapshot_computer_env: a computer env's files, kept as a FileArtifactUniverse by their paths, which load_artifact +"""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 -import base64 -import logging -from pathlib import PurePosixPath from typing import ClassVar, Optional -from agentenv_protocol import uploaded_object_path -from agentenv_protocol.types import FileWithBytes from agentenv_protocol.transfers import RelativePath from pydantic import TypeAdapter, ValidationError -from agent_env.a2a_agent.object_transfer import ObjectLimits -from agent_env.artifact import FileArtifact -from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse from agent_env.entity.refs import EntityRef, RefRole -from agent_env.env import legacy_protocol 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 derive_id, is_local_id, validate_local_id +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, - _snapshot_upload, - get_data, - registered_upload, -) - -logger = logging.getLogger(__name__) +from agent_env.task_step.task_steps.snapshot_env import DEFAULT_EXPORT_TIMEOUT_SECONDS, SnapshotEnvTaskStep _PATHS = TypeAdapter(list[RelativePath]) -#: What a computer may upload through its snapshot 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 its files are written to the store at once. -_WRITES_IN_FLIGHT = 8 class SnapshotComputerEnvTaskStep(SnapshotEnvTaskStep): - """Snapshot a deployed computer env's files into a new ``FileArtifactUniverse``, each under its path relative to - the session user's home: those under ``file_paths``, by default the folders its card declares. The computer - uploads each file through a grant when the store's grants reach it, and returns it inline otherwise. - ``load_artifact`` loads the universe back into a computer, through ``data/add``. + """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``. - Picks the deployment and names the snapshot as ``snapshot_env`` does. Output: - ``context.metadata["env_snapshotted_universes"][self.id]`` = ``{"id", "version"}``.""" + 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="file_artifact_universe"), + EntityRef.artifact("snapshot_id", role=RefRole.OUTPUT, artifact_type="environment_universe"), ) def __init__( @@ -117,50 +94,14 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: env = await asyncio.to_thread(Env.get, deployed.env_id, deployed.env_version) if not isinstance(env, ComputerEnv): raise RuntimeError(_not_a_computer(env)) - files = await _files(deployed, env, snapshot_id, self.file_paths, self.export_timeout_seconds) - if not files: - under = ", ".join(self.file_paths) if self.file_paths is not None else "the folders its card declares" - raise RuntimeError(f"computer env {env.id} has no files under {under} to snapshot") - universe = await asyncio.to_thread(FileArtifactUniverse.put, snapshot_id, file_artifacts=files) - logger.info(f"snapshot_computer_env: {env.id} -> {universe.id} v{universe.version} ({len(files)} files)") - context.metadata.setdefault("env_snapshotted_universes", {})[self.id] = {"id": universe.id, - "version": universe.version} + 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 -async def _files(deployed, env: ComputerEnv, snapshot_id: str, paths: Optional[list[str]], - timeout_seconds: float) -> dict[str, FileArtifact]: - """Each file ``env``'s own data/get returns under ``paths``, by its path, written to the store as ``__file-``: an upload under the grant in place, without a download, else its inline bytes.""" - name = env.environment_name - base_url = await legacy_protocol.v1_base_url(deployed, None, name) - _, card = await legacy_protocol.child_env_card(deployed, None, name) - if base_url is None: - raise RuntimeError(f"computer env {env.id}'s recorded card names no environment {name!r}") - upload = await asyncio.to_thread(_snapshot_upload, card or {}, timeout_seconds, deployed.sandbox_type, - FILE_TREE_LIMITS) - resp, upload = await get_data(base_url, name, timeout_seconds, upload, - **({} if paths is None else {"paths": paths})) - writing = asyncio.Semaphore(_WRITES_IN_FLIGHT) - - async def written(index: int, part) -> tuple[str, FileArtifact]: - path = uploaded_object_path(part) - if path is None: - raise RuntimeError(f"computer env {env.id}'s data/get returned a {part.kind} part naming no path") - file_id, filename = derive_id(snapshot_id, f"file-{index}"), PurePosixPath(path).name - content_type, description = part.file.mimeType or "application/octet-stream", f"{path} of computer env {env.id}" - async with writing: - if isinstance(part.file, FileWithBytes): - return path, await asyncio.to_thread(FileArtifact.put_bytes, file_id, description=description, - filename=filename, content=base64.b64decode(part.file.bytes), - content_type=content_type) - if upload is None: - raise RuntimeError(f"computer env {env.id}'s data/get returned {path} neither inline nor uploaded") - return path, await asyncio.to_thread(registered_upload, upload.object_url(path), file_id, filename, - content_type, description) - - return dict(await asyncio.gather(*(written(index, part) for index, part in enumerate(resp.parts)))) - - 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 44822df77..6bd8d2055 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -12,6 +12,7 @@ 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 @@ -21,12 +22,13 @@ 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 @@ -50,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" @@ -162,8 +169,8 @@ class SnapshotEnvTaskStep(TaskStep): ``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 refused, since - one file can't keep it: ``snapshot_computer_env`` snapshots a computer env's files. + 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" @@ -469,33 +476,36 @@ def _enumerate_environments(env) -> list[str]: @staticmethod async def _export_environment_to_file( - gateway_url: Optional[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`` → object url (no download) / base64 ``bytes`` / streamed relative-or-same-origin ``uri`` → ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON - guard. Errors are caught by the caller (fails just this service): a file tree, which one file can't - keep, among them.""" + guard. Errors are caught by the caller (fails just this service).""" from agent_env.env import legacy_protocol 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 {}): - raise ValueError(f"'{environment_name}' returns a file tree from data/get, which one file can't keep; " - "snapshot_computer_env snapshots a computer env's files") + 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) - resp, upload = await get_data(base_url, environment_name, timeout_seconds, upload) + 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) @@ -578,9 +588,11 @@ async def snapshot_env_state( 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, or without one (``gateway_url`` None) through - ``deployed``'s recorded card, into a new EnvironmentUniverseArtifact ``snapshot_id``. + ``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). @@ -633,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 = "" @@ -649,6 +661,19 @@ 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), @@ -768,6 +793,35 @@ async def get_data(base_url: str, environment_name: str, timeout_seconds: float, 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.""" diff --git a/tst/integration/computer/computer_env_test.py b/tst/integration/computer/computer_env_test.py index 3a9bb115e..0f0467849 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -4,7 +4,7 @@ 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, is snapshotted by snapshot_computer_env, its files under the paths it names -kept by their paths after its upload under the grant fails, comes back from a reset fresh, takes the snapshot back +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.""" @@ -22,7 +22,7 @@ from agentenv_protocol.computer.conformance.client import Driver from agent_env.artifact import FileArtifact -from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +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 @@ -126,13 +126,11 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an 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) - taken = context.metadata["env_snapshotted_universes"]["snapshot"] - snapshot = {path: file.load() for path, file in FileArtifactUniverse.get( - taken["id"], version=taken["version"]).get_file_artifacts().items()} + 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", destination_path=".").execute(context) + 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) @@ -161,9 +159,7 @@ async def test_an_open_computer_uploads_each_file_under_the_snapshots_grant_and_ "&& printf two > Downloads/b.txt") context = await SnapshotComputerEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, snapshot_id=f"{desktop.id}-snapshot").execute(context) - taken = context.metadata["env_snapshotted_universes"]["snapshot"] - snapshot = {path: (file.object_url, file.load()) for path, file in FileArtifactUniverse.get( - taken["id"], version=taken["version"]).get_file_artifacts().items()} + 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() @@ -215,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 87e14ad87..36c807f52 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 eb8896ff4..d3f60293e 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 index 000e6b35a..0fb423c77 100644 --- a/tst/unit/task_step/snapshot_computer_env_test.py +++ b/tst/unit/task_step/snapshot_computer_env_test.py @@ -1,7 +1,7 @@ -"""snapshot_computer_env: a computer env's files kept as a FileArtifactUniverse by their paths, 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.""" +"""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 @@ -10,7 +10,7 @@ import pytest from agentenv_protocol import uploaded_file_part -from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +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 @@ -32,16 +32,19 @@ class _Mail: id, version, type, environment_name = "mail", 1, "mcp_server", "mail" -async def _snapshot(monkeypatch, card: dict, env=None, **fields) -> FileArtifactUniverse: +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"] - return FileArtifactUniverse.get(snapshotted["id"], version=snapshotted["version"]) + (desktop,) = EnvironmentUniverseArtifact.get(snapshotted["id"], version=snapshotted["version"]).get_environment_artifacts() + assert desktop.environment_name == "desktop" + return desktop -def _kept(universe: FileArtifactUniverse) -> dict[str, bytes]: - return {path: file.load() for path, file in universe.get_file_artifacts().items()} +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(): @@ -82,10 +85,10 @@ async def test_each_file_returned_inline_is_kept_by_its_path(local_stores, monke sent = _serve(monkeypatch, {"/agentenv": lambda params: [_inline("Documents/notes.txt", b"hello"), _inline("Documents/b/c.bin", b"\x00\x01")]}) - universe = await _snapshot(monkeypatch, _computer(), file_paths=file_paths) + desktop = await _snapshot(monkeypatch, _computer(), file_paths=file_paths) - assert _kept(universe) == {"Documents/notes.txt": b"hello", "Documents/b/c.bin": b"\x00\x01"} - assert universe.get_file_artifacts()["Documents/b/c.bin"].filename == "c.bin" + 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)] @@ -105,12 +108,12 @@ def uploading(params): sent = _serve(monkeypatch, {"/agentenv": uploading}) - universe = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) + 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 universe.get_file_artifacts().items()} == { + 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()} @@ -120,16 +123,16 @@ async def test_a_computer_that_fails_under_the_grant_is_asked_again_without_one( sent = _serve(monkeypatch, {"/agentenv": lambda params: UNREACHABLE if "write_namespace" in params else [_inline("Desktop/left.txt", b"x")]}) - universe = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) + desktop = await _snapshot(monkeypatch, _computer(DATA_OBJECTS)) - assert _kept(universe) == {"Desktop/left.txt": b"x"} + 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: []}, "computer env desk has no files under Documents to snapshot"), + (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): diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py index f964221d4..d201c1114 100644 --- a/tst/unit/task_step/snapshot_env_without_gateway_test.py +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -1,7 +1,7 @@ """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, refused for snapshot_computer_env; and an env that fails under the upload -grant asked again without it.""" +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 @@ -105,13 +105,18 @@ async def test_the_env_itself_is_reached_when_its_card_has_no_children(local_sto @pytest.mark.asyncio -async def test_a_file_tree_is_refused_for_snapshot_computer_env_before_its_data_get(local_stores, monkeypatch): - sent = _serve(monkeypatch, {}) +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)) - with pytest.raises(RuntimeError, match="'desktop' returns a file tree from data/get, which one file can't keep; " - "snapshot_computer_env snapshots a computer env's files"): - await _snapshot(_computer()) - assert sent == [] + (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 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 a4757eed1..f102e4d75 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 6e70635b1..b118bc393 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -557,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: @@ -601,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" @@ -665,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") @@ -903,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)) @@ -925,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)) @@ -942,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" @@ -1038,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") @@ -1131,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"