import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif public enum AntigravityRemoteFetchError: LocalizedError, Sendable, Equatable {    case notLoggedIn    case permissionDenied(String)    case apiError(String)    case parseFailed(String)     public var errorDescription: String? {        switch self {        case .notLoggedIn:            "Antigravity Google auth not found. Use Antigravity login to authenticate."        case let .permissionDenied(message):            "Antigravity remote API permission denied: \(message)"        case let .apiError(message):            "Antigravity remote API error: \(message)"        case let .parseFailed(message):            "Could not parse Antigravity remote usage: \(message)"        }    }} public struct AntigravityRemoteUsageFetcher: Sendable {    public var timeout: TimeInterval = 10.0    public var homeDirectory: String    public var environment: [String: String]    public var dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)    public var oauthClientResolver: @Sendable () -> AntigravityOAuthClient?    public var credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void     private static let log = CodexBarLog.logger(LogCategories.antigravity)    private static let userAgent = "antigravity"    private static let baseURL = "https://cloudcode-pa.googleapis.com"    private static let loadCodeAssistEndpoint = "\(baseURL)/v1internal:loadCodeAssist"    private static let onboardUserEndpoint = "\(baseURL)/v1internal:onboardUser"    private static let fetchAvailableModelsEndpoint = "\(baseURL)/v1internal:fetchAvailableModels"    private static let retrieveUserQuotaEndpoint = "\(baseURL)/v1internal:retrieveUserQuota"    private static let refreshSafetyWindow: TimeInterval = 60     private struct FetchContext {        let timeout: TimeInterval        let store: AntigravityOAuthCredentialsStore?        let dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)        let oauthClientResolver: @Sendable () -> AntigravityOAuthClient?        let credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void         func persistCredentials(_ credentials: AntigravityOAuthCredentials) async throws {            if let store {                try store.save(credentials)            } else {                try await self.credentialsUpdateHandler(credentials)            }        }    }     public init(        timeout: TimeInterval = 10.0,        homeDirectory: String = NSHomeDirectory(),        environment: [String: String] = ProcessInfo.processInfo.environment,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse) = { request in            try await ProviderHTTPClient.shared.data(for: request)        },        oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient? = {            AntigravityOAuthConfig.resolvedClient()        },        credentialsUpdateHandler: @escaping @Sendable (AntigravityOAuthCredentials) async throws -> Void = { _ in })    {        self.timeout = timeout        self.homeDirectory = homeDirectory        self.environment = environment        self.dataLoader = dataLoader        self.oauthClientResolver = oauthClientResolver        self.credentialsUpdateHandler = credentialsUpdateHandler    }     public func fetch() async throws -> AntigravityStatusSnapshot {        let source = try Self.resolveCredentialSource(homeDirectory: self.homeDirectory, environment: self.environment)        guard let credentials = source.credentials else {            throw AntigravityRemoteFetchError.notLoggedIn        }        return try await self.fetchSnapshot(            using: credentials,            store: source.store)    }     private func fetchSnapshot(        using initialCredentials: AntigravityOAuthCredentials,        store: AntigravityOAuthCredentialsStore?) async throws        -> AntigravityStatusSnapshot    {        guard let storedAccessToken = initialCredentials.accessToken?.trimmedNonEmpty else {            throw AntigravityRemoteFetchError.notLoggedIn        }         var credentials = initialCredentials        var accessToken = storedAccessToken        let context = FetchContext(            timeout: self.timeout,            store: store,            dataLoader: self.dataLoader,            oauthClientResolver: self.oauthClientResolver,            credentialsUpdateHandler: self.credentialsUpdateHandler)        if Self.shouldRefresh(expiryDate: credentials.expiryDate, now: Date()) {            guard let refreshToken = credentials.refreshToken?.trimmedNonEmpty else {                throw AntigravityRemoteFetchError.notLoggedIn            }            let refreshed = try await Self.refreshAccessToken(                credentials: credentials,                refreshToken: refreshToken,                context: context)            accessToken = refreshed.accessToken            credentials = refreshed.credentials            if let store {                credentials = try store.load() ?? credentials            }            credentials.accessToken = credentials.accessToken?.trimmedNonEmpty ?? accessToken        }         let claims = Self.extractClaims(from: credentials)        let codeAssist = try await Self.loadCodeAssist(            accessToken: accessToken,            timeout: self.timeout,            dataLoader: self.dataLoader)        let projectId = try await Self.resolveProjectID(            accessToken: accessToken,            storedProjectID: credentials.projectID?.trimmedNonEmpty,            initialResponse: codeAssist,            context: context)        if let projectId, credentials.projectID?.trimmedNonEmpty != projectId {            credentials.projectID = projectId            do {                try await context.persistCredentials(credentials)            } catch {                Self.log.warning("Could not persist Antigravity project ID: \(error.localizedDescription)")            }        }        let models = try await Self.fetchModelQuotas(            accessToken: accessToken,            projectId: projectId,            timeout: self.timeout,            dataLoader: self.dataLoader)         return AntigravityStatusSnapshot(            modelQuotas: models,            accountEmail: claims.email,            accountPlan: Self.resolvePlan(response: codeAssist, claims: claims),            source: .remote)    }     private static func shouldRefresh(expiryDate: Date?, now: Date) -> Bool {        guard let expiryDate else { return false }        return expiryDate.timeIntervalSince(now) <= Self.refreshSafetyWindow    }     private static func loadCodeAssist(        accessToken: String,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> CodeAssistResponse    {        let body = [            "metadata": [                "ideType": "ANTIGRAVITY",                "platform": "PLATFORM_UNSPECIFIED",                "pluginType": "GEMINI",            ],        ]        return try await Self.sendRequest(            endpoint: Self.loadCodeAssistEndpoint,            accessToken: accessToken,            body: body,            timeout: timeout,            dataLoader: dataLoader)    }     private static func fetchAvailableModels(        accessToken: String,        projectId: String?,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> FetchAvailableModelsResponse    {        let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty {            ["project": projectId]        } else {            [:]        }        return try await Self.sendRequest(            endpoint: Self.fetchAvailableModelsEndpoint,            accessToken: accessToken,            body: body,            timeout: timeout,            dataLoader: dataLoader)    }     private static func fetchModelQuotas(        accessToken: String,        projectId: String?,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> [AntigravityModelQuota]    {        do {            let response = try await Self.fetchAvailableModels(                accessToken: accessToken,                projectId: projectId,                timeout: timeout,                dataLoader: dataLoader)            let modelQuotas = try Self.parseModelQuotas(response)            if Self.shouldVerifyFullRemoteQuotas(modelQuotas) {                let quotaBuckets = try await Self.fetchQuotaBucketsIfPermitted(                    accessToken: accessToken,                    projectId: projectId,                    timeout: timeout,                    dataLoader: dataLoader)                guard let quotaBuckets, Self.hasQuotaFractionData(quotaBuckets) else {                    return []                }                return Self.mergeVerifiedQuotas(modelQuotas: modelQuotas, verifiedQuotas: quotaBuckets)            }            return modelQuotas        } catch let error as AntigravityRemoteFetchError {            guard case .permissionDenied = error else {                throw error            }            Self.log.info("Falling back to retrieveUserQuota for Antigravity remote usage")            return try await Self.fetchQuotaBucketsIfPermitted(                accessToken: accessToken,                projectId: projectId,                timeout: timeout,                dataLoader: dataLoader) ?? []        }    }     private static func mergeVerifiedQuotas(        modelQuotas: [AntigravityModelQuota],        verifiedQuotas: [AntigravityModelQuota]) -> [AntigravityModelQuota]    {        var verifiedByModelID = Dictionary(            uniqueKeysWithValues: verifiedQuotas.map { (Self.quotaKey($0), $0) })        var merged = modelQuotas.compactMap { modelQuota -> AntigravityModelQuota? in            guard let verifiedQuota = verifiedByModelID.removeValue(forKey: Self.quotaKey(modelQuota)) else {                return nil            }            let resetTime = verifiedQuota.resetTime ?? modelQuota.resetTime            return AntigravityModelQuota(                label: modelQuota.label,                modelId: modelQuota.modelId,                remainingFraction: verifiedQuota.remainingFraction ?? modelQuota.remainingFraction,                resetTime: resetTime,                resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })        }        let unmatchedVerifiedQuotas = verifiedByModelID.values            .filter { $0.remainingFraction != nil }            .sorted { lhs, rhs in                lhs.modelId.localizedCaseInsensitiveCompare(rhs.modelId) == .orderedAscending            }        merged.append(contentsOf: unmatchedVerifiedQuotas)        return merged    }     private static func quotaKey(_ quota: AntigravityModelQuota) -> String {        quota.modelId.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()    }     private static func shouldVerifyFullRemoteQuotas(_ quotas: [AntigravityModelQuota]) -> Bool {        guard !quotas.isEmpty else { return false }        return quotas.allSatisfy { quota in            guard let remaining = quota.remainingFraction else { return false }            return remaining >= 0.999        }    }     private static func hasQuotaFractionData(_ quotas: [AntigravityModelQuota]) -> Bool {        quotas.contains { quota in            quota.remainingFraction != nil        }    }     private static func fetchQuotaBucketsIfPermitted(        accessToken: String,        projectId: String?,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> [AntigravityModelQuota]?    {        do {            let response = try await Self.retrieveUserQuota(                accessToken: accessToken,                projectId: projectId,                timeout: timeout,                dataLoader: dataLoader)            return try Self.parseQuotaBuckets(response)        } catch let quotaError as AntigravityRemoteFetchError {            guard case .permissionDenied = quotaError else {                throw quotaError            }            Self.log.info("Antigravity remote quota endpoint is not permitted for this account")            return nil        }    }     private static func retrieveUserQuota(        accessToken: String,        projectId: String?,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> RetrieveUserQuotaResponse    {        let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty {            ["project": projectId]        } else {            [:]        }        return try await Self.sendRequest(            endpoint: Self.retrieveUserQuotaEndpoint,            accessToken: accessToken,            body: body,            timeout: timeout,            dataLoader: dataLoader)    }     private static func resolveProjectID(        accessToken: String,        storedProjectID: String?,        initialResponse: CodeAssistResponse,        context: FetchContext) async throws        -> String?    {        if let storedProjectID {            return storedProjectID        }         if let projectID = initialResponse.projectID {            return projectID        }         guard let tierID = Self.pickOnboardTier(from: initialResponse) else {            return nil        }         let onboardBody: [String: Any] = [            "tierId": tierID,            "metadata": [                "ideType": "ANTIGRAVITY",                "platform": "PLATFORM_UNSPECIFIED",                "pluginType": "GEMINI",            ],        ]         do {            let onboardResponse: OnboardResponse = try await Self.sendRequest(                endpoint: Self.onboardUserEndpoint,                accessToken: accessToken,                body: onboardBody,                timeout: context.timeout,                dataLoader: context.dataLoader)            if let projectID = onboardResponse.projectID {                return projectID            }        } catch {            Self.log.warning("Antigravity onboarding request failed", metadata: [                "error": "\(error.localizedDescription)",            ])        }         for _ in 0..<5 {            try? await Task.sleep(for: .milliseconds(2000))            let refreshed = try await Self.loadCodeAssist(                accessToken: accessToken,                timeout: context.timeout,                dataLoader: context.dataLoader)            if let projectID = refreshed.projectID {                return projectID            }        }         return nil    }     private static func sendRequest<Response: Decodable>(        endpoint: String,        accessToken: String,        body: [String: Any],        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> Response    {        guard let url = URL(string: endpoint) else {            throw AntigravityRemoteFetchError.apiError("Invalid endpoint URL")        }         var request = URLRequest(url: url)        request.httpMethod = "POST"        request.timeoutInterval = timeout        request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")        request.setValue("application/json", forHTTPHeaderField: "Content-Type")        request.setValue(Self.userAgent, forHTTPHeaderField: "User-Agent")        request.httpBody = try JSONSerialization.data(withJSONObject: body)         let httpResponse = try await ProviderHTTPTransportHandler(dataLoader).response(for: request)         switch httpResponse.statusCode {        case 200:            break        case 401:            throw AntigravityRemoteFetchError.notLoggedIn        case 403:            let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty ?? "HTTP 403"            throw AntigravityRemoteFetchError.permissionDenied(message)        default:            let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty                ?? "HTTP \(httpResponse.statusCode)"            throw AntigravityRemoteFetchError.apiError("HTTP \(httpResponse.statusCode): \(message)")        }         do {            return try JSONDecoder().decode(Response.self, from: httpResponse.data)        } catch {            throw AntigravityRemoteFetchError.parseFailed(error.localizedDescription)        }    }     private static func parseModelQuotas(_ response: FetchAvailableModelsResponse) throws -> [AntigravityModelQuota] {        let models = response.models ?? [:]        return models.compactMap { modelID, model in            guard let quotaInfo = model.quotaInfo else { return nil }            let resetTime = quotaInfo.resetTime.flatMap(Self.parseResetTime(_:))            let label = model.displayName?.trimmedNonEmpty                ?? model.label?.trimmedNonEmpty                ?? modelID            return AntigravityModelQuota(                label: label,                modelId: modelID,                remainingFraction: quotaInfo.remainingFraction,                resetTime: resetTime,                resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })        }    }     private static func parseQuotaBuckets(_ response: RetrieveUserQuotaResponse) throws -> [AntigravityModelQuota] {        guard let buckets = response.buckets else {            throw AntigravityRemoteFetchError.parseFailed("No quota buckets in response")        }        guard !buckets.isEmpty else { return [] }         var modelQuotaMap: [String: (fraction: Double?, resetTime: String?)] = [:]        for bucket in buckets {            guard let modelID = bucket.modelId?.trimmedNonEmpty else { continue }            let next = (bucket.remainingFraction, bucket.resetTime)            if let existing = modelQuotaMap[modelID] {                let existingValue = existing.fraction ?? Double.greatestFiniteMagnitude                let nextValue = next.0 ?? Double.greatestFiniteMagnitude                if nextValue < existingValue {                    modelQuotaMap[modelID] = next                }            } else {                modelQuotaMap[modelID] = next            }        }         return modelQuotaMap.keys.sorted().compactMap { modelID in            guard let info = modelQuotaMap[modelID] else { return nil }            let resetTime = info.resetTime.flatMap(Self.parseResetTime(_:))            return AntigravityModelQuota(                label: modelID,                modelId: modelID,                remainingFraction: info.fraction,                resetTime: resetTime,                resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })        }    }     private static func resolvePlan(response: CodeAssistResponse, claims: TokenClaims) -> String? {        if let planType = response.planInfo?.planType?.trimmedNonEmpty {            return planType        }         switch (response.currentTier?.id?.trimmedNonEmpty, claims.hostedDomain) {        case ("standard-tier", _):            return "Paid"        case ("free-tier", .some):            return "Workspace"        case ("free-tier", .none):            return "Free"        case ("legacy-tier", _):            return "Legacy"        default:            return response.currentTier?.name?.trimmedNonEmpty        }    }     private static func pickOnboardTier(from response: CodeAssistResponse) -> String? {        if let defaultTier = response.allowedTiers?            .first(where: { $0.isDefault == true && $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty        {            return defaultTier

Lines 1–500 of 777. Load more