spb/zyquo-local Public MIT

1//2//  DownloadManager.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import Observation1112/// Queued, resumable model downloads.13///14/// - Files download to `<name>.partial` and are moved into place atomically;15///   a model is valid only when every file is present and size-verified.16/// - Resume uses HTTP Range from the partial file's size — surviving pause,17///   cancel-restart, and app relaunch (manifest + partials persist on disk).18/// - Max 2 concurrent file downloads per model.19@MainActor20@Observable21final class DownloadManager {22    private(set) var tasks: [DownloadTask] = []23    /// Live transfer speed per repoID, bytes/sec.24    private(set) var speeds: [String: Double] = [:]2526    var hub: HubService27    private let store: ModelStore28    private var workers: [String: Task<Void, Never>] = [:]29    /// Speed window per repoID: last sample time and cumulative bytes.30    private var speedSamples: [String: (time: ContinuousClock.Instant, bytes: Int64)] = [:]3132    /// Active + queued downloads (for the sidebar badge).33    var activeCount: Int {34        tasks.filter { $0.state == .downloading || $0.state == .queued }.count35    }3637    var overallFraction: Double {38        let active = tasks.filter { $0.state == .downloading || $0.state == .paused }39        guard !active.isEmpty else { return 0 }40        return active.reduce(0.0) { $0 + $1.fractionCompleted } / Double(active.count)41    }4243    init(hub: HubService, store: ModelStore) {44        self.hub = hub45        self.store = store46        restoreManifest()47    }4849    enum DownloadError: LocalizedError {50        case insufficientDiskSpace(needed: Int64, available: Int64)51        case sizeMismatch(file: String)52        case interrupted5354        var errorDescription: String? {55            switch self {56            case .insufficientDiskSpace(let needed, let available):57                "Not enough disk space: \(ByteCountFormatter.string(fromByteCount: needed, countStyle: .file)) needed, \(ByteCountFormatter.string(fromByteCount: available, countStyle: .file)) available."58            case .sizeMismatch(let file):59                "Downloaded file \(file) does not match its expected size. The download will retry from scratch for this file."60            case .interrupted:61                "The download was interrupted. It can be resumed from where it left off."62            }63        }64    }6566    // MARK: - Public API6768    /// Starts (or restarts) downloading all required files of a repo.69    func download(repoID: String) async {70        if let existing = task(for: repoID),71            existing.state == .downloading || existing.state == .queued72        {73            return74        }75        do {76            let (files, totalBytes) = try await hub.requiredFiles(of: repoID)77            try checkDiskSpace(needed: totalBytes)78            let dir = store.directory(for: repoID)79            try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)8081            var progress: [DownloadTask.FileProgress] = files.map { file in82                let partial = dir.appendingPathComponent(file.path + ".partial")83                let final = dir.appendingPathComponent(file.path)84                let already: Int6485                let done: Bool86                if FileManager.default.fileExists(atPath: final.path) {87                    already = file.size ?? 088                    done = true89                } else {90                    already = (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? 091                    done = false92                }93                return DownloadTask.FileProgress(94                    path: file.path,95                    totalBytes: file.size ?? 0,96                    receivedBytes: already,97                    sha256: file.lfs?.oid,98                    completed: done99                )100            }101            progress.sort { $0.totalBytes < $1.totalBytes }  // small config files first102            upsert(DownloadTask(103                repoID: repoID, state: .downloading, files: progress,104                errorDescription: nil, startedAt: Date()105            ))106            startWorker(repoID: repoID)107        } catch {108            upsert(DownloadTask(109                repoID: repoID, state: .failed, files: task(for: repoID)?.files ?? [],110                errorDescription: error.localizedDescription, startedAt: Date()111            ))112        }113    }114115    func pause(repoID: String) {116        workers[repoID]?.cancel()117        workers[repoID] = nil118        mutate(repoID) { $0.state = .paused }119    }120121    func resume(repoID: String) {122        guard var t = task(for: repoID), t.state == .paused || t.state == .failed else { return }123        t.state = .downloading124        t.errorDescription = nil125        upsert(t)126        startWorker(repoID: repoID)127    }128129    /// Cancels and removes all partial data.130    func cancel(repoID: String) {131        workers[repoID]?.cancel()132        workers[repoID] = nil133        let dir = store.directory(for: repoID)134        if let t = task(for: repoID) {135            for f in t.files where !f.completed {136                try? FileManager.default.removeItem(at: dir.appendingPathComponent(f.path + ".partial"))137            }138        }139        // Remove the directory when nothing complete remains in it.140        if let contents = try? FileManager.default.contentsOfDirectory(atPath: dir.path), contents.isEmpty {141            try? FileManager.default.removeItem(at: dir)142        }143        tasks.removeAll { $0.repoID == repoID }144        speeds[repoID] = nil145        speedSamples[repoID] = nil146        persistManifest()147    }148149    func task(for repoID: String) -> DownloadTask? {150        tasks.first { $0.repoID == repoID }151    }152153    /// ETA in seconds based on current speed.154    func eta(for repoID: String) -> TimeInterval? {155        guard let t = task(for: repoID), let speed = speeds[t.repoID], speed > 1 else { return nil }156        return Double(t.totalBytes - t.receivedBytes) / speed157    }158159    // MARK: - Worker160161    private func startWorker(repoID: String) {162        workers[repoID]?.cancel()163        workers[repoID] = Task { [weak self] in164            await self?.runWorker(repoID: repoID)165        }166    }167168    private func runWorker(repoID: String) async {169        let dir = store.directory(for: repoID)170        guard let snapshot = task(for: repoID) else { return }171        let pending = snapshot.files.filter { !$0.completed }172173        do {174            try await withThrowingTaskGroup(of: Void.self) { group in175                var iterator = pending.makeIterator()176                var inFlight = 0177                func addNext(_ group: inout ThrowingTaskGroup<Void, Error>) {178                    guard let file = iterator.next() else { return }179                    inFlight += 1180                    group.addTask { [weak self] in181                        try await self?.downloadFile(repoID: repoID, file: file, directory: dir)182                    }183                }184                addNext(&group)185                addNext(&group)  // 2 concurrent file downloads max186                while inFlight > 0 {187                    try await group.next()188                    inFlight -= 1189                    addNext(&group)190                }191            }192            mutate(repoID) { $0.state = .verifying }193            try verify(repoID: repoID, directory: dir)194            mutate(repoID) { $0.state = .completed }195            speeds[repoID] = nil196            store.rescan()197        } catch is CancellationError {198            // pause() already set the state; on app quit partials just persist.199        } catch {200            if Task.isCancelled { return }201            mutate(repoID) {202                $0.state = .failed203                $0.errorDescription = error.localizedDescription204            }205            speeds[repoID] = nil206        }207        workers[repoID] = nil208    }209210    /// Downloads one file with Range-resume into `<path>.partial`, then moves211    /// it into place. Retries transient network drops from the partial offset.212    /// Throws CancellationError on pause.213    private func downloadFile(repoID: String, file: DownloadTask.FileProgress, directory: URL) async throws {214        let partial = directory.appendingPathComponent(file.path + ".partial")215        let final = directory.appendingPathComponent(file.path)216        try FileManager.default.createDirectory(217            at: partial.deletingLastPathComponent(), withIntermediateDirectories: true)218        if !FileManager.default.fileExists(atPath: partial.path) {219            FileManager.default.createFile(atPath: partial.path, contents: nil)220        }221222        var request = URLRequest(url: HubService.resolveURL(repoID: repoID, path: file.path))223        request.setValue("ZyquoLocal/1.0.0", forHTTPHeaderField: "User-Agent")224        if let token = hub.token, !token.isEmpty {225            request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")226        }227228        let path = file.path229        let transfer = FileTransfer(partial: partial) { [weak self] received in230            Task { @MainActor [weak self] in231                self?.updateProgress(repoID: repoID, path: path, received: received)232            }233        }234235        var attempts = 0236        while true {237            attempts += 1238            let offset: Int64 =239                (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? 0240            do {241                _ = try await transfer.run(request: request, offset: offset)242                break243            } catch let error as FileTransfer.TransferError {244                switch error {245                case .badStatus(401), .badStatus(403):246                    throw HubService.HubError.gatedOrMissing(repoID)247                case .badStatus(416):248                    // Range beyond EOF: complete if the size matches, else restart.249                    let size = (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? 0250                    if size == file.totalBytes { break }251                    try? FileManager.default.removeItem(at: partial)252                    FileManager.default.createFile(atPath: partial.path, contents: nil)253                    if attempts >= 3 { throw DownloadError.sizeMismatch(file: file.path) }254                    continue255                case .badStatus(let code):256                    throw HubService.HubError.http(code, repoID)257                }258            } catch is CancellationError {259                throw CancellationError()260            } catch {261                // Transient drop (connection lost, timeout…): retry with backoff262                // from whatever the partial already holds.263                guard attempts < 4 else { throw error }264                try await Task.sleep(for: .seconds(Double(attempts)))265                try Task.checkCancellation()266                continue267            }268        }269        try finishFile(partial: partial, final: final, repoID: repoID, path: file.path, expected: file.totalBytes)270    }271272    /// Size check + atomic move into place.273    private func finishFile(partial: URL, final: URL, repoID: String, path: String, expected: Int64) throws {274        let size = (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? -1275        guard expected == 0 || size == expected else {276            try? FileManager.default.removeItem(at: partial)277            updateProgress(repoID: repoID, path: path, received: 0)278            throw DownloadError.sizeMismatch(file: path)279        }280        if FileManager.default.fileExists(atPath: final.path) {281            try FileManager.default.removeItem(at: final)282        }283        try FileManager.default.moveItem(at: partial, to: final)284        mutate(repoID) { t in285            if let i = t.files.firstIndex(where: { $0.path == path }) {286                t.files[i].completed = true287                t.files[i].receivedBytes = t.files[i].totalBytes288            }289        }290    }291292    /// Final integrity pass: every required file present with expected size.293    private func verify(repoID: String, directory: URL) throws {294        guard let t = task(for: repoID) else { return }295        for f in t.files {296            let url = directory.appendingPathComponent(f.path)297            let size = (try? FileManager.default.attributesOfItem(atPath: url.path)[.size] as? Int64) ?? -1298            guard f.totalBytes == 0 || size == f.totalBytes else {299                throw DownloadError.sizeMismatch(file: f.path)300            }301        }302    }303304    private func checkDiskSpace(needed: Int64) throws {305        let values = try? store.modelsRoot.resourceValues(forKeys: [.volumeAvailableCapacityForImportantUsageKey])306        let available = values?.volumeAvailableCapacityForImportantUsage ?? 0307        // 5 % headroom so a download never fills the disk completely.308        if available < Int64(Double(needed) * 1.05) {309            throw DownloadError.insufficientDiskSpace(needed: needed, available: available)310        }311    }312313    // MARK: - State & manifest314315    private func upsert(_ t: DownloadTask) {316        if let i = tasks.firstIndex(where: { $0.repoID == t.repoID }) {317            tasks[i] = t318        } else {319            tasks.append(t)320        }321        persistManifest()322    }323324    private func mutate(_ repoID: String, _ change: (inout DownloadTask) -> Void) {325        guard let i = tasks.firstIndex(where: { $0.repoID == repoID }) else { return }326        change(&tasks[i])327        persistManifest()328    }329330    private func updateProgress(repoID: String, path: String, received: Int64) {331        guard let i = tasks.firstIndex(where: { $0.repoID == repoID }) else { return }332        if let j = tasks[i].files.firstIndex(where: { $0.path == path }) {333            tasks[i].files[j].receivedBytes = received334        }335        // Smoothed speed over a sliding window of the whole repo's bytes.336        let totalReceived = tasks[i].receivedBytes337        let now = ContinuousClock.now338        if let sample = speedSamples[repoID] {339            let d = sample.time.duration(to: now)340            let elapsed = Double(d.components.seconds) + Double(d.components.attoseconds) * 1e-18341            if elapsed >= 1.0 {342                let instant = Double(totalReceived - sample.bytes) / elapsed343                let previous = speeds[repoID] ?? instant344                speeds[repoID] = previous * 0.6 + instant * 0.4345                speedSamples[repoID] = (now, totalReceived)346            }347        } else {348            speedSamples[repoID] = (now, totalReceived)349        }350    }351352    private var manifestURL: URL {353        PersistenceService.appSupportDirectory.appendingPathComponent("downloads.json")354    }355356    private func persistManifest() {357        let persistable = tasks.filter { $0.state != .completed && $0.state != .cancelled }358        try? JSONEncoder().encode(persistable).write(to: manifestURL, options: .atomic)359    }360361    /// Restores paused/interrupted downloads after relaunch.362    private func restoreManifest() {363        guard let data = try? Data(contentsOf: manifestURL),364            var restored = try? JSONDecoder().decode([DownloadTask].self, from: data)365        else { return }366        for i in restored.indices where restored[i].state == .downloading || restored[i].state == .queued {367            restored[i].state = .paused368        }369        tasks = restored370    }371}