diff --git a/core/download/api.py b/core/download/api.py index 2b390d4..5056ff3 100644 --- a/core/download/api.py +++ b/core/download/api.py @@ -56,12 +56,15 @@ except ImportError: log = create_module_logger("core.downloader") SPEED_HISTORY_SIZE = 5 -CHUNK_SIZE = 1024 * 1024 +CHUNK_SIZE = 4 * 1024 * 1024 CLI_LOG_INTERVAL = 5 ARIA2_RPC_TIMEOUT = (2, 5) ARIA2_STATUS_RPC_RETRIES = 4 ARIA2_STATUS_RPC_RETRY_DELAY = 0.15 ARIA2_IDLE_STOP_SECONDS = 5 * 60 +MAX_ARIA2_URL_RECOVERY_ATTEMPTS = 2 +PYTHON_DOWNLOAD_MAX_RETRIES = 2 +PYTHON_DOWNLOAD_RETRY_DELAY = 0.5 # Keep validation aligned with the install location used by the aria2 installer. # ARIA2_INSTALL_ROOT resolves the custom node package first, so paths reached # through a shared ComfyUI/custom_nodes symlink remain valid while external @@ -95,10 +98,14 @@ context.ARIA2_RPC_TIMEOUT = ARIA2_RPC_TIMEOUT context.ARIA2_STATUS_RPC_RETRIES = ARIA2_STATUS_RPC_RETRIES context.ARIA2_STATUS_RPC_RETRY_DELAY = ARIA2_STATUS_RPC_RETRY_DELAY context.ARIA2_IDLE_STOP_SECONDS = ARIA2_IDLE_STOP_SECONDS +context.MAX_ARIA2_URL_RECOVERY_ATTEMPTS = MAX_ARIA2_URL_RECOVERY_ATTEMPTS +context.PYTHON_DOWNLOAD_MAX_RETRIES = PYTHON_DOWNLOAD_MAX_RETRIES +context.PYTHON_DOWNLOAD_RETRY_DELAY = PYTHON_DOWNLOAD_RETRY_DELAY context.MANAGED_ARIA2_ROOT = MANAGED_ARIA2_ROOT context.HF_XET_ARIA2_AUTH_HOSTS = HF_XET_ARIA2_AUTH_HOSTS context.aria2_process = None +context.aria2_recent_stderr = deque(maxlen=50) context.aria2_rpc_url = "" context.aria2_rpc_secret = "" context.aria2_rpc_lock = threading.Lock() @@ -247,6 +254,16 @@ context._aria2_action_error_is_ok = aria2_backend.aria2_action_error_is_ok context._is_aria2_missing_control_file_error = ( aria2_backend.is_aria2_missing_control_file_error ) +context._is_aria2_important_stderr = aria2_backend.is_aria2_important_stderr +context.is_aria2_important_stderr = aria2_backend.is_aria2_important_stderr +context._is_aria2_recoverable_url_error = ( + aria2_backend.is_aria2_recoverable_url_error +) +context.is_aria2_recoverable_url_error = ( + aria2_backend.is_aria2_recoverable_url_error +) +context._drain_aria2_stderr = _bind(aria2_backend.drain_aria2_stderr) +context.drain_aria2_stderr = context._drain_aria2_stderr context.Aria2Error = aria2_backend.Aria2Error context._find_free_port = aria2_backend.find_free_port context._parse_aria2_int = aria2_backend.parse_aria2_int diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index 4299272..a53b54c 100644 --- a/core/download/aria2_backend.py +++ b/core/download/aria2_backend.py @@ -87,6 +87,84 @@ def parse_aria2_int(value: Any) -> int: return 0 +ARIA2_IMPORTANT_STDERR_MARKERS = ( + "write disk cache flush failure", + "error when trying to flush write cache", + "failed to write into the file", + "failed to open the file", + "failed to seek the file", + "no space left on device", + "not enough space on the disk", + "input/output error", + "permission denied", + "access is denied", + "disk quota exceeded", +) + +MAX_ARIA2_URL_RECOVERY_ATTEMPTS = 2 + + +def is_aria2_important_stderr(line: str) -> bool: + """Identify critical disk or permission error messages in aria2's stderr.""" + lowered = line.lower() + return any(marker in lowered for marker in ARIA2_IMPORTANT_STDERR_MARKERS) + + +def is_aria2_recoverable_url_error(error: Any) -> bool: + """Identify aria2 transfer errors that may be resolved by refreshing a signed URL.""" + message = str(error or "").lower() + return ( + "no uri available" in message + or "403 forbidden" in message + or "401 unauthorized" in message + or "410 gone" in message + or "status code 403" in message + or "status code 401" in message + or "bad or unexpected: 403" in message + or "bad or unexpected: 401" in message + ) + + +def drain_aria2_stderr( + process: Any, + recent_lines: Any = None, + *, + dependencies: Any = None, +) -> None: + """Continuously drain aria2's stderr pipe so it never saturates Windows pipe buffers.""" + if not process or not getattr(process, "stderr", None): + return + stderr = process.stderr + try: + while True: + raw_line = stderr.readline() + if not raw_line: + break + if isinstance(raw_line, bytes): + line = raw_line.decode("utf-8", errors="replace").rstrip() + else: + line = str(raw_line).rstrip() + if not line: + continue + if recent_lines is not None: + try: + recent_lines.append(line) + except Exception: + pass + if is_aria2_important_stderr(line): + log.warning(f"aria2: {line}") + else: + log.debug(f"aria2 stderr: {line}") + except Exception: + pass + finally: + try: + if stderr and not getattr(stderr, "closed", False): + stderr.close() + except Exception: + pass + + def is_aria2_missing_control_file_error(error: Any) -> bool: """Identify aria2's stale-output error caused by a missing .aria2 file.""" message = str(error or "").lower() @@ -440,6 +518,12 @@ def stop_aria2_daemon( except Exception: pass + try: + if getattr(process, "stderr", None) and not getattr(process.stderr, "closed", False): + process.stderr.close() + except Exception: + pass + facade.aria2_process = None facade.aria2_rpc_url = "" facade.aria2_rpc_secret = "" @@ -557,6 +641,15 @@ def ensure_aria2_daemon( process.terminate() except Exception: pass + if ( + process is not None + and getattr(process, "stderr", None) + and not getattr(process.stderr, "closed", False) + ): + try: + process.stderr.close() + except Exception: + pass facade.aria2_process = None facade.aria2_process_started_by_resolver = False @@ -592,6 +685,7 @@ def ensure_aria2_daemon( creationflags = getattr(facade.subprocess, "CREATE_NO_WINDOW", 0) log.info(f"Starting aria2 RPC daemon from {executable}") + facade.aria2_recent_stderr = facade.deque(maxlen=50) facade.aria2_process = facade.subprocess.Popen( command, stdout=facade.subprocess.DEVNULL, @@ -600,18 +694,22 @@ def ensure_aria2_daemon( ) facade.aria2_process_started_by_resolver = True + drain_thread = facade.threading.Thread( + target=facade._drain_aria2_stderr, + args=(facade.aria2_process, facade.aria2_recent_stderr), + daemon=True, + name="aria2-stderr-drain", + ) + drain_thread.start() + start_time = facade.time.time() last_error = "" while facade.time.time() - start_time < 10: process = facade.aria2_process if process.poll() is not None: - stderr = "" - try: - stderr = ( - process.stderr.read() if process.stderr else b"" - ).decode("utf-8", errors="replace") - except Exception: - stderr = "" + facade.time.sleep(0.05) + recent = getattr(facade, "aria2_recent_stderr", None) + stderr = "\n".join(recent) if recent else "" raise Aria2Error( "aria2 RPC process exited early with code " f"{process.returncode}: {stderr.strip()}" @@ -736,6 +834,12 @@ def download_file_with_aria2( start_time = facade.time.time() filename = facade.get_filename_from_path(dest_path) expected_sha256 = facade._extract_expected_sha256(metadata) + max_url_recovery = getattr( + facade, + "MAX_ARIA2_URL_RECOVERY_ATTEMPTS", + MAX_ARIA2_URL_RECOVERY_ATTEMPTS, + ) + url_recovery_attempts = 0 with facade.download_lock: facade.download_progress[download_id] = create_initial_progress( @@ -757,45 +861,60 @@ def download_file_with_aria2( headers, ) - options: Dict[str, Any] = { - "dir": facade.os.path.dirname(dest_path), - "out": filename, - "continue": "true", - "max-connection-per-server": "4", - "split": "4", - "min-split-size": "1M", - "allow-overwrite": "true", - "auto-file-renaming": "false", - "file-allocation": "none", - "no-want-digest-header": "true", - # Redirects were already resolved and validated above. Keeping them - # disabled prevents sensitive headers from reaching another host. - "max-redirect": "0", - } - # Checksum and check-integrity are deliberately omitted from aria2 options. - # On Windows, aria2c's single-threaded, unaccelerated software checksum - # verification blocks download completion and freezes the process at 100% - # for many minutes on 20+ GB models. Instead, verification is handled - # asynchronously in the background via hardware-accelerated Python hashlib. - user_agent = facade._get_header_value(request_headers, "User-Agent") - referer = facade._get_header_value(request_headers, "Referer") - if user_agent: - options["user-agent"] = user_agent - if referer: - options["referer"] = referer + def build_transfer_options(target_headers: Dict[str, str]) -> Dict[str, Any]: + opts: Dict[str, Any] = { + "dir": facade.os.path.dirname(dest_path), + "out": filename, + "continue": "true", + "max-connection-per-server": "4", + "split": "4", + "min-split-size": "1M", + "allow-overwrite": "true", + "auto-file-renaming": "false", + "file-allocation": "none", + "no-want-digest-header": "true", + # Redirects were already resolved and validated above. Keeping them + # disabled prevents sensitive headers from reaching another host. + "max-redirect": "0", + } + # Checksum and check-integrity are deliberately omitted from aria2 options. + # On Windows, aria2c's single-threaded, unaccelerated software checksum + # verification blocks download completion and freezes the process at 100% + # for many minutes on 20+ GB models. Instead, verification is handled + # asynchronously in the background via hardware-accelerated Python hashlib. + user_agent = facade._get_header_value(target_headers, "User-Agent") + referer = facade._get_header_value(target_headers, "Referer") + if user_agent: + opts["user-agent"] = user_agent + if referer: + opts["referer"] = referer - header_values = [ - f"{key}: {value}" - for key, value in request_headers.items() - if str(key).lower() not in {"user-agent", "referer"} - ] - if header_values: - options["header"] = header_values + header_values = [ + f"{key}: {value}" + for key, value in target_headers.items() + if str(key).lower() not in {"user-agent", "referer"} + ] + if header_values: + opts["header"] = header_values + return opts + + options = build_transfer_options(request_headers) + + def add_aria2_transfer( + target_url: str = "", + target_options: Optional[Dict[str, Any]] = None, + *, + preserve_progress: bool = False, + ) -> str: + eff_url = target_url or aria2_url + eff_options = target_options if target_options is not None else options + + if download_id in facade.cancelled_downloads: + raise Aria2Error("Download was cancelled") - def add_aria2_transfer() -> str: new_gid = facade._aria2_rpc( "aria2.addUri", - [[aria2_url], options], + [[eff_url], eff_options], ) if not isinstance(new_gid, str) or not new_gid: raise Aria2Error("aria2 did not return a download gid") @@ -805,18 +924,41 @@ def download_file_with_aria2( "gid": new_gid, "path": dest_path, } + + if download_id in facade.cancelled_downloads: + try: + facade._aria2_rpc("aria2.forceRemove", [new_gid]) + except Exception: + pass + raise Aria2Error("Download was cancelled") + with facade.download_lock: - facade.download_progress[download_id].update( - { - "aria2_gid": new_gid, - "status": "downloading", - "error": None, - "progress": 0, - "downloaded": 0, - "total_size": 0, - "speed": 0, - } + progress_dict = facade.download_progress.setdefault( + download_id, + create_initial_progress( + url=url, + path=dest_path, + filename=filename, + directory=facade.os.path.dirname(dest_path), + download_backend="aria2", + start_time=start_time, + ), ) + updates: Dict[str, Any] = { + "aria2_gid": new_gid, + "status": "downloading", + "error": None, + } + if not preserve_progress: + updates.update( + { + "progress": 0, + "downloaded": 0, + "total_size": 0, + "speed": 0, + } + ) + progress_dict.update(updates) return new_gid recovery_attempted = False @@ -1000,8 +1142,49 @@ def download_file_with_aria2( settings, ) if recovered: - gid = add_aria2_transfer() + gid = add_aria2_transfer( + aria2_url, + options, + preserve_progress=True, + ) continue + + if ( + url_recovery_attempts < max_url_recovery + and download_id not in facade.cancelled_downloads + and facade._is_aria2_recoverable_url_error(error_msg) + ): + url_recovery_attempts += 1 + log.warning( + f"aria2 transfer for {filename} failed with '{error_msg}'; " + f"refreshing signed URL and resuming partial download " + f"(attempt {url_recovery_attempts}/{max_url_recovery})" + ) + try: + facade._aria2_rpc("aria2.removeDownloadResult", [gid]) + except Exception: + pass + facade.time.sleep(1.0) + try: + facade._ensure_aria2_daemon(settings) + aria2_url, request_headers = ( + facade._resolve_download_url_for_aria2( + url, + headers, + ) + ) + options = build_transfer_options(request_headers) + gid = add_aria2_transfer( + aria2_url, + options, + preserve_progress=True, + ) + continue + except Exception as refresh_exc: + log.warning( + f"Could not refresh signed URL for {filename}: {refresh_exc}" + ) + with facade.download_lock: facade.download_progress[download_id]["status"] = "error" facade.download_progress[download_id]["error"] = error_msg @@ -1022,6 +1205,11 @@ def download_file_with_aria2( facade.time.sleep(0.5) except Exception as exc: + if download_id in facade.cancelled_downloads: + facade._delete_partial_download_files(dest_path) + facade.cancelled_downloads.discard(download_id) + result["error"] = "Download cancelled" + return result error_msg = facade._sanitize_download_error(exc) with facade.download_lock: if download_id in facade.download_progress: diff --git a/core/download/orchestrator.py b/core/download/orchestrator.py index fc13701..81af031 100644 --- a/core/download/orchestrator.py +++ b/core/download/orchestrator.py @@ -10,6 +10,29 @@ class DownloadCancelled(Exception): """The caller cancelled before the download response became available.""" +def parse_retry_after(value: Optional[str]) -> Optional[int]: + """Parse HTTP Retry-After header into integer seconds.""" + if not value or not str(value).strip(): + return None + val = str(value).strip() + try: + return max(1, int(val)) + except ValueError: + pass + try: + from datetime import datetime, timezone + from email.utils import parsedate_to_datetime + + parsed = parsedate_to_datetime(val) + now = datetime.now(timezone.utc) + if parsed.tzinfo is None: + now = datetime.now(timezone.utc).replace(tzinfo=None) + delta = (parsed - now).total_seconds() + return max(1, int(delta)) + except Exception: + return None + + def request_download_response(facade: Any, download_id: str, url: str, **kwargs: Any): """Allow cancellation while waiting for HTTP headers without opening a file.""" finished = threading.Event() @@ -240,6 +263,37 @@ def download_file( is_resumed = bool( existing_size > 0 and getattr(response, "status_code", None) == 206 ) + if not is_resumed and existing_size > 0 and final_url and final_url != url: + log.info( + "Range request was not honored across redirect (HTTP 200). " + f"Retrying Range request directly on target URL for {filename}..." + ) + try: + response.close() + except Exception: + pass + try: + retry_resp, retry_final_url, _ = request_download_response( + facade, + download_id, + final_url, + headers=request_headers, + stream=True, + timeout=(20, 60), + ) + retry_resp.raise_for_status() + if getattr(retry_resp, "status_code", None) == 206: + response = retry_resp + final_url = retry_final_url + is_resumed = True + log.info( + f"Successfully resumed download from target URL at byte {existing_size}" + ) + else: + response = retry_resp + except Exception as retry_exc: + log.debug(f"Direct Range retry on target URL failed: {retry_exc}") + if not is_resumed and existing_size > 0: log.info( f"Server returned HTTP {getattr(response, 'status_code', None)} without partial content support. " @@ -286,51 +340,135 @@ def download_file( break sha256_hasher.update(existing_chunk) + max_stream_retries = getattr(facade, "PYTHON_DOWNLOAD_MAX_RETRIES", 2) + stream_retry_delay = getattr(facade, "PYTHON_DOWNLOAD_RETRY_DELAY", 0.5) + stream_attempts = 0 + cancelled = False file_mode = "ab" if is_resumed else "wb" with open(partial_path, file_mode) as file_handle: - for chunk in response.iter_content(chunk_size=chunk_size): - if download_id in cancelled_downloads: - cancelled = True + while True: + try: + for chunk in response.iter_content(chunk_size=chunk_size): + if download_id in cancelled_downloads: + cancelled = True + break + + if chunk: + file_handle.write(chunk) + if sha256_hasher is not None: + sha256_hasher.update(chunk) + downloaded += len(chunk) + + current_time = time.time() + time_delta = current_time - last_speed_update + if time_delta >= 0.5: + bytes_delta = downloaded - last_downloaded + instant_speed = bytes_delta / time_delta if time_delta > 0 else 0 + speed_history.append(instant_speed) + smoothed_speed = sum(speed_history) / len(speed_history) if speed_history else 0 + last_speed_update = current_time + last_downloaded = downloaded + + with download_lock: + download_progress[download_id]["downloaded"] = downloaded + download_progress[download_id]["speed"] = int(smoothed_speed) + if total_size > 0: + download_progress[download_id]["progress"] = int((downloaded / total_size) * 100) + + if current_time - last_cli_log >= cli_log_interval: + last_cli_log = current_time + progress_pct = int((downloaded / total_size) * 100) if total_size > 0 else 0 + downloaded_str = facade.format_bytes(downloaded) + total_str = facade.format_bytes(total_size) if total_size > 0 else "?" + speed_str = facade.format_bytes(int(smoothed_speed)) + "/s" + log.info(f"Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}") + else: + with download_lock: + download_progress[download_id]["downloaded"] = downloaded + if total_size > 0: + download_progress[download_id]["progress"] = int((downloaded / total_size) * 100) + + if progress_callback: + progress_callback(downloaded, total_size) + + if cancelled or download_id in cancelled_downloads: + break + + # Check if stream closed prematurely + if total_size > 0 and downloaded < total_size: + raise requests.exceptions.RequestException( + f"Connection closed prematurely ({facade.format_bytes(downloaded)} / {total_size_str})" + ) + + # Finished transferring all data break - if chunk: - file_handle.write(chunk) - if sha256_hasher is not None: - sha256_hasher.update(chunk) - downloaded += len(chunk) + except (requests.exceptions.RequestException, OSError, TimeoutError) as stream_exc: + if cancelled or download_id in cancelled_downloads: + cancelled = True + break - current_time = time.time() - time_delta = current_time - last_speed_update - if time_delta >= 0.5: - bytes_delta = downloaded - last_downloaded - instant_speed = bytes_delta / time_delta if time_delta > 0 else 0 - speed_history.append(instant_speed) - smoothed_speed = sum(speed_history) / len(speed_history) if speed_history else 0 - last_speed_update = current_time - last_downloaded = downloaded + if stream_attempts < max_stream_retries: + stream_attempts += 1 + file_handle.flush() + try: + response.close() + except Exception: + pass - with download_lock: - download_progress[download_id]["downloaded"] = downloaded - download_progress[download_id]["speed"] = int(smoothed_speed) - if total_size > 0: - download_progress[download_id]["progress"] = int((downloaded / total_size) * 100) + log.warning( + f"Download stream interrupted for {filename} ({facade._sanitize_download_error(stream_exc)}); " + f"auto-reconnecting and resuming from byte {downloaded} " + f"(attempt {stream_attempts}/{max_stream_retries})..." + ) + time.sleep(stream_retry_delay) - if current_time - last_cli_log >= cli_log_interval: - last_cli_log = current_time - progress_pct = int((downloaded / total_size) * 100) if total_size > 0 else 0 - downloaded_str = facade.format_bytes(downloaded) - total_str = facade.format_bytes(total_size) if total_size > 0 else "?" - speed_str = facade.format_bytes(int(smoothed_speed)) + "/s" - log.info(f"Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}") - else: - with download_lock: - download_progress[download_id]["downloaded"] = downloaded - if total_size > 0: - download_progress[download_id]["progress"] = int((downloaded / total_size) * 100) + resume_headers = facade.build_download_headers(url, headers) + if downloaded > 0: + resume_headers["Range"] = f"bytes={downloaded}-" - if progress_callback: - progress_callback(downloaded, total_size) + reconnected = False + if final_url and final_url != url: + try: + resume_resp, resume_final_url, _ = request_download_response( + facade, + download_id, + final_url, + headers=resume_headers, + stream=True, + timeout=(20, 60), + ) + resume_resp.raise_for_status() + if getattr(resume_resp, "status_code", None) == 206: + response = resume_resp + final_url = resume_final_url + reconnected = True + except Exception as rec_exc: + log.debug(f"Direct reconnect on target URL failed: {rec_exc}") + + if not reconnected: + try: + resume_resp, resume_final_url, _ = request_download_response( + facade, + download_id, + url, + headers=resume_headers, + stream=True, + timeout=(20, 60), + ) + resume_resp.raise_for_status() + if getattr(resume_resp, "status_code", None) == 206 or downloaded == 0: + response = resume_resp + final_url = resume_final_url + reconnected = True + except Exception as rec_exc: + log.debug(f"Origin reconnect failed: {rec_exc}") + + if reconnected: + continue + + raise if cancelled or download_id in cancelled_downloads: with download_lock: @@ -457,6 +595,14 @@ def download_file( error_msg = f"Unauthorized (HTTP {status_code}): Authentication required." elif status_code == 404: error_msg = "Model not found (HTTP 404): The file may have been moved or deleted." + elif status_code == 429: + resp = getattr(exc, "response", None) + retry_hdr = resp.headers.get("Retry-After") if resp is not None else None + retry_sec = parse_retry_after(retry_hdr) if retry_hdr else None + if retry_sec: + error_msg = f"Rate limit exceeded (HTTP 429): please retry after {retry_sec}s." + else: + error_msg = "Rate limit exceeded (HTTP 429): too many requests, please try again later." with download_lock: download_progress[download_id]["status"] = "error" diff --git a/tests/test_downloader_boundaries.py b/tests/test_downloader_boundaries.py index b229a5f..c526018 100644 --- a/tests/test_downloader_boundaries.py +++ b/tests/test_downloader_boundaries.py @@ -1,19 +1,28 @@ +import io import os import tempfile import threading import unittest +from collections import deque from types import SimpleNamespace from unittest.mock import MagicMock, call, patch from core.download.api import context as downloader from core.download.aria2_backend import ( download_file_with_aria2, + drain_aria2_stderr, + is_aria2_important_stderr, + is_aria2_recoverable_url_error, pause_download, recover_aria2_missing_control_file, resume_download, start_background_sha256_verification, ) -from core.download.orchestrator import download_file, find_active_download_for_path +from core.download.orchestrator import ( + download_file, + find_active_download_for_path, + parse_retry_after, +) from core.path_utils import normalize_metadata_file_path @@ -483,6 +492,240 @@ class DownloaderBoundaryTests(unittest.TestCase): sidecar_meta = downloader.read_json_safe(sidecar_path, {}) self.assertFalse(sidecar_meta.get("sha256_verified")) + def test_is_aria2_important_stderr(self): + self.assertTrue(is_aria2_important_stderr("Error: Write disk cache flush failure")) + self.assertTrue(is_aria2_important_stderr("cause: No space left on device")) + self.assertTrue(is_aria2_important_stderr("failed to write into the file: Access is denied.")) + self.assertFalse(is_aria2_important_stderr("aria2 progress: 50% - 20 MB/s")) + self.assertFalse(is_aria2_important_stderr("Connecting to server 127.0.0.1:8000")) + + def test_drain_aria2_stderr_reads_lines_without_blocking(self): + mock_process = MagicMock() + mock_process.stderr = io.BytesIO( + b"connecting to host\n" + b"no space left on device\n" + b"closing connection\n" + ) + recent_lines = deque(maxlen=10) + drain_aria2_stderr(mock_process, recent_lines) + self.assertEqual( + list(recent_lines), + [ + "connecting to host", + "no space left on device", + "closing connection", + ], + ) + + def test_is_aria2_recoverable_url_error(self): + self.assertTrue(is_aria2_recoverable_url_error("No URI available. See --max-tries option.")) + self.assertTrue(is_aria2_recoverable_url_error("HTTP response header was bad or unexpected: 403 Forbidden")) + self.assertTrue(is_aria2_recoverable_url_error("Failed with status code 401")) + self.assertTrue(is_aria2_recoverable_url_error("Remote server returned 410 Gone")) + self.assertFalse(is_aria2_recoverable_url_error("Checksum mismatch error")) + self.assertFalse(is_aria2_recoverable_url_error("File already exists and cannot be overwritten")) + + def test_aria2_download_recovers_from_expired_signed_url(self): + with tempfile.TemporaryDirectory() as temp_dir: + dest_path = os.path.join(temp_dir, "model_resumed.safetensors") + with open(dest_path, "wb") as f: + f.write(b"partially downloaded bytes") + + mock_resolve = MagicMock(side_effect=[ + ("https://cdn.example.com/expired-url", {"User-Agent": "test-ua"}), + ("https://cdn.example.com/fresh-url", {"User-Agent": "test-ua"}), + ]) + mock_rpc = MagicMock(side_effect=[ + "gid-1", # 1st addUri + "OK", # removeDownloadResult for gid-1 + "gid-2", # 2nd addUri + ]) + mock_status = MagicMock(side_effect=[ + {"status": "error", "errorMessage": "No URI available. See --max-tries option."}, + {"status": "complete", "completedLength": "500", "totalLength": "500"}, + ]) + + with patch.object(downloader, "_ensure_aria2_daemon"), \ + patch.object(downloader, "_resolve_download_url_for_aria2", mock_resolve), \ + patch.object(downloader, "_aria2_rpc", mock_rpc), \ + patch.object(downloader, "_aria2_tell_status", mock_status), \ + patch.object(downloader, "write_model_resolver_metadata", return_value=""), \ + patch.object(downloader, "_schedule_aria2_idle_stop"), \ + patch.object(downloader, "invalidate_model_caches"): + + result = download_file_with_aria2( + "https://civitai.com/api/download/models/999", + dest_path, + "url-refresh-test", + dependencies=downloader, + ) + + self.assertTrue(result["success"]) + self.assertEqual(mock_resolve.call_count, 2) + self.assertEqual(mock_rpc.call_args_list[0].args[0], "aria2.addUri") + self.assertEqual(mock_rpc.call_args_list[0].args[1][0], ["https://cdn.example.com/expired-url"]) + self.assertEqual(mock_rpc.call_args_list[1].args, ("aria2.removeDownloadResult", ["gid-1"])) + self.assertEqual(mock_rpc.call_args_list[2].args[0], "aria2.addUri") + self.assertEqual(mock_rpc.call_args_list[2].args[1][0], ["https://cdn.example.com/fresh-url"]) + + def test_aria2_cancellation_shielding_prevents_orphaned_transfer(self): + with tempfile.TemporaryDirectory() as temp_dir: + dest_path = os.path.join(temp_dir, "cancelled.safetensors") + download_id = "test-shielded-cancel" + downloader.cancelled_downloads.add(download_id) + + mock_rpc = MagicMock() + mock_resolve = MagicMock(return_value=("https://example.com/file.safetensors", {})) + + with patch.object(downloader, "_ensure_aria2_daemon"), \ + patch.object(downloader, "_resolve_download_url_for_aria2", mock_resolve), \ + patch.object(downloader, "_aria2_rpc", mock_rpc), \ + patch.object(downloader, "_schedule_aria2_idle_stop"): + + result = download_file_with_aria2( + "https://example.com/file.safetensors", + dest_path, + download_id, + dependencies=downloader, + ) + + self.assertFalse(result["success"]) + self.assertEqual(result["error"], "Download cancelled") + self.assertEqual(mock_rpc.call_count, 0) + + def test_parse_retry_after(self): + self.assertEqual(parse_retry_after("120"), 120) + self.assertEqual(parse_retry_after("0"), 1) + self.assertIsNone(parse_retry_after("")) + self.assertIsNone(parse_retry_after(None)) + self.assertIsNone(parse_retry_after("invalid-text")) + + def test_python_download_handles_rate_limit_429_with_retry_after(self): + import requests + + with tempfile.TemporaryDirectory() as temp_dir: + dest_path = os.path.join(temp_dir, "model.safetensors") + mock_resp = MagicMock() + mock_resp.status_code = 429 + mock_resp.headers = {"Retry-After": "45"} + http_err = requests.exceptions.HTTPError(response=mock_resp) + + with patch.object(downloader, "_download_backend_from_settings", return_value="python"), \ + patch.object(downloader, "request_public_url", side_effect=http_err): + + result = download_file( + "https://civitai.com/api/download/models/12345", + dest_path, + "rate-limit-test", + dependencies=downloader, + ) + + self.assertFalse(result["success"]) + self.assertIn("Rate limit exceeded (HTTP 429)", result["error"]) + self.assertIn("45s", result["error"]) + + def test_python_download_resumes_via_direct_target_retry_when_redirect_drops_range(self): + with tempfile.TemporaryDirectory() as temp_dir: + dest_path = os.path.join(temp_dir, "model.safetensors") + partial_path = f"{dest_path}.part" + initial_data = b"0123456789" + with open(partial_path, "wb") as f: + f.write(initial_data) + + # Response 1: from origin redirect, drops range -> returns 200 + resp_redirected = MagicMock() + resp_redirected.status_code = 200 + resp_redirected.headers = {"Content-Length": "20"} + + # Response 2: from direct target URL with range -> returns 206 + remaining_data = b"ABCDEFGHIJ" + total_data = initial_data + remaining_data + resp_direct = MagicMock() + resp_direct.status_code = 206 + resp_direct.headers = { + "Content-Range": f"bytes 10-19/{len(total_data)}", + "Content-Length": "10", + } + resp_direct.iter_content.return_value = [remaining_data] + + mock_request = MagicMock(side_effect=[ + (resp_redirected, "https://cdn.example.com/direct-file", {}), + (resp_direct, "https://cdn.example.com/direct-file", {}), + ]) + + with patch.object(downloader, "_download_backend_from_settings", return_value="python"), \ + patch.object(downloader, "request_public_url", mock_request), \ + patch.object(downloader, "write_model_resolver_metadata", return_value=""): + + result = download_file( + "https://civitai.com/api/download/models/100", + dest_path, + "range-redirect-test", + dependencies=downloader, + ) + + self.assertTrue(result["success"]) + self.assertEqual(len(total_data), result["size"]) + with open(dest_path, "rb") as f: + self.assertEqual(total_data, f.read()) + self.assertEqual(mock_request.call_count, 2) + # Second call was direct to the CDN target + self.assertEqual(mock_request.call_args_list[1].args[1], "https://cdn.example.com/direct-file") + + def test_python_download_auto_reconnects_on_stream_interruption(self): + import requests + + with tempfile.TemporaryDirectory() as temp_dir: + dest_path = os.path.join(temp_dir, "model_reconnect.safetensors") + chunk1 = b"AAAAA" + chunk2 = b"BBBBB" + total_data = chunk1 + chunk2 + + # Stream 1: yields chunk1 then drops connection + resp1 = MagicMock() + resp1.status_code = 200 + resp1.headers = {"Content-Length": str(len(total_data))} + + def failing_stream(*args, **kwargs): + yield chunk1 + raise requests.exceptions.ConnectionError("Connection reset by peer") + + resp1.iter_content.side_effect = failing_stream + + # Stream 2: resumes with Range 206, yields chunk2 + resp2 = MagicMock() + resp2.status_code = 206 + resp2.headers = { + "Content-Range": f"bytes 5-9/{len(total_data)}", + "Content-Length": "5", + } + resp2.iter_content.return_value = [chunk2] + + mock_request = MagicMock(side_effect=[ + (resp1, "https://example.com/model.safetensors", {}), + (resp2, "https://example.com/model.safetensors", {}), + ]) + + with patch.object(downloader, "_download_backend_from_settings", return_value="python"), \ + patch.object(downloader, "request_public_url", mock_request), \ + patch.object(downloader, "write_model_resolver_metadata", return_value=""): + + result = download_file( + "https://example.com/model.safetensors", + dest_path, + "stream-reconnect-test", + dependencies=downloader, + ) + + self.assertTrue(result["success"]) + self.assertEqual(len(total_data), result["size"]) + with open(dest_path, "rb") as f: + self.assertEqual(total_data, f.read()) + self.assertEqual(mock_request.call_count, 2) + # Reconnected call should have requested Range from byte 5 + second_call_headers = mock_request.call_args_list[1].kwargs.get("headers", {}) + self.assertEqual(second_call_headers.get("Range"), "bytes=5-") + if __name__ == "__main__":