diff --git a/Sources/CodexBar/IconRenderer.swift b/Sources/CodexBar/IconRenderer.swift index 15e103da1d..f5864c1626 100644 --- a/Sources/CodexBar/IconRenderer.swift +++ b/Sources/CodexBar/IconRenderer.swift @@ -938,10 +938,10 @@ enum IconRenderer { switch indicator { case .minor, .maintenance: - let size: CGFloat = 4 + let size: CGFloat = 2 let rect = Self.snapRect( - x: Self.baseSize.width - size - 2, - y: 2, + x: Self.baseSize.width - size - 1, + y: 7, width: size, height: size) let path = NSBezierPath(ovalIn: rect) @@ -949,19 +949,19 @@ enum IconRenderer { path.fill() case .major, .critical, .unknown: let lineRect = Self.snapRect( - x: Self.baseSize.width - 6, - y: 4, - width: 2.0, - height: 6) - let linePath = NSBezierPath(roundedRect: lineRect, xRadius: 1, yRadius: 1) + x: Self.baseSize.width - 2.5, + y: 7.4, + width: 1.5, + height: 1.8) + let linePath = NSBezierPath(roundedRect: lineRect, xRadius: 0.75, yRadius: 0.75) color.setFill() linePath.fill() let dotRect = Self.snapRect( - x: Self.baseSize.width - 6, - y: 2, - width: 2.0, - height: 2.0) + x: Self.baseSize.width - 2.45, + y: 6.6, + width: 1.4, + height: 1.4) NSBezierPath(ovalIn: dotRect).fill() case .none: break diff --git a/Sources/CodexBar/InlineUsageDashboardContent.swift b/Sources/CodexBar/InlineUsageDashboardContent.swift index e1ec45057d..7b2366faec 100644 --- a/Sources/CodexBar/InlineUsageDashboardContent.swift +++ b/Sources/CodexBar/InlineUsageDashboardContent.swift @@ -141,6 +141,24 @@ extension UsageMenuCardView.Model { return nil } + static func inlineUsageDashboardLoading( + input: Input, + tokenUsageSnapshot: CostUsageTokenSnapshot?, + inlineUsageDashboard: InlineUsageDashboardModel?) + -> Bool + { + guard inlineUsageDashboard == nil else { return false } + guard input.tokenCostUsageEnabled else { return false } + guard tokenUsageSnapshot == nil else { return false } + guard input.tokenError?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty ?? true else { return false } + guard input.lastError?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty ?? true else { return false } + if self.usesProviderCostHistoryAsPrimaryDashboard(input.provider) { + return input.isRefreshing || input.tokenRefreshInFlight || input.tokenRefreshQueued + } + guard [.codex, .claude, .vertexai, .bedrock].contains(input.provider) else { return false } + return input.tokenRefreshInFlight || input.tokenRefreshQueued + } + static func usesProviderCostHistoryAsPrimaryDashboard(_ provider: UsageProvider) -> Bool { provider == .openai || provider == .mistral } @@ -659,3 +677,101 @@ struct InlineUsageDashboardContent: View { } } } + +struct InlineUsageDashboardLoadingContent: View { + private let barHeights: [CGFloat] = [ + 26, + 19, + 10, + 48, + 12, + 4, + 4, + 4, + 4, + 8, + 13, + 12, + 15, + 9, + 36, + 33, + 39, + 7, + 22, + 16, + 24, + 26, + 27, + 18, + 13, + ] + @Environment(\.menuItemHighlighted) private var isHighlighted + + var body: some View { + VStack(alignment: .leading, spacing: 10) { + LazyVGrid( + columns: [ + GridItem(.flexible(minimum: 118), alignment: .leading), + GridItem(.flexible(minimum: 100), alignment: .leading), + ], + alignment: .leading, + spacing: 6) + { + LoadingKPI(title: L("Today"), width: 68) + LoadingKPI(title: L("30d cost"), width: 78) + LoadingKPI(title: L("30d tokens"), width: 58) + LoadingKPI(title: L("Latest tokens"), width: 62) + } + HStack(alignment: .bottom, spacing: 2) { + ForEach(Array(self.barHeights.enumerated()), id: \.offset) { _, height in + RoundedRectangle(cornerRadius: 1.5, style: .continuous) + .fill(self.skeletonFill.opacity(0.85)) + .frame(maxWidth: .infinity) + .frame(height: height) + } + } + .frame(maxWidth: .infinity) + .frame(height: 58, alignment: .bottom) + .overlay(alignment: .bottomLeading) { + Rectangle() + .fill(MenuHighlightStyle.secondary(self.isHighlighted).opacity(0.18)) + .frame(height: 1) + } + VStack(alignment: .leading, spacing: 5) { + self.skeletonLine(width: 138) + self.skeletonLine(width: 228) + } + } + .accessibilityLabel("Loading cost history") + } + + private var skeletonFill: Color { + MenuHighlightStyle.secondary(self.isHighlighted).opacity(self.isHighlighted ? 0.34 : 0.2) + } + + private func skeletonLine(width: CGFloat, height: CGFloat = 10) -> some View { + RoundedRectangle(cornerRadius: height / 2, style: .continuous) + .fill(self.skeletonFill) + .frame(width: width, height: height) + } + + private struct LoadingKPI: View { + let title: String + let width: CGFloat + @Environment(\.menuItemHighlighted) private var isHighlighted + + var body: some View { + VStack(alignment: .leading, spacing: 4) { + Text(self.title) + .font(.caption2) + .foregroundStyle(MenuHighlightStyle.secondary(self.isHighlighted)) + .lineLimit(1) + RoundedRectangle(cornerRadius: 6, style: .continuous) + .fill(MenuHighlightStyle.secondary(self.isHighlighted).opacity(self.isHighlighted ? 0.34 : 0.2)) + .frame(width: self.width, height: 18) + } + .frame(maxWidth: .infinity, alignment: .leading) + } + } +} diff --git a/Sources/CodexBar/MainRunLoopStallMonitor.swift b/Sources/CodexBar/MainRunLoopStallMonitor.swift new file mode 100644 index 0000000000..90a2f2766a --- /dev/null +++ b/Sources/CodexBar/MainRunLoopStallMonitor.swift @@ -0,0 +1,66 @@ +import CodexBarCore +import Foundation + +@MainActor +final class MainRunLoopStallMonitor { + private let logger: CodexBarLogger + private let interval: TimeInterval + private let threshold: TimeInterval + private let metadataProvider: @MainActor () -> [String: String] + private var timer: DispatchSourceTimer? + private var lastTick: TimeInterval = ProcessInfo.processInfo.systemUptime + + init( + logger: CodexBarLogger, + interval: TimeInterval = 0.1, + threshold: TimeInterval = 0.35, + metadataProvider: @escaping @MainActor () -> [String: String]) + { + self.logger = logger + self.interval = interval + self.threshold = threshold + self.metadataProvider = metadataProvider + } + + func start() { + guard self.timer == nil else { return } + + self.lastTick = ProcessInfo.processInfo.systemUptime + let timer = DispatchSource.makeTimerSource(queue: .main) + timer.schedule( + deadline: .now() + self.interval, + repeating: self.interval, + leeway: .milliseconds(20)) + timer.setEventHandler { [weak self] in + Task { @MainActor [weak self] in + self?.tick() + } + } + self.timer = timer + timer.resume() + } + + func stop() { + self.timer?.cancel() + self.timer = nil + } + + private func tick() { + let now = ProcessInfo.processInfo.systemUptime + let elapsed = now - self.lastTick + self.lastTick = now + + let gap = max(0, elapsed - self.interval) + guard gap >= self.threshold else { return } + + var metadata = self.metadataProvider() + metadata["elapsedMs"] = String(format: "%.1f", elapsed * 1000) + metadata["gapMs"] = String(format: "%.1f", gap * 1000) + metadata["thresholdMs"] = String(format: "%.1f", self.threshold * 1000) + self.logger.warning("main runloop stall detected", metadata: metadata) + } + + deinit { + self.timer?.cancel() + } +} diff --git a/Sources/CodexBar/MenuCardView+ModelHelpers.swift b/Sources/CodexBar/MenuCardView+ModelHelpers.swift index 87dfd512aa..7e844d1894 100644 --- a/Sources/CodexBar/MenuCardView+ModelHelpers.swift +++ b/Sources/CodexBar/MenuCardView+ModelHelpers.swift @@ -15,6 +15,7 @@ extension UsageMenuCardView.Model { self.usageNotes.isEmpty && self.openAIAPIUsage == nil && self.inlineUsageDashboard == nil && + !self.inlineUsageDashboardLoading && self.creditsRemaining == nil && self.providerCost == nil && self.tokenUsage == nil && @@ -26,6 +27,7 @@ extension UsageMenuCardView.Model { !self.usageNotes.isEmpty || self.openAIAPIUsage != nil || self.inlineUsageDashboard != nil || + self.inlineUsageDashboardLoading || self.placeholder != nil } diff --git a/Sources/CodexBar/MenuCardView.swift b/Sources/CodexBar/MenuCardView.swift index 072810518b..82c51d3812 100644 --- a/Sources/CodexBar/MenuCardView.swift +++ b/Sources/CodexBar/MenuCardView.swift @@ -111,6 +111,7 @@ struct UsageMenuCardView: View { let creditsHintCopyText: String? let providerCost: ProviderCostSection? let tokenUsage: TokenUsageSection? + let inlineUsageDashboardLoading: Bool let placeholder: String? let progressColor: Color } @@ -137,6 +138,8 @@ struct UsageMenuCardView: View { if self.model.metrics.isEmpty { if let dashboard = self.model.inlineUsageDashboard { InlineUsageDashboardContent(model: dashboard) + } else if self.model.inlineUsageDashboardLoading { + InlineUsageDashboardLoadingContent() } else if !self.model.usageNotes.isEmpty { UsageNotesContent(notes: self.model.usageNotes) } else if let placeholder = self.model.placeholder { @@ -161,6 +164,8 @@ struct UsageMenuCardView: View { } if let dashboard = self.model.inlineUsageDashboard { InlineUsageDashboardContent(model: dashboard) + } else if self.model.inlineUsageDashboardLoading { + InlineUsageDashboardLoadingContent() } else if !self.model.usageNotes.isEmpty { UsageNotesContent(notes: self.model.usageNotes) } @@ -484,6 +489,8 @@ struct UsageMenuCardUsageSectionView: View { if self.model.metrics.isEmpty { if let dashboard = self.model.inlineUsageDashboard { InlineUsageDashboardContent(model: dashboard) + } else if self.model.inlineUsageDashboardLoading { + InlineUsageDashboardLoadingContent() } else if !self.model.usageNotes.isEmpty { UsageNotesContent(notes: self.model.usageNotes) } else if let placeholder = self.model.placeholder { @@ -500,6 +507,8 @@ struct UsageMenuCardUsageSectionView: View { } if let dashboard = self.model.inlineUsageDashboard { InlineUsageDashboardContent(model: dashboard) + } else if self.model.inlineUsageDashboardLoading { + InlineUsageDashboardLoadingContent() } else if !self.model.usageNotes.isEmpty { UsageNotesContent(notes: self.model.usageNotes) } @@ -685,6 +694,8 @@ extension UsageMenuCardView.Model { let dashboardError: String? let tokenSnapshot: CostUsageTokenSnapshot? let tokenError: String? + let tokenRefreshInFlight: Bool + let tokenRefreshQueued: Bool let account: AccountInfo let isRefreshing: Bool let lastError: String? @@ -711,6 +722,8 @@ extension UsageMenuCardView.Model { dashboardError: String?, tokenSnapshot: CostUsageTokenSnapshot?, tokenError: String?, + tokenRefreshInFlight: Bool = false, + tokenRefreshQueued: Bool = false, account: AccountInfo, isRefreshing: Bool, lastError: String?, @@ -736,6 +749,8 @@ extension UsageMenuCardView.Model { self.dashboardError = dashboardError self.tokenSnapshot = tokenSnapshot self.tokenError = tokenError + self.tokenRefreshInFlight = tokenRefreshInFlight + self.tokenRefreshQueued = tokenRefreshQueued self.account = account self.isRefreshing = isRefreshing self.lastError = lastError @@ -761,7 +776,12 @@ extension UsageMenuCardView.Model { metadata: input.metadata) let metrics = Self.metrics(input: input) let openAIAPIUsage = input.snapshot?.openAIAPIUsage + let tokenUsageSnapshot = Self.tokenUsageSnapshot(input: input) let inlineUsageDashboard = Self.inlineUsageDashboard(input: input) + let inlineUsageDashboardLoading = Self.inlineUsageDashboardLoading( + input: input, + tokenUsageSnapshot: tokenUsageSnapshot, + inlineUsageDashboard: inlineUsageDashboard) let usageNotes = Self.usageNotes(input: input) let creditsText: String? = if input.provider == .openrouter { nil @@ -783,7 +803,6 @@ extension UsageMenuCardView.Model { } else { Self.providerCostSection(provider: input.provider, cost: input.snapshot?.providerCost) } - let tokenUsageSnapshot = Self.tokenUsageSnapshot(input: input) let tokenUsage = Self.tokenUsageSection( provider: input.provider, enabled: input.tokenCostUsageEnabled, @@ -814,6 +833,7 @@ extension UsageMenuCardView.Model { creditsHintCopyText: redacted.creditsHintCopyText, providerCost: providerCost, tokenUsage: tokenUsage, + inlineUsageDashboardLoading: inlineUsageDashboardLoading, placeholder: placeholder, progressColor: Self.progressColor(for: input.provider)) } diff --git a/Sources/CodexBar/Providers/Codex/CodexConsumerProjection.swift b/Sources/CodexBar/Providers/Codex/CodexConsumerProjection.swift index 287db61f4c..46e87af75d 100644 --- a/Sources/CodexBar/Providers/Codex/CodexConsumerProjection.swift +++ b/Sources/CodexBar/Providers/Codex/CodexConsumerProjection.swift @@ -15,6 +15,10 @@ struct CodexUIErrorMapper { return trimmed } + if self.looksCancellation(lower: lower) { + return nil + } + if let cachedMessage = self.cachedMessage(raw: trimmed, lower: lower) { return cachedMessage } @@ -125,6 +129,14 @@ struct CodexUIErrorMapper { || lower.contains("returned invalid data") } + private static func looksCancellation(lower: String) -> Bool { + lower == "cancelled" + || lower.contains("network error: cancelled") + || lower.contains("nsurlerrordomain code=-999") + || lower.contains("nsurlerrordomain error -999") + || lower.contains("cancellationerror") + } + private static func looksOpenAIWebTimeout(lower: String) -> Bool { lower.contains("nsurlerrordomain") && (lower.contains("timed out") || lower.contains("error -1001")) diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift index 0952a47ae9..8234042dab 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift @@ -63,12 +63,22 @@ extension UsageStore { guard previousGuard != nil, previousGuard != currentGuard else { return false } - self.snapshots.removeValue(forKey: .codex) - self.errors[.codex] = nil - self.lastSourceLabels.removeValue(forKey: .codex) + let cachedSelection = self.cachedCodexVisibleAccountSnapshotForActiveSelection() + if let cachedSnapshot = cachedSelection?.snapshot { + self.snapshots[.codex] = cachedSnapshot + if let sourceLabel = cachedSelection?.sourceLabel { + self.lastSourceLabels[.codex] = sourceLabel + } else { + self.lastSourceLabels.removeValue(forKey: .codex) + } + self.errors[.codex] = nil + } else { + self.snapshots.removeValue(forKey: .codex) + self.errors[.codex] = cachedSelection?.error + self.lastSourceLabels.removeValue(forKey: .codex) + } self.lastFetchAttempts.removeValue(forKey: .codex) self.accountSnapshots.removeValue(forKey: .codex) - self.codexAccountSnapshots = [] self.failureGates[.codex]?.reset() self.lastKnownSessionRemaining.removeValue(forKey: .codex) self.lastKnownSessionWindowSource.removeValue(forKey: .codex) diff --git a/Sources/CodexBar/StatusItemController+HostedSubmenus.swift b/Sources/CodexBar/StatusItemController+HostedSubmenus.swift index 65bad39ddf..bc8fe9629f 100644 --- a/Sources/CodexBar/StatusItemController+HostedSubmenus.swift +++ b/Sources/CodexBar/StatusItemController+HostedSubmenus.swift @@ -37,6 +37,40 @@ extension StatusItemController { return submenu } + func makeStorageBreakdownSubmenu(provider: UsageProvider, width: CGFloat? = nil) -> NSMenu? { + guard self.store.storageFootprint(for: provider)?.components.isEmpty == false else { return nil } + if let width { + return self.makeHostedSubviewPlaceholderMenu( + chartID: Self.storageBreakdownID, + provider: provider, + width: width) + } + return self.makeHostedSubviewPlaceholderMenu(chartID: Self.storageBreakdownID, provider: provider) + } + + func isOpenAIWebSubviewMenu(_ menu: NSMenu) -> Bool { + let ids: Set = [ + Self.usageBreakdownChartID, + Self.creditsHistoryChartID, + ] + return menu.items.contains { item in + guard let id = item.representedObject as? String else { return false } + return ids.contains(id) + } + } + + func refreshHostedSubviewHeights(in menu: NSMenu) { + let width = self.renderedMenuWidth(for: menu) + + for item in menu.items { + guard let view = item.view else { continue } + view.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + view.layoutSubtreeIfNeeded() + let height = view.fittingSize.height + view.frame = NSRect(origin: .zero, size: NSSize(width: width, height: height)) + } + } + func hydrateHostedSubviewMenuIfNeeded(_ menu: NSMenu, width requestedWidth: CGFloat? = nil) { guard let placeholder = menu.items.first, menu.items.count == 1, @@ -46,8 +80,22 @@ extension StatusItemController { return } + let startedAt = Date() let width = requestedWidth ?? self.renderedMenuWidth(for: menu.supermenu ?? menu) menu.removeAllItems() + func logHydration(didHydrate: Bool) { + let elapsedMs = Date().timeIntervalSince(startedAt) * 1000 + if elapsedMs >= 16 { + self.menuLogger.info( + "hosted submenu hydration finished", + metadata: [ + "chartID": chartID, + "didHydrate": "\(didHydrate)", + "elapsedMs": String(format: "%.1f", elapsedMs), + "width": String(format: "%.0f", width), + ]) + } + } let didHydrate: Bool = switch chartID { case Self.usageBreakdownChartID: @@ -90,13 +138,17 @@ extension StatusItemController { false } - guard !didHydrate else { return } + guard !didHydrate else { + logHydration(didHydrate: true) + return + } let unavailableItem = NSMenuItem(title: L("No data available"), action: nil, keyEquivalent: "") unavailableItem.isEnabled = false unavailableItem.representedObject = chartID unavailableItem.toolTip = placeholder.toolTip menu.addItem(unavailableItem) + logHydration(didHydrate: false) } @discardableResult @@ -115,8 +167,8 @@ extension StatusItemController { let chartView = UsageBreakdownChartMenuView(breakdown: breakdown, width: width) let hosting = MenuHostingView(rootView: chartView) - let controller = NSHostingController(rootView: chartView) - let size = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + let size = hosting.fittingSize hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) let chartItem = NSMenuItem() @@ -142,8 +194,8 @@ extension StatusItemController { let chartView = CreditsHistoryChartMenuView(breakdown: breakdown, width: width) let hosting = MenuHostingView(rootView: chartView) - let controller = NSHostingController(rootView: chartView) - let size = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + let size = hosting.fittingSize hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) let chartItem = NSMenuItem() @@ -180,8 +232,8 @@ extension StatusItemController { windowLabel: tokenSnapshot.historyLabel, width: width) let hosting = MenuHostingView(rootView: chartView) - let controller = NSHostingController(rootView: chartView) - let size = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + let size = hosting.fittingSize hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) let chartItem = NSMenuItem() @@ -215,8 +267,8 @@ extension StatusItemController { let maxHeight = self.storageBreakdownMenuMaxHeight() let view = StorageBreakdownMenuView(footprint: footprint, width: width, maxHeight: maxHeight) let hosting = MenuHostingView(rootView: view) - let controller = NSHostingController(rootView: view) - let size = controller.sizeThatFits(in: CGSize(width: width, height: maxHeight)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: maxHeight)) + let size = hosting.fittingSize hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) let item = NSMenuItem() @@ -255,8 +307,8 @@ extension StatusItemController { let chartView = ZaiHourlyUsageChartMenuView(modelUsage: modelUsage, width: width) let hosting = MenuHostingView(rootView: chartView) - let controller = NSHostingController(rootView: chartView) - let size = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + let size = hosting.fittingSize hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) let chartItem = NSMenuItem() diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 5ff987741c..aecc807f25 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -6,11 +6,17 @@ import SwiftUI // MARK: - NSMenu construction +enum MenuPopulateMode { + case full + case preserveExistingSwitcherWidth +} + extension StatusItemController { static let menuCardBaseWidth: CGFloat = 310 private static let maxOverviewProviders = SettingsStore.mergedOverviewProviderLimit private static let overviewRowIdentifierPrefix = "overviewRow-" private static let defaultMenuOpenRefreshDelay: Duration = .seconds(1.2) + private static let openMenuContentRefreshDelay: Duration = .milliseconds(140) #if DEBUG private static var menuOpenRefreshDelayForTesting: Duration = .seconds(1.2) static func setMenuOpenRefreshDelayForTesting(_ delay: Duration) { @@ -49,22 +55,6 @@ extension StatusItemController { } } - private func menuCardWidth( - for providers: [UsageProvider], - sections: [MenuDescriptor.Section]) -> CGFloat - { - _ = providers - let baselineWidth = Self.menuCardBaseWidth - return max(baselineWidth, self.measuredStandardMenuWidth(for: sections, baseWidth: baselineWidth)) - } - - private func measuredStandardMenuWidth(for sections: [MenuDescriptor.Section], baseWidth: CGFloat) -> CGFloat { - let measuringMenu = NSMenu() - measuringMenu.autoenablesItems = false - self.addActionableSections(sections, to: measuringMenu, width: baseWidth) - return ceil(measuringMenu.size.width) - } - func makeMenu() -> NSMenu { guard self.shouldMergeIcons else { return self.makeMenu(for: nil) @@ -73,6 +63,22 @@ extension StatusItemController { } func menuWillOpen(_ menu: NSMenu) { + let startedAt = Date() + func logMenuWillOpen(kind: String, didRefresh: Bool) { + let elapsedMs = Date().timeIntervalSince(startedAt) * 1000 + if elapsedMs >= 16 { + self.menuLogger.info( + "menu will open processed", + metadata: [ + "didRefresh": "\(didRefresh)", + "elapsedMs": String(format: "%.1f", elapsedMs), + "items": "\(menu.items.count)", + "kind": kind, + ]) + } + } + self.beginMenuTokenCostDeferral(menu, reason: "menu-open") + if self.isHostedSubviewMenu(menu) { self.hydrateHostedSubviewMenuIfNeeded(menu) self.refreshHostedSubviewHeights(in: menu) @@ -85,6 +91,7 @@ extension StatusItemController { self.openMenus[ObjectIdentifier(menu)] = menu } // Removed redundant async refresh - single pass is sufficient after initial layout + logMenuWillOpen(kind: "hosted", didRefresh: true) return } @@ -107,11 +114,20 @@ extension StatusItemController { } } - let didRefresh = self.menuNeedsRefresh(menu) - if didRefresh { + let needsRefresh = self.menuNeedsRefresh(menu) + let canUseCachedMenu = needsRefresh && + !menu.items.isEmpty && + Self.menuRefreshEnabled && + self.isPersistentParentMenu(menu) && + self.canOpenStaleParentMenuFromCache(menu, provider: provider) + if needsRefresh, !canUseCachedMenu { self.populateMenu(menu, provider: provider) self.markMenuFresh(menu) // Heights are already set during populateMenu, no need to remeasure + } else if canUseCachedMenu { + self.menuLogger.debug( + "stale parent menu opened from cache", + metadata: ["provider": provider?.rawValue ?? "overview"]) } if Self.menuRefreshEnabled { // Intentionally skip open-menu tracking when refresh is disabled (tests). @@ -121,6 +137,44 @@ extension StatusItemController { // Only schedule refresh after menu is registered as open - refreshNow is called async self.scheduleOpenMenuRefresh(for: menu) } + logMenuWillOpen(kind: "parent", didRefresh: needsRefresh && !canUseCachedMenu) + } + + private func canOpenStaleParentMenuFromCache(_ menu: NSMenu, provider: UsageProvider?) -> Bool { + let enabledProviders = self.store.enabledProvidersForDisplay() + let includesOverview = self.includesOverviewTab(enabledProviders: enabledProviders) + let switcherSelection = self.shouldMergeIcons && enabledProviders.count > 1 + ? self.resolvedSwitcherSelection( + enabledProviders: enabledProviders, + includesOverview: includesOverview) + : nil + let selectedProvider = switcherSelection == .overview + ? nil + : (switcherSelection?.provider ?? provider ?? enabledProviders.first) + + let hasProviderSwitcher = menu.items.first?.view is ProviderSwitcherView + let wantsProviderSwitcher = self.shouldMergeIcons && enabledProviders.count > 1 + guard wantsProviderSwitcher, hasProviderSwitcher else { return false } + guard enabledProviders == self.lastSwitcherProviders else { return false } + guard self.settings.usageBarsShowUsed == self.lastSwitcherUsageBarsShowUsed else { return false } + guard includesOverview == self.lastSwitcherIncludesOverview else { return false } + guard switcherSelection == self.lastMergedSwitcherSelection else { return false } + + let hasCodexSwitcher = menu.items.contains { $0.view is CodexAccountSwitcherView } + let wantsCodexSwitcher = selectedProvider + .map { self.codexAccountMenuDisplay(for: $0)?.showSwitcher == true } ?? false + guard hasCodexSwitcher == wantsCodexSwitcher else { return false } + + let hasTokenSwitcher = menu.items.contains { $0.view is TokenAccountSwitcherView } + let wantsTokenSwitcher = selectedProvider + .map { self.tokenAccountMenuDisplay(for: $0)?.showSwitcher == true } ?? false + return hasTokenSwitcher == wantsTokenSwitcher + } + + private func isPersistentParentMenu(_ menu: NSMenu) -> Bool { + menu === self.mergedMenu || + menu === self.fallbackMenu || + self.providerMenus.values.contains { $0 === menu } } func menuDidClose(_ menu: NSMenu) { @@ -134,46 +188,63 @@ extension StatusItemController { func forgetClosedMenu(_ menu: NSMenu) { let key = ObjectIdentifier(menu) + self.endMenuTokenCostDeferral(menu, reason: "menu-close") + if key == self.providerSwitcherShortcutMenuID { self.removeProviderSwitcherShortcutMonitor() } self.openMenus.removeValue(forKey: key) + self.openMenuRefreshTokens.removeValue(forKey: key) self.menuRefreshTasks.removeValue(forKey: key)?.cancel() - self.openMenuRebuildTasks.removeValue(forKey: key)?.cancel() - self.openMenuRebuildTokens.removeValue(forKey: key) - self.openMenuRebuildsClosingHostedSubviewMenus.remove(key) - if let highlightedView = self.highlightedMenuItems.removeValue(forKey: key)?.view { - (highlightedView as? MenuCardHighlighting)?.setHighlighted(false) + if !self.closingHostedSubviewMenusForParentSwitch { + self.finishProviderSwitcherSettingsSuppression(nil, menuWasRebuilt: false) + self.applyDeferredSettingsChangeAfterMenuTrackingIfNeeded() + self.scheduleDeferredRefreshStoreAfterMenuTrackingIfNeeded(reason: "menu-close") } - let isPersistentMenu = menu === self.mergedMenu || - menu === self.fallbackMenu || - self.providerMenus.values.contains { $0 === menu } + let isPersistentMenu = self.isPersistentParentMenu(menu) if !isPersistentMenu { self.menuProviders.removeValue(forKey: key) self.menuVersions.removeValue(forKey: key) + } else if self.menuNeedsRefresh(menu) { + self.rebuildClosedMenuIfNeeded(menu) } - } - - func menu(_ menu: NSMenu, willHighlight item: NSMenuItem?) { - let key = ObjectIdentifier(menu) - let previous = self.highlightedMenuItems[key] - guard previous !== item else { return } - - if let previous { - (previous.view as? MenuCardHighlighting)?.setHighlighted(false) + self.highlightedMenuItems.removeValue(forKey: key) + for menuItem in menu.items { + (menuItem.view as? MenuCardHighlighting)?.setHighlighted(false) } + } - if let item, item.isEnabled { - self.highlightedMenuItems[key] = item - (item.view as? MenuCardHighlighting)?.setHighlighted(true) - } else { - self.highlightedMenuItems.removeValue(forKey: key) + private func rebuildClosedMenuIfNeeded(_ menu: NSMenu) { + let provider = self.menuProvider(for: menu) + Task { @MainActor [weak self, weak menu] in + await Task.yield() + guard let self, let menu else { return } + guard self.openMenus[ObjectIdentifier(menu)] == nil else { return } + guard self.menuNeedsRefresh(menu) else { return } + let startedAt = Date() + self.populateMenu(menu, provider: provider) + self.markMenuFresh(menu) + let elapsedMs = Date().timeIntervalSince(startedAt) * 1000 + if elapsedMs >= 16 { + self.menuLogger.info( + "closed parent menu rebuilt", + metadata: [ + "elapsedMs": String(format: "%.1f", elapsedMs), + "provider": provider?.rawValue ?? "overview", + ]) + } } } - func populateMenu(_ menu: NSMenu, provider: UsageProvider?) { + @discardableResult + func populateMenu( + _ menu: NSMenu, + provider: UsageProvider?, + mode: MenuPopulateMode = .full) + -> Bool + { let enabledProviders = self.store.enabledProvidersForDisplay() let includesOverview = self.includesOverviewTab(enabledProviders: enabledProviders) let switcherSelection = self.shouldMergeIcons && enabledProviders.count > 1 @@ -188,6 +259,25 @@ extension StatusItemController { switcherSelection?.provider ?? provider } let currentProvider = selectedProvider ?? enabledProviders.first ?? .codex + let populateStartedAt = Date() + var populatePath = "unknown" + var populateResult = true + defer { + let elapsedMs = Date().timeIntervalSince(populateStartedAt) * 1000 + if elapsedMs >= 16 { + self.menuLogger.info( + "populateMenu finished", + metadata: [ + "currentProvider": currentProvider.rawValue, + "elapsedMs": String(format: "%.1f", elapsedMs), + "items": "\(menu.items.count)", + "open": "\(self.openMenus[ObjectIdentifier(menu)] != nil)", + "path": populatePath, + "result": "\(populateResult)", + "selectedProvider": selectedProvider?.rawValue ?? "overview", + ]) + } + } let rawCodexAccountDisplay = isOverviewSelected ? nil : self.codexAccountMenuDisplay(for: currentProvider) let codexAccountDisplay = isOverviewSelected ? nil @@ -209,39 +299,21 @@ extension StatusItemController { codexAccountPromotionCoordinator: self.codexAccountPromotionCoordinator, updateReady: self.updater.updateStatus.isUpdateReady, includeContextualActions: !isOverviewSelected) - let menuWidth = self.menuCardWidth(for: enabledProviders, sections: descriptor.sections) + let menuWidth = self.resolvedMenuWidth( + enabledProviders: enabledProviders, + sections: descriptor.sections, + menu: menu, + mode: mode) - let hasTokenSwitcher = menu.items.contains { $0.view is TokenAccountSwitcherView } - let hasCodexSwitcher = menu.items.contains { $0.view is CodexAccountSwitcherView } - let switcherProvidersMatch = enabledProviders == self.lastSwitcherProviders - let switcherUsageBarsShowUsedMatch = self.settings.usageBarsShowUsed == self.lastSwitcherUsageBarsShowUsed - let switcherSelectionMatches = switcherSelection == self.lastMergedSwitcherSelection - let switcherOverviewAvailabilityMatches = includesOverview == self.lastSwitcherIncludesOverview - let tokenSwitcherCompatible = tokenAccountDisplay == self.lastTokenAccountMenuDisplay && - ((tokenAccountDisplay?.showSwitcher == true && hasTokenSwitcher) || - (tokenAccountDisplay?.showSwitcher != true && !hasTokenSwitcher)) - let codexSwitcherCompatible = codexAccountDisplay == self.lastCodexAccountMenuDisplay && - ((codexAccountDisplay?.showSwitcher == true && hasCodexSwitcher) || - (codexAccountDisplay?.showSwitcher != true && !hasCodexSwitcher)) - let reusableRowWidthsMatch = self.reusableFixedWidthRows(in: menu).allSatisfy { item in - guard let view = item.view else { return false } - return abs(view.frame.width - menuWidth) <= 0.5 - } - let providerSwitcherWidthMatches = (menu.items.first?.view as? ProviderSwitcherView).map { view in - abs(view.frame.width - menuWidth) <= 0.5 - } ?? false - let canSmartUpdate = self.shouldMergeIcons && - enabledProviders.count > 1 && - !isOverviewSelected && - switcherProvidersMatch && - switcherUsageBarsShowUsedMatch && - switcherSelectionMatches && - switcherOverviewAvailabilityMatches && - tokenSwitcherCompatible && - codexSwitcherCompatible && - reusableRowWidthsMatch && - !menu.items.isEmpty && - menu.items.first?.view is ProviderSwitcherView + let compatibility = self.menuPopulationCompatibility(.init( + menu: menu, + enabledProviders: enabledProviders, + includesOverview: includesOverview, + switcherSelection: switcherSelection, + codexAccountDisplay: codexAccountDisplay, + tokenAccountDisplay: tokenAccountDisplay, + menuWidth: menuWidth, + isOverviewSelected: isOverviewSelected)) #if DEBUG if self.openMenus[ObjectIdentifier(menu)] != nil { @@ -251,11 +323,12 @@ extension StatusItemController { "available=\(self.store.enabledProviders().map(\.rawValue)) " + "selection=\(String(describing: switcherSelection)) " + "last=\(String(describing: self.lastMergedSwitcherSelection)) " + - "smart=\(canSmartUpdate)") + "smart=\(compatibility.canSmartUpdate)") } #endif - if canSmartUpdate { + if compatibility.canSmartUpdate { + populatePath = "smart" self.updateMenuContentPreservingSwitcher( menu, context: MenuUpdateContext( @@ -266,27 +339,19 @@ extension StatusItemController { codexAccountDisplay: codexAccountDisplay, tokenAccountDisplay: tokenAccountDisplay, openAIContext: openAIContext)) - return + return true } - let canPreserveProviderSwitcher = self.shouldMergeIcons && - enabledProviders.count > 1 && - switcherProvidersMatch && - switcherUsageBarsShowUsedMatch && - switcherOverviewAvailabilityMatches && - providerSwitcherWidthMatches && - !menu.items.isEmpty && - menu.items.first?.view is ProviderSwitcherView - #if DEBUG if self.openMenus[ObjectIdentifier(menu)] != nil { self.menuLogger.debug( - "populateMenu(open): preserveSwitcher=\(canPreserveProviderSwitcher) " + - "widthMatch=\(providerSwitcherWidthMatches)") + "populateMenu(open): preserveSwitcher=\(compatibility.canPreserveProviderSwitcher) " + + "widthMatch=\(compatibility.providerSwitcherWidthMatches)") } #endif - if canPreserveProviderSwitcher { + if compatibility.canPreserveProviderSwitcher { + populatePath = "preserve-switcher" self.updateMenuContentPreservingSwitcher( menu, context: MenuUpdateContext( @@ -297,7 +362,14 @@ extension StatusItemController { codexAccountDisplay: codexAccountDisplay, tokenAccountDisplay: tokenAccountDisplay, openAIContext: openAIContext)) - return + return true + } + + if mode == .preserveExistingSwitcherWidth { + populatePath = "switcher-only" + populateResult = false + self.updateProviderSwitcherOnly(in: menu, selection: switcherSelection) + return false } #if DEBUG @@ -305,6 +377,7 @@ extension StatusItemController { self.menuLogger.debug("populateMenu(open): rebuilding whole menu and replacing provider switcher") } #endif + populatePath = "rebuild" self.rebuildMenuContent( menu, context: MenuRebuildContext( @@ -318,28 +391,19 @@ extension StatusItemController { tokenAccountDisplay: tokenAccountDisplay, openAIContext: openAIContext, descriptor: descriptor)) + return true } - private func reusableFixedWidthRows(in menu: NSMenu) -> [NSMenuItem] { - guard !menu.items.isEmpty else { return [] } - - var reusableRows: [NSMenuItem] = [] - var index = self.providerSwitcherContentStartIndex(in: menu) - if index > 0 { - reusableRows.append(menu.items[0]) - } - if menu.items.count > index, - menu.items[index].view is CodexAccountSwitcherView - { - reusableRows.append(menu.items[index]) - index += 2 - } - if menu.items.count > index, - menu.items[index].view is TokenAccountSwitcherView - { - reusableRows.append(menu.items[index]) - } - return reusableRows + private func updateProviderSwitcherOnly( + in menu: NSMenu, + selection: ProviderSwitcherSelection?) + { + guard let switcherView = menu.items.first?.view as? ProviderSwitcherView, + let selection + else { return } + switcherView.updateSelection(selection) + switcherView.updateQuotaIndicators() + self.lastMergedSwitcherSelection = selection } /// Smart update: rebuild everything below the provider switcher while keeping the switcher view intact. @@ -924,125 +988,7 @@ extension StatusItemController { return menu } - private func makeProviderSwitcherItem( - providers: [UsageProvider], - includesOverview: Bool, - selected: ProviderSwitcherSelection, - menu: NSMenu, - width: CGFloat) -> NSMenuItem - { - let view = ProviderSwitcherView( - providers: providers, - selected: selected, - includesOverview: includesOverview, - width: width, - showsIcons: self.settings.switcherShowsIcons, - iconProvider: { [weak self] provider in - self?.switcherIcon(for: provider) ?? NSImage() - }, - weeklyRemainingProvider: { [weak self] provider in - self?.switcherWeeklyRemaining(for: provider) - }, - onSelect: { [weak self, weak menu] selection in - guard let self, let menu else { return } - let provider: UsageProvider? - switch selection { - case .overview: - self.settings.mergedMenuLastSelectedWasOverview = true - provider = self.resolvedMenuProvider() - case let .provider(selectedProvider): - self.settings.mergedMenuLastSelectedWasOverview = false - self.selectedMenuProvider = selectedProvider - provider = selectedProvider - } - switch selection { - case .overview: - self.lastMenuProvider = provider ?? .codex - case let .provider(provider): - self.lastMenuProvider = provider - } - self.lastMergedSwitcherSelection = selection - self.deferSwitcherMenuRebuildIfStillVisible(menu, provider: provider) - }) - let item = NSMenuItem() - item.view = view - item.isEnabled = false - return item - } - - private func makeTokenAccountSwitcherItem( - display: TokenAccountMenuDisplay, - menu: NSMenu, - width: CGFloat) -> NSMenuItem - { - let view = TokenAccountSwitcherView( - accounts: display.accounts, - selectedIndex: display.activeIndex, - width: width, - onSelect: { [weak self, weak menu] index -> Task? in - guard let self, let menu else { return nil } - self.settings.setActiveTokenAccountIndex(index, for: display.provider) - self.applyIcon(phase: nil) - self.deferSwitcherMenuRebuildIfStillVisible(menu, provider: display.provider) - return Task { @MainActor [weak self, weak menu] in - guard let self else { return } - await ProviderInteractionContext.$current.withValue(.userInitiated) { - await self.store.refreshProvider(display.provider) - } - guard let menu else { return } - self.refreshOpenMenuIfStillVisible(menu, provider: display.provider) - } - }) - let item = NSMenuItem() - item.view = view - item.isEnabled = false - return item - } - - private func makeCodexAccountSwitcherItem( - display: CodexAccountMenuDisplay, - menu: NSMenu, - width: CGFloat) -> NSMenuItem - { - let view = CodexAccountSwitcherView( - accounts: display.accounts, - selectedAccountID: display.activeVisibleAccountID, - width: width, - onSelect: { [weak self, weak menu] account in - guard let self else { return } - self.handleCodexVisibleAccountSelection(account, menu: menu) - }) - let item = NSMenuItem() - item.view = view - item.isEnabled = false - return item - } - - @discardableResult - private func handleCodexVisibleAccountSelection(_ account: CodexVisibleAccount, menu: NSMenu?) -> Bool { - let visibleAccountID = account.id - self.settings.selectDisplayedCodexVisibleAccount(account) - if self.store.prepareCodexAccountScopedRefreshIfNeeded(), let menu { - self.deferSwitcherMenuRebuildIfStillVisible(menu, provider: .codex) - } - Task { @MainActor in - await ProviderInteractionContext.$current.withValue(.userInitiated) { - await self.store.refreshCodexAccountScopedState( - allowDisabled: true, - phaseDidChange: { [weak self, weak menu] _ in - guard let self, let menu else { return } - guard self.settings.codexVisibleAccountProjection.activeVisibleAccountID == visibleAccountID - else { - return - } - self.refreshOpenMenuIfStillVisible(menu, provider: .codex) - }) - } - } - return true - } - - private func resolvedMenuProvider(enabledProviders: [UsageProvider]? = nil) -> UsageProvider? { + func resolvedMenuProvider(enabledProviders: [UsageProvider]? = nil) -> UsageProvider? { let enabled = enabledProviders ?? self.store.enabledProvidersForDisplay() if enabled.isEmpty { return .codex } if let selected = self.selectedMenuProvider, enabled.contains(selected) { @@ -1097,25 +1043,83 @@ extension StatusItemController { } func refreshOpenMenuIfStillVisible(_ menu: NSMenu, provider: UsageProvider?) { - self.scheduleOpenMenuRebuildIfStillVisible(menu, provider: provider) + guard Self.menuRefreshEnabled else { return } + guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } + let menuID = ObjectIdentifier(menu) + self.openMenuRefreshGeneration &+= 1 + let token = self.openMenuRefreshGeneration + self.openMenuRefreshTokens[menuID] = token + let scheduledAt = Date() + self.menuLogger.debug( + "open menu refresh scheduled", + metadata: [ + "provider": provider?.rawValue ?? "overview", + "token": "\(token)", + ]) + Task { @MainActor [weak self, weak menu] in + guard let self, let menu else { return } + #if DEBUG + if let override = self._test_openMenuRefreshYieldOverride { + await override() + } else { + try? await Task.sleep(for: Self.openMenuContentRefreshDelay) + } + #else + try? await Task.sleep(for: Self.openMenuContentRefreshDelay) + #endif + guard !Task.isCancelled else { return } + guard self.openMenuRefreshTokens[menuID] == token else { + self.menuLogger.debug( + "open menu refresh skipped stale", + metadata: [ + "currentToken": "\(self.openMenuRefreshTokens[menuID] ?? -1)", + "token": "\(token)", + ]) + return + } + guard self.openMenus[menuID] != nil else { + self.menuLogger.debug( + "open menu refresh skipped closed menu", + metadata: ["token": "\(token)"]) + return + } + self.openMenuRefreshTokens.removeValue(forKey: menuID) + let rebuildStartedAt = Date() + self.rebuildOpenMenuIfStillVisible(menu, provider: provider) + let rebuildMs = Date().timeIntervalSince(rebuildStartedAt) * 1000 + let totalMs = Date().timeIntervalSince(scheduledAt) * 1000 + self.menuLogger.info( + "open menu refresh finished", + metadata: [ + "provider": provider?.rawValue ?? "overview", + "rebuildMs": String(format: "%.1f", rebuildMs), + "token": "\(token)", + "totalMs": String(format: "%.1f", totalMs), + ]) + } } func rebuildOpenMenuIfStillVisible(_ menu: NSMenu, provider: UsageProvider?) { - guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } - guard self.isHostedSubviewMenu(menu) || !self.hasOpenHostedSubviewMenu() else { return } - self.populateMenu(menu, provider: provider) - self.markMenuFresh(menu) - self.applyIcon(phase: nil) - #if DEBUG - self._test_openMenuRebuildObserver?(menu) - #endif + if self.isHostedSubviewMenu(menu) { + guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } + self.populateMenu(menu, provider: provider) + self.markMenuFresh(menu) + self.applyIcon(phase: nil) + #if DEBUG + self._test_openMenuRebuildObserver?(menu) + #endif + } else { + self.updateTrackedOpenParentMenuContentIfPossible(menu, provider: provider, reason: "open-menu-rebuild") + } } private func scheduleOpenMenuRefresh(for menu: NSMenu) { - // Kick off a refresh on open (non-forced) and re-check after a delay. - // NEVER block menu opening with network requests. - if !self.store.isRefreshing { - self.refreshStore(forceTokenUsage: false, refreshOpenMenusWhenComplete: false) + if self.menuNeedsRefresh(menu) { + self.refreshOpenMenuIfStillVisible(menu, provider: self.menuProvider(for: menu)) + } + // Keep menu tracking responsive: render cached menu state while open and refresh after interaction ends. + if !self.store.isRefreshing, self.menuNeedsDelayedRefreshRetry(for: menu) { + self.deferRefreshStoreUntilMenuInteractionEnds(reason: "menu-open-stale") } let key = ObjectIdentifier(menu) self.menuRefreshTasks[key]?.cancel() @@ -1134,7 +1138,7 @@ extension StatusItemController { let retryMissingSnapshotCount = retryProviders.count { self.store.snapshot(for: $0) == nil } let willRetryRefresh = retryStaleProviderCount > 0 || retryMissingSnapshotCount > 0 guard willRetryRefresh else { return } - self.refreshStore(forceTokenUsage: false, refreshOpenMenusWhenComplete: false) + self.deferRefreshStoreUntilMenuInteractionEnds(reason: "menu-open-delayed-stale") } } @@ -1216,10 +1220,23 @@ extension StatusItemController { { view } + let startedAt = Date() let hosting = MenuCardItemHostingView(rootView: wrapped, highlightState: highlightState, onClick: onClick) // Set frame with target width immediately let height = self.menuCardHeight(for: hosting, width: width) hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: height)) + let elapsedMs = Date().timeIntervalSince(startedAt) * 1000 + if elapsedMs >= 16 { + self.menuLogger.info( + "menu card item built", + metadata: [ + "elapsedMs": String(format: "%.1f", elapsedMs), + "height": String(format: "%.0f", height), + "id": id, + "submenu": "\(submenu != nil)", + "width": String(format: "%.0f", width), + ]) + } let item = NSMenuItem() item.view = hosting item.isEnabled = true @@ -1363,7 +1380,7 @@ extension StatusItemController { return true } - private func switcherIcon(for provider: UsageProvider) -> NSImage { + func switcherIcon(for provider: UsageProvider) -> NSImage { if let brand = ProviderBrandIcon.image(for: provider) { return brand } @@ -1578,40 +1595,6 @@ extension StatusItemController { provider == .openai && self.tokenSnapshotForCostHistorySubmenu(provider: provider)?.daily.isEmpty == false } - func makeStorageBreakdownSubmenu(provider: UsageProvider, width: CGFloat? = nil) -> NSMenu? { - guard self.store.storageFootprint(for: provider)?.components.isEmpty == false else { return nil } - if let width { - return self.makeHostedSubviewPlaceholderMenu( - chartID: Self.storageBreakdownID, - provider: provider, - width: width) - } - return self.makeHostedSubviewPlaceholderMenu(chartID: Self.storageBreakdownID, provider: provider) - } - - private func isOpenAIWebSubviewMenu(_ menu: NSMenu) -> Bool { - let ids: Set = [ - Self.usageBreakdownChartID, - Self.creditsHistoryChartID, - ] - return menu.items.contains { item in - guard let id = item.representedObject as? String else { return false } - return ids.contains(id) - } - } - - func refreshHostedSubviewHeights(in menu: NSMenu) { - let width = self.renderedMenuWidth(for: menu) - - for item in menu.items { - guard let view = item.view else { continue } - view.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) - view.layoutSubtreeIfNeeded() - let height = view.fittingSize.height - view.frame = NSRect(origin: .zero, size: NSSize(width: width, height: height)) - } - } - @objc private func menuCardNoOp(_ sender: NSMenuItem) { _ = sender } diff --git a/Sources/CodexBar/StatusItemController+MenuCardModel.swift b/Sources/CodexBar/StatusItemController+MenuCardModel.swift index 1518d12235..9a52cdb87d 100644 --- a/Sources/CodexBar/StatusItemController+MenuCardModel.swift +++ b/Sources/CodexBar/StatusItemController+MenuCardModel.swift @@ -25,6 +25,9 @@ extension StatusItemController { } else { snapshotOverride ?? self.store.snapshot(for: target) } + if surface == .liveCard, snapshotOverride == nil { + self.store.ensureTokenCostSnapshotScheduled(for: target, reason: "menu-card") + } let projectedTokenSnapshot = self.store.tokenSnapshot(fromProviderSnapshot: snapshot, provider: target) let storedTokenSnapshot = UsageStore.tokenCostRequiresProviderSnapshot(target) ? nil @@ -96,6 +99,8 @@ extension StatusItemController { dashboardError: dashboardError, tokenSnapshot: tokenSnapshot, tokenError: tokenError, + tokenRefreshInFlight: self.store.isTokenRefreshInFlight(for: target), + tokenRefreshQueued: self.store.isTokenRefreshQueued(for: target), account: accountOverride ?? self.store.accountInfo(for: target), isRefreshing: self.store.shouldShowRefreshingMenuCard(for: target), lastError: errorOverride diff --git a/Sources/CodexBar/StatusItemController+MenuHighlighting.swift b/Sources/CodexBar/StatusItemController+MenuHighlighting.swift new file mode 100644 index 0000000000..0cf585f126 --- /dev/null +++ b/Sources/CodexBar/StatusItemController+MenuHighlighting.swift @@ -0,0 +1,20 @@ +import AppKit + +extension StatusItemController { + func menu(_ menu: NSMenu, willHighlight item: NSMenuItem?) { + let key = ObjectIdentifier(menu) + let previous = self.highlightedMenuItems[key] + guard previous !== item else { return } + + if let previous { + (previous.view as? MenuCardHighlighting)?.setHighlighted(false) + } + + if let item, item.isEnabled { + self.highlightedMenuItems[key] = item + (item.view as? MenuCardHighlighting)?.setHighlighted(true) + } else { + self.highlightedMenuItems.removeValue(forKey: key) + } + } +} diff --git a/Sources/CodexBar/StatusItemController+MenuPopulation.swift b/Sources/CodexBar/StatusItemController+MenuPopulation.swift new file mode 100644 index 0000000000..2c4dea5e29 --- /dev/null +++ b/Sources/CodexBar/StatusItemController+MenuPopulation.swift @@ -0,0 +1,92 @@ +import AppKit +import CodexBarCore + +extension StatusItemController { + struct MenuPopulationCompatibility { + let canSmartUpdate: Bool + let canPreserveProviderSwitcher: Bool + let providerSwitcherWidthMatches: Bool + } + + struct MenuPopulationInput { + let menu: NSMenu + let enabledProviders: [UsageProvider] + let includesOverview: Bool + let switcherSelection: ProviderSwitcherSelection? + let codexAccountDisplay: CodexAccountMenuDisplay? + let tokenAccountDisplay: TokenAccountMenuDisplay? + let menuWidth: CGFloat + let isOverviewSelected: Bool + } + + func menuPopulationCompatibility(_ input: MenuPopulationInput) -> MenuPopulationCompatibility { + let hasTokenSwitcher = input.menu.items.contains { $0.view is TokenAccountSwitcherView } + let hasCodexSwitcher = input.menu.items.contains { $0.view is CodexAccountSwitcherView } + let switcherProvidersMatch = input.enabledProviders == self.lastSwitcherProviders + let switcherUsageBarsShowUsedMatch = self.settings.usageBarsShowUsed == self.lastSwitcherUsageBarsShowUsed + let switcherSelectionMatches = input.switcherSelection == self.lastMergedSwitcherSelection + let switcherOverviewAvailabilityMatches = input.includesOverview == self.lastSwitcherIncludesOverview + let tokenSwitcherCompatible = input.tokenAccountDisplay == self.lastTokenAccountMenuDisplay && + ((input.tokenAccountDisplay?.showSwitcher == true && hasTokenSwitcher) || + (input.tokenAccountDisplay?.showSwitcher != true && !hasTokenSwitcher)) + let codexSwitcherCompatible = input.codexAccountDisplay == self.lastCodexAccountMenuDisplay && + ((input.codexAccountDisplay?.showSwitcher == true && hasCodexSwitcher) || + (input.codexAccountDisplay?.showSwitcher != true && !hasCodexSwitcher)) + let reusableRowWidthsMatch = self.reusableFixedWidthRows(in: input.menu).allSatisfy { item in + guard let view = item.view else { return false } + return abs(view.frame.width - input.menuWidth) <= 0.5 + } + let providerSwitcherWidthMatches = (input.menu.items.first?.view as? ProviderSwitcherView).map { view in + abs(view.frame.width - input.menuWidth) <= 0.5 + } ?? false + + let canSmartUpdate = self.shouldMergeIcons && + input.enabledProviders.count > 1 && + !input.isOverviewSelected && + switcherProvidersMatch && + switcherUsageBarsShowUsedMatch && + switcherSelectionMatches && + switcherOverviewAvailabilityMatches && + tokenSwitcherCompatible && + codexSwitcherCompatible && + reusableRowWidthsMatch && + !input.menu.items.isEmpty && + input.menu.items.first?.view is ProviderSwitcherView + + let canPreserveProviderSwitcher = self.shouldMergeIcons && + input.enabledProviders.count > 1 && + switcherProvidersMatch && + switcherUsageBarsShowUsedMatch && + switcherOverviewAvailabilityMatches && + providerSwitcherWidthMatches && + !input.menu.items.isEmpty && + input.menu.items.first?.view is ProviderSwitcherView + + return MenuPopulationCompatibility( + canSmartUpdate: canSmartUpdate, + canPreserveProviderSwitcher: canPreserveProviderSwitcher, + providerSwitcherWidthMatches: providerSwitcherWidthMatches) + } + + func reusableFixedWidthRows(in menu: NSMenu) -> [NSMenuItem] { + guard !menu.items.isEmpty else { return [] } + + var reusableRows: [NSMenuItem] = [] + var index = self.providerSwitcherContentStartIndex(in: menu) + if index > 0 { + reusableRows.append(menu.items[0]) + } + if menu.items.count > index, + menu.items[index].view is CodexAccountSwitcherView + { + reusableRows.append(menu.items[index]) + index += 2 + } + if menu.items.count > index, + menu.items[index].view is TokenAccountSwitcherView + { + reusableRows.append(menu.items[index]) + } + return reusableRows + } +} diff --git a/Sources/CodexBar/StatusItemController+MenuPresentation.swift b/Sources/CodexBar/StatusItemController+MenuPresentation.swift index f1097a50b5..913f29051b 100644 --- a/Sources/CodexBar/StatusItemController+MenuPresentation.swift +++ b/Sources/CodexBar/StatusItemController+MenuPresentation.swift @@ -127,9 +127,8 @@ final class MenuCardItemHostingView: NSHostingView, Menu } func measuredHeight(width: CGFloat) -> CGFloat { - let controller = NSHostingController(rootView: self.rootView) - let measured = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) - return measured.height + self.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + return self.fittingSize.height } func setHighlighted(_ highlighted: Bool) { diff --git a/Sources/CodexBar/StatusItemController+MenuPrewarming.swift b/Sources/CodexBar/StatusItemController+MenuPrewarming.swift new file mode 100644 index 0000000000..f3519b8d85 --- /dev/null +++ b/Sources/CodexBar/StatusItemController+MenuPrewarming.swift @@ -0,0 +1,140 @@ +import AppKit +import CodexBarCore +import Foundation +import ObjectiveC + +@MainActor +private final class PersistentMenuPrewarmState { + var task: Task? + var scheduledVersion: Int? + var completedVersion: Int? +} + +private enum PersistentMenuPrewarmAssociation { + nonisolated(unsafe) static var key: UInt8 = 0 +} + +extension StatusItemController { + private var persistentMenuPrewarmState: PersistentMenuPrewarmState { + if let state = objc_getAssociatedObject( + self, + &PersistentMenuPrewarmAssociation.key) as? PersistentMenuPrewarmState + { + return state + } + + let state = PersistentMenuPrewarmState() + objc_setAssociatedObject( + self, + &PersistentMenuPrewarmAssociation.key, + state, + .OBJC_ASSOCIATION_RETAIN_NONATOMIC) + return state + } + + func startResponsiveMenuSupportIfNeeded() { + self.startMainRunLoopStallMonitorIfNeeded() + self.schedulePersistentMenuPrewarm(reason: "startup") + } + + func schedulePersistentMenuPrewarm(reason: String) { + guard Self.menuRefreshEnabled else { return } + guard !SettingsStore.isRunningTests else { return } + + let version = self.menuContentVersion + let state = self.persistentMenuPrewarmState + guard state.completedVersion != version else { return } + if state.scheduledVersion == version, state.task != nil { + return + } + + state.task?.cancel() + state.scheduledVersion = version + let delay: Duration = reason == "startup" ? .milliseconds(250) : .milliseconds(900) + state.task = Task { @MainActor [weak self] in + try? await Task.sleep(for: delay) + guard let self, !Task.isCancelled else { return } + let state = self.persistentMenuPrewarmState + guard state.scheduledVersion == version else { return } + state.task = nil + guard !self.hasPreparedForAppShutdown else { return } + guard self.openMenus.isEmpty else { return } + guard self.persistentMenusNeedPrewarm(reason: reason) else { + state.completedVersion = version + return + } + self.prewarmPersistentMenus(reason: reason) + state.completedVersion = version + } + } + + private func persistentMenusNeedPrewarm(reason: String) -> Bool { + if reason == "startup" { + return true + } + if self.shouldMergeIcons { + return self.mergedMenu.map { self.menuNeedsRefresh($0) } ?? false + } + + for menu in self.providerMenus.values where self.menuNeedsRefresh(menu) { + return true + } + return self.fallbackMenu.map { self.menuNeedsRefresh($0) } ?? false + } + + private func prewarmPersistentMenus(reason: String) { + guard Self.menuRefreshEnabled else { return } + guard self.openMenus.isEmpty else { return } + + let startedAt = Date() + let count: Int = if self.shouldMergeIcons { + self.prewarmMergedMenu() + } else { + self.prewarmSplitProviderMenus() + } + + let elapsedMs = Date().timeIntervalSince(startedAt) * 1000 + guard count > 0 || elapsedMs >= 16 else { return } + self.menuLogger.info( + "persistent menus prewarmed", + metadata: [ + "elapsedMs": String(format: "%.1f", elapsedMs), + "menus": "\(count)", + "reason": reason, + ]) + } + + private func prewarmMergedMenu() -> Int { + let menu = self.mergedMenu ?? self.makeMenu() + self.mergedMenu = menu + if self.statusItem.menu !== menu { + self.statusItem.menu = menu + } + return self.populateClosedMenuIfNeeded(menu, provider: self.resolvedMenuProvider()) ? 1 : 0 + } + + private func prewarmSplitProviderMenus() -> Int { + var count = 0 + let fallback = self.fallbackProvider + for provider in self.settings.orderedProviders() { + if self.isEnabled(provider) { + let menu = self.providerMenus[provider] ?? self.makeMenu(for: provider) + self.providerMenus[provider] = menu + count += self.populateClosedMenuIfNeeded(menu, provider: provider) ? 1 : 0 + } else if fallback == provider { + let menu = self.fallbackMenu ?? self.makeMenu(for: nil) + self.fallbackMenu = menu + count += self.populateClosedMenuIfNeeded(menu, provider: nil) ? 1 : 0 + } + } + return count + } + + private func populateClosedMenuIfNeeded(_ menu: NSMenu, provider: UsageProvider?) -> Bool { + guard self.openMenus[ObjectIdentifier(menu)] == nil else { return false } + guard self.menuNeedsRefresh(menu) else { return false } + self.populateMenu(menu, provider: provider) + self.markMenuFresh(menu) + return true + } +} diff --git a/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift b/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift index 594aa483d4..3266416777 100644 --- a/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift +++ b/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift @@ -10,65 +10,298 @@ extension StatusItemController { updates() } - func deferSwitcherMenuRebuildIfStillVisible(_ menu: NSMenu, provider: UsageProvider?) { + @discardableResult + func beginProviderSwitcherSettingsSuppression() -> Int { + self.providerSwitcherSettingsSuppressionGeneration &+= 1 + let token = self.providerSwitcherSettingsSuppressionGeneration + self.activeProviderSwitcherSettingsSuppressionGeneration = token + self.providerSwitcherDeferredIconUpdatePending = false + return token + } + + func finishProviderSwitcherSettingsSuppression( + _ token: Int?, + menuWasRebuilt: Bool, + menuWasInvalidated: Bool = false) + { + guard let active = self.activeProviderSwitcherSettingsSuppressionGeneration else { return } + if let token, token != active { return } + + self.activeProviderSwitcherSettingsSuppressionGeneration = nil + self.providerSwitcherSettingsObservationSuppressionsRemaining = max( + self.providerSwitcherSettingsObservationSuppressionsRemaining, + 2) + if !menuWasRebuilt, !menuWasInvalidated { + self.menuContentVersion &+= 1 + } + guard self.providerSwitcherDeferredIconUpdatePending else { return } + self.providerSwitcherDeferredIconUpdatePending = false + self.applyIcon(phase: nil) + } + + func deferSettingsChangeDuringMenuTracking( + reason: String, + needsStatusItemRebuild: Bool, + needsOpenMenuRefresh: Bool) + { + self.settingsChangeDeferredDuringMenuTracking = true + self.deferredSettingsChangeNeedsStatusItemRebuild = + self.deferredSettingsChangeNeedsStatusItemRebuild || needsStatusItemRebuild + self.deferredSettingsChangeNeedsOpenMenuRefresh = + self.deferredSettingsChangeNeedsOpenMenuRefresh || needsOpenMenuRefresh + self.menuLogger.debug( + "settings change deferred during menu tracking", + metadata: [ + "needsOpenMenuRefresh": "\(needsOpenMenuRefresh)", + "needsStatusItemRebuild": "\(needsStatusItemRebuild)", + "openMenus": "\(self.openMenus.count)", + "reason": reason, + ]) + } + + func applyDeferredSettingsChangeAfterMenuTrackingIfNeeded() { + guard self.openMenus.isEmpty else { return } + let storeChangeDeferred = self.storeChangeDeferredDuringMenuTracking + let storeChangeDeferredVersion = self.storeChangeDeferredMenuContentVersion + let storeChangeObservationCount = self.storeChangeDeferredObservationCount + self.storeChangeDeferredDuringMenuTracking = false + self.storeChangeDeferredMenuContentVersion = nil + self.storeChangeDeferredObservationCount = 0 + if storeChangeDeferred { + self.observeStoreChanges() + } + guard self.settingsChangeDeferredDuringMenuTracking || storeChangeDeferred else { return } + + let needsStatusItemRebuild = self.deferredSettingsChangeNeedsStatusItemRebuild + let needsOpenMenuRefresh = self.deferredSettingsChangeNeedsOpenMenuRefresh + let settingsChangeDeferred = self.settingsChangeDeferredDuringMenuTracking + self.settingsChangeDeferredDuringMenuTracking = false + self.deferredSettingsChangeNeedsStatusItemRebuild = false + self.deferredSettingsChangeNeedsOpenMenuRefresh = false + let shouldInvalidateForStoreChange = storeChangeDeferred && + storeChangeDeferredVersion == self.menuContentVersion + + self.menuLogger.info( + "deferred menu change applied after menu tracking", + metadata: [ + "needsOpenMenuRefresh": "\(needsOpenMenuRefresh)", + "needsStatusItemRebuild": "\(needsStatusItemRebuild)", + "storeChangeDeferred": "\(storeChangeDeferred)", + "storeObservationCount": "\(storeChangeObservationCount)", + ]) + if settingsChangeDeferred || shouldInvalidateForStoreChange { + self.invalidateMenus(refreshOpenMenus: needsOpenMenuRefresh) + } + if needsStatusItemRebuild { + self.rebuildProviderStatusItems() + } + self.updateVisibility() + self.updateIcons() + } + + func deferRefreshStoreUntilMenuInteractionEnds(reason: String) { + self.deferredMenuInteractionRefreshPending = true + self.deferredMenuInteractionRefreshGeneration &+= 1 + self.deferredMenuInteractionRefreshTask?.cancel() + self.deferredMenuInteractionRefreshTask = nil + self.menuLogger.debug( + "store refresh deferred until menu interaction ends", + metadata: [ + "openMenus": "\(self.openMenus.count)", + "reason": reason, + ]) + } + + func scheduleDeferredRefreshStoreAfterMenuTrackingIfNeeded(reason: String) { + guard self.openMenus.isEmpty else { return } + guard self.deferredMenuInteractionRefreshPending else { return } + guard self.deferredMenuInteractionRefreshTask == nil else { return } + + self.deferredMenuInteractionRefreshGeneration &+= 1 + let generation = self.deferredMenuInteractionRefreshGeneration + self.deferredMenuInteractionRefreshTask = Task { @MainActor [weak self] in + try? await Task.sleep(for: .seconds(1)) + guard let self, !Task.isCancelled else { return } + guard self.deferredMenuInteractionRefreshGeneration == generation else { return } + self.deferredMenuInteractionRefreshTask = nil + guard self.openMenus.isEmpty else { return } + self.deferredMenuInteractionRefreshPending = false + guard !self.store.isRefreshing else { return } + self.menuLogger.debug( + "store refresh resumed after menu interaction", + metadata: ["reason": reason]) + self.refreshStore(forceTokenUsage: false, refreshOpenMenusWhenComplete: false) + } + } + + func noteProviderSwitcherInteraction() { + self.lastProviderSwitcherInteractionAt = Date() + } + + func deferSwitcherMenuRebuildIfStillVisible( + _ menu: NSMenu, + provider: UsageProvider?, + settingsSuppressionToken: Int? = nil) + { self.providerSwitcherUpdateToken &+= 1 let updateToken = self.providerSwitcherUpdateToken - self.scheduleOpenMenuRebuildIfStillVisible( + let updatedMenu = self.updateOpenMenuForSwitcherSelectionImmediately( + menu, + provider: provider, + token: updateToken) + if updatedMenu { + self.closeHostedSubviewMenusForParentSwitch() + self.finishProviderSwitcherSettingsSuppression( + settingsSuppressionToken, + menuWasRebuilt: true) + return + } + + let updatedSwitcher = self.updateOpenMenuSwitcherSelectionIfPossible(menu) + if updatedSwitcher { + self.closeHostedSubviewMenusForParentSwitch() + let menuWasInvalidated = self.deferOpenParentMenuMutationIfTracking( + menu, + provider: provider, + reason: "provider-switch-selection") + self.finishProviderSwitcherSettingsSuppression( + settingsSuppressionToken, + menuWasRebuilt: false, + menuWasInvalidated: menuWasInvalidated) + return + } + + self.menuLogger.debug( + "provider switch parent rebuild deferred during menu tracking", + metadata: [ + "provider": provider?.rawValue ?? "overview", + "token": "\(updateToken)", + ]) + let menuWasInvalidated = self.deferOpenParentMenuMutationIfTracking( menu, provider: provider, - closeHostedSubviewMenusBeforeRebuild: true) - { [weak self] in - guard let self else { return false } - return self.providerSwitcherUpdateToken == updateToken + reason: "provider-switch") + if menuWasInvalidated { + self.closeHostedSubviewMenusForParentSwitch() } + self.finishProviderSwitcherSettingsSuppression( + settingsSuppressionToken, + menuWasRebuilt: false, + menuWasInvalidated: menuWasInvalidated) } - func scheduleOpenMenuRebuildIfStillVisible( + private func updateOpenMenuSwitcherSelectionIfPossible(_ menu: NSMenu) -> Bool { + let menuID = ObjectIdentifier(menu) + guard self.openMenus[menuID] != nil else { return false } + guard !self.isHostedSubviewMenu(menu) else { return false } + guard let switcherView = menu.items.first?.view as? ProviderSwitcherView else { return false } + guard let selection = self.lastMergedSwitcherSelection else { return false } + + switcherView.updateSelection(selection) + switcherView.updateQuotaIndicators() + self.applyIcon(phase: nil) + return true + } + + private func updateOpenMenuForSwitcherSelectionImmediately( _ menu: NSMenu, provider: UsageProvider?, - closeHostedSubviewMenusBeforeRebuild: Bool = false, - beforeRebuild: (@MainActor () -> Bool)? = nil) + token: Int) + -> Bool { - let key = ObjectIdentifier(menu) - if closeHostedSubviewMenusBeforeRebuild { - self.openMenuRebuildsClosingHostedSubviewMenus.insert(key) + let menuID = ObjectIdentifier(menu) + guard self.openMenus[menuID] != nil else { return false } + guard !self.isHostedSubviewMenu(menu) else { return false } + guard menu.items.first?.view is ProviderSwitcherView else { return false } + + let startedAt = Date() + let didUpdate = self.populateMenu( + menu, + provider: provider, + mode: .preserveExistingSwitcherWidth) + guard didUpdate else { return false } + + self.markMenuFresh(menu) + self.applyIcon(phase: nil) + let updateMs = Date().timeIntervalSince(startedAt) * 1000 + self.menuLogger.info( + "provider switch content updated in tracked menu", + metadata: [ + "provider": provider?.rawValue ?? "overview", + "token": "\(token)", + "updateMs": String(format: "%.1f", updateMs), + ]) + return true + } + + @discardableResult + func updateTrackedOpenParentMenuContentIfPossible( + _ menu: NSMenu, + provider: UsageProvider?, + reason: String) + -> Bool + { + let menuID = ObjectIdentifier(menu) + guard self.openMenus[menuID] != nil else { return false } + guard !self.isHostedSubviewMenu(menu) else { return false } + guard !self.hasOpenHostedSubviewMenu() else { return false } + + let startedAt = Date() + let didUpdate = self.populateMenu(menu, provider: provider, mode: .preserveExistingSwitcherWidth) + guard didUpdate else { + _ = self.deferOpenParentMenuMutationIfTracking( + menu, + provider: provider, + reason: "\(reason)-incompatible") + self.applyIcon(phase: nil) + return false } - let shouldCloseHostedSubviewMenus = self.openMenuRebuildsClosingHostedSubviewMenus.contains(key) - self.openMenuRebuildTokenCounter &+= 1 - let rebuildToken = self.openMenuRebuildTokenCounter - self.openMenuRebuildTokens[key] = rebuildToken - self.openMenuRebuildTasks[key]?.cancel() - self.openMenuRebuildTasks[key] = Task { @MainActor [weak self, weak menu] in - guard let self, let menu else { return } - #if DEBUG - if let override = self._test_openMenuRefreshYieldOverride { - await override() - } else { - await Task.yield() - } - #else - await Task.yield() - #endif - guard !Task.isCancelled else { return } - guard self.openMenuRebuildTokens[key] == rebuildToken else { return } - defer { - if self.openMenuRebuildTokens[key] == rebuildToken { - self.openMenuRebuildTasks.removeValue(forKey: key) - self.openMenuRebuildTokens.removeValue(forKey: key) - self.openMenuRebuildsClosingHostedSubviewMenus.remove(key) - } - } - guard self.openMenus[key] != nil else { return } - guard beforeRebuild?() ?? true else { return } - if shouldCloseHostedSubviewMenus { - self.closeHostedSubviewMenusForParentSwitch() - } - self.rebuildOpenMenuIfStillVisible(menu, provider: provider) + + self.markMenuFresh(menu) + self.applyIcon(phase: nil) + let updateMs = Date().timeIntervalSince(startedAt) * 1000 + self.menuLogger.info( + "tracked parent menu content updated", + metadata: [ + "provider": provider?.rawValue ?? "overview", + "reason": reason, + "updateMs": String(format: "%.1f", updateMs), + ]) + #if DEBUG + self._test_openMenuRebuildObserver?(menu) + #endif + return true + } + + @discardableResult + func deferOpenParentMenuMutationIfTracking( + _ menu: NSMenu, + provider: UsageProvider?, + reason: String) + -> Bool + { + let menuID = ObjectIdentifier(menu) + guard self.openMenus[menuID] != nil else { return false } + guard !self.isHostedSubviewMenu(menu) else { return false } + + self.openMenuRefreshTokens.removeValue(forKey: menuID) + if self.menuVersions[menuID] == self.menuContentVersion { + self.menuContentVersion &+= 1 } + self.menuLogger.debug( + "open parent menu mutation deferred until next open", + metadata: [ + "provider": provider?.rawValue ?? "overview", + "reason": reason, + ]) + return true } private func closeHostedSubviewMenusForParentSwitch() { let hostedMenus = self.openMenus.values.filter { self.isHostedSubviewMenu($0) } + guard !hostedMenus.isEmpty else { return } + self.closingHostedSubviewMenusForParentSwitch = true + defer { self.closingHostedSubviewMenusForParentSwitch = false } for hostedMenu in hostedMenus { hostedMenu.cancelTrackingWithoutAnimation() self.forgetClosedMenu(hostedMenu) diff --git a/Sources/CodexBar/StatusItemController+MenuSizing.swift b/Sources/CodexBar/StatusItemController+MenuSizing.swift new file mode 100644 index 0000000000..f3f0b5e7fb --- /dev/null +++ b/Sources/CodexBar/StatusItemController+MenuSizing.swift @@ -0,0 +1,64 @@ +import AppKit +import CodexBarCore + +extension StatusItemController { + func menuCardWidth( + for providers: [UsageProvider], + sections: [MenuDescriptor.Section]) -> CGFloat + { + _ = providers + let baselineWidth = Self.menuCardBaseWidth + return max(baselineWidth, self.measuredStandardMenuWidth(for: sections, baseWidth: baselineWidth)) + } + + func resolvedMenuWidth( + enabledProviders: [UsageProvider], + sections: [MenuDescriptor.Section], + menu: NSMenu, + mode: MenuPopulateMode) + -> CGFloat + { + if mode == .preserveExistingSwitcherWidth, + let switcherView = menu.items.first?.view as? ProviderSwitcherView, + switcherView.frame.width > 0 + { + return switcherView.frame.width + } + return self.menuCardWidth(for: enabledProviders, sections: sections) + } + + private func measuredStandardMenuWidth(for sections: [MenuDescriptor.Section], baseWidth: CGFloat) -> CGFloat { + let menuFont = NSFont.menuFont(ofSize: NSFont.systemFontSize) + let secondaryFont = NSFont.menuFont(ofSize: NSFont.smallSystemFontSize) + let persistentActionPadding: CGFloat = 88 + let standardActionPadding: CGFloat = 56 + let textPadding: CGFloat = 36 + var measuredWidth = baseWidth + + for section in sections { + for entry in section.entries { + switch entry { + case let .action(title, action): + let titleWidth = self.textWidth(L(title), font: menuFont) + let padding = self.usesPersistentMenuActionItem(for: action) + ? persistentActionPadding + : standardActionPadding + measuredWidth = max(measuredWidth, titleWidth + padding) + case let .submenu(title, _, _): + measuredWidth = max(measuredWidth, self.textWidth(title, font: menuFont) + standardActionPadding) + case let .text(text, style): + let font = style == .secondary ? secondaryFont : menuFont + measuredWidth = max(measuredWidth, min(self.textWidth(text, font: font) + textPadding, baseWidth)) + case .divider: + continue + } + } + } + + return ceil(measuredWidth) + } + + private func textWidth(_ text: String, font: NSFont) -> CGFloat { + ceil((text as NSString).size(withAttributes: [.font: font]).width) + } +} diff --git a/Sources/CodexBar/StatusItemController+MenuTracking.swift b/Sources/CodexBar/StatusItemController+MenuTracking.swift index d210bff624..a440605841 100644 --- a/Sources/CodexBar/StatusItemController+MenuTracking.swift +++ b/Sources/CodexBar/StatusItemController+MenuTracking.swift @@ -22,7 +22,16 @@ extension StatusItemController { self.refreshOpenMenusIfNeeded(allowsParentRebuild: true) } - private func refreshOpenMenusIfNeeded(allowsParentRebuild: Bool) { + func refreshOpenMenusForTokenCostHydration() { + guard Self.menuRefreshEnabled else { return } + guard !self.openMenus.isEmpty else { return } + self.refreshOpenMenusIfNeeded(allowsParentRebuild: true, allowsTrackedParentSmartUpdate: true) + } + + private func refreshOpenMenusIfNeeded( + allowsParentRebuild: Bool, + allowsTrackedParentSmartUpdate: Bool = false) + { var orphanedKeys: [ObjectIdentifier] = [] let hasOpenHostedSubviewMenu = self.hasOpenHostedSubviewMenu() for (key, menu) in self.openMenus { @@ -33,6 +42,7 @@ extension StatusItemController { self.refreshOpenMenuIfNeeded( menu, allowsParentRebuild: allowsParentRebuild, + allowsTrackedParentSmartUpdate: allowsTrackedParentSmartUpdate, hasOpenHostedSubviewMenu: hasOpenHostedSubviewMenu) } self.removeOrphanedOpenMenuEntries(orphanedKeys) @@ -41,6 +51,7 @@ extension StatusItemController { private func refreshOpenMenuIfNeeded( _ menu: NSMenu, allowsParentRebuild: Bool, + allowsTrackedParentSmartUpdate: Bool, hasOpenHostedSubviewMenu: Bool) { if self.isHostedSubviewMenu(menu) { @@ -52,7 +63,19 @@ extension StatusItemController { guard self.menuNeedsRefresh(menu) else { return } let provider = self.menuProvider(for: menu) - self.scheduleOpenMenuRebuildIfStillVisible(menu, provider: provider) + if allowsTrackedParentSmartUpdate { + self.updateTrackedOpenParentMenuContentIfPossible(menu, provider: provider, reason: "tracked-menu-refresh") + return + } + if self.deferOpenParentMenuMutationIfTracking( + menu, + provider: provider, + reason: "tracked-menu-refresh") + { + return + } + self.populateMenu(menu, provider: provider) + self.markMenuFresh(menu) } private func removeOrphanedOpenMenuEntries(_ keys: [ObjectIdentifier]) { diff --git a/Sources/CodexBar/StatusItemController+Shutdown.swift b/Sources/CodexBar/StatusItemController+Shutdown.swift index 7dcfdcff8d..b621fe6ad6 100644 --- a/Sources/CodexBar/StatusItemController+Shutdown.swift +++ b/Sources/CodexBar/StatusItemController+Shutdown.swift @@ -28,6 +28,12 @@ extension StatusItemController { self.loginTask = nil self.screenChangeVisibilityTask?.cancel() self.screenChangeVisibilityTask = nil + self.mainRunLoopStallMonitor?.stop() + self.mainRunLoopStallMonitor = nil + self.codexAccountSelectionRefreshTask?.cancel() + self.codexAccountSelectionRefreshTask = nil + self.tokenCostMenuHydrationRefreshTask?.cancel() + self.tokenCostMenuHydrationRefreshTask = nil self.pendingScreenChangePreviousCount = nil self.animationDriver?.stop() self.animationDriver = nil diff --git a/Sources/CodexBar/StatusItemController+StallMonitoring.swift b/Sources/CodexBar/StatusItemController+StallMonitoring.swift new file mode 100644 index 0000000000..850c428ece --- /dev/null +++ b/Sources/CodexBar/StatusItemController+StallMonitoring.swift @@ -0,0 +1,54 @@ +import CodexBarCore +import Foundation +import ObjectiveC + +private enum MainRunLoopStallMonitorAssociation { + nonisolated(unsafe) static var key: UInt8 = 0 +} + +extension StatusItemController { + var mainRunLoopStallMonitor: MainRunLoopStallMonitor? { + get { + objc_getAssociatedObject( + self, + &MainRunLoopStallMonitorAssociation.key) as? MainRunLoopStallMonitor + } + set { + objc_setAssociatedObject( + self, + &MainRunLoopStallMonitorAssociation.key, + newValue, + .OBJC_ASSOCIATION_RETAIN_NONATOMIC) + } + } + + func startMainRunLoopStallMonitorIfNeeded() { + guard !SettingsStore.isRunningTests else { return } + guard self.mainRunLoopStallMonitor == nil else { return } + + let monitor = MainRunLoopStallMonitor(logger: self.menuLogger) { [weak self] in + guard let self else { return [:] } + + var metadata: [String: String] = [ + "menuContentVersion": "\(self.menuContentVersion)", + "openMenus": "\(self.openMenus.count)", + "providerSwitcherToken": "\(self.providerSwitcherUpdateToken)", + "selectedProvider": self.selectedMenuProvider?.rawValue ?? "nil", + "storeRefreshing": self.store.isRefreshing ? "1" : "0", + ] + if let lastMenuProvider = self.lastMenuProvider { + metadata["lastMenuProvider"] = lastMenuProvider.rawValue + } + if let selection = self.lastMergedSwitcherSelection { + metadata["switcherSelection"] = selection.logValue + } + if let lastInteraction = self.lastProviderSwitcherInteractionAt { + let ageMs = Date().timeIntervalSince(lastInteraction) * 1000 + metadata["lastSwitcherInteractionAgeMs"] = String(format: "%.1f", ageMs) + } + return metadata + } + self.mainRunLoopStallMonitor = monitor + monitor.start() + } +} diff --git a/Sources/CodexBar/StatusItemController+SwitcherMenuItems.swift b/Sources/CodexBar/StatusItemController+SwitcherMenuItems.swift new file mode 100644 index 0000000000..ebe663511e --- /dev/null +++ b/Sources/CodexBar/StatusItemController+SwitcherMenuItems.swift @@ -0,0 +1,143 @@ +import AppKit +import CodexBarCore + +extension StatusItemController { + func makeProviderSwitcherItem( + providers: [UsageProvider], + includesOverview: Bool, + selected: ProviderSwitcherSelection, + menu: NSMenu, + width: CGFloat) -> NSMenuItem + { + let view = ProviderSwitcherView( + providers: providers, + selected: selected, + includesOverview: includesOverview, + width: width, + showsIcons: self.settings.switcherShowsIcons, + iconProvider: { [weak self] provider in + self?.switcherIcon(for: provider) ?? NSImage() + }, + weeklyRemainingProvider: { [weak self] provider in + self?.switcherWeeklyRemaining(for: provider) + }, + onSelect: { [weak self, weak menu] selection in + guard let self, let menu else { return } + self.noteProviderSwitcherInteraction() + if selection == self.lastMergedSwitcherSelection { + self.menuLogger.debug( + "provider switch selection ignored", + metadata: ["selection": selection.logValue]) + return + } + let settingsSuppressionToken = self.beginProviderSwitcherSettingsSuppression() + let provider: UsageProvider? + switch selection { + case .overview: + self.settings.mergedMenuLastSelectedWasOverview = true + provider = self.resolvedMenuProvider() + case let .provider(selectedProvider): + self.settings.mergedMenuLastSelectedWasOverview = false + self.selectedMenuProvider = selectedProvider + provider = selectedProvider + } + switch selection { + case .overview: + self.lastMenuProvider = provider ?? .codex + case let .provider(provider): + self.lastMenuProvider = provider + } + self.lastMergedSwitcherSelection = selection + self.deferSwitcherMenuRebuildIfStillVisible( + menu, + provider: provider, + settingsSuppressionToken: settingsSuppressionToken) + }) + let item = NSMenuItem() + item.view = view + item.isEnabled = false + return item + } + + func makeTokenAccountSwitcherItem( + display: TokenAccountMenuDisplay, + menu: NSMenu, + width: CGFloat) -> NSMenuItem + { + let view = TokenAccountSwitcherView( + accounts: display.accounts, + selectedIndex: display.activeIndex, + width: width, + onSelect: { [weak self, weak menu] index -> Task? in + guard let self, let menu else { return nil } + self.noteProviderSwitcherInteraction() + self.settings.setActiveTokenAccountIndex(index, for: display.provider) + self.applyIcon(phase: nil) + self.deferSwitcherMenuRebuildIfStillVisible(menu, provider: display.provider) + return Task { @MainActor [weak self, weak menu] in + guard let self else { return } + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await self.store.refreshProvider(display.provider) + } + guard let menu else { return } + self.refreshOpenMenuIfStillVisible(menu, provider: display.provider) + } + }) + let item = NSMenuItem() + item.view = view + item.isEnabled = false + return item + } + + func makeCodexAccountSwitcherItem( + display: CodexAccountMenuDisplay, + menu: NSMenu, + width: CGFloat) -> NSMenuItem + { + let view = CodexAccountSwitcherView( + accounts: display.accounts, + selectedAccountID: display.activeVisibleAccountID, + width: width, + onSelect: { [weak self, weak menu] account in + guard let self else { return } + self.handleCodexVisibleAccountSelection(account, menu: menu) + }) + let item = NSMenuItem() + item.view = view + item.isEnabled = false + return item + } + + @discardableResult + private func handleCodexVisibleAccountSelection(_ account: CodexVisibleAccount, menu: NSMenu?) -> Bool { + let visibleAccountID = account.id + self.noteProviderSwitcherInteraction() + self.codexAccountSelectionRefreshGeneration &+= 1 + let generation = self.codexAccountSelectionRefreshGeneration + self.codexAccountSelectionRefreshTask?.cancel() + self.settings.selectDisplayedCodexVisibleAccount(account) + if self.store.prepareCodexAccountScopedRefreshIfNeeded(), let menu { + self.deferSwitcherMenuRebuildIfStillVisible(menu, provider: .codex) + } + self.codexAccountSelectionRefreshTask = Task { @MainActor [weak self, weak menu] in + guard let self else { return } + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await self.store.refreshCodexAccountScopedState( + allowDisabled: true, + phaseDidChange: { [weak self, weak menu] _ in + guard let self, let menu else { return } + guard self.codexAccountSelectionRefreshGeneration == generation else { return } + guard self.settings.codexVisibleAccountProjection.activeVisibleAccountID == visibleAccountID + else { + return + } + self.refreshOpenMenuIfStillVisible(menu, provider: .codex) + }) + } + if self.codexAccountSelectionRefreshGeneration == generation { + self.codexAccountSelectionRefreshTask = nil + } + } + return true + } +} diff --git a/Sources/CodexBar/StatusItemController+SwitcherViews.swift b/Sources/CodexBar/StatusItemController+SwitcherViews.swift index a7fe4c684e..48cb27d062 100644 --- a/Sources/CodexBar/StatusItemController+SwitcherViews.swift +++ b/Sources/CodexBar/StatusItemController+SwitcherViews.swift @@ -5,6 +5,15 @@ import QuartzCore enum ProviderSwitcherSelection: Equatable { case overview case provider(UsageProvider) + + var logValue: String { + switch self { + case .overview: + "overview" + case let .provider(provider): + provider.rawValue + } + } } final class ProviderSwitcherView: NSView { @@ -25,6 +34,7 @@ final class ProviderSwitcherView: NSView { private let onSelect: (ProviderSwitcherSelection) -> Void private let showsIcons: Bool private let weeklyRemainingProvider: (UsageProvider) -> Double? + private let reservesQuotaIndicatorSpace: Bool private var buttons: [NSButton] = [] private var quotaIndicators: [ObjectIdentifier: QuotaIndicator] = [:] private var hoverTrackingArea: NSTrackingArea? @@ -60,32 +70,17 @@ final class ProviderSwitcherView: NSView { onSelect: @escaping (ProviderSwitcherSelection) -> Void) { let minimumGap: CGFloat = 1 - var segments = providers.map { provider in - let fullTitle = Self.switcherTitle(for: provider) - let icon = iconProvider(provider) - icon.isTemplate = true - // Avoid any resampling: we ship exact 16pt/32px assets for crisp rendering. - icon.size = NSSize(width: 16, height: 16) - return Segment( - selection: .provider(provider), - image: icon, - title: fullTitle) - } - if includesOverview { - let overviewIcon = Self.overviewIcon() - overviewIcon.isTemplate = true - overviewIcon.size = NSSize(width: 16, height: 16) - segments.insert( - Segment( - selection: .overview, - image: overviewIcon, - title: L("Overview")), - at: 0) - } + let segments = Self.makeSegments( + providers: providers, + includesOverview: includesOverview, + iconProvider: iconProvider) self.segments = segments self.onSelect = onSelect self.showsIcons = showsIcons self.weeklyRemainingProvider = weeklyRemainingProvider + self.reservesQuotaIndicatorSpace = Self.reservesQuotaIndicatorSpace( + segments: segments, + weeklyRemainingProvider: weeklyRemainingProvider) self.stackedIcons = showsIcons && self.segments.count > 3 let initialOuterPadding = Self.switcherOuterPadding( for: width, @@ -169,6 +164,9 @@ final class ProviderSwitcherView: NSView { case .overview: nil } + if remaining == nil, self.reservesQuotaIndicatorSpace { + Self.applyQuotaBarContentInset(to: button) + } self.addQuotaIndicator(to: button, selection: segment.selection, remainingPercent: remaining) button.bezelStyle = .regularSquare button.isBordered = false @@ -182,6 +180,7 @@ final class ProviderSwitcherView: NSView { button.state = (selected == segment.selection) ? .on : .off button.toolTip = nil button.translatesAutoresizingMaskIntoConstraints = false + button.heightAnchor.constraint(equalToConstant: self.rowHeight).isActive = true self.buttons.append(button) return button } @@ -622,9 +621,13 @@ final class ProviderSwitcherView: NSView { self.addQuotaIndicator(to: button, selection: segment.selection, remainingPercent: remaining) } } else if let indicator = self.quotaIndicators.removeValue(forKey: key) { - Self.applyQuotaBarContentInset(to: button, height: 0) + let reservedHeight = self.reservesQuotaIndicatorSpace + ? Self.quotaIndicatorReservedHeight + : 0 + Self.applyQuotaBarContentInset(to: button, height: reservedHeight) indicator.track.removeFromSuperview() - continue + } else if self.reservesQuotaIndicatorSpace { + Self.applyQuotaBarContentInset(to: button) } self.updateQuotaIndicatorVisibility(for: button) } @@ -667,6 +670,10 @@ final class ProviderSwitcherView: NSView { self.buttons.map(\.fittingSize) } + func _test_buttonIntrinsicSizes() -> [NSSize] { + self.buttons.map(\.intrinsicContentSize) + } + func _test_rowCount() -> Int { self.rowCount } @@ -897,6 +904,50 @@ final class ProviderSwitcherView: NSView { } extension ProviderSwitcherView { + private static func reservesQuotaIndicatorSpace( + segments: [Segment], + weeklyRemainingProvider: (UsageProvider) -> Double?) -> Bool + { + segments.contains { segment in + switch segment.selection { + case let .provider(provider): + weeklyRemainingProvider(provider) != nil + case .overview: + false + } + } + } + + private static func makeSegments( + providers: [UsageProvider], + includesOverview: Bool, + iconProvider: (UsageProvider) -> NSImage) -> [Segment] + { + var segments = providers.map { provider in + let fullTitle = self.switcherTitle(for: provider) + let icon = iconProvider(provider) + icon.isTemplate = true + // Avoid any resampling: we ship exact 16pt/32px assets for crisp rendering. + icon.size = NSSize(width: 16, height: 16) + return Segment( + selection: .provider(provider), + image: icon, + title: fullTitle) + } + if includesOverview { + let overviewIcon = self.overviewIcon() + overviewIcon.isTemplate = true + overviewIcon.size = NSSize(width: 16, height: 16) + segments.insert( + Segment( + selection: .overview, + image: overviewIcon, + title: L("Overview")), + at: 0) + } + return segments + } + private func addQuotaIndicator(to view: NSView, selection: ProviderSwitcherSelection, remainingPercent: Double?) { guard let remainingPercent else { return } Self.applyQuotaBarContentInset(to: view) diff --git a/Sources/CodexBar/StatusItemController+TokenCostDeferral.swift b/Sources/CodexBar/StatusItemController+TokenCostDeferral.swift new file mode 100644 index 0000000000..3cb6435978 --- /dev/null +++ b/Sources/CodexBar/StatusItemController+TokenCostDeferral.swift @@ -0,0 +1,16 @@ +import AppKit + +extension StatusItemController { + func beginMenuTokenCostDeferral(_ menu: NSMenu, reason: String) { + let key = ObjectIdentifier(menu) + guard !self.tokenCostDeferredMenuIDs.contains(key) else { return } + self.tokenCostDeferredMenuIDs.insert(key) + self.store.beginInteractiveMenuTokenCostDeferral(reason: reason) + } + + func endMenuTokenCostDeferral(_ menu: NSMenu, reason: String) { + let key = ObjectIdentifier(menu) + guard self.tokenCostDeferredMenuIDs.remove(key) != nil else { return } + self.store.endInteractiveMenuTokenCostDeferral(reason: reason) + } +} diff --git a/Sources/CodexBar/StatusItemController+TokenCostHydration.swift b/Sources/CodexBar/StatusItemController+TokenCostHydration.swift new file mode 100644 index 0000000000..12eb737362 --- /dev/null +++ b/Sources/CodexBar/StatusItemController+TokenCostHydration.swift @@ -0,0 +1,33 @@ +import Observation + +extension StatusItemController { + func observeTokenCostMenuHydrationChanges() { + withObservationTracking { + _ = self.store.tokenSnapshots + _ = self.store.tokenErrors + } onChange: { [weak self] in + Task { @MainActor [weak self] in + guard let self else { return } + self.observeTokenCostMenuHydrationChanges() + self.scheduleOpenMenuTokenCostHydrationRefreshIfNeeded() + } + } + } + + private func scheduleOpenMenuTokenCostHydrationRefreshIfNeeded() { + guard Self.menuRefreshEnabled else { return } + guard !self.openMenus.isEmpty else { return } + self.tokenCostMenuHydrationRefreshTask?.cancel() + self.tokenCostMenuHydrationRefreshTask = Task { @MainActor [weak self] in + try? await Task.sleep(for: .milliseconds(180)) + guard let self, !Task.isCancelled else { return } + self.tokenCostMenuHydrationRefreshTask = nil + guard !self.openMenus.isEmpty else { return } + self.menuContentVersion &+= 1 + self.menuLogger.debug( + "token cost menu hydration refresh", + metadata: ["openMenus": "\(self.openMenus.count)"]) + self.refreshOpenMenusForTokenCostHydration() + } + } +} diff --git a/Sources/CodexBar/StatusItemController+UsageHistoryMenu.swift b/Sources/CodexBar/StatusItemController+UsageHistoryMenu.swift index 73233cb38d..9070bf38ab 100644 --- a/Sources/CodexBar/StatusItemController+UsageHistoryMenu.swift +++ b/Sources/CodexBar/StatusItemController+UsageHistoryMenu.swift @@ -48,7 +48,9 @@ extension StatusItemController { provider: UsageProvider, width: CGFloat) -> Bool { - let histories = self.store.planUtilizationHistory(for: provider) + let startedAt = Date() + let histories = self.store.planUtilizationHistoryForMenu(for: provider) + let historyMs = Date().timeIntervalSince(startedAt) * 1000 let snapshot = self.store.snapshot(for: provider) if !Self.menuCardRenderingEnabled { @@ -64,9 +66,13 @@ extension StatusItemController { histories: histories, snapshot: snapshot, width: width) + let hostingStartedAt = Date() let hosting = UsageHistoryMenuHostingView(rootView: chartView) - let controller = NSHostingController(rootView: chartView) - let size = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: 1)) + let hostingMs = Date().timeIntervalSince(hostingStartedAt) * 1000 + let sizeStartedAt = Date() + let size = hosting.fittingSize + let sizeMs = Date().timeIntervalSince(sizeStartedAt) * 1000 hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) let chartItem = NSMenuItem() @@ -74,6 +80,20 @@ extension StatusItemController { chartItem.isEnabled = true chartItem.representedObject = Self.usageHistoryChartID submenu.addItem(chartItem) + let totalMs = Date().timeIntervalSince(startedAt) * 1000 + if totalMs >= 16 { + self.menuLogger.info( + "usage history submenu chart built", + metadata: [ + "entries": "\(histories.reduce(0) { $0 + $1.entries.count })", + "historyMs": String(format: "%.1f", historyMs), + "hostingMs": String(format: "%.1f", hostingMs), + "provider": provider.rawValue, + "sizeMs": String(format: "%.1f", sizeMs), + "totalMs": String(format: "%.1f", totalMs), + "width": String(format: "%.0f", width), + ]) + } return true } } diff --git a/Sources/CodexBar/StatusItemController.swift b/Sources/CodexBar/StatusItemController.swift index 09499730dc..7fa56d16b9 100644 --- a/Sources/CodexBar/StatusItemController.swift +++ b/Sources/CodexBar/StatusItemController.swift @@ -21,6 +21,20 @@ extension StatusItemControlling { func prepareForAppShutdown() {} } +private enum StatusItemIdentity { + case merged + case provider(UsageProvider) + + var accessibilityIdentifier: String { + switch self { + case .merged: + StatusItemController.statusItemAccessibilityIdentifierPrefix + case let .provider(provider): + "\(StatusItemController.statusItemAccessibilityIdentifierPrefix).\(provider.rawValue)" + } + } +} + @MainActor final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControlling { // Disable SwiftUI menu cards + menu refresh work in tests to avoid swiftpm-testing-helper crashes. @@ -29,21 +43,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin private(set) static var menuRefreshEnabled = !SettingsStore.isRunningTests static let quotaWarningFlashDuration: TimeInterval = 60 private nonisolated static let statusItemAccessibilityTitle = "CodexBar" - private nonisolated static let statusItemAccessibilityIdentifierPrefix = "CodexBar.StatusItem" - - private enum StatusItemIdentity { - case merged - case provider(UsageProvider) - - var accessibilityIdentifier: String { - switch self { - case .merged: - StatusItemController.statusItemAccessibilityIdentifierPrefix - case let .provider(provider): - "\(StatusItemController.statusItemAccessibilityIdentifierPrefix).\(provider.rawValue)" - } - } - } + fileprivate nonisolated static let statusItemAccessibilityIdentifierPrefix = "CodexBar.StatusItem" #if DEBUG static func setMenuRefreshEnabledForTesting(_ enabled: Bool) { @@ -114,6 +114,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin var highlightedMenuItems: [ObjectIdentifier: NSMenuItem] = [:] var providerSwitcherShortcutEventMonitor: ProviderSwitcherShortcutEventMonitor? var providerSwitcherShortcutMenuID: ObjectIdentifier? + var tokenCostDeferredMenuIDs: Set = [] var hasPreparedForAppShutdown = false #if DEBUG var onDelayedMenuRefreshAttemptForTesting: (() -> Void)? @@ -179,6 +180,26 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin var lastTokenAccountMenuDisplay: TokenAccountMenuDisplay? /// Monotonic token used to ignore stale deferred provider-switcher menu rebuilds. var providerSwitcherUpdateToken = 0 + var providerSwitcherSettingsSuppressionGeneration = 0 + var activeProviderSwitcherSettingsSuppressionGeneration: Int? + var providerSwitcherDeferredIconUpdatePending = false + var closingHostedSubviewMenusForParentSwitch = false + var providerSwitcherSettingsObservationSuppressionsRemaining = 0 + var lastProviderSwitcherInteractionAt: Date? + var codexAccountSelectionRefreshGeneration = 0 + var codexAccountSelectionRefreshTask: Task? + var openMenuRefreshGeneration = 0 + var openMenuRefreshTokens: [ObjectIdentifier: Int] = [:] + var settingsChangeDeferredDuringMenuTracking = false + var deferredSettingsChangeNeedsStatusItemRebuild = false + var deferredSettingsChangeNeedsOpenMenuRefresh = false + var deferredMenuInteractionRefreshPending = false + var deferredMenuInteractionRefreshGeneration = 0 + var deferredMenuInteractionRefreshTask: Task? + var tokenCostMenuHydrationRefreshTask: Task? + var storeChangeDeferredDuringMenuTracking = false + var storeChangeDeferredMenuContentVersion: Int? + var storeChangeDeferredObservationCount = 0 var lastAppliedMergedIconRenderSignature: String? var lastAppliedProviderIconRenderSignatures: [UsageProvider: String] = [:] var lastObservedStoreIconWorkSignature: String? @@ -355,6 +376,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin selector: #selector(self.handleScreenParametersDidChange(_:)), name: NSApplication.didChangeScreenParametersNotification, object: nil) + self.startResponsiveMenuSupportIfNeeded() } convenience init( @@ -380,6 +402,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin private func wireBindings() { self.observeStoreChanges() + self.observeTokenCostMenuHydrationChanges() self.observeStoreIconChanges() self.observeIconPerfRefreshCycleChanges() self.observeDebugForceAnimation() @@ -388,12 +411,24 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin self.observeManagedCodexCoordinatorChanges() } - private func observeStoreChanges() { + func observeStoreChanges() { withObservationTracking { _ = self.store.menuObservationToken } onChange: { [weak self] in Task { @MainActor [weak self] in guard let self else { return } + if !self.openMenus.isEmpty { + if !self.storeChangeDeferredDuringMenuTracking { + self.storeChangeDeferredMenuContentVersion = self.menuContentVersion + self.storeChangeDeferredObservationCount = 0 + self.menuLogger.debug( + "store menu observations deferred during menu tracking", + metadata: ["openMenus": "\(self.openMenus.count)"]) + } + self.storeChangeDeferredDuringMenuTracking = true + self.storeChangeDeferredObservationCount += 1 + return + } self.observeStoreChanges() self.invalidateMenus() } @@ -624,6 +659,34 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin let configChanged = self.settings.configRevision != self.lastConfigRevision let orderChanged = self.settings.providerOrder != self.lastProviderOrder let shouldRefreshOpenMenus = self.shouldRefreshOpenMenusForProviderSwitcher() + let hasPendingProviderSwitcherObservation = reason == "observation" && + self.providerSwitcherSettingsObservationSuppressionsRemaining > 0 + let shouldSuppressProviderSwitcherObservation = + self.activeProviderSwitcherSettingsSuppressionGeneration != nil || + hasPendingProviderSwitcherObservation + if shouldSuppressProviderSwitcherObservation, + !configChanged, + !orderChanged, + !shouldRefreshOpenMenus + { + if self.activeProviderSwitcherSettingsSuppressionGeneration != nil { + self.providerSwitcherDeferredIconUpdatePending = true + } else if hasPendingProviderSwitcherObservation { + self.providerSwitcherSettingsObservationSuppressionsRemaining -= 1 + } + self.menuLogger.debug("provider switch settings observation suppressed") + return + } + if hasPendingProviderSwitcherObservation { + self.providerSwitcherSettingsObservationSuppressionsRemaining = 0 + } + if !self.openMenus.isEmpty, reason.hasPrefix("config:notification:provider-") { + self.deferSettingsChangeDuringMenuTracking( + reason: reason, + needsStatusItemRebuild: orderChanged || configChanged, + needsOpenMenuRefresh: shouldRefreshOpenMenus) + return + } self.invalidateMenus() if orderChanged || configChanged { self.rebuildProviderStatusItems() @@ -635,7 +698,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin } } - private func updateIcons() { + func updateIcons() { #if DEBUG guard !self.isReleasedForTesting else { return } #endif @@ -698,7 +761,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin self.updateIcons() } - private func updateVisibility() { + func updateVisibility() { #if DEBUG guard !self.isReleasedForTesting else { return } #endif @@ -792,7 +855,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin } } - private func rebuildProviderStatusItems() { + func rebuildProviderStatusItems() { #if DEBUG guard !self.isReleasedForTesting else { return } #endif @@ -838,19 +901,6 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin self.settings.mergeIcons && self.store.enabledProvidersForDisplay().count > 1 } - func switchAccountSubtitle(for target: UsageProvider) -> String? { - guard self.loginTask != nil, let provider = self.activeLoginProvider, provider == target - else { return nil } - let base: String - switch self.loginPhase { - case .idle: return nil - case .requesting: base = "Requesting login…" - case .waitingBrowser: base = "Waiting in browser…" - } - let prefix = ProviderDescriptorRegistry.descriptor(for: provider).metadata.displayName - return "\(prefix): \(base)" - } - deinit { let animationDriver = self.animationDriver Task { @MainActor in @@ -865,6 +915,19 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin } extension StatusItemController { + func switchAccountSubtitle(for target: UsageProvider) -> String? { + guard self.loginTask != nil, let provider = self.activeLoginProvider, provider == target + else { return nil } + let base: String + switch self.loginPhase { + case .idle: return nil + case .requesting: base = "Requesting login…" + case .waitingBrowser: base = "Waiting in browser…" + } + let prefix = ProviderDescriptorRegistry.descriptor(for: provider).metadata.displayName + return "\(prefix): \(base)" + } + func refreshExistingStatusItemsForVisibilityRecovery() { #if DEBUG guard !self.isReleasedForTesting else { return } diff --git a/Sources/CodexBar/UsageStore+PlanUtilization.swift b/Sources/CodexBar/UsageStore+PlanUtilization.swift index 82b81708b5..e6e61395d0 100644 --- a/Sources/CodexBar/UsageStore+PlanUtilization.swift +++ b/Sources/CodexBar/UsageStore+PlanUtilization.swift @@ -53,6 +53,16 @@ extension UsageStore { return providerBuckets.histories(for: accountKey) } + func planUtilizationHistoryForMenu(for provider: UsageProvider) -> [PlanUtilizationSeriesHistory] { + let providerBuckets = self.planUtilizationHistory[provider] ?? PlanUtilizationHistoryBuckets() + let snapshot = self.snapshots[provider] + let accountKey = self.planUtilizationMenuAccountKey( + provider: provider, + snapshot: snapshot, + providerBuckets: providerBuckets) + return providerBuckets.histories(for: accountKey) + } + func shouldShowRefreshingMenuCard(for provider: UsageProvider) -> Bool { let isRefreshing = self.isRefreshing || self.refreshingProviders.contains(provider) return isRefreshing @@ -495,6 +505,79 @@ extension UsageStore { return resolvedSnapshot.flatMap { Self.planUtilizationIdentityAccountKey(provider: provider, snapshot: $0) } } + private func planUtilizationMenuAccountKey( + provider: UsageProvider, + snapshot: UsageSnapshot?, + providerBuckets: PlanUtilizationHistoryBuckets) -> String? + { + if provider == .codex { + return self.codexPlanUtilizationMenuAccountKey( + snapshot: snapshot, + providerBuckets: providerBuckets) + } + + if let tokenAccountKey = Self.planUtilizationAccountKey( + provider: provider, + account: self.settings.selectedTokenAccount(for: provider)), + providerBuckets.histories(for: tokenAccountKey).isEmpty == false + { + return tokenAccountKey + } + + if let snapshot, + let identityAccountKey = Self.planUtilizationIdentityAccountKey(provider: provider, snapshot: snapshot), + providerBuckets.histories(for: identityAccountKey).isEmpty == false + { + return identityAccountKey + } + + if provider == .claude, + let snapshot, + let legacyAccountKey = Self.legacyClaudePlanUtilizationEmailAccountKey(snapshot: snapshot), + providerBuckets.histories(for: legacyAccountKey).isEmpty == false + { + return legacyAccountKey + } + + return self.stickyPlanUtilizationAccountKey(providerBuckets: providerBuckets) + } + + private func codexPlanUtilizationMenuAccountKey( + snapshot: UsageSnapshot?, + providerBuckets: PlanUtilizationHistoryBuckets) -> String? + { + let ownership = self.codexOwnershipContext(snapshot: snapshot, includeDashboardFallback: true) + let preferredKeys = [ + ownership.canonicalKey, + ownership.canonicalEmailHashKey, + ownership.planUtilizationLegacyEmailHash, + ].compactMap(\.self) + if let preferredKey = preferredKeys.first(where: { providerBuckets.histories(for: $0).isEmpty == false }) { + return preferredKey + } + + if let canonicalKey = ownership.canonicalKey { + let matchingKeys = providerBuckets.accounts.keys.filter { rawKey in + let owner = CodexHistoryOwnership.classifyPersistedKey( + rawKey, + legacyEmailHash: ownership.planUtilizationLegacyEmailHash) + return CodexHistoryOwnership.belongsToTargetContinuity( + owner, + targetCanonicalKey: canonicalKey, + canonicalEmailHashKey: ownership.canonicalEmailHashKey) + } + if let latestMatchingKey = matchingKeys.max(by: { lhs, rhs in + let lhsDate = providerBuckets.accounts[lhs]?.compactMap(\.latestCapturedAt).max() ?? .distantPast + let rhsDate = providerBuckets.accounts[rhs]?.compactMap(\.latestCapturedAt).max() ?? .distantPast + return lhsDate < rhsDate + }) { + return latestMatchingKey + } + } + + return self.stickyPlanUtilizationAccountKey(providerBuckets: providerBuckets) + } + private nonisolated static func planUtilizationAccountKey( provider: UsageProvider, account: ProviderTokenAccount?) -> String? diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index bf73bb84e6..1400320a0f 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -36,8 +36,6 @@ extension UsageStore { if provider == .codex, self.shouldFetchAllCodexVisibleAccounts() { await self.refreshCodexVisibleAccountsForMenu() return - } else if provider == .codex { - self.codexAccountSnapshots = [] } if provider == .kilo, self.shouldFanOutKiloScopes() { @@ -124,6 +122,7 @@ extension UsageStore { if provider == .codex { self.rememberLiveSystemCodexEmailIfNeeded(scoped.accountEmail(for: .codex)) self.seedCodexAccountScopedRefreshGuard(accountEmail: scoped.accountEmail(for: .codex)) + self.rememberActiveCodexVisibleAccountSnapshot(backfilled, sourceLabel: result.sourceLabel) } return backfilled } @@ -290,6 +289,13 @@ extension UsageStore { } private func handleProviderFetchFailure(provider: UsageProvider, error: Error) async { + if Self.errorIsCancellation(error) { + await MainActor.run { + self.errors[provider] = nil + } + return + } + let shouldNotifyPermissionPrompt = Self.isPermissionPromptWaiting(error) await MainActor.run { let hadPriorData = self.snapshots[provider] != nil @@ -330,8 +336,8 @@ extension UsageStore { } private static func shouldPreservePriorSnapshot(after error: Error, hadPriorData: Bool) -> Bool { + if self.errorIsCancellation(error) { return true } guard hadPriorData else { return false } - if error is CancellationError { return true } if self.isPreservableNetworkTransportError(error) { return true } let message = error.localizedDescription.lowercased() diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 4a5a96d4e7..ec275db98e 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -215,7 +215,7 @@ extension UsageStore { return false } - private static func errorIsCancellation(_ error: any Error) -> Bool { + nonisolated static func errorIsCancellation(_ error: any Error) -> Bool { if error is CancellationError { return true } @@ -607,6 +607,37 @@ extension UsageStore { } } + func rememberActiveCodexVisibleAccountSnapshot(_ snapshot: UsageSnapshot, sourceLabel: String?) { + let projection = self.settings.codexVisibleAccountProjection + guard let activeID = projection.activeVisibleAccountID, + let account = projection.visibleAccounts.first(where: { $0.id == activeID }) + else { + return + } + + let labeled = self.applyCodexVisibleAccountLabel(snapshot.scoped(to: .codex), account: account) + let accountSnapshot = CodexAccountUsageSnapshot( + account: account, + snapshot: labeled, + error: nil, + sourceLabel: sourceLabel) + self.upsertCodexAccountUsageSnapshot(accountSnapshot) + } + + func cachedCodexVisibleAccountSnapshotForActiveSelection() -> CodexAccountUsageSnapshot? { + guard let activeID = self.settings.codexVisibleAccountProjection.activeVisibleAccountID else { return nil } + return self.codexAccountSnapshots.first { $0.id == activeID } + } + + private func upsertCodexAccountUsageSnapshot(_ snapshot: CodexAccountUsageSnapshot) { + if let index = self.codexAccountSnapshots.firstIndex(where: { $0.id == snapshot.id }) { + self.codexAccountSnapshots[index] = snapshot + } else { + self.codexAccountSnapshots.append(snapshot) + } + self.codexAccountUsageSnapshotStore?.store(self.codexAccountSnapshots) + } + func applySelectedOutcome( _ outcome: ProviderFetchOutcome, provider: UsageProvider, diff --git a/Sources/CodexBar/UsageStore+TokenCost.swift b/Sources/CodexBar/UsageStore+TokenCost.swift index 13c162ed9e..327c1df489 100644 --- a/Sources/CodexBar/UsageStore+TokenCost.swift +++ b/Sources/CodexBar/UsageStore+TokenCost.swift @@ -18,6 +18,10 @@ extension UsageStore { self.tokenRefreshInFlight.contains(provider) } + func isTokenRefreshQueued(for provider: UsageProvider) -> Bool { + self.tokenRefreshQueuedProviders.contains(provider) + } + func tokenCostScope(for provider: UsageProvider) -> (codexHomePath: String?, signature: String) { guard provider == .codex else { return (nil, provider.rawValue) @@ -66,4 +70,93 @@ extension UsageStore { nonisolated static func tokenCostNoDataMessage(for provider: UsageProvider) -> String { ProviderDescriptorRegistry.descriptor(for: provider).tokenCost.noDataMessage() } + + func ensureTokenCostSnapshotScheduled(for provider: UsageProvider, reason: String) { + guard ProviderDescriptorRegistry.descriptor(for: provider).tokenCost.supportsTokenCost else { return } + guard self.settings.isCostUsageEffectivelyEnabled(for: provider) else { return } + guard !Self.tokenCostRequiresProviderSnapshot(provider) else { return } + guard self.isEnabled(provider) else { return } + guard self.tokenSnapshots[provider] == nil else { return } + guard self.tokenErrors[provider] == nil else { return } + guard !self.tokenRefreshQueuedProviders.contains(provider) else { return } + guard !self.tokenRefreshInFlight.contains(provider) else { return } + + if !self.hydrateCachedTokenCostSnapshotIfNeeded(for: provider, reason: reason) { + self.scheduleTokenRefresh( + force: false, + providers: [provider], + allowDuringMenuInteraction: true, + reason: reason) + } + } + + func shouldRunTokenCostRefreshDuringMenuInteraction(_ provider: UsageProvider) -> Bool { + guard ProviderDescriptorRegistry.descriptor(for: provider).tokenCost.supportsTokenCost else { return false } + guard self.settings.isCostUsageEffectivelyEnabled(for: provider) else { return false } + guard !Self.tokenCostRequiresProviderSnapshot(provider) else { return false } + guard self.isEnabled(provider) else { return false } + return self.tokenSnapshots[provider] == nil && self.tokenErrors[provider] == nil + } + + @discardableResult + private func hydrateCachedTokenCostSnapshotIfNeeded(for provider: UsageProvider, reason: String) -> Bool { + guard self.tokenCacheHydrationTasks[provider] == nil else { return true } + let now = Date() + if let lastAttempt = self.lastTokenCacheHydrationAttemptAt[provider], + now.timeIntervalSince(lastAttempt) < self.tokenCacheHydrationRetryInterval + { + return false + } + self.lastTokenCacheHydrationAttemptAt[provider] = now + self.tokenRefreshQueuedProviders.insert(provider) + + let fetcher = self.costUsageFetcher + let environment = self.environmentBase + let historyDays = self.settings.costUsageHistoryDays + let costScope = self.tokenCostScope(for: provider) + let allowVertexClaudeFallback = !self.isEnabled(.claude) + let override = self._test_cachedTokenUsageLoadOverride + self.tokenCacheHydrationTasks[provider] = Task(priority: .utility) { @MainActor [weak self] in + let snapshot: CostUsageTokenSnapshot? = if let override { + await override(provider) + } else { + await Task.detached(priority: .utility) { + fetcher.loadCachedTokenSnapshot( + provider: provider, + environment: environment, + now: now, + allowVertexClaudeFallback: allowVertexClaudeFallback, + codexHomePath: costScope.codexHomePath, + historyDays: historyDays) + }.value + } + + guard let self else { return } + self.tokenCacheHydrationTasks[provider] = nil + if Task.isCancelled { + self.tokenRefreshQueuedProviders.remove(provider) + return + } + guard let snapshot, !snapshot.daily.isEmpty else { + self.scheduleTokenRefresh( + force: false, + providers: [provider], + allowDuringMenuInteraction: true, + reason: "\(reason)-cache-miss") + return + } + self.tokenRefreshQueuedProviders.remove(provider) + guard self.tokenSnapshots[provider] == nil else { return } + self.tokenCostLogger.debug( + "cost usage hydrated cached snapshot provider=\(provider.rawValue) reason=\(reason)") + self.tokenSnapshots[provider] = snapshot + self.tokenErrors[provider] = nil + self.tokenFailureGates[provider]?.recordSuccess() + self.scheduleTokenRefresh( + force: false, + providers: [provider], + reason: "\(reason)-revalidate") + } + return true + } } diff --git a/Sources/CodexBar/UsageStore+TokenCostDeferral.swift b/Sources/CodexBar/UsageStore+TokenCostDeferral.swift new file mode 100644 index 0000000000..d463e329b5 --- /dev/null +++ b/Sources/CodexBar/UsageStore+TokenCostDeferral.swift @@ -0,0 +1,166 @@ +import Foundation + +extension UsageStore { + func beginInteractiveMenuTokenCostDeferral(reason: String) { + self.tokenCostInteractiveDeferralDepth += 1 + self.tokenCostInteractiveDeferralGeneration &+= 1 + self.tokenCostDeferredUntil = nil + self.tokenRefreshResumeTask?.cancel() + self.tokenRefreshResumeTask = nil + self.cancelScheduledTokenRefreshForMenuInteraction(reason: reason) + } + + func endInteractiveMenuTokenCostDeferral(reason: String) { + guard self.tokenCostInteractiveDeferralDepth > 0 else { return } + self.tokenCostInteractiveDeferralDepth -= 1 + self.tokenCostInteractiveDeferralGeneration &+= 1 + guard self.tokenCostInteractiveDeferralDepth == 0 else { return } + + let delay = max(0, self.tokenCostInteractionResumeDelay) + self.tokenCostDeferredUntil = Date().addingTimeInterval(delay) + self.scheduleTokenRefreshResumeAfterInteraction(reason: reason, delay: delay) + } + + func deferRefreshDuringMenuInteractionIfNeeded(forceTokenUsage: Bool) -> Bool { + if self.tokenCostInteractiveDeferralDepth > 0 { + self.refreshDeferredDuringMenuInteraction = true + self.refreshDeferredForceTokenUsage = self.refreshDeferredForceTokenUsage || forceTokenUsage + self.tokenCostLogger.debug( + "store refresh deferred for menu interaction " + + "activeMenuDepth=\(self.tokenCostInteractiveDeferralDepth) " + + "forceTokenUsage=\(forceTokenUsage)") + return true + } + + guard let deferredUntil = self.tokenCostDeferredUntil else { return false } + + let remaining = deferredUntil.timeIntervalSinceNow + guard remaining > 0 else { + self.tokenCostDeferredUntil = nil + return false + } + + self.refreshDeferredDuringMenuInteraction = true + self.refreshDeferredForceTokenUsage = self.refreshDeferredForceTokenUsage || forceTokenUsage + self.scheduleTokenRefreshResumeAfterInteraction(reason: "refresh-deferred-during-menu", delay: remaining) + return true + } + + func shouldDeferScheduledTokenRefresh(reason: String) -> Bool { + self.shouldDeferTokenRefreshForMenuInteraction(reason: reason) + } + + func shouldDeferTokenRefreshForMenuInteraction(reason: String) -> Bool { + if self.tokenCostInteractiveDeferralDepth > 0 { + self.tokenCostLogger + .debug("cost usage deferred reason=\(reason) activeMenuDepth=\(self.tokenCostInteractiveDeferralDepth)") + return true + } + + guard let deferredUntil = self.tokenCostDeferredUntil else { return false } + + let remaining = deferredUntil.timeIntervalSinceNow + guard remaining > 0 else { + self.tokenCostDeferredUntil = nil + return false + } + + self.scheduleTokenRefreshResumeAfterInteraction(reason: reason, delay: remaining) + return true + } + + private func scheduleTokenRefreshResumeAfterInteraction(reason: String, delay: TimeInterval) { + self.tokenRefreshResumeTask?.cancel() + let generation = self.tokenCostInteractiveDeferralGeneration + self.tokenRefreshResumeTask = Task { @MainActor [weak self] in + try? await Task.sleep(for: .seconds(delay)) + guard let self, !Task.isCancelled else { return } + guard self.tokenCostInteractiveDeferralDepth == 0 else { return } + guard self.tokenCostInteractiveDeferralGeneration == generation else { return } + self.tokenCostDeferredUntil = nil + self.tokenRefreshResumeTask = nil + self.tokenCostLogger + .debug("cost usage resuming after menu interaction reason=\(reason) delaySeconds=\(Int(delay))") + let shouldRunRefresh = self.refreshDeferredDuringMenuInteraction + let refreshForceTokenUsage = self.refreshDeferredForceTokenUsage + let regularTokenRefresh = self.tokenCostDeferredRefreshPending + let forceTokenRefresh = self.tokenCostDeferredForceRefreshPending + self.refreshDeferredDuringMenuInteraction = false + self.refreshDeferredForceTokenUsage = false + self.tokenCostDeferredRefreshPending = false + self.tokenCostDeferredForceRefreshPending = false + + if shouldRunRefresh { + await self.refresh(forceTokenUsage: refreshForceTokenUsage) + if forceTokenRefresh, !refreshForceTokenUsage { + self.scheduleTokenRefresh(force: true) + } else if regularTokenRefresh, !refreshForceTokenUsage { + self.scheduleTokenRefresh(force: false) + } + } else if forceTokenRefresh { + self.scheduleTokenRefresh(force: true) + } else if regularTokenRefresh { + self.scheduleTokenRefresh(force: false) + } else { + self.scheduleTokenRefresh(force: false) + } + } + } + + private func cancelScheduledTokenRefreshForMenuInteraction(reason: String) { + guard let task = self.tokenRefreshSequenceTask else { return } + let protectedProviders = self.tokenRefreshMenuAllowedProviders + .filter { self.shouldRunTokenCostRefreshDuringMenuInteraction($0) } + let inFlightProviders = Set(self.tokenRefreshInFlightStartedAt.keys) + let shouldKeepProtectedJob = !protectedProviders.isEmpty && + (inFlightProviders.isEmpty || !inFlightProviders.isDisjoint(with: protectedProviders)) + if shouldKeepProtectedJob { + let providersText = protectedProviders.map(\.rawValue).sorted().joined(separator: ",") + self.tokenCostLogger + .debug("cost usage kept running during menu interaction reason=\(reason) providers=\(providersText)") + return + } + + let inFlightDetails = self.tokenRefreshInFlightStartedAt + .sorted { $0.key.rawValue < $1.key.rawValue } + .map { provider, startedAt in + let elapsedMs = Int(Date().timeIntervalSince(startedAt) * 1000) + return "inFlightProvider=\(provider.rawValue) inFlightMs=\(elapsedMs)" + } + .joined(separator: " ") + + let suffix = inFlightDetails.isEmpty ? "inFlightProvider=none" : inFlightDetails + self.tokenCostLogger + .info("cost usage cancelled for menu interaction reason=\(reason) \(suffix)") + task.cancel() + } + + func clearCostUsageCache() async -> String? { + let errorMessage: String? = await Task.detached(priority: .utility) { + let fm = FileManager.default + let cacheDirs = [ + Self.costUsageCacheDirectory(fileManager: fm), + ] + + for cacheDir in cacheDirs { + do { + try fm.removeItem(at: cacheDir) + } catch let error as NSError { + if error.domain == NSCocoaErrorDomain, error.code == NSFileNoSuchFileError { continue } + return error.localizedDescription + } + } + return nil + }.value + + guard errorMessage == nil else { return errorMessage } + + self.tokenSnapshots.removeAll() + self.tokenErrors.removeAll() + self.lastTokenFetchAt.removeAll() + self.lastTokenFetchScope.removeAll() + self.tokenFailureGates[.codex]?.reset() + self.tokenFailureGates[.claude]?.reset() + return nil + } +} diff --git a/Sources/CodexBar/UsageStore+TokenCostScheduling.swift b/Sources/CodexBar/UsageStore+TokenCostScheduling.swift new file mode 100644 index 0000000000..f6c93fd5ec --- /dev/null +++ b/Sources/CodexBar/UsageStore+TokenCostScheduling.swift @@ -0,0 +1,115 @@ +import CodexBarCore +import Foundation + +extension UsageStore { + func scheduleTokenRefresh( + force: Bool, + providers: [UsageProvider]? = nil, + allowDuringMenuInteraction: Bool = false, + reason: String = "scheduled-refresh") + { + let refreshProviders = providers ?? self.enabledProvidersForBackgroundWork() + guard !refreshProviders.isEmpty else { return } + if !allowDuringMenuInteraction, + self + .shouldDeferTokenRefreshForMenuInteraction(reason: force ? "forced-scheduled-refresh" : reason) + { + if force { + self.tokenCostDeferredForceRefreshPending = true + } else { + self.tokenCostDeferredRefreshPending = true + } + return + } + + self.tokenRefreshQueuedProviders.formUnion(refreshProviders) + if allowDuringMenuInteraction { + self.tokenRefreshMenuAllowedProviders.formUnion(refreshProviders) + } + if force { + self.tokenRefreshSequenceTask?.cancel() + self.tokenRefreshSequenceTask = nil + } else if self.tokenRefreshSequenceTask != nil { + self.pendingTokenRefreshProviders.formUnion(refreshProviders) + self.pendingTokenRefreshAllowsMenuInteraction = + self.pendingTokenRefreshAllowsMenuInteraction || allowDuringMenuInteraction + return + } + + self.tokenRefreshSequenceTask = Task(priority: .utility) { [weak self] in + guard let self else { return } + defer { + Task { @MainActor [weak self] in + self?.tokenRefreshSequenceTask = nil + } + } + await self.refreshTokenUsageSequenceWorker( + force: force, + providers: refreshProviders, + allowDuringMenuInteraction: allowDuringMenuInteraction) + } + } + + func refreshTokenUsageSequenceNow(force: Bool) async { + if force, let existing = self.tokenRefreshSequenceTask { + existing.cancel() + await existing.value + self.tokenRefreshSequenceTask = nil + } + + await self.refreshTokenUsageSequence(force: force, providers: self.enabledProvidersForBackgroundWork()) + } + + private func refreshTokenUsageSequenceWorker( + force: Bool, + providers: [UsageProvider], + allowDuringMenuInteraction: Bool) async + { + var nextProviders = providers + var nextAllowsMenuInteraction = allowDuringMenuInteraction + var nextForce = force + defer { + if Task.isCancelled { + self.tokenRefreshQueuedProviders.subtract(nextProviders) + self.tokenRefreshQueuedProviders.subtract(self.pendingTokenRefreshProviders) + self.tokenRefreshMenuAllowedProviders.subtract(nextProviders) + self.tokenRefreshMenuAllowedProviders.subtract(self.pendingTokenRefreshProviders) + self.pendingTokenRefreshProviders = [] + self.pendingTokenRefreshAllowsMenuInteraction = false + } + } + + while !Task.isCancelled { + await self.refreshTokenUsageSequence( + force: nextForce, + providers: nextProviders, + allowDuringMenuInteraction: nextAllowsMenuInteraction) + guard !Task.isCancelled else { break } + guard !self.pendingTokenRefreshProviders.isEmpty else { break } + + nextProviders = Array(self.pendingTokenRefreshProviders) + nextAllowsMenuInteraction = self.pendingTokenRefreshAllowsMenuInteraction + nextForce = false + self.pendingTokenRefreshProviders = [] + self.pendingTokenRefreshAllowsMenuInteraction = false + } + } + + private func refreshTokenUsageSequence( + force: Bool, + providers: [UsageProvider], + allowDuringMenuInteraction: Bool = false) async + { + for provider in providers { + if Task.isCancelled { break } + if !force, + !allowDuringMenuInteraction, + !self.shouldRunTokenCostRefreshDuringMenuInteraction(provider), + self.shouldDeferScheduledTokenRefresh(reason: "sequence-refresh") + { + break + } + await self.refreshTokenUsage(provider, force: force) + } + } +} diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index d2c35695d2..fd01739025 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -18,6 +18,7 @@ extension UsageStore { _ = self.kiloScopeSnapshots _ = self.tokenSnapshots _ = self.tokenErrors + _ = self.tokenRefreshQueuedProviders _ = self.tokenRefreshInFlight _ = self.credits _ = self.lastCreditsError @@ -148,6 +149,7 @@ final class UsageStore { var kiloScopeSnapshots: [KiloScopeSnapshot] = [] var tokenSnapshots: [UsageProvider: CostUsageTokenSnapshot] = [:] var tokenErrors: [UsageProvider: String] = [:] + var tokenRefreshQueuedProviders: Set = [] var tokenRefreshInFlight: Set = [] var credits: CreditsSnapshot? var lastCreditsError: String? @@ -200,11 +202,13 @@ final class UsageStore { @ObservationIgnored var _test_widgetSnapshotSaveOverride: (@MainActor (WidgetSnapshot) async -> Void)? @ObservationIgnored var _test_providerRefreshOverride: (@MainActor (UsageProvider) async -> Void)? @ObservationIgnored var _test_tokenUsageRefreshOverride: (@MainActor (UsageProvider, Bool) async -> Void)? + @ObservationIgnored var _test_cachedTokenUsageLoadOverride: + (@MainActor (UsageProvider) async -> CostUsageTokenSnapshot?)? @ObservationIgnored var widgetSnapshotPersistTask: Task? @ObservationIgnored let codexFetcher: UsageFetcher @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching - @ObservationIgnored private let costUsageFetcher: CostUsageFetcher + @ObservationIgnored let costUsageFetcher: CostUsageFetcher @ObservationIgnored let browserDetection: BrowserDetection @ObservationIgnored private let registry: ProviderRegistry @ObservationIgnored let settings: SettingsStore @@ -212,7 +216,7 @@ final class UsageStore { @ObservationIgnored private let sessionQuotaNotifier: any SessionQuotaNotifying @ObservationIgnored private let sessionQuotaLogger = CodexBarLog.logger(LogCategories.sessionQuota) @ObservationIgnored let openAIWebLogger = CodexBarLog.logger(LogCategories.openAIWeb) - @ObservationIgnored private let tokenCostLogger = CodexBarLog.logger(LogCategories.tokenCost) + @ObservationIgnored let tokenCostLogger = CodexBarLog.logger(LogCategories.tokenCost) @ObservationIgnored let augmentLogger = CodexBarLog.logger(LogCategories.augment) @ObservationIgnored let providerLogger = CodexBarLog.logger(LogCategories.providers) @ObservationIgnored var openAIWebDebugLines: [String] = [] @@ -224,7 +228,20 @@ final class UsageStore { @ObservationIgnored private var providerAvailabilityCache: [UsageProvider: ProviderAvailabilityCacheEntry] = [:] @ObservationIgnored private var timerTask: Task? @ObservationIgnored private var tokenTimerTask: Task? - @ObservationIgnored private var tokenRefreshSequenceTask: Task? + @ObservationIgnored var tokenRefreshSequenceTask: Task? + @ObservationIgnored var tokenRefreshResumeTask: Task? + @ObservationIgnored var pendingTokenRefreshProviders: Set = [] + @ObservationIgnored var pendingTokenRefreshAllowsMenuInteraction = false + @ObservationIgnored var tokenRefreshMenuAllowedProviders: Set = [] + @ObservationIgnored var tokenCacheHydrationTasks: [UsageProvider: Task] = [:] + @ObservationIgnored var lastTokenCacheHydrationAttemptAt: [UsageProvider: Date] = [:] + @ObservationIgnored var tokenCostInteractiveDeferralDepth = 0 + @ObservationIgnored var tokenCostInteractiveDeferralGeneration: UInt64 = 0 + @ObservationIgnored var tokenCostDeferredUntil: Date? + @ObservationIgnored var refreshDeferredDuringMenuInteraction = false + @ObservationIgnored var refreshDeferredForceTokenUsage = false + @ObservationIgnored var tokenCostDeferredRefreshPending = false + @ObservationIgnored var tokenCostDeferredForceRefreshPending = false @ObservationIgnored var storageRefreshTask: Task? @ObservationIgnored var storageRefreshGeneration: UInt64 = 0 @ObservationIgnored var storageRefreshInFlightSignature: String? @@ -245,12 +262,15 @@ final class UsageStore { @ObservationIgnored var lastPermissionPromptNotificationAt: [UsageProvider: Date] = [:] @ObservationIgnored var lastTokenFetchAt: [UsageProvider: Date] = [:] @ObservationIgnored var lastTokenFetchScope: [UsageProvider: String] = [:] + @ObservationIgnored var tokenRefreshInFlightStartedAt: [UsageProvider: Date] = [:] @ObservationIgnored var planUtilizationHistory: [UsageProvider: PlanUtilizationHistoryBuckets] = [:] @ObservationIgnored var weeklyLimitResetDetectorStates: [String: WeeklyLimitResetDetectorState] = [:] @ObservationIgnored private var hasCompletedInitialRefresh: Bool = false @ObservationIgnored private let providerAvailabilityCacheTTL: TimeInterval = 1 @ObservationIgnored private let tokenFetchTTL: TimeInterval = 60 * 60 @ObservationIgnored private let tokenFetchTimeout: TimeInterval = 10 * 60 + @ObservationIgnored var tokenCacheHydrationRetryInterval: TimeInterval = 30 + @ObservationIgnored var tokenCostInteractionResumeDelay: TimeInterval = 2 * 60 @ObservationIgnored private let startupBehavior: StartupBehavior @ObservationIgnored let planUtilizationPersistenceCoordinator: PlanUtilizationHistoryPersistenceCoordinator @@ -568,7 +588,11 @@ final class UsageStore { } if forceTokenUsage { - await self.refreshTokenUsageSequenceNow(force: true) + if self.shouldDeferTokenRefreshForMenuInteraction(reason: "forced-refresh") { + self.tokenCostDeferredForceRefreshPending = true + } else { + await self.refreshTokenUsageSequenceNow(force: true) + } } else { // Token-cost usage can be slow; run it outside regular/menu-open refreshes so we don't block UI. self.scheduleTokenRefresh(force: false) @@ -658,46 +682,14 @@ final class UsageStore { } } - private func scheduleTokenRefresh(force: Bool) { - if force { - self.tokenRefreshSequenceTask?.cancel() - self.tokenRefreshSequenceTask = nil - } else if self.tokenRefreshSequenceTask != nil { - return - } - - self.tokenRefreshSequenceTask = Task(priority: .utility) { [weak self] in - guard let self else { return } - defer { - Task { @MainActor [weak self] in - self?.tokenRefreshSequenceTask = nil - } - } - await self.refreshTokenUsageSequence(force: force) - } - } - - private func refreshTokenUsageSequenceNow(force: Bool) async { - if force, let existing = self.tokenRefreshSequenceTask { - existing.cancel() - await existing.value - self.tokenRefreshSequenceTask = nil - } - - await self.refreshTokenUsageSequence(force: force) - } - - private func refreshTokenUsageSequence(force: Bool) async { - for provider in self.enabledProvidersForBackgroundWork() { - if Task.isCancelled { break } - await self.refreshTokenUsage(provider, force: force) - } - } - deinit { self.timerTask?.cancel() self.tokenTimerTask?.cancel() self.tokenRefreshSequenceTask?.cancel() + self.tokenRefreshResumeTask?.cancel() + for task in self.tokenCacheHydrationTasks.values { + task.cancel() + } self.storageRefreshTask?.cancel() self.codexPlanHistoryBackfillTask?.cancel() } @@ -1451,36 +1443,11 @@ extension UsageStore { } } - func clearCostUsageCache() async -> String? { - let errorMessage: String? = await Task.detached(priority: .utility) { - let fm = FileManager.default - let cacheDirs = [ - Self.costUsageCacheDirectory(fileManager: fm), - ] - - for cacheDir in cacheDirs { - do { - try fm.removeItem(at: cacheDir) - } catch let error as NSError { - if error.domain == NSCocoaErrorDomain, error.code == NSFileNoSuchFileError { continue } - return error.localizedDescription - } - } - return nil - }.value - - guard errorMessage == nil else { return errorMessage } - - self.tokenSnapshots.removeAll() - self.tokenErrors.removeAll() - self.lastTokenFetchAt.removeAll() - self.lastTokenFetchScope.removeAll() - self.tokenFailureGates[.codex]?.reset() - self.tokenFailureGates[.claude]?.reset() - return nil - } - - private func refreshTokenUsage(_ provider: UsageProvider, force: Bool) async { + func refreshTokenUsage(_ provider: UsageProvider, force: Bool) async { + defer { + self.tokenRefreshQueuedProviders.remove(provider) + self.tokenRefreshMenuAllowedProviders.remove(provider) + } guard ProviderDescriptorRegistry.descriptor(for: provider).tokenCost.supportsTokenCost else { self.tokenSnapshots.removeValue(forKey: provider) self.tokenErrors[provider] = nil @@ -1543,9 +1510,14 @@ extension UsageStore { self.lastTokenFetchAt[provider] = now self.lastTokenFetchScope[provider] = costScopeSignature self.tokenRefreshInFlight.insert(provider) - defer { self.tokenRefreshInFlight.remove(provider) } let startedAt = Date() + self.tokenRefreshInFlightStartedAt[provider] = startedAt + defer { + self.tokenRefreshInFlight.remove(provider) + self.tokenRefreshInFlightStartedAt.removeValue(forKey: provider) + } + let providerText = provider.rawValue self.tokenCostLogger .debug("cost usage start provider=\(providerText) force=\(force)") @@ -1584,6 +1556,7 @@ extension UsageStore { guard let snapshot = try await group.next() else { throw CancellationError() } return snapshot } + try Task.checkCancellation() guard !snapshot.daily.isEmpty else { self.tokenSnapshots.removeValue(forKey: provider) @@ -1608,8 +1581,15 @@ extension UsageStore { self.tokenFailureGates[provider]?.recordSuccess() self.persistWidgetSnapshot(reason: "token-usage") } catch { - if error is CancellationError { return } let duration = Date().timeIntervalSince(startedAt) + if error is CancellationError { + let durationText = String(format: "%.2f", duration) + self.tokenCostLogger + .info("cost usage cancelled provider=\(providerText) duration=\(durationText)s") + self.lastTokenFetchAt.removeValue(forKey: provider) + self.lastTokenFetchScope.removeValue(forKey: provider) + return + } let msg = error.localizedDescription let durationText = String(format: "%.2f", duration) let message = "cost usage failed provider=\(providerText) duration=\(durationText)s error=\(msg)" diff --git a/Sources/CodexBarCore/CostUsageFetcher.swift b/Sources/CodexBarCore/CostUsageFetcher.swift index f9413e6516..025408707f 100644 --- a/Sources/CodexBarCore/CostUsageFetcher.swift +++ b/Sources/CodexBarCore/CostUsageFetcher.swift @@ -41,6 +41,23 @@ public struct CostUsageFetcher: Sendable { refreshPricingInBackground: refreshPricingInBackground) } + public func loadCachedTokenSnapshot( + provider: UsageProvider, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date(), + allowVertexClaudeFallback: Bool = false, + codexHomePath: String? = nil, + historyDays: Int = 30) -> CostUsageTokenSnapshot? + { + Self.loadCachedTokenSnapshot( + provider: provider, + environment: environment, + now: now, + allowVertexClaudeFallback: allowVertexClaudeFallback, + codexHomePath: codexHomePath, + historyDays: historyDays) + } + static func loadTokenSnapshot( provider: UsageProvider, environment: [String: String] = ProcessInfo.processInfo.environment, @@ -98,12 +115,14 @@ public struct CostUsageFetcher: Sendable { if forceRefresh { options.refreshMinIntervalSeconds = 0 } + try Task.checkCancellation() var daily = CostUsageScanner.loadDailyReport( provider: provider, since: since, until: until, now: now, options: options) + try Task.checkCancellation() if provider == .vertexai, !allowVertexClaudeFallback, @@ -118,6 +137,7 @@ public struct CostUsageFetcher: Sendable { until: until, now: now, options: fallback) + try Task.checkCancellation() } if provider == .codex || provider == .claude { @@ -134,12 +154,74 @@ public struct CostUsageFetcher: Sendable { until: until, now: now, options: piOptions) + try Task.checkCancellation() daily = CostUsageDailyReport.merged([daily, piReport]) } return Self.tokenSnapshot(from: daily, now: now, historyDays: clampedHistoryDays) } + static func loadCachedTokenSnapshot( + provider: UsageProvider, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date(), + allowVertexClaudeFallback: Bool = false, + codexHomePath: String? = nil, + historyDays: Int = 30, + scannerOptions overrideScannerOptions: CostUsageScanner.Options? = nil, + piScannerOptions overridePiScannerOptions: PiSessionCostScanner + .Options? = nil) -> CostUsageTokenSnapshot? + { + guard provider == .codex || provider == .claude || provider == .vertexai else { + return nil + } + + _ = environment + let until = now + let clampedHistoryDays = max(1, min(365, historyDays)) + let since = Calendar.current.date(byAdding: .day, value: -(clampedHistoryDays - 1), to: now) ?? now + + var options = overrideScannerOptions ?? CostUsageScanner.Options(cacheOnly: true) + options.cacheOnly = true + if provider == .codex, + let codexHomePath = codexHomePath?.trimmingCharacters(in: .whitespacesAndNewlines), + !codexHomePath.isEmpty + { + options.codexSessionsRoot = URL(fileURLWithPath: codexHomePath, isDirectory: true) + .appendingPathComponent("sessions", isDirectory: true) + } + if provider == .vertexai { + options.claudeLogProviderFilter = allowVertexClaudeFallback ? .all : .vertexAIOnly + } else if provider == .claude { + options.claudeLogProviderFilter = .excludeVertexAI + } + + var daily = CostUsageScanner.loadDailyReport( + provider: provider, + since: since, + until: until, + now: now, + options: options) + + if provider == .codex || provider == .claude { + var piOptions = overridePiScannerOptions ?? PiSessionCostScanner.Options(cacheOnly: true) + piOptions.cacheOnly = true + if piOptions.cacheRoot == nil { + piOptions.cacheRoot = options.cacheRoot + } + let piReport = PiSessionCostScanner.loadDailyReport( + provider: provider, + since: since, + until: until, + now: now, + options: piOptions) + daily = CostUsageDailyReport.merged([daily, piReport]) + } + + let snapshot = Self.tokenSnapshot(from: daily, now: now, historyDays: clampedHistoryDays) + return snapshot.daily.isEmpty ? nil : snapshot + } + private static func loadBedrockDailyReport( environment: [String: String], since: Date, diff --git a/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift b/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift index 714759e3ca..86cecf019b 100644 --- a/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift +++ b/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift @@ -1,5 +1,5 @@ // Generated by Scripts/regenerate-codex-parser-hash.sh. Do not edit by hand. enum CodexParserHash { - static let value = "5387f958e9e06f7d" + static let value = "8dfbffcff47827cf" } diff --git a/Sources/CodexBarCore/PiSessionCostScanner.swift b/Sources/CodexBarCore/PiSessionCostScanner.swift index 2fcce0c816..2fa7b36b14 100644 --- a/Sources/CodexBarCore/PiSessionCostScanner.swift +++ b/Sources/CodexBarCore/PiSessionCostScanner.swift @@ -5,17 +5,20 @@ enum PiSessionCostScanner { var piSessionsRoot: URL? var cacheRoot: URL? var refreshMinIntervalSeconds: TimeInterval = 60 + var cacheOnly: Bool = false var forceRescan: Bool = false init( piSessionsRoot: URL? = nil, cacheRoot: URL? = nil, refreshMinIntervalSeconds: TimeInterval = 60, + cacheOnly: Bool = false, forceRescan: Bool = false) { self.piSessionsRoot = piSessionsRoot self.cacheRoot = cacheRoot self.refreshMinIntervalSeconds = refreshMinIntervalSeconds + self.cacheOnly = cacheOnly self.forceRescan = forceRescan } } @@ -53,17 +56,18 @@ enum PiSessionCostScanner { let range = CostUsageScanner.CostUsageDayRange(since: since, until: until) var cache = PiSessionCostCacheIO.load(cacheRoot: options.cacheRoot) + let cacheBeforeRefresh = cache let nowMs = Int64(now.timeIntervalSince1970 * 1000) let refreshMs = Int64(max(0, options.refreshMinIntervalSeconds) * 1000) let pricingContext = ModelsDevPricingContext( catalog: CostUsagePricing.modelsDevCatalog(now: now, cacheRoot: options.cacheRoot), cacheRoot: options.cacheRoot) let windowExpanded = self.requestedWindowExpandsCache(range: range, cache: cache) - let shouldRefresh = options.forceRescan + let shouldRefresh = !options.cacheOnly && (options.forceRescan || windowExpanded || refreshMs == 0 || cache.lastScanUnixMs == 0 - || nowMs - cache.lastScanUnixMs > refreshMs + || nowMs - cache.lastScanUnixMs > refreshMs) if shouldRefresh { let root = self.defaultPiSessionsRoot(options: options) @@ -71,29 +75,37 @@ enum PiSessionCostScanner { let files = self.listPiSessionFiles(root: root, startCutoffLocal: startCutoff) let filePathsInScan = Set(files.map(\.path)) - for fileURL in files { - self.scanPiSessionFile( - fileURL: fileURL, - range: range, - forceRescan: options.forceRescan || windowExpanded, - pricingContext: pricingContext, - cache: &cache) - } + do { + for fileURL in files { + try Task.checkCancellation() + try self.scanPiSessionFile( + fileURL: fileURL, + range: range, + forceRescan: options.forceRescan || windowExpanded, + pricingContext: pricingContext, + cache: &cache) + } - for key in cache.files.keys where !filePathsInScan.contains(key) { - if let old = cache.files[key] { - self.applyContributions( - daysByProvider: &cache.daysByProvider, - contributions: old.contributions, - sign: -1) + try Task.checkCancellation() + for key in cache.files.keys where !filePathsInScan.contains(key) { + if let old = cache.files[key] { + self.applyContributions( + daysByProvider: &cache.daysByProvider, + contributions: old.contributions, + sign: -1) + } + cache.files.removeValue(forKey: key) } - cache.files.removeValue(forKey: key) - } - cache.scanSinceKey = range.scanSinceKey - cache.scanUntilKey = range.scanUntilKey - cache.lastScanUnixMs = nowMs - PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot) + cache.scanSinceKey = range.scanSinceKey + cache.scanUntilKey = range.scanUntilKey + cache.lastScanUnixMs = nowMs + PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot) + } catch is CancellationError { + cache = cacheBeforeRefresh + } catch { + cache = cacheBeforeRefresh + } } return self.buildReport( @@ -180,6 +192,7 @@ enum PiSessionCostScanner { forceRescan: Bool, pricingContext: ModelsDevPricingContext, cache: inout PiSessionCostCache) + throws { let path = fileURL.path let attrs = (try? FileManager.default.attributesOfItem(atPath: path)) ?? [:] @@ -206,7 +219,7 @@ enum PiSessionCostScanner { cached.parsedBytes > 0, cached.parsedBytes <= size { - let delta = self.parsePiSessionFile( + let delta = try self.parsePiSessionFile( fileURL: fileURL, range: range, startOffset: cached.parsedBytes, @@ -235,7 +248,7 @@ enum PiSessionCostScanner { sign: -1) } - let parsed = self.parsePiSessionFile( + let parsed = try self.parsePiSessionFile( fileURL: fileURL, range: range, pricingContext: pricingContext) @@ -256,7 +269,7 @@ enum PiSessionCostScanner { range: CostUsageScanner.CostUsageDayRange, startOffset: Int64 = 0, initialModelContext: PiModelContext? = nil, - pricingContext: ModelsDevPricingContext? = nil) -> ParseResult + pricingContext: ModelsDevPricingContext? = nil) throws -> ParseResult { var currentModelContext = initialModelContext var contributions: [String: [String: [String: PiPackedUsage]]] = [:] @@ -292,41 +305,48 @@ enum PiSessionCostScanner { } } - let parsedBytes = (try? CostUsageJsonl.scan( - fileURL: fileURL, - offset: startOffset, - maxLineBytes: Self.maxLineBytes, - prefixBytes: Self.maxLineBytes, - onLine: { line in - guard !line.bytes.isEmpty, !line.wasTruncated else { return } - autoreleasepool { - guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any] - else { return } - guard let type = object["type"] as? String else { return } - - if type == "model_change" { - currentModelContext = self.modelContext(from: object) - return + let parsedBytes: Int64 + do { + parsedBytes = try CostUsageJsonl.scan( + fileURL: fileURL, + offset: startOffset, + maxLineBytes: Self.maxLineBytes, + prefixBytes: Self.maxLineBytes, + onLine: { line in + guard !line.bytes.isEmpty, !line.wasTruncated else { return } + autoreleasepool { + guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any] + else { return } + guard let type = object["type"] as? String else { return } + + if type == "model_change" { + currentModelContext = self.modelContext(from: object) + return + } + + guard type == "message", let message = object["message"] as? [String: Any] else { return } + guard (message["role"] as? String) == "assistant" else { return } + + let identity = self.resolveAssistantIdentity( + entry: object, + message: message, + fallback: currentModelContext) + guard let identity else { return } + guard let date = self.timestampDate(entry: object, message: message) else { return } + let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date) + let usage = self.extractUsage( + provider: identity.provider, + modelName: identity.modelName, + message: message, + pricingContext: pricingContext) + add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage) } - - guard type == "message", let message = object["message"] as? [String: Any] else { return } - guard (message["role"] as? String) == "assistant" else { return } - - let identity = self.resolveAssistantIdentity( - entry: object, - message: message, - fallback: currentModelContext) - guard let identity else { return } - guard let date = self.timestampDate(entry: object, message: message) else { return } - let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date) - let usage = self.extractUsage( - provider: identity.provider, - modelName: identity.modelName, - message: message, - pricingContext: pricingContext) - add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage) - } - })) ?? startOffset + }) + } catch is CancellationError { + throw CancellationError() + } catch { + parsedBytes = startOffset + } return ParseResult( contributions: contributions, diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageJsonl.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageJsonl.swift index 2bf74841c0..e844b3ce7d 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageJsonl.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageJsonl.swift @@ -28,6 +28,7 @@ enum CostUsageJsonl { var lineBytes = 0 var truncated = false var bytesRead: Int64 = 0 + var scannedLines = 0 func appendSegment(_ bytes: UnsafePointer, count: Int) { guard count > 0 else { return } @@ -43,31 +44,36 @@ enum CostUsageJsonl { } } - func flushLine() { + func flushLine() throws { guard lineBytes > 0 else { return } let line = Line(bytes: current, wasTruncated: truncated) onLine(line) + scannedLines += 1 + if scannedLines.isMultiple(of: 128) { + try Task.checkCancellation() + } current.removeAll(keepingCapacity: true) lineBytes = 0 truncated = false } while true { + try Task.checkCancellation() let chunk = try handle.read(upToCount: 256 * 1024) ?? Data() if chunk.isEmpty { - flushLine() + try flushLine() break } bytesRead += Int64(chunk.count) - chunk.withUnsafeBytes { rawBuffer in + try chunk.withUnsafeBytes { rawBuffer in guard let base = rawBuffer.bindMemory(to: UInt8.self).baseAddress else { return } var segmentStart = 0 var index = 0 while index < rawBuffer.count { if base[index] == 0x0A { appendSegment(base.advanced(by: segmentStart), count: index - segmentStart) - flushLine() + try flushLine() segmentStart = index + 1 } index += 1 diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift index 6e56a88591..e47793d186 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift @@ -567,15 +567,16 @@ extension CostUsageScanner { options: Options) -> CostUsageDailyReport { var cache = CostUsageCacheIO.load(provider: provider, cacheRoot: options.cacheRoot) + let cacheBeforeRefresh = cache let nowMs = Int64(now.timeIntervalSince1970 * 1000) let refreshMs = Int64(max(0, options.refreshMinIntervalSeconds) * 1000) let windowExpanded = Self.requestedWindowExpandsCache(range: range, cache: cache) - let shouldRefresh = options.forceRescan + let shouldRefresh = !options.cacheOnly && (options.forceRescan || windowExpanded || refreshMs == 0 || cache.lastScanUnixMs == 0 - || nowMs - cache.lastScanUnixMs > refreshMs + || nowMs - cache.lastScanUnixMs > refreshMs) let roots = self.defaultClaudeProjectsRoots(options: options) let providerFilter = options.claudeLogProviderFilter @@ -583,6 +584,16 @@ extension CostUsageScanner { var touched: Set = [] if shouldRefresh { + func cancelledReport() -> CostUsageDailyReport { + Self.buildClaudeReportFromCache( + cache: cacheBeforeRefresh, + range: range, + modelsDevCatalog: nil, + modelsDevCacheRoot: options.cacheRoot) + } + + if Task.isCancelled { return cancelledReport() } + if options.forceRescan { cache = CostUsageCache() } @@ -596,9 +607,11 @@ extension CostUsageScanner { modelsDevCacheRoot: options.cacheRoot) for root in roots { + if Task.isCancelled { return cancelledReport() } Self.scanClaudeRoot( root: root, state: scanState) + if Task.isCancelled { return cancelledReport() } } cache = scanState.cache @@ -613,6 +626,7 @@ extension CostUsageScanner { Self.pruneDays(cache: &cache, sinceKey: range.scanSinceKey, untilKey: range.scanUntilKey) cache.scanSinceKey = range.scanSinceKey cache.scanUntilKey = range.scanUntilKey + if Task.isCancelled { return cancelledReport() } cache.lastScanUnixMs = nowMs CostUsageCacheIO.save(provider: provider, cache: cache, cacheRoot: options.cacheRoot) } diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index 55602a594b..5b5dc42232 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -24,6 +24,7 @@ enum CostUsageScanner { var codexTraceDatabaseURL: URL? var refreshMinIntervalSeconds: TimeInterval = 60 var claudeLogProviderFilter: ClaudeLogProviderFilter = .all + var cacheOnly: Bool = false /// Force a full rescan, ignoring per-file cache and incremental offsets. var forceRescan: Bool = false @@ -33,6 +34,7 @@ enum CostUsageScanner { cacheRoot: URL? = nil, codexTraceDatabaseURL: URL? = nil, claudeLogProviderFilter: ClaudeLogProviderFilter = .all, + cacheOnly: Bool = false, forceRescan: Bool = false) { self.codexSessionsRoot = codexSessionsRoot @@ -40,6 +42,7 @@ enum CostUsageScanner { self.cacheRoot = cacheRoot self.codexTraceDatabaseURL = codexTraceDatabaseURL self.claudeLogProviderFilter = claudeLogProviderFilter + self.cacheOnly = cacheOnly self.forceRescan = forceRescan } } @@ -1562,7 +1565,7 @@ enum CostUsageScanner { newKeys: priorityTurnKeys, range: range) : [] - let shouldRefresh = options.forceRescan + let shouldRefresh = !options.cacheOnly && (options.forceRescan || windowExpanded || rootsChanged || needsCostCacheMigration @@ -1572,7 +1575,7 @@ enum CostUsageScanner { || priorityTurnsChanged || refreshMs == 0 || cache.lastScanUnixMs == 0 - || nowMs - cache.lastScanUnixMs > refreshMs + || nowMs - cache.lastScanUnixMs > refreshMs) return CodexRefreshPlan( refreshMs: refreshMs, @@ -1596,12 +1599,32 @@ enum CostUsageScanner { shouldRefresh: shouldRefresh) } + private static func clearMovedCodexCacheForCacheOnlyLoad( + cache: inout CostUsageCache, + options: Options, + plan: CodexRefreshPlan) + { + guard options.cacheOnly, plan.rootsChanged else { return } + cache = CostUsageCache() + } + private static func loadCodexDaily(range: CostUsageDayRange, now: Date, options: Options) -> CostUsageDailyReport { var cache = CostUsageCacheIO.load(provider: .codex, cacheRoot: options.cacheRoot) + let cacheBeforeRefresh = cache let nowMs = Int64(now.timeIntervalSince1970 * 1000) let plan = Self.makeCodexRefreshPlan(cache: cache, range: range, now: now, nowMs: nowMs, options: options) + Self.clearMovedCodexCacheForCacheOnlyLoad(cache: &cache, options: options, plan: plan) if plan.shouldRefresh { + func cancelledReport() -> CostUsageDailyReport { + Self.buildCodexReportFromCache( + cache: cacheBeforeRefresh, + range: range, + modelsDevCatalog: plan.modelsDevCatalog, + modelsDevCacheRoot: options.cacheRoot, + priorityTurns: plan.priorityTurns) + } + if options.forceRescan { cache = CostUsageCache() } @@ -1613,7 +1636,7 @@ enum CostUsageScanner { .map { Calendar.current.startOfDay(for: $0) } var seenPaths: Set = [] var files: [URL] = [] - for root in plan.roots { + for root in plan.roots where !Task.isCancelled { let rootFiles = Self.listCodexSessionFiles( root: root, scanSinceKey: range.scanSinceKey, @@ -1638,6 +1661,7 @@ enum CostUsageScanner { } } } + guard !Task.isCancelled else { return cancelledReport() } for fileURL in Self.cachedCodexSessionFiles(cache: cache, range: range, roots: plan.roots) .sorted(by: { $0.path < $1.path }) @@ -1661,7 +1685,7 @@ enum CostUsageScanner { modelsDevCatalog: plan.modelsDevCatalog, modelsDevCacheRoot: options.cacheRoot, priorityTurns: plan.priorityTurns) - for fileURL in files { + for fileURL in files where !Task.isCancelled { Self.scanCodexFile( fileURL: fileURL, context: CodexFileScanContext( @@ -1677,6 +1701,7 @@ enum CostUsageScanner { cache: &cache, state: &scanState) } + guard !Task.isCancelled else { return cancelledReport() } Self.pruneForceRescanFilesOutsideWindow( cache: &cache, @@ -1734,6 +1759,7 @@ enum CostUsageScanner { retainedSinceKey: retainedSinceKey, retainedUntilKey: retainedUntilKey) } + guard !Task.isCancelled else { return cancelledReport() } cache.lastScanUnixMs = nowMs CostUsageCacheIO.save(provider: .codex, cache: cache, cacheRoot: options.cacheRoot) } diff --git a/Tests/CodexBarTests/CodexAccountScopedRefreshCacheTests.swift b/Tests/CodexBarTests/CodexAccountScopedRefreshCacheTests.swift new file mode 100644 index 0000000000..ff1cb3782e --- /dev/null +++ b/Tests/CodexBarTests/CodexAccountScopedRefreshCacheTests.swift @@ -0,0 +1,75 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +extension CodexAccountScopedRefreshTests { + @Test + func `cancelled codex usage refresh preserves last good snapshot without surfacing error`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-cancelled-refresh") + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + let cached = self.codexSnapshot(email: "alpha@example.com", usedPercent: 17) + store._setSnapshotForTesting(cached, provider: .codex) + self.installFailingCodexProvider(on: store, error: TestRefreshError(message: "Network error: cancelled")) + + await store.refreshProvider(.codex, allowDisabled: true) + + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "alpha@example.com") + #expect(store.snapshots[.codex]?.primary?.usedPercent == 17) + #expect(store.errors[.codex] == nil) + } + + @Test + func `account transition reuses cached snapshot for selected visible codex account`() throws { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-cached-selected-account") + settings.refreshFrequency = .manual + settings.codexCookieSource = .off + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + settings.codexActiveSource = .liveSystem + + let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let managedAccount = ManagedCodexAccount( + id: managedAccountID, + email: "beta@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) + defer { + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + try? FileManager.default.removeItem(at: storeURL) + } + settings._test_managedCodexAccountStoreURL = storeURL + + let store = self.makeUsageStore(settings: settings) + store._setSnapshotForTesting(self.codexSnapshot(email: "alpha@example.com", usedPercent: 9), provider: .codex) + store.lastCodexAccountScopedRefreshGuard = store + .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) + + let betaAccount = try #require(settings.codexVisibleAccountProjection.visibleAccounts + .first { $0.email == "beta@example.com" }) + let betaSnapshot = self.codexSnapshot(email: "beta@example.com", usedPercent: 41) + store.codexAccountSnapshots = [ + CodexAccountUsageSnapshot( + account: betaAccount, + snapshot: betaSnapshot, + error: nil, + sourceLabel: "cached"), + ] + + settings.selectDisplayedCodexVisibleAccount(betaAccount) + let didInvalidate = store.prepareCodexAccountScopedRefreshIfNeeded() + + #expect(didInvalidate) + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "beta@example.com") + #expect(store.snapshots[.codex]?.primary?.usedPercent == 41) + #expect(store.lastSourceLabels[.codex] == "cached") + #expect(store.errors[.codex] == nil) + #expect(store.codexAccountSnapshots.count == 1) + } +} diff --git a/Tests/CodexBarTests/CodexUserFacingErrorTests.swift b/Tests/CodexBarTests/CodexUserFacingErrorTests.swift index 4fb88fc3c0..cae468eb34 100644 --- a/Tests/CodexBarTests/CodexUserFacingErrorTests.swift +++ b/Tests/CodexBarTests/CodexUserFacingErrorTests.swift @@ -60,6 +60,14 @@ struct CodexUserFacingErrorTests { #expect(store.userFacingError(for: .codex) == "Codex usage is temporarily unavailable. Try refreshing.") } + @Test + func `cancelled codex refresh is hidden from user facing errors`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-cancelled") + store.errors[.codex] = "Network error: cancelled" + + #expect(store.userFacingError(for: .codex) == nil) + } + @Test func `decode mismatch codex error is sanitized`() { let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-decode-mismatch") diff --git a/Tests/CodexBarTests/CostUsageFetcherTests.swift b/Tests/CodexBarTests/CostUsageFetcherTests.swift index d58c63a378..0d79d13a1c 100644 --- a/Tests/CodexBarTests/CostUsageFetcherTests.swift +++ b/Tests/CodexBarTests/CostUsageFetcherTests.swift @@ -865,3 +865,74 @@ struct CostUsageFetcherTests { ]).write(to: url, atomically: true, encoding: .utf8) } } + +extension CostUsageFetcherTests { + @Test + func `fetcher can hydrate codex snapshot from cache without scanning`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 4, day: 8) + try Self.writeCodexSessionFile( + homeRoot: env.codexHomeRoot, + env: env, + day: day, + filename: "cached.jsonl", + tokens: 100) + + let options = CostUsageScanner.Options(cacheRoot: env.cacheRoot) + let piOptions = PiSessionCostScanner.Options(piSessionsRoot: env.piSessionsRoot, cacheRoot: env.cacheRoot) + _ = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + codexHomePath: env.codexHomeRoot.path, + scannerOptions: options, + piScannerOptions: piOptions) + + try FileManager.default.removeItem( + at: env.codexHomeRoot.appendingPathComponent("sessions", isDirectory: true)) + + let cached = CostUsageFetcher.loadCachedTokenSnapshot( + provider: .codex, + now: day.addingTimeInterval(120), + codexHomePath: env.codexHomeRoot.path, + scannerOptions: options, + piScannerOptions: piOptions) + + #expect(cached?.sessionTokens == 100) + } + + @Test + func `fetcher does not hydrate codex snapshot from another home cache`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 4, day: 8) + let managedHome = env.root.appendingPathComponent("managed-codex-home", isDirectory: true) + try Self.writeCodexSessionFile( + homeRoot: env.codexHomeRoot, + env: env, + day: day, + filename: "ambient.jsonl", + tokens: 100) + try Self.writeCodexSessionFile(homeRoot: managedHome, env: env, day: day, filename: "managed.jsonl", tokens: 10) + + let options = CostUsageScanner.Options(cacheRoot: env.cacheRoot) + let piOptions = PiSessionCostScanner.Options(piSessionsRoot: env.piSessionsRoot, cacheRoot: env.cacheRoot) + _ = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + codexHomePath: env.codexHomeRoot.path, + scannerOptions: options, + piScannerOptions: piOptions) + + let cachedManaged = CostUsageFetcher.loadCachedTokenSnapshot( + provider: .codex, + now: day.addingTimeInterval(120), + codexHomePath: managedHome.path, + scannerOptions: options, + piScannerOptions: piOptions) + + #expect(cachedManaged == nil) + } +} diff --git a/Tests/CodexBarTests/MenuCardModelTests.swift b/Tests/CodexBarTests/MenuCardModelTests.swift index dc823cea5c..de3d9ab0ee 100644 --- a/Tests/CodexBarTests/MenuCardModelTests.swift +++ b/Tests/CodexBarTests/MenuCardModelTests.swift @@ -233,6 +233,67 @@ struct ProviderInlineDashboardModelTests { #expect(model.tokenUsage?.monthLine.contains("$0.37") == true) } + @Test + func `local cost history shows loading dashboard while snapshot is pending`() throws { + let now = Date(timeIntervalSince1970: 1_700_179_200) + let metadata = try #require(ProviderDefaults.metadata[.codex]) + let model = UsageMenuCardView.Model.make(.init( + provider: .codex, + metadata: metadata, + snapshot: UsageSnapshot( + primary: RateWindow(usedPercent: 10, windowMinutes: nil, resetsAt: nil, resetDescription: nil), + secondary: nil, + updatedAt: now), + credits: nil, + creditsError: nil, + dashboard: nil, + dashboardError: nil, + tokenSnapshot: nil, + tokenError: nil, + tokenRefreshInFlight: true, + account: AccountInfo(email: nil, plan: nil), + isRefreshing: false, + lastError: nil, + usageBarsShowUsed: false, + resetTimeDisplayStyle: .countdown, + tokenCostUsageEnabled: true, + showOptionalCreditsAndExtraUsage: true, + hidePersonalInfo: false, + now: now)) + + #expect(model.inlineUsageDashboard == nil) + #expect(model.inlineUsageDashboardLoading == true) + #expect(model.hasUsageContent == true) + } + + @Test + func `primary cost history provider keeps error visible instead of loading forever`() throws { + let now = Date(timeIntervalSince1970: 1_700_179_200) + let metadata = try #require(ProviderDefaults.metadata[.openai]) + let model = UsageMenuCardView.Model.make(.init( + provider: .openai, + metadata: metadata, + snapshot: nil, + credits: nil, + creditsError: nil, + dashboard: nil, + dashboardError: nil, + tokenSnapshot: nil, + tokenError: nil, + account: AccountInfo(email: nil, plan: nil), + isRefreshing: false, + lastError: "Missing OpenAI API key.", + usageBarsShowUsed: false, + resetTimeDisplayStyle: .countdown, + tokenCostUsageEnabled: true, + showOptionalCreditsAndExtraUsage: true, + hidePersonalInfo: false, + now: now)) + + #expect(model.inlineUsageDashboard == nil) + #expect(model.inlineUsageDashboardLoading == false) + } + @Test func `mistral daily buckets get inline dashboard`() throws { let now = Date(timeIntervalSince1970: 1_700_179_200) diff --git a/Tests/CodexBarTests/PiSessionCostScannerTests.swift b/Tests/CodexBarTests/PiSessionCostScannerTests.swift index 1a063ce78f..92656f0064 100644 --- a/Tests/CodexBarTests/PiSessionCostScannerTests.swift +++ b/Tests/CodexBarTests/PiSessionCostScannerTests.swift @@ -227,6 +227,86 @@ struct PiSessionCostScannerTests { #expect(abs((secondReport.data.first?.costUSD ?? 0) - secondExpectedCost) < 0.000001) } + @Test + func `pi scanner cancellation keeps existing cache refreshable`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 4, day: 4) + let firstTimestamp = Int(day.timeIntervalSince1970 * 1000) + let secondTimestamp = Int(day.addingTimeInterval(60).timeIntervalSince1970 * 1000) + + let firstAssistant: [String: Any] = [ + "type": "message", + "timestamp": env.isoString(for: day), + "message": [ + "role": "assistant", + "provider": "openai-codex", + "model": "openai/gpt-5.4", + "timestamp": firstTimestamp, + "usage": [ + "input": 10, + "output": 5, + "totalTokens": 15, + ], + ], + ] + let secondAssistant: [String: Any] = [ + "type": "message", + "timestamp": env.isoString(for: day), + "message": [ + "role": "assistant", + "provider": "openai-codex", + "model": "gpt-5.4", + "timestamp": secondTimestamp, + "usage": [ + "input": 20, + "output": 10, + "totalTokens": 30, + ], + ], + ] + + let url = try env.writePiSessionFile( + relativePath: "2026-04-04T10-00-00-000Z_test.jsonl", + contents: env.jsonl([firstAssistant])) + let options = PiSessionCostScanner.Options( + piSessionsRoot: env.piSessionsRoot, + cacheRoot: env.cacheRoot, + refreshMinIntervalSeconds: 0) + let firstReport = PiSessionCostScanner.loadDailyReport( + provider: .codex, + since: day, + until: day, + now: day, + options: options) + #expect(firstReport.data.first?.totalTokens == 15) + + try env.jsonl([firstAssistant, secondAssistant]).write(to: url, atomically: true, encoding: .utf8) + + let cancelledRefresh = Task { () -> CostUsageDailyReport in + while !Task.isCancelled { + await Task.yield() + } + return PiSessionCostScanner.loadDailyReport( + provider: .codex, + since: day, + until: day, + now: day.addingTimeInterval(60), + options: options) + } + cancelledRefresh.cancel() + _ = await cancelledRefresh.value + + let refreshedReport = PiSessionCostScanner.loadDailyReport( + provider: .codex, + since: day, + until: day, + now: day.addingTimeInterval(120), + options: options) + #expect(refreshedReport.data.first?.totalTokens == 45) + } + @Test func `pi scanner ignores explicit unsupported provider even with fallback context`() throws { let env = try CostUsageTestEnvironment() diff --git a/Tests/CodexBarTests/ProviderStorageFootprintTests.swift b/Tests/CodexBarTests/ProviderStorageFootprintTests.swift index 0a903e95e1..786fda3461 100644 --- a/Tests/CodexBarTests/ProviderStorageFootprintTests.swift +++ b/Tests/CodexBarTests/ProviderStorageFootprintTests.swift @@ -384,6 +384,7 @@ struct ProviderStorageFootprintTests { browserDetection: BrowserDetection(cacheTTL: 0), settings: settings, environmentBase: ["CODEX_HOME": codexHome.path]) + store.managedCodexAccountsForStorageOverride = [] settings.providerStorageFootprintsEnabled = true store.storageRefreshGeneration = 41 store.storageRefreshInFlightSignature = "codex=\(codexHome.path)" diff --git a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift index 4ad6ce8033..c0efab7976 100644 --- a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift +++ b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift @@ -998,7 +998,7 @@ extension StatusMenuCodexSwitcherTests { @MainActor extension StatusMenuCodexSwitcherTests { @Test - func `codex account switch defers open menu rebuild until after switcher action`() async throws { + func `codex account switch updates tracked menu content before refresh completes`() async throws { self.disableMenuCardsForTesting() StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.setMenuRefreshEnabledForTesting(false) } @@ -1072,10 +1072,8 @@ extension StatusMenuCodexSwitcherTests { switcher._test_selectAccount(id: managedVisibleAccount.id) #expect(rebuildCount == 0) - for _ in 0..<20 where rebuildCount == 0 { - await Task.yield() - } - #expect(rebuildCount == 1) + #expect(controller.menuNeedsRefresh(menu) == false) + #expect(settings.codexVisibleAccountProjection.activeVisibleAccountID == managedVisibleAccount.id) await blocker.waitUntilStarted() await blocker.resume(with: .success(self.snapshot(email: "managed@example.com", percent: 17))) diff --git a/Tests/CodexBarTests/StatusMenuHighlightTests.swift b/Tests/CodexBarTests/StatusMenuHighlightTests.swift index 96fccbb561..b47209e631 100644 --- a/Tests/CodexBarTests/StatusMenuHighlightTests.swift +++ b/Tests/CodexBarTests/StatusMenuHighlightTests.swift @@ -52,4 +52,32 @@ extension StatusMenuTests { #expect(secondView.states == [true]) #expect(thirdView.states.isEmpty) } + + @Test + func `menu highlight clears tracked state on close so reopen rehighlights same row`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + + let menu = NSMenu() + let view = HighlightProbeView() + let item = NSMenuItem() + item.view = view + item.isEnabled = true + menu.addItem(item) + + controller.menu(menu, willHighlight: item) + controller.forgetClosedMenu(menu) + controller.menu(menu, willHighlight: item) + + #expect(view.states == [true, false, true]) + } } diff --git a/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift b/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift index 39d01eeaa0..d51163e365 100644 --- a/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift +++ b/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift @@ -5,7 +5,7 @@ import Testing extension StatusMenuTests { @Test - func `store observation marks open menu stale without rebuilding during tracking`() async { + func `store observation defers parent menu invalidation until tracking ends`() async { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false @@ -53,53 +53,46 @@ extension StatusMenuTests { loginMethod: "Plus Plan")), provider: .codex) - for _ in 0..<20 where controller.menuContentVersion == openedVersion { + for _ in 0..<20 where !controller.storeChangeDeferredDuringMenuTracking { await Task.yield() } - #expect(controller.menuContentVersion != openedVersion) + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow( + usedPercent: 44, + windowMinutes: 300, + resetsAt: now.addingTimeInterval(2400), + resetDescription: nil), + secondary: nil, + tertiary: nil, + updatedAt: now, + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: "codex@example.com", + accountOrganization: nil, + loginMethod: "Plus Plan")), + provider: .codex) + await Task.yield() + + controller.refreshOpenMenusIfNeeded() + + #expect(controller.menuContentVersion == openedVersion) #expect(controller.menuVersions[key] == openedVersion) + #expect(controller.storeChangeDeferredDuringMenuTracking) + #expect(controller.storeChangeDeferredObservationCount == 1) #expect(rebuildCount == 0) - } - @Test - func `explicit store actions refresh a visible open menu`() async { - self.disableMenuCardsForTesting() - let settings = self.makeSettings() - settings.statusChecksEnabled = false - settings.refreshFrequency = .manual - settings.mergeIcons = false + controller.menuDidClose(menu) + let staleVersion = controller.menuContentVersion + #expect(staleVersion != openedVersion) - let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) - let controller = StatusItemController( - store: store, - settings: settings, - account: UsageFetcher().loadAccountInfo(), - updater: DisabledUpdaterController(), - preferencesSelection: PreferencesSelection(), - statusBar: self.makeStatusBarForTesting()) - defer { controller.releaseStatusItemsForTesting() } - - let menu = controller.makeMenu() controller.menuWillOpen(menu) - let key = ObjectIdentifier(menu) - controller.openMenus[key] = menu - StatusItemController.setMenuRefreshEnabledForTesting(true) - defer { StatusItemController.resetMenuRefreshEnabledForTesting() } - - let openedVersion = controller.menuVersions[key] - - controller.refreshOpenMenusAfterExplicitStoreAction() - for _ in 0..<20 where controller.menuVersions[key] != controller.menuContentVersion { - await Task.yield() - } - - #expect(controller.menuContentVersion != openedVersion) - #expect(controller.menuVersions[key] == controller.menuContentVersion) + #expect(controller.menuVersions[key] == staleVersion) } @Test - func `repeated explicit store actions coalesce to one open menu rebuild`() async { + func `explicit store actions leave tracked parent menu stale until next open`() { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false @@ -123,70 +116,12 @@ extension StatusMenuTests { StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.resetMenuRefreshEnabledForTesting() } - var rebuildCount = 0 - controller._test_openMenuRebuildObserver = { _ in - rebuildCount += 1 - } - defer { controller._test_openMenuRebuildObserver = nil } + let openedVersion = controller.menuVersions[key] controller.refreshOpenMenusAfterExplicitStoreAction() - controller.refreshOpenMenusAfterExplicitStoreAction() - controller.refreshOpenMenusAfterExplicitStoreAction() - - for _ in 0..<20 where rebuildCount == 0 { - await Task.yield() - } - - #expect(rebuildCount == 1) - #expect(controller.menuVersions[key] == controller.menuContentVersion) - } - @Test - func `plain open menu refresh preserves pending switcher hosted submenu cleanup`() async { - self.disableMenuCardsForTesting() - let settings = self.makeSettings() - settings.statusChecksEnabled = false - settings.refreshFrequency = .manual - settings.mergeIcons = false - - let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) - let controller = StatusItemController( - store: store, - settings: settings, - account: UsageFetcher().loadAccountInfo(), - updater: DisabledUpdaterController(), - preferencesSelection: PreferencesSelection(), - statusBar: self.makeStatusBarForTesting()) - defer { controller.releaseStatusItemsForTesting() } - - let menu = controller.makeMenu() - controller.menuWillOpen(menu) - let menuKey = ObjectIdentifier(menu) - controller.openMenus[menuKey] = menu - - let submenu = controller.makeHostedSubviewPlaceholderMenu( - chartID: StatusItemController.usageBreakdownChartID, - provider: .codex) - let submenuKey = ObjectIdentifier(submenu) - controller.openMenus[submenuKey] = submenu - StatusItemController.setMenuRefreshEnabledForTesting(true) - defer { StatusItemController.resetMenuRefreshEnabledForTesting() } - - var rebuildCount = 0 - controller._test_openMenuRebuildObserver = { _ in - rebuildCount += 1 - } - defer { controller._test_openMenuRebuildObserver = nil } - - controller.deferSwitcherMenuRebuildIfStillVisible(menu, provider: .codex) - controller.refreshOpenMenuIfStillVisible(menu, provider: .codex) - - for _ in 0..<20 where rebuildCount == 0 { - await Task.yield() - } - - #expect(controller.openMenus[submenuKey] == nil) - #expect(rebuildCount == 1) - #expect(controller.menuVersions[menuKey] == controller.menuContentVersion) + #expect(controller.menuContentVersion != openedVersion) + #expect(controller.menuVersions[key] == openedVersion) + #expect(controller.menuNeedsRefresh(menu)) } } diff --git a/Tests/CodexBarTests/StatusMenuSwitcherClickTests.swift b/Tests/CodexBarTests/StatusMenuSwitcherClickTests.swift index 9fe93b78d1..2e491ad34f 100644 --- a/Tests/CodexBarTests/StatusMenuSwitcherClickTests.swift +++ b/Tests/CodexBarTests/StatusMenuSwitcherClickTests.swift @@ -130,7 +130,7 @@ struct StatusMenuSwitcherClickTests { } @Test - func `merged switcher switches provider while overview chart submenu is open`() async throws { + func `merged switcher switches provider without rebuilding tracked parent menu`() async throws { let previousMenuCardRendering = StatusItemController.menuCardRenderingEnabled let previousMenuRefresh = StatusItemController.menuRefreshEnabled StatusItemController.menuCardRenderingEnabled = false @@ -202,15 +202,13 @@ struct StatusMenuSwitcherClickTests { let switcher = try #require(menu.items.first?.view as? ProviderSwitcherView) #expect(switcher._test_simulateRuntimeClick(buttonTag: 2)) - for _ in 0..<100 where rebuildCount == 0 { - await Task.yield() - try? await Task.sleep(for: .milliseconds(10)) - } + await Task.yield() #expect(settings.mergedMenuLastSelectedWasOverview == false) #expect(settings.selectedMenuProvider == .claude) - #expect(rebuildCount == 1) + #expect(rebuildCount == 0) #expect(controller.openMenus[ObjectIdentifier(submenu)] == nil) + #expect(controller.menuNeedsRefresh(menu) == false) let ids = menu.items.compactMap { $0.representedObject as? String } #expect(ids.contains("menuCard")) @@ -525,16 +523,39 @@ struct StatusMenuSwitcherClickTests { }, onSelect: { _ in }) #expect(view._test_quotaIndicatorFillRatios().count == 2) - let noQuotaHeight = try #require(noQuotaView._test_buttonFittingSizes().last?.height) - let quotaHeight = try #require(view._test_buttonFittingSizes().last?.height) + let noQuotaHeight = try #require(noQuotaView._test_buttonIntrinsicSizes().last?.height) + let quotaHeight = try #require(view._test_buttonIntrinsicSizes().last?.height) #expect(quotaHeight > noQuotaHeight) grokRemaining = nil view.updateQuotaIndicators() #expect(view._test_quotaIndicatorFillRatios().count == 1) - let removedQuotaHeight = try #require(view._test_buttonFittingSizes().last?.height) - #expect(removedQuotaHeight == noQuotaHeight) + let removedQuotaHeight = try #require(view._test_buttonIntrinsicSizes().last?.height) + #expect(removedQuotaHeight == quotaHeight) + } + + @Test + func `overview tab reserves quota bar space beside provider tabs`() throws { + let view = ProviderSwitcherView( + providers: [.codex, .claude, .gemini], + selected: .overview, + includesOverview: true, + width: 560, + showsIcons: true, + iconProvider: { _ in NSImage(size: NSSize(width: 16, height: 16)) }, + weeklyRemainingProvider: { _ in 50 }, + onSelect: { _ in }) + view.updateConstraintsForSubtreeIfNeeded() + view.layoutSubtreeIfNeeded() + + let frames = view._test_buttonFrames() + #expect(frames.count == 4) + let rowHeight = view._test_rowHeight() + let overviewHeight = try #require(frames.first?.height) + let providerHeights = frames.dropFirst().map(\.height) + #expect(overviewHeight == rowHeight) + #expect(providerHeights.allSatisfy { $0 == rowHeight }) } @Test @@ -559,8 +580,8 @@ struct StatusMenuSwitcherClickTests { weeklyRemainingProvider: { _ in 50 }, onSelect: { _ in }) - let withoutQuotaHeight = try #require(textOnlyWithoutQuota._test_buttonFittingSizes().first?.height) - let withQuotaHeight = try #require(textOnlyWithQuota._test_buttonFittingSizes().first?.height) + let withoutQuotaHeight = try #require(textOnlyWithoutQuota._test_buttonIntrinsicSizes().first?.height) + let withQuotaHeight = try #require(textOnlyWithQuota._test_buttonIntrinsicSizes().first?.height) #expect(withQuotaHeight > withoutQuotaHeight) } diff --git a/Tests/CodexBarTests/StatusMenuSwitcherRefreshTests.swift b/Tests/CodexBarTests/StatusMenuSwitcherRefreshTests.swift index 4d71f58e1e..9722cc04d7 100644 --- a/Tests/CodexBarTests/StatusMenuSwitcherRefreshTests.swift +++ b/Tests/CodexBarTests/StatusMenuSwitcherRefreshTests.swift @@ -7,7 +7,7 @@ import Testing @Suite(.serialized) struct StatusMenuSwitcherRefreshTests { @Test - func `merged provider switch rebuilds stale width switcher rows`() async throws { + func `merged provider switch updates tracked parent body from cache`() async throws { let previousMenuCardRendering = StatusItemController.menuCardRenderingEnabled StatusItemController.menuCardRenderingEnabled = false StatusItemController.setMenuRefreshEnabledForTesting(true) @@ -47,8 +47,11 @@ struct StatusMenuSwitcherRefreshTests { defer { controller.releaseStatusItemsForTesting() } let menu = controller.makeMenu() + controller.mergedMenu = menu controller.menuWillOpen(menu) - #expect(controller.openMenus[ObjectIdentifier(menu)] === menu) + let key = ObjectIdentifier(menu) + #expect(controller.openMenus[key] === menu) + let openedVersion = controller.menuVersions[key] let initialSwitcher = try #require(menu.items.first?.view as? ProviderSwitcherView) initialSwitcher.frame.size.width = 250 @@ -62,17 +65,140 @@ struct StatusMenuSwitcherRefreshTests { let nextProviderButton = try #require(Self.switcherButtons(in: menu).first { $0.state == .off }) #expect(initialSwitcher._test_simulateRuntimeClick(buttonTag: nextProviderButton.tag) == true) - for _ in 0..<100 where rebuildCount == 0 { - await Task.yield() - try? await Task.sleep(for: .milliseconds(10)) - } + await Task.yield() - #expect(rebuildCount == 1) + #expect(rebuildCount == 0) + #expect(controller.menuVersions[key] == openedVersion) + #expect(controller.menuNeedsRefresh(menu) == false) let updatedSwitcher = try #require(menu.items.first?.view as? ProviderSwitcherView) - #expect(updatedSwitcher.frame.width == 310) + #expect(updatedSwitcher.frame.width == 250) #expect(Self.switcherButtons(in: menu).first { $0.tag == nextProviderButton.tag }?.state == .on) } + @Test + func `stale cached merged parent menu hydrates after opening`() async throws { + let previousMenuCardRendering = StatusItemController.menuCardRenderingEnabled + StatusItemController.menuCardRenderingEnabled = false + StatusItemController.setMenuRefreshEnabledForTesting(true) + defer { + StatusItemController.menuCardRenderingEnabled = previousMenuCardRendering + StatusItemController.resetMenuRefreshEnabledForTesting() + } + + let settings = Self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + settings.selectedMenuProvider = .codex + Self.enableCodexAndClaude(settings) + + let activeProviders: [UsageProvider] = [.codex, .claude] + _ = settings.setMergedOverviewProviderSelection( + provider: .codex, + isSelected: false, + activeProviders: activeProviders) + _ = settings.setMergedOverviewProviderSelection( + provider: .claude, + isSelected: false, + activeProviders: activeProviders) + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: .system) + defer { controller.releaseStatusItemsForTesting() } + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.menuWillOpen(menu) + let key = ObjectIdentifier(menu) + let openedVersion = try #require(controller.menuVersions[key]) + controller.menuDidClose(menu) + + controller.menuContentVersion &+= 1 + let staleVersion = controller.menuContentVersion + var updateCount = 0 + controller._test_openMenuRefreshYieldOverride = {} + controller._test_openMenuRebuildObserver = { _ in + updateCount += 1 + } + defer { + controller._test_openMenuRefreshYieldOverride = nil + controller._test_openMenuRebuildObserver = nil + } + + controller.menuWillOpen(menu) + #expect(controller.menuVersions[key] == openedVersion) + + for _ in 0..<10 where updateCount == 0 { + await Task.yield() + } + + #expect(updateCount == 1) + #expect(controller.menuVersions[key] == staleVersion) + #expect(controller.menuNeedsRefresh(menu) == false) + } + + @Test + func `stale merged parent menu rebuilds immediately when provider list changes`() throws { + let previousMenuCardRendering = StatusItemController.menuCardRenderingEnabled + StatusItemController.menuCardRenderingEnabled = false + StatusItemController.setMenuRefreshEnabledForTesting(true) + defer { + StatusItemController.menuCardRenderingEnabled = previousMenuCardRendering + StatusItemController.resetMenuRefreshEnabledForTesting() + } + + let settings = Self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + settings.selectedMenuProvider = .codex + Self.enableCodexAndClaude(settings) + + let activeProviders: [UsageProvider] = [.codex, .claude] + _ = settings.setMergedOverviewProviderSelection( + provider: .codex, + isSelected: false, + activeProviders: activeProviders) + _ = settings.setMergedOverviewProviderSelection( + provider: .claude, + isSelected: false, + activeProviders: activeProviders) + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: .system) + defer { controller.releaseStatusItemsForTesting() } + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.menuWillOpen(menu) + let key = ObjectIdentifier(menu) + controller.menuDidClose(menu) + + let geminiMetadata = try #require(ProviderRegistry.shared.metadata[.gemini]) + settings.setProviderEnabled(provider: .gemini, metadata: geminiMetadata, enabled: true) + controller.menuContentVersion &+= 1 + let changedVersion = controller.menuContentVersion + + controller.menuWillOpen(menu) + + #expect(controller.menuVersions[key] == changedVersion) + #expect(controller.menuNeedsRefresh(menu) == false) + } + private static func makeSettings() -> SettingsStore { let suite = "StatusMenuSwitcherRefreshTests-\(UUID().uuidString)" let defaults = UserDefaults(suiteName: suite)! diff --git a/Tests/CodexBarTests/StatusMenuTests.swift b/Tests/CodexBarTests/StatusMenuTests.swift index cd4e5e9130..0910f75414 100644 --- a/Tests/CodexBarTests/StatusMenuTests.swift +++ b/Tests/CodexBarTests/StatusMenuTests.swift @@ -277,64 +277,6 @@ struct StatusMenuTests { #expect(controller.closeOpenMenusFromShortcutIfNeeded() == false) } - @Test - func `open menu defers store data refresh until next open`() async { - self.disableMenuCardsForTesting() - let settings = self.makeSettings() - settings.statusChecksEnabled = false - settings.refreshFrequency = .manual - settings.mergeIcons = false - - let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) - let controller = StatusItemController( - store: store, - settings: settings, - account: UsageFetcher().loadAccountInfo(), - updater: DisabledUpdaterController(), - preferencesSelection: PreferencesSelection(), - statusBar: self.makeStatusBarForTesting()) - - let menu = controller.makeMenu() - controller.menuWillOpen(menu) - let key = ObjectIdentifier(menu) - controller.openMenus[key] = menu - StatusItemController.setMenuRefreshEnabledForTesting(true) - defer { StatusItemController.resetMenuRefreshEnabledForTesting() } - let openedVersion = controller.menuVersions[key] - - let now = Date() - store._setSnapshotForTesting( - UsageSnapshot( - primary: RateWindow( - usedPercent: 11, - windowMinutes: 300, - resetsAt: now.addingTimeInterval(1800), - resetDescription: nil), - secondary: nil, - tertiary: nil, - updatedAt: now, - identity: ProviderIdentitySnapshot( - providerID: .codex, - accountEmail: "codex@example.com", - accountOrganization: nil, - loginMethod: "Plus Plan")), - provider: .codex) - - for _ in 0..<50 where controller.menuContentVersion == openedVersion { - await Task.yield() - } - - let staleVersion = controller.menuContentVersion - controller.refreshOpenMenusIfNeeded() - - #expect(controller.menuContentVersion != openedVersion) - #expect(controller.menuVersions[key] == openedVersion) - - controller.menuDidClose(menu) - controller.menuWillOpen(menu) - #expect(controller.menuVersions[key] == staleVersion) - } - @Test func `merged menu refresh uses resolved enabled provider when selection is cleared`() { self.disableMenuCardsForTesting() @@ -440,6 +382,45 @@ struct StatusMenuTests { #expect(delayedRefreshWakeCount == 0) } + @Test + func `menu open defers automatic provider refresh until tracking ends`() async { + StatusItemController.menuCardRenderingEnabled = false + StatusItemController.setMenuRefreshEnabledForTesting(true) + StatusItemController.setMenuOpenRefreshDelayForTesting(.milliseconds(50)) + defer { + StatusItemController.resetMenuOpenRefreshDelayForTesting() + StatusItemController.resetMenuRefreshEnabledForTesting() + } + + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + var refreshCount = 0 + store._test_providerRefreshOverride = { _ in + refreshCount += 1 + } + + await withStatusItemControllerForTesting( + store: store, + settings: settings, + fetcher: fetcher, + statusBar: self.makeStatusBarForTesting()) + { controller in + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + try? await Task.sleep(for: .milliseconds(180)) + #expect(refreshCount == 0) + #expect(controller.deferredMenuInteractionRefreshPending) + controller.menuDidClose(menu) + controller.deferredMenuInteractionRefreshPending = false + controller.deferredMenuInteractionRefreshTask?.cancel() + controller.deferredMenuInteractionRefreshTask = nil + } + } + @Test func `login state callbacks do not attach menus after release`() { self.disableMenuCardsForTesting() @@ -521,7 +502,7 @@ struct StatusMenuTests { } @Test - func `open merged menu rebuilds switcher when usage bars mode changes`() async { + func `open merged menu marks switcher stale when usage bars mode changes`() { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false @@ -552,9 +533,11 @@ struct StatusMenuTests { let menu = controller.makeMenu() controller.menuWillOpen(menu) - controller.openMenus[ObjectIdentifier(menu)] = menu + let key = ObjectIdentifier(menu) + controller.openMenus[key] = menu StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.resetMenuRefreshEnabledForTesting() } + let openedVersion = controller.menuVersions[key] let initialSwitcher = menu.items.first?.view as? ProviderSwitcherView #expect(initialSwitcher != nil) @@ -562,17 +545,15 @@ struct StatusMenuTests { settings.usageBarsShowUsed = true controller.handleProviderConfigChange(reason: "usageBarsShowUsed") - for _ in 0..<20 - where initialSwitcherID == (menu.items.first?.view as? ProviderSwitcherView).map(ObjectIdentifier.init) - { - await Task.yield() - } let updatedSwitcher = menu.items.first?.view as? ProviderSwitcherView #expect(updatedSwitcher != nil) if let initialSwitcherID, let updatedSwitcher { - #expect(initialSwitcherID != ObjectIdentifier(updatedSwitcher)) + #expect(initialSwitcherID == ObjectIdentifier(updatedSwitcher)) } + #expect(controller.menuContentVersion != openedVersion) + #expect(controller.menuVersions[key] == openedVersion) + #expect(controller.menuNeedsRefresh(menu)) } @Test diff --git a/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift b/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift index 49053cc937..89ccc3d3a2 100644 --- a/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift +++ b/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift @@ -293,7 +293,7 @@ final class StatusMenuTokenAccountSwitcherTests: XCTestCase { ["menuCard-0", "menuCard-1", "menuCard-2", "menuCard-3", "menuCard-4", "menuCard-5"]) } - func test_tokenAccountSwitchDefersOpenMenuRebuildUntilAfterSwitcherAction() async throws { + func test_tokenAccountSwitchUpdatesTrackedMenuContentBeforeRefreshCompletes() async throws { self.disableMenuCardsForTesting() StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.setMenuRefreshEnabledForTesting(false) } @@ -342,10 +342,8 @@ final class StatusMenuTokenAccountSwitcherTests: XCTestCase { let selectionTask = try XCTUnwrap(switcher._test_select(index: 1)) XCTAssertEqual(rebuildCount, 0) - for _ in 0..<20 where rebuildCount == 0 { - await Task.yield() - } - XCTAssertEqual(rebuildCount, 1) + XCTAssertFalse(controller.menuNeedsRefresh(menu)) + XCTAssertEqual(settings.tokenAccountsData(for: .claude)?.clampedActiveIndex(), 1) await blocker.waitUntilStarted(count: 1) await blocker.resumeAll(with: .success(self.snapshot(percent: 17))) diff --git a/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift b/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift index 014ce6716f..a89a937e1c 100644 --- a/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift +++ b/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift @@ -67,6 +67,23 @@ private actor TokenRefreshRecorder { func record(provider: UsageProvider, force: Bool) { self.calls.append((provider, force)) } + + func waitUntilCalls(count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let start = ContinuousClock.now + while self.calls.count < count { + if start.duration(to: .now) >= timeout { return false } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } +} + +private actor ProviderRefreshRecorder { + private(set) var calls: [UsageProvider] = [] + + func record(provider: UsageProvider) { + self.calls.append(provider) + } } @MainActor @@ -172,6 +189,180 @@ struct UsageStoreManualTokenRefreshTests { } } + @Test + func `menu interaction defers regular token-cost refresh until idle`() async { + let store = Self.makeStore() + let recorder = TokenRefreshRecorder() + store.tokenCostInteractionResumeDelay = 0.01 + store._test_providerRefreshOverride = { _ in } + store._test_tokenUsageRefreshOverride = { provider, force in + await recorder.record(provider: provider, force: force) + } + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + await store.refresh(forceTokenUsage: false) + try? await Task.sleep(for: .milliseconds(50)) + #expect(await recorder.calls.isEmpty) + + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + #expect(await recorder.waitUntilCalls(count: 1)) + + #expect(await recorder.calls.map(\.provider) == [.codex]) + #expect(await recorder.calls.map(\.force) == [false]) + } + + @Test + func `menu interaction starts missing token-cost snapshot scan so skeleton can hydrate`() async { + let store = Self.makeStore() + let recorder = TokenRefreshRecorder() + store.tokenCostInteractionResumeDelay = 0.01 + store._test_cachedTokenUsageLoadOverride = { _ in nil } + store._test_tokenUsageRefreshOverride = { provider, force in + await recorder.record(provider: provider, force: force) + } + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + store.ensureTokenCostSnapshotScheduled(for: .codex, reason: "test-missing-chart") + + #expect(await recorder.waitUntilCalls(count: 1)) + #expect(await recorder.calls.map(\.provider) == [.codex]) + #expect(await recorder.calls.map(\.force) == [false]) + + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + #expect(await recorder.calls.map(\.provider) == [.codex]) + #expect(await recorder.calls.map(\.force) == [false]) + } + + @Test + func `menu interaction hydrates cached token-cost snapshot before revalidation`() async { + let store = Self.makeStore() + let recorder = TokenRefreshRecorder() + store.tokenCostInteractionResumeDelay = 0.01 + store._test_cachedTokenUsageLoadOverride = { _ in Self.tokenSnapshot() } + store._test_tokenUsageRefreshOverride = { provider, force in + await recorder.record(provider: provider, force: force) + } + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + store.ensureTokenCostSnapshotScheduled(for: .codex, reason: "test-cached-chart") + + #expect(await Self.waitUntil(timeout: .seconds(1)) { + store.tokenSnapshot(for: .codex)?.sessionTokens == 150 + }) + #expect(await recorder.calls.isEmpty) + + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + #expect(await recorder.waitUntilCalls(count: 1)) + #expect(await recorder.calls.map(\.provider) == [.codex]) + #expect(await recorder.calls.map(\.force) == [false]) + } + + @Test + func `menu interaction defers stale token-cost refresh when cached data exists`() async { + let store = Self.makeStore() + let recorder = TokenRefreshRecorder() + store.tokenCostInteractionResumeDelay = 0.01 + store._setTokenSnapshotForTesting(Self.tokenSnapshot(), provider: .codex) + store._test_tokenUsageRefreshOverride = { provider, force in + await recorder.record(provider: provider, force: force) + } + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + store.scheduleTokenRefresh(force: false, reason: "test-stale-chart") + try? await Task.sleep(for: .milliseconds(50)) + #expect(await recorder.calls.isEmpty) + + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + #expect(await recorder.waitUntilCalls(count: 1)) + #expect(await recorder.calls.map(\.provider) == [.codex]) + #expect(await recorder.calls.map(\.force) == [false]) + } + + @Test + func `menu interaction cancels token refresh even when unrelated provider is queued`() { + let store = Self.makeStore() + let task = Task { + while !Task.isCancelled { + try? await Task.sleep(for: .milliseconds(20)) + } + } + store.tokenRefreshSequenceTask = task + store.tokenRefreshQueuedProviders = [.gemini] + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + + #expect(task.isCancelled) + store.tokenRefreshSequenceTask = nil + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + } + + @Test + func `menu interaction keeps protected missing chart refresh running`() { + let store = Self.makeStore() + let task = Task { + while !Task.isCancelled { + try? await Task.sleep(for: .milliseconds(20)) + } + } + store.tokenRefreshSequenceTask = task + store.tokenRefreshQueuedProviders = [.codex] + store.tokenRefreshMenuAllowedProviders = [.codex] + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + + #expect(!task.isCancelled) + task.cancel() + store.tokenRefreshSequenceTask = nil + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + } + + @Test + func `menu interaction cancels stale in-flight refresh even with protected pending work`() { + let store = Self.makeStore() + let task = Task { + while !Task.isCancelled { + try? await Task.sleep(for: .milliseconds(20)) + } + } + store.tokenRefreshSequenceTask = task + store.tokenRefreshQueuedProviders = [.claude, .codex] + store.tokenRefreshMenuAllowedProviders = [.codex] + store.tokenRefreshInFlightStartedAt[.claude] = Date() + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + + #expect(task.isCancelled) + store.tokenRefreshSequenceTask = nil + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + } + + @Test + func `menu interaction refreshes providers immediately but defers forced token-cost refresh until idle`() async { + let store = Self.makeStore() + let providerRecorder = ProviderRefreshRecorder() + let tokenRecorder = TokenRefreshRecorder() + store.tokenCostInteractionResumeDelay = 0.01 + store._test_providerRefreshOverride = { provider in + await providerRecorder.record(provider: provider) + } + store._test_tokenUsageRefreshOverride = { provider, force in + await tokenRecorder.record(provider: provider, force: force) + } + + store.beginInteractiveMenuTokenCostDeferral(reason: "test-menu-open") + await store.refresh(forceTokenUsage: true) + + #expect(await providerRecorder.calls == [.codex]) + #expect(await tokenRecorder.calls.isEmpty) + + store.endInteractiveMenuTokenCostDeferral(reason: "test-menu-close") + #expect(await tokenRecorder.waitUntilCalls(count: 1)) + + #expect(await providerRecorder.calls == [.codex]) + #expect(await tokenRecorder.calls.map(\.provider) == [.codex]) + #expect(await tokenRecorder.calls.map(\.force) == [true]) + } + private static func makeStore() -> UsageStore { let suite = "UsageStoreManualTokenRefreshTests-\(UUID().uuidString)" let defaults = UserDefaults(suiteName: suite)! @@ -201,4 +392,37 @@ struct UsageStoreManualTokenRefreshTests { startupBehavior: .testing, environmentBase: [:]) } + + private static func tokenSnapshot() -> CostUsageTokenSnapshot { + CostUsageTokenSnapshot( + sessionTokens: 150, + sessionCostUSD: 0.12, + last30DaysTokens: 150, + last30DaysCostUSD: 0.12, + daily: [ + CostUsageDailyReport.Entry( + date: "2026-05-28", + inputTokens: 100, + outputTokens: 50, + totalTokens: 150, + costUSD: 0.12, + modelsUsed: ["gpt-5"], + modelBreakdowns: [ + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5", + costUSD: 0.12, + totalTokens: 150), + ]), + ], + updatedAt: Date()) + } + + private static func waitUntil(timeout: Duration, predicate: @escaping @MainActor () -> Bool) async -> Bool { + let start = ContinuousClock.now + while !predicate() { + if start.duration(to: .now) >= timeout { return false } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } } diff --git a/Tests/CodexBarTests/UsageStorePlanUtilizationTests.swift b/Tests/CodexBarTests/UsageStorePlanUtilizationTests.swift index aa3207f6dd..5c2b7d2cf0 100644 --- a/Tests/CodexBarTests/UsageStorePlanUtilizationTests.swift +++ b/Tests/CodexBarTests/UsageStorePlanUtilizationTests.swift @@ -539,6 +539,35 @@ struct UsageStorePlanUtilizationTests { #expect(bobHistory == [bobWeekly]) } + @MainActor + @Test + func `plan history menu lookup does not mutate account buckets`() throws { + let store = Self.makeStore() + let aliceSnapshot = Self.makeSnapshot(provider: .codex, email: "alice@example.com") + let aliceKey = try #require( + UsageStore._planUtilizationAccountKeyForTesting( + provider: .codex, + snapshot: aliceSnapshot)) + let bootstrap = planSeries(name: .session, windowMinutes: 300, entries: [ + planEntry(at: Date(timeIntervalSince1970: 1_699_913_600), usedPercent: 90), + ]) + let aliceWeekly = planSeries(name: .weekly, windowMinutes: 10080, entries: [ + planEntry(at: Date(timeIntervalSince1970: 1_700_000_000), usedPercent: 20), + ]) + let buckets = PlanUtilizationHistoryBuckets( + unscoped: [bootstrap], + accounts: [ + aliceKey: [aliceWeekly], + ]) + store.planUtilizationHistory[.codex] = buckets + store._setSnapshotForTesting(aliceSnapshot, provider: .codex) + + let history = store.planUtilizationHistoryForMenu(for: .codex) + + #expect(history == [aliceWeekly]) + #expect(store.planUtilizationHistory[.codex] == buckets) + } + @MainActor @Test func `plan utilization menu hides while refreshing without current snapshot`() throws {