feat(collect_artifacts): collect from framework computer envs

.github/scripts/check_plugin_api.py

@@ -57,6 +57,10 @@
     "agent_env.env.envs.website.WebsiteEnv.frontend_docker_image_artifact", "agent_env.env.envs.website.WebsiteEnv.environment_name",
     "agent_env.env.envs.multi_env.MultiEnv.mcp_server_envs", "agent_env.env.envs.multi_env.MultiEnv.website_envs",
     "agent_env.env.envs.multi_env.MultiEnv.name", "agent_env.artifact.artifacts.docker_image.DockerImageArtifact.image_name",
+    # What a plugin step runs and reads inside a deployed computer, as its session user, through its driver.
+    "agent_env.env.envs.computer.utils.inside.run", "agent_env.env.envs.computer.utils.inside.read_file",
+    "agent_env.env.envs.computer.utils.inside.max_timeout_s",
+    "agent_env.env.envs.computer.utils.inside.max_output_bytes",
 )
 SURFACE = BASES + USED
 

src/agent_env/env/envs/computer/utils/inside.py

@@ -46,6 +46,11 @@ def max_timeout_s(deployed: DeployedSandboxEnv) -> float:
     return protocol_v1.extension_params(deployed.environment_card, COMPUTER_EXEC)["max_timeout_s"]
 
 
+def max_output_bytes(deployed: DeployedSandboxEnv) -> int:
+    """The most bytes of stdout, and of stderr, one command returns whole, as the card's computer-exec/v1 says."""
+    return protocol_v1.extension_params(deployed.environment_card, COMPUTER_EXEC)["max_output_bytes"]
+
+
 def loopback(deployed: DeployedSandboxEnv) -> dict:
     """The env's id, address and MCP endpoint as code inside the computer reaches them, the driver on loopback, and
     the ``headers`` its framework-only routes take: code agent-env runs in the computer gets the computer's token."""

src/agent_env/task_step/task_steps/collect_artifacts.py

@@ -12,6 +12,9 @@
     Required for CUA tasks, where deliverables live on the desktop VM (e.g.
     /home/docker/Desktop) which is NOT mounted into the a2a `agent-api`
     container — so the docker-exec path reports "File not found".
+  - Computer (`env_id` names a `computer` env, one whose card serves computer-exec/v1
+    ``run``): the same reads, run inside the computer as its session user through
+    computer-exec, in ranges that fit the card's ``max_output_bytes``.
 
 Artifact paths come from one of these sources (in priority order):
   1. Manifest (`manifest_step_id`): a prior prompt step's declared artifacts from
@@ -77,6 +80,8 @@
 import time
 from typing import ClassVar, Optional
 
+from agentenv_protocol.computer import COMPUTER_EXEC
+
 from agent_env.config import get_config
 from agent_env.env.gateway.constants import EXT_STEP_URI
 from agent_env.store.ids import derive_path_id, is_local_id, validate_local_id
@@ -99,6 +104,8 @@
 # cua_bash so no single call has to move the whole file.
 _CUA_SINGLE_SHOT_MAX = 16 * 1024 * 1024  # 16 MB
 _CUA_CHUNK_BYTES = 8 * 1024 * 1024       # 8 MB per ranged read
+# How long one command inside a computer may run when collect reads a file through computer-exec.
+_COMPUTER_EXEC_TIMEOUT_S = 600
 
 # artifact_paths entries with a URL scheme are not VM files — they're gold/source
 # URLs that downstream consumers (evaluator, artifact previewer) fetch directly.
@@ -192,6 +199,15 @@ def _cua_bash_error(text: str) -> str:
     return str(payload.get("error") or "")
 
 
+def _is_computer(deployed_env) -> bool:
+    """Whether a deployed env is a computer collect reads through computer-exec/v1 ``run``, not a CUA desktop VM: no
+    CUA marker in its metadata, and a card that serves ``run``."""
+    if _CUA_ENV_METADATA_MARKER in (deployed_env.metadata or {}):
+        return False
+    supports = getattr(deployed_env, "supports", None)
+    return bool(supports and supports(COMPUTER_EXEC, "run"))
+
+
 def _looks_absent(stderr: str) -> bool:
     """Whether a stat/wc failure message means the file simply isn't there. Match the
     missing-file text specifically — a permission or I/O error also says 'cannot stat',
@@ -540,7 +556,10 @@ def upload() -> str:
                 _remove(local_path)
 
     async def _controller_call(self, deployed_env, tool_name: str, arguments: dict) -> str:
-        """Call a CUA MCP tool through the gateway's step/v1, as the env card declares it. Returns the text content."""
+        """Call a CUA MCP tool through the gateway's step/v1, as the env card declares it. Returns the text content.
+        On a computer env, ``cua_bash`` and ``cua_get_file`` run inside the computer instead (``_computer_call``)."""
+        if _is_computer(deployed_env):
+            return await self._computer_call(deployed_env, tool_name, arguments)
         data = await deployed_env.invoke(
             EXT_STEP_URI, "step", {"action": "call_tool", "tool_name": tool_name, "arguments": arguments}, timeout=600,
         )
