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 a42d1f80..31481ba6 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -89,6 +89,14 @@ # a step's `env_id` resolved to a CUA env before issuing cua_* controller calls. _CUA_ENV_METADATA_MARKER = "cua_vm_sandbox_id" +# cua_get_file returns a file's whole base64 body in one gateway response. That buffer-and-return +# can't carry a large file within one request — a ~500MB deliverable blew past the step's 600s +# controller timeout (and could exceed the gateway body limit) and was silently dropped. Files at +# or under this size take that single-shot path; larger ones stream in fixed byte ranges via +# 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 + # artifact_paths entries with a URL scheme are not VM files — they're gold/source # URLs that downstream consumers (evaluator, artifact previewer) fetch directly. # Collect skips them rather than joining them to base_path (which would mangle @@ -172,6 +180,26 @@ def _cua_bash_stdout(text: str) -> str: return payload.get("output") or "" +def _cua_bash_error(text: str) -> str: + """The `error`/stderr field of a `cua_bash` envelope, or "" if absent/unwrapped. + Used to tell a genuinely-missing file from a size-probe that failed for another + reason — the former is a lenient skip, the latter a dropped deliverable.""" + try: + payload = json.loads(text) + except (TypeError, ValueError): + return "" + if not isinstance(payload, dict): + return "" + return str(payload.get("error") or "") + + +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', + and treating that as absent would skip an existing deliverable.""" + return "No such file" in stderr + + def _remove(path: str) -> None: try: os.unlink(path) @@ -179,21 +207,19 @@ def _remove(path: str) -> None: pass -def _write_and_upload(store, content: bytes, artifact_id: str, version: int, object_name: str, content_type: str) -> str: - """Stage ``content`` in a temp file, upload it as a collected artifact and remove the file.""" - with tempfile.NamedTemporaryFile(delete=False, suffix=f"-{object_name.replace('/', '_')}") as tmp: - tmp.write(content) +def _upload_file(store, file_path: str, artifact_id: str, version: int, object_name: str, content_type: str) -> str: + """Upload an already-staged local file as a collected artifact and remove it.""" try: return store.put_object_file( artifact_type="collected_artifacts", id=artifact_id, version=version, object_name=object_name, - file_path=tmp.name, + file_path=file_path, content_type=content_type, ) finally: - _remove(tmp.name) + _remove(file_path) class CollectArtifactsTaskStep(TaskStep): @@ -453,12 +479,16 @@ async def _discover_container(self, sandbox) -> str: ) async def _get_file_size(self, sandbox, container: Optional[str], source_path: str) -> int: - """Get file size on the agent's filesystem. Returns -1 if the file doesn't exist. + """Get file size on the agent's filesystem. Returns -1 only if the file is genuinely + absent; a `stat` that fails for any other reason RAISES so a failed size check is a + hard failure rather than a silent 'missing' skip that drops an existing deliverable. VM-mode wraps in `docker exec `; container-mode runs `stat` directly.""" args = _exec_args(sandbox, container, ("stat", "-c", "%s", source_path)) - exit_code, stdout, _ = await sandbox.exec_with_output(*args) + exit_code, stdout, stderr = await sandbox.exec_with_output(*args) if exit_code != 0: - return -1 + if _looks_absent(stderr): + return -1 + raise RuntimeError(f"could not size {source_path} (stat exit {exit_code}): {stderr[:200]}") return int(stdout.strip()) async def _collect_file( @@ -555,12 +585,83 @@ 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 _controller_get_file(self, deployed_env, path: str) -> bytes: - """Fetch a file from the CUA VM as raw bytes (controller returns base64).""" + 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 + returning -1, so a failed size check becomes a hard failure instead of masquerading + as a missing file and silently dropping an existing deliverable. ``wc -c`` is + portable across the Ubuntu and macOS desktop VMs; ``stat``'s flags are not. A + transport exception from ``_controller_call`` propagates unchanged.""" + out = await self._controller_call( + deployed_env, "cua_bash", {"script": f"wc -c < {shlex.quote(path)}"} + ) + # A present file (even 0 bytes) prints its size to stdout; empty stdout means wc failed. + text = _cua_bash_stdout(out).strip() + if text: + try: + return int(text.split()[0]) + except (ValueError, IndexError): + raise RuntimeError(f"unexpected `wc -c` output sizing {path}: {text!r}") + err = _cua_bash_error(out) + if _looks_absent(err): + return -1 + raise RuntimeError(f"could not size {path} on the CUA VM (`wc -c` failed: {err or 'no output'})") + + async def _controller_get_file(self, deployed_env, path: str) -> str: + """Download a CUA VM file to a local temp file and return its path. The caller owns + the returned file (uploads then removes it). + + Raises ``FileNotFoundError`` when the file is absent (the caller treats that as a + lenient skip). Small files come back in one ``cua_get_file`` call; large ones stream + in byte ranges straight to disk, so neither the transfer nor this process ever holds + the whole file in memory.""" import base64 - b64 = await self._controller_call(deployed_env, "cua_get_file", {"path": path}) - return base64.b64decode(b64) + size = await self._controller_file_size(deployed_env, path) + if size < 0: + raise FileNotFoundError(f"file not found on CUA VM: {path}") + with tempfile.NamedTemporaryFile(delete=False, suffix=f"-{os.path.basename(path)}") as tmp: + local_path = tmp.name + try: + if size > _CUA_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}) + with open(local_path, "wb") as fh: + fh.write(base64.b64decode(b64)) + return local_path + except BaseException: + _remove(local_path) # nothing downstream owns it yet, so it's ours to clean up + raise + + async def _controller_stream_to_file(self, deployed_env, path: str, size: int, local_path: str) -> None: + """Stream ``path`` off the CUA VM into ``local_path`` in ``_CUA_CHUNK_BYTES`` ranges. + + ``tail -c +N | head -c M | base64`` is portable (no GNU ``dd iflag``) and seeks rather + than scans. Each range is decoded and written straight to the file so memory stays + bounded to one chunk. base64 may wrap, so whitespace is stripped before decoding. A + short read fails loudly rather than leaving a truncated artifact.""" + import base64 + + logger.info( + f"Streaming {path} ({size} bytes) off the CUA VM in {_CUA_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" + 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: + break # no forward progress — stop rather than loop forever + fh.write(chunk) + written += len(chunk) + if written != size: + raise RuntimeError( + f"streamed {written} of {size} bytes for {path} (chunked read stalled or " + f"truncated) — refusing to upload a partial artifact" + ) async def _controller_list_dir(self, deployed_env) -> list[str]: """Recursively enumerate regular files under base_path via cua_bash. @@ -625,13 +726,13 @@ async def _collect_via_controller(self, context, store, artifact_id, version): ) if self._configured_entries(context): logger.info(f"collect '{self.id}': every configured entry is a URL; nothing to read from the VM") - return {}, {} + return {}, {}, [] enumerated = self._drop_excluded(await self._controller_list_dir(deployed_env)) if not enumerated: logger.warning( f"No artifact filenames resolved and {self.base_path} is empty — nothing to collect" ) - return {}, {} + return {}, {}, [] items = [(f, f"{self.base_path}/{f}", f) for f in enumerated] logger.info( f"No manifest/artifacts_key/artifact_paths configured; enumerated " @@ -642,34 +743,45 @@ async def _collect_via_controller(self, context, store, artifact_id, version): collected: dict[str, str] = {} file_artifacts: dict[str, FileArtifact] = {} + hard_failures: list[str] = [] for key, source_path, object_name in items: try: - content = await self._controller_get_file(deployed_env, source_path) + local_path = await self._controller_get_file(deployed_env, source_path) + except FileNotFoundError: + # The file isn't on the VM — a listed-but-not-produced artifact. Stay lenient: + # only a file that EXISTS but we failed to move counts as a hard failure. + logger.warning(f"File not found on the VM, skipping: {source_path}") + continue except Exception as e: logger.warning(f"Failed to read {source_path} via controller: {e}") + hard_failures.append(source_path) continue ext = os.path.splitext(object_name)[1] content_type = _CONTENT_TYPES.get(ext, "application/octet-stream") + size = os.path.getsize(local_path) try: + # The upload owns the file from here: a cancel must not remove it under the upload. s3_url = await finish_on_thread( functools.partial( - _write_and_upload, store, content, artifact_id, version, object_name, content_type, + _upload_file, store, local_path, artifact_id, version, object_name, content_type, ), f"Uploading collected {source_path}", + if_never_run=functools.partial(_remove, local_path), ) collected[key] = s3_url self._register_file_artifact( store, file_artifacts, artifact_id, key, object_name, content_type, s3_url, context, ) - logger.info(f"Collected {source_path} -> {s3_url} ({len(content)} bytes) via controller") + logger.info(f"Collected {source_path} -> {s3_url} ({size} bytes) via controller") except Exception as e: logger.warning(f"Failed to collect {source_path} via controller: {e}") + hard_failures.append(source_path) logger.info(f"Collected {len(collected)}/{len(items)} artifacts via controller") - return collected, file_artifacts + return collected, file_artifacts, hard_failures async def _collect_via_agent_container(self, context, store, artifact_id, version): from agent_env.providers.sandbox_providers.sandbox_provider import ( @@ -803,7 +915,7 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto if self._configured_entries(context): logger.info(f"collect '{self.id}': every configured entry is a URL; nothing to read from the VM") await provider.close() - return {}, {} + return {}, {}, [] # Neither seed-driven nor static list was configured. Fall back # to enumerating base_path on the VM — useful when the agent's # output set is dynamic (e.g. "everything the solver wrote to @@ -814,7 +926,7 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto f"No artifact filenames resolved and {self.base_path} is empty — nothing to collect" ) await provider.close() - return {}, {} + return {}, {}, [] items = [(f, f"{self.base_path}/{f}", f) for f in enumerated] logger.info( f"No manifest/artifacts_key/artifact_paths configured; enumerated " @@ -825,7 +937,11 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto collected: dict[str, str] = {} file_artifacts: dict[str, FileArtifact] = {} + hard_failures: list[str] = [] for key, source_path, object_name in items: + # In the guarded block so a size-probe error is caught as a hard failure (not a + # silent skip) and provider.close() below still runs; only a genuinely-absent + # file (size -1) is the lenient skip. try: file_size = await self._get_file_size(sandbox, container, source_path) if file_size < 0: @@ -847,6 +963,7 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto except Exception as e: logger.warning(f"Failed to collect {source_path}: {e}") + hard_failures.append(source_path) if not collected and items: # Surface what the agent actually left in base_path so a re-collect is a pick, not a guess. @@ -857,7 +974,7 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto } logger.info(f"Collected {len(collected)}/{len(items)} artifacts") await provider.close() - return collected, file_artifacts + return collected, file_artifacts, hard_failures def _register_file_artifact(self, store, file_artifacts, artifact_id, key, object_name, content_type, s3_url, context): """Register a FileArtifact document pointing at the already-uploaded S3 @@ -902,13 +1019,13 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: version = int(time.time()) if self.env_id: - collected, file_artifacts = await self._collect_via_controller(context, store, artifact_id, version) + collected, file_artifacts, hard_failures = await self._collect_via_controller(context, store, artifact_id, version) elif self.container_name: - collected, file_artifacts = await self._collect_via_sandbox_container(context, store, artifact_id, version) + collected, file_artifacts, hard_failures = await self._collect_via_sandbox_container(context, store, artifact_id, version) elif self.sandbox_name: - collected, file_artifacts = await self._collect_via_vm_host(context, store, artifact_id, version) + collected, file_artifacts, hard_failures = await self._collect_via_vm_host(context, store, artifact_id, version) else: - collected, file_artifacts = await self._collect_via_agent_container(context, store, artifact_id, version) + collected, file_artifacts, hard_failures = await self._collect_via_agent_container(context, store, artifact_id, version) # `artifacts` and `file_artifact_universe` (set below) are legacy flat keys # that mirror the latest collect step; they can be removed once all @@ -935,6 +1052,14 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # After the metadata write, so tolerant (fail_task_on_error=False) dependents still see # present-but-empty. Enumeration fallback resolves to empty, so it stays lenient. + # A file that EXISTS on the VM but could not be collected (read or upload failure) fails + # the step even when others succeeded — a partial collect used to drop such files with only + # a warning, silently losing a deliverable. Absent files stay lenient (not in hard_failures). + if hard_failures: + raise RuntimeError( + f"collect: {len(hard_failures)} artifact(s) exist on the VM but could not be collected " + f"(see per-file warnings for the cause): {hard_failures}" + ) expected = self._resolve_items(context) if expected and collected_entry["file_artifact_universe"] is None: raise RuntimeError( diff --git a/tst/unit/task_step/test_collect_artifacts.py b/tst/unit/task_step/test_collect_artifacts.py index cc2e70fd..0af64b66 100644 --- a/tst/unit/task_step/test_collect_artifacts.py +++ b/tst/unit/task_step/test_collect_artifacts.py @@ -10,6 +10,7 @@ import asyncio import json +import os import tempfile import threading from unittest.mock import AsyncMock, MagicMock, patch @@ -29,6 +30,23 @@ def _run(coro): return asyncio.run(coro) +def _staged(data: bytes) -> str: + """A real temp file holding `data` — `_controller_get_file` now downloads to a temp file + and returns its path (not bytes), so stubs of it hand back a staged file the uploader reads.""" + with tempfile.NamedTemporaryFile(delete=False) as f: + f.write(data) + return f.name + + +def _read_staged(path: str) -> bytes: + """Read and delete a file `_controller_get_file` downloaded (it returns a path now).""" + try: + with open(path, "rb") as fh: + return fh.read() + finally: + os.unlink(path) + + def _cua_context(): return TaskStepContext(instance_id="inst-1", deployed_envs=[_cua_record()]) @@ -103,7 +121,7 @@ def test_reads_route_to_gateway_mcp_not_controller_sidecar(self): store.put_object_file.return_value = "s3://bucket/report.docx" store.next_version.return_value = 1 with patch.object(step, "_controller_get_file", new_callable=AsyncMock, - return_value=b"data") as get_file, \ + side_effect=lambda *a, **k: _staged(b"data")) as get_file, \ 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) @@ -121,7 +139,7 @@ def test_collects_via_controller_when_env_id_set(self): store.next_version.return_value = 1 with patch.object(step, "_controller_get_file", new_callable=AsyncMock, - return_value=b"%PDF-1.4 data") as get_file, \ + side_effect=lambda *a, **k: _staged(b"%PDF-1.4 data")) as get_file, \ 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) @@ -145,7 +163,7 @@ def test_the_controller_path_uploads_off_the_event_loop(self): store = MagicMock() store.put_object_file.side_effect = lambda **kw: on_loop.append(on_event_loop()) or "s3://bucket/report.pdf" store.next_version.return_value = 1 - with patch.object(step, "_controller_get_file", new_callable=AsyncMock, return_value=b"%PDF"), \ + with patch.object(step, "_controller_get_file", new_callable=AsyncMock, side_effect=lambda *a, **k: _staged(b"%PDF")), \ 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) @@ -183,7 +201,8 @@ async def run(): return await asyncio.sleep(0.02) - with patch.object(step, "_controller_get_file", new_callable=AsyncMock, return_value=b"%PDF-1.4"), \ + with patch.object(step, "_controller_get_file", new_callable=AsyncMock, + side_effect=lambda *a, **k: _staged(b"%PDF-1.4")), \ patch("agent_env.artifact.store.get_artifact_store", return_value=store), \ patch("agent_env.artifact.FileArtifactUniverse"), \ caplog.at_level("INFO", logger="agent_env.task_step.thread_work"): @@ -246,7 +265,7 @@ def test_absolute_paths_bypass_base_path(self): store.next_version.return_value = 1 with patch.object(step, "_controller_get_file", new_callable=AsyncMock, - return_value=b"%PDF data") as get_file, \ + side_effect=lambda *a, **k: _staged(b"%PDF data")) as get_file, \ 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) @@ -261,7 +280,9 @@ def test_absolute_paths_bypass_base_path(self): # S3 object name preserves dir structure with the leading slash stripped assert store.put_object_file.call_args.kwargs["object_name"] == "tmp/output/result.pdf" - def test_missing_file_is_skipped_not_fatal(self): + def test_absent_file_is_skipped_not_fatal(self): + """A listed file that isn't on the VM (FileNotFoundError) is skipped; a file that IS + present but fails to transfer is fatal — see test_present_file_that_fails_is_fatal.""" step = CollectArtifactsTaskStep( id="collect-artifacts", version=None, env_id="ubuntu-cua", base_path="/home/docker/Desktop", artifact_paths=["there.pdf", "gone.pdf"], @@ -272,8 +293,8 @@ def test_missing_file_is_skipped_not_fatal(self): async def fake_get(deployed_env, path): if path.endswith("gone.pdf"): - raise RuntimeError("cua_get_file failed: file not found") - return b"data" + raise FileNotFoundError(f"file not found on CUA VM: {path}") + return _staged(b"data") with patch.object(step, "_controller_get_file", side_effect=fake_get), \ patch("agent_env.artifact.store.get_artifact_store", return_value=store), \ @@ -283,7 +304,37 @@ async def fake_get(deployed_env, path): assert ctx.metadata["artifacts"] == {"there.pdf": "s3://bucket/there.pdf"} - def test_s3_upload_failure_is_skipped_not_fatal(self): + def test_present_file_that_fails_to_read_is_fatal(self): + """A file that EXISTS on the VM but can't be read (a non-FileNotFoundError, e.g. a + timeout/transport error — the ~500MB-deliverable failure mode) fails the step even + when other files collected, instead of being silently dropped.""" + step = CollectArtifactsTaskStep( + id="collect-artifacts", version=None, env_id="ubuntu-cua", + base_path="/home/docker/Desktop", artifact_paths=["ok.pdf", "big.mp4"], + ) + store = MagicMock() + store.put_object_file.return_value = "s3://bucket/ok.pdf" + store.next_version.return_value = 1 + + async def fake_get(deployed_env, path): + if path.endswith("big.mp4"): + raise RuntimeError("ReadTimeout") # present, but the read timed out + return _staged(b"data") + + with patch.object(step, "_controller_get_file", side_effect=fake_get), \ + 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) + try: + _run(step.execute(_cua_context())) + assert False, "expected RuntimeError when a present file fails to collect" + except RuntimeError as e: + assert "could not be collected" in str(e) + assert "big.mp4" in str(e) + + def test_s3_upload_failure_is_fatal(self): + """A present file we fetched but failed to upload is a hard failure (retryable), not a + silent drop — otherwise a flaky S3 write loses a deliverable.""" step = CollectArtifactsTaskStep( id="collect-artifacts", version=None, env_id="ubuntu-cua", base_path="/home/docker/Desktop", artifact_paths=["bad.pdf", "ok.pdf"], @@ -298,24 +349,27 @@ def put(**kw): store.put_object_file.side_effect = put - with patch.object(step, "_controller_get_file", new_callable=AsyncMock, return_value=b"data"), \ + with patch.object(step, "_controller_get_file", new_callable=AsyncMock, + side_effect=lambda *a, **k: _staged(b"data")), \ 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(_cua_context())) - - # bad.pdf's S3 failure is non-fatal; ok.pdf still collected - assert ctx.metadata["artifacts"] == {"ok.pdf": "s3://bucket/ok.pdf"} + try: + _run(step.execute(_cua_context())) + assert False, "expected RuntimeError when a present file fails to upload" + except RuntimeError as e: + assert "could not be collected" in str(e) + assert "bad.pdf" in str(e) - def test_all_resolved_missing_is_fatal(self): - """Named artifacts that ALL fail to collect -> loud failure, not a silent empty run.""" + def test_all_resolved_absent_is_fatal(self): + """Named artifacts that are ALL absent -> loud failure (no universe), not a silent empty run.""" step = CollectArtifactsTaskStep( id="collect-artifacts", version=None, env_id="ubuntu-cua", base_path="/home/docker/Desktop", artifact_paths=["gone1.pdf", "gone2.pdf"], ) store = MagicMock() with patch.object(step, "_controller_get_file", new_callable=AsyncMock, - side_effect=RuntimeError("cua_get_file failed: file not found")), \ + side_effect=FileNotFoundError("file not found on CUA VM")), \ patch("agent_env.artifact.store.get_artifact_store", return_value=store): try: _run(step.execute(_cua_context())) @@ -360,11 +414,11 @@ def test_url_only_list_skips_enumeration_on_the_container_path(self): provider.close = AsyncMock() ctx = TaskStepContext() with patch.object(step, "_list_base_directory", new_callable=AsyncMock) as list_dir: - collected, file_artifacts = _run(step._collect_items( + collected, file_artifacts, hard_failures = _run(step._collect_items( provider, MagicMock(), None, step._resolve_items(ctx), ctx, MagicMock(), "aid", 1, )) list_dir.assert_not_awaited() - assert (collected, file_artifacts) == ({}, {}) + assert (collected, file_artifacts, hard_failures) == ({}, {}, []) provider.close.assert_awaited_once() def test_collected_but_no_universe_is_fatal(self): @@ -377,7 +431,8 @@ def test_collected_but_no_universe_is_fatal(self): store.put_object_file.return_value = "s3://bucket/present.pdf" store.next_version.return_value = 1 store.put_document.side_effect = RuntimeError("doc write failed") # registration swallowed -> no universe - with patch.object(step, "_controller_get_file", new_callable=AsyncMock, return_value=b"data"), \ + with patch.object(step, "_controller_get_file", new_callable=AsyncMock, + side_effect=lambda *a, **k: _staged(b"data")), \ patch("agent_env.artifact.store.get_artifact_store", return_value=store): try: _run(step.execute(_cua_context())) @@ -399,7 +454,7 @@ def test_param_override_redirects_collection(self): "step_params": {"collect-artifacts": {"artifact_paths": ["right.pdf"]}} } with patch.object(step, "_controller_get_file", new_callable=AsyncMock, - return_value=b"data") as get_file, \ + side_effect=lambda *a, **k: _staged(b"data")) as get_file, \ 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) @@ -437,6 +492,29 @@ def test_non_cua_env_refuses_controller_calls(self): class TestAgentContainerPath: + def test_get_file_size_absent_returns_minus_one(self): + """A genuinely-missing file (stat says 'No such file') is the lenient -1 skip.""" + step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver") + sandbox = MagicMock() + sandbox.mode = "container" + sandbox.exec_with_output = AsyncMock( + return_value=(1, "", "stat: cannot stat '/x': No such file or directory") + ) + assert _run(step._get_file_size(sandbox, None, "/x")) == -1 + + def test_get_file_size_other_error_raises(self): + """A stat failure that is NOT 'missing' raises, so it becomes a hard failure rather than a + silent 'missing' skip that drops an existing deliverable.""" + step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver") + sandbox = MagicMock() + sandbox.mode = "container" + sandbox.exec_with_output = AsyncMock(return_value=(1, "", "stat: '/x': Permission denied")) + try: + _run(step._get_file_size(sandbox, None, "/x")) + assert False, "expected RuntimeError, not a -1 'absent' result" + except RuntimeError as e: + assert "could not size" in str(e) + def test_unreachable_sandbox_is_clear_error(self): """A re-run targeting an expired/unreachable sandbox fails with a clear message.""" step = CollectArtifactsTaskStep( @@ -709,10 +787,10 @@ def test_collects_through_the_named_container(self): provider.close = AsyncMock() with patch.object(step, "_resolve_live_sandbox", AsyncMock(return_value=sandbox)), \ - patch.object(step, "_collect_items", AsyncMock(return_value=({"page.html": "s3://x"}, {}))) as collect, \ + patch.object(step, "_collect_items", AsyncMock(return_value=({"page.html": "s3://x"}, {}, []))) as collect, \ patch("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider): - collected, _ = _run( + collected, _, _ = _run( step._collect_via_sandbox_container(self._context(), MagicMock(), "aid", 1) ) @@ -748,6 +826,180 @@ def test_a_card_without_step_raises_before_any_read(self, monkeypatch): assert sent == [] +class TestControllerChunkedRead: + """`_controller_get_file` reads small files in one cua_get_file call and streams large + ones in byte ranges via cua_bash — so no single controller call has to move a ~500MB + file (the deliverable that silently vanished under the old single-shot read).""" + + @staticmethod + def _fake_controller(data: bytes): + """Answer the three calls the reader makes against an in-VM file of `data`: + `wc -c` (size), `cua_get_file` (whole base64 body), and the `tail|head|base64` + ranged reads. Missing file -> empty `wc` output (size unknowable).""" + import base64 as b64 + import re as _re + + async def call(deployed_env, tool_name, arguments): + if tool_name == "cua_get_file": + return b64.b64encode(data).decode() + script = arguments["script"] + if "wc -c" in script: + out = "" if data is None else f"{len(data)}\n" + return json.dumps({"output": out, "error": "", "returncode": 0, "status": "success"}) + m = _re.search(r"tail -c \+(\d+) .*head -c (\d+)", script) + start, length = int(m.group(1)) - 1, int(m.group(2)) + # GNU base64 wraps at 76 cols; exercise that the reader strips it before decoding. + wrapped = "\n".join( + __import__("textwrap").wrap(b64.b64encode(data[start:start + length]).decode(), 76) + ) + return json.dumps({"output": wrapped, "error": "", "returncode": 0}) + + return call + + def test_small_file_uses_single_shot(self, monkeypatch): + import agent_env.task_step.task_steps.collect_artifacts as mod + monkeypatch.setattr(mod, "_CUA_SINGLE_SHOT_MAX", 16) + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + calls = [] + + async def call(deployed_env, tool_name, arguments): + calls.append(tool_name) + if "wc -c" in arguments.get("script", ""): + return json.dumps({"output": "5\n", "returncode": 0}) + import base64 as b64 + return b64.b64encode(b"hello").decode() + + with patch.object(step, "_controller_call", side_effect=call): + out = _run(step._controller_get_file(_cua_record(), "/x/small.bin")) + assert _read_staged(out) == b"hello" + assert "cua_get_file" in calls # single-shot, not chunked + assert not any("tail -c" in c for c in calls if isinstance(c, str)) + + def test_large_file_streams_and_reassembles(self, monkeypatch): + import agent_env.task_step.task_steps.collect_artifacts as mod + monkeypatch.setattr(mod, "_CUA_SINGLE_SHOT_MAX", 16) + monkeypatch.setattr(mod, "_CUA_CHUNK_BYTES", 7) # force several ranges over the payload + data = bytes(range(256)) * 5 # 1280 bytes, well over the 16-byte single-shot ceiling + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + with patch.object(step, "_controller_call", side_effect=self._fake_controller(data)): + out = _run(step._controller_get_file(_cua_record(), "/x/big.mp4")) + assert _read_staged(out) == data # every byte-range reassembled in order, base64 wrapping stripped + + def test_large_file_is_uploaded_from_a_path_not_held_in_memory(self, monkeypatch): + """End-to-end: a large file streams to a temp file and the uploader reads it from that + path (bytes intact), rather than the whole file being buffered in memory.""" + import agent_env.task_step.task_steps.collect_artifacts as mod + monkeypatch.setattr(mod, "_CUA_SINGLE_SHOT_MAX", 16) + monkeypatch.setattr(mod, "_CUA_CHUNK_BYTES", 8) + data = bytes(range(256)) * 4 # 1024 bytes, over the 16-byte single-shot ceiling + step = CollectArtifactsTaskStep( + id="collect-artifacts", version=None, env_id="ubuntu-cua", + base_path="/home/docker/Desktop", artifact_paths=["big.mp4"], + ) + uploaded = {} + store = MagicMock() + store.next_version.return_value = 1 + + def put(**kw): + with open(kw["file_path"], "rb") as fh: # the upload reads a real on-disk file + uploaded["bytes"] = fh.read() + return "s3://bucket/big.mp4" + + store.put_object_file.side_effect = put + with patch.object(step, "_controller_call", side_effect=self._fake_controller(data)), \ + 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(_cua_context())) + assert uploaded["bytes"] == data + assert ctx.metadata["artifacts"] == {"big.mp4": "s3://bucket/big.mp4"} + + def test_truncated_stream_is_fatal(self, monkeypatch): + """If the ranged reads stop short of the stat size, fail loudly rather than upload a partial file.""" + import agent_env.task_step.task_steps.collect_artifacts as mod + monkeypatch.setattr(mod, "_CUA_SINGLE_SHOT_MAX", 4) + monkeypatch.setattr(mod, "_CUA_CHUNK_BYTES", 4) + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + + async def call(deployed_env, tool_name, arguments): + script = arguments["script"] + if "wc -c" in script: + return json.dumps({"output": "100\n", "returncode": 0}) # claims 100 bytes + return json.dumps({"output": "", "returncode": 0}) # but every range reads empty + + with patch.object(step, "_controller_call", side_effect=call): + try: + _run(step._controller_get_file(_cua_record(), "/x/big.mp4")) + assert False, "expected RuntimeError on a short read" + except RuntimeError as e: + assert "refusing to upload a partial artifact" in str(e) + + def test_absent_file_raises_filenotfound(self): + """A missing file (wc errors with 'No such file') is reported absent so the caller + skips it leniently.""" + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + + async def call(deployed_env, tool_name, arguments): + return json.dumps({"output": "", "error": "wc: /x: No such file or directory", "returncode": 1}) + + with patch.object(step, "_controller_call", side_effect=call): + try: + _run(step._controller_get_file(_cua_record(), "/x/missing.bin")) + assert False, "expected FileNotFoundError for an absent file" + except FileNotFoundError: + pass + + def test_zero_byte_file_is_present_not_absent(self): + """`wc -c` prints 0 for an empty file; that's a present (size-0) file, not a missing one.""" + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + + async def call(deployed_env, tool_name, arguments): + if "wc -c" in arguments.get("script", ""): + return json.dumps({"output": "0\n", "error": "", "returncode": 0}) + return "" # cua_get_file for an empty file + + with patch.object(step, "_controller_call", side_effect=call): + assert _read_staged(_run(step._controller_get_file(_cua_record(), "/x/empty.bin"))) == b"" + + def test_size_probe_error_raises_not_absent(self): + """A `wc` failure that is NOT 'file missing' (e.g. a permission/IO error) raises — it must + not be mistaken for an absent file, which would silently drop an existing deliverable.""" + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + + async def call(deployed_env, tool_name, arguments): + return json.dumps({"output": "", "error": "wc: /x: Permission denied", "returncode": 1}) + + with patch.object(step, "_controller_call", side_effect=call): + try: + _run(step._controller_file_size(_cua_record(), "/x/locked.bin")) + assert False, "expected RuntimeError, not a -1 'absent' result" + except RuntimeError as e: + assert "could not size" in str(e) + except FileNotFoundError: + assert False, "a non-missing wc error must not look like an absent file" + + def test_size_probe_transport_error_is_fatal_not_a_missing_skip(self): + """If the controller times out running the size probe, the file may well exist — the step + fails (retryable) instead of skipping it as 'missing' and reporting success without it.""" + step = CollectArtifactsTaskStep( + id="collect-artifacts", version=None, env_id="ubuntu-cua", + base_path="/home/docker/Desktop", artifact_paths=["deliverable.mp4"], + ) + store = MagicMock() + + async def call(deployed_env, tool_name, arguments): + raise RuntimeError("ReadTimeout") # the wc -c size probe times out + + with patch.object(step, "_controller_call", side_effect=call), \ + patch("agent_env.artifact.store.get_artifact_store", return_value=store): + try: + _run(step.execute(_cua_context())) + assert False, "expected RuntimeError when the size probe errors" + except RuntimeError as e: + assert "could not be collected" in str(e) + assert "deliverable.mp4" in str(e) + + def _gateway_card(extensions: list = GATEWAY_EXTENSIONS) -> dict: """The gateway's own card, advertising `extensions`.""" return {"name": "gw", "capabilities": {"extensions": extensions}} diff --git a/tst/unit/task_step/test_collect_artifacts_vm_host.py b/tst/unit/task_step/test_collect_artifacts_vm_host.py index a8b771ab..ef6a6425 100644 --- a/tst/unit/task_step/test_collect_artifacts_vm_host.py +++ b/tst/unit/task_step/test_collect_artifacts_vm_host.py @@ -241,7 +241,7 @@ async def test_agent_target_still_discovers_its_container(self, monkeypatch): async def _fake_agent_collect(self, context, store, artifact_id, version): seen["called"] = True - return {}, {} + return {}, {}, [] monkeypatch.setattr( CollectArtifactsTaskStep, "_collect_via_agent_container", _fake_agent_collect