Merge 81882da444 into 851743adcf

strix/interface/cloud/source_scan.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import re
 import sys
+import time
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 from urllib.parse import quote
@@ -27,6 +28,25 @@ if TYPE_CHECKING:
 
 _SHA256 = re.compile(r"^[0-9a-fA-F]{64}$")
 
+_UPLOAD_PROCESSING_POLL_S = 5
+_UPLOAD_PROCESSING_TIMEOUT_S = 300
+
+
+class _UploadStillProcessing(http.CloudError):
+    """The platform is still processing an upload past the polling bound.
+
+    Unlike other completion failures this is not a rejection: the upload
+    is retained so the scan can be started from it once it is ready.
+    """
+
+    def __init__(self, upload_id: str) -> None:
+        super().__init__(
+            f"the platform is still processing source upload {upload_id} after "
+            f"{_UPLOAD_PROCESSING_TIMEOUT_S:g} seconds; the upload was kept. "
+            f"Start the scan with `strix cloud scans start --upload-ids {upload_id}` "
+            "once processing finishes."
+        )
+
 
 @dataclass
 class LocalSourceScan:
@@ -251,15 +271,9 @@ def _upload_scan_source(bundle: SourceBundle, *, token: str | None) -> str:
         raise error
     try:
         http.upload_file(cast("str", signed_url), cast("str", upload_token), bundle.archive_path)
-        completed = http.check(
-            http.request(
-                "POST",
-                "/uploads/complete",
-                token=token,
-                body={"upload_id": upload_id},
-            )
-        )
-        _validate_completed_upload(completed, expected_id=cast("str", upload_id))
+        _complete_upload(cast("str", upload_id), token=token)
+    except _UploadStillProcessing:
+        raise
     except BaseException as error:
         try:
             _delete_upload(cast("str", upload_id), token=token)
@@ -274,6 +288,35 @@ def _upload_scan_source(bundle: SourceBundle, *, token: str | None) -> str:
     return cast("str", upload_id)
 
 
+def _complete_upload(upload_id: str, *, token: str | None) -> None:
+    """Complete an upload, polling while the platform is still processing it.
+
+    The platform answers ``202`` with ``status: "processing"`` until the
+    upload record is ready, so a single completion call is not enough for
+    larger sources. Poll on a short bounded loop; past the bound, raise
+    :class:`_UploadStillProcessing` so the upload is kept for a later
+    ``--upload-ids`` launch instead of being deleted as a rejection.
+    """
+    deadline = time.monotonic() + _UPLOAD_PROCESSING_TIMEOUT_S
+    while True:
+        completed = http.check(
+            http.request(
+                "POST",
+                "/uploads/complete",
+                token=token,
+                body={"upload_id": upload_id},
+            )
+        )
+        fields = cast("dict[str, Any]", completed) if isinstance(completed, dict) else {}
+        if str(fields.get("status", "")).lower() != "processing":
+            _validate_completed_upload(completed, expected_id=upload_id)
+            return
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise _UploadStillProcessing(upload_id)
+        time.sleep(min(_UPLOAD_PROCESSING_POLL_S, remaining))
+
+
 def _validate_completed_upload(completed: Any, *, expected_id: str) -> None:
     fields = cast("dict[str, Any]", completed) if isinstance(completed, dict) else {}
     if fields.get("id") != expected_id:

tests/test_cloud_source_upload.py

@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any
 import pytest
 
 from strix.interface import cloud
-from strix.interface.cloud import http, source_upload
+from strix.interface.cloud import http, source_scan, source_upload
 
 
 if TYPE_CHECKING:
