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
     )

tests/test_download_sound.mjs

@@ -0,0 +1,58 @@
+import test from 'node:test';
+import assert from 'node:assert/strict';
+
+test('success sound unlocks on interaction and schedules the supplied chord', async () => {
+    const listeners = {};
+    const oscillators = [];
+    const gains = [];
+    globalThis.document = { addEventListener: (name, callback) => { listeners[name] = callback; } };
+    globalThis.AudioContext = class {
+        state = 'suspended';
+        currentTime = 10;
+        destination = {};
+        async resume() { this.state = 'running'; }
+        createOscillator() {
+            const oscillator = {
+                frequency: { setValueAtTime(value, time) { this.value = value; this.time = time; } }, connect() {}, disconnect() {},
+                start(time) { this.started = time; },
+                stop(time) { this.stopped = time; },
+            };
+            oscillators.push(oscillator);
+            return oscillator;
+        }
+        createGain() {
+            const events = [];
+            gains.push(events);
+            return { gain: {
+                setValueAtTime(value, time) { events.push(['set', value, time]); },
+                exponentialRampToValueAtTime(value, time) { events.push(['exponential', value, time]); },
+            }, connect() {}, disconnect() {} };
+        }
+    };
+    try {
+        const { playDownloadSuccessSound } = await import('../web/utils/download_sound.js');
+        playDownloadSuccessSound();
+        assert.equal(oscillators.length, 0);
+        listeners.pointerdown();
+        playDownloadSuccessSound();
+        assert.deepEqual(oscillators.map(item => item.frequency.value), [523, 659, 784, 1046]);
+        assert.deepEqual(oscillators.map(item => item.started), [10, 10.15, 10.30, 10.47]);
+        oscillators.forEach((item, index) => {
+            const duration = [0.27, 0.33, 0.48, 0.40][index];
+            assert.ok(Math.abs(item.stopped - item.started - duration - 0.02) < 1e-10);
+            assert.deepEqual(gains[index], [
+                ['set', 0.0001, item.started],
+                ['exponential', [0.055, 0.060, 0.055, 0.022][index], item.started + 0.012],
+                ['exponential', 0.0001, item.started + duration],
+            ]);
+            item.onended();
+        });
+        const { safeStorage } = await import('../web/resolver/utils/html_utils.js');
+        safeStorage.setItem('ModelResolver.downloadSoundEnabled', 'false');
+        playDownloadSuccessSound();
+        assert.equal(oscillators.length, 4, 'disabled sound must not schedule more notes');
+    } finally {
+        delete globalThis.document;
+        delete globalThis.AudioContext;
+    }
+});

tests/test_download_state.py

@@ -1,4 +1,78 @@
 import importlib
+import threading
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+
+def test_cancel_does_not_wait_for_http_headers_and_closes_late_response():
+    orchestrator = importlib.import_module("core.download.orchestrator")
+    entered = threading.Event()
+    release = threading.Event()
+    closed = threading.Event()
+    cancelled = threading.Event()
+    response = SimpleNamespace(close=closed.set)
+
+    def request(*args, **kwargs):
+        entered.set()
+        assert release.wait(2)
+        return response, "https://example.com/model", {}
+
+    facade = SimpleNamespace(request_public_url=request, cancelled_downloads=set())
+
+    def run():
+        with pytest.raises(orchestrator.DownloadCancelled):
+            orchestrator.request_download_response(facade, "download", "https://example.com/model")
+        cancelled.set()
+
+    worker = threading.Thread(target=run)
+    worker.start()
+    try:
+        assert entered.wait(1)
+        facade.cancelled_downloads.add("download")
+        assert cancelled.wait(1), "Cancellation must finish before HTTP headers arrive"
+        assert not closed.is_set()
+    finally:
+        release.set()
+        worker.join(2)
+    assert closed.wait(1)
+
+
+def test_download_response_preserves_result_and_network_errors():
+    orchestrator = importlib.import_module("core.download.orchestrator")
+    result = (Mock(), "https://example.com/model", {})
+    facade = SimpleNamespace(request_public_url=Mock(return_value=result), cancelled_downloads=set())
+    assert orchestrator.request_download_response(facade, "download", result[1]) == result
+    facade.request_public_url.side_effect = ValueError("network error")
+    with pytest.raises(ValueError, match="network error"):
+        orchestrator.request_download_response(facade, "download", result[1])
+
+
+def test_cancel_is_idempotent_and_does_not_revive_terminal_downloads():
+    state_module = importlib.import_module("core.download.state")
+    progress = {"download": {"status": "starting", "speed": 5}}
+    state = SimpleNamespace(
+        download_progress=progress,
+        download_lock=threading.Lock(),
+        cancelled_downloads=set(),
+        aria2_lock=threading.Lock(),
+        aria2_desired_states={},
+        aria2_transfers={},
+        xet_transfers_lock=threading.Lock(),
+        xet_transfers={},
+    )
+    assert state_module.cancel_download("download", state)
+    assert progress["download"] == {"status": "cancelling", "speed": 0}
+    assert "download" in state.cancelled_downloads
+    for status in ("cancelled", "completed", "error"):
+        progress["download"]["status"] = status
+        state.cancelled_downloads.clear()
+        assert state_module.cancel_download("download", state)
+        assert progress["download"]["status"] == status
+        assert not state.cancelled_downloads
+    assert state_module.cancel_download("unknown", state)
+    assert "unknown" not in state.cancelled_downloads
 
 
 def test_initial_download_progress_preserves_the_shared_contract():

