From 05fb99d9af47451de862f6fbe8c96239af3f42f2 Mon Sep 17 00:00:00 2001 From: iamtoruk Date: Mon, 24 Aug 2026 08:43:47 -0700 Subject: [PATCH 01/12] =?UTF-8?q?feat(menubar):=20live=20Gemini=20quota=20?= =?UTF-8?q?=E2=80=94=20service,=20Settings=20pane,=20popover=20Plan=20tab?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the desktop app's Gemini Code Assist quota flow (loadCodeAssist tier + retrieveUserQuota buckets, ~/.gemini/oauth_creds.json read-only, in-memory-only refresh) to the menubar. Plan tab and empty-state gating now treat Gemini as plan-capable. --- mac/Sources/CodeBurnMenubar/AppStore.swift | 131 ++++++ mac/Sources/CodeBurnMenubar/CodeBurnApp.swift | 29 +- .../Data/GeminiQuotaPresentation.swift | 51 +++ .../Data/GeminiSubscriptionService.swift | 376 ++++++++++++++++++ .../Views/HeatmapSection.swift | 102 ++++- .../Views/MenuBarContent.swift | 2 +- .../CodeBurnMenubar/Views/SettingsView.swift | 136 ++++++- .../GeminiQuotaTests.swift | 299 ++++++++++++++ 8 files changed, 1111 insertions(+), 15 deletions(-) create mode 100644 mac/Sources/CodeBurnMenubar/Data/GeminiQuotaPresentation.swift create mode 100644 mac/Sources/CodeBurnMenubar/Data/GeminiSubscriptionService.swift create mode 100644 mac/Tests/CodeBurnMenubarTests/GeminiQuotaTests.swift diff --git a/mac/Sources/CodeBurnMenubar/AppStore.swift b/mac/Sources/CodeBurnMenubar/AppStore.swift index f8ff5bab..f2a22b90 100644 --- a/mac/Sources/CodeBurnMenubar/AppStore.swift +++ b/mac/Sources/CodeBurnMenubar/AppStore.swift @@ -145,12 +145,20 @@ final class AppStore { // first refresh tick. var kimiLoadState: SubscriptionLoadState = KimiSubscriptionService.hasCredential ? .dormant : .notBootstrapped + var geminiUsage: GeminiUsage? + var geminiError: String? + // Same file-based activation as Kimi — reading ~/.gemini/oauth_creds.json + // is prompt-free, so we start dormant and auto-activate on the first + // refresh tick. + var geminiLoadState: SubscriptionLoadState = GeminiSubscriptionService.hasCredential ? .dormant : .notBootstrapped + /// Generation tokens for the in-flight refresh tasks. Incremented on every /// disconnect / reset so a fetch that started before the disconnect cannot /// resume after the await and re-populate the freshly-cleared state. private var claudeRefreshGen: Int = 0 private var codexRefreshGen: Int = 0 private var kimiRefreshGen: Int = 0 + private var geminiRefreshGen: Int = 0 private var cache: [PayloadCacheKey: CachedPayload] = [:] private var cacheDate: String = "" @@ -1252,6 +1260,91 @@ final class AppStore { } } + // MARK: - Gemini + + /// Same prompt-free activation as Kimi: reading the CLI's credential file + /// needs no keychain, so the first refresh tick activates dormant state. + func bootstrapGemini() async { + // Capture the generation before the await so a disconnect that lands + // mid-fetch cannot be resurrected into .loaded when the fetch returns. + let gen = geminiRefreshGen + geminiLoadState = .bootstrapping + do { + let usage = try await GeminiSubscriptionService.refresh() + guard gen == geminiRefreshGen else { return } + geminiUsage = usage + geminiError = nil + geminiLoadState = .loaded + } catch let err as GeminiSubscriptionService.FetchError { + guard gen == geminiRefreshGen else { return } + applyGeminiFetchError(err) + } catch { + guard gen == geminiRefreshGen else { return } + geminiError = sanitizeForUI(error.localizedDescription) + geminiLoadState = .failed + } + } + + func refreshGemini() async { + _ = await refreshGeminiReportingSuccess() + } + + @discardableResult + func refreshGeminiReportingSuccess() async -> Bool { + if case .dormant = geminiLoadState { + await bootstrapGemini() + return geminiLoadState == .loaded + } + guard GeminiSubscriptionService.hasCredential else { + if geminiLoadState != .notBootstrapped { geminiLoadState = .notBootstrapped } + return false + } + let gen = geminiRefreshGen + if geminiUsage == nil { geminiLoadState = .loading } + do { + let usage = try await GeminiSubscriptionService.refresh() + guard gen == geminiRefreshGen else { return false } + geminiUsage = usage + geminiError = nil + geminiLoadState = .loaded + return true + } catch let err as GeminiSubscriptionService.FetchError { + guard gen == geminiRefreshGen else { return false } + applyGeminiFetchError(err) + return false + } catch { + guard gen == geminiRefreshGen else { return false } + geminiError = sanitizeForUI(error.localizedDescription) + geminiLoadState = .failed + return false + } + } + + func disconnectGemini() { + GeminiSubscriptionService.disconnect() + geminiRefreshGen &+= 1 + geminiUsage = nil + geminiError = nil + geminiLoadState = .notBootstrapped + NotificationCenter.default.post(name: .codeBurnSubscriptionDisconnected, object: nil) + } + + private func applyGeminiFetchError(_ err: GeminiSubscriptionService.FetchError) { + let sanitized = sanitizeForUI(err.errorDescription) + geminiError = sanitized + if case .noCredentials = err { + geminiLoadState = .noCredentials + } else if err.isTerminal { + geminiLoadState = .terminalFailure(reason: sanitized) + } else if let retryAt = err.rateLimitRetryAt { + geminiLoadState = .transientFailure(retryAt: retryAt) + } else { + // 5xx / network blips back off automatically, mirroring the + // Electron provider's transientFailure mapping. + geminiLoadState = .transientFailure(retryAt: nil) + } + } + private func applyFetchError(_ err: ClaudeSubscriptionService.FetchError) { let sanitized = sanitizeForUI(err.errorDescription) subscriptionError = sanitized @@ -1324,6 +1417,10 @@ final class AppStore { let worst = max(usage.primary?.usedPercent ?? 0, usage.details.map(\.usedPercent).max() ?? 0) if worst > 0 { providers.append(("Kimi Code", worst)) } } + if let usage = geminiUsage, shouldIncludeCachedQuota(loadState: geminiLoadState) { + let worst = usage.details.map(\.usedPercent).max() ?? 0 + if worst > 0 { providers.append(("Gemini", worst)) } + } let worst = providers.map(\.percent).max() ?? 0 let severity = QuotaSummary.severity(for: worst / 100) let sorted = providers.sorted { $0.percent > $1.percent } @@ -1345,6 +1442,7 @@ final class AppStore { case .claude: return claudeQuotaSummary(filter: filter) case .codex: return codexQuotaSummary(filter: filter) case .kimiCode: return kimiQuotaSummary(filter: filter) + case .gemini: return geminiQuotaSummary(filter: filter) default: return nil } } @@ -1506,6 +1604,39 @@ final class AppStore { return QuotaSummary(providerFilter: filter, connection: connection, primary: primary, details: details, planLabel: kimiUsage?.plan ?? "Kimi Code", footerLines: []) } + private func geminiQuotaSummary(filter: ProviderFilter) -> QuotaSummary? { + if case .notBootstrapped = geminiLoadState { return nil } + if case .bootstrapping = geminiLoadState { return nil } + if case .noCredentials = geminiLoadState { return nil } + + let connection: QuotaSummary.Connection = { + switch geminiLoadState { + case .notBootstrapped, .dormant, .bootstrapping, .noCredentials: return .disconnected + case .loading: return geminiUsage == nil ? .loading : .stale + case .loaded: return .connected + case .failed: return geminiUsage == nil ? .loading : .stale + // An expired Gemini login without a refresh path sits terminal + // until the user runs the CLI. Keep the last-known bars (marked + // stale) instead of flapping the chip to a reconnect card. + case let .terminalFailure(reason): return geminiUsage == nil ? .terminalFailure(reason: reason) : .stale + case .transientFailure: return .transientFailure + } + }() + + var primary: QuotaSummary.Window? + var details: [QuotaSummary.Window] = [] + if let usage = geminiUsage { + // details is sorted most-constrained first, so the first row is + // the headline bar. + for w in usage.details { + let row = QuotaSummary.Window(label: w.label, percent: w.usedPercent / 100, resetsAt: w.resetsAt) + if primary == nil { primary = row } + details.append(row) + } + } + return QuotaSummary(providerFilter: filter, connection: connection, primary: primary, details: details, planLabel: geminiUsage?.plan ?? "Gemini", footerLines: []) + } + /// Persist one snapshot per window so we can answer "what did the prior cycle end at?" /// when the current window has just reset and projection from current data isn't meaningful. /// Also computes the effective_tokens consumed inside each 7-day window from local history, diff --git a/mac/Sources/CodeBurnMenubar/CodeBurnApp.swift b/mac/Sources/CodeBurnMenubar/CodeBurnApp.swift index fa42d332..33fb6b6f 100644 --- a/mac/Sources/CodeBurnMenubar/CodeBurnApp.swift +++ b/mac/Sources/CodeBurnMenubar/CodeBurnApp.swift @@ -530,6 +530,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSPopoverDelegate, NSM fileprivate var lastSubscriptionRefreshAt: Date? fileprivate var lastCodexRefreshAt: Date? fileprivate var lastKimiRefreshAt: Date? + fileprivate var lastGeminiRefreshAt: Date? private var claudeQuotaFailureCount = 0 private var nextClaudeQuotaRefreshAt: Date? @@ -625,23 +626,32 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSPopoverDelegate, NSM if let task = codexQuotaRefreshTask { return await task.value } - // Kimi Code rides the same tick, but with its own cadence anchor: - // when Codex is not connected its refresh returns false immediately, - // so lastCodexRefreshAt never advances — anchoring Kimi on it would - // poll api.kimi.com on every payload tick instead of the configured - // quota cadence. Anchor on attempt (not success) so a failing Kimi - // endpoint also respects the cadence. + // Kimi Code and Gemini ride the same tick, each with its own cadence + // anchor: when Codex is not connected its refresh returns false + // immediately, so lastCodexRefreshAt never advances — anchoring them + // on it would poll the quota endpoints on every payload tick instead + // of the configured quota cadence. Anchor on attempt (not success) + // so a failing endpoint also respects the cadence. let kimiDue: Bool = { let cadence = SubscriptionRefreshCadence.current guard cadence != .manual else { return false } return Date().timeIntervalSince(lastKimiRefreshAt ?? .distantPast) >= TimeInterval(cadence.rawValue) }() if kimiDue { lastKimiRefreshAt = Date() } - let task = Task { [store, kimiDue] in + let geminiDue: Bool = { + let cadence = SubscriptionRefreshCadence.current + guard cadence != .manual else { return false } + return Date().timeIntervalSince(lastGeminiRefreshAt ?? .distantPast) >= TimeInterval(cadence.rawValue) + }() + if geminiDue { lastGeminiRefreshAt = Date() } + let task = Task { [store, kimiDue, geminiDue] in async let codex = store.refreshCodexReportingSuccess() if kimiDue { _ = await store.refreshKimiReportingSuccess() } + if geminiDue { + _ = await store.refreshGeminiReportingSuccess() + } return await codex } codexQuotaRefreshTask = task @@ -861,6 +871,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSPopoverDelegate, NSM lastSubscriptionRefreshAt = nil lastCodexRefreshAt = nil lastKimiRefreshAt = nil + lastGeminiRefreshAt = nil claudeQuotaFailureCount = 0 nextClaudeQuotaRefreshAt = nil } @@ -892,11 +903,13 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSPopoverDelegate, NSM // relying on payload/menubarPayload happening to touch the same cache. _ = self.store.isOverDailyBudget // Track the live-quota state too so the flame icon re-tints on - // every subscription / codex usage update, not just every 30s. + // every connected provider's usage update, not just every 30s. _ = self.store.subscription _ = self.store.subscriptionLoadState _ = self.store.codexUsage _ = self.store.codexLoadState + _ = self.store.geminiUsage + _ = self.store.geminiLoadState } onChange: { [weak self] in DispatchQueue.main.async { guard let self else { return } diff --git a/mac/Sources/CodeBurnMenubar/Data/GeminiQuotaPresentation.swift b/mac/Sources/CodeBurnMenubar/Data/GeminiQuotaPresentation.swift new file mode 100644 index 00000000..45972ce4 --- /dev/null +++ b/mac/Sources/CodeBurnMenubar/Data/GeminiQuotaPresentation.swift @@ -0,0 +1,51 @@ +import Foundation + +/// Pure display-decision helpers for the Gemini quota surfaces. +/// +/// Gemini OAuth access tokens expire (~1h) and CodeBurn only refreshes them +/// in memory when the CLI's documented env overrides are set, so an expired +/// login without a refresh path sits in `.terminalFailure` until the user +/// runs the Gemini CLI once. Like Kimi, the always-visible surfaces keep +/// showing the last good snapshot with a quiet caption instead of flapping +/// to a reconnect screen; the reconnect screen is reserved for the no-data +/// case, where there is genuinely nothing to show. +enum GeminiQuotaPresentation { + /// Which Plan-tab subview to render, given the load state and whether a + /// last-known snapshot exists. + enum PlanContent: Equatable { + case noCredentials + case loading + case failed + case transientFailed + case reconnect(reason: String?) + /// Render the loaded usage bars. `idle` is true when the login has + /// gone terminal but a snapshot is still on hand — the caller stamps a + /// quiet "run the CLI" caption instead of hiding the data. + case usage(idle: Bool) + } + + static func planContent(loadState: SubscriptionLoadState, hasUsage: Bool) -> PlanContent { + switch loadState { + case .notBootstrapped, .noCredentials: + return .noCredentials + case .dormant, .bootstrapping: + return .loading + case .loading, .loaded: + return hasUsage ? .usage(idle: false) : .loading + case .failed: + return .failed + case .transientFailure: + return hasUsage ? .usage(idle: false) : .transientFailed + case .terminalFailure(let reason): + return hasUsage ? .usage(idle: true) : .reconnect(reason: reason) + } + } + + /// Snapshot age past which a loaded view stamps an "as of