agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/task_step/task_steps/run_docker_container.py

@@ -0,0 +1,380 @@
+"""Run a Docker container task step — builds a user-provided Dockerfile on a
+pre-deployed VM sandbox and runs the resulting container in the background."""
+
+from __future__ import annotations
+
+import logging
+import posixpath
+import shlex
+from typing import ClassVar, Optional
+from urllib.parse import urlparse
+
+from agent_env.task_step.context import TaskStepContext
+from agent_env.entity_refs import EntityRef
+from agent_env.task_step.task_step import TaskStep, TaskStepDependency
+from agent_env.utils.paths import validate_relative_filename
+
+logger = logging.getLogger(__name__)
+
+
+class RunDockerContainerTaskStep(TaskStep):
+    type: ClassVar[str] = "run_docker_container"
+    entity_refs = (
+        EntityRef.artifact(
+            "docker_context_artifact_id",
+            version_field="docker_context_artifact_version",
+            artifact_type="file_artifact_universe",
+        ),
+    )
+
+    def __init__(
+        self,
+        id: str,
+        version: Optional[int],
+        sandbox_name: str,
+        docker_context_artifact_id: Optional[str] = None,
+        docker_context_artifact_version: Optional[int] = None,
+        docker_context_url: Optional[str] = None,
+        dockerfile_path: str = "Dockerfile",
+        container_name: str = "task-container",
+        image_tag: Optional[str] = None,
+        ports: Optional[list[int]] = None,
+        env_vars: Optional[dict[str, str]] = None,
+        build_args: Optional[dict[str, str]] = None,
+        command_override: Optional[str] = None,
+        keep_alive_with_base_command: bool = False,
+        network: Optional[str] = None,
+        extra_networks: Optional[list[str]] = None,
+        ready_command: Optional[str] = None,
+        devices: Optional[list[str]] = None,
+        cap_add: Optional[list[str]] = None,
+        privileged: bool = False,
+        volumes: Optional[list[str]] = None,
+        shm_size: Optional[str] = None,
+        depends_on: Optional[list[TaskStepDependency]] = None,
+        fail_task_on_error: bool = True,
+    ):
+        super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        modes = sum(x is not None for x in (docker_context_artifact_id, docker_context_url))
+        if modes != 1:
+            raise ValueError(
+                "Must pass exactly one of `docker_context_artifact_id` or `docker_context_url`"
+            )
+        if command_override and keep_alive_with_base_command:
+            raise ValueError(
+                "Pass either `command_override` or `keep_alive_with_base_command`, not both"
+            )
+        self.sandbox_name = sandbox_name
+        self.docker_context_artifact_id = docker_context_artifact_id
+        self.docker_context_artifact_version = docker_context_artifact_version
+        self.docker_context_url = docker_context_url
+        self.dockerfile_path = dockerfile_path
+        self.container_name = container_name
+        self.image_tag = image_tag
+        self.ports = ports
+        self.env_vars = env_vars
+        self.build_args = build_args
+        self.command_override = command_override
+        self.keep_alive_with_base_command = keep_alive_with_base_command
+        self.network = network
+        self.extra_networks = list(extra_networks or [])
+        # host / none / container:<name> replace the container's whole netns, so they can't be attached as extra networks.
+        netns_extra = [n for n in self.extra_networks if n in ("host", "none") or n.startswith("container:")]
+        if netns_extra:
+            raise ValueError(f"extra_networks cannot contain netns-replacing modes {netns_extra}; use user-defined bridges")
+        self.ready_command = ready_command
+        self.devices = list(devices or [])
+        self.cap_add = list(cap_add or [])
+        self.privileged = privileged
+        self.volumes = list(volumes or [])
+        self.shm_size = shm_size
+
+    def to_dict(self) -> dict:
+        base = super().to_dict()
+        base["sandbox_name"] = self.sandbox_name
+        base["docker_context_artifact_id"] = self.docker_context_artifact_id
+        base["docker_context_artifact_version"] = self.docker_context_artifact_version
+        base["docker_context_url"] = self.docker_context_url
+        base["dockerfile_path"] = self.dockerfile_path
+        base["container_name"] = self.container_name
+        base["image_tag"] = self.image_tag
+        base["ports"] = self.ports
+        base["env_vars"] = self.env_vars
+        base["build_args"] = self.build_args
+        base["command_override"] = self.command_override
+        base["keep_alive_with_base_command"] = self.keep_alive_with_base_command
+        base["network"] = self.network
+        base["extra_networks"] = self.extra_networks
+        base["ready_command"] = self.ready_command
+        base["devices"] = self.devices
+        base["cap_add"] = self.cap_add
+        base["privileged"] = self.privileged
+        base["volumes"] = self.volumes
+        base["shm_size"] = self.shm_size
+        return base
+
+    @classmethod
+    def from_dict(cls, data: dict) -> RunDockerContainerTaskStep:
+        return cls(
+            **cls._base_from_dict(data),
+            sandbox_name=data["sandbox_name"],
+            docker_context_artifact_id=data.get("docker_context_artifact_id"),
+            docker_context_artifact_version=data.get("docker_context_artifact_version"),
+            docker_context_url=data.get("docker_context_url"),
+            dockerfile_path=data.get("dockerfile_path", "Dockerfile"),
+            container_name=data.get("container_name", "task-container"),
+            image_tag=data.get("image_tag"),
+            ports=data.get("ports"),
+            env_vars=data.get("env_vars"),
+            build_args=data.get("build_args"),
+            command_override=data.get("command_override"),
+            keep_alive_with_base_command=data.get("keep_alive_with_base_command", False),
+            network=data.get("network"),
+            extra_networks=data.get("extra_networks"),
+            ready_command=data.get("ready_command"),
+            devices=data.get("devices"),
+            cap_add=data.get("cap_add"),
+            privileged=data.get("privileged", False),
+            volumes=data.get("volumes"),
+            shm_size=data.get("shm_size"),
+        )
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        from agent_env.providers.sandbox_providers.sandbox_provider import (
+            SANDBOX_MODE_VM,
+            build_sandbox_provider,
+            get_sandbox_provider,
+        )
+
+        ds = next(
+            (s for s in context.deployed_sandboxes if s.sandbox_name == self.sandbox_name),
+            None,
+        )
+        if ds is None:
+            raise RuntimeError(
+                f"Sandbox '{self.sandbox_name}' not found in context.deployed_sandboxes; "
+                f"deploy it with DeploySandboxTaskStep before this step"
+            )
+        if ds.sandbox_mode != SANDBOX_MODE_VM:
+            raise RuntimeError(
+                f"Cannot run docker container on sandbox '{self.sandbox_name}' "
+                f"(mode={ds.sandbox_mode!r}); only VM-mode sandboxes support docker build/run"
+            )
+
+        provider = (
+            build_sandbox_provider(ds.sandbox_type) if ds.sandbox_type else get_sandbox_provider()
+        )
+        sandbox = await provider.get_sandbox(ds.sandbox_id)
+
+        image_tag = self.image_tag or f"{self.container_name}:latest"
+
+        work_dir = f"/tmp/docker-context-{self.container_name}"
+        await sandbox.exec_script(
+            f"rm -rf {shlex.quote(work_dir)} && mkdir -p {shlex.quote(work_dir)}"
+        )
+
+        if self.docker_context_artifact_id is not None:
+            await self._stage_from_universe(sandbox, ds.sandbox_id, work_dir)
+        else:
+            await self._stage_zip_from_url(sandbox, self.docker_context_url, work_dir)
+
+        dockerfile_abs = posixpath.join(work_dir, self.dockerfile_path)
+        await sandbox.exec_script(
+            f"test -f {shlex.quote(dockerfile_abs)} "
+            f"|| (echo 'Dockerfile not found at {self.dockerfile_path} inside the build context' >&2; exit 1)"
+        )
+
+        build_arg_flags = ""
+        if self.build_args:
+            build_arg_flags = " ".join(
+                f"--build-arg {shlex.quote(f'{k}={v}')}" for k, v in self.build_args.items()
+            )
+        build_cmd = (
+            f"cd {shlex.quote(work_dir)} && "
+            f"docker build --platform linux/amd64 "
+            f"-f {shlex.quote(self.dockerfile_path)} "
+            f"-t {shlex.quote(image_tag)} "
+            f"{build_arg_flags} ."
+        )
+        logger.info(
+            f"Building docker image '{image_tag}' on sandbox {ds.sandbox_id} "
+            f"(dockerfile={self.dockerfile_path})"
+        )
+        await sandbox.exec_script(build_cmd)
+
+        network_flag = ""
+        if self.network:
+            network_flag = f"--network {await self._ensure_network(sandbox, self.network)} "
+
+        port_flags = " ".join(f"-p {p}:{p}" for p in (self.ports or []))
+        env_flag_parts = []
+        if self.env_vars:
+            for k, v in self.env_vars.items():
+                escaped = v.replace("'", "'\\''")
+                env_flag_parts.append(f"-e {k}='{escaped}'")
+        env_flags = " ".join(env_flag_parts)
+        device_flags = " ".join(f"--device {shlex.quote(d)}" for d in self.devices)
+        cap_flags = " ".join(f"--cap-add {shlex.quote(c)}" for c in self.cap_add)
+        volume_flags = " ".join(f"-v {shlex.quote(v)}" for v in self.volumes)
+        extra_flags = " ".join(f for f in (
+            "--privileged" if self.privileged else "",
+            f"--shm-size {shlex.quote(self.shm_size)}" if self.shm_size else "",
+            device_flags, cap_flags, volume_flags,
+        ) if f)
+        entrypoint_flag = ""
+        if self.command_override:
+            command_tail = f" {self.command_override}"
+        elif self.keep_alive_with_base_command:
+            # Keep the container alive for install/exec WITHOUT discarding the
+            # image's own startup (a `sleep infinity` override would stop a base
+            # CMD that launches a DB/server the task needs). Replay the recorded
+            # ENTRYPOINT+CMD in the background, then hold PID 1 open.
+            base_cmd = await self._inspect_base_command(sandbox, image_tag)
+            if base_cmd:
+                base_str = " ".join(shlex.quote(part) for part in base_cmd)
+                script = f"{base_str} & exec sleep infinity"
+                entrypoint_flag = "--entrypoint /bin/sh "
+                command_tail = f" -c {shlex.quote(script)}"
+            else:
+                command_tail = " sleep infinity"
+        else:
+            command_tail = ""
+        run_cmd = (
+            f"docker run -d --name {shlex.quote(self.container_name)} "
+            f"{network_flag}{entrypoint_flag}{extra_flags} {port_flags} {env_flags} "
+            f"{shlex.quote(image_tag)}{command_tail}"
+        )
+        logger.info(
+            f"Starting container '{self.container_name}' from image '{image_tag}' "
+            f"on sandbox {ds.sandbox_id}"
+        )
+        await sandbox.exec_script(run_cmd)
+
+        # Dual-home the container onto additional networks (e.g. a jump host straddling segments).
+        # Dedup and skip the primary: `docker network connect` errors on an already-attached endpoint.
+        for extra in dict.fromkeys(self.extra_networks):
+            if extra == self.network:
+                continue
+            enet = await self._ensure_network(sandbox, extra)
+            await sandbox.exec_script(f"docker network connect {enet} {shlex.quote(self.container_name)}")
+            logger.info(f"Connected container '{self.container_name}' to extra network '{extra}'")
+
+        if self.ready_command:
+            ready = shlex.quote(self.ready_command)
+            logger.info(f"Waiting for container '{self.container_name}' readiness: {self.ready_command}")
+            await sandbox.exec_script(
+                f"for i in $(seq 1 60); do "
+                f"if docker exec {shlex.quote(self.container_name)} sh -c {ready} >/dev/null 2>&1; then exit 0; fi; "
+                f"sleep 2; done; "
+                f"echo 'container {self.container_name} failed readiness check within 120s' >&2; exit 1"
+            )
+
+        context.metadata.setdefault("deployed_docker_containers", []).append({
+            "container_name": self.container_name,
+            "sandbox_name": self.sandbox_name,
+            "sandbox_id": ds.sandbox_id,
+            "image_tag": image_tag,
+            "ports": self.ports,
+            "network": self.network,
+        })
+        return context
+
+    async def _ensure_network(self, sandbox, name: str) -> str:
+        net = shlex.quote(name)
+        # host / none / bridge / container:<name> are Docker built-in modes, not user bridges -- nothing to create.
+        if name in ("host", "none", "bridge") or name.startswith("container:"):
+            return net
+        await sandbox.exec_script(
+            f"docker network inspect {net} >/dev/null 2>&1 "    # already exists -> reuse it
+            f"|| docker network create {net} >/dev/null 2>&1 "  # else create the user bridge
+            f"|| docker network inspect {net} >/dev/null"       # lost a create race -> confirm it exists
+        )
+        return net
+
+    async def _inspect_base_command(self, sandbox, image_tag: str) -> list[str]:
+        """Return the built image's ENTRYPOINT+CMD as a single argv list (or [])."""
+        import json
+
+        out = await sandbox.exec_script(
+            "docker inspect --format "
+            "'{{json .Config.Entrypoint}}|{{json .Config.Cmd}}' "
+            f"{shlex.quote(image_tag)}"
+        )
+        ep_raw, _, cmd_raw = out.strip().partition("|")
+        entrypoint = json.loads(ep_raw) if ep_raw and ep_raw != "null" else []
+        cmd = json.loads(cmd_raw) if cmd_raw and cmd_raw != "null" else []
+        return list(entrypoint or []) + list(cmd or [])
+
+    async def _stage_from_universe(self, sandbox, sandbox_id: str, work_dir: str) -> None:
+        from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse
+
+        universe = FileArtifactUniverse.get(
+            self.docker_context_artifact_id, self.docker_context_artifact_version,
+        )
+        file_artifacts = universe.get_file_artifacts()
+        if not file_artifacts:
+            raise RuntimeError(
+                f"FileArtifactUniverse '{universe.id}' v{universe.version} has no files; "
+                f"nothing to use as docker build context"
+            )
+
+        dirs_to_make: set[str] = set()
+        loaded: dict[str, str] = {}
+        for filename in file_artifacts:
+            validate_relative_filename(filename)
+            dest_path = posixpath.join(work_dir, filename)
+            parent = posixpath.dirname(dest_path)
+            if parent and parent != work_dir:
+                dirs_to_make.add(parent)
+            loaded[filename] = dest_path
+        if dirs_to_make:
+            mkdir_cmd = " && ".join(f"mkdir -p {shlex.quote(d)}" for d in sorted(dirs_to_make))
+            await sandbox.exec_script(mkdir_cmd)
+
+        total = len(file_artifacts)
+        logger.info(
+            f"Loading docker context universe '{universe.id}' v{universe.version} "
+            f"({total} file(s)) onto sandbox {sandbox_id} at {work_dir}"
+        )
+        for idx, (filename, fa) in enumerate(file_artifacts.items(), 1):
+            dest_path = loaded[filename]
+            logger.info(f"  [{idx}/{total}] {fa.object_url} -> {dest_path}")
+            await sandbox.load_s3_file(fa.object_url, dest_path)
+
+    @staticmethod
+    async def _stage_zip_from_url(sandbox, url: str, work_dir: str) -> None:
+        """http(s) urls are fetched with curl; any other url is read through the configured
+        object store, which reaches whatever its backend can."""
+        parsed = urlparse(url)
+        if not parsed.scheme:
+            raise ValueError(
+                f"docker_context_url={url!r} must be an object store url or an http(s) url"
+            )
+        from_store = parsed.scheme not in ("http", "https")
+        path = url.rsplit("/", 1)[-1] if from_store else parsed.path
+        if path.endswith(".zip"):
+            pass
+        elif path.endswith((".tar.gz", ".tgz", ".tar.bz2", ".tar.xz", ".tar")):
+            raise NotImplementedError(
+                f"tar archives are not yet supported (URL: {url}); only .zip is supported today"
+            )
+        else:
+            raise ValueError(
+                f"Unsupported archive format for docker_context_url={url!r}; "
+                f"only .zip is supported"
+            )
+
+        archive_path = posixpath.join(work_dir, "_context.zip")
+        logger.info(f"Downloading docker context from {url} into {work_dir}")
+        if from_store:
+            await sandbox.load_s3_file(url, archive_path)
+        else:
+            await sandbox.exec_script(
+                f"curl -fsSL {shlex.quote(url)} -o {shlex.quote(archive_path)}"
+            )
+
+        await sandbox.exec_script(
+            f"unzip -q -o {shlex.quote(archive_path)} -d {shlex.quote(work_dir)} && "
+            f"rm -f {shlex.quote(archive_path)}"
+        )
+        logger.info(f"Extracted {url} into {work_dir}")