@@ -551,6 +570,41 @@ async def _controller_call(self, deployed_env, tool_name: str, arguments: dict)
             raise RuntimeError(f"{tool_name} returned no content: {data}")
         return content[0].get("text", "")
 
+    async def _computer_call(self, deployed_env, tool_name: str, arguments: dict) -> str:
+        """``cua_bash`` or ``cua_get_file`` answered as the CUA MCP server answers them, by a command run inside a
+        computer through computer-exec/v1 ``run``: ``cua_bash`` as its JSON envelope, ``cua_get_file`` as the file's
+        bare base64. Output computer-exec cut short raises rather than handing back a partial file or listing."""
+        from agent_env.env.envs.computer.utils.inside import run
+
+        if tool_name == "cua_bash":
+            script = arguments["script"]
+        elif tool_name == "cua_get_file":
+            script = f"base64 < {shlex.quote(arguments['path'])} | tr -d '\\n'"
+        else:
+            raise RuntimeError(f"env '{self.env_id}' is a computer: collect has no {tool_name} there")
+        answer = await run(deployed_env, ["/bin/sh", "-c", script], timeout_s=_COMPUTER_EXEC_TIMEOUT_S)
+        if answer.get("truncated") or answer.get("timed_out"):
+            raise RuntimeError(f"{tool_name} on computer env '{self.env_id}' was cut short "
+                               f"({'timed out' if answer.get('timed_out') else 'output truncated'})")
+        if tool_name == "cua_get_file":
+            if answer.get("exit_code"):
+                raise RuntimeError(f"cua_get_file failed on computer env '{self.env_id}': "
+                                   f"{answer.get('stderr', '')[:300]}")
+            return answer.get("stdout", "")
+        return json.dumps({"output": answer.get("stdout", ""), "error": answer.get("stderr", ""),
+                           "returncode": answer.get("exit_code")})
+
+    def _read_limits(self, deployed_env) -> tuple[int, int]:
+        """The largest file read in one call, and the bytes each ranged read moves: on a computer, both fit the base64
+        of a range in the output its card's computer-exec/v1 returns whole."""
+        if not _is_computer(deployed_env):
+            return _CUA_SINGLE_SHOT_MAX, _CUA_CHUNK_BYTES
+        from agent_env.env.envs.computer.utils.inside import max_output_bytes
+
+        # base64 is 4 bytes out per 3 in, and GNU base64 wraps a ranged read's lines at 76 columns.
+        chunk = max(3, (max_output_bytes(deployed_env) * 3 // 4) * 76 // 77 // 3 * 3)
+        return chunk, chunk
+
     async def _controller_file_size(self, deployed_env, path: str) -> int:
         """Size of ``path`` on the CUA VM in bytes, or -1 only if the file is genuinely
         absent. A controller/transport error (e.g. a ``wc`` timeout) RAISES rather than
@@ -589,7 +643,8 @@ async def _controller_get_file(self, deployed_env, path: str) -> str:
         with tempfile.NamedTemporaryFile(delete=False, suffix=f"-{os.path.basename(path)}") as tmp:
             local_path = tmp.name
         try:
-            if size > _CUA_SINGLE_SHOT_MAX:
+            single_shot_max, _ = self._read_limits(deployed_env)
+            if size > single_shot_max:
                 await self._controller_stream_to_file(deployed_env, path, size, local_path)
             else:
                 b64 = await self._controller_call(deployed_env, "cua_get_file", {"path": path})
@@ -609,14 +664,15 @@ async def _controller_stream_to_file(self, deployed_env, path: str, size: int, l
         short read fails loudly rather than leaving a truncated artifact."""
         import base64
 
+        _, chunk_bytes = self._read_limits(deployed_env)
         logger.info(
-            f"Streaming {path} ({size} bytes) off the CUA VM in {_CUA_CHUNK_BYTES}-byte chunks"
+            f"Streaming {path} ({size} bytes) off the CUA VM in {chunk_bytes}-byte chunks"
         )
         quoted = shlex.quote(path)
         written = 0
         with open(local_path, "wb") as fh:
             while written < size:
-                script = f"tail -c +{written + 1} {quoted} | head -c {_CUA_CHUNK_BYTES} | base64"
+                script = f"tail -c +{written + 1} {quoted} | head -c {chunk_bytes} | base64"
                 out = await self._controller_call(deployed_env, "cua_bash", {"script": script})
                 chunk = base64.b64decode(_cua_bash_stdout(out).encode().translate(None, b"\r\n\t "))
                 if not chunk:
@@ -633,7 +689,8 @@ async def _controller_list_dir(self, deployed_env) -> list[str]:
         """Recursively enumerate regular files under base_path via cua_bash.
         Returns paths relative to base_path. Fallback when no explicit list
         was configured."""
-        script = f"find {shlex.quote(self.base_path)} -type f -printf '%P\\n'"
+        # Portable across GNU and BSD find (macOS computers): BSD find has no -printf.
+        script = f"cd {shlex.quote(self.base_path)} && find . -type f | sed 's|^\\./||'"
         try:
             out = await self._controller_call(deployed_env, "cua_bash", {"script": script})
         except Exception as e:
@@ -642,7 +699,7 @@ async def _controller_list_dir(self, deployed_env) -> list[str]:
         return [f.strip() for f in _cua_bash_stdout(out).splitlines() if f.strip()]
 
     def _resolve_cua_env(self, context):
-        """Resolve self.env_id to a deployed CUA env, refusing non-CUA envs.
+        """Resolve self.env_id to a deployed CUA env or computer env, refusing any other env.
 
         cua_get_file / cua_bash are CUA-controller tools; firing them at an
         arbitrary env's gateway would hang or error confusingly. A desktop-VM env's
@@ -653,12 +710,15 @@ def _resolve_cua_env(self, context):
             raise RuntimeError(
                 f"Env '{self.env_id}' not found in context.deployed_envs — cannot collect artifacts via controller"
             )
+        if _CUA_ENV_METADATA_MARKER not in (deployed_env.metadata or {}) and _is_computer(deployed_env):
+            return deployed_env
         if _CUA_ENV_METADATA_MARKER not in (deployed_env.metadata or {}):
             raise RuntimeError(
                 f"Env '{self.env_id}' is not a CUA environment "
                 f"(no '{_CUA_ENV_METADATA_MARKER}' in its deployed metadata) — refusing to issue "
                 f"cua_* controller calls against a non-CUA gateway. collect_artifacts with `env_id` "
-                f"set is only valid for CUA envs; omit `env_id` to use the agent-container path."
+                f"set is only valid for CUA envs and computer envs serving {COMPUTER_EXEC} run; omit `env_id` to use "
+                f"the agent-container path."
             )
         # Up front: enumeration swallows errors into [], so a missing step/v1 would become a silent empty collect.
         deployed_env.require(EXT_STEP_URI, "step")
@@ -681,7 +741,8 @@ async def _collect_via_controller(self, context, store, artifact_id, version):
         # gateway's step/v1. These are MCP tools registered on the gateway on both
         # Ubuntu and macOS — NOT endpoints on the macOS controller sidecar, whose
         # /step only accepts a computer-use action and 500s on a call_tool payload.
-        logger.info(f"Collecting artifacts via CUA MCP server (env_id={self.env_id}, gateway={deployed_env.environment_url})")
+        via = "computer-exec" if _is_computer(deployed_env) else "CUA MCP server"
+        logger.info(f"Collecting artifacts via {via} (env_id={self.env_id}, gateway={deployed_env.environment_url})")
 
         items = self._resolve_items(context)
         if not items:

tst/unit/task_step/test_collect_artifacts.py

@@ -1047,3 +1047,112 @@ async def exec_with_output(self, *args):
 
     monkeypatch.setattr(LocalSandbox, "exec_with_output", exec_with_output)
     return _run(LocalSandboxProvider().get_sandbox("local-agent1")), calls
+
+
+class _Computer:
+    """A deployed `computer` env as collect sees it: no CUA marker, a card serving computer-exec/v1 ``run`` with an
+    output cap, and ``run`` answered by running the command for real, so the shell collect sends is exercised."""
+
+    def __init__(self, max_output_bytes: int = 1 << 20):
+        from agentenv_protocol.computer import COMPUTER_EXEC
+
+        self.env_id = "desk"
+        self.environment_url = "http://desk"
+        self.metadata = {"computer": {"port": 8080}}
+        self.environment_card = {"capabilities": {"extensions": [
+            {"uri": COMPUTER_EXEC, "params": {"max_output_bytes": max_output_bytes, "max_timeout_s": 3600}}]}}
+        self.max_output_bytes = max_output_bytes
+        self.calls: list[list[str]] = []
+        self._uri = COMPUTER_EXEC
+
+    def supports(self, uri, method):
+        return uri == self._uri and method == "run"
+
+    async def run(self, deployed, argv, *, timeout_s, env=None):
+        import subprocess
+
+        self.calls.append(argv)
+        done = subprocess.run(argv, capture_output=True, timeout=timeout_s)
+        cap = self.max_output_bytes
+        return {"exit_code": done.returncode, "stdout": done.stdout.decode()[:cap], "stderr": done.stderr.decode()[:cap],
+                "truncated": len(done.stdout) > cap or len(done.stderr) > cap, "timed_out": False}
+
+
+def _computer_context(computer):
+    ctx = TaskStepContext(instance_id="inst-1")
+    ctx.deployed_envs = [computer]
+    return ctx
+
+
+class TestComputerPath:
+    """`env_id` naming a framework `computer` env: collect reads inside the computer through computer-exec/v1 ``run``,
+    in ranges whose base64 fits the card's ``max_output_bytes``, instead of refusing it as a non-CUA env."""
+
+    def _collect(self, computer, tmp_path, entries):
+        step = CollectArtifactsTaskStep(id="collect", version=None, env_id="desk", base_path=str(tmp_path),
+                                        artifact_paths=entries)
+        uploaded = {}
+
+        def put(local_path, *a, **k):
+            with open(local_path, "rb") as fh:
+                uploaded[os.path.basename(local_path)] = fh.read()
+            return "s3://bucket/x"
+
+        store = MagicMock()
+        store.put_object_file.side_effect = lambda **k: put(k.get("file_path") or k.get("local_path"))
+        with patch("agent_env.env.envs.computer.utils.inside.run", new=computer.run), \
+             patch("agent_env.task_step.task_steps.collect_artifacts._upload_file",
+                   side_effect=lambda store, file_path, *a: put(file_path)), \
+             patch("agent_env.artifact.store.get_artifact_store", return_value=store), \
+             patch("agent_env.artifact.FileArtifactUniverse") as universe:
+            universe.put.return_value = MagicMock(id="inst-1", version=1)
+            ctx = _run(step.execute(_computer_context(computer)))
+        return ctx, uploaded
+
+    def test_a_small_file_is_read_in_one_call(self, tmp_path):
+        (tmp_path / "report.txt").write_bytes(b"hello computer")
+        computer = _Computer()
+        ctx, uploaded = self._collect(computer, tmp_path, ["report.txt"])
+        assert list(uploaded.values()) == [b"hello computer"]
+        assert ctx.metadata["artifacts"] == {"report.txt": "s3://bucket/x"}
+        assert all(argv[:2] == ["/bin/sh", "-c"] for argv in computer.calls)
+        assert not any("tail -c" in argv[2] for argv in computer.calls)
+
+    def test_a_file_past_the_output_cap_streams_in_ranges(self, tmp_path):
+        data = os.urandom(5000)
+        (tmp_path / "big.bin").write_bytes(data)
+        computer = _Computer(max_output_bytes=1024)
+        _, uploaded = self._collect(computer, tmp_path, ["big.bin"])
+        assert list(uploaded.values()) == [data]
+        ranged = [argv for argv in computer.calls if "tail -c" in argv[2]]
+        assert len(ranged) > 1  # several ranges, each under the cap
+        assert not any(argv for argv in computer.calls if argv[2].startswith("base64 <"))
+
+    def test_an_absent_file_is_skipped(self, tmp_path):
+        (tmp_path / "here.txt").write_bytes(b"x")
+        _, uploaded = self._collect(_Computer(), tmp_path, ["here.txt", "gone.txt"])
+        assert list(uploaded.values()) == [b"x"]
+
+    def test_no_list_enumerates_the_base_path(self, tmp_path):
+        (tmp_path / "a.txt").write_bytes(b"a")
+        (tmp_path / "sub").mkdir()
+        (tmp_path / "sub" / "b.txt").write_bytes(b"b")
+        ctx, uploaded = self._collect(_Computer(), tmp_path, None)
+        assert sorted(uploaded.values()) == [b"a", b"b"]
+        assert sorted(ctx.metadata["artifacts"]) == ["a.txt", "sub/b.txt"]
+
+    def test_truncated_output_is_an_error_not_a_partial_file(self):
+        computer = _Computer()
+
+        async def truncated(deployed, argv, *, timeout_s, env=None):
+            return {"exit_code": 0, "stdout": "aGk=", "stderr": "", "truncated": True, "timed_out": False}
+
+        step = CollectArtifactsTaskStep(id="collect", version=None, env_id="desk")
+        with patch("agent_env.env.envs.computer.utils.inside.run", new=truncated):
+            with pytest.raises(RuntimeError, match="cut short"):
+                _run(step._controller_call(computer, "cua_get_file", {"path": "/x"}))
+
+    def test_a_cua_env_never_takes_the_computer_path(self):
+        cua = _cua_record()
+        step = CollectArtifactsTaskStep(id="collect", version=None, env_id="ubuntu-cua")
+        assert step._read_limits(cua) == (16 * 1024 * 1024, 8 * 1024 * 1024)