fix: make model integrity downloads reliable (#49)

Author: Mr-Sunglasses

README.md

@@ -454,11 +454,10 @@ model can execute code in those runtimes.
 
 ### Coverage
 
-40 of 58 catalog models are pinned, covering every Hugging Face source. The
-remainder cannot be pinned from published metadata:
+Every catalog source that publishes usable integrity metadata is pinned,
+including all Hugging Face sources and the Moonshine asset manifests. The
+current exceptions cannot be pinned from published metadata:
 
-- 13 Moonshine models are downloaded by the `moonshine_voice` library rather
-  than by the gateway, so their transfer is outside our control.
 - 3 Handy-mirrored models on `blob.handy.computer` return a multipart S3 ETag,
   which is not a digest of the file content.
 - 2 sherpa-onnx release tarballs on GitHub publish no checksum.
@@ -476,17 +475,23 @@ is silently downgraded.
 
 ### Refreshing pins
 
-When upstream legitimately re-uploads a model, its pinned download starts
-failing until the pin is updated. That is intentional: the change becomes a
-reviewable diff instead of a silent swap.
+The harvester is incremental by default: after adding a catalog model, it pins
+only entries missing from `app/model_pins.json`. Existing records remain byte
+for byte unchanged, so adding one model cannot silently refresh every model in
+the catalog. Use `--refresh` only when intentionally reviewing every upstream
+change; `--only` explicitly refreshes the matching model or family.
 
 ```sh
-uv run scripts/harvest-model-pins.py                    # all free sources
-uv run scripts/harvest-model-pins.py --only whisperkit:  # one family
+uv run scripts/harvest-model-pins.py                     # newly added models
+uv run scripts/harvest-model-pins.py --only whisperkit:  # refresh one family
+uv run scripts/harvest-model-pins.py --refresh           # refresh everything
 ```
 
-Review the resulting diff as carefully as code. A changed digest means the
-upstream bytes changed, and the commit message should say why.
+Each revision and its digests are written as one snapshot. If the complete
+snapshot cannot be collected, the command fails and preserves the previous
+record rather than combining a new revision with stale hashes. Review the
+resulting diff as carefully as code. A changed digest means the upstream bytes
+changed, and the commit message should say why.
 
 ## Engine selection
 

app/model_manager.py

@@ -8,7 +8,7 @@
 import tarfile
 import threading
 import time
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 from functools import partial
 from http import client as http_client
@@ -51,10 +51,7 @@
 
 _MAX_NETWORK_ATTEMPTS = 3
 _RETRY_BACKOFF_SECONDS = 2.0
-# Transport-level failures only. A hash mismatch is deliberately excluded: it
-# is retried by the user, not automatically, since a mismatch can mean the
-# host is compromised and silently re-trying could mask that rather than
-# surface it.
+MAX_PARALLEL_FILE_DOWNLOADS = 4
 _RETRYABLE_NETWORK_ERRORS = (
     TimeoutError,
     ConnectionError,
@@ -69,6 +66,8 @@ def _call_with_retries[ActionResult](
     *,
     on_retry: Callable[[], None] | None = None,
     attempts: int = _MAX_NETWORK_ATTEMPTS,
+    retryable_errors: tuple[type[BaseException], ...] = _RETRYABLE_NETWORK_ERRORS,
+    retry_delay: Callable[[BaseException, int], float] | None = None,
 ) -> ActionResult:
     """Call *action* up to *attempts* times, retrying transport failures.
 
@@ -80,12 +79,15 @@ def _call_with_retries[ActionResult](
     for attempt in range(1, attempts + 1):
         try:
             return action()
-        except _RETRYABLE_NETWORK_ERRORS:
+        except retryable_errors as error:
             if on_retry is not None:
                 on_retry()
             if attempt == attempts:
                 raise
-            time.sleep(_RETRY_BACKOFF_SECONDS * attempt)
+            delay = _RETRY_BACKOFF_SECONDS * attempt
+            if retry_delay is not None:
+                delay = retry_delay(error, attempt)
+            time.sleep(delay)
     raise AssertionError("unreachable")  # pragma: no cover
 
 
@@ -149,6 +151,50 @@ class _DownloadHandle:
     task: asyncio.Task[None] | None = field(default=None)
 
 
+@dataclass(slots=True)
+class _DownloadBatch:
+    downloads: tuple[Callable[[], Awaitable[None]], ...]
+    download_handle: _DownloadHandle
+    semaphore: asyncio.Semaphore = field(
+        default_factory=lambda: asyncio.Semaphore(MAX_PARALLEL_FILE_DOWNLOADS)
+    )
+
+    async def run(self) -> None:
+        """Wait for every worker to stop before reporting a folder failure."""
+        outcomes = await asyncio.gather(
+            *(self._run_one(download) for download in self.downloads),
+            return_exceptions=True,
+        )
+        failure = self._first_failure(outcomes)
+        if failure is not None:
+            raise failure
+        if any(isinstance(outcome, DownloadCancelled) for outcome in outcomes):
+            raise DownloadCancelled
+
+    async def _run_one(self, download: Callable[[], Awaitable[None]]) -> None:
+        try:
+            await self._run_bounded(download)
+        except Exception:
+            self.download_handle.cancel.set()
+            raise
+
+    async def _run_bounded(self, download: Callable[[], Awaitable[None]]) -> None:
+        async with self.semaphore:
+            if self.download_handle.cancel.is_set():
+                raise DownloadCancelled
+            await download()
+
+    def _first_failure(self, outcomes: list[None | BaseException]) -> Exception | None:
+        return next(
+            (
+                outcome
+                for outcome in outcomes
+                if isinstance(outcome, Exception) and not isinstance(outcome, DownloadCancelled)
+            ),
+            None,
+        )
+
+
 class ModelManager:
     """Downloads, lists, and deletes local speech models."""
 
@@ -383,7 +429,7 @@ async def _run_archive_download(
         partial_dir = final_dir.with_name(f"{final_dir.name}{PARTIAL_DIRECTORY_SUFFIX}")
         extraction_dir = final_dir.with_name(f"{final_dir.name}.extracting")
         archive_path = final_dir.with_name(f"{final_dir.name}.download")
-        _clear_staging_directory(partial_dir)
+        shutil.rmtree(partial_dir, ignore_errors=True)
         _remove_tree(extraction_dir)
         archive_path.unlink(missing_ok=True)
         try:
@@ -465,14 +511,18 @@ async def _run_sherpa_huggingface_download(
                 for name in model.required_files
                 if (entry := available.get(name)) is not None
             )
-            await asyncio.gather(
-                *(
-                    self._download_required_file(
-                        model, partial_dir, download_handle, name, available
-                    )
-                    for name in model.required_files
+            downloads = tuple(
+                partial(
+                    self._download_required_file,
+                    model,
+                    partial_dir,
+                    download_handle,
+                    name,
+                    available,
                 )
+                for name in model.required_files
             )
+            await _DownloadBatch(downloads, download_handle).run()
             missing = [name for name in model.required_files if not (partial_dir / name).is_file()]
             if missing:
                 raise RuntimeError(_missing_model_files_message(missing))
@@ -590,12 +640,11 @@ async def _run_huggingface_download(
         _remove_tree(partial_dir)
         partial_dir.mkdir(parents=True, exist_ok=True)
         try:
-            await asyncio.gather(
-                *(
-                    self._download_repo_file(model, partial_dir, download_handle, entry)
-                    for entry in files
-                )
+            downloads = tuple(
+                partial(self._download_repo_file, model, partial_dir, download_handle, entry)
+                for entry in files
             )
+            await _DownloadBatch(downloads, download_handle).run()
             if download_handle.cancel.is_set():
                 raise DownloadCancelled
             final_dir.parent.mkdir(parents=True, exist_ok=True)
@@ -856,13 +905,17 @@ def _download_file(
     before raising, so a rejected download cannot be left behind for an engine
     to load later.
 
-    A dropped connection or read timeout mid-stream is retried a few times
-    (see `_call_with_retries`) rather than surfaced immediately, since it can
-    otherwise produce a truncated file that fails SHA-256 verification for a
-    reason that has nothing to do with the source's integrity.
+    A dropped connection, short response, or checksum mismatch is retried a
+    few times against the exact same expected digest. A persistent mismatch is
+    still rejected; retrying can only accept bytes that satisfy the pin.
     """
     download = _DownloadAttempt(url, destination, state, cancel, display_name, expected_sha256)
-    return _call_with_retries(download.run, on_retry=download.rollback)
+    return _call_with_retries(
+        download.run,
+        on_retry=download.rollback,
+        retryable_errors=(*_RETRYABLE_NETWORK_ERRORS, ModelIntegrityError),
+        retry_delay=download.retry_delay,
+    )
 
 
 @dataclass
@@ -874,9 +927,11 @@ class _DownloadAttempt:
     display_name: str
     expected_sha256: str | None
     bytes_this_attempt: int = 0
+    response_size: int | None = None
 
     def run(self) -> str:
         self.bytes_this_attempt = 0
+        self.response_size = None
         request = urllib_request.Request(self.url, headers={"User-Agent": USER_AGENT})
         digest = hashlib.sha256()
         with urllib_request.urlopen(request, timeout=60) as response:
@@ -885,17 +940,27 @@ def run(self) -> str:
             self.state.current_file = self.display_name or self.destination.name
             with self.destination.open("wb") as output:
                 self._write_response(response, output, digest)
+        self._verify_response_size()
         actual = digest.hexdigest()
         self._verify_digest(actual)
         return actual
 
     def rollback(self) -> None:
-        self.state.downloaded_bytes -= self.bytes_this_attempt
+        self.destination.unlink(missing_ok=True)
+        self.state.downloaded_bytes = max(0, self.state.downloaded_bytes - self.bytes_this_attempt)
+
+    def retry_delay(self, error: BaseException, attempt: int) -> float:
+        """Retry a checksum mismatch immediately; back off for network failures."""
+        if isinstance(error, ModelIntegrityError):
+            return 0
+        return _RETRY_BACKOFF_SECONDS * attempt
 
     def _set_total_bytes(self, response: Any) -> None:
         length = response.headers.get("Content-Length")
-        if self.state.total_bytes is None and length and length.isdigit():
-            self.state.total_bytes = int(length)
+        if length and length.isdigit():
+            self.response_size = int(length)
+            if self.state.total_bytes is None:
+                self.state.total_bytes = self.response_size
 
     def _write_response(self, response: Any, output: Any, digest: Any) -> None:
         chunk = response.read(CHUNK_SIZE)
@@ -908,6 +973,12 @@ def _write_response(self, response: Any, output: Any, digest: Any) -> None:
             self.state.downloaded_bytes += len(chunk)
             chunk = response.read(CHUNK_SIZE)
 
+    def _verify_response_size(self) -> None:
+        if self.response_size is None or self.bytes_this_attempt == self.response_size:
+            return
+        missing = max(0, self.response_size - self.bytes_this_attempt)
+        raise http_client.IncompleteRead(b"", missing)
+
     def _verify_digest(self, actual: str) -> None:
         if self.expected_sha256 is None or actual == self.expected_sha256:
             return
@@ -926,10 +997,6 @@ def _remove_tree(path: Path) -> None:
     shutil.rmtree(path, ignore_errors=True)
 
 
-def _clear_staging_directory(path: Path) -> None:
-    _remove_tree(path)
-
-
 def _missing_model_files_message(missing: list[str]) -> str:
     names = ", ".join(missing)
     return f"Downloaded model is missing: {names}."

docs/troubleshooting.md

@@ -320,9 +320,11 @@ includes the bearer token, recordings, transcripts, or session identifiers.
 ## A model download fails SHA-256 verification
 
 The gateway pins the expected digest of every catalog model it can
-(`app/model_pins.json`) and discards a download that does not match, so the
-failure means the bytes served differ from the bytes this release was built
-against. The partial file is already deleted; nothing unverified is kept.
+(`app/model_pins.json`) and discards a download that does not match. It retries
+short responses and checksum mismatches three times against that same pin; it
+never changes or bypasses the expected digest. A reported failure therefore
+means the mismatch persisted. The partial file is already deleted; nothing
+unverified is kept.
 
 Two very different causes look identical here, so check which one it is before
 retrying:
@@ -337,10 +339,9 @@ retrying:
    ```
 
 2. **The bytes were altered in transit or at the source.** A corrupted proxy or
-   mirror, or a compromised upstream. Retry once — a corrupted transfer usually
-   will not reproduce, while an altered source will fail identically every
-   time. Do not "fix" a reproducible mismatch by refreshing the pin unless the
-   upstream commit genuinely changed.
+   mirror, or a compromised upstream. Because the gateway already retried the
+   transfer, do not keep clicking Retry. Do not "fix" a reproducible mismatch
+   by refreshing the pin unless the upstream commit genuinely changed.
 
 Never work around this by deleting the pin. The check is the only thing
 standing between a swapped model file and an ONNX/GGUF/Core ML runtime that