Add configurable download sound and improve cancellation retries

core/download/orchestrator.py

@@ -1,10 +1,53 @@
 """High-level download orchestration used by the downloader facade."""
 
+import threading
 from typing import Any, Callable, Dict, Optional
 
 from .state import create_initial_progress
 
 
+class DownloadCancelled(Exception):
+    """The caller cancelled before the download response became available."""
+
+
+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()
+    guard = threading.Lock()
+    outcome = {}
+    abandoned = False
+
+    def request():
+        try:
+            result = facade.request_public_url("GET", url, **kwargs)
+            with guard:
+                if abandoned:
+                    result[0].close()
+                else:
+                    outcome["result"] = result
+        except Exception as exc:
+            with guard:
+                outcome["error"] = exc
+        finally:
+            finished.set()
+
+    if download_id in facade.cancelled_downloads:
+        raise DownloadCancelled()
+    threading.Thread(target=request, daemon=True).start()
+    while True:
+        finished.wait(0.1)
+        with guard:
+            if download_id in facade.cancelled_downloads:
+                abandoned = True
+                if "result" in outcome:
+                    outcome.pop("result")[0].close()
+                raise DownloadCancelled()
+            if finished.is_set():
+                if "error" in outcome:
+                    raise outcome["error"]
+                return outcome["result"]
+
+
 def find_active_download_for_path(
     download_progress: Dict[str, Dict[str, Any]],
     target_path: str,
@@ -147,8 +190,9 @@ def download_file(
         log.info(f"URL: {facade._strip_sensitive_url_params(url)}")
 
         request_headers = facade.build_download_headers(url, headers)
-        response, final_url, _final_headers = facade.request_public_url(
-            "GET",
+        response, final_url, _final_headers = request_download_response(
+            facade,
+            download_id,
             url,
             headers=request_headers,
             stream=True,
@@ -318,6 +362,13 @@ def download_file(
         )
         facade.invalidate_model_caches()
 
+    except DownloadCancelled:
+        facade._delete_python_partial_download_file(partial_path)
+        result["error"] = "Download cancelled"
+        with download_lock:
+            download_progress[download_id].update(status="cancelled", speed=0)
+        cancelled_downloads.discard(download_id)
+
     except requests.exceptions.RequestException as exc:
         error_msg = facade._sanitize_download_error(exc)
         if hasattr(exc, "response") and exc.response is not None:
@@ -622,6 +673,11 @@ def start_background_download(
 
     def run_download() -> None:
         try:
+            with download_lock:
+                if download_id in facade.cancelled_downloads:
+                    download_progress[download_id].update(status="cancelled", speed=0)
+                    facade.cancelled_downloads.discard(download_id)
+                    return
             result = facade.download_model(
                 url,
                 filename,
@@ -635,7 +691,12 @@ def start_background_download(
             if not result.get("success"):
                 with download_lock:
                     if download_id in download_progress:
-                        download_progress[download_id]["status"] = "error"
+                        cancelled = (
+                            download_id in facade.cancelled_downloads
+                            or download_progress[download_id].get("status") == "cancelled"
+                        )
+                        download_progress[download_id]["status"] = "cancelled" if cancelled else "error"
+                        facade.cancelled_downloads.discard(download_id)
                         download_progress[download_id]["error"] = result.get(
                             "error",
                             "Download failed",

core/download/state.py

@@ -78,10 +78,14 @@ def get_all_progress(state: DownloadStateDependencies) -> Dict[str, Dict[str, An
 
 def cancel_download(download_id: str, state: DownloadStateDependencies) -> bool:
     """Cancel a download in progress."""
-    state.cancelled_downloads.add(download_id)
+    with state.download_lock:
+        progress = state.download_progress.get(download_id)
+        if not progress or progress.get("status") in {"completed", "error", "cancelled"}:
+            return True
+        state.cancelled_downloads.add(download_id)
+        progress.update(status="cancelling", speed=0)
     with state.aria2_lock:
         state.aria2_desired_states.pop(download_id, None)
-    state.set_download_progress_status(download_id, "cancelling", speed=0)
 
     transfer = state.aria2_transfers.get(download_id)
     if transfer and transfer.get("gid"):

core/settings.py

@@ -36,6 +36,7 @@ SETTINGS_SCHEMA: List[Dict[str, Any]] = [
     { "serverKey": "download_backend", "localKey": "ModelResolver.downloadBackend", "type": "backend", "default": "python" },
     { "serverKey": "aria2c_path", "localKey": "ModelResolver.aria2cPath", "type": "string", "default": "" },
     { "serverKey": "aria2_auto_stop_daemon", "localKey": "ModelResolver.aria2AutoStopDaemon", "type": "boolean", "default": True },
+    { "serverKey": "download_sound_enabled", "localKey": "ModelResolver.downloadSoundEnabled", "type": "boolean", "default": True },
     { "serverKey": "download_path_mode", "localKey": "ModelResolver.downloadPathMode", "type": "pathMode", "default": "suggested" },
     { "serverKey": "download_path_templates", "localKey": "ModelResolver.downloadPathTemplates", "type": "json", "default": {} },
     { "serverKey": "base_model_path_mappings", "localKey": "ModelResolver.baseModelPathMappings", "type": "json", "default": {} },
@@ -196,6 +197,9 @@ def normalize_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
     data["aria2_auto_stop_daemon"] = bool_setting(
         data.get("aria2_auto_stop_daemon"), True
     )
+    data["download_sound_enabled"] = bool_setting(
+        data.get("download_sound_enabled"), True
+    )
     data["auto_refresh_comfy_models_after_apply"] = bool_setting(
         data.get("auto_refresh_comfy_models_after_apply"), True
     )