diff --git a/core/download/api.py b/core/download/api.py index 7a747da..60ce26e 100644 --- a/core/download/api.py +++ b/core/download/api.py @@ -230,6 +230,7 @@ context._run_huggingface_xet_transfer = _bind( ) context.download_file_with_aria2 = _bind(aria2_backend.download_file_with_aria2) +context.delete_badsha_download = _bind(aria2_backend.delete_badsha_download) context.verify_download_sha256_with_progress = _bind( aria2_backend.verify_download_sha256_with_progress ) @@ -355,10 +356,19 @@ def clear_completed_downloads(): return state.clear_completed_downloads(_state_dependencies()) +def delete_badsha_download(download_id: str = "", path: Optional[str] = None) -> Dict[str, Any]: + return aria2_backend.delete_badsha_download( + download_id=download_id, + path=path, + dependencies=context, + ) + + context.get_progress = get_progress context.get_all_progress = get_all_progress context.cancel_download = cancel_download context.clear_completed_downloads = clear_completed_downloads +context.delete_badsha_download = delete_badsha_download class _BoundHuggingFaceXetProgressAdapter( @@ -409,6 +419,7 @@ get_existing_model_preview_path = context.get_existing_model_preview_path __all__ = [ "cancel_download", "clear_completed_downloads", + "delete_badsha_download", "download_file", "download_model", "get_all_progress", diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index 9c2a6fa..d02e000 100644 --- a/core/download/aria2_backend.py +++ b/core/download/aria2_backend.py @@ -301,6 +301,102 @@ def delete_xet_partial_file( ) +def delete_badsha_download( + download_id: str = "", + path: Optional[str] = None, + *, + dependencies: Any = None, +) -> Dict[str, Any]: + """Delete a corrupted download with .badsha suffix and its associated metadata.""" + facade = require_download_dependencies(dependencies, "aria2 backend") + target_bad_path = "" + + if download_id: + with facade.download_lock: + progress = facade.download_progress.get(download_id, {}) + target_bad_path = progress.get("bad_path") or "" + if not target_bad_path: + candidate = progress.get("path") or "" + if candidate.lower().endswith(".badsha"): + target_bad_path = candidate + elif candidate and facade.os.path.isfile(f"{candidate}.badsha"): + target_bad_path = f"{candidate}.badsha" + + if not target_bad_path and path: + target_bad_path = str(path).strip() + + if not target_bad_path: + return {"success": False, "error": "No .badsha file path found"} + + normalized = facade.os.path.normpath(target_bad_path) + if not normalized.lower().endswith(".badsha"): + return { + "success": False, + "error": "Only files ending with .badsha can be deleted through this operation", + } + + check_roots = getattr(facade, "is_path_in_configured_model_roots", None) + if callable(check_roots): + in_model_roots = check_roots(normalized) + if not in_model_roots: + import tempfile + temp_dir = facade.os.path.normpath(tempfile.gettempdir()).lower() + if not normalized.lower().startswith(temp_dir): + return { + "success": False, + "error": "Access denied: file is outside configured model directories", + } + + removed_files = [] + candidates = [ + normalized, + f"{normalized}.aria2", + f"{normalized}.modelresolver.json", + ] + if normalized.lower().endswith(".badsha"): + original_base = normalized[:-7] + sidecar_orig = facade.get_model_resolver_sidecar_path(original_base) + if sidecar_orig and facade.os.path.isfile(sidecar_orig): + meta = facade.read_json_safe(sidecar_orig, {}) + if isinstance(meta, dict) and meta.get("sha256_verified") is False: + candidates.append(sidecar_orig) + + for p in candidates: + if facade.os.path.isfile(p): + try: + facade.os.remove(p) + removed_files.append(p) + log.info(f"Deleted corrupted download artifact: {p}") + except OSError as exc: + log.warning(f"Could not remove {p}: {exc}") + + if download_id: + with facade.download_lock: + if download_id in facade.download_progress: + facade.download_progress.pop(download_id, None) + facade.cancelled_downloads.discard(download_id) + + remove_persisted = getattr(facade, "remove_persisted_download", None) + if callable(remove_persisted) and download_id: + try: + remove_persisted(download_id) + except Exception: + pass + + invalidate_fn = getattr(facade, "invalidate_model_caches", None) + if callable(invalidate_fn): + try: + invalidate_fn() + except Exception: + pass + + return { + "success": True, + "removed_files": removed_files, + "path": normalized, + } + + def aria2_action_error_is_ok(status: str, message: str) -> bool: """Recognize idempotent aria2 pause/resume errors.""" lowered = str(message or "").lower() diff --git a/core/download/orchestrator.py b/core/download/orchestrator.py index 6217ee8..b598bec 100644 --- a/core/download/orchestrator.py +++ b/core/download/orchestrator.py @@ -680,6 +680,8 @@ def download_file( download_progress[download_id]["status"] = "error" download_progress[download_id]["error"] = error_msg download_progress[download_id]["sha256_verified"] = False + if bad_path: + download_progress[download_id]["bad_path"] = bad_path result.update( { "error": error_msg, diff --git a/core/routes/downloads.py b/core/routes/downloads.py index e212933..88c69df 100644 --- a/core/routes/downloads.py +++ b/core/routes/downloads.py @@ -58,6 +58,22 @@ def register_download_routes(context: RouteContext): operation=service.clear_completed_downloads, return_success_on_error=True, ) + register_service_route( + context, + path="/model_resolver/delete_bad_download/{download_id}", + methods=("post",), + error_prefix="delete bad download", + operation=service.delete_bad_download, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/delete_bad_download", + methods=("post",), + error_prefix="delete bad download", + operation=service.delete_bad_download, + return_success_on_error=True, + ) register_service_route( context, path="/model_resolver/aria2/status", diff --git a/core/routes/registry.py b/core/routes/registry.py index c484232..5277452 100644 --- a/core/routes/registry.py +++ b/core/routes/registry.py @@ -155,6 +155,7 @@ def register_routes(self): from ..download.api import ( cancel_download, clear_completed_downloads, + delete_badsha_download, get_all_progress, get_aria2_status, get_download_directory, diff --git a/core/services/download_service.py b/core/services/download_service.py index 632da22..e7c8baa 100644 --- a/core/services/download_service.py +++ b/core/services/download_service.py @@ -26,6 +26,7 @@ class DownloadService: self.clear_completed_downloads_fn = context.require( "clear_completed_downloads" ) + self.delete_badsha_download_fn = context.require("delete_badsha_download") self.get_all_progress_fn = context.require("get_all_progress") self.get_aria2_status_fn = context.require("get_aria2_status") self.get_civarchive_model_details = context.require( @@ -382,6 +383,20 @@ class DownloadService: self.clear_completed_downloads_fn() return self.web.json_response({"success": True}) + async def delete_bad_download(self, request): + """Delete a corrupted (.badsha) download file and its metadata.""" + download_id = request.match_info.get("download_id") or "" + path = None + if request.can_read_body: + try: + data = await request.json() + if isinstance(data, dict): + path = data.get("path") or data.get("bad_path") + except Exception: + pass + result = self.delete_badsha_download_fn(download_id=download_id, path=path) + return self._json_result_with_success_status(result) + async def aria2_status(self, request): """Report aria2 availability using optional request settings.""" try: diff --git a/tests/test_downloader_boundaries.py b/tests/test_downloader_boundaries.py index 9363042..9fdcdae 100644 --- a/tests/test_downloader_boundaries.py +++ b/tests/test_downloader_boundaries.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock, call, patch from core.download.api import context as downloader from core.download.aria2_backend import ( + delete_badsha_download, download_file_with_aria2, drain_aria2_stderr, is_aria2_important_stderr, @@ -738,6 +739,49 @@ class DownloaderBoundaryTests(unittest.TestCase): progress = downloader.download_progress.get(download_id, {}) self.assertEqual("cancelled", progress.get("status")) + def test_delete_badsha_download_removes_badsha_file_and_sidecar(self): + with tempfile.TemporaryDirectory() as temp_dir: + bad_path = os.path.join(temp_dir, "test_corrupt.safetensors.badsha") + with open(bad_path, "wb") as f: + f.write(b"corrupt data") + + aria2_path = f"{bad_path}.aria2" + with open(aria2_path, "wb") as f: + f.write(b"aria2 state") + + bad_sidecar = f"{bad_path}.modelresolver.json" + downloader.write_json_atomic(bad_sidecar, {"sha256_verified": False}) + + orig_sidecar = os.path.join(temp_dir, "test_corrupt.safetensors.modelresolver.json") + downloader.write_json_atomic(orig_sidecar, {"sha256_verified": False}) + + download_id = "test-delete-badsha-1" + with downloader.download_lock: + downloader.download_progress[download_id] = { + "status": "error", + "bad_path": bad_path, + } + + result = delete_badsha_download(download_id=download_id, dependencies=downloader) + self.assertTrue(result["success"]) + self.assertFalse(os.path.exists(bad_path)) + self.assertFalse(os.path.exists(aria2_path)) + self.assertFalse(os.path.exists(bad_sidecar)) + self.assertFalse(os.path.exists(orig_sidecar)) + with downloader.download_lock: + self.assertNotIn(download_id, downloader.download_progress) + + def test_delete_badsha_download_rejects_non_badsha_file(self): + with tempfile.TemporaryDirectory() as temp_dir: + safe_model = os.path.join(temp_dir, "safe_model.safetensors") + with open(safe_model, "wb") as f: + f.write(b"safe model data") + + result = delete_badsha_download(path=safe_model, dependencies=downloader) + self.assertFalse(result["success"]) + self.assertIn("only files ending with .badsha", result["error"].lower()) + self.assertTrue(os.path.exists(safe_model)) + def test_parse_retry_after(self): self.assertEqual(parse_retry_after("120"), 120) self.assertEqual(parse_retry_after("0"), 1) diff --git a/web/css/resolver-shell.css b/web/css/resolver-shell.css index 8c56d42..9e03617 100644 --- a/web/css/resolver-shell.css +++ b/web/css/resolver-shell.css @@ -979,6 +979,20 @@ color: var(--mr-color-danger-bright); } +#model-resolver-modal .mr-download-queue-delete-badsha { + color: var(--mr-color-danger-bright, #ff4d4f); +} + +#model-resolver-modal .mr-download-queue-delete-badsha:hover { + color: #ff7875; + background: rgba(239, 68, 68, 0.16); +} + +.mr-status .mr-delete-badsha-btn { + margin-left: auto; + flex-shrink: 0; +} + .mr-download-queue-progress { gap: var(--mr-spacing-8); margin-top: var(--mr-spacing-3); diff --git a/web/resolver/actions/queue_methods.js b/web/resolver/actions/queue_methods.js index acced26..a01f2d4 100644 --- a/web/resolver/actions/queue_methods.js +++ b/web/resolver/actions/queue_methods.js @@ -1161,6 +1161,7 @@ export const queueMethods = { pause: '.mr-download-queue-pause', resume: '.mr-download-queue-resume', retry: '.mr-download-queue-retry', + deleteBadsha: '.mr-download-queue-delete-badsha', openFolder: '.mr-download-queue-open-folder', switchWorkflow: '.mr-download-queue-switch-workflow', more: '.mr-download-queue-more', @@ -1171,6 +1172,7 @@ export const queueMethods = { onPause: downloadId => this.pauseDownload(downloadId), onResume: downloadId => this.resumeDownload(downloadId), onRetry: downloadId => this.retryDownload?.(downloadId), + onDeleteBadsha: downloadId => this.deleteBadDownload?.(downloadId), onOpenFolder: context => this.openContainingFolder?.(context), onSwitchWorkflow: context => this.switchToDownloadWorkflow?.(context), moreBinding: 'click', @@ -1293,6 +1295,9 @@ export const queueMethods = { html += `
`; if (isError) { html += actionButton('mr-download-queue-retry', 'Retry download', 'play'); + if (isShaMismatch || Boolean(progress.bad_path)) { + html += actionButton('mr-download-queue-delete-badsha', 'Delete corrupted (.badsha) file', 'trash2'); + } } else { html += actionButton( isPaused ? 'mr-download-queue-resume' : 'mr-download-queue-pause', diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index f3a6a47..9b0eb05 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -392,9 +392,19 @@ export const resolveDownloadMethods = { renderDownloadStatusMessage(message, type, progress = {}, info = {}) { const contextMenuModel = this.getDownloadFolderContext(progress, info); + const errorMessage = String(progress.error || info.statusSnapshot?.message || message || '').toLowerCase(); + const hasBadsha = Boolean(progress.bad_path || errorMessage.includes('sha256') || errorMessage.includes('.badsha')); + const downloadId = info.downloadId || progress.download_id || ''; + let extraHtml = ''; + if (type === 'error' && hasBadsha && downloadId) { + const safeDownloadId = this.escapeHtml(downloadId); + const deleteLabel = 'Delete corrupted (.badsha) file'; + extraHtml = ``; + } return this.renderStatusMessage(message, type, { contextMenuModel, - contextMenuTooltip: 'Right-click to open download folder' + contextMenuTooltip: 'Right-click to open download folder', + extraHtml, }); }, @@ -1112,6 +1122,47 @@ export const resolveDownloadMethods = { return { progressDiv, downloadBtn }; }, + async deleteBadDownload(downloadId) { + if (!downloadId) return; + const info = this.activeDownloads?.[downloadId]; + const badPath = info?.lastProgress?.bad_path || info?.statusSnapshot?.progress?.bad_path || ''; + try { + const res = await this.fetchJson( + `/model_resolver/delete_bad_download/${encodeURIComponent(downloadId)}`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ path: badPath }) + }, + 'Delete corrupted download' + ); + if (res && res.success) { + this.showNotification('Corrupted (.badsha) file deleted', 'info'); + delete this.activeDownloads?.[downloadId]; + this.removeActiveDownloadRecovery?.(downloadId); + this.updateDownloadAllButtonState?.(); + this.updateQueuePanel?.(); + if (info?.missing) { + const elements = this.resolveDownloadUiElements(info); + if (elements.progressDiv) { + elements.progressDiv.classList.add('mr-is-hidden'); + elements.progressDiv.classList.remove('mr-is-visible'); + elements.progressDiv.innerHTML = ''; + } + if (elements.downloadBtn) { + this.restoreDownloadButtonReadyState(elements.downloadBtn); + } + this.refreshLocalMatchesUiForMissing?.(info.missing); + } + } else { + this.showNotification(res?.error || 'Failed to delete corrupted file', 'error'); + } + } catch (err) { + console.error('Model Resolver: Failed to delete corrupted download:', err); + this.showNotification('Failed to delete corrupted file: ' + err.message, 'error'); + } + }, + attachDownloadActionHandlers(progressDiv, downloadId) { if (!progressDiv) return; const getDownloadContext = (button) => { @@ -1133,6 +1184,7 @@ export const resolveDownloadMethods = { pause: '.pause-download-btn, .mr-download-queue-pause', resume: '.resume-download-btn, .mr-download-queue-resume', retry: '.retry-download-btn, .mr-download-queue-retry', + deleteBadsha: '.mr-download-queue-delete-badsha, .mr-delete-badsha-btn', openFolder: '.mr-download-queue-open-folder', switchWorkflow: '.mr-download-queue-switch-workflow', more: '.mr-download-queue-more', @@ -1142,6 +1194,7 @@ export const resolveDownloadMethods = { onPause: targetDownloadId => this.pauseDownload(targetDownloadId), onResume: targetDownloadId => this.resumeDownload(targetDownloadId), onRetry: targetDownloadId => this.retryDownload?.(targetDownloadId), + onDeleteBadsha: targetDownloadId => this.deleteBadDownload?.(targetDownloadId), onOpenFolder: context => this.openContainingFolder?.(context), onSwitchWorkflow: context => this.switchToDownloadWorkflow?.(context), onMore: (event, button) => { diff --git a/web/resolver/utils/download_action_handlers.js b/web/resolver/utils/download_action_handlers.js index 00c32a1..e8ce6a5 100644 --- a/web/resolver/utils/download_action_handlers.js +++ b/web/resolver/utils/download_action_handlers.js @@ -28,6 +28,7 @@ export function bindDownloadActionHandlers(root, { onPause = null, onResume = null, onRetry = null, + onDeleteBadsha = null, onOpenFolder = null, onSwitchWorkflow = null, onMore = null, @@ -49,6 +50,7 @@ export function bindDownloadActionHandlers(root, { bindIdAction(selectors.pause, onPause); bindIdAction(selectors.resume, onResume); bindIdAction(selectors.retry, onRetry); + bindIdAction(selectors.deleteBadsha, onDeleteBadsha); bindActionButtons(root, selectors.openFolder, (button) => { const context = resolveContext(button); diff --git a/web/resolver/utils/render_format_methods.js b/web/resolver/utils/render_format_methods.js index 8e68de4..211ecbe 100644 --- a/web/resolver/utils/render_format_methods.js +++ b/web/resolver/utils/render_format_methods.js @@ -141,11 +141,13 @@ export const renderFormatMethods = { const className = contextMenuAttrs ? `mr-status mr-status-${type} mr-download-folder-context` : `mr-status mr-status-${type}`; + const extraHtml = options?.extraHtml ? `${options.extraHtml}` : ''; return `
${icon} - ${message} + ${message} + ${extraHtml}
`; },