From 52e671638b3bf339f8073c2861c6edf1ca5c8137 Mon Sep 17 00:00:00 2001 From: pixelpaws Date: Tue, 14 Oct 2025 22:18:57 +0800 Subject: [PATCH] feat(example-images): add stop control for download panel --- locales/de.json | 2 + locales/en.json | 2 + locales/es.json | 2 + locales/fr.json | 2 + locales/he.json | 2 + locales/ja.json | 2 + locales/ko.json | 2 + locales/ru.json | 2 + locales/zh-CN.json | 2 + locales/zh-TW.json | 2 + py/routes/example_images_route_registrar.py | 1 + py/routes/handlers/example_images_handlers.py | 8 + py/utils/example_images_download_manager.py | 176 ++++++++++++++---- static/js/managers/ExampleImagesManager.js | 148 ++++++++++++--- templates/components/progress_panel.html | 3 + ...example_images_route_registrar_handlers.py | 13 ++ tests/routes/test_example_images_routes.py | 20 +- ...st_example_images_download_manager_unit.py | 57 ++++++ 18 files changed, 389 insertions(+), 57 deletions(-) diff --git a/locales/de.json b/locales/de.json index d96cba50..a19e7047 100644 --- a/locales/de.json +++ b/locales/de.json @@ -1245,6 +1245,8 @@ "pauseFailed": "Fehler beim Pausieren des Downloads: {error}", "downloadResumed": "Download fortgesetzt", "resumeFailed": "Fehler beim Fortsetzen des Downloads: {error}", + "downloadStopped": "Download abgebrochen", + "stopFailed": "Download konnte nicht abgebrochen werden: {error}", "deleted": "Beispielbild gelöscht", "deleteFailed": "Fehler beim Löschen des Beispielbilds", "setPreviewFailed": "Fehler beim Setzen des Vorschaubilds" diff --git a/locales/en.json b/locales/en.json index 3cb82402..3785588d 100644 --- a/locales/en.json +++ b/locales/en.json @@ -1245,6 +1245,8 @@ "pauseFailed": "Failed to pause download: {error}", "downloadResumed": "Download resumed", "resumeFailed": "Failed to resume download: {error}", + "downloadStopped": "Download cancelled", + "stopFailed": "Failed to cancel download: {error}", "deleted": "Example image deleted", "deleteFailed": "Failed to delete example image", "setPreviewFailed": "Failed to set preview image" diff --git a/locales/es.json b/locales/es.json index bf706254..a47d171e 100644 --- a/locales/es.json +++ b/locales/es.json @@ -1245,6 +1245,8 @@ "pauseFailed": "Error al pausar descarga: {error}", "downloadResumed": "Descarga reanudada", "resumeFailed": "Error al reanudar descarga: {error}", + "downloadStopped": "Descarga cancelada", + "stopFailed": "Error al cancelar descarga: {error}", "deleted": "Imagen de ejemplo eliminada", "deleteFailed": "Error al eliminar imagen de ejemplo", "setPreviewFailed": "Error al establecer imagen de vista previa" diff --git a/locales/fr.json b/locales/fr.json index c64de794..93556d9f 100644 --- a/locales/fr.json +++ b/locales/fr.json @@ -1245,6 +1245,8 @@ "pauseFailed": "Échec de la mise en pause du téléchargement : {error}", "downloadResumed": "Téléchargement repris", "resumeFailed": "Échec de la reprise du téléchargement : {error}", + "downloadStopped": "Téléchargement annulé", + "stopFailed": "Échec de l'annulation du téléchargement : {error}", "deleted": "Image d'exemple supprimée", "deleteFailed": "Échec de la suppression de l'image d'exemple", "setPreviewFailed": "Échec de la définition de l'image d'aperçu" diff --git a/locales/he.json b/locales/he.json index ecc2ade6..79fdd7e9 100644 --- a/locales/he.json +++ b/locales/he.json @@ -1245,6 +1245,8 @@ "pauseFailed": "השהיית ההורדה נכשלה: {error}", "downloadResumed": "ההורדה חודשה", "resumeFailed": "חידוש ההורדה נכשל: {error}", + "downloadStopped": "ההורדה בוטלה", + "stopFailed": "נכשל בביטול ההורדה: {error}", "deleted": "תמונת הדוגמה נמחקה", "deleteFailed": "מחיקת תמונת הדוגמה נכשלה", "setPreviewFailed": "הגדרת תמונת התצוגה המקדימה נכשלה" diff --git a/locales/ja.json b/locales/ja.json index 76a71313..d39fccaa 100644 --- a/locales/ja.json +++ b/locales/ja.json @@ -1245,6 +1245,8 @@ "pauseFailed": "ダウンロードの一時停止に失敗しました：{error}", "downloadResumed": "ダウンロードが再開されました", "resumeFailed": "ダウンロードの再開に失敗しました：{error}", + "downloadStopped": "ダウンロードをキャンセルしました", + "stopFailed": "ダウンロードのキャンセルに失敗しました：{error}", "deleted": "例画像が削除されました", "deleteFailed": "例画像の削除に失敗しました", "setPreviewFailed": "プレビュー画像の設定に失敗しました" diff --git a/locales/ko.json b/locales/ko.json index 7fb88db8..a6565c10 100644 --- a/locales/ko.json +++ b/locales/ko.json @@ -1245,6 +1245,8 @@ "pauseFailed": "다운로드 일시정지 실패: {error}", "downloadResumed": "다운로드가 재개되었습니다", "resumeFailed": "다운로드 재개 실패: {error}", + "downloadStopped": "다운로드가 취소되었습니다", + "stopFailed": "다운로드 취소 실패: {error}", "deleted": "예시 이미지가 삭제되었습니다", "deleteFailed": "예시 이미지 삭제 실패", "setPreviewFailed": "미리보기 이미지 설정 실패" diff --git a/locales/ru.json b/locales/ru.json index 4b014df9..5a3b496e 100644 --- a/locales/ru.json +++ b/locales/ru.json @@ -1245,6 +1245,8 @@ "pauseFailed": "Не удалось приостановить загрузку: {error}", "downloadResumed": "Загрузка возобновлена", "resumeFailed": "Не удалось возобновить загрузку: {error}", + "downloadStopped": "Загрузка отменена", + "stopFailed": "Не удалось отменить загрузку: {error}", "deleted": "Пример изображения удален", "deleteFailed": "Не удалось удалить пример изображения", "setPreviewFailed": "Не удалось установить превью изображение" diff --git a/locales/zh-CN.json b/locales/zh-CN.json index a1454a81..474b61ca 100644 --- a/locales/zh-CN.json +++ b/locales/zh-CN.json @@ -1245,6 +1245,8 @@ "pauseFailed": "暂停下载失败：{error}", "downloadResumed": "下载已恢复", "resumeFailed": "恢复下载失败：{error}", + "downloadStopped": "下载已取消", + "stopFailed": "取消下载失败：{error}", "deleted": "示例图片已删除", "deleteFailed": "删除示例图片失败", "setPreviewFailed": "设置预览图片失败" diff --git a/locales/zh-TW.json b/locales/zh-TW.json index 1a7f98fa..9e1385fd 100644 --- a/locales/zh-TW.json +++ b/locales/zh-TW.json @@ -1245,6 +1245,8 @@ "pauseFailed": "暫停下載失敗：{error}", "downloadResumed": "下載已恢復", "resumeFailed": "恢復下載失敗：{error}", + "downloadStopped": "下載已取消", + "stopFailed": "取消下載失敗：{error}", "deleted": "範例圖片已刪除", "deleteFailed": "刪除範例圖片失敗", "setPreviewFailed": "設定預覽圖片失敗" diff --git a/py/routes/example_images_route_registrar.py b/py/routes/example_images_route_registrar.py index aa12c3b1..63bac34b 100644 --- a/py/routes/example_images_route_registrar.py +++ b/py/routes/example_images_route_registrar.py @@ -22,6 +22,7 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = ( RouteDefinition("GET", "/api/lm/example-images-status", "get_example_images_status"), RouteDefinition("POST", "/api/lm/pause-example-images", "pause_example_images"), RouteDefinition("POST", "/api/lm/resume-example-images", "resume_example_images"), + RouteDefinition("POST", "/api/lm/stop-example-images", "stop_example_images"), RouteDefinition("POST", "/api/lm/open-example-images-folder", "open_example_images_folder"), RouteDefinition("GET", "/api/lm/example-image-files", "get_example_image_files"), RouteDefinition("GET", "/api/lm/has-example-images", "has_example_images"), diff --git a/py/routes/handlers/example_images_handlers.py b/py/routes/handlers/example_images_handlers.py index 111013e2..f7acf5ef 100644 --- a/py/routes/handlers/example_images_handlers.py +++ b/py/routes/handlers/example_images_handlers.py @@ -68,6 +68,13 @@ class ExampleImagesDownloadHandler: except DownloadNotRunningError as exc: return web.json_response({'success': False, 'error': str(exc)}, status=400) + async def stop_example_images(self, request: web.Request) -> web.StreamResponse: + try: + result = await self._download_manager.stop_download(request) + return web.json_response(result) + except DownloadNotRunningError as exc: + return web.json_response({'success': False, 'error': str(exc)}, status=400) + async def force_download_example_images(self, request: web.Request) -> web.StreamResponse: try: payload = await request.json() @@ -149,6 +156,7 @@ class ExampleImagesHandlerSet: "get_example_images_status": self.download.get_example_images_status, "pause_example_images": self.download.pause_example_images, "resume_example_images": self.download.resume_example_images, + "stop_example_images": self.download.stop_example_images, "force_download_example_images": self.download.force_download_example_images, "import_example_images": self.management.import_example_images, "delete_example_image": self.management.delete_example_image, diff --git a/py/utils/example_images_download_manager.py b/py/utils/example_images_download_manager.py index 4726e81b..85a70b6c 100644 --- a/py/utils/example_images_download_manager.py +++ b/py/utils/example_images_download_manager.py @@ -105,6 +105,7 @@ class DownloadManager: self._progress = _DownloadProgress() self._ws_manager = ws_manager self._state_lock = state_lock or asyncio.Lock() + self._stop_requested = False def _resolve_output_dir(self, library_name: str | None = None) -> str: base_path = get_settings_manager().get('example_images_path') @@ -145,6 +146,7 @@ class DownloadManager: raise DownloadConfigurationError('Example images path not configured in settings') self._progress.reset() + self._stop_requested = False self._progress['status'] = 'running' self._progress['start_time'] = time.time() self._progress['end_time'] = None @@ -267,6 +269,27 @@ class DownloadManager: 'success': True, 'message': 'Download resumed' } + + async def stop_download(self, request): + """Stop the example images download after the current model completes.""" + + async with self._state_lock: + if not self._is_downloading: + raise DownloadNotRunningError() + + if self._progress['status'] in {'completed', 'error', 'stopped'}: + raise DownloadNotRunningError() + + if self._progress['status'] != 'stopping': + self._stop_requested = True + self._progress['status'] = 'stopping' + + await self._broadcast_progress(status='stopping') + + return { + 'success': True, + 'message': 'Download stopping' + } async def _download_all_example_images( self, @@ -311,6 +334,12 @@ class DownloadManager: # Process each model for i, (scanner_type, model, scanner) in enumerate(all_models): + async with self._state_lock: + current_status = self._progress['status'] + + if current_status not in {'running', 'paused', 'stopping'}: + break + # Main logic for processing model is here, but actual operations are delegated to other classes was_remote_download = await self._process_model( scanner_type, @@ -321,24 +350,59 @@ class DownloadManager: downloader, library_name, ) - + # Update progress self._progress['completed'] += 1 - await self._broadcast_progress(status='running') - + + async with self._state_lock: + current_status = self._progress['status'] + should_stop = self._stop_requested and current_status == 'stopping' + + broadcast_status = 'running' if current_status == 'running' else current_status + await self._broadcast_progress(status=broadcast_status) + + if should_stop: + break + # Only add delay after remote download of models, and not after processing the last model - if was_remote_download and i < len(all_models) - 1 and self._progress['status'] == 'running': + if ( + was_remote_download + and i < len(all_models) - 1 + and current_status == 'running' + ): await asyncio.sleep(delay) - - # Mark as completed - self._progress['status'] = 'completed' - self._progress['end_time'] = time.time() - logger.debug( - "Example images download completed: %s/%s models processed", - self._progress['completed'], - self._progress['total'], - ) - await self._broadcast_progress(status='completed') + + async with self._state_lock: + if self._stop_requested and self._progress['status'] == 'stopping': + self._progress['status'] = 'stopped' + self._progress['end_time'] = time.time() + self._stop_requested = False + final_status = 'stopped' + elif self._progress['status'] not in {'error', 'stopped'}: + self._progress['status'] = 'completed' + self._progress['end_time'] = time.time() + self._stop_requested = False + final_status = 'completed' + else: + final_status = self._progress['status'] + self._stop_requested = False + if self._progress['end_time'] is None: + self._progress['end_time'] = time.time() + + if final_status == 'completed': + logger.debug( + "Example images download completed: %s/%s models processed", + self._progress['completed'], + self._progress['total'], + ) + elif final_status == 'stopped': + logger.debug( + "Example images download stopped: %s/%s models processed", + self._progress['completed'], + self._progress['total'], + ) + + await self._broadcast_progress(status=final_status) except Exception as e: error_msg = f"Error during example images download: {str(e)}" @@ -360,6 +424,7 @@ class DownloadManager: async with self._state_lock: self._is_downloading = False self._download_task = None + self._stop_requested = False async def _process_model( self, @@ -378,7 +443,7 @@ class DownloadManager: await asyncio.sleep(1) # Check if download should continue - if self._progress['status'] != 'running': + if self._progress['status'] not in {'running', 'stopping'}: logger.info(f"Download stopped: {self._progress['status']}") return False # Return False to indicate no remote download happened @@ -567,6 +632,7 @@ class DownloadManager: raise DownloadConfigurationError('Example images path not configured in settings') self._progress.reset() + self._stop_requested = False self._progress['total'] = len(model_hashes) self._progress['status'] = 'running' self._progress['start_time'] = time.time() @@ -588,10 +654,15 @@ class DownloadManager: async with self._state_lock: self._is_downloading = False + final_status = self._progress['status'] + + message = 'Force download completed' + if final_status == 'stopped': + message = 'Force download stopped' return { 'success': True, - 'message': 'Force download completed', + 'message': message, 'result': result } @@ -649,6 +720,12 @@ class DownloadManager: # Process each model success_count = 0 for i, (scanner_type, model, scanner) in enumerate(models_to_process): + async with self._state_lock: + current_status = self._progress['status'] + + if current_status not in {'running', 'paused', 'stopping'}: + break + # Force process this model regardless of previous status was_successful = await self._process_specific_model( scanner_type, @@ -659,32 +736,65 @@ class DownloadManager: downloader, library_name, ) - + if was_successful: success_count += 1 - + # Update progress self._progress['completed'] += 1 + async with self._state_lock: + current_status = self._progress['status'] + should_stop = self._stop_requested and current_status == 'stopping' + + broadcast_status = 'running' if current_status == 'running' else current_status # Send progress update via WebSocket - await self._broadcast_progress(status='running') - + await self._broadcast_progress(status=broadcast_status) + + if should_stop: + break + # Only add delay after remote download, and not after processing the last model - if was_successful and i < len(models_to_process) - 1 and self._progress['status'] == 'running': + if ( + was_successful + and i < len(models_to_process) - 1 + and current_status == 'running' + ): await asyncio.sleep(delay) - - # Mark as completed - self._progress['status'] = 'completed' - self._progress['end_time'] = time.time() - logger.debug( - "Forced example images download completed: %s/%s models processed", - self._progress['completed'], - self._progress['total'], - ) + + async with self._state_lock: + if self._stop_requested and self._progress['status'] == 'stopping': + self._progress['status'] = 'stopped' + self._progress['end_time'] = time.time() + self._stop_requested = False + final_status = 'stopped' + elif self._progress['status'] not in {'error', 'stopped'}: + self._progress['status'] = 'completed' + self._progress['end_time'] = time.time() + self._stop_requested = False + final_status = 'completed' + else: + final_status = self._progress['status'] + self._stop_requested = False + if self._progress['end_time'] is None: + self._progress['end_time'] = time.time() + + if final_status == 'completed': + logger.debug( + "Forced example images download completed: %s/%s models processed", + self._progress['completed'], + self._progress['total'], + ) + elif final_status == 'stopped': + logger.debug( + "Forced example images download stopped: %s/%s models processed", + self._progress['completed'], + self._progress['total'], + ) # Send final progress via WebSocket - await self._broadcast_progress(status='completed') - + await self._broadcast_progress(status=final_status) + return { 'total': self._progress['total'], 'processed': self._progress['completed'], @@ -726,7 +836,7 @@ class DownloadManager: await asyncio.sleep(1) # Check if download should continue - if self._progress['status'] != 'running': + if self._progress['status'] not in {'running', 'stopping'}: logger.info(f"Download stopped: {self._progress['status']}") return False diff --git a/static/js/managers/ExampleImagesManager.js b/static/js/managers/ExampleImagesManager.js index d2a05fef..162f0f4f 100644 --- a/static/js/managers/ExampleImagesManager.js +++ b/static/js/managers/ExampleImagesManager.js @@ -13,8 +13,10 @@ export class ExampleImagesManager { this.progressPanel = null; this.isProgressPanelCollapsed = false; this.pauseButton = null; // Store reference to the pause button + this.stopButton = null; this.isMigrating = false; // Track migration state separately from downloading this.hasShownCompletionToast = false; // Flag to track if completion toast has been shown + this.isStopping = false; // Auto download properties this.autoDownloadInterval = null; @@ -52,11 +54,16 @@ export class ExampleImagesManager { // Initialize progress panel button handlers this.pauseButton = document.getElementById('pauseExampleDownloadBtn'); + this.stopButton = document.getElementById('stopExampleDownloadBtn'); const collapseBtn = document.getElementById('collapseProgressBtn'); - + if (this.pauseButton) { this.pauseButton.onclick = () => this.pauseDownload(); } + + if (this.stopButton) { + this.stopButton.onclick = () => this.stopDownload(); + } if (collapseBtn) { collapseBtn.onclick = () => this.toggleProgressPanel(); @@ -210,10 +217,14 @@ export class ExampleImagesManager { updateDownloadButtonText() { const btnTextElement = document.getElementById('exampleDownloadBtnText'); if (btnTextElement) { - if (this.isDownloading && this.isPaused) { + if (this.isStopping) { + btnTextElement.textContent = "Stopping..."; + } else if (this.isDownloading && this.isPaused) { btnTextElement.textContent = "Resume"; } else if (!this.isDownloading) { btnTextElement.textContent = "Download"; + } else { + btnTextElement.textContent = "Download"; } } } @@ -239,18 +250,22 @@ export class ExampleImagesManager { }); const data = await response.json(); - + if (data.success) { this.isDownloading = true; this.isPaused = false; + this.isStopping = false; this.hasShownCompletionToast = false; // Reset toast flag when starting new download this.startTime = new Date(); this.updateUI(data.status); this.showProgressPanel(); this.startProgressUpdates(); this.updateDownloadButtonText(); + if (this.stopButton) { + this.stopButton.disabled = false; + } showToast('toast.exampleImages.downloadStarted', {}, 'success'); - + // Close settings modal modalManager.closeModal('settingsModal'); } else { @@ -263,7 +278,7 @@ export class ExampleImagesManager { } async pauseDownload() { - if (!this.isDownloading || this.isPaused) { + if (!this.isDownloading || this.isPaused || this.isStopping) { return; } @@ -299,21 +314,21 @@ export class ExampleImagesManager { } async resumeDownload() { - if (!this.isDownloading || !this.isPaused) { + if (!this.isDownloading || !this.isPaused || this.isStopping) { return; } - + try { const response = await fetch('/api/lm/resume-example-images', { method: 'POST' }); - + const data = await response.json(); - + if (data.success) { this.isPaused = false; document.getElementById('downloadStatusText').textContent = 'Downloading'; - + // Only update the icon element, not the entire innerHTML if (this.pauseButton) { const iconElement = this.pauseButton.querySelector('i'); @@ -322,7 +337,7 @@ export class ExampleImagesManager { } this.pauseButton.onclick = () => this.pauseDownload(); } - + this.updateDownloadButtonText(); showToast('toast.exampleImages.downloadResumed', {}, 'success'); } else { @@ -333,6 +348,60 @@ export class ExampleImagesManager { showToast('toast.exampleImages.resumeFailed', {}, 'error'); } } + + async stopDownload() { + if (this.isStopping) { + return; + } + + if (!this.isDownloading) { + this.hideProgressPanel(); + return; + } + + this.isStopping = true; + this.isPaused = false; + this.updateDownloadButtonText(); + + if (this.stopButton) { + this.stopButton.disabled = true; + } + + try { + const response = await fetch('/api/lm/stop-example-images', { + method: 'POST' + }); + + let data; + try { + data = await response.json(); + } catch (parseError) { + data = { success: false, error: 'Invalid server response' }; + } + + if (response.ok && data.success) { + showToast('toast.exampleImages.downloadStopped', {}, 'info'); + this.hideProgressPanel(); + } else { + this.isStopping = false; + if (this.stopButton) { + this.stopButton.disabled = false; + } + const errorMessage = data && data.error ? data.error : 'Unknown error'; + showToast('toast.exampleImages.stopFailed', { error: errorMessage }, 'error'); + } + } catch (error) { + console.error('Failed to stop download:', error); + this.isStopping = false; + if (this.stopButton) { + this.stopButton.disabled = false; + } + const errorMessage = error && error.message ? error.message : 'Unknown error'; + showToast('toast.exampleImages.stopFailed', { error: errorMessage }, 'error'); + } finally { + this.updateDownloadButtonText(); + } + } startProgressUpdates() { // Clear any existing interval @@ -352,21 +421,36 @@ export class ExampleImagesManager { const data = await response.json(); if (data.success) { + const currentStatus = data.status.status; this.isDownloading = data.is_downloading; - this.isPaused = data.status.status === 'paused'; + this.isPaused = currentStatus === 'paused'; this.isMigrating = data.is_migrating || false; - + + if (currentStatus === 'stopping') { + this.isStopping = true; + } else if ( + !data.is_downloading || + currentStatus === 'stopped' || + currentStatus === 'completed' || + currentStatus === 'error' + ) { + this.isStopping = false; + } + // Update download button text this.updateDownloadButtonText(); - + if (this.isDownloading) { this.updateUI(data.status); } else { // Download completed or failed clearInterval(this.progressUpdateInterval); this.progressUpdateInterval = null; - - if (data.status.status === 'completed' && !this.hasShownCompletionToast) { + if (this.stopButton) { + this.stopButton.disabled = true; + } + + if (currentStatus === 'completed' && !this.hasShownCompletionToast) { const actionType = this.isMigrating ? 'migration' : 'download'; showToast('toast.downloads.imagesCompleted', { action: actionType }, 'success'); // Mark as shown to prevent duplicate toasts @@ -375,10 +459,13 @@ export class ExampleImagesManager { this.isMigrating = false; // Hide the panel after a delay setTimeout(() => this.hideProgressPanel(), 5000); - } else if (data.status.status === 'error') { + } else if (currentStatus === 'error') { const actionType = this.isMigrating ? 'migration' : 'download'; showToast('toast.downloads.imagesFailed', { action: actionType }, 'error'); this.isMigrating = false; + } else if (currentStatus === 'stopped') { + this.hideProgressPanel(); + this.isMigrating = false; } } } @@ -434,7 +521,11 @@ export class ExampleImagesManager { if (!this.pauseButton) { this.pauseButton = document.getElementById('pauseExampleDownloadBtn'); } - + + if (!this.stopButton) { + this.stopButton = document.getElementById('stopExampleDownloadBtn'); + } + if (this.pauseButton) { // Check if the button already has the SVG elements let hasProgressElements = !!this.pauseButton.querySelector('.mini-progress-circle'); @@ -456,12 +547,14 @@ export class ExampleImagesManager { iconElement.className = status.status === 'paused' ? 'fas fa-play' : 'fas fa-pause'; } } - + // Update click handler - this.pauseButton.onclick = status.status === 'paused' - ? () => this.resumeDownload() + this.pauseButton.onclick = status.status === 'paused' + ? () => this.resumeDownload() : () => this.pauseDownload(); - + + this.pauseButton.disabled = ['completed', 'error', 'stopped'].includes(status.status) || status.status === 'stopping'; + // Update progress immediately const progressBar = document.getElementById('downloadProgressBar'); if (progressBar) { @@ -469,6 +562,15 @@ export class ExampleImagesManager { this.updateMiniProgress(progressPercent); } } + + if (this.stopButton) { + if (status.status === 'stopping' || this.isStopping) { + this.stopButton.disabled = true; + } else { + const canStop = ['running', 'paused'].includes(status.status); + this.stopButton.disabled = !canStop; + } + } // Update title text const titleElement = document.querySelector('.progress-panel-title'); @@ -584,6 +686,8 @@ export class ExampleImagesManager { case 'paused': return 'Paused'; case 'completed': return 'Completed'; case 'error': return 'Error'; + case 'stopping': return 'Stopping'; + case 'stopped': return 'Stopped'; default: return 'Initializing'; } } diff --git a/templates/components/progress_panel.html b/templates/components/progress_panel.html index eaee35f9..ad51a621 100644 --- a/templates/components/progress_panel.html +++ b/templates/components/progress_panel.html @@ -13,6 +13,9 @@ + diff --git a/tests/routes/test_example_images_route_registrar_handlers.py b/tests/routes/test_example_images_route_registrar_handlers.py index b4441f0b..1dceffc3 100644 --- a/tests/routes/test_example_images_route_registrar_handlers.py +++ b/tests/routes/test_example_images_route_registrar_handlers.py @@ -41,9 +41,11 @@ class StubDownloadManager: def __init__(self) -> None: self.pause_calls = 0 self.resume_calls = 0 + self.stop_calls = 0 self.force_payloads: list[dict[str, Any]] = [] self.pause_error: Exception | None = None self.resume_error: Exception | None = None + self.stop_error: Exception | None = None self.force_error: Exception | None = None async def get_status(self, request: web.Request) -> dict[str, Any]: @@ -61,6 +63,12 @@ class StubDownloadManager: raise self.resume_error return {"success": True, "message": "resumed"} + async def stop_download(self, request: web.Request) -> dict[str, Any]: + self.stop_calls += 1 + if self.stop_error: + raise self.stop_error + return {"success": True, "message": "stopping"} + async def start_force_download(self, payload: dict[str, Any]) -> dict[str, Any]: self.force_payloads.append(payload) if self.force_error: @@ -193,17 +201,22 @@ async def test_pause_and_resume_return_client_errors_when_not_running(): async with registrar_app() as harness: harness.download_manager.pause_error = DownloadNotRunningError() harness.download_manager.resume_error = DownloadNotRunningError("Stopped") + harness.download_manager.stop_error = DownloadNotRunningError("Not running") pause_response = await harness.client.post("/api/lm/pause-example-images") resume_response = await harness.client.post("/api/lm/resume-example-images") + stop_response = await harness.client.post("/api/lm/stop-example-images") assert pause_response.status == 400 assert resume_response.status == 400 + assert stop_response.status == 400 pause_body = await _json(pause_response) resume_body = await _json(resume_response) + stop_body = await _json(stop_response) assert pause_body == {"success": False, "error": "No download in progress"} assert resume_body == {"success": False, "error": "Stopped"} + assert stop_body == {"success": False, "error": "Not running"} async def test_import_route_returns_validation_errors(): diff --git a/tests/routes/test_example_images_routes.py b/tests/routes/test_example_images_routes.py index 1cc57eb0..10b305c1 100644 --- a/tests/routes/test_example_images_routes.py +++ b/tests/routes/test_example_images_routes.py @@ -51,6 +51,10 @@ class StubDownloadManager: self.calls.append(("resume_download", None)) return {"operation": "resume_download"} + async def stop_download(self, request: web.Request) -> dict: + self.calls.append(("stop_download", None)) + return {"operation": "stop_download"} + async def start_force_download(self, payload: Any) -> dict: self.calls.append(("start_force_download", payload)) return {"operation": "start_force_download", "payload": payload} @@ -195,19 +199,23 @@ async def test_status_route_returns_manager_payload(): assert harness.download_manager.calls == [("get_status", {"detail": "true"})] -async def test_pause_and_resume_routes_delegate(): +async def test_pause_resume_and_stop_routes_delegate(): async with example_images_app() as harness: pause_response = await harness.client.post("/api/lm/pause-example-images") resume_response = await harness.client.post("/api/lm/resume-example-images") + stop_response = await harness.client.post("/api/lm/stop-example-images") assert pause_response.status == 200 assert await pause_response.json() == {"operation": "pause_download"} assert resume_response.status == 200 assert await resume_response.json() == {"operation": "resume_download"} + assert stop_response.status == 200 + assert await stop_response.json() == {"operation": "stop_download"} - assert harness.download_manager.calls[-2:] == [ + assert harness.download_manager.calls[-3:] == [ ("pause_download", None), ("resume_download", None), + ("stop_download", None), ] @@ -309,6 +317,10 @@ async def test_download_handler_methods_delegate() -> None: self.calls.append(("resume_download", request)) return {"status": "running"} + async def stop_download(self, request) -> dict: + self.calls.append(("stop_download", request)) + return {"status": "stopping"} + async def start_force_download(self, payload) -> dict: self.calls.append(("start_force_download", payload)) return {"status": "force", "payload": payload} @@ -342,6 +354,8 @@ async def test_download_handler_methods_delegate() -> None: assert json.loads(pause_response.text) == {"status": "paused"} resume_response = await handler.resume_example_images(request) assert json.loads(resume_response.text) == {"status": "running"} + stop_response = await handler.stop_example_images(request) + assert json.loads(stop_response.text) == {"status": "stopping"} force_response = await handler.force_download_example_images(request) assert json.loads(force_response.text) == {"status": "force", "payload": {"foo": "bar"}} @@ -350,6 +364,7 @@ async def test_download_handler_methods_delegate() -> None: ("get_status", request), ("pause_download", request), ("resume_download", request), + ("stop_download", request), ("start_force_download", {"foo": "bar"}), ] @@ -460,6 +475,7 @@ def test_handler_set_route_mapping_includes_all_handlers() -> None: "get_example_images_status", "pause_example_images", "resume_example_images", + "stop_example_images", "force_download_example_images", "import_example_images", "delete_example_image", diff --git a/tests/utils/test_example_images_download_manager_unit.py b/tests/utils/test_example_images_download_manager_unit.py index 42d88bf8..f6617ba6 100644 --- a/tests/utils/test_example_images_download_manager_unit.py +++ b/tests/utils/test_example_images_download_manager_unit.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import time from typing import Any, Dict import pytest @@ -128,6 +129,59 @@ async def test_pause_and_resume_flow(monkeypatch: pytest.MonkeyPatch, tmp_path) await asyncio.wait_for(task, timeout=1) +async def test_stop_download_transitions_to_stopped(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + settings_manager = get_settings_manager() + settings_manager.settings["example_images_path"] = str(tmp_path) + settings_manager.settings["libraries"] = {"default": {}} + settings_manager.settings["active_library"] = "default" + + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + started = asyncio.Event() + release = asyncio.Event() + + async def fake_download(self, *_args): + started.set() + await release.wait() + async with self._state_lock: + if self._stop_requested and self._progress['status'] == 'stopping': + self._progress['status'] = 'stopped' + else: + self._progress['status'] = 'completed' + self._progress['end_time'] = time.time() + self._stop_requested = False + await self._broadcast_progress(status=self._progress['status']) + async with self._state_lock: + self._is_downloading = False + self._download_task = None + + monkeypatch.setattr( + download_module.DownloadManager, + "_download_all_example_images", + fake_download, + ) + + await manager.start_download({}) + await asyncio.wait_for(started.wait(), timeout=1) + + stop_response = await manager.stop_download(object()) + assert stop_response == {"success": True, "message": "Download stopping"} + assert manager._progress["status"] == "stopping" + + task = manager._download_task + assert task is not None + release.set() + await asyncio.wait_for(task, timeout=1) + + assert manager._progress["status"] == "stopped" + assert manager._is_downloading is False + assert manager._stop_requested is False + statuses = [payload["status"] for payload in ws_manager.payloads] + assert "stopping" in statuses + assert "stopped" in statuses + + async def test_pause_or_resume_without_running_download(monkeypatch: pytest.MonkeyPatch) -> None: manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager()) @@ -136,3 +190,6 @@ async def test_pause_or_resume_without_running_download(monkeypatch: pytest.Monk with pytest.raises(download_module.DownloadNotRunningError): await manager.resume_download(object()) + + with pytest.raises(download_module.DownloadNotRunningError): + await manager.stop_download(object())