diff --git a/core/network_utils.py b/core/network_utils.py index 9b4f417..9fa68c3 100644 --- a/core/network_utils.py +++ b/core/network_utils.py @@ -199,3 +199,89 @@ def request_public_url( current_headers = next_headers raise UnsafeUrlError("The download URL has too many redirects") + + +def request_source_response( + url: str, + method: str = "GET", + headers: Optional[Dict[str, str]] = None, + params: Optional[Dict[str, Any]] = None, + timeout: int = 20, + max_attempts: int = 2, + log_name: str = "Source API", +) -> Optional[requests.Response]: + """Perform a GET/POST request with retry logic for rate limit status (HTTP 429).""" + import time + from .log_system import create_module_logger + log = create_module_logger(__name__) + + request_params = {k: v for k, v in (params or {}).items() if v is not None} + response = None + + for attempt in range(max_attempts): + try: + if method.upper() == "POST": + response = requests.post(url, json=request_params, headers=headers, timeout=timeout) + else: + response = requests.get(url, params=request_params, headers=headers, timeout=timeout) + except Exception as e: + log.warning(f"{log_name} request failed: url={url}, error={e}") + return None + + if response.status_code != 429 or attempt == max_attempts - 1: + break + + retry_after = response.headers.get("Retry-After") + try: + delay = float(retry_after) if retry_after else 1.2 + except (TypeError, ValueError): + delay = 1.2 + time.sleep(max(0.5, min(delay, 3.0))) + + return response + + +def request_source_json( + url: str, + method: str = "GET", + headers: Optional[Dict[str, str]] = None, + params: Optional[Dict[str, Any]] = None, + timeout: int = 20, + max_attempts: int = 2, + log_name: str = "Source API", + raise_on_error: bool = False, +) -> Optional[Dict[str, Any]]: + """Perform a GET/POST request and return JSON with retry logic for HTTP 429.""" + response = request_source_response( + url, + method=method, + headers=headers, + params=params, + timeout=timeout, + max_attempts=max_attempts, + log_name=log_name, + ) + if response is None: + if raise_on_error: + raise ValueError(f"{log_name} request failed: no response") + return None + + if response.status_code != 200: + from .log_system import create_module_logger + log = create_module_logger(__name__) + log.debug(f"{log_name} returned HTTP {response.status_code}: url={url}") + if raise_on_error: + response.raise_for_status() + return None + + try: + return response.json() + except Exception as e: + from .log_system import create_module_logger + log = create_module_logger(__name__) + log.warning(f"{log_name} JSON parse failed: url={url}, error={e}") + if raise_on_error: + raise + return None + + diff --git a/core/path_templates.py b/core/path_templates.py index f814fdb..09902f7 100644 --- a/core/path_templates.py +++ b/core/path_templates.py @@ -64,11 +64,8 @@ def _base_model_token_variants(value: Any) -> List[str]: def _split_relative_folder_segments(relative_path: Any) -> List[str]: - parts = [ - part.strip() - for part in str(relative_path or "").replace("\\", "/").split("/") - if part.strip() - ] + from .path_utils import split_path_segments + parts = split_path_segments(relative_path) return parts[:-1] if len(parts) > 1 else [] diff --git a/core/path_utils.py b/core/path_utils.py index 178f9ac..c0cb9f6 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -1079,3 +1079,22 @@ def find_metadata_sidecar_path(model_path: str) -> str: return path return "" + + +def split_path_segments(path_value: Any, filter_dots: bool = True) -> list[str]: + """Split path into standard segments, replacing backslashes and filtering empty segments.""" + text = str(path_value or "").replace("\\", "/").strip() + if not text: + return [] + + parts = [] + for raw_part in text.split("/"): + part = raw_part.strip() + if not part: + continue + if filter_dots and part in {".", ".."}: + continue + parts.append(part) + + return parts + diff --git a/core/settings.py b/core/settings.py index abd21f2..e0429b1 100644 --- a/core/settings.py +++ b/core/settings.py @@ -85,28 +85,14 @@ def sanitize_folder_name(value: Any, fallback: str = "") -> str: def normalize_relative_subfolder(path_value: Any) -> str: - text = str(path_value or "").replace("\\", "/").strip() - if not text: - return "" - parts = [] - for raw_part in text.split("/"): - part = sanitize_folder_name(raw_part) - if part and part not in {".", ".."}: - parts.append(part) - return "/".join(parts) + from .path_utils import split_path_segments + parts = [sanitize_folder_name(part) for part in split_path_segments(path_value)] + return "/".join([part for part in parts if part]) def normalize_download_path_template(template: Any) -> str: - text = str(template or "").replace("\\", "/").strip() - if not text: - return "" - parts = [] - for raw_part in text.split("/"): - part = raw_part.strip() - if not part or part in {".", ".."}: - continue - parts.append(part) - return "/".join(parts) + from .path_utils import split_path_segments + return "/".join(split_path_segments(template)) def _read_settings_file() -> Dict[str, Any]: diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 123a869..d8f7b7e 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -43,6 +43,7 @@ from ..type_utils import ( extract_file_size, normalize_model_file_info, build_search_result, + normalize_sha256, ) from ..progress import report_progress, get_progress_reporter from ..log_system import create_module_logger @@ -125,43 +126,16 @@ def _request_json( params: Optional[Dict[str, Any]] = None, timeout: int = 20, ) -> Optional[Dict[str, Any]]: + from ..network_utils import request_source_json url = f"{CIVARCHIVE_API_URL}{path}" - response = None - request_params = {k: v for k, v in (params or {}).items() if v is not None} - for attempt in range(2): - try: - response = requests.get( - url, - params=request_params, - headers=REQUEST_HEADERS, - timeout=timeout, - ) - except Exception as e: - log.warning(f"CivArchive API request failed: path={path}, error={e}") - return None + return request_source_json( + url, + params=params, + headers=REQUEST_HEADERS, + timeout=timeout, + log_name="CivArchive API" + ) - if response.status_code != 429 or attempt == 1: - break - - retry_after = response.headers.get("Retry-After") - try: - delay = float(retry_after) if retry_after else 1.2 - except (TypeError, ValueError): - delay = 1.2 - time.sleep(max(0.5, min(delay, 3.0))) - - if response is None or response.status_code != 200: - status = response.status_code if response is not None else "no-response" - log.debug( - f"CivArchive API returned {status}: path={path}, params={params}" - ) - return None - - try: - return response.json() - except Exception as e: - log.warning(f"CivArchive API JSON parse failed: path={path}, error={e}") - return None def _normalize_embedded_html_text(html_text: str) -> str: @@ -478,30 +452,24 @@ def _search_page( def parse_civarchive_url(url: str) -> Optional[Dict[str, Any]]: """Parse CivArchive URLs into sha256 or model/version identifiers.""" + from ..type_utils import parse_provider_model_url if not url: return None - - parsed = urlparse(urljoin(CIVARCHIVE_BASE_URL, url)) - if parsed.hostname and not host_matches_domain(parsed.hostname, "civarchive.com"): - return None - - sha_match = re.search(r"/sha256/([a-fA-F0-9]{64})", parsed.path) - if sha_match: - return {"sha256": sha_match.group(1).lower()} - - return parse_civitai_model_path(parsed.path, parsed.query) + full_url = urljoin(CIVARCHIVE_BASE_URL, url) + return parse_provider_model_url(full_url, ["civarchive.com"]) def _extract_sha256(value: str) -> Optional[str]: if not value: return None - direct = re.fullmatch(r"[a-fA-F0-9]{64}", value.strip()) - if direct: - return direct.group(0).lower() + normalized = normalize_sha256(value) + if normalized: + return normalized parsed = parse_civarchive_url(value) return parsed.get("sha256") if parsed else None + def _extract_model_context( payload: Dict[str, Any], ) -> Tuple[Dict[str, Any], Dict[str, Any], List[Dict[str, Any]]]: diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 0fb7535..2ec7af6 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -9,6 +9,7 @@ import re import json import hashlib import requests +from ..network_utils import request_source_response, request_source_json from typing import Dict, Any, Optional, List, Callable from urllib.parse import urlparse, parse_qs, quote, urlencode @@ -445,10 +446,9 @@ def _search_civitai_trpc_candidates( f"CivitAI tRPC search start: filename={filename}, model_type={model_type}, type_filter={type_filter or 'none'}, session_token={'yes' if session_token else 'no'}, url={url}" ) - try: - response = requests.get(url, headers=headers, timeout=timeout) - except Exception as e: - log.warning(f"CivitAI tRPC request failed for filename={filename}: {e}") + response = request_source_response(url, headers=headers, timeout=timeout, log_name="CivitAI tRPC") + if response is None: + log.warning(f"CivitAI tRPC request failed for filename={filename}: no response") return [] text_preview = response.text[:800].replace("\n", " ").replace("\r", " ") @@ -497,12 +497,13 @@ def _search_civitai_red_candidates( headers["Cookie"] = build_civitai_session_cookie(session_token) log.info(f"CivitAI.red search start: filename={filename}, url={search_url}") - response = requests.get(search_url, headers=headers, timeout=timeout) - if response.status_code != 200: - if response.status_code in {401, 403}: + response = request_source_response(search_url, headers=headers, timeout=timeout, log_name="CivitAI.red") + if response is None or response.status_code != 200: + if response is not None and response.status_code in {401, 403}: return [] + status = response.status_code if response is not None else "no-response" log.warning( - f"CivitAI.red search returned {response.status_code} for filename={filename}" + f"CivitAI.red search returned {status} for filename={filename}" ) return [] @@ -582,15 +583,15 @@ def _search_civitai_public_api_candidates( f"CivitAI public API search start: filename={filename}, model_type={model_type}, api_key={'yes' if api_key else 'no'}, session_token={'yes' if session_token else 'no'}" ) - try: - response = requests.get( - f"{CIVITAI_API_URL}/models", - params=params, - headers=headers, - timeout=timeout, - ) - except Exception as e: - log.warning(f"CivitAI public API request failed for filename={filename}: {e}") + response = request_source_response( + f"{CIVITAI_API_URL}/models", + params=params, + headers=headers, + timeout=timeout, + log_name="CivitAI public API" + ) + if response is None: + log.warning(f"CivitAI public API request failed for filename={filename}: no response") return [] text_preview = response.text[:800].replace("\n", " ").replace("\r", " ") @@ -686,19 +687,15 @@ def _find_civitai_file_in_model( best_resolved_result = None best_resolved_confidence = 0.0 - headers = {} if api_key: headers["Authorization"] = f"Bearer {api_key}" - response = requests.get( - f"{CIVITAI_API_URL}/models/{model_id}", headers=headers, timeout=15 + data = request_source_json( + f"{CIVITAI_API_URL}/models/{model_id}", headers=headers, timeout=15, log_name="CivitAI model lookup" ) - if response.status_code != 200: - log.warning(f"CivitAI model lookup returned {response.status_code} for model_id={model_id}") + if data is None: return None - - data = response.json() versions = data.get("modelVersions", []) if allow_model_title_match: @@ -796,18 +793,8 @@ def parse_civitai_url(url: str) -> Optional[Dict[str, Any]]: """ Parse a CivitAI URL to extract model/version info. """ - if not isinstance(url, str) or not url.strip(): - return None - parsed = urlparse(url) - if not _is_civitai_host(parsed.hostname): - return None - - if "/api/download/models/" in parsed.path: - match = re.search(r"/api/download/models/(\d+)", parsed.path) - if match: - return {"version_id": int(match.group(1))} - - return parse_civitai_model_path(parsed.path, parsed.query) + from ..type_utils import parse_provider_model_url + return parse_provider_model_url(url, ["civitai.com", "civitai.red"]) def get_civitai_download_url(version_id: int, api_key: Optional[str] = None) -> str: @@ -1085,12 +1072,11 @@ def search_civitai( if api_key: headers["Authorization"] = f"Bearer {api_key}" - response = requests.get( - f"{CIVITAI_API_URL}/models", params=params, headers=headers, timeout=15 + data = request_source_json( + f"{CIVITAI_API_URL}/models", params=params, headers=headers, timeout=15, log_name="CivitAI search models" ) - if response.status_code == 200: - data = response.json() + if data: for model in data.get("items", []): model_id = model.get("id") @@ -1174,10 +1160,11 @@ def get_civitai_model_details( if api_key: headers["Authorization"] = f"Bearer {api_key}" - response = requests.get( + response = request_source_response( f"{CIVITAI_API_URL}/models/{model_id}", headers=headers, timeout=20, + log_name="CivitAI model details" ) if response.status_code != 200: log.warning(f"CivitAI model details returned {response.status_code}: model_id={model_id}") @@ -1332,7 +1319,7 @@ def resolve_urn( url = f"{CIVITAI_API_URL}/models/{model_id}" params = {"modelVersionId": version_id} - response = requests.get(url, headers=headers, params=params, timeout=15) + response = request_source_response(url, headers=headers, params=params, timeout=15, log_name="CivitAI URN resolve") if response.status_code != 200: log.warning(f"CivitAI URN resolve failed: {response.status_code}") @@ -1443,7 +1430,7 @@ def get_model_info_by_hash( if api_key: headers["Authorization"] = f"Bearer {api_key}" - response = requests.get(api_url, headers=headers, timeout=15) + response = request_source_response(api_url, headers=headers, timeout=15, log_name="CivitAI by-hash") if response.status_code == 200: data = response.json() diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 0d1e019..27d51e0 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -18,7 +18,7 @@ from ..log_system import create_module_logger log = create_module_logger(__name__) from ..path_utils import write_json_atomic, METADATA_DIR, read_json_safe -from ..network_utils import host_matches_domain +from ..network_utils import host_matches_domain, request_source_response, request_source_json from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache, extract_file_size from ..matcher import build_filename_search_queries, clean_filename_for_search @@ -188,11 +188,12 @@ def _build_author_index_from_models( def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200): try: - response = requests.get( + response = request_source_response( f"{HF_API_URL}/models", params={"author": author, "limit": limit, "full": "true"}, headers=headers, timeout=15, + log_name="HuggingFace author index" ) if response.status_code != 200: log.debug( @@ -389,11 +390,12 @@ def _get_repo_tree( files_url = f"{HF_API_URL}/models/{repo_id}/tree/{quote(branch, safe='')}" try: - response = requests.get( + response = request_source_response( files_url, params={"recursive": "1"}, headers=headers, timeout=15, + log_name="HuggingFace tree" ) if response.status_code != 200: log.debug( @@ -528,11 +530,12 @@ def _get_repos_by_author( author: str, headers: Dict[str, str], limit: int = 200 ) -> List[str]: try: - response = requests.get( + response = request_source_response( f"{HF_API_URL}/models", params={"author": author, "limit": limit}, headers=headers, timeout=15, + log_name="HuggingFace repos by author" ) if response.status_code != 200: log.debug( @@ -569,8 +572,8 @@ def _search_brave_for_huggingface_candidates( seen = set() try: - response = requests.get( - BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=15 + response = request_source_response( + BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=15, log_name="Brave search" ) if response.status_code != 200: log.warning( @@ -689,7 +692,7 @@ def search_huggingface_for_file( query_count=total_queries, ) search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20" - response = requests.get(search_url, headers=headers, timeout=10) + response = request_source_response(search_url, headers=headers, timeout=10, log_name="HuggingFace models search") if response.status_code != 200: log.warning( f"HuggingFace API status={response.status_code} query={search_query}" @@ -969,12 +972,12 @@ def search_huggingface( params = {"search": query, "limit": limit, "full": "true"} - response = requests.get( - f"{HF_API_URL}/models", params=params, headers=headers, timeout=15 + data = request_source_json( + f"{HF_API_URL}/models", params=params, headers=headers, timeout=15, log_name="HuggingFace general search" ) - if response.status_code == 200: - models = response.json() + if data: + models = data for model in models: repo_id = model.get("id", "") diff --git a/core/sources/model_list.py b/core/sources/model_list.py index 080f90b..794c37d 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -27,6 +27,7 @@ MODEL_LIST_SOURCE_URL = ( MODEL_LIST_GITHUB_API_URL = ( "https://api.github.com/repos/Comfy-Org/ComfyUI-Manager/contents/model-list.json?ref=main" ) +from ..network_utils import request_source_json HTTP_TIMEOUT = 30 # Cache for loaded data @@ -42,9 +43,14 @@ def _fetch_json_url(url: str) -> Dict[str, Any]: "Accept": "application/json", "User-Agent": "comfyui-model-resolver", } - response = requests.get(url, headers=headers, timeout=HTTP_TIMEOUT) - response.raise_for_status() - return response.json() + return request_source_json( + url, + headers=headers, + timeout=HTTP_TIMEOUT, + log_name="ModelList", + raise_on_error=True, + ) + diff --git a/core/type_utils.py b/core/type_utils.py index 1caee5e..d647017 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -1072,3 +1072,38 @@ def build_search_result( } result.update(extra) return result + + +def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[Dict[str, Any]]: + """Parse a provider model URL (e.g. CivitAI or mirror site) to extract model/version info.""" + from urllib.parse import urlparse + import re + + if not isinstance(url, str) or not url.strip(): + return None + + parsed = urlparse(url) + hostname = parsed.hostname + if hostname: + hostname = hostname.lower().strip(".") + matched = False + for domain in allowed_domains: + if hostname == domain or hostname.endswith("." + domain): + matched = True + break + if not matched: + return None + + # Standard CivArchive sha256 mirrors pattern + sha_match = re.search(r"/sha256/([a-fA-F0-9]{64})", parsed.path) + if sha_match: + return {"sha256": sha_match.group(1).lower()} + + # Standard CivitAI download pattern + if "/api/download/models/" in parsed.path: + match = re.search(r"/api/download/models/(\d+)", parsed.path) + if match: + return {"version_id": int(match.group(1))} + + return parse_civitai_model_path(parsed.path, parsed.query) + diff --git a/requirements.txt b/requirements.txt index 8b13721..069b09a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ requests aiohttp rapidfuzz +huggingface-hub \ No newline at end of file diff --git a/tests/test_unified_helpers.py b/tests/test_unified_helpers.py index 1d67272..35d18e8 100644 --- a/tests/test_unified_helpers.py +++ b/tests/test_unified_helpers.py @@ -1128,3 +1128,64 @@ class UnifiedHelpersTests(unittest.TestCase): # Verify new data is written self.assertEqual(read_json_safe(data_file), new_data) self.assertEqual(read_json_safe(meta_file), new_meta) + + @patch("requests.get") + def test_request_source_json_success(self, mock_get): + from core.network_utils import request_source_json + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"success": True} + mock_get.return_value = mock_response + + res = request_source_json("https://example.com/api") + self.assertEqual(res, {"success": True}) + + @patch("requests.get") + @patch("time.sleep") + def test_request_source_json_429_retry(self, mock_sleep, mock_get): + from core.network_utils import request_source_json + mock_response_429 = MagicMock() + mock_response_429.status_code = 429 + mock_response_429.headers = {"Retry-After": "1.5"} + + mock_response_200 = MagicMock() + mock_response_200.status_code = 200 + mock_response_200.json.return_value = {"ok": True} + + mock_get.side_effect = [mock_response_429, mock_response_200] + + res = request_source_json("https://example.com/api", max_attempts=2) + self.assertEqual(res, {"ok": True}) + mock_sleep.assert_called_once_with(1.5) + + def test_split_path_segments(self): + from core.path_utils import split_path_segments + self.assertEqual(split_path_segments("a/b/c"), ["a", "b", "c"]) + self.assertEqual(split_path_segments("a\\b\\c"), ["a", "b", "c"]) + self.assertEqual(split_path_segments("a/./b/../c"), ["a", "b", "c"]) + self.assertEqual(split_path_segments("a/./b/../c", filter_dots=False), ["a", ".", "b", "..", "c"]) + self.assertEqual(split_path_segments(""), []) + self.assertEqual(split_path_segments(None), []) + + def test_parse_provider_model_url(self): + from core.type_utils import parse_provider_model_url + allowed = ["civitai.com", "civitai.red"] + + # Valid url download format + res = parse_provider_model_url("https://civitai.com/api/download/models/12345", allowed) + self.assertEqual(res, {"version_id": 12345}) + + # Valid model page format + res = parse_provider_model_url("https://civitai.red/models/9876?modelVersionId=54321", allowed) + self.assertEqual(res, {"model_id": 9876, "version_id": 54321}) + + # Non-matching hostname + res = parse_provider_model_url("https://example.com/api/download/models/12345", allowed) + self.assertIsNone(res) + + # Mirror sha256 path pattern + res = parse_provider_model_url("https://civarchive.com/sha256/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2", ["civarchive.com"]) + self.assertEqual(res, {"sha256": "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2"}) + + + diff --git a/web/resolver/utils/html_utils.js b/web/resolver/utils/html_utils.js index 11c513c..c9749c7 100644 --- a/web/resolver/utils/html_utils.js +++ b/web/resolver/utils/html_utils.js @@ -119,3 +119,43 @@ export function sanitizeDescriptionHtml(html) { Array.from(root.childNodes).forEach(child => sanitizeNode(child)); return root.innerHTML; } + +/** + * Generic background task progress poller with token cancellation checks. + */ +export async function pollBackgroundTask({ + endpoint, + tokenCheck, + onProgress, + isTerminal, + onTerminal, + onError, + intervalMs = 250, + fetchJson, + filterIgnoredStatus +}) { + while (tokenCheck()) { + try { + const data = await fetchJson(endpoint, { silent: true }, 'Poll background task'); + + if (filterIgnoredStatus && filterIgnoredStatus(data)) { + await new Promise(resolve => setTimeout(resolve, intervalMs)); + continue; + } + + if (tokenCheck()) { + onProgress(data); + if (isTerminal(data)) { + onTerminal(data); + return; + } + } + } catch (error) { + if (tokenCheck()) { + onError(error); + return; + } + } + await new Promise(resolve => setTimeout(resolve, intervalMs)); + } +} diff --git a/web/resolver/utils/render_format_methods.js b/web/resolver/utils/render_format_methods.js index 84af6ab..323890c 100644 --- a/web/resolver/utils/render_format_methods.js +++ b/web/resolver/utils/render_format_methods.js @@ -2,7 +2,7 @@ import { app } from "../../../../../scripts/app.js"; import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; -import { escapeHtml, escapeJsString } from "./html_utils.js"; +import { escapeHtml, escapeJsString, pollBackgroundTask } from "./html_utils.js"; export const renderFormatMethods = { encodeContextMenuModel(context = null) { @@ -248,69 +248,54 @@ export const renderFormatMethods = { }, async pollAnalysisProgress(analysisId, token) { - while (this._analysisProgressToken === token) { - try { - const progress = await this.fetchJson(`/model_resolver/analyze-progress/${analysisId}`, { - silent: true - }, 'Poll analysis progress'); - if ( - progress && - this.contentElement && - this._analysisProgressToken === token && - this.activeTab === 'missing' - ) { + await pollBackgroundTask({ + endpoint: `/model_resolver/analyze-progress/${analysisId}`, + tokenCheck: () => this._analysisProgressToken === token, + onProgress: (progress) => { + if (this.contentElement && this.activeTab === 'missing') { this.contentElement.innerHTML = this.renderAnalysisProgress(progress); - if (progress.status === 'completed' || progress.status === 'error') { - if (this._analysisProgressToken === token) { - this._analysisProgressToken = null; - } - return; - } } - } catch (error) { + }, + isTerminal: (progress) => progress.status === 'completed' || progress.status === 'error', + onTerminal: () => { + if (this._analysisProgressToken === token) { + this._analysisProgressToken = null; + } + }, + onError: (error) => { console.warn('Model Resolver: analysis progress polling failed', error); - } - - await new Promise(resolve => setTimeout(resolve, 250)); - } + }, + intervalMs: 250, + fetchJson: this.fetchJson.bind(this) + }); }, async pollLoadedModelsProgress(loadedId, token) { - while (this._loadedModelsProgressToken === token) { - try { - const progress = await this.fetchJson(`/model_resolver/loaded-progress/${encodeURIComponent(loadedId)}`, { - silent: true - }, 'Poll loaded models progress'); - if (progress?.status === 'unknown') { - await new Promise(resolve => setTimeout(resolve, 250)); - continue; - } - if ( - progress && - this.contentElement && - this._loadedModelsProgressToken === token && - this.activeTab === 'loaded' - ) { + await pollBackgroundTask({ + endpoint: `/model_resolver/loaded-progress/${encodeURIComponent(loadedId)}`, + tokenCheck: () => this._loadedModelsProgressToken === token, + filterIgnoredStatus: (progress) => progress?.status === 'unknown', + onProgress: (progress) => { + if (this.contentElement && this.activeTab === 'loaded') { this.contentElement.innerHTML = this.renderLoadedModelsProgress(progress); - if (progress.status === 'completed' || progress.status === 'error') { - if (this._loadedModelsProgressToken === token) { - this._loadedModelsProgressToken = null; - } - return; - } } - } catch (error) { + }, + isTerminal: (progress) => progress.status === 'completed' || progress.status === 'error', + onTerminal: () => { + if (this._loadedModelsProgressToken === token) { + this._loadedModelsProgressToken = null; + } + }, + onError: (error) => { console.warn('Model Resolver: loaded models progress polling failed', error); - } - - await new Promise(resolve => setTimeout(resolve, 250)); - } + }, + intervalMs: 250, + fetchJson: this.fetchJson.bind(this) + }); }, /** * Get display information for the best local match of a missing model - * @param {object} missing - The missing model object - * @returns {object} Match display metadata */ getLocalMatchDisplayInfo(missing) { const bestMatch = this.getBestLocalMatch?.(missing, 70) || null; diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index e082143..9b334c0 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -4,6 +4,7 @@ import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; import { LOG_LEVEL as DEFAULT_FRONTEND_LOG_LEVEL } from "../../log_system/config.js"; import { logger as frontendLogger } from "../../log_system/logger.js"; +import { pollBackgroundTask } from "../utils/html_utils.js"; const SETTINGS_MAP = [ { serverKey: 'civitai_key', localKey: 'ModelResolver.civitaiApiKey', type: 'string', default: '' }, { serverKey: 'civitai_session_token', localKey: 'ModelResolver.civitaiSessionToken', type: 'string', default: '' }, @@ -2067,34 +2068,31 @@ export const optionsMethods = { const pollMetadataBuildProgress = async () => { const progressId = this.metadataBuildProgressId; if (!progressId) return; - try { - const data = await this.fetchJson( - `/model_resolver/metadata-build/progress/${encodeURIComponent(progressId)}`, - { silent: true }, - 'Metadata build progress' - ); - appendMetadataBuildHistory(data); - updateMetadataBuildSummary(data); - updateMetadataBuildProgressView(data); - renderMetadataBuildResults(data); - if (isMetadataBuildTerminal(data)) { - stopMetadataBuildPolling(); + stopMetadataBuildPolling(); + await pollBackgroundTask({ + endpoint: `/model_resolver/metadata-build/progress/${encodeURIComponent(progressId)}`, + tokenCheck: () => this.metadataBuildProgressId === progressId, + onProgress: (data) => { + appendMetadataBuildHistory(data); + updateMetadataBuildSummary(data); + updateMetadataBuildProgressView(data); + renderMetadataBuildResults(data); + }, + isTerminal: (data) => isMetadataBuildTerminal(data), + onTerminal: (data) => { finishMetadataBuild(data); if (this.metadataBuildProgressId === progressId) { this.metadataBuildProgressId = ''; } - return; - } - stopMetadataBuildPolling(); - this.metadataBuildPollTimer = window.setTimeout(() => { - pollMetadataBuildProgress(); - }, 700); - } catch (error) { - stopMetadataBuildPolling(); - this.metadataBuildProgressId = ''; - setMetadataBuildBusy(false); - setMetadataBuildStatus(error.message || 'Metadata build progress failed.', 'is-invalid'); - } + }, + onError: (error) => { + this.metadataBuildProgressId = ''; + setMetadataBuildBusy(false); + setMetadataBuildStatus(error.message || 'Metadata build progress failed.', 'is-invalid'); + }, + intervalMs: 700, + fetchJson: this.fetchJson.bind(this) + }); }; const runMetadataBuild = async () => {