369 lines · 13.9 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif public struct DoubaoUsageSnapshot: Sendable {    public let remainingRequests: Int    public let limitRequests: Int    public let resetTime: Date?    public let updatedAt: Date    public let apiKeyValid: Bool    public let totalTokens: Int?    public let requestLimitsReliable: Bool    public init(        remainingRequests: Int,        limitRequests: Int,        resetTime: Date?,        updatedAt: Date,        apiKeyValid: Bool = false,        totalTokens: Int? = nil,        requestLimitsReliable: Bool = true)    {        self.remainingRequests = remainingRequests        self.limitRequests = limitRequests        self.resetTime = resetTime        self.updatedAt = updatedAt        self.apiKeyValid = apiKeyValid        self.totalTokens = totalTokens        self.requestLimitsReliable = requestLimitsReliable    }     public func toUsageSnapshot() -> UsageSnapshot {        let primary: RateWindow?        if self.limitRequests > 0, self.requestLimitsReliable {            let used = max(0, self.limitRequests - self.remainingRequests)            primary = RateWindow(                usedPercent: min(100, max(0, Double(used) / Double(self.limitRequests) * 100)),                windowMinutes: nil,                resetsAt: self.resetTime,                resetDescription: "\(used)/\(self.limitRequests) requests")        } else if self.apiKeyValid {            // Ark can return successful requests without a trustworthy request-limit window.            // Omitting the window prevents the UI from presenting unknown usage as 100% left.            primary = nil        } else {            primary = RateWindow(                usedPercent: 0,                windowMinutes: nil,                resetsAt: self.resetTime,                resetDescription: "No usage data")        }         let identity = ProviderIdentitySnapshot(            providerID: .doubao,            accountEmail: nil,            accountOrganization: nil,            loginMethod: nil)         return UsageSnapshot(            primary: primary,            secondary: nil,            tertiary: nil,            providerCost: nil,            updatedAt: self.updatedAt,            identity: identity)    }} public enum DoubaoUsageError: LocalizedError, Sendable {    case missingCredentials    case networkError(String)    case apiError(Int, String)    case parseFailed(String)     public var errorDescription: String? {        switch self {        case .missingCredentials:            "Missing Doubao API key (ARK_API_KEY)."        case let .networkError(message):            "Doubao network error: \(message)"        case let .apiError(code, message):            "Doubao API error (\(code)): \(message)"        case let .parseFailed(message):            "Failed to parse Doubao response: \(message)"        }    }} public struct DoubaoUsageFetcher: Sendable {    private static let log = CodexBarLog.logger(LogCategories.doubaoUsage)    private static let apiURL = URL(string: "https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions")!     /// Models to probe, ordered by likelihood. We try multiple models because    /// different key types may not have access to every model.    private static let probeModels = [        "doubao-seed-2.0-code",        "doubao-1.5-pro-32k",        "doubao-lite-32k",    ]     private struct ProbeResult {        let snapshot: DoubaoUsageSnapshot        let statusCode: Int         var hasAmbiguousZeroRemaining: Bool {            self.statusCode == 200                && self.snapshot.requestLimitsReliable                && self.snapshot.limitRequests > 0                && self.snapshot.remainingRequests == 0        }    }     public static func fetchUsage(        apiKey: String,        session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> DoubaoUsageSnapshot    {        guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {            throw DoubaoUsageError.missingCredentials        }         var lastError: Error?        for model in self.probeModels {            do {                let result = try await self.probe(apiKey: apiKey, model: model, transport: transport)                guard result.hasAmbiguousZeroRemaining else {                    return result.snapshot                }                 return try await self.confirmAmbiguousZeroRemaining(                    initial: result,                    apiKey: apiKey,                    model: model,                    transport: transport)            } catch let error as DoubaoUsageError {                if case let .apiError(code, _) = error, code == 404 || code == 403 {                    Self.log.debug("Doubao probe model \(model) unavailable (\(code)), trying next")                    lastError = error                    continue                }                throw error            }        }        throw lastError ?? DoubaoUsageError.apiError(0, "All probe models failed")    }     private static func confirmAmbiguousZeroRemaining(        initial: ProbeResult,        apiKey: String,        model: String,        transport: any ProviderHTTPTransport) async throws -> DoubaoUsageSnapshot    {        do {            let confirmation = try await self.probe(apiKey: apiKey, model: model, transport: transport)            // This path starts only after a complete HTTP 200 request-limit pair            // reported zero. An immediate 429 confirms that exhausted state even            // when Ark omits the headers from the throttle response.            if confirmation.statusCode == 429 {                return confirmation.snapshot.requestLimitsReliable                    ? confirmation.snapshot                    : initial.snapshot            }            guard confirmation.hasAmbiguousZeroRemaining else {                return confirmation.snapshot            }             Self.log.warning(                """                Doubao Ark returned limit=\(confirmation.snapshot.limitRequests) remaining=0 \                with HTTP 200 twice; treating request-limit headers as unreliable.                """)            return DoubaoUsageSnapshot(                remainingRequests: confirmation.snapshot.remainingRequests,                limitRequests: confirmation.snapshot.limitRequests,                resetTime: confirmation.snapshot.resetTime,                updatedAt: confirmation.snapshot.updatedAt,                apiKeyValid: confirmation.snapshot.apiKeyValid,                totalTokens: confirmation.snapshot.totalTokens,                requestLimitsReliable: false)        } catch {            if error is CancellationError || (error as? URLError)?.code == .cancelled {                throw error            }            self.log.warning(                """                Doubao zero-remaining confirmation failed; preserving the initial exhausted state: \                \(error.localizedDescription)                """)            return initial.snapshot        }    }     private static func probe(        apiKey: String,        model: String,        transport: any ProviderHTTPTransport) async throws -> ProbeResult    {        var request = URLRequest(url: self.apiURL)        request.httpMethod = "POST"        request.timeoutInterval = 15        request.setValue("application/json", forHTTPHeaderField: "Content-Type")        request.setValue("application/json", forHTTPHeaderField: "Accept")        request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")         let body: [String: Any] = [            "model": model,            "max_tokens": 1,            "messages": [                ["role": "user", "content": "hi"],            ] as [[String: Any]],        ]         request.httpBody = try JSONSerialization.data(withJSONObject: body)         let response = try await transport.response(for: request)        let data = response.data         // Accept both 200 (success) and 429 (rate limited) – both carry rate limit headers.        guard response.statusCode == 200 || response.statusCode == 429 else {            let summary = Self.apiErrorSummary(statusCode: response.statusCode, data: data)            Self.log.error("Doubao API returned \(response.statusCode): \(summary)")            throw DoubaoUsageError.apiError(response.statusCode, summary)        }         let headers = response.response.allHeaderFields        let remaining = Self.intHeader(headers, "x-ratelimit-remaining-requests")        let limit = Self.intHeader(headers, "x-ratelimit-limit-requests")        let resetString = Self.stringHeader(headers, "x-ratelimit-reset-requests")         let resetTime: Date? = resetString.flatMap(Self.parseResetTime)         var totalTokens: Int?        if remaining == nil, limit == nil,           let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],           let usage = json["usage"] as? [String: Any]        {            totalTokens = usage["total_tokens"] as? Int        }         // 429 means the key is valid but rate-limited; treat it as valid so the UI        // shows "Active" instead of "No usage data" when headers are absent.        let keyValid = response.statusCode == 200 || response.statusCode == 429        // A request-limit header on 429 identifies request-bucket exhaustion even        // when Ark omits remaining. A bare 429 may describe another throttle.        let requestLimitsReliable = response.statusCode == 429            ? limit != nil            : limit != nil && remaining != nil         let snapshot = DoubaoUsageSnapshot(            remainingRequests: remaining ?? 0,            limitRequests: limit ?? 0,            resetTime: resetTime,            updatedAt: Date(),            apiKeyValid: keyValid,            totalTokens: totalTokens,            requestLimitsReliable: requestLimitsReliable)         Self.log.debug(            """            Doubao usage parsed remaining=\(snapshot.remainingRequests) \            limit=\(snapshot.limitRequests) valid=\(snapshot.apiKeyValid)            """)         return ProbeResult(snapshot: snapshot, statusCode: response.statusCode)    }     private static func stringHeader(_ headers: [AnyHashable: Any], _ name: String) -> String? {        if let value = headers[name] as? String { return value }        for (key, val) in headers {            if let keyStr = key as? String,               keyStr.caseInsensitiveCompare(name) == .orderedSame,               let valStr = val as? String            {                return valStr            }        }        return nil    }     private static func intHeader(_ headers: [AnyHashable: Any], _ name: String) -> Int? {        if let value = headers[name] as? String, let int = Int(value) {            return int        }        if let value = headers[name.lowercased()] as? String, let int = Int(value) {            return int        }        for (key, val) in headers {            if let keyStr = key as? String,               keyStr.lowercased() == name.lowercased(),               let valStr = val as? String,               let int = Int(valStr)            {                return int            }        }        return nil    }     private static func parseResetTime(_ value: String) -> Date? {        let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)        if trimmed.isEmpty { return nil }         let isoFormatter = ISO8601DateFormatter()        isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]        if let date = isoFormatter.date(from: trimmed) { return date }        let isoFallback = ISO8601DateFormatter()        isoFallback.formatOptions = [.withInternetDateTime]        if let date = isoFallback.date(from: trimmed) { return date }         var seconds: TimeInterval = 0        let pattern = /(\d+)([dhms])/        for match in trimmed.matches(of: pattern) {            guard let num = Double(match.1) else { continue }            switch match.2 {            case "d": seconds += num * 86400            case "h": seconds += num * 3600            case "m": seconds += num * 60            case "s": seconds += num            default: break            }        }        if seconds > 0 {            return Date().addingTimeInterval(seconds)        }         if let secs = TimeInterval(trimmed) {            return Date().addingTimeInterval(secs)        }         return nil    }     private static func apiErrorSummary(statusCode: Int, data: Data) -> String {        guard let root = try? JSONSerialization.jsonObject(with: data),              let json = root as? [String: Any]        else {            if let text = String(data: data, encoding: .utf8)?                .trimmingCharacters(in: .whitespacesAndNewlines),                !text.isEmpty            {                return self.compactText(text)            }            return "Unexpected response body (\(data.count) bytes)."        }         if let error = json["error"] as? [String: Any],           let message = error["message"] as? String        {            let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines)            if !trimmed.isEmpty { return Self.compactText(trimmed) }        }         if let message = json["message"] as? String {            let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines)            if !trimmed.isEmpty { return Self.compactText(trimmed) }        }         return "HTTP \(statusCode) (\(data.count) bytes)."    }     private static func compactText(_ text: String, maxLength: Int = 200) -> String {        let collapsed = text            .components(separatedBy: .newlines)            .joined(separator: " ")            .trimmingCharacters(in: .whitespacesAndNewlines)        if collapsed.count <= maxLength { return collapsed }        let limitIndex = collapsed.index(collapsed.startIndex, offsetBy: maxLength)        return "\(collapsed[..<limitIndex])..."    }}