wehub-resource-sync
778 lines
29 KiB
Swift
778 lines
29 KiB
Swift
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
|
|
}
|
|
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
|
|
}
|
|
}
|