From 7658c64678864d3e2b9011c8c6ba6116766e3509 Mon Sep 17 00:00:00 2001 From: Shreeraman A K <16458670+shreeraman96@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:44:20 -0700 Subject: [PATCH 1/2] fix: eliminate HuggingFace rate-limit failures in model downloads Replace the per-directory tree walk with a single recursive listing (with Link pagination), honor ratelimit/Retry-After headers on 429/503 with capped backoff, stop wiping the model cache on rate-limit and network errors so partial downloads stay resumable, fall back to the HuggingFace CLI token file for authenticated quotas, and add a URLProtocol-based regression suite. --- Sources/FluidAudio/DownloadUtils.swift | 567 +++++++++++++----- .../Shared/DownloadUtilsRateLimitTests.swift | 390 ++++++++++++ 2 files changed, 791 insertions(+), 166 deletions(-) create mode 100644 Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift diff --git a/Sources/FluidAudio/DownloadUtils.swift b/Sources/FluidAudio/DownloadUtils.swift index 7d3c8f02b..b65e0b154 100644 --- a/Sources/FluidAudio/DownloadUtils.swift +++ b/Sources/FluidAudio/DownloadUtils.swift @@ -10,15 +10,40 @@ public class DownloadUtils { /// Shared URLSession with registry and proxy configuration public static let sharedSession: URLSession = ModelRegistry.configuredSession() - /// Get HuggingFace token from environment if available. + /// Test-only override for the session used by internal request helpers. + /// `nonisolated(unsafe)` is acceptable here because it is a test seam: set once + /// before any concurrent access (in `setUp`), cleared in `tearDown`, never mutated + /// while requests are in flight. + nonisolated(unsafe) internal static var sessionOverride: URLSession? + + /// Session used for all internal request/download traffic. Prefer this over + /// `sharedSession` inside the type so tests can substitute a stub session. + private static var session: URLSession { sessionOverride ?? sharedSession } + + /// Get HuggingFace token from environment or the HF CLI token file if available. /// Supports multiple env vars for compatibility with different HuggingFace tools: /// - HF_TOKEN: Official HuggingFace CLI /// - HUGGING_FACE_HUB_TOKEN: Python huggingface_hub library /// - HUGGINGFACEHUB_API_TOKEN: LangChain and older integrations + /// - ~/.cache/huggingface/token: HF CLI login token file (used when no env var is set) private static var huggingFaceToken: String? { - ProcessInfo.processInfo.environment["HF_TOKEN"] + if let token = ProcessInfo.processInfo.environment["HF_TOKEN"] ?? ProcessInfo.processInfo.environment["HUGGING_FACE_HUB_TOKEN"] ?? ProcessInfo.processInfo.environment["HUGGINGFACEHUB_API_TOKEN"] + { + return token + } + return tokenFromCLICacheFile + } + + /// Best-effort read of the HF CLI's cached login token. Never throws — sandboxed + /// apps without access to the home directory simply get `nil`. + private static var tokenFromCLICacheFile: String? { + let tokenFileURL = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".cache/huggingface/token") + guard let contents = try? String(contentsOf: tokenFileURL, encoding: .utf8) else { return nil } + let trimmed = contents.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed } /// Create a URLRequest with optional auth header and timeout @@ -26,6 +51,7 @@ public class DownloadUtils { url: URL, timeout: TimeInterval = DownloadConfig.default.timeout ) -> URLRequest { var request = URLRequest(url: url, timeoutInterval: timeout) + request.setValue("FluidAudio-Swift", forHTTPHeaderField: "User-Agent") if let token = huggingFaceToken { request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization") } @@ -36,7 +62,7 @@ public class DownloadUtils { /// Use this for API calls that need auth tokens for private repos or higher rate limits public static func fetchWithAuth(from url: URL) async throws -> (Data, URLResponse) { let request = authorizedRequest(url: url) - return try await sharedSession.data(for: request) + return try await session.data(for: request) } /// Validate that response data is JSON, not HTML error page @@ -51,6 +77,321 @@ public class DownloadUtils { } } + // MARK: - Rate-limit-aware retry helpers + + /// Whether an HTTP status code indicates HuggingFace rate limiting. + internal static func isRateLimitedStatus(_ code: Int) -> Bool { + code == 429 || code == 503 + } + + /// Compute the delay before the next retry attempt. + /// + /// Priority: the `ratelimit` response header's `t=` field (HuggingFace's + /// window-reset hint) → `Retry-After` → exponential backoff. Result is clamped to + /// `[0.5, 300]` seconds. + internal static func retryDelaySeconds( + from response: HTTPURLResponse?, attempt: Int, minBackoff: TimeInterval + ) -> TimeInterval { + let delay: TimeInterval + if let seconds = rateLimitResetSeconds(from: response) { + delay = seconds + } else if let retryAfter = response?.value(forHTTPHeaderField: "Retry-After"), + let seconds = TimeInterval(retryAfter.trimmingCharacters(in: .whitespaces)) + { + delay = seconds + } else { + delay = pow(2.0, Double(max(attempt - 1, 0))) * minBackoff + } + return min(max(delay, 0.5), 300) + } + + /// Parse the `t=` field out of the HuggingFace `ratelimit` header, e.g. + /// `"api";r=0;t=280`. + private static func rateLimitResetSeconds(from response: HTTPURLResponse?) -> TimeInterval? { + guard let header = response?.value(forHTTPHeaderField: "ratelimit") else { return nil } + for component in header.split(separator: ";") { + let trimmed = component.trimmingCharacters(in: .whitespaces) + guard trimmed.hasPrefix("t=") else { continue } + return TimeInterval(trimmed.dropFirst(2)) + } + return nil + } + + /// Parse the RFC 5988 `Link` response header for a `rel="next"` pagination URL. + internal static func nextLinkURL(from response: HTTPURLResponse) -> URL? { + guard let linkHeader = response.value(forHTTPHeaderField: "Link") else { return nil } + for part in linkHeader.split(separator: ",") { + let segments = part.split(separator: ";").map { $0.trimmingCharacters(in: .whitespaces) } + guard let urlSegment = segments.first, urlSegment.hasPrefix("<"), urlSegment.hasSuffix(">") else { + continue + } + let isNext = segments.dropFirst().contains { segment in + let normalized = segment.lowercased().replacingOccurrences(of: "\"", with: "") + return normalized == "rel=next" + } + guard isNext else { continue } + let urlString = String(urlSegment.dropFirst().dropLast()) + return URL(string: urlString) + } + return nil + } + + /// Perform a data request, retrying on 429/503 (rate limit) and transient network + /// errors with an appropriate backoff. Throws `HuggingFaceDownloadError.rateLimited` + /// once attempts are exhausted while rate limited. + private static func dataWithRetry( + request: URLRequest, + description: String, + maxAttempts: Int = 4, + minBackoff: TimeInterval = 1.0 + ) async throws -> (Data, HTTPURLResponse) { + var attempt = 1 + while true { + do { + let (data, response) = try await session.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw HuggingFaceDownloadError.invalidResponse + } + + if isRateLimitedStatus(httpResponse.statusCode) { + if attempt < maxAttempts { + let delay = retryDelaySeconds(from: httpResponse, attempt: attempt, minBackoff: minBackoff) + logger.warning( + "Rate limited (\(httpResponse.statusCode)) while \(description), attempt \(attempt)/\(maxAttempts). Retrying in \(String(format: "%.1f", delay))s." + ) + try Task.checkCancellation() + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + attempt += 1 + continue + } else { + throw HuggingFaceDownloadError.rateLimited( + statusCode: httpResponse.statusCode, + message: "Rate limited while \(description)") + } + } + + return (data, httpResponse) + } catch is CancellationError { + throw CancellationError() + } catch let error as URLError where error.code == .cancelled { + throw error + } catch let error as HuggingFaceDownloadError { + throw error + } catch { + if attempt < maxAttempts { + let delay = pow(2.0, Double(max(attempt - 1, 0))) * minBackoff + logger.warning( + "Request failed while \(description), attempt \(attempt)/\(maxAttempts): \(error.localizedDescription). Retrying in \(String(format: "%.1f", delay))s." + ) + try Task.checkCancellation() + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + attempt += 1 + continue + } + throw error + } + } + } + + /// Perform a download request, retrying on 429/503 (rate limit) and transient + /// network errors with an appropriate backoff. + private static func downloadWithRetry( + request: URLRequest, + description: String, + onProgress: (@Sendable (Int64, Int64) -> Void)?, + maxAttempts: Int = 4, + minBackoff: TimeInterval = 1.0 + ) async throws -> (URL, HTTPURLResponse) { + var attempt = 1 + while true { + do { + let tempFileURL: URL + let httpResponse: HTTPURLResponse + if let onProgress { + (tempFileURL, httpResponse) = try await downloadWithProgress( + request: request, onProgress: onProgress) + } else { + let (url, response) = try await session.download(for: request) + guard let resp = response as? HTTPURLResponse else { + throw HuggingFaceDownloadError.invalidResponse + } + tempFileURL = url + httpResponse = resp + } + + if isRateLimitedStatus(httpResponse.statusCode) { + if attempt < maxAttempts { + let delay = retryDelaySeconds(from: httpResponse, attempt: attempt, minBackoff: minBackoff) + logger.warning( + "Rate limited (\(httpResponse.statusCode)) while downloading \(description), attempt \(attempt)/\(maxAttempts). Retrying in \(String(format: "%.1f", delay))s." + ) + // Discard the (empty/error-body) temp file from this attempt before + // sleeping and retrying — nothing downstream will ever consume it. + try? FileManager.default.removeItem(at: tempFileURL) + try Task.checkCancellation() + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + attempt += 1 + continue + } else { + try? FileManager.default.removeItem(at: tempFileURL) + throw HuggingFaceDownloadError.rateLimited( + statusCode: httpResponse.statusCode, + message: "Rate limited while downloading \(description)") + } + } + + return (tempFileURL, httpResponse) + } catch is CancellationError { + throw CancellationError() + } catch let error as URLError where error.code == .cancelled { + throw error + } catch let error as HuggingFaceDownloadError { + throw error + } catch { + if attempt < maxAttempts { + let delay = pow(2.0, Double(max(attempt - 1, 0))) * minBackoff + logger.warning( + "Download failed while downloading \(description), attempt \(attempt)/\(maxAttempts): \(error.localizedDescription). Retrying in \(String(format: "%.1f", delay))s." + ) + try Task.checkCancellation() + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + attempt += 1 + continue + } + throw error + } + } + } + + /// Fetch and parse a full (paginated) HuggingFace repo tree listing in as few API + /// calls as possible, using `recursive=true` instead of walking directories one at + /// a time. + private static func listRepoTree( + remotePath: String, path: String + ) async throws -> [(path: String, size: Int, type: String)] { + let apiPath = path.isEmpty ? "tree/main" : "tree/main/\(path)" + let baseURL = try ModelRegistry.apiModels(remotePath, apiPath) + + guard var components = URLComponents(url: baseURL, resolvingAgainstBaseURL: false) else { + throw HuggingFaceDownloadError.invalidResponse + } + var queryItems = components.queryItems ?? [] + queryItems.append(URLQueryItem(name: "recursive", value: "true")) + components.queryItems = queryItems + + guard var nextURL = components.url else { + throw HuggingFaceDownloadError.invalidResponse + } + + // Guards against a misbehaving/malicious server sending an unbounded or looping + // pagination chain. + let maxPages = 100 + var seenURLs: Set = [] + + var results: [(path: String, size: Int, type: String)] = [] + var page = 0 + while true { + page += 1 + guard page <= maxPages else { + logger.error("Aborting pagination for \(path.isEmpty ? "root" : path) after \(maxPages) pages") + throw HuggingFaceDownloadError.invalidResponse + } + seenURLs.insert(nextURL) + + let request = authorizedRequest(url: nextURL) + let (data, httpResponse) = try await dataWithRetry(request: request, description: "listing files") + + try validateJSONResponse(data, path: path) + + guard let items = try JSONSerialization.jsonObject(with: data) as? [[String: Any]] else { + throw HuggingFaceDownloadError.invalidResponse + } + + for item in items { + guard let itemPath = item["path"] as? String, + let itemType = item["type"] as? String + else { continue } + let itemSize = item["size"] as? Int ?? -1 + results.append((path: itemPath, size: itemSize, type: itemType)) + } + + guard let next = nextLinkURL(from: httpResponse), !seenURLs.contains(next) else { break } + nextURL = next + } + + return results + } + + /// Pure file-selection logic for `downloadRepo`. + /// + /// Given a flat repo tree listing (as returned by `listRepoTree`), the required-model + /// filter `patterns`, and the optional `subPath`, returns the files that should be + /// downloaded. Reproduces two effects the old per-directory recursive walk had: + /// + /// - Directory pruning: a file is only included if every ancestor directory strictly + /// deeper than the listing root (`subPath`, or the repo root when `subPath` is nil) + /// passes `directoryPasses`. + /// - File inclusion: the `subPath` branch requires the file be inside `subPath` and + /// either match a pattern or look like metadata (`.json`/`.model`/`.bin`); the + /// non-`subPath` branch matches a pattern or a `.json`/`.txt` suffix. + /// + /// Extracted as a standalone pure function (no I/O) so this selection logic can be + /// unit tested directly against synthetic tree listings. + internal static func selectFilesToDownload( + tree: [(path: String, size: Int, type: String)], + patterns: [String], + subPath: String? + ) -> [(path: String, size: Int)] { + // Whether a directory path is allowed to be descended into / must have all its + // files reachable. Mirrors the pruning a per-directory recursive walk used to do. + func directoryPasses(_ dirPath: String) -> Bool { + if let sub = subPath { + return dirPath == sub || dirPath.hasPrefix("\(sub)/") + || patterns.contains { dirPath.hasPrefix($0) || $0.hasPrefix(dirPath + "/") } + } + return patterns.isEmpty || patterns.contains { dirPath.hasPrefix($0) || $0.hasPrefix(dirPath + "/") } + } + + // Progressive ancestor-directory prefixes of a file path, excluding the file itself. + func ancestorDirectories(of filePath: String) -> [String] { + let components = filePath.split(separator: "/").map(String.init) + guard components.count > 1 else { return [] } + return (1.. rootDepth else { return false } + return !directoryPasses(ancestor) + } + guard !isPruned else { continue } + + // For subPath repos, only include files within the subPath + let shouldInclude: Bool + if let sub = subPath { + let isInSubPath = entry.path.hasPrefix("\(sub)/") + let matchesPattern = + patterns.isEmpty || patterns.contains { entry.path.hasPrefix($0) } + let isMetadata = + entry.path.hasSuffix(".json") || entry.path.hasSuffix(".model") || entry.path.hasSuffix(".bin") + shouldInclude = isInSubPath && (matchesPattern || isMetadata) + } else { + shouldInclude = + patterns.isEmpty || patterns.contains { entry.path.hasPrefix($0) } + || entry.path.hasSuffix(".json") || entry.path.hasSuffix(".txt") + } + if shouldInclude { + filesToDownload.append((path: entry.path, size: entry.size)) + } + } + return filesToDownload + } + public enum HuggingFaceDownloadError: LocalizedError { case invalidResponse case rateLimited(statusCode: Int, message: String) @@ -130,6 +471,13 @@ public class DownloadUtils { directory: directory, computeUnits: computeUnits, variant: variant, progressHandler: progressHandler) } catch { + guard !isTransientAndUnwipeable(error) else { + logger.warning( + "First load failed with a transient/network error, not wiping cache: \(error.localizedDescription)" + ) + throw error + } + logger.warning("First load failed: \(error.localizedDescription)") logger.info("Deleting cache and re-downloading…") let repoPath = directory.appendingPathComponent(repo.folderName) @@ -142,6 +490,27 @@ public class DownloadUtils { } } + /// Errors where wiping the cache and retrying would just discard a resumable + /// partial download for no benefit: cancellation, rate limiting, HTML error pages + /// (usually rate-limit related), and any network-layer failure (offline, timeout, …). + private static func isTransientAndUnwipeable(_ error: Error) -> Bool { + if error is CancellationError { + return true + } + if error is URLError { + return true + } + if let downloadError = error as? HuggingFaceDownloadError { + switch downloadError { + case .rateLimited, .htmlErrorResponse: + return true + case .invalidResponse, .downloadFailed, .modelNotFound: + return false + } + } + return false + } + public static func clearModelCache(forRepo repo: Repo, directory: URL) { let repoPath = directory.appendingPathComponent(repo.folderName) try? FileManager.default.removeItem(at: repoPath) @@ -289,76 +658,12 @@ public class DownloadUtils { } } - // Get all files recursively using HuggingFace API - var filesToDownload: [(path: String, size: Int)] = [] - - func listDirectory(path: String) async throws { - let apiPath = path.isEmpty ? "tree/main" : "tree/main/\(path)" - let dirURL = try ModelRegistry.apiModels(repo.remotePath, apiPath) - let request = authorizedRequest(url: dirURL) - - let (dirData, response) = try await sharedSession.data(for: request) - - if let httpResponse = response as? HTTPURLResponse { - if httpResponse.statusCode == 429 || httpResponse.statusCode == 503 { - throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, message: "Rate limited while listing files") - } - } - - // Validate that response is JSON, not HTML error page - try validateJSONResponse(dirData, path: path) - - guard let items = try JSONSerialization.jsonObject(with: dirData) as? [[String: Any]] else { - throw HuggingFaceDownloadError.invalidResponse - } - - for item in items { - guard let itemPath = item["path"] as? String, - let itemType = item["type"] as? String - else { continue } - - if itemType == "directory" { - // For subPath repos, only process paths within the subPath - let shouldProcess: Bool - if let sub = subPath { - shouldProcess = - itemPath == sub || itemPath.hasPrefix("\(sub)/") - || patterns.contains { itemPath.hasPrefix($0) || $0.hasPrefix(itemPath + "/") } - } else { - shouldProcess = - patterns.isEmpty - || patterns.contains { itemPath.hasPrefix($0) || $0.hasPrefix(itemPath + "/") } - } - if shouldProcess { - try await listDirectory(path: itemPath) - } - } else if itemType == "file" { - // For subPath repos, only include files within the subPath - let shouldInclude: Bool - if let sub = subPath { - let isInSubPath = itemPath.hasPrefix("\(sub)/") - let matchesPattern = - patterns.isEmpty || patterns.contains { itemPath.hasPrefix($0) } - let isMetadata = - itemPath.hasSuffix(".json") || itemPath.hasSuffix(".model") || itemPath.hasSuffix(".bin") - shouldInclude = isInSubPath && (matchesPattern || isMetadata) - } else { - shouldInclude = - patterns.isEmpty || patterns.contains { itemPath.hasPrefix($0) } - || itemPath.hasSuffix(".json") || itemPath.hasSuffix(".txt") - } - if shouldInclude { - let fileSize = item["size"] as? Int ?? -1 - filesToDownload.append((path: itemPath, size: fileSize)) - } - } - } - } - - // Start listing from subPath if specified, otherwise from root + // Get all files recursively in a single (paginated) API call, then reproduce the + // pruning + inclusion semantics of the old per-directory recursive walk. progressHandler?(DownloadProgress(fractionCompleted: 0.0, phase: .listing)) - try await listDirectory(path: subPath ?? "") + let tree = try await listRepoTree(remotePath: repo.remotePath, path: subPath ?? "") + + let filesToDownload = selectFilesToDownload(tree: tree, patterns: patterns, subPath: subPath) logger.info("Found \(filesToDownload.count) files to download") // Compute total known bytes for byte-weighted progress. @@ -406,8 +711,9 @@ public class DownloadUtils { let baseBytes = completedBytes let fileCount = filesToDownload.count let totalBytesSnapshot = totalBytes - (tempFileURL, httpResponse) = try await downloadWithProgress( + (tempFileURL, httpResponse) = try await downloadWithRetry( request: request, + description: file.path, onProgress: { bytesWritten, _ in guard totalBytesSnapshot > 0 else { return } let current = baseBytes + bytesWritten @@ -421,19 +727,8 @@ public class DownloadUtils { } ) } else { - let (url, response) = try await sharedSession.download(for: request) - guard let resp = response as? HTTPURLResponse else { - throw HuggingFaceDownloadError.invalidResponse - } - tempFileURL = url - httpResponse = resp - } - - // Validate response - if httpResponse.statusCode == 429 || httpResponse.statusCode == 503 { - throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, - message: "Rate limited while downloading \(file.path)") + (tempFileURL, httpResponse) = try await downloadWithRetry( + request: request, description: file.path, onProgress: nil) } guard (200..<300).contains(httpResponse.statusCode) else { @@ -497,14 +792,14 @@ public class DownloadUtils { onProgress: onProgress, completion: { continuation.resume(with: $0) } ) - let session = URLSession( - configuration: sharedSession.configuration, + let delegateSession = URLSession( + configuration: session.configuration, delegate: delegate, delegateQueue: nil ) - delegate.session = session + delegate.session = delegateSession - let task = session.downloadTask(with: request) + let task = delegateSession.downloadTask(with: request) taskHolder.setTask(task) task.resume() } @@ -528,40 +823,11 @@ public class DownloadUtils { subdirectory: String, to repoDirectory: URL ) async throws { - var filesToDownload: [(path: String, size: Int)] = [] - - func listFiles(at path: String) async throws { - let dirURL = try ModelRegistry.apiModels(repo.remotePath, "tree/main/\(path)") - let (dirData, response) = try await fetchWithAuth(from: dirURL) - if let httpResponse = response as? HTTPURLResponse, - httpResponse.statusCode == 429 || httpResponse.statusCode == 503 - { - throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, - message: "Rate limited while listing files in \(path)") - } - - // Validate that response is JSON, not HTML error page - try validateJSONResponse(dirData, path: path) - - guard let items = try JSONSerialization.jsonObject(with: dirData) as? [[String: Any]] else { - throw HuggingFaceDownloadError.invalidResponse - } - for item in items { - guard let itemPath = item["path"] as? String, - let itemType = item["type"] as? String - else { continue } - - if itemType == "directory" { - try await listFiles(at: itemPath) - } else if itemType == "file" { - let fileSize = item["size"] as? Int ?? -1 - filesToDownload.append((path: itemPath, size: fileSize)) - } - } - } - - try await listFiles(at: subdirectory) + let tree = try await listRepoTree(remotePath: repo.remotePath, path: subdirectory) + let filesToDownload: [(path: String, size: Int)] = + tree + .filter { $0.type == "file" } + .map { (path: $0.path, size: $0.size) } logger.info("Found \(filesToDownload.count) files in \(subdirectory)") for (index, file) in filesToDownload.enumerated() { @@ -586,16 +852,8 @@ public class DownloadUtils { let fileURL = try ModelRegistry.resolveModel(repo.remotePath, encodedPath) let request = authorizedRequest(url: fileURL) - let (tempURL, response) = try await sharedSession.download(for: request) - guard let httpResponse = response as? HTTPURLResponse else { - throw HuggingFaceDownloadError.invalidResponse - } - - if httpResponse.statusCode == 429 || httpResponse.statusCode == 503 { - throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, - message: "Rate limited while downloading \(file.path)") - } + let (tempURL, httpResponse) = try await downloadWithRetry( + request: request, description: file.path, onProgress: nil) guard (200..<300).contains(httpResponse.statusCode) else { throw HuggingFaceDownloadError.downloadFailed( @@ -618,49 +876,26 @@ public class DownloadUtils { } /// Fetch a single file from HuggingFace with retry + /// + /// Rate-limit (429/503) retries and transient network-error retries both happen + /// inside `dataWithRetry`, sharing the full `maxAttempts` budget so 429s use + /// HuggingFace's header-provided delay instead of blind exponential backoff. public static func fetchHuggingFaceFile( from url: URL, description: String, maxAttempts: Int = 4, minBackoff: TimeInterval = 1.0 ) async throws -> Data { - var lastError: Error? let request = authorizedRequest(url: url) - for attempt in 1...maxAttempts { - do { - let (data, response) = try await sharedSession.data(for: request) - - guard let httpResponse = response as? HTTPURLResponse else { - throw HuggingFaceDownloadError.invalidResponse - } + let (data, httpResponse) = try await dataWithRetry( + request: request, description: description, maxAttempts: maxAttempts, minBackoff: minBackoff) - if httpResponse.statusCode == 429 || httpResponse.statusCode == 503 { - throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, - message: "HTTP \(httpResponse.statusCode)" - ) - } - - guard (200..<300).contains(httpResponse.statusCode) else { - throw HuggingFaceDownloadError.invalidResponse - } - - return data - - } catch { - lastError = error - if attempt < maxAttempts { - let backoffSeconds = pow(2.0, Double(attempt - 1)) * minBackoff - logger.warning( - "Download attempt \(attempt) for \(description) failed: \(error.localizedDescription). Retrying in \(String(format: "%.1f", backoffSeconds))s." - ) - try await Task.sleep(nanoseconds: UInt64(backoffSeconds * 1_000_000_000)) - } - } + guard (200..<300).contains(httpResponse.statusCode) else { + throw HuggingFaceDownloadError.invalidResponse } - throw lastError ?? HuggingFaceDownloadError.invalidResponse + return data } } diff --git a/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift b/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift new file mode 100644 index 000000000..88575525e --- /dev/null +++ b/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift @@ -0,0 +1,390 @@ +import Foundation +import XCTest + +@testable import FluidAudio + +/// A stubbed `URLProtocol` that serves queued responses in order (falling back to a +/// default response once the queue is exhausted) and records every request it saw. +/// Used to exercise `DownloadUtils`' rate-limit retry/pagination logic without any +/// real network traffic. +final class StubURLProtocol: URLProtocol { + + struct StubResponse { + let statusCode: Int + let headers: [String: String] + let data: Data + /// When set, the protocol reports a transport-level failure (e.g. a timeout) + /// instead of an HTTP response. + let failureCode: URLError.Code? + + init(statusCode: Int, headers: [String: String] = [:], data: Data = Data()) { + self.statusCode = statusCode + self.headers = headers + self.data = data + self.failureCode = nil + } + + init(failureCode: URLError.Code) { + self.statusCode = 0 + self.headers = [:] + self.data = Data() + self.failureCode = failureCode + } + } + + private static let lock = NSLock() + nonisolated(unsafe) private static var queuedResponses: [StubResponse] = [] + nonisolated(unsafe) private static var fallbackResponse = StubResponse(statusCode: 200) + nonisolated(unsafe) private static var _recordedRequests: [URLRequest] = [] + + static var recordedRequests: [URLRequest] { + lock.withLock { _recordedRequests } + } + + /// Queue responses to be returned in order, one per incoming request. Once the + /// queue is exhausted, `fallback` is returned for any further requests. + static func enqueue(_ responses: [StubResponse], fallback: StubResponse = StubResponse(statusCode: 200)) { + lock.withLock { + queuedResponses = responses + fallbackResponse = fallback + } + } + + static func reset() { + lock.withLock { + queuedResponses = [] + fallbackResponse = StubResponse(statusCode: 200) + _recordedRequests = [] + } + } + + override class func canInit(with request: URLRequest) -> Bool { true } + + override class func canonicalRequest(for request: URLRequest) -> URLRequest { request } + + override func startLoading() { + let stub = Self.lock.withLock { () -> StubResponse in + Self._recordedRequests.append(self.request) + guard !Self.queuedResponses.isEmpty else { return Self.fallbackResponse } + return Self.queuedResponses.removeFirst() + } + + if let failureCode = stub.failureCode { + client?.urlProtocol(self, didFailWithError: URLError(failureCode)) + return + } + + guard let url = request.url, + let httpResponse = HTTPURLResponse( + url: url, statusCode: stub.statusCode, httpVersion: "HTTP/1.1", headerFields: stub.headers) + else { + client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse)) + return + } + + client?.urlProtocol(self, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + client?.urlProtocol(self, didLoad: stub.data) + client?.urlProtocolDidFinishLoading(self) + } + + override func stopLoading() {} +} + +final class DownloadUtilsRateLimitTests: XCTestCase { + + override func setUp() { + super.setUp() + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [StubURLProtocol.self] + DownloadUtils.sessionOverride = URLSession(configuration: configuration) + } + + override func tearDown() { + DownloadUtils.sessionOverride = nil + StubURLProtocol.reset() + super.tearDown() + } + + private func makeTempDirectory() -> URL { + let dir = FileManager.default.temporaryDirectory + .appendingPathComponent("DownloadUtilsRateLimitTests-\(UUID().uuidString)") + try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + return dir + } + + private func jsonData(_ items: [[String: Any]]) -> Data { + (try? JSONSerialization.data(withJSONObject: items)) ?? Data() + } + + private func treeRequests() -> [URLRequest] { + StubURLProtocol.recordedRequests.filter { $0.url?.path.contains("/tree/") ?? false } + } + + // MARK: - retryDelaySeconds + + func testRetryDelayUsesRateLimitHeaderTValue() { + let response = HTTPURLResponse( + url: URL(string: "https://huggingface.co/")!, + statusCode: 429, + httpVersion: nil, + headerFields: ["ratelimit": "\"api\";r=0;t=7"] + ) + let delay = DownloadUtils.retryDelaySeconds(from: response, attempt: 1, minBackoff: 1.0) + XCTAssertEqual(delay, 7, accuracy: 0.001) + } + + func testRetryDelayFallsBackToExponentialBackoffWithoutHeaders() { + let delay = DownloadUtils.retryDelaySeconds(from: nil, attempt: 3, minBackoff: 1.0) + XCTAssertEqual(delay, 4, accuracy: 0.001) // pow(2, 2) * 1.0 + } + + func testRetryDelayClampsToFloor() { + let response = HTTPURLResponse( + url: URL(string: "https://huggingface.co/")!, + statusCode: 429, + httpVersion: nil, + headerFields: ["ratelimit": "\"api\";r=0;t=0"] + ) + let delay = DownloadUtils.retryDelaySeconds(from: response, attempt: 1, minBackoff: 1.0) + XCTAssertEqual(delay, 0.5, accuracy: 0.001) + } + + func testRetryDelayUsesRetryAfterHeaderWhenNoRateLimitHeader() { + let response = HTTPURLResponse( + url: URL(string: "https://huggingface.co/")!, + statusCode: 429, + httpVersion: nil, + headerFields: ["Retry-After": "12"] + ) + let delay = DownloadUtils.retryDelaySeconds(from: response, attempt: 1, minBackoff: 1.0) + XCTAssertEqual(delay, 12, accuracy: 0.001) + } + + // MARK: - nextLinkURL + + func testNextLinkURLParsesRelNext() { + let response = HTTPURLResponse( + url: URL(string: "https://huggingface.co/api/models/x/tree/main")!, + statusCode: 200, + httpVersion: nil, + headerFields: ["Link": "; rel=\"next\""] + )! + XCTAssertEqual(DownloadUtils.nextLinkURL(from: response), URL(string: "https://x/api?cursor=abc")) + } + + func testNextLinkURLReturnsNilWhenAbsent() { + let response = HTTPURLResponse( + url: URL(string: "https://huggingface.co/")!, + statusCode: 200, + httpVersion: nil, + headerFields: nil + )! + XCTAssertNil(DownloadUtils.nextLinkURL(from: response)) + } + + // MARK: - Listing retries on rate limit, then succeeds + + func testDownloadSubdirectoryRetriesOnRateLimitThenSucceeds() async throws { + StubURLProtocol.enqueue([ + .init(statusCode: 429, headers: ["ratelimit": "\"api\";r=0;t=0"]), + .init(statusCode: 200, data: jsonData([["type": "file", "path": "a/b.bin", "size": 3]])), + .init(statusCode: 200, data: Data([1, 2, 3])), + ]) + + let repoDirectory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: repoDirectory) } + + try await DownloadUtils.downloadSubdirectory(.parakeet, subdirectory: "extra", to: repoDirectory) + + XCTAssertEqual(treeRequests().count, 2) + XCTAssertTrue( + FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("a/b.bin").path)) + } + + // MARK: - Single-call recursive listing (no per-directory walk) + + func testListingMakesExactlyOneTreeRequestForNestedPaths() async throws { + StubURLProtocol.enqueue( + [ + .init( + statusCode: 200, + data: jsonData([ + ["type": "directory", "path": "sub"], + ["type": "file", "path": "sub/nested/deep.bin", "size": 0], + ["type": "file", "path": "top.bin", "size": 0], + ])) + ], fallback: .init(statusCode: 200)) + + let repoDirectory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: repoDirectory) } + + try await DownloadUtils.downloadSubdirectory(.parakeet, subdirectory: "extra", to: repoDirectory) + + XCTAssertEqual(treeRequests().count, 1) + XCTAssertTrue( + FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("sub/nested/deep.bin").path)) + XCTAssertTrue( + FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("top.bin").path)) + } + + // MARK: - Pagination follows Link rel="next" and unions results + + func testListingFollowsPaginationAndUnionsResults() async throws { + StubURLProtocol.enqueue([ + .init( + statusCode: 200, + headers: ["Link": "; rel=\"next\""], + data: jsonData([["type": "file", "path": "p1.bin", "size": 0]]) + ), + .init(statusCode: 200, data: jsonData([["type": "file", "path": "p2.bin", "size": 0]])), + ]) + + let repoDirectory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: repoDirectory) } + + try await DownloadUtils.downloadSubdirectory(.parakeet, subdirectory: "extra", to: repoDirectory) + + XCTAssertEqual(treeRequests().count, 2) + XCTAssertTrue(FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("p1.bin").path)) + XCTAssertTrue(FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("p2.bin").path)) + } + + // MARK: - Exhausted retries surface a rateLimited error + + func testListingExhaustedRetriesThrowsRateLimited() async { + StubURLProtocol.enqueue([], fallback: .init(statusCode: 429, headers: ["ratelimit": "\"api\";r=0;t=0"])) + + let repoDirectory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: repoDirectory) } + + do { + try await DownloadUtils.downloadSubdirectory(.parakeet, subdirectory: "extra", to: repoDirectory) + XCTFail("Expected rateLimited error") + } catch let DownloadUtils.HuggingFaceDownloadError.rateLimited(statusCode, message) { + XCTAssertEqual(statusCode, 429) + XCTAssertTrue(message.contains("Rate limited while listing files")) + } catch { + XCTFail("Expected rateLimited error, got \(error)") + } + } + + // MARK: - loadModels does not wipe cache on rate limit / network errors + + func testLoadModelsDoesNotWipeCacheOnRateLimit() async { + StubURLProtocol.enqueue([], fallback: .init(statusCode: 429, headers: ["ratelimit": "\"api\";r=0;t=0"])) + + let directory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: directory) } + + let repoPath = directory.appendingPathComponent(Repo.parakeet.folderName) + try? FileManager.default.createDirectory(at: repoPath, withIntermediateDirectories: true) + let markerPath = repoPath.appendingPathComponent("marker.txt") + FileManager.default.createFile(atPath: markerPath.path, contents: Data("marker".utf8)) + + do { + _ = try await DownloadUtils.loadModels(.parakeet, modelNames: [], directory: directory) + XCTFail("Expected rateLimited error") + } catch is DownloadUtils.HuggingFaceDownloadError { + // expected + } catch { + XCTFail("Expected rateLimited error, got \(error)") + } + + XCTAssertTrue( + FileManager.default.fileExists(atPath: markerPath.path), + "Cache should not be wiped when the failure is a rate limit") + } + + func testLoadModelsDoesNotWipeCacheOnNetworkTimeout() async { + StubURLProtocol.enqueue([], fallback: .init(failureCode: .timedOut)) + + let directory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: directory) } + + let repoPath = directory.appendingPathComponent(Repo.parakeet.folderName) + try? FileManager.default.createDirectory(at: repoPath, withIntermediateDirectories: true) + let markerPath = repoPath.appendingPathComponent("marker.txt") + FileManager.default.createFile(atPath: markerPath.path, contents: Data("marker".utf8)) + + do { + _ = try await DownloadUtils.loadModels(.parakeet, modelNames: [], directory: directory) + XCTFail("Expected a URLError") + } catch let error as URLError { + XCTAssertEqual(error.code, .timedOut) + } catch { + XCTFail("Expected URLError.timedOut, got \(error)") + } + + XCTAssertTrue( + FileManager.default.fileExists(atPath: markerPath.path), + "Cache should not be wiped when the failure is a network error") + } + + // MARK: - selectFilesToDownload (pure file-selection logic) + + func testSelectFilesToDownloadNonSubPathPatternsAndRootMetadata() { + let tree: [(path: String, size: Int, type: String)] = [ + ("model_a/weights.mlmodelc", 10, "file"), + ("model_a", 0, "directory"), + ("model_b/weights.mlmodelc", 10, "file"), + ("config.json", 5, "file"), + ("README.txt", 3, "file"), + ("notes.md", 3, "file"), + ] + let patterns = ["model_a/"] + + let selected = DownloadUtils.selectFilesToDownload(tree: tree, patterns: patterns, subPath: nil) + let paths = Set(selected.map { $0.path }) + + XCTAssertEqual( + paths, ["model_a/weights.mlmodelc", "config.json", "README.txt"], + "Only pattern-matched files plus root .json/.txt metadata should be selected") + } + + func testSelectFilesToDownloadSubPathMetadataRule() { + let tree: [(path: String, size: Int, type: String)] = [ + ("160ms/encoder.mlmodelc/model.mlmodel", 10, "file"), + ("160ms/vocab.json", 2, "file"), + ("160ms/tokenizer.model", 2, "file"), + ("160ms/extra.bin", 2, "file"), + ("160ms/readme.md", 2, "file"), + ("320ms/encoder.mlmodelc/model.mlmodel", 10, "file"), + ] + let patterns = ["160ms/encoder.mlmodelc/"] + + let selected = DownloadUtils.selectFilesToDownload(tree: tree, patterns: patterns, subPath: "160ms") + let paths = Set(selected.map { $0.path }) + + XCTAssertEqual( + paths, + [ + "160ms/encoder.mlmodelc/model.mlmodel", "160ms/vocab.json", "160ms/tokenizer.model", + "160ms/extra.bin", + ], + "Files inside subPath matching a pattern or .json/.model/.bin metadata should be selected;" + + " files outside subPath or without a matching suffix should not") + XCTAssertFalse(paths.contains("160ms/readme.md")) + XCTAssertFalse(paths.contains("320ms/encoder.mlmodelc/model.mlmodel")) + } + + func testSelectFilesToDownloadPrunesFilesUnderRejectedAncestorDirectory() { + // "other_model/config.json" matches the root-level .json metadata rule by suffix, + // but its ancestor directory "other_model" does not pass directoryPasses (it isn't + // a required-model pattern), so the old per-directory walk would never have + // descended into it. The flat-tree selector must reproduce that pruning. + let tree: [(path: String, size: Int, type: String)] = [ + ("required_model/weights.mlmodelc", 10, "file"), + ("other_model/config.json", 5, "file"), + ("top_level.json", 5, "file"), + ] + let patterns = ["required_model/"] + + let selected = DownloadUtils.selectFilesToDownload(tree: tree, patterns: patterns, subPath: nil) + let paths = Set(selected.map { $0.path }) + + XCTAssertEqual( + paths, ["required_model/weights.mlmodelc", "top_level.json"], + "A file whose suffix matches but whose ancestor directory is pruned must be excluded") + XCTAssertFalse(paths.contains("other_model/config.json")) + } +} From e77f5b73af87f2798895a2ada06cd2ba38356b37 Mon Sep 17 00:00:00 2001 From: Shreeraman A K <16458670+shreeraman96@users.noreply.github.com> Date: Sat, 25 Jul 2026 12:47:55 -0700 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20address=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20retry=20transient=205xx,=20monotonic=20progress,=20?= =?UTF-8?q?cached=20token?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - dataWithRetry/downloadWithRetry now also retry 500/502/504 with backoff, restoring the retry coverage fetchHuggingFaceFile had before the refactor; exhausted 5xx retries surface through the caller's status validation. - Download progress is clamped monotonic so a retried file transfer cannot rewind the reported fraction. - HF token resolved once per process instead of re-reading the CLI token file on every request. - Drop the prose-message assertion from the exhausted-retries test; add a 502-then-success retry regression test. --- Sources/FluidAudio/DownloadUtils.swift | 48 ++++++++++++++----- .../Shared/DownloadUtilsRateLimitTests.swift | 20 +++++++- 2 files changed, 53 insertions(+), 15 deletions(-) diff --git a/Sources/FluidAudio/DownloadUtils.swift b/Sources/FluidAudio/DownloadUtils.swift index b65e0b154..1dd71b1e3 100644 --- a/Sources/FluidAudio/DownloadUtils.swift +++ b/Sources/FluidAudio/DownloadUtils.swift @@ -1,6 +1,7 @@ import CoreML import Foundation import OSLog +import os /// HuggingFace model downloader using URLSession public class DownloadUtils { @@ -26,7 +27,8 @@ public class DownloadUtils { /// - HUGGING_FACE_HUB_TOKEN: Python huggingface_hub library /// - HUGGINGFACEHUB_API_TOKEN: LangChain and older integrations /// - ~/.cache/huggingface/token: HF CLI login token file (used when no env var is set) - private static var huggingFaceToken: String? { + /// Resolved once per process: a token added after launch is picked up on next start. + private static let huggingFaceToken: String? = { if let token = ProcessInfo.processInfo.environment["HF_TOKEN"] ?? ProcessInfo.processInfo.environment["HUGGING_FACE_HUB_TOKEN"] ?? ProcessInfo.processInfo.environment["HUGGINGFACEHUB_API_TOKEN"] @@ -34,7 +36,7 @@ public class DownloadUtils { return token } return tokenFromCLICacheFile - } + }() /// Best-effort read of the HF CLI's cached login token. Never throws — sandboxed /// apps without access to the home directory simply get `nil`. @@ -84,6 +86,11 @@ public class DownloadUtils { code == 429 || code == 503 } + /// Whether an HTTP status code indicates a transient server-side failure worth retrying. + internal static func isTransientServerStatus(_ code: Int) -> Bool { + code == 500 || code == 502 || code == 504 + } + /// Compute the delay before the next retry attempt. /// /// Priority: the `ratelimit` response header's `t=` field (HuggingFace's @@ -153,21 +160,24 @@ public class DownloadUtils { throw HuggingFaceDownloadError.invalidResponse } - if isRateLimitedStatus(httpResponse.statusCode) { + let status = httpResponse.statusCode + if isRateLimitedStatus(status) || isTransientServerStatus(status) { if attempt < maxAttempts { let delay = retryDelaySeconds(from: httpResponse, attempt: attempt, minBackoff: minBackoff) logger.warning( - "Rate limited (\(httpResponse.statusCode)) while \(description), attempt \(attempt)/\(maxAttempts). Retrying in \(String(format: "%.1f", delay))s." + "HTTP \(status) while \(description), attempt \(attempt)/\(maxAttempts). Retrying in \(String(format: "%.1f", delay))s." ) try Task.checkCancellation() try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) attempt += 1 continue - } else { + } else if isRateLimitedStatus(status) { throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, + statusCode: status, message: "Rate limited while \(description)") } + // Exhausted retries on a transient 5xx: return the response and let the + // caller's status validation surface the failure. } return (data, httpResponse) @@ -219,11 +229,12 @@ public class DownloadUtils { httpResponse = resp } - if isRateLimitedStatus(httpResponse.statusCode) { + let status = httpResponse.statusCode + if isRateLimitedStatus(status) || isTransientServerStatus(status) { if attempt < maxAttempts { let delay = retryDelaySeconds(from: httpResponse, attempt: attempt, minBackoff: minBackoff) logger.warning( - "Rate limited (\(httpResponse.statusCode)) while downloading \(description), attempt \(attempt)/\(maxAttempts). Retrying in \(String(format: "%.1f", delay))s." + "HTTP \(status) while downloading \(description), attempt \(attempt)/\(maxAttempts). Retrying in \(String(format: "%.1f", delay))s." ) // Discard the (empty/error-body) temp file from this attempt before // sleeping and retrying — nothing downstream will ever consume it. @@ -232,12 +243,14 @@ public class DownloadUtils { try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) attempt += 1 continue - } else { + } else if isRateLimitedStatus(status) { try? FileManager.default.removeItem(at: tempFileURL) throw HuggingFaceDownloadError.rateLimited( - statusCode: httpResponse.statusCode, + statusCode: status, message: "Rate limited while downloading \(description)") } + // Exhausted retries on a transient 5xx: return the response and let the + // caller's status validation surface the failure. } return (tempFileURL, httpResponse) @@ -670,6 +683,9 @@ public class DownloadUtils { // Files with unknown sizes (size == -1) are treated as 0 for weighting. let totalBytes: Int64 = filesToDownload.reduce(0) { $0 + Int64(max(0, $1.size)) } var completedBytes: Int64 = 0 + // Highest download-phase fraction reported so far; keeps progress monotonic + // across per-file retries (a retried attempt restarts its byte count at zero). + let maxReportedFraction = OSAllocatedUnfairLock(initialState: 0.0) // Download each file for (index, file) in filesToDownload.enumerated() { @@ -717,11 +733,17 @@ public class DownloadUtils { onProgress: { bytesWritten, _ in guard totalBytesSnapshot > 0 else { return } let current = baseBytes + bytesWritten - // Download phase occupies 0.0–0.5 of the overall range. - let fraction = 0.5 * Double(current) / Double(totalBytesSnapshot) + // Download phase occupies 0.0–0.5 of the overall range. Clamp to the + // highest fraction reported so far: a retried attempt restarts its + // byte count at zero and must not rewind the visible progress. + let fraction = min(0.5 * Double(current) / Double(totalBytesSnapshot), 0.5) + let monotonic = maxReportedFraction.withLock { maxSoFar in + maxSoFar = max(maxSoFar, fraction) + return maxSoFar + } handler( DownloadProgress( - fractionCompleted: min(fraction, 0.5), + fractionCompleted: monotonic, phase: .downloading(completedFiles: index, totalFiles: fileCount) )) } diff --git a/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift b/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift index 88575525e..447107f2b 100644 --- a/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift +++ b/Tests/FluidAudioTests/Shared/DownloadUtilsRateLimitTests.swift @@ -201,6 +201,23 @@ final class DownloadUtilsRateLimitTests: XCTestCase { FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("a/b.bin").path)) } + func testListingRetriesOnTransientServerErrorThenSucceeds() async throws { + StubURLProtocol.enqueue([ + .init(statusCode: 502), + .init(statusCode: 200, data: jsonData([["type": "file", "path": "a/b.bin", "size": 3]])), + .init(statusCode: 200, data: Data([1, 2, 3])), + ]) + + let repoDirectory = makeTempDirectory() + defer { try? FileManager.default.removeItem(at: repoDirectory) } + + try await DownloadUtils.downloadSubdirectory(.parakeet, subdirectory: "extra", to: repoDirectory) + + XCTAssertEqual(treeRequests().count, 2) + XCTAssertTrue( + FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent("a/b.bin").path)) + } + // MARK: - Single-call recursive listing (no per-directory walk) func testListingMakesExactlyOneTreeRequestForNestedPaths() async throws { @@ -260,9 +277,8 @@ final class DownloadUtilsRateLimitTests: XCTestCase { do { try await DownloadUtils.downloadSubdirectory(.parakeet, subdirectory: "extra", to: repoDirectory) XCTFail("Expected rateLimited error") - } catch let DownloadUtils.HuggingFaceDownloadError.rateLimited(statusCode, message) { + } catch let DownloadUtils.HuggingFaceDownloadError.rateLimited(statusCode, _) { XCTAssertEqual(statusCode, 429) - XCTAssertTrue(message.contains("Rate limited while listing files")) } catch { XCTFail("Expected rateLimited error, got \(error)") }