From 2a2e5c3592e5a0616674572c2710a492be8d9cb3 Mon Sep 17 00:00:00 2001 From: jam-cai Date: Thu, 23 Apr 2026 14:56:35 -0400 Subject: [PATCH] feat: Add cancel control to in-flight model downloads (#6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires Swift Task cancellation through to URLSession so users can stop a long model download from the UI. Until now, Task.cancel() only flipped Task.isCancelled — the URLSession download kept running until natural completion, wasting bytes and silently ignoring the user's intent. Changes: - ModelDownloadManager.cancel(filename:) — public, idempotent. - ModelDownloadSessionDelegate.download(from:) wraps the continuation in withTaskCancellationHandler. The onCancel block calls .cancel() on the retained URLSessionDownloadTask, which fires didCompleteWithError with URLError.cancelled and resumes the continuation. - New pure helper DownloadOutcomeClassifier distinguishes user-cancel (CancellationError or URLError.cancelled) from real failures, so performDownload's catch routes the two correctly: - cancel -> restore .idle (or .downloaded if a prior copy is on disk) - failure -> .failed(message), shown to the user - DownloadableModelRow renders an xmark.circle.fill cancel button next to the percentage during .downloading; tooltip explains. Tests (7 new, 32 total): - DownloadOutcomeClassifierTests covers cancel surfaces (CancellationError, URLError.cancelled), real failure surfaces (timedOut, notConnected, badServerResponse, generic NSError), and the LlamaRuntimeError case (must NOT be misclassified as cancel — would silently roll back a real failure to .idle). Verified locally: xcodebuild test -project tabby.xcodeproj -scheme tabby \ -destination 'platform=macOS' CODE_SIGNING_ALLOWED=NO # ** TEST SUCCEEDED ** 32 tests 0 failures First of three stacked PRs against #6: 1. (this PR) cancel + cleanup 2. size + SHA-256 validation 3. retry UX polish Refs #6 --- tabby.xcodeproj/project.pbxproj | 13 ++-- .../Utilities/DownloadOutcomeClassifier.swift | 32 +++++++++ .../Utilities/ModelDownloadManager.swift | 65 ++++++++++++++++--- tabby/UI/DownloadableModelCatalogView.swift | 37 ++++++++--- .../DownloadOutcomeClassifierTests.swift | 50 ++++++++++++++ 5 files changed, 168 insertions(+), 29 deletions(-) create mode 100644 tabby/Services/Utilities/DownloadOutcomeClassifier.swift create mode 100644 tabbyTests/DownloadOutcomeClassifierTests.swift diff --git a/tabby.xcodeproj/project.pbxproj b/tabby.xcodeproj/project.pbxproj index 7470fb39..f31fcb7c 100644 --- a/tabby.xcodeproj/project.pbxproj +++ b/tabby.xcodeproj/project.pbxproj @@ -7,6 +7,7 @@ objects = { /* Begin PBXBuildFile section */ + 8B6282F0C1CCA0746D96B914 /* DownloadOutcomeClassifierTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 562F89255AF340C15A0554BE /* DownloadOutcomeClassifierTests.swift */; }; A1C3E0012F90000100AAA001 /* LlamaSwift in Frameworks */ = {isa = PBXBuildFile; productRef = A1C3E0002F90000100AAA001 /* LlamaSwift */; }; A1C3E0112F90000100AAA001 /* Sparkle in Frameworks */ = {isa = PBXBuildFile; productRef = A1C3E0102F90000100AAA001 /* Sparkle */; }; A404828463CADB2ECDAE7AF3 /* LlamaPromptRendererTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = F29623C5C0A67B992D383A3C /* LlamaPromptRendererTests.swift */; }; @@ -27,6 +28,7 @@ /* Begin PBXFileReference section */ 193741492F81DE7000BEC04F /* tabby.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = tabby.app; sourceTree = BUILT_PRODUCTS_DIR; }; 3FBFA92FA44AA317135426FB /* tabbyTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = tabbyTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; }; + 562F89255AF340C15A0554BE /* DownloadOutcomeClassifierTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = DownloadOutcomeClassifierTests.swift; sourceTree = ""; }; 5A39EC1A44E9160E13E2AE77 /* SuggestionAvailabilityEvaluatorTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionAvailabilityEvaluatorTests.swift; sourceTree = ""; }; 8B1121FFDAD30C7F62E15289 /* SuggestionRequestFactoryTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionRequestFactoryTests.swift; sourceTree = ""; }; F29623C5C0A67B992D383A3C /* LlamaPromptRendererTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = LlamaPromptRendererTests.swift; sourceTree = ""; }; @@ -67,7 +69,6 @@ children = ( 1937414B2F81DE7000BEC04F /* tabby */, 1937414A2F81DE7000BEC04F /* Products */, - 98425D3A1DB21E10B0F3BFA6 /* Frameworks */, 38079AF58386B5C2E9373742 /* tabbyTests */, ); sourceTree = ""; @@ -87,18 +88,11 @@ F29623C5C0A67B992D383A3C /* LlamaPromptRendererTests.swift */, 5A39EC1A44E9160E13E2AE77 /* SuggestionAvailabilityEvaluatorTests.swift */, 8B1121FFDAD30C7F62E15289 /* SuggestionRequestFactoryTests.swift */, + 562F89255AF340C15A0554BE /* DownloadOutcomeClassifierTests.swift */, ); - name = tabbyTests; path = tabbyTests; sourceTree = ""; }; - 98425D3A1DB21E10B0F3BFA6 /* Frameworks */ = { - isa = PBXGroup; - children = ( - ); - name = Frameworks; - sourceTree = ""; - }; /* End PBXGroup section */ /* Begin PBXNativeTarget section */ @@ -215,6 +209,7 @@ A404828463CADB2ECDAE7AF3 /* LlamaPromptRendererTests.swift in Sources */, B053492719F6106E48290C32 /* SuggestionAvailabilityEvaluatorTests.swift in Sources */, ADEFEE12C197DB6C990E3812 /* SuggestionRequestFactoryTests.swift in Sources */, + 8B6282F0C1CCA0746D96B914 /* DownloadOutcomeClassifierTests.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; diff --git a/tabby/Services/Utilities/DownloadOutcomeClassifier.swift b/tabby/Services/Utilities/DownloadOutcomeClassifier.swift new file mode 100644 index 00000000..ba67d303 --- /dev/null +++ b/tabby/Services/Utilities/DownloadOutcomeClassifier.swift @@ -0,0 +1,32 @@ +import Foundation + +/// File overview: +/// Classifies download errors into "user pressed cancel" vs "something genuinely +/// went wrong." Used by `ModelDownloadManager` to decide whether a failed +/// download should restore the prior state (cancel) or surface as `.failed`. +/// +/// Why this is its own type: +/// Swift's `Task.cancel()` triggers two distinct error shapes by the time we +/// catch them downstream: +/// +/// 1. `CancellationError` — when `Task.checkCancellation()` runs *before* the +/// URLSession download even starts. +/// 2. `URLError(.cancelled)` — when the URLSession download is in flight and +/// our `withTaskCancellationHandler` aborts it via +/// `URLSessionDownloadTask.cancel()`. +/// +/// Without this classification, case (2) would route through the catch-all +/// failure path and the user would see "The operation couldn't be completed" +/// despite having pressed Cancel themselves. This helper makes the +/// distinction testable in isolation, with no URLSession or Task setup. +enum DownloadOutcomeClassifier { + static func isUserCancellation(_ error: Error) -> Bool { + if error is CancellationError { + return true + } + if let urlError = error as? URLError, urlError.code == .cancelled { + return true + } + return false + } +} diff --git a/tabby/Services/Utilities/ModelDownloadManager.swift b/tabby/Services/Utilities/ModelDownloadManager.swift index 0a4387fb..f9faf525 100644 --- a/tabby/Services/Utilities/ModelDownloadManager.swift +++ b/tabby/Services/Utilities/ModelDownloadManager.swift @@ -126,6 +126,27 @@ final class ModelDownloadManager: ObservableObject { downloadTasks[model.filename] = task } + /// User-initiated cancel of an in-flight model download. Idempotent — + /// calling it on a filename that isn't downloading is a safe no-op. + /// + /// Cancellation flow: + /// 1. `Task.cancel()` flips `Task.isCancelled` and triggers the + /// `withTaskCancellationHandler` block in the delegate. + /// 2. That block calls `URLSessionDownloadTask.cancel()`, which aborts + /// the in-flight download. + /// 3. The delegate receives `didCompleteWithError(URLError.cancelled)` + /// and resumes the continuation throwing. + /// 4. `performDownload`'s catch routes the error through + /// `DownloadOutcomeClassifier`, sees a user cancel, and restores + /// `.idle` (or `.downloaded` if a prior copy is on disk) — never + /// `.failed`, since the user pressed Cancel deliberately. + func cancel(filename: String) { + guard let task = downloadTasks[filename] else { + return + } + task.cancel() + } + func openModelsDirectory() { do { try ensureRuntimeDirectoryExists() @@ -191,14 +212,18 @@ final class ModelDownloadManager: ObservableObject { modelStates[model.filename] = .downloaded onModelDirectoryChanged?() - } catch is CancellationError { - if isInstalled(model: model) { - modelStates[model.filename] = .downloaded + } catch { + // A user-initiated cancel surfaces here as either CancellationError + // (cancelled before URLSession ran) or URLError.cancelled + // (cancelled while in flight). Both should restore the prior + // visible state, not show a failure — the user already knows what + // they did. DownloadOutcomeClassifier owns the discrimination so + // the rule is unit-tested in isolation. + if DownloadOutcomeClassifier.isUserCancellation(error) { + modelStates[model.filename] = isInstalled(model: model) ? .downloaded : .idle } else { - modelStates[model.filename] = .idle + modelStates[model.filename] = .failed(error.localizedDescription) } - } catch { - modelStates[model.filename] = .failed(error.localizedDescription) } } @@ -250,6 +275,11 @@ private final class ModelDownloadSessionDelegate: NSObject, URLSessionDownloadDe private var downloadedFileURL: URL? private var response: URLResponse? private var hasCompleted = false + // Held so `withTaskCancellationHandler` can call .cancel() on it when the + // surrounding Swift Task is cancelled. Without this, Task.cancel() would + // only flip Task.isCancelled — the URLSession download would keep running + // until natural completion, wasting bytes and ignoring the user's intent. + private var activeDownloadTask: URLSessionDownloadTask? init(progressHandler: @escaping @Sendable (Double?) -> Void) { self.progressHandler = progressHandler @@ -260,10 +290,25 @@ private final class ModelDownloadSessionDelegate: NSObject, URLSessionDownloadDe configuration.requestCachePolicy = .reloadIgnoringLocalCacheData let session = URLSession(configuration: configuration, delegate: self, delegateQueue: nil) - return try await withCheckedThrowingContinuation { continuation in - self.continuation = continuation - let task = session.downloadTask(with: url) - task.resume() + // withTaskCancellationHandler bridges Swift Task cancellation into the + // URLSession world. When `Task.cancel()` runs upstream (e.g., from + // ModelDownloadManager.cancel(filename:)), the onCancel block fires and + // aborts the URLSession download task. The delegate then receives + // didCompleteWithError(URLError.cancelled), which resumes the + // continuation throwing — and the catch in performDownload routes it + // through DownloadOutcomeClassifier as a user cancel rather than a + // hard failure. + return try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + self.continuation = continuation + let task = session.downloadTask(with: url) + self.activeDownloadTask = task + task.resume() + } + } onCancel: { [weak self] in + // URLSessionDownloadTask.cancel() is thread-safe by Apple's docs, + // so calling it from arbitrary cancellation contexts is fine. + self?.activeDownloadTask?.cancel() } } diff --git a/tabby/UI/DownloadableModelCatalogView.swift b/tabby/UI/DownloadableModelCatalogView.swift index a39b68c9..74bfb8f2 100644 --- a/tabby/UI/DownloadableModelCatalogView.swift +++ b/tabby/UI/DownloadableModelCatalogView.swift @@ -18,7 +18,8 @@ struct DownloadableModelCatalogView: View { DownloadableModelRow( model: model, state: modelDownloadManager.state(for: model), - onDownload: { modelDownloadManager.download(model) } + onDownload: { modelDownloadManager.download(model) }, + onCancel: { modelDownloadManager.cancel(filename: model.filename) } ) } @@ -58,6 +59,7 @@ private struct DownloadableModelRow: View { let model: DownloadableRuntimeModel let state: ModelDownloadState let onDownload: () -> Void + let onCancel: () -> Void var body: some View { VStack(alignment: .leading, spacing: 8) { @@ -100,15 +102,30 @@ private struct DownloadableModelRow: View { .buttonStyle(.borderedProminent) .controlSize(.small) case .downloading(let progress): - if let progress { - Text("\(Int((progress * 100).rounded()))%") - .font(.system(size: 11, weight: .medium, design: .monospaced)) - .foregroundStyle(.blue) - .frame(width: 40, alignment: .trailing) - } else { - ProgressView() - .controlSize(.small) - .frame(width: 40) + HStack(spacing: 6) { + if let progress { + Text("\(Int((progress * 100).rounded()))%") + .font(.system(size: 11, weight: .medium, design: .monospaced)) + .foregroundStyle(.blue) + .frame(width: 40, alignment: .trailing) + } else { + ProgressView() + .controlSize(.small) + .frame(width: 40) + } + // Plain SF Symbol button keeps the row compact and matches + // the "Get"/"Retry" button affordance scale. Cancel is the + // affirmative action while downloading, so it gets the same + // visual weight as Get does in the idle state. + Button { + onCancel() + } label: { + Image(systemName: "xmark.circle.fill") + .font(.system(size: 16)) + .foregroundStyle(.secondary) + } + .buttonStyle(.plain) + .help("Cancel download") } case .downloaded: Image(systemName: "checkmark.circle.fill") diff --git a/tabbyTests/DownloadOutcomeClassifierTests.swift b/tabbyTests/DownloadOutcomeClassifierTests.swift new file mode 100644 index 00000000..ad6f1f12 --- /dev/null +++ b/tabbyTests/DownloadOutcomeClassifierTests.swift @@ -0,0 +1,50 @@ +import XCTest +@testable import tabby + +/// Tests for the rule that decides whether a download error is "the user +/// pressed Cancel" or "something went genuinely wrong." +/// +/// This classification matters because the two cases produce different errors +/// at runtime (CancellationError vs URLError.cancelled depending on whether +/// the URLSession download had started yet) but both should restore the prior +/// state, never surface as a user-visible failure. +final class DownloadOutcomeClassifierTests: XCTestCase { + + // MARK: - cancellation surfaces + + func test_isUserCancellation_trueForCancellationError() { + XCTAssertTrue(DownloadOutcomeClassifier.isUserCancellation(CancellationError())) + } + + func test_isUserCancellation_trueForURLErrorCancelled() { + XCTAssertTrue(DownloadOutcomeClassifier.isUserCancellation(URLError(.cancelled))) + } + + // MARK: - real failures + + func test_isUserCancellation_falseForURLErrorTimedOut() { + XCTAssertFalse(DownloadOutcomeClassifier.isUserCancellation(URLError(.timedOut))) + } + + func test_isUserCancellation_falseForURLErrorNotConnected() { + XCTAssertFalse(DownloadOutcomeClassifier.isUserCancellation(URLError(.notConnectedToInternet))) + } + + func test_isUserCancellation_falseForURLErrorBadServerResponse() { + XCTAssertFalse(DownloadOutcomeClassifier.isUserCancellation(URLError(.badServerResponse))) + } + + func test_isUserCancellation_falseForGenericNSError() { + let error = NSError(domain: "TestDomain", code: 42, userInfo: nil) + XCTAssertFalse(DownloadOutcomeClassifier.isUserCancellation(error)) + } + + /// Important: a domain-level runtime error (model unavailable, etc.) must + /// NOT be misclassified as a cancellation. Otherwise a real failure would + /// silently roll back to .idle and the user would never see the problem. + func test_isUserCancellation_falseForLlamaRuntimeError() { + XCTAssertFalse(DownloadOutcomeClassifier.isUserCancellation( + LlamaRuntimeError.unavailable("Model download failed with status code 500.") + )) + } +}