wehub-resource-sync
270 lines
11 KiB
Swift
270 lines
11 KiB
Swift
import CodexBarCore
|
|
import Foundation
|
|
|
|
@MainActor
|
|
extension UsageStore {
|
|
nonisolated static let codexSnapshotWaitTimeoutSeconds: TimeInterval = 6
|
|
nonisolated static let codexRefreshStartGraceSeconds: TimeInterval = 0.25
|
|
nonisolated static let codexSnapshotPollIntervalNanoseconds: UInt64 = 100_000_000
|
|
|
|
func codexCreditsFetcher() -> UsageFetcher {
|
|
// Credits are remote Codex account state, so they need the same managed-home routing as the
|
|
// primary Codex usage fetch. Local token-cost scanning intentionally stays ambient-system scoped.
|
|
self.makeFetchContext(provider: .codex, override: nil).fetcher
|
|
}
|
|
|
|
func scheduleCreditsRefreshIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) {
|
|
let refreshKey = self.codexCreditsRefreshKey(
|
|
expectedGuard: self.freshCodexAccountScopedRefreshGuard())
|
|
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 {
|
|
if self.creditsRefreshTaskKey == refreshKey {
|
|
self.creditsRefreshTask = nil
|
|
self.creditsRefreshTaskKey = nil
|
|
}
|
|
}
|
|
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:
|
|
"live"
|
|
case let .managedAccount(id):
|
|
"managed:\(id.uuidString)"
|
|
case let .profileHome(path):
|
|
"profile:\(path)"
|
|
}
|
|
|
|
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",
|
|
"auth:\(expectedGuard.authFingerprint ?? "nil")",
|
|
].joined(separator: "|")
|
|
}
|
|
|
|
func refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) async {
|
|
guard self.isEnabled(.codex) else { return }
|
|
var expectedGuard = self.freshCodexAccountScopedRefreshGuard()
|
|
if expectedGuard.identity == .unresolved,
|
|
let minimumSnapshotUpdatedAt,
|
|
case .liveSystem = expectedGuard.source
|
|
{
|
|
_ = await self.waitForCodexSnapshotOrRefreshCompletion(minimumUpdatedAt: minimumSnapshotUpdatedAt)
|
|
expectedGuard = self.freshCodexAccountScopedRefreshGuard()
|
|
}
|
|
guard expectedGuard.identity != .unresolved,
|
|
expectedGuard.accountKey != nil
|
|
else {
|
|
return
|
|
}
|
|
do {
|
|
let credits = try await self.loadLatestCodexCredits()
|
|
guard !Task.isCancelled else { return }
|
|
guard let applyGuard = self.codexScopedNonUsageSuccessApplyGuard(
|
|
expectedGuard: expectedGuard) else { return }
|
|
self.reconcileCodexPublishedUsageOwner(with: applyGuard)
|
|
await MainActor.run {
|
|
self.credits = credits
|
|
self.lastCreditsError = nil
|
|
self.lastCreditsSnapshot = credits
|
|
self.lastCreditsSnapshotAccountKey = applyGuard.accountKey
|
|
self.lastCreditsSource = .api
|
|
self.creditsFailureStreak = 0
|
|
self.lastCodexAccountScopedRefreshGuard = applyGuard
|
|
}
|
|
let codexSnapshot = await MainActor.run {
|
|
self.snapshots[.codex]
|
|
}
|
|
if let minimumSnapshotUpdatedAt,
|
|
codexSnapshot == nil || codexSnapshot?.updatedAt ?? .distantPast < minimumSnapshotUpdatedAt
|
|
{
|
|
self.scheduleCodexPlanHistoryBackfill(
|
|
minimumSnapshotUpdatedAt: minimumSnapshotUpdatedAt)
|
|
return
|
|
}
|
|
|
|
self.cancelCodexPlanHistoryBackfill()
|
|
guard let codexSnapshot else { return }
|
|
await self.recordPlanUtilizationHistorySample(
|
|
provider: .codex,
|
|
snapshot: codexSnapshot,
|
|
now: codexSnapshot.updatedAt)
|
|
} catch {
|
|
guard !Task.isCancelled else { return }
|
|
let message = error.localizedDescription
|
|
if message.localizedCaseInsensitiveContains("data not available yet") {
|
|
guard self.shouldApplyCodexScopedNonUsageFailure(expectedGuard: expectedGuard) else { return }
|
|
self.reconcileCodexPublishedUsageOwner(with: expectedGuard)
|
|
await MainActor.run {
|
|
if let cached = self.lastCreditsSnapshot,
|
|
self.lastCreditsSnapshotAccountKey == expectedGuard.accountKey
|
|
{
|
|
self.credits = cached
|
|
self.lastCreditsError = nil
|
|
self.lastCodexAccountScopedRefreshGuard = expectedGuard
|
|
} else {
|
|
self.credits = nil
|
|
self.lastCreditsSource = .none
|
|
self.lastCreditsError = L("Codex credits are still loading; will retry shortly.")
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
guard self.shouldApplyCodexScopedNonUsageFailure(expectedGuard: expectedGuard) else { return }
|
|
self.reconcileCodexPublishedUsageOwner(with: expectedGuard)
|
|
await MainActor.run {
|
|
self.creditsFailureStreak += 1
|
|
if let cached = self.lastCreditsSnapshot,
|
|
self.lastCreditsSnapshotAccountKey == expectedGuard.accountKey
|
|
{
|
|
self.credits = cached
|
|
let stamp = cached.updatedAt.formatted(date: .abbreviated, time: .shortened)
|
|
self.lastCreditsError =
|
|
"Last Codex credits refresh failed: \(message). Cached values from \(stamp)."
|
|
self.lastCodexAccountScopedRefreshGuard = expectedGuard
|
|
} else {
|
|
self.lastCreditsError = message
|
|
self.credits = nil
|
|
self.lastCreditsSource = .none
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private func loadLatestCodexCredits() async throws -> CreditsSnapshot {
|
|
if let override = self._test_codexCreditsLoaderOverride {
|
|
return try await override()
|
|
}
|
|
let descriptor = self.providerSpecs[.codex]?.descriptor ?? ProviderDescriptorRegistry.descriptor(for: .codex)
|
|
let context = self.makeFetchContext(provider: .codex, override: nil, includeCredits: true)
|
|
let strategies = await descriptor.fetchPlan.pipeline.resolveStrategies(context)
|
|
var lastAvailableError: Error?
|
|
|
|
for strategy in strategies {
|
|
guard await strategy.isAvailable(context) else { continue }
|
|
do {
|
|
let result = try await strategy.fetch(context)
|
|
if let credits = result.credits {
|
|
return credits
|
|
}
|
|
lastAvailableError = UsageError.noRateLimitsFound
|
|
guard context.sourceMode == .auto else { break }
|
|
} catch {
|
|
lastAvailableError = error
|
|
guard strategy.shouldFallback(on: error, context: context) else { break }
|
|
}
|
|
}
|
|
throw lastAvailableError ?? ProviderFetchError.noAvailableStrategy(.codex)
|
|
}
|
|
|
|
func waitForCodexSnapshot(minimumUpdatedAt: Date) async -> UsageSnapshot? {
|
|
let deadline = Date().addingTimeInterval(Self.codexSnapshotWaitTimeoutSeconds)
|
|
|
|
while Date() < deadline {
|
|
if Task.isCancelled {
|
|
return nil
|
|
}
|
|
if let snapshot = await MainActor.run(body: { self.snapshots[.codex] }),
|
|
snapshot.updatedAt >= minimumUpdatedAt
|
|
{
|
|
return snapshot
|
|
}
|
|
try? await Task.sleep(nanoseconds: Self.codexSnapshotPollIntervalNanoseconds)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func waitForCodexSnapshotOrRefreshCompletion(minimumUpdatedAt: Date) async -> UsageSnapshot? {
|
|
let deadline = Date().addingTimeInterval(Self.codexSnapshotWaitTimeoutSeconds)
|
|
let refreshStartDeadline = Date().addingTimeInterval(Self.codexRefreshStartGraceSeconds)
|
|
|
|
while Date() < deadline {
|
|
if Task.isCancelled {
|
|
return nil
|
|
}
|
|
let state = await MainActor.run {
|
|
(
|
|
snapshot: self.snapshots[.codex],
|
|
isRefreshing: self.refreshingProviders.contains(.codex),
|
|
hasAttempts: !(self.lastFetchAttempts[.codex] ?? []).isEmpty,
|
|
hasError: self.errors[.codex] != nil)
|
|
}
|
|
if let snapshot = state.snapshot, snapshot.updatedAt >= minimumUpdatedAt {
|
|
return snapshot
|
|
}
|
|
if !state.isRefreshing, state.hasAttempts || state.hasError {
|
|
return nil
|
|
}
|
|
if !state.isRefreshing,
|
|
!state.hasAttempts,
|
|
!state.hasError,
|
|
Date() >= refreshStartDeadline
|
|
{
|
|
return nil
|
|
}
|
|
try? await Task.sleep(nanoseconds: Self.codexSnapshotPollIntervalNanoseconds)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func scheduleCodexPlanHistoryBackfill(
|
|
minimumSnapshotUpdatedAt: Date)
|
|
{
|
|
self.cancelCodexPlanHistoryBackfill()
|
|
self.codexPlanHistoryBackfillTask = Task { @MainActor [weak self] in
|
|
guard let self else { return }
|
|
guard let snapshot = await self.waitForCodexSnapshot(minimumUpdatedAt: minimumSnapshotUpdatedAt) else {
|
|
return
|
|
}
|
|
await self.recordPlanUtilizationHistorySample(
|
|
provider: .codex,
|
|
snapshot: snapshot,
|
|
now: snapshot.updatedAt)
|
|
self.codexPlanHistoryBackfillTask = nil
|
|
}
|
|
}
|
|
|
|
func cancelCodexPlanHistoryBackfill() {
|
|
self.codexPlanHistoryBackfillTask?.cancel()
|
|
self.codexPlanHistoryBackfillTask = nil
|
|
}
|
|
}
|