Skip to content

Commit 060a9fd

Browse files
steipetesteipete
committed
fix: keep provider cost widgets fresh
1 parent 3e137e3 commit 060a9fd

4 files changed

Lines changed: 54 additions & 48 deletions

File tree

‎Sources/CodexBar/UsageStore+WidgetSnapshot.swift‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,8 @@ extension UsageStore {
3636
private func makeWidgetEntry(for provider: UsageProvider) -> WidgetSnapshot.ProviderEntry? {
3737
guard let snapshot = self.snapshots[provider] else { return nil }
3838

39-
let tokenSnapshot = self.tokenSnapshots[provider]
39+
let tokenSnapshot = self.tokenSnapshot(fromProviderSnapshot: snapshot, provider: provider) ?? self
40+
.tokenSnapshots[provider]
4041
let dailyUsage = tokenSnapshot?.daily.map { entry in
4142
WidgetSnapshot.DailyUsagePoint(
4243
dayKey: entry.date,

‎Sources/CodexBar/UsageStore.swift‎

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1495,6 +1495,20 @@ extension UsageStore {
14951495
return
14961496
}
14971497

1498+
if Self.tokenCostRequiresProviderSnapshot(provider) {
1499+
if let snapshot = self.tokenSnapshot(fromProviderSnapshot: self.snapshots[provider], provider: provider) {
1500+
self.tokenSnapshots[provider] = snapshot
1501+
self.tokenErrors[provider] = nil
1502+
self.tokenFailureGates[provider]?.recordSuccess()
1503+
self.persistWidgetSnapshot(reason: "token-usage")
1504+
} else {
1505+
self.tokenSnapshots.removeValue(forKey: provider)
1506+
self.tokenErrors[provider] = nil
1507+
self.tokenFailureGates[provider]?.reset()
1508+
}
1509+
return
1510+
}
1511+
14981512
guard self.settings.costUsageEnabled else {
14991513
self.tokenSnapshots.removeValue(forKey: provider)
15001514
self.tokenErrors[provider] = nil
@@ -1531,21 +1545,6 @@ extension UsageStore {
15311545
self.tokenRefreshInFlight.insert(provider)
15321546
defer { self.tokenRefreshInFlight.remove(provider) }
15331547

1534-
if let snapshot = self.tokenSnapshot(fromProviderSnapshot: self.snapshots[provider], provider: provider) {
1535-
self.tokenSnapshots[provider] = snapshot
1536-
self.tokenErrors[provider] = nil
1537-
self.tokenFailureGates[provider]?.recordSuccess()
1538-
self.persistWidgetSnapshot(reason: "token-usage")
1539-
return
1540-
}
1541-
1542-
if Self.tokenCostRequiresProviderSnapshot(provider) {
1543-
self.tokenSnapshots.removeValue(forKey: provider)
1544-
self.tokenErrors[provider] = Self.tokenCostNoDataMessage(for: provider)
1545-
self.tokenFailureGates[provider]?.recordSuccess()
1546-
return
1547-
}
1548-
15491548
let startedAt = Date()
15501549
let providerText = provider.rawValue
15511550
self.tokenCostLogger

‎Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift‎

Lines changed: 0 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -820,38 +820,6 @@ struct CodexAccountScopedRefreshTests {
820820
#expect(codexEntry.codeReviewRemainingPercent == 88)
821821
}
822822

823-
@Test
824-
func `widget snapshot preserves token usage currency`() async throws {
825-
let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-widget-token-currency")
826-
settings.refreshFrequency = .manual
827-
828-
let store = self.makeUsageStore(settings: settings)
829-
store._setSnapshotForTesting(UsageSnapshot(primary: nil, secondary: nil, updatedAt: Date()), provider: .mistral)
830-
store._setTokenSnapshotForTesting(
831-
CostUsageTokenSnapshot(
832-
sessionTokens: 10,
833-
sessionCostUSD: 1.2,
834-
last30DaysTokens: 100,
835-
last30DaysCostUSD: 9.0,
836-
currencyCode: "eur",
837-
historyLabel: "This month",
838-
daily: [],
839-
updatedAt: Date()),
840-
provider: .mistral)
841-
842-
var widgetSnapshots: [WidgetSnapshot] = []
843-
store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) }
844-
defer { store._test_widgetSnapshotSaveOverride = nil }
845-
846-
store.persistWidgetSnapshot(reason: "token-currency")
847-
await store.widgetSnapshotPersistTask?.value
848-
849-
let mistralEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .mistral })
850-
#expect(mistralEntry.tokenUsage?.currencyCode == "EUR")
851-
#expect(mistralEntry.tokenUsage?.sessionLabel == "Latest billing day")
852-
#expect(mistralEntry.tokenUsage?.last30DaysLabel == "This month")
853-
}
854-
855823
@Test
856824
func `codex account refresh reports usage and credits phases before completion`() async {
857825
let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-phases")

‎Tests/CodexBarTests/UsageStoreCoverageTests.swift‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,44 @@ struct UsageStoreCoverageTests {
352352
#expect(store.enabledProvidersForBackgroundWork().isEmpty)
353353
}
354354

355+
@Test
356+
func `widget snapshot projects provider derived token usage`() async throws {
357+
let settings = Self.makeSettingsStore(suite: "UsageStoreCoverageTests-widget-provider-cost")
358+
let store = Self.makeUsageStore(settings: settings)
359+
let day = MistralDailyUsageBucket(
360+
day: "2026-05-26",
361+
cost: 1.2,
362+
inputTokens: 10,
363+
cachedTokens: 0,
364+
outputTokens: 5,
365+
models: [])
366+
store._setSnapshotForTesting(MistralUsageSnapshot(
367+
totalCost: 9,
368+
currency: "eur",
369+
currencySymbol: "€",
370+
totalInputTokens: 10,
371+
totalOutputTokens: 5,
372+
totalCachedTokens: 0,
373+
modelCount: 1,
374+
daily: [day],
375+
startDate: nil,
376+
endDate: nil,
377+
updatedAt: Date()).toUsageSnapshot(), provider: .mistral)
378+
379+
var widgetSnapshots: [WidgetSnapshot] = []
380+
store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) }
381+
defer { store._test_widgetSnapshotSaveOverride = nil }
382+
383+
store.persistWidgetSnapshot(reason: "provider-cost")
384+
await store.widgetSnapshotPersistTask?.value
385+
386+
let mistralEntry = try #require(widgetSnapshots.last?.entries.first { $0.provider == .mistral })
387+
#expect(mistralEntry.tokenUsage?.currencyCode == "EUR")
388+
#expect(mistralEntry.tokenUsage?.sessionLabel == "Latest billing day")
389+
#expect(mistralEntry.tokenUsage?.last30DaysLabel == "This month")
390+
#expect(mistralEntry.tokenUsage?.last30DaysCostUSD == 9)
391+
}
392+
355393
@Test
356394
func `unavailable provider with only cached status gets single cleanup pass`() async throws {
357395
let settings = Self.makeSettingsStore(suite: "UsageStoreCoverageTests-background-status-cleanup")

0 commit comments

Comments
 (0)