wehub-resource-sync
122 lines
4.5 KiB
Swift
122 lines
4.5 KiB
Swift
import CodexBarCore
|
|
import Foundation
|
|
|
|
struct KiloScopeSnapshot: Identifiable, Equatable {
|
|
let id: String // KiloUsageScope.scopeIdentifier
|
|
let scope: KiloUsageScope
|
|
let snapshot: UsageSnapshot?
|
|
let errorMessage: String?
|
|
let sourceLabel: String?
|
|
|
|
static func == (lhs: KiloScopeSnapshot, rhs: KiloScopeSnapshot) -> Bool {
|
|
lhs.id == rhs.id
|
|
&& lhs.snapshot?.updatedAt == rhs.snapshot?.updatedAt
|
|
&& lhs.errorMessage == rhs.errorMessage
|
|
&& lhs.sourceLabel == rhs.sourceLabel
|
|
}
|
|
}
|
|
|
|
extension UsageStore {
|
|
var kiloEnabledScopes: [KiloUsageScope] {
|
|
var scopes: [KiloUsageScope] = [.personal]
|
|
let enabled = self.settings.kiloEnabledOrganizationIDs
|
|
guard !enabled.isEmpty else { return scopes }
|
|
let knownByID = Dictionary(
|
|
uniqueKeysWithValues: self.settings.kiloKnownOrganizations.map { ($0.id, $0) })
|
|
for id in enabled {
|
|
if let org = knownByID[id] {
|
|
scopes.append(.organization(id: org.id, name: org.name))
|
|
}
|
|
}
|
|
return scopes
|
|
}
|
|
|
|
func shouldFanOutKiloScopes() -> Bool {
|
|
self.kiloEnabledScopes.count > 1
|
|
}
|
|
|
|
func refreshKiloScopes(generation: UInt64? = nil) async {
|
|
let scopes = self.kiloEnabledScopes
|
|
guard scopes.count > 1 else {
|
|
await MainActor.run { self.kiloScopeSnapshots = [] }
|
|
return
|
|
}
|
|
let env = ProcessInfo.processInfo.environment
|
|
let resolved: KiloResolvedBearerToken
|
|
do {
|
|
resolved = try KiloBearerTokenResolver.resolve(
|
|
source: self.settings.kiloUsageDataSource,
|
|
apiKey: self.settings.configSnapshot.providerConfig(for: .kilo)?.sanitizedAPIKey,
|
|
environment: env)
|
|
} catch {
|
|
let message = (error as? LocalizedError)?.errorDescription
|
|
?? error.localizedDescription
|
|
await MainActor.run {
|
|
self.kiloScopeSnapshots = scopes.map {
|
|
KiloScopeSnapshot(
|
|
id: $0.scopeIdentifier,
|
|
scope: $0,
|
|
snapshot: nil,
|
|
errorMessage: message,
|
|
sourceLabel: nil)
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
let results: [KiloScopeSnapshot] = await withTaskGroup(of: KiloScopeSnapshot.self) { group in
|
|
for scope in scopes {
|
|
group.addTask {
|
|
do {
|
|
let raw = try await KiloUsageFetcher.fetchUsage(
|
|
apiKey: resolved.token,
|
|
scope: scope,
|
|
environment: env)
|
|
let snapshot = raw.toUsageSnapshot()
|
|
.withAccountOrganization(scope.displayName)
|
|
return KiloScopeSnapshot(
|
|
id: scope.scopeIdentifier,
|
|
scope: scope,
|
|
snapshot: snapshot,
|
|
errorMessage: nil,
|
|
sourceLabel: resolved.sourceLabel)
|
|
} catch {
|
|
return KiloScopeSnapshot(
|
|
id: scope.scopeIdentifier,
|
|
scope: scope,
|
|
snapshot: nil,
|
|
errorMessage: (error as? LocalizedError)?.errorDescription
|
|
?? error.localizedDescription,
|
|
sourceLabel: nil)
|
|
}
|
|
}
|
|
}
|
|
var collected: [KiloScopeSnapshot] = []
|
|
for await result in group {
|
|
collected.append(result)
|
|
}
|
|
return collected
|
|
}
|
|
|
|
let resultByID = Dictionary(uniqueKeysWithValues: results.map { ($0.id, $0) })
|
|
let ordered = scopes.compactMap { resultByID[$0.scopeIdentifier] }
|
|
|
|
await MainActor.run {
|
|
guard self.isCurrentProviderRefreshGeneration(.kilo, generation: generation) else { return }
|
|
self.kiloScopeSnapshots = ordered
|
|
}
|
|
}
|
|
}
|
|
|
|
extension UsageSnapshot {
|
|
fileprivate func withAccountOrganization(_ org: String) -> UsageSnapshot {
|
|
let baseIdentity = self.identity
|
|
let newIdentity = ProviderIdentitySnapshot(
|
|
providerID: baseIdentity?.providerID ?? .kilo,
|
|
accountEmail: baseIdentity?.accountEmail,
|
|
accountOrganization: org,
|
|
loginMethod: baseIdentity?.loginMethod)
|
|
return self.withIdentity(newIdentity)
|
|
}
|
|
}
|