diff --git a/.github/scripts/check_plugin_api.py b/.github/scripts/check_plugin_api.py index 464f49c85..e9bee3b38 100644 --- a/.github/scripts/check_plugin_api.py +++ b/.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 diff --git a/src/agent_env/env/envs/computer/utils/inside.py b/src/agent_env/env/envs/computer/utils/inside.py index c03ccecaa..e15e8e9fa 100644 --- a/src/agent_env/env/envs/computer/utils/inside.py +++ b/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.""" diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index c55923a6a..1f133edd7 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/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: diff --git a/tst/unit/task_step/test_collect_artifacts.py b/tst/unit/task_step/test_collect_artifacts.py index 31bc1ad76..6cc38479a 100644 --- a/tst/unit/task_step/test_collect_artifacts.py +++ b/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)