import CodexBarCoreimport 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) }}
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)
}
}