web/resolver/actions/resolve_download_methods.js

@@ -9,6 +9,7 @@ import { bindDownloadActionHandlers } from "../utils/download_action_handlers.js
 import { matchesLocalModelDownload } from "../utils/local_match_utils.js";
 import { matchesSearchText } from "../utils/search_utils.js";
 import { normalizeSourceKey } from "../utils/source_labels.js";
+import { playDownloadSuccessSound } from "../../utils/download_sound.js";
 const log = createModuleLogger('resolve_download_methods');
 
 export const resolveDownloadMethods = {
@@ -2225,6 +2226,7 @@ export const resolveDownloadMethods = {
                 this.updateDownloadAllButtonState();
                 this.updateQueuePanel?.();
                 this.refreshLocalMatchesUiForMissing?.(missing);
+                if (!alreadyExists) playDownloadSuccessSound();
                 this.showNotification(alreadyExists
                     ? `Already downloaded: ${progress.filename}`
                     : `Downloaded: ${progress.filename}`, 'success', {
@@ -2333,16 +2335,10 @@ export const resolveDownloadMethods = {
         this.fetchJson(`/model_resolver/cancel/${downloadId}`, {
             method: 'POST'
         }, 'Cancel download').then(() => {
-            if (!info) return;
-            this.finalizeCancelledDownloadFrontend?.(downloadId, info, {
-                ...(info.lastProgress || {}),
-                status: 'cancelled',
-                filename: info.lastProgress?.filename || info.filename || '',
-                path: info.lastProgress?.path || info.downloadPath || '',
-                directory: info.lastProgress?.directory || info.downloadDirectory || ''
-            });
-            this.showNotification('Download cancelled', 'info');
+            // The request only acknowledges cancellation. Progress polling
+            // releases the file and enables retry when the worker has stopped.
         }).catch((error) => {
+            this.clearPendingDownloadStatus(info);
             console.error('Model Resolver: Cancel error:', error);
             this.showNotification('Failed to cancel download', 'error');
         });

web/resolver/views/options_methods.js

@@ -295,6 +295,15 @@ export const optionsMethods = {
                                                 ${renderDownloadBackendOptions(tokens.download_backend)}
                                             </select>
                                         </div>
+                                        <label class="mr-options-toggle-row mr-options-compact-toggle-row">
+                                            <div class="mr-options-toggle-copy">
+                                                <span class="mr-options-toggle-title">Download sound notifications <span class="mr-tooltip-badge" data-tooltip="Play a sound when a model finishes downloading successfully.">?</span></span>
+                                            </div>
+                                            <span class="mr-options-toggle-control">
+                                                <input id="mr-options-download-sound" class="mr-options-switch-input" type="checkbox" ${tokens.download_sound_enabled !== false ? 'checked' : ''}>
+                                                <span class="mr-options-switch"></span>
+                                            </span>
+                                        </label>
                                         <div class="mr-options-dependent-block mr-options-aria2-setting">
                                             <div class="mr-options-number-row mr-options-wide-row">
                                                 <div class="mr-options-number-copy">
@@ -969,6 +978,7 @@ export const optionsMethods = {
         const aria2AvailabilityEl = this.contentElement.querySelector('#mr-options-aria2-availability');
         const aria2DaemonEl = this.contentElement.querySelector('#mr-options-aria2-daemon');
         const aria2AutoStopInput = this.contentElement.querySelector('#mr-options-aria2-auto-stop');
+        const downloadSoundInput = this.contentElement.querySelector('#mr-options-download-sound');
         const aria2SettingEls = Array.from(this.contentElement.querySelectorAll('.mr-options-aria2-setting'));
         const downloadPathModeInput = this.contentElement.querySelector('#mr-options-download-path-mode');
         const defaultRootSelectInputs = Array.from(this.contentElement.querySelectorAll('.mr-options-default-root'));
@@ -1091,6 +1101,7 @@ export const optionsMethods = {
             downloadBackendInput,
             aria2cPathInput,
             aria2AutoStopInput,
+            downloadSoundInput,
             downloadPathModeInput,
             ...defaultRootSelectInputs,
             ...templatePresetInputs,
@@ -3508,6 +3519,7 @@ export const optionsMethods = {
                     download_backend: downloadBackend,
                     aria2c_path: aria2cPathInput?.value?.trim() || '',
                     aria2_auto_stop_daemon: Boolean(aria2AutoStopInput?.checked),
+                    download_sound_enabled: Boolean(downloadSoundInput?.checked),
                     download_path_mode: downloadPathMode,
                     download_path_templates: downloadPathTemplates,
                     base_model_path_mappings: baseModelPathMappings,

web/utils/download_sound.js

@@ -0,0 +1,53 @@
+import { safeStorage } from '../resolver/utils/html_utils.js';
+
+const notes = [
+    { freq: 523, start: 0, duration: 0.27, volume: 0.055 },
+    { freq: 659, start: 0.15, duration: 0.33, volume: 0.060 },
+    { freq: 784, start: 0.30, duration: 0.48, volume: 0.055 },
+    { freq: 1046, start: 0.47, duration: 0.40, volume: 0.022 },
+];
+
+let context;
+
+function unlockAudio() {
+    try {
+        const AudioContext = globalThis.AudioContext || globalThis.webkitAudioContext;
+        if (!AudioContext) return;
+        context ||= new AudioContext();
+        if (context.state === 'suspended') void context.resume().catch(() => {});
+    } catch {
+        // Audio availability must not affect downloads.
+    }
+}
+
+// Browsers require a user gesture before allowing notification audio.
+globalThis.document?.addEventListener('pointerdown', unlockAudio, { capture: true, passive: true });
+globalThis.document?.addEventListener('keydown', unlockAudio, { capture: true, passive: true });
+
+export function playDownloadSuccessSound() {
+    try {
+        if (safeStorage.getItem('ModelResolver.downloadSoundEnabled') === 'false' || context?.state !== 'running') return;
+        const now = context.currentTime;
+        for (const note of notes) {
+            const oscillator = context.createOscillator();
+            const gain = context.createGain();
+            const start = now + note.start;
+            const end = start + note.duration;
+            oscillator.type = 'sine';
+            oscillator.frequency.setValueAtTime(note.freq, start);
+            gain.gain.setValueAtTime(0.0001, start);
+            gain.gain.exponentialRampToValueAtTime(note.volume, start + 0.012);
+            gain.gain.exponentialRampToValueAtTime(0.0001, end);
+            oscillator.connect(gain);
+            gain.connect(context.destination);
+            oscillator.onended = () => {
+                oscillator.disconnect();
+                gain.disconnect();
+            };
+            oscillator.start(start);
+            oscillator.stop(end + 0.02);
+        }
+    } catch {
+        // Notification audio is optional when the browser cannot play it.
+    }
+}