diff --git a/Sources/CodexBar/PreferencesGeneralPane.swift b/Sources/CodexBar/PreferencesGeneralPane.swift index 39a95a55f..72b13d6a1 100644 --- a/Sources/CodexBar/PreferencesGeneralPane.swift +++ b/Sources/CodexBar/PreferencesGeneralPane.swift @@ -96,6 +96,13 @@ struct GeneralPane: View { subtitle: "Notifies when the 5-hour session quota hits 0% and when it becomes " + "available again.", binding: self.$settings.sessionQuotaNotificationsEnabled) + PreferenceToggleRow( + title: "Quota warning notifications", + subtitle: "Warns before quota runs out. Fires once per threshold per window.", + binding: self.$settings.quotaWarningNotificationsEnabled) + if self.settings.quotaWarningNotificationsEnabled { + QuotaWarningThresholdPicker(settings: self.settings) + } } Divider() @@ -163,3 +170,35 @@ struct GeneralPane: View { .foregroundStyle(.tertiary) } } + +@MainActor +private struct QuotaWarningThresholdPicker: View { + @Bindable var settings: SettingsStore + + private static let availableThresholds = [80, 50, 20, 10] + + var body: some View { + VStack(alignment: .leading, spacing: 4) { + Text("Warn at these remaining levels:") + .font(.footnote) + .foregroundStyle(.secondary) + HStack(spacing: 8) { + ForEach(Self.availableThresholds, id: \.self) { threshold in + Toggle("\(threshold)%", isOn: Binding( + get: { self.settings.quotaWarningThresholds.contains(threshold) }, + set: { isOn in + var current = self.settings.quotaWarningThresholds + if isOn { + if !current.contains(threshold) { current.append(threshold) } + } else { + current.removeAll { $0 == threshold } + } + self.settings.quotaWarningThresholds = current.sorted(by: >) + })) + .toggleStyle(.checkbox) + } + } + } + .padding(.leading, 20) + } +} diff --git a/Sources/CodexBar/QuotaWarningNotifications.swift b/Sources/CodexBar/QuotaWarningNotifications.swift new file mode 100644 index 000000000..b80edb8fb --- /dev/null +++ b/Sources/CodexBar/QuotaWarningNotifications.swift @@ -0,0 +1,67 @@ +import CodexBarCore +import Foundation + +enum QuotaWarningWindow: String, Sendable { + case primary + case secondary +} + +struct QuotaWarningEvent: Equatable, Sendable { + let threshold: Int + let window: QuotaWarningWindow + let currentRemaining: Double +} + +enum QuotaWarningNotificationLogic { + /// Returns thresholds that were newly crossed downward. + /// A threshold is considered crossed when `previousRemaining > threshold` and + /// `currentRemaining <= threshold`, and the threshold has not already fired. + static func crossedThresholds( + previousRemaining: Double?, + currentRemaining: Double, + thresholds: [Int], + alreadyFired: Set + ) -> [Int] { + guard let previousRemaining else { return [] } + + var crossed: [Int] = [] + for threshold in thresholds { + let t = Double(threshold) + if previousRemaining > t, + currentRemaining <= t, + !alreadyFired.contains(threshold) + { + crossed.append(threshold) + } + } + return crossed + } + + /// Returns thresholds that should be cleared from `alreadyFired` because + /// the remaining percentage has risen back above them. + static func restoredThresholds( + currentRemaining: Double, + alreadyFired: Set + ) -> Set { + Set(alreadyFired.filter { Double($0) < currentRemaining }) + } +} + +@MainActor +final class QuotaWarningNotifier { + private let logger = CodexBarLog.logger(LogCategories.quotaWarningNotifications) + + init() {} + + func post(event: QuotaWarningEvent, provider: UsageProvider) { + let providerName = ProviderDescriptorRegistry.descriptor(for: provider).metadata.displayName + let windowLabel = event.window == .primary ? "session" : "weekly" + + let title = "\(providerName) \(windowLabel) quota low" + let body = "\(event.threshold)% remaining – currently at \(Int(event.currentRemaining))%." + + let idPrefix = "quota-warning-\(provider.rawValue)-\(event.window.rawValue)-\(event.threshold)" + self.logger.info("enqueuing", metadata: ["prefix": idPrefix]) + AppNotifications.shared.post(idPrefix: idPrefix, title: title, body: body) + } +} diff --git a/Sources/CodexBar/SettingsStore+Defaults.swift b/Sources/CodexBar/SettingsStore+Defaults.swift index 44d83a023..1e0f3d1c7 100644 --- a/Sources/CodexBar/SettingsStore+Defaults.swift +++ b/Sources/CodexBar/SettingsStore+Defaults.swift @@ -101,6 +101,22 @@ extension SettingsStore { } } + var quotaWarningNotificationsEnabled: Bool { + get { self.defaultsState.quotaWarningNotificationsEnabled } + set { + self.defaultsState.quotaWarningNotificationsEnabled = newValue + self.userDefaults.set(newValue, forKey: "quotaWarningNotificationsEnabled") + } + } + + var quotaWarningThresholds: [Int] { + get { self.defaultsState.quotaWarningThresholdsRaw } + set { + self.defaultsState.quotaWarningThresholdsRaw = newValue + self.userDefaults.set(newValue, forKey: "quotaWarningThresholds") + } + } + var usageBarsShowUsed: Bool { get { self.defaultsState.usageBarsShowUsed } set { diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index 09f3e3caa..0dd818e6d 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -184,6 +184,10 @@ extension SettingsStore { if sessionQuotaDefault == nil { userDefaults.set(true, forKey: "sessionQuotaNotificationsEnabled") } + let quotaWarningNotificationsEnabled = userDefaults.object( + forKey: "quotaWarningNotificationsEnabled") as? Bool ?? false + let quotaWarningThresholdsRaw = userDefaults.array( + forKey: "quotaWarningThresholds") as? [Int] ?? [50, 20] let usageBarsShowUsed = userDefaults.object(forKey: "usageBarsShowUsed") as? Bool ?? false let resetTimesShowAbsolute = userDefaults.object(forKey: "resetTimesShowAbsolute") as? Bool ?? false let menuBarShowsBrandIconWithPercent = userDefaults.object( @@ -235,6 +239,8 @@ extension SettingsStore { debugKeepCLISessionsAlive: debugKeepCLISessionsAlive, statusChecksEnabled: statusChecksEnabled, sessionQuotaNotificationsEnabled: sessionQuotaNotificationsEnabled, + quotaWarningNotificationsEnabled: quotaWarningNotificationsEnabled, + quotaWarningThresholdsRaw: quotaWarningThresholdsRaw, usageBarsShowUsed: usageBarsShowUsed, resetTimesShowAbsolute: resetTimesShowAbsolute, menuBarShowsBrandIconWithPercent: menuBarShowsBrandIconWithPercent, diff --git a/Sources/CodexBar/SettingsStoreState.swift b/Sources/CodexBar/SettingsStoreState.swift index ff3d640bd..256e520a0 100644 --- a/Sources/CodexBar/SettingsStoreState.swift +++ b/Sources/CodexBar/SettingsStoreState.swift @@ -11,6 +11,8 @@ struct SettingsDefaultsState: Sendable { var debugKeepCLISessionsAlive: Bool var statusChecksEnabled: Bool var sessionQuotaNotificationsEnabled: Bool + var quotaWarningNotificationsEnabled: Bool + var quotaWarningThresholdsRaw: [Int] var usageBarsShowUsed: Bool var resetTimesShowAbsolute: Bool var menuBarShowsBrandIconWithPercent: Bool diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index 1e60dfb37..1dc4c1ee0 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -25,6 +25,8 @@ extension UsageStore { self.statuses.removeValue(forKey: provider) self.lastKnownSessionRemaining.removeValue(forKey: provider) self.lastKnownSessionWindowSource.removeValue(forKey: provider) + self.lastKnownSecondaryRemaining.removeValue(forKey: provider) + self.quotaWarningFiredThresholds.removeValue(forKey: provider) self.lastTokenFetchAt.removeValue(forKey: provider) } return @@ -79,6 +81,7 @@ extension UsageStore { case let .success(result): let scoped = result.usage.scoped(to: provider) await MainActor.run { + self.handleQuotaWarningTransition(provider: provider, snapshot: scoped) self.handleSessionQuotaTransition(provider: provider, snapshot: scoped) self.snapshots[provider] = scoped self.lastSourceLabels[provider] = result.sourceLabel diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 3e55ffa9f..fb6d69fce 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -156,6 +156,11 @@ extension UsageStore { scoped } await MainActor.run { + if let account, self.lastWarningAccountID[provider] != account.id { + self.resetQuotaWarningState(for: provider) + self.lastWarningAccountID[provider] = account.id + } + self.handleQuotaWarningTransition(provider: provider, snapshot: labeled) self.handleSessionQuotaTransition(provider: provider, snapshot: labeled) self.snapshots[provider] = labeled self.lastSourceLabels[provider] = result.sourceLabel diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index e3f0bcae4..7d84a807b 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -39,6 +39,8 @@ extension UsageStore { _ = self.settings.refreshFrequency _ = self.settings.statusChecksEnabled _ = self.settings.sessionQuotaNotificationsEnabled + _ = self.settings.quotaWarningNotificationsEnabled + _ = self.settings.quotaWarningThresholds _ = self.settings.usageBarsShowUsed _ = self.settings.costUsageEnabled _ = self.settings.randomBlinkEnabled @@ -152,6 +154,11 @@ final class UsageStore { @ObservationIgnored var codexHistoricalDatasetAccountKey: String? @ObservationIgnored var lastKnownSessionRemaining: [UsageProvider: Double] = [:] @ObservationIgnored var lastKnownSessionWindowSource: [UsageProvider: SessionQuotaWindowSource] = [:] + @ObservationIgnored var lastKnownSecondaryRemaining: [UsageProvider: Double] = [:] + @ObservationIgnored var quotaWarningFiredThresholds: [UsageProvider: [QuotaWarningWindow: Set]] = [:] + @ObservationIgnored var lastWarningAccountID: [UsageProvider: UUID] = [:] + @ObservationIgnored private let quotaWarningNotifier: QuotaWarningNotifier + @ObservationIgnored private let quotaWarningLogger = CodexBarLog.logger(LogCategories.quotaWarning) @ObservationIgnored var lastTokenFetchAt: [UsageProvider: Date] = [:] @ObservationIgnored private var hasCompletedInitialRefresh: Bool = false @ObservationIgnored private let tokenFetchTTL: TimeInterval = 60 * 60 @@ -167,6 +174,7 @@ final class UsageStore { registry: ProviderRegistry = .shared, historicalUsageHistoryStore: HistoricalUsageHistoryStore = HistoricalUsageHistoryStore(), sessionQuotaNotifier: any SessionQuotaNotifying = SessionQuotaNotifier(), + quotaWarningNotifier: QuotaWarningNotifier = QuotaWarningNotifier(), startupBehavior: StartupBehavior = .automatic) { self.codexFetcher = fetcher @@ -177,6 +185,7 @@ final class UsageStore { self.registry = registry self.historicalUsageHistoryStore = historicalUsageHistoryStore self.sessionQuotaNotifier = sessionQuotaNotifier + self.quotaWarningNotifier = quotaWarningNotifier self.startupBehavior = startupBehavior.resolved(isRunningTests: Self.isRunningTestsProcess()) self.providerMetadata = registry.metadata self @@ -589,6 +598,87 @@ final class UsageStore { self.sessionQuotaNotifier.post(transition: transition, provider: provider, badge: nil) } + func handleQuotaWarningTransition(provider: UsageProvider, snapshot: UsageSnapshot) { + // Always update tracking state so baselines stay current even when + // notifications are disabled. This prevents retroactive alerts when + // the user re-enables notifications after quota has dropped. + defer { + if let secondary = snapshot.secondary { + self.lastKnownSecondaryRemaining[provider] = secondary.remainingPercent + } + } + + guard self.settings.quotaWarningNotificationsEnabled else { return } + let thresholds = self.settings.quotaWarningThresholds.sorted(by: >) + guard !thresholds.isEmpty else { return } + + if let primary = snapshot.primary { + self.processWarningWindow( + provider: provider, + window: .primary, + currentRemaining: primary.remainingPercent, + previousRemaining: self.lastKnownSessionRemaining[provider], + thresholds: thresholds) + } + + if let secondary = snapshot.secondary { + self.processWarningWindow( + provider: provider, + window: .secondary, + currentRemaining: secondary.remainingPercent, + previousRemaining: self.lastKnownSecondaryRemaining[provider], + thresholds: thresholds) + } + } + + /// Resets quota warning baselines and fired thresholds for a provider. + /// Call when the selected token account changes to avoid false alerts + /// caused by a different account's remaining quota. + func resetQuotaWarningState(for provider: UsageProvider) { + self.lastKnownSecondaryRemaining.removeValue(forKey: provider) + self.quotaWarningFiredThresholds.removeValue(forKey: provider) + } + + private func processWarningWindow( + provider: UsageProvider, + window: QuotaWarningWindow, + currentRemaining: Double, + previousRemaining: Double?, + thresholds: [Int]) + { + var firedSet = self.quotaWarningFiredThresholds[provider]?[window] ?? [] + + let restored = QuotaWarningNotificationLogic.restoredThresholds( + currentRemaining: currentRemaining, alreadyFired: firedSet) + firedSet.subtract(restored) + + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: previousRemaining, + currentRemaining: currentRemaining, + thresholds: thresholds, + alreadyFired: firedSet) + + for threshold in crossed { + let event = QuotaWarningEvent( + threshold: threshold, window: window, currentRemaining: currentRemaining) + self.quotaWarningLogger.info( + "threshold crossed", + metadata: [ + "provider": provider.rawValue, + "window": window.rawValue, + "threshold": "\(threshold)", + "remaining": "\(currentRemaining)", + ]) + self.quotaWarningNotifier.post(event: event, provider: provider) + firedSet.insert(threshold) + } + + if self.quotaWarningFiredThresholds[provider] == nil { + self.quotaWarningFiredThresholds[provider] = [:] + } + self.quotaWarningFiredThresholds[provider]?[window] = firedSet + } + private func refreshStatus(_ provider: UsageProvider) async { guard self.settings.statusChecksEnabled else { return } guard let meta = self.providerMetadata[provider] else { return } diff --git a/Sources/CodexBarCore/Logging/LogCategories.swift b/Sources/CodexBarCore/Logging/LogCategories.swift index 37a7726ef..352839f62 100644 --- a/Sources/CodexBarCore/Logging/LogCategories.swift +++ b/Sources/CodexBarCore/Logging/LogCategories.swift @@ -45,6 +45,8 @@ public enum LogCategories { public static let openRouterUsage = "openrouter-usage" public static let providerDetection = "provider-detection" public static let providers = "providers" + public static let quotaWarning = "quotaWarning" + public static let quotaWarningNotifications = "quotaWarningNotifications" public static let sessionQuota = "sessionQuota" public static let sessionQuotaNotifications = "sessionQuotaNotifications" public static let settings = "settings" diff --git a/Tests/CodexBarTests/QuotaWarningNotificationLogicTests.swift b/Tests/CodexBarTests/QuotaWarningNotificationLogicTests.swift new file mode 100644 index 000000000..052a68e37 --- /dev/null +++ b/Tests/CodexBarTests/QuotaWarningNotificationLogicTests.swift @@ -0,0 +1,93 @@ +import Testing +@testable import CodexBar + +@Suite +struct QuotaWarningNotificationLogicTests { + @Test + func doesNothingWithoutPreviousValue() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: nil, currentRemaining: 40, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed.isEmpty) + } + + @Test + func detectsSingleThresholdCrossing() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 55, currentRemaining: 45, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed == [50]) + } + + @Test + func detectsMultipleThresholdCrossings() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 55, currentRemaining: 15, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed == [50, 20]) + } + + @Test + func doesNotRefireAlreadyFiredThreshold() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 55, currentRemaining: 45, + thresholds: [80, 50, 20], alreadyFired: [50]) + #expect(crossed.isEmpty) + } + + @Test + func clearsRestoredThresholds() { + let restored = QuotaWarningNotificationLogic.restoredThresholds( + currentRemaining: 60, alreadyFired: [80, 50, 20]) + #expect(restored == [50, 20]) + } + + @Test + func doesNotCrossOnUpwardMovement() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 45, currentRemaining: 55, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed.isEmpty) + } + + @Test + func handlesExactThresholdBoundary() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 50.001, currentRemaining: 50.0, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed == [50]) + } + + @Test + func handlesEmptyThresholds() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 80, currentRemaining: 10, + thresholds: [], alreadyFired: []) + #expect(crossed.isEmpty) + } + + @Test + func doesNotCrossWhenStayingAboveThreshold() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 90, currentRemaining: 85, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed.isEmpty) + } + + @Test + func doesNotRestoreThresholdAboveCurrentRemaining() { + let restored = QuotaWarningNotificationLogic.restoredThresholds( + currentRemaining: 60, alreadyFired: [80, 50, 20]) + #expect(!restored.contains(80)) + #expect(restored.contains(50)) + #expect(restored.contains(20)) + } + + @Test + func crossesAllThresholdsOnLargeDrop() { + let crossed = QuotaWarningNotificationLogic.crossedThresholds( + previousRemaining: 100, currentRemaining: 5, + thresholds: [80, 50, 20], alreadyFired: []) + #expect(crossed == [80, 50, 20]) + } +}