@@ -393,6 +393,114 @@ def test_source_upload_is_completed_and_attached_to_scan(
     assert uploaded_path is not None and not uploaded_path.exists()
 
 
+def test_source_upload_polls_while_platform_is_processing(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: Any
+) -> None:
+    (tmp_path / "app.py").write_text("print('safe')\n", encoding="utf-8")
+    monkeypatch.setattr(source_scan, "_UPLOAD_PROCESSING_POLL_S", 0)
+    calls: list[tuple[str, str]] = []
+    complete_calls = 0
+
+    def fake_request(method: str, path: str, **_kwargs: Any) -> FakeResponse:
+        nonlocal complete_calls
+        calls.append((method, path))
+        if path == "/uploads/request":
+            return FakeResponse(
+                {
+                    "upload_id": "upload-1",
+                    "signed_url": "https://storage.test/object",
+                    "token": "signed",
+                }
+            )
+        if path == "/uploads/complete":
+            complete_calls += 1
+            if complete_calls <= 2:
+                return FakeResponse({"status": "processing"}, status_code=202)
+            return FakeResponse({"id": "upload-1"})
+        if path == "/scans":
+            return FakeResponse({"scan_id": "scan-1", "status": "pending"})
+        raise AssertionError(path)
+
+    monkeypatch.setattr(http, "request", fake_request)
+    monkeypatch.setattr(http, "upload_file", lambda *_args, **_kwargs: None)
+
+    assert cloud.run_cloud(["scans", "start", "--source", str(tmp_path), "--yes", "--json"]) == 0
+    payload = json.loads(capsys.readouterr().out)
+    assert payload["upload_id"] == "upload-1"
+    assert payload["scan"]["scan_id"] == "scan-1"
+    assert complete_calls == 3
+    assert ("DELETE", "/uploads/upload-1") not in calls
+
+
+def test_source_upload_processing_poll_stops_at_deadline(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(source_scan, "_UPLOAD_PROCESSING_POLL_S", 5)
+    monkeypatch.setattr(source_scan, "_UPLOAD_PROCESSING_TIMEOUT_S", 12)
+    now = 0.0
+    sleeps: list[float] = []
+    complete_calls = 0
+
+    def monotonic() -> float:
+        return now
+
+    def sleep(seconds: float) -> None:
+        nonlocal now
+        sleeps.append(seconds)
+        now += seconds
+
+    def fake_request(method: str, path: str, **_kwargs: Any) -> FakeResponse:
+        nonlocal complete_calls
+        assert method == "POST"
+        assert path == "/uploads/complete"
+        complete_calls += 1
+        return FakeResponse({"status": "processing"}, status_code=202)
+
+    monkeypatch.setattr("strix.interface.cloud.source_scan.time.monotonic", monotonic)
+    monkeypatch.setattr("strix.interface.cloud.source_scan.time.sleep", sleep)
+    monkeypatch.setattr(http, "request", fake_request)
+
+    with pytest.raises(http.CloudError, match=r"still processing.*after 12 seconds"):
+        source_scan._complete_upload("upload-1", token=None)
+
+    assert complete_calls == 4
+    assert sleeps == [5, 5, 2]
+
+
+def test_source_upload_kept_when_processing_outlasts_poll_bound(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: Any
+) -> None:
+    (tmp_path / "app.py").write_text("print('safe')\n", encoding="utf-8")
+    monkeypatch.setattr(source_scan, "_UPLOAD_PROCESSING_TIMEOUT_S", 0)
+    calls: list[tuple[str, str]] = []
+
+    def fake_request(method: str, path: str, **_kwargs: Any) -> FakeResponse:
+        calls.append((method, path))
+        if path == "/uploads/request":
+            return FakeResponse(
+                {
+                    "upload_id": "upload-1",
+                    "signed_url": "https://storage.test/object",
+                    "token": "signed",
+                }
+            )
+        if path == "/uploads/complete":
+            return FakeResponse({"status": "processing"}, status_code=202)
+        raise AssertionError(path)
+
+    monkeypatch.setattr(http, "request", fake_request)
+    monkeypatch.setattr(http, "upload_file", lambda *_args, **_kwargs: None)
+
+    assert (
+        cloud.run_cloud(["scans", "start", "--source", str(tmp_path), "--yes", "--json"])
+        == http.EXIT_ERROR
+    )
+    error = json.loads(capsys.readouterr().out)["error"]
+    assert "still processing" in error
+    assert "--upload-ids upload-1" in error
+    assert ("DELETE", "/uploads/upload-1") not in calls
+
+
 def test_source_upload_with_domain_is_a_live_test(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: Any
 ) -> None: