feat(snapshot): snapshot envs without a gateway, keep a file tree's files, and load a computer's snapshot back

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

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.

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,

src/agent_env/task_step/task_steps/computer/__init__.py

@@ -0,0 +1 @@
+"""Task steps for computer envs."""

src/agent_env/task_step/task_steps/computer/snapshot_computer_env.py

@@ -0,0 +1,107 @@
+"""snapshot_computer_env: a computer env's files, kept by their paths in its environment artifact, which load_artifact
+loads back into a computer."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import ClassVar, Optional
+
+from agentenv_protocol.transfers import RelativePath
+from pydantic import TypeAdapter, ValidationError
+
+from agent_env.entity.refs import EntityRef, RefRole
+from agent_env.env.env import Env
+from agent_env.env.envs.computer.env import ComputerEnv
+from agent_env.store.base import NotFoundError
+from agent_env.store.ids import is_local_id, validate_local_id
+from agent_env.task_step.context import TaskStepContext
+from agent_env.task_step.task_step import TaskStep, TaskStepDependency
+from agent_env.task_step.task_steps.snapshot_env import DEFAULT_EXPORT_TIMEOUT_SECONDS, SnapshotEnvTaskStep
+
+_PATHS = TypeAdapter(list[RelativePath])
+
+
+class SnapshotComputerEnvTaskStep(SnapshotEnvTaskStep):
+    """Snapshot a deployed computer env's files into a new environment universe, as ``snapshot_env`` snapshots an env,
+    each file under its path relative to the session user's home in the computer's environment artifact: those under
+    ``file_paths``, by default the folders its card declares. ``load_artifact`` loads the snapshot back into a
+    computer, through ``data/add``.
+
+    Output: ``context.metadata["env_snapshotted_universes"][self.id]`` = ``{"id", "version"}``."""
+
+    type: ClassVar[str] = "snapshot_computer_env"
+    entity_refs = (
+        EntityRef.env("env_id"),
+        EntityRef.artifact("snapshot_id", role=RefRole.OUTPUT, artifact_type="environment_universe"),
+    )
+
+    def __init__(
+        self,
+        id: str,
+        version: Optional[int],
+        env_id: Optional[str] = None,
+        env_instance_id: Optional[str] = None,
+        env_step_id: Optional[str] = None,
+        snapshot_id: Optional[str] = None,
+        file_paths: Optional[list[str]] = None,
+        export_timeout_seconds: int = DEFAULT_EXPORT_TIMEOUT_SECONDS,
+        depends_on: Optional[list[TaskStepDependency]] = None,
+        fail_task_on_error: bool = True,
+    ):
+        super().__init__(id, version, env_id=env_id, env_instance_id=env_instance_id, env_step_id=env_step_id,
+                         snapshot_id=snapshot_id, export_timeout_seconds=export_timeout_seconds,
+                         depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        self.file_paths = file_paths
+
+    def to_dict(self) -> dict:
+        return {**TaskStep.to_dict(self), "env_id": self.env_id, "env_instance_id": self.env_instance_id,
+                "env_step_id": self.env_step_id, "snapshot_id": self.snapshot_id, "file_paths": self.file_paths,
+                "export_timeout_seconds": self.export_timeout_seconds}
+
+    @classmethod
+    def from_dict(cls, data: dict) -> SnapshotComputerEnvTaskStep:
+        return cls(**cls._base_from_dict(data), env_id=data.get("env_id"), env_instance_id=data.get("env_instance_id"),
+                   env_step_id=data.get("env_step_id"), snapshot_id=data.get("snapshot_id"),
+                   file_paths=data.get("file_paths"),
+                   export_timeout_seconds=data.get("export_timeout_seconds", DEFAULT_EXPORT_TIMEOUT_SECONDS))
+
+    def config_problems(self) -> list[str]:
+        problems = super().config_problems()
+        if self.file_paths is not None:
+            try:
+                _PATHS.validate_python(self.file_paths)
+            except ValidationError:
+                problems.append(f"file_paths are normalized paths relative to the computer's session user's home, "
+                                f"not {self.file_paths!r}")
+        return problems
+
+    def preflight(self) -> list[str]:
+        """The env this step names when it isn't a computer env."""
+        problems = super().preflight()
+        if self.env_id is None:
+            return problems
+        try:
+            env = Env.get(self.env_id)
+        except NotFoundError:  # what the run reports
+            return problems
+        return problems if isinstance(env, ComputerEnv) else [*problems, _not_a_computer(env)]
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        snapshot_id = self._derive_snapshot_id(context)
+        if is_local_id(snapshot_id):
+            validate_local_id(snapshot_id)
+        deployed = self._resolve_deployed_env(context)
+        env = await asyncio.to_thread(Env.get, deployed.env_id, deployed.env_version)
+        if not isinstance(env, ComputerEnv):
+            raise RuntimeError(_not_a_computer(env))
+        result = await self.snapshot_env_state(
+            env=env, gateway_url=None, snapshot_id=snapshot_id, deployed=deployed,
+            export_timeout_seconds=self.export_timeout_seconds,
+            data_get_params=None if self.file_paths is None else {"paths": self.file_paths})
+        context.metadata.setdefault("env_snapshotted_universes", {})[self.id] = {
+            "id": result.environment_universe_artifact_id, "version": result.environment_universe_artifact_version}
+        return context
+
+
+def _not_a_computer(env: Env) -> str:
+    return f"snapshot_computer_env snapshots a computer env, and {env.id} is a {env.type} env"