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( 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 } if let firstTier = response.allowedTiers? .first(where: { $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty { return firstTier } if let paidTier = response.paidTier?.id?.trimmedNonEmpty { return paidTier } if let currentTier = response.currentTier?.id?.trimmedNonEmpty { return currentTier } return nil } private static func parseResetTime(_ value: String) -> Date? { let formatter = ISO8601DateFormatter() formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] if let date = formatter.date(from: value) { return date } formatter.formatOptions = [.withInternetDateTime] return formatter.date(from: value) } private static func credentialsStore(homeDirectory: String) -> AntigravityOAuthCredentialsStore { let homeURL = URL(fileURLWithPath: homeDirectory, isDirectory: true) return AntigravityOAuthCredentialsStore(fileURL: AntigravityOAuthCredentialsStore.defaultURL(home: homeURL)) } private static func resolveCredentialSource( homeDirectory: String, environment: [String: String]) throws -> ( credentials: AntigravityOAuthCredentials?, store: AntigravityOAuthCredentialsStore?) { let primaryStore = Self.credentialsStore(homeDirectory: homeDirectory) if let tokenValue = environment[AntigravityOAuthCredentialsStore.environmentCredentialsKey] { guard let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: tokenValue) else { throw AntigravityRemoteFetchError.parseFailed("Could not decode selected account credentials.") } return (credentials, nil) } return try (primaryStore.load(), primaryStore) } private struct RefreshResult { let accessToken: String let credentials: AntigravityOAuthCredentials } private static func refreshAccessToken( credentials: AntigravityOAuthCredentials, refreshToken: String, context: FetchContext) async throws -> RefreshResult { let oauthClient = try Self.refreshOAuthClient( from: credentials, oauthClientResolver: context.oauthClientResolver) var request = URLRequest(url: AntigravityOAuthConfig.tokenURL) request.httpMethod = "POST" request.timeoutInterval = context.timeout request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") request.httpBody = Self.formBody([ "client_id": oauthClient.clientID, "client_secret": oauthClient.clientSecret, "refresh_token": refreshToken, "grant_type": "refresh_token", ]) let httpResponse = try await ProviderHTTPTransportHandler(context.dataLoader).response(for: request) guard httpResponse.statusCode == 200 else { throw AntigravityRemoteFetchError.notLoggedIn } guard let json = try JSONSerialization.jsonObject(with: httpResponse.data) as? [String: Any], let accessToken = json["access_token"] as? String else { throw AntigravityRemoteFetchError.parseFailed("Could not parse refresh response") } let updatedCredentials = Self.updatedCredentials(credentials, refreshResponse: json) try await context.persistCredentials(updatedCredentials) return RefreshResult(accessToken: accessToken, credentials: updatedCredentials) } private static func refreshOAuthClient( from credentials: AntigravityOAuthCredentials, oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient?) throws -> AntigravityOAuthClient { if let clientID = credentials.clientID?.trimmedNonEmpty, let clientSecret = credentials.clientSecret?.trimmedNonEmpty { return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret) } guard let client = oauthClientResolver() else { throw AntigravityRemoteFetchError.apiError(AntigravityOAuthConfig.missingCredentialsMessage) } return client } private static func updatedCredentials( _ credentials: AntigravityOAuthCredentials, refreshResponse: [String: Any]) -> AntigravityOAuthCredentials { var credentials = credentials if let accessToken = refreshResponse["access_token"] as? String { credentials.accessToken = accessToken } if let expiresIn = refreshResponse["expires_in"] as? Double { credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + expiresIn) * 1000 } if let expiresIn = refreshResponse["expires_in"] as? Int { credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + Double(expiresIn)) * 1000 } if let idToken = refreshResponse["id_token"] as? String { credentials.idToken = idToken } return credentials } private static func formBody(_ values: [String: String]) -> Data? { var components = URLComponents() components.queryItems = values.map { key, value in URLQueryItem(name: key, value: value) } return components.query?.data(using: .utf8) } private struct TokenClaims { let email: String? let hostedDomain: String? } private static func extractClaims(from credentials: AntigravityOAuthCredentials) -> TokenClaims { let tokenClaims = Self.extractClaimsFromToken(credentials.idToken) return TokenClaims( email: tokenClaims.email ?? credentials.email?.trimmedNonEmpty, hostedDomain: tokenClaims.hostedDomain) } private static func extractClaimsFromToken(_ idToken: String?) -> TokenClaims { guard let idToken else { return TokenClaims(email: nil, hostedDomain: nil) } let parts = idToken.components(separatedBy: ".") guard parts.count >= 2 else { return TokenClaims(email: nil, hostedDomain: nil) } var payload = parts[1] .replacingOccurrences(of: "-", with: "+") .replacingOccurrences(of: "_", with: "/") let remainder = payload.count % 4 if remainder > 0 { payload += String(repeating: "=", count: 4 - remainder) } guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters), let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { return TokenClaims(email: nil, hostedDomain: nil) } return TokenClaims( email: (json["email"] as? String)?.trimmedNonEmpty, hostedDomain: (json["hd"] as? String)?.trimmedNonEmpty) } } extension String { fileprivate var trimmedNonEmpty: String? { let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines) return trimmed.isEmpty ? nil : trimmed } } private struct ProjectReference: Decodable { let value: String? init(from decoder: Decoder) throws { let single = try decoder.singleValueContainer() if let stringValue = try? single.decode(String.self) { self.value = stringValue return } let keyed = try decoder.container(keyedBy: CodingKeys.self) self.value = try keyed.decodeIfPresent(String.self, forKey: .id) ?? keyed.decodeIfPresent(String.self, forKey: .projectID) } private enum CodingKeys: String, CodingKey { case id case projectID = "projectId" } } private struct CodeAssistResponse: Decodable { let planInfo: CodeAssistPlanInfo? let currentTier: TierInfo? let paidTier: TierInfo? let allowedTiers: [AllowedTier]? let cloudaicompanionProject: ProjectReference? var projectID: String? { self.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty } } private struct CodeAssistPlanInfo: Decodable { let planType: String? } private struct TierInfo: Decodable { let id: String? let name: String? } private struct AllowedTier: Decodable { let id: String? let isDefault: Bool? } private struct OnboardResponse: Decodable { let response: OnboardInnerResponse? var projectID: String? { self.response?.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty } } private struct OnboardInnerResponse: Decodable { let cloudaicompanionProject: ProjectReference? } private struct FetchAvailableModelsResponse: Decodable { let models: [String: AntigravityRemoteModel]? } private struct RetrieveUserQuotaResponse: Decodable { let buckets: [RetrieveUserQuotaBucket]? } private struct RetrieveUserQuotaBucket: Decodable { let modelId: String? let remainingFraction: Double? let resetTime: String? } private struct AntigravityRemoteModel: Decodable { let displayName: String? let label: String? let quotaInfo: AntigravityRemoteQuotaInfo? } private struct AntigravityRemoteQuotaInfo: Decodable { let remainingFraction: Double? let resetTime: String? } extension String? { fileprivate var trimmedNonEmpty: String? { self?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty } } extension String { fileprivate var nilIfEmpty: String? { self.isEmpty ? nil : self } }