agent_env/task_step/task_steps/sandbox_utils/sandbox_utils.py

@@ -0,0 +1,79 @@
+"""Sandbox diagnostics shared by the steps that drive an A2A agent."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_sandbox_provider
+
+logger = logging.getLogger(__name__)
+
+# Room for the agent's own account of a failure (its protocol ``error_message``).
+_AGENT_ERROR_TEXT_LIMIT = 2000
+# Containers tailed when none is the agent's own, e.g. an agent installed into an env's containers.
+_FALLBACK_CONTAINER_LIMIT = 3
+# Core names the agent containers it starts ``agent-api``, ``agent-<id>`` or ``a2a-agent-*``.
+_AGENT_CONTAINER_MARKER = "agent"
+# Per stream, so a long stdout cannot push the stderr error out of the window.
+_STREAM_TAIL_CHARS = 500
+
+
+def agent_error_text(error_message: Optional[str], response_text: Optional[str]) -> str:
+    """The agent's own failure text, bounded: its ``error_message``, else its response."""
+    return (error_message or response_text or "")[:_AGENT_ERROR_TEXT_LIMIT]
+
+
+async def fetch_container_logs(agent, tail: int = 500) -> Optional[str]:
+    """Tail the docker logs of the agent's container, so a failed A2A task stays diagnosable
+    after its sandbox is reaped.
+
+    The agent's own container is the one named ``sandbox.container_name``. When none is, the
+    containers are tailed up to ``_FALLBACK_CONTAINER_LIMIT``, agent-named ones first. Best
+    effort: None on any failure, including a sandbox with no docker, so it never masks the
+    task's own failure.
+    """
+    sandbox_id = getattr(agent, "sandbox_id", None)
+    if not sandbox_id:
+        return None
+    try:
+        sandbox_type = getattr(agent, "sandbox_type", None)
+        provider = (
+            build_sandbox_provider(sandbox_type) if sandbox_type else get_sandbox_provider()
+        )
+        sandbox = await provider.get_sandbox(sandbox_id)
+        chunks: list[str] = []
+        try:
+            _, ps_out, _ = await sandbox.exec_with_output(
+                "sudo", "docker", "ps", "-a", "--format", "{{.ID}}|{{.Names}}",
+            )
+            rows = [
+                (cid.strip(), name.strip())
+                for cid, _, name in (line.partition("|") for line in (ps_out or "").splitlines())
+                if cid.strip()
+            ]
+            own = [cid for cid, name in rows if name == sandbox.container_name]
+            by_agent_name = sorted(rows, key=lambda row: _AGENT_CONTAINER_MARKER not in row[1])
+            cids = own or [cid for cid, _ in by_agent_name][:_FALLBACK_CONTAINER_LIMIT]
+            for cid in cids:
+                _, out, err = await sandbox.exec_with_output(
+                    "sudo", "docker", "logs", "--tail", str(tail), cid,
+                )
+                parts = []
+                out_tail = (out or "").strip()[-_STREAM_TAIL_CHARS:]
+                err_tail = (err or "").strip()[-_STREAM_TAIL_CHARS:]
+                if out_tail:
+                    parts.append(f"[stdout]\n{out_tail}")
+                if err_tail:
+                    parts.append(f"[stderr]\n{err_tail}")
+                if parts:
+                    chunks.append(f"[container {cid}]\n" + "\n".join(parts))
+        except Exception as e:
+            logger.debug("docker logs unavailable on sandbox %s: %s", sandbox_id, e)
+        return "\n\n".join(chunks) or None
+    except Exception as e:
+        logger.warning(
+            "Could not fetch agent container logs from sandbox %s: %s",
+            sandbox_id, e, exc_info=True,
+        )
+        return None

