1669 lines · 63.9 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif public struct AntigravityModelQuota: Sendable {    public let label: String    public let modelId: String    public let remainingFraction: Double?    public let resetTime: Date?    public let resetDescription: String?     public init(        label: String,        modelId: String,        remainingFraction: Double?,        resetTime: Date?,        resetDescription: String?)    {        self.label = label        self.modelId = modelId        self.remainingFraction = remainingFraction        self.resetTime = resetTime        self.resetDescription = resetDescription    }     public var remainingPercent: Double {        guard let remainingFraction else { return 0 }        return max(0, min(100, remainingFraction * 100))    }} private enum AntigravityModelFamily {    case claude    case gpt    case geminiPro    case geminiFlash    case unknown} private enum AntigravityUsagePool: Hashable {    case gemini    case claudeGPT     var id: String {        switch self {        case .gemini: "antigravity-gemini"        case .claudeGPT: "antigravity-claude-gpt"        }    }     var title: String {        switch self {        case .gemini: "Gemini Models"        case .claudeGPT: "Claude and GPT models"        }    }     var sortRank: Int {        switch self {        case .gemini: 0        case .claudeGPT: 1        }    }} private struct AntigravityModelVersion: Comparable {    let major: Int    let minor: Int     static func < (lhs: AntigravityModelVersion, rhs: AntigravityModelVersion) -> Bool {        if lhs.major != rhs.major { return lhs.major < rhs.major }        return lhs.minor < rhs.minor    }} private struct AntigravityNormalizedModel {    let quota: AntigravityModelQuota    let family: AntigravityModelFamily    let selectionPriority: Int?    let isImage: Bool    let isLite: Bool    let isAutocomplete: Bool    let version: AntigravityModelVersion?    let tier: Int} public enum AntigravityModelQuotaSource: Sendable {    case local    case remote} public struct AntigravityStatusSnapshot: Sendable {    public let modelQuotas: [AntigravityModelQuota]    public let accountEmail: String?    public let accountPlan: String?    public let source: AntigravityModelQuotaSource    let quotaSummary: AntigravityQuotaSummary?     public init(        modelQuotas: [AntigravityModelQuota],        accountEmail: String?,        accountPlan: String?,        source: AntigravityModelQuotaSource = .remote)    {        self.modelQuotas = modelQuotas        self.accountEmail = accountEmail        self.accountPlan = accountPlan        self.source = source        self.quotaSummary = nil    }     init(        quotaSummary: AntigravityQuotaSummary,        accountEmail: String?,        accountPlan: String?,        source: AntigravityModelQuotaSource = .local)    {        self.modelQuotas = []        self.accountEmail = accountEmail        self.accountPlan = accountPlan        self.source = source        self.quotaSummary = quotaSummary    }     public func toUsageSnapshot() throws -> UsageSnapshot {        if let quotaSummary {            return try Self.usageSnapshot(                from: quotaSummary,                accountEmail: self.accountEmail,                accountPlan: self.accountPlan)        }         guard !self.modelQuotas.isEmpty else {            throw AntigravityStatusProbeError.parseFailed("No quota models available")        }         let normalized = Self.normalizedModels(self.modelQuotas)        let summaryCandidates = normalized.filter(Self.isSummaryCandidate)        let primaryQuota = Self.representative(for: .gemini, in: summaryCandidates)        let secondaryQuota = Self.representative(for: .claudeGPT, in: summaryCandidates)        let fallbackQuota: AntigravityModelQuota? = if primaryQuota == nil, secondaryQuota == nil {            switch self.source {            case .local:                Self.fallbackRepresentative(in: normalized.filter {                    $0.family == .unknown &&                        Self.isSelectableTextModel($0) &&                        $0.quota.remainingFraction != nil                })            case .remote:                nil            }        } else {            nil        }         let primary = primaryQuota.map(Self.rateWindow(for:))        let secondary = secondaryQuota.map(Self.rateWindow(for:))        let extraWindows = Self.extraRateWindows(            from: normalized,            summaryCandidates: summaryCandidates,            compactFallbackModelID: fallbackQuota?.modelId,            representedPools: Set([                primaryQuota.map { _ in AntigravityUsagePool.gemini },                secondaryQuota.map { _ in AntigravityUsagePool.claudeGPT },            ].compactMap(\.self)))         let identity = ProviderIdentitySnapshot(            providerID: .antigravity,            accountEmail: self.accountEmail,            accountOrganization: nil,            loginMethod: self.accountPlan)        return UsageSnapshot(            primary: primary,            secondary: secondary,            tertiary: nil,            extraRateWindows: extraWindows.isEmpty ? nil : extraWindows,            updatedAt: Date(),            identity: identity)    }     func withIdentity(from snapshot: AntigravityStatusSnapshot?) -> AntigravityStatusSnapshot {        guard let snapshot else { return self }        let accountEmail = snapshot.accountEmail ?? self.accountEmail        let accountPlan = snapshot.accountPlan ?? self.accountPlan        if let quotaSummary {            return AntigravityStatusSnapshot(                quotaSummary: quotaSummary,                accountEmail: accountEmail,                accountPlan: accountPlan,                source: self.source)        }        return AntigravityStatusSnapshot(            modelQuotas: self.modelQuotas,            accountEmail: accountEmail,            accountPlan: accountPlan,            source: self.source)    }     private static func usageSnapshot(        from quotaSummary: AntigravityQuotaSummary,        accountEmail: String?,        accountPlan: String?) throws -> UsageSnapshot    {        let namedWindows = Self.quotaSummaryWindows(from: quotaSummary)        guard !namedWindows.isEmpty else {            throw AntigravityStatusProbeError.parseFailed("No quota buckets available")        }         let primary = Self.quotaSummaryRepresentative(            matching: { $0.lowercased().contains("gemini") },            in: namedWindows)        let secondary = Self.quotaSummaryRepresentative(            matching: { $0.lowercased().contains("claude") || $0.lowercased().contains("gpt") },            in: namedWindows)         let identity = ProviderIdentitySnapshot(            providerID: .antigravity,            accountEmail: accountEmail,            accountOrganization: nil,            loginMethod: accountPlan)        return UsageSnapshot(            primary: primary,            secondary: secondary,            tertiary: nil,            extraRateWindows: namedWindows,            updatedAt: Date(),            identity: identity)    }     private static func quotaSummaryRepresentative(        matching predicate: (String) -> Bool,        in windows: [NamedRateWindow]) -> RateWindow?    {        windows            .filter { $0.usageKnown && predicate($0.title) }            .max { lhs, rhs in                if lhs.window.usedPercent != rhs.window.usedPercent {                    return lhs.window.usedPercent < rhs.window.usedPercent                }                return lhs.title.localizedCaseInsensitiveCompare(rhs.title) == .orderedDescending            }?            .window    }     private static func quotaSummaryWindows(from quotaSummary: AntigravityQuotaSummary) -> [NamedRateWindow] {        let sortedGroups = quotaSummary.groups.enumerated().sorted { lhs, rhs in            let lhsRank = Self.quotaGroupSortRank(lhs.element)            let rhsRank = Self.quotaGroupSortRank(rhs.element)            if lhsRank != rhsRank {                return lhsRank < rhsRank            }            return lhs.offset < rhs.offset        }.map(\.element)         return sortedGroups.flatMap { group in            let groupTitle = Self.displayTitle(forQuotaGroup: group)            let sortedBuckets = group.buckets.enumerated().sorted { lhs, rhs in                let lhsRank = Self.quotaBucketSortRank(lhs.element)                let rhsRank = Self.quotaBucketSortRank(rhs.element)                if lhsRank != rhsRank {                    return lhsRank < rhsRank                }                return lhs.offset < rhs.offset            }.map(\.element)             return sortedBuckets.map { bucket in                let bucketTitle = Self.displayTitle(forQuotaBucket: bucket)                let remainingPercent = Self.remainingPercent(from: bucket.remainingFraction)                let usedPercent = remainingPercent.map { 100 - $0 } ?? 0                let window = RateWindow(                    usedPercent: usedPercent,                    windowMinutes: Self.windowMinutes(forQuotaBucket: bucket),                    resetsAt: bucket.resetTime,                    resetDescription: bucket.resetDescription)                return NamedRateWindow(                    id: Self.quotaSummaryWindowID(for: bucket),                    title: "\(groupTitle) \(bucketTitle)",                    window: window,                    usageKnown: !bucket.disabled && bucket.remainingFraction != nil)            }        }    }     static func isQuotaSummaryWindowID(_ id: String) -> Bool {        id.hasPrefix(self.quotaSummaryWindowIDPrefix)    }     private static let quotaSummaryWindowIDPrefix = "antigravity-quota-summary-"     private static func quotaSummaryWindowID(for bucket: AntigravityQuotaSummaryBucket) -> String {        self.quotaSummaryWindowIDPrefix + bucket.bucketId    }     private static func remainingPercent(from remainingFraction: Double?) -> Double? {        guard let remainingFraction else { return nil }        return max(0, min(100, remainingFraction * 100))    }     private static func displayTitle(forQuotaGroup group: AntigravityQuotaSummaryGroup) -> String {        let title = group.displayName.trimmingCharacters(in: .whitespacesAndNewlines)        return title.isEmpty ? "Quota" : title    }     private static func displayTitle(forQuotaBucket bucket: AntigravityQuotaSummaryBucket) -> String {        bucket.displayName    }     private static func windowMinutes(forQuotaBucket bucket: AntigravityQuotaSummaryBucket) -> Int? {        switch self.quotaBucketKind(for: bucket) {        case .session:            300        case .weekly:            10080        case .other:            nil        }    }     private enum QuotaBucketKind {        case session        case weekly        case other    }     private static func quotaGroupSortRank(_ group: AntigravityQuotaSummaryGroup) -> Int {        let title = group.displayName.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()        if title.contains("gemini") {            return 0        }        if title.contains("claude") || title.contains("gpt") {            return 1        }        return 2    }     private static func quotaBucketSortRank(_ bucket: AntigravityQuotaSummaryBucket) -> Int {        switch self.quotaBucketKind(for: bucket) {        case .session:            0        case .weekly:            1        case .other:            2        }    }     private static func quotaBucketKind(for bucket: AntigravityQuotaSummaryBucket) -> QuotaBucketKind {        let combined = "\(bucket.bucketId) \(bucket.displayName)".lowercased()        if combined.contains("5h") || combined.contains("5-hour") || combined.contains("five hour") {            return .session        }        if combined.contains("weekly") {            return .weekly        }        return .other    }     static func quotaDisplayLabel(_ quota: AntigravityModelQuota) -> String {        let trimmed = quota.label.trimmingCharacters(in: .whitespacesAndNewlines)        guard trimmed.isEmpty || trimmed == quota.modelId else { return quota.label }        return Self.humanizedModelID(quota.modelId)    }     static func humanizedModelID(_ modelId: String) -> String {        let parts = modelId.split(separator: "-").map(String.init)        var words: [String] = []        var index = 0         while index < parts.count {            let part = parts[index]            if Self.isSingleDigit(part), index + 1 < parts.count, Self.isSingleDigit(parts[index + 1]) {                var versionParts = [part]                index += 1                while index < parts.count, Self.isSingleDigit(parts[index]) {                    versionParts.append(parts[index])                    index += 1                }                words.append(versionParts.joined(separator: "."))                continue            }             if part.allSatisfy({ $0.isNumber || $0 == "." }) {                words.append(part)            } else if Self.modelLabelAcronyms.contains(part.lowercased()) {                words.append(part.uppercased())            } else {                words.append(part.prefix(1).uppercased() + part.dropFirst())            }            index += 1        }         return words.joined(separator: " ")    }     private static let modelLabelAcronyms: Set<String> = ["ai", "api", "gpt", "oss"]     private static func isSingleDigit(_ value: String) -> Bool {        value.count == 1 && value.allSatisfy(\.isNumber)    }     private static func rateWindow(for quota: AntigravityModelQuota) -> RateWindow {        RateWindow(            usedPercent: 100 - quota.remainingPercent,            windowMinutes: nil,            resetsAt: quota.resetTime,            resetDescription: quota.resetDescription)    }     private static func modelOrderPrecedes(        _ lhs: AntigravityNormalizedModel,        _ rhs: AntigravityNormalizedModel) -> Bool    {        // 1. Family rank: claude=0, geminiPro=1, geminiFlash=2, unknown=3        let lhsFamilyRank = Self.familyRank(lhs.family)        let rhsFamilyRank = Self.familyRank(rhs.family)        if lhsFamilyRank != rhsFamilyRank {            return lhsFamilyRank < rhsFamilyRank        }         // 2. Version descending (newer first); nil version sorts after non-nil        switch (lhs.version, rhs.version) {        case let (.some(lv), .some(rv)):            if lv != rv {                return lv > rv            }        case (.some, .none):            return true        case (.none, .some):            return false        case (.none, .none):            break        }         // 3. Tier ascending: High(0) < Medium(1) < Low(2)        if lhs.tier != rhs.tier {            return lhs.tier < rhs.tier        }         // 4. Label tiebreaker        return lhs.quota.label.localizedCaseInsensitiveCompare(rhs.quota.label) == .orderedAscending    }     private static func familyRank(_ family: AntigravityModelFamily) -> Int {        switch family {        case .claude: 0        case .gpt: 1        case .geminiPro: 2        case .geminiFlash: 3        case .unknown: 4        }    }     private static func isSummaryCandidate(_ model: AntigravityNormalizedModel) -> Bool {        self.usagePool(for: model) != nil && self.isSelectableTextModel(model)    }     private static func isSelectableTextModel(_ model: AntigravityNormalizedModel) -> Bool {        !model.isLite && !model.isAutocomplete && !model.isImage    }     private static func normalizedModels(_ models: [AntigravityModelQuota]) -> [AntigravityNormalizedModel] {        models.map { self.normalizeModel($0) }    }     private static func normalizeModel(_ quota: AntigravityModelQuota) -> AntigravityNormalizedModel {        let modelId = quota.modelId.lowercased()        let label = quota.label.lowercased()        let family = Self.family(forModelID: modelId, label: label)         let isLite = modelId.contains("lite") || label.contains("lite")        let isAutocomplete = modelId.contains("autocomplete") || label.contains("autocomplete") || modelId            .hasPrefix("tab_")        let isImage = modelId.contains("image") || label.contains("image")        let isSelectableTextModel = !isLite && !isAutocomplete && !isImage        let isLowPriorityGeminiPro = modelId.contains("pro-low")            || (label.contains("pro") && label.contains("low"))         let selectionPriority: Int? = switch family {        case .claude, .gpt:            0        case .geminiPro:            if isLowPriorityGeminiPro, isSelectableTextModel {                0            } else if isSelectableTextModel {                1            } else {                nil            }        case .geminiFlash:            isSelectableTextModel ? 0 : nil        case .unknown:            nil        }         let version = Self.parseVersion(from: label)        let tier = Self.parseTier(from: label, modelId: modelId)         return AntigravityNormalizedModel(            quota: quota,

Lines 1–500 of 1669.