From ce4a110c0cbc9bb2473b4eda9206d4f49f7355ab Mon Sep 17 00:00:00 2001 From: joshuavial Date: Thu, 11 Jun 2026 03:19:20 +0800 Subject: [PATCH 1/2] Fix menu open refresh delay --- Sources/CodexBar/MenuCardView.swift | 2 +- .../Kilo/UsageStore+KiloOrgRefresh.swift | 3 +- .../StatusItemController+Actions.swift | 3 +- .../CodexBar/StatusItemController+Menu.swift | 55 +- .../StatusItemController+MenuCardModel.swift | 2 +- ...temController+MenuInteractionRefresh.swift | 14 +- ...ItemController+MenuRefreshScheduling.swift | 5 + .../StatusItemController+MenuTracking.swift | 125 +- Sources/CodexBar/StatusItemController.swift | 13 +- .../CodexBar/UsageStore+PlanUtilization.swift | 5 + Sources/CodexBar/UsageStore+Refresh.swift | 213 ++- .../CodexBar/UsageStore+TokenAccounts.swift | 53 +- Sources/CodexBar/UsageStore.swift | 18 +- Sources/CodexBar/UsageStoreSupport.swift | 72 + .../CodexBarTests/MenuCardSubtitleTests.swift | 45 + .../StatusMenuClosedPreparationTests.swift | 8 + .../StatusMenuInstantOpenTests.swift | 1540 +++++++++++++++++ .../StatusMenuOpenRefreshTests.swift | 125 +- Tests/CodexBarTests/StatusMenuTests.swift | 13 +- .../StatusMenuTokenAccountSwitcherTests.swift | 7 +- 20 files changed, 2144 insertions(+), 177 deletions(-) create mode 100644 Tests/CodexBarTests/StatusMenuInstantOpenTests.swift diff --git a/Sources/CodexBar/MenuCardView.swift b/Sources/CodexBar/MenuCardView.swift index 89c3d2204e..1c3532da14 100644 --- a/Sources/CodexBar/MenuCardView.swift +++ b/Sources/CodexBar/MenuCardView.swift @@ -1047,7 +1047,7 @@ extension UsageMenuCardView.Model { return (lastError.trimmingCharacters(in: .whitespacesAndNewlines), .error) } - if isRefreshing, snapshot == nil { + if isRefreshing { return ("\(L("Refreshing"))…", .loading) } diff --git a/Sources/CodexBar/Providers/Kilo/UsageStore+KiloOrgRefresh.swift b/Sources/CodexBar/Providers/Kilo/UsageStore+KiloOrgRefresh.swift index ed1f7129e2..59e76dc722 100644 --- a/Sources/CodexBar/Providers/Kilo/UsageStore+KiloOrgRefresh.swift +++ b/Sources/CodexBar/Providers/Kilo/UsageStore+KiloOrgRefresh.swift @@ -35,7 +35,7 @@ extension UsageStore { self.kiloEnabledScopes.count > 1 } - func refreshKiloScopes() async { + func refreshKiloScopes(generation: UInt64? = nil) async { let scopes = self.kiloEnabledScopes guard scopes.count > 1 else { await MainActor.run { self.kiloScopeSnapshots = [] } @@ -102,6 +102,7 @@ extension UsageStore { let ordered = scopes.compactMap { resultByID[$0.scopeIdentifier] } await MainActor.run { + guard self.isCurrentProviderRefreshGeneration(.kilo, generation: generation) else { return } self.kiloScopeSnapshots = ordered } } diff --git a/Sources/CodexBar/StatusItemController+Actions.swift b/Sources/CodexBar/StatusItemController+Actions.swift index 9ce0d39ddd..d0607c57c8 100644 --- a/Sources/CodexBar/StatusItemController+Actions.swift +++ b/Sources/CodexBar/StatusItemController+Actions.swift @@ -44,8 +44,7 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { func refreshOpenMenusAfterExplicitStoreAction() { self.invalidateMenus( refreshOpenMenus: true, - deferOpenParentMenuRebuild: true, - allowStaleContentDuringDataRefresh: true) + deferOpenParentMenuRebuild: true) } @objc func refreshNow() { diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 4740f03360..9dfebcf3b9 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -1135,8 +1135,11 @@ extension StatusItemController { // provider fetch failed and needs a retry; periodic freshness is handled by the refresh timer. // AppKit menu tracking is modal, so starting provider refreshes while it is active can make the menu // feel frozen and can block keyboard focus from returning. - if self.menuNeedsDelayedRefreshRetry(for: menu) { - self.deferMenuInteractionRefreshIfNeeded() + let providersNeedingRetry = self.delayedRefreshRetryProviders(for: menu).filter { + self.store.isStale(provider: $0) || self.store.snapshot(for: $0) == nil + } + if !providersNeedingRetry.isEmpty { + self.deferMenuInteractionRefreshIfNeeded(providers: providersNeedingRetry) } let key = ObjectIdentifier(menu) self.menuRefreshTasks[key]?.cancel() @@ -1149,13 +1152,43 @@ extension StatusItemController { self.onDelayedMenuRefreshAttemptForTesting?() #endif guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } - guard !self.store.isRefreshing else { return } - let retryProviders = self.delayedRefreshRetryProviders(for: menu) - let retryStaleProviderCount = retryProviders.count { self.store.isStale(provider: $0) } - let retryMissingSnapshotCount = retryProviders.count { self.store.snapshot(for: $0) == nil } - let willRetryRefresh = retryStaleProviderCount > 0 || retryMissingSnapshotCount > 0 - guard willRetryRefresh else { return } - self.deferMenuInteractionRefreshIfNeeded() + let availableProviders = Set(self.store.enabledProvidersForBackgroundWork()) + let retryProviders = self.delayedRefreshRetryProviders(for: menu).filter { + availableProviders.contains($0) && + (self.store.refreshingProviders.contains($0) || + self.store.isStale(provider: $0) || + self.store.snapshot(for: $0) == nil) + } + guard !retryProviders.isEmpty else { + self.clearSatisfiedDeferredMenuInteractionRefreshes( + for: self.delayedRefreshRetryProviders(for: menu)) + if self.menuNeedsRefresh(menu) { + self.scheduleOpenMenuRebuildIfStillVisible( + menu, + provider: self.menuProvider(for: menu), + resyncReadinessBaselineAfterRebuild: self.openMenus.count == 1) + } + return + } + self.deferMenuInteractionRefreshIfNeeded(providers: retryProviders) + await ProviderInteractionContext.$current.withValue(.background) { + for provider in retryProviders { + guard !Task.isCancelled else { return } + await self.store.refreshProvider(provider, coalesceIfRefreshing: true) + } + } + let stillNeedsRetry = retryProviders.contains { + self.store.isStale(provider: $0) || self.store.snapshot(for: $0) == nil + } + if !stillNeedsRetry { + self.clearSatisfiedDeferredMenuInteractionRefreshes(for: retryProviders) + } + guard !Task.isCancelled else { return } + guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } + self.invalidateMenus( + refreshOpenMenus: true, + deferOpenParentMenuRebuild: false, + allowStaleContentDuringDataRefresh: true) } } @@ -1168,6 +1201,10 @@ extension StatusItemController { } private func delayedRefreshRetryProviders(for menu: NSMenu) -> [UsageProvider] { + self.renderedProviders(for: menu) + } + + func renderedProviders(for menu: NSMenu) -> [UsageProvider] { let enabledProviders = self.store.enabledProvidersForDisplay() guard !enabledProviders.isEmpty else { return [] } let includesOverview = self.includesOverviewTab(enabledProviders: enabledProviders) diff --git a/Sources/CodexBar/StatusItemController+MenuCardModel.swift b/Sources/CodexBar/StatusItemController+MenuCardModel.swift index 7c17a65593..eebc4b830f 100644 --- a/Sources/CodexBar/StatusItemController+MenuCardModel.swift +++ b/Sources/CodexBar/StatusItemController+MenuCardModel.swift @@ -101,7 +101,7 @@ extension StatusItemController { tokenSnapshot: tokenSnapshot, tokenError: tokenError, account: fallbackAccount, - isRefreshing: self.store.shouldShowRefreshingMenuCard(for: target), + isRefreshing: self.store.shouldShowRefreshingMenuCardIndicator(for: target), lastError: errorOverride ?? codexProjection?.userFacingErrors.usage ?? self.store.userFacingError(for: target), diff --git a/Sources/CodexBar/StatusItemController+MenuInteractionRefresh.swift b/Sources/CodexBar/StatusItemController+MenuInteractionRefresh.swift index 6c27e63671..5cfcf76560 100644 --- a/Sources/CodexBar/StatusItemController+MenuInteractionRefresh.swift +++ b/Sources/CodexBar/StatusItemController+MenuInteractionRefresh.swift @@ -58,9 +58,17 @@ extension StatusItemController { ]) } - func deferMenuInteractionRefreshIfNeeded() { + func deferMenuInteractionRefreshIfNeeded(providers: [UsageProvider]) { guard !self.store.isRefreshing else { return } - self.deferredMenuInteractionRefreshPending = true + self.deferredMenuInteractionRefreshProviders.formUnion(providers) + } + + func clearSatisfiedDeferredMenuInteractionRefreshes(for providers: [UsageProvider]) { + for provider in providers + where !self.store.isStale(provider: provider) && self.store.snapshot(for: provider) != nil + { + self.deferredMenuInteractionRefreshProviders.remove(provider) + } } func deferOpenAIDashboardRefreshUntilMenuCloses(reason: String) { @@ -110,7 +118,7 @@ extension StatusItemController { return } self.deferredMenuInteractionRefreshTask = nil - self.deferredMenuInteractionRefreshPending = false + self.deferredMenuInteractionRefreshProviders.removeAll() self.deferredOpenAIDashboardRefreshReason = nil #if DEBUG self.onDeferredMenuInteractionRefreshForTesting?() diff --git a/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift b/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift index 4980a019be..a6b21ed34f 100644 --- a/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift +++ b/Sources/CodexBar/StatusItemController+MenuRefreshScheduling.swift @@ -120,6 +120,7 @@ extension StatusItemController { [ provider.rawValue, "token=\(tokenSignature)", + "refreshing=\(self.store.shouldShowRefreshingMenuCardIndicator(for: provider) ? "1" : "0")", "usageHistory=\(usageHistoryVisible ? "1" : "0")", ].joined(separator: ":")) } @@ -214,6 +215,7 @@ extension StatusItemController { _ menu: NSMenu, provider: UsageProvider?, closeHostedSubviewMenusBeforeRebuild: Bool = false, + resyncReadinessBaselineAfterRebuild: Bool = false, debounceNanoseconds: UInt64 = 0, beforeRebuild: (@MainActor () -> Bool)? = nil) { @@ -255,6 +257,9 @@ extension StatusItemController { self.closeHostedSubviewMenusForParentSwitch() } self.rebuildOpenMenuIfStillVisible(menu, provider: provider) + if resyncReadinessBaselineAfterRebuild, !self.menuNeedsRefresh(menu) { + self.resyncMenuAdjunctReadinessBaseline() + } } } diff --git a/Sources/CodexBar/StatusItemController+MenuTracking.swift b/Sources/CodexBar/StatusItemController+MenuTracking.swift index 9d437f01bc..037edcdd86 100644 --- a/Sources/CodexBar/StatusItemController+MenuTracking.swift +++ b/Sources/CodexBar/StatusItemController+MenuTracking.swift @@ -37,6 +37,11 @@ extension StatusItemController { self.clearMergedSwitcherContentCaches() } self.pruneVersionScopedMenuCardHeightCache() + if allowStaleContentDuringDataRefresh { + self.latestDataOnlyMenuContentVersion = self.menuContentVersion + } else { + self.latestStructuralMenuContentVersion = self.menuContentVersion + } if !allowStaleContentDuringDataRefresh, !preservesMergedSwitcherContentCaches { self.latestRequiredMenuRebuildVersion = self.menuContentVersion } @@ -49,31 +54,56 @@ extension StatusItemController { deferParentRebuildDuringTracking: deferOpenParentMenuRebuild) return } + if allowStaleContentDuringDataRefresh { + if !self.cancelNonRequiredClosedMenuPreparation() { + self.prepareAttachedClosedMenusIfNeeded() + } + return + } self.prepareAttachedClosedMenusIfNeeded() } + @discardableResult + private func cancelNonRequiredClosedMenuPreparation() -> Bool { + let menus = self.attachedMenusForClosedPreparation() + let hasRequiredClosedMenu = self.latestRequiredMenuRebuildVersion > 0 && menus.contains { menu in + let key = ObjectIdentifier(menu) + return (self.menuVersions[key] ?? -1) < self.latestRequiredMenuRebuildVersion + } + guard !hasRequiredClosedMenu else { return false } + self.cancelAllClosedMenuRebuilds() + for menu in menus { + self.closedMenusDeferredUntilNextOpen.remove(ObjectIdentifier(menu)) + } + return true + } + func prepareAttachedClosedMenusIfNeeded() { guard self.isMenuRefreshEnabled else { return } guard self.openMenus.isEmpty else { return } guard !self.isMenuDataRefreshInFlight else { return } let menus = self.attachedMenusForClosedPreparation() let requiredClosedPreparationVersion: Int? - if self.menuContentVersion > self.latestRequiredMenuRebuildVersion { - guard self.latestRequiredMenuRebuildVersion > 0 else { return } - let hasRequiredClosedMenu = menus.contains { menu in - let key = ObjectIdentifier(menu) - return (self.menuVersions[key] ?? -1) < self.latestRequiredMenuRebuildVersion - } - guard hasRequiredClosedMenu else { return } + if self.latestRequiredMenuRebuildVersion > 0, + menus.contains(where: { menu in + let key = ObjectIdentifier(menu) + return (self.menuVersions[key] ?? -1) < self.latestRequiredMenuRebuildVersion + }) + { requiredClosedPreparationVersion = self.latestRequiredMenuRebuildVersion + } else if self.menuContentVersion > self.latestRequiredMenuRebuildVersion { + guard self.latestRequiredMenuRebuildVersion > 0 else { return } + return } else { requiredClosedPreparationVersion = nil } for menu in menus { let key = ObjectIdentifier(menu) - guard !self.closedMenusDeferredUntilNextOpen.contains(key) else { continue } if let requiredClosedPreparationVersion { + self.closedMenusDeferredUntilNextOpen.remove(key) guard (self.menuVersions[key] ?? -1) < requiredClosedPreparationVersion else { continue } + } else { + guard !self.closedMenusDeferredUntilNextOpen.contains(key) else { continue } } // Pre-warming the merged menu while it is closed runs a full main-thread populateMenu // (incl. SwiftUI hosting-view layout) that menuWillOpen redoes synchronously on display @@ -96,6 +126,7 @@ extension StatusItemController { self.menuProviders.removeValue(forKey: key) self.menuVersions.removeValue(forKey: key) self.menuReadinessSignatures.removeValue(forKey: key) + self.menuIdentitySignatures.removeValue(forKey: key) self.closedMenusDeferredUntilNextOpen.remove(key) } @@ -112,28 +143,36 @@ extension StatusItemController { func refreshMenuForOpenIfNeeded(_ menu: NSMenu, provider: UsageProvider?) { self.closedMenusDeferredUntilNextOpen.remove(ObjectIdentifier(menu)) guard self.menuNeedsRefresh(menu) else { return } - if self.canPreserveStaleMenuContentDuringRefresh(menu) { + if self.canPreserveStaleMenuContentForInstantOpen(menu) { #if DEBUG self.menuLogger.debug( - "menu open kept existing content during refresh", + "menu open kept existing content for instant render", metadata: [ "items": "\(menu.items.count)", "provider": provider?.rawValue ?? "nil", "storeRefreshing": self.store.isRefreshing ? "1" : "0", ]) #endif - self.deferMenuInteractionRefreshIfNeeded() + if self.isMenuRefreshEnabled, !self.isMenuDataRefreshInFlight { + self.scheduleOpenMenuRebuildIfStillVisible( + menu, + provider: provider, + resyncReadinessBaselineAfterRebuild: self.openMenus.isEmpty) + } return } self.populateMenu(menu, provider: provider) self.markMenuFresh(menu) } - private func canPreserveStaleMenuContentDuringRefresh(_ menu: NSMenu) -> Bool { - guard self.isMenuDataRefreshInFlight, !menu.items.isEmpty else { return false } + private func canPreserveStaleMenuContentForInstantOpen(_ menu: NSMenu) -> Bool { + guard !menu.items.isEmpty else { return false } let key = ObjectIdentifier(menu) guard let menuVersion = self.menuVersions[key] else { return false } - return menuVersion >= self.latestRequiredMenuRebuildVersion + return self.menuContentVersion == self.latestDataOnlyMenuContentVersion && + menuVersion >= self.latestStructuralMenuContentVersion && + self.menuIdentitySignatures[key] == self.menuIdentitySignature( + for: self.renderedProviders(for: menu)) } private func attachedMenusForClosedPreparation() -> [NSMenu] { @@ -233,6 +272,64 @@ extension StatusItemController { let key = ObjectIdentifier(menu) self.menuVersions[key] = self.menuContentVersion self.menuReadinessSignatures[key] = self.menuAdjunctReadinessSignature() + self.menuIdentitySignatures[key] = self.menuIdentitySignature( + for: self.renderedProviders(for: menu)) + } + + private func menuIdentitySignature(for providers: [UsageProvider]) -> String { + var parts: [String] = [] + for target in providers { + parts.append(target.rawValue) + parts.append(self.providerIdentitySignature(self.store.snapshot(for: target)?.identity(for: target))) + + if self.store.metadata(for: target).usesAccountFallback { + let account = self.store.accountInfo(for: target) + parts.append(Self.menuIdentityField(account.email)) + parts.append(Self.menuIdentityField(account.plan)) + } + + for accountSnapshot in self.store.accountSnapshots[target] ?? [] { + parts.append(accountSnapshot.account.id.uuidString) + parts.append(Self.menuIdentityField(accountSnapshot.account.label)) + parts.append(self.providerIdentitySignature(accountSnapshot.snapshot?.identity(for: target))) + } + + if target == .codex { + for account in self.settings.codexVisibleAccountProjection.visibleAccounts { + parts.append(Self.menuIdentityField(account.id)) + parts.append(Self.menuIdentityField(account.email)) + parts.append(Self.menuIdentityField(account.workspaceLabel)) + parts.append(account.isActive ? "active" : "inactive") + parts.append(account.isLive ? "live" : "stored") + } + for accountSnapshot in self.store.codexAccountSnapshots { + parts.append(Self.menuIdentityField(accountSnapshot.id)) + parts.append(self.providerIdentitySignature(accountSnapshot.snapshot?.identity(for: target))) + } + } + + if target == .kilo { + for scopeSnapshot in self.store.kiloScopeSnapshots { + parts.append(Self.menuIdentityField(scopeSnapshot.id)) + parts.append(self.providerIdentitySignature(scopeSnapshot.snapshot?.identity(for: target))) + } + } + } + return parts.joined(separator: "|") + } + + private func providerIdentitySignature(_ identity: ProviderIdentitySnapshot?) -> String { + [ + identity?.providerID?.rawValue ?? "", + Self.menuIdentityField(identity?.accountEmail), + Self.menuIdentityField(identity?.accountOrganization), + Self.menuIdentityField(identity?.loginMethod), + ].joined(separator: ":") + } + + private static func menuIdentityField(_ value: String?) -> String { + let value = value ?? "" + return "\(value.utf8.count):\(value)" } func hasOpenHostedSubviewMenu() -> Bool { diff --git a/Sources/CodexBar/StatusItemController.swift b/Sources/CodexBar/StatusItemController.swift index 23208fcf55..b1e57e6ea3 100644 --- a/Sources/CodexBar/StatusItemController.swift +++ b/Sources/CodexBar/StatusItemController.swift @@ -116,6 +116,8 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin var menuProviders: [ObjectIdentifier: UsageProvider] = [:] var menuContentVersion: Int = 0 var latestRequiredMenuRebuildVersion: Int = 0 + var latestDataOnlyMenuContentVersion: Int = 0 + var latestStructuralMenuContentVersion: Int = 0 var menuVersions: [ObjectIdentifier: Int] = [:] var menuReadinessSignatures: [ObjectIdentifier: String] = [:] let hostedSubviewRenderSignatures = NSMapTable.weakToStrongObjects() @@ -136,9 +138,14 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin var openMenuRebuildTasks: [ObjectIdentifier: Task] = [:] var openMenuRebuildTokens: [ObjectIdentifier: Int] = [:] var openMenuRebuildTokenCounter = 0 + var menuIdentitySignatures: [ObjectIdentifier: String] = [:] var openMenuRebuildsClosingHostedSubviewMenus: Set = [] var parentMenuRebuildsDeferredDuringTracking: Set = [] - var deferredMenuInteractionRefreshPending = false + var deferredMenuInteractionRefreshProviders: Set = [] + var deferredMenuInteractionRefreshPending: Bool { + !self.deferredMenuInteractionRefreshProviders.isEmpty + } + var deferredOpenAIDashboardRefreshReason: String? var deferredMenuInteractionRefreshTask: Task? var highlightedMenuItems: [ObjectIdentifier: NSMenuItem] = [:] @@ -649,6 +656,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin #endif let configChanged = self.settings.configRevision != self.lastConfigRevision let orderChanged = self.settings.providerOrder != self.lastProviderOrder + let localizationChanged = self.menuLocalizationSignature() != self.lastMenuLocalizationSignature let shouldRefreshOpenMenus = self.shouldRefreshOpenMenusForProviderSwitcher() self.invalidateMenus() if orderChanged || configChanged { @@ -657,7 +665,8 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin self.updateVisibility() self.updateIcons() if shouldRefreshOpenMenus { - self.refreshOpenMenusForStructureChange() + self.refreshOpenMenusAllowingParentRebuild( + deferParentRebuildDuringTracking: !localizationChanged) } } diff --git a/Sources/CodexBar/UsageStore+PlanUtilization.swift b/Sources/CodexBar/UsageStore+PlanUtilization.swift index 29406373ac..988e252571 100644 --- a/Sources/CodexBar/UsageStore+PlanUtilization.swift +++ b/Sources/CodexBar/UsageStore+PlanUtilization.swift @@ -94,6 +94,11 @@ extension UsageStore { && self.error(for: provider) == nil } + func shouldShowRefreshingMenuCardIndicator(for provider: UsageProvider) -> Bool { + let isRefreshing = self.isRefreshing || self.refreshingProviders.contains(provider) + return isRefreshing && self.error(for: provider) == nil + } + func shouldHidePlanUtilizationMenuItem(for provider: UsageProvider) -> Bool { guard self.supportsPlanUtilizationHistory(for: provider) else { return true } return self.shouldShowRefreshingMenuCard(for: provider) diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index eabe08c780..ed33360cf0 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -2,6 +2,19 @@ import CodexBarCore import Foundation extension UsageStore { + private struct ProviderRefreshOutcomeContext { + let generation: UInt64 + let codexExpectedGuard: CodexAccountScopedRefreshGuard? + let claudeCredentialsChanged: Bool + let shouldConsumeClaudeKeychainFingerprint: Bool + } + + func refreshForSettingsChange() async { + await self.runRefresh( + startupConnectivityRetryAttempt: nil, + coalesceProviderRefreshesOverride: false) + } + func prepareRefreshState(for provider: UsageProvider? = nil) { guard provider == nil || provider == .codex else { return } _ = self.settings.persistResolvedCodexActiveSourceCorrectionIfNeeded() @@ -20,9 +33,140 @@ extension UsageStore { return self.providerSpecs[provider] } - func refreshProvider(_ provider: UsageProvider, allowDisabled: Bool = false) async { + func refreshProvider( + _ provider: UsageProvider, + allowDisabled: Bool = false, + coalesceIfRefreshing: Bool = false) async + { + while coalesceIfRefreshing, + let states = self.providerRefreshTasks[provider], + let latestGeneration = self.latestProviderRefreshGenerations[provider], + let existingState = states.last(where: { $0.generation == latestGeneration }) + { + await self.waitForProviderRefresh(provider, state: existingState) + if Task.isCancelled { return } + if existingState.shouldRetry { + self.removeProviderRefreshTask(provider, state: existingState) + continue + } + return + } + + self.providerRefreshTaskGeneration &+= 1 + let generation = self.providerRefreshTaskGeneration + let predecessorStates = self.providerRefreshTasks[provider] ?? [] + for predecessorState in predecessorStates { + predecessorState.cancelTask() + } + self.latestProviderRefreshGenerations[provider] = generation + let state = ProviderRefreshTaskState(generation: generation) + let task = Task { @MainActor [weak self] in + guard let self else { return } + var snapshotUpdatedAtBeforeRefresh: Date? + var didStartRefresh = false + for predecessorState in predecessorStates { + await predecessorState.waitForTaskCompletion() + } + if !Task.isCancelled, self.isCurrentProviderRefreshGeneration(provider, generation: generation) { + snapshotUpdatedAtBeforeRefresh = self.snapshot(for: provider)?.updatedAt + didStartRefresh = true + await self.refreshProviderTracked( + provider, + allowDisabled: allowDisabled, + generation: generation) + } + let publishedNewSnapshot = didStartRefresh && + self.snapshot(for: provider)?.updatedAt != snapshotUpdatedAtBeforeRefresh + let retryRequired = Task.isCancelled && !publishedNewSnapshot + self.providerRefreshDidComplete(provider, state: state, retryRequired: retryRequired) + } + state.install(task: task) + self.providerRefreshTasks[provider, default: []].append(state) + await self.waitForProviderRefresh(provider, state: state) + } + + private func waitForProviderRefresh(_ provider: UsageProvider, state: ProviderRefreshTaskState) async { + self.providerRefreshWaiterGeneration &+= 1 + let waiterID = self.providerRefreshWaiterGeneration + guard let task = state.addWaiter(waiterID) else { return } + await withTaskCancellationHandler { + await task.value + } onCancel: { + state.cancelWaiter(waiterID) + } + state.finishWaiter(waiterID) + if state.canRemove { + self.scheduleProviderRefreshTaskRemoval(provider, state: state) + } + } + + private func providerRefreshDidComplete( + _ provider: UsageProvider, + state: ProviderRefreshTaskState, + retryRequired: Bool) + { + state.markCompleted(retryRequired: retryRequired) + self.scheduleProviderRefreshTaskRemoval(provider, state: state) + } + + private func removeProviderRefreshTask(_ provider: UsageProvider, state: ProviderRefreshTaskState) { + guard var states = self.providerRefreshTasks[provider] else { return } + states.removeAll { $0 === state } + if states.isEmpty { + self.providerRefreshTasks.removeValue(forKey: provider) + } else { + self.providerRefreshTasks[provider] = states + } + } + + private func scheduleProviderRefreshTaskRemoval(_ provider: UsageProvider, state: ProviderRefreshTaskState) { + Task { @MainActor [weak self] in + await Task.yield() + guard let self, + self.providerRefreshTasks[provider]?.contains(where: { $0 === state }) == true, + state.canRemove + else { + return + } + self.removeProviderRefreshTask(provider, state: state) + } + } + + func isCurrentProviderRefreshGeneration(_ provider: UsageProvider, generation: UInt64?) -> Bool { + guard let generation else { return true } + return self.latestProviderRefreshGenerations[provider] == generation + } + + private func refreshProviderTracked( + _ provider: UsageProvider, + allowDisabled: Bool, + generation: UInt64) async + { + self.providerRefreshCounts[provider, default: 0] += 1 + self.refreshingProviders.insert(provider) + defer { + let remaining = max(0, self.providerRefreshCounts[provider, default: 1] - 1) + if remaining == 0 { + self.providerRefreshCounts.removeValue(forKey: provider) + self.refreshingProviders.remove(provider) + } else { + self.providerRefreshCounts[provider] = remaining + } + } + await self.refreshProviderNow( + provider, + allowDisabled: allowDisabled, + generation: generation) + } + + private func refreshProviderNow( + _ provider: UsageProvider, + allowDisabled: Bool, + generation: UInt64) async + { self.prepareRefreshState(for: provider) guard let spec = await self.providerRefreshSpec(provider) else { return } + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } let codexExpectedGuard = provider == .codex ? self.freshCodexAccountScopedRefreshGuard() : nil if !spec.isEnabled(), !allowDisabled { @@ -30,18 +174,16 @@ extension UsageStore { return } - self.refreshingProviders.insert(provider) - defer { self.refreshingProviders.remove(provider) } - if provider == .codex, self.shouldFetchAllCodexVisibleAccounts() { - await self.refreshCodexVisibleAccountsForMenu() + await self.refreshCodexVisibleAccountsForMenu(generation: generation) return } else if provider == .codex { self.codexAccountSnapshots = [] } if provider == .kilo, self.shouldFanOutKiloScopes() { - await self.refreshKiloScopes() + await self.refreshKiloScopes(generation: generation) + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } // Continue to also fetch the personal snapshot through the regular path // so the existing single-card render keeps working when only personal is shown. // The presence of multi-element kiloScopeSnapshots triggers stacked rendering. @@ -51,7 +193,10 @@ extension UsageStore { let tokenAccounts = self.tokenAccounts(for: provider) if self.shouldFetchAllTokenAccounts(provider: provider, accounts: tokenAccounts) { - await self.refreshTokenAccounts(provider: provider, accounts: tokenAccounts) + await self.refreshTokenAccounts( + provider: provider, + accounts: tokenAccounts, + generation: generation) return } else { _ = await MainActor.run { @@ -62,7 +207,7 @@ extension UsageStore { let claudeAuthStateBeforeFetch = provider == .claude ? await Self.captureClaudeRefreshAuthState(invalidateCredentialsFile: true) : nil - let fetchContext = spec.makeFetchContext() + let fetchContext = self.makeFetchContext(provider: provider, override: nil) let descriptor = spec.descriptor // Keep provider fetch work off MainActor so slow keychain/process reads don't stall menu/UI responsiveness. let outcome = await withTaskGroup( @@ -74,6 +219,7 @@ extension UsageStore { } return await group.next()! } + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } let claudeAuthFingerprintAfterFetch = provider == .claude ? await Self.captureClaudeAuthFingerprintToken() : nil @@ -88,6 +234,22 @@ extension UsageStore { let shouldConsumeClaudeKeychainFingerprint = Self.shouldConsumeClaudeKeychainFingerprintChange( beforeFetch: claudeAuthStateBeforeFetch, changedDuringFetch: claudeAuthChangedDuringFetch) + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } + await self.applyProviderRefreshOutcome( + provider: provider, + outcome: outcome, + context: ProviderRefreshOutcomeContext( + generation: generation, + codexExpectedGuard: codexExpectedGuard, + claudeCredentialsChanged: claudeCredentialsChanged, + shouldConsumeClaudeKeychainFingerprint: shouldConsumeClaudeKeychainFingerprint)) + } + + private func applyProviderRefreshOutcome( + provider: UsageProvider, + outcome: ProviderFetchOutcome, + context: ProviderRefreshOutcomeContext) async + { await MainActor.run { self.lastFetchAttempts[provider] = outcome.attempts } @@ -96,17 +258,20 @@ extension UsageStore { case let .success(result): let scoped = result.usage.scoped(to: provider) if provider == .codex, - let codexExpectedGuard, + let codexExpectedGuard = context.codexExpectedGuard, !self.shouldApplyCodexUsageResult(expectedGuard: codexExpectedGuard, usage: scoped) { return } - let backfilled = await MainActor.run { - if claudeCredentialsChanged { + let backfilled = await MainActor.run { () -> UsageSnapshot? in + guard self.isCurrentProviderRefreshGeneration(provider, generation: context.generation) else { + return nil + } + if context.claudeCredentialsChanged { self.clearClaudeCredentialDerivedStateForCredentialSwapNow() } let resetBackfillSource = provider == .codex - ? self.codexLastKnownResetSnapshot(matching: codexExpectedGuard) + ? self.codexLastKnownResetSnapshot(matching: context.codexExpectedGuard) : self.lastKnownResetSnapshots[provider] let backfilled = scoped.backfillingResetTimes(from: resetBackfillSource) self.handleQuotaWarningTransitions(provider: provider, snapshot: backfilled) @@ -130,12 +295,14 @@ extension UsageStore { } return backfilled } - if shouldConsumeClaudeKeychainFingerprint { + guard let backfilled else { return } + if context.shouldConsumeClaudeKeychainFingerprint { _ = await Self.consumeClaudeKeychainFingerprintChangeWithoutPrompt() } await self.recordPlanUtilizationHistorySample( provider: provider, snapshot: backfilled) + guard self.isCurrentProviderRefreshGeneration(provider, generation: context.generation) else { return } if let runtime = self.providerRuntimes[provider] { let context = ProviderRuntimeContext( provider: provider, settings: self.settings, store: self) @@ -146,19 +313,23 @@ extension UsageStore { } case let .failure(error): if provider == .codex, - let codexExpectedGuard, + let codexExpectedGuard = context.codexExpectedGuard, !self.shouldApplyCodexScopedFailure(expectedGuard: codexExpectedGuard) { return } + guard self.isCurrentProviderRefreshGeneration(provider, generation: context.generation) else { return } self.recordStartupConnectivityRetryableFailure(error) - if claudeCredentialsChanged { + if context.claudeCredentialsChanged { await self.clearClaudeCredentialDerivedStateForCredentialSwap() } - if shouldConsumeClaudeKeychainFingerprint { + if context.shouldConsumeClaudeKeychainFingerprint { _ = await Self.consumeClaudeKeychainFingerprintChangeWithoutPrompt() } - await self.handleProviderFetchFailure(provider: provider, error: error) + await self.handleProviderFetchFailure( + provider: provider, + error: error, + generation: context.generation) } } @@ -293,9 +464,14 @@ extension UsageStore { self.lastTokenFetchAt.removeValue(forKey: .claude) } - private func handleProviderFetchFailure(provider: UsageProvider, error: Error) async { + private func handleProviderFetchFailure( + provider: UsageProvider, + error: Error, + generation: UInt64) async + { let shouldNotifyPermissionPrompt = Self.isPermissionPromptWaiting(error) await MainActor.run { + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } let hadPriorData = self.snapshots[provider] != nil let preservesPriorData = Self.shouldPreservePriorSnapshot( after: error, @@ -336,6 +512,7 @@ extension UsageStore { self.postPermissionPromptNotificationIfNeeded(provider: provider, error: error) } } + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } if let runtime = self.providerRuntimes[provider] { let context = ProviderRuntimeContext( provider: provider, settings: self.settings, store: self) diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 5f0476c7ce..aa7d4cada1 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -75,7 +75,7 @@ extension UsageStore { projection.visibleAccounts.count > 1 } - func refreshCodexVisibleAccountsForMenu() async { + func refreshCodexVisibleAccountsForMenu(generation: UInt64? = nil) async { let projection = self.freshCodexVisibleAccountProjectionForAccountRefresh() let accounts = self.limitedCodexVisibleAccounts( projection.visibleAccounts, @@ -131,6 +131,7 @@ extension UsageStore { let currentProjection = self.freshCodexVisibleAccountProjectionForAccountRefresh( requireLiveManagedAuthFor: managedAccountIDsWithReadableAuthAtStart) + guard self.isCurrentProviderRefreshGeneration(.codex, generation: generation) else { return } let currentSnapshots = snapshots.compactMap { snapshot -> CodexAccountUsageSnapshot? in guard let currentAccount = Self.currentCodexVisibleAccount( matching: snapshot.account, @@ -188,7 +189,8 @@ extension UsageStore { selectedOutcome, account: currentSelectedAccount, snapshot: currentSelectedSnapshot, - sourceLabel: selectedSourceLabel) + sourceLabel: selectedSourceLabel, + generation: generation) } } else { _ = self.prepareCodexAccountScopedRefreshIfNeeded() @@ -386,7 +388,11 @@ extension UsageStore { } } - func refreshTokenAccounts(provider: UsageProvider, accounts: [ProviderTokenAccount]) async { + func refreshTokenAccounts( + provider: UsageProvider, + accounts: [ProviderTokenAccount], + generation: UInt64? = nil) async + { let selectedAccount = self.settings.selectedTokenAccount(for: provider) let limitedAccounts = self.limitedTokenAccounts(accounts, selected: selectedAccount) let effectiveSelected = selectedAccount ?? limitedAccounts.first @@ -405,6 +411,7 @@ extension UsageStore { var sawAnyNonCancellationOutcome = false let results = await self.fetchTokenAccountOutcomes(provider: provider, accounts: limitedAccounts) + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } for result in results { let account = result.account let outcome = result.outcome @@ -445,9 +452,11 @@ extension UsageStore { selectedOutcome, provider: provider, account: effectiveSelected, - fallbackSnapshot: selectedSnapshot) + fallbackSnapshot: selectedSnapshot, + generation: generation) } + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } await self.recordFetchedTokenAccountPlanUtilizationHistory( provider: provider, samples: historySamples, @@ -632,6 +641,9 @@ extension UsageStore { codexActiveSourceOverride: codexActiveSourceOverride) let fetcher = ProviderRegistry.makeFetcher(base: self.codexFetcher, provider: provider, env: env) let verbose = self.settings.isVerboseLoggingEnabled + let contextProvider = provider + let originalAccountToken = account?.token + let originalManualToken = provider == .stepfun ? self.settings.stepfunToken : nil return ProviderFetchContext( runtime: .app, sourceMode: sourceMode, @@ -646,19 +658,26 @@ extension UsageStore { claudeFetcher: self.claudeFetcher, browserDetection: self.browserDetection, selectedTokenAccountID: account?.id, - tokenAccountTokenUpdater: { [weak settings = self.settings] provider, accountID, token in + tokenAccountTokenUpdater: { [weak self] provider, accountID, token in await MainActor.run { - settings?.updateTokenAccount( + guard let self, provider == contextProvider, + self.settings.tokenAccounts(for: provider) + .first(where: { $0.id == accountID })?.token == originalAccountToken + else { + return + } + self.settings.updateTokenAccount( provider: provider, accountID: accountID, token: token) } }, - providerManualTokenUpdater: { [weak settings = self.settings] provider, token in + providerManualTokenUpdater: { [weak self] provider, token in await MainActor.run { - if provider == .stepfun { - settings?.stepfunToken = token - } + guard let self, provider == .stepfun, + self.settings.stepfunToken == originalManualToken + else { return } + self.settings.stepfunToken = token } }, costUsageHistoryDays: self.settings.costUsageHistoryDays) @@ -1141,8 +1160,10 @@ extension UsageStore { _ outcome: ProviderFetchOutcome, account: CodexVisibleAccount, snapshot: UsageSnapshot?, - sourceLabel: String?) async + sourceLabel: String?, + generation: UInt64? = nil) async { + guard self.isCurrentProviderRefreshGeneration(.codex, generation: generation) else { return } self.lastFetchAttempts[.codex] = outcome.attempts switch outcome.result { case .success: @@ -1159,6 +1180,7 @@ extension UsageStore { self.rememberLiveSystemCodexEmailIfNeeded(snapshot.accountEmail(for: .codex)) self.seedCodexAccountScopedRefreshGuard(accountEmail: account.email) await self.recordPlanUtilizationHistorySample(provider: .codex, snapshot: snapshot) + guard self.isCurrentProviderRefreshGeneration(.codex, generation: generation) else { return } self.recordCodexHistoricalSampleIfNeeded(snapshot: snapshot) case let .failure(error): guard let message = self.tokenAccountErrorMessage(error) else { @@ -1182,11 +1204,14 @@ extension UsageStore { _ outcome: ProviderFetchOutcome, provider: UsageProvider, account: ProviderTokenAccount?, - fallbackSnapshot: UsageSnapshot?) async + fallbackSnapshot: UsageSnapshot?, + generation: UInt64? = nil) async { await MainActor.run { + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } self.lastFetchAttempts[provider] = outcome.attempts } + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } switch outcome.result { case let .success(result): let scoped = result.usage.scoped(to: provider) @@ -1196,6 +1221,9 @@ extension UsageStore { scoped } let backfilled = await MainActor.run { + guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { + return nil as UsageSnapshot? + } let backfilled = labeled.backfillingResetTimes(from: self.lastKnownResetSnapshots[provider]) self.handleQuotaWarningTransitions(provider: provider, snapshot: backfilled) self.handleSessionQuotaTransition(provider: provider, snapshot: backfilled) @@ -1206,6 +1234,7 @@ extension UsageStore { self.failureGates[provider]?.recordSuccess() return backfilled } + guard let backfilled else { return } await self.recordPlanUtilizationHistorySample( provider: provider, snapshot: backfilled, diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index bae9cef7e0..e3ea41a79d 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -64,7 +64,7 @@ extension UsageStore { self.startTimer() self.updateProviderRuntimes() await self.refreshHistoricalDatasetIfNeeded() - await self.refresh() + await self.refreshForSettingsChange() } } } @@ -227,6 +227,11 @@ final class UsageStore { @ObservationIgnored var providerSpecs: [UsageProvider: ProviderSpec] = [:] @ObservationIgnored let providerMetadata: [UsageProvider: ProviderMetadata] @ObservationIgnored var providerRuntimes: [UsageProvider: any ProviderRuntime] = [:] + @ObservationIgnored var providerRefreshTasks: [UsageProvider: [ProviderRefreshTaskState]] = [:] + @ObservationIgnored var providerRefreshTaskGeneration: UInt64 = 0 + @ObservationIgnored var providerRefreshWaiterGeneration: UInt64 = 0 + @ObservationIgnored var latestProviderRefreshGenerations: [UsageProvider: UInt64] = [:] + @ObservationIgnored var providerRefreshCounts: [UsageProvider: Int] = [:] @ObservationIgnored private var providerAvailabilityCache: [UsageProvider: ProviderAvailabilityCacheEntry] = [:] @ObservationIgnored var accountInfoCache: [UsageProvider: AccountInfoCacheEntry] = [:] @ObservationIgnored private var timerTask: Task? @@ -549,8 +554,8 @@ final class UsageStore { func runRefresh( forceTokenUsage: Bool = false, - startupConnectivityRetryAttempt: Int?) - async + startupConnectivityRetryAttempt: Int?, + coalesceProviderRefreshesOverride: Bool? = nil) async { guard !self.isRefreshing else { return } self.prepareRefreshState() @@ -583,7 +588,12 @@ final class UsageStore { await withTaskGroup(of: Void.self) { group in for provider in refreshProviders { - group.addTask { await self.refreshProvider(provider) } + group.addTask { + await self.refreshProvider( + provider, + coalesceIfRefreshing: coalesceProviderRefreshesOverride ?? + (ProviderInteractionContext.current == .background)) + } if availableRefreshProviders.contains(provider) { group.addTask { await self.refreshStatus(provider) } } diff --git a/Sources/CodexBar/UsageStoreSupport.swift b/Sources/CodexBar/UsageStoreSupport.swift index 3ac92fc59e..6f9aa532cd 100644 --- a/Sources/CodexBar/UsageStoreSupport.swift +++ b/Sources/CodexBar/UsageStoreSupport.swift @@ -1,6 +1,78 @@ import CodexBarCore import Foundation +final class ProviderRefreshTaskState: @unchecked Sendable { + let generation: UInt64 + + private let lock = NSLock() + private var task: Task? + private var waiterIDs: Set = [] + private var completed = false + private var retryRequired = false + + init(generation: UInt64) { + self.generation = generation + } + + func install(task: Task) { + self.lock.withLock { + self.task = task + } + } + + func addWaiter(_ waiterID: UInt64) -> Task? { + self.lock.withLock { + self.waiterIDs.insert(waiterID) + return self.task + } + } + + func cancelWaiter(_ waiterID: UInt64) { + let taskToCancel = self.lock.withLock { + guard self.waiterIDs.remove(waiterID) != nil else { return nil as Task? } + return self.waiterIDs.isEmpty && !self.completed ? self.task : nil + } + taskToCancel?.cancel() + } + + func finishWaiter(_ waiterID: UInt64) { + _ = self.lock.withLock { + self.waiterIDs.remove(waiterID) + } + } + + func markCompleted(retryRequired: Bool) { + self.lock.withLock { + self.completed = true + self.retryRequired = retryRequired + } + } + + func cancelTask() { + let task = self.lock.withLock { + self.completed ? nil : self.task + } + task?.cancel() + } + + func waitForTaskCompletion() async { + let task = self.lock.withLock { self.task } + await task?.value + } + + var isCompleted: Bool { + self.lock.withLock { self.completed } + } + + var shouldRetry: Bool { + self.lock.withLock { self.retryRequired } + } + + var canRemove: Bool { + self.lock.withLock { self.completed && self.waiterIDs.isEmpty } + } +} + enum ProviderStatusIndicator: String { case none case minor diff --git a/Tests/CodexBarTests/MenuCardSubtitleTests.swift b/Tests/CodexBarTests/MenuCardSubtitleTests.swift index 8407c1c10d..6cda8e0d51 100644 --- a/Tests/CodexBarTests/MenuCardSubtitleTests.swift +++ b/Tests/CodexBarTests/MenuCardSubtitleTests.swift @@ -46,4 +46,49 @@ struct MenuCardSubtitleTests { #expect(model.subtitleText == UsageFormatter.updatedString(from: updatedAt, now: now)) } + + @Test + func `subtitle shows refreshing while cached snapshot remains visible`() throws { + let updatedAt = Date(timeIntervalSinceReferenceDate: 0) + let now = updatedAt.addingTimeInterval(5 * 3600) + let snapshot = UsageSnapshot( + primary: RateWindow( + usedPercent: 22, + windowMinutes: 300, + resetsAt: now.addingTimeInterval(3000), + resetDescription: nil), + secondary: nil, + tertiary: nil, + updatedAt: updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: "codex@example.com", + accountOrganization: nil, + loginMethod: "Plus Plan")) + let metadata = try #require(ProviderDefaults.metadata[.codex]) + + let model = UsageMenuCardView.Model.make(.init( + provider: .codex, + metadata: metadata, + snapshot: snapshot, + credits: nil, + creditsError: nil, + dashboard: nil, + dashboardError: nil, + tokenSnapshot: nil, + tokenError: nil, + account: AccountInfo(email: "codex@example.com", plan: "Plus Plan"), + isRefreshing: true, + lastError: nil, + usageBarsShowUsed: false, + resetTimeDisplayStyle: .countdown, + tokenCostUsageEnabled: false, + showOptionalCreditsAndExtraUsage: true, + hidePersonalInfo: false, + now: now)) + + #expect(model.subtitleText == "Refreshing…") + #expect(model.subtitleStyle == .loading) + #expect(!model.metrics.isEmpty) + } } diff --git a/Tests/CodexBarTests/StatusMenuClosedPreparationTests.swift b/Tests/CodexBarTests/StatusMenuClosedPreparationTests.swift index e2e1fbaf37..92ebb005f8 100644 --- a/Tests/CodexBarTests/StatusMenuClosedPreparationTests.swift +++ b/Tests/CodexBarTests/StatusMenuClosedPreparationTests.swift @@ -32,6 +32,9 @@ extension StatusMenuTests { defer { StatusItemController.resetClosedMenuPreparationDelayForTesting() } controller.menuRefreshEnabledOverrideForTesting = true + for _ in 0..<20 { + await Task.yield() + } let menu = controller.makeMenu() // Simulate a closed menu that was attached by an icon update but has never been opened. controller.fallbackMenu = menu @@ -80,6 +83,9 @@ extension StatusMenuTests { defer { StatusItemController.resetClosedMenuPreparationDelayForTesting() } controller.menuRefreshEnabledOverrideForTesting = true + for _ in 0..<20 { + await Task.yield() + } let menu = controller.makeMenu() controller.fallbackMenu = menu controller.statusItem.menu = menu @@ -101,6 +107,8 @@ extension StatusMenuTests { #expect(controller.menuVersions[key] == openedVersion) store.isRefreshing = false + controller.fallbackMenu = menu + controller.statusItem.menu = menu controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) for _ in 0..<40 where controller.menuVersions[key] == openedVersion { await Task.yield() diff --git a/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift b/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift new file mode 100644 index 0000000000..00dd8dbb4e --- /dev/null +++ b/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift @@ -0,0 +1,1540 @@ +import AppKit +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +extension StatusMenuTests { + @Test + func `opening fresh menu does not schedule deferred refresh`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + var providerRefreshCount = 0 + var refreshInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + refreshInteractions.append(ProviderInteractionContext.current) + providerRefreshCount += 1 + } + defer { store._test_providerRefreshOverride = nil } + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + StatusItemController.setClosedMenuPreparationDelayForTesting(.zero) + defer { StatusItemController.resetClosedMenuPreparationDelayForTesting() } + + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.zero) + defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + + for _ in 0..<20 { + await Task.yield() + } + #expect(providerRefreshCount == 0) + #expect(!controller.deferredMenuInteractionRefreshPending) + + controller.menuDidClose(menu) + for _ in 0..<40 { + await Task.yield() + } + + #expect(providerRefreshCount == 0) + #expect(refreshInteractions.isEmpty) + } + + @Test + func `menu open with missing data refreshes asynchronously while tracking`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + var providerRefreshCount = 0 + var refreshInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + refreshInteractions.append(ProviderInteractionContext.current) + providerRefreshCount += 1 + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow( + usedPercent: 25, + windowMinutes: 300, + resetsAt: Date().addingTimeInterval(3600), + resetDescription: nil), + secondary: nil, + updatedAt: Date()), + provider: .codex) + } + defer { store._test_providerRefreshOverride = nil } + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + StatusItemController.setClosedMenuPreparationDelayForTesting(.zero) + defer { StatusItemController.resetClosedMenuPreparationDelayForTesting() } + + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.zero) + defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + #expect(controller.deferredMenuInteractionRefreshPending) + + for _ in 0..<40 where providerRefreshCount == 0 { + await Task.yield() + } + + #expect(providerRefreshCount == 1) + #expect(refreshInteractions == [.background]) + for _ in 0..<40 where controller.deferredMenuInteractionRefreshPending { + await Task.yield() + } + #expect(!controller.deferredMenuInteractionRefreshPending) + controller.menuDidClose(menu) + #expect(!controller.deferredMenuInteractionRefreshPending) + } + + @Test + func `menu open renders cached data immediately after data only invalidation`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + var providerRefreshCount = 0 + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + providerRefreshCount += 1 + } + defer { store._test_providerRefreshOverride = nil } + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + + controller.menuRefreshEnabledOverrideForTesting = true + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let key = ObjectIdentifier(menu) + let cachedItemCount = menu.items.count + let cachedVersion = controller.menuVersions[key] + controller.lastMenuAdjunctReadinessSignature = "stale-baseline" + + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + let dataOnlyVersion = controller.menuContentVersion + var asyncRebuilds = 0 + controller._test_openMenuRebuildObserver = { _ in + asyncRebuilds += 1 + } + defer { controller._test_openMenuRebuildObserver = nil } + + controller.menuWillOpen(menu) + + #expect(cachedVersion != dataOnlyVersion) + #expect(menu.items.count == cachedItemCount) + #expect(controller.menuVersions[key] == cachedVersion) + #expect(asyncRebuilds == 0) + #expect(!controller.deferredMenuInteractionRefreshPending) + + for _ in 0..<40 where asyncRebuilds == 0 { + await Task.yield() + } + + #expect(asyncRebuilds == 1) + #expect(controller.menuVersions[key] == controller.menuContentVersion) + #expect(!controller.didMenuAdjunctReadinessChange()) + controller.menuDidClose(menu) + for _ in 0..<40 { + await Task.yield() + } + #expect(providerRefreshCount == 0) + } + + @Test + func `closing before cached menu rebuild keeps next open stale`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + self.enableOnlyCodexForInstantOpenTesting(settings) + + 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() } + controller.menuRefreshEnabledOverrideForTesting = true + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let key = ObjectIdentifier(menu) + let cachedVersion = controller.menuVersions[key] + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + + var rebuildGateEntries = 0 + var rebuildGate: CheckedContinuation? + controller._test_openMenuRefreshYieldOverride = { + rebuildGateEntries += 1 + await withCheckedContinuation { continuation in + rebuildGate = continuation + } + } + defer { + rebuildGate?.resume() + controller._test_openMenuRefreshYieldOverride = nil + } + + controller.menuWillOpen(menu) + for _ in 0..<40 where rebuildGateEntries == 0 { + await Task.yield() + } + + #expect(rebuildGateEntries == 1) + #expect(controller.menuVersions[key] == cachedVersion) + controller.menuDidClose(menu) + #expect(controller.menuNeedsRefresh(menu)) + + rebuildGate?.resume() + rebuildGate = nil + controller._test_openMenuRefreshYieldOverride = nil + for _ in 0..<20 { + await Task.yield() + } + + controller.menuWillOpen(menu) + for _ in 0..<40 where controller.menuNeedsRefresh(menu) { + await Task.yield() + } + + #expect(!controller.menuNeedsRefresh(menu)) + controller.menuDidClose(menu) + } + + @Test + func `menu open rebuilds synchronously after provider identity changes`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "old@example.com"), + provider: .codex) + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let cachedVersion = controller.menuVersions[ObjectIdentifier(menu)] + + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "new@example.com"), + provider: .codex) + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + + #expect(controller.menuVersions[ObjectIdentifier(menu)] != cachedVersion) + #expect(controller.menuVersions[ObjectIdentifier(menu)] == controller.menuContentVersion) + } + + @Test + func `overview menu rebuilds synchronously after secondary provider identity changes`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + settings.mergedMenuLastSelectedWasOverview = true + self.enableProvidersForInstantOpenTesting([.codex, .claude], settings: settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "codex@example.com"), + provider: .codex) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "old@example.com", provider: .claude), + provider: .claude) + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + controller.selectedMenuProvider = .codex + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let cachedVersion = controller.menuVersions[ObjectIdentifier(menu)] + + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "new@example.com", provider: .claude), + provider: .claude) + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + + #expect(controller.menuVersions[ObjectIdentifier(menu)] != cachedVersion) + #expect(controller.menuVersions[ObjectIdentifier(menu)] == controller.menuContentVersion) + } + + @Test + func `stacked Codex menu rebuilds synchronously after secondary account identity changes`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + settings.multiAccountMenuLayout = .stacked + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "selected@example.com"), + provider: .codex) + let selectedAccount = CodexVisibleAccount( + id: "selected", + email: "selected@example.com", + workspaceLabel: nil, + workspaceAccountID: nil, + authFingerprint: nil, + storedAccountID: nil, + selectionSource: .liveSystem, + isActive: true, + isLive: true, + canReauthenticate: false, + canRemove: false) + let secondaryAccount = CodexVisibleAccount( + id: "secondary", + email: "secondary@example.com", + workspaceLabel: nil, + workspaceAccountID: nil, + authFingerprint: nil, + storedAccountID: nil, + selectionSource: .liveSystem, + isActive: false, + isLive: false, + canReauthenticate: false, + canRemove: false) + store.codexAccountSnapshots = [ + CodexAccountUsageSnapshot( + account: selectedAccount, + snapshot: self.instantOpenSnapshot(email: "selected@example.com"), + error: nil, + sourceLabel: "test"), + CodexAccountUsageSnapshot( + account: secondaryAccount, + snapshot: self.instantOpenSnapshot(email: "old@example.com"), + error: nil, + sourceLabel: "test"), + ] + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let cachedVersion = controller.menuVersions[ObjectIdentifier(menu)] + + store.codexAccountSnapshots[1] = CodexAccountUsageSnapshot( + account: secondaryAccount, + snapshot: self.instantOpenSnapshot(email: "new@example.com"), + error: nil, + sourceLabel: "test") + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + + #expect(controller.menuVersions[ObjectIdentifier(menu)] != cachedVersion) + #expect(controller.menuVersions[ObjectIdentifier(menu)] == controller.menuContentVersion) + } + + @Test + func `cache preserving structural invalidation rebuilds synchronously on open`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + self.enableOnlyCodexForInstantOpenTesting(settings) + + 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() } + controller.menuRefreshEnabledOverrideForTesting = true + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let key = ObjectIdentifier(menu) + let cachedVersion = controller.menuVersions[key] + + controller.preservingMergedSwitcherContentCachesDuringInvalidation { + controller.invalidateMenus() + } + #expect(controller.menuVersions[key] == cachedVersion) + #expect(controller.menuContentVersion != controller.latestDataOnlyMenuContentVersion) + + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + + #expect(controller.menuVersions[key] == controller.menuContentVersion) + #expect(controller.openMenuRebuildTasks[key] == nil) + } + + @Test + func `data invalidation after cache preserving structural invalidation still rebuilds synchronously`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + self.enableOnlyCodexForInstantOpenTesting(settings) + + 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() } + controller.menuRefreshEnabledOverrideForTesting = true + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + let key = ObjectIdentifier(menu) + let cachedVersion = controller.menuVersions[key] + + controller.preservingMergedSwitcherContentCachesDuringInvalidation { + controller.invalidateMenus() + } + let structuralVersion = controller.menuContentVersion + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + + #expect(controller.menuVersions[key] == cachedVersion) + #expect(controller.latestStructuralMenuContentVersion == structuralVersion) + #expect(controller.menuContentVersion == controller.latestDataOnlyMenuContentVersion) + + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + + #expect(controller.menuVersions[key] == controller.menuContentVersion) + #expect(controller.openMenuRebuildTasks[key] == nil) + } + + @Test + func `menu open does not overlap provider specific refresh`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + let refreshGate = BlockingInstantOpenProviderRefresh() + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + await refreshGate.run() + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow( + usedPercent: 25, + windowMinutes: 300, + resetsAt: Date().addingTimeInterval(3600), + resetDescription: nil), + secondary: nil, + updatedAt: Date()), + provider: .codex) + } + defer { store._test_providerRefreshOverride = nil } + let existingRefreshTask = Task { + await store.refreshProvider(.codex) + } + await refreshGate.waitUntilStarted(count: 1) + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + for _ in 0..<40 { + await Task.yield() + } + + #expect(await refreshGate.startCount == 1) + #expect(store.refreshingProviders.contains(.codex)) + await refreshGate.releaseFirst() + await existingRefreshTask.value + for _ in 0..<40 where controller.deferredMenuInteractionRefreshPending { + await Task.yield() + } + #expect(!store.isRefreshing) + #expect(!controller.deferredMenuInteractionRefreshPending) + } + + @Test + func `cached menu rebuilds after active provider refresh completes`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + let refreshGate = BlockingInstantOpenProviderRefresh() + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + await refreshGate.run() + } + defer { store._test_providerRefreshOverride = nil } + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.mergedMenu = menu + controller.statusItem.menu = menu + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + + let existingRefreshTask = Task { + await store.refreshProvider(.codex) + } + await refreshGate.waitUntilStarted(count: 1) + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + #expect(controller.menuNeedsRefresh(menu)) + + await refreshGate.releaseFirst() + await existingRefreshTask.value + for _ in 0..<80 where controller.menuNeedsRefresh(menu) { + await Task.yield() + } + + #expect(await refreshGate.startCount == 1) + #expect(!controller.menuNeedsRefresh(menu)) + } + + @Test + func `menu rebuilds after displayed provider completes while another provider refreshes`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableProvidersForInstantOpenTesting([.claude, .codex], settings: settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "old@example.com"), + provider: .codex) + store.refreshingProviders = [.claude] + defer { store.refreshingProviders = [] } + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu(for: .codex) + controller.populateMenu(menu, provider: .codex) + controller.markMenuFresh(menu) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "new@example.com"), + provider: .codex) + controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) + + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + for _ in 0..<80 where controller.menuNeedsRefresh(menu) { + await Task.yield() + } + + #expect(!controller.menuNeedsRefresh(menu)) + } + + @Test + func `user refresh supersedes background provider refresh`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshGate = BlockingInstantOpenProviderRefresh() + var refreshInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + refreshInteractions.append(ProviderInteractionContext.current) + await refreshGate.run() + } + defer { store._test_providerRefreshOverride = nil } + + let backgroundRefreshTask = Task { + await ProviderInteractionContext.$current.withValue(.background) { + await store.refreshProvider(.codex, coalesceIfRefreshing: true) + } + } + await refreshGate.waitUntilStarted(count: 1) + + let userRefreshTask = Task { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh() + } + } + for _ in 0..<40 { + await Task.yield() + } + #expect(await refreshGate.startCount == 1) + + await refreshGate.releaseFirst() + await refreshGate.waitUntilStarted(count: 2) + await userRefreshTask.value + await backgroundRefreshTask.value + #expect(await refreshGate.startCount == 2) + #expect(refreshInteractions == [.background, .userInitiated]) + } + + @Test + func `settings refresh supersedes background provider refresh without becoming user initiated`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshGate = BlockingInstantOpenProviderRefresh() + var refreshInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + refreshInteractions.append(ProviderInteractionContext.current) + await refreshGate.run() + } + defer { store._test_providerRefreshOverride = nil } + + let backgroundRefreshTask = Task { + await store.refreshProvider(.codex, coalesceIfRefreshing: true) + } + await refreshGate.waitUntilStarted(count: 1) + + let settingsRefreshTask = Task { + await store.refreshForSettingsChange() + } + for _ in 0..<40 { + await Task.yield() + } + #expect(await refreshGate.startCount == 1) + + await refreshGate.releaseFirst() + await refreshGate.waitUntilStarted(count: 2) + await settingsRefreshTask.value + await backgroundRefreshTask.value + #expect(await refreshGate.startCount == 2) + #expect(refreshInteractions == [.background, .background]) + } + + @Test + func `superseded provider refresh drains before newer result`() async throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableProvidersForInstantOpenTesting([.claude], settings: settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshes = OrderedInstantOpenProviderRefresh() + let baseSpec = try #require(store.providerSpecs[.claude]) + let baseDescriptor = baseSpec.descriptor + let strategy = InstantOpenProviderFetchStrategy { + await refreshes.awaitSnapshot() + } + store.providerSpecs[.claude] = ProviderSpec( + style: baseSpec.style, + isEnabled: baseSpec.isEnabled, + descriptor: ProviderDescriptor( + id: .claude, + metadata: baseDescriptor.metadata, + branding: baseDescriptor.branding, + tokenCost: baseDescriptor.tokenCost, + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli, .oauth], + pipeline: ProviderFetchPipeline { _ in [strategy] }), + cli: baseDescriptor.cli), + makeFetchContext: baseSpec.makeFetchContext) + + let olderTask = Task { + await store.refreshProvider(.claude) + } + await refreshes.waitUntilStarted(count: 1) + let newerTask = Task { + await store.refreshProvider(.claude) + } + for _ in 0..<40 { + await Task.yield() + } + #expect(await refreshes.startCount == 1) + + await refreshes.resume( + call: 1, + snapshot: self.instantOpenSnapshot( + email: "old@example.com", + provider: .claude, + percent: 10)) + await refreshes.waitUntilStarted(count: 2) + await refreshes.resume( + call: 2, + snapshot: self.instantOpenSnapshot( + email: "new@example.com", + provider: .claude, + percent: 80)) + await newerTask.value + await olderTask.value + + #expect(store.snapshot(for: .claude)?.primary?.usedPercent == 80) + #expect(store.snapshot(for: .claude)?.accountEmail(for: .claude) == "new@example.com") + } + + @Test + func `superseded provider refresh cannot overwrite manually changed token`() async throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableProvidersForInstantOpenTesting([.stepfun], settings: settings) + settings.stepfunToken = "initial-token" + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshes = OrderedInstantOpenProviderMutation() + let baseSpec = try #require(store.providerSpecs[.stepfun]) + let baseDescriptor = baseSpec.descriptor + let strategy = InstantOpenProviderMutationFetchStrategy(mutations: refreshes) + store.providerSpecs[.stepfun] = ProviderSpec( + style: baseSpec.style, + isEnabled: baseSpec.isEnabled, + descriptor: ProviderDescriptor( + id: .stepfun, + metadata: baseDescriptor.metadata, + branding: baseDescriptor.branding, + tokenCost: baseDescriptor.tokenCost, + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline { _ in [strategy] }), + cli: baseDescriptor.cli), + makeFetchContext: baseSpec.makeFetchContext) + + let olderTask = Task { + await store.refreshProvider(.stepfun) + } + await refreshes.waitUntilStarted(count: 1) + let newerTask = Task { + await store.refreshProvider(.stepfun) + } + for _ in 0..<40 { + await Task.yield() + } + #expect(await refreshes.startCount == 1) + + settings.stepfunToken = "user-token" + await refreshes.resume(call: 1, token: "old-token") + await refreshes.waitUntilStarted(count: 2) + #expect(settings.stepfunToken == "user-token") + await refreshes.resume(call: 2, token: "new-token") + await newerTask.value + await olderTask.value + + #expect(settings.stepfunToken == "new-token") + } + + @Test + func `superseded provider refresh preserves rotated token when credential is unchanged`() async throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableProvidersForInstantOpenTesting([.stepfun], settings: settings) + settings.stepfunToken = "initial-token" + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshes = OrderedInstantOpenProviderMutation() + let baseSpec = try #require(store.providerSpecs[.stepfun]) + let baseDescriptor = baseSpec.descriptor + let strategy = InstantOpenProviderMutationFetchStrategy(mutations: refreshes) + store.providerSpecs[.stepfun] = ProviderSpec( + style: baseSpec.style, + isEnabled: baseSpec.isEnabled, + descriptor: ProviderDescriptor( + id: .stepfun, + metadata: baseDescriptor.metadata, + branding: baseDescriptor.branding, + tokenCost: baseDescriptor.tokenCost, + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline { _ in [strategy] }), + cli: baseDescriptor.cli), + makeFetchContext: baseSpec.makeFetchContext) + + let olderTask = Task { + await store.refreshProvider(.stepfun) + } + await refreshes.waitUntilStarted(count: 1) + let newerTask = Task { + await store.refreshProvider(.stepfun) + } + + await refreshes.resume(call: 1, token: "rotated-token") + await refreshes.waitUntilStarted(count: 2) + #expect(settings.stepfunToken == "rotated-token") + await refreshes.resume(call: 2, token: "newer-token") + await newerTask.value + await olderTask.value + + #expect(settings.stepfunToken == "newer-token") + } + + @Test + func `canceling provider refresh cancels its owned probe task`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshGate = BlockingInstantOpenProviderRefresh() + var refreshWasCancelled = false + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + await refreshGate.run() + refreshWasCancelled = Task.isCancelled + } + defer { store._test_providerRefreshOverride = nil } + + let refreshTask = Task { + await store.refreshProvider(.codex) + } + await refreshGate.waitUntilStarted(count: 1) + refreshTask.cancel() + await refreshGate.releaseFirst() + await refreshTask.value + + #expect(refreshWasCancelled) + } + + @Test + func `canceling refresh owner keeps shared provider probe alive`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + let refreshGate = BlockingInstantOpenProviderRefresh() + var refreshWasCancelled = false + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + await refreshGate.run() + refreshWasCancelled = Task.isCancelled + } + defer { store._test_providerRefreshOverride = nil } + + let ownerTask = Task { + await store.refreshProvider(.codex) + } + await refreshGate.waitUntilStarted(count: 1) + let sharedWaiterTask = Task { + await store.refreshProvider(.codex, coalesceIfRefreshing: true) + } + for _ in 0..<40 { + await Task.yield() + } + + ownerTask.cancel() + await refreshGate.releaseFirst() + await ownerTask.value + await sharedWaiterTask.value + + #expect(!refreshWasCancelled) + #expect(await refreshGate.startCount == 1) + } + + @Test + func `background refresh retries canceled provider probe with cached data`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "cached@example.com"), + provider: .codex) + let refreshGate = BlockingInstantOpenProviderRefresh() + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + await refreshGate.run() + } + defer { store._test_providerRefreshOverride = nil } + + let ownerTask = Task { + await store.refreshProvider(.codex) + } + await refreshGate.waitUntilStarted(count: 1) + ownerTask.cancel() + let backgroundTask = Task { + await store.refreshProvider(.codex, coalesceIfRefreshing: true) + } + for _ in 0..<40 { + await Task.yield() + } + await refreshGate.releaseFirst() + await ownerTask.value + await backgroundTask.value + + #expect(await refreshGate.startCount == 2) + } + + @Test + func `menu open refresh only retries the displayed provider`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableProvidersForInstantOpenTesting([.claude, .codex], settings: settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + store.refreshingProviders.insert(.claude) + defer { store.refreshingProviders.remove(.claude) } + var refreshedProviders: [UsageProvider] = [] + store._test_providerRefreshOverride = { provider in + refreshedProviders.append(provider) + } + defer { store._test_providerRefreshOverride = nil } + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuProviders[ObjectIdentifier(menu)] = .codex + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + for _ in 0..<40 where refreshedProviders.isEmpty { + await Task.yield() + } + + #expect(refreshedProviders == [.codex]) + } + + @Test + func `opening fresh split menu preserves another provider deferred retry`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableProvidersForInstantOpenTesting([.claude, .codex], settings: settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "claude@example.com", provider: .claude), + provider: .claude) + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.seconds(60)) + StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.seconds(60)) + defer { + StatusItemController.resetMenuOpenRefreshDelayForTesting() + StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() + } + + let codexMenu = controller.makeMenu(for: .codex) + controller.menuWillOpen(codexMenu) + controller.menuDidClose(codexMenu) + #expect(controller.deferredMenuInteractionRefreshProviders == [.codex]) + + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + let claudeMenu = controller.makeMenu(for: .claude) + controller.menuWillOpen(claudeMenu) + defer { controller.menuDidClose(claudeMenu) } + for _ in 0..<40 { + await Task.yield() + } + + #expect(controller.deferredMenuInteractionRefreshProviders == [.codex]) + #expect(controller.deferredMenuInteractionRefreshPending) + } + + @Test + func `overview defers only providers that need retry`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + settings.mergedMenuLastSelectedWasOverview = true + self.enableProvidersForInstantOpenTesting([.claude, .codex], settings: settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + store._setSnapshotForTesting( + self.instantOpenSnapshot(email: "claude@example.com", provider: .claude), + provider: .claude) + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.seconds(60)) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + + #expect(controller.deferredMenuInteractionRefreshProviders == [.codex]) + } + + @Test + func `closing overview menu stops before refreshing another provider`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = true + settings.mergedMenuLastSelectedWasOverview = true + self.enableProvidersForInstantOpenTesting([.codex, .openai], settings: settings) + settings.updateProviderConfig(provider: .openai) { config in + config.apiKey = "test-openai-key" + } + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + store._setSnapshotForTesting(nil, provider: .openai) + let refreshGate = BlockingInstantOpenProviderRefresh() + store._test_providerRefreshOverride = { _ in + await refreshGate.run() + } + defer { store._test_providerRefreshOverride = nil } + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.seconds(60)) + defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + await refreshGate.waitUntilStarted(count: 1) + controller.menuDidClose(menu) + await refreshGate.releaseFirst() + for _ in 0..<80 { + await Task.yield() + } + + #expect(await refreshGate.startCount == 1) + #expect(controller.deferredMenuInteractionRefreshPending) + } + + @Test + func `closing menu during missing data refresh preserves deferred retry`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + let refreshGate = BlockingInstantOpenProviderRefresh() + var refreshInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + refreshInteractions.append(ProviderInteractionContext.current) + await refreshGate.run() + } + defer { store._test_providerRefreshOverride = nil } + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.zero) + defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + await refreshGate.waitUntilStarted(count: 1) + #expect(controller.deferredMenuInteractionRefreshPending) + #expect(store.refreshingProviders.contains(.codex)) + + let periodicRefreshTask = Task { + await store.refresh() + } + for _ in 0..<40 { + await Task.yield() + } + #expect(await refreshGate.startCount == 1) + + controller.menuDidClose(menu) + #expect(controller.deferredMenuInteractionRefreshPending) + await refreshGate.releaseFirst() + await periodicRefreshTask.value + for _ in 0..<40 where store.isRefreshing { + await Task.yield() + } + for _ in 0..<40 { + await Task.yield() + } + #expect(controller.deferredMenuInteractionRefreshPending) + + controller.scheduleDeferredMenuInteractionRefreshIfNeeded(delay: .zero) + await refreshGate.waitUntilStarted(count: 2) + for _ in 0..<40 where controller.deferredMenuInteractionRefreshPending { + await Task.yield() + } + + #expect(await refreshGate.startCount == 2) + #expect(refreshInteractions == [.background, .background]) + #expect(!controller.deferredMenuInteractionRefreshPending) + } + + @Test + func `closing menu during successful missing data refresh clears deferred retry`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodexForInstantOpenTesting(settings) + + let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) + store._setSnapshotForTesting(nil, provider: .codex) + let refreshGate = BlockingInstantOpenProviderRefresh() + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + await refreshGate.run() + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow( + usedPercent: 25, + windowMinutes: 300, + resetsAt: Date().addingTimeInterval(3600), + resetDescription: nil), + secondary: nil, + updatedAt: Date()), + provider: .codex) + } + defer { store._test_providerRefreshOverride = nil } + + let controller = StatusItemController( + store: store, + settings: settings, + account: UsageFetcher().loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + defer { controller.releaseStatusItemsForTesting() } + controller.menuRefreshEnabledOverrideForTesting = true + StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.zero) + defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } + StatusItemController.setMenuOpenRefreshDelayForTesting(.zero) + defer { StatusItemController.resetMenuOpenRefreshDelayForTesting() } + + let menu = controller.makeMenu() + controller.menuWillOpen(menu) + await refreshGate.waitUntilStarted(count: 1) + controller.menuDidClose(menu) + await refreshGate.releaseFirst() + + for _ in 0..<80 where controller.deferredMenuInteractionRefreshPending { + await Task.yield() + } + for _ in 0..<40 { + await Task.yield() + } + + #expect(await refreshGate.startCount == 1) + #expect(!controller.deferredMenuInteractionRefreshPending) + } + + private func enableOnlyCodexForInstantOpenTesting(_ settings: SettingsStore) { + self.enableProvidersForInstantOpenTesting([.codex], settings: settings) + } + + private func instantOpenSnapshot( + email: String, + provider: UsageProvider = .codex, + percent: Double = 25) -> UsageSnapshot + { + UsageSnapshot( + primary: RateWindow( + usedPercent: percent, + windowMinutes: 300, + resetsAt: Date().addingTimeInterval(3600), + resetDescription: nil), + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: provider, + accountEmail: email, + accountOrganization: nil, + loginMethod: "ChatGPT")) + } + + private func enableProvidersForInstantOpenTesting( + _ enabledProviders: Set, + settings: SettingsStore) + { + let registry = ProviderRegistry.shared + for provider in UsageProvider.allCases { + guard let metadata = registry.metadata[provider] else { continue } + settings.setProviderEnabled( + provider: provider, + metadata: metadata, + enabled: enabledProviders.contains(provider)) + } + } +} + +private struct InstantOpenProviderFetchStrategy: ProviderFetchStrategy { + let loader: @Sendable () async -> UsageSnapshot + + var id: String { + "instant-open-provider-refresh-test" + } + + var kind: ProviderFetchKind { + .cli + } + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { + let usage = await self.loader() + return self.makeResult(usage: usage, sourceLabel: self.id) + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +private struct InstantOpenProviderMutationFetchStrategy: ProviderFetchStrategy { + let mutations: OrderedInstantOpenProviderMutation + + let id = "instant-open-provider-mutation-test" + let kind: ProviderFetchKind = .web + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let token = await self.mutations.awaitToken() + await context.providerManualTokenUpdater?(.stepfun, token) + let usage = UsageSnapshot( + primary: RateWindow( + usedPercent: 25, + windowMinutes: 300, + resetsAt: Date().addingTimeInterval(3600), + resetDescription: nil), + secondary: nil, + updatedAt: Date()) + return self.makeResult(usage: usage, sourceLabel: self.id) + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +private actor OrderedInstantOpenProviderRefresh { + private var started = 0 + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var continuations: [Int: CheckedContinuation] = [:] + + var startCount: Int { + self.started + } + + func awaitSnapshot() async -> UsageSnapshot { + self.started += 1 + let call = self.started + self.resumeReadyStartWaiters() + return await withCheckedContinuation { continuation in + self.continuations[call] = continuation + } + } + + func waitUntilStarted(count: Int) async { + if self.started >= count { return } + await withCheckedContinuation { continuation in + self.startWaiters.append((count: count, continuation: continuation)) + } + } + + func resume(call: Int, snapshot: UsageSnapshot) { + self.continuations.removeValue(forKey: call)?.resume(returning: snapshot) + } + + private func resumeReadyStartWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.startWaiters { + if self.started >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.startWaiters = remaining + } +} + +private actor OrderedInstantOpenProviderMutation { + private var started = 0 + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var continuations: [Int: CheckedContinuation] = [:] + + var startCount: Int { + self.started + } + + func awaitToken() async -> String { + self.started += 1 + let call = self.started + self.resumeReadyStartWaiters() + return await withCheckedContinuation { continuation in + self.continuations[call] = continuation + } + } + + func waitUntilStarted(count: Int) async { + if self.started >= count { return } + await withCheckedContinuation { continuation in + self.startWaiters.append((count: count, continuation: continuation)) + } + } + + func resume(call: Int, token: String) { + self.continuations.removeValue(forKey: call)?.resume(returning: token) + } + + private func resumeReadyStartWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.startWaiters { + if self.started >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.startWaiters = remaining + } +} + +private actor BlockingInstantOpenProviderRefresh { + private var started = 0 + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var firstReleaseWaiters: [CheckedContinuation] = [] + private var firstReleased = false + + var startCount: Int { + self.started + } + + func run() async { + self.started += 1 + self.resumeReadyStartWaiters() + guard self.started == 1, !self.firstReleased else { return } + await withCheckedContinuation { continuation in + self.firstReleaseWaiters.append(continuation) + } + } + + func waitUntilStarted(count: Int) async { + if self.started >= count { return } + await withCheckedContinuation { continuation in + self.startWaiters.append((count: count, continuation: continuation)) + } + } + + func releaseFirst() { + self.firstReleased = true + let waiters = self.firstReleaseWaiters + self.firstReleaseWaiters.removeAll() + waiters.forEach { $0.resume() } + } + + private func resumeReadyStartWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.startWaiters { + if self.started >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.startWaiters = remaining + } +} diff --git a/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift b/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift index e25496bbab..4a523d593a 100644 --- a/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift +++ b/Tests/CodexBarTests/StatusMenuOpenRefreshTests.swift @@ -5,112 +5,6 @@ import Testing @testable import CodexBar extension StatusMenuTests { - @Test - func `opening fresh menu does not schedule deferred refresh`() async { - self.disableMenuCardsForTesting() - let settings = self.makeSettings() - settings.statusChecksEnabled = false - settings.refreshFrequency = .manual - settings.mergeIcons = false - self.enableOnlyCodex(settings) - - let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) - var providerRefreshCount = 0 - var refreshInteractions: [ProviderInteraction] = [] - store._test_providerRefreshOverride = { provider in - guard provider == .codex else { return } - refreshInteractions.append(ProviderInteractionContext.current) - providerRefreshCount += 1 - } - defer { store._test_providerRefreshOverride = nil } - - let controller = StatusItemController( - store: store, - settings: settings, - account: UsageFetcher().loadAccountInfo(), - updater: DisabledUpdaterController(), - preferencesSelection: PreferencesSelection(), - statusBar: self.makeStatusBarForTesting()) - defer { controller.releaseStatusItemsForTesting() } - StatusItemController.setClosedMenuPreparationDelayForTesting(.zero) - defer { StatusItemController.resetClosedMenuPreparationDelayForTesting() } - - controller.menuRefreshEnabledOverrideForTesting = true - StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.zero) - defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } - - let menu = controller.makeMenu() - controller.menuWillOpen(menu) - - for _ in 0..<20 { - await Task.yield() - } - #expect(providerRefreshCount == 0) - #expect(!controller.deferredMenuInteractionRefreshPending) - - controller.menuDidClose(menu) - for _ in 0..<40 { - await Task.yield() - } - - #expect(providerRefreshCount == 0) - #expect(refreshInteractions.isEmpty) - } - - @Test - func `menu open with missing data defers automatic refresh until tracking ends`() async { - self.disableMenuCardsForTesting() - let settings = self.makeSettings() - settings.statusChecksEnabled = false - settings.refreshFrequency = .manual - settings.mergeIcons = false - self.enableOnlyCodex(settings) - - let store = self.makeCodexStore(settings: settings, dashboardAuthorized: false) - store._setSnapshotForTesting(nil, provider: .codex) - var providerRefreshCount = 0 - var refreshInteractions: [ProviderInteraction] = [] - store._test_providerRefreshOverride = { provider in - guard provider == .codex else { return } - refreshInteractions.append(ProviderInteractionContext.current) - providerRefreshCount += 1 - } - defer { store._test_providerRefreshOverride = nil } - - let controller = StatusItemController( - store: store, - settings: settings, - account: UsageFetcher().loadAccountInfo(), - updater: DisabledUpdaterController(), - preferencesSelection: PreferencesSelection(), - statusBar: self.makeStatusBarForTesting()) - defer { controller.releaseStatusItemsForTesting() } - StatusItemController.setClosedMenuPreparationDelayForTesting(.zero) - defer { StatusItemController.resetClosedMenuPreparationDelayForTesting() } - - controller.menuRefreshEnabledOverrideForTesting = true - StatusItemController.setDeferredMenuInteractionRefreshDelayForTesting(.zero) - defer { StatusItemController.resetDeferredMenuInteractionRefreshDelayForTesting() } - - let menu = controller.makeMenu() - controller.menuWillOpen(menu) - - for _ in 0..<20 { - await Task.yield() - } - #expect(providerRefreshCount == 0) - #expect(controller.deferredMenuInteractionRefreshPending) - - controller.menuDidClose(menu) - for _ in 0..<40 where providerRefreshCount == 0 { - await Task.yield() - } - - #expect(providerRefreshCount == 1) - #expect(refreshInteractions == [.background]) - #expect(!controller.deferredMenuInteractionRefreshPending) - } - @Test func `store observation marks open menu stale without rebuilding during tracking`() async { self.disableMenuCardsForTesting() @@ -193,10 +87,20 @@ extension StatusMenuTests { let menu = controller.makeMenu() controller.mergedMenu = menu controller.statusItem.menu = menu + for _ in 0..<20 { + await Task.yield() + } controller.populateMenu(menu, provider: nil) controller.markMenuFresh(menu) let key = ObjectIdentifier(menu) + for _ in 0..<40 { + await Task.yield() + } + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + controller.cancelAllClosedMenuRebuilds() + controller.closedMenusDeferredUntilNextOpen.removeAll(keepingCapacity: false) let openedVersion = controller.menuVersions[key] // Background data-refresh tick (stale allowed): closed prep is skipped entirely, so @@ -256,6 +160,12 @@ extension StatusMenuTests { controller.populateMenu(menu, provider: nil) controller.markMenuFresh(menu) let key = ObjectIdentifier(menu) + for _ in 0..<40 { + await Task.yield() + } + controller.populateMenu(menu, provider: nil) + controller.markMenuFresh(menu) + controller.cancelAllClosedMenuRebuilds() let openedVersion = controller.menuVersions[key] controller.invalidateMenus(allowStaleContentDuringDataRefresh: true) @@ -270,6 +180,9 @@ extension StatusMenuTests { controller.menuWillOpen(menu) defer { controller.menuDidClose(menu) } + for _ in 0..<40 where controller.menuVersions[key] != controller.menuContentVersion { + await Task.yield() + } #expect(controller.menuVersions[key] == controller.menuContentVersion) } diff --git a/Tests/CodexBarTests/StatusMenuTests.swift b/Tests/CodexBarTests/StatusMenuTests.swift index d59357daa6..73160c41dc 100644 --- a/Tests/CodexBarTests/StatusMenuTests.swift +++ b/Tests/CodexBarTests/StatusMenuTests.swift @@ -563,12 +563,19 @@ struct StatusMenuTests { settings.usageBarsShowUsed = true controller.handleProviderConfigChange(reason: "usageBarsShowUsed") - for _ in 0..<20 - where initialSwitcherID == (menu.items.first?.view as? ProviderSwitcherView).map(ObjectIdentifier.init) - { + for _ in 0..<20 { await Task.yield() } + #expect(controller.parentMenuRebuildsDeferredDuringTracking.contains(ObjectIdentifier(menu))) + if let initialSwitcherID, let currentSwitcher = menu.items.first?.view as? ProviderSwitcherView { + #expect(initialSwitcherID == ObjectIdentifier(currentSwitcher)) + } + + controller.menuDidClose(menu) + controller.menuWillOpen(menu) + defer { controller.menuDidClose(menu) } + let updatedSwitcher = menu.items.first?.view as? ProviderSwitcherView #expect(updatedSwitcher != nil) if let initialSwitcherID, let updatedSwitcher { diff --git a/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift b/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift index 49053cc937..1e505ce00e 100644 --- a/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift +++ b/Tests/CodexBarTests/StatusMenuTokenAccountSwitcherTests.swift @@ -131,10 +131,15 @@ final class StatusMenuTokenAccountSwitcherTests: XCTestCase { let switcher = try XCTUnwrap(menu.items.compactMap { $0.view as? TokenAccountSwitcherView }.first) let selectionTask = try XCTUnwrap(switcher._test_select(index: 1)) - await blocker.waitUntilStarted(count: 2) XCTAssertEqual(settings.tokenAccountsData(for: .claude)?.clampedActiveIndex(), 1) + for _ in 0..<40 { + await Task.yield() + } + let startedBeforeDrain = await blocker.startedCallCount() + XCTAssertEqual(startedBeforeDrain, 1) await blocker.resumeAll(with: .success(self.snapshot(percent: 17))) + await blocker.waitUntilStarted(count: 2) await selectionTask.value await refreshTask.value let startedCallCount = await blocker.startedCallCount() From 01748970b5a9208240c08fa39d87affba5df6259 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Wed, 10 Jun 2026 21:12:22 +0100 Subject: [PATCH 2/2] docs: note instant cached menu opening --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 80e3f43641..c7636ecd9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ - Doubao: confirm zero-remaining HTTP 200 request limits before falling back, preserving genuine exhaustion and avoiding false 100% usage (#1383). Thanks @LeoLin990405 and @foobra! - Menu bar: defer pasteboard writes and copy feedback outside the `NSMenu` tracking callback so in-menu copy buttons no longer beachball on macOS 26 (#1388). Thanks @LeoLin990405! - Menu bar: defer Overview-row provider transitions out of AppKit's click callback so opening provider detail no longer performs a full synchronous menu rebuild (#1325). +- Menu bar: open cached menus immediately after data-only invalidations, then refresh missing or stale provider data asynchronously without queuing redundant work on close (#1398). Thanks @joshuavial! - Menu bar: recycle SwiftUI card hosting views across data refreshes and provider switches, and reconcile matching menu rows in place instead of removing and reinserting every row, cutting open-click, switch, and idle rebuild cost (#1394). Thanks @bcssewl! - Development: disable Keychain access for unbundled executables to avoid repeated password prompts while preserving packaged app behavior (#1271). Thanks @Yuxin-Qiao! - Antigravity: exclude model quotas without a remaining fraction from family summaries so they no longer mask tracked usage in the automatic menu-bar metric (#1369). Thanks @Martin-Hausleitner!