feat(civarchive): prioritize known mirrors for remote size probes

core/sources/civarchive.py

@@ -9,7 +9,7 @@ import json
 import re
 import time
 from typing import Any, Callable, Dict, List, Optional, Tuple
-from urllib.parse import urljoin
+from urllib.parse import urljoin, urlparse
 
 from ..log_system import create_module_logger
 from ..matcher import (
@@ -27,7 +27,7 @@ from ..matcher import (
 from ..matcher import (
     normalize_base_model as _normalize_base_model,
 )
-from ..network_utils import request_source_response
+from ..network_utils import host_matches_domain, request_source_response
 from ..path_utils import get_filename_from_path
 from ..progress import get_progress_reporter
 from ..type_utils import (
@@ -64,6 +64,22 @@ DEFAULT_CIVARCHIVE_CANDIDATE_LIMIT = 10
 MAX_CIVARCHIVE_CANDIDATE_LIMIT = 30
 SEARCH_RESULT_DETAIL_LIMIT = 5
 HASH_PAGE_MODEL_LINK_LIMIT = 5
+CIVARCHIVE_SIZE_PROBE_DOMAINS = (
+    "huggingface.co",
+    "civarchive.com",
+    "civitai.com",
+    "civitai.red",
+    "modelscope.ai",
+    "modelscope.cn",
+)
+CIVARCHIVE_SIZE_PROBE_PRIORITY = (
+    "huggingface.co",
+    "civarchive.com",
+    "civitai.com",
+    "civitai.red",
+    "modelscope.ai",
+    "modelscope.cn",
+)
 
 
 
@@ -540,12 +556,27 @@ def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]:
     normalized = _normalize_download_url(url)
     if not normalized:
         return None
-    probe_url = prepare_remote_size_probe_url(normalized, ["civarchive.com", "civitai.com", "civitai.red"])
+    probe_url = prepare_remote_size_probe_url(
+        normalized,
+        list(CIVARCHIVE_SIZE_PROBE_DOMAINS),
+    )
     if not probe_url:
         return None
     return fetch_remote_file_size_cached(probe_url, headers=REQUEST_HEADERS, timeout=timeout)
 
 
+def _remote_size_probe_priority(url: Any) -> int:
+    try:
+        host = urlparse(str(url or "")).hostname
+    except ValueError:
+        return len(CIVARCHIVE_SIZE_PROBE_PRIORITY)
+
+    for priority, domain in enumerate(CIVARCHIVE_SIZE_PROBE_PRIORITY):
+        if host_matches_domain(host, domain):
+            return priority
+    return len(CIVARCHIVE_SIZE_PROBE_PRIORITY)
+
+
 
 def _resolve_file_size_bytes(
     file_info: Dict[str, Any],
@@ -556,7 +587,7 @@ def _resolve_file_size_bytes(
         return size
 
     urls = download_urls if download_urls is not None else _collect_normalized_download_urls(file_info)
-    for url in urls:
+    for url in sorted(urls, key=_remote_size_probe_priority):
         size = _fetch_remote_file_size_bytes(url)
         if size:
             return size

tests/test_source_internals.py

@@ -10,8 +10,8 @@ extract_trained_words), or where a refactored unified function changes
 its output contract.
 """
 
-import os
 import json
+import os
 import tempfile
 import unittest
 from unittest.mock import MagicMock, patch
@@ -22,38 +22,39 @@ import requests
 # civarchive internals
 # ---------------------------------------------------------------------------
 from core.sources.civarchive import (
-    _normalize_download_url,
     _archive_link_is_dead,
-    _download_url_looks_like_model_file,
     _collect_download_urls,
-    _collect_normalized_download_urls,
     _collect_download_urls_unified,
+    _collect_normalized_download_urls,
+    _fetch_remote_file_size_bytes,
     _normalize_archive_version,
+    _normalize_download_url,
+    _resolve_file_size_bytes,
     parse_civarchive_url,
 )
-from core.type_utils import prepare_remote_size_probe_url
 
 # ---------------------------------------------------------------------------
 # civitai internals
 # ---------------------------------------------------------------------------
 from core.sources.civitai import (
-    clear_search_cache,
-    build_civitai_session_cookie,
-    parse_civitai_url,
-    _normalize_civitai_file,
-    _extract_model_images,
     _build_civitai_result_from_version,
     _enrich_model_info_with_details,
-    _find_matching_file_in_versions,
+    _extract_model_images,
     _extract_public_api_model_candidates,
     _extract_trpc_model_candidates,
+    _find_matching_file_in_versions,
+    _normalize_civitai_file,
     _search_civitai_public_api_candidates,
     _search_civitai_red_candidates,
     _search_civitai_trpc_candidates,
+    build_civitai_session_cookie,
+    clear_search_cache,
     get_model_info_by_hash,
     get_model_info_for_file,
+    parse_civitai_url,
     search_civitai_for_file,
 )
+from core.type_utils import prepare_remote_size_probe_url
 
 # ---------------------------------------------------------------------------
 # workflow_updater
@@ -535,6 +536,60 @@ class PrepareSizeProbeUrlTests(unittest.TestCase):
         self.assertIsNone(self._prepare_size_probe_url(None))
 
 
+class CivArchiveMirrorSizeProbeTests(unittest.TestCase):
+
+    @patch(
+        "core.sources.civarchive.fetch_remote_file_size_cached",
+        return_value=364853140,
+    )
+    def test_huggingface_mirror_is_allowed(self, mock_fetch):
+        url = (
+            "https://huggingface.co/unsloth/LTX-2.3-GGUF/resolve/main/vae/"
+            "ltx-2.3-22b-distilled_audio_vae.safetensors"
+        )
+
+        self.assertEqual(_fetch_remote_file_size_bytes(url), 364853140)
+        mock_fetch.assert_called_once()
+        self.assertEqual(mock_fetch.call_args.args[0], url)
+
+    @patch(
+        "core.sources.civarchive.fetch_remote_file_size_cached",
+        return_value=364853140,
+    )
+    def test_modelscope_mirror_is_allowed(self, mock_fetch):
+        url = (
+            "https://www.modelscope.ai/unsloth/LTX-2.3-GGUF/resolve/master/vae/"
+            "ltx-2.3-22b-distilled_audio_vae.safetensors"
+        )
+
+        self.assertEqual(_fetch_remote_file_size_bytes(url), 364853140)
+        mock_fetch.assert_called_once()
+        self.assertEqual(mock_fetch.call_args.args[0], url)
+
+    @patch("core.sources.civarchive._fetch_remote_file_size_bytes")
+    def test_huggingface_is_probed_before_modelscope(self, mock_fetch):
+        modelscope_url = "https://www.modelscope.ai/org/repo/resolve/master/model.safetensors"
+        huggingface_url = "https://huggingface.co/org/repo/resolve/main/model.safetensors"
+        mock_fetch.side_effect = lambda url: 123456 if url == huggingface_url else None
+
+        size = _resolve_file_size_bytes(
+            {},
+            [modelscope_url, huggingface_url],
+        )
+
+        self.assertEqual(size, 123456)
+        mock_fetch.assert_called_once_with(huggingface_url)
+
+    @patch("core.sources.civarchive.fetch_remote_file_size_cached")
+    def test_untrusted_mirror_is_not_probed(self, mock_fetch):
+        self.assertIsNone(
+            _fetch_remote_file_size_bytes(
+                "https://example.com/models/model.safetensors"
+            )
+        )
+        mock_fetch.assert_not_called()
+
+
 
 # ===========================================================================
 # civitai - parse_civitai_url