1154 lines · 43.5 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif#if canImport(Darwin)import Darwin#elseif canImport(Glibc)import Glibc#elseif canImport(Musl)import Musl#endif public struct GeminiModelQuota: Sendable {    public let modelId: String    public let percentLeft: Double    public let resetTime: Date?    public let resetDescription: String?} public struct GeminiStatusSnapshot: Sendable {    public let modelQuotas: [GeminiModelQuota]    public let rawText: String    public let accountEmail: String?    public let accountPlan: String?     // Convenience: lowest quota across all models (for icon display)    public var lowestPercentLeft: Double? {        self.modelQuotas.min(by: { $0.percentLeft < $1.percentLeft })?.percentLeft    }     /// Legacy compatibility    public var dailyPercentLeft: Double? {        self.lowestPercentLeft    }     public var resetDescription: String? {        self.modelQuotas.min(by: { $0.percentLeft < $1.percentLeft })?.resetDescription    }     /// Converts Gemini quotas to a unified UsageSnapshot.    /// Groups quotas by tier: Pro (24h window) as primary, Flash (24h window) as secondary,    /// Flash Lite (24h window) as tertiary.    public func toUsageSnapshot() -> UsageSnapshot {        let lower = self.modelQuotas.map { ($0.modelId.lowercased(), $0) }        let flashLiteQuotas = lower.filter { Self.isFlashLiteModel(id: $0.0) }.map(\.1)        let flashQuotas = lower.filter { Self.isFlashModel(id: $0.0) }.map(\.1)        let proQuotas = lower.filter { Self.isProModel(id: $0.0) }.map(\.1)         let flashLiteMin = flashLiteQuotas.min(by: { $0.percentLeft < $1.percentLeft })        let flashMin = flashQuotas.min(by: { $0.percentLeft < $1.percentLeft })        let proMin = proQuotas.min(by: { $0.percentLeft < $1.percentLeft })         let primary = RateWindow(            usedPercent: proMin.map { 100 - $0.percentLeft } ?? 0,            windowMinutes: 1440,            resetsAt: proMin?.resetTime,            resetDescription: proMin?.resetDescription)         let secondary: RateWindow? = flashMin.map {            RateWindow(                usedPercent: 100 - $0.percentLeft,                windowMinutes: 1440,                resetsAt: $0.resetTime,                resetDescription: $0.resetDescription)        }        let tertiary: RateWindow? = flashLiteMin.map {            RateWindow(                usedPercent: 100 - $0.percentLeft,                windowMinutes: 1440,                resetsAt: $0.resetTime,                resetDescription: $0.resetDescription)        }         let identity = ProviderIdentitySnapshot(            providerID: .gemini,            accountEmail: self.accountEmail,            accountOrganization: nil,            loginMethod: self.accountPlan)        return UsageSnapshot(            primary: primary,            secondary: secondary,            tertiary: tertiary,            updatedAt: Date(),            identity: identity)    }     private static func isFlashLiteModel(id: String) -> Bool {        id.contains("flash-lite")    }     private static func isFlashModel(id: String) -> Bool {        id.contains("flash") && !self.isFlashLiteModel(id: id)    }     private static func isProModel(id: String) -> Bool {        id.contains("pro")    }} public enum GeminiStatusProbeError: LocalizedError, Sendable, Equatable {    case geminiNotInstalled    case notLoggedIn    case unsupportedAuthType(String)    case parseFailed(String)    case timedOut    case apiError(String)     public var errorDescription: String? {        switch self {        case .geminiNotInstalled:            "Gemini CLI is not installed or not on PATH."        case .notLoggedIn:            "Not logged in to Gemini. Run 'gemini' in Terminal to authenticate."        case let .unsupportedAuthType(authType):            "Gemini \(authType) auth not supported. Use Google account (OAuth) instead."        case let .parseFailed(msg):            "Could not parse Gemini usage: \(msg)"        case .timedOut:            "Gemini quota API request timed out."        case let .apiError(msg):            "Gemini API error: \(msg)"        }    }} public enum GeminiAuthType: String, Sendable {    case oauthPersonal = "oauth-personal"    case apiKey = "gemini-api-key"    case vertexAI = "vertex-ai"    case unknown} /// User tier IDs returned from the Cloud Code Private API (loadCodeAssist)./// Maps to: google3/cloud/developer_experience/cloudcode/pa/service/usertier.gopublic enum GeminiUserTierId: String, Sendable {    case free = "free-tier"    case legacy = "legacy-tier"    case standard = "standard-tier"} public struct GeminiStatusProbe: Sendable {    public var timeout: TimeInterval = 10.0    public var homeDirectory: String    public var dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)    private static let log = CodexBarLog.logger(LogCategories.geminiProbe)    private static let quotaEndpoint = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"    private static let loadCodeAssistEndpoint = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"    private static let projectsEndpoint = "https://cloudresourcemanager.googleapis.com/v1/projects"    private static let credentialsPath = "/.gemini/oauth_creds.json"    private static let settingsPath = "/.gemini/settings.json"    private static let tokenRefreshEndpoint = "https://oauth2.googleapis.com/token"     public init(        timeout: TimeInterval = 10.0,        homeDirectory: String = NSHomeDirectory(),        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse) = Self.defaultDataLoader)    {        self.timeout = timeout        self.homeDirectory = homeDirectory        self.dataLoader = dataLoader    }     /// Reads the current Gemini auth type from settings.json    public static func currentAuthType(homeDirectory: String = NSHomeDirectory()) -> GeminiAuthType {        let settingsURL = URL(fileURLWithPath: homeDirectory + Self.settingsPath)         guard let data = try? Data(contentsOf: settingsURL),              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],              let security = json["security"] as? [String: Any],              let auth = security["auth"] as? [String: Any],              let selectedType = auth["selectedType"] as? String        else {            return .unknown        }         if selectedType == "api-key" {            return .apiKey        }        return GeminiAuthType(rawValue: selectedType) ?? .unknown    }     public func fetch() async throws -> GeminiStatusSnapshot {        // Block explicitly unsupported auth types; allow unknown to try OAuth creds        let authType = Self.currentAuthType(homeDirectory: self.homeDirectory)        switch authType {        case .apiKey:            throw GeminiStatusProbeError.unsupportedAuthType("API key")        case .vertexAI:            throw GeminiStatusProbeError.unsupportedAuthType("Vertex AI")        case .oauthPersonal, .unknown:            break        }         let snap = try await Self.fetchViaAPI(            timeout: self.timeout,            homeDirectory: self.homeDirectory,            dataLoader: self.dataLoader)         Self.log.info("Gemini API fetch ok", metadata: [            "dailyPercentLeft": "\(snap.dailyPercentLeft ?? -1)",        ])        return snap    }     // MARK: - Direct API approach     private static func fetchViaAPI(        timeout: TimeInterval,        homeDirectory: String,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> GeminiStatusSnapshot    {        let creds = try Self.loadCredentials(homeDirectory: homeDirectory)         let expiryStr = creds.expiryDate.map { "\($0)" } ?? "nil"        let hasRefresh = creds.refreshToken != nil        Self.log.debug("Token check", metadata: [            "expiry": expiryStr,            "hasRefresh": hasRefresh ? "1" : "0",            "now": "\(Date())",        ])         var accessToken = creds.accessToken?.isEmpty == false ? creds.accessToken : nil        var idToken = creds.idToken        let needsRefresh = accessToken == nil || creds.expiryDate.map { $0 < Date() } == true        if needsRefresh {            if accessToken == nil {                Self.log.info("No access token found; attempting refresh from stored Gemini credentials")            } else if let expiry = creds.expiryDate {                Self.log.info("Token expired; attempting refresh", metadata: [                    "expiry": "\(expiry)",                ])            }             guard let refreshToken = creds.refreshToken, !refreshToken.isEmpty else {                Self.log.error("No refresh token available")                throw GeminiStatusProbeError.notLoggedIn            }            accessToken = try await Self.refreshAccessToken(                refreshToken: refreshToken,                timeout: timeout,                homeDirectory: homeDirectory,                dataLoader: dataLoader)            idToken = (try? Self.loadCredentials(homeDirectory: homeDirectory).idToken) ?? idToken        }        guard let accessToken else {            Self.log.error("No access token found")            throw GeminiStatusProbeError.notLoggedIn        }         // Extract account info from JWT        let claims = Self.extractClaimsFromToken(idToken)         // Load Code Assist status to get project ID and tier (aligned with CLI setupUser logic)        let caStatus = await Self.loadCodeAssistStatus(            accessToken: accessToken,            timeout: timeout,            dataLoader: dataLoader)         // Determine the project ID to use for quota fetching.        // Priority:        // 1. Project ID returned by loadCodeAssist (e.g. managed project for free tier)        // 2. Discovered project ID from cloud resource manager (e.g. user's own project)        var projectId = caStatus.projectId        if projectId == nil {            projectId = try? await Self.discoverGeminiProjectId(                accessToken: accessToken,                timeout: timeout,                dataLoader: dataLoader)        }         guard let url = URL(string: Self.quotaEndpoint) else {            throw GeminiStatusProbeError.apiError("Invalid endpoint URL")        }         var request = URLRequest(url: url)        request.httpMethod = "POST"        request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")        request.setValue("application/json", forHTTPHeaderField: "Content-Type")         // Include project ID for accurate quota        if let projectId {            request.httpBody = Data("{\"project\": \"\(projectId)\"}".utf8)        } else {            request.httpBody = Data("{}".utf8)        }        request.timeoutInterval = timeout         let (data, response) = try await dataLoader(request)         guard let httpResponse = response as? HTTPURLResponse else {            throw GeminiStatusProbeError.apiError("Invalid response")        }         if httpResponse.statusCode == 401 {            throw GeminiStatusProbeError.notLoggedIn        }         guard httpResponse.statusCode == 200 else {            throw GeminiStatusProbeError.apiError("HTTP \(httpResponse.statusCode)")        }         let snapshot = try Self.parseAPIResponse(data, email: claims.email)         // Plan display strings with tier mapping:        // - standard-tier: Paid subscription (AI Pro, AI Ultra, Code Assist        //   Standard/Enterprise, Developer Program Premium)        // - free-tier + hd claim: Workspace account (Gemini included free since Jan 2025)        // - free-tier: Personal free account (1000 req/day limit)        // - legacy-tier: Unknown legacy/grandfathered tier        // - nil (API failed): Leave blank (no display)        let plan: String? = switch (caStatus.tier, claims.hostedDomain) {        case (.standard, _):            "Paid"        case let (.free, .some(domain)):            { Self.log.info("Workspace account detected", metadata: ["domain": domain]); return "Workspace" }()        case (.free, .none):            { Self.log.info("Personal free account"); return "Free" }()        case (.legacy, _):            "Legacy"        case (.none, _):            { Self.log.info("Tier detection failed, leaving plan blank"); return nil }()        }         return GeminiStatusSnapshot(            modelQuotas: snapshot.modelQuotas,            rawText: snapshot.rawText,            accountEmail: snapshot.accountEmail,            accountPlan: plan)    }     private static func discoverGeminiProjectId(        accessToken: String,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws        -> String?    {        guard let url = URL(string: projectsEndpoint) else { return nil }         var request = URLRequest(url: url)        request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")        request.timeoutInterval = timeout         let (data, response) = try await dataLoader(request)         guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else {            return nil        }         guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],              let projects = json["projects"] as? [[String: Any]]        else {            return nil        }         // Look for Gemini API project (has "generative-language" label or "gen-lang-client" prefix)        for project in projects {            guard let projectId = project["projectId"] as? String else { continue }             // Check for gen-lang-client prefix (Gemini CLI projects)            if projectId.hasPrefix("gen-lang-client") {                return projectId            }             // Check for generative-language label            if let labels = project["labels"] as? [String: String],               labels["generative-language"] != nil            {                return projectId            }        }         return nil    }     private struct CodeAssistStatus {        let tier: GeminiUserTierId?        let projectId: String?         static let empty = CodeAssistStatus(tier: nil, projectId: nil)    }     private static func loadCodeAssistStatus(        accessToken: String,        timeout: TimeInterval,        dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async -> CodeAssistStatus    {        guard let url = URL(string: loadCodeAssistEndpoint) else {            self.log.warning("loadCodeAssist: invalid endpoint URL")            return .empty        }         var request = URLRequest(url: url)        request.httpMethod = "POST"        request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")        request.setValue("application/json", forHTTPHeaderField: "Content-Type")        request.httpBody = Data("{\"metadata\":{\"ideType\":\"GEMINI_CLI\",\"pluginType\":\"GEMINI\"}}".utf8)        request.timeoutInterval = timeout         let data: Data        let response: URLResponse        do {            (data, response) = try await dataLoader(request)        } catch {            Self.log.warning("loadCodeAssist: request failed", metadata: ["error": "\(error)"])            return .empty        }         guard let httpResponse = response as? HTTPURLResponse else {            Self.log.warning("loadCodeAssist: invalid response type")            return .empty        }         guard httpResponse.statusCode == 200 else {            Self.log.warning("loadCodeAssist: HTTP error", metadata: [                "statusCode": "\(httpResponse.statusCode)",                "body": String(data: data, encoding: .utf8) ?? "<binary>",            ])            return .empty        }         guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {            Self.log.warning("loadCodeAssist: failed to parse JSON", metadata: [                "body": String(data: data, encoding: .utf8) ?? "<binary>",            ])            return .empty        }         let rawProjectId: String? = {            if let project = json["cloudaicompanionProject"] as? String {                return project            }            if let project = json["cloudaicompanionProject"] as? [String: Any] {                if let projectId = project["id"] as? String {                    return projectId                }                if let projectId = project["projectId"] as? String {                    return projectId                }            }            return nil        }()        let trimmedProjectId = rawProjectId?.trimmingCharacters(in: .whitespacesAndNewlines)        let projectId = trimmedProjectId?.isEmpty == true ? nil : trimmedProjectId        if let projectId {            Self.log.info("loadCodeAssist: project detected", metadata: ["projectId": projectId])        }         let tierId = (json["currentTier"] as? [String: Any])?["id"] as? String        guard let tierId else {            Self.log.warning("loadCodeAssist: no currentTier.id in response", metadata: [                "json": "\(json)",            ])            return CodeAssistStatus(tier: nil, projectId: projectId)        }         guard let tier = GeminiUserTierId(rawValue: tierId) else {            Self.log.warning("loadCodeAssist: unknown tier ID", metadata: ["tierId": tierId])            return CodeAssistStatus(tier: nil, projectId: projectId)        }         Self.log.info("loadCodeAssist: success", metadata: ["tier": tierId, "projectId": projectId ?? "nil"])        return CodeAssistStatus(tier: tier, projectId: projectId)    }     private struct OAuthCredentials {        let accessToken: String?        let idToken: String?        let refreshToken: String?        let expiryDate: Date?    }     private struct OAuthClientCredentials {        let clientId: String        let clientSecret: String    }     private static func extractOAuthCredentials() -> OAuthClientCredentials? {        let env = ProcessInfo.processInfo.environment         // Find the gemini binary        guard let geminiPath = BinaryLocator.resolveGeminiBinary(            env: env,            loginPATH: LoginShellPathCache.shared.current)            ?? TTYCommandRunner.which("gemini")        else {            return nil        }         // Resolve symlinks to find the actual installation        let resolvedGeminiPath = URL(fileURLWithPath: geminiPath).resolvingSymlinksInPath().path         // Try the legacy layouts first — they're cheap file reads and cover the common cases        // (Homebrew, npm/bun sibling, Nix) without spawning subprocesses or walking the tree.        if let credentials = Self.extractOAuthCredentialsFromLegacyPaths(realGeminiPath: resolvedGeminiPath) {            return credentials        }         // For fnm-managed installs, ask fnm where the package lives        if Self.isLikelyFnmManagedPath(geminiPath) || Self.isLikelyFnmManagedPath(resolvedGeminiPath),

Lines 1–500 of 1154.