agent_env/task_step/task_steps/snapshot_agent_state.py

@@ -0,0 +1,205 @@
+"""Snapshot a deployed A2A agent's conversation state to a FileArtifactUniverse."""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import ClassVar, Optional
+
+from agent_env.task_step.context import TaskStepContext
+from agent_env.entity_refs import EntityRef, RefRole
+from agent_env.task_step.task_step import TaskStep, TaskStepDependency
+from agent_env.task_step.snapshot_utils import agent_state_capture
+
+logger = logging.getLogger(__name__)
+
+
+class SnapshotAgentStateTaskStep(TaskStep):
+    """Capture a deployed agent's conversation state and persist it as a
+    ``FileArtifactUniverse``.
+
+    Looks up the deployed agent (by ``agent_name``) and the source ``context_id``
+    via a previously-recorded ``PromptResponse`` (by ``prompt_id``). Calls the
+    agent's ``/ext/snapshot`` extension to write the conversation transcript to
+    S3, then wraps that S3 prefix as a ``FileArtifactUniverse`` via
+    ``put_existing`` (no download/re-upload). Appends an entry recording
+    ``{id, version, bundle_object_url, source_agent_name, source_context_id}`` to
+    ``context.metadata['agent_snapshots']`` (a list of such entries).
+
+    When ``env_id`` is set, the step ALSO dumps each MCP service's current
+    state by hitting ``{gateway_url}/svc/mcp-{service}/export-state`` on the
+    deployed env, uploads each as JSON alongside the workspace tarball, and
+    writes a JSON-stringified ``{service: s3_url}`` (unsigned ``s3://`` object
+    references) to ``context.metadata['snapshot_json_url']``. A downstream
+    verifier re-signs each value before handing the map to its test script, so
+    the script can fetch each service's state and assert against it. Without this, only the
+    agent's *workspace files* are visible to the verifier — not per-MCP-service
+    DB state.
+    """
+
+    type: ClassVar[str] = "snapshot_agent_state"
+    entity_refs = (
+        EntityRef.artifact("artifact_id", role=RefRole.OUTPUT, artifact_type="file_artifact_universe"),
+        EntityRef.env("env_id"),
+        EntityRef.artifact("universe_artifact_id", artifact_type="environment_universe"),
+    )
+    DEFAULT_TIMEOUT_SECONDS: ClassVar[int] = 120
+
+    def __init__(
+        self,
+        id: str,
+        version: Optional[int],
+        artifact_id: str,
+        prompt_id: str,
+        agent_name: Optional[str] = None,
+        env_id: Optional[str] = None,
+        universe_artifact_id: Optional[str] = None,
+        timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
+        depends_on: Optional[list[TaskStepDependency]] = None,
+        fail_task_on_error: bool = True,
+    ):
+        super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        self.artifact_id = artifact_id
+        self.prompt_id = prompt_id
+        self.agent_name = agent_name or TaskStep.DEFAULT_AGENT_NAME
+        # Set both `env_id` (which deployed env to query for /export-state) AND
+        # `universe_artifact_id` (to enumerate the env's service names) to
+        # enable per-MCP-service capture. Leaving either unset captures only
+        # the agent's workspace (legacy behaviour).
+        self.env_id = env_id
+        self.universe_artifact_id = universe_artifact_id
+        self.timeout_seconds = timeout_seconds
+
+    def to_dict(self) -> dict:
+        base = super().to_dict()
+        base["artifact_id"] = self.artifact_id
+        base["prompt_id"] = self.prompt_id
+        base["agent_name"] = self.agent_name
+        base["env_id"] = self.env_id
+        base["universe_artifact_id"] = self.universe_artifact_id
+        base["timeout_seconds"] = self.timeout_seconds
+        return base
+
+    @classmethod
+    def from_dict(cls, data: dict) -> SnapshotAgentStateTaskStep:
+        return cls(
+            **cls._base_from_dict(data),
+            artifact_id=data["artifact_id"],
+            prompt_id=data["prompt_id"],
+            agent_name=data.get("agent_name"),
+            env_id=data.get("env_id"),
+            universe_artifact_id=data.get("universe_artifact_id"),
+            timeout_seconds=data.get("timeout_seconds", cls.DEFAULT_TIMEOUT_SECONDS),
+        )
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        deployed = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None)
+        if deployed is None:
+            raise RuntimeError(f"No deployed agent named '{self.agent_name}' found in context")
+        a2a_url = deployed.a2a_url or deployed.api_url
+        if not a2a_url:
+            raise RuntimeError(f"Deployed agent '{self.agent_name}' has no a2a_url")
+
+        prompt = next((p for p in context.prompt_responses if p.prompt_id == self.prompt_id), None)
+        if prompt is None:
+            raise RuntimeError(f"No PromptResponse with prompt_id '{self.prompt_id}' in context")
+        a2a_context_id = prompt.a2a_context_id
+        if not a2a_context_id:
+            raise RuntimeError(
+                f"PromptResponse for '{self.prompt_id}' has no a2a_context_id "
+                "(prompt step did not capture one — make sure the agent is using a recent gateway)"
+            )
+
+        workspace = await agent_state_capture.capture_workspace(
+            a2a_url=a2a_url,
+            a2a_card=deployed.a2a_card or {},
+            agent_name=self.agent_name,
+            a2a_context_id=a2a_context_id,
+            artifact_id=self.artifact_id,
+            timeout_seconds=self.timeout_seconds,
+        )
+
+        snapshots = context.metadata.setdefault("agent_snapshots", [])
+        snapshots.append({
+            "id": workspace.universe_id,
+            "version": workspace.universe_version,
+            "bundle_object_url": workspace.bundle_object_url,
+            "source_agent_name": self.agent_name,
+            "source_context_id": a2a_context_id,
+        })
+
+        # Per-MCP-service state dump. Only runs when env_id +
+        # universe_artifact_id are both configured. Mirrors the export-state
+        # pattern in `multienv_validator/verify_universe_roundtrip.py:_export_all`.
+        # Result lands in `context.metadata['snapshot_json_url']` as a
+        # JSON-stringified `{service: s3_url}` (unsigned `s3://` references).
+        # A downstream verifier re-signs each value and hands the map to its test
+        # script, which accepts either a single URL or this map.
+        if self.env_id and self.universe_artifact_id:
+            await self._capture_universe_state(context, workspace.capture_prefix)
+
+        return context
+
+    async def _capture_universe_state(
+        self, context: TaskStepContext, capture_prefix: str
+    ) -> None:
+        """GET /svc/mcp-<name>/export-state for each service, upload, publish `s3://` refs."""
+        from agent_env.artifact import EnvironmentUniverseArtifact
+        from agent_env.config import get_config
+
+        deployed_env = next(
+            (d for d in context.deployed_envs if d.env_id == self.env_id), None
+        )
+        if deployed_env is None:
+            logger.warning(
+                "snapshot_agent_state: env_id=%s not found in deployed_envs (%s); "
+                "skipping per-service capture",
+                self.env_id,
+                [d.env_id for d in context.deployed_envs],
+            )
+            return
+
+        try:
+            universe = EnvironmentUniverseArtifact.get(self.universe_artifact_id)
+            environment_names = [sa.environment_name for sa in universe.get_environment_artifacts()]
+        except Exception as exc:
+            logger.warning(
+                "snapshot_agent_state: could not enumerate services from universe %s: %s",
+                self.universe_artifact_id, exc,
+            )
+            return
+
+        # Reuse the workspace capture's S3 prefix — services land in a
+        # `services/` subdir alongside the workspace tarball.
+        store = get_config().get_object_store()
+        key_prefix = store.get_object_key(capture_prefix).rstrip("/") + "/"
+
+        from agent_env.env import legacy_protocol
+        from agent_env.env.env import gateway_url_of
+        from agentenv_protocol import client as protocol_v1
+        urls: dict[str, str] = {}
+        for name in environment_names:
+            try:
+                base_url = await legacy_protocol.v1_base_url(deployed_env, gateway_url_of(deployed_env), name, mcp=True)
+                if base_url is not None:
+                    resp = await protocol_v1.get_data(base_url)
+                    state = resp.parts[0].data if resp.parts else {}
+                else:
+                    state = await legacy_protocol.export_state(gateway_url_of(deployed_env), name)
+                urls[name] = await asyncio.to_thread(
+                    store.put, f"{key_prefix}services/{name}.json", json.dumps(state).encode(),
+                    content_type="application/json", allow_overwrite=True,
+                )
+            except Exception as exc:
+                logger.warning(
+                    "snapshot_agent_state: capturing %s state failed: %s — skipping",
+                    name, exc,
+                )
+                continue
+
+        if urls:
+            context.metadata["snapshot_json_url"] = json.dumps(urls)
+            logger.info(
+                "snapshot_agent_state: captured %d service states → snapshot_json_url",
+                len(urls),
+            )