From 2dbbed6f2e45e9402e72ff91070d6884a3bace15 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Tue, 27 Jan 2026 16:45:20 +0300 Subject: [PATCH 01/19] feat: add Antigravity OAuth support with manual token and local import Add OAuth-based authentication for Antigravity provider with multiple credential sources: - OAuth flow with token refresh capability - Manual token input in settings - Local credential import from SQLite database - Cloud Code API client for quota fetching - Settings snapshot support for Antigravity provider --- .../AntigravityProviderImplementation.swift | 5 + .../AntigravitySettingsStore.swift | 52 +++ .../CodexBarCore/Config/CodexBarConfig.swift | 8 +- .../AntigravityAuthorizedFetchStrategy.swift | 120 +++++++ .../AntigravityCloudCodeClient.swift | 233 +++++++++++++ .../AntigravityLocalImporter.swift | 212 ++++++++++++ .../AntigravityOAuthCredentials.swift | 240 ++++++++++++++ .../AntigravityOAuthFlow.swift | 306 ++++++++++++++++++ .../AntigravityTokenRefresher.swift | 114 +++++++ .../AntigravityUsageSource.swift | 29 ++ .../AntigravityProviderDescriptor.swift | 25 +- .../Providers/ProviderSettingsSnapshot.swift | 31 +- .../CodexBarTests/AntigravityOAuthTests.swift | 203 ++++++++++++ docs/antigravity.md | 64 +++- 14 files changed, 1629 insertions(+), 13 deletions(-) create mode 100644 Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityTokenRefresher.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift create mode 100644 Tests/CodexBarTests/AntigravityOAuthTests.swift diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index 88492f07..177b12d4 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -10,6 +10,11 @@ struct AntigravityProviderImplementation: ProviderImplementation { await AntigravityStatusProbe.detectVersion() } + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .antigravity(context.settings.antigravitySettingsSnapshot()) + } + @MainActor func runLoginFlow(context: ProviderLoginContext) async -> Bool { await context.controller.runAntigravityLoginFlow() diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift new file mode 100644 index 00000000..1138d70f --- /dev/null +++ b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift @@ -0,0 +1,52 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var antigravityUsageSource: AntigravityUsageSource { + get { + guard let rawValue = self.configSnapshot.providerConfig(for: .antigravity)?.usageSource else { + return .auto + } + return AntigravityUsageSource(rawValue: rawValue) ?? .auto + } + set { + self.updateProviderConfig(provider: .antigravity) { entry in + entry.usageSource = newValue.rawValue + } + self.logProviderModeChange(provider: .antigravity, field: "usageSource", value: newValue.rawValue) + } + } + + var antigravityManualToken: String { + get { self.configSnapshot.providerConfig(for: .antigravity)?.manualToken ?? "" } + set { + self.updateProviderConfig(provider: .antigravity) { entry in + entry.manualToken = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .antigravity, field: "manualToken", value: newValue) + } + } + + var antigravityAccountEmail: String? { + AntigravityOAuthCredentialsStore.load()?.email + } + + var antigravityHasCredentials: Bool { + if let credentials = AntigravityOAuthCredentialsStore.load() { + return !credentials.accessToken.isEmpty || credentials.isRefreshable + } + return false + } + + func clearAntigravityCredentials() { + AntigravityOAuthCredentialsStore.clear() + } +} + +extension SettingsStore { + func antigravitySettingsSnapshot() -> ProviderSettingsSnapshot.AntigravityProviderSettings { + ProviderSettingsSnapshot.AntigravityProviderSettings( + usageSource: self.antigravityUsageSource, + manualToken: self.antigravityManualToken) + } +} diff --git a/Sources/CodexBarCore/Config/CodexBarConfig.swift b/Sources/CodexBarCore/Config/CodexBarConfig.swift index c4bbbc2c..f151b070 100644 --- a/Sources/CodexBarCore/Config/CodexBarConfig.swift +++ b/Sources/CodexBarCore/Config/CodexBarConfig.swift @@ -82,6 +82,8 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { public var region: String? public var workspaceID: String? public var tokenAccounts: ProviderTokenAccountData? + public var usageSource: String? + public var manualToken: String? public init( id: UsageProvider, @@ -92,7 +94,9 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { cookieSource: ProviderCookieSource? = nil, region: String? = nil, workspaceID: String? = nil, - tokenAccounts: ProviderTokenAccountData? = nil) + tokenAccounts: ProviderTokenAccountData? = nil, + usageSource: String? = nil, + manualToken: String? = nil) { self.id = id self.enabled = enabled @@ -103,6 +107,8 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { self.region = region self.workspaceID = workspaceID self.tokenAccounts = tokenAccounts + self.usageSource = usageSource + self.manualToken = manualToken } public var sanitizedAPIKey: String? { diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift new file mode 100644 index 00000000..22cf10b0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift @@ -0,0 +1,120 @@ +import Foundation + +public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { + public let id: String = "antigravity.authorized" + public let kind: ProviderFetchKind = .oauth + + private static let log = CodexBarLog.logger(LogCategories.antigravity) + + public init() {} + + public func isAvailable(_ context: ProviderFetchContext) async -> Bool { + if let credentials = AntigravityOAuthCredentialsStore.load() { + if !credentials.accessToken.isEmpty { return true } + if credentials.isRefreshable { return true } + } + + if let manualToken = context.settings?.antigravityManualToken, + !manualToken.isEmpty + { + return true + } + + return false + } + + public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let credentials = try await self.resolveCredentials(context: context) + + let accessToken: String + if credentials.needsRefresh, credentials.isRefreshable { + Self.log.info("Antigravity credentials need refresh") + let refreshed = try await self.refreshCredentials(credentials) + accessToken = refreshed.accessToken + } else if credentials.accessToken.isEmpty, credentials.isRefreshable { + Self.log.info("Antigravity credentials have no access token, refreshing") + let refreshed = try await self.refreshCredentials(credentials) + accessToken = refreshed.accessToken + } else { + accessToken = credentials.accessToken + } + + let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: accessToken) + let snapshot = AntigravityStatusSnapshot( + modelQuotas: quota.models, + accountEmail: credentials.email ?? quota.email, + accountPlan: nil) + + let usage = try snapshot.toUsageSnapshot() + return self.makeResult(usage: usage, sourceLabel: "authorized") + } + + public func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool { + if let oauthError = error as? AntigravityOAuthCredentialsError { + switch oauthError { + case .invalidGrant: + return true + case .notFound: + return true + default: + return false + } + } + return false + } + + private func resolveCredentials(context: ProviderFetchContext) async throws -> AntigravityOAuthCredentials { + if let cached = AntigravityOAuthCredentialsStore.load() { + return cached + } + + if let manualToken = context.settings?.antigravityManualToken, + !manualToken.isEmpty + { + if let parsed = AntigravityOAuthCredentialsStore.parseManualToken(manualToken) { + if parsed.isRefreshable, parsed.accessToken.isEmpty { + let refreshed = try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( + refreshToken: parsed.refreshToken!) + AntigravityOAuthCredentialsStore.save(refreshed) + return refreshed + } + return parsed + } + } + + if AntigravityLocalImporter.isAvailable() { + let localInfo = try await AntigravityLocalImporter.importCredentials() + if let refreshToken = localInfo.refreshToken, !refreshToken.isEmpty { + let credentials = try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( + refreshToken: refreshToken, + fallbackEmail: localInfo.email) + AntigravityOAuthCredentialsStore.save(credentials) + return credentials + } + if let accessToken = localInfo.accessToken, !accessToken.isEmpty { + let credentials = AntigravityOAuthCredentials( + accessToken: accessToken, + refreshToken: nil, + expiresAt: nil, + email: localInfo.email, + scopes: AntigravityOAuthConfig.scopes) + AntigravityOAuthCredentialsStore.save(credentials) + return credentials + } + } + + throw AntigravityOAuthCredentialsError.notFound + } + + private func refreshCredentials(_ credentials: AntigravityOAuthCredentials) async throws -> AntigravityOAuthCredentials { + guard let refreshToken = credentials.refreshToken else { + throw AntigravityOAuthCredentialsError.invalidGrant + } + + let refreshed = try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( + refreshToken: refreshToken, + fallbackEmail: credentials.email) + AntigravityOAuthCredentialsStore.save(refreshed) + return refreshed + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift new file mode 100644 index 00000000..507ec419 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift @@ -0,0 +1,233 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum AntigravityCloudCodeConfig { + public static let baseURLs = [ + "https://daily-cloudcode-pa.googleapis.com", + "https://cloudcode-pa.googleapis.com", + "https://daily-cloudcode-pa.sandbox.googleapis.com", + ] + + public static let fetchAvailableModelsPath = "/v1internal:fetchAvailableModels" + public static let loadCodeAssistPath = "/v1internal:loadCodeAssist" + public static let onboardUserPath = "/v1internal:onboardUser" + public static let userAgent = "antigravity" + + public static let metadata: [String: String] = [ + "ideType": "ANTIGRAVITY", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + ] + + public static let defaultAttempts = 2 + public static let backoffBaseMs = 500 + public static let backoffMaxMs = 4000 +} + +public struct AntigravityCloudCodeQuota: Sendable { + public let models: [AntigravityModelQuota] + public let email: String? + public let projectId: String? +} + +public struct AntigravityProjectInfo: Sendable { + public let projectId: String? + public let tierId: String? +} + +public enum AntigravityCloudCodeClient { + private static let log = CodexBarLog.logger(LogCategories.antigravity) + private static let httpTimeout: TimeInterval = 15.0 + + public static func fetchQuota(accessToken: String, projectId: String? = nil) async throws -> AntigravityCloudCodeQuota { + try await self.requestWithRetry { baseURL in + try await self.fetchQuotaFromEndpoint( + baseURL: baseURL, + accessToken: accessToken, + projectId: projectId) + } + } + + public static func loadProjectInfo(accessToken: String) async throws -> AntigravityProjectInfo { + try await self.requestWithRetry { baseURL in + try await self.loadProjectInfoFromEndpoint(baseURL: baseURL, accessToken: accessToken) + } + } + + private static func requestWithRetry( + _ operation: (String) async throws -> T) async throws -> T + { + var lastError: Error? + + for attempt in 1...AntigravityCloudCodeConfig.defaultAttempts { + if attempt > 1 { + let delay = self.getBackoffDelay(attempt: attempt) + self.log.info("Cloud Code retry round \(attempt)/\(AntigravityCloudCodeConfig.defaultAttempts) in \(delay)ms") + try await Task.sleep(nanoseconds: UInt64(delay) * 1_000_000) + } + + for baseURL in AntigravityCloudCodeConfig.baseURLs { + do { + return try await operation(baseURL) + } catch let error as AntigravityOAuthCredentialsError { + if case .invalidGrant = error { + throw error + } + lastError = error + self.log.debug("Cloud Code request failed (\(baseURL)): \(error.localizedDescription)") + } catch { + lastError = error + self.log.debug("Cloud Code request failed (\(baseURL)): \(error.localizedDescription)") + } + } + } + + throw lastError ?? AntigravityOAuthCredentialsError.networkError("All Cloud Code endpoints failed") + } + + private static func getBackoffDelay(attempt: Int) -> Int { + let raw = AntigravityCloudCodeConfig.backoffBaseMs * Int(pow(2.0, Double(attempt - 2))) + let jitter = Int.random(in: 0..<100) + return min(raw + jitter, AntigravityCloudCodeConfig.backoffMaxMs) + } + + private static func fetchQuotaFromEndpoint( + baseURL: String, + accessToken: String, + projectId: String?) async throws -> AntigravityCloudCodeQuota + { + let urlString = baseURL + AntigravityCloudCodeConfig.fetchAvailableModelsPath + guard let url = URL(string: urlString) else { + throw AntigravityOAuthCredentialsError.networkError("Invalid Cloud Code URL") + } + + var requestBody: [String: Any] = [:] + if let projectId { + requestBody["project"] = projectId + } + + let data = try await self.makeRequest(url: url, body: requestBody, accessToken: accessToken) + return try self.parseQuotaResponse(data: data) + } + + private static func loadProjectInfoFromEndpoint( + baseURL: String, + accessToken: String) async throws -> AntigravityProjectInfo + { + let urlString = baseURL + AntigravityCloudCodeConfig.loadCodeAssistPath + guard let url = URL(string: urlString) else { + throw AntigravityOAuthCredentialsError.networkError("Invalid Cloud Code URL") + } + + let requestBody: [String: Any] = ["metadata": AntigravityCloudCodeConfig.metadata] + let data = try await self.makeRequest(url: url, body: requestBody, accessToken: accessToken) + return try self.parseProjectInfoResponse(data: data) + } + + private static func makeRequest( + url: URL, + body: [String: Any], + accessToken: String) async throws -> Data + { + let bodyData = try JSONSerialization.data(withJSONObject: body) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = bodyData + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue(AntigravityCloudCodeConfig.userAgent, forHTTPHeaderField: "User-Agent") + request.timeoutInterval = self.httpTimeout + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let http = response as? HTTPURLResponse else { + throw AntigravityOAuthCredentialsError.networkError("Invalid response") + } + + if http.statusCode == 401 || http.statusCode == 403 { + throw AntigravityOAuthCredentialsError.invalidGrant + } + + guard http.statusCode == 200 else { + let errorBody = String(data: data, encoding: .utf8) ?? "" + throw AntigravityOAuthCredentialsError.networkError("HTTP \(http.statusCode): \(errorBody)") + } + + return data + } + + private static func parseProjectInfoResponse(data: Data) throws -> AntigravityProjectInfo { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid Cloud Code response JSON") + } + + let projectId = self.extractProjectId(from: json["cloudaicompanionProject"]) + let tierId: String? + if let paidTier = json["paidTier"] as? [String: Any], let id = paidTier["id"] as? String { + tierId = id + } else if let currentTier = json["currentTier"] as? [String: Any], let id = currentTier["id"] as? String { + tierId = id + } else { + tierId = nil + } + + return AntigravityProjectInfo(projectId: projectId, tierId: tierId) + } + + private static func extractProjectId(from project: Any?) -> String? { + if let projectString = project as? String, !projectString.isEmpty { + return projectString + } + if let projectDict = project as? [String: Any], let id = projectDict["id"] as? String, !id.isEmpty { + return id + } + return nil + } + + private static func parseQuotaResponse(data: Data) throws -> AntigravityCloudCodeQuota { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid Cloud Code response JSON") + } + + var models: [AntigravityModelQuota] = [] + + if let modelsDict = json["models"] as? [String: [String: Any]] { + for (modelKey, modelData) in modelsDict { + if let quota = self.parseModelFromDict(key: modelKey, data: modelData) { + models.append(quota) + } + } + } + + return AntigravityCloudCodeQuota(models: models, email: nil, projectId: nil) + } + + private static func parseModelFromDict(key: String, data: [String: Any]) -> AntigravityModelQuota? { + let displayName = (data["displayName"] as? String) ?? key + let modelId = (data["model"] as? String) ?? key + + var remainingFraction: Double? + var resetTime: Date? + + if let quotaInfo = data["quotaInfo"] as? [String: Any] { + remainingFraction = quotaInfo["remainingFraction"] as? Double + + if let resetTimeStr = quotaInfo["resetTime"] as? String { + resetTime = ISO8601DateFormatter().date(from: resetTimeStr) + if resetTime == nil, let seconds = Double(resetTimeStr) { + resetTime = Date(timeIntervalSince1970: seconds) + } + } + } + + return AntigravityModelQuota( + label: displayName, + modelId: modelId, + remainingFraction: remainingFraction, + resetTime: resetTime, + resetDescription: nil) + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift new file mode 100644 index 00000000..9311bed6 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift @@ -0,0 +1,212 @@ +import Foundation +import SQLite3 + +public enum AntigravityLocalImporter { + private static let log = CodexBarLog.logger(LogCategories.antigravity) + + public struct LocalCredentialInfo: Sendable { + public let accessToken: String? + public let refreshToken: String? + public let email: String? + public let name: String? + + public var hasAccessToken: Bool { + guard let accessToken else { return false } + return !accessToken.isEmpty + } + + public var hasRefreshToken: Bool { + guard let refreshToken else { return false } + return !refreshToken.isEmpty + } + } + + public static func stateDbPath() -> URL { + let home = FileManager.default.homeDirectoryForCurrentUser + return home + .appendingPathComponent("Library") + .appendingPathComponent("Application Support") + .appendingPathComponent("Antigravity") + .appendingPathComponent("User") + .appendingPathComponent("globalStorage") + .appendingPathComponent("state.vscdb") + } + + public static func importCredentials() async throws -> LocalCredentialInfo { + let dbPath = self.stateDbPath() + guard FileManager.default.fileExists(atPath: dbPath.path) else { + throw AntigravityOAuthCredentialsError.notFound + } + + var refreshToken: String? + if let protoInfo = try? self.readProtoTokenInfo(dbPath: dbPath) { + refreshToken = protoInfo.refreshToken + } + + if let authStatus = try? self.readAuthStatus(dbPath: dbPath) { + return LocalCredentialInfo( + accessToken: authStatus.apiKey, + refreshToken: refreshToken, + email: authStatus.email, + name: authStatus.name) + } + + if let refreshToken, !refreshToken.isEmpty { + return LocalCredentialInfo( + accessToken: nil, + refreshToken: refreshToken, + email: nil, + name: nil) + } + + throw AntigravityOAuthCredentialsError.notFound + } + + public static func isAvailable() -> Bool { + FileManager.default.fileExists(atPath: self.stateDbPath().path) + } + + private struct AuthStatus { + let apiKey: String? + let email: String? + let name: String? + } + + private struct ProtoTokenInfo { + let accessToken: String? + let refreshToken: String? + let tokenType: String? + let expirySeconds: Int? + } + + private static func readAuthStatus(dbPath: URL) throws -> AuthStatus { + let json = try self.readStateValue(dbPath: dbPath, key: "antigravityAuthStatus") + guard let data = json.data(using: .utf8), + let dict = try JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid antigravityAuthStatus JSON") + } + + let apiKey = (dict["apiKey"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let email = (dict["email"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let name = (dict["name"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + + return AuthStatus(apiKey: apiKey, email: email, name: name) + } + + private static func readProtoTokenInfo(dbPath: URL) throws -> ProtoTokenInfo { + let base64 = try self.readStateValue(dbPath: dbPath, key: "jetskiStateSync.agentManagerInitState") + guard let data = Data(base64Encoded: base64) else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid base64 in agentManagerInitState") + } + return try self.parseProtoTokenInfo(data: data) + } + + private static func readStateValue(dbPath: URL, key: String) throws -> String { + var db: OpaquePointer? + let openStatus = sqlite3_open_v2(dbPath.path, &db, SQLITE_OPEN_READONLY, nil) + guard openStatus == SQLITE_OK, let db else { + throw AntigravityOAuthCredentialsError.decodeFailed("Failed to open state.vscdb: \(openStatus)") + } + defer { sqlite3_close(db) } + + let query = "SELECT value FROM ItemTable WHERE key = ?" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK, let stmt else { + throw AntigravityOAuthCredentialsError.decodeFailed("Failed to prepare query") + } + defer { sqlite3_finalize(stmt) } + + sqlite3_bind_text(stmt, 1, key, -1, nil) + + guard sqlite3_step(stmt) == SQLITE_ROW else { + throw AntigravityOAuthCredentialsError.notFound + } + + guard let cValue = sqlite3_column_text(stmt, 0) else { + throw AntigravityOAuthCredentialsError.decodeFailed("Empty value for key: \(key)") + } + + let value = String(cString: cValue).trimmingCharacters(in: .whitespacesAndNewlines) + guard !value.isEmpty else { + throw AntigravityOAuthCredentialsError.decodeFailed("Empty value for key: \(key)") + } + + return value + } + + private static func parseProtoTokenInfo(data: Data) throws -> ProtoTokenInfo { + var accessToken: String? + var refreshToken: String? + var tokenType: String? + var expirySeconds: Int? + + var offset = 0 + while offset < data.count { + let (fieldTag, newOffset) = try self.readVarint(data: data, offset: offset) + offset = newOffset + + let fieldNumber = fieldTag >> 3 + let wireType = fieldTag & 0x07 + + switch wireType { + case 0: + let (value, nextOffset) = try self.readVarint(data: data, offset: offset) + offset = nextOffset + if fieldNumber == 4 { + expirySeconds = value + } + case 2: + let (length, lengthOffset) = try self.readVarint(data: data, offset: offset) + offset = lengthOffset + let endOffset = offset + length + guard endOffset <= data.count else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid protobuf length") + } + let stringData = data[offset.. (Int, Int) { + var result = 0 + var shift = 0 + var pos = offset + + while pos < data.count { + let byte = Int(data[pos]) + result |= (byte & 0x7F) << shift + pos += 1 + if (byte & 0x80) == 0 { + return (result, pos) + } + shift += 7 + if shift > 63 { + throw AntigravityOAuthCredentialsError.decodeFailed("Varint too long") + } + } + + throw AntigravityOAuthCredentialsError.decodeFailed("Incomplete varint") + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift new file mode 100644 index 00000000..763961ff --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift @@ -0,0 +1,240 @@ +import Foundation +#if os(macOS) +import Security +#endif + +public struct AntigravityOAuthCredentials: Sendable, Codable { + public let accessToken: String + public let refreshToken: String? + public let expiresAt: Date? + public let email: String? + public let scopes: [String] + + public init( + accessToken: String, + refreshToken: String?, + expiresAt: Date?, + email: String?, + scopes: [String] = AntigravityOAuthConfig.scopes) + { + self.accessToken = accessToken + self.refreshToken = refreshToken + self.expiresAt = expiresAt + self.email = email + self.scopes = scopes + } + + public var isExpired: Bool { + guard let expiresAt else { return false } + return Date() >= expiresAt + } + + public var expiresIn: TimeInterval? { + guard let expiresAt else { return nil } + return expiresAt.timeIntervalSinceNow + } + + public var isRefreshable: Bool { + guard let refreshToken else { return false } + return !refreshToken.isEmpty + } + + public var needsRefresh: Bool { + guard let expiresAt else { return self.isRefreshable } + let bufferTime: TimeInterval = 5 * 60 + return expiresAt.timeIntervalSinceNow < bufferTime + } +} + +public enum AntigravityOAuthConfig { + public static let clientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" + public static let clientSecret = "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" + public static let scopes = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", + "https://www.googleapis.com/auth/experimentsandconfigs", + ] + public static let tokenURL = "https://oauth2.googleapis.com/token" + public static let authURL = "https://accounts.google.com/o/oauth2/auth" + public static let userInfoURL = "https://www.googleapis.com/oauth2/v2/userinfo" + public static let callbackHost = "127.0.0.1" + public static let callbackPortStart = 11451 + public static let callbackPortRange = 100 +} + +public enum AntigravityOAuthCredentialsError: LocalizedError, Sendable { + case notFound + case decodeFailed(String) + case missingAccessToken + case refreshFailed(String) + case invalidGrant + case networkError(String) + case keychainError(Int) + + public var errorDescription: String? { + switch self { + case .notFound: + "Antigravity credentials not found. Authorize or import from Antigravity app." + case let .decodeFailed(message): + "Failed to decode Antigravity credentials: \(message)" + case .missingAccessToken: + "Antigravity access token is missing." + case let .refreshFailed(message): + "Failed to refresh Antigravity token: \(message)" + case .invalidGrant: + "Antigravity refresh token is invalid. Please re-authorize." + case let .networkError(message): + "Antigravity network error: \(message)" + case let .keychainError(status): + "Antigravity keychain error: \(status)" + } + } +} + +public enum AntigravityOAuthCredentialsStore { + private static let cacheKey = KeychainCacheStore.Key.oauth(provider: .antigravity) + private static let log = CodexBarLog.logger(LogCategories.antigravity) + public static let environmentTokenKey = "CODEXBAR_ANTIGRAVITY_TOKEN" + + struct CacheEntry: Codable, Sendable { + let credentials: AntigravityOAuthCredentials + let storedAt: Date + } + + private nonisolated(unsafe) static var cachedCredentials: AntigravityOAuthCredentials? + private nonisolated(unsafe) static var cacheTimestamp: Date? + private static let memoryCacheValidityDuration: TimeInterval = 1800 + + public static func load( + environment: [String: String] = ProcessInfo.processInfo.environment) -> AntigravityOAuthCredentials? + { + if let credentials = self.loadFromEnvironment(environment) { + return credentials + } + + if let cached = self.cachedCredentials, + let timestamp = self.cacheTimestamp, + Date().timeIntervalSince(timestamp) < self.memoryCacheValidityDuration, + !cached.isExpired + { + return cached + } + + switch KeychainCacheStore.load(key: self.cacheKey, as: CacheEntry.self) { + case let .found(entry): + if entry.credentials.isExpired, !entry.credentials.isRefreshable { + self.log.debug("Antigravity cached credentials expired and not refreshable") + return entry.credentials + } + self.cachedCredentials = entry.credentials + self.cacheTimestamp = Date() + return entry.credentials + case .invalid: + KeychainCacheStore.clear(key: self.cacheKey) + return nil + case .missing: + return nil + } + } + + public static func save(_ credentials: AntigravityOAuthCredentials) { + let entry = CacheEntry(credentials: credentials, storedAt: Date()) + KeychainCacheStore.store(key: self.cacheKey, entry: entry) + self.cachedCredentials = credentials + self.cacheTimestamp = Date() + self.log.info("Antigravity credentials saved", metadata: [ + "email": credentials.email ?? "unknown", + "hasRefreshToken": "\(credentials.isRefreshable)", + ]) + } + + public static func clear() { + KeychainCacheStore.clear(key: self.cacheKey) + self.cachedCredentials = nil + self.cacheTimestamp = nil + self.log.info("Antigravity credentials cleared") + } + + public static func invalidateCache() { + self.cachedCredentials = nil + self.cacheTimestamp = nil + } + + private static func loadFromEnvironment(_ environment: [String: String]) -> AntigravityOAuthCredentials? { + guard let token = environment[self.environmentTokenKey]?.trimmingCharacters(in: .whitespacesAndNewlines), + !token.isEmpty + else { + return nil + } + + return AntigravityOAuthCredentials( + accessToken: token, + refreshToken: nil, + expiresAt: Date.distantFuture, + email: nil, + scopes: AntigravityOAuthConfig.scopes) + } +} + +extension AntigravityOAuthCredentialsStore { + public static func parseManualToken(_ input: String) -> AntigravityOAuthCredentials? { + let trimmed = input.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + + if let jsonCredentials = self.parseJSONInput(trimmed) { + return jsonCredentials + } + + if trimmed.hasPrefix("ya29.") || trimmed.hasPrefix("1//") { + let isRefreshToken = trimmed.hasPrefix("1//") + if isRefreshToken { + return AntigravityOAuthCredentials( + accessToken: "", + refreshToken: trimmed, + expiresAt: nil, + email: nil, + scopes: AntigravityOAuthConfig.scopes) + } else { + return AntigravityOAuthCredentials( + accessToken: trimmed, + refreshToken: nil, + expiresAt: nil, + email: nil, + scopes: AntigravityOAuthConfig.scopes) + } + } + + return nil + } + + private static func parseJSONInput(_ input: String) -> AntigravityOAuthCredentials? { + guard let data = input.data(using: .utf8) else { return nil } + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return nil + } + + let apiKey = (json["apiKey"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let accessToken = (json["accessToken"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let refreshToken = (json["refreshToken"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let email = (json["email"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + + let token = apiKey ?? accessToken + guard let token, !token.isEmpty else { return nil } + + var expiresAt: Date? + if let expiresAtString = json["expiresAt"] as? String { + expiresAt = ISO8601DateFormatter().date(from: expiresAtString) + } else if let expiresAtMillis = json["expiresAt"] as? Double { + expiresAt = Date(timeIntervalSince1970: expiresAtMillis / 1000.0) + } + + return AntigravityOAuthCredentials( + accessToken: token, + refreshToken: refreshToken, + expiresAt: expiresAt, + email: email, + scopes: AntigravityOAuthConfig.scopes) + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift new file mode 100644 index 00000000..e1b5ba10 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -0,0 +1,306 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif +#if os(macOS) +import AppKit +#endif + +public actor AntigravityOAuthFlow { + private static let log = CodexBarLog.logger(LogCategories.antigravity) + private static let httpTimeout: TimeInterval = 15.0 + + private var callbackServer: CallbackServer? + private var pendingState: String? + private var pendingContinuation: CheckedContinuation? + + public init() {} + + public func startAuthorization() async throws -> AntigravityOAuthCredentials { + let port = try await self.startCallbackServer() + let redirectUri = "http://\(AntigravityOAuthConfig.callbackHost):\(port)" + let state = self.generateState() + self.pendingState = state + + let authURL = self.buildAuthURL(redirectUri: redirectUri, state: state) + + Self.log.info("Opening Antigravity OAuth authorization URL") + #if os(macOS) + if let url = URL(string: authURL) { + await MainActor.run { + NSWorkspace.shared.open(url) + } + } + #endif + + do { + let code = try await withCheckedThrowingContinuation { continuation in + self.pendingContinuation = continuation + } + + let credentials = try await self.exchangeCodeForToken(code: code, redirectUri: redirectUri) + return credentials + } catch { + throw error + } + } + + public func cancelAuthorization() { + self.pendingContinuation?.resume(throwing: CancellationError()) + self.pendingContinuation = nil + self.pendingState = nil + self.stopCallbackServer() + } + + private func startCallbackServer() async throws -> Int { + var port = AntigravityOAuthConfig.callbackPortStart + var attempts = 0 + + while attempts < AntigravityOAuthConfig.callbackPortRange { + do { + let server = try CallbackServer(port: port) { code, state in + Task { [weak self] in + await self?.handleCallback(code: code, state: state) + } + } + self.callbackServer = server + Self.log.info("Antigravity OAuth callback server started on port \(port)") + return port + } catch { + port += 1 + attempts += 1 + } + } + + throw AntigravityOAuthCredentialsError.networkError("No available port for OAuth callback") + } + + private func stopCallbackServer() { + self.callbackServer?.stop() + self.callbackServer = nil + } + + private func handleCallback(code: String?, state: String?) { + defer { + self.stopCallbackServer() + } + + guard let code, let state else { + self.pendingContinuation?.resume(throwing: AntigravityOAuthCredentialsError.decodeFailed("Missing code or state in callback")) + self.pendingContinuation = nil + return + } + + guard state == self.pendingState else { + self.pendingContinuation?.resume(throwing: AntigravityOAuthCredentialsError.decodeFailed("State mismatch")) + self.pendingContinuation = nil + return + } + + self.pendingContinuation?.resume(returning: code) + self.pendingContinuation = nil + self.pendingState = nil + } + + private func generateState() -> String { + let chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + return String((0..<32).map { _ in chars.randomElement()! }) + } + + private func buildAuthURL(redirectUri: String, state: String) -> String { + var components = URLComponents(string: AntigravityOAuthConfig.authURL)! + components.queryItems = [ + URLQueryItem(name: "client_id", value: AntigravityOAuthConfig.clientID), + URLQueryItem(name: "redirect_uri", value: redirectUri), + URLQueryItem(name: "response_type", value: "code"), + URLQueryItem(name: "scope", value: AntigravityOAuthConfig.scopes.joined(separator: " ")), + URLQueryItem(name: "state", value: state), + URLQueryItem(name: "access_type", value: "offline"), + URLQueryItem(name: "prompt", value: "consent"), + URLQueryItem(name: "include_granted_scopes", value: "true"), + ] + return components.url!.absoluteString + } + + private func exchangeCodeForToken(code: String, redirectUri: String) async throws -> AntigravityOAuthCredentials { + let params = [ + "client_id": AntigravityOAuthConfig.clientID, + "client_secret": AntigravityOAuthConfig.clientSecret, + "code": code, + "redirect_uri": redirectUri, + "grant_type": "authorization_code", + ] + + let body = params + .map { "\($0.key)=\($0.value.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) ?? $0.value)" } + .joined(separator: "&") + + guard let url = URL(string: AntigravityOAuthConfig.tokenURL) else { + throw AntigravityOAuthCredentialsError.networkError("Invalid token URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = body.data(using: .utf8) + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.timeoutInterval = Self.httpTimeout + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { + let statusCode = (response as? HTTPURLResponse)?.statusCode ?? 0 + let errorBody = String(data: data, encoding: .utf8) ?? "" + throw AntigravityOAuthCredentialsError.refreshFailed("Token exchange failed: HTTP \(statusCode) - \(errorBody)") + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let accessToken = json["access_token"] as? String, + let refreshToken = json["refresh_token"] as? String, + let expiresIn = json["expires_in"] as? Int + else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid token exchange response") + } + + let expiresAt = Date(timeIntervalSinceNow: TimeInterval(expiresIn)) + let email = try? await AntigravityTokenRefresher.fetchUserEmail(accessToken: accessToken) + + Self.log.info("Antigravity OAuth authorization successful", metadata: [ + "email": email ?? "unknown", + ]) + + return AntigravityOAuthCredentials( + accessToken: accessToken, + refreshToken: refreshToken, + expiresAt: expiresAt, + email: email, + scopes: AntigravityOAuthConfig.scopes) + } +} + +private final class CallbackServer: @unchecked Sendable { + private var serverSocket: Int32 = -1 + private var isRunning = false + private let callback: @Sendable (String?, String?) -> Void + private let queue = DispatchQueue(label: "com.codexbar.antigravity.oauth.callback") + + init(port: Int, callback: @escaping @Sendable (String?, String?) -> Void) throws { + self.callback = callback + try self.start(port: port) + } + + private func start(port: Int) throws { + self.serverSocket = socket(AF_INET, SOCK_STREAM, 0) + guard self.serverSocket >= 0 else { + throw AntigravityOAuthCredentialsError.networkError("Failed to create socket") + } + + var opt: Int32 = 1 + setsockopt(self.serverSocket, SOL_SOCKET, SO_REUSEADDR, &opt, socklen_t(MemoryLayout.size)) + + var addr = sockaddr_in() + addr.sin_family = sa_family_t(AF_INET) + addr.sin_port = UInt16(port).bigEndian + addr.sin_addr.s_addr = inet_addr("127.0.0.1") + + let bindResult = withUnsafePointer(to: &addr) { ptr in + ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { sockPtr in + bind(self.serverSocket, sockPtr, socklen_t(MemoryLayout.size)) + } + } + + guard bindResult >= 0 else { + close(self.serverSocket) + throw AntigravityOAuthCredentialsError.networkError("Failed to bind to port \(port)") + } + + guard listen(self.serverSocket, 1) >= 0 else { + close(self.serverSocket) + throw AntigravityOAuthCredentialsError.networkError("Failed to listen on port \(port)") + } + + self.isRunning = true + self.queue.async { [weak self] in + self?.acceptConnections() + } + } + + func stop() { + self.isRunning = false + if self.serverSocket >= 0 { + close(self.serverSocket) + self.serverSocket = -1 + } + } + + private func acceptConnections() { + while self.isRunning { + var clientAddr = sockaddr_in() + var addrLen = socklen_t(MemoryLayout.size) + + let clientSocket = withUnsafeMutablePointer(to: &clientAddr) { ptr in + ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { sockPtr in + accept(self.serverSocket, sockPtr, &addrLen) + } + } + + guard clientSocket >= 0 else { continue } + + var buffer = [UInt8](repeating: 0, count: 4096) + let bytesRead = recv(clientSocket, &buffer, buffer.count, 0) + + if bytesRead > 0 { + let request = String(bytes: buffer[0.. + Authorization Successful + +

Authorization Successful!

+

You can close this window and return to CodexBar.

+ + + + """ + } else { + responseHTML = """ + + Authorization Failed + +

Authorization Failed

+

Please close this window and try again.

+ + + """ + } + + let response = "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nConnection: close\r\n\r\n\(responseHTML)" + _ = response.withCString { ptr in + send(clientSocket, ptr, strlen(ptr), 0) + } + + close(clientSocket) + self.callback(code, state) + break + } + + close(clientSocket) + } + } + + private func parseOAuthCallback(request: String) -> (code: String?, state: String?) { + guard let firstLine = request.split(separator: "\r\n").first else { return (nil, nil) } + let parts = firstLine.split(separator: " ") + guard parts.count >= 2 else { return (nil, nil) } + + let path = String(parts[1]) + guard let components = URLComponents(string: "http://localhost\(path)") else { return (nil, nil) } + + let code = components.queryItems?.first(where: { $0.name == "code" })?.value + let state = components.queryItems?.first(where: { $0.name == "state" })?.value + + return (code, state) + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityTokenRefresher.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityTokenRefresher.swift new file mode 100644 index 00000000..843f337e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityTokenRefresher.swift @@ -0,0 +1,114 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum AntigravityTokenRefresher { + private static let log = CodexBarLog.logger(LogCategories.antigravity) + private static let httpTimeout: TimeInterval = 15.0 + + public struct RefreshResult: Sendable { + public let accessToken: String + public let expiresAt: Date + public let email: String? + } + + public static func refreshAccessToken(refreshToken: String) async throws -> RefreshResult { + let params = [ + "client_id": AntigravityOAuthConfig.clientID, + "client_secret": AntigravityOAuthConfig.clientSecret, + "refresh_token": refreshToken, + "grant_type": "refresh_token", + ] + + let body = params + .map { "\($0.key)=\($0.value.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) ?? $0.value)" } + .joined(separator: "&") + + guard let url = URL(string: AntigravityOAuthConfig.tokenURL) else { + throw AntigravityOAuthCredentialsError.networkError("Invalid token URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = body.data(using: .utf8) + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.timeoutInterval = self.httpTimeout + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let http = response as? HTTPURLResponse else { + throw AntigravityOAuthCredentialsError.networkError("Invalid response") + } + + guard http.statusCode == 200 else { + let errorBody = String(data: data, encoding: .utf8) ?? "" + if errorBody.lowercased().contains("invalid_grant") { + self.log.warning("Antigravity refresh token is invalid (invalid_grant)") + throw AntigravityOAuthCredentialsError.invalidGrant + } + throw AntigravityOAuthCredentialsError.refreshFailed("HTTP \(http.statusCode): \(errorBody)") + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let accessToken = json["access_token"] as? String, + let expiresIn = json["expires_in"] as? Int + else { + throw AntigravityOAuthCredentialsError.decodeFailed("Invalid token response") + } + + let expiresAt = Date(timeIntervalSinceNow: TimeInterval(expiresIn)) + + var email: String? + do { + email = try await self.fetchUserEmail(accessToken: accessToken) + } catch { + self.log.debug("Failed to fetch user email during refresh: \(error.localizedDescription)") + } + + self.log.info("Antigravity access token refreshed", metadata: [ + "expiresIn": "\(expiresIn)s", + "email": email ?? "unknown", + ]) + + return RefreshResult(accessToken: accessToken, expiresAt: expiresAt, email: email) + } + + public static func buildCredentialsFromRefreshToken( + refreshToken: String, + fallbackEmail: String? = nil) async throws -> AntigravityOAuthCredentials + { + let result = try await self.refreshAccessToken(refreshToken: refreshToken) + return AntigravityOAuthCredentials( + accessToken: result.accessToken, + refreshToken: refreshToken, + expiresAt: result.expiresAt, + email: result.email ?? fallbackEmail, + scopes: AntigravityOAuthConfig.scopes) + } + + public static func fetchUserEmail(accessToken: String) async throws -> String { + guard let url = URL(string: AntigravityOAuthConfig.userInfoURL) else { + throw AntigravityOAuthCredentialsError.networkError("Invalid userinfo URL") + } + + var request = URLRequest(url: url) + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.timeoutInterval = self.httpTimeout + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { + let statusCode = (response as? HTTPURLResponse)?.statusCode ?? 0 + throw AntigravityOAuthCredentialsError.networkError("Failed to fetch user info: HTTP \(statusCode)") + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let email = json["email"] as? String + else { + throw AntigravityOAuthCredentialsError.decodeFailed("Missing email in userinfo response") + } + + return email + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift new file mode 100644 index 00000000..274a9d16 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift @@ -0,0 +1,29 @@ +import Foundation + +public enum AntigravityUsageSource: String, CaseIterable, Sendable, Codable { + case auto + case authorized + case local + + public var displayName: String { + switch self { + case .auto: + "Auto" + case .authorized: + "Authorized (OAuth)" + case .local: + "Local Server" + } + } + + public var description: String { + switch self { + case .auto: + "Use authorized credentials if available, fallback to local server" + case .authorized: + "Use OAuth credentials to fetch quota from Cloud Code API" + case .local: + "Use local Antigravity language server" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift index 027be879..f8743497 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift @@ -33,19 +33,34 @@ public enum AntigravityProviderDescriptor { supportsTokenCost: false, noDataMessage: { "Antigravity cost summary is not supported." }), fetchPlan: ProviderFetchPlan( - sourceModes: [.auto, .cli], - pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [AntigravityStatusFetchStrategy()] })), + sourceModes: [.auto, .oauth, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: Self.resolveStrategies)), cli: ProviderCLIConfig( name: "antigravity", versionDetector: nil)) } + + private static func resolveStrategies(_ context: ProviderFetchContext) -> [any ProviderFetchStrategy] { + let usageSource = context.settings?.antigravity?.usageSource ?? .auto + + switch usageSource { + case .auto: + return [AntigravityAuthorizedFetchStrategy(), AntigravityLocalFetchStrategy()] + case .authorized: + return [AntigravityAuthorizedFetchStrategy()] + case .local: + return [AntigravityLocalFetchStrategy()] + } + } } -struct AntigravityStatusFetchStrategy: ProviderFetchStrategy { +struct AntigravityLocalFetchStrategy: ProviderFetchStrategy { let id: String = "antigravity.local" let kind: ProviderFetchKind = .localProbe - func isAvailable(_: ProviderFetchContext) async -> Bool { true } + func isAvailable(_: ProviderFetchContext) async -> Bool { + await AntigravityStatusProbe.isRunning() + } func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { let probe = AntigravityStatusProbe() @@ -57,6 +72,6 @@ struct AntigravityStatusFetchStrategy: ProviderFetchStrategy { } func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { - false + true } } diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index f83cb9fd..e43f96f3 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -15,7 +15,8 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: KimiProviderSettings? = nil, augment: AugmentProviderSettings? = nil, amp: AmpProviderSettings? = nil, - jetbrains: JetBrainsProviderSettings? = nil) -> ProviderSettingsSnapshot + jetbrains: JetBrainsProviderSettings? = nil, + antigravity: AntigravityProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot( debugMenuEnabled: debugMenuEnabled, @@ -31,7 +32,8 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: kimi, augment: augment, amp: amp, - jetbrains: jetbrains) + jetbrains: jetbrains, + antigravity: antigravity) } public struct CodexProviderSettings: Sendable { @@ -167,6 +169,16 @@ public struct ProviderSettingsSnapshot: Sendable { } } + public struct AntigravityProviderSettings: Sendable { + public let usageSource: AntigravityUsageSource + public let manualToken: String? + + public init(usageSource: AntigravityUsageSource, manualToken: String?) { + self.usageSource = usageSource + self.manualToken = manualToken + } + } + public let debugMenuEnabled: Bool public let debugKeepCLISessionsAlive: Bool public let codex: CodexProviderSettings? @@ -181,11 +193,16 @@ public struct ProviderSettingsSnapshot: Sendable { public let augment: AugmentProviderSettings? public let amp: AmpProviderSettings? public let jetbrains: JetBrainsProviderSettings? + public let antigravity: AntigravityProviderSettings? public var jetbrainsIDEBasePath: String? { self.jetbrains?.ideBasePath } + public var antigravityManualToken: String? { + self.antigravity?.manualToken + } + public init( debugMenuEnabled: Bool, debugKeepCLISessionsAlive: Bool, @@ -200,7 +217,8 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: KimiProviderSettings?, augment: AugmentProviderSettings?, amp: AmpProviderSettings?, - jetbrains: JetBrainsProviderSettings? = nil) + jetbrains: JetBrainsProviderSettings? = nil, + antigravity: AntigravityProviderSettings? = nil) { self.debugMenuEnabled = debugMenuEnabled self.debugKeepCLISessionsAlive = debugKeepCLISessionsAlive @@ -216,6 +234,7 @@ public struct ProviderSettingsSnapshot: Sendable { self.augment = augment self.amp = amp self.jetbrains = jetbrains + self.antigravity = antigravity } } @@ -232,6 +251,7 @@ public enum ProviderSettingsSnapshotContribution: Sendable { case augment(ProviderSettingsSnapshot.AugmentProviderSettings) case amp(ProviderSettingsSnapshot.AmpProviderSettings) case jetbrains(ProviderSettingsSnapshot.JetBrainsProviderSettings) + case antigravity(ProviderSettingsSnapshot.AntigravityProviderSettings) } public struct ProviderSettingsSnapshotBuilder: Sendable { @@ -249,6 +269,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { public var augment: ProviderSettingsSnapshot.AugmentProviderSettings? public var amp: ProviderSettingsSnapshot.AmpProviderSettings? public var jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? + public var antigravity: ProviderSettingsSnapshot.AntigravityProviderSettings? public init(debugMenuEnabled: Bool = false, debugKeepCLISessionsAlive: Bool = false) { self.debugMenuEnabled = debugMenuEnabled @@ -269,6 +290,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { case let .augment(value): self.augment = value case let .amp(value): self.amp = value case let .jetbrains(value): self.jetbrains = value + case let .antigravity(value): self.antigravity = value } } @@ -287,6 +309,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { kimi: self.kimi, augment: self.augment, amp: self.amp, - jetbrains: self.jetbrains) + jetbrains: self.jetbrains, + antigravity: self.antigravity) } } diff --git a/Tests/CodexBarTests/AntigravityOAuthTests.swift b/Tests/CodexBarTests/AntigravityOAuthTests.swift new file mode 100644 index 00000000..6388e17b --- /dev/null +++ b/Tests/CodexBarTests/AntigravityOAuthTests.swift @@ -0,0 +1,203 @@ +import Foundation +import Testing +@testable import CodexBarCore + +@Suite("AntigravityOAuthCredentials") +struct AntigravityOAuthCredentialsTests { + @Test("Credentials are not expired when expiresAt is in the future") + func test_credentialsNotExpired() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: Date().addingTimeInterval(3600), + email: "test@example.com") + #expect(!creds.isExpired) + } + + @Test("Credentials are expired when expiresAt is in the past") + func test_credentialsExpired() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: Date().addingTimeInterval(-3600), + email: "test@example.com") + #expect(creds.isExpired) + } + + @Test("Credentials with nil expiresAt are not expired") + func test_credentialsNilExpiresAtNotExpired() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: nil, + email: "test@example.com") + #expect(!creds.isExpired) + } + + @Test("Credentials are refreshable when refresh token is present") + func test_credentialsIsRefreshable() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: nil, + email: nil) + #expect(creds.isRefreshable) + } + + @Test("Credentials are not refreshable when refresh token is nil") + func test_credentialsNotRefreshableNil() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: nil, + expiresAt: nil, + email: nil) + #expect(!creds.isRefreshable) + } + + @Test("Credentials are not refreshable when refresh token is empty") + func test_credentialsNotRefreshableEmpty() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "", + expiresAt: nil, + email: nil) + #expect(!creds.isRefreshable) + } + + @Test("Credentials need refresh when close to expiry") + func test_credentialsNeedRefresh() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: Date().addingTimeInterval(120), + email: nil) + #expect(creds.needsRefresh) + } + + @Test("Credentials don't need refresh when not close to expiry") + func test_credentialsDontNeedRefresh() { + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: Date().addingTimeInterval(3600), + email: nil) + #expect(!creds.needsRefresh) + } +} + +@Suite("AntigravityManualTokenParsing") +struct AntigravityManualTokenParsingTests { + @Test("Parses access token starting with ya29.") + func test_parseAccessToken() { + let token = "ya29.a0ARrdaM..." + let creds = AntigravityOAuthCredentialsStore.parseManualToken(token) + #expect(creds != nil) + #expect(creds?.accessToken == token) + #expect(creds?.refreshToken == nil) + } + + @Test("Parses refresh token starting with 1//") + func test_parseRefreshToken() { + let token = "1//0gXyz..." + let creds = AntigravityOAuthCredentialsStore.parseManualToken(token) + #expect(creds != nil) + #expect(creds?.accessToken.isEmpty == true) + #expect(creds?.refreshToken == token) + } + + @Test("Parses JSON with apiKey") + func test_parseJSONWithApiKey() { + let json = """ + {"apiKey": "ya29.test", "email": "test@example.com"} + """ + let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) + #expect(creds != nil) + #expect(creds?.accessToken == "ya29.test") + #expect(creds?.email == "test@example.com") + } + + @Test("Parses JSON with accessToken") + func test_parseJSONWithAccessToken() { + let json = """ + {"accessToken": "ya29.test", "refreshToken": "1//refresh"} + """ + let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) + #expect(creds != nil) + #expect(creds?.accessToken == "ya29.test") + #expect(creds?.refreshToken == "1//refresh") + } + + @Test("Parses JSON with expiresAt as string") + func test_parseJSONWithExpiresAtString() { + let json = """ + {"apiKey": "ya29.test", "expiresAt": "2025-01-01T00:00:00Z"} + """ + let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) + #expect(creds != nil) + #expect(creds?.expiresAt != nil) + } + + @Test("Parses JSON with expiresAt as milliseconds") + func test_parseJSONWithExpiresAtMillis() { + let json = """ + {"apiKey": "ya29.test", "expiresAt": 1735689600000} + """ + let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) + #expect(creds != nil) + #expect(creds?.expiresAt != nil) + } + + @Test("Returns nil for empty input") + func test_parseEmptyInput() { + let creds = AntigravityOAuthCredentialsStore.parseManualToken("") + #expect(creds == nil) + } + + @Test("Returns nil for whitespace-only input") + func test_parseWhitespaceOnlyInput() { + let creds = AntigravityOAuthCredentialsStore.parseManualToken(" \n\t ") + #expect(creds == nil) + } + + @Test("Returns nil for invalid token format") + func test_parseInvalidTokenFormat() { + let creds = AntigravityOAuthCredentialsStore.parseManualToken("not-a-valid-token") + #expect(creds == nil) + } + + @Test("Returns nil for JSON without apiKey or accessToken") + func test_parseJSONWithoutToken() { + let json = """ + {"email": "test@example.com"} + """ + let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) + #expect(creds == nil) + } +} + +@Suite("AntigravityUsageSource") +struct AntigravityUsageSourceTests { + @Test("UsageSource has expected cases") + func test_usageSourceCases() { + let allCases = AntigravityUsageSource.allCases + #expect(allCases.contains(.auto)) + #expect(allCases.contains(.authorized)) + #expect(allCases.contains(.local)) + #expect(allCases.count == 3) + } + + @Test("UsageSource rawValue matches expected strings") + func test_usageSourceRawValues() { + #expect(AntigravityUsageSource.auto.rawValue == "auto") + #expect(AntigravityUsageSource.authorized.rawValue == "authorized") + #expect(AntigravityUsageSource.local.rawValue == "local") + } + + @Test("UsageSource can be initialized from rawValue") + func test_usageSourceFromRawValue() { + #expect(AntigravityUsageSource(rawValue: "auto") == .auto) + #expect(AntigravityUsageSource(rawValue: "authorized") == .authorized) + #expect(AntigravityUsageSource(rawValue: "local") == .local) + #expect(AntigravityUsageSource(rawValue: "invalid") == nil) + } +} diff --git a/docs/antigravity.md b/docs/antigravity.md index ab99af30..02baf1e9 100644 --- a/docs/antigravity.md +++ b/docs/antigravity.md @@ -1,16 +1,55 @@ --- -summary: "Antigravity provider notes: local LSP probing, port discovery, quota parsing, and UI mapping." +summary: "Antigravity provider notes: OAuth credentials, local LSP probing, port discovery, quota parsing, and UI mapping." read_when: - Adding or modifying the Antigravity provider - Debugging Antigravity port detection or quota parsing - Adjusting Antigravity menu labels or model mapping + - Working with Antigravity OAuth credentials --- # Antigravity provider -Antigravity is a local-only provider. We talk directly to the Antigravity language server running on the same machine. +Antigravity supports both OAuth-authorized API access and local language server probing. -## Data sources + fallback order +## Usage source modes + +- **Auto** (default): Try authorized credentials first, fallback to local server +- **Authorized**: Use OAuth credentials to fetch quota from Cloud Code API +- **Local**: Use local Antigravity language server only + +## Credential acquisition + +### Fallback order (Auto mode) + +1. **Keychain OAuth credentials** - Previously saved credentials from OAuth flow or import +2. **Manual token** - Token pasted in settings (refresh token or access token) +3. **Local DB import** - Import from `state.vscdb` (Antigravity app's local storage) +4. **Local server** - Direct probe to running Antigravity language server + +### OAuth browser flow + +Opens Google OAuth in browser with local callback server (port 11451+). Returns refresh token for long-term storage. + +### Local DB import (`state.vscdb`) + +Path: `~/Library/Application Support/Antigravity/User/globalStorage/state.vscdb` + +- **Refresh token**: `jetskiStateSync.agentManagerInitState` (protobuf field 6) +- **Access token + email**: `antigravityAuthStatus` JSON (`apiKey`, `email` fields) + +### Manual token import + +Accepts: +- Raw access token (`ya29.xxx`) +- Raw refresh token (`1//xxx`) +- JSON payload: `{"apiKey": "...", "email": "...", "refreshToken": "..."}` + +### Cloud Code API endpoints + +- Primary: `https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` +- Fallback: `https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` + +## Local server data sources + fallback order 1) **Process detection** - Command: `ps -ax -o pid=,command=`. @@ -71,5 +110,24 @@ Antigravity is a local-only provider. We talk directly to the Antigravity langua - Local HTTPS uses a self-signed cert; the probe allows insecure TLS. ## Key files + +### OAuth/Credentials +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityTokenRefresher.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift` + +### Local Server - `Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift` + +### App Integration - `Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift` +- `Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift` + +### Tests +- `Tests/CodexBarTests/AntigravityOAuthTests.swift` +- `Tests/CodexBarTests/AntigravityStatusProbeTests.swift` From fdc1aeff9a457a0e1b56daadabbfb99bb35d97d8 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 17:31:44 +0300 Subject: [PATCH 02/19] feat: enhance Antigravity provider with improved token management and UI updates - Added support for two-field token entry in the Antigravity settings. - Updated the UI for selecting log providers to include Antigravity. - Enhanced the settings snapshot to include new token management features. - Implemented credential refresh handling and improved error messaging for Antigravity OAuth. - Refactored related components for better maintainability and clarity. --- Sources/CodexBar/PreferencesDebugPane.swift | 21 +- .../PreferencesProviderSettingsRows.swift | 87 ++++- .../PreferencesProvidersPane+Testing.swift | 7 +- .../CodexBar/PreferencesProvidersPane.swift | 124 +++++++- .../Antigravity/AntigravityLoginFlow.swift | 196 +++++++++++- .../AntigravityProviderImplementation.swift | 61 +++- .../AntigravitySettingsStore.swift | 134 ++++++-- .../Shared/ProviderSettingsDescriptors.swift | 12 +- .../CodexBar/UsageStore+TokenAccounts.swift | 51 ++- Sources/CodexBar/UsageStore.swift | 59 +++- Sources/CodexBarCLI/CLIHelpers.swift | 3 + Sources/CodexBarCLI/TokenAccountCLI.swift | 16 +- .../CodexBarCore/Config/CodexBarConfig.swift | 5 +- .../AntigravityAuthorizedFetchStrategy.swift | 134 ++++---- .../AntigravityLocalImporter.swift | 9 +- .../AntigravityOAuthCredentials.swift | 198 ++++++------ .../AntigravityOAuthFlow.swift | 15 +- .../AntigravityUsageSource.swift | 21 +- .../AntigravityProviderDescriptor.swift | 4 +- .../Providers/ProviderFetchPlan.swift | 5 +- .../Providers/ProviderSettingsSnapshot.swift | 16 +- .../TokenAccountSupportCatalog+Data.swift | 7 + .../CodexBarTests/AntigravityOAuthTests.swift | 298 ++++++++---------- docs/antigravity.md | 48 +-- 24 files changed, 1074 insertions(+), 457 deletions(-) diff --git a/Sources/CodexBar/PreferencesDebugPane.swift b/Sources/CodexBar/PreferencesDebugPane.swift index 1b808340..79d8b3b1 100644 --- a/Sources/CodexBar/PreferencesDebugPane.swift +++ b/Sources/CodexBar/PreferencesDebugPane.swift @@ -100,15 +100,20 @@ struct DebugPane: View { title: "Probe logs", caption: "Fetch the latest probe output for debugging; Copy keeps the full text.") { - Picker("Provider", selection: self.$currentLogProvider) { - Text("Codex").tag(UsageProvider.codex) - Text("Claude").tag(UsageProvider.claude) - Text("Cursor").tag(UsageProvider.cursor) - Text("Augment").tag(UsageProvider.augment) - Text("Amp").tag(UsageProvider.amp) + ScrollView(.horizontal, showsIndicators: false) { + Picker("Provider", selection: self.$currentLogProvider) { + Text("Codex").tag(UsageProvider.codex) + Text("Claude").tag(UsageProvider.claude) + Text("Cursor").tag(UsageProvider.cursor) + Text("Augment").tag(UsageProvider.augment) + Text("Amp").tag(UsageProvider.amp) + Text("Antigravity").tag(UsageProvider.antigravity) + } + .pickerStyle(.segmented) + .controlSize(.small) + .fixedSize() } - .pickerStyle(.segmented) - .frame(width: 460) + .frame(maxWidth: .infinity, alignment: .leading) HStack(spacing: 12) { Button { self.loadLog(self.currentLogProvider) } label: { diff --git a/Sources/CodexBar/PreferencesProviderSettingsRows.swift b/Sources/CodexBar/PreferencesProviderSettingsRows.swift index 5d7abde9..4cba7137 100644 --- a/Sources/CodexBar/PreferencesProviderSettingsRows.swift +++ b/Sources/CodexBar/PreferencesProviderSettingsRows.swift @@ -203,6 +203,7 @@ struct ProviderSettingsTokenAccountsRowView: View { let descriptor: ProviderSettingsTokenAccountsDescriptor @State private var newLabel: String = "" @State private var newToken: String = "" + @State private var newToken2: String = "" var body: some View { VStack(alignment: .leading, spacing: 8) { @@ -243,25 +244,79 @@ struct ProviderSettingsTokenAccountsRowView: View { .controlSize(.small) } - HStack(spacing: 8) { - TextField("Label", text: self.$newLabel) - .textFieldStyle(.roundedBorder) - .font(.footnote) - SecureField(self.descriptor.placeholder, text: self.$newToken) - .textFieldStyle(.roundedBorder) - .font(.footnote) - Button("Add") { - let label = self.newLabel.trimmingCharacters(in: .whitespacesAndNewlines) - let token = self.newToken.trimmingCharacters(in: .whitespacesAndNewlines) - guard !label.isEmpty, !token.isEmpty else { return } - self.descriptor.addAccount(label, token) - self.newLabel = "" - self.newToken = "" + if self.descriptor.supportsManualEntry { + let trimmedLabel = self.newLabel.trimmingCharacters(in: .whitespacesAndNewlines) + let trimmedToken = self.newToken.trimmingCharacters(in: .whitespacesAndNewlines) + let trimmedToken2 = self.newToken2.trimmingCharacters(in: .whitespacesAndNewlines) + let addDisabled = trimmedLabel.isEmpty || trimmedToken.isEmpty + + if self.descriptor.supportsTwoFieldEntry { + VStack(alignment: .leading, spacing: 8) { + HStack(spacing: 8) { + TextField("Label", text: self.$newLabel) + .textFieldStyle(.roundedBorder) + .font(.footnote) + .frame(maxWidth: .infinity) + } + HStack(spacing: 8) { + SecureField("Access Token (ya29...)", text: self.$newToken) + .textFieldStyle(.roundedBorder) + .font(.footnote) + SecureField("Refresh Token - optional (1//...)", text: self.$newToken2) + .textFieldStyle(.roundedBorder) + .font(.footnote) + } + Button("Add") { + guard !addDisabled else { return } + self.descriptor.addAccount(trimmedLabel, trimmedToken, trimmedToken2) + self.newLabel = "" + self.newToken = "" + self.newToken2 = "" + } + .buttonStyle(.bordered) + .controlSize(.small) + .disabled(addDisabled) + } + } else { + HStack(spacing: 8) { + TextField("Label", text: self.$newLabel) + .textFieldStyle(.roundedBorder) + .font(.footnote) + SecureField(self.descriptor.placeholder, text: self.$newToken) + .textFieldStyle(.roundedBorder) + .font(.footnote) + Button("Add") { + guard !addDisabled else { return } + self.descriptor.addAccount(trimmedLabel, trimmedToken, "") + self.newLabel = "" + self.newToken = "" + self.newToken2 = "" + } + .buttonStyle(.bordered) + .controlSize(.small) + .disabled(addDisabled) + } + } + } + + if let addAction = self.descriptor.addAction { + Button(self.descriptor.addActionTitle ?? "Add account") { + Task { @MainActor in + await addAction() + } + } + .buttonStyle(.bordered) + .controlSize(.small) + } + + if let importAction = self.descriptor.importAction { + Button(importAction.title) { + Task { @MainActor in + await importAction.action() + } } .buttonStyle(.bordered) .controlSize(.small) - .disabled(self.newLabel.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty || - self.newToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty) } HStack(spacing: 10) { diff --git a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift index e2dce0a7..77c67e18 100644 --- a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift +++ b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift @@ -166,12 +166,17 @@ enum ProvidersPaneTestHarness { title: "Accounts", subtitle: "Accounts subtitle", placeholder: "Token", + supportsManualEntry: true, + supportsTwoFieldEntry: false, + addActionTitle: nil, + addAction: nil, + importAction: nil, provider: .codex, isVisible: { true }, accounts: { [] }, activeIndex: { 0 }, setActiveIndex: { _ in }, - addAccount: { _, _ in }, + addAccount: { _, _, _ in }, removeAccount: { _ in }, openConfigFile: {}, reloadFromDisk: {}) diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index d4b29d86..de858710 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -164,11 +164,36 @@ struct ProvidersPane: View { func tokenAccountDescriptor(for provider: UsageProvider) -> ProviderSettingsTokenAccountsDescriptor? { guard let support = TokenAccountSupportCatalog.support(for: provider) else { return nil } let context = self.makeSettingsContext(provider: provider) + + let isAntigravity = provider == .antigravity + let keychainEnabled = !KeychainAccessGate.isDisabled + + let supportsTwoFieldEntry = isAntigravity + let supportsManualEntry = !isAntigravity || !keychainEnabled + let addActionTitle = (isAntigravity && keychainEnabled) ? "Sign in with Google" : nil + let addAction: (() async -> Void)? = (isAntigravity && keychainEnabled) + ? { + _ = await AntigravityLoginFlow.runOAuthFlow( + settings: self.settings, + store: self.store) + } + : nil + let importAction: ProviderSettingsTokenAccountsDescriptor.ImportAction? = (isAntigravity && !keychainEnabled) + ? ProviderSettingsTokenAccountsDescriptor.ImportAction( + title: "Import from Antigravity app", + action: { await self.importAntigravityCredentials() }) + : nil + return ProviderSettingsTokenAccountsDescriptor( id: "token-accounts-\(provider.rawValue)", title: support.title, subtitle: support.subtitle, placeholder: support.placeholder, + supportsManualEntry: supportsManualEntry, + supportsTwoFieldEntry: supportsTwoFieldEntry, + addActionTitle: addActionTitle, + addAction: addAction, + importAction: importAction, provider: provider, isVisible: { ProviderCatalog.implementation(for: provider)? @@ -187,14 +212,25 @@ struct ProvidersPane: View { await self.store.refreshProvider(provider, allowDisabled: true) } }, - addAccount: { label, token in - self.settings.addTokenAccount(provider: provider, label: label, token: token) + addAccount: { label, token, token2 in + if isAntigravity { + _ = self.settings.addManualAntigravityTokenAccount( + label: label, + accessToken: token, + refreshToken: token2.isEmpty ? nil : token2) + } else { + self.settings.addTokenAccount(provider: provider, label: label, token: token) + } Task { @MainActor in await self.store.refreshProvider(provider, allowDisabled: true) } }, removeAccount: { accountID in - self.settings.removeTokenAccount(provider: provider, accountID: accountID) + if isAntigravity { + self.settings.removeAntigravityTokenAccount(accountID: accountID) + } else { + self.settings.removeTokenAccount(provider: provider, accountID: accountID) + } Task { @MainActor in await self.store.refreshProvider(provider, allowDisabled: true) } @@ -210,6 +246,88 @@ struct ProvidersPane: View { }) } + @MainActor + private func importAntigravityCredentials() async { + do { + let credentials = try await AntigravityLocalImporter.importCredentials() + guard let accessToken = credentials.accessToken, !accessToken.isEmpty else { + self.presentAlert( + title: "Import Failed", + message: "No access token found in Antigravity database.") + return + } + + let label = credentials.email ?? credentials.name ?? "Imported Account" + guard let account = self.settings.addManualAntigravityTokenAccount( + label: label, + accessToken: accessToken, + refreshToken: credentials.refreshToken) else { + self.presentAlert( + title: "Import Failed", + message: "Unable to save imported credentials.") + return + } + + await self.store.refreshProvider(.antigravity, allowDisabled: true) + + let alert = NSAlert() + alert.messageText = "Import Successful" + alert.informativeText = "Imported account: \(account.label)" + alert.runModal() + + } catch let error as AntigravityOAuthCredentialsError { + switch error { + case .notFound: + if AntigravityLocalImporter.isAvailable() { + self.presentAlert( + title: "No Credentials Found", + message: "Antigravity database found, but no credentials were found inside.\n\nPlease ensure:\n1. You are signed in to Antigravity IDE (check the Account menu)\n2. Try restarting Antigravity IDE if you just signed in\n3. If the issue persists, paste your tokens manually below") + } else { + self.presentAlert( + title: "Import Failed", + message: "Antigravity database not found. Ensure Antigravity app is installed.") + } + default: + self.presentAlert( + title: "Import Failed", + message: error.localizedDescription) + } + } catch { + let nsError = error as NSError + if nsError.domain == NSPOSIXErrorDomain && nsError.code == 1 { + self.presentFullDiskAccessAlert() + } else { + self.presentAlert( + title: "Import Failed", + message: error.localizedDescription) + } + } + } + + @MainActor + private func presentFullDiskAccessAlert() { + let alert = NSAlert() + alert.messageText = "Full Disk Access Required" + alert.informativeText = "Full Disk Access is required to read the Antigravity database. Please grant access in System Settings." + alert.addButton(withTitle: "Open System Settings") + alert.addButton(withTitle: "Cancel") + + if alert.runModal() == .alertFirstButtonReturn { + if let url = URL(string: "x-apple.systempreferences:com.apple.preference.security?Privacy_AllFiles") { + NSWorkspace.shared.open(url) + } + } + } + + @MainActor + private func presentAlert(title: String, message: String) { + let alert = NSAlert() + alert.messageText = title + alert.informativeText = message + alert.alertStyle = .warning + alert.runModal() + } + private func makeSettingsContext(provider: UsageProvider) -> ProviderSettingsContext { ProviderSettingsContext( provider: provider, diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift index e41aa8fe..3cd020ba 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift @@ -1,11 +1,199 @@ +import AppKit import CodexBarCore +@MainActor +struct AntigravityLoginFlow { + static func runOAuthFlow(settings: SettingsStore, store: UsageStore? = nil) async -> Bool { + if KeychainAccessGate.isDisabled { + Self.presentAlert( + title: "Authorization Unavailable", + message: "Keychain access is disabled. Enable it to store Antigravity OAuth accounts.") + return false + } + let flow = AntigravityOAuthFlow() + + let waitingAlert = NSAlert() + waitingAlert.messageText = "Waiting for Authentication..." + waitingAlert.informativeText = """ + Please complete the sign-in in your browser. + This window will close automatically when finished. + """ + waitingAlert.addButton(withTitle: "Cancel") + let parentWindow = Self.resolveWaitingParentWindow() + let hostWindow = parentWindow ?? Self.makeWaitingHostWindow() + let shouldCloseHostWindow = parentWindow == nil + + let waitTask = Task { @MainActor in + let response = await Self.presentWaitingAlert(waitingAlert, parentWindow: hostWindow) + if response == .alertFirstButtonReturn { + await flow.cancelAuthorization() + } + return response + } + await Task.yield() + + let authTask = Task.detached(priority: .userInitiated) { + try await flow.startAuthorization() + } + + let authResult: Result + do { + let credentials = try await authTask.value + authResult = .success(credentials) + } catch { + authResult = .failure(error) + } + + Self.dismissWaitingAlert(waitingAlert, parentWindow: hostWindow, closeHost: shouldCloseHostWindow) + let waitResponse = await waitTask.value + if waitResponse == .alertFirstButtonReturn { + return false + } + + switch authResult { + case let .success(credentials): + guard let accountLabel = Self.persistCredentials(credentials, settings: settings) else { + Self.presentAlert(title: "Authorization Failed", message: "Unable to store Antigravity credentials.") + return false + } + if let store { + await store.refreshProvider(.antigravity, allowDisabled: true) + } + + let success = NSAlert() + success.messageText = "Authorization Successful" + var info: [String] = [] + info.append("Signed in as \(accountLabel).") + if !info.isEmpty { + success.informativeText = info.joined(separator: " ") + } + success.runModal() + return true + case let .failure(error): + guard !(error is CancellationError) else { return false } + Self.presentAlert(title: "Authorization Failed", message: error.localizedDescription) + return false + } + } + + private static func persistCredentials( + _ credentials: AntigravityOAuthCredentials, + settings: SettingsStore) -> String? + { + guard let accountLabel = Self.resolveAccountLabel(credentials: credentials, settings: settings) else { + return nil + } + guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: accountLabel) else { + return nil + } + _ = settings.upsertAntigravityTokenAccount(label: accountLabel) + settings.setProviderEnabled( + provider: .antigravity, + metadata: ProviderRegistry.shared.metadata[.antigravity]!, + enabled: true) + return accountLabel + } + + private static func resolveAccountLabel( + credentials: AntigravityOAuthCredentials, + settings: SettingsStore) -> String? + { + if let email = credentials.email, + let normalized = AntigravityOAuthCredentialsStore.normalizedLabel(email) + { + return normalized + } + + let existingLabels = Set(settings.tokenAccounts(for: .antigravity).map { $0.label.lowercased() }) + var index = 1 + while existingLabels.contains("account \(index)") { + index += 1 + } + return "account \(index)" + } + + private static func presentAlert(title: String, message: String) { + let alert = NSAlert() + alert.messageText = title + alert.informativeText = message + alert.alertStyle = .warning + alert.runModal() + } + + @MainActor + private static func presentWaitingAlert( + _ alert: NSAlert, + parentWindow: NSWindow) async -> NSApplication.ModalResponse + { + await withCheckedContinuation { continuation in + alert.beginSheetModal(for: parentWindow) { response in + continuation.resume(returning: response) + } + } + } + + @MainActor + private static func dismissWaitingAlert( + _ alert: NSAlert, + parentWindow: NSWindow, + closeHost: Bool) + { + let alertWindow = alert.window + if alertWindow.sheetParent != nil { + parentWindow.endSheet(alertWindow) + } else { + alertWindow.orderOut(nil) + } + + guard closeHost else { return } + parentWindow.orderOut(nil) + parentWindow.close() + } + + @MainActor + private static func resolveWaitingParentWindow() -> NSWindow? { + if let window = NSApp.keyWindow ?? NSApp.mainWindow { + return window + } + if let window = NSApp.windows.first(where: { $0.isVisible && !$0.ignoresMouseEvents }) { + return window + } + return NSApp.windows.first + } + + @MainActor + private static func makeWaitingHostWindow() -> NSWindow { + let window = NSWindow( + contentRect: NSRect(x: 0, y: 0, width: 420, height: 1), + styleMask: [.titled, .fullSizeContentView], + backing: .buffered, + defer: false) + window.isReleasedWhenClosed = false + window.titleVisibility = .hidden + window.titlebarAppearsTransparent = true + window.standardWindowButton(.closeButton)?.isHidden = true + window.standardWindowButton(.miniaturizeButton)?.isHidden = true + window.standardWindowButton(.zoomButton)?.isHidden = true + window.backgroundColor = .clear + window.isOpaque = false + window.hasShadow = false + window.level = .floating + window.collectionBehavior = [.canJoinAllSpaces, .fullScreenAuxiliary] + window.center() + window.makeKeyAndOrderFront(nil) + return window + } +} + @MainActor extension StatusItemController { - func runAntigravityLoginFlow() async { + func runAntigravityLoginFlow() async -> Bool { + self.loginPhase = .waitingBrowser + let success = await AntigravityLoginFlow.runOAuthFlow(settings: self.settings) self.loginPhase = .idle - self.presentLoginAlert( - title: "Antigravity login is managed in the app", - message: "Open Antigravity to sign in, then refresh CodexBar.") + if success { + self.postLoginNotification(for: .antigravity) + } + return success } } diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index 177b12d4..ecc34318 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -1,10 +1,12 @@ import CodexBarCore import CodexBarMacroSupport import Foundation +import SwiftUI @ProviderImplementationRegistration struct AntigravityProviderImplementation: ProviderImplementation { let id: UsageProvider = .antigravity + let supportsLoginFlow: Bool = true func detectVersion(context _: ProviderVersionContext) async -> String? { await AntigravityStatusProbe.detectVersion() @@ -12,12 +14,67 @@ struct AntigravityProviderImplementation: ProviderImplementation { @MainActor func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { - .antigravity(context.settings.antigravitySettingsSnapshot()) + .antigravity(context.settings.antigravitySettingsSnapshot(tokenOverride: context.tokenOverride)) } @MainActor func runLoginFlow(context: ProviderLoginContext) async -> Bool { await context.controller.runAntigravityLoginFlow() - return false + } + + @MainActor + func sourceMode(context: ProviderSourceModeContext) -> ProviderSourceMode { + switch context.settings.antigravityUsageSource { + case .auto: .auto + case .authorized: .oauth + case .local: .cli + } + } + + @MainActor + func settingsPickers(context: ProviderSettingsContext) -> [ProviderSettingsPickerDescriptor] { + let usageBinding = Binding( + get: { context.settings.antigravityUsageSource.rawValue }, + set: { + if let source = AntigravityUsageSource(rawValue: $0) { + context.settings.antigravityUsageSource = source + } + }) + + let usageOptions = AntigravityUsageSource.allCases.map { + ProviderSettingsPickerOption(id: $0.rawValue, title: $0.displayName) + } + + return [ + ProviderSettingsPickerDescriptor( + id: "antigravity-usage-source", + title: "Usage source", + subtitle: "Choose how to fetch Antigravity usage data.", + dynamicSubtitle: { + switch context.settings.antigravityUsageSource { + case .auto: + return "Auto: Try OAuth/manual first, fallback to local server" + case .authorized: + return "OAuth: Use OAuth account or manual tokens only" + case .local: + return "Local: Use local Antigravity language server only" + } + }, + binding: usageBinding, + options: usageOptions, + isVisible: nil, + onChange: nil, + trailingText: { + let label = context.store.sourceLabel(for: .antigravity) + return label.isEmpty ? nil : label + }), + ] + } + + @MainActor + func tokenAccountsVisibility(context: ProviderSettingsContext, support: TokenAccountSupport) -> Bool { + _ = context + _ = support + return true } } diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift index 1138d70f..818b39cf 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift @@ -16,37 +16,129 @@ extension SettingsStore { self.logProviderModeChange(provider: .antigravity, field: "usageSource", value: newValue.rawValue) } } +} - var antigravityManualToken: String { - get { self.configSnapshot.providerConfig(for: .antigravity)?.manualToken ?? "" } - set { +extension SettingsStore { + func antigravitySettingsSnapshot( + tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.AntigravityProviderSettings + { + let account = ProviderTokenAccountSelection.selectedAccount( + provider: .antigravity, + settings: self, + override: tokenOverride) + let tokenAccounts = self.tokenAccountsData(for: .antigravity) + return ProviderSettingsSnapshot.AntigravityProviderSettings( + usageSource: self.antigravityUsageSource, + accountLabel: account?.label, + tokenAccounts: tokenAccounts) + } + + func upsertAntigravityTokenAccount(label: String) -> ProviderTokenAccount? { + guard let normalized = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return nil } + let tokenValue = normalized + let existing = self.tokenAccountsData(for: .antigravity) + var accounts = existing?.accounts ?? [] + if let index = accounts.firstIndex(where: { $0.label.lowercased() == normalized }) { + let current = accounts[index] + let updated = ProviderTokenAccount( + id: current.id, + label: normalized, + token: tokenValue, + addedAt: current.addedAt, + lastUsed: current.lastUsed) + accounts[index] = updated + let updatedData = ProviderTokenAccountData( + version: existing?.version ?? 1, + accounts: accounts, + activeIndex: index) self.updateProviderConfig(provider: .antigravity) { entry in - entry.manualToken = self.normalizedConfigValue(newValue) + entry.tokenAccounts = updatedData } - self.logSecretUpdate(provider: .antigravity, field: "manualToken", value: newValue) + return updated } + + let account = ProviderTokenAccount( + id: UUID(), + label: normalized, + token: tokenValue, + addedAt: Date().timeIntervalSince1970, + lastUsed: nil) + let updatedData = ProviderTokenAccountData( + version: existing?.version ?? 1, + accounts: accounts + [account], + activeIndex: accounts.count) + self.updateProviderConfig(provider: .antigravity) { entry in + entry.tokenAccounts = updatedData + } + return account } - var antigravityAccountEmail: String? { - AntigravityOAuthCredentialsStore.load()?.email + func removeAntigravityTokenAccount(accountID: UUID) { + guard let data = self.tokenAccountsData(for: .antigravity) else { return } + guard let removed = data.accounts.first(where: { $0.id == accountID }) else { return } + self.removeTokenAccount(provider: .antigravity, accountID: accountID) + AntigravityOAuthCredentialsStore.clear(accountLabel: removed.label) } - var antigravityHasCredentials: Bool { - if let credentials = AntigravityOAuthCredentialsStore.load() { - return !credentials.accessToken.isEmpty || credentials.isRefreshable + func addManualAntigravityTokenAccount( + label: String, + accessToken: String, + refreshToken: String? = nil + ) -> ProviderTokenAccount? { + guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return nil } + + if !KeychainAccessGate.isDisabled { + let credentials = AntigravityOAuthCredentials( + accessToken: accessToken, + refreshToken: refreshToken, + expiresAt: nil, + email: normalizedLabel, + scopes: [] + ) + guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: normalizedLabel) else { + return nil + } } - return false - } - func clearAntigravityCredentials() { - AntigravityOAuthCredentialsStore.clear() - } -} + let tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( + accessToken: accessToken, + refreshToken: refreshToken) -extension SettingsStore { - func antigravitySettingsSnapshot() -> ProviderSettingsSnapshot.AntigravityProviderSettings { - ProviderSettingsSnapshot.AntigravityProviderSettings( - usageSource: self.antigravityUsageSource, - manualToken: self.antigravityManualToken) + let existing = self.tokenAccountsData(for: .antigravity) + var accounts = existing?.accounts ?? [] + + if let index = accounts.firstIndex(where: { $0.label.lowercased() == normalizedLabel }) { + let current = accounts[index] + let updated = ProviderTokenAccount( + id: current.id, + label: normalizedLabel, + token: tokenValue, + addedAt: current.addedAt, + lastUsed: current.lastUsed) + accounts[index] = updated + let updatedData = ProviderTokenAccountData( + version: existing?.version ?? 1, + accounts: accounts, + activeIndex: index) + self.updateProviderConfig(provider: .antigravity) { entry in + entry.tokenAccounts = updatedData + } + return updated + } + + let account = ProviderTokenAccount( + id: UUID(), + label: normalizedLabel, + token: tokenValue, + addedAt: Date().timeIntervalSince1970, + lastUsed: nil) + let updatedData = ProviderTokenAccountData( + version: existing?.version ?? 1, + accounts: accounts + [account], + activeIndex: accounts.count) + self.updateProviderConfig(provider: .antigravity) { entry in + entry.tokenAccounts = updatedData + } + return account } } diff --git a/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift b/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift index c624a671..eb2642e2 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift @@ -87,16 +87,26 @@ struct ProviderSettingsFieldDescriptor: Identifiable { /// Shared token account descriptor rendered in the Providers settings pane. @MainActor struct ProviderSettingsTokenAccountsDescriptor: Identifiable { + struct ImportAction { + let title: String + let action: () async -> Void + } + let id: String let title: String let subtitle: String let placeholder: String + let supportsManualEntry: Bool + let supportsTwoFieldEntry: Bool + let addActionTitle: String? + let addAction: (() async -> Void)? + let importAction: ImportAction? let provider: UsageProvider let isVisible: (() -> Bool)? let accounts: () -> [ProviderTokenAccount] let activeIndex: () -> Int let setActiveIndex: (Int) -> Void - let addAccount: (_ label: String, _ token: String) -> Void + let addAccount: (_ label: String, _ token: String, _ token2: String) -> Void let removeAccount: (_ accountID: UUID) -> Void let openConfigFile: () -> Void let reloadFromDisk: () -> Void diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 7127a123..128eb54c 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -87,6 +87,14 @@ extension UsageStore { settings: self.settings, tokenOverride: override) let verbose = self.settings.isVerboseLoggingEnabled + + let onCredentialsRefreshed: (@Sendable (UsageProvider, String, String?) -> Void) = { [weak self] provider, accountLabel, newAccessToken in + Task { @MainActor in + guard let self = self else { return } + self.saveRefreshedCredentialsToConfig(provider: provider, accountLabel: accountLabel, newAccessToken: newAccessToken) + } + } + let context = ProviderFetchContext( runtime: .app, sourceMode: sourceMode, @@ -98,7 +106,8 @@ extension UsageStore { settings: snapshot, fetcher: self.codexFetcher, claudeFetcher: self.claudeFetcher, - browserDetection: self.browserDetection) + browserDetection: self.browserDetection, + onCredentialsRefreshed: onCredentialsRefreshed) return await descriptor.fetchOutcome(context: context) } @@ -202,4 +211,44 @@ extension UsageStore { updatedAt: snapshot.updatedAt, identity: identity) } + + @MainActor + private func saveRefreshedCredentialsToConfig(provider: UsageProvider, accountLabel: String, newAccessToken: String?) { + guard provider == .antigravity else { return } + guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(accountLabel) else { return } + + let tokenAccounts = self.settings.tokenAccountsData(for: .antigravity) + guard let account = tokenAccounts?.accounts.first(where: { $0.label.lowercased() == normalizedLabel }) else { return } + guard let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: account.token) else { return } + + let accessToken = newAccessToken ?? payload.accessToken + let tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( + accessToken: accessToken, + refreshToken: payload.refreshToken) + + guard var accounts = tokenAccounts?.accounts else { return } + guard let index = accounts.firstIndex(where: { $0.id == account.id }) else { return } + + let updatedAccount = ProviderTokenAccount( + id: account.id, + label: account.label, + token: tokenValue, + addedAt: account.addedAt, + lastUsed: Date().timeIntervalSince1970) + accounts[index] = updatedAccount + + let updatedData = ProviderTokenAccountData( + version: tokenAccounts?.version ?? 1, + accounts: accounts, + activeIndex: tokenAccounts?.activeIndex ?? 0) + + self.settings.updateProviderConfig(provider: .antigravity) { entry in + entry.tokenAccounts = updatedData + } + + self.providerLogger.info("Saved refreshed credentials to JSON config", metadata: [ + "provider": "antigravity", + "account": normalizedLabel, + ]) + } } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 69491056..c96b0784 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -61,6 +61,59 @@ extension UsageStore { } } } + + private func debugAntigravityLog( + usageSource: AntigravityUsageSource, + accountLabel: String?) async -> String + { + let tokenAccounts = await MainActor.run { self.settings.tokenAccountsData(for: .antigravity) } + + return await self.runWithTimeout(seconds: 15) { + var lines: [String] = [] + + let trimmedLabel = accountLabel?.trimmingCharacters(in: .whitespacesAndNewlines) + let hasAccountLabel = !(trimmedLabel?.isEmpty ?? true) + + let keychainCreds = hasAccountLabel + ? AntigravityOAuthCredentialsStore.load(accountLabel: trimmedLabel ?? "") + : nil + let hasKeychainAccessToken = !(keychainCreds?.accessToken.isEmpty ?? true) + let hasKeychainRefreshToken = keychainCreds?.isRefreshable ?? false + let isKeychainExpired = keychainCreds?.isExpired ?? false + + let normalizedLabel = trimmedLabel?.lowercased() ?? "" + let manualAccount = tokenAccounts?.accounts.first { $0.label.lowercased() == normalizedLabel } + let manualPayload = manualAccount.flatMap { + AntigravityOAuthCredentialsStore.manualTokenPayload(from: $0.token) + } + let hasManualAccessToken = manualPayload?.accessToken.isEmpty == false + + let serverRunning = await AntigravityStatusProbe.isRunning() + + let present = { $0 ? "present" : "missing" } + let available = { $0 ? "available" : "unavailable" } + + lines.append("usageSource=\(usageSource.rawValue)") + lines.append("accountLabel=\(hasAccountLabel ? (trimmedLabel ?? "") : "none")") + lines.append("keychainCredentials=\(present(keychainCreds != nil))") + lines.append("keychainAccessToken=\(present(hasKeychainAccessToken))") + lines.append("keychainRefreshToken=\(present(hasKeychainRefreshToken))") + lines.append("keychainExpired=\(isKeychainExpired)") + if let email = keychainCreds?.email, !email.isEmpty { + lines.append("keychainEmail=\(email)") + } + lines.append("manualCredentials=\(present(manualPayload != nil))") + lines.append("manualAccessToken=\(present(hasManualAccessToken))") + lines.append("localServer=\(serverRunning ? "running" : "not_running")") + + lines.append("") + let isAuthorizedAvailable = hasAccountLabel && (hasKeychainAccessToken || hasKeychainRefreshToken || manualPayload != nil) + lines.append("authorizedStrategy=\(available(isAuthorizedAvailable))") + lines.append("localStrategy=\(available(serverRunning))") + + return lines.joined(separator: "\n") + } + } } enum ProviderStatusIndicator: String { @@ -1134,6 +1187,8 @@ extension UsageStore { let keepCLISessionsAlive = self.settings.debugKeepCLISessionsAlive let cursorCookieSource = self.settings.cursorCookieSource let cursorCookieHeader = self.settings.cursorCookieHeader + let antigravityUsageSource = self.settings.antigravityUsageSource + let antigravityAccountLabel = self.settings.selectedTokenAccount(for: .antigravity)?.label return await Task.detached(priority: .utility) { () -> String in switch provider { case .codex: @@ -1168,7 +1223,9 @@ extension UsageStore { await MainActor.run { self.probeLogs[.gemini] = text } return text case .antigravity: - let text = "Antigravity debug log not yet implemented" + let text = await self.debugAntigravityLog( + usageSource: antigravityUsageSource, + accountLabel: antigravityAccountLabel) await MainActor.run { self.probeLogs[.antigravity] = text } return text case .cursor: diff --git a/Sources/CodexBarCLI/CLIHelpers.swift b/Sources/CodexBarCLI/CLIHelpers.swift index 00d25058..67f80fb4 100644 --- a/Sources/CodexBarCLI/CLIHelpers.swift +++ b/Sources/CodexBarCLI/CLIHelpers.swift @@ -197,6 +197,9 @@ extension CodexBarCLI { return .web } guard let raw = values.options["source"]?.last?.lowercased() else { return nil } + if raw == "local" { + return .cli + } return ProviderSourceMode(rawValue: raw) } diff --git a/Sources/CodexBarCLI/TokenAccountCLI.swift b/Sources/CodexBarCLI/TokenAccountCLI.swift index a0a88371..bbc4bfa7 100644 --- a/Sources/CodexBarCLI/TokenAccountCLI.swift +++ b/Sources/CodexBarCLI/TokenAccountCLI.swift @@ -147,7 +147,15 @@ struct TokenAccountCLIContext { return self.makeSnapshot( jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings( ideBasePath: nil)) - case .gemini, .antigravity, .copilot, .kiro, .vertexai, .kimik2, .synthetic: + case .antigravity: + let usageSourceRaw = config?.usageSource ?? "" + let usageSource = AntigravityUsageSource(rawValue: usageSourceRaw) ?? .auto + return self.makeSnapshot( + antigravity: ProviderSettingsSnapshot.AntigravityProviderSettings( + usageSource: usageSource, + accountLabel: account?.label, + tokenAccounts: config?.tokenAccounts)) + case .gemini, .copilot, .kiro, .vertexai, .kimik2, .synthetic: return nil } } @@ -163,7 +171,8 @@ struct TokenAccountCLIContext { kimi: ProviderSettingsSnapshot.KimiProviderSettings? = nil, augment: ProviderSettingsSnapshot.AugmentProviderSettings? = nil, amp: ProviderSettingsSnapshot.AmpProviderSettings? = nil, - jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? = nil) -> ProviderSettingsSnapshot + jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? = nil, + antigravity: ProviderSettingsSnapshot.AntigravityProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot.make( codex: codex, @@ -176,7 +185,8 @@ struct TokenAccountCLIContext { kimi: kimi, augment: augment, amp: amp, - jetbrains: jetbrains) + jetbrains: jetbrains, + antigravity: antigravity) } func environment( diff --git a/Sources/CodexBarCore/Config/CodexBarConfig.swift b/Sources/CodexBarCore/Config/CodexBarConfig.swift index f151b070..cbd11cb4 100644 --- a/Sources/CodexBarCore/Config/CodexBarConfig.swift +++ b/Sources/CodexBarCore/Config/CodexBarConfig.swift @@ -83,7 +83,6 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { public var workspaceID: String? public var tokenAccounts: ProviderTokenAccountData? public var usageSource: String? - public var manualToken: String? public init( id: UsageProvider, @@ -95,8 +94,7 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { region: String? = nil, workspaceID: String? = nil, tokenAccounts: ProviderTokenAccountData? = nil, - usageSource: String? = nil, - manualToken: String? = nil) + usageSource: String? = nil) { self.id = id self.enabled = enabled @@ -108,7 +106,6 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { self.workspaceID = workspaceID self.tokenAccounts = tokenAccounts self.usageSource = usageSource - self.manualToken = manualToken } public var sanitizedAPIKey: String? { diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift index 22cf10b0..98b00f1c 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift @@ -9,52 +9,59 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { public init() {} public func isAvailable(_ context: ProviderFetchContext) async -> Bool { - if let credentials = AntigravityOAuthCredentialsStore.load() { - if !credentials.accessToken.isEmpty { return true } - if credentials.isRefreshable { return true } + guard let accountLabel = context.settings?.antigravity?.accountLabel else { + return false } - if let manualToken = context.settings?.antigravityManualToken, - !manualToken.isEmpty - { - return true + if let manualCredentials = self.loadManualCredentials(accountLabel: accountLabel, context: context) { + return !manualCredentials.accessToken.isEmpty } - return false + guard let credentials = AntigravityOAuthCredentialsStore.load(accountLabel: accountLabel) else { + return false + } + if !credentials.accessToken.isEmpty { return true } + return credentials.isRefreshable } public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { - let credentials = try await self.resolveCredentials(context: context) + let resolved = try await self.resolveCredentials(context: context) + let accountLabel = resolved.accountLabel + let credentials = resolved.credentials + let sourceLabel = resolved.sourceLabel - let accessToken: String - if credentials.needsRefresh, credentials.isRefreshable { - Self.log.info("Antigravity credentials need refresh") + var refreshedCredentials: AntigravityOAuthCredentials? + if credentials.needsRefresh || (credentials.accessToken.isEmpty && credentials.isRefreshable) { let refreshed = try await self.refreshCredentials(credentials) - accessToken = refreshed.accessToken - } else if credentials.accessToken.isEmpty, credentials.isRefreshable { - Self.log.info("Antigravity credentials have no access token, refreshing") - let refreshed = try await self.refreshCredentials(credentials) - accessToken = refreshed.accessToken - } else { - accessToken = credentials.accessToken + refreshedCredentials = refreshed + if KeychainAccessGate.isDisabled { + context.onCredentialsRefreshed?(.antigravity, accountLabel, refreshed.accessToken) + } else { + _ = AntigravityOAuthCredentialsStore.save(refreshed, accountLabel: accountLabel) + } } - let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: accessToken) + let activeCredentials = refreshedCredentials ?? credentials + let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: activeCredentials.accessToken) let snapshot = AntigravityStatusSnapshot( modelQuotas: quota.models, - accountEmail: credentials.email ?? quota.email, + accountEmail: activeCredentials.email ?? quota.email, accountPlan: nil) let usage = try snapshot.toUsageSnapshot() - return self.makeResult(usage: usage, sourceLabel: "authorized") + return self.makeResult(usage: usage, sourceLabel: sourceLabel) } - public func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool { + public func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + let usageSource = context.settings?.antigravity?.usageSource ?? .auto + + if usageSource == .auto { + return true + } + if let oauthError = error as? AntigravityOAuthCredentialsError { switch oauthError { - case .invalidGrant: - return true - case .notFound: + case .invalidGrant, .notFound: return true default: return false @@ -63,47 +70,28 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { return false } - private func resolveCredentials(context: ProviderFetchContext) async throws -> AntigravityOAuthCredentials { - if let cached = AntigravityOAuthCredentialsStore.load() { - return cached + private func resolveCredentials(context: ProviderFetchContext) async throws + -> (accountLabel: String, credentials: AntigravityOAuthCredentials, sourceLabel: String) + { + guard let accountLabel = context.settings?.antigravity?.accountLabel, + let normalized = AntigravityOAuthCredentialsStore.normalizedLabel(accountLabel) + else { + throw AntigravityOAuthCredentialsError.notFound } - if let manualToken = context.settings?.antigravityManualToken, - !manualToken.isEmpty - { - if let parsed = AntigravityOAuthCredentialsStore.parseManualToken(manualToken) { - if parsed.isRefreshable, parsed.accessToken.isEmpty { - let refreshed = try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( - refreshToken: parsed.refreshToken!) - AntigravityOAuthCredentialsStore.save(refreshed) - return refreshed - } - return parsed - } + if let manualCredentials = self.loadManualCredentials(accountLabel: accountLabel, context: context) { + return (normalized, manualCredentials, "Manual") } - if AntigravityLocalImporter.isAvailable() { - let localInfo = try await AntigravityLocalImporter.importCredentials() - if let refreshToken = localInfo.refreshToken, !refreshToken.isEmpty { - let credentials = try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( - refreshToken: refreshToken, - fallbackEmail: localInfo.email) - AntigravityOAuthCredentialsStore.save(credentials) - return credentials - } - if let accessToken = localInfo.accessToken, !accessToken.isEmpty { - let credentials = AntigravityOAuthCredentials( - accessToken: accessToken, - refreshToken: nil, - expiresAt: nil, - email: localInfo.email, - scopes: AntigravityOAuthConfig.scopes) - AntigravityOAuthCredentialsStore.save(credentials) - return credentials - } + guard let cached = AntigravityOAuthCredentialsStore.load(accountLabel: normalized) else { + throw AntigravityOAuthCredentialsError.notFound + } + + if cached.accessToken.isEmpty, !cached.isRefreshable { + throw AntigravityOAuthCredentialsError.notFound } - throw AntigravityOAuthCredentialsError.notFound + return (normalized, cached, "OAuth") } private func refreshCredentials(_ credentials: AntigravityOAuthCredentials) async throws -> AntigravityOAuthCredentials { @@ -111,10 +99,28 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { throw AntigravityOAuthCredentialsError.invalidGrant } - let refreshed = try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( + return try await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( refreshToken: refreshToken, fallbackEmail: credentials.email) - AntigravityOAuthCredentialsStore.save(refreshed) - return refreshed + } + + private func loadManualCredentials( + accountLabel: String, + context: ProviderFetchContext + ) -> AntigravityOAuthCredentials? { + guard let normalized = AntigravityOAuthCredentialsStore.normalizedLabel(accountLabel) else { return nil } + + let tokenAccounts = context.settings?.antigravity?.tokenAccounts + guard let account = tokenAccounts?.accounts.first(where: { $0.label.lowercased() == normalized }) else { + return nil + } + + guard let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: account.token) else { return nil } + return AntigravityOAuthCredentials( + accessToken: payload.accessToken, + refreshToken: payload.refreshToken, + expiresAt: nil, + email: account.label, + scopes: []) } } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift index 9311bed6..7bcf4ecf 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift @@ -2,8 +2,6 @@ import Foundation import SQLite3 public enum AntigravityLocalImporter { - private static let log = CodexBarLog.logger(LogCategories.antigravity) - public struct LocalCredentialInfo: Sendable { public let accessToken: String? public let refreshToken: String? @@ -117,7 +115,12 @@ public enum AntigravityLocalImporter { } defer { sqlite3_finalize(stmt) } - sqlite3_bind_text(stmt, 1, key, -1, nil) + let keyCString = key.cString(using: .utf8) + guard let keyCString else { + throw AntigravityOAuthCredentialsError.decodeFailed("Failed to convert key to UTF-8: \(key)") + } + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + sqlite3_bind_text(stmt, 1, keyCString, -1, transient) guard sqlite3_step(stmt) == SQLITE_ROW else { throw AntigravityOAuthCredentialsError.notFound diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift index 763961ff..541492de 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift @@ -1,7 +1,4 @@ import Foundation -#if os(macOS) -import Security -#endif public struct AntigravityOAuthCredentials: Sendable, Codable { public let accessToken: String @@ -76,7 +73,7 @@ public enum AntigravityOAuthCredentialsError: LocalizedError, Sendable { public var errorDescription: String? { switch self { case .notFound: - "Antigravity credentials not found. Authorize or import from Antigravity app." + "Antigravity credentials not found. Sign in with Google to add an OAuth account." case let .decodeFailed(message): "Failed to decode Antigravity credentials: \(message)" case .missingAccessToken: @@ -94,147 +91,134 @@ public enum AntigravityOAuthCredentialsError: LocalizedError, Sendable { } public enum AntigravityOAuthCredentialsStore { - private static let cacheKey = KeychainCacheStore.Key.oauth(provider: .antigravity) + public static let manualTokenPrefix = "manual:" private static let log = CodexBarLog.logger(LogCategories.antigravity) - public static let environmentTokenKey = "CODEXBAR_ANTIGRAVITY_TOKEN" + public static let environmentAccountKey = "CODEXBAR_ANTIGRAVITY_ACCOUNT" + private static let cacheCategory = "oauth.antigravity" + + public struct ManualTokenPayload: Sendable { + public let accessToken: String + public let refreshToken: String? + + public init(accessToken: String, refreshToken: String?) { + self.accessToken = accessToken + self.refreshToken = refreshToken + } + } struct CacheEntry: Codable, Sendable { let credentials: AntigravityOAuthCredentials let storedAt: Date } - private nonisolated(unsafe) static var cachedCredentials: AntigravityOAuthCredentials? - private nonisolated(unsafe) static var cacheTimestamp: Date? + private nonisolated(unsafe) static var cachedCredentialsByLabel: [String: AntigravityOAuthCredentials] = [:] + private nonisolated(unsafe) static var cacheTimestampByLabel: [String: Date] = [:] private static let memoryCacheValidityDuration: TimeInterval = 1800 - public static func load( - environment: [String: String] = ProcessInfo.processInfo.environment) -> AntigravityOAuthCredentials? - { - if let credentials = self.loadFromEnvironment(environment) { - return credentials + public static func normalizedLabel(_ label: String) -> String? { + let trimmed = label.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + return trimmed.lowercased() + } + + public static func manualTokenPayload(from token: String) -> ManualTokenPayload? { + guard token.hasPrefix(self.manualTokenPrefix) else { return nil } + let content = String(token.dropFirst(self.manualTokenPrefix.count)) + if let jsonData = content.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: jsonData) as? [String: String], + let accessToken = json["access"], + !accessToken.isEmpty + { + return ManualTokenPayload(accessToken: accessToken, refreshToken: json["refresh"]) + } + guard !content.isEmpty else { return nil } + return ManualTokenPayload(accessToken: content, refreshToken: nil) + } + + public static func manualTokenValue(accessToken: String, refreshToken: String?) -> String { + let trimmedAccess = accessToken.trimmingCharacters(in: .whitespacesAndNewlines) + let trimmedRefresh = refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines) + if let refresh = trimmedRefresh, !refresh.isEmpty { + let tokenData = ["access": trimmedAccess, "refresh": refresh] + if let jsonData = try? JSONSerialization.data(withJSONObject: tokenData), + let jsonString = String(data: jsonData, encoding: .utf8) + { + return "\(self.manualTokenPrefix)\(jsonString)" + } } + return "\(self.manualTokenPrefix)\(trimmedAccess)" + } + + public static func load(accountLabel: String) -> AntigravityOAuthCredentials? { + guard let normalized = self.normalizedLabel(accountLabel) else { return nil } + guard !KeychainAccessGate.isDisabled else { return nil } - if let cached = self.cachedCredentials, - let timestamp = self.cacheTimestamp, + if let cached = self.cachedCredentialsByLabel[normalized], + let timestamp = self.cacheTimestampByLabel[normalized], Date().timeIntervalSince(timestamp) < self.memoryCacheValidityDuration, !cached.isExpired { return cached } - switch KeychainCacheStore.load(key: self.cacheKey, as: CacheEntry.self) { + switch KeychainCacheStore.load(key: self.key(for: normalized), as: CacheEntry.self) { case let .found(entry): if entry.credentials.isExpired, !entry.credentials.isRefreshable { self.log.debug("Antigravity cached credentials expired and not refreshable") return entry.credentials } - self.cachedCredentials = entry.credentials - self.cacheTimestamp = Date() + self.cachedCredentialsByLabel[normalized] = entry.credentials + self.cacheTimestampByLabel[normalized] = Date() return entry.credentials case .invalid: - KeychainCacheStore.clear(key: self.cacheKey) + KeychainCacheStore.clear(key: self.key(for: normalized)) + self.cachedCredentialsByLabel.removeValue(forKey: normalized) + self.cacheTimestampByLabel.removeValue(forKey: normalized) return nil case .missing: return nil } } - public static func save(_ credentials: AntigravityOAuthCredentials) { - let entry = CacheEntry(credentials: credentials, storedAt: Date()) - KeychainCacheStore.store(key: self.cacheKey, entry: entry) - self.cachedCredentials = credentials - self.cacheTimestamp = Date() - self.log.info("Antigravity credentials saved", metadata: [ - "email": credentials.email ?? "unknown", - "hasRefreshToken": "\(credentials.isRefreshable)", - ]) - } + public static func save(_ credentials: AntigravityOAuthCredentials, accountLabel: String) -> Bool { + guard let normalized = self.normalizedLabel(accountLabel) else { return false } + guard !KeychainAccessGate.isDisabled else { + self.log.error("Antigravity OAuth save failed: keychain access disabled") + return false + } - public static func clear() { - KeychainCacheStore.clear(key: self.cacheKey) - self.cachedCredentials = nil - self.cacheTimestamp = nil - self.log.info("Antigravity credentials cleared") + self.saveToKeychain(credentials, normalizedLabel: normalized) + return true } - public static func invalidateCache() { - self.cachedCredentials = nil - self.cacheTimestamp = nil + public static func clear(accountLabel: String) { + guard let normalized = self.normalizedLabel(accountLabel) else { return } + KeychainCacheStore.clear(key: self.key(for: normalized)) + self.cachedCredentialsByLabel.removeValue(forKey: normalized) + self.cacheTimestampByLabel.removeValue(forKey: normalized) + self.log.info("Antigravity credentials cleared", metadata: [ + "label": normalized, + ]) } - private static func loadFromEnvironment(_ environment: [String: String]) -> AntigravityOAuthCredentials? { - guard let token = environment[self.environmentTokenKey]?.trimmingCharacters(in: .whitespacesAndNewlines), - !token.isEmpty - else { - return nil - } - - return AntigravityOAuthCredentials( - accessToken: token, - refreshToken: nil, - expiresAt: Date.distantFuture, - email: nil, - scopes: AntigravityOAuthConfig.scopes) + public static func invalidateCache() { + self.cachedCredentialsByLabel.removeAll() + self.cacheTimestampByLabel.removeAll() } -} - -extension AntigravityOAuthCredentialsStore { - public static func parseManualToken(_ input: String) -> AntigravityOAuthCredentials? { - let trimmed = input.trimmingCharacters(in: .whitespacesAndNewlines) - guard !trimmed.isEmpty else { return nil } - - if let jsonCredentials = self.parseJSONInput(trimmed) { - return jsonCredentials - } - - if trimmed.hasPrefix("ya29.") || trimmed.hasPrefix("1//") { - let isRefreshToken = trimmed.hasPrefix("1//") - if isRefreshToken { - return AntigravityOAuthCredentials( - accessToken: "", - refreshToken: trimmed, - expiresAt: nil, - email: nil, - scopes: AntigravityOAuthConfig.scopes) - } else { - return AntigravityOAuthCredentials( - accessToken: trimmed, - refreshToken: nil, - expiresAt: nil, - email: nil, - scopes: AntigravityOAuthConfig.scopes) - } - } - return nil + private static func saveToKeychain(_ credentials: AntigravityOAuthCredentials, normalizedLabel: String) { + let entry = CacheEntry(credentials: credentials, storedAt: Date()) + KeychainCacheStore.store(key: self.key(for: normalizedLabel), entry: entry) + self.cachedCredentialsByLabel[normalizedLabel] = credentials + self.cacheTimestampByLabel[normalizedLabel] = Date() + self.log.info("Antigravity credentials saved", metadata: [ + "label": normalizedLabel, + "email": credentials.email ?? "unknown", + "hasRefreshToken": "\(credentials.isRefreshable)", + ]) } - private static func parseJSONInput(_ input: String) -> AntigravityOAuthCredentials? { - guard let data = input.data(using: .utf8) else { return nil } - guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { - return nil - } - - let apiKey = (json["apiKey"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) - let accessToken = (json["accessToken"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) - let refreshToken = (json["refreshToken"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) - let email = (json["email"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) - - let token = apiKey ?? accessToken - guard let token, !token.isEmpty else { return nil } - - var expiresAt: Date? - if let expiresAtString = json["expiresAt"] as? String { - expiresAt = ISO8601DateFormatter().date(from: expiresAtString) - } else if let expiresAtMillis = json["expiresAt"] as? Double { - expiresAt = Date(timeIntervalSince1970: expiresAtMillis / 1000.0) - } - - return AntigravityOAuthCredentials( - accessToken: token, - refreshToken: refreshToken, - expiresAt: expiresAt, - email: email, - scopes: AntigravityOAuthConfig.scopes) + private static func key(for normalizedLabel: String) -> KeychainCacheStore.Key { + KeychainCacheStore.Key(category: self.cacheCategory, identifier: normalizedLabel) } } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift index e1b5ba10..c8557ee8 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -27,22 +27,17 @@ public actor AntigravityOAuthFlow { Self.log.info("Opening Antigravity OAuth authorization URL") #if os(macOS) if let url = URL(string: authURL) { - await MainActor.run { + _ = await MainActor.run { NSWorkspace.shared.open(url) } } #endif - do { - let code = try await withCheckedThrowingContinuation { continuation in - self.pendingContinuation = continuation - } - - let credentials = try await self.exchangeCodeForToken(code: code, redirectUri: redirectUri) - return credentials - } catch { - throw error + let code = try await withCheckedThrowingContinuation { continuation in + self.pendingContinuation = continuation } + + return try await self.exchangeCodeForToken(code: code, redirectUri: redirectUri) } public func cancelAuthorization() { diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift index 274a9d16..ecaf8c10 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift @@ -5,12 +5,25 @@ public enum AntigravityUsageSource: String, CaseIterable, Sendable, Codable { case authorized case local + public init?(rawValue: String) { + switch rawValue { + case "auto": + self = .auto + case "authorized": + self = .authorized + case "local", "cli": + self = .local + default: + return nil + } + } + public var displayName: String { switch self { case .auto: "Auto" case .authorized: - "Authorized (OAuth)" + "OAuth" case .local: "Local Server" } @@ -19,11 +32,11 @@ public enum AntigravityUsageSource: String, CaseIterable, Sendable, Codable { public var description: String { switch self { case .auto: - "Use authorized credentials if available, fallback to local server" + "Try OAuth/manual tokens first, fallback to local server" case .authorized: - "Use OAuth credentials to fetch quota from Cloud Code API" + "Use OAuth account or manual tokens only" case .local: - "Use local Antigravity language server" + "Use local Antigravity language server only" } } } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift index f8743497..bfea3874 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift @@ -59,7 +59,7 @@ struct AntigravityLocalFetchStrategy: ProviderFetchStrategy { let kind: ProviderFetchKind = .localProbe func isAvailable(_: ProviderFetchContext) async -> Bool { - await AntigravityStatusProbe.isRunning() + true } func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { @@ -68,7 +68,7 @@ struct AntigravityLocalFetchStrategy: ProviderFetchStrategy { let usage = try snap.toUsageSnapshot() return self.makeResult( usage: usage, - sourceLabel: "local") + sourceLabel: "Local Server") } func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { diff --git a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift index cadbdc81..77aa1f43 100644 --- a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift +++ b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift @@ -29,6 +29,7 @@ public struct ProviderFetchContext: Sendable { public let fetcher: UsageFetcher public let claudeFetcher: any ClaudeUsageFetching public let browserDetection: BrowserDetection + public let onCredentialsRefreshed: (@Sendable (UsageProvider, String, String?) -> Void)? public init( runtime: ProviderRuntime, @@ -41,7 +42,8 @@ public struct ProviderFetchContext: Sendable { settings: ProviderSettingsSnapshot?, fetcher: UsageFetcher, claudeFetcher: any ClaudeUsageFetching, - browserDetection: BrowserDetection) + browserDetection: BrowserDetection, + onCredentialsRefreshed: (@Sendable (UsageProvider, String, String?) -> Void)? = nil) { self.runtime = runtime self.sourceMode = sourceMode @@ -54,6 +56,7 @@ public struct ProviderFetchContext: Sendable { self.fetcher = fetcher self.claudeFetcher = claudeFetcher self.browserDetection = browserDetection + self.onCredentialsRefreshed = onCredentialsRefreshed } } diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index e43f96f3..0f21f6fe 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -171,11 +171,17 @@ public struct ProviderSettingsSnapshot: Sendable { public struct AntigravityProviderSettings: Sendable { public let usageSource: AntigravityUsageSource - public let manualToken: String? + public let accountLabel: String? + public let tokenAccounts: ProviderTokenAccountData? - public init(usageSource: AntigravityUsageSource, manualToken: String?) { + public init( + usageSource: AntigravityUsageSource, + accountLabel: String?, + tokenAccounts: ProviderTokenAccountData? = nil) + { self.usageSource = usageSource - self.manualToken = manualToken + self.accountLabel = accountLabel + self.tokenAccounts = tokenAccounts } } @@ -199,10 +205,6 @@ public struct ProviderSettingsSnapshot: Sendable { self.jetbrains?.ideBasePath } - public var antigravityManualToken: String? { - self.antigravity?.manualToken - } - public init( debugMenuEnabled: Bool, debugKeepCLISessionsAlive: Bool, diff --git a/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift b/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift index b2bc65d4..dc3e752c 100644 --- a/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift +++ b/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift @@ -9,6 +9,13 @@ extension TokenAccountSupportCatalog { injection: .cookieHeader, requiresManualCookieSource: true, cookieName: "sessionKey"), + .antigravity: TokenAccountSupport( + title: "OAuth accounts", + subtitle: "Sign in with Google or paste tokens manually to add Antigravity accounts.", + placeholder: "OAuth account", + injection: .environment(key: AntigravityOAuthCredentialsStore.environmentAccountKey), + requiresManualCookieSource: true, + cookieName: nil), .zai: TokenAccountSupport( title: "API tokens", subtitle: "Stored in the CodexBar config file.", diff --git a/Tests/CodexBarTests/AntigravityOAuthTests.swift b/Tests/CodexBarTests/AntigravityOAuthTests.swift index 6388e17b..253c96f9 100644 --- a/Tests/CodexBarTests/AntigravityOAuthTests.swift +++ b/Tests/CodexBarTests/AntigravityOAuthTests.swift @@ -2,202 +2,176 @@ import Foundation import Testing @testable import CodexBarCore -@Suite("AntigravityOAuthCredentials") +@Suite struct AntigravityOAuthCredentialsTests { - @Test("Credentials are not expired when expiresAt is in the future") - func test_credentialsNotExpired() { - let creds = AntigravityOAuthCredentials( - accessToken: "ya29.test", - refreshToken: "1//refresh", - expiresAt: Date().addingTimeInterval(3600), - email: "test@example.com") - #expect(!creds.isExpired) - } - - @Test("Credentials are expired when expiresAt is in the past") - func test_credentialsExpired() { - let creds = AntigravityOAuthCredentials( + @Test + func isExpired() { + let expired = AntigravityOAuthCredentials( accessToken: "ya29.test", refreshToken: "1//refresh", expiresAt: Date().addingTimeInterval(-3600), email: "test@example.com") - #expect(creds.isExpired) - } - - @Test("Credentials with nil expiresAt are not expired") - func test_credentialsNilExpiresAtNotExpired() { - let creds = AntigravityOAuthCredentials( + let valid = AntigravityOAuthCredentials( accessToken: "ya29.test", refreshToken: "1//refresh", - expiresAt: nil, + expiresAt: Date().addingTimeInterval(3600), email: "test@example.com") - #expect(!creds.isExpired) + #expect(expired.isExpired) + #expect(!valid.isExpired) } - @Test("Credentials are refreshable when refresh token is present") - func test_credentialsIsRefreshable() { + @Test + func needsRefreshWhenExpiringSoon() { let creds = AntigravityOAuthCredentials( accessToken: "ya29.test", refreshToken: "1//refresh", - expiresAt: nil, + expiresAt: Date().addingTimeInterval(120), email: nil) - #expect(creds.isRefreshable) + #expect(creds.needsRefresh) } +} - @Test("Credentials are not refreshable when refresh token is nil") - func test_credentialsNotRefreshableNil() { - let creds = AntigravityOAuthCredentials( - accessToken: "ya29.test", - refreshToken: nil, - expiresAt: nil, - email: nil) - #expect(!creds.isRefreshable) +@Suite(.serialized) +struct AntigravityOAuthCredentialsStoreTests { + @Test + func normalizesLabel() { + let normalized = AntigravityOAuthCredentialsStore.normalizedLabel(" User@Example.com \n") + #expect(normalized == "user@example.com") } - @Test("Credentials are not refreshable when refresh token is empty") - func test_credentialsNotRefreshableEmpty() { - let creds = AntigravityOAuthCredentials( - accessToken: "ya29.test", - refreshToken: "", - expiresAt: nil, - email: nil) - #expect(!creds.isRefreshable) - } + @Test + func saveAndLoadByLabel() { + KeychainCacheStore.setTestStoreForTesting(true) + defer { KeychainCacheStore.setTestStoreForTesting(false) } + let previousKeychainDisabled = KeychainAccessGate.isDisabled + KeychainAccessGate.isDisabled = false + defer { KeychainAccessGate.isDisabled = previousKeychainDisabled } + AntigravityOAuthCredentialsStore.invalidateCache() - @Test("Credentials need refresh when close to expiry") - func test_credentialsNeedRefresh() { - let creds = AntigravityOAuthCredentials( - accessToken: "ya29.test", - refreshToken: "1//refresh", - expiresAt: Date().addingTimeInterval(120), - email: nil) - #expect(creds.needsRefresh) - } - - @Test("Credentials don't need refresh when not close to expiry") - func test_credentialsDontNeedRefresh() { let creds = AntigravityOAuthCredentials( accessToken: "ya29.test", refreshToken: "1//refresh", expiresAt: Date().addingTimeInterval(3600), - email: nil) - #expect(!creds.needsRefresh) - } -} - -@Suite("AntigravityManualTokenParsing") -struct AntigravityManualTokenParsingTests { - @Test("Parses access token starting with ya29.") - func test_parseAccessToken() { - let token = "ya29.a0ARrdaM..." - let creds = AntigravityOAuthCredentialsStore.parseManualToken(token) - #expect(creds != nil) - #expect(creds?.accessToken == token) - #expect(creds?.refreshToken == nil) - } - - @Test("Parses refresh token starting with 1//") - func test_parseRefreshToken() { - let token = "1//0gXyz..." - let creds = AntigravityOAuthCredentialsStore.parseManualToken(token) - #expect(creds != nil) - #expect(creds?.accessToken.isEmpty == true) - #expect(creds?.refreshToken == token) - } - - @Test("Parses JSON with apiKey") - func test_parseJSONWithApiKey() { - let json = """ - {"apiKey": "ya29.test", "email": "test@example.com"} - """ - let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) - #expect(creds != nil) - #expect(creds?.accessToken == "ya29.test") - #expect(creds?.email == "test@example.com") - } - - @Test("Parses JSON with accessToken") - func test_parseJSONWithAccessToken() { - let json = """ - {"accessToken": "ya29.test", "refreshToken": "1//refresh"} - """ - let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) - #expect(creds != nil) - #expect(creds?.accessToken == "ya29.test") - #expect(creds?.refreshToken == "1//refresh") - } - - @Test("Parses JSON with expiresAt as string") - func test_parseJSONWithExpiresAtString() { - let json = """ - {"apiKey": "ya29.test", "expiresAt": "2025-01-01T00:00:00Z"} - """ - let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) - #expect(creds != nil) - #expect(creds?.expiresAt != nil) - } + email: "user@example.com") + #expect(AntigravityOAuthCredentialsStore.save(creds, accountLabel: "User@Example.com")) - @Test("Parses JSON with expiresAt as milliseconds") - func test_parseJSONWithExpiresAtMillis() { - let json = """ - {"apiKey": "ya29.test", "expiresAt": 1735689600000} - """ - let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) - #expect(creds != nil) - #expect(creds?.expiresAt != nil) - } - - @Test("Returns nil for empty input") - func test_parseEmptyInput() { - let creds = AntigravityOAuthCredentialsStore.parseManualToken("") - #expect(creds == nil) - } - - @Test("Returns nil for whitespace-only input") - func test_parseWhitespaceOnlyInput() { - let creds = AntigravityOAuthCredentialsStore.parseManualToken(" \n\t ") - #expect(creds == nil) + let loaded = AntigravityOAuthCredentialsStore.load(accountLabel: "user@example.com") + #expect(loaded?.accessToken == "ya29.test") + #expect(loaded?.refreshToken == "1//refresh") } +} - @Test("Returns nil for invalid token format") - func test_parseInvalidTokenFormat() { - let creds = AntigravityOAuthCredentialsStore.parseManualToken("not-a-valid-token") - #expect(creds == nil) +@Suite +struct AntigravityManualTokenPayloadTests { + @Test + func parsesJSONPayload() { + let token = AntigravityOAuthCredentialsStore.manualTokenValue( + accessToken: "ya29.test", + refreshToken: "1//refresh") + let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) + #expect(payload?.accessToken == "ya29.test") + #expect(payload?.refreshToken == "1//refresh") } - @Test("Returns nil for JSON without apiKey or accessToken") - func test_parseJSONWithoutToken() { - let json = """ - {"email": "test@example.com"} - """ - let creds = AntigravityOAuthCredentialsStore.parseManualToken(json) - #expect(creds == nil) + @Test + func parsesLegacyPayload() { + let token = "\(AntigravityOAuthCredentialsStore.manualTokenPrefix)ya29.test" + let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) + #expect(payload?.accessToken == "ya29.test") + #expect(payload?.refreshToken == nil) } } -@Suite("AntigravityUsageSource") +@Suite struct AntigravityUsageSourceTests { - @Test("UsageSource has expected cases") - func test_usageSourceCases() { - let allCases = AntigravityUsageSource.allCases - #expect(allCases.contains(.auto)) - #expect(allCases.contains(.authorized)) - #expect(allCases.contains(.local)) - #expect(allCases.count == 3) + @Test + func parsesRawValue() { + #expect(AntigravityUsageSource(rawValue: "cli") == .local) + #expect(AntigravityUsageSource(rawValue: "authorized") == .authorized) } +} - @Test("UsageSource rawValue matches expected strings") - func test_usageSourceRawValues() { - #expect(AntigravityUsageSource.auto.rawValue == "auto") - #expect(AntigravityUsageSource.authorized.rawValue == "authorized") - #expect(AntigravityUsageSource.local.rawValue == "local") - } +@Suite(.serialized) +struct AntigravityAuthorizedFetchStrategyTests { + private func makeContext(usageSource: AntigravityUsageSource, accountLabel: String?) -> ProviderFetchContext { + let browserDetection = BrowserDetection(cacheTTL: 0) + let settings = ProviderSettingsSnapshot( + debugMenuEnabled: false, + debugKeepCLISessionsAlive: false, + codex: nil, + claude: nil, + cursor: nil, + opencode: nil, + factory: nil, + minimax: nil, + zai: nil, + copilot: nil, + kimi: nil, + augment: nil, + amp: nil, + jetbrains: nil, + antigravity: .init(usageSource: usageSource, accountLabel: accountLabel)) + return ProviderFetchContext( + runtime: .cli, + sourceMode: .auto, + includeCredits: false, + webTimeout: 60, + webDebugDumpHTML: false, + verbose: false, + env: [:], + settings: settings, + fetcher: UsageFetcher(), + claudeFetcher: ClaudeUsageFetcher(browserDetection: browserDetection), + browserDetection: browserDetection) + } + + @Test + func unavailableWithoutAccountLabel() async { + let strategy = AntigravityAuthorizedFetchStrategy() + let context = self.makeContext(usageSource: .auto, accountLabel: nil) + let available = await strategy.isAvailable(context) + #expect(!available) + } + + @Test + func availableWithStoredCredentials() async { + KeychainCacheStore.setTestStoreForTesting(true) + defer { KeychainCacheStore.setTestStoreForTesting(false) } + let previousKeychainDisabled = KeychainAccessGate.isDisabled + KeychainAccessGate.isDisabled = false + defer { KeychainAccessGate.isDisabled = previousKeychainDisabled } + AntigravityOAuthCredentialsStore.invalidateCache() - @Test("UsageSource can be initialized from rawValue") - func test_usageSourceFromRawValue() { - #expect(AntigravityUsageSource(rawValue: "auto") == .auto) - #expect(AntigravityUsageSource(rawValue: "authorized") == .authorized) - #expect(AntigravityUsageSource(rawValue: "local") == .local) - #expect(AntigravityUsageSource(rawValue: "invalid") == nil) + let creds = AntigravityOAuthCredentials( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: Date().addingTimeInterval(3600), + email: "user@example.com") + _ = AntigravityOAuthCredentialsStore.save(creds, accountLabel: "user@example.com") + + let strategy = AntigravityAuthorizedFetchStrategy() + let context = self.makeContext(usageSource: .auto, accountLabel: "user@example.com") + let available = await strategy.isAvailable(context) + #expect(available) + } + + @Test + func fallbackInAutoMode() { + let strategy = AntigravityAuthorizedFetchStrategy() + let context = self.makeContext(usageSource: .auto, accountLabel: "user@example.com") + let shouldFallback = strategy.shouldFallback( + on: AntigravityOAuthCredentialsError.networkError("test"), + context: context) + #expect(shouldFallback) + } + + @Test + func noFallbackOnNetworkErrorInOAuthMode() { + let strategy = AntigravityAuthorizedFetchStrategy() + let context = self.makeContext(usageSource: .authorized, accountLabel: "user@example.com") + let shouldFallback = strategy.shouldFallback( + on: AntigravityOAuthCredentialsError.networkError("test"), + context: context) + #expect(!shouldFallback) } } diff --git a/docs/antigravity.md b/docs/antigravity.md index 02baf1e9..0286bdcd 100644 --- a/docs/antigravity.md +++ b/docs/antigravity.md @@ -9,45 +9,28 @@ read_when: # Antigravity provider -Antigravity supports both OAuth-authorized API access and local language server probing. +Antigravity supports OAuth-authorized Cloud Code quota and local language server probing. ## Usage source modes -- **Auto** (default): Try authorized credentials first, fallback to local server -- **Authorized**: Use OAuth credentials to fetch quota from Cloud Code API -- **Local**: Use local Antigravity language server only +- **Auto** (default): OAuth/manual first, fallback to local server +- **OAuth**: OAuth/manual only +- **Local**: local Antigravity language server only -## Credential acquisition +## OAuth credentials -### Fallback order (Auto mode) +- **Keychain**: stored from the OAuth browser flow +- **Manual tokens**: stored in token accounts with `manual:` prefix (access `ya29.` + optional refresh `1//`) +- **Local import**: `~/Library/Application Support/Antigravity/User/globalStorage/state.vscdb` + - refresh token: `jetskiStateSync.agentManagerInitState` (protobuf field 6) + - access token/email: `antigravityAuthStatus` JSON (`apiKey`, `email`) +- OAuth callback server listens on `http://127.0.0.1:11451+` -1. **Keychain OAuth credentials** - Previously saved credentials from OAuth flow or import -2. **Manual token** - Token pasted in settings (refresh token or access token) -3. **Local DB import** - Import from `state.vscdb` (Antigravity app's local storage) -4. **Local server** - Direct probe to running Antigravity language server +## Cloud Code API endpoints -### OAuth browser flow - -Opens Google OAuth in browser with local callback server (port 11451+). Returns refresh token for long-term storage. - -### Local DB import (`state.vscdb`) - -Path: `~/Library/Application Support/Antigravity/User/globalStorage/state.vscdb` - -- **Refresh token**: `jetskiStateSync.agentManagerInitState` (protobuf field 6) -- **Access token + email**: `antigravityAuthStatus` JSON (`apiKey`, `email` fields) - -### Manual token import - -Accepts: -- Raw access token (`ya29.xxx`) -- Raw refresh token (`1//xxx`) -- JSON payload: `{"apiKey": "...", "email": "...", "refreshToken": "..."}` - -### Cloud Code API endpoints - -- Primary: `https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` -- Fallback: `https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` +- `POST https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` +- `POST https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` +- `POST https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist` ## Local server data sources + fallback order @@ -127,6 +110,7 @@ Accepts: ### App Integration - `Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift` - `Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift` +- `Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift` ### Tests - `Tests/CodexBarTests/AntigravityOAuthTests.swift` From 54e3e10d3b5224f18eea7dacbfb200befe74db97 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 17:42:25 +0300 Subject: [PATCH 03/19] feat: add SQLite3 installation step to CI workflow --- .github/workflows/ci.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3338bcdf..e829d02c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -61,6 +61,9 @@ jobs: uname -a uname -m + - name: Install SQLite3 + run: sudo apt-get update && sudo apt-get install -y libsqlite3-dev + - name: Setup Swift 6.2.1 uses: swift-actions/setup-swift@v3 with: From 48cc9e7b4a758d56abd076a493c9d5c12ad14bda Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 17:50:28 +0300 Subject: [PATCH 04/19] feat: add conditional compilation for macOS in AntigravityLocalImporter --- .../AntigravityOAuth/AntigravityLocalImporter.swift | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift index 7bcf4ecf..8131d568 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift @@ -1,4 +1,6 @@ import Foundation + +#if os(macOS) import SQLite3 public enum AntigravityLocalImporter { @@ -213,3 +215,4 @@ public enum AntigravityLocalImporter { throw AntigravityOAuthCredentialsError.decodeFailed("Incomplete varint") } } +#endif From 6f7c4814c65872563a446dde1d114abed2b15609 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 17:51:36 +0300 Subject: [PATCH 05/19] feat: add documentation comments for AntigravityOAuthConfig client credentials --- .../AntigravityOAuth/AntigravityOAuthCredentials.swift | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift index 541492de..8ea88c87 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift @@ -44,6 +44,10 @@ public struct AntigravityOAuthCredentials: Sendable, Codable { } public enum AntigravityOAuthConfig { + // Public OAuth client credentials for Antigravity service. + // These are meant to be embedded in client applications and are distributed with the app. + // They are visible in OAuth flows to users anyway; the actual security comes from + // the refresh/access tokens stored in Keychain, not these public client identifiers. public static let clientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" public static let clientSecret = "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" public static let scopes = [ From efc1f2b7ecc5ce1fa12b0955e9dee5a5ee8294c9 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 17:53:20 +0300 Subject: [PATCH 06/19] feat: add GitGuardian configuration for Antigravity OAuth client credentials --- .gitguardian.yml | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 .gitguardian.yml diff --git a/.gitguardian.yml b/.gitguardian.yml new file mode 100644 index 00000000..ea53df08 --- /dev/null +++ b/.gitguardian.yml @@ -0,0 +1,14 @@ +version: 2 + +# GitGuardian configuration +# See: https://docs.gitguardian.com/ggshield-docs/configuration + +secret: + ignored_matches: + # Antigravity Google OAuth public client credentials + # These are public OAuth client IDs meant to be embedded in the app + # They are distributed with the client and visible in OAuth flows + - match: 1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com + name: Antigravity OAuth Client ID + - match: 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 + name: Antigravity OAuth Client Secret From 8853c754dc899c94c2672626c78b6b066ebcf154 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 17:58:09 +0300 Subject: [PATCH 07/19] fix: correct socket type casting in CallbackServer --- .../Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift index c8557ee8..cea298bb 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -184,7 +184,7 @@ private final class CallbackServer: @unchecked Sendable { } private func start(port: Int) throws { - self.serverSocket = socket(AF_INET, SOCK_STREAM, 0) + self.serverSocket = socket(AF_INET, Int32(SOCK_STREAM), 0) guard self.serverSocket >= 0 else { throw AntigravityOAuthCredentialsError.networkError("Failed to create socket") } From 6efd60575e4d5035a4f1fcc25c980fbf7ab104ab Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 18:05:13 +0300 Subject: [PATCH 08/19] fix: correct socket type casting in CallbackServer --- .../Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift index cea298bb..1673b24d 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -184,7 +184,7 @@ private final class CallbackServer: @unchecked Sendable { } private func start(port: Int) throws { - self.serverSocket = socket(AF_INET, Int32(SOCK_STREAM), 0) + self.serverSocket = socket(AF_INET, CInt(SOCK_STREAM), 0) guard self.serverSocket >= 0 else { throw AntigravityOAuthCredentialsError.networkError("Failed to create socket") } From a9686424cacdd40357d675bd7e44adc24d27242e Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 18:43:40 +0300 Subject: [PATCH 09/19] fix: adjust socket creation for Linux compatibility in CallbackServer --- .../Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift index 1673b24d..8d128b27 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -184,7 +184,11 @@ private final class CallbackServer: @unchecked Sendable { } private func start(port: Int) throws { - self.serverSocket = socket(AF_INET, CInt(SOCK_STREAM), 0) + #if os(Linux) + self.serverSocket = socket(AF_INET, Int32(SOCK_STREAM.rawValue), 0) + #else + self.serverSocket = socket(AF_INET, SOCK_STREAM, 0) + #endif guard self.serverSocket >= 0 else { throw AntigravityOAuthCredentialsError.networkError("Failed to create socket") } From e35005a12bf78aebb521cf3828c1bb8f9caa2a8d Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 23:45:29 +0300 Subject: [PATCH 10/19] chore: clean up unused parameter pattern --- .../Antigravity/AntigravityProviderImplementation.swift | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index ecc34318..9e003e4c 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -72,9 +72,7 @@ struct AntigravityProviderImplementation: ProviderImplementation { } @MainActor - func tokenAccountsVisibility(context: ProviderSettingsContext, support: TokenAccountSupport) -> Bool { - _ = context - _ = support - return true + func tokenAccountsVisibility(context _: ProviderSettingsContext, support _: TokenAccountSupport) -> Bool { + true } } From 9c7ec4d567d4ab488cadaf2c3cfbdf419fd05999 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Wed, 28 Jan 2026 23:59:50 +0300 Subject: [PATCH 11/19] feat: enhance presentation method to handle unrecognized source labels --- .../AntigravityProviderImplementation.swift | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index 9e003e4c..5ce5a8bb 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -8,6 +8,21 @@ struct AntigravityProviderImplementation: ProviderImplementation { let id: UsageProvider = .antigravity let supportsLoginFlow: Bool = true + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { context in + let sourceLabel = context.store.sourceLabel(for: .antigravity) + switch sourceLabel.lowercased() { + case "oauth", "manual": + return "oauth" + case "local server": + return "local" + default: + return "not detected" + } + } + } + func detectVersion(context _: ProviderVersionContext) async -> String? { await AntigravityStatusProbe.detectVersion() } From 4f76db344659388aa3cb8beabf01028977035d11 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 00:02:51 +0300 Subject: [PATCH 12/19] fix: simplify success alert message in AntigravityLoginFlow --- .../Providers/Antigravity/AntigravityLoginFlow.swift | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift index 3cd020ba..29c8aa48 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift @@ -62,11 +62,7 @@ struct AntigravityLoginFlow { let success = NSAlert() success.messageText = "Authorization Successful" - var info: [String] = [] - info.append("Signed in as \(accountLabel).") - if !info.isEmpty { - success.informativeText = info.joined(separator: " ") - } + success.informativeText = "Signed in as \(accountLabel)." success.runModal() return true case let .failure(error): From 44a96524db3fed0f4dae8f0603949f8d62e5f95a Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 00:18:12 +0300 Subject: [PATCH 13/19] style: remove unnecessary blank line in ProvidersPane view --- Sources/CodexBar/PreferencesProvidersPane.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index de858710..b8e26f2a 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -303,7 +303,7 @@ struct ProvidersPane: View { } } } - + @MainActor private func presentFullDiskAccessAlert() { let alert = NSAlert() From e3f86ac5823fe03bd4cceddc76dd74a043f790fb Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 01:45:26 +0300 Subject: [PATCH 14/19] feat: integrate Swift Protobuf for Antigravity OAuth state management and enhance logging --- Package.resolved | 11 +- Package.swift | 3 + .../CodexBar/PreferencesProvidersPane.swift | 14 +- .../Antigravity/AntigravityLoginFlow.swift | 10 + .../AntigravityAuthorizedFetchStrategy.swift | 16 ++ .../AntigravityLocalImporter.swift | 129 +++++------ .../AntigravityOAuthFlow.swift | 13 +- .../antigravity_state.pb.swift | 202 ++++++++++++++++++ .../AntigravityOAuth/antigravity_state.proto | 30 +++ docs/antigravity.md | 5 +- 10 files changed, 358 insertions(+), 75 deletions(-) create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift create mode 100644 Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.proto diff --git a/Package.resolved b/Package.resolved index f84c0c21..d4c83370 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "74bd6f3ab6e0b0cb0c2cddb00f2167c2ab0a1c00cd54ffc1a2899c7ef8c56367", + "originHash" : "7cde0aee9a244cdd6b664c0e4f76d9a9e7cbddbc2e5bcee2c71bb39243cffb66", "pins" : [ { "identity" : "commander", @@ -46,6 +46,15 @@ "version" : "1.9.1" } }, + { + "identity" : "swift-protobuf", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-protobuf.git", + "state" : { + "revision" : "c169a5744230951031770e27e475ff6eefe51f9d", + "version" : "1.33.3" + } + }, { "identity" : "swift-syntax", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index 83cbfca6..52d5559e 100644 --- a/Package.swift +++ b/Package.swift @@ -22,6 +22,7 @@ let package = Package( .package(url: "https://github.com/apple/swift-log", from: "1.9.1"), .package(url: "https://github.com/apple/swift-syntax", from: "600.0.1"), .package(url: "https://github.com/sindresorhus/KeyboardShortcuts", from: "2.4.0"), + .package(url: "https://github.com/apple/swift-protobuf.git", from: "1.21.0"), sweetCookieKitDependency, ], targets: { @@ -32,7 +33,9 @@ let package = Package( "CodexBarMacroSupport", .product(name: "Logging", package: "swift-log"), .product(name: "SweetCookieKit", package: "SweetCookieKit"), + .product(name: "SwiftProtobuf", package: "swift-protobuf"), ], + exclude: ["Providers/Antigravity/AntigravityOAuth/antigravity_state.proto"], swiftSettings: [ .enableUpcomingFeature("StrictConcurrency"), ]), diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index b8e26f2a..6c16be96 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -178,7 +178,7 @@ struct ProvidersPane: View { store: self.store) } : nil - let importAction: ProviderSettingsTokenAccountsDescriptor.ImportAction? = (isAntigravity && !keychainEnabled) + let importAction: ProviderSettingsTokenAccountsDescriptor.ImportAction? = isAntigravity ? ProviderSettingsTokenAccountsDescriptor.ImportAction( title: "Import from Antigravity app", action: { await self.importAntigravityCredentials() }) @@ -248,9 +248,15 @@ struct ProvidersPane: View { @MainActor private func importAntigravityCredentials() async { + let log = CodexBarLog.logger(LogCategories.antigravity) + log.debug("Starting credentials import from Antigravity DB") + do { let credentials = try await AntigravityLocalImporter.importCredentials() + log.debug("Import successful - email: \(credentials.email ?? "none"), hasAccessToken: \(credentials.hasAccessToken), hasRefreshToken: \(credentials.hasRefreshToken)") + guard let accessToken = credentials.accessToken, !accessToken.isEmpty else { + log.debug("Import failed: no access token found") self.presentAlert( title: "Import Failed", message: "No access token found in Antigravity database.") @@ -258,16 +264,20 @@ struct ProvidersPane: View { } let label = credentials.email ?? credentials.name ?? "Imported Account" + log.debug("Creating manual token account with label: \(label)") + guard let account = self.settings.addManualAntigravityTokenAccount( label: label, accessToken: accessToken, refreshToken: credentials.refreshToken) else { + log.debug("Import failed: unable to save imported credentials") self.presentAlert( title: "Import Failed", message: "Unable to save imported credentials.") return } + log.debug("Account created successfully: \(account.label)") await self.store.refreshProvider(.antigravity, allowDisabled: true) let alert = NSAlert() @@ -276,6 +286,7 @@ struct ProvidersPane: View { alert.runModal() } catch let error as AntigravityOAuthCredentialsError { + log.debug("Import failed with AntigravityOAuthCredentialsError: \(error)") switch error { case .notFound: if AntigravityLocalImporter.isAvailable() { @@ -293,6 +304,7 @@ struct ProvidersPane: View { message: error.localizedDescription) } } catch { + log.debug("Import failed with error: \(error)") let nsError = error as NSError if nsError.domain == NSPOSIXErrorDomain && nsError.code == 1 { self.presentFullDiskAccessAlert() diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift index 29c8aa48..6012e352 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift @@ -3,7 +3,12 @@ import CodexBarCore @MainActor struct AntigravityLoginFlow { + private static let log = CodexBarLog.logger(LogCategories.antigravity) + static func runOAuthFlow(settings: SettingsStore, store: UsageStore? = nil) async -> Bool { + Self.log.debug("Starting Antigravity OAuth login flow") + Self.log.debug("Keychain access disabled: \(KeychainAccessGate.isDisabled)") + if KeychainAccessGate.isDisabled { Self.presentAlert( title: "Authorization Unavailable", @@ -77,16 +82,21 @@ struct AntigravityLoginFlow { settings: SettingsStore) -> String? { guard let accountLabel = Self.resolveAccountLabel(credentials: credentials, settings: settings) else { + Self.log.debug("Failed to resolve account label") return nil } + Self.log.debug("Persisting credentials for account: \(accountLabel)") guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: accountLabel) else { + Self.log.debug("Failed to save credentials to Keychain") return nil } + Self.log.debug("Saved credentials to Keychain successfully") _ = settings.upsertAntigravityTokenAccount(label: accountLabel) settings.setProviderEnabled( provider: .antigravity, metadata: ProviderRegistry.shared.metadata[.antigravity]!, enabled: true) + Self.log.debug("Provider enabled and token account created") return accountLabel } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift index 98b00f1c..4e76fa49 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift @@ -10,30 +10,44 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { public func isAvailable(_ context: ProviderFetchContext) async -> Bool { guard let accountLabel = context.settings?.antigravity?.accountLabel else { + Self.log.debug("Authorized strategy not available: no account label") return false } + Self.log.debug("Checking authorized strategy availability for account: \(accountLabel)") + if let manualCredentials = self.loadManualCredentials(accountLabel: accountLabel, context: context) { + Self.log.debug("Manual credentials found") return !manualCredentials.accessToken.isEmpty } guard let credentials = AntigravityOAuthCredentialsStore.load(accountLabel: accountLabel) else { + Self.log.debug("Keychain credentials not found") return false } + + Self.log.debug("Keychain credentials found - hasAccessToken: \(!credentials.accessToken.isEmpty), isRefreshable: \(credentials.isRefreshable)") + if !credentials.accessToken.isEmpty { return true } return credentials.isRefreshable } public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + Self.log.debug("Fetching with authorized strategy") + let resolved = try await self.resolveCredentials(context: context) let accountLabel = resolved.accountLabel let credentials = resolved.credentials let sourceLabel = resolved.sourceLabel + Self.log.debug("Resolved credentials - source: \(sourceLabel), needsRefresh: \(credentials.needsRefresh), isRefreshable: \(credentials.isRefreshable)") + var refreshedCredentials: AntigravityOAuthCredentials? if credentials.needsRefresh || (credentials.accessToken.isEmpty && credentials.isRefreshable) { + Self.log.debug("Credentials need refresh, refreshing token...") let refreshed = try await self.refreshCredentials(credentials) refreshedCredentials = refreshed + Self.log.debug("Token refresh successful") if KeychainAccessGate.isDisabled { context.onCredentialsRefreshed?(.antigravity, accountLabel, refreshed.accessToken) } else { @@ -43,6 +57,8 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { let activeCredentials = refreshedCredentials ?? credentials let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: activeCredentials.accessToken) + Self.log.debug("Successfully fetched quota from Cloud Code API") + let snapshot = AntigravityStatusSnapshot( modelQuotas: quota.models, accountEmail: activeCredentials.email ?? quota.email, diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift index 8131d568..0582a5c2 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift @@ -1,4 +1,5 @@ import Foundation +import SwiftProtobuf #if os(macOS) import SQLite3 @@ -21,6 +22,8 @@ public enum AntigravityLocalImporter { } } + private static let log = CodexBarLog.logger(LogCategories.antigravity) + public static func stateDbPath() -> URL { let home = FileManager.default.homeDirectoryForCurrentUser return home @@ -33,32 +36,53 @@ public enum AntigravityLocalImporter { } public static func importCredentials() async throws -> LocalCredentialInfo { + Self.log.debug("Starting Antigravity DB import") + let dbPath = self.stateDbPath() + Self.log.debug("Database path: \(dbPath.path)") + guard FileManager.default.fileExists(atPath: dbPath.path) else { + Self.log.debug("Database file not found at path") throw AntigravityOAuthCredentialsError.notFound } var refreshToken: String? - if let protoInfo = try? self.readProtoTokenInfo(dbPath: dbPath) { + var accessToken: String? + + do { + let protoInfo = try self.readProtoTokenInfo(dbPath: dbPath) refreshToken = protoInfo.refreshToken + accessToken = protoInfo.accessToken + Self.log.debug("Extracted OAuth token info - access_token present: \(accessToken != nil && !accessToken!.isEmpty), refresh_token present: \(refreshToken != nil && !refreshToken!.isEmpty)") + } catch { + Self.log.debug("Failed to read proto token info: \(error)") } if let authStatus = try? self.readAuthStatus(dbPath: dbPath) { + Self.log.debug("Read auth status - email: \(authStatus.email ?? "none"), apiKey present: \(authStatus.apiKey != nil && !authStatus.apiKey!.isEmpty)") + + // Prefer accessToken from proto, fallback to apiKey from auth status + let finalAccessToken = accessToken ?? authStatus.apiKey + + Self.log.debug("Import result - email: \(authStatus.email ?? "none"), hasAccessToken: \(finalAccessToken != nil && !finalAccessToken!.isEmpty), hasRefreshToken: \(refreshToken != nil && !refreshToken!.isEmpty)") + return LocalCredentialInfo( - accessToken: authStatus.apiKey, + accessToken: finalAccessToken, refreshToken: refreshToken, email: authStatus.email, name: authStatus.name) } if let refreshToken, !refreshToken.isEmpty { + Self.log.debug("Using refresh token only (no auth status found)") return LocalCredentialInfo( - accessToken: nil, + accessToken: accessToken, refreshToken: refreshToken, email: nil, name: nil) } + Self.log.debug("No credentials found in database") throw AntigravityOAuthCredentialsError.notFound } @@ -80,6 +104,7 @@ public enum AntigravityLocalImporter { } private static func readAuthStatus(dbPath: URL) throws -> AuthStatus { + Self.log.debug("Reading antigravityAuthStatus from DB") let json = try self.readStateValue(dbPath: dbPath, key: "antigravityAuthStatus") guard let data = json.data(using: .utf8), let dict = try JSONSerialization.jsonObject(with: data) as? [String: Any] @@ -95,10 +120,15 @@ public enum AntigravityLocalImporter { } private static func readProtoTokenInfo(dbPath: URL) throws -> ProtoTokenInfo { + Self.log.debug("Reading jetskiStateSync.agentManagerInitState from DB") let base64 = try self.readStateValue(dbPath: dbPath, key: "jetskiStateSync.agentManagerInitState") - guard let data = Data(base64Encoded: base64) else { + Self.log.debug("Read base64 value, length: \(base64.count)") + + guard let data = Data(base64Encoded: base64.trimmingCharacters(in: .whitespacesAndNewlines)) else { throw AntigravityOAuthCredentialsError.decodeFailed("Invalid base64 in agentManagerInitState") } + Self.log.debug("Decoded base64, data length: \(data.count)") + return try self.parseProtoTokenInfo(data: data) } @@ -141,78 +171,35 @@ public enum AntigravityLocalImporter { } private static func parseProtoTokenInfo(data: Data) throws -> ProtoTokenInfo { - var accessToken: String? - var refreshToken: String? - var tokenType: String? - var expirySeconds: Int? - - var offset = 0 - while offset < data.count { - let (fieldTag, newOffset) = try self.readVarint(data: data, offset: offset) - offset = newOffset - - let fieldNumber = fieldTag >> 3 - let wireType = fieldTag & 0x07 - - switch wireType { - case 0: - let (value, nextOffset) = try self.readVarint(data: data, offset: offset) - offset = nextOffset - if fieldNumber == 4 { - expirySeconds = value - } - case 2: - let (length, lengthOffset) = try self.readVarint(data: data, offset: offset) - offset = lengthOffset - let endOffset = offset + length - guard endOffset <= data.count else { - throw AntigravityOAuthCredentialsError.decodeFailed("Invalid protobuf length") - } - let stringData = data[offset.. (Int, Int) { - var result = 0 - var shift = 0 - var pos = offset - - while pos < data.count { - let byte = Int(data[pos]) - result |= (byte & 0x7F) << shift - pos += 1 - if (byte & 0x80) == 0 { - return (result, pos) + guard state.hasOauthToken else { + Self.log.debug("No oauth_token field (field 6) found in protobuf") + throw AntigravityOAuthCredentialsError.decodeFailed("No oauth_token field found") } - shift += 7 - if shift > 63 { - throw AntigravityOAuthCredentialsError.decodeFailed("Varint too long") + + let oauthToken = state.oauthToken + Self.log.debug("Found OAuthTokenInfo - access_token length: \(oauthToken.accessToken.count), refresh_token length: \(oauthToken.refreshToken.count)") + + var expirySeconds: Int? + if oauthToken.hasExpiry { + expirySeconds = Int(oauthToken.expiry.seconds) + Self.log.debug("Token expiry: \(expirySeconds!) seconds since epoch") } - } - throw AntigravityOAuthCredentialsError.decodeFailed("Incomplete varint") + return ProtoTokenInfo( + accessToken: oauthToken.accessToken.isEmpty ? nil : oauthToken.accessToken, + refreshToken: oauthToken.refreshToken.isEmpty ? nil : oauthToken.refreshToken, + tokenType: oauthToken.tokenType.isEmpty ? nil : oauthToken.tokenType, + expirySeconds: expirySeconds) + } catch { + Self.log.debug("Protobuf parsing failed: \(error)") + throw AntigravityOAuthCredentialsError.decodeFailed("Failed to parse protobuf: \(error)") + } } } #endif diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift index 8d128b27..1b8fe1fe 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -17,13 +17,17 @@ public actor AntigravityOAuthFlow { public init() {} public func startAuthorization() async throws -> AntigravityOAuthCredentials { + Self.log.debug("Starting OAuth authorization flow") + let port = try await self.startCallbackServer() let redirectUri = "http://\(AntigravityOAuthConfig.callbackHost):\(port)" let state = self.generateState() self.pendingState = state let authURL = self.buildAuthURL(redirectUri: redirectUri, state: state) - + + Self.log.debug("Requesting scopes: \(AntigravityOAuthConfig.scopes.joined(separator: ", "))") + Self.log.debug("Using access_type=offline and prompt=consent to ensure refresh token") Self.log.info("Opening Antigravity OAuth authorization URL") #if os(macOS) if let url = URL(string: authURL) { @@ -159,6 +163,13 @@ public actor AntigravityOAuthFlow { let expiresAt = Date(timeIntervalSinceNow: TimeInterval(expiresIn)) let email = try? await AntigravityTokenRefresher.fetchUserEmail(accessToken: accessToken) + Self.log.debug("Received OAuth response - access_token length: \(accessToken.count)") + Self.log.debug("Received OAuth response - refresh_token length: \(refreshToken.count)") + Self.log.debug("Token expires in: \(expiresIn) seconds") + if let email = email { + Self.log.debug("Resolved account email: \(email)") + } + Self.log.info("Antigravity OAuth authorization successful", metadata: [ "email": email ?? "unknown", ]) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift new file mode 100644 index 00000000..c41d5647 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift @@ -0,0 +1,202 @@ +// DO NOT EDIT. +// swift-format-ignore-file +// swiftlint:disable all +// +// Generated by the Swift generator plugin for the protocol buffer compiler. +// Source: antigravity_state.proto +// +// For information on using the generated types, please see the documentation: +// https://github.com/apple/swift-protobuf/ + +import SwiftProtobuf + +// If the compiler emits an error on this type, it is because this file +// was generated by a version of the `protoc` Swift plug-in that is +// incompatible with the version of SwiftProtobuf to which you are linking. +// Please ensure that you are building against the same version of the API +// that was used to generate this file. +fileprivate struct _GeneratedWithProtocGenSwiftVersion: SwiftProtobuf.ProtobufAPIVersionCheck { + struct _2: SwiftProtobuf.ProtobufAPIVersion_2 {} + typealias Version = _2 +} + +struct AgentManagerInitState: Sendable { + // SwiftProtobuf.Message conformance is added in an extension below. See the + // `Message` and `Message+*Additions` files in the SwiftProtobuf library for + // methods supported on all messages. + + /// OAuth token information is stored in field 6 + var oauthToken: OAuthTokenInfo { + get {return _oauthToken ?? OAuthTokenInfo()} + set {_oauthToken = newValue} + } + /// Returns true if `oauthToken` has been explicitly set. + var hasOauthToken: Bool {return self._oauthToken != nil} + /// Clears the value of `oauthToken`. Subsequent reads from it will return its default value. + mutating func clearOauthToken() {self._oauthToken = nil} + + var unknownFields = SwiftProtobuf.UnknownStorage() + + init() {} + + fileprivate var _oauthToken: OAuthTokenInfo? = nil +} + +struct OAuthTokenInfo: Sendable { + // SwiftProtobuf.Message conformance is added in an extension below. See the + // `Message` and `Message+*Additions` files in the SwiftProtobuf library for + // methods supported on all messages. + + /// Field 1: Access token (short-lived API token) + var accessToken: String = String() + + /// Field 2: Token type (typically "Bearer") + var tokenType: String = String() + + /// Field 3: Refresh token (long-lived, used to get new access tokens) + var refreshToken: String = String() + + /// Field 4: Token expiry timestamp + var expiry: Timestamp { + get {return _expiry ?? Timestamp()} + set {_expiry = newValue} + } + /// Returns true if `expiry` has been explicitly set. + var hasExpiry: Bool {return self._expiry != nil} + /// Clears the value of `expiry`. Subsequent reads from it will return its default value. + mutating func clearExpiry() {self._expiry = nil} + + var unknownFields = SwiftProtobuf.UnknownStorage() + + init() {} + + fileprivate var _expiry: Timestamp? = nil +} + +struct Timestamp: Sendable { + // SwiftProtobuf.Message conformance is added in an extension below. See the + // `Message` and `Message+*Additions` files in the SwiftProtobuf library for + // methods supported on all messages. + + /// Unix timestamp in seconds + var seconds: Int64 = 0 + + var unknownFields = SwiftProtobuf.UnknownStorage() + + init() {} +} + +// MARK: - Code below here is support for the SwiftProtobuf runtime. + +extension AgentManagerInitState: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding { + static let protoMessageName: String = "AgentManagerInitState" + static let _protobuf_nameMap = SwiftProtobuf._NameMap(bytecode: "\0\u{4}\u{6}oauth_token\0") + + mutating func decodeMessage(decoder: inout D) throws { + while let fieldNumber = try decoder.nextFieldNumber() { + // The use of inline closures is to circumvent an issue where the compiler + // allocates stack space for every case branch when no optimizations are + // enabled. https://github.com/apple/swift-protobuf/issues/1034 + switch fieldNumber { + case 6: try { try decoder.decodeSingularMessageField(value: &self._oauthToken) }() + default: break + } + } + } + + func traverse(visitor: inout V) throws { + // The use of inline closures is to circumvent an issue where the compiler + // allocates stack space for every if/case branch local when no optimizations + // are enabled. https://github.com/apple/swift-protobuf/issues/1034 and + // https://github.com/apple/swift-protobuf/issues/1182 + try { if let v = self._oauthToken { + try visitor.visitSingularMessageField(value: v, fieldNumber: 6) + } }() + try unknownFields.traverse(visitor: &visitor) + } + + static func ==(lhs: AgentManagerInitState, rhs: AgentManagerInitState) -> Bool { + if lhs._oauthToken != rhs._oauthToken {return false} + if lhs.unknownFields != rhs.unknownFields {return false} + return true + } +} + +extension OAuthTokenInfo: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding { + static let protoMessageName: String = "OAuthTokenInfo" + static let _protobuf_nameMap = SwiftProtobuf._NameMap(bytecode: "\0\u{3}access_token\0\u{3}token_type\0\u{3}refresh_token\0\u{1}expiry\0") + + mutating func decodeMessage(decoder: inout D) throws { + while let fieldNumber = try decoder.nextFieldNumber() { + // The use of inline closures is to circumvent an issue where the compiler + // allocates stack space for every case branch when no optimizations are + // enabled. https://github.com/apple/swift-protobuf/issues/1034 + switch fieldNumber { + case 1: try { try decoder.decodeSingularStringField(value: &self.accessToken) }() + case 2: try { try decoder.decodeSingularStringField(value: &self.tokenType) }() + case 3: try { try decoder.decodeSingularStringField(value: &self.refreshToken) }() + case 4: try { try decoder.decodeSingularMessageField(value: &self._expiry) }() + default: break + } + } + } + + func traverse(visitor: inout V) throws { + // The use of inline closures is to circumvent an issue where the compiler + // allocates stack space for every if/case branch local when no optimizations + // are enabled. https://github.com/apple/swift-protobuf/issues/1034 and + // https://github.com/apple/swift-protobuf/issues/1182 + if !self.accessToken.isEmpty { + try visitor.visitSingularStringField(value: self.accessToken, fieldNumber: 1) + } + if !self.tokenType.isEmpty { + try visitor.visitSingularStringField(value: self.tokenType, fieldNumber: 2) + } + if !self.refreshToken.isEmpty { + try visitor.visitSingularStringField(value: self.refreshToken, fieldNumber: 3) + } + try { if let v = self._expiry { + try visitor.visitSingularMessageField(value: v, fieldNumber: 4) + } }() + try unknownFields.traverse(visitor: &visitor) + } + + static func ==(lhs: OAuthTokenInfo, rhs: OAuthTokenInfo) -> Bool { + if lhs.accessToken != rhs.accessToken {return false} + if lhs.tokenType != rhs.tokenType {return false} + if lhs.refreshToken != rhs.refreshToken {return false} + if lhs._expiry != rhs._expiry {return false} + if lhs.unknownFields != rhs.unknownFields {return false} + return true + } +} + +extension Timestamp: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding { + static let protoMessageName: String = "Timestamp" + static let _protobuf_nameMap = SwiftProtobuf._NameMap(bytecode: "\0\u{1}seconds\0") + + mutating func decodeMessage(decoder: inout D) throws { + while let fieldNumber = try decoder.nextFieldNumber() { + // The use of inline closures is to circumvent an issue where the compiler + // allocates stack space for every case branch when no optimizations are + // enabled. https://github.com/apple/swift-protobuf/issues/1034 + switch fieldNumber { + case 1: try { try decoder.decodeSingularInt64Field(value: &self.seconds) }() + default: break + } + } + } + + func traverse(visitor: inout V) throws { + if self.seconds != 0 { + try visitor.visitSingularInt64Field(value: self.seconds, fieldNumber: 1) + } + try unknownFields.traverse(visitor: &visitor) + } + + static func ==(lhs: Timestamp, rhs: Timestamp) -> Bool { + if lhs.seconds != rhs.seconds {return false} + if lhs.unknownFields != rhs.unknownFields {return false} + return true + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.proto b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.proto new file mode 100644 index 00000000..29adf94f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.proto @@ -0,0 +1,30 @@ +syntax = "proto3"; + +// Antigravity IDE state database protobuf definitions +// These definitions match the structure stored in: +// ~/Library/Application Support/Antigravity/User/globalStorage/state.vscdb +// Key: jetskiStateSync.agentManagerInitState (base64 encoded) + +message AgentManagerInitState { + // OAuth token information is stored in field 6 + OAuthTokenInfo oauth_token = 6; +} + +message OAuthTokenInfo { + // Field 1: Access token (short-lived API token) + string access_token = 1; + + // Field 2: Token type (typically "Bearer") + string token_type = 2; + + // Field 3: Refresh token (long-lived, used to get new access tokens) + string refresh_token = 3; + + // Field 4: Token expiry timestamp + Timestamp expiry = 4; +} + +message Timestamp { + // Unix timestamp in seconds + int64 seconds = 1; +} diff --git a/docs/antigravity.md b/docs/antigravity.md index 0286bdcd..f2bfd190 100644 --- a/docs/antigravity.md +++ b/docs/antigravity.md @@ -22,8 +22,9 @@ Antigravity supports OAuth-authorized Cloud Code quota and local language server - **Keychain**: stored from the OAuth browser flow - **Manual tokens**: stored in token accounts with `manual:` prefix (access `ya29.` + optional refresh `1//`) - **Local import**: `~/Library/Application Support/Antigravity/User/globalStorage/state.vscdb` - - refresh token: `jetskiStateSync.agentManagerInitState` (protobuf field 6) + - refresh token: `jetskiStateSync.agentManagerInitState` (base64 protobuf, field 6 contains nested OAuthTokenInfo) - access token/email: `antigravityAuthStatus` JSON (`apiKey`, `email`) + - Import button always visible; storage adapts to Keychain setting (Keychain when enabled, config.json when disabled) - OAuth callback server listens on `http://127.0.0.1:11451+` ## Cloud Code API endpoints @@ -98,6 +99,8 @@ Antigravity supports OAuth-authorized Cloud Code quota and local language server - `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift` - `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityTokenRefresher.swift` - `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift` +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.proto` (protobuf definition) +- `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift` (generated Swift) - `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift` - `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift` - `Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift` From 0b724af143bc63bd9761691c4ffcb0d79d3d8c50 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 13:45:58 +0300 Subject: [PATCH 15/19] refactor: streamline Antigravity login flow and settings management by conditionally handling Keychain access --- .../PreferencesProviderSettingsRows.swift | 15 +++------- .../CodexBar/PreferencesProvidersPane.swift | 4 +-- .../Antigravity/AntigravityLoginFlow.swift | 30 ++++++++++++------- .../AntigravitySettingsStore.swift | 17 ++++++----- 4 files changed, 35 insertions(+), 31 deletions(-) diff --git a/Sources/CodexBar/PreferencesProviderSettingsRows.swift b/Sources/CodexBar/PreferencesProviderSettingsRows.swift index 4cba7137..c61e6566 100644 --- a/Sources/CodexBar/PreferencesProviderSettingsRows.swift +++ b/Sources/CodexBar/PreferencesProviderSettingsRows.swift @@ -319,18 +319,11 @@ struct ProviderSettingsTokenAccountsRowView: View { .controlSize(.small) } - HStack(spacing: 10) { - Button("Open token file") { - self.descriptor.openConfigFile() - } - .buttonStyle(.link) - .controlSize(.small) - Button("Reload") { - self.descriptor.reloadFromDisk() - } - .buttonStyle(.link) - .controlSize(.small) + Button("Open config file") { + self.descriptor.openConfigFile() } + .buttonStyle(.link) + .controlSize(.small) } } } diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index 6c16be96..417d4135 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -170,8 +170,8 @@ struct ProvidersPane: View { let supportsTwoFieldEntry = isAntigravity let supportsManualEntry = !isAntigravity || !keychainEnabled - let addActionTitle = (isAntigravity && keychainEnabled) ? "Sign in with Google" : nil - let addAction: (() async -> Void)? = (isAntigravity && keychainEnabled) + let addActionTitle = isAntigravity ? "Sign in with Google" : nil + let addAction: (() async -> Void)? = isAntigravity ? { _ = await AntigravityLoginFlow.runOAuthFlow( settings: self.settings, diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift index 6012e352..9b53f08e 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift @@ -9,12 +9,6 @@ struct AntigravityLoginFlow { Self.log.debug("Starting Antigravity OAuth login flow") Self.log.debug("Keychain access disabled: \(KeychainAccessGate.isDisabled)") - if KeychainAccessGate.isDisabled { - Self.presentAlert( - title: "Authorization Unavailable", - message: "Keychain access is disabled. Enable it to store Antigravity OAuth accounts.") - return false - } let flow = AntigravityOAuthFlow() let waitingAlert = NSAlert() @@ -86,12 +80,26 @@ struct AntigravityLoginFlow { return nil } Self.log.debug("Persisting credentials for account: \(accountLabel)") - guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: accountLabel) else { - Self.log.debug("Failed to save credentials to Keychain") - return nil + + if !KeychainAccessGate.isDisabled { + guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: accountLabel) else { + Self.log.debug("Failed to save credentials to Keychain") + return nil + } + Self.log.debug("Saved credentials to Keychain") + _ = settings.upsertAntigravityTokenAccount(label: accountLabel) + } else { + Self.log.debug("Keychain disabled, storing tokens in config") + guard settings.addManualAntigravityTokenAccount( + label: accountLabel, + accessToken: credentials.accessToken, + refreshToken: credentials.refreshToken) != nil + else { + Self.log.debug("Failed to save credentials to config") + return nil + } } - Self.log.debug("Saved credentials to Keychain successfully") - _ = settings.upsertAntigravityTokenAccount(label: accountLabel) + settings.setProviderEnabled( provider: .antigravity, metadata: ProviderRegistry.shared.metadata[.antigravity]!, diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift index 818b39cf..fac6ee37 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift @@ -77,7 +77,9 @@ extension SettingsStore { guard let data = self.tokenAccountsData(for: .antigravity) else { return } guard let removed = data.accounts.first(where: { $0.id == accountID }) else { return } self.removeTokenAccount(provider: .antigravity, accountID: accountID) - AntigravityOAuthCredentialsStore.clear(accountLabel: removed.label) + if !KeychainAccessGate.isDisabled { + AntigravityOAuthCredentialsStore.clear(accountLabel: removed.label) + } } func addManualAntigravityTokenAccount( @@ -87,23 +89,24 @@ extension SettingsStore { ) -> ProviderTokenAccount? { guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return nil } + let tokenValue: String if !KeychainAccessGate.isDisabled { let credentials = AntigravityOAuthCredentials( accessToken: accessToken, refreshToken: refreshToken, expiresAt: nil, email: normalizedLabel, - scopes: [] - ) + scopes: []) guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: normalizedLabel) else { return nil } + tokenValue = normalizedLabel + } else { + tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( + accessToken: accessToken, + refreshToken: refreshToken) } - let tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( - accessToken: accessToken, - refreshToken: refreshToken) - let existing = self.tokenAccountsData(for: .antigravity) var accounts = existing?.accounts ?? [] From accb33064b6a1e2a109e1857fe81698f21a9e856 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 14:41:11 +0300 Subject: [PATCH 16/19] feat: enhance Antigravity OAuth credentials management by adding expiresAt support and updating related methods --- .../CodexBar/PreferencesProvidersPane.swift | 3 +- .../AntigravitySettingsStore.swift | 36 ++++++++++- .../CodexBar/UsageStore+TokenAccounts.swift | 24 ++++---- .../AntigravityAuthorizedFetchStrategy.swift | 61 ++++++++++++++----- .../AntigravityLocalImporter.swift | 29 ++++----- .../AntigravityOAuthCredentials.swift | 54 +++++++++------- .../Providers/ProviderFetchPlan.swift | 6 +- .../CodexBarTests/AntigravityOAuthTests.swift | 29 +++++++-- 8 files changed, 168 insertions(+), 74 deletions(-) diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index 417d4135..9aa7c0d3 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -269,7 +269,8 @@ struct ProvidersPane: View { guard let account = self.settings.addManualAntigravityTokenAccount( label: label, accessToken: accessToken, - refreshToken: credentials.refreshToken) else { + refreshToken: credentials.refreshToken, + expiresAt: credentials.expiresAt) else { log.debug("Import failed: unable to save imported credentials") self.presentAlert( title: "Import Failed", diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift index fac6ee37..8398f929 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift @@ -85,7 +85,8 @@ extension SettingsStore { func addManualAntigravityTokenAccount( label: String, accessToken: String, - refreshToken: String? = nil + refreshToken: String? = nil, + expiresAt: Date? = nil ) -> ProviderTokenAccount? { guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return nil } @@ -94,7 +95,7 @@ extension SettingsStore { let credentials = AntigravityOAuthCredentials( accessToken: accessToken, refreshToken: refreshToken, - expiresAt: nil, + expiresAt: expiresAt, email: normalizedLabel, scopes: []) guard AntigravityOAuthCredentialsStore.save(credentials, accountLabel: normalizedLabel) else { @@ -104,7 +105,8 @@ extension SettingsStore { } else { tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( accessToken: accessToken, - refreshToken: refreshToken) + refreshToken: refreshToken, + expiresAt: expiresAt) } let existing = self.tokenAccountsData(for: .antigravity) @@ -126,6 +128,10 @@ extension SettingsStore { self.updateProviderConfig(provider: .antigravity) { entry in entry.tokenAccounts = updatedData } + self.triggerBackgroundRefreshIfNeeded( + label: label, + refreshToken: refreshToken, + expiresAt: expiresAt) return updated } @@ -142,6 +148,30 @@ extension SettingsStore { self.updateProviderConfig(provider: .antigravity) { entry in entry.tokenAccounts = updatedData } + self.triggerBackgroundRefreshIfNeeded( + label: label, + refreshToken: refreshToken, + expiresAt: expiresAt) return account } + + private func triggerBackgroundRefreshIfNeeded( + label: String, + refreshToken: String?, + expiresAt: Date? + ) { + guard let refreshToken, !refreshToken.isEmpty, expiresAt == nil else { return } + guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return } + + Task { + guard let refreshed = try? await AntigravityTokenRefresher.buildCredentialsFromRefreshToken( + refreshToken: refreshToken, + fallbackEmail: normalizedLabel) else { return } + _ = self.addManualAntigravityTokenAccount( + label: label, + accessToken: refreshed.accessToken, + refreshToken: refreshed.refreshToken, + expiresAt: refreshed.expiresAt) + } + } } diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 128eb54c..b0713d3e 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -88,13 +88,13 @@ extension UsageStore { tokenOverride: override) let verbose = self.settings.isVerboseLoggingEnabled - let onCredentialsRefreshed: (@Sendable (UsageProvider, String, String?) -> Void) = { [weak self] provider, accountLabel, newAccessToken in + let onAntigravityCredentialsRefreshed: (@Sendable (String, AntigravityOAuthCredentials) -> Void) = { [weak self] accountLabel, credentials in Task { @MainActor in - guard let self = self else { return } - self.saveRefreshedCredentialsToConfig(provider: provider, accountLabel: accountLabel, newAccessToken: newAccessToken) + guard let self else { return } + self.saveRefreshedAntigravityCredentials(accountLabel: accountLabel, credentials: credentials) } } - + let context = ProviderFetchContext( runtime: .app, sourceMode: sourceMode, @@ -107,7 +107,7 @@ extension UsageStore { fetcher: self.codexFetcher, claudeFetcher: self.claudeFetcher, browserDetection: self.browserDetection, - onCredentialsRefreshed: onCredentialsRefreshed) + onAntigravityCredentialsRefreshed: onAntigravityCredentialsRefreshed) return await descriptor.fetchOutcome(context: context) } @@ -213,18 +213,16 @@ extension UsageStore { } @MainActor - private func saveRefreshedCredentialsToConfig(provider: UsageProvider, accountLabel: String, newAccessToken: String?) { - guard provider == .antigravity else { return } + private func saveRefreshedAntigravityCredentials(accountLabel: String, credentials: AntigravityOAuthCredentials) { guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(accountLabel) else { return } let tokenAccounts = self.settings.tokenAccountsData(for: .antigravity) guard let account = tokenAccounts?.accounts.first(where: { $0.label.lowercased() == normalizedLabel }) else { return } - guard let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: account.token) else { return } - let accessToken = newAccessToken ?? payload.accessToken let tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( - accessToken: accessToken, - refreshToken: payload.refreshToken) + accessToken: credentials.accessToken, + refreshToken: credentials.refreshToken, + expiresAt: credentials.expiresAt) guard var accounts = tokenAccounts?.accounts else { return } guard let index = accounts.firstIndex(where: { $0.id == account.id }) else { return } @@ -246,9 +244,9 @@ extension UsageStore { entry.tokenAccounts = updatedData } - self.providerLogger.info("Saved refreshed credentials to JSON config", metadata: [ - "provider": "antigravity", + self.providerLogger.info("Saved refreshed Antigravity credentials with expiresAt", metadata: [ "account": normalizedLabel, + "hasExpiresAt": "\(credentials.expiresAt != nil)", ]) } } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift index 4e76fa49..e045e48a 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift @@ -34,40 +34,71 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { Self.log.debug("Fetching with authorized strategy") - + let resolved = try await self.resolveCredentials(context: context) let accountLabel = resolved.accountLabel - let credentials = resolved.credentials + var credentials = resolved.credentials let sourceLabel = resolved.sourceLabel + var didRefresh = false Self.log.debug("Resolved credentials - source: \(sourceLabel), needsRefresh: \(credentials.needsRefresh), isRefreshable: \(credentials.isRefreshable)") - var refreshedCredentials: AntigravityOAuthCredentials? if credentials.needsRefresh || (credentials.accessToken.isEmpty && credentials.isRefreshable) { Self.log.debug("Credentials need refresh, refreshing token...") - let refreshed = try await self.refreshCredentials(credentials) - refreshedCredentials = refreshed + credentials = try await self.refreshAndSave(credentials, accountLabel: accountLabel, context: context) + didRefresh = true Self.log.debug("Token refresh successful") - if KeychainAccessGate.isDisabled { - context.onCredentialsRefreshed?(.antigravity, accountLabel, refreshed.accessToken) - } else { - _ = AntigravityOAuthCredentialsStore.save(refreshed, accountLabel: accountLabel) - } } - let activeCredentials = refreshedCredentials ?? credentials - let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: activeCredentials.accessToken) + do { + return try await self.fetchQuotaAndMakeResult(credentials: credentials, sourceLabel: sourceLabel) + } catch AntigravityOAuthCredentialsError.invalidGrant where credentials.isRefreshable && !didRefresh { + Self.log.info("API returned invalidGrant, attempting refresh-and-retry") + credentials = try await self.refreshAndSave(credentials, accountLabel: accountLabel, context: context) + return try await self.fetchQuotaAndMakeResult(credentials: credentials, sourceLabel: sourceLabel) + } + } + + private func fetchQuotaAndMakeResult( + credentials: AntigravityOAuthCredentials, + sourceLabel: String + ) async throws -> ProviderFetchResult { + let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: credentials.accessToken) Self.log.debug("Successfully fetched quota from Cloud Code API") - + let snapshot = AntigravityStatusSnapshot( modelQuotas: quota.models, - accountEmail: activeCredentials.email ?? quota.email, + accountEmail: credentials.email ?? quota.email, accountPlan: nil) let usage = try snapshot.toUsageSnapshot() return self.makeResult(usage: usage, sourceLabel: sourceLabel) } + private func refreshAndSave( + _ credentials: AntigravityOAuthCredentials, + accountLabel: String, + context: ProviderFetchContext + ) async throws -> AntigravityOAuthCredentials { + do { + let refreshed = try await self.refreshCredentials(credentials) + + if KeychainAccessGate.isDisabled { + context.onAntigravityCredentialsRefreshed?(accountLabel, refreshed) + } else { + _ = AntigravityOAuthCredentialsStore.save(refreshed, accountLabel: accountLabel) + } + + return refreshed + } catch AntigravityOAuthCredentialsError.invalidGrant { + Self.log.warning("Refresh token invalid, clearing keychain credentials") + if !KeychainAccessGate.isDisabled { + AntigravityOAuthCredentialsStore.clear(accountLabel: accountLabel) + } + throw AntigravityOAuthCredentialsError.invalidGrant + } + } + public func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { let usageSource = context.settings?.antigravity?.usageSource ?? .auto @@ -135,7 +166,7 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { return AntigravityOAuthCredentials( accessToken: payload.accessToken, refreshToken: payload.refreshToken, - expiresAt: nil, + expiresAt: payload.expiresAt, email: account.label, scopes: []) } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift index 0582a5c2..80340199 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift @@ -10,6 +10,7 @@ public enum AntigravityLocalImporter { public let refreshToken: String? public let email: String? public let name: String? + public let expiresAt: Date? public var hasAccessToken: Bool { guard let accessToken else { return false } @@ -48,29 +49,28 @@ public enum AntigravityLocalImporter { var refreshToken: String? var accessToken: String? - - do { - let protoInfo = try self.readProtoTokenInfo(dbPath: dbPath) + var expiresAt: Date? + + if let protoInfo = try? self.readProtoTokenInfo(dbPath: dbPath) { refreshToken = protoInfo.refreshToken accessToken = protoInfo.accessToken - Self.log.debug("Extracted OAuth token info - access_token present: \(accessToken != nil && !accessToken!.isEmpty), refresh_token present: \(refreshToken != nil && !refreshToken!.isEmpty)") - } catch { - Self.log.debug("Failed to read proto token info: \(error)") + if let expiry = protoInfo.expirySeconds { + expiresAt = Date(timeIntervalSince1970: TimeInterval(expiry)) + } + Self.log.debug("Extracted OAuth token info - access_token present: \(accessToken?.isEmpty == false), refresh_token present: \(refreshToken?.isEmpty == false)") } if let authStatus = try? self.readAuthStatus(dbPath: dbPath) { - Self.log.debug("Read auth status - email: \(authStatus.email ?? "none"), apiKey present: \(authStatus.apiKey != nil && !authStatus.apiKey!.isEmpty)") - - // Prefer accessToken from proto, fallback to apiKey from auth status + Self.log.debug("Read auth status - email: \(authStatus.email ?? "none"), apiKey present: \(authStatus.apiKey?.isEmpty == false)") let finalAccessToken = accessToken ?? authStatus.apiKey - - Self.log.debug("Import result - email: \(authStatus.email ?? "none"), hasAccessToken: \(finalAccessToken != nil && !finalAccessToken!.isEmpty), hasRefreshToken: \(refreshToken != nil && !refreshToken!.isEmpty)") - + Self.log.debug("Import result - email: \(authStatus.email ?? "none"), hasAccessToken: \(finalAccessToken?.isEmpty == false), hasRefreshToken: \(refreshToken?.isEmpty == false)") + return LocalCredentialInfo( accessToken: finalAccessToken, refreshToken: refreshToken, email: authStatus.email, - name: authStatus.name) + name: authStatus.name, + expiresAt: expiresAt) } if let refreshToken, !refreshToken.isEmpty { @@ -79,7 +79,8 @@ public enum AntigravityLocalImporter { accessToken: accessToken, refreshToken: refreshToken, email: nil, - name: nil) + name: nil, + expiresAt: expiresAt) } Self.log.debug("No credentials found in database") diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift index 8ea88c87..b15b52d6 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift @@ -37,7 +37,7 @@ public struct AntigravityOAuthCredentials: Sendable, Codable { } public var needsRefresh: Bool { - guard let expiresAt else { return self.isRefreshable } + guard let expiresAt else { return false } let bufferTime: TimeInterval = 5 * 60 return expiresAt.timeIntervalSinceNow < bufferTime } @@ -103,10 +103,12 @@ public enum AntigravityOAuthCredentialsStore { public struct ManualTokenPayload: Sendable { public let accessToken: String public let refreshToken: String? + public let expiresAt: Date? - public init(accessToken: String, refreshToken: String?) { + public init(accessToken: String, refreshToken: String?, expiresAt: Date?) { self.accessToken = accessToken self.refreshToken = refreshToken + self.expiresAt = expiresAt } } @@ -128,29 +130,39 @@ public enum AntigravityOAuthCredentialsStore { public static func manualTokenPayload(from token: String) -> ManualTokenPayload? { guard token.hasPrefix(self.manualTokenPrefix) else { return nil } let content = String(token.dropFirst(self.manualTokenPrefix.count)) - if let jsonData = content.data(using: .utf8), - let json = try? JSONSerialization.jsonObject(with: jsonData) as? [String: String], - let accessToken = json["access"], - !accessToken.isEmpty - { - return ManualTokenPayload(accessToken: accessToken, refreshToken: json["refresh"]) - } - guard !content.isEmpty else { return nil } - return ManualTokenPayload(accessToken: content, refreshToken: nil) + + guard let jsonData = content.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: jsonData) as? [String: Any], + let accessToken = json["accessToken"] as? String, + !accessToken.isEmpty + else { return nil } + + let refreshToken = json["refreshToken"] as? String + let expiresAt: Date? = (json["expiresAt"] as? TimeInterval).map { Date(timeIntervalSince1970: $0) } + + return ManualTokenPayload(accessToken: accessToken, refreshToken: refreshToken, expiresAt: expiresAt) } - public static func manualTokenValue(accessToken: String, refreshToken: String?) -> String { + public static func manualTokenValue( + accessToken: String, + refreshToken: String?, + expiresAt: Date? = nil + ) -> String { let trimmedAccess = accessToken.trimmingCharacters(in: .whitespacesAndNewlines) - let trimmedRefresh = refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines) - if let refresh = trimmedRefresh, !refresh.isEmpty { - let tokenData = ["access": trimmedAccess, "refresh": refresh] - if let jsonData = try? JSONSerialization.data(withJSONObject: tokenData), - let jsonString = String(data: jsonData, encoding: .utf8) - { - return "\(self.manualTokenPrefix)\(jsonString)" - } + var tokenData: [String: Any] = ["accessToken": trimmedAccess] + + if let refresh = refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines), !refresh.isEmpty { + tokenData["refreshToken"] = refresh } - return "\(self.manualTokenPrefix)\(trimmedAccess)" + if let expiresAt { + tokenData["expiresAt"] = expiresAt.timeIntervalSince1970 + } + + guard let jsonData = try? JSONSerialization.data(withJSONObject: tokenData), + let jsonString = String(data: jsonData, encoding: .utf8) + else { return "\(self.manualTokenPrefix){\"accessToken\":\"\(trimmedAccess)\"}" } + + return "\(self.manualTokenPrefix)\(jsonString)" } public static func load(accountLabel: String) -> AntigravityOAuthCredentials? { diff --git a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift index 77aa1f43..14f35ecc 100644 --- a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift +++ b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift @@ -29,7 +29,7 @@ public struct ProviderFetchContext: Sendable { public let fetcher: UsageFetcher public let claudeFetcher: any ClaudeUsageFetching public let browserDetection: BrowserDetection - public let onCredentialsRefreshed: (@Sendable (UsageProvider, String, String?) -> Void)? + public let onAntigravityCredentialsRefreshed: (@Sendable (String, AntigravityOAuthCredentials) -> Void)? public init( runtime: ProviderRuntime, @@ -43,7 +43,7 @@ public struct ProviderFetchContext: Sendable { fetcher: UsageFetcher, claudeFetcher: any ClaudeUsageFetching, browserDetection: BrowserDetection, - onCredentialsRefreshed: (@Sendable (UsageProvider, String, String?) -> Void)? = nil) + onAntigravityCredentialsRefreshed: (@Sendable (String, AntigravityOAuthCredentials) -> Void)? = nil) { self.runtime = runtime self.sourceMode = sourceMode @@ -56,7 +56,7 @@ public struct ProviderFetchContext: Sendable { self.fetcher = fetcher self.claudeFetcher = claudeFetcher self.browserDetection = browserDetection - self.onCredentialsRefreshed = onCredentialsRefreshed + self.onAntigravityCredentialsRefreshed = onAntigravityCredentialsRefreshed } } diff --git a/Tests/CodexBarTests/AntigravityOAuthTests.swift b/Tests/CodexBarTests/AntigravityOAuthTests.swift index 253c96f9..1667c8a6 100644 --- a/Tests/CodexBarTests/AntigravityOAuthTests.swift +++ b/Tests/CodexBarTests/AntigravityOAuthTests.swift @@ -67,18 +67,39 @@ struct AntigravityManualTokenPayloadTests { func parsesJSONPayload() { let token = AntigravityOAuthCredentialsStore.manualTokenValue( accessToken: "ya29.test", - refreshToken: "1//refresh") + refreshToken: "1//refresh", + expiresAt: nil) let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) #expect(payload?.accessToken == "ya29.test") #expect(payload?.refreshToken == "1//refresh") + #expect(payload?.expiresAt == nil) } @Test - func parsesLegacyPayload() { - let token = "\(AntigravityOAuthCredentialsStore.manualTokenPrefix)ya29.test" + func parsesJSONPayloadWithExpiresAt() { + let expiresAt = Date(timeIntervalSince1970: 1_738_160_000) + let token = AntigravityOAuthCredentialsStore.manualTokenValue( + accessToken: "ya29.test", + refreshToken: "1//refresh", + expiresAt: expiresAt) let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) #expect(payload?.accessToken == "ya29.test") - #expect(payload?.refreshToken == nil) + #expect(payload?.refreshToken == "1//refresh") + #expect(payload?.expiresAt?.timeIntervalSince1970 == 1_738_160_000) + } + + @Test + func rejectsLegacyPayload() { + let token = "\(AntigravityOAuthCredentialsStore.manualTokenPrefix)ya29.test" + let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) + #expect(payload == nil) + } + + @Test + func rejectsOldJSONFormat() { + let token = "\(AntigravityOAuthCredentialsStore.manualTokenPrefix){\"access\":\"ya29.test\",\"refresh\":\"1//refresh\"}" + let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) + #expect(payload == nil) } } From 3372d8deebe044109fbabd36abf036f1374d16fb Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 15:18:43 +0300 Subject: [PATCH 17/19] fix: correct label text by removing redundant "language" from local server description --- .../Antigravity/AntigravityProviderImplementation.swift | 2 +- .../Antigravity/AntigravityOAuth/AntigravityUsageSource.swift | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index 5ce5a8bb..297daeef 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -72,7 +72,7 @@ struct AntigravityProviderImplementation: ProviderImplementation { case .authorized: return "OAuth: Use OAuth account or manual tokens only" case .local: - return "Local: Use local Antigravity language server only" + return "Local: Use local Antigravity local server only" } }, binding: usageBinding, diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift index ecaf8c10..17aa3d5f 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift @@ -36,7 +36,7 @@ public enum AntigravityUsageSource: String, CaseIterable, Sendable, Codable { case .authorized: "Use OAuth account or manual tokens only" case .local: - "Use local Antigravity language server only" + "Use local Antigravity local server only" } } } From 5dad68b2c8a7e1aefe909b46f15e075803cf588d Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Thu, 29 Jan 2026 15:56:27 +0300 Subject: [PATCH 18/19] fix: update local server description in AntigravityProviderImplementation and documentation for clarity --- .../Antigravity/AntigravityProviderImplementation.swift | 2 +- .../AntigravityOAuth/AntigravityUsageSource.swift | 2 +- docs/antigravity.md | 5 +++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index 297daeef..a27d0384 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -72,7 +72,7 @@ struct AntigravityProviderImplementation: ProviderImplementation { case .authorized: return "OAuth: Use OAuth account or manual tokens only" case .local: - return "Local: Use local Antigravity local server only" + return "Local: Use Antigravity local server only" } }, binding: usageBinding, diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift index 17aa3d5f..8128cce0 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityUsageSource.swift @@ -36,7 +36,7 @@ public enum AntigravityUsageSource: String, CaseIterable, Sendable, Codable { case .authorized: "Use OAuth account or manual tokens only" case .local: - "Use local Antigravity local server only" + "Use Antigravity local server only" } } } diff --git a/docs/antigravity.md b/docs/antigravity.md index f2bfd190..c95c79ee 100644 --- a/docs/antigravity.md +++ b/docs/antigravity.md @@ -15,12 +15,12 @@ Antigravity supports OAuth-authorized Cloud Code quota and local language server - **Auto** (default): OAuth/manual first, fallback to local server - **OAuth**: OAuth/manual only -- **Local**: local Antigravity language server only +- **Local**: Antigravity local server only ## OAuth credentials - **Keychain**: stored from the OAuth browser flow -- **Manual tokens**: stored in token accounts with `manual:` prefix (access `ya29.` + optional refresh `1//`) +- **Manual tokens**: stored in token accounts with `manual:` prefix (access `ya29.` + optional refresh `1//`) when Keychain is disabled; otherwise saved to Keychain under the account label. - **Local import**: `~/Library/Application Support/Antigravity/User/globalStorage/state.vscdb` - refresh token: `jetskiStateSync.agentManagerInitState` (base64 protobuf, field 6 contains nested OAuthTokenInfo) - access token/email: `antigravityAuthStatus` JSON (`apiKey`, `email`) @@ -31,6 +31,7 @@ Antigravity supports OAuth-authorized Cloud Code quota and local language server - `POST https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` - `POST https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels` +- `POST https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:fetchAvailableModels` (fallback) - `POST https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist` ## Local server data sources + fallback order From 3a75c695d61f83d0706aca6ef466823e9b5493b3 Mon Sep 17 00:00:00 2001 From: Mert Can Demir Date: Tue, 3 Feb 2026 13:55:31 +0300 Subject: [PATCH 19/19] chore: linting Add antigravity_state.pb.swift to exclusion lists in .swiftformat and .swiftlint.yml, and apply SwiftFormat baseline formatting changes across the codebase. --- .swiftformat | 2 +- .swiftlint.yml | 2 + .../CodexBar/PreferencesProvidersPane.swift | 23 +++++++--- .../Antigravity/AntigravityLoginFlow.swift | 8 ++-- .../AntigravityProviderImplementation.swift | 6 +-- .../AntigravitySettingsStore.swift | 8 ++-- .../CodexBar/UsageStore+TokenAccounts.swift | 14 +++--- Sources/CodexBar/UsageStore.swift | 3 +- .../AntigravityAuthorizedFetchStrategy.swift | 31 ++++++++----- .../AntigravityCloudCodeClient.swift | 18 +++++--- .../AntigravityLocalImporter.swift | 43 +++++++++++++------ .../AntigravityOAuthCredentials.swift | 4 +- .../AntigravityOAuthFlow.swift | 22 +++++----- .../AntigravityProviderDescriptor.swift | 2 +- .../CodexBarTests/AntigravityOAuthTests.swift | 3 +- 15 files changed, 120 insertions(+), 69 deletions(-) diff --git a/.swiftformat b/.swiftformat index 4f3218d9..0b214dc9 100644 --- a/.swiftformat +++ b/.swiftformat @@ -47,4 +47,4 @@ --allman false # Exclusions ---exclude .build,.swiftpm,DerivedData,node_modules,dist,coverage,xcuserdata,Core/PeekabooCore/Sources/PeekabooCore/Extensions/NSArray+Extensions.swift +--exclude .build,.swiftpm,DerivedData,node_modules,dist,coverage,xcuserdata,Core/PeekabooCore/Sources/PeekabooCore/Extensions/NSArray+Extensions.swift,Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift diff --git a/.swiftlint.yml b/.swiftlint.yml index ceae70f1..a90d939d 100644 --- a/.swiftlint.yml +++ b/.swiftlint.yml @@ -24,6 +24,8 @@ excluded: - "*.playground" # Exclude specific files that should not be linted/formatted - "Core/PeekabooCore/Sources/PeekabooCore/Extensions/NSArray+Extensions.swift" + # Protobuf generated file + - "Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/antigravity_state.pb.swift" # Analyzer rules (require compilation) analyzer_rules: diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index d9ff9fbe..91b5a989 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -255,7 +255,11 @@ struct ProvidersPane: View { do { let credentials = try await AntigravityLocalImporter.importCredentials() - log.debug("Import successful - email: \(credentials.email ?? "none"), hasAccessToken: \(credentials.hasAccessToken), hasRefreshToken: \(credentials.hasRefreshToken)") + log.debug( + """ + Import successful - email: \(credentials.email ?? "none"), \ + hasAccessToken: \(credentials.hasAccessToken), hasRefreshToken: \(credentials.hasRefreshToken) + """) guard let accessToken = credentials.accessToken, !accessToken.isEmpty else { log.debug("Import failed: no access token found") @@ -272,7 +276,8 @@ struct ProvidersPane: View { label: label, accessToken: accessToken, refreshToken: credentials.refreshToken, - expiresAt: credentials.expiresAt) else { + expiresAt: credentials.expiresAt) + else { log.debug("Import failed: unable to save imported credentials") self.presentAlert( title: "Import Failed", @@ -295,7 +300,14 @@ struct ProvidersPane: View { if AntigravityLocalImporter.isAvailable() { self.presentAlert( title: "No Credentials Found", - message: "Antigravity database found, but no credentials were found inside.\n\nPlease ensure:\n1. You are signed in to Antigravity IDE (check the Account menu)\n2. Try restarting Antigravity IDE if you just signed in\n3. If the issue persists, paste your tokens manually below") + message: """ + Antigravity database found, but no credentials were found inside. + + Please ensure: + 1. You are signed in to Antigravity IDE (check the Account menu) + 2. Try restarting Antigravity IDE if you just signed in + 3. If the issue persists, paste your tokens manually below + """) } else { self.presentAlert( title: "Import Failed", @@ -309,7 +321,7 @@ struct ProvidersPane: View { } catch { log.debug("Import failed with error: \(error)") let nsError = error as NSError - if nsError.domain == NSPOSIXErrorDomain && nsError.code == 1 { + if nsError.domain == NSPOSIXErrorDomain, nsError.code == 1 { self.presentFullDiskAccessAlert() } else { self.presentAlert( @@ -323,7 +335,8 @@ struct ProvidersPane: View { private func presentFullDiskAccessAlert() { let alert = NSAlert() alert.messageText = "Full Disk Access Required" - alert.informativeText = "Full Disk Access is required to read the Antigravity database. Please grant access in System Settings." + alert.informativeText = + "Full Disk Access is required to read the Antigravity database. Please grant access in System Settings." alert.addButton(withTitle: "Open System Settings") alert.addButton(withTitle: "Cancel") diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift index 9b53f08e..575d4689 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityLoginFlow.swift @@ -6,8 +6,8 @@ struct AntigravityLoginFlow { private static let log = CodexBarLog.logger(LogCategories.antigravity) static func runOAuthFlow(settings: SettingsStore, store: UsageStore? = nil) async -> Bool { - Self.log.debug("Starting Antigravity OAuth login flow") - Self.log.debug("Keychain access disabled: \(KeychainAccessGate.isDisabled)") + self.log.debug("Starting Antigravity OAuth login flow") + self.log.debug("Keychain access disabled: \(KeychainAccessGate.isDisabled)") let flow = AntigravityOAuthFlow() @@ -75,8 +75,8 @@ struct AntigravityLoginFlow { _ credentials: AntigravityOAuthCredentials, settings: SettingsStore) -> String? { - guard let accountLabel = Self.resolveAccountLabel(credentials: credentials, settings: settings) else { - Self.log.debug("Failed to resolve account label") + guard let accountLabel = resolveAccountLabel(credentials: credentials, settings: settings) else { + self.log.debug("Failed to resolve account label") return nil } Self.log.debug("Persisting credentials for account: \(accountLabel)") diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift index a27d0384..4449997f 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravityProviderImplementation.swift @@ -68,11 +68,11 @@ struct AntigravityProviderImplementation: ProviderImplementation { dynamicSubtitle: { switch context.settings.antigravityUsageSource { case .auto: - return "Auto: Try OAuth/manual first, fallback to local server" + "Auto: Try OAuth/manual first, fallback to local server" case .authorized: - return "OAuth: Use OAuth account or manual tokens only" + "OAuth: Use OAuth account or manual tokens only" case .local: - return "Local: Use Antigravity local server only" + "Local: Use Antigravity local server only" } }, binding: usageBinding, diff --git a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift index 8398f929..686a1de3 100644 --- a/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift +++ b/Sources/CodexBar/Providers/Antigravity/AntigravitySettingsStore.swift @@ -86,8 +86,8 @@ extension SettingsStore { label: String, accessToken: String, refreshToken: String? = nil, - expiresAt: Date? = nil - ) -> ProviderTokenAccount? { + expiresAt: Date? = nil) -> ProviderTokenAccount? + { guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return nil } let tokenValue: String @@ -158,8 +158,8 @@ extension SettingsStore { private func triggerBackgroundRefreshIfNeeded( label: String, refreshToken: String?, - expiresAt: Date? - ) { + expiresAt: Date?) + { guard let refreshToken, !refreshToken.isEmpty, expiresAt == nil else { return } guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(label) else { return } diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index b0713d3e..cd4282bb 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -88,12 +88,13 @@ extension UsageStore { tokenOverride: override) let verbose = self.settings.isVerboseLoggingEnabled - let onAntigravityCredentialsRefreshed: (@Sendable (String, AntigravityOAuthCredentials) -> Void) = { [weak self] accountLabel, credentials in - Task { @MainActor in - guard let self else { return } - self.saveRefreshedAntigravityCredentials(accountLabel: accountLabel, credentials: credentials) + let onAntigravityCredentialsRefreshed: (@Sendable (String, AntigravityOAuthCredentials) -> Void) = + { [weak self] accountLabel, credentials in + Task { @MainActor in + guard let self else { return } + self.saveRefreshedAntigravityCredentials(accountLabel: accountLabel, credentials: credentials) + } } - } let context = ProviderFetchContext( runtime: .app, @@ -217,7 +218,8 @@ extension UsageStore { guard let normalizedLabel = AntigravityOAuthCredentialsStore.normalizedLabel(accountLabel) else { return } let tokenAccounts = self.settings.tokenAccountsData(for: .antigravity) - guard let account = tokenAccounts?.accounts.first(where: { $0.label.lowercased() == normalizedLabel }) else { return } + guard let account = tokenAccounts?.accounts.first(where: { $0.label.lowercased() == normalizedLabel }) + else { return } let tokenValue = AntigravityOAuthCredentialsStore.manualTokenValue( accessToken: credentials.accessToken, diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index c96b0784..f9088549 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -107,7 +107,8 @@ extension UsageStore { lines.append("localServer=\(serverRunning ? "running" : "not_running")") lines.append("") - let isAuthorizedAvailable = hasAccountLabel && (hasKeychainAccessToken || hasKeychainRefreshToken || manualPayload != nil) + let isAuthorizedAvailable = hasAccountLabel && + (hasKeychainAccessToken || hasKeychainRefreshToken || manualPayload != nil) lines.append("authorizedStrategy=\(available(isAuthorizedAvailable))") lines.append("localStrategy=\(available(serverRunning))") diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift index e045e48a..1b94c00a 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityAuthorizedFetchStrategy.swift @@ -25,9 +25,13 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { Self.log.debug("Keychain credentials not found") return false } - - Self.log.debug("Keychain credentials found - hasAccessToken: \(!credentials.accessToken.isEmpty), isRefreshable: \(credentials.isRefreshable)") - + + Self.log.debug( + """ + Keychain credentials found - hasAccessToken: \(!credentials.accessToken.isEmpty), \ + isRefreshable: \(credentials.isRefreshable) + """) + if !credentials.accessToken.isEmpty { return true } return credentials.isRefreshable } @@ -41,7 +45,11 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { let sourceLabel = resolved.sourceLabel var didRefresh = false - Self.log.debug("Resolved credentials - source: \(sourceLabel), needsRefresh: \(credentials.needsRefresh), isRefreshable: \(credentials.isRefreshable)") + Self.log.debug( + """ + Resolved credentials - source: \(sourceLabel), needsRefresh: \(credentials.needsRefresh), \ + isRefreshable: \(credentials.isRefreshable) + """) if credentials.needsRefresh || (credentials.accessToken.isEmpty && credentials.isRefreshable) { Self.log.debug("Credentials need refresh, refreshing token...") @@ -61,8 +69,8 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { private func fetchQuotaAndMakeResult( credentials: AntigravityOAuthCredentials, - sourceLabel: String - ) async throws -> ProviderFetchResult { + sourceLabel: String) async throws -> ProviderFetchResult + { let quota = try await AntigravityCloudCodeClient.fetchQuota(accessToken: credentials.accessToken) Self.log.debug("Successfully fetched quota from Cloud Code API") @@ -78,8 +86,8 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { private func refreshAndSave( _ credentials: AntigravityOAuthCredentials, accountLabel: String, - context: ProviderFetchContext - ) async throws -> AntigravityOAuthCredentials { + context: ProviderFetchContext) async throws -> AntigravityOAuthCredentials + { do { let refreshed = try await self.refreshCredentials(credentials) @@ -141,7 +149,8 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { return (normalized, cached, "OAuth") } - private func refreshCredentials(_ credentials: AntigravityOAuthCredentials) async throws -> AntigravityOAuthCredentials { + private func refreshCredentials(_ credentials: AntigravityOAuthCredentials) async throws + -> AntigravityOAuthCredentials { guard let refreshToken = credentials.refreshToken else { throw AntigravityOAuthCredentialsError.invalidGrant } @@ -153,8 +162,8 @@ public struct AntigravityAuthorizedFetchStrategy: ProviderFetchStrategy { private func loadManualCredentials( accountLabel: String, - context: ProviderFetchContext - ) -> AntigravityOAuthCredentials? { + context: ProviderFetchContext) -> AntigravityOAuthCredentials? + { guard let normalized = AntigravityOAuthCredentialsStore.normalizedLabel(accountLabel) else { return nil } let tokenAccounts = context.settings?.antigravity?.tokenAccounts diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift index 507ec419..ae961d41 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityCloudCodeClient.swift @@ -41,7 +41,10 @@ public enum AntigravityCloudCodeClient { private static let log = CodexBarLog.logger(LogCategories.antigravity) private static let httpTimeout: TimeInterval = 15.0 - public static func fetchQuota(accessToken: String, projectId: String? = nil) async throws -> AntigravityCloudCodeQuota { + public static func fetchQuota( + accessToken: String, + projectId: String? = nil) async throws -> AntigravityCloudCodeQuota + { try await self.requestWithRetry { baseURL in try await self.fetchQuotaFromEndpoint( baseURL: baseURL, @@ -64,7 +67,9 @@ public enum AntigravityCloudCodeClient { for attempt in 1...AntigravityCloudCodeConfig.defaultAttempts { if attempt > 1 { let delay = self.getBackoffDelay(attempt: attempt) - self.log.info("Cloud Code retry round \(attempt)/\(AntigravityCloudCodeConfig.defaultAttempts) in \(delay)ms") + self.log + .info( + "Cloud Code retry round \(attempt)/\(AntigravityCloudCodeConfig.defaultAttempts) in \(delay)ms") try await Task.sleep(nanoseconds: UInt64(delay) * 1_000_000) } @@ -165,13 +170,12 @@ public enum AntigravityCloudCodeClient { } let projectId = self.extractProjectId(from: json["cloudaicompanionProject"]) - let tierId: String? - if let paidTier = json["paidTier"] as? [String: Any], let id = paidTier["id"] as? String { - tierId = id + let tierId: String? = if let paidTier = json["paidTier"] as? [String: Any], let id = paidTier["id"] as? String { + id } else if let currentTier = json["currentTier"] as? [String: Any], let id = currentTier["id"] as? String { - tierId = id + id } else { - tierId = nil + nil } return AntigravityProjectInfo(projectId: projectId, tierId: tierId) diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift index 80340199..83f6d335 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityLocalImporter.swift @@ -37,11 +37,11 @@ public enum AntigravityLocalImporter { } public static func importCredentials() async throws -> LocalCredentialInfo { - Self.log.debug("Starting Antigravity DB import") - + self.log.debug("Starting Antigravity DB import") + let dbPath = self.stateDbPath() Self.log.debug("Database path: \(dbPath.path)") - + guard FileManager.default.fileExists(atPath: dbPath.path) else { Self.log.debug("Database file not found at path") throw AntigravityOAuthCredentialsError.notFound @@ -57,13 +57,26 @@ public enum AntigravityLocalImporter { if let expiry = protoInfo.expirySeconds { expiresAt = Date(timeIntervalSince1970: TimeInterval(expiry)) } - Self.log.debug("Extracted OAuth token info - access_token present: \(accessToken?.isEmpty == false), refresh_token present: \(refreshToken?.isEmpty == false)") + Self.log.debug( + """ + Extracted OAuth token info - access_token present: \(accessToken?.isEmpty == false), \ + refresh_token present: \(refreshToken?.isEmpty == false) + """) } if let authStatus = try? self.readAuthStatus(dbPath: dbPath) { - Self.log.debug("Read auth status - email: \(authStatus.email ?? "none"), apiKey present: \(authStatus.apiKey?.isEmpty == false)") + Self.log.debug( + """ + Read auth status - email: \(authStatus.email ?? "none"), \ + apiKey present: \(authStatus.apiKey?.isEmpty == false) + """) let finalAccessToken = accessToken ?? authStatus.apiKey - Self.log.debug("Import result - email: \(authStatus.email ?? "none"), hasAccessToken: \(finalAccessToken?.isEmpty == false), hasRefreshToken: \(refreshToken?.isEmpty == false)") + Self.log.debug( + """ + Import result - email: \(authStatus.email ?? "none"), \ + hasAccessToken: \(finalAccessToken?.isEmpty == false), hasRefreshToken: \(refreshToken? + .isEmpty == false) + """) return LocalCredentialInfo( accessToken: finalAccessToken, @@ -105,7 +118,7 @@ public enum AntigravityLocalImporter { } private static func readAuthStatus(dbPath: URL) throws -> AuthStatus { - Self.log.debug("Reading antigravityAuthStatus from DB") + self.log.debug("Reading antigravityAuthStatus from DB") let json = try self.readStateValue(dbPath: dbPath, key: "antigravityAuthStatus") guard let data = json.data(using: .utf8), let dict = try JSONSerialization.jsonObject(with: data) as? [String: Any] @@ -121,15 +134,15 @@ public enum AntigravityLocalImporter { } private static func readProtoTokenInfo(dbPath: URL) throws -> ProtoTokenInfo { - Self.log.debug("Reading jetskiStateSync.agentManagerInitState from DB") + self.log.debug("Reading jetskiStateSync.agentManagerInitState from DB") let base64 = try self.readStateValue(dbPath: dbPath, key: "jetskiStateSync.agentManagerInitState") Self.log.debug("Read base64 value, length: \(base64.count)") - + guard let data = Data(base64Encoded: base64.trimmingCharacters(in: .whitespacesAndNewlines)) else { throw AntigravityOAuthCredentialsError.decodeFailed("Invalid base64 in agentManagerInitState") } Self.log.debug("Decoded base64, data length: \(data.count)") - + return try self.parseProtoTokenInfo(data: data) } @@ -172,7 +185,7 @@ public enum AntigravityLocalImporter { } private static func parseProtoTokenInfo(data: Data) throws -> ProtoTokenInfo { - Self.log.debug("Parsing protobuf data using swift-protobuf") + self.log.debug("Parsing protobuf data using swift-protobuf") do { let state = try AgentManagerInitState(serializedBytes: data) @@ -184,7 +197,11 @@ public enum AntigravityLocalImporter { } let oauthToken = state.oauthToken - Self.log.debug("Found OAuthTokenInfo - access_token length: \(oauthToken.accessToken.count), refresh_token length: \(oauthToken.refreshToken.count)") + Self.log.debug( + """ + Found OAuthTokenInfo - access_token length: \(oauthToken.accessToken.count), \ + refresh_token length: \(oauthToken.refreshToken.count) + """) var expirySeconds: Int? if oauthToken.hasExpiry { @@ -198,7 +215,7 @@ public enum AntigravityLocalImporter { tokenType: oauthToken.tokenType.isEmpty ? nil : oauthToken.tokenType, expirySeconds: expirySeconds) } catch { - Self.log.debug("Protobuf parsing failed: \(error)") + self.log.debug("Protobuf parsing failed: \(error)") throw AntigravityOAuthCredentialsError.decodeFailed("Failed to parse protobuf: \(error)") } } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift index b15b52d6..4d6159f4 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthCredentials.swift @@ -146,8 +146,8 @@ public enum AntigravityOAuthCredentialsStore { public static func manualTokenValue( accessToken: String, refreshToken: String?, - expiresAt: Date? = nil - ) -> String { + expiresAt: Date? = nil) -> String + { let trimmedAccess = accessToken.trimmingCharacters(in: .whitespacesAndNewlines) var tokenData: [String: Any] = ["accessToken": trimmedAccess] diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift index 1b8fe1fe..403ca602 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuth/AntigravityOAuthFlow.swift @@ -18,14 +18,14 @@ public actor AntigravityOAuthFlow { public func startAuthorization() async throws -> AntigravityOAuthCredentials { Self.log.debug("Starting OAuth authorization flow") - + let port = try await self.startCallbackServer() let redirectUri = "http://\(AntigravityOAuthConfig.callbackHost):\(port)" let state = self.generateState() self.pendingState = state let authURL = self.buildAuthURL(redirectUri: redirectUri, state: state) - + Self.log.debug("Requesting scopes: \(AntigravityOAuthConfig.scopes.joined(separator: ", "))") Self.log.debug("Using access_type=offline and prompt=consent to ensure refresh token") Self.log.info("Opening Antigravity OAuth authorization URL") @@ -85,7 +85,8 @@ public actor AntigravityOAuthFlow { } guard let code, let state else { - self.pendingContinuation?.resume(throwing: AntigravityOAuthCredentialsError.decodeFailed("Missing code or state in callback")) + self.pendingContinuation? + .resume(throwing: AntigravityOAuthCredentialsError.decodeFailed("Missing code or state in callback")) self.pendingContinuation = nil return } @@ -149,7 +150,8 @@ public actor AntigravityOAuthFlow { guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { let statusCode = (response as? HTTPURLResponse)?.statusCode ?? 0 let errorBody = String(data: data, encoding: .utf8) ?? "" - throw AntigravityOAuthCredentialsError.refreshFailed("Token exchange failed: HTTP \(statusCode) - \(errorBody)") + throw AntigravityOAuthCredentialsError + .refreshFailed("Token exchange failed: HTTP \(statusCode) - \(errorBody)") } guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], @@ -166,7 +168,7 @@ public actor AntigravityOAuthFlow { Self.log.debug("Received OAuth response - access_token length: \(accessToken.count)") Self.log.debug("Received OAuth response - refresh_token length: \(refreshToken.count)") Self.log.debug("Token expires in: \(expiresIn) seconds") - if let email = email { + if let email { Self.log.debug("Resolved account email: \(email)") } @@ -262,9 +264,8 @@ private final class CallbackServer: @unchecked Sendable { let request = String(bytes: buffer[0.. Authorization Successful @@ -275,7 +276,7 @@ private final class CallbackServer: @unchecked Sendable { """ } else { - responseHTML = """ + """ Authorization Failed @@ -286,7 +287,8 @@ private final class CallbackServer: @unchecked Sendable { """ } - let response = "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nConnection: close\r\n\r\n\(responseHTML)" + let headers = "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nConnection: close\r\n\r\n" + let response = headers + responseHTML _ = response.withCString { ptr in send(clientSocket, ptr, strlen(ptr), 0) } diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift index bfea3874..d7588922 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift @@ -34,7 +34,7 @@ public enum AntigravityProviderDescriptor { noDataMessage: { "Antigravity cost summary is not supported." }), fetchPlan: ProviderFetchPlan( sourceModes: [.auto, .oauth, .cli], - pipeline: ProviderFetchPipeline(resolveStrategies: Self.resolveStrategies)), + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), cli: ProviderCLIConfig( name: "antigravity", versionDetector: nil)) diff --git a/Tests/CodexBarTests/AntigravityOAuthTests.swift b/Tests/CodexBarTests/AntigravityOAuthTests.swift index 1667c8a6..49d6af67 100644 --- a/Tests/CodexBarTests/AntigravityOAuthTests.swift +++ b/Tests/CodexBarTests/AntigravityOAuthTests.swift @@ -97,7 +97,8 @@ struct AntigravityManualTokenPayloadTests { @Test func rejectsOldJSONFormat() { - let token = "\(AntigravityOAuthCredentialsStore.manualTokenPrefix){\"access\":\"ya29.test\",\"refresh\":\"1//refresh\"}" + let token = + "\(AntigravityOAuthCredentialsStore.manualTokenPrefix){\"access\":\"ya29.test\",\"refresh\":\"1//refresh\"}" let payload = AntigravityOAuthCredentialsStore.manualTokenPayload(from: token) #expect(payload == nil) }