diff --git a/core/download/orchestrator.py b/core/download/orchestrator.py index d93ee39..3ad506f 100644 --- a/core/download/orchestrator.py +++ b/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", diff --git a/core/download/state.py b/core/download/state.py index 610139b..d1b5176 100644 --- a/core/download/state.py +++ b/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"): diff --git a/core/settings.py b/core/settings.py index 22c87db..0ee54e2 100644 --- a/core/settings.py +++ b/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 ) diff --git a/tests/test_download_sound.mjs b/tests/test_download_sound.mjs new file mode 100644 index 0000000..c51f20d --- /dev/null +++ b/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; + } +}); diff --git a/tests/test_download_state.py b/tests/test_download_state.py index 28a4e0e..7f24e11 100644 --- a/tests/test_download_state.py +++ b/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(): diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index 38cf2bb..44e1e53 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/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'); }); diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index e1adbc0..78ef458 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -295,6 +295,15 @@ export const optionsMethods = { ${renderDownloadBackendOptions(tokens.download_backend)} +
@@ -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, diff --git a/web/utils/download_sound.js b/web/utils/download_sound.js new file mode 100644 index 0000000..a422ef2 --- /dev/null +++ b/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. + } +}