Preserve Claude widget quota on partial refresh (#2480 by @ChenZiHong-Gavin)

Sources/CodexBar/UsageStore+BackgroundRefresh.swift

@@ -50,6 +50,7 @@ extension UsageStore {
             self.kiloScopeSnapshots = []
         }
         if provider == .claude {
+            self.widgetUsagePreservationBlockedProviders.insert(provider)
             self.clearClaudeSwapAccountState()
         }
         self.clearTokenSnapshot(for: provider)

Sources/CodexBar/UsageStore+Refresh.swift

@@ -556,6 +556,7 @@ extension UsageStore {
             }
             self.lastKnownResetSnapshots[provider] = backfilled
             self.snapshots[provider] = backfilled
+            self.widgetUsagePreservationBlockedProviders.remove(provider)
             if provider == .deepseek {
                 self.clearDeepSeekProfileTransition()
             }
@@ -969,6 +970,7 @@ extension UsageStore {
     }
 
     private func clearClaudeCredentialDerivedStateForCredentialSwapNow() {
+        self.widgetUsagePreservationBlockedProviders.insert(.claude)
         self.snapshots.removeValue(forKey: .claude)
         self.lastKnownResetSnapshots.removeValue(forKey: .claude)
         self.errors[.claude] = nil
@@ -1049,6 +1051,7 @@ extension UsageStore {
                 self.lastSourceLabels.removeValue(forKey: provider)
                 self.errors[provider] = nil
                 self.knownLimitsAvailabilityByProvider[provider] = .unavailable
+                self.widgetUsagePreservationBlockedProviders.insert(provider)
                 self.failureGates[provider]?.reset()
                 return
             }

Sources/CodexBar/UsageStore+TokenAccounts.swift

@@ -1554,6 +1554,7 @@ extension UsageStore {
                     accountDiscriminatorOverride: provider == .claude ? warningAccountDiscriminator : nil)
                 self.lastKnownResetSnapshots[provider] = backfilled
                 self.snapshots[provider] = backfilled
+                self.widgetUsagePreservationBlockedProviders.remove(provider)
                 if provider == .deepseek {
                     self.clearDeepSeekProfileTransition()
                 }

Sources/CodexBar/UsageStore+WidgetSnapshot.swift

@@ -6,7 +6,17 @@ import WidgetKit
 
 extension UsageStore {
     func persistWidgetSnapshot(reason: String) {
-        let snapshot = self.makeWidgetSnapshot()
+        // A fresh process has token-cost data before a user-authorized Claude OAuth refresh can run.
+        // Keep the last queued snapshot in memory so back-to-back writes cannot race the on-disk cache.
+        let previousSnapshot = self.lastQueuedWidgetSnapshot ?? {
+            #if DEBUG
+            // Snapshot-save overrides must stay isolated from a developer's real app-group data.
+            guard self._test_widgetSnapshotSaveOverride == nil else { return nil }
+            #endif
+            return WidgetSnapshotStore.load()
+        }()
+        let snapshot = self.makeWidgetSnapshot(previousSnapshot: previousSnapshot)
+        self.lastQueuedWidgetSnapshot = snapshot
         let previousTask = self.widgetSnapshotPersistTask
         self.widgetSnapshotPersistTask = Task { @MainActor in
             _ = await previousTask?.result
@@ -25,11 +35,14 @@ extension UsageStore {
         }
     }
 
-    private func makeWidgetSnapshot() -> WidgetSnapshot {
+    private func makeWidgetSnapshot(previousSnapshot: WidgetSnapshot?) -> WidgetSnapshot {
         let now = Date()
         let enabledProviders = self.enabledProviders()
         let entries = UsageProvider.allCases.compactMap { provider in
-            self.makeWidgetEntry(for: provider, now: now)
+            self.makeWidgetEntry(
+                for: provider,
+                now: now,
+                previousEntry: previousSnapshot?.entries.first { $0.provider == provider })
         }
         return WidgetSnapshot(
             entries: entries,
@@ -38,10 +51,36 @@ extension UsageStore {
             generatedAt: now)
     }
 
-    private func makeWidgetEntry(for provider: UsageProvider, now: Date) -> WidgetSnapshot.ProviderEntry? {
+    private func makeWidgetEntry(
+        for provider: UsageProvider,
+        now: Date,
+        previousEntry: WidgetSnapshot.ProviderEntry?) -> WidgetSnapshot.ProviderEntry?
+    {
         let snapshot = self.snapshots[provider]
         let storedTokenSnapshot = self.tokenSnapshotForCurrentProviderConfig(for: provider)?.snapshot
-        guard snapshot != nil || (provider == .claude && storedTokenSnapshot != nil) else { return nil }
+        let claudeQuotaOwnerKey: String? = if provider == .claude {
+            self.claudeWidgetQuotaOwnerKey()
+        } else {
+            nil
+        }
+        let preservedClaudeUsage: PreservedClaudeWidgetUsage? = if provider == .claude,
+                                                                   snapshot == nil,
+                                                                   !self.widgetUsagePreservationBlockedProviders
+                                                                       .contains(provider),
+                                                                       self.knownLimitsAvailabilityByProvider[provider]?
+                                                                           .isUnavailable != true
+        {
+            Self.preservedClaudeWidgetUsage(
+                from: previousEntry,
+                expectedQuotaOwnerKey: claudeQuotaOwnerKey)
+        } else {
+            nil
+        }
+        guard snapshot != nil ||
+            (provider == .claude && (storedTokenSnapshot != nil || preservedClaudeUsage != nil))
+        else {
+            return nil
+        }
 
         let tokenSnapshot = storedTokenSnapshot
         let dailyUsage = tokenSnapshot?.daily.map { entry in
@@ -52,7 +91,9 @@ extension UsageStore {
         } ?? []
 
         let tokenUsage = Self.widgetTokenUsageSummary(from: tokenSnapshot, provider: provider)
-        let usageRows = snapshot.map { self.widgetUsageRows(provider: provider, snapshot: $0, now: now) } ?? []
+        let usageRows = snapshot.map {
+            self.widgetUsageRows(provider: provider, snapshot: $0, now: now)
+        } ?? preservedClaudeUsage?.usageRows ?? []
 
         let creditsRemaining: Double?
         let codeReviewRemaining: Double?
@@ -75,19 +116,82 @@ extension UsageStore {
         } else {
             nil
         }
+        let quotaOwnerKey: String? = if provider == .claude {
+            snapshot != nil ? claudeQuotaOwnerKey : preservedClaudeUsage?.quotaOwnerKey
+        } else {
+            nil
+        }
 
         return WidgetSnapshot.ProviderEntry(
             provider: provider,
-            updatedAt: snapshot?.updatedAt ?? tokenSnapshot?.updatedAt ?? now,
-            primary: snapshot?.primary,
-            secondary: snapshot?.secondary,
-            tertiary: snapshot?.tertiary,
+            updatedAt: snapshot?.updatedAt ?? preservedClaudeUsage?.updatedAt ?? tokenSnapshot?.updatedAt ?? now,
+            primary: snapshot?.primary ?? preservedClaudeUsage?.primary,
+            secondary: snapshot?.secondary ?? preservedClaudeUsage?.secondary,
+            tertiary: snapshot?.tertiary ?? preservedClaudeUsage?.tertiary,
             usageRows: usageRows,
             creditsRemaining: creditsRemaining,
             codeReviewRemainingPercent: codeReviewRemaining,
             tokenUsage: tokenUsage,
             dailyUsage: dailyUsage,
-            providerCost: providerCost)
+            providerCost: providerCost,
+            quotaOwnerKey: quotaOwnerKey)
+    }
+
+    private struct PreservedClaudeWidgetUsage {
+        let updatedAt: Date
+        let primary: RateWindow?
+        let secondary: RateWindow?
+        let tertiary: RateWindow?
+        let usageRows: [WidgetSnapshot.WidgetUsageRowSnapshot]?
+        let quotaOwnerKey: String?
+    }
+
+    private func claudeWidgetQuotaOwnerKey() -> String {
+        if let account = self.settings.effectiveSelectedTokenAccount(for: .claude) {
+            return self.tokenAccountSnapshotCacheKey(provider: .claude, account: account)
+        }
+        let environment = ProviderRegistry.makeEnvironment(
+            base: self.environmentBase,
+            provider: .claude,
+            settings: self.settings,
+            tokenOverride: nil)
+        return ClaudeOAuthCredentialsStore.credentialsProfileIdentifier(environment: environment)
+    }
+
+    private nonisolated static func preservedClaudeWidgetUsage(
+        from entry: WidgetSnapshot.ProviderEntry?,
+        expectedQuotaOwnerKey: String?) -> PreservedClaudeWidgetUsage?
+    {
+        guard let entry, entry.provider == .claude else { return nil }
+        guard let expectedQuotaOwnerKey,
+              let quotaOwnerKey = entry.quotaOwnerKey,
+              quotaOwnerKey == expectedQuotaOwnerKey
+        else {
+            return nil
+        }
+
+        let primary = entry.primary?.isSyntheticPlaceholder == true ? nil : entry.primary
+        let secondary = entry.secondary?.isSyntheticPlaceholder == true ? nil : entry.secondary
+        let tertiary = entry.tertiary?.isSyntheticPlaceholder == true ? nil : entry.tertiary
+        let usageRows = entry.usageRows?.filter { row in
+            guard row.window?.isSyntheticPlaceholder != true else { return false }
+            return switch row.id {
+            case "primary": primary != nil
+            case "secondary": secondary != nil
+            case "tertiary": tertiary != nil
+            default: row.percentLeft != nil
+            }
+        }
+        guard primary != nil || secondary != nil || tertiary != nil || usageRows?.isEmpty == false else {
+            return nil
+        }
+        return PreservedClaudeWidgetUsage(
+            updatedAt: entry.updatedAt,
+            primary: primary,
+            secondary: secondary,
+            tertiary: tertiary,
+            usageRows: usageRows,
+            quotaOwnerKey: quotaOwnerKey)
     }
 
     nonisolated static func widgetTokenUsageSummary(

Sources/CodexBar/UsageStore.swift

@@ -268,6 +268,8 @@ final class UsageStore {
     @ObservationIgnored var _test_startupConnectivityRetrySleepOverride: (@MainActor (
         TimeInterval) async throws -> Void)?
     @ObservationIgnored var widgetSnapshotPersistTask: Task<Void, Never>?
+    @ObservationIgnored var lastQueuedWidgetSnapshot: WidgetSnapshot?
+    @ObservationIgnored var widgetUsagePreservationBlockedProviders: Set<UsageProvider> = []
 
     @ObservationIgnored let codexFetcher: UsageFetcher
     @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching

Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift

@@ -2911,7 +2911,7 @@ public enum ClaudeOAuthCredentialsStore {
         return ClaudeConfigPaths.credentialsURL(environment: environment)
     }
 
-    static func credentialsProfileIdentifier(environment: [String: String]) -> String {
+    public static func credentialsProfileIdentifier(environment: [String: String]) -> String {
         #if DEBUG
         if let override = self.taskCredentialsProfileIdentifierOverride {
             return override

Sources/CodexBarCore/WidgetSnapshot.swift

@@ -27,6 +27,7 @@ public struct WidgetSnapshot: Codable, Sendable {
         public let tokenUsage: TokenUsageSummary?
         public let dailyUsage: [DailyUsagePoint]
         public let providerCost: ProviderCostSnapshot?
+        public let quotaOwnerKey: String?
 
         public init(
             provider: UsageProvider,
@@ -39,7 +40,8 @@ public struct WidgetSnapshot: Codable, Sendable {
             codeReviewRemainingPercent: Double?,
             tokenUsage: TokenUsageSummary?,
             dailyUsage: [DailyUsagePoint],
-            providerCost: ProviderCostSnapshot? = nil)
+            providerCost: ProviderCostSnapshot? = nil,
+            quotaOwnerKey: String? = nil)
         {
             self.provider = provider
             self.updatedAt = updatedAt
@@ -52,6 +54,7 @@ public struct WidgetSnapshot: Codable, Sendable {
             self.tokenUsage = tokenUsage
             self.dailyUsage = dailyUsage
             self.providerCost = providerCost
+            self.quotaOwnerKey = quotaOwnerKey
         }
     }