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() } 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() } 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()))) 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) } 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) } 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 alphaRefreshTask.value await betaRefreshTask.value await staleCreditsTask.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 regularCompletion = RefreshCompletionProbe() let regularRefreshTask = Task { await store.refresh(forceTokenUsage: false) await regularCompletion.markCompleted() } 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) } 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 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( 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 } await store.refresh(forceTokenUsage: false) store._test_openAIDashboardLoaderOverride = { _, _, _, _ in try await blocker.awaitResult() } defer { store._test_openAIDashboardLoaderOverride = nil } let firstRefreshTask = Task { await store.refresh(forceTokenUsage: false) await firstCompletion.markCompleted() } 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() } 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) 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 `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) } 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")]))) let imported = await importTask.value #expect(imported == nil) #expect(store.openAIDashboard == nil) #expect(store.openAIDashboardCookieImportStatus == nil) #expect(store.openAIDashboardRequiresLogin == false) } @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 settings.openAIWebAccessEnabled = true settings.codexCookieSource = .auto return settings } 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, environmentBase: environment) } 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))." } } 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 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 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, 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 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 cancel(id: UUID) { guard self.continuations.contains(where: { $0.id == id }), self.cancelledIDs.insert(id).inserted else { return } self.cancellations += 1 } }