diff --git a/Sources/CodexBar/UsageStore+BackgroundRefresh.swift b/Sources/CodexBar/UsageStore+BackgroundRefresh.swift index 2a87229bdd..5fefcd1fe8 100644 --- a/Sources/CodexBar/UsageStore+BackgroundRefresh.swift +++ b/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) diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index 9b01f15b66..a133454949 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/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 } diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index e973092515..219251ec02 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/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() } diff --git a/Sources/CodexBar/UsageStore+WidgetSnapshot.swift b/Sources/CodexBar/UsageStore+WidgetSnapshot.swift index 2b5cdfe981..d45afc0e28 100644 --- a/Sources/CodexBar/UsageStore+WidgetSnapshot.swift +++ b/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( diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 9223f3cdc1..0e5ae769de 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -268,6 +268,8 @@ final class UsageStore { @ObservationIgnored var _test_startupConnectivityRetrySleepOverride: (@MainActor ( TimeInterval) async throws -> Void)? @ObservationIgnored var widgetSnapshotPersistTask: Task? + @ObservationIgnored var lastQueuedWidgetSnapshot: WidgetSnapshot? + @ObservationIgnored var widgetUsagePreservationBlockedProviders: Set = [] @ObservationIgnored let codexFetcher: UsageFetcher @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift index 560dc9db65..7f51a1a991 100644 --- a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift +++ b/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 diff --git a/Sources/CodexBarCore/WidgetSnapshot.swift b/Sources/CodexBarCore/WidgetSnapshot.swift index c30e407c47..1080856a10 100644 --- a/Sources/CodexBarCore/WidgetSnapshot.swift +++ b/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 } } diff --git a/Tests/CodexBarTests/UsageStoreWidgetSnapshotAccountTests.swift b/Tests/CodexBarTests/UsageStoreWidgetSnapshotAccountTests.swift new file mode 100644 index 0000000000..6b25694f72 --- /dev/null +++ b/Tests/CodexBarTests/UsageStoreWidgetSnapshotAccountTests.swift @@ -0,0 +1,322 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +@MainActor +struct UsageStoreWidgetSnapshotAccountTests { + @Test + func `legacy ownerless Claude quota is not preserved after upgrade`() async throws { + let suite = "UsageStoreWidgetSnapshotTests-claude-legacy-ownerless-drops-quota" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings.statusChecksEnabled = false + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings) + let quotaUpdatedAt = Date(timeIntervalSince1970: 1_800_000_000) + let tokenUpdatedAt = quotaUpdatedAt.addingTimeInterval(60) + let quota = RateWindow( + usedPercent: 28, + windowMinutes: 300, + resetsAt: quotaUpdatedAt.addingTimeInterval(3600), + resetDescription: nil) + store.lastQueuedWidgetSnapshot = WidgetSnapshot( + entries: [ + WidgetSnapshot.ProviderEntry( + provider: .claude, + updatedAt: quotaUpdatedAt, + primary: quota, + secondary: nil, + tertiary: nil, + usageRows: [ + WidgetSnapshot.WidgetUsageRowSnapshot( + id: "primary", + title: "Session", + percentLeft: 72), + ], + creditsRemaining: nil, + codeReviewRemainingPercent: nil, + tokenUsage: nil, + dailyUsage: []), + ], + enabledProviders: [.claude], + generatedAt: quotaUpdatedAt) + store._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 4300, + sessionCostUSD: 1.50, + last30DaysTokens: 43000, + last30DaysCostUSD: 13.50, + daily: [], + updatedAt: tokenUpdatedAt), + provider: .claude) + + var widgetSnapshots: [WidgetSnapshot] = [] + store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) } + defer { store._test_widgetSnapshotSaveOverride = nil } + + store.persistWidgetSnapshot(reason: "claude-legacy-ownerless-drops-quota-test") + await store.widgetSnapshotPersistTask?.value + + let entry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(entry.primary == nil) + #expect(entry.usageRows?.isEmpty == true) + #expect(entry.quotaOwnerKey == nil) + #expect(entry.tokenUsage?.sessionTokens == 4300) + } + + @Test + func `widget snapshot does not preserve Claude quota across OAuth profiles`() async throws { + let suiteA = "UsageStoreWidgetSnapshotTests-claude-oauth-profile-a" + let defaultsA = try #require(UserDefaults(suiteName: suiteA)) + defaultsA.removePersistentDomain(forName: suiteA) + let settingsA = SettingsStore( + userDefaults: defaultsA, + configStore: testConfigStore(suiteName: suiteA), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settingsA.statusChecksEnabled = false + + let environmentA = ["CLAUDE_CONFIG_DIR": "/tmp/codexbar-widget-profile-a"] + let storeA = UsageStore( + fetcher: UsageFetcher(environment: environmentA), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settingsA, + environmentBase: environmentA) + let quotaUpdatedAt = Date(timeIntervalSince1970: 1_800_000_000) + let tokenUpdatedAt = quotaUpdatedAt.addingTimeInterval(60) + let quota = RateWindow( + usedPercent: 28, + windowMinutes: 300, + resetsAt: quotaUpdatedAt.addingTimeInterval(3600), + resetDescription: nil) + storeA._setSnapshotForTesting( + UsageSnapshot(primary: quota, secondary: nil, updatedAt: quotaUpdatedAt), + provider: .claude) + + var profileASnapshots: [WidgetSnapshot] = [] + storeA._test_widgetSnapshotSaveOverride = { profileASnapshots.append($0) } + defer { storeA._test_widgetSnapshotSaveOverride = nil } + ClaudeOAuthCredentialsStore.withEnvironmentCredentialsURLForTesting { + storeA.persistWidgetSnapshot(reason: "claude-oauth-profile-a-primes-quota-test") + } + await storeA.widgetSnapshotPersistTask?.value + let profileAEntry = try #require(profileASnapshots.last?.entries.first { $0.provider == .claude }) + let profileAOwner = try #require(profileAEntry.quotaOwnerKey) + + let suiteB = "UsageStoreWidgetSnapshotTests-claude-oauth-profile-b" + let defaultsB = try #require(UserDefaults(suiteName: suiteB)) + defaultsB.removePersistentDomain(forName: suiteB) + let settingsB = SettingsStore( + userDefaults: defaultsB, + configStore: testConfigStore(suiteName: suiteB), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settingsB.statusChecksEnabled = false + + let environmentB = ["CLAUDE_CONFIG_DIR": "/tmp/codexbar-widget-profile-b"] + let profileBOwner = ClaudeOAuthCredentialsStore.withEnvironmentCredentialsURLForTesting { + ClaudeOAuthCredentialsStore.credentialsProfileIdentifier(environment: environmentB) + } + #expect(profileAOwner != profileBOwner) + + let storeB = UsageStore( + fetcher: UsageFetcher(environment: environmentB), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settingsB, + environmentBase: environmentB) + storeB.lastQueuedWidgetSnapshot = WidgetSnapshot( + entries: [profileAEntry], + enabledProviders: [.claude], + generatedAt: quotaUpdatedAt) + storeB._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 4300, + sessionCostUSD: 1.50, + last30DaysTokens: 43000, + last30DaysCostUSD: 13.50, + daily: [], + updatedAt: tokenUpdatedAt), + provider: .claude) + + var profileBSnapshots: [WidgetSnapshot] = [] + storeB._test_widgetSnapshotSaveOverride = { profileBSnapshots.append($0) } + defer { storeB._test_widgetSnapshotSaveOverride = nil } + ClaudeOAuthCredentialsStore.withEnvironmentCredentialsURLForTesting { + storeB.persistWidgetSnapshot(reason: "claude-oauth-profile-change-drops-quota-test") + } + await storeB.widgetSnapshotPersistTask?.value + + let entry = try #require(profileBSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(entry.primary == nil) + #expect(entry.usageRows?.isEmpty == true) + #expect(entry.quotaOwnerKey == nil) + #expect(entry.tokenUsage?.sessionTokens == 4300) + } + + @Test + func `widget snapshot does not preserve Claude quota after selected account changes`() async throws { + let suite = "UsageStoreWidgetSnapshotTests-claude-account-change-drops-quota" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings.statusChecksEnabled = false + settings.addTokenAccount(provider: .claude, label: "Primary", token: "primary-token") + settings.addTokenAccount(provider: .claude, label: "Secondary", token: "secondary-token") + settings.setActiveTokenAccountIndex(0, for: .claude) + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings) + let accounts = settings.tokenAccounts(for: .claude) + let primaryAccount = try #require(accounts.first) + let quotaOwnerKey = store.tokenAccountSnapshotCacheKey(provider: .claude, account: primaryAccount) + let quotaUpdatedAt = Date(timeIntervalSince1970: 1_800_000_000) + let tokenUpdatedAt = quotaUpdatedAt.addingTimeInterval(60) + let quota = RateWindow( + usedPercent: 28, + windowMinutes: 300, + resetsAt: quotaUpdatedAt.addingTimeInterval(3600), + resetDescription: nil) + store.lastQueuedWidgetSnapshot = WidgetSnapshot( + entries: [ + WidgetSnapshot.ProviderEntry( + provider: .claude, + updatedAt: quotaUpdatedAt, + primary: quota, + secondary: nil, + tertiary: nil, + usageRows: [ + WidgetSnapshot.WidgetUsageRowSnapshot( + id: "primary", + title: "Session", + percentLeft: 72), + ], + creditsRemaining: nil, + codeReviewRemainingPercent: nil, + tokenUsage: nil, + dailyUsage: [], + quotaOwnerKey: quotaOwnerKey), + ], + enabledProviders: [.claude], + generatedAt: quotaUpdatedAt) + + settings.setActiveTokenAccountIndex(1, for: .claude) + store._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 4300, + sessionCostUSD: 1.50, + last30DaysTokens: 43000, + last30DaysCostUSD: 13.50, + daily: [], + updatedAt: tokenUpdatedAt), + provider: .claude) + + var widgetSnapshots: [WidgetSnapshot] = [] + store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) } + defer { store._test_widgetSnapshotSaveOverride = nil } + + store.persistWidgetSnapshot(reason: "claude-account-change-drops-quota-test") + await store.widgetSnapshotPersistTask?.value + + let entry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(entry.primary == nil) + #expect(entry.usageRows?.isEmpty == true) + #expect(entry.quotaOwnerKey == nil) + #expect(entry.tokenUsage?.sessionTokens == 4300) + } + + @Test + func `stacked Claude account success re-enables quota preservation`() async throws { + let suite = "UsageStoreWidgetSnapshotTests-claude-stacked-success-unblocks-quota" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings.statusChecksEnabled = false + settings.multiAccountMenuLayout = .stacked + settings.addTokenAccount(provider: .claude, label: "Primary", token: "primary-token") + settings.addTokenAccount(provider: .claude, label: "Secondary", token: "secondary-token") + settings.setActiveTokenAccountIndex(0, for: .claude) + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings) + let account = try #require(settings.effectiveSelectedTokenAccount(for: .claude)) + let quotaUpdatedAt = Date(timeIntervalSince1970: 1_800_000_000) + let tokenUpdatedAt = quotaUpdatedAt.addingTimeInterval(60) + let quota = RateWindow( + usedPercent: 28, + windowMinutes: 300, + resetsAt: quotaUpdatedAt.addingTimeInterval(3600), + resetDescription: nil) + let outcome = ProviderFetchOutcome( + result: .success(ProviderFetchResult( + usage: UsageSnapshot( + primary: quota, + secondary: nil, + updatedAt: quotaUpdatedAt), + credits: nil, + dashboard: nil, + sourceLabel: "fixture", + strategyID: "fixture.api-token", + strategyKind: .apiToken)), + attempts: []) + + store.widgetUsagePreservationBlockedProviders.insert(.claude) + await store.applySelectedOutcome( + outcome, + provider: .claude, + account: account, + fallbackSnapshot: nil) + #expect(!store.widgetUsagePreservationBlockedProviders.contains(.claude)) + + var widgetSnapshots: [WidgetSnapshot] = [] + store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) } + defer { store._test_widgetSnapshotSaveOverride = nil } + + store.persistWidgetSnapshot(reason: "claude-stacked-success-primes-quota-test") + await store.widgetSnapshotPersistTask?.value + let freshEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(freshEntry.primary == quota) + #expect(freshEntry.quotaOwnerKey != nil) + + store.snapshots.removeValue(forKey: .claude) + store._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 4300, + sessionCostUSD: 1.50, + last30DaysTokens: 43000, + last30DaysCostUSD: 13.50, + daily: [], + updatedAt: tokenUpdatedAt), + provider: .claude) + store.persistWidgetSnapshot(reason: "claude-stacked-success-preserves-quota-test") + await store.widgetSnapshotPersistTask?.value + + let preservedEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(preservedEntry.primary == quota) + #expect(preservedEntry.quotaOwnerKey == freshEntry.quotaOwnerKey) + #expect(preservedEntry.tokenUsage?.sessionTokens == 4300) + } +} diff --git a/Tests/CodexBarTests/UsageStoreWidgetSnapshotTests.swift b/Tests/CodexBarTests/UsageStoreWidgetSnapshotTests.swift index bf05168caf..d7bb0566a2 100644 --- a/Tests/CodexBarTests/UsageStoreWidgetSnapshotTests.swift +++ b/Tests/CodexBarTests/UsageStoreWidgetSnapshotTests.swift @@ -369,6 +369,155 @@ struct UsageStoreWidgetSnapshotTests { #expect(entry.tokenUsage?.last30DaysTokens == 42000) } + @Test + func `widget snapshot preserves prior Claude quota rows during token only refresh`() async throws { + let suite = "UsageStoreWidgetSnapshotTests-claude-token-only-preserves-quota" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings.statusChecksEnabled = false + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings) + let quotaUpdatedAt = Date(timeIntervalSince1970: 1_800_000_000) + let tokenUpdatedAt = quotaUpdatedAt.addingTimeInterval(60) + let primary = RateWindow( + usedPercent: 28, + windowMinutes: 300, + resetsAt: quotaUpdatedAt.addingTimeInterval(3600), + resetDescription: nil) + let secondary = RateWindow( + usedPercent: 12, + windowMinutes: 10080, + resetsAt: quotaUpdatedAt.addingTimeInterval(86400), + resetDescription: nil) + store._setSnapshotForTesting( + UsageSnapshot( + primary: primary, + secondary: secondary, + updatedAt: quotaUpdatedAt), + provider: .claude) + + var widgetSnapshots: [WidgetSnapshot] = [] + store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) } + defer { store._test_widgetSnapshotSaveOverride = nil } + + store.persistWidgetSnapshot(reason: "claude-pre-token-preserves-quota-test") + await store.widgetSnapshotPersistTask?.value + + let preTokenEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(preTokenEntry.updatedAt == quotaUpdatedAt) + #expect(preTokenEntry.primary == primary) + #expect(preTokenEntry.secondary == secondary) + #expect(preTokenEntry.usageRows?.map(\.id) == ["primary", "secondary"]) + #expect(preTokenEntry.tokenUsage == nil) + let quotaOwnerKey = try #require(preTokenEntry.quotaOwnerKey) + + store.snapshots.removeValue(forKey: .claude) + store._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 4300, + sessionCostUSD: 1.50, + last30DaysTokens: 43000, + last30DaysCostUSD: 13.50, + daily: [], + updatedAt: tokenUpdatedAt), + provider: .claude) + store.persistWidgetSnapshot(reason: "claude-token-only-preserves-quota-test") + await store.widgetSnapshotPersistTask?.value + + let entry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(entry.updatedAt == quotaUpdatedAt) + #expect(entry.primary == primary) + #expect(entry.secondary == secondary) + #expect(entry.usageRows?.map(\.id) == ["primary", "secondary"]) + #expect(entry.usageRows?.compactMap(\.percentLeft) == [72, 88]) + #expect(entry.tokenUsage?.updatedAt == tokenUpdatedAt) + #expect(entry.tokenUsage?.sessionTokens == 4300) + + store.lastQueuedWidgetSnapshot = WidgetSnapshot( + entries: [ + WidgetSnapshot.ProviderEntry( + provider: .claude, + updatedAt: quotaUpdatedAt, + primary: primary, + secondary: secondary, + tertiary: nil, + usageRows: [ + WidgetSnapshot.WidgetUsageRowSnapshot( + id: "primary", + title: "Session", + percentLeft: 72), + WidgetSnapshot.WidgetUsageRowSnapshot( + id: "secondary", + title: "Weekly", + percentLeft: 88), + ], + creditsRemaining: nil, + codeReviewRemainingPercent: nil, + tokenUsage: nil, + dailyUsage: [], + quotaOwnerKey: quotaOwnerKey), + ], + enabledProviders: [.claude], + generatedAt: quotaUpdatedAt) + store.widgetUsagePreservationBlockedProviders.insert(.claude) + + store.persistWidgetSnapshot(reason: "claude-token-only-credential-change-test") + await store.widgetSnapshotPersistTask?.value + + let blockedEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(blockedEntry.updatedAt == tokenUpdatedAt) + #expect(blockedEntry.primary == nil) + #expect(blockedEntry.secondary == nil) + #expect(blockedEntry.usageRows?.isEmpty == true) + + let placeholder = RateWindow( + usedPercent: 0, + windowMinutes: 300, + resetsAt: nil, + resetDescription: nil, + isSyntheticPlaceholder: true) + store.widgetUsagePreservationBlockedProviders.remove(.claude) + store.lastQueuedWidgetSnapshot = WidgetSnapshot( + entries: [ + WidgetSnapshot.ProviderEntry( + provider: .claude, + updatedAt: quotaUpdatedAt, + primary: placeholder, + secondary: nil, + tertiary: nil, + usageRows: [ + WidgetSnapshot.WidgetUsageRowSnapshot( + id: "primary", + title: "Session", + percentLeft: 100), + ], + creditsRemaining: nil, + codeReviewRemainingPercent: nil, + tokenUsage: nil, + dailyUsage: [], + quotaOwnerKey: quotaOwnerKey), + ], + enabledProviders: [.claude], + generatedAt: quotaUpdatedAt) + + store.persistWidgetSnapshot(reason: "claude-token-only-drops-placeholder-test") + await store.widgetSnapshotPersistTask?.value + + let filteredEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .claude }) + #expect(filteredEntry.updatedAt == tokenUpdatedAt) + #expect(filteredEntry.primary == nil) + #expect(filteredEntry.usageRows?.isEmpty == true) + } + @Test func `widget snapshot uses Claude enterprise spend limit instead of placeholder quota`() async throws { let suite = "UsageStoreWidgetSnapshotTests-claude-enterprise-spend-limit" diff --git a/Tests/CodexBarTests/WidgetSnapshotTests.swift b/Tests/CodexBarTests/WidgetSnapshotTests.swift index 9264145c19..8ccf40f138 100644 --- a/Tests/CodexBarTests/WidgetSnapshotTests.swift +++ b/Tests/CodexBarTests/WidgetSnapshotTests.swift @@ -48,7 +48,8 @@ struct WidgetSnapshotTests { last30DaysLabel: "This month"), dailyUsage: [ WidgetSnapshot.DailyUsagePoint(dayKey: "2025-12-20", totalTokens: 1200, costUSD: 12.3), - ]) + ], + quotaOwnerKey: "claude-account-cache-key") let snapshot = WidgetSnapshot( entries: [entry], @@ -71,6 +72,7 @@ struct WidgetSnapshotTests { #expect(decoded.entries.first?.tokenUsage?.sessionLabel == "Latest billing day") #expect(decoded.entries.first?.tokenUsage?.last30DaysLabel == "This month") #expect(decoded.entries.first?.usageRows?.map(\.id) == ["session", "weekly"]) + #expect(decoded.entries.first?.quotaOwnerKey == "claude-account-cache-key") #expect(decoded.enabledProviders == [.codex, .claude]) #expect(decoded.usageBarsShowUsed) } @@ -188,6 +190,7 @@ struct WidgetSnapshotTests { #expect(decoded.entries.count == 1) #expect(decoded.entries.first?.usageRows == nil) + #expect(decoded.entries.first?.quotaOwnerKey == nil) #expect(decoded.entries.first?.secondary?.usedPercent == 25) #expect(!decoded.usageBarsShowUsed) }