spb/zyquo-mlx Public MIT

1//2//  DownloadManager.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// Progress of one repo download.12struct DownloadProgress: Sendable {13    var repo: String14    var completedBytes: Int6415    var totalBytes: Int6416    var currentFile: String17    var fraction: Double { totalBytes > 0 ? Double(completedBytes) / Double(totalBytes) : 0 }18}1920enum DownloadEvent: Sendable {21    case progress(DownloadProgress)22    case finished(LocalModel)23    case failed(String)24}2526/// Downloads Hub repos into the model library with resumable, Range-based27/// transfers (contract verified in docs/MODELS.md §1.3: fresh resolve URL +28/// `Range: bytes=N-` → 206; CDN URLs expire so we never persist them).29actor DownloadManager {3031    static let shared = DownloadManager()3233    private var activeTasks: [String: Task<Void, Never>] = [:]3435    var activeRepos: [String] { Array(activeTasks.keys) }3637    /// Download a full repo. Partially-downloaded files resume from their38    /// current byte count.39    func download(repo: String) -> AsyncStream<DownloadEvent> {40        let (stream, continuation) = AsyncStream.makeStream(of: DownloadEvent.self)4142        let task = Task {43            let directory = PersistenceService.modelsDirectory44                .appendingPathComponent(repo.replacingOccurrences(of: "/", with: "--"), isDirectory: true)45            do {46                try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)47                let files = try await HubService.shared.files(repo: repo)48                    .filter { !$0.path.hasPrefix(".") }49                let totalBytes = files.reduce(0) { $0 + $1.size }50                var completed: Int64 = 05152                for file in files {53                    try Task.checkCancellation()54                    let destination = directory.appendingPathComponent(file.path)55                    try FileManager.default.createDirectory(56                        at: destination.deletingLastPathComponent(),57                        withIntermediateDirectories: true)5859                    let existing = (try? destination.resourceValues(forKeys: [.fileSizeKey]))?60                        .fileSize.map(Int64.init) ?? 061                    if existing == file.size, file.size > 0 {62                        completed += file.size63                        continuation.yield(.progress(DownloadProgress(64                            repo: repo, completedBytes: completed,65                            totalBytes: totalBytes, currentFile: file.path)))66                        continue67                    }6869                    let completedSoFar = completed70                    try await downloadFile(71                        repo: repo, path: file.path, to: destination,72                        resumeFrom: existing < file.size ? existing : 073                    ) { bytesSoFar in74                        continuation.yield(.progress(DownloadProgress(75                            repo: repo, completedBytes: completedSoFar + bytesSoFar,76                            totalBytes: totalBytes, currentFile: file.path)))77                    }78                    completed += file.size79                }8081                let model = try await ModelStore.shared.describe(directory: directory)82                continuation.yield(.finished(model))83            } catch is CancellationError {84                continuation.yield(.failed("Download paused — it will resume from where it stopped."))85            } catch {86                continuation.yield(.failed(error.localizedDescription))87            }88            self.clearTask(repo: repo)89            continuation.finish()90        }91        activeTasks[repo] = task92        continuation.onTermination = { termination in93            if case .cancelled = termination { task.cancel() }94        }95        return stream96    }9798    func cancel(repo: String) {99        activeTasks[repo]?.cancel()100    }101102    private func clearTask(repo: String) {103        activeTasks[repo] = nil104    }105106    /// Stream one file to disk with Range resume and periodic progress.107    private func downloadFile(108        repo: String, path: String, to destination: URL,109        resumeFrom: Int64,110        progress: @Sendable (Int64) -> Void111    ) async throws {112        var request = URLRequest(url: HubService.resolveURL(repo: repo, path: path))113        if resumeFrom > 0 {114            request.setValue("bytes=\(resumeFrom)-", forHTTPHeaderField: "Range")115        }116        if let token = HFTokenStore.token, !token.isEmpty {117            request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")118        }119120        let (bytes, response) = try await URLSession.shared.bytes(for: request)121        guard let http = response as? HTTPURLResponse,122            http.statusCode == 200 || http.statusCode == 206123        else {124            throw HubServiceError.badResponse((response as? HTTPURLResponse)?.statusCode ?? 0)125        }126127        let appending = http.statusCode == 206 && resumeFrom > 0128        if !appending {129            FileManager.default.createFile(atPath: destination.path, contents: nil)130        }131        let handle = try FileHandle(forWritingTo: destination)132        defer { try? handle.close() }133        if appending {134            try handle.seekToEnd()135        } else {136            try handle.truncate(atOffset: 0)137        }138139        var written: Int64 = appending ? resumeFrom : 0140        var buffer = Data(capacity: 1 << 20)141        var lastReport = Date()142        for try await byte in bytes {143            buffer.append(byte)144            if buffer.count >= 1 << 20 {145                try handle.write(contentsOf: buffer)146                written += Int64(buffer.count)147                buffer.removeAll(keepingCapacity: true)148                if Date().timeIntervalSince(lastReport) > 0.2 {149                    progress(written)150                    lastReport = Date()151                }152                try Task.checkCancellation()153            }154        }155        if !buffer.isEmpty {156            try handle.write(contentsOf: buffer)157            written += Int64(buffer.count)158        }159        progress(written)160    }161}162