From 4477ec469cc0fe28fb3b2bd670d6092a7d3f99f6 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:37:09 +0800 Subject: [PATCH 1/4] fix: keep manual refresh responsive --- .../StatusItemController+Actions.swift | 2 +- .../CodexBar/UsageStore+AdaptiveRefresh.swift | 15 +++++++++++++++ Sources/CodexBar/UsageStore.swift | 14 ++++++++++++-- .../AdaptiveRefreshTimerTests.swift | 19 +++++++++++++++++++ .../StatusMenuPersistentRefreshTests.swift | 8 ++++---- 5 files changed, 51 insertions(+), 7 deletions(-) diff --git a/Sources/CodexBar/StatusItemController+Actions.swift b/Sources/CodexBar/StatusItemController+Actions.swift index 4bf5d99cf2..3d616a2aec 100644 --- a/Sources/CodexBar/StatusItemController+Actions.swift +++ b/Sources/CodexBar/StatusItemController+Actions.swift @@ -179,7 +179,7 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { interaction: .userInitiated) } else { await self.performStoreRefresh( - forceTokenUsage: true, + forceTokenUsage: false, refreshOpenMenusWhenComplete: true, interaction: .userInitiated) } diff --git a/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift b/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift index 670ef113d7..4a4c85befb 100644 --- a/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift +++ b/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift @@ -34,6 +34,21 @@ extension UsageStore { return candidate < scheduledAt } + /// Advances a fixed timer from the last scheduled tick instead of the refresh completion time. + /// Missed ticks are skipped so a refresh that runs longer than its interval does not create + /// overlapping catch-up refreshes. + nonisolated static func nextFixedTimerScheduledAt( + previousScheduledAt: ContinuousClock.Instant, + completedAt: ContinuousClock.Instant, + interval: Duration) -> ContinuousClock.Instant + { + var scheduledAt = previousScheduledAt + interval + while scheduledAt <= completedAt { + scheduledAt += interval + } + return scheduledAt + } + func logAdaptiveRefreshDecision(_ decision: AdaptiveRefreshPolicy.Decision) { // Reason and delay only; never provider/account/email/path/credential/response data. // No "adaptive refresh: " prefix — the adaptiveRefresh log category already identifies the source. diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index abc3937ff7..6a0cbf91f6 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -769,13 +769,23 @@ final class UsageStore { guard let wait = frequency.seconds else { return } // Background poller so the menu stays responsive; canceled when settings change or store deallocates. - // `self` is only briefly borrowed to read the (DEBUG-only) sleep override, never held across the sleep. + // Fixed cadence is anchored to the scheduled tick time, not refresh completion, so slow provider + // work doesn't permanently stretch a two-minute interval into "refresh duration + two minutes". self.timerTask = Task.detached(priority: .utility) { [weak self] in + let interval = Duration.seconds(wait) + let clock = ContinuousClock() + var scheduledAt = clock.now + interval while !Task.isCancelled { - let sleepDuration = await self?.effectiveTimerSleepDuration(.seconds(wait)) ?? .seconds(wait) + let now = clock.now + let computedSleep = now >= scheduledAt ? .zero : scheduledAt - now + let sleepDuration = await self?.effectiveTimerSleepDuration(computedSleep) ?? computedSleep try? await Task.sleep(for: sleepDuration) guard !Task.isCancelled else { return } await self?.refresh() + scheduledAt = Self.nextFixedTimerScheduledAt( + previousScheduledAt: scheduledAt, + completedAt: clock.now, + interval: interval) } } } diff --git a/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift b/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift index 84f3aedffe..456773c338 100644 --- a/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift +++ b/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift @@ -130,6 +130,25 @@ struct AdaptiveRefreshTimerTests { #expect(store.completedRefreshCountForTesting >= 2) } + @Test + func `fixed cadence advances from scheduled tick instead of refresh completion`() { + let interval = Duration.milliseconds(100) + let start = ContinuousClock.now + let firstScheduledAt = start + interval + + let nextAfterSlowRefresh = UsageStore.nextFixedTimerScheduledAt( + previousScheduledAt: firstScheduledAt, + completedAt: firstScheduledAt + .milliseconds(60), + interval: interval) + #expect(nextAfterSlowRefresh == start + .milliseconds(200)) + + let nextAfterMissedTicks = UsageStore.nextFixedTimerScheduledAt( + previousScheduledAt: firstScheduledAt, + completedAt: firstScheduledAt + .milliseconds(260), + interval: interval) + #expect(nextAfterMissedTicks == start + .milliseconds(400)) + } + @Test func `adaptive mode keeps recomputing and refreshing across menu-open changes`() async throws { let settings = Self.makeSettingsStore(suite: "AdaptiveRefreshTimerTests-adaptive", frequency: .adaptive) diff --git a/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift b/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift index 76a6348326..06cee511a8 100644 --- a/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift +++ b/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift @@ -1201,17 +1201,17 @@ extension StatusMenuPersistentRefreshTests { } controller.refreshNow() + let manualTask = try? #require(controller.manualRefreshTasks[.global]) await tokenRefreshStarted.wait() + await manualTask?.value - #expect(controller.store.isRefreshing) + #expect(!controller.store.isRefreshing) #expect(controller.store.refreshingProviders.isEmpty) - #expect(monitor.isManualRefreshInFlight) + #expect(!monitor.isManualRefreshInFlight) #expect(!monitor.isManualRefreshInFlight(for: .codex)) #expect(monitor.subtitle(for: .codex, fallback: fallback).style == .info) releaseTokenRefresh.resume() - await controller.manualRefreshTasks[.global]?.value - #expect(!monitor.isManualRefreshInFlight) } @Test From 47f93632c0c33def44ab09fed7d18bb61b1dcebd Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 10 Jul 2026 15:13:18 -0700 Subject: [PATCH 2/4] fix: serialize manual refresh enrichment Co-authored-by: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> --- CHANGELOG.md | 1 + .../CodexBar/ProviderRefreshCoordinator.swift | 8 +- .../Codex/UsageStore+CodexAccountState.swift | 1 + .../StatusItemController+Actions.swift | 33 +- ...temController+MenuInteractionRefresh.swift | 5 +- ...tatusItemController+MenuPresentation.swift | 5 + ...ItemController+PersistentMenuActions.swift | 4 + .../StatusItemController+Shutdown.swift | 2 + .../CodexBar/UsageStore+AdaptiveRefresh.swift | 29 + Sources/CodexBar/UsageStore+OpenAIWeb.swift | 4 + .../CodexBar/UsageStore+ProviderRuntime.swift | 21 + Sources/CodexBar/UsageStore+Refresh.swift | 3 +- .../UsageStore+RefreshEnrichment.swift | 325 ++++ .../UsageStore+ResetBoundaryRefresh.swift | 5 +- .../UsageStore+StartupConnectivityRetry.swift | 4 +- .../UsageStore+TokenRefreshSequence.swift | 122 ++ Sources/CodexBar/UsageStore.swift | 287 ++-- .../BrowserCookieAccessGate.swift | 22 + Sources/CodexBarCore/KeychainAccessGate.swift | 1 + .../KeychainAccessPreflight.swift | 1 + .../AdaptiveRefreshTimerTests.swift | 73 +- ...CodexAccountScopedRefreshTestSupport.swift | 33 +- ...odexBackgroundRefreshCoalescingTests.swift | 1321 ++++++++++++++++- .../CodexManagedOpenAIWebRefreshTests.swift | 126 +- ...erredMenuInteractionRefreshTailTests.swift | 177 +++ .../PersistentRefreshAccessibilityTests.swift | 31 + .../ProviderRefreshCoordinatorTests.swift | 12 + .../RequiredRefreshCoalescingTests.swift | 487 ++++++ .../StatusItemControllerShutdownTests.swift | 192 ++- .../StatusMenuInstantOpenTests.swift | 12 +- .../StatusMenuPersistentRefreshTests.swift | 263 +++- .../UsageStoreManualTokenRefreshTests.swift | 63 +- .../UsageStoreResetBoundaryRefreshTests.swift | 109 ++ 33 files changed, 3473 insertions(+), 309 deletions(-) create mode 100644 Sources/CodexBar/UsageStore+ProviderRuntime.swift create mode 100644 Sources/CodexBar/UsageStore+RefreshEnrichment.swift create mode 100644 Sources/CodexBar/UsageStore+TokenRefreshSequence.swift create mode 100644 Tests/CodexBarTests/DeferredMenuInteractionRefreshTailTests.swift create mode 100644 Tests/CodexBarTests/PersistentRefreshAccessibilityTests.swift create mode 100644 Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index b4ab322d31..90f0bc2a94 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ - Quota warnings: add opt-in predictive pace alerts for Codex and Claude session and weekly limits, with one alert per risk episode. Thanks @vincent-peng! ### Fixed +- Refresh: keep all-provider manual refresh responsive while forced cost, credit, and dashboard enrichment finishes in a serialized background tail, and keep fixed intervals anchored to scheduled ticks. Thanks @Yuxin-Qiao! - Agent Sessions: keep local and remote session discovery off by default while preserving the General setting for explicit opt-in. - Menus: stop completed provider cards and plan-utilization rows from remaining in “Refreshing…” while unrelated provider or token-cost work is still running. Thanks @Yuxin-Qiao! - Menu: keep native hover highlights aligned by deferring geometry-changing open-menu rebuilds until the pointer leaves the row. Thanks @Zihao-Qi! diff --git a/Sources/CodexBar/ProviderRefreshCoordinator.swift b/Sources/CodexBar/ProviderRefreshCoordinator.swift index 2ebc29c116..fc7ed25a27 100644 --- a/Sources/CodexBar/ProviderRefreshCoordinator.swift +++ b/Sources/CodexBar/ProviderRefreshCoordinator.swift @@ -22,7 +22,9 @@ final class ProviderRefreshCoordinator { func coalescingState(for key: Key) -> ProviderRefreshTaskState? { guard let latestGeneration = self.latestGenerations[key] else { return nil } - return self.states[key]?.last { $0.generation == latestGeneration } + return self.states[key]?.last { state in + state.generation == latestGeneration && !state.isCompleted + } } func beginReplacingRequest(for key: Key) -> Request { @@ -177,6 +179,10 @@ final class ProviderRefreshTaskState: @unchecked Sendable { self.lock.withLock { self.retryRequired } } + fileprivate var isCompleted: Bool { + self.lock.withLock { self.completed } + } + var canRemove: Bool { self.lock.withLock { self.completed && self.waiterIDs.isEmpty } } diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift index 6f034179fc..1197c505db 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift @@ -51,6 +51,7 @@ extension UsageStore { await self.refreshOpenAIDashboardIfNeeded( force: true, expectedGuard: expectedGuard, + bypassCoalescing: true, allowCodexUsageBackfill: true) phaseDidChange?(.dashboard) } diff --git a/Sources/CodexBar/StatusItemController+Actions.swift b/Sources/CodexBar/StatusItemController+Actions.swift index 3d616a2aec..6f614252c4 100644 --- a/Sources/CodexBar/StatusItemController+Actions.swift +++ b/Sources/CodexBar/StatusItemController+Actions.swift @@ -51,12 +51,31 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { } } + func performStoreRefresh( + enrichmentMode: UsageStore.RefreshEnrichmentMode, + refreshOpenMenusWhenComplete: Bool, + interaction: ProviderInteraction) async + { + await self.withProviderInteraction(interaction) { + await self.store.refresh(enrichmentMode: enrichmentMode) + guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return } + self.store.scheduleStorageFootprintRefreshForOverview(force: true) + if refreshOpenMenusWhenComplete { + self.refreshOpenMenusAfterExplicitStoreAction() + } else { + self.invalidateMenus() + } + } + } + func performStoreRefresh( for provider: UsageProvider, refreshOpenMenusWhenComplete: Bool, interaction: ProviderInteraction) async { await self.withProviderInteraction(interaction) { + await self.store.awaitForcedRefreshEnrichment() + guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return } let refreshStartedAt = Date() await self.store.refreshProvider(provider) guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return } @@ -69,7 +88,8 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return } await self.store.refreshOpenAIDashboardIfNeeded( force: true, - expectedGuard: self.store.freshCodexOpenAIWebRefreshGuard()) + expectedGuard: self.store.freshCodexOpenAIWebRefreshGuard(), + bypassCoalescing: true) guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return } if self.store.openAIDashboardRequiresLogin { await self.store.refreshProvider(.codex) @@ -151,6 +171,7 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { guard !self.hasPreparedForAppShutdown, self.manualRefreshTasks[scope] == nil, !conflictsWithOtherScope, + !self.store.hasForcedRefreshEnrichmentInFlight, !self.store.isRefreshing, !scopedRefreshInFlight else { return } @@ -179,7 +200,7 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { interaction: .userInitiated) } else { await self.performStoreRefresh( - forceTokenUsage: false, + enrichmentMode: .forcedBackground, refreshOpenMenusWhenComplete: true, interaction: .userInitiated) } @@ -785,8 +806,12 @@ extension StatusItemController: StatusItemMenuPersistentActionDelegate { private func trimmedLoginOutput(_ text: String) -> String { let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) let limit = 600 - if trimmed.isEmpty { return L("No output captured.") } - if trimmed.count <= limit { return trimmed } + if trimmed.isEmpty { + return L("No output captured.") + } + if trimmed.count <= limit { + return trimmed + } let idx = trimmed.index(trimmed.startIndex, offsetBy: limit) return "\(trimmed[.. Bool { + self.isRowEnabled + } + override func accessibilityPerformPress() -> Bool { + guard self.isRowEnabled else { return false } guard let onClick = self.onClick else { return super.accessibilityPerformPress() } diff --git a/Sources/CodexBar/StatusItemController+PersistentMenuActions.swift b/Sources/CodexBar/StatusItemController+PersistentMenuActions.swift index f04dc52b51..4a03556182 100644 --- a/Sources/CodexBar/StatusItemController+PersistentMenuActions.swift +++ b/Sources/CodexBar/StatusItemController+PersistentMenuActions.swift @@ -16,6 +16,10 @@ extension StatusItemController { } func isRefreshActionInFlight(for menu: NSMenu) -> Bool { + if self.store.hasForcedRefreshEnrichmentInFlight { + return true + } + // An all-providers manual refresh (⌘R / overview) legitimately busies every row. if self.manualRefreshTasks[.global] != nil { return true diff --git a/Sources/CodexBar/StatusItemController+Shutdown.swift b/Sources/CodexBar/StatusItemController+Shutdown.swift index 1ca78b57a9..223ea7e362 100644 --- a/Sources/CodexBar/StatusItemController+Shutdown.swift +++ b/Sources/CodexBar/StatusItemController+Shutdown.swift @@ -33,6 +33,8 @@ extension StatusItemController { task.cancel() } self.manualRefreshTasks.removeAll() + self.store.cancelForcedRefreshEnrichment() + self.store.cancelRequiredRefresh() self.menuCardRefreshMonitor.resetManualRefresh() self.screenChangeVisibilityTask?.cancel() self.screenChangeVisibilityTask = nil diff --git a/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift b/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift index 4a4c85befb..05ab36cd57 100644 --- a/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift +++ b/Sources/CodexBar/UsageStore+AdaptiveRefresh.swift @@ -42,6 +42,7 @@ extension UsageStore { completedAt: ContinuousClock.Instant, interval: Duration) -> ContinuousClock.Instant { + precondition(interval > .zero) var scheduledAt = previousScheduledAt + interval while scheduledAt <= completedAt { scheduledAt += interval @@ -49,6 +50,34 @@ extension UsageStore { return scheduledAt } + nonisolated static func runFixedRefreshTimer( + interval: Duration, + sleepOverride: Duration? = nil, + now: @escaping @Sendable () async -> ContinuousClock.Instant = { ContinuousClock.now }, + sleep: @escaping @Sendable (Duration) async throws -> Void = { duration in + try await Task.sleep(for: duration) + }, + refresh: @escaping @Sendable () async -> Void) async + { + precondition(interval > .zero) + var scheduledAt = await now() + interval + while !Task.isCancelled { + let current = await now() + let computedSleep = current >= scheduledAt ? .zero : scheduledAt - current + do { + try await sleep(sleepOverride ?? computedSleep) + } catch { + return + } + guard !Task.isCancelled else { return } + await refresh() + scheduledAt = await self.nextFixedTimerScheduledAt( + previousScheduledAt: scheduledAt, + completedAt: now(), + interval: interval) + } + } + func logAdaptiveRefreshDecision(_ decision: AdaptiveRefreshPolicy.Decision) { // Reason and delay only; never provider/account/email/path/credential/response data. // No "adaptive refresh: " prefix — the adaptiveRefresh log category already identifies the source. diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index 4ed5a410f5..0abc974835 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -430,6 +430,10 @@ extension UsageStore { await task.value return } + if bypassCoalescing { + self.openAIDashboardBackgroundRefreshTask?.cancel() + self.openAIDashboardRefreshTask?.cancel() + } self.handleOpenAIWebTargetEmailChangeIfNeeded( targetEmail: targetEmail, targetScope: self.codexCookieCacheScopeForOpenAIWeb()) diff --git a/Sources/CodexBar/UsageStore+ProviderRuntime.swift b/Sources/CodexBar/UsageStore+ProviderRuntime.swift new file mode 100644 index 0000000000..860a4705ee --- /dev/null +++ b/Sources/CodexBar/UsageStore+ProviderRuntime.swift @@ -0,0 +1,21 @@ +import CodexBarCore + +extension UsageStore { + func performRuntimeAction(_ action: ProviderRuntimeAction, for provider: UsageProvider) async { + guard let runtime = self.providerRuntimes[provider] else { return } + let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self) + await runtime.perform(action: action, context: context) + } + + func updateProviderRuntimes() { + for (provider, runtime) in self.providerRuntimes { + let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self) + if self.isEnabled(provider) { + runtime.start(context: context) + } else { + runtime.stop(context: context) + } + runtime.settingsDidChange(context: context) + } + } +} diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index 134d54ac5e..f9541a7c78 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -35,7 +35,8 @@ extension UsageStore { func refreshForSettingsChange() async { await self.runRefresh( startupConnectivityRetryAttempt: nil, - coalesceProviderRefreshesOverride: false) + coalesceProviderRefreshesOverride: false, + waitForRefreshAvailability: true) } func prepareRefreshState(for provider: UsageProvider? = nil) { diff --git a/Sources/CodexBar/UsageStore+RefreshEnrichment.swift b/Sources/CodexBar/UsageStore+RefreshEnrichment.swift new file mode 100644 index 0000000000..c755db1dcd --- /dev/null +++ b/Sources/CodexBar/UsageStore+RefreshEnrichment.swift @@ -0,0 +1,325 @@ +import CodexBarCore +import Foundation + +extension UsageStore { + enum RefreshEnrichmentMode: Equatable, Sendable { + case automatic + case forcedForeground + case forcedBackground + } + + struct RequiredRefreshRequest: Sendable { + var throughGeneration: UInt64 + var startupConnectivityRetryAttempt: Int? + var coalesceProviderRefreshes: Bool + var interaction: ProviderInteraction + + mutating func merge(_ newer: Self) { + self.throughGeneration = max(self.throughGeneration, newer.throughGeneration) + if let newerAttempt = newer.startupConnectivityRetryAttempt { + self.startupConnectivityRetryAttempt = max( + self.startupConnectivityRetryAttempt ?? newerAttempt, + newerAttempt) + } + // Replacement is the stronger policy: a settings change must not join work started + // with the old configuration merely because another required refresh arrived first. + self.coalesceProviderRefreshes = self.coalesceProviderRefreshes && newer.coalesceProviderRefreshes + if newer.interaction == .userInitiated { + self.interaction = .userInitiated + } + } + } + + func refresh(forceTokenUsage: Bool = false) async { + if forceTokenUsage { + await self.refresh(enrichmentMode: .forcedForeground) + } else { + await self.runRefresh( + startupConnectivityRetryAttempt: nil, + waitForRefreshAvailability: true) + } + } + + private struct ForcedRefreshEnrichmentRequest: Sendable { + let generation: UInt64 + let refreshStartedAt: Date + let openAIWebRefreshPhase: ProviderRefreshPhase + } + + func refresh(enrichmentMode: RefreshEnrichmentMode) async { + if enrichmentMode == .forcedForeground { + await self.cancelForcedRefreshEnrichmentAndWait() + } + await self.runRefresh( + enrichmentMode: enrichmentMode, + startupConnectivityRetryAttempt: nil) + } + + func enqueueRequiredRefresh( + startupConnectivityRetryAttempt: Int?, + coalesceProviderRefreshesOverride: Bool?) async -> Bool + { + self.requiredRefreshRequestGeneration &+= 1 + let interaction = ProviderInteractionContext.current + let request = RequiredRefreshRequest( + throughGeneration: self.requiredRefreshRequestGeneration, + startupConnectivityRetryAttempt: startupConnectivityRetryAttempt, + coalesceProviderRefreshes: coalesceProviderRefreshesOverride ?? (interaction == .background), + interaction: interaction) + if var pending = self.pendingRequiredRefreshRequest { + pending.merge(request) + self.pendingRequiredRefreshRequest = pending + } else { + self.pendingRequiredRefreshRequest = request + } + + if let task = self.requiredRefreshTask { + return await task.value + } + + let token = UUID() + self.requiredRefreshTaskToken = token + let task = Task { @MainActor [weak self] in + guard let self else { return false } + let didRefresh = await self.drainRequiredRefreshRequests() + self.completeRequiredRefreshTask(token: token) + return didRefresh + } + self.requiredRefreshTask = task + return await task.value + } + + func cancelRequiredRefresh() { + self.pendingRequiredRefreshRequest = nil + self.requiredRefreshTaskToken = nil + let task = self.requiredRefreshTask + self.requiredRefreshTask = nil + task?.cancel() + } + + private func drainRequiredRefreshRequests() async -> Bool { + var completedAnyRefresh = false + while !Task.isCancelled { + guard await self.waitForRequiredRefreshAvailability(), + let request = self.pendingRequiredRefreshRequest + else { + break + } + self.pendingRequiredRefreshRequest = nil + + let didRefresh = await ProviderInteractionContext.$current.withValue(request.interaction) { + await self.runRefresh( + startupConnectivityRetryAttempt: request.startupConnectivityRetryAttempt, + coalesceProviderRefreshesOverride: request.coalesceProviderRefreshes) + } + if didRefresh { + completedAnyRefresh = true + self.requiredRefreshCompletedGeneration = max( + self.requiredRefreshCompletedGeneration, + request.throughGeneration) + } else if !Task.isCancelled { + var retry = request + if let pending = self.pendingRequiredRefreshRequest { + retry.merge(pending) + } + self.pendingRequiredRefreshRequest = retry + } + } + return completedAnyRefresh + } + + private func waitForRequiredRefreshAvailability() async -> Bool { + while self.isRefreshing || self.hasForcedRefreshEnrichmentInFlight { + guard !Task.isCancelled else { return false } + if self.hasForcedRefreshEnrichmentInFlight { + await self.awaitForcedRefreshEnrichment() + } else { + do { + try await Task.sleep(for: .milliseconds(20)) + } catch { + return false + } + } + } + return !Task.isCancelled + } + + private func completeRequiredRefreshTask(token: UUID) { + guard self.requiredRefreshTaskToken == token else { return } + self.requiredRefreshTask = nil + self.requiredRefreshTaskToken = nil + } + + func enqueueForcedRefreshEnrichment( + generation: UInt64, + refreshStartedAt: Date, + openAIWebRefreshPhase: ProviderRefreshPhase) + { + let request = ForcedRefreshEnrichmentRequest( + generation: generation, + refreshStartedAt: refreshStartedAt, + openAIWebRefreshPhase: openAIWebRefreshPhase) + if let predecessor = self.forcedRefreshEnrichmentTask { + self.replacePendingForcedRefreshEnrichment(request, predecessor: predecessor) + } else { + self.startForcedRefreshEnrichment(request) + } + } + + func awaitForcedRefreshEnrichment() async { + var reportedWait = false + while !Task.isCancelled { + guard let task = self.pendingForcedRefreshEnrichmentTask ?? self.forcedRefreshEnrichmentTask else { + return + } + if !reportedWait { + self._test_forcedRefreshEnrichmentWaitObserver?() + reportedWait = true + } + await task.value + } + } + + func cancelForcedRefreshEnrichment() { + _ = self.cancelForcedRefreshEnrichmentTasks() + } + + private func cancelForcedRefreshEnrichmentAndWait() async { + let tasks = self.cancelForcedRefreshEnrichmentTasks() + for task in tasks { + await task.value + } + } + + private func startForcedRefreshEnrichment(_ request: ForcedRefreshEnrichmentRequest) { + let token = UUID() + self.forcedRefreshEnrichmentToken = token + self.hasForcedRefreshEnrichmentInFlight = true + self.forcedRefreshEnrichmentTask = Task(priority: .utility) { @MainActor [weak self] in + guard let self else { return } + await self.runForcedRefreshEnrichment(request) + self.completeForcedRefreshEnrichment(token: token) + } + } + + private func replacePendingForcedRefreshEnrichment( + _ request: ForcedRefreshEnrichmentRequest, + predecessor: Task) + { + self.pendingForcedRefreshEnrichmentTask?.cancel() + let token = UUID() + self.pendingForcedRefreshEnrichmentToken = token + self.hasForcedRefreshEnrichmentInFlight = true + self.pendingForcedRefreshEnrichmentTask = Task(priority: .utility) { @MainActor [weak self] in + await predecessor.value + guard !Task.isCancelled, + let self, + self.pendingForcedRefreshEnrichmentToken == token, + let promotedTask = self.pendingForcedRefreshEnrichmentTask + else { return } + + self.pendingForcedRefreshEnrichmentTask = nil + self.pendingForcedRefreshEnrichmentToken = nil + self.forcedRefreshEnrichmentTask = promotedTask + self.forcedRefreshEnrichmentToken = token + await self.runForcedRefreshEnrichment(request) + self.completeForcedRefreshEnrichment(token: token) + } + } + + private func completeForcedRefreshEnrichment(token: UUID) { + guard self.forcedRefreshEnrichmentToken == token else { return } + // Keep the completed predecessor installed until its latest pending waiter promotes itself. + // This avoids an actor-reentrancy gap where a new request could otherwise start beside it. + guard self.pendingForcedRefreshEnrichmentTask == nil else { return } + self.forcedRefreshEnrichmentTask = nil + self.forcedRefreshEnrichmentToken = nil + self.hasForcedRefreshEnrichmentInFlight = false + } + + private func cancelForcedRefreshEnrichmentTasks() -> [Task] { + let tasks = [ + self.forcedRefreshEnrichmentTask, + self.pendingForcedRefreshEnrichmentTask, + self.openAIDashboardBackgroundRefreshTask, + self.openAIDashboardRefreshTask, + ].compactMap(\.self) + + self.forcedRefreshEnrichmentTask = nil + self.forcedRefreshEnrichmentToken = nil + self.pendingForcedRefreshEnrichmentTask = nil + self.pendingForcedRefreshEnrichmentToken = nil + self.hasForcedRefreshEnrichmentInFlight = false + tasks.forEach { $0.cancel() } + self.invalidateOpenAIDashboardRefreshTask() + return tasks + } + + private func runForcedRefreshEnrichment(_ request: ForcedRefreshEnrichmentRequest) async { + await withTaskGroup(of: Void.self) { group in + group.addTask { + await self.refreshCreditsNow(minimumSnapshotUpdatedAt: request.refreshStartedAt) + } + group.addTask { + await self.refreshTokenUsageSequenceNow(force: true) + } + } + guard !Task.isCancelled else { return } + + await self.refreshOpenAIWebAfterProviderRefresh( + force: true, + refreshPhase: request.openAIWebRefreshPhase) + guard !Task.isCancelled else { return } + + if self.openAIDashboardRequiresLogin, + request.generation == self.forcedRefreshEnrichmentGeneration + { + // Join a newer in-flight Codex request rather than replacing it. A newer accepted all-provider + // pass owns reconciliation even before it can enqueue its tail, so recheck generation afterward. + await self.refreshProvider(.codex, coalesceIfRefreshing: true) + guard !Task.isCancelled else { return } + if request.generation == self.forcedRefreshEnrichmentGeneration { + await self.refreshCreditsNow(minimumSnapshotUpdatedAt: request.refreshStartedAt) + guard !Task.isCancelled else { return } + } + } + + self.persistWidgetSnapshot(reason: "forced-refresh-enrichment") + } + + func refreshOpenAIWebAfterProviderRefresh( + force: Bool, + refreshPhase: ProviderRefreshPhase) async + { + self.syncOpenAIWebState() + let refreshPolicy = OpenAIWebRefreshPolicyContext( + accessEnabled: self.isEnabled(.codex) && + self.settings.openAIWebAccessEnabled && + self.settings.codexCookieSource.isEnabled, + batterySaverEnabled: self.settings.openAIWebBatterySaverEnabled, + force: force, + refreshPhase: refreshPhase) + let shouldRefreshOpenAIWeb = Self.shouldRunOpenAIWebRefresh(refreshPolicy) + self.openAIWebLogger.debug( + "OpenAI web refresh gate", + metadata: [ + "allowed": shouldRefreshOpenAIWeb ? "1" : "0", + "accessEnabled": refreshPolicy.accessEnabled ? "1" : "0", + "batterySaverEnabled": refreshPolicy.batterySaverEnabled ? "1" : "0", + "force": refreshPolicy.force ? "1" : "0", + "interaction": ProviderInteractionContext.current == .userInitiated ? "user" : "background", + "phase": refreshPhase == .startup ? "startup" : "regular", + ]) + guard shouldRefreshOpenAIWeb, !Task.isCancelled else { return } + + let codexDashboardGuard = self.freshCodexOpenAIWebRefreshGuard() + if force { + await self.refreshOpenAIDashboardIfNeeded( + force: true, + expectedGuard: codexDashboardGuard, + bypassCoalescing: true) + } else { + self.scheduleOpenAIDashboardRefreshIfNeeded(expectedGuard: codexDashboardGuard) + } + } +} diff --git a/Sources/CodexBar/UsageStore+ResetBoundaryRefresh.swift b/Sources/CodexBar/UsageStore+ResetBoundaryRefresh.swift index db58480a03..bd447f88ce 100644 --- a/Sources/CodexBar/UsageStore+ResetBoundaryRefresh.swift +++ b/Sources/CodexBar/UsageStore+ResetBoundaryRefresh.swift @@ -42,8 +42,11 @@ extension UsageStore { self.resetBoundaryRefreshTask = nil self.scheduledResetBoundaryRefreshAt = nil guard Self.shouldRecordResetBoundaryAttempt(isRefreshing: self.isRefreshing) else { return } + // Mark the boundary before the pass so runRefresh cannot schedule the same stale boundary again. self.recordAttemptedResetBoundaryRefresh(boundaryRefreshAt) - await self.refresh() + await self.runRefresh( + startupConnectivityRetryAttempt: nil, + waitForRefreshAvailability: true) } private func recordAttemptedResetBoundaryRefresh(_ refreshAt: Date) { diff --git a/Sources/CodexBar/UsageStore+StartupConnectivityRetry.swift b/Sources/CodexBar/UsageStore+StartupConnectivityRetry.swift index aceceb48d9..6101f11b54 100644 --- a/Sources/CodexBar/UsageStore+StartupConnectivityRetry.swift +++ b/Sources/CodexBar/UsageStore+StartupConnectivityRetry.swift @@ -61,7 +61,9 @@ extension UsageStore { do { try await self.sleepForStartupConnectivityRetry(delay) guard !Task.isCancelled else { return } - await self.runRefresh(startupConnectivityRetryAttempt: attempt) + await self.runRefresh( + startupConnectivityRetryAttempt: attempt, + waitForRefreshAvailability: true) } catch { return } diff --git a/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift b/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift new file mode 100644 index 0000000000..e7df6d1e9f --- /dev/null +++ b/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift @@ -0,0 +1,122 @@ +import CodexBarCore +import Foundation + +extension UsageStore { + private enum TokenRefreshSequenceScope: Sendable { + case all + case provider(UsageProvider) + } + + func startTokenTimer() { + self.tokenTimerTask?.cancel() + let wait = self.tokenFetchTTL + self.tokenTimerTask = Task.detached(priority: .utility) { [weak self] in + while !Task.isCancelled { + try? await Task.sleep(for: .seconds(wait)) + await self?.scheduleTokenRefresh() + } + } + } + + func scheduleTokenRefresh() { + guard self.tokenRefreshSequenceTask == nil, !self.hasForcedRefreshEnrichmentInFlight else { return } + self.startTokenRefreshSequence(force: false, scope: .all) + } + + func refreshTokenUsageSequenceNow(force: Bool) async { + guard let task = await self.serializedTokenRefreshTask(force: force, scope: .all) else { return } + await self.awaitTokenRefreshSequence(task) + } + + func refreshTokenUsageNow(for provider: UsageProvider, force: Bool) async { + if force, + self.tokenRefreshSequenceTask != nil, + let activeProvider = self.tokenRefreshSequenceProvider, + activeProvider != provider + { + // A scoped user refresh can run beside unrelated scheduled work. The scheduled + // sequence still owns the shared slot, so the timer cannot introduce a third pass. + await self.refreshTokenUsage(provider, force: true) + self.scheduleMemoryPressureRelief() + return + } + guard let task = await self.serializedTokenRefreshTask(force: force, scope: .provider(provider)) else { + return + } + await self.awaitTokenRefreshSequence(task) + } + + private func serializedTokenRefreshTask( + force: Bool, + scope: TokenRefreshSequenceScope) async -> Task? + { + if force { + while let existing = self.tokenRefreshSequenceTask { + existing.cancel() + await existing.value + guard !Task.isCancelled else { return nil } + } + } else if let existing = self.tokenRefreshSequenceTask { + return existing + } + return self.startTokenRefreshSequence(force: force, scope: scope) + } + + @discardableResult + private func startTokenRefreshSequence( + force: Bool, + scope: TokenRefreshSequenceScope) -> Task + { + let token = UUID() + self.tokenRefreshSequenceToken = token + let task = Task(priority: .utility) { @MainActor [weak self] in + guard let self else { return } + switch scope { + case .all: + await self.refreshTokenUsageSequence(force: force) + case let .provider(provider): + self.tokenRefreshSequenceProvider = provider + await self.refreshTokenUsage(provider, force: force) + self.tokenRefreshSequenceProvider = nil + self.scheduleMemoryPressureRelief() + } + self.completeTokenRefreshSequence(token: token) + } + self.tokenRefreshSequenceTask = task + return task + } + + private func awaitTokenRefreshSequence(_ task: Task) async { + await withTaskCancellationHandler { + await task.value + } onCancel: { + task.cancel() + } + } + + private func completeTokenRefreshSequence(token: UUID) { + guard self.tokenRefreshSequenceToken == token else { return } + self.tokenRefreshSequenceTask = nil + self.tokenRefreshSequenceToken = nil + self.tokenRefreshSequenceProvider = nil + } + + private func refreshTokenUsageSequence(force: Bool) async { + defer { self.tokenRefreshSequenceProvider = nil } + for provider in self.enabledProvidersForBackgroundWork() { + if Task.isCancelled { + break + } + self.tokenRefreshSequenceProvider = provider + await self.refreshTokenUsage(provider, force: force) + self.tokenRefreshSequenceProvider = nil + } + self.scheduleMemoryPressureRelief() + } + + #if DEBUG + func scheduleTokenRefreshForTesting() { + self.scheduleTokenRefresh() + } + #endif +} diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 6a0cbf91f6..fc85dd2c69 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -31,6 +31,7 @@ extension UsageStore { _ = self.openAIDashboardAttachmentRevision _ = self.versions _ = self.isRefreshing + _ = self.hasForcedRefreshEnrichmentInFlight _ = self.refreshingProviders _ = self.pathDebugInfo _ = self.statuses @@ -86,9 +87,15 @@ extension UsageStore { private static func isRunningTestsProcess() -> Bool { let environment = ProcessInfo.processInfo.environment - if environment["XCTestConfigurationFilePath"] != nil { return true } - if environment["XCTestSessionIdentifier"] != nil { return true } - if environment["SWIFT_TESTING_ENABLED"] != nil { return true } + if environment["XCTestConfigurationFilePath"] != nil { + return true + } + if environment["XCTestSessionIdentifier"] != nil { + return true + } + if environment["SWIFT_TESTING_ENABLED"] != nil { + return true + } return CommandLine.arguments.contains { argument in argument.contains("xctest") || argument.contains("swift-testing") } @@ -113,7 +120,9 @@ extension UsageStore { var preferredSnapshot: UsageSnapshot? { for provider in self.enabledProviders() { - if let snap = self.snapshots[provider] { return snap } + if let snap = self.snapshots[provider] { + return snap + } } return nil } @@ -179,6 +188,7 @@ final class UsageStore { var openAIDashboardCookieImportDebugLog: String? var versions: [UsageProvider: String] = [:] var isRefreshing = false + var hasForcedRefreshEnrichmentInFlight = false var refreshingProviders: Set = [] var debugForceAnimation = false var pathDebugInfo: PathDebugSnapshot = .empty @@ -235,6 +245,7 @@ final class UsageStore { @ObservationIgnored var _test_tokenUsageRefreshOverride: (@MainActor (UsageProvider, Bool) async -> Void)? @ObservationIgnored var _test_providerStatusFetchOverride: (@MainActor ( UsageProvider) async throws -> ProviderStatus)? + @ObservationIgnored var _test_forcedRefreshEnrichmentWaitObserver: (@MainActor () -> Void)? @ObservationIgnored var _test_startupConnectivityRetryScheduled: (@MainActor (Int, TimeInterval) -> Void)? @ObservationIgnored var _test_startupConnectivityRetrySleepOverride: (@MainActor ( TimeInterval) async throws -> Void)? @@ -267,9 +278,20 @@ final class UsageStore { /// In-memory only; resets on every launch. @ObservationIgnored private(set) var lastMenuOpenAt: Date? @ObservationIgnored var adaptiveRefreshScheduledAt: Date? - @ObservationIgnored private var tokenTimerTask: Task? - @ObservationIgnored private var tokenRefreshSequenceTask: Task? - @ObservationIgnored private var tokenRefreshSequenceProvider: UsageProvider? + @ObservationIgnored var tokenTimerTask: Task? + @ObservationIgnored var tokenRefreshSequenceTask: Task? + @ObservationIgnored var tokenRefreshSequenceToken: UUID? + @ObservationIgnored var tokenRefreshSequenceProvider: UsageProvider? + @ObservationIgnored var forcedRefreshEnrichmentTask: Task? + @ObservationIgnored var forcedRefreshEnrichmentToken: UUID? + @ObservationIgnored var pendingForcedRefreshEnrichmentTask: Task? + @ObservationIgnored var pendingForcedRefreshEnrichmentToken: UUID? + @ObservationIgnored var forcedRefreshEnrichmentGeneration: UInt64 = 0 + @ObservationIgnored var requiredRefreshTask: Task? + @ObservationIgnored var requiredRefreshTaskToken: UUID? + @ObservationIgnored var pendingRequiredRefreshRequest: RequiredRefreshRequest? + @ObservationIgnored var requiredRefreshRequestGeneration: UInt64 = 0 + @ObservationIgnored var requiredRefreshCompletedGeneration: UInt64 = 0 @ObservationIgnored var memoryPressureReliefTask: Task? @ObservationIgnored var startupConnectivityRetryTask: Task? @ObservationIgnored var startupConnectivityRetryNeeded = false @@ -391,14 +413,16 @@ final class UsageStore { Task { @MainActor [weak self] in await self?.refreshHistoricalDatasetIfNeeded() } - Task { await self.refresh() } + Task { await self.refresh(enrichmentMode: .automatic) } self.startTimer() self.startTokenTimer() } var iconStyle: IconStyle { let enabled = self.enabledProviders() - if enabled.count > 1 { return .combined } + if enabled.count > 1 { + return .combined + } if let provider = enabled.first { return self.style(for: provider) } @@ -546,55 +570,32 @@ final class UsageStore { self.providerAvailabilityCache.removeAll(keepingCapacity: true) } - func performRuntimeAction(_ action: ProviderRuntimeAction, for provider: UsageProvider) async { - guard let runtime = self.providerRuntimes[provider] else { return } - let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self) - await runtime.perform(action: action, context: context) - } - - private func updateProviderRuntimes() { - for (provider, runtime) in self.providerRuntimes { - let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self) - if self.isEnabled(provider) { - runtime.start(context: context) - } else { - runtime.stop(context: context) - } - runtime.settingsDidChange(context: context) - } - } - - func refresh(forceTokenUsage: Bool = false) async { - await self.runRefresh(forceTokenUsage: forceTokenUsage, startupConnectivityRetryAttempt: nil) - } - - func noteMenuOpened(at date: Date = Date()) { - self.lastMenuOpenAt = date - guard self.settings.refreshFrequency == .adaptive else { return } - - let decision = Self.adaptiveRefreshDecision( - now: date, - lastMenuOpenAt: date, - lowPowerModeEnabled: ProcessInfo.processInfo.isLowPowerModeEnabled, - thermalState: ProcessInfo.processInfo.thermalState) - let candidate = date.addingTimeInterval(TimeInterval(decision.delay.components.seconds)) - guard Self.shouldAdvanceAdaptiveTimer( - scheduledAt: self.adaptiveRefreshScheduledAt, - candidate: candidate) - else { return } - self.startTimer(preservingResetBoundaryRefresh: true) - } - #if DEBUG @ObservationIgnored private(set) var completedRefreshCountForTesting = 0 #endif + @discardableResult func runRefresh( - forceTokenUsage: Bool = false, + enrichmentMode: RefreshEnrichmentMode = .automatic, startupConnectivityRetryAttempt: Int?, - coalesceProviderRefreshesOverride: Bool? = nil) async + coalesceProviderRefreshesOverride: Bool? = nil, + waitForRefreshAvailability: Bool = false) async -> Bool { - guard !self.isRefreshing else { return } + if enrichmentMode == .automatic, waitForRefreshAvailability { + return await self.enqueueRequiredRefresh( + startupConnectivityRetryAttempt: startupConnectivityRetryAttempt, + coalesceProviderRefreshesOverride: coalesceProviderRefreshesOverride) + } + + guard !self.isRefreshing else { return false } + guard enrichmentMode != .automatic || !self.hasForcedRefreshEnrichmentInFlight else { return false } + let forcedBackgroundGeneration: UInt64? + if enrichmentMode == .forcedBackground { + self.forcedRefreshEnrichmentGeneration &+= 1 + forcedBackgroundGeneration = self.forcedRefreshEnrichmentGeneration + } else { + forcedBackgroundGeneration = nil + } self.prepareRefreshState() let refreshPhase = Self.refreshPhase(hasCompletedInitialRefresh: self.hasCompletedInitialRefresh) let openAIWebRefreshPhase = Self.openAIWebRefreshPhase( @@ -609,7 +610,7 @@ final class UsageStore { let availableRefreshProviders = Set(self.enabledProviders()) let refreshStartedAt = Date() - await ProviderRefreshContext.$current.withValue(refreshPhase) { + let completedRefresh = await ProviderRefreshContext.$current.withValue(refreshPhase) { self.isRefreshing = true defer { self.isRefreshing = false @@ -635,62 +636,52 @@ final class UsageStore { group.addTask { await self.refreshProviderStatus(provider) } } } - if forceTokenUsage { + if enrichmentMode == .forcedForeground { group.addTask { await self.refreshCreditsNow(minimumSnapshotUpdatedAt: refreshStartedAt) } } } + guard !Task.isCancelled else { return false } - if !forceTokenUsage { + if enrichmentMode == .automatic { self.scheduleCreditsRefreshIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) } - if forceTokenUsage { + if enrichmentMode == .forcedForeground { await self.refreshTokenUsageSequenceNow(force: true) - } else { + } else if enrichmentMode == .automatic { // Token-cost usage can be slow; run it outside regular/menu-open refreshes so we don't block UI. - self.scheduleTokenRefresh(force: false) + self.scheduleTokenRefresh() } // OpenAI web scrape depends on the current Codex account email (which can change after login/account // switch). Run this after Codex usage refresh so we don't accidentally scrape with stale credentials. - self.syncOpenAIWebState() - let refreshPolicy = OpenAIWebRefreshPolicyContext( - accessEnabled: self.isEnabled(.codex) && - self.settings.openAIWebAccessEnabled && - self.settings.codexCookieSource.isEnabled, - batterySaverEnabled: self.settings.openAIWebBatterySaverEnabled, - force: forceTokenUsage, - refreshPhase: openAIWebRefreshPhase) - let shouldRefreshOpenAIWeb = Self.shouldRunOpenAIWebRefresh(refreshPolicy) - self.openAIWebLogger.debug( - "OpenAI web refresh gate", - metadata: [ - "allowed": shouldRefreshOpenAIWeb ? "1" : "0", - "accessEnabled": refreshPolicy.accessEnabled ? "1" : "0", - "batterySaverEnabled": refreshPolicy.batterySaverEnabled ? "1" : "0", - "force": refreshPolicy.force ? "1" : "0", - "interaction": ProviderInteractionContext.current == .userInitiated ? "user" : "background", - "phase": openAIWebRefreshPhase == .startup ? "startup" : "regular", - ]) - if shouldRefreshOpenAIWeb { - let codexDashboardGuard = self.freshCodexOpenAIWebRefreshGuard() - if forceTokenUsage { - await self.refreshOpenAIDashboardIfNeeded( - force: true, - expectedGuard: codexDashboardGuard) - } else { - self.scheduleOpenAIDashboardRefreshIfNeeded(expectedGuard: codexDashboardGuard) - } + if enrichmentMode == .forcedBackground { + // Account ownership must fail closed before the responsive foreground pass returns; + // only the expensive dashboard fetch belongs in the deferred enrichment tail. + self.syncOpenAIWebState() + } else { + await self.refreshOpenAIWebAfterProviderRefresh( + force: enrichmentMode == .forcedForeground, + refreshPhase: openAIWebRefreshPhase) } - if forceTokenUsage, self.openAIDashboardRequiresLogin { + if enrichmentMode == .forcedForeground, self.openAIDashboardRequiresLogin { await self.refreshProvider(.codex) await self.refreshCreditsNow(minimumSnapshotUpdatedAt: refreshStartedAt) } self.persistWidgetSnapshot(reason: "refresh") + if let forcedBackgroundGeneration { + self.enqueueForcedRefreshEnrichment( + generation: forcedBackgroundGeneration, + refreshStartedAt: refreshStartedAt, + openAIWebRefreshPhase: openAIWebRefreshPhase) + } + return true } + guard completedRefresh else { return false } + self.scheduleResetBoundaryRefreshIfNeeded( normalRefreshInterval: self.normalRefreshIntervalForHeuristics()) @@ -703,6 +694,7 @@ final class UsageStore { #if DEBUG self.completedRefreshCountForTesting += 1 #endif + return true } /// For demo/testing: drop the snapshot so the loading animation plays, then restore the last snapshot. @@ -760,100 +752,42 @@ final class UsageStore { guard let sleepDuration = await Self.nextAdaptiveTimerSleepDuration(for: self) else { return } try? await Task.sleep(for: sleepDuration) guard !Task.isCancelled else { return } - await self?.refresh() + await self?.refresh(enrichmentMode: .automatic) } } return } guard let wait = frequency.seconds else { return } + #if DEBUG + let fixedTimerSleepOverride = self.refreshTimerSleepOverrideForTesting + #else + let fixedTimerSleepOverride: Duration? = nil + #endif // Background poller so the menu stays responsive; canceled when settings change or store deallocates. // Fixed cadence is anchored to the scheduled tick time, not refresh completion, so slow provider // work doesn't permanently stretch a two-minute interval into "refresh duration + two minutes". self.timerTask = Task.detached(priority: .utility) { [weak self] in - let interval = Duration.seconds(wait) - let clock = ContinuousClock() - var scheduledAt = clock.now + interval - while !Task.isCancelled { - let now = clock.now - let computedSleep = now >= scheduledAt ? .zero : scheduledAt - now - let sleepDuration = await self?.effectiveTimerSleepDuration(computedSleep) ?? computedSleep - try? await Task.sleep(for: sleepDuration) - guard !Task.isCancelled else { return } - await self?.refresh() - scheduledAt = Self.nextFixedTimerScheduledAt( - previousScheduledAt: scheduledAt, - completedAt: clock.now, - interval: interval) - } - } - } - - private func startTokenTimer() { - self.tokenTimerTask?.cancel() - let wait = self.tokenFetchTTL - self.tokenTimerTask = Task.detached(priority: .utility) { [weak self] in - while !Task.isCancelled { - try? await Task.sleep(for: .seconds(wait)) - await self?.scheduleTokenRefresh(force: false) - } - } - } - - private func scheduleTokenRefresh(force: Bool) { - if force { - self.tokenRefreshSequenceTask?.cancel() - self.tokenRefreshSequenceTask = nil - } else if self.tokenRefreshSequenceTask != nil { - return - } - - self.tokenRefreshSequenceTask = Task(priority: .utility) { [weak self] in - guard let self else { return } - defer { - Task { @MainActor [weak self] in - self?.tokenRefreshSequenceTask = nil - } - } - await self.refreshTokenUsageSequence(force: force) + await Self.runFixedRefreshTimer( + interval: .seconds(wait), + sleepOverride: fixedTimerSleepOverride, + refresh: { [weak self] in + await self?.refresh(enrichmentMode: .automatic) + }) } } - private func refreshTokenUsageSequenceNow(force: Bool) async { - await self.drainScheduledTokenRefreshIfNeeded(force: force, scopedTo: nil) - await self.refreshTokenUsageSequence(force: force) - } - - func refreshTokenUsageNow(for provider: UsageProvider, force: Bool) async { - await self.drainScheduledTokenRefreshIfNeeded(force: force, scopedTo: provider) - await self.refreshTokenUsage(provider, force: force) - self.scheduleMemoryPressureRelief() - } - - private func drainScheduledTokenRefreshIfNeeded(force: Bool, scopedTo provider: UsageProvider?) async { - guard force, let existing = self.tokenRefreshSequenceTask else { return } - if let provider, self.tokenRefreshSequenceProvider != provider { return } - existing.cancel() - await existing.value - self.tokenRefreshSequenceTask = nil - } - - private func refreshTokenUsageSequence(force: Bool) async { - defer { self.tokenRefreshSequenceProvider = nil } - for provider in self.enabledProvidersForBackgroundWork() { - if Task.isCancelled { break } - self.tokenRefreshSequenceProvider = provider - await self.refreshTokenUsage(provider, force: force) - self.tokenRefreshSequenceProvider = nil - } - self.scheduleMemoryPressureRelief() - } - deinit { self.timerTask?.cancel() self.tokenTimerTask?.cancel() self.tokenRefreshSequenceTask?.cancel() + self.forcedRefreshEnrichmentTask?.cancel() + self.pendingForcedRefreshEnrichmentTask?.cancel() + self.requiredRefreshTask?.cancel() + self.creditsRefreshTask?.cancel() + self.openAIDashboardBackgroundRefreshTask?.cancel() + self.openAIDashboardRefreshTask?.cancel() self.memoryPressureReliefTask?.cancel() self.startupConnectivityRetryTask?.cancel() self.storageRefreshTask?.cancel() @@ -920,7 +854,9 @@ final class UsageStore { return } guard let sessionWindow = self.sessionQuotaWindow(provider: provider, snapshot: snapshot) else { - if provider == .commandcode, snapshot.commandCodeSubscriptionEnrichmentUnavailable { return } + if provider == .commandcode, snapshot.commandCodeSubscriptionEnrichmentUnavailable { + return + } self.lastKnownSessionRemaining.removeValue(forKey: provider) self.lastKnownSessionWindowSource.removeValue(forKey: provider) return @@ -1490,7 +1426,7 @@ extension UsageStore { } } - private func refreshTokenUsage(_ provider: UsageProvider, force: Bool) async { + func refreshTokenUsage(_ provider: UsageProvider, force: Bool) async { guard ProviderDescriptorRegistry.descriptor(for: provider).tokenCost.supportsTokenCost else { self.tokenSnapshots.removeValue(forKey: provider) self.tokenErrors[provider] = nil @@ -1652,7 +1588,28 @@ extension UsageStore { /// Fast failures may retry on the next scheduled pass instead of waiting out the fetch /// TTL; timed-out scans keep the TTL so a slow corpus cannot thrash back-to-back rescans. nonisolated static func tokenFetchFailureAllowsEarlyRetry(_ error: Error) -> Bool { - if case CostUsageError.timedOut = error { return false } + if case CostUsageError.timedOut = error { + return false + } return true } } + +extension UsageStore { + func noteMenuOpened(at date: Date = Date()) { + self.lastMenuOpenAt = date + guard self.settings.refreshFrequency == .adaptive else { return } + + let decision = Self.adaptiveRefreshDecision( + now: date, + lastMenuOpenAt: date, + lowPowerModeEnabled: ProcessInfo.processInfo.isLowPowerModeEnabled, + thermalState: ProcessInfo.processInfo.thermalState) + let candidate = date.addingTimeInterval(TimeInterval(decision.delay.components.seconds)) + guard Self.shouldAdvanceAdaptiveTimer( + scheduledAt: self.adaptiveRefreshScheduledAt, + candidate: candidate) + else { return } + self.startTimer(preservingResetBoundaryRefresh: true) + } +} diff --git a/Sources/CodexBarCore/BrowserCookieAccessGate.swift b/Sources/CodexBarCore/BrowserCookieAccessGate.swift index d896eb6fb0..436d316565 100644 --- a/Sources/CodexBarCore/BrowserCookieAccessGate.swift +++ b/Sources/CodexBarCore/BrowserCookieAccessGate.swift @@ -59,6 +59,7 @@ public enum BrowserCookieAccessGate { private static let cooldownInterval: TimeInterval = 60 * 60 * 6 private static let log = CodexBarLog.logger(LogCategories.browserCookieGate) @TaskLocal private static var explicitRetryScope: ExplicitRetryScope? + @TaskLocal private static var deniedBrowsersForTesting: [Browser]? static let allowTestCookieAccessEnvironmentKey = "CODEXBAR_ALLOW_TEST_BROWSER_COOKIE_ACCESS" @@ -81,6 +82,9 @@ public enum BrowserCookieAccessGate { public static func shouldAttempt(_ browser: Browser, now: Date = Date()) -> Bool { guard browser.usesKeychainForCookieDecryption else { return true } guard !KeychainAccessGate.isDisabled else { return false } + if self.deniedBrowsersForTesting?.contains(browser) == true { + return self.isExplicitRetryAllowed(for: browser) + } let shouldCheckKeychain = self.lock.withLock { state in self.loadIfNeeded(&state) if let blockedUntil = state.deniedUntilByBrowser[browser.rawValue] { @@ -146,6 +150,16 @@ public enum BrowserCookieAccessGate { } } + static func withDeniedBrowsersForTesting( + _ browsers: [Browser], + isolation _: isolated (any Actor)? = #isolation, + operation: () async throws -> T) async rethrows -> T + { + try await self.$deniedBrowsersForTesting.withValue(browsers) { + try await operation() + } + } + static func operationPreservingAccessContext( _ operation: @escaping @Sendable () throws -> T) -> @Sendable () throws -> T { @@ -326,6 +340,14 @@ public enum BrowserCookieAccessGate { try await operation() } + static func withDeniedBrowsersForTesting( + _: [Browser], + isolation _: isolated (any Actor)? = #isolation, + operation: () async throws -> T) async rethrows -> T + { + try await operation() + } + public static func recordIfNeeded(_ error: Error, now: Date = Date()) {} public static func recordDenied(for browser: Browser, now: Date = Date()) {} public static func resetForTesting() {} diff --git a/Sources/CodexBarCore/KeychainAccessGate.swift b/Sources/CodexBarCore/KeychainAccessGate.swift index a25e19fc4b..f39865b0ba 100644 --- a/Sources/CodexBarCore/KeychainAccessGate.swift +++ b/Sources/CodexBarCore/KeychainAccessGate.swift @@ -74,6 +74,7 @@ public enum KeychainAccessGate { static func withTaskOverrideForTesting( _ disabled: Bool?, + isolation _: isolated (any Actor)? = #isolation, operation: () async throws -> T) async rethrows -> T { try await self.$taskOverrideValue.withValue(disabled) { diff --git a/Sources/CodexBarCore/KeychainAccessPreflight.swift b/Sources/CodexBarCore/KeychainAccessPreflight.swift index d003bf2249..d9523abf6f 100644 --- a/Sources/CodexBarCore/KeychainAccessPreflight.swift +++ b/Sources/CodexBarCore/KeychainAccessPreflight.swift @@ -117,6 +117,7 @@ public enum KeychainAccessPreflight { static func withCheckGenericPasswordOverrideForTesting( _ override: ((String, String?) -> Outcome)?, + isolation _: isolated (any Actor)? = #isolation, operation: () async throws -> T) async rethrows -> T { try await self.$taskCheckGenericPasswordOverrideStore.withValue( diff --git a/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift b/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift index 456773c338..5630fa60ee 100644 --- a/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift +++ b/Tests/CodexBarTests/AdaptiveRefreshTimerTests.swift @@ -91,12 +91,12 @@ struct AdaptiveRefreshTimerTests { } @Test - func `refresh call is a no-op while another refresh is already in flight`() async { + func `opportunistic timer refresh is a no-op while another refresh is already in flight`() async { let settings = Self.makeSettingsStore(suite: "AdaptiveRefreshTimerTests-coalesce", frequency: .manual) let store = Self.makeUsageStore(settings: settings, startupBehavior: .testing) store.isRefreshing = true - await store.refresh() + await store.refresh(enrichmentMode: .automatic) // The guard at the top of runRefresh() returned immediately: no completion was recorded and the // flag was left untouched by this call. This is the invariant every timer tick (fixed or @@ -136,6 +136,24 @@ struct AdaptiveRefreshTimerTests { let start = ContinuousClock.now let firstScheduledAt = start + interval + let nextAfterExactTick = UsageStore.nextFixedTimerScheduledAt( + previousScheduledAt: firstScheduledAt, + completedAt: firstScheduledAt, + interval: interval) + #expect(nextAfterExactTick == start + .milliseconds(200)) + + let nextJustBeforeFollowingTick = UsageStore.nextFixedTimerScheduledAt( + previousScheduledAt: firstScheduledAt, + completedAt: firstScheduledAt + .milliseconds(100) - .nanoseconds(1), + interval: interval) + #expect(nextJustBeforeFollowingTick == start + .milliseconds(200)) + + let nextAtFollowingTick = UsageStore.nextFixedTimerScheduledAt( + previousScheduledAt: firstScheduledAt, + completedAt: firstScheduledAt + .milliseconds(100), + interval: interval) + #expect(nextAtFollowingTick == start + .milliseconds(300)) + let nextAfterSlowRefresh = UsageStore.nextFixedTimerScheduledAt( previousScheduledAt: firstScheduledAt, completedAt: firstScheduledAt + .milliseconds(60), @@ -149,6 +167,20 @@ struct AdaptiveRefreshTimerTests { #expect(nextAfterMissedTicks == start + .milliseconds(400)) } + @Test + func `fixed timer loop stays interval aligned after a slow refresh`() async { + let harness = FixedTimerLoopHarness() + + await UsageStore.runFixedRefreshTimer( + interval: .milliseconds(100), + now: { await harness.now() }, + sleep: { duration in await harness.sleep(for: duration) }, + refresh: { await harness.refresh() }) + + #expect(await harness.recordedStarts() == [.milliseconds(100), .milliseconds(300)]) + #expect(await harness.maximumConcurrentRefreshes() == 1) + } + @Test func `adaptive mode keeps recomputing and refreshing across menu-open changes`() async throws { let settings = Self.makeSettingsStore(suite: "AdaptiveRefreshTimerTests-adaptive", frequency: .adaptive) @@ -332,3 +364,40 @@ struct AdaptiveRefreshTimerTests { environmentBase: [:]) } } + +private actor FixedTimerLoopHarness { + private let origin = ContinuousClock.now + private var elapsed = Duration.zero + private var starts: [Duration] = [] + private var activeRefreshes = 0 + private var maximumActiveRefreshes = 0 + + func now() -> ContinuousClock.Instant { + self.origin + self.elapsed + } + + func sleep(for duration: Duration) { + self.elapsed += duration + } + + func refresh() { + self.activeRefreshes += 1 + self.maximumActiveRefreshes = max(self.maximumActiveRefreshes, self.activeRefreshes) + self.starts.append(self.elapsed) + if self.starts.count == 1 { + self.elapsed += .milliseconds(160) + } + self.activeRefreshes -= 1 + if self.starts.count == 2 { + withUnsafeCurrentTask { $0?.cancel() } + } + } + + func recordedStarts() -> [Duration] { + self.starts + } + + func maximumConcurrentRefreshes() -> Int { + self.maximumActiveRefreshes + } +} diff --git a/Tests/CodexBarTests/CodexAccountScopedRefreshTestSupport.swift b/Tests/CodexBarTests/CodexAccountScopedRefreshTestSupport.swift index 48e114c133..0e4eca805d 100644 --- a/Tests/CodexBarTests/CodexAccountScopedRefreshTestSupport.swift +++ b/Tests/CodexBarTests/CodexAccountScopedRefreshTestSupport.swift @@ -68,12 +68,25 @@ extension CodexAccountScopedRefreshTests { } func makeUsageStore(settings: SettingsStore, environmentBase: [String: String] = [:]) -> UsageStore { - UsageStore( - fetcher: UsageFetcher(environment: [:]), - browserDetection: BrowserDetection(cacheTTL: 0), + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-tests", isDirectory: true) + .appendingPathComponent(UUID().uuidString, isDirectory: true) + var environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + "XDG_CONFIG_HOME": root.appendingPathComponent(".config", isDirectory: true).path, + ] + if let reconciliationEnvironment = settings._test_codexReconciliationEnvironment { + environment.merge(reconciliationEnvironment) { _, override in override } + } + environment.merge(environmentBase) { _, override in override } + settings._test_codexReconciliationEnvironment = environment + return UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(homeDirectory: environment["HOME"] ?? root.path, cacheTTL: 0), settings: settings, startupBehavior: .testing, - environmentBase: environmentBase) + environmentBase: environment) } func liveAccount(email: String, identity: CodexIdentity = .unresolved) -> ObservedSystemCodexAccount { @@ -402,7 +415,9 @@ actor BlockingCodexFetchStrategy { } func waitUntilStarted() async { - if self.didStart { return } + if self.didStart { + return + } await withCheckedContinuation { continuation in self.startedWaiters.append(continuation) } @@ -430,7 +445,9 @@ actor BlockingOpenAIDashboardLoader { } func waitUntilStarted() async { - if self.didStart { return } + if self.didStart { + return + } await withCheckedContinuation { continuation in self.startedWaiters.append(continuation) } @@ -457,7 +474,9 @@ actor BlockingWidgetSnapshotSaver { } func waitUntilStarted(count: Int) async { - if self.snapshots.count >= count { return } + if self.snapshots.count >= count { + return + } await withCheckedContinuation { continuation in self.startedWaiters.append(continuation) } diff --git a/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift b/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift index 6b0aa0ab03..664e375524 100644 --- a/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift +++ b/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift @@ -32,15 +32,33 @@ struct CodexBackgroundRefreshCoalescingTests { await store.refresh(forceTokenUsage: false) await firstCompletion.markCompleted() } - await blocker.waitUntilStarted(count: 1) - #expect(await firstCompletion.waitUntilCompleted() == true) + let didStartFirstCreditsRefresh = await blocker.waitUntilStartedWithin(count: 1) + #expect(didStartFirstCreditsRefresh) + guard didStartFirstCreditsRefresh else { + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [firstRefreshTask]) + return + } + let didCompleteFirstRefresh = await firstCompletion.waitUntilCompleted() + #expect(didCompleteFirstRefresh) + guard didCompleteFirstRefresh else { + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [firstRefreshTask]) + return + } let secondRefreshTask = Task { await store.refresh(forceTokenUsage: false) await secondCompletion.markCompleted() } - #expect(await secondCompletion.waitUntilCompleted() == true) + let didCompleteSecondRefresh = await secondCompletion.waitUntilCompleted() + #expect(didCompleteSecondRefresh) + guard didCompleteSecondRefresh else { + await self.cancelCreditsWork( + store: store, + blocker: blocker, + tasks: [firstRefreshTask, secondRefreshTask]) + return + } #expect(await blocker.startedCount() == 1) await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) @@ -77,20 +95,44 @@ struct CodexBackgroundRefreshCoalescingTests { let alphaRefreshTask = Task { await store.refresh(forceTokenUsage: false) } - await blocker.waitUntilStarted(count: 1) + let didStartAlphaRefresh = await blocker.waitUntilStartedWithin(count: 1) + #expect(didStartAlphaRefresh) + guard didStartAlphaRefresh else { + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [alphaRefreshTask]) + return + } + let staleCreditsTask = try #require(store.creditsRefreshTask) settings._test_activeManagedCodexAccount = betaAccount settings.codexActiveSource = .managedAccount(id: betaAccount.id) let betaRefreshTask = Task { await store.refresh(forceTokenUsage: false) } - await blocker.waitUntilStarted(count: 2) + let didStartBetaRefresh = await blocker.waitUntilStartedWithin(count: 2) + #expect(didStartBetaRefresh) + guard didStartBetaRefresh else { + await self.cancelCreditsWork( + store: store, + blocker: blocker, + tasks: [alphaRefreshTask, betaRefreshTask]) + return + } + let didCancelAlphaRefresh = await blocker.waitUntilCancellationCount(1) + #expect(didCancelAlphaRefresh) + guard didCancelAlphaRefresh else { + await self.cancelCreditsWork( + store: store, + blocker: blocker, + tasks: [alphaRefreshTask, betaRefreshTask]) + return + } + await blocker.resumeLast(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 10, events: [], updatedAt: Date()))) - await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) await alphaRefreshTask.value await betaRefreshTask.value + await staleCreditsTask.value await store.creditsRefreshTask?.value #expect(await blocker.startedCount() == 2) @@ -123,24 +165,114 @@ struct CodexBackgroundRefreshCoalescingTests { await store.refresh(forceTokenUsage: false) await regularCompletion.markCompleted() } - await blocker.waitUntilStarted(count: 1) - #expect(await regularCompletion.waitUntilCompleted() == true) + let didStartRegularCreditsRefresh = await blocker.waitUntilStartedWithin(count: 1) + #expect(didStartRegularCreditsRefresh) + guard didStartRegularCreditsRefresh else { + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [regularRefreshTask]) + return + } + let didCompleteRegularRefresh = await regularCompletion.waitUntilCompleted() + #expect(didCompleteRegularRefresh) + guard didCompleteRegularRefresh else { + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [regularRefreshTask]) + return + } + let staleCreditsTask = try #require(store.creditsRefreshTask) let forceRefreshTask = Task { await store.refresh(forceTokenUsage: true) } - await blocker.waitUntilStarted(count: 2) + let didStartForcedCreditsRefresh = await blocker.waitUntilStartedWithin(count: 2) + #expect(didStartForcedCreditsRefresh) + guard didStartForcedCreditsRefresh else { + await self.cancelCreditsWork( + store: store, + blocker: blocker, + tasks: [regularRefreshTask, forceRefreshTask]) + return + } + let didCancelStaleCreditsRefresh = await blocker.waitUntilCancellationCount(1) + #expect(didCancelStaleCreditsRefresh) + guard didCancelStaleCreditsRefresh else { + await self.cancelCreditsWork( + store: store, + blocker: blocker, + tasks: [regularRefreshTask, forceRefreshTask]) + return + } + await blocker.resumeLast(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 10, events: [], updatedAt: Date()))) - await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) await regularRefreshTask.value await forceRefreshTask.value + await staleCreditsTask.value #expect(await blocker.startedCount() == 2) + #expect(await blocker.cancellationCount() == 1) #expect(store.credits?.remaining == 25) } + @Test + func `forced background tail replaces stale scheduled Codex credits fetch`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-credits-tail-cancels-background") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + settings.openAIWebAccessEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let blocker = BlockingCreditsLoader() + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + } + + await store.refresh(forceTokenUsage: false) + let didStartScheduledCreditsRefresh = await blocker.waitUntilStartedWithin(count: 1) + #expect(didStartScheduledCreditsRefresh) + guard didStartScheduledCreditsRefresh else { + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: []) + return + } + let staleCreditsTask = try #require(store.creditsRefreshTask) + + await store.refresh(enrichmentMode: .forcedBackground) + let tailTask = try #require(store.forcedRefreshEnrichmentTask) + let didStartForcedCreditsRefresh = await blocker.waitUntilStartedWithin(count: 2) + #expect(didStartForcedCreditsRefresh) + guard didStartForcedCreditsRefresh else { + tailTask.cancel() + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [tailTask]) + return + } + + let didCancelStaleCreditsRefresh = await blocker.waitUntilCancellationCount(1) + #expect(didCancelStaleCreditsRefresh) + guard didCancelStaleCreditsRefresh else { + tailTask.cancel() + await self.cancelCreditsWork(store: store, blocker: blocker, tasks: [tailTask]) + return + } + await blocker.resumeLast(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) + await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 10, events: [], updatedAt: Date()))) + await tailTask.value + await staleCreditsTask.value + + #expect(await blocker.startedCount() == 2) + #expect(await blocker.cancellationCount() == 1) + #expect(store.credits?.remaining == 25) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + @Test func `rapid regular refreshes coalesce concurrent OpenAI dashboard fetches`() async throws { let settings = try self.makeSettingsStore( @@ -172,15 +304,33 @@ struct CodexBackgroundRefreshCoalescingTests { await store.refresh(forceTokenUsage: false) await firstCompletion.markCompleted() } - await blocker.waitUntilStarted(count: 1) - #expect(await firstCompletion.waitUntilCompleted() == true) + let didStartDashboardRefresh = await blocker.waitUntilStartedWithin(count: 1) + #expect(didStartDashboardRefresh) + guard didStartDashboardRefresh else { + await self.cancelDashboardWork(store: store, blocker: blocker, tasks: [firstRefreshTask]) + return + } + let didCompleteFirstRefresh = await firstCompletion.waitUntilCompleted() + #expect(didCompleteFirstRefresh) + guard didCompleteFirstRefresh else { + await self.cancelDashboardWork(store: store, blocker: blocker, tasks: [firstRefreshTask]) + return + } let secondRefreshTask = Task { await store.refresh(forceTokenUsage: false) await secondCompletion.markCompleted() } - #expect(await secondCompletion.waitUntilCompleted() == true) + let didCompleteSecondRefresh = await secondCompletion.waitUntilCompleted() + #expect(didCompleteSecondRefresh) + guard didCompleteSecondRefresh else { + await self.cancelDashboardWork( + store: store, + blocker: blocker, + tasks: [firstRefreshTask, secondRefreshTask]) + return + } #expect(await blocker.startedCount() == 1) let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask) @@ -224,8 +374,22 @@ struct CodexBackgroundRefreshCoalescingTests { targetEmail: managedAccount.email, force: true) } - await importBlocker.waitUntilStarted() + let didStartImport = await importBlocker.waitUntilStarted() + #expect(didStartImport) + guard didStartImport else { + importTask.cancel() + await importBlocker.cancelAll() + _ = await importTask.value + return + } importTask.cancel() + let didObserveCancellation = await importBlocker.waitUntilCancellationCount(1) + #expect(didObserveCancellation) + guard didObserveCancellation else { + await importBlocker.cancelAll() + _ = await importTask.value + return + } await importBlocker.resumeNext(with: .failure( OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount( found: [.init(sourceLabel: "Chrome", email: "other@example.com")]))) @@ -237,16 +401,786 @@ struct CodexBackgroundRefreshCoalescingTests { #expect(store.openAIDashboardRequiresLogin == false) } - private func makeSettingsStore(suite: String) throws -> SettingsStore { - let defaults = UserDefaults(suiteName: suite)! - defaults.removePersistentDomain(forName: suite) - defaults.set(true, forKey: "providerDetectionCompleted") - let configStore = testConfigStore(suiteName: suite) - let settings = SettingsStore( - userDefaults: defaults, - configStore: configStore, - zaiTokenStore: NoopZaiTokenStore(), - syntheticTokenStore: NoopSyntheticTokenStore()) + @Test + func `settings refresh waits for forced enrichment instead of being dropped`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-settings-waits-for-tail") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let tokenGate = BlockingForcedTokenRefresh() + var providerInteractions: [ProviderInteraction] = [] + var didObserveWait = false + store._test_providerRefreshOverride = { _ in + providerInteractions.append(ProviderInteractionContext.current) + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + guard force else { return } + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + store._test_forcedRefreshEnrichmentWaitObserver = { + didObserveWait = true + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + store._test_forcedRefreshEnrichmentWaitObserver = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + let didStartTokenTail = await tokenGate.waitUntilStarted(count: 1) + #expect(didStartTokenTail) + guard didStartTokenTail else { + await self.cancelForcedEnrichmentWork(store: store) + return + } + settings.costUsageEnabled = false + + let completion = RefreshCompletionProbe() + var settingsRefreshes: [Task] = [] + for expectedGeneration in 1...3 { + let task = Task { @MainActor in + await store.refreshForSettingsChange() + await completion.markCompleted() + } + settingsRefreshes.append(task) + for _ in 0..<100 where store.requiredRefreshRequestGeneration < expectedGeneration { + await Task.yield() + } + #expect(store.requiredRefreshRequestGeneration == expectedGeneration) + } + for _ in 0..<100 where !didObserveWait { + await Task.yield() + } + + #expect(didObserveWait) + #expect(providerInteractions == [.background]) + #expect(await completion.isCompleted == false) + + await tokenGate.resumeNext() + for task in settingsRefreshes { + await task.value + } + + #expect(providerInteractions == [.background, .background]) + #expect(await completion.isCompleted) + #expect(store.requiredRefreshCompletedGeneration == 3) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `required post-action refresh waits for forced enrichment`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-required-refresh-waits-for-tail") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let tokenGate = BlockingForcedTokenRefresh() + var providerInteractions: [ProviderInteraction] = [] + var didObserveWait = false + store._test_providerRefreshOverride = { _ in + providerInteractions.append(ProviderInteractionContext.current) + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + guard force else { return } + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + store._test_forcedRefreshEnrichmentWaitObserver = { + didObserveWait = true + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + store._test_forcedRefreshEnrichmentWaitObserver = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + let didStartTokenTail = await tokenGate.waitUntilStarted(count: 1) + #expect(didStartTokenTail) + guard didStartTokenTail else { + await self.cancelForcedEnrichmentWork(store: store) + return + } + settings.costUsageEnabled = false + + let completion = RefreshCompletionProbe() + let requiredRefresh = Task { @MainActor in + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh() + } + await completion.markCompleted() + } + for _ in 0..<100 where !didObserveWait { + await Task.yield() + } + + #expect(didObserveWait) + #expect(providerInteractions == [.background]) + #expect(await completion.isCompleted == false) + + await tokenGate.resumeNext() + await requiredRefresh.value + + #expect(providerInteractions == [.background, .userInitiated]) + #expect(await completion.isCompleted) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `startup retry waits for forced enrichment and completes its retry pass`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-startup-retry-waits-for-tail") + settings.refreshFrequency = .manual + settings.statusChecksEnabled = true + settings.costUsageEnabled = false + settings.openAIWebAccessEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let retrySleep = ForcedRefreshRetrySleepGate() + let tokenGate = BlockingForcedTokenRefresh() + var statusAttempts = 0 + var didObserveWait = false + store._test_providerRefreshOverride = { _ in } + store._test_providerStatusFetchOverride = { _ in + statusAttempts += 1 + if statusAttempts == 1 { + throw URLError(.cannotFindHost) + } + return ProviderStatus(indicator: .none, description: "Operational", updatedAt: Date()) + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + guard force else { return } + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + store._test_startupConnectivityRetrySleepOverride = { delay in + try await retrySleep.sleep(delay) + } + store._test_forcedRefreshEnrichmentWaitObserver = { + didObserveWait = true + } + defer { + store._test_providerRefreshOverride = nil + store._test_providerStatusFetchOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + store._test_startupConnectivityRetrySleepOverride = nil + store._test_forcedRefreshEnrichmentWaitObserver = nil + store.startupConnectivityRetryTask?.cancel() + store.startupConnectivityRetryTask = nil + } + + await store.refresh() + let didStartRetrySleep = await retrySleep.waitUntilSleeping() + #expect(didStartRetrySleep) + guard didStartRetrySleep else { + store.startupConnectivityRetryTask?.cancel() + return + } + let retryTask = try #require(store.startupConnectivityRetryTask) + + settings.costUsageEnabled = true + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh(enrichmentMode: .forcedBackground) + } + let didStartTokenTail = await tokenGate.waitUntilStarted(count: 1) + #expect(didStartTokenTail) + guard didStartTokenTail else { + await self.cancelForcedEnrichmentWork(store: store) + retryTask.cancel() + await retryTask.value + return + } + settings.costUsageEnabled = false + + await retrySleep.resume() + for _ in 0..<100 where !didObserveWait { + await Task.yield() + } + + #expect(didObserveWait) + #expect(statusAttempts == 2) + + await tokenGate.resumeNext() + await retryTask.value + + #expect(statusAttempts == 3) + #expect(store.statuses[.codex]?.indicator == ProviderStatusIndicator.none) + #expect(store.startupConnectivityRetryTask == nil) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } +} + +extension CodexBackgroundRefreshCoalescingTests { + @Test + func `forced enrichment keeps one active and the latest contextual follow-up`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-enrichment-latest") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let gate = BlockingForcedTokenRefresh() + let deniedAt = Date() + var providerRefreshCount = 0 + store._test_providerRefreshOverride = { _ in + providerRefreshCount += 1 + } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + defer { store._test_codexCreditsLoaderOverride = nil } + store._test_tokenUsageRefreshOverride = { provider, force in + let retryAllowed = BrowserCookieAccessGate.shouldAttempt( + .arc, + now: deniedAt.addingTimeInterval(1)) + await gate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: retryAllowed) + } + defer { store._test_tokenUsageRefreshOverride = nil } + + await BrowserCookieAccessGate.withDeniedBrowsersForTesting([.arc]) { + await KeychainAccessGate.withTaskOverrideForTesting(false) { + let firstRefresh = Task { @MainActor in + await ProviderInteractionContext.$current.withValue(.background) { + await store.refresh(enrichmentMode: .forcedBackground) + } + } + let didStartFirstTail = await gate.waitUntilStarted(count: 1) + #expect(didStartFirstTail) + guard didStartFirstTail else { + firstRefresh.cancel() + await self.cancelForcedEnrichmentWork(store: store) + await firstRefresh.value + return + } + await firstRefresh.value + + await store.refresh(enrichmentMode: .automatic) + #expect(providerRefreshCount == 1) + + await ProviderInteractionContext.$current.withValue(.background) { + await store.refresh(enrichmentMode: .forcedBackground) + } + let preflightOverride: (String, String?) -> KeychainAccessPreflight.Outcome = { _, _ in .allowed } + await KeychainAccessPreflight.withCheckGenericPasswordOverrideForTesting(preflightOverride) { + await BrowserCookieAccessGate.withExplicitRetry { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh(enrichmentMode: .forcedBackground) + } + } + } + + #expect(store.forcedRefreshEnrichmentTask != nil) + #expect(store.pendingForcedRefreshEnrichmentTask != nil) + #expect(store.hasForcedRefreshEnrichmentInFlight) + + await gate.resumeNext() + let didStartLatestTail = await gate.waitUntilStarted(count: 2) + #expect(didStartLatestTail) + guard didStartLatestTail else { + await self.cancelForcedEnrichmentWork(store: store) + return + } + + let calls = await gate.recordedCalls() + #expect(calls.map(\.provider) == [.codex, .codex]) + #expect(calls.map(\.force) == [true, true]) + #expect(calls.map(\.interaction) == [.background, .userInitiated]) + #expect(calls.map(\.refreshPhase) == [.startup, .regular]) + #expect(calls.map(\.browserRetryAllowed) == [false, true]) + + await gate.resumeNext() + await store.awaitForcedRefreshEnrichment() + #expect(!store.hasForcedRefreshEnrichmentInFlight) + #expect(store.forcedRefreshEnrichmentTask == nil) + #expect(store.pendingForcedRefreshEnrichmentTask == nil) + } + } + } + + @Test + func `forced background login failure reconciles provider and credits once`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-login-reconciliation") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + var providerRefreshes = 0 + var creditsRefreshes = 0 + var dashboardLoads = 0 + store._test_providerRefreshOverride = { provider in + #expect(provider == .codex) + providerRefreshes += 1 + } + store._test_codexCreditsLoaderOverride = { + creditsRefreshes += 1 + return CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_openAIDashboardLoaderOverride = { _, _, _, _ in + dashboardLoads += 1 + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + store._test_openAIDashboardCookieImportOverride = { targetEmail, _, _, _, _ in + OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Fixture", + cookieCount: 2, + signedInEmail: targetEmail, + matchesCodexEmail: true) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_openAIDashboardLoaderOverride = nil + store._test_openAIDashboardCookieImportOverride = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + await store.awaitForcedRefreshEnrichment() + + #expect(dashboardLoads == 2) + #expect(providerRefreshes == 2) + #expect(creditsRefreshes == 2) + #expect(store.openAIDashboardRequiresLogin) + } + + @Test + func `older login reconciliation yields to an already pending forced tail`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-login-pending-generation") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let tokenGate = BlockingForcedTokenRefresh() + let dashboardLoader = LoginThenSuccessDashboardLoader(email: managedAccount.email) + var providerInteractions: [ProviderInteraction] = [] + var creditsInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { _ in + providerInteractions.append(ProviderInteractionContext.current) + } + store._test_codexCreditsLoaderOverride = { + creditsInteractions.append(ProviderInteractionContext.current) + return CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + store._test_openAIDashboardLoaderOverride = { _, _, _, _ in + try await dashboardLoader.load() + } + store._test_openAIDashboardCookieImportOverride = { targetEmail, _, _, _, _ in + OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Fixture", + cookieCount: 2, + signedInEmail: targetEmail, + matchesCodexEmail: true) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + store._test_openAIDashboardLoaderOverride = nil + store._test_openAIDashboardCookieImportOverride = nil + } + + await ProviderInteractionContext.$current.withValue(.background) { + await store.refresh(enrichmentMode: .forcedBackground) + } + let didStartOlderTail = await tokenGate.waitUntilStarted(count: 1) + #expect(didStartOlderTail) + guard didStartOlderTail else { + let activeTask = store.forcedRefreshEnrichmentTask + store.cancelForcedRefreshEnrichment() + await activeTask?.value + return + } + + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh(enrichmentMode: .forcedBackground) + } + #expect(store.pendingForcedRefreshEnrichmentTask != nil) + + await tokenGate.resumeNext() + let didStartNewerTail = await tokenGate.waitUntilStarted(count: 2) + #expect(didStartNewerTail) + guard didStartNewerTail else { + let tasks = [ + store.forcedRefreshEnrichmentTask, + store.pendingForcedRefreshEnrichmentTask, + ].compactMap(\.self) + store.cancelForcedRefreshEnrichment() + for task in tasks { + await task.value + } + return + } + + #expect(await dashboardLoader.callCount() == 2) + #expect(store.openAIDashboardRequiresLogin) + #expect(providerInteractions == [.background, .userInitiated]) + #expect(creditsInteractions == [.background, .userInitiated]) + + await tokenGate.resumeNext() + await store.awaitForcedRefreshEnrichment() + + #expect(await dashboardLoader.callCount() == 3) + #expect(!store.openAIDashboardRequiresLogin) + #expect(providerInteractions.count == 2) + #expect(creditsInteractions.count == 2) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `older login reconciliation does not replace newer forced provider work`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-login-inflight-generation") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let providerGate = BlockingSecondProviderRefresh() + let tokenGate = BlockingForcedTokenRefresh() + let dashboardLoader = LoginThenSuccessDashboardLoader(email: managedAccount.email) + var creditsInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { _ in + await providerGate.run(interaction: ProviderInteractionContext.current) + } + store._test_codexCreditsLoaderOverride = { + creditsInteractions.append(ProviderInteractionContext.current) + return CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + store._test_openAIDashboardLoaderOverride = { _, _, _, _ in + try await dashboardLoader.load() + } + store._test_openAIDashboardCookieImportOverride = { targetEmail, _, _, _, _ in + OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Fixture", + cookieCount: 2, + signedInEmail: targetEmail, + matchesCodexEmail: true) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + store._test_openAIDashboardLoaderOverride = nil + store._test_openAIDashboardCookieImportOverride = nil + } + + await ProviderInteractionContext.$current.withValue(.background) { + await store.refresh(enrichmentMode: .forcedBackground) + } + let didStartOlderTail = await tokenGate.waitUntilStarted(count: 1) + #expect(didStartOlderTail) + guard didStartOlderTail else { + let activeTask = store.forcedRefreshEnrichmentTask + store.cancelForcedRefreshEnrichment() + await activeTask?.value + return + } + let olderTail = try #require(store.forcedRefreshEnrichmentTask) + + let newerRefresh = Task { @MainActor in + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh(enrichmentMode: .forcedBackground) + } + } + let didStartNewerProviderRefresh = await providerGate.waitUntilStarted(count: 2) + #expect(didStartNewerProviderRefresh) + guard didStartNewerProviderRefresh else { + newerRefresh.cancel() + await providerGate.releaseBlockedCall() + store.cancelForcedRefreshEnrichment() + await tokenGate.resumeNext() + await newerRefresh.value + return + } + + let olderTailCompletion = RefreshCompletionProbe() + let olderTailWaiter = Task { + await olderTail.value + await olderTailCompletion.markCompleted() + } + await tokenGate.resumeNext() + let didCompleteOlderTail = await olderTailCompletion.waitUntilCompleted() + #expect(didCompleteOlderTail) + guard didCompleteOlderTail else { + await providerGate.releaseBlockedCall() + await newerRefresh.value + let tasks = [ + store.forcedRefreshEnrichmentTask, + store.pendingForcedRefreshEnrichmentTask, + ].compactMap(\.self) + store.cancelForcedRefreshEnrichment() + for task in tasks { + await task.value + } + await olderTailWaiter.value + return + } + await olderTailWaiter.value + + #expect(await dashboardLoader.callCount() == 2) + #expect(store.openAIDashboardRequiresLogin) + #expect(await providerGate.wasBlockedCallCancelled() == false) + #expect(await providerGate.recordedInteractions() == [.background, .userInitiated]) + #expect(creditsInteractions == [.background]) + + await providerGate.releaseBlockedCall() + await newerRefresh.value + let didStartNewerTail = await tokenGate.waitUntilStarted(count: 2) + #expect(didStartNewerTail) + guard didStartNewerTail else { + let tasks = [ + store.forcedRefreshEnrichmentTask, + store.pendingForcedRefreshEnrichmentTask, + ].compactMap(\.self) + store.cancelForcedRefreshEnrichment() + for task in tasks { + await task.value + } + return + } + await tokenGate.resumeNext() + await store.awaitForcedRefreshEnrichment() + + #expect(await dashboardLoader.callCount() == 3) + #expect(!store.openAIDashboardRequiresLogin) + #expect(await providerGate.wasBlockedCallCancelled() == false) + #expect(await providerGate.recordedInteractions() == [.background, .userInitiated]) + #expect(creditsInteractions == [.background, .userInitiated]) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `cancelling forced enrichment cancels its real dashboard child`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-dashboard-child-cancellation") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let dashboardLoader = CancellationAwareOpenAIDashboardLoader(email: managedAccount.email) + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_openAIDashboardLoaderOverride = { _, _, _, _ in + try await dashboardLoader.load() + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_openAIDashboardLoaderOverride = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + let capturedEnrichmentTask = store.forcedRefreshEnrichmentTask + let didStartDashboardRefresh = await dashboardLoader.waitUntilStarted() + #expect(didStartDashboardRefresh) + guard didStartDashboardRefresh else { + store.cancelForcedRefreshEnrichment() + await capturedEnrichmentTask?.value + return + } + let enrichmentTask = try #require(capturedEnrichmentTask) + let dashboardTask = try #require(store.openAIDashboardRefreshTask) + + store.cancelForcedRefreshEnrichment() + await enrichmentTask.value + await dashboardTask.value + + #expect(await dashboardLoader.wasCancelled()) + #expect(dashboardTask.isCancelled) + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardError == nil) + #expect(!store.openAIDashboardRequiresLogin) + #expect(store.openAIDashboardRefreshTask == nil) + #expect(store.openAIDashboardBackgroundRefreshTask == nil) + #expect(store.forcedRefreshEnrichmentTask == nil) + #expect(store.pendingForcedRefreshEnrichmentTask == nil) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `forced background enrichment runs dashboard under battery saver with user context`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-dashboard-battery") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + settings.openAIWebBatterySaverEnabled = true + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + var dashboardInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + defer { store._test_codexCreditsLoaderOverride = nil } + store._test_openAIDashboardLoaderOverride = { _, _, _, _ in + dashboardInteractions.append(ProviderInteractionContext.current) + return OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()) + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + await BrowserCookieAccessGate.withExplicitRetry { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh(enrichmentMode: .forcedBackground) + } + } + await store.awaitForcedRefreshEnrichment() + + #expect(dashboardInteractions == [.userInitiated]) + #expect(store.openAIDashboard?.signedInEmail == managedAccount.email) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } +} + +extension CodexBackgroundRefreshCoalescingTests { + private func cancelForcedEnrichmentWork(store: UsageStore) async { + let tasks = [ + store.forcedRefreshEnrichmentTask, + store.pendingForcedRefreshEnrichmentTask, + ].compactMap(\.self) + store.cancelForcedRefreshEnrichment() + for task in tasks { + await task.value + } + } + + private func cancelCreditsWork( + store: UsageStore, + blocker: BlockingCreditsLoader, + tasks: [Task]) async + { + let creditsTask = store.creditsRefreshTask + tasks.forEach { $0.cancel() } + creditsTask?.cancel() + await blocker.cancelAll() + for task in tasks { + await task.value + } + await creditsTask?.value + } + + private func cancelDashboardWork( + store: UsageStore, + blocker: BlockingManagedOpenAIDashboardLoader, + tasks: [Task]) async + { + let dashboardTasks = [ + store.openAIDashboardBackgroundRefreshTask, + store.openAIDashboardRefreshTask, + ].compactMap(\.self) + tasks.forEach { $0.cancel() } + store.invalidateOpenAIDashboardRefreshTask() + await blocker.cancelAll() + for task in tasks { + await task.value + } + for task in dashboardTasks { + await task.value + } + } + + func makeSettingsStore(suite: String) throws -> SettingsStore { + let settings = testSettingsStore(suiteName: suite) let codexMetadata = try #require(ProviderDescriptorRegistry.metadata[.codex]) settings.setProviderEnabled(provider: .codex, metadata: codexMetadata, enabled: true) settings.providerDetectionCompleted = true @@ -255,15 +1189,24 @@ struct CodexBackgroundRefreshCoalescingTests { return settings } - private func makeStore(settings: SettingsStore) -> UsageStore { - UsageStore( - fetcher: UsageFetcher(environment: [:]), + func makeStore(settings: SettingsStore) -> UsageStore { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-tests", isDirectory: true) + .appendingPathComponent(UUID().uuidString, isDirectory: true) + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + "XDG_CONFIG_HOME": root.appendingPathComponent(".config", isDirectory: true).path, + ] + return UsageStore( + fetcher: UsageFetcher(environment: environment), browserDetection: BrowserDetection(cacheTTL: 0), settings: settings, - startupBehavior: .testing) + startupBehavior: .testing, + environmentBase: environment) } - private static func installManagedAccount( + static func installManagedAccount( email: String, settings: SettingsStore) throws -> ManagedCodexAccount { @@ -321,44 +1264,316 @@ struct CodexBackgroundRefreshCoalescingTests { } } +actor BlockingForcedTokenRefresh { + struct Call: Sendable { + let provider: UsageProvider + let force: Bool + let interaction: ProviderInteraction + let refreshPhase: ProviderRefreshPhase + let browserRetryAllowed: Bool + } + + private var calls: [Call] = [] + private var continuations: [(id: UUID, continuation: CheckedContinuation)] = [] + + func run( + provider: UsageProvider, + force: Bool, + interaction: ProviderInteraction, + refreshPhase: ProviderRefreshPhase, + browserRetryAllowed: Bool) async + { + let id = UUID() + self.calls.append(Call( + provider: provider, + force: force, + interaction: interaction, + refreshPhase: refreshPhase, + browserRetryAllowed: browserRetryAllowed)) + await withTaskCancellationHandler { + await withCheckedContinuation { (continuation: CheckedContinuation) in + if Task.isCancelled { + continuation.resume() + } else { + self.continuations.append((id: id, continuation: continuation)) + } + } + } onCancel: { + Task { await self.cancel(id: id) } + } + } + + @discardableResult + func waitUntilStarted(count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.calls.count < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func resumeNext() { + guard !self.continuations.isEmpty else { return } + self.continuations.removeFirst().continuation.resume() + } + + func recordedCalls() -> [Call] { + self.calls + } + + private func cancel(id: UUID) { + guard let index = self.continuations.firstIndex(where: { $0.id == id }) else { return } + self.continuations.remove(at: index).continuation.resume() + } +} + +private actor LoginThenSuccessDashboardLoader { + private let email: String + private var calls = 0 + + init(email: String) { + self.email = email + } + + func load() throws -> OpenAIDashboardSnapshot { + self.calls += 1 + if self.calls <= 2 { + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + return OpenAIDashboardSnapshot( + signedInEmail: self.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()) + } + + func callCount() -> Int { + self.calls + } +} + +private actor ForcedRefreshRetrySleepGate { + private var continuation: CheckedContinuation? + private var cancelled = false + + func sleep(_ delay: TimeInterval) async throws { + #expect(delay == 15) + try await withTaskCancellationHandler(operation: { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + if self.cancelled || Task.isCancelled { + continuation.resume(throwing: CancellationError()) + } else { + self.continuation = continuation + } + } + }, onCancel: { + Task { await self.cancel() } + }) + } + + @discardableResult + func waitUntilSleeping(timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.continuation == nil { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func resume() { + self.continuation?.resume() + self.continuation = nil + } + + private func cancel() { + self.cancelled = true + self.continuation?.resume(throwing: CancellationError()) + self.continuation = nil + } +} + +private actor BlockingSecondProviderRefresh { + private var interactions: [ProviderInteraction] = [] + private var blockedContinuation: CheckedContinuation? + private var blockedCallCancelled = false + private var blockedCallReleased = false + + func run(interaction: ProviderInteraction) async { + self.interactions.append(interaction) + guard self.interactions.count == 2 else { return } + + await withTaskCancellationHandler { + await withCheckedContinuation { (continuation: CheckedContinuation) in + if self.blockedCallCancelled || self.blockedCallReleased || Task.isCancelled { + if Task.isCancelled { + self.blockedCallCancelled = true + } + continuation.resume() + } else { + self.blockedContinuation = continuation + } + } + } onCancel: { + Task { await self.cancelBlockedCall() } + } + } + + func waitUntilStarted(count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.interactions.count < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func releaseBlockedCall() { + self.blockedCallReleased = true + self.blockedContinuation?.resume() + self.blockedContinuation = nil + } + + func wasBlockedCallCancelled() -> Bool { + self.blockedCallCancelled + } + + func recordedInteractions() -> [ProviderInteraction] { + self.interactions + } + + private func cancelBlockedCall() { + self.blockedCallCancelled = true + self.blockedContinuation?.resume() + self.blockedContinuation = nil + } +} + +private actor CancellationAwareOpenAIDashboardLoader { + private let email: String + private var started = false + private var cancelled = false + + init(email: String) { + self.email = email + } + + func load() async throws -> OpenAIDashboardSnapshot { + self.started = true + + do { + try await Task.sleep(for: .seconds(30)) + } catch is CancellationError { + self.cancelled = true + throw CancellationError() + } + + return OpenAIDashboardSnapshot( + signedInEmail: self.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()) + } + + func waitUntilStarted(timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while !self.started { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func wasCancelled() -> Bool { + self.cancelled + } +} + private actor BlockingOpenAIDashboardCookieImport { - private var continuations: [ - CheckedContinuation, Never> - ] = [] - private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private typealias ImportResult = OpenAIDashboardBrowserCookieImporter.ImportResult + private typealias ResultContinuation = CheckedContinuation, Never> + + private var continuations: [(id: UUID, continuation: ResultContinuation)] = [] private var started = 0 + private var cancellations = 0 + private var cancelledIDs: Set = [] + private var rejectsNewCalls = false func awaitResult() async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult { - let result = await withCheckedContinuation { continuation in - self.continuations.append(continuation) - self.started += 1 - self.resumeReadyStartWaiters() + let id = UUID() + let result = await withTaskCancellationHandler { + await withCheckedContinuation { (continuation: ResultContinuation) in + if self.rejectsNewCalls || Task.isCancelled { + continuation.resume(returning: .failure(CancellationError())) + } else { + self.continuations.append((id: id, continuation: continuation)) + self.started += 1 + } + } + } onCancel: { + Task { await self.cancel(id: id) } } return try result.get() } - func waitUntilStarted(count: Int = 1) async { - if self.started >= count { return } - await withCheckedContinuation { continuation in - self.startWaiters.append((count: count, continuation: continuation)) + func waitUntilStarted(count: Int = 1, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.started < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) } + return true + } + + func waitUntilCancellationCount(_ count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.cancellations < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true } func resumeNext(with result: Result) { guard !self.continuations.isEmpty else { return } - let continuation = self.continuations.removeFirst() - continuation.resume(returning: result) + let record = self.continuations.removeFirst() + self.cancelledIDs.remove(record.id) + record.continuation.resume(returning: result) } - 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 + func cancelAll() { + self.rejectsNewCalls = true + let continuations = self.continuations + self.continuations.removeAll() + self.cancelledIDs.removeAll() + continuations.forEach { $0.continuation.resume(returning: .failure(CancellationError())) } + } + + private func cancel(id: UUID) { + guard self.continuations.contains(where: { $0.id == id }), self.cancelledIDs.insert(id).inserted else { return } + self.cancellations += 1 } } diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index a5ed210fca..25fb0fcc74 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -648,15 +648,7 @@ struct CodexManagedOpenAIWebRefreshTests { } private func makeSettingsStore(suite: String) throws -> SettingsStore { - let defaults = UserDefaults(suiteName: suite)! - defaults.removePersistentDomain(forName: suite) - defaults.set(true, forKey: "providerDetectionCompleted") - let configStore = testConfigStore(suiteName: suite) - let settings = SettingsStore( - userDefaults: defaults, - configStore: configStore, - zaiTokenStore: NoopZaiTokenStore(), - syntheticTokenStore: NoopSyntheticTokenStore()) + let settings = testSettingsStore(suiteName: suite) let codexMetadata = try #require(ProviderDescriptorRegistry.metadata[.codex]) settings.setProviderEnabled(provider: .codex, metadata: codexMetadata, enabled: true) settings.providerDetectionCompleted = true @@ -747,15 +739,28 @@ actor RefreshCompletionProbe { } actor BlockingManagedOpenAIDashboardLoader { - private var continuations: [CheckedContinuation, Never>] = [] + private typealias ResultContinuation = CheckedContinuation, Never> + + private var continuations: [(id: UUID, continuation: ResultContinuation)] = [] private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] private var started: Int = 0 + private var cancelledIDs: Set = [] + private var rejectsNewCalls = false func awaitResult() async throws -> OpenAIDashboardSnapshot { - let result = await withCheckedContinuation { continuation in - self.continuations.append(continuation) - self.started += 1 - self.resumeReadyStartWaiters() + let id = UUID() + let result = await withTaskCancellationHandler { + await withCheckedContinuation { (continuation: ResultContinuation) in + if self.rejectsNewCalls || Task.isCancelled { + continuation.resume(returning: .failure(CancellationError())) + } else { + self.continuations.append((id: id, continuation: continuation)) + self.started += 1 + self.resumeReadyStartWaiters() + } + } + } onCancel: { + Task { await self.cancel(id: id) } } return try result.get() } @@ -784,8 +789,17 @@ actor BlockingManagedOpenAIDashboardLoader { func resumeNext(with result: Result) { guard !self.continuations.isEmpty else { return } - let continuation = self.continuations.removeFirst() - continuation.resume(returning: result) + let record = self.continuations.removeFirst() + self.cancelledIDs.remove(record.id) + record.continuation.resume(returning: result) + } + + func cancelAll() { + self.rejectsNewCalls = true + let continuations = self.continuations + self.continuations.removeAll() + self.cancelledIDs.removeAll() + continuations.forEach { $0.continuation.resume(returning: .failure(CancellationError())) } } private func resumeReadyStartWaiters() { @@ -799,18 +813,37 @@ actor BlockingManagedOpenAIDashboardLoader { } self.startWaiters = remaining } + + private func cancel(id: UUID) { + guard self.continuations.contains(where: { $0.id == id }) else { return } + _ = self.cancelledIDs.insert(id) + } } actor BlockingCreditsLoader { - private var continuations: [CheckedContinuation, Never>] = [] + private typealias ResultContinuation = CheckedContinuation, Never> + + private var continuations: [(id: UUID, continuation: ResultContinuation)] = [] private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] private var started = 0 + private var cancellations = 0 + private var cancelledIDs: Set = [] + private var rejectsNewCalls = false func awaitResult() async throws -> CreditsSnapshot { - let result = await withCheckedContinuation { continuation in - self.continuations.append(continuation) - self.started += 1 - self.resumeReadyStartWaiters() + let id = UUID() + let result = await withTaskCancellationHandler { + await withCheckedContinuation { (continuation: ResultContinuation) in + if self.rejectsNewCalls || Task.isCancelled { + continuation.resume(returning: .failure(CancellationError())) + } else { + self.continuations.append((id: id, continuation: continuation)) + self.started += 1 + self.resumeReadyStartWaiters() + } + } + } onCancel: { + Task { await self.cancel(id: id) } } return try result.get() } @@ -822,14 +855,56 @@ actor BlockingCreditsLoader { } } + func waitUntilStartedWithin(count: Int = 1, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.started < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(50)) + } + return true + } + func startedCount() -> Int { self.started } + func cancellationCount() -> Int { + self.cancellations + } + + func waitUntilCancellationCount(_ count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.cancellations < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(50)) + } + return true + } + func resumeNext(with result: Result) { guard !self.continuations.isEmpty else { return } - let continuation = self.continuations.removeFirst() - continuation.resume(returning: result) + let record = self.continuations.removeFirst() + self.cancelledIDs.remove(record.id) + record.continuation.resume(returning: result) + } + + func resumeLast(with result: Result) { + guard !self.continuations.isEmpty else { return } + let record = self.continuations.removeLast() + self.cancelledIDs.remove(record.id) + record.continuation.resume(returning: result) + } + + func cancelAll() { + self.rejectsNewCalls = true + let continuations = self.continuations + self.continuations.removeAll() + self.cancelledIDs.removeAll() + continuations.forEach { $0.continuation.resume(returning: .failure(CancellationError())) } } private func resumeReadyStartWaiters() { @@ -843,6 +918,11 @@ actor BlockingCreditsLoader { } self.startWaiters = remaining } + + private func cancel(id: UUID) { + guard self.continuations.contains(where: { $0.id == id }), self.cancelledIDs.insert(id).inserted else { return } + self.cancellations += 1 + } } private actor OpenAIDashboardImportCallTracker { diff --git a/Tests/CodexBarTests/DeferredMenuInteractionRefreshTailTests.swift b/Tests/CodexBarTests/DeferredMenuInteractionRefreshTailTests.swift new file mode 100644 index 0000000000..11cc484a06 --- /dev/null +++ b/Tests/CodexBarTests/DeferredMenuInteractionRefreshTailTests.swift @@ -0,0 +1,177 @@ +import AppKit +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@MainActor +@Suite(.serialized) +struct DeferredMenuInteractionRefreshTailTests { + @Test + func `repeated scheduling during forced enrichment produces one deferred refresh`() async { + let settings = testSettingsStore( + suiteName: "DeferredMenuInteractionRefreshTailTests-forced-tail") + settings.providerDetectionCompleted = true + settings.refreshFrequency = .manual + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + settings.codexCookieSource = .off + for provider in UsageProvider.allCases { + guard let metadata = ProviderRegistry.shared.metadata[provider] else { continue } + settings.setProviderEnabled( + provider: provider, + metadata: metadata, + enabled: provider == .codex) + } + + let isolatedRoot = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-deferred-refresh-tests", isDirectory: true) + .appendingPathComponent(UUID().uuidString, isDirectory: true) + let environment = [ + "HOME": isolatedRoot.path, + "CODEX_HOME": isolatedRoot.appendingPathComponent(".codex", isDirectory: true).path, + ] + let fetcher = UsageFetcher(environment: environment) + let store = UsageStore( + fetcher: fetcher, + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + let tokenTail = DeferredMenuRefreshTokenTailBlocker() + var providerRefreshCount = 0 + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + providerRefreshCount += 1 + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + guard provider == .codex, force else { return } + await tokenTail.run() + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + } + + let controller = StatusItemController( + store: store, + settings: settings, + account: AccountInfo(email: nil, plan: nil), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: testStatusBar(), + menuCardRenderingEnabled: false, + menuRefreshEnabled: false) + defer { + controller.cancelDeferredMenuInteractionRefreshTask() + controller.releaseStatusItemsForTesting() + } + var deferredRefreshCount = 0 + controller.onDeferredMenuInteractionRefreshForTesting = { + deferredRefreshCount += 1 + } + + await store.refresh(enrichmentMode: .forcedBackground) + let enrichmentTask = store.forcedRefreshEnrichmentTask + let didStartTail = await tokenTail.waitUntilStarted() + #expect(didStartTail) + guard didStartTail else { + store.cancelForcedRefreshEnrichment() + await enrichmentTask?.value + return + } + #expect(providerRefreshCount == 1) + #expect(store.hasForcedRefreshEnrichmentInFlight) + + // The follow-up automatic refresh should not start another token-cost tail. + settings.costUsageEnabled = false + controller.deferMenuInteractionRefreshIfNeeded(providers: [.codex]) + for _ in 0..<3 { + controller.scheduleDeferredMenuInteractionRefreshIfNeeded(delay: .zero) + try? await Task.sleep(for: .milliseconds(30)) + } + + #expect(deferredRefreshCount == 0) + #expect(providerRefreshCount == 1) + #expect(controller.deferredMenuInteractionRefreshProviders == [.codex]) + #expect(controller.deferredMenuInteractionRefreshTask != nil) + + await tokenTail.release() + await enrichmentTask?.value + controller.scheduleDeferredMenuInteractionRefreshIfNeeded(delay: .zero) + + let completedExactlyOnce = await Self.waitUntil { + deferredRefreshCount == 1 && + providerRefreshCount == 2 && + !controller.deferredMenuInteractionRefreshPending + } + #expect(completedExactlyOnce) + #expect(deferredRefreshCount == 1) + #expect(providerRefreshCount == 2) + #expect(controller.deferredMenuInteractionRefreshProviders.isEmpty) + } + + private static func waitUntil( + timeout: Duration = .seconds(2), + condition: @escaping @MainActor () -> Bool) async -> Bool + { + let deadline = ContinuousClock.now + timeout + while !condition() { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } +} + +private actor DeferredMenuRefreshTokenTailBlocker { + private var started = 0 + private var released = false + private var waiter: (id: UUID, continuation: CheckedContinuation)? + + func run() async { + let id = UUID() + self.started += 1 + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + if self.released || Task.isCancelled { + continuation.resume() + } else { + self.waiter = (id: id, continuation: continuation) + } + } + } onCancel: { + Task { await self.cancel(id: id) } + } + } + + func waitUntilStarted(timeout: Duration = .seconds(2)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.started == 0 { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func release() { + self.released = true + self.waiter?.continuation.resume() + self.waiter = nil + } + + private func cancel(id: UUID) { + guard self.waiter?.id == id else { return } + self.waiter?.continuation.resume() + self.waiter = nil + } +} diff --git a/Tests/CodexBarTests/PersistentRefreshAccessibilityTests.swift b/Tests/CodexBarTests/PersistentRefreshAccessibilityTests.swift new file mode 100644 index 0000000000..67a61504e8 --- /dev/null +++ b/Tests/CodexBarTests/PersistentRefreshAccessibilityTests.swift @@ -0,0 +1,31 @@ +import AppKit +import Testing +@testable import CodexBar + +@Suite(.serialized) +@MainActor +struct PersistentRefreshAccessibilityTests { + @Test + func `disabled refresh row rejects accessibility press`() { + var pressCount = 0 + let view = PersistentRefreshMenuView( + title: "Refresh", + systemImageName: "arrow.clockwise", + shortcutText: "⌘ R", + onClick: { pressCount += 1 }) + + #expect(view.isAccessibilityEnabled()) + #expect(view.accessibilityPerformPress()) + #expect(pressCount == 1) + + view.setEnabled(false) + #expect(!view.isAccessibilityEnabled()) + #expect(!view.accessibilityPerformPress()) + #expect(pressCount == 1) + + view.setEnabled(true) + #expect(view.isAccessibilityEnabled()) + #expect(view.accessibilityPerformPress()) + #expect(pressCount == 2) + } +} diff --git a/Tests/CodexBarTests/ProviderRefreshCoordinatorTests.swift b/Tests/CodexBarTests/ProviderRefreshCoordinatorTests.swift index 9829410b79..b14aab9aa0 100644 --- a/Tests/CodexBarTests/ProviderRefreshCoordinatorTests.swift +++ b/Tests/CodexBarTests/ProviderRefreshCoordinatorTests.swift @@ -77,6 +77,18 @@ struct ProviderRefreshCoordinatorTests { #expect(result == .retryRequired) } + @Test + func `completed request is not offered for coalescing before deferred removal`() { + let coordinator = ProviderRefreshCoordinator() + let request = coordinator.beginReplacingRequest(for: "codex") + let task = Task {} + request.state.install(task: task) + + coordinator.complete(request.state, for: "codex", retryRequired: false) + + #expect(coordinator.coalescingState(for: "codex") == nil) + } + @Test func `completion removal and activity counts are key scoped`() async { let coordinator = ProviderRefreshCoordinator() diff --git a/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift b/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift new file mode 100644 index 0000000000..57f0596c0f --- /dev/null +++ b/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift @@ -0,0 +1,487 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +extension CodexBackgroundRefreshCoalescingTests { + @Test + func `required refresh requests during a pass share one follow-up`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-required-refresh-follow-up") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + settings.openAIWebAccessEnabled = false + let store = self.makeStore(settings: settings) + let providerGate = BlockingRequiredProviderRefresh() + store._test_providerRefreshOverride = { _ in + await providerGate.run(interaction: ProviderInteractionContext.current) + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + } + + let firstRefresh = Task { @MainActor in + await store.refreshForSettingsChange() + } + let didStartFirstPass = await providerGate.waitUntilStarted(count: 1) + #expect(didStartFirstPass) + guard didStartFirstPass else { + firstRefresh.cancel() + store.cancelRequiredRefresh() + await providerGate.cancelAll() + await firstRefresh.value + return + } + + var laterRefreshes: [Task] = [] + for expectedGeneration in 2...4 { + let task = Task { @MainActor in + await store.refreshForSettingsChange() + } + laterRefreshes.append(task) + for _ in 0..<100 where store.requiredRefreshRequestGeneration < expectedGeneration { + await Task.yield() + } + #expect(store.requiredRefreshRequestGeneration == expectedGeneration) + } + #expect(await providerGate.startedCount() == 1) + + await providerGate.resumeNext() + let didStartFollowUp = await providerGate.waitUntilStarted(count: 2) + #expect(didStartFollowUp) + guard didStartFollowUp else { + firstRefresh.cancel() + laterRefreshes.forEach { $0.cancel() } + store.cancelRequiredRefresh() + await providerGate.cancelAll() + await firstRefresh.value + for task in laterRefreshes { + await task.value + } + return + } + #expect(store.requiredRefreshCompletedGeneration == 1) + + await providerGate.resumeNext() + await firstRefresh.value + for task in laterRefreshes { + await task.value + } + try await Task.sleep(for: .milliseconds(50)) + + #expect(await providerGate.startedCount() == 2) + #expect(await providerGate.recordedInteractions() == [.background, .background]) + #expect(store.requiredRefreshCompletedGeneration == 4) + #expect(store.requiredRefreshTask == nil) + #expect(store.pendingRequiredRefreshRequest == nil) + } + + @Test + func `forced dashboard enrichment supersedes weaker background request`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-dashboard-supersedes-background") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let dashboardLoader = BlockingManagedOpenAIDashboardLoader() + var allowNavigationTimeoutRetries: [Bool] = [] + var dashboardInteractions: [ProviderInteraction] = [] + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_openAIDashboardLoaderOverride = { _, _, allowNavigationTimeoutRetry, _ in + allowNavigationTimeoutRetries.append(allowNavigationTimeoutRetry) + dashboardInteractions.append(ProviderInteractionContext.current) + return try await dashboardLoader.awaitResult() + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_openAIDashboardLoaderOverride = nil + } + + store.scheduleOpenAIDashboardRefreshIfNeeded( + expectedGuard: store.freshCodexOpenAIWebRefreshGuard()) + let didStartBackground = await dashboardLoader.waitUntilStartedWithin(count: 1) + #expect(didStartBackground) + guard didStartBackground else { + store.invalidateOpenAIDashboardRefreshTask() + await dashboardLoader.cancelAll() + return + } + let staleDashboardTask = store.openAIDashboardRefreshTask + let backgroundTask = store.openAIDashboardBackgroundRefreshTask + + let forcedRefresh = Task { @MainActor in + await BrowserCookieAccessGate.withExplicitRetry { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refresh(enrichmentMode: .forcedBackground) + await store.awaitForcedRefreshEnrichment() + } + } + } + let didStartForced = await dashboardLoader.waitUntilStartedWithin(count: 2) + #expect(didStartForced) + guard didStartForced else { + forcedRefresh.cancel() + store.cancelForcedRefreshEnrichment() + store.invalidateOpenAIDashboardRefreshTask() + await dashboardLoader.cancelAll() + await staleDashboardTask?.value + await backgroundTask?.value + await forcedRefresh.value + return + } + #expect(staleDashboardTask?.isCancelled == true) + #expect(backgroundTask?.isCancelled == true) + + await dashboardLoader.resumeNext(with: .failure(URLError(.timedOut))) + await staleDashboardTask?.value + await backgroundTask?.value + #expect(store.lastOpenAIDashboardError == nil) + + await dashboardLoader.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + await forcedRefresh.value + + #expect(allowNavigationTimeoutRetries == [false, true]) + #expect(dashboardInteractions == [.background, .userInitiated]) + #expect(store.openAIDashboard?.creditsRemaining == 25) + #expect(store.lastOpenAIDashboardError == nil) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `account scoped refresh supersedes weaker background dashboard request`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-account-dashboard-supersedes-background") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let dashboardLoader = BlockingManagedOpenAIDashboardLoader() + var allowNavigationTimeoutRetries: [Bool] = [] + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_openAIDashboardLoaderOverride = { _, _, allowNavigationTimeoutRetry, _ in + allowNavigationTimeoutRetries.append(allowNavigationTimeoutRetry) + return try await dashboardLoader.awaitResult() + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_openAIDashboardLoaderOverride = nil + } + + store.scheduleOpenAIDashboardRefreshIfNeeded( + expectedGuard: store.freshCodexOpenAIWebRefreshGuard()) + let didStartBackground = await dashboardLoader.waitUntilStartedWithin(count: 1) + #expect(didStartBackground) + guard didStartBackground else { + store.invalidateOpenAIDashboardRefreshTask() + await dashboardLoader.cancelAll() + return + } + let staleDashboardTask = store.openAIDashboardRefreshTask + let backgroundTask = store.openAIDashboardBackgroundRefreshTask + + let accountRefresh = Task { @MainActor in + await BrowserCookieAccessGate.withExplicitRetry { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await store.refreshCodexAccountScopedState() + } + } + } + let didStartForced = await dashboardLoader.waitUntilStartedWithin(count: 2) + #expect(didStartForced) + guard didStartForced else { + accountRefresh.cancel() + store.invalidateOpenAIDashboardRefreshTask() + await dashboardLoader.cancelAll() + await staleDashboardTask?.value + await backgroundTask?.value + await accountRefresh.value + return + } + #expect(staleDashboardTask?.isCancelled == true) + #expect(backgroundTask?.isCancelled == true) + + await dashboardLoader.resumeNext(with: .failure(URLError(.timedOut))) + await staleDashboardTask?.value + await backgroundTask?.value + #expect(store.lastOpenAIDashboardError == nil) + + await dashboardLoader.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + await accountRefresh.value + + #expect(allowNavigationTimeoutRetries == [false, true]) + #expect(store.openAIDashboard?.creditsRemaining == 25) + #expect(store.lastOpenAIDashboardError == nil) + } + + @Test + func `forced background refresh detaches stale dashboard before its tail`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-dashboard-detaches-account") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + let alphaAccount = try Self.installManagedAccount( + email: "alpha@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: alphaAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + store.syncOpenAIWebState() + let alphaDashboard = OpenAIDashboardSnapshot( + signedInEmail: alphaAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()) + store.openAIDashboard = alphaDashboard + store.openAIDashboardAttachmentAuthorized = true + store.lastOpenAIDashboardSnapshot = alphaDashboard + store.lastOpenAIDashboardAttachmentAuthorized = true + + let betaAccount = ManagedCodexAccount( + id: UUID(), + email: "beta@example.com", + managedHomePath: "/tmp/codexbar-managed-beta", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let tokenGate = BlockingForcedTokenRefresh() + store._test_providerRefreshOverride = { provider in + guard provider == .codex else { return } + settings._test_activeManagedCodexAccount = betaAccount + settings.codexActiveSource = .managedAccount(id: betaAccount.id) + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + guard force else { return } + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + defer { + settings._test_activeManagedCodexAccount = nil + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + #expect(await tokenGate.waitUntilStarted(count: 1)) + + #expect(store.openAIDashboard == nil) + #expect(!store.openAIDashboardAttachmentAuthorized) + #expect(store.lastOpenAIDashboardSnapshot == nil) + #expect(!store.lastOpenAIDashboardAttachmentAuthorized) + #expect(store.openAIDashboardRequiresLogin) + #expect(store.hasForcedRefreshEnrichmentInFlight) + + let enrichmentTask = store.forcedRefreshEnrichmentTask + store.cancelForcedRefreshEnrichment() + await tokenGate.resumeNext() + await enrichmentTask?.value + } + + @Test + func `forced token tail excludes periodic token sequence`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-token-excludes-timer") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + let store = self.makeStore(settings: settings) + let tokenGate = BlockingForcedTokenRefresh() + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + #expect(await tokenGate.waitUntilStarted(count: 1)) + + store.scheduleTokenRefreshForTesting() + try await Task.sleep(for: .milliseconds(100)) + #expect(await tokenGate.recordedCalls().count == 1) + + await tokenGate.resumeNext() + await store.awaitForcedRefreshEnrichment() + + let calls = await tokenGate.recordedCalls() + #expect(calls.count == 1) + #expect(calls.first?.force == true) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + func `forced enrichment excludes timer after token child completes`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-tail-excludes-token-timer") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + let tokenGate = BlockingForcedTokenRefresh() + let creditsGate = BlockingCreditsLoader() + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + try await creditsGate.awaitResult() + } + store._test_tokenUsageRefreshOverride = { provider, force in + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + #expect(await tokenGate.waitUntilStarted(count: 1)) + #expect(await creditsGate.waitUntilStartedWithin(count: 1)) + + await tokenGate.resumeNext() + for _ in 0..<100 where store.tokenRefreshSequenceTask != nil { + await Task.yield() + } + #expect(store.tokenRefreshSequenceTask == nil) + #expect(store.hasForcedRefreshEnrichmentInFlight) + + store.scheduleTokenRefreshForTesting() + try await Task.sleep(for: .milliseconds(100)) + #expect(store.tokenRefreshSequenceTask == nil) + #expect(await tokenGate.recordedCalls().count == 1) + + await creditsGate.resumeNext(with: .success(CreditsSnapshot( + remaining: 25, + events: [], + updatedAt: Date()))) + await store.awaitForcedRefreshEnrichment() + + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } +} + +private actor BlockingRequiredProviderRefresh { + private var interactions: [ProviderInteraction] = [] + private var continuations: [(id: UUID, continuation: CheckedContinuation)] = [] + + func run(interaction: ProviderInteraction) async { + let id = UUID() + self.interactions.append(interaction) + await withTaskCancellationHandler { + await withCheckedContinuation { (continuation: CheckedContinuation) in + if Task.isCancelled { + continuation.resume() + } else { + self.continuations.append((id: id, continuation: continuation)) + } + } + } onCancel: { + Task { await self.cancel(id: id) } + } + } + + func waitUntilStarted(count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.interactions.count < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func startedCount() -> Int { + self.interactions.count + } + + func recordedInteractions() -> [ProviderInteraction] { + self.interactions + } + + func resumeNext() { + guard !self.continuations.isEmpty else { return } + self.continuations.removeFirst().continuation.resume() + } + + func cancelAll() { + let continuations = self.continuations + self.continuations.removeAll() + continuations.forEach { $0.continuation.resume() } + } + + private func cancel(id: UUID) { + guard let index = self.continuations.firstIndex(where: { $0.id == id }) else { return } + self.continuations.remove(at: index).continuation.resume() + } +} diff --git a/Tests/CodexBarTests/StatusItemControllerShutdownTests.swift b/Tests/CodexBarTests/StatusItemControllerShutdownTests.swift index 69e185a9df..7f4745ef81 100644 --- a/Tests/CodexBarTests/StatusItemControllerShutdownTests.swift +++ b/Tests/CodexBarTests/StatusItemControllerShutdownTests.swift @@ -27,12 +27,18 @@ struct StatusItemControllerShutdownTests { settings.setProviderEnabled(provider: .claude, metadata: claudeMetadata, enabled: true) } - let fetcher = UsageFetcher() - let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let environment = Self.isolatedEnvironment() + let fetcher = UsageFetcher(environment: environment) + let store = UsageStore( + fetcher: fetcher, + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) let controller = StatusItemController( store: store, settings: settings, - account: fetcher.loadAccountInfo(), + account: AccountInfo(email: nil, plan: nil), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: .system) @@ -102,6 +108,117 @@ struct StatusItemControllerShutdownTests { #expect(didTerminate) } + @Test + func `app shutdown cancels forced enrichment`() async { + let controller = self.makeController() + defer { + StatusItemController.menuCardRenderingEnabled = !SettingsStore.isRunningTests + StatusItemController.resetMenuRefreshEnabledForTesting() + } + controller.settings.statusChecksEnabled = false + controller.settings.costUsageEnabled = true + controller.settings.openAIWebAccessEnabled = false + controller.settings.codexCookieSource = .off + let tokenTail = CancellationAwareTokenTail() + + controller.store._test_providerRefreshOverride = { _ in } + controller.store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + controller.store._test_tokenUsageRefreshOverride = { _, _ in + await tokenTail.run() + } + defer { + controller.store._test_providerRefreshOverride = nil + controller.store._test_codexCreditsLoaderOverride = nil + controller.store._test_tokenUsageRefreshOverride = nil + } + + controller.refreshNow() + let didStartTokenTail = await tokenTail.waitUntilStarted() + #expect(didStartTokenTail) + guard didStartTokenTail else { + controller.prepareForAppShutdown() + return + } + await controller.manualRefreshTasks[.global]?.value + let enrichmentTask = controller.store.forcedRefreshEnrichmentTask + let requiredRefresh = Task { @MainActor in + await controller.store.refreshForSettingsChange() + } + for _ in 0..<100 where controller.store.requiredRefreshTask == nil { + await Task.yield() + } + + #expect(controller.store.hasForcedRefreshEnrichmentInFlight) + #expect(controller.store.requiredRefreshTask != nil) + controller.prepareForAppShutdown() + await enrichmentTask?.value + await requiredRefresh.value + + #expect(await tokenTail.wasCancelled()) + #expect(!controller.store.hasForcedRefreshEnrichmentInFlight) + #expect(controller.store.forcedRefreshEnrichmentTask == nil) + #expect(controller.store.pendingForcedRefreshEnrichmentTask == nil) + #expect(controller.store.requiredRefreshTask == nil) + #expect(controller.store.pendingRequiredRefreshRequest == nil) + #expect(controller.store.openAIDashboardRefreshTask == nil) + #expect(controller.store.tokenRefreshInFlight.isEmpty) + } + + @Test + func `app shutdown cancels active and pending forced enrichment without promotion`() async { + let controller = self.makeController() + defer { + StatusItemController.menuCardRenderingEnabled = !SettingsStore.isRunningTests + StatusItemController.resetMenuRefreshEnabledForTesting() + } + controller.settings.statusChecksEnabled = false + controller.settings.costUsageEnabled = true + controller.settings.openAIWebAccessEnabled = false + controller.settings.codexCookieSource = .off + let tokenTail = CancellationAwareTokenTail() + + controller.store._test_providerRefreshOverride = { _ in } + controller.store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + controller.store._test_tokenUsageRefreshOverride = { _, _ in + await tokenTail.run() + } + defer { + controller.store._test_providerRefreshOverride = nil + controller.store._test_codexCreditsLoaderOverride = nil + controller.store._test_tokenUsageRefreshOverride = nil + } + + controller.refreshNow() + let didStartTokenTail = await tokenTail.waitUntilStarted(count: 1) + #expect(didStartTokenTail) + guard didStartTokenTail else { + controller.prepareForAppShutdown() + return + } + await controller.manualRefreshTasks[.global]?.value + + await controller.store.refresh(enrichmentMode: .forcedBackground) + let activeTask = controller.store.forcedRefreshEnrichmentTask + let pendingTask = controller.store.pendingForcedRefreshEnrichmentTask + #expect(activeTask != nil) + #expect(pendingTask != nil) + + controller.prepareForAppShutdown() + await activeTask?.value + await pendingTask?.value + + #expect(await tokenTail.startedCount() == 1) + #expect(await tokenTail.cancelledCount() == 1) + #expect(pendingTask?.isCancelled == true) + #expect(!controller.store.hasForcedRefreshEnrichmentInFlight) + #expect(controller.store.forcedRefreshEnrichmentTask == nil) + #expect(controller.store.pendingForcedRefreshEnrichmentTask == nil) + } + private func makeController() -> StatusItemController { StatusItemController.menuCardRenderingEnabled = false StatusItemController.setMenuRefreshEnabledForTesting(true) @@ -114,25 +231,72 @@ struct StatusItemControllerShutdownTests { settings.setProviderEnabled(provider: .codex, metadata: codexMetadata, enabled: true) } - let fetcher = UsageFetcher() - let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let environment = Self.isolatedEnvironment() + let fetcher = UsageFetcher(environment: environment) + let store = UsageStore( + fetcher: fetcher, + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) return StatusItemController( store: store, settings: settings, - account: fetcher.loadAccountInfo(), + account: AccountInfo(email: nil, plan: nil), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: .system) } private func makeSettings() -> SettingsStore { - let suite = "StatusItemControllerShutdownTests-\(UUID().uuidString)" - let defaults = UserDefaults(suiteName: suite)! - defaults.removePersistentDomain(forName: suite) - return SettingsStore( - userDefaults: defaults, - configStore: testConfigStore(suiteName: suite), - zaiTokenStore: NoopZaiTokenStore(), - syntheticTokenStore: NoopSyntheticTokenStore()) + testSettingsStore(suiteName: "StatusItemControllerShutdownTests") + } + + private static func isolatedEnvironment() -> [String: String] { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-tests", isDirectory: true) + .appendingPathComponent(UUID().uuidString, isDirectory: true) + return [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + "XDG_CONFIG_HOME": root.appendingPathComponent(".config", isDirectory: true).path, + ] + } +} + +private actor CancellationAwareTokenTail { + private var started = 0 + private var cancelled = 0 + + func run() async { + self.started += 1 + do { + try await Task.sleep(for: .seconds(30)) + } catch is CancellationError { + self.cancelled += 1 + } catch {} + } + + func waitUntilStarted(count: Int = 1, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.started < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } + + func wasCancelled() -> Bool { + self.cancelled > 0 + } + + func startedCount() -> Int { + self.started + } + + func cancelledCount() -> Int { + self.cancelled } } diff --git a/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift b/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift index 9a4afe5b7b..17e4e5c785 100644 --- a/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift +++ b/Tests/CodexBarTests/StatusMenuInstantOpenTests.swift @@ -686,16 +686,7 @@ extension StatusMenuTests { 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) + store._setSnapshotForTesting(self.instantOpenSnapshot(email: "refreshed@example.com"), provider: .codex) } defer { store._test_providerRefreshOverride = nil } let existingRefreshTask = Task { @@ -748,6 +739,7 @@ extension StatusMenuTests { store._test_providerRefreshOverride = { provider in guard provider == .codex else { return } await refreshGate.run() + store._setSnapshotForTesting(self.instantOpenSnapshot(email: "refreshed@example.com"), provider: .codex) } defer { store._test_providerRefreshOverride = nil } let controller = StatusItemController( diff --git a/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift b/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift index 06cee511a8..ef9ca20e86 100644 --- a/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift +++ b/Tests/CodexBarTests/StatusMenuPersistentRefreshTests.swift @@ -64,21 +64,30 @@ private final class ManualRefreshGate { self.isOpen = true } } + + func waitUntilSignaled(timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while !self.isOpen { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + self.isOpen = false + return true + } +} + +enum BlockingEnrichmentStage: Sendable { + case credits + case dashboard } @MainActor @Suite(.serialized) struct StatusMenuPersistentRefreshTests { private func makeSettings() -> SettingsStore { - let suite = "StatusMenuPersistentRefreshTests-\(UUID().uuidString)" - let defaults = UserDefaults(suiteName: suite)! - defaults.removePersistentDomain(forName: suite) - let configStore = testConfigStore(suiteName: suite) - return SettingsStore( - userDefaults: defaults, - configStore: configStore, - zaiTokenStore: NoopZaiTokenStore(), - syntheticTokenStore: NoopSyntheticTokenStore()) + testSettingsStore(suiteName: "StatusMenuPersistentRefreshTests") } private func makeController( @@ -86,8 +95,14 @@ struct StatusMenuPersistentRefreshTests { updater: UpdaterProviding = DisabledUpdaterController(), account: AccountInfo? = nil) -> StatusItemController { - let fetcher = UsageFetcher() - let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let environment = Self.isolatedEnvironment() + let fetcher = UsageFetcher(environment: environment) + let store = UsageStore( + fetcher: fetcher, + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) if let account { store.accountInfoCache[.codex] = UsageStore.AccountInfoCacheEntry( account: account, @@ -97,12 +112,23 @@ struct StatusMenuPersistentRefreshTests { return StatusItemController( store: store, settings: settings, - account: account ?? fetcher.loadAccountInfo(), + account: account ?? AccountInfo(email: nil, plan: nil), updater: updater, preferencesSelection: PreferencesSelection(), statusBar: .system) } + private static func isolatedEnvironment() -> [String: String] { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-tests", isDirectory: true) + .appendingPathComponent(UUID().uuidString, isDirectory: true) + return [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + "XDG_CONFIG_HOME": root.appendingPathComponent(".config", isDirectory: true).path, + ] + } + private func enableOnly(_ providers: Set, settings: SettingsStore) { for provider in UsageProvider.allCases { guard let metadata = ProviderRegistry.shared.metadata[provider] else { continue } @@ -565,7 +591,9 @@ struct StatusMenuPersistentRefreshTests { #expect(inFlight.metrics.isEmpty) #expect(inFlight.creditsText == fallback.creditsText) } +} +extension StatusMenuPersistentRefreshTests { @Test func `manual refresh uses fallback when empty quota gains a placeholder`() throws { let settings = self.makeSettings() @@ -1187,31 +1215,234 @@ extension StatusMenuPersistentRefreshTests { let controller = self.makeController(settings: settings) let tokenRefreshStarted = ManualRefreshGate() let releaseTokenRefresh = ManualRefreshGate() + defer { releaseTokenRefresh.resume() } let monitor = controller.menuCardRefreshMonitor let fallback = MenuCardLiveSubtitle(text: "Idle", style: .info) + let menu = controller.makeMenu() + var tokenRefreshWasForced: Bool? controller.store._test_providerRefreshOverride = { _ in } - controller.store._test_tokenUsageRefreshOverride = { _, _ in + controller.store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + controller.store._test_tokenUsageRefreshOverride = { _, force in + tokenRefreshWasForced = force tokenRefreshStarted.resume() await releaseTokenRefresh.wait() } defer { controller.store._test_providerRefreshOverride = nil + controller.store._test_codexCreditsLoaderOverride = nil controller.store._test_tokenUsageRefreshOverride = nil } controller.refreshNow() - let manualTask = try? #require(controller.manualRefreshTasks[.global]) - await tokenRefreshStarted.wait() - await manualTask?.value + let manualTask = controller.manualRefreshTasks[.global] + #expect(await tokenRefreshStarted.waitUntilSignaled()) + let manualCompletion = RefreshCompletionProbe() + let manualCompletionTask = Task { + await manualTask?.value + await manualCompletion.markCompleted() + } + #expect(await manualCompletion.waitUntilCompleted()) + let tailTask = controller.store.forcedRefreshEnrichmentTask #expect(!controller.store.isRefreshing) + #expect(controller.store.hasForcedRefreshEnrichmentInFlight) #expect(controller.store.refreshingProviders.isEmpty) + #expect(tokenRefreshWasForced == true) + #expect(controller.isRefreshActionInFlight(for: menu)) #expect(!monitor.isManualRefreshInFlight) #expect(!monitor.isManualRefreshInFlight(for: .codex)) #expect(monitor.subtitle(for: .codex, fallback: fallback).style == .info) releaseTokenRefresh.resume() + await manualCompletionTask.value + await tailTask?.value + #expect(!controller.store.hasForcedRefreshEnrichmentInFlight) + #expect(!controller.isRefreshActionInFlight(for: menu)) + } + + @Test + func `credit tail does not keep completed provider card refreshing`() async { + await self.verifyCompletedProviderCardStopsRefreshing(whileBlocking: .credits) + } + + @Test + func `dashboard tail does not keep completed provider card refreshing`() async { + await self.verifyCompletedProviderCardStopsRefreshing(whileBlocking: .dashboard) + } + + private func verifyCompletedProviderCardStopsRefreshing(whileBlocking stage: BlockingEnrichmentStage) async { + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + settings.openAIWebAccessEnabled = stage == .dashboard + settings.codexCookieSource = stage == .dashboard ? .auto : .off + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "fixture@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date(), + identity: .emailOnly(normalizedEmail: "fixture@example.com")) + settings.codexActiveSource = .liveSystem + self.enableOnly([.codex], settings: settings) + let account = AccountInfo(email: "fixture@example.com", plan: "pro") + let controller = self.makeController(settings: settings, account: account) + controller.store.accountInfoCache[.codex] = UsageStore.AccountInfoCacheEntry( + account: account, + configRevision: settings.configRevision, + expiresAt: .distantFuture) + let stageStarted = ManualRefreshGate() + let releaseStage = ManualRefreshGate() + defer { + releaseStage.resume() + controller.prepareForAppShutdown() + settings._test_liveSystemCodexAccount = nil + } + let monitor = controller.menuCardRefreshMonitor + let fallback = MenuCardLiveSubtitle(text: "Idle", style: .info) + let menu = controller.makeMenu() + var creditsLoaderCalls = 0 + var dashboardLoaderCalls = 0 + + controller.store._test_providerRefreshOverride = { _ in } + controller.store._test_codexCreditsLoaderOverride = { + creditsLoaderCalls += 1 + if stage == .credits, creditsLoaderCalls == 1 { + stageStarted.resume() + await releaseStage.wait() + } + return CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + controller.store._test_openAIDashboardLoaderOverride = { _, _, _, _ in + dashboardLoaderCalls += 1 + if stage == .dashboard, dashboardLoaderCalls == 1 { + stageStarted.resume() + await releaseStage.wait() + } + return OpenAIDashboardSnapshot( + signedInEmail: account.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()) + } + defer { + controller.store._test_providerRefreshOverride = nil + controller.store._test_codexCreditsLoaderOverride = nil + controller.store._test_openAIDashboardLoaderOverride = nil + } + + controller.refreshNow() + let manualTask = controller.manualRefreshTasks[.global] + let didStartStage = await stageStarted.waitUntilSignaled() + #expect(didStartStage) + guard didStartStage else { return } + let manualCompletion = RefreshCompletionProbe() + let manualCompletionTask = Task { + await manualTask?.value + await manualCompletion.markCompleted() + } + let didCompleteManualRefresh = await manualCompletion.waitUntilCompleted() + #expect(didCompleteManualRefresh) + guard didCompleteManualRefresh else { return } + let tailTask = controller.store.forcedRefreshEnrichmentTask + + #expect(!controller.store.isRefreshing) + #expect(controller.store.hasForcedRefreshEnrichmentInFlight) + #expect(controller.store.refreshingProviders.isEmpty) + #expect(controller.isRefreshActionInFlight(for: menu)) + #expect(!monitor.isManualRefreshInFlight) + #expect(!monitor.isManualRefreshInFlight(for: .codex)) + #expect(monitor.subtitle(for: .codex, fallback: fallback).style == .info) + + releaseStage.resume() + await manualCompletionTask.value + let tailCompletion = RefreshCompletionProbe() + let tailCompletionTask = Task { + await tailTask?.value + await tailCompletion.markCompleted() + } + let didCompleteTail = await tailCompletion.waitUntilCompleted() + #expect(didCompleteTail) + guard didCompleteTail else { return } + await tailCompletionTask.value + + #expect(creditsLoaderCalls >= 1) + #expect(dashboardLoaderCalls == (stage == .dashboard ? 1 : 0)) + #expect(!controller.store.hasForcedRefreshEnrichmentInFlight) + #expect(!controller.isRefreshActionInFlight(for: menu)) + } + + @Test + func `provider scoped refresh waits for global forced enrichment`() async { + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + settings.codexCookieSource = .off + self.enableOnly([.codex], settings: settings) + let controller = self.makeController(settings: settings) + let tokenRefreshStarted = ManualRefreshGate() + let releaseTokenRefresh = ManualRefreshGate() + let scopedWaitStarted = ManualRefreshGate() + defer { releaseTokenRefresh.resume() } + var providerRefreshCalls = 0 + var tokenRefreshCalls = 0 + + controller.store._test_providerRefreshOverride = { _ in + providerRefreshCalls += 1 + } + controller.store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + controller.store._test_tokenUsageRefreshOverride = { _, _ in + tokenRefreshCalls += 1 + if tokenRefreshCalls == 1 { + tokenRefreshStarted.resume() + await releaseTokenRefresh.wait() + } + } + controller.store._test_forcedRefreshEnrichmentWaitObserver = { + scopedWaitStarted.resume() + } + defer { + controller.store._test_providerRefreshOverride = nil + controller.store._test_codexCreditsLoaderOverride = nil + controller.store._test_tokenUsageRefreshOverride = nil + controller.store._test_forcedRefreshEnrichmentWaitObserver = nil + } + + controller.refreshNow() + let globalTask = controller.manualRefreshTasks[.global] + #expect(await tokenRefreshStarted.waitUntilSignaled()) + let globalCompletion = RefreshCompletionProbe() + let globalCompletionTask = Task { + await globalTask?.value + await globalCompletion.markCompleted() + } + #expect(await globalCompletion.waitUntilCompleted()) + #expect(providerRefreshCalls == 1) + + let scopedTask = Task { @MainActor in + await controller.performStoreRefresh( + for: .codex, + refreshOpenMenusWhenComplete: false, + interaction: .userInitiated) + } + #expect(await scopedWaitStarted.waitUntilSignaled()) + #expect(providerRefreshCalls == 1) + + releaseTokenRefresh.resume() + await globalCompletionTask.value + await scopedTask.value + #expect(providerRefreshCalls == 2) + #expect(tokenRefreshCalls == 2) + #expect(!controller.store.hasForcedRefreshEnrichmentInFlight) } @Test diff --git a/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift b/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift index 2268cc3431..b770f8f6d1 100644 --- a/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift +++ b/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift @@ -78,6 +78,17 @@ private actor TokenRefreshRecorder { func record(provider: UsageProvider, force: Bool) { self.calls.append((provider, force)) } + + func waitForCallCount(_ count: Int, timeout: Duration = .seconds(5)) async -> Bool { + let deadline = ContinuousClock.now + timeout + while self.calls.count < count { + if ContinuousClock.now >= deadline { + return false + } + try? await Task.sleep(for: .milliseconds(10)) + } + return true + } } @MainActor @@ -273,15 +284,39 @@ struct UsageStoreManualTokenRefreshTests { } } + @Test + func `forced background refresh bypasses a fresh token cache`() async { + let store = Self.makeStore() + let recorder = TokenRefreshRecorder() + store._test_providerRefreshOverride = { _ in } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + await recorder.record(provider: provider, force: force) + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + } + + await store.refresh(forceTokenUsage: false) + let didRecordScheduledRefresh = await recorder.waitForCallCount(1) + #expect(didRecordScheduledRefresh) + guard didRecordScheduledRefresh else { + store.cancelForcedRefreshEnrichment() + return + } + await store.refresh(enrichmentMode: .forcedBackground) + await store.awaitForcedRefreshEnrichment() + + #expect(await recorder.calls.map(\.provider) == [.codex, .codex]) + #expect(await recorder.calls.map(\.force) == [false, true]) + } + private static func makeStore(enabledProviders: Set = [.codex]) -> UsageStore { - let suite = "UsageStoreManualTokenRefreshTests-\(UUID().uuidString)" - let defaults = UserDefaults(suiteName: suite)! - defaults.removePersistentDomain(forName: suite) - let settings = SettingsStore( - userDefaults: defaults, - configStore: testConfigStore(suiteName: suite), - zaiTokenStore: NoopZaiTokenStore(), - syntheticTokenStore: NoopSyntheticTokenStore()) + let settings = testSettingsStore(suiteName: "UsageStoreManualTokenRefreshTests") settings.refreshFrequency = .manual settings.statusChecksEnabled = false settings.costUsageEnabled = true @@ -298,11 +333,19 @@ struct UsageStoreManualTokenRefreshTests { enabled: enabledProviders.contains(provider)) } + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-tests", isDirectory: true) + .appendingPathComponent(UUID().uuidString, isDirectory: true) + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + "XDG_CONFIG_HOME": root.appendingPathComponent(".config", isDirectory: true).path, + ] return UsageStore( - fetcher: UsageFetcher(), + fetcher: UsageFetcher(environment: environment), browserDetection: BrowserDetection(cacheTTL: 0), settings: settings, startupBehavior: .testing, - environmentBase: [:]) + environmentBase: environment) } } diff --git a/Tests/CodexBarTests/UsageStoreResetBoundaryRefreshTests.swift b/Tests/CodexBarTests/UsageStoreResetBoundaryRefreshTests.swift index b199270e42..1e11a59493 100644 --- a/Tests/CodexBarTests/UsageStoreResetBoundaryRefreshTests.swift +++ b/Tests/CodexBarTests/UsageStoreResetBoundaryRefreshTests.swift @@ -113,6 +113,115 @@ struct UsageStoreResetBoundaryRefreshTests { #expect(store.scheduledResetBoundaryRefreshAt == retryAt) } + @Test + @MainActor + func `boundary refresh records its attempt before waiting for forced enrichment`() async { + let settings = testSettingsStore(suiteName: "UsageStoreResetBoundaryRefreshTests-waits-for-tail") + settings.statusChecksEnabled = false + settings.costUsageEnabled = true + settings.openAIWebAccessEnabled = false + settings.codexCookieSource = .off + for provider in UsageProvider.allCases { + guard let metadata = ProviderRegistry.shared.metadata[provider] else { continue } + settings.setProviderEnabled(provider: provider, metadata: metadata, enabled: provider == .codex) + } + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let tokenGate = BlockingForcedTokenRefresh() + var providerRefreshes = 0 + var didObserveWait = false + store._test_providerRefreshOverride = { _ in + providerRefreshes += 1 + } + store._test_codexCreditsLoaderOverride = { + CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()) + } + store._test_tokenUsageRefreshOverride = { provider, force in + guard force else { return } + await tokenGate.run( + provider: provider, + force: force, + interaction: ProviderInteractionContext.current, + refreshPhase: ProviderRefreshContext.current, + browserRetryAllowed: false) + } + store._test_forcedRefreshEnrichmentWaitObserver = { + didObserveWait = true + } + defer { + store._test_providerRefreshOverride = nil + store._test_codexCreditsLoaderOverride = nil + store._test_tokenUsageRefreshOverride = nil + store._test_forcedRefreshEnrichmentWaitObserver = nil + } + + await store.refresh(enrichmentMode: .forcedBackground) + #expect(await tokenGate.waitUntilStarted(count: 1)) + settings.costUsageEnabled = false + + let boundaryRefreshAt = Date(timeIntervalSince1970: 12345) + let boundaryTask = Task { @MainActor in + await store.runResetBoundaryRefresh(boundaryRefreshAt: boundaryRefreshAt) + } + for _ in 0..<100 where !didObserveWait { + await Task.yield() + } + + #expect(didObserveWait) + #expect(providerRefreshes == 1) + #expect(store.attemptedResetBoundaryRefreshes.contains(boundaryRefreshAt)) + + await tokenGate.resumeNext() + await boundaryTask.value + + #expect(providerRefreshes == 2) + #expect(store.attemptedResetBoundaryRefreshes.contains(boundaryRefreshAt)) + #expect(!store.hasForcedRefreshEnrichmentInFlight) + } + + @Test + @MainActor + func `boundary refresh does not reschedule unchanged stale snapshot`() async { + let settings = testSettingsStore(suiteName: "UsageStoreResetBoundaryRefreshTests-no-duplicate") + settings.refreshFrequency = .oneMinute + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + settings.openAIWebAccessEnabled = false + settings.codexCookieSource = .off + for provider in UsageProvider.allCases { + guard let metadata = ProviderRegistry.shared.metadata[provider] else { continue } + settings.setProviderEnabled(provider: provider, metadata: metadata, enabled: provider == .codex) + } + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let now = Date() + let boundaryRefreshAt = now.addingTimeInterval(1) + store.snapshots[.codex] = Self.snapshot( + updatedAt: now.addingTimeInterval(-60), + primaryResetsAt: boundaryRefreshAt.addingTimeInterval(-UsageStore.resetBoundaryRefreshGraceSeconds)) + var providerRefreshes = 0 + store._test_providerRefreshOverride = { _ in + providerRefreshes += 1 + } + defer { + store._test_providerRefreshOverride = nil + store.cancelResetBoundaryRefresh() + } + + await store.runResetBoundaryRefresh(boundaryRefreshAt: boundaryRefreshAt) + + #expect(providerRefreshes == 1) + #expect(store.attemptedResetBoundaryRefreshes.contains(boundaryRefreshAt)) + #expect(store.resetBoundaryRefreshTask == nil) + #expect(store.scheduledResetBoundaryRefreshAt == nil) + } + @Test func `ignores reset boundary after normal poll`() { let now = Date(timeIntervalSince1970: 3000) From 57cd28cf6e5a67c3ab35b85e8e273884a758aa6e Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 10 Jul 2026 15:26:17 -0700 Subject: [PATCH 3/4] fix: preserve queued token refresh ownership --- .../UsageStore+TokenRefreshSequence.swift | 23 ++++++++++--------- .../UsageStoreManualTokenRefreshTests.swift | 23 +++++++++++++++++++ 2 files changed, 35 insertions(+), 11 deletions(-) diff --git a/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift b/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift index e7df6d1e9f..d6d230e74b 100644 --- a/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift +++ b/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift @@ -67,19 +67,20 @@ extension UsageStore { force: Bool, scope: TokenRefreshSequenceScope) -> Task { + let providers: [UsageProvider] = switch scope { + case .all: + self.enabledProvidersForBackgroundWork() + case let .provider(provider): + [provider] + } let token = UUID() self.tokenRefreshSequenceToken = token + // Publish the first owner before installing the task. A scoped forced refresh can arrive + // before the task gets its first MainActor turn and must not mistake this slot for unknown work. + self.tokenRefreshSequenceProvider = providers.first let task = Task(priority: .utility) { @MainActor [weak self] in guard let self else { return } - switch scope { - case .all: - await self.refreshTokenUsageSequence(force: force) - case let .provider(provider): - self.tokenRefreshSequenceProvider = provider - await self.refreshTokenUsage(provider, force: force) - self.tokenRefreshSequenceProvider = nil - self.scheduleMemoryPressureRelief() - } + await self.refreshTokenUsageSequence(providers: providers, force: force) self.completeTokenRefreshSequence(token: token) } self.tokenRefreshSequenceTask = task @@ -101,9 +102,9 @@ extension UsageStore { self.tokenRefreshSequenceProvider = nil } - private func refreshTokenUsageSequence(force: Bool) async { + private func refreshTokenUsageSequence(providers: [UsageProvider], force: Bool) async { defer { self.tokenRefreshSequenceProvider = nil } - for provider in self.enabledProvidersForBackgroundWork() { + for provider in providers { if Task.isCancelled { break } diff --git a/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift b/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift index b770f8f6d1..eb28ffd4fd 100644 --- a/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift +++ b/Tests/CodexBarTests/UsageStoreManualTokenRefreshTests.swift @@ -260,6 +260,29 @@ struct UsageStoreManualTokenRefreshTests { await scheduledGate.waitForFinish() } + @Test + func `scoped manual refresh preserves an unrelated token sequence before it starts`() async { + let store = Self.makeStore(enabledProviders: [.claude, .codex]) + let recorder = TokenRefreshRecorder() + store._test_tokenUsageRefreshOverride = { provider, force in + await recorder.record(provider: provider, force: force) + } + + // Do not yield between installing the scheduled slot and starting the scoped refresh. This + // exercises the window before the scheduled task receives its first MainActor turn. + store.scheduleTokenRefreshForTesting() + await store.refreshTokenUsageNow(for: .claude, force: true) + + let recordedBothRefreshes = await recorder.waitForCallCount(2) + #expect(recordedBothRefreshes) + let scheduledTask = store.tokenRefreshSequenceTask + await scheduledTask?.value + + let calls = await recorder.calls + #expect(calls.contains { $0.provider == .codex && !$0.force }) + #expect(calls.contains { $0.provider == .claude && $0.force }) + } + @Test func `regular refresh schedules token-cost refresh without waiting`() async { let store = Self.makeStore() From cbface08793c6bc7f71ebf5e6333ab19ba0e9dde Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 10 Jul 2026 16:26:02 -0700 Subject: [PATCH 4/4] fix: cancel stale dashboard schedulers --- Sources/CodexBar/UsageStore+OpenAIWeb.swift | 1 + .../RequiredRefreshCoalescingTests.swift | 68 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index 0abc974835..9c91d812a0 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -507,6 +507,7 @@ extension UsageStore { } } + guard !Task.isCancelled else { return } await self.refreshOpenAIDashboardIfNeeded(force: false, expectedGuard: expectedGuard) guard !Task.isCancelled else { return } self.persistWidgetSnapshot(reason: "dashboard") diff --git a/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift b/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift index 57f0596c0f..f0694de8c1 100644 --- a/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift +++ b/Tests/CodexBarTests/RequiredRefreshCoalescingTests.swift @@ -80,6 +80,74 @@ extension CodexBackgroundRefreshCoalescingTests { #expect(store.pendingRequiredRefreshRequest == nil) } + @Test + func `forced dashboard refresh stops a queued stale scheduler before it starts`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-forced-dashboard-prestart-cancellation") + settings.statusChecksEnabled = false + settings.costUsageEnabled = false + let managedAccount = try Self.installManagedAccount( + email: "managed@example.com", + settings: settings) + defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) } + + let store = self.makeStore(settings: settings) + var allowNavigationTimeoutRetries: [Bool] = [] + store._test_openAIDashboardCookieImportOverride = { targetEmail, _, _, _, _ in + OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Fixture", + cookieCount: 2, + signedInEmail: targetEmail, + matchesCodexEmail: true) + } + store._test_openAIDashboardLoaderOverride = { _, _, allowNavigationTimeoutRetry, _ in + allowNavigationTimeoutRetries.append(allowNavigationTimeoutRetry) + return OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()) + } + defer { + store._test_openAIDashboardCookieImportOverride = nil + store._test_openAIDashboardLoaderOverride = nil + } + + let currentGuard = store.freshCodexOpenAIWebRefreshGuard() + let backgroundGuard = CodexAccountScopedRefreshGuard( + source: currentGuard.source, + identity: currentGuard.identity, + accountKey: currentGuard.accountKey, + authFingerprint: "background-token-material") + let forcedGuard = CodexAccountScopedRefreshGuard( + source: currentGuard.source, + identity: currentGuard.identity, + accountKey: currentGuard.accountKey, + authFingerprint: "forced-token-material") + store.openAIWebAccountDidChange = true + + // Keep both calls on this MainActor turn so the forced request cancels the scheduler + // before its task body starts. + store.scheduleOpenAIDashboardRefreshIfNeeded(expectedGuard: backgroundGuard) + let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask) + await store.refreshOpenAIDashboardIfNeeded( + force: true, + expectedGuard: forcedGuard, + bypassCoalescing: true, + allowCodexUsageBackfill: false) + await backgroundTask.value + + #expect(backgroundTask.isCancelled) + #expect(allowNavigationTimeoutRetries == [true]) + #expect(store.openAIDashboardBackgroundRefreshTask == nil) + #expect(store.openAIDashboardRefreshTask == nil) + } + @Test func `forced dashboard enrichment supersedes weaker background request`() async throws { let settings = try self.makeSettingsStore(