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