feat(load_artifact): allow urls entries to set the saved filename

src/agent_env/task_step/task_steps/load_artifact.py

@@ -28,13 +28,90 @@
     ("env_snapshotted_universes", lambda entry: entry),
 )
 
+# A `urls` entry: a bare URL, saved under the name its path ends in, or
+# {"url": ..., "filename": ...}, saved under that filename verbatim.
+UrlEntry = str | dict[str, str]
+_URL_ENTRY_KEYS = frozenset({"url", "filename"})
+
+
+def _named_download(step_id: str, entry: dict) -> tuple[str, str]:
+    """The ``(url, filename)`` of an object entry, rejecting anything but a plain file name for ``filename``."""
+    unknown = sorted(set(entry) - _URL_ENTRY_KEYS)
+    if unknown:
+        raise ValueError(
+            f"load_artifact '{step_id}': urls entry {entry!r} has unknown key(s) {unknown}; "
+            f"an object entry takes only 'url' and 'filename'"
+        )
+    url, filename = entry.get("url"), entry.get("filename")
+    if not isinstance(url, str) or not url:
+        raise ValueError(f"load_artifact '{step_id}': urls entry {entry!r} needs a non-empty string 'url'")
+    if not isinstance(filename, str) or not filename:
+        raise ValueError(
+            f"load_artifact '{step_id}': urls entry {entry!r} needs a non-empty string 'filename' "
+            f"(a bare URL string is saved under the name its path ends in)"
+        )
+    try:
+        validate_relative_filename(filename)
+    except ValueError as e:
+        raise ValueError(f"load_artifact '{step_id}': urls entry {entry!r}: {e}") from e
+    if "/" in filename or "\\" in filename or filename == ".":
+        raise ValueError(
+            f"load_artifact '{step_id}': urls entry {entry!r}: filename must be a plain file name, not a path"
+        )
+    return url, filename
+
+
+def _url_downloads(step_id: str, urls: list[UrlEntry]) -> list[tuple[str, str]]:
+    """Each ``urls`` entry as ``(url, filename)``, the name it is saved under in the destination.
+
+    A bare URL takes its last path segment, a repeat suffixed ``-1``, ``-2``, ...; an explicit
+    filename is never renamed, so one given twice, or matching a bare URL's name, raises.
+    """
+    explicit: dict[str, str] = {}
+    for entry in urls:
+        if isinstance(entry, dict):
+            url, filename = _named_download(step_id, entry)
+            if filename in explicit:
+                raise ValueError(
+                    f"load_artifact '{step_id}': filename {filename!r} is given to more than one urls entry "
+                    f"({explicit[filename]!r} and {url!r})"
+                )
+            explicit[filename] = url
+        elif not isinstance(entry, str):
+            raise ValueError(
+                f"load_artifact '{step_id}': urls entry {entry!r} must be a URL string or an object "
+                f"with 'url' and 'filename'"
+            )
+
+    downloads: list[tuple[str, str]] = []
+    seen: dict[str, int] = {}
+    for entry in urls:
+        if isinstance(entry, dict):
+            downloads.append((entry["url"], entry["filename"]))
+            continue
+        base = unquote(os.path.basename(urlparse(entry).path)) or "downloaded"
+        count = seen.get(base, 0)
+        seen[base] = count + 1
+        if count == 0:
+            filename = base
+        else:
+            stem, ext = os.path.splitext(base)
+            filename = f"{stem}-{count}{ext}"
+        if filename in explicit:
+            raise ValueError(
+                f"load_artifact '{step_id}': {entry!r} would be saved as {filename!r}, the filename given to "
+                f"{explicit[filename]!r}; give it a 'filename' of its own"
+            )
+        downloads.append((entry, filename))
+    return downloads
+
 
 @dataclass
 class _LoadInputs:
     """What a load_artifact step loads and where, after per-run overrides."""
 
     artifacts: list[dict]
-    urls: list[str]
+    downloads: list[tuple[str, str]]
     destination_path: Optional[str]
 
 
@@ -55,7 +132,7 @@ def __init__(
         artifacts: Optional[list[dict]] = None,
         artifact_id: Optional[str] = None,
         artifact_version: Optional[int] = None,
-        urls: Optional[list[str]] = None,
+        urls: Optional[list[UrlEntry]] = None,
         collected_artifacts_step_id: Optional[str] = None,
         artifact_from_step_id: Optional[str] = None,
         destination_path: Optional[str] = None,
@@ -107,6 +184,8 @@ def __init__(
         self.env_step_id = env_step_id
         self.artifacts = artifacts or []
         self.urls = list(urls or [])
+        # Stored urls fail here, at construction; a per-run `urls` override is checked in `_resolve_inputs`.
+        _url_downloads(self.id, self.urls)
         # Superseded by the producer-agnostic `artifact_from_step_id`, but kept
         # readable/writable: it is persisted on existing task-step documents.
         self.collected_artifacts_step_id = collected_artifacts_step_id
@@ -223,10 +302,11 @@ def _deployed_env(self, context: TaskStepContext):
         return matches[0] if matches else None
 
     def _resolve_inputs(self, context: TaskStepContext) -> _LoadInputs:
-        """Resolve ``(artifacts, urls, destination_path)``, honoring per-run
+        """Resolve ``(artifacts, downloads, destination_path)``, honoring per-run
         ``step_overrides``. An override wins over the stored value; an explicit
         artifact override also wins over ``artifact_from_step_id`` /
-        ``collected_artifacts_step_id`` wiring.
+        ``collected_artifacts_step_id`` wiring. A ``urls`` override is named and
+        validated here, before anything is loaded.
         Topology (env_id/agent_name/container_name) is not overridable here.
         """
         overrides = self.step_param_overrides(context)
@@ -257,11 +337,12 @@ def _resolve_inputs(self, context: TaskStepContext) -> _LoadInputs:
 
         # `or []` so a `{"urls": null}` override cleanly clears rather than crashing on list(None).
         urls = list(overrides["urls"] or []) if "urls" in overrides else list(self.urls)
+        downloads = _url_downloads(self.id, urls)
         destination_path = overrides.get("destination_path", self.destination_path)
 
         if overrides:
             logger.info("load '%s': applying step_overrides for %s", self.id, sorted(overrides))
-        return _LoadInputs(artifacts=resolved_artifacts, urls=urls, destination_path=destination_path)
+        return _LoadInputs(artifacts=resolved_artifacts, downloads=downloads, destination_path=destination_path)
 
     async def _load_url_onto_vm(self, sandbox, url: str, destination_path: str) -> None:
         """Host counterpart of ``sandbox.load_object_file``; ``write_file_from_url`` targets a container instead."""
@@ -302,7 +383,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
 
         inputs = self._resolve_inputs(context)
         resolved_artifacts = inputs.artifacts
-        urls = inputs.urls
+        downloads = inputs.downloads
         destination_path = inputs.destination_path
 
         # With container_name it hosts that container; alone it is the target itself.
@@ -500,7 +581,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
             else:
                 raise ValueError(f"Unsupported artifact type for loading: '{artifact.type}'")
 
-        if urls:
+        if downloads:
             if self.sandbox_name is not None:
                 sandbox = target_sandbox
             else:
@@ -514,19 +595,6 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
                 )
                 sandbox = await provider.get_sandbox(agent.sandbox_id)
             destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/"
-
-            files: list[str] = []
-            seen: dict[str, int] = {}
-            for url in urls:
-                base = unquote(os.path.basename(urlparse(url).path)) or "downloaded"
-                count = seen.get(base, 0)
-                seen[base] = count + 1
-                if count == 0:
-                    files.append(base)
-                else:
-                    stem, ext = os.path.splitext(base)
-                    files.append(f"{stem}-{count}{ext}")
-
             semaphore = asyncio.Semaphore(8)
 
             async def _load_one(url: str, filename: str) -> None:
@@ -539,13 +607,13 @@ async def _load_one(url: str, filename: str) -> None:
                     target_desc = f"VM sandbox '{self.sandbox_name}'" if onto_vm_host else "agent"
                     logger.info(f"Loaded URL into {target_desc}: {url} -> {dest}")
 
-            await asyncio.gather(*(_load_one(u, f) for u, f in zip(urls, files)))
+            await asyncio.gather(*(_load_one(u, f) for u, f in downloads))
             context.metadata.setdefault("loaded_urls", []).append({
                 "step_id": self.id,
                 "agent_name": None if self.sandbox_name is not None else self.agent_name,
                 "sandbox_name": self.sandbox_name,
                 "destination_path": destination,
-                "files": files,
+                "files": [filename for _, filename in downloads],
             })
 
         return context

tst/unit/task_step/test_load_artifact_container_sandbox.py

@@ -107,3 +107,70 @@ async def test_urls_go_into_a_container_sandbox(monkeypatch):
 
     sandbox.write_file_from_url.assert_awaited_once_with("https://example.com/data.zip", "/work/inputs/data.zip")
     sandbox.exec_script.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_bare_urls_are_named_after_their_path_and_repeats_are_suffixed(monkeypatch):
+    sandbox = _container_sandbox(vm_backed=False)
+    step = LoadArtifactTaskStep(
+        id="load", version=None, sandbox_name="bundler", destination_path="/work/inputs",
+        urls=[
+            "https://a.example/x/data.csv", "https://b.example/data.csv?sig=1", "https://c.example/",
+            "https://d.example/my%20notes.txt",
+        ],
+    )
+
+    ctx = await step.execute(_context(sandbox, monkeypatch))
+
+    assert sorted(c.args for c in sandbox.write_file_from_url.await_args_list) == [
+        ("https://a.example/x/data.csv", "/work/inputs/data.csv"),
+        ("https://b.example/data.csv?sig=1", "/work/inputs/data-1.csv"),
+        ("https://c.example/", "/work/inputs/downloaded"),
+        ("https://d.example/my%20notes.txt", "/work/inputs/my notes.txt"),
+    ]
+    (loaded,) = ctx.metadata["loaded_urls"]
+    assert loaded["files"] == ["data.csv", "data-1.csv", "downloaded", "my notes.txt"]
+    assert loaded["destination_path"] == "/work/inputs"
+
+
+@pytest.mark.asyncio
+async def test_a_mixed_list_names_each_entry_its_own_way(monkeypatch):
+    sandbox = _container_sandbox(vm_backed=False)
+    step = LoadArtifactTaskStep(
+        id="load", version=None, sandbox_name="bundler", destination_path="/tmp/data",
+        urls=[
+            "https://a.example/data.csv",
+            {"url": "https://files.example/objects/obj-4f9c2a", "filename": "report.txt"},
+            "https://b.example/data.csv",
+        ],
+    )
+
+    ctx = await step.execute(_context(sandbox, monkeypatch))
+
+    assert sorted(c.args for c in sandbox.write_file_from_url.await_args_list) == [
+        ("https://a.example/data.csv", "/tmp/data/data.csv"),
+        ("https://b.example/data.csv", "/tmp/data/data-1.csv"),
+        ("https://files.example/objects/obj-4f9c2a", "/tmp/data/report.txt"),
+    ]
+    (loaded,) = ctx.metadata["loaded_urls"]
+    assert loaded["files"] == ["data.csv", "report.txt", "data-1.csv"]
+
+
+@pytest.mark.asyncio
+async def test_an_invalid_urls_override_fails_before_anything_is_loaded(universe, monkeypatch):
+    sandbox = _container_sandbox(vm_backed=False)
+    step = LoadArtifactTaskStep(
+        id="load", version=None, sandbox_name="bundler", artifact_id=UNIVERSE_ID,
+        urls=["https://a.example/data.csv"], destination_path="/work/inputs",
+    )
+    ctx = _context(sandbox, monkeypatch)
+    ctx.metadata["user_overrides"] = {"step_params": {"load": {"urls": [
+        "https://a.example/data.csv", {"url": "https://files.example/objects/obj-7d01", "filename": "data.csv"},
+    ]}}}
+
+    with pytest.raises(ValueError, match="would be saved as 'data.csv'"):
+        await step.execute(ctx)
+
+    sandbox.write_file_from_object.assert_not_awaited()
+    sandbox.write_file_from_url.assert_not_awaited()
+    assert "loaded_urls" not in ctx.metadata

tst/unit/task_step/test_load_artifact_overrides.py

@@ -42,7 +42,7 @@ class TestResolveInputs:
     def test_no_override_uses_stored_values(self):
         res = _load_step()._resolve_inputs(_ctx())
         assert res.artifacts == [{"id": "stored-artifact", "version": None}]
-        assert res.urls == []
+        assert res.downloads == []
         assert res.destination_path == "/dest"
 
     def test_override_artifacts_list(self):
@@ -69,14 +69,27 @@ def test_override_urls(self):
         )
         ctx = _ctx({"load": {"urls": ["https://new/x", "https://new/y"]}})
         res = step._resolve_inputs(ctx)
-        assert res.urls == ["https://new/x", "https://new/y"]
+        assert res.downloads == [("https://new/x", "x"), ("https://new/y", "y")]
         assert res.artifacts == []
 
     def test_override_urls_none_clears_list(self):
         # {"urls": null} in the JSON override must clear the list, not crash on list(None)
         step = LoadArtifactTaskStep(id="load", version=None, agent_name="agent", urls=["https://old/a"])
         res = step._resolve_inputs(_ctx({"load": {"urls": None}}))
-        assert res.urls == []
+        assert res.downloads == []
+
+    def test_override_urls_take_named_entries(self):
+        step = LoadArtifactTaskStep(id="load", version=None, agent_name="agent", urls=["https://old/a"])
+        ctx = _ctx({"load": {"urls": [
+            {"url": "https://files.example/objects/obj-4f9c2a", "filename": "sheet.xlsx"},
+            "https://new/notes.txt",
+        ]}})
+        res = step._resolve_inputs(ctx)
+        assert res.downloads == [
+            ("https://files.example/objects/obj-4f9c2a", "sheet.xlsx"),
+            ("https://new/notes.txt", "notes.txt"),
+        ]
+        assert step.urls == ["https://old/a"]
 
     def test_override_artifact_id_null_raises(self):
         # {"artifact_id": null} is malformed — fail clear, not with an opaque Artifact.get(None) error

tst/unit/task_step/test_load_artifact_url_filenames.py

@@ -0,0 +1,46 @@
+"""LoadArtifactTaskStep `urls` entries: an object entry ``{"url", "filename"}`` is saved under its
+filename verbatim, survives a round-trip, and a bad or colliding entry fails at construction."""
+
+from __future__ import annotations
+
+import pytest
+
+from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep
+
+OPAQUE_URL = "https://files.example/objects/obj-4f9c2a"
+
+
+def _step(urls) -> LoadArtifactTaskStep:
+    return LoadArtifactTaskStep(id="load", version=None, sandbox_name="box", urls=urls)
+
+
+def test_object_entries_survive_a_round_trip():
+    urls = ["https://a.example/data.csv", {"url": OPAQUE_URL, "filename": "report.txt"}]
+    data = _step(urls).to_dict()
+    assert data["urls"] == urls
+    assert LoadArtifactTaskStep.from_dict(data).urls == urls
+
+
+@pytest.mark.parametrize(
+    "entry, match",
+    [
+        ({"url": OPAQUE_URL}, "non-empty string 'filename'"),
+        ({"url": OPAQUE_URL, "name": "a.txt"}, r"unknown key\(s\) \['name'\]"),
+        ({"url": OPAQUE_URL, "filename": "../a.txt"}, r"must not contain '\.\.' segments"),
+        ({"url": OPAQUE_URL, "filename": "/etc/passwd"}, "must be relative"),
+        ({"url": OPAQUE_URL, "filename": "sub/a.txt"}, "plain file name, not a path"),
+    ],
+)
+def test_a_bad_object_entry_fails_at_construction(entry, match):
+    with pytest.raises(ValueError, match=match):
+        _step([entry])
+
+
+def test_a_filename_given_twice_fails():
+    with pytest.raises(ValueError, match="'a.txt' is given to more than one urls entry"):
+        _step([{"url": OPAQUE_URL, "filename": "a.txt"}, {"url": OPAQUE_URL + "2", "filename": "a.txt"}])
+
+
+def test_a_bare_url_named_like_an_explicit_filename_fails():
+    with pytest.raises(ValueError, match=r"'https://a.example/data.csv' would be saved as 'data.csv'"):
+        _step(["https://a.example/data.csv", {"url": OPAQUE_URL, "filename": "data.csv"}])

tst/unit/task_step/test_load_artifact_vm_target.py

@@ -164,6 +164,20 @@ async def test_urls_use_the_host_downloader(self, vm):
         assert not [c for c in calls if c[0] == "container_url"]
         assert ctx.metadata["loaded_urls"][0]["sandbox_name"] == "mk"
 
+    @pytest.mark.asyncio
+    async def test_a_named_url_is_curled_to_its_filename(self, vm):
+        ctx, calls = vm
+        step = LoadArtifactTaskStep(
+            id="stage", version=None, sandbox_name="mk", destination_path="/app/seeds",
+            urls=[{"url": "https://files.example/objects/obj-4f9c2a", "filename": "repo.tar.gz"}],
+        )
+        ctx = await step.execute(ctx)
+
+        curls = [c[1] for c in calls if c[0] == "exec" and c[1].startswith("curl -fsSL")]
+        assert len(curls) == 1, curls
+        assert "https://files.example/objects/obj-4f9c2a" in curls[0] and curls[0].endswith("-o /app/seeds/repo.tar.gz")
+        assert ctx.metadata["loaded_urls"][0]["files"] == ["repo.tar.gz"]
+
 
 class TestContainerPathUnchanged:
     @pytest.mark.asyncio