fix(collect): stream large CUA artifacts and fail on dropped deliverables

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,28 +180,46 @@ 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)
     except FileNotFoundError:
         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>`; 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(