feat: add internal response passthrough and sanitize download URLs

README.md

@@ -74,13 +74,20 @@ curl --request POST --url http://localhost:8000/inference  --header 'Content-Typ
     --data '{"input_1": "Hello"}'
 ```
 
-Workers return a `WorkerResponse` with three fields:
+Workers return a `WorkerResponse`:
 
 - `billable_seconds`: the billable duration as a number of seconds, including
   fractional seconds when available, or `null` when it cannot be determined.
 - `stats`: numeric worker measurements. Use `duration` for total worker
   processing time; additional keys may report individual phases.
 - `result`: the worker-specific JSON object.
+- `internal`: an optional object passed through without inspection, for data the
+  deployment's own response consumer reads and then removes before answering its
+  callers. Its contents are entirely up to that deployment; this library gives it
+  no meaning and does not validate it. Leave it unset if nothing consumes it.
+
+Any other field a worker returns is dropped. `internal` is the one exception, so
+a worker cannot reach a caller with a field nobody agreed to.
 
 The `/health` endpoint reports the worker artifact version supplied through
 `WORKER_VERSION` and available GPU metadata in addition to `ok`. Deployments

maestro_worker_python/download_file.py

@@ -9,10 +9,11 @@
 import requests
 
 from .response import ValidationError
+from .sanitize_url import sanitize_url
 
 
 def download_file(url: str, filename: str | None = None) -> str:
-    logging.info(f"Downloading input: {url}")
+    logging.info("Downloading input: %s", sanitize_url(url))
     response = requests.get(url, allow_redirects=True, timeout=300)
 
     try:
@@ -39,7 +40,7 @@ def download_files_manager(*urls: str) -> Iterator[None | str | list[str]]:
         list_objects = []
         for url in urls:
             filename = tempfile.NamedTemporaryFile()
-            logging.info("Downloading file from url -> %s, filename -> %s", url, filename.name)
+            logging.info("Downloading file from url -> %s, filename -> %s", sanitize_url(url), filename.name)
             with ThreadPoolExecutor(max_workers=20) as exe:
                 thread_list.append(exe.submit(download_file, url, filename.name))
             list_objects.append(filename)

maestro_worker_python/response.py

@@ -7,6 +7,8 @@ class WorkerResponse(BaseModel):
     billable_seconds: float | None
     stats: dict[str, float]
     result: dict[str, Any]
+    # Opaque passthrough for the deployment; other undeclared fields are dropped.
+    internal: dict[str, Any] | None = None
 
 
 class ValidationError(Exception):

maestro_worker_python/sanitize_url.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from urllib.parse import urlsplit, urlunsplit
+
+# Ports implied by their scheme carry no information, and keeping them would make
+# ``https://host:443/x`` and ``https://host/x`` read as two different locations.
+_DEFAULT_PORTS = {
+    "ftp": 21,
+    "http": 80,
+    "https": 443,
+    "ws": 80,
+    "wss": 443,
+}
+
+# Substituted for a URL that cannot be parsed. An unparseable string is the one
+# case where there is no way to show it holds no credential, so none of it is kept.
+UNPARSEABLE_URL = "<unparseable-url>"
+
+
+def sanitize_url(url: str) -> str:
+    """Strip credentials and signing material from a URL while keeping it identifiable.
+
+    Removes userinfo, query, and fragment, lowercases scheme and host, and drops
+    the scheme's default port. The path is passed through byte-for-byte: decoding
+    it would let percent-encoded data be reinterpreted as structure, and
+    re-encoding it would change which object the URL names.
+
+    The result is safe to log but is not anonymous — it still names an object.
+    """
+    try:
+        parts = urlsplit(url)
+        port = parts.port
+    except ValueError:
+        return UNPARSEABLE_URL
+
+    # hostname drops userinfo and lowercases in one step, unlike netloc.
+    host = (parts.hostname or "").lower()
+
+    if ":" in host:
+        # An IPv6 literal loses the brackets it needs to round-trip.
+        host = f"[{host}]"
+
+    if port is not None and port != _DEFAULT_PORTS.get(parts.scheme.lower()):
+        host = f"{host}:{port}"
+
+    return urlunsplit((parts.scheme.lower(), host, parts.path, "", ""))

pyproject.toml

@@ -26,6 +26,7 @@ maestro-upload-server = "maestro_worker_python.run_upload_server:main"
 
 [dependency-groups]
 dev = [
+  "httpx2>=2.9.1",
   "pytest>=9.0.3,<10",
   "pytest-httpserver>=1.0.6,<2",
   "ruff>=0.14.9,<0.15",

tests/test_download_file.py

@@ -1,10 +1,27 @@
+import logging
 import os
 
 import pytest
 
 from maestro_worker_python.download_file import download_file, download_files_manager
 from maestro_worker_python.response import ValidationError
 
+SIGNED_QUERY = "X-Goog-Signature=deadbeefcafe"
+
+
+def _assert_no_signing_material_logged(caplog, httpserver):
+    # The in-process test server logs the request line it received, query included.
+    # That is the fixture's own access log, not the library under test.
+    records = [record for record in caplog.records if record.name != "werkzeug"]
+    assert records, "nothing was logged, so this proves nothing"
+    for record in records:
+        message = record.getMessage()
+        assert "deadbeefcafe" not in message
+        assert "X-Goog-Signature" not in message
+    # The sanitized location must still be logged, or redaction has cost the
+    # operator the ability to tell which input a download refers to.
+    assert any(httpserver.url_for("/test") in record.getMessage() for record in records)
+
 
 def test_download_file(httpserver):
     httpserver.expect_request("/test").respond_with_data("hello")
@@ -38,6 +55,26 @@ def test_download_files_manager(httpserver):
     assert all(files_content)
 
 
+def test_download_file_keeps_signing_material_out_of_the_logs(httpserver, caplog):
+    httpserver.expect_request("/test").respond_with_data("hello")
+    url = httpserver.url_for(f"/test?{SIGNED_QUERY}")
+
+    with caplog.at_level(logging.INFO):
+        download_file(url)
+
+    _assert_no_signing_material_logged(caplog, httpserver)
+
+
+def test_download_files_manager_keeps_signing_material_out_of_the_logs(httpserver, caplog):
+    httpserver.expect_request("/test").respond_with_data("hello")
+    url = httpserver.url_for(f"/test?{SIGNED_QUERY}")
+
+    with caplog.at_level(logging.INFO), download_files_manager(url) as downloaded:
+        assert downloaded is not None
+
+    _assert_no_signing_material_logged(caplog, httpserver)
+
+
 def test_download_files_manager_delete(httpserver):
     httpserver.expect_request("/test").respond_with_data("hello")
     url = httpserver.url_for("/test?foo=bar")

tests/test_sanitize_url.py

@@ -0,0 +1,37 @@
+import pytest
+
+from maestro_worker_python.sanitize_url import UNPARSEABLE_URL, sanitize_url
+
+
+@pytest.mark.parametrize(
+    ("url", "expected"),
+    [
+        # A signed GCS read: the signature is the credential.
+        (
+            "https://storage.googleapis.com/bucket/song.mp3"
+            "?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Signature=deadbeef",
+            "https://storage.googleapis.com/bucket/song.mp3",
+        ),
+        ("https://user:hunter2@host/path", "https://host/path"),
+        ("https://host/path#fragment", "https://host/path"),
+        ("HTTPS://Storage.GoogleAPIs.COM/Bucket/Song.mp3", "https://storage.googleapis.com/Bucket/Song.mp3"),
+        ("https://host:443/path", "https://host/path"),
+        ("http://host:80/path", "http://host/path"),
+        ("https://host:8443/path", "https://host:8443/path"),
+        ("gs://My-Bucket/path/to/obj.wav?x=1", "gs://my-bucket/path/to/obj.wav"),
+        ("hf://meta-llama/Llama-3/model.safetensors", "hf://meta-llama/Llama-3/model.safetensors"),
+        ("file:///tmp/local.wav", "file:///tmp/local.wav"),
+        ("http://[::1]:8080/path", "http://[::1]:8080/path"),
+    ],
+)
+def test_sanitize_url_removes_credentials_and_normalizes_the_location(url, expected):
+    assert sanitize_url(url) == expected
+
+
+def test_sanitize_url_passes_the_path_through_without_decoding_it():
+    """Decoding would let %2F be reread as a separator, naming a different object."""
+    assert sanitize_url("https://host/a%2Fb/My%20Song%20(1).mp3?sig=x") == "https://host/a%2Fb/My%20Song%20(1).mp3"
+
+
+def test_sanitize_url_drops_a_url_it_cannot_parse():
+    assert sanitize_url("https://host:not-a-port/path") == UNPARSEABLE_URL

tests/test_serve.py

@@ -4,9 +4,30 @@
 from pathlib import Path
 
 import pytest
+from fastapi.testclient import TestClient
 
 from maestro_worker_python.config import settings
 
+# Deliberately mixed value types: `internal` is typed only as an object, so
+# coercion of anything inside it would be a silent corruption of a channel this
+# library is supposed to be agnostic about.
+INTERNAL_PAYLOAD = {
+    "file_transfers": [
+        {"direction": "input", "field": "input_url", "url": "https://host/in.mp3", "bytes": 0},
+        {"direction": "output", "field": "vocals", "url": "https://host/out%2Fa.wav", "bytes": 1234},
+    ],
+    "nested": {"flag": True, "count": 7, "ratio": 0.5, "absent": None},
+}
+
+
+def _write_worker(worker_path: Path, returned: str) -> None:
+    worker_path.write_text(
+        f"INTERNAL = {INTERNAL_PAYLOAD!r}\n"
+        "class MoisesWorker:\n"
+        "    def inference(self, input_data):\n"
+        f"        return {returned}\n"
+    )
+
 
 def _import_serve(monkeypatch: pytest.MonkeyPatch, worker_path: Path):
     monkeypatch.setattr(settings, "model_path", str(worker_path))
@@ -55,3 +76,50 @@ async def run():
 
     asyncio.run(run())
     assert not ready_path.exists()
+
+
+def test_inference_passes_internal_content_through_uninspected(tmp_path, monkeypatch):
+    worker_path = tmp_path / "worker.py"
+    _write_worker(
+        worker_path,
+        '{"billable_seconds": 1.0, "stats": {"duration": 0.5}, "result": {"ok": True}, "internal": INTERNAL}',
+    )
+    serve_module = _import_serve(monkeypatch, worker_path)
+
+    body = TestClient(serve_module.app).post("/inference", json={}).json()
+
+    assert body["internal"] == INTERNAL_PAYLOAD
+    nested = body["internal"]["nested"]
+    # Equality alone would accept 7 becoming 7.0 or True becoming 1.
+    assert isinstance(nested["count"], int) and not isinstance(nested["count"], bool)
+    assert isinstance(nested["flag"], bool)
+    assert body["result"] == {"ok": True}
+
+
+def test_inference_drops_a_field_the_response_model_does_not_declare(tmp_path, monkeypatch):
+    """The single declared passthrough is what keeps the envelope closed."""
+    worker_path = tmp_path / "worker.py"
+    _write_worker(
+        worker_path,
+        '{"billable_seconds": 1.0, "stats": {}, "result": {}, "internal": {}, "leaked": "not for callers"}',
+    )
+    serve_module = _import_serve(monkeypatch, worker_path)
+
+    body = TestClient(serve_module.app).post("/inference", json={}).json()
+
+    assert "leaked" not in body
+
+
+def test_inference_distinguishes_an_empty_internal_object_from_no_internal_object(tmp_path, monkeypatch):
+    """Producers signal support with the object itself, so the two must not collapse."""
+    worker_path = tmp_path / "worker.py"
+    _write_worker(worker_path, '{"billable_seconds": 1.0, "stats": {}, "result": {}, "internal": {}}')
+    serve_module = _import_serve(monkeypatch, worker_path)
+    empty = TestClient(serve_module.app).post("/inference", json={}).json()
+
+    _write_worker(worker_path, '{"billable_seconds": 1.0, "stats": {}, "result": {}}')
+    serve_module = _import_serve(monkeypatch, worker_path)
+    unset = TestClient(serve_module.app).post("/inference", json={}).json()
+
+    assert empty["internal"] == {}
+    assert unset["internal"] is None