From 40db997bbf81842f6d2975a4dcd657bf493b0d99 Mon Sep 17 00:00:00 2001 From: Perry Story Date: Thu, 21 May 2026 05:16:09 -0400 Subject: [PATCH 1/7] perf: background credits and dashboard fetch for regular refreshes Move refreshCreditsIfNeeded() and refreshOpenAIDashboardIfNeeded() into background tasks for non-forced refreshes. Add Task?-based coalescing guard for credits refresh to prevent unbounded stacking. Forced refreshes (manual) still await inline. Split from #1073. --- .../Codex/UsageStore+CodexRefresh.swift | 15 + Sources/CodexBar/UsageStore.swift | 23 +- .../CodexManagedOpenAIWebRefreshTests.swift | 283 ++++++++++++++++++ 3 files changed, 317 insertions(+), 4 deletions(-) diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift index 192f7f344..05ca4fc04 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift @@ -13,6 +13,21 @@ extension UsageStore { self.makeFetchContext(provider: .codex, override: nil).fetcher } + func scheduleCreditsRefreshIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) { + if let existing = self.creditsRefreshTask { + guard existing.isCancelled else { return } + self.creditsRefreshTask = nil + } + + self.creditsRefreshTask = Task(priority: .utility) { @MainActor [weak self] in + guard let self else { return } + defer { + self.creditsRefreshTask = nil + } + await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: minimumSnapshotUpdatedAt) + } + } + func refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) async { guard self.isEnabled(.codex) else { return } var expectedGuard = self.currentCodexAccountScopedRefreshGuard() diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 608d779d9..a48b75cbe 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -179,6 +179,7 @@ final class UsageStore { @ObservationIgnored var lastCodexAccountScopedRefreshGuard: CodexAccountScopedRefreshGuard? @ObservationIgnored var lastKnownLiveSystemCodexEmail: String? @ObservationIgnored var openAIWebAccountDidChange: Bool = false + @ObservationIgnored var creditsRefreshTask: Task? @ObservationIgnored var openAIDashboardRefreshTask: Task? @ObservationIgnored var openAIDashboardRefreshTaskKey: String? @ObservationIgnored var openAIDashboardRefreshTaskToken: UUID? @@ -554,7 +555,13 @@ final class UsageStore { group.addTask { await self.refreshStatus(provider) } } } - group.addTask { await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) } + if forceTokenUsage { + group.addTask { await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) } + } + } + + if !forceTokenUsage { + self.scheduleCreditsRefreshIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) } if forceTokenUsage { @@ -586,9 +593,17 @@ final class UsageStore { ]) if shouldRefreshOpenAIWeb { let codexDashboardGuard = self.currentCodexOpenAIWebRefreshGuard() - await self.refreshOpenAIDashboardIfNeeded( - force: forceTokenUsage, - expectedGuard: codexDashboardGuard) + if forceTokenUsage { + await self.refreshOpenAIDashboardIfNeeded( + force: true, + expectedGuard: codexDashboardGuard) + } else { + Task { @MainActor [weak self] in + await self?.refreshOpenAIDashboardIfNeeded( + force: false, + expectedGuard: codexDashboardGuard) + } + } } if forceTokenUsage, self.openAIDashboardRequiresLogin { diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index 213d965f8..53f1b4f4e 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -6,12 +6,205 @@ import Testing @Suite(.serialized) @MainActor struct CodexManagedOpenAIWebRefreshTests { + @Test + func `regular refresh does not await OpenAI web scrape`() async throws { + let settings = try self.makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-regular-refresh-nonblocking") + settings.statusChecksEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingManagedOpenAIDashboardLoader() + let completion = RefreshCompletionProbe() + 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 + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let refreshTask = Task { + await store.refresh(forceTokenUsage: false) + await completion.markCompleted() + } + + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await blocker.startedCount() == 1) + #expect(await completion.isCompleted == true) + + await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + + await refreshTask.value + } + + @Test + func `regular refresh does not await Codex credits fetch`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-regular-refresh-nonblocking-credits") + settings.statusChecksEnabled = false + settings.openAIWebAccessEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingCreditsLoader() + let completion = RefreshCompletionProbe() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let refreshTask = Task { + await store.refresh(forceTokenUsage: false) + await completion.markCompleted() + } + + await blocker.waitUntilStarted(count: 1) + + #expect(await blocker.startedCount() == 1) + #expect(await completion.isCompleted == true) + + await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) + + await refreshTask.value + } + + @Test + func `rapid regular refreshes coalesce concurrent Codex credits fetches`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-credits-coalescing") + settings.statusChecksEnabled = false + settings.openAIWebAccessEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingCreditsLoader() + let firstCompletion = RefreshCompletionProbe() + let secondCompletion = RefreshCompletionProbe() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let firstRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await firstCompletion.markCompleted() + } + await blocker.waitUntilStarted(count: 1) + #expect(await firstCompletion.isCompleted == true) + + let secondRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await secondCompletion.markCompleted() + } + + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await blocker.startedCount() == 1) + #expect(await secondCompletion.isCompleted == true) + + await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) + + await firstRefreshTask.value + await secondRefreshTask.value + } + @Test func `manual cookie import bypasses same account refresh coalescing`() async throws { let settings = try self.makeSettingsStore( suite: "CodexManagedOpenAIWebRefreshTests-manual-import-bypass-coalesce") let managedHome = FileManager.default.temporaryDirectory .appendingPathComponent("codex-managed-openai-web-refresh-\(UUID().uuidString)", isDirectory: true) + try Self.writeCodexAuthFile( + homeURL: managedHome, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHome) } let managedAccount = ManagedCodexAccount( id: UUID(), email: "managed@example.com", @@ -312,6 +505,44 @@ struct CodexManagedOpenAIWebRefreshTests { settings.codexCookieSource = .auto return settings } + + private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String, accountId: String? = nil) throws { + try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true) + var tokens: [String: Any] = [ + "accessToken": "access-token", + "refreshToken": "refresh-token", + "idToken": Self.fakeJWT(email: email, plan: plan, accountId: accountId), + ] + if let accountId { + tokens["accountId"] = accountId + } + let data = try JSONSerialization.data(withJSONObject: ["tokens": tokens], options: [.sortedKeys]) + try data.write(to: homeURL.appendingPathComponent("auth.json")) + } + + private static func fakeJWT(email: String, plan: String, accountId: String? = nil) -> String { + let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data() + var authClaims: [String: Any] = [ + "chatgpt_plan_type": plan, + ] + if let accountId { + authClaims["chatgpt_account_id"] = accountId + } + let payload = (try? JSONSerialization.data(withJSONObject: [ + "email": email, + "chatgpt_plan_type": plan, + "https://api.openai.com/auth": authClaims, + ])) ?? Data() + + func base64URL(_ data: Data) -> String { + data.base64EncodedString() + .replacingOccurrences(of: "=", with: "") + .replacingOccurrences(of: "+", with: "-") + .replacingOccurrences(of: "/", with: "_") + } + + return "\(base64URL(header)).\(base64URL(payload))." + } } private enum ManagedDashboardTestError: LocalizedError { @@ -325,6 +556,14 @@ private enum ManagedDashboardTestError: LocalizedError { } } +private actor RefreshCompletionProbe { + private(set) var isCompleted = false + + func markCompleted() { + self.isCompleted = true + } +} + private actor BlockingManagedOpenAIDashboardLoader { private var continuations: [CheckedContinuation, Never>] = [] private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] @@ -369,6 +608,50 @@ private actor BlockingManagedOpenAIDashboardLoader { } } +private actor BlockingCreditsLoader { + private var continuations: [CheckedContinuation, Never>] = [] + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var started = 0 + + func awaitResult() async throws -> CreditsSnapshot { + self.started += 1 + self.resumeReadyStartWaiters() + let result = await withCheckedContinuation { continuation in + self.continuations.append(continuation) + } + 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 startedCount() -> Int { + self.started + } + + func resumeNext(with result: Result) { + guard !self.continuations.isEmpty else { return } + let continuation = self.continuations.removeFirst() + 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 + } +} + private actor OpenAIDashboardImportCallTracker { private var calls: Int = 0 private var waiters: [(count: Int, continuation: CheckedContinuation)] = [] From 4be86eb2be93ea5e4cae6595a07904b5ddcd336d Mon Sep 17 00:00:00 2001 From: Perry Story Date: Fri, 22 May 2026 19:17:16 +0100 Subject: [PATCH 2/7] fix: persist widget snapshots after background refreshes --- .../Codex/UsageStore+CodexRefresh.swift | 1 + Sources/CodexBar/UsageStore.swift | 1 + .../CodexManagedOpenAIWebRefreshTests.swift | 167 ++++++++++++++++++ 3 files changed, 169 insertions(+) diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift index 05ca4fc04..a4ba4a787 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift @@ -25,6 +25,7 @@ extension UsageStore { self.creditsRefreshTask = nil } await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: minimumSnapshotUpdatedAt) + self.persistWidgetSnapshot(reason: "credits") } } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index a48b75cbe..36e1adb77 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -602,6 +602,7 @@ final class UsageStore { await self?.refreshOpenAIDashboardIfNeeded( force: false, expectedGuard: codexDashboardGuard) + self?.persistWidgetSnapshot(reason: "dashboard") } } } diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index 53f1b4f4e..9dd148c8f 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -194,6 +194,161 @@ struct CodexManagedOpenAIWebRefreshTests { await secondRefreshTask.value } + @Test + func `background credits refresh persists updated widget snapshot after refresh returns`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-widget-background-credits") + settings.statusChecksEnabled = false + settings.openAIWebAccessEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store.snapshots[.codex] = Self.codexSnapshot(email: managedAccount.email, usedPercent: 18) + + let creditsBlocker = BlockingCreditsLoader() + let saver = BlockingWidgetSnapshotSaver() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await creditsBlocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + store._test_widgetSnapshotSaveOverride = { snapshot in + await saver.save(snapshot) + } + defer { store._test_widgetSnapshotSaveOverride = nil } + + let refreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + + await refreshTask.value + await saver.waitUntilStarted(count: 1) + + let firstSnapshots = await saver.savedSnapshots() + let firstCodexEntry = try #require(firstSnapshots.first?.entries.first { $0.provider == .codex }) + #expect(firstCodexEntry.creditsRemaining == nil) + + await saver.resumeNext() + await creditsBlocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await saver.startedCount() == 2) + let secondSnapshots = await saver.savedSnapshots() + let secondCodexEntry = try #require(secondSnapshots.last?.entries.first { $0.provider == .codex }) + #expect(secondCodexEntry.creditsRemaining == 25) + + await saver.resumeNext() + await store.widgetSnapshotPersistTask?.value + } + + @Test + func `background dashboard refresh persists updated widget snapshot after refresh returns`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-widget-background-dashboard") + settings.statusChecksEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store.snapshots[.codex] = Self.codexSnapshot(email: managedAccount.email, usedPercent: 18) + store.creditsRefreshTask = Task {} + + let dashboardBlocker = BlockingManagedOpenAIDashboardLoader() + let saver = BlockingWidgetSnapshotSaver() + 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 + try await dashboardBlocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + store._test_widgetSnapshotSaveOverride = { snapshot in + await saver.save(snapshot) + } + defer { store._test_widgetSnapshotSaveOverride = nil } + + let refreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + + await refreshTask.value + await saver.waitUntilStarted(count: 1) + + let firstSnapshots = await saver.savedSnapshots() + let firstCodexEntry = try #require(firstSnapshots.first?.entries.first { $0.provider == .codex }) + #expect(firstCodexEntry.codeReviewRemainingPercent == nil) + + await saver.resumeNext() + await dashboardBlocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await saver.startedCount() == 2) + let secondSnapshots = await saver.savedSnapshots() + let secondCodexEntry = try #require(secondSnapshots.last?.entries.first { $0.provider == .codex }) + #expect(secondCodexEntry.codeReviewRemainingPercent == 95) + + await saver.resumeNext() + await store.widgetSnapshotPersistTask?.value + } + @Test func `manual cookie import bypasses same account refresh coalescing`() async throws { let settings = try self.makeSettingsStore( @@ -543,6 +698,18 @@ struct CodexManagedOpenAIWebRefreshTests { return "\(base64URL(header)).\(base64URL(payload))." } + + private static func codexSnapshot(email: String, usedPercent: Double) -> UsageSnapshot { + UsageSnapshot( + primary: RateWindow(usedPercent: usedPercent, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: email, + accountOrganization: nil, + loginMethod: "Pro")) + } } private enum ManagedDashboardTestError: LocalizedError { From 8a1b17b9305577367428e3f884417195352d4787 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 22 May 2026 19:28:46 +0100 Subject: [PATCH 3/7] fix: scope background codex refresh coalescing --- CHANGELOG.md | 1 + .../Codex/UsageStore+CodexRefresh.swift | 41 ++++- Sources/CodexBar/UsageStore+OpenAIWeb.swift | 31 ++++ Sources/CodexBar/UsageStore.swift | 10 +- .../CodexManagedOpenAIWebRefreshTests.swift | 166 +++++++++++++++++- 5 files changed, 238 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f5e36bc8..d96ee94aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ ### Fixed - Menu bar: show extra-usage spend as currency text for Claude and Cursor when that metric is selected (#1107). Thanks @Yuxin-Qiao! +- Codex: run regular credits and OpenAI dashboard refreshes in the background while coalescing overlapping refresh work (#1078). Thanks @ptstory! ## 0.29.0 — 2026-05-22 diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift index a4ba4a787..a85b0ca2b 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift @@ -14,21 +14,54 @@ extension UsageStore { } func scheduleCreditsRefreshIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) { - if let existing = self.creditsRefreshTask { - guard existing.isCancelled else { return } - self.creditsRefreshTask = nil + let refreshKey = self.codexCreditsRefreshKey( + expectedGuard: self.currentCodexAccountScopedRefreshGuard()) + if let existing = self.creditsRefreshTask, + !existing.isCancelled, + self.creditsRefreshTaskKey == refreshKey + { + return } + self.creditsRefreshTask?.cancel() + self.creditsRefreshTaskKey = refreshKey self.creditsRefreshTask = Task(priority: .utility) { @MainActor [weak self] in guard let self else { return } defer { - self.creditsRefreshTask = nil + if self.creditsRefreshTaskKey == refreshKey { + self.creditsRefreshTask = nil + self.creditsRefreshTaskKey = nil + } } await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: minimumSnapshotUpdatedAt) self.persistWidgetSnapshot(reason: "credits") } } + func codexCreditsRefreshKey(expectedGuard: CodexAccountScopedRefreshGuard) -> String { + let sourceKey = switch expectedGuard.source { + case .liveSystem: + "live" + case let .managedAccount(id): + "managed:\(id.uuidString)" + } + + let identityKey = switch expectedGuard.identity { + case let .providerAccount(id): + "provider:\(id)" + case let .emailOnly(normalizedEmail): + "email:\(normalizedEmail)" + case .unresolved: + "unresolved" + } + + return [ + sourceKey, + identityKey, + expectedGuard.accountKey ?? "account:nil", + ].joined(separator: "|") + } + func refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) async { guard self.isEnabled(.codex) else { return } var expectedGuard = self.currentCodexAccountScopedRefreshGuard() diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index c02e06c52..3e4471ff3 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -409,6 +409,37 @@ extension UsageStore { } } + func scheduleOpenAIDashboardRefreshIfNeeded(expectedGuard: CodexAccountScopedRefreshGuard? = nil) { + self.syncOpenAIWebState() + let allowCurrentSnapshotFallback = expectedGuard?.source == .liveSystem && expectedGuard? + .identity == .unresolved + let targetEmail = self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: allowCurrentSnapshotFallback, + allowLastKnownLiveFallback: expectedGuard?.identity != .unresolved) + let refreshKey = self.openAIDashboardRefreshKey(targetEmail: targetEmail, expectedGuard: expectedGuard) + if let task = self.openAIDashboardBackgroundRefreshTask, + !task.isCancelled, + self.openAIDashboardBackgroundRefreshTaskKey == refreshKey + { + return + } + + self.openAIDashboardBackgroundRefreshTask?.cancel() + self.openAIDashboardBackgroundRefreshTaskKey = refreshKey + self.openAIDashboardBackgroundRefreshTask = Task(priority: .utility) { @MainActor [weak self] in + guard let self else { return } + defer { + if self.openAIDashboardBackgroundRefreshTaskKey == refreshKey { + self.openAIDashboardBackgroundRefreshTask = nil + self.openAIDashboardBackgroundRefreshTaskKey = nil + } + } + + await self.refreshOpenAIDashboardIfNeeded(force: false, expectedGuard: expectedGuard) + self.persistWidgetSnapshot(reason: "dashboard") + } + } + private func performOpenAIDashboardRefreshIfNeeded(_ context: OpenAIDashboardRefreshContext) async { self.openAIDashboardCookieImportStatus = nil var latestCookieImportStatus: String? diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 36e1adb77..508ee0dbe 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -180,6 +180,9 @@ final class UsageStore { @ObservationIgnored var lastKnownLiveSystemCodexEmail: String? @ObservationIgnored var openAIWebAccountDidChange: Bool = false @ObservationIgnored var creditsRefreshTask: Task? + @ObservationIgnored var creditsRefreshTaskKey: String? + @ObservationIgnored var openAIDashboardBackgroundRefreshTask: Task? + @ObservationIgnored var openAIDashboardBackgroundRefreshTaskKey: String? @ObservationIgnored var openAIDashboardRefreshTask: Task? @ObservationIgnored var openAIDashboardRefreshTaskKey: String? @ObservationIgnored var openAIDashboardRefreshTaskToken: UUID? @@ -598,12 +601,7 @@ final class UsageStore { force: true, expectedGuard: codexDashboardGuard) } else { - Task { @MainActor [weak self] in - await self?.refreshOpenAIDashboardIfNeeded( - force: false, - expectedGuard: codexDashboardGuard) - self?.persistWidgetSnapshot(reason: "dashboard") - } + self.scheduleOpenAIDashboardRefreshIfNeeded(expectedGuard: codexDashboardGuard) } } diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index 9dd148c8f..f66b1b5f7 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -8,7 +8,8 @@ import Testing struct CodexManagedOpenAIWebRefreshTests { @Test func `regular refresh does not await OpenAI web scrape`() async throws { - let settings = try self.makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-regular-refresh-nonblocking") + let settings = try self + .makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-regular-refresh-nonblocking") settings.statusChecksEnabled = false if let codexMeta = ProviderRegistry.shared.metadata[.codex] { settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) @@ -194,6 +195,167 @@ struct CodexManagedOpenAIWebRefreshTests { await secondRefreshTask.value } + @Test + func `regular credits refresh reschedules when Codex account changes`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-credits-account-switch") + settings.statusChecksEnabled = false + settings.openAIWebAccessEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let alphaHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + let betaHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: alphaHomeURL, + email: "alpha@example.com", + plan: "Pro") + try? Self.writeCodexAuthFile( + homeURL: betaHomeURL, + email: "beta@example.com", + plan: "Pro") + defer { + try? FileManager.default.removeItem(at: alphaHomeURL) + try? FileManager.default.removeItem(at: betaHomeURL) + } + let alphaAccount = ManagedCodexAccount( + id: UUID(), + email: "alpha@example.com", + managedHomePath: alphaHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let betaAccount = ManagedCodexAccount( + id: UUID(), + email: "beta@example.com", + managedHomePath: betaHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = alphaAccount + settings.codexActiveSource = .managedAccount(id: alphaAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingCreditsLoader() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let alphaRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + await blocker.waitUntilStarted(count: 1) + + settings._test_activeManagedCodexAccount = betaAccount + settings.codexActiveSource = .managedAccount(id: betaAccount.id) + let betaRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + await blocker.waitUntilStarted(count: 2) + + 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 store.creditsRefreshTask?.value + + #expect(await blocker.startedCount() == 2) + #expect(store.lastCreditsSnapshotAccountKey == "beta@example.com") + #expect(store.credits?.remaining == 25) + } + + @Test + func `rapid regular refreshes coalesce concurrent OpenAI dashboard fetches`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-dashboard-coalescing") + settings.statusChecksEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingManagedOpenAIDashboardLoader() + let firstCompletion = RefreshCompletionProbe() + let secondCompletion = RefreshCompletionProbe() + 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 + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let firstRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await firstCompletion.markCompleted() + } + await blocker.waitUntilStarted(count: 1) + #expect(await firstCompletion.isCompleted == true) + + let secondRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await secondCompletion.markCompleted() + } + + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await blocker.startedCount() == 1) + #expect(await secondCompletion.isCompleted == true) + + let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask) + await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + + await firstRefreshTask.value + await secondRefreshTask.value + await backgroundTask.value + + #expect(store.openAIDashboard?.creditsRemaining == 25) + } + @Test func `background credits refresh persists updated widget snapshot after refresh returns`() async throws { let settings = try self.makeSettingsStore( @@ -298,6 +460,8 @@ struct CodexManagedOpenAIWebRefreshTests { startupBehavior: .testing) store.snapshots[.codex] = Self.codexSnapshot(email: managedAccount.email, usedPercent: 18) store.creditsRefreshTask = Task {} + store.creditsRefreshTaskKey = store.codexCreditsRefreshKey( + expectedGuard: store.currentCodexAccountScopedRefreshGuard()) let dashboardBlocker = BlockingManagedOpenAIDashboardLoader() let saver = BlockingWidgetSnapshotSaver() From bc5e5c007c2ceb366978ba49813fab962556e6eb Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 22 May 2026 19:32:40 +0100 Subject: [PATCH 4/7] fix: cancel scheduled codex credits before force refresh --- .../Codex/UsageStore+CodexRefresh.swift | 14 +++++ Sources/CodexBar/UsageStore.swift | 4 +- .../CodexManagedOpenAIWebRefreshTests.swift | 60 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift index a85b0ca2b..dc0608017 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift @@ -34,10 +34,22 @@ extension UsageStore { } } await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: minimumSnapshotUpdatedAt) + guard !Task.isCancelled else { return } self.persistWidgetSnapshot(reason: "credits") } } + func cancelScheduledCreditsRefresh() { + self.creditsRefreshTask?.cancel() + self.creditsRefreshTask = nil + self.creditsRefreshTaskKey = nil + } + + func refreshCreditsNow(minimumSnapshotUpdatedAt: Date? = nil) async { + self.cancelScheduledCreditsRefresh() + await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: minimumSnapshotUpdatedAt) + } + func codexCreditsRefreshKey(expectedGuard: CodexAccountScopedRefreshGuard) -> String { let sourceKey = switch expectedGuard.source { case .liveSystem: @@ -79,6 +91,7 @@ extension UsageStore { } do { let credits = try await self.loadLatestCodexCredits() + guard !Task.isCancelled else { return } guard self.shouldApplyCodexScopedNonUsageResult(expectedGuard: expectedGuard) else { return } await MainActor.run { self.credits = credits @@ -107,6 +120,7 @@ extension UsageStore { snapshot: codexSnapshot, now: codexSnapshot.updatedAt) } catch { + guard !Task.isCancelled else { return } let message = error.localizedDescription if message.localizedCaseInsensitiveContains("data not available yet") { guard self.shouldApplyCodexScopedNonUsageResult(expectedGuard: expectedGuard) else { return } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 508ee0dbe..0bc6f514f 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -559,7 +559,7 @@ final class UsageStore { } } if forceTokenUsage { - group.addTask { await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) } + group.addTask { await self.refreshCreditsNow(minimumSnapshotUpdatedAt: refreshStartedAt) } } } @@ -607,7 +607,7 @@ final class UsageStore { if forceTokenUsage, self.openAIDashboardRequiresLogin { await self.refreshProvider(.codex) - await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) + await self.refreshCreditsNow(minimumSnapshotUpdatedAt: refreshStartedAt) } self.persistWidgetSnapshot(reason: "refresh") diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index f66b1b5f7..e81ac1ed4 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -275,6 +275,66 @@ struct CodexManagedOpenAIWebRefreshTests { #expect(store.credits?.remaining == 25) } + @Test + func `force refresh cancels stale background Codex credits fetch`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexManagedOpenAIWebRefreshTests-credits-force-cancels-background") + settings.statusChecksEnabled = false + settings.openAIWebAccessEnabled = false + if let codexMeta = ProviderRegistry.shared.metadata[.codex] { + settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) + } + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try? Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: "managed@example.com", + plan: "Pro") + defer { try? FileManager.default.removeItem(at: managedHomeURL) } + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingCreditsLoader() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let regularRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + await blocker.waitUntilStarted(count: 1) + + let forceRefreshTask = Task { + await store.refresh(forceTokenUsage: true) + } + await blocker.waitUntilStarted(count: 2) + + 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 + + #expect(await blocker.startedCount() == 2) + #expect(store.credits?.remaining == 25) + } + @Test func `rapid regular refreshes coalesce concurrent OpenAI dashboard fetches`() async throws { let settings = try self.makeSettingsStore( From 7500f8cd791bb8da561b7ad651729b7e1872397e Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 22 May 2026 19:35:55 +0100 Subject: [PATCH 5/7] test: split codex background refresh coalescing coverage --- ...odexBackgroundRefreshCoalescingTests.swift | 286 +++++++++++++++++ .../CodexManagedOpenAIWebRefreshTests.swift | 292 +----------------- 2 files changed, 289 insertions(+), 289 deletions(-) create mode 100644 Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift diff --git a/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift b/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift new file mode 100644 index 000000000..fe1f12702 --- /dev/null +++ b/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift @@ -0,0 +1,286 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +@Suite(.serialized) +@MainActor +struct CodexBackgroundRefreshCoalescingTests { + @Test + func `rapid regular refreshes coalesce concurrent Codex credits fetches`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-credits-coalescing") + settings.statusChecksEnabled = 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() + let firstCompletion = RefreshCompletionProbe() + let secondCompletion = RefreshCompletionProbe() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let firstRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await firstCompletion.markCompleted() + } + await blocker.waitUntilStarted(count: 1) + #expect(await firstCompletion.isCompleted == true) + + let secondRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await secondCompletion.markCompleted() + } + + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await blocker.startedCount() == 1) + #expect(await secondCompletion.isCompleted == true) + + await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) + + await firstRefreshTask.value + await secondRefreshTask.value + } + + @Test + func `regular credits refresh reschedules when Codex account changes`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-credits-account-switch") + settings.statusChecksEnabled = false + settings.openAIWebAccessEnabled = false + let alphaAccount = try Self.makeManagedAccount(email: "alpha@example.com") + let betaAccount = try Self.makeManagedAccount(email: "beta@example.com") + defer { + try? FileManager.default.removeItem(atPath: alphaAccount.managedHomePath) + try? FileManager.default.removeItem(atPath: betaAccount.managedHomePath) + } + settings._test_activeManagedCodexAccount = alphaAccount + settings.codexActiveSource = .managedAccount(id: alphaAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = self.makeStore(settings: settings) + let blocker = BlockingCreditsLoader() + store._test_providerRefreshOverride = { _ in } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let alphaRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + await blocker.waitUntilStarted(count: 1) + + settings._test_activeManagedCodexAccount = betaAccount + settings.codexActiveSource = .managedAccount(id: betaAccount.id) + let betaRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + await blocker.waitUntilStarted(count: 2) + + 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 store.creditsRefreshTask?.value + + #expect(await blocker.startedCount() == 2) + #expect(store.lastCreditsSnapshotAccountKey == "beta@example.com") + #expect(store.credits?.remaining == 25) + } + + @Test + func `force refresh cancels stale background Codex credits fetch`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-credits-force-cancels-background") + settings.statusChecksEnabled = 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 } + defer { store._test_providerRefreshOverride = nil } + store._test_codexCreditsLoaderOverride = { + try await blocker.awaitResult() + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let regularRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + } + await blocker.waitUntilStarted(count: 1) + + let forceRefreshTask = Task { + await store.refresh(forceTokenUsage: true) + } + await blocker.waitUntilStarted(count: 2) + + 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 + + #expect(await blocker.startedCount() == 2) + #expect(store.credits?.remaining == 25) + } + + @Test + func `rapid regular refreshes coalesce concurrent OpenAI dashboard fetches`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-dashboard-coalescing") + settings.statusChecksEnabled = 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 = BlockingManagedOpenAIDashboardLoader() + let firstCompletion = RefreshCompletionProbe() + let secondCompletion = RefreshCompletionProbe() + 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 + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let firstRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await firstCompletion.markCompleted() + } + await blocker.waitUntilStarted(count: 1) + #expect(await firstCompletion.isCompleted == true) + + let secondRefreshTask = Task { + await store.refresh(forceTokenUsage: false) + await secondCompletion.markCompleted() + } + + try? await Task.sleep(for: .milliseconds(200)) + + #expect(await blocker.startedCount() == 1) + #expect(await secondCompletion.isCompleted == true) + + let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask) + await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + + await firstRefreshTask.value + await secondRefreshTask.value + await backgroundTask.value + + #expect(store.openAIDashboard?.creditsRemaining == 25) + } + + 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 codexMetadata = try #require(ProviderDescriptorRegistry.metadata[.codex]) + settings.setProviderEnabled(provider: .codex, metadata: codexMetadata, enabled: true) + settings.providerDetectionCompleted = true + settings.openAIWebAccessEnabled = true + settings.codexCookieSource = .auto + return settings + } + + private func makeStore(settings: SettingsStore) -> UsageStore { + UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + } + + private static func installManagedAccount( + email: String, + settings: SettingsStore) throws -> ManagedCodexAccount + { + let account = try Self.makeManagedAccount(email: email) + settings._test_activeManagedCodexAccount = account + settings.codexActiveSource = .managedAccount(id: account.id) + return account + } + + private static func makeManagedAccount(email: String) throws -> ManagedCodexAccount { + let managedHomeURL = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try Self.writeCodexAuthFile( + homeURL: managedHomeURL, + email: email, + plan: "Pro") + return ManagedCodexAccount( + id: UUID(), + email: email, + managedHomePath: managedHomeURL.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + } + + private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String) throws { + try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true) + let tokens: [String: Any] = [ + "accessToken": "access-token", + "refreshToken": "refresh-token", + "idToken": Self.fakeJWT(email: email, plan: plan), + ] + let data = try JSONSerialization.data(withJSONObject: ["tokens": tokens], options: [.sortedKeys]) + try data.write(to: homeURL.appendingPathComponent("auth.json")) + } + + private static func fakeJWT(email: String, plan: String) -> String { + let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data() + let payload = (try? JSONSerialization.data(withJSONObject: [ + "email": email, + "chatgpt_plan_type": plan, + "https://api.openai.com/auth": [ + "chatgpt_plan_type": plan, + ], + ])) ?? Data() + + func base64URL(_ data: Data) -> String { + data.base64EncodedString() + .replacingOccurrences(of: "=", with: "") + .replacingOccurrences(of: "+", with: "-") + .replacingOccurrences(of: "/", with: "_") + } + + return "\(base64URL(header)).\(base64URL(payload))." + } +} diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index e81ac1ed4..f44eb9bcd 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -130,292 +130,6 @@ struct CodexManagedOpenAIWebRefreshTests { await refreshTask.value } - @Test - func `rapid regular refreshes coalesce concurrent Codex credits fetches`() async throws { - let settings = try self.makeSettingsStore( - suite: "CodexManagedOpenAIWebRefreshTests-credits-coalescing") - settings.statusChecksEnabled = false - settings.openAIWebAccessEnabled = false - if let codexMeta = ProviderRegistry.shared.metadata[.codex] { - settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) - } - let managedHomeURL = FileManager.default.temporaryDirectory - .appendingPathComponent(UUID().uuidString, isDirectory: true) - try? Self.writeCodexAuthFile( - homeURL: managedHomeURL, - email: "managed@example.com", - plan: "Pro") - defer { try? FileManager.default.removeItem(at: managedHomeURL) } - let managedAccount = ManagedCodexAccount( - id: UUID(), - email: "managed@example.com", - managedHomePath: managedHomeURL.path, - createdAt: 1, - updatedAt: 1, - lastAuthenticatedAt: 1) - settings._test_activeManagedCodexAccount = managedAccount - settings.codexActiveSource = .managedAccount(id: managedAccount.id) - defer { settings._test_activeManagedCodexAccount = nil } - - let store = UsageStore( - fetcher: UsageFetcher(environment: [:]), - browserDetection: BrowserDetection(cacheTTL: 0), - settings: settings, - startupBehavior: .testing) - let blocker = BlockingCreditsLoader() - let firstCompletion = RefreshCompletionProbe() - let secondCompletion = RefreshCompletionProbe() - store._test_providerRefreshOverride = { _ in } - defer { store._test_providerRefreshOverride = nil } - store._test_codexCreditsLoaderOverride = { - try await blocker.awaitResult() - } - defer { store._test_codexCreditsLoaderOverride = nil } - - let firstRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - await firstCompletion.markCompleted() - } - await blocker.waitUntilStarted(count: 1) - #expect(await firstCompletion.isCompleted == true) - - let secondRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - await secondCompletion.markCompleted() - } - - try? await Task.sleep(for: .milliseconds(200)) - - #expect(await blocker.startedCount() == 1) - #expect(await secondCompletion.isCompleted == true) - - await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) - - await firstRefreshTask.value - await secondRefreshTask.value - } - - @Test - func `regular credits refresh reschedules when Codex account changes`() async throws { - let settings = try self.makeSettingsStore( - suite: "CodexManagedOpenAIWebRefreshTests-credits-account-switch") - settings.statusChecksEnabled = false - settings.openAIWebAccessEnabled = false - if let codexMeta = ProviderRegistry.shared.metadata[.codex] { - settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) - } - let alphaHomeURL = FileManager.default.temporaryDirectory - .appendingPathComponent(UUID().uuidString, isDirectory: true) - let betaHomeURL = FileManager.default.temporaryDirectory - .appendingPathComponent(UUID().uuidString, isDirectory: true) - try? Self.writeCodexAuthFile( - homeURL: alphaHomeURL, - email: "alpha@example.com", - plan: "Pro") - try? Self.writeCodexAuthFile( - homeURL: betaHomeURL, - email: "beta@example.com", - plan: "Pro") - defer { - try? FileManager.default.removeItem(at: alphaHomeURL) - try? FileManager.default.removeItem(at: betaHomeURL) - } - let alphaAccount = ManagedCodexAccount( - id: UUID(), - email: "alpha@example.com", - managedHomePath: alphaHomeURL.path, - createdAt: 1, - updatedAt: 1, - lastAuthenticatedAt: 1) - let betaAccount = ManagedCodexAccount( - id: UUID(), - email: "beta@example.com", - managedHomePath: betaHomeURL.path, - createdAt: 1, - updatedAt: 1, - lastAuthenticatedAt: 1) - settings._test_activeManagedCodexAccount = alphaAccount - settings.codexActiveSource = .managedAccount(id: alphaAccount.id) - defer { settings._test_activeManagedCodexAccount = nil } - - let store = UsageStore( - fetcher: UsageFetcher(environment: [:]), - browserDetection: BrowserDetection(cacheTTL: 0), - settings: settings, - startupBehavior: .testing) - let blocker = BlockingCreditsLoader() - store._test_providerRefreshOverride = { _ in } - defer { store._test_providerRefreshOverride = nil } - store._test_codexCreditsLoaderOverride = { - try await blocker.awaitResult() - } - defer { store._test_codexCreditsLoaderOverride = nil } - - let alphaRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - } - await blocker.waitUntilStarted(count: 1) - - settings._test_activeManagedCodexAccount = betaAccount - settings.codexActiveSource = .managedAccount(id: betaAccount.id) - let betaRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - } - await blocker.waitUntilStarted(count: 2) - - 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 store.creditsRefreshTask?.value - - #expect(await blocker.startedCount() == 2) - #expect(store.lastCreditsSnapshotAccountKey == "beta@example.com") - #expect(store.credits?.remaining == 25) - } - - @Test - func `force refresh cancels stale background Codex credits fetch`() async throws { - let settings = try self.makeSettingsStore( - suite: "CodexManagedOpenAIWebRefreshTests-credits-force-cancels-background") - settings.statusChecksEnabled = false - settings.openAIWebAccessEnabled = false - if let codexMeta = ProviderRegistry.shared.metadata[.codex] { - settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) - } - let managedHomeURL = FileManager.default.temporaryDirectory - .appendingPathComponent(UUID().uuidString, isDirectory: true) - try? Self.writeCodexAuthFile( - homeURL: managedHomeURL, - email: "managed@example.com", - plan: "Pro") - defer { try? FileManager.default.removeItem(at: managedHomeURL) } - let managedAccount = ManagedCodexAccount( - id: UUID(), - email: "managed@example.com", - managedHomePath: managedHomeURL.path, - createdAt: 1, - updatedAt: 1, - lastAuthenticatedAt: 1) - settings._test_activeManagedCodexAccount = managedAccount - settings.codexActiveSource = .managedAccount(id: managedAccount.id) - defer { settings._test_activeManagedCodexAccount = nil } - - let store = UsageStore( - fetcher: UsageFetcher(environment: [:]), - browserDetection: BrowserDetection(cacheTTL: 0), - settings: settings, - startupBehavior: .testing) - let blocker = BlockingCreditsLoader() - store._test_providerRefreshOverride = { _ in } - defer { store._test_providerRefreshOverride = nil } - store._test_codexCreditsLoaderOverride = { - try await blocker.awaitResult() - } - defer { store._test_codexCreditsLoaderOverride = nil } - - let regularRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - } - await blocker.waitUntilStarted(count: 1) - - let forceRefreshTask = Task { - await store.refresh(forceTokenUsage: true) - } - await blocker.waitUntilStarted(count: 2) - - 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 - - #expect(await blocker.startedCount() == 2) - #expect(store.credits?.remaining == 25) - } - - @Test - func `rapid regular refreshes coalesce concurrent OpenAI dashboard fetches`() async throws { - let settings = try self.makeSettingsStore( - suite: "CodexManagedOpenAIWebRefreshTests-dashboard-coalescing") - settings.statusChecksEnabled = false - if let codexMeta = ProviderRegistry.shared.metadata[.codex] { - settings.setProviderEnabled(provider: .codex, metadata: codexMeta, enabled: true) - } - let managedHomeURL = FileManager.default.temporaryDirectory - .appendingPathComponent(UUID().uuidString, isDirectory: true) - try? Self.writeCodexAuthFile( - homeURL: managedHomeURL, - email: "managed@example.com", - plan: "Pro") - defer { try? FileManager.default.removeItem(at: managedHomeURL) } - let managedAccount = ManagedCodexAccount( - id: UUID(), - email: "managed@example.com", - managedHomePath: managedHomeURL.path, - createdAt: 1, - updatedAt: 1, - lastAuthenticatedAt: 1) - settings._test_activeManagedCodexAccount = managedAccount - settings.codexActiveSource = .managedAccount(id: managedAccount.id) - defer { settings._test_activeManagedCodexAccount = nil } - - let store = UsageStore( - fetcher: UsageFetcher(environment: [:]), - browserDetection: BrowserDetection(cacheTTL: 0), - settings: settings, - startupBehavior: .testing) - let blocker = BlockingManagedOpenAIDashboardLoader() - let firstCompletion = RefreshCompletionProbe() - let secondCompletion = RefreshCompletionProbe() - 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 - try await blocker.awaitResult() - } - defer { store._test_openAIDashboardLoaderOverride = nil } - - let firstRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - await firstCompletion.markCompleted() - } - await blocker.waitUntilStarted(count: 1) - #expect(await firstCompletion.isCompleted == true) - - let secondRefreshTask = Task { - await store.refresh(forceTokenUsage: false) - await secondCompletion.markCompleted() - } - - try? await Task.sleep(for: .milliseconds(200)) - - #expect(await blocker.startedCount() == 1) - #expect(await secondCompletion.isCompleted == true) - - let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask) - await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( - signedInEmail: managedAccount.email, - codeReviewRemainingPercent: 95, - creditEvents: [], - dailyBreakdown: [], - usageBreakdown: [], - creditsPurchaseURL: nil, - creditsRemaining: 25, - accountPlan: "Pro", - updatedAt: Date()))) - - await firstRefreshTask.value - await secondRefreshTask.value - await backgroundTask.value - - #expect(store.openAIDashboard?.creditsRemaining == 25) - } - @Test func `background credits refresh persists updated widget snapshot after refresh returns`() async throws { let settings = try self.makeSettingsStore( @@ -947,7 +661,7 @@ private enum ManagedDashboardTestError: LocalizedError { } } -private actor RefreshCompletionProbe { +actor RefreshCompletionProbe { private(set) var isCompleted = false func markCompleted() { @@ -955,7 +669,7 @@ private actor RefreshCompletionProbe { } } -private actor BlockingManagedOpenAIDashboardLoader { +actor BlockingManagedOpenAIDashboardLoader { private var continuations: [CheckedContinuation, Never>] = [] private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] private var started: Int = 0 @@ -999,7 +713,7 @@ private actor BlockingManagedOpenAIDashboardLoader { } } -private actor BlockingCreditsLoader { +actor BlockingCreditsLoader { private var continuations: [CheckedContinuation, Never>] = [] private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] private var started = 0 From a9a8d1ddbc7d93c7aae000ac4a0dd9e1f9ffb4c4 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 22 May 2026 20:06:38 +0100 Subject: [PATCH 6/7] fix: stop cancelled dashboard refresh side effects --- Sources/CodexBar/UsageStore+OpenAIWeb.swift | 135 +++++++++++++----- ...odexBackgroundRefreshCoalescingTests.swift | 77 ++++++++++ .../CodexManagedOpenAIWebRefreshTests.swift | 52 ++----- 3 files changed, 190 insertions(+), 74 deletions(-) diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index 3e4471ff3..3b75b446c 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -28,6 +28,15 @@ extension UsageStore { let allowCodexUsageBackfill: Bool } + private struct OpenAIDashboardCookieImportRequest { + let normalizedTarget: String? + let allowAnyAccount: Bool + let cookieSource: ProviderCookieSource + let cacheScope: CookieHeaderCache.Scope? + let preferCachedCookieHeader: Bool? + let force: Bool + } + private static let openAIWebRefreshMultiplier: TimeInterval = 5 private static let openAIWebPrimaryFetchTimeout: TimeInterval = 25 private static let openAIWebRetryFetchTimeout: TimeInterval = 8 @@ -424,6 +433,11 @@ extension UsageStore { return } + if self.openAIDashboardBackgroundRefreshTaskKey != nil, + self.openAIDashboardBackgroundRefreshTaskKey != refreshKey + { + self.invalidateOpenAIDashboardRefreshTask() + } self.openAIDashboardBackgroundRefreshTask?.cancel() self.openAIDashboardBackgroundRefreshTaskKey = refreshKey self.openAIDashboardBackgroundRefreshTask = Task(priority: .utility) { @MainActor [weak self] in @@ -436,11 +450,13 @@ extension UsageStore { } await self.refreshOpenAIDashboardIfNeeded(force: false, expectedGuard: expectedGuard) + guard !Task.isCancelled else { return } self.persistWidgetSnapshot(reason: "dashboard") } } private func performOpenAIDashboardRefreshIfNeeded(_ context: OpenAIDashboardRefreshContext) async { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } self.openAIDashboardCookieImportStatus = nil var latestCookieImportStatus: String? if self.openAIWebDebugLines.isEmpty { @@ -471,6 +487,7 @@ extension UsageStore { let imported = await self.importOpenAIDashboardCookiesIfNeeded( targetEmail: targetEmail, force: true) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } didImportCookiesForRefresh = true latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() if await self.abortOpenAIDashboardRetryAfterImportFailure( @@ -493,12 +510,14 @@ extension UsageStore { accountEmail: effectiveEmail, logger: log, timeout: Self.openAIWebDashboardFetchTimeout(didImportCookies: didImportCookiesForRefresh)) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } if self.dashboardEmailMismatch(expected: normalized, actual: dash.signedInEmail) { if let imported = await self.importOpenAIDashboardCookiesIfNeeded( targetEmail: context.targetEmail, force: true) { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } effectiveEmail = imported } latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() @@ -506,6 +525,7 @@ extension UsageStore { accountEmail: effectiveEmail, logger: log, timeout: Self.openAIWebRetryDashboardFetchTimeout(afterCookieImport: true)) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } } await self.applyOpenAIDashboard( @@ -515,17 +535,20 @@ extension UsageStore { refreshTaskToken: context.refreshTaskToken, allowCodexUsageBackfill: context.allowCodexUsageBackfill) } catch let OpenAIDashboardFetcher.FetchError.noDashboardData(body) { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } await self.retryOpenAIDashboardAfterNoData( body: body, context: context, latestCookieImportStatus: &latestCookieImportStatus, logger: log) } catch OpenAIDashboardFetcher.FetchError.loginRequired { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } await self.retryOpenAIDashboardAfterLoginRequired( context: context, latestCookieImportStatus: &latestCookieImportStatus, logger: log) } catch { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } if Self.isOpenAIDashboardTimeout(error) { await self.retryOpenAIDashboardAfterTimeout( context: context, @@ -558,6 +581,7 @@ extension UsageStore { targetEmail: targetEmail, force: true, preferCachedCookieHeader: true) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() if await self.abortOpenAIDashboardRetryAfterImportFailure( importedEmail: imported, @@ -576,6 +600,7 @@ extension UsageStore { accountEmail: effectiveEmail, logger: logger, timeout: Self.openAIWebRetryDashboardFetchTimeout(afterCookieImport: true)) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } await self.applyOpenAIDashboard( dash, targetEmail: effectiveEmail, @@ -583,6 +608,7 @@ extension UsageStore { refreshTaskToken: context.refreshTaskToken, allowCodexUsageBackfill: context.allowCodexUsageBackfill) } catch { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } let message = self.preferredOpenAIDashboardFailureMessage( error: error, targetEmail: targetEmail, @@ -606,6 +632,7 @@ extension UsageStore { allowLastKnownLiveFallback: context.expectedGuard?.identity != .unresolved) var effectiveEmail = targetEmail let imported = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() if await self.abortOpenAIDashboardRetryAfterImportFailure( importedEmail: imported, @@ -624,6 +651,7 @@ extension UsageStore { accountEmail: effectiveEmail, logger: logger, timeout: Self.openAIWebRetryDashboardFetchTimeout(afterCookieImport: true)) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } await self.applyOpenAIDashboard( dash, targetEmail: effectiveEmail, @@ -631,6 +659,7 @@ extension UsageStore { refreshTaskToken: context.refreshTaskToken, allowCodexUsageBackfill: context.allowCodexUsageBackfill) } catch let OpenAIDashboardFetcher.FetchError.noDashboardData(retryBody) { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } let finalBody = retryBody.isEmpty ? body : retryBody let message = self.openAIDashboardFriendlyError( body: finalBody, @@ -643,6 +672,7 @@ extension UsageStore { refreshTaskToken: context.refreshTaskToken, routingTargetEmail: targetEmail) } catch { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } let message = self.preferredOpenAIDashboardFailureMessage( error: error, targetEmail: targetEmail, @@ -665,6 +695,7 @@ extension UsageStore { allowLastKnownLiveFallback: context.expectedGuard?.identity != .unresolved) var effectiveEmail = targetEmail let imported = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() if await self.abortOpenAIDashboardRetryAfterImportFailure( importedEmail: imported, @@ -683,6 +714,7 @@ extension UsageStore { accountEmail: effectiveEmail, logger: logger, timeout: Self.openAIWebRetryDashboardFetchTimeout(afterCookieImport: true)) + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } await self.applyOpenAIDashboard( dash, targetEmail: effectiveEmail, @@ -690,11 +722,13 @@ extension UsageStore { refreshTaskToken: context.refreshTaskToken, allowCodexUsageBackfill: context.allowCodexUsageBackfill) } catch OpenAIDashboardFetcher.FetchError.loginRequired { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } await self.applyOpenAIDashboardLoginRequiredFailure( expectedGuard: context.expectedGuard, refreshTaskToken: context.refreshTaskToken, routingTargetEmail: targetEmail) } catch { + guard self.shouldContinueOpenAIDashboardRefresh(token: context.refreshTaskToken) else { return } let message = self.preferredOpenAIDashboardFailureMessage( error: error, targetEmail: targetEmail, @@ -875,6 +909,10 @@ extension UsageStore { return self.openAIDashboardRefreshTaskToken == token } + private func shouldContinueOpenAIDashboardRefresh(token: UUID?) -> Bool { + !Task.isCancelled && self.shouldApplyOpenAIDashboardRefreshTask(token: token) + } + func invalidateOpenAIDashboardRefreshTask() { self.openAIDashboardRefreshTask?.cancel() self.openAIDashboardRefreshTask = nil @@ -951,14 +989,62 @@ extension UsageStore { return false } + private func openAIDashboardCookieImportResult( + request: OpenAIDashboardCookieImportRequest, + logger: @escaping (String) -> Void) async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult + { + if let override = self._test_openAIDashboardCookieImportOverride { + return try await override( + request.normalizedTarget, + request.allowAnyAccount, + request.cookieSource, + request.cacheScope, + logger) + } + + let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: self.browserDetection) + switch request.cookieSource { + case .manual: + self.settings.ensureCodexCookieLoaded() + // Manual OpenAI cookies still come from one provider-level setting. Auto-imported cookies are + // isolated per managed account, but a manual header is an explicit override owned by settings, + // so switching managed accounts does not currently swap it underneath the user. + let manualHeader = self.settings.codexCookieHeader + guard CookieHeaderNormalizer.normalize(manualHeader) != nil else { + throw OpenAIDashboardBrowserCookieImporter.ImportError.manualCookieHeaderInvalid + } + return try await importer.importManualCookies( + cookieHeader: manualHeader, + intoAccountEmail: request.normalizedTarget, + allowAnyAccount: request.allowAnyAccount, + cacheScope: request.cacheScope, + logger: logger) + case .auto: + return try await importer.importBestCookies( + intoAccountEmail: request.normalizedTarget, + allowAnyAccount: request.allowAnyAccount, + preferCachedCookieHeader: request.preferCachedCookieHeader ?? !request.force, + cacheScope: request.cacheScope, + logger: logger) + case .off: + return OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Off", + cookieCount: 0, + signedInEmail: request.normalizedTarget, + matchesCodexEmail: true) + } + } + func importOpenAIDashboardCookiesIfNeeded( targetEmail: String?, force: Bool, preferCachedCookieHeader: Bool? = nil) async -> String? { + guard !Task.isCancelled else { return nil } if await self.openAIWebCookieImportShouldFailClosed() { return nil } + guard !Task.isCancelled else { return nil } let normalizedTarget = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) let allowAnyAccount = normalizedTarget == nil || normalizedTarget?.isEmpty == true @@ -996,42 +1082,17 @@ extension UsageStore { self.logOpenAIWeb(message) } - let result: OpenAIDashboardBrowserCookieImporter.ImportResult - if let override = self._test_openAIDashboardCookieImportOverride { - result = try await override(normalizedTarget, allowAnyAccount, cookieSource, cacheScope, log) - } else { - let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: self.browserDetection) - switch cookieSource { - case .manual: - self.settings.ensureCodexCookieLoaded() - // Manual OpenAI cookies still come from one provider-level setting. Auto-imported cookies are - // isolated per managed account, but a manual header is an explicit override owned by settings, - // so switching managed accounts does not currently swap it underneath the user. - let manualHeader = self.settings.codexCookieHeader - guard CookieHeaderNormalizer.normalize(manualHeader) != nil else { - throw OpenAIDashboardBrowserCookieImporter.ImportError.manualCookieHeaderInvalid - } - result = try await importer.importManualCookies( - cookieHeader: manualHeader, - intoAccountEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, - cacheScope: cacheScope, - logger: log) - case .auto: - result = try await importer.importBestCookies( - intoAccountEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, - preferCachedCookieHeader: preferCachedCookieHeader ?? !force, - cacheScope: cacheScope, - logger: log) - case .off: - result = OpenAIDashboardBrowserCookieImporter.ImportResult( - sourceLabel: "Off", - cookieCount: 0, - signedInEmail: normalizedTarget, - matchesCodexEmail: true) - } - } + let request = OpenAIDashboardCookieImportRequest( + normalizedTarget: normalizedTarget, + allowAnyAccount: allowAnyAccount, + cookieSource: cookieSource, + cacheScope: cacheScope, + preferCachedCookieHeader: preferCachedCookieHeader, + force: force) + let result = try await self.openAIDashboardCookieImportResult( + request: request, + logger: log) + guard !Task.isCancelled else { return nil } let effectiveEmail = result.signedInEmail? .trimmingCharacters(in: .whitespacesAndNewlines) .isEmpty == false @@ -1067,6 +1128,7 @@ extension UsageStore { } return effectiveEmail } catch let err as OpenAIDashboardBrowserCookieImporter.ImportError { + guard !Task.isCancelled else { return nil } switch err { case let .noMatchingAccount(found): let foundText: String = if found.isEmpty { @@ -1104,6 +1166,7 @@ extension UsageStore { } } } catch { + guard !Task.isCancelled else { return nil } self.logOpenAIWeb("[\(stamp)] import failed: \(error.localizedDescription)") await MainActor.run { self.openAIDashboardCookieImportStatus = diff --git a/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift b/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift index fe1f12702..33f25c47b 100644 --- a/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift +++ b/Tests/CodexBarTests/CodexBackgroundRefreshCoalescingTests.swift @@ -201,6 +201,41 @@ struct CodexBackgroundRefreshCoalescingTests { #expect(store.openAIDashboard?.creditsRemaining == 25) } + @Test + func `cancelled background dashboard import does not publish stale account status`() async throws { + let settings = try self.makeSettingsStore( + suite: "CodexBackgroundRefreshCoalescingTests-dashboard-cancelled-import") + settings.statusChecksEnabled = 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 importBlocker = BlockingOpenAIDashboardCookieImport() + store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in + try await importBlocker.awaitResult() + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + + let importTask = Task { @MainActor in + await store.importOpenAIDashboardCookiesIfNeeded( + targetEmail: managedAccount.email, + force: true) + } + await importBlocker.waitUntilStarted() + importTask.cancel() + await importBlocker.resumeNext(with: .failure( + OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount( + found: [.init(sourceLabel: "Chrome", email: "other@example.com")]))) + + let imported = await importTask.value + #expect(imported == nil) + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardCookieImportStatus == nil) + #expect(store.openAIDashboardRequiresLogin == false) + } + private func makeSettingsStore(suite: String) throws -> SettingsStore { let defaults = UserDefaults(suiteName: suite)! defaults.removePersistentDomain(forName: suite) @@ -284,3 +319,45 @@ struct CodexBackgroundRefreshCoalescingTests { return "\(base64URL(header)).\(base64URL(payload))." } } + +private actor BlockingOpenAIDashboardCookieImport { + private var continuations: [ + CheckedContinuation, Never> + ] = [] + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var started = 0 + + func awaitResult() async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult { + let result = await withCheckedContinuation { continuation in + self.continuations.append(continuation) + self.started += 1 + self.resumeReadyStartWaiters() + } + 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 resumeNext(with result: Result) { + guard !self.continuations.isEmpty else { return } + let continuation = self.continuations.removeFirst() + 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 + } +} diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift index f44eb9bcd..894341f78 100644 --- a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -189,8 +189,10 @@ struct CodexManagedOpenAIWebRefreshTests { #expect(firstCodexEntry.creditsRemaining == nil) await saver.resumeNext() + let backgroundTask = try #require(store.creditsRefreshTask) await creditsBlocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date()))) - try? await Task.sleep(for: .milliseconds(200)) + await backgroundTask.value + await saver.waitUntilStarted(count: 2) #expect(await saver.startedCount() == 2) let secondSnapshots = await saver.savedSnapshots() @@ -266,6 +268,7 @@ struct CodexManagedOpenAIWebRefreshTests { #expect(firstCodexEntry.codeReviewRemainingPercent == nil) await saver.resumeNext() + let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask) await dashboardBlocker.resumeNext(with: .success(OpenAIDashboardSnapshot( signedInEmail: managedAccount.email, codeReviewRemainingPercent: 95, @@ -276,7 +279,8 @@ struct CodexManagedOpenAIWebRefreshTests { creditsRemaining: 25, accountPlan: "Pro", updatedAt: Date()))) - try? await Task.sleep(for: .milliseconds(200)) + await backgroundTask.value + await saver.waitUntilStarted(count: 2) #expect(await saver.startedCount() == 2) let secondSnapshots = await saver.savedSnapshots() @@ -532,43 +536,15 @@ struct CodexManagedOpenAIWebRefreshTests { browserDetection: BrowserDetection(cacheTTL: 0), settings: settings, startupBehavior: .testing) - let blocker = BlockingManagedOpenAIDashboardLoader() - let importTracker = OpenAIDashboardImportCallTracker() + store.openAIDashboardCookieImportStatus = + "OpenAI cookies are for other@example.com, not managed@example.com." store._test_openAIDashboardLoaderOverride = { _, _, _ in - try await blocker.awaitResult() + throw ManagedDashboardTestError.networkTimeout } defer { store._test_openAIDashboardLoaderOverride = nil } - store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in - let call = await importTracker.recordCall() - if call == 1 { - return OpenAIDashboardBrowserCookieImporter.ImportResult( - sourceLabel: "Chrome", - cookieCount: 2, - signedInEmail: managedAccount.email, - matchesCodexEmail: true) - } - throw OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount( - found: [.init(sourceLabel: "Chrome", email: "other@example.com")]) - } - defer { store._test_openAIDashboardCookieImportOverride = nil } let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() - let firstTask = Task { - await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) - } - await blocker.waitUntilStarted(count: 1) - - let secondTask = Task { - await store.importOpenAIDashboardBrowserCookiesNow() - } - await blocker.waitUntilStarted(count: 2) - - await blocker.resumeNext(with: .failure(OpenAIDashboardFetcher.FetchError.loginRequired)) - await importTracker.waitUntilCalls(count: 2) - await blocker.resumeNext(with: .failure(ManagedDashboardTestError.networkTimeout)) - - await firstTask.value - await secondTask.value + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) #expect(store.lastOpenAIDashboardError == ManagedDashboardTestError.networkTimeout.localizedDescription) } @@ -675,10 +651,10 @@ actor BlockingManagedOpenAIDashboardLoader { private var started: Int = 0 func awaitResult() async throws -> OpenAIDashboardSnapshot { - self.started += 1 - self.resumeReadyStartWaiters() let result = await withCheckedContinuation { continuation in self.continuations.append(continuation) + self.started += 1 + self.resumeReadyStartWaiters() } return try result.get() } @@ -719,10 +695,10 @@ actor BlockingCreditsLoader { private var started = 0 func awaitResult() async throws -> CreditsSnapshot { - self.started += 1 - self.resumeReadyStartWaiters() let result = await withCheckedContinuation { continuation in self.continuations.append(continuation) + self.started += 1 + self.resumeReadyStartWaiters() } return try result.get() } From 68895e19006a267d054eab700744348c86bd1309 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 22 May 2026 20:34:45 +0100 Subject: [PATCH 7/7] test: await background codex credits refresh --- Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift b/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift index 3d64d4232..6eb336799 100644 --- a/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift +++ b/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift @@ -871,10 +871,12 @@ struct CodexAccountScopedRefreshTests { await blocker.waitUntilStarted() await blocker.resume(with: .success(self.codexSnapshot(email: "alpha@example.com", usedPercent: 12))) await refreshTask.value + #expect(store.lastCodexAccountScopedRefreshGuard?.accountKey == "alpha@example.com") + + await store.creditsRefreshTask?.value #expect(store.credits?.remaining == 55) #expect(store.lastCreditsSource == .api) - #expect(store.lastCodexAccountScopedRefreshGuard?.accountKey == "alpha@example.com") } @Test