fix(snapshot): stream a service's export only from its own address

src/agent_env/task_step/task_steps/snapshot_env.py

@@ -469,7 +469,7 @@ async def _export_environment_to_file(
 
         v1 ``get_data``, handed a namespace grant when the service takes the data-objects
         form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` →
-        object url (no download) / base64 ``bytes`` / streamed relative-or-http ``uri`` →
+        object url (no download) / base64 ``bytes`` / streamed relative-or-same-origin ``uri`` →
         ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON
         guard. Errors are caught by the caller (fails just this service)."""
         from agent_env.env import legacy_protocol
@@ -510,9 +510,14 @@ async def _export_environment_to_file(
                         return uri
                     # A relative uri (e.g. "export-snapshot") means "stream it from
                     # my service endpoint" — resolve against the service base_url so
-                    # multi-GB bundles never ride inline in the JSON-RPC response.
+                    # multi-GB bundles never ride inline in the JSON-RPC response. An
+                    # absolute one must be on that same origin: this process fetches it.
                     if not scheme:
                         uri = f"{base_url.rstrip('/')}/{uri.lstrip('/')}"
+                    elif urlparse(uri)[:2] != urlparse(base_url)[:2]:
+                        raise ValueError(
+                            f"FilePart for '{environment_name}' names {uri!r}, which is not on its service's address {base_url}"
+                        )
                     async with httpx.AsyncClient() as client:
                         async with client.stream(
                             "GET",

tst/unit/task_step/test_snapshot_env.py

@@ -758,21 +758,24 @@ async def fake_get_card(base_url, timeout=10):
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize("absolute", [False, True], ids=["relative", "same-origin"])
 async def test_export_service_to_file_resolves_relative_uri_and_streams(
-    monkeypatch, tmp_path
+    monkeypatch, tmp_path, absolute
 ):
     """A FilePart with a RELATIVE uri (e.g. "export-snapshot") is resolved against
     the service base_url and streamed to disk — so multi-GB bundles never ride
-    inline in the JSON-RPC response."""
+    inline in the JSON-RPC response. An absolute uri on that address streams the same."""
     from agentenv_protocol import FilePart
     from agentenv_protocol.client import GetDataResponse
     from agentenv_protocol import client as protocol_v1
     from agent_env.env import legacy_protocol
     from agent_env.task_step.task_steps import snapshot_env as mod
 
     zip_bytes = b"PK\x03\x04streamed-from-export-snapshot"
+    expected_base = legacy_protocol.environment_base_url("https://gw", "gdrive", mcp=True)
+    uri = f"{expected_base.rstrip('/')}/export-snapshot" if absolute else "export-snapshot"
     part = FilePart(
-        file={"uri": "export-snapshot", "name": "gdrive.zip", "mimeType": "application/zip"}
+        file={"uri": uri, "name": "gdrive.zip", "mimeType": "application/zip"}
     )
 
     async def fake_supports_v1(base_url):
@@ -826,7 +829,6 @@ def stream(self, method, url, timeout=None):
     assert suffix == ".zip"
     with open(out, "rb") as f:
         assert f.read() == zip_bytes
-    expected_base = legacy_protocol.environment_base_url("https://gw", "gdrive", mcp=True)
     assert captured["method"] == "GET"
     assert captured["url"] == f"{expected_base.rstrip('/')}/export-snapshot"
 
@@ -868,6 +870,25 @@ async def test_export_service_to_file_returns_an_object_url_as_the_services_uplo
     assert not out.exists()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("uri", [
+    "http://127.0.0.1:8234/api/runs",
+    "http://169.254.169.254/latest/meta-data/",
+    "https://internal.example/export-snapshot",
+    "https://gw@internal.example/export-snapshot",
+    "https://gw:8443/export-snapshot",
+    "http://gw/export-snapshot",
+])
+async def test_export_service_to_file_refuses_an_absolute_uri_off_its_service_address(monkeypatch, tmp_path, uri):
+    """This process makes the request, from where it sits, so a service names only its own address."""
+    _serve_filepart_uri(monkeypatch, uri)
+
+    out = tmp_path / "export-tmp"
+    with pytest.raises(ValueError, match="not on its service's address"):
+        await SnapshotEnvTaskStep._export_environment_to_file("https://gw", "gdrive", str(out), 30)
+    assert not out.exists()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("upload_url", ["fake://home/snapshots/gdrive#v2.zip", "fake://shared/snapshots/gdrive#v2.zip"])
 async def test_core_registers_a_service_upload_in_place_and_it_loads_back(local_stores, monkeypatch, tmp_path, upload_url):