Files
steipete--codexbar/Sources/CodexBar/UsageStore+RefreshEnrichment.swift
2026-07-13 12:22:33 +08:00

326 lines
13 KiB
Swift

import CodexBarCore
import Foundation
extension UsageStore {
enum RefreshEnrichmentMode: Equatable, Sendable {
case automatic
case forcedForeground
case forcedBackground
}
struct RequiredRefreshRequest: Sendable {
var throughGeneration: UInt64
var startupConnectivityRetryAttempt: Int?
var coalesceProviderRefreshes: Bool
var interaction: ProviderInteraction
mutating func merge(_ newer: Self) {
self.throughGeneration = max(self.throughGeneration, newer.throughGeneration)
if let newerAttempt = newer.startupConnectivityRetryAttempt {
self.startupConnectivityRetryAttempt = max(
self.startupConnectivityRetryAttempt ?? newerAttempt,
newerAttempt)
}
// Replacement is the stronger policy: a settings change must not join work started
// with the old configuration merely because another required refresh arrived first.
self.coalesceProviderRefreshes = self.coalesceProviderRefreshes && newer.coalesceProviderRefreshes
if newer.interaction == .userInitiated {
self.interaction = .userInitiated
}
}
}
func refresh(forceTokenUsage: Bool = false) async {
if forceTokenUsage {
await self.refresh(enrichmentMode: .forcedForeground)
} else {
await self.runRefresh(
startupConnectivityRetryAttempt: nil,
waitForRefreshAvailability: true)
}
}
private struct ForcedRefreshEnrichmentRequest: Sendable {
let generation: UInt64
let refreshStartedAt: Date
let openAIWebRefreshPhase: ProviderRefreshPhase
}
func refresh(enrichmentMode: RefreshEnrichmentMode) async {
if enrichmentMode == .forcedForeground {
await self.cancelForcedRefreshEnrichmentAndWait()
}
await self.runRefresh(
enrichmentMode: enrichmentMode,
startupConnectivityRetryAttempt: nil)
}
func enqueueRequiredRefresh(
startupConnectivityRetryAttempt: Int?,
coalesceProviderRefreshesOverride: Bool?) async -> Bool
{
self.requiredRefreshRequestGeneration &+= 1
let interaction = ProviderInteractionContext.current
let request = RequiredRefreshRequest(
throughGeneration: self.requiredRefreshRequestGeneration,
startupConnectivityRetryAttempt: startupConnectivityRetryAttempt,
coalesceProviderRefreshes: coalesceProviderRefreshesOverride ?? (interaction == .background),
interaction: interaction)
if var pending = self.pendingRequiredRefreshRequest {
pending.merge(request)
self.pendingRequiredRefreshRequest = pending
} else {
self.pendingRequiredRefreshRequest = request
}
if let task = self.requiredRefreshTask {
return await task.value
}
let token = UUID()
self.requiredRefreshTaskToken = token
let task = Task { @MainActor [weak self] in
guard let self else { return false }
let didRefresh = await self.drainRequiredRefreshRequests()
self.completeRequiredRefreshTask(token: token)
return didRefresh
}
self.requiredRefreshTask = task
return await task.value
}
func cancelRequiredRefresh() {
self.pendingRequiredRefreshRequest = nil
self.requiredRefreshTaskToken = nil
let task = self.requiredRefreshTask
self.requiredRefreshTask = nil
task?.cancel()
}
private func drainRequiredRefreshRequests() async -> Bool {
var completedAnyRefresh = false
while !Task.isCancelled {
guard await self.waitForRequiredRefreshAvailability(),
let request = self.pendingRequiredRefreshRequest
else {
break
}
self.pendingRequiredRefreshRequest = nil
let didRefresh = await ProviderInteractionContext.$current.withValue(request.interaction) {
await self.runRefresh(
startupConnectivityRetryAttempt: request.startupConnectivityRetryAttempt,
coalesceProviderRefreshesOverride: request.coalesceProviderRefreshes)
}
if didRefresh {
completedAnyRefresh = true
self.requiredRefreshCompletedGeneration = max(
self.requiredRefreshCompletedGeneration,
request.throughGeneration)
} else if !Task.isCancelled {
var retry = request
if let pending = self.pendingRequiredRefreshRequest {
retry.merge(pending)
}
self.pendingRequiredRefreshRequest = retry
}
}
return completedAnyRefresh
}
private func waitForRequiredRefreshAvailability() async -> Bool {
while self.isRefreshing || self.hasForcedRefreshEnrichmentInFlight {
guard !Task.isCancelled else { return false }
if self.hasForcedRefreshEnrichmentInFlight {
await self.awaitForcedRefreshEnrichment()
} else {
do {
try await Task.sleep(for: .milliseconds(20))
} catch {
return false
}
}
}
return !Task.isCancelled
}
private func completeRequiredRefreshTask(token: UUID) {
guard self.requiredRefreshTaskToken == token else { return }
self.requiredRefreshTask = nil
self.requiredRefreshTaskToken = nil
}
func enqueueForcedRefreshEnrichment(
generation: UInt64,
refreshStartedAt: Date,
openAIWebRefreshPhase: ProviderRefreshPhase)
{
let request = ForcedRefreshEnrichmentRequest(
generation: generation,
refreshStartedAt: refreshStartedAt,
openAIWebRefreshPhase: openAIWebRefreshPhase)
if let predecessor = self.forcedRefreshEnrichmentTask {
self.replacePendingForcedRefreshEnrichment(request, predecessor: predecessor)
} else {
self.startForcedRefreshEnrichment(request)
}
}
func awaitForcedRefreshEnrichment() async {
var reportedWait = false
while !Task.isCancelled {
guard let task = self.pendingForcedRefreshEnrichmentTask ?? self.forcedRefreshEnrichmentTask else {
return
}
if !reportedWait {
self._test_forcedRefreshEnrichmentWaitObserver?()
reportedWait = true
}
await task.value
}
}
func cancelForcedRefreshEnrichment() {
_ = self.cancelForcedRefreshEnrichmentTasks()
}
private func cancelForcedRefreshEnrichmentAndWait() async {
let tasks = self.cancelForcedRefreshEnrichmentTasks()
for task in tasks {
await task.value
}
}
private func startForcedRefreshEnrichment(_ request: ForcedRefreshEnrichmentRequest) {
let token = UUID()
self.forcedRefreshEnrichmentToken = token
self.hasForcedRefreshEnrichmentInFlight = true
self.forcedRefreshEnrichmentTask = Task(priority: .utility) { @MainActor [weak self] in
guard let self else { return }
await self.runForcedRefreshEnrichment(request)
self.completeForcedRefreshEnrichment(token: token)
}
}
private func replacePendingForcedRefreshEnrichment(
_ request: ForcedRefreshEnrichmentRequest,
predecessor: Task<Void, Never>)
{
self.pendingForcedRefreshEnrichmentTask?.cancel()
let token = UUID()
self.pendingForcedRefreshEnrichmentToken = token
self.hasForcedRefreshEnrichmentInFlight = true
self.pendingForcedRefreshEnrichmentTask = Task(priority: .utility) { @MainActor [weak self] in
await predecessor.value
guard !Task.isCancelled,
let self,
self.pendingForcedRefreshEnrichmentToken == token,
let promotedTask = self.pendingForcedRefreshEnrichmentTask
else { return }
self.pendingForcedRefreshEnrichmentTask = nil
self.pendingForcedRefreshEnrichmentToken = nil
self.forcedRefreshEnrichmentTask = promotedTask
self.forcedRefreshEnrichmentToken = token
await self.runForcedRefreshEnrichment(request)
self.completeForcedRefreshEnrichment(token: token)
}
}
private func completeForcedRefreshEnrichment(token: UUID) {
guard self.forcedRefreshEnrichmentToken == token else { return }
// Keep the completed predecessor installed until its latest pending waiter promotes itself.
// This avoids an actor-reentrancy gap where a new request could otherwise start beside it.
guard self.pendingForcedRefreshEnrichmentTask == nil else { return }
self.forcedRefreshEnrichmentTask = nil
self.forcedRefreshEnrichmentToken = nil
self.hasForcedRefreshEnrichmentInFlight = false
}
private func cancelForcedRefreshEnrichmentTasks() -> [Task<Void, Never>] {
let tasks = [
self.forcedRefreshEnrichmentTask,
self.pendingForcedRefreshEnrichmentTask,
self.openAIDashboardBackgroundRefreshTask,
self.openAIDashboardRefreshTask,
].compactMap(\.self)
self.forcedRefreshEnrichmentTask = nil
self.forcedRefreshEnrichmentToken = nil
self.pendingForcedRefreshEnrichmentTask = nil
self.pendingForcedRefreshEnrichmentToken = nil
self.hasForcedRefreshEnrichmentInFlight = false
tasks.forEach { $0.cancel() }
self.invalidateOpenAIDashboardRefreshTask()
return tasks
}
private func runForcedRefreshEnrichment(_ request: ForcedRefreshEnrichmentRequest) async {
await withTaskGroup(of: Void.self) { group in
group.addTask {
await self.refreshCreditsNow(minimumSnapshotUpdatedAt: request.refreshStartedAt)
}
group.addTask {
await self.refreshTokenUsageSequenceNow(force: true)
}
}
guard !Task.isCancelled else { return }
await self.refreshOpenAIWebAfterProviderRefresh(
force: true,
refreshPhase: request.openAIWebRefreshPhase)
guard !Task.isCancelled else { return }
if self.openAIDashboardRequiresLogin,
request.generation == self.forcedRefreshEnrichmentGeneration
{
// Join a newer in-flight Codex request rather than replacing it. A newer accepted all-provider
// pass owns reconciliation even before it can enqueue its tail, so recheck generation afterward.
await self.refreshProvider(.codex, coalesceIfRefreshing: true)
guard !Task.isCancelled else { return }
if request.generation == self.forcedRefreshEnrichmentGeneration {
await self.refreshCreditsNow(minimumSnapshotUpdatedAt: request.refreshStartedAt)
guard !Task.isCancelled else { return }
}
}
self.persistWidgetSnapshot(reason: "forced-refresh-enrichment")
}
func refreshOpenAIWebAfterProviderRefresh(
force: Bool,
refreshPhase: ProviderRefreshPhase) async
{
self.syncOpenAIWebState()
let refreshPolicy = OpenAIWebRefreshPolicyContext(
accessEnabled: self.isEnabled(.codex) &&
self.settings.openAIWebAccessEnabled &&
self.settings.codexCookieSource.isEnabled,
batterySaverEnabled: self.settings.openAIWebBatterySaverEnabled,
force: force,
refreshPhase: refreshPhase)
let shouldRefreshOpenAIWeb = Self.shouldRunOpenAIWebRefresh(refreshPolicy)
self.openAIWebLogger.debug(
"OpenAI web refresh gate",
metadata: [
"allowed": shouldRefreshOpenAIWeb ? "1" : "0",
"accessEnabled": refreshPolicy.accessEnabled ? "1" : "0",
"batterySaverEnabled": refreshPolicy.batterySaverEnabled ? "1" : "0",
"force": refreshPolicy.force ? "1" : "0",
"interaction": ProviderInteractionContext.current == .userInitiated ? "user" : "background",
"phase": refreshPhase == .startup ? "startup" : "regular",
])
guard shouldRefreshOpenAIWeb, !Task.isCancelled else { return }
let codexDashboardGuard = self.freshCodexOpenAIWebRefreshGuard()
if force {
await self.refreshOpenAIDashboardIfNeeded(
force: true,
expectedGuard: codexDashboardGuard,
bypassCoalescing: true)
} else {
self.scheduleOpenAIDashboardRefreshIfNeeded(expectedGuard: codexDashboardGuard)
}
}
}