diff --git a/Sources/CodexBar/PlanUtilizationHistoryStore.swift b/Sources/CodexBar/PlanUtilizationHistoryStore.swift index 621749d503..ef7f0d83a9 100644 --- a/Sources/CodexBar/PlanUtilizationHistoryStore.swift +++ b/Sources/CodexBar/PlanUtilizationHistoryStore.swift @@ -40,6 +40,12 @@ struct PlanUtilizationSeriesHistory: Codable, Equatable, Sendable { let windowMinutes: Int let entries: [PlanUtilizationHistoryEntry] + private enum CodingKeys: String, CodingKey { + case name + case windowMinutes + case entries + } + init(name: PlanUtilizationSeriesName, windowMinutes: Int, entries: [PlanUtilizationHistoryEntry]) { self.name = name self.windowMinutes = windowMinutes @@ -56,6 +62,14 @@ struct PlanUtilizationSeriesHistory: Codable, Equatable, Sendable { } } + init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let name = try container.decode(PlanUtilizationSeriesName.self, forKey: .name) + let windowMinutes = try container.decode(Int.self, forKey: .windowMinutes) + let entries = try container.decode([PlanUtilizationHistoryEntry].self, forKey: .entries) + self.init(name: name, windowMinutes: windowMinutes, entries: entries) + } + var latestCapturedAt: Date? { self.entries.last?.capturedAt } diff --git a/Sources/CodexBar/UsageStore+PlanUtilization.swift b/Sources/CodexBar/UsageStore+PlanUtilization.swift index 6ac3d3dc76..a3eb54d74e 100644 --- a/Sources/CodexBar/UsageStore+PlanUtilization.swift +++ b/Sources/CodexBar/UsageStore+PlanUtilization.swift @@ -373,8 +373,10 @@ extension UsageStore { let canonicalWindowMinutes = sample.name.canonicalWindowMinutes(sample.windowMinutes) let key = PlanUtilizationSeriesKey(name: sample.name, windowMinutes: canonicalWindowMinutes) if let existingHistory = historiesByKey[key] { - guard let updatedEntries = self.updatedPlanUtilizationEntries( - existingEntries: existingHistory.entries, + var updatedEntries = existingHistory.entries + self.assertPlanUtilizationEntriesSorted(updatedEntries) + guard self.updatedPlanUtilizationEntries( + existingEntries: &updatedEntries, entry: sample.entry) else { continue @@ -411,28 +413,72 @@ extension UsageStore { } private nonisolated static func updatedPlanUtilizationEntries( - existingEntries: [PlanUtilizationHistoryEntry], - entry: PlanUtilizationHistoryEntry) -> [PlanUtilizationHistoryEntry]? + existingEntries: inout [PlanUtilizationHistoryEntry], + entry: PlanUtilizationHistoryEntry) -> Bool { - var entries = existingEntries - let insertionIndex = entries.firstIndex(where: { $0.capturedAt > entry.capturedAt }) ?? entries.endIndex + let insertionIndex = self.planUtilizationEntryInsertionIndex( + entries: existingEntries, + capturedAt: entry.capturedAt) let sampleHourBucket = self.planUtilizationHourBucket(for: entry.capturedAt) let sameHourRange = self.planUtilizationHourRange( - entries: entries, + entries: existingEntries, insertionIndex: insertionIndex, hourBucket: sampleHourBucket) - let existingHourEntries = Array(entries[sameHourRange]) + let existingHourEntries = Array(existingEntries[sameHourRange]) let canonicalHourEntries = self.canonicalPlanUtilizationHourEntries( existingHourEntries: existingHourEntries, incomingEntry: entry) - guard canonicalHourEntries != existingHourEntries else { return nil } - entries.replaceSubrange(sameHourRange, with: canonicalHourEntries) + guard canonicalHourEntries != existingHourEntries else { return false } + existingEntries.replaceSubrange(sameHourRange, with: canonicalHourEntries) - if entries.count > self.planUtilizationMaxSamples { - entries.removeFirst(entries.count - self.planUtilizationMaxSamples) + if existingEntries.count > self.planUtilizationMaxSamples { + existingEntries.removeFirst(existingEntries.count - self.planUtilizationMaxSamples) } - return entries + return true + } + + /// The insertion search assumes `capturedAt` order. Checked once per merged series rather than per entry — + /// a per-entry check would reintroduce the O(n) scan this insertion path exists to remove. + private nonisolated static func assertPlanUtilizationEntriesSorted( + _ entries: [PlanUtilizationHistoryEntry]) + { + #if DEBUG + assert( + zip(entries, entries.dropFirst()).allSatisfy { $0.capturedAt <= $1.capturedAt }, + "plan-utilization entries must be sorted by capturedAt before insertion") + #endif + } + + private nonisolated static func planUtilizationEntryInsertionIndex( + entries: [PlanUtilizationHistoryEntry], + capturedAt: Date) -> Int + { + guard let lastCapturedAt = entries.last?.capturedAt else { + return entries.endIndex + } + if lastCapturedAt <= capturedAt { + return entries.endIndex + } + return self.planUtilizationEntryUpperBound(entries: entries, capturedAt: capturedAt) + } + + /// First index i such that entries[i].capturedAt > capturedAt (endIndex if none). + private nonisolated static func planUtilizationEntryUpperBound( + entries: [PlanUtilizationHistoryEntry], + capturedAt: Date) -> Int + { + var low = entries.startIndex + var high = entries.endIndex + while low < high { + let mid = low + ((high - low) / 2) + if entries[mid].capturedAt > capturedAt { + high = mid + } else { + low = mid + 1 + } + } + return low } #if DEBUG @@ -440,7 +486,18 @@ extension UsageStore { existingEntries: [PlanUtilizationHistoryEntry], entry: PlanUtilizationHistoryEntry) -> [PlanUtilizationHistoryEntry]? { - self.updatedPlanUtilizationEntries(existingEntries: existingEntries, entry: entry) + var entries = existingEntries + guard self.updatedPlanUtilizationEntries(existingEntries: &entries, entry: entry) else { + return nil + } + return entries + } + + nonisolated static func _planUtilizationEntryUpperBoundForTesting( + entries: [PlanUtilizationHistoryEntry], + capturedAt: Date) -> Int + { + self.planUtilizationEntryUpperBound(entries: entries, capturedAt: capturedAt) } nonisolated static func _updatedPlanUtilizationHistoriesForTesting( @@ -1202,6 +1259,7 @@ extension UsageStore { var historiesToMerge: [[PlanUtilizationSeriesHistory]] = [] let scopedRawKeys = Array(providerBuckets.accounts.keys) var legacyRawKeysToRemove: [String] = [] + var hasForeignHistoryToMerge = false for rawKey in scopedRawKeys { let owner = CodexHistoryOwnership.classifyPersistedKey( @@ -1219,6 +1277,7 @@ extension UsageStore { historiesToMerge.append(accountHistories) if rawKey != canonicalKey { legacyRawKeysToRemove.append(rawKey) + hasForeignHistoryToMerge = true } } } @@ -1232,6 +1291,7 @@ extension UsageStore { { historiesToMerge.append(opaqueHistories) legacyRawKeysToRemove.append(recoverableOpaqueRawKey) + hasForeignHistoryToMerge = true } if shouldAdoptUnscopedHistory, @@ -1245,9 +1305,14 @@ extension UsageStore { { historiesToMerge.append(providerBuckets.unscoped) providerBuckets.unscoped = [] + hasForeignHistoryToMerge = true } - guard !historiesToMerge.isEmpty else { return canonicalKey } + // Canonical-only contribution is not a migration; re-merging it with itself is quadratic. + // Skipping also skips incidental re-canonicalization of already-canonical series (duplicate + // (name, windowMinutes) collapse and per-hour peak replay). Repairing legacy data belongs + // at load time, not on every refresh and menu open. + guard hasForeignHistoryToMerge else { return canonicalKey } for rawKey in legacyRawKeysToRemove { providerBuckets.accounts.removeValue(forKey: rawKey) } @@ -1523,29 +1588,29 @@ extension UsageStore { provider _: UsageProvider, histories: [[PlanUtilizationSeriesHistory]]) -> [PlanUtilizationSeriesHistory] { - var mergedByKey: [PlanUtilizationSeriesKey: PlanUtilizationSeriesHistory] = [:] + var mergedEntriesByKey: [PlanUtilizationSeriesKey: [PlanUtilizationHistoryEntry]] = [:] for historyGroup in histories { for history in historyGroup { let key = PlanUtilizationSeriesKey(name: history.name, windowMinutes: history.windowMinutes) - let existingEntries = mergedByKey[key]?.entries ?? [] - var mergedEntries = existingEntries + var mergedEntries = mergedEntriesByKey[key] ?? [] + self.assertPlanUtilizationEntriesSorted(mergedEntries) for entry in history.entries.sorted(by: { $0.capturedAt < $1.capturedAt }) { - if let updatedEntries = self.updatedPlanUtilizationEntries( - existingEntries: mergedEntries, + _ = self.updatedPlanUtilizationEntries( + existingEntries: &mergedEntries, entry: entry) - { - mergedEntries = updatedEntries - } } - mergedByKey[key] = PlanUtilizationSeriesHistory( - name: history.name, - windowMinutes: history.windowMinutes, - entries: mergedEntries) + mergedEntriesByKey[key] = mergedEntries } } - return mergedByKey.values.sorted { lhs, rhs in + return mergedEntriesByKey.map { key, entries in + PlanUtilizationSeriesHistory( + name: key.name, + windowMinutes: key.windowMinutes, + entries: entries) + } + .sorted { lhs, rhs in if lhs.windowMinutes != rhs.windowMinutes { return lhs.windowMinutes < rhs.windowMinutes } diff --git a/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift b/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift index a189eecbaa..61a629fc29 100644 --- a/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift +++ b/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift @@ -2736,7 +2736,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact app-runtime bridge coordinates provider-owned state through the shared controller."), AllowedProviderConstruct( path: "Sources/CodexBar/UsageStore+PlanUtilization.swift", - line: 846, + line: 903, anchor: "if provider == .claude {", expectedProviderIDs: ["claude"], expectedReferenceCount: 1, @@ -2744,7 +2744,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact app-runtime bridge coordinates provider-owned state through the shared controller."), AllowedProviderConstruct( path: "Sources/CodexBar/UsageStore+PlanUtilization.swift", - line: 861, + line: 918, anchor: "guard let identity = snapshot.identity(for: .claude) else { return nil }", expectedProviderIDs: ["claude"], expectedReferenceCount: 3, @@ -2752,7 +2752,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact app-runtime bridge coordinates provider-owned state through the shared controller."), AllowedProviderConstruct( path: "Sources/CodexBar/UsageStore+PlanUtilization.swift", - line: 888, + line: 945, anchor: "key.hasPrefix(\"\\(UsageProvider.claude.rawValue):\")", expectedProviderIDs: ["claude"], expectedReferenceCount: 1, @@ -2760,7 +2760,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact app-runtime bridge coordinates provider-owned state through the shared controller."), AllowedProviderConstruct( path: "Sources/CodexBar/UsageStore+PlanUtilization.swift", - line: 1318, + line: 1383, anchor: "if ![UsageProvider.codex, .claude, .antigravity].contains(provider) {", expectedProviderIDs: ["antigravity", "claude", "codex"], expectedReferenceCount: 3, diff --git a/Tests/CodexBarTests/SessionEquivalentForecastTests.swift b/Tests/CodexBarTests/SessionEquivalentForecastTests.swift index f2bcd14f09..2bd1d817c9 100644 --- a/Tests/CodexBarTests/SessionEquivalentForecastTests.swift +++ b/Tests/CodexBarTests/SessionEquivalentForecastTests.swift @@ -341,7 +341,7 @@ struct SessionEquivalentForecastTests { } @Test - func `rejects hostile dates percentages and unsorted history`() throws { + func `rejects hostile dates and percentages and normalizes unsorted history`() throws { let now = Date(timeIntervalSince1970: 1_900_000_000) let session = RateWindow( usedPercent: 20, @@ -380,15 +380,18 @@ struct SessionEquivalentForecastTests { let encodedSession = try JSONEncoder().encode(fixture.histories[0]) var sessionJSON = try #require(JSONSerialization.jsonObject(with: encodedSession) as? [String: Any]) let entriesJSON = try #require(sessionJSON["entries"] as? [[String: Any]]) + // Decoding sorts entries (PlanUtilizationSeriesHistory.init(from:)), so a reversed payload + // normalizes to the same series and yields the same estimate as the sorted fixture above. sessionJSON["entries"] = Array(entriesJSON.reversed()) let shuffledData = try JSONSerialization.data(withJSONObject: sessionJSON) let shuffledSession = try JSONDecoder().decode(PlanUtilizationSeriesHistory.self, from: shuffledData) - #expect((shuffledSession.entries.first?.capturedAt ?? .distantPast) - > (shuffledSession.entries.last?.capturedAt ?? .distantFuture)) - #expect(SessionEquivalentBurnEstimator.estimate( + #expect(shuffledSession == fixture.histories[0]) + let shuffledEstimate = try #require(SessionEquivalentBurnEstimator.estimate( histories: [shuffledSession, fixture.histories[1]], currentSessionResetsAt: fixture.currentSessionReset, - now: fixture.currentSessionReset.addingTimeInterval(-3600)) == nil) + now: fixture.currentSessionReset.addingTimeInterval(-3600))) + #expect(shuffledEstimate.sampleCount == 3) + #expect(shuffledEstimate.medianWeeklyPercentPerWindow == 6) let huge = UsagePaceText.sessionEquivalentDetail(forecast: SessionEquivalentForecast( estimatedWindowsToExhaustWeekly: .greatestFiniteMagnitude, diff --git a/Tests/CodexBarTests/UsageStorePlanUtilizationCodexMergeTests.swift b/Tests/CodexBarTests/UsageStorePlanUtilizationCodexMergeTests.swift new file mode 100644 index 0000000000..0562b631f0 --- /dev/null +++ b/Tests/CodexBarTests/UsageStorePlanUtilizationCodexMergeTests.swift @@ -0,0 +1,220 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +struct UsageStorePlanUtilizationCodexMergeTests { + @MainActor + @Test + func `codex materialize leaves canonical-only history untouched`() throws { + let store = UsageStorePlanUtilizationTests.makeStore() + let snapshot = UsageStorePlanUtilizationTests.makeSnapshot(provider: .codex, email: "alice@example.com") + let canonicalKey = try #require( + UsageStore._planUtilizationAccountKeyForTesting( + provider: .codex, + snapshot: snapshot)) + let hourStart = Date(timeIntervalSince1970: 1_699_999_200) + let session = planSeries(name: .session, windowMinutes: 300, entries: [ + planEntry(at: hourStart.addingTimeInterval(3600), usedPercent: 40), + planEntry(at: hourStart, usedPercent: 10), + ]) + let weekly = planSeries(name: .weekly, windowMinutes: 10080, entries: [ + planEntry(at: hourStart.addingTimeInterval(7200), usedPercent: 15), + ]) + // Direct assignment keeps this unsorted; a self-merge would reorder by windowMinutes. + let originalHistories = [weekly, session] + let originalBuckets = PlanUtilizationHistoryBuckets( + preferredAccountKey: canonicalKey, + unscoped: [], + accounts: [ + canonicalKey: originalHistories, + ]) + store.planUtilizationHistory[.codex] = originalBuckets + store._setSnapshotForTesting(snapshot, provider: .codex) + + let revisionBefore = store.planUtilizationHistoryRevision + let history = store.planUtilizationHistory(for: .codex) + let buckets = try #require(store.planUtilizationHistory[.codex]) + + #expect(history == originalHistories) + #expect(buckets == originalBuckets) + #expect(buckets.accounts[canonicalKey] == originalHistories) + #expect(buckets.accounts.keys.sorted() == [canonicalKey]) + #expect(buckets.unscoped.isEmpty) + #expect(store.planUtilizationHistoryRevision == revisionBefore) + } + + @MainActor + @Test + func `codex materialize merge matches reference for overlapping hours out of order entries and distinct series`() + throws + { + let store = UsageStorePlanUtilizationTests.makeStore() + let snapshot = UsageStorePlanUtilizationTests.makeSnapshot(provider: .codex, email: "alice@example.com") + let canonicalKey = try #require( + UsageStore._planUtilizationAccountKeyForTesting( + provider: .codex, + snapshot: snapshot)) + let legacyEmailHash = UsageStore._codexLegacyPlanUtilizationEmailHashKeyForTesting( + normalizedEmail: "alice@example.com") + let hourStart = Date(timeIntervalSince1970: 1_699_999_200) + + let canonicalSession = planSeries(name: .session, windowMinutes: 300, entries: [ + planEntry(at: hourStart.addingTimeInterval(3 * 3600), usedPercent: 40), + planEntry(at: hourStart.addingTimeInterval(5 * 60), usedPercent: 10), + ]) + let canonicalWeekly = planSeries(name: .weekly, windowMinutes: 10080, entries: [ + planEntry(at: hourStart, usedPercent: 15), + ]) + let legacySession = planSeries(name: .session, windowMinutes: 300, entries: [ + planEntry(at: hourStart.addingTimeInterval(3600), usedPercent: 20), + planEntry(at: hourStart.addingTimeInterval(25 * 60), usedPercent: 35), + ]) + let legacyWeekly = planSeries(name: .weekly, windowMinutes: 10080, entries: [ + planEntry(at: hourStart.addingTimeInterval(3600), usedPercent: 25), + ]) + let unscopedSession = planSeries(name: .session, windowMinutes: 300, entries: [ + planEntry(at: hourStart.addingTimeInterval(2 * 3600), usedPercent: 30), + ]) + + store.planUtilizationHistory[.codex] = PlanUtilizationHistoryBuckets( + unscoped: [unscopedSession], + accounts: [ + canonicalKey: [canonicalWeekly, canonicalSession], + legacyEmailHash: [legacySession, legacyWeekly], + ]) + store._setSnapshotForTesting(snapshot, provider: .codex) + + let history = store.planUtilizationHistory(for: .codex) + let buckets = try #require(store.planUtilizationHistory[.codex]) + + let expectedSession = planSeries(name: .session, windowMinutes: 300, entries: [ + planEntry(at: hourStart.addingTimeInterval(25 * 60), usedPercent: 35), + planEntry(at: hourStart.addingTimeInterval(3600), usedPercent: 20), + planEntry(at: hourStart.addingTimeInterval(2 * 3600), usedPercent: 30), + planEntry(at: hourStart.addingTimeInterval(3 * 3600), usedPercent: 40), + ]) + let expectedWeekly = planSeries(name: .weekly, windowMinutes: 10080, entries: [ + planEntry(at: hourStart, usedPercent: 15), + planEntry(at: hourStart.addingTimeInterval(3600), usedPercent: 25), + ]) + let expectedHistories = [expectedSession, expectedWeekly] + + #expect(history == expectedHistories) + #expect(buckets.accounts[canonicalKey] == expectedHistories) + #expect(buckets.accounts[legacyEmailHash] == nil) + #expect(buckets.unscoped.isEmpty) + #expect(findSeries(history, name: .session, windowMinutes: 300)?.entries.map(\.usedPercent) == [ + 35, 20, 30, 40, + ]) + #expect(findSeries(history, name: .weekly, windowMinutes: 10080)?.entries.map(\.usedPercent) == [15, 25]) + } + + @MainActor + @Test + func `codex materialize merge matches reference when retention trimming is exercised`() throws { + let store = UsageStorePlanUtilizationTests.makeStore() + let snapshot = UsageStorePlanUtilizationTests.makeSnapshot(provider: .codex, email: "alice@example.com") + let canonicalKey = try #require( + UsageStore._planUtilizationAccountKeyForTesting( + provider: .codex, + snapshot: snapshot)) + let legacyEmailHash = UsageStore._codexLegacyPlanUtilizationEmailHashKeyForTesting( + normalizedEmail: "alice@example.com") + let hourStart = Date(timeIntervalSince1970: 1_699_999_200) + let maxSamples = UsageStore._planUtilizationMaxSamplesForTesting + let overflow = 12 + let totalSessionEntries = maxSamples + overflow + let legacySessionEntries = (0..