From ecd46b86fcc40f87f9ec13af369c2d7e1733a9ca Mon Sep 17 00:00:00 2001 From: Tejas Polakam Date: Wed, 7 Oct 2026 16:38:53 -0400 Subject: [PATCH] feat(load_artifact): allow urls entries to set the saved filename A `urls` entry may now be {"url": ..., "filename": ...} as well as a bare URL string. A bare string keeps today's naming exactly (last path segment, percent-decoded, "downloaded" when empty, repeats suffixed -1, -2, ...); an object entry is saved as / verbatim, so a URL with no usable name (an opaque object key) can land under the name the prompt refers to. Entries are validated and named by one function, called at construction for the stored urls and in _resolve_inputs for a per-run `urls` override, so a bad override fails before anything is loaded. An object entry needs a non-empty string url and filename and no other key; the filename must pass validate_relative_filename and be a plain file name (no path separators, not "."). An explicit filename is never renamed: one given twice, or one a bare URL's derived name would match, raises instead of being suffixed. loaded_urls[].files records the names as written. to_dict/from_dict round- trip the entries unchanged: strings stay strings, objects stay objects. Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/load_artifact.py | 112 ++++++++++++++---- .../test_load_artifact_container_sandbox.py | 67 +++++++++++ .../task_step/test_load_artifact_overrides.py | 19 ++- .../test_load_artifact_url_filenames.py | 46 +++++++ .../task_step/test_load_artifact_vm_target.py | 14 +++ 5 files changed, 233 insertions(+), 25 deletions(-) create mode 100644 tst/unit/task_step/test_load_artifact_url_filenames.py diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index b1564c83..84899af6 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/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 diff --git a/tst/unit/task_step/test_load_artifact_container_sandbox.py b/tst/unit/task_step/test_load_artifact_container_sandbox.py index 24e68243..7973ce0e 100644 --- a/tst/unit/task_step/test_load_artifact_container_sandbox.py +++ b/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 diff --git a/tst/unit/task_step/test_load_artifact_overrides.py b/tst/unit/task_step/test_load_artifact_overrides.py index 6db9f070..bc8bd98c 100644 --- a/tst/unit/task_step/test_load_artifact_overrides.py +++ b/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 diff --git a/tst/unit/task_step/test_load_artifact_url_filenames.py b/tst/unit/task_step/test_load_artifact_url_filenames.py new file mode 100644 index 00000000..23234ce9 --- /dev/null +++ b/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"}]) diff --git a/tst/unit/task_step/test_load_artifact_vm_target.py b/tst/unit/task_step/test_load_artifact_vm_target.py index 30ec4d85..528c487d 100644 --- a/tst/unit/task_step/test_load_artifact_vm_target.py +++ b/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