1460 lines · 58.0 KB · Swift
RawDownload
import CodexBarCoreimport Commanderimport Foundation struct ServeOptions: CommanderParsable {    @Flag(names: [.short("v"), .long("verbose")], help: "Enable verbose logging")    var verbose: Bool = false     @Flag(name: .long("json-output"), help: "Emit machine-readable logs")    var jsonOutput: Bool = false     @Option(name: .long("log-level"), help: "Set log level (trace|verbose|debug|info|warning|error|critical)")    var logLevel: String?     @Option(name: .long("port"), help: "Local HTTP port (default: 8080)")    var port: Int?     @Option(name: .long("host"), help: "IPv4 bind address or localhost (default: 127.0.0.1)")    var host: String?     @Option(name: .long("refresh-interval"), help: "Response cache TTL in seconds (default: 60)")    var refreshInterval: Double?     @Option(        name: .long("request-timeout"),        help: "Total per-request deadline in seconds; 0 disables (default: 30)")    var requestTimeout: Double?     @Option(        name: .long("dashboard-token"),        help: "Bearer token for /dashboard/v1/snapshot (prefer CODEXBAR_DASHBOARD_TOKEN)")    var dashboardBearer: String?     @Flag(        name: .long("allow-plain-http"),        help: "Accept sending the dashboard token over cleartext HTTP on a non-loopback host")    var allowPlainHTTP: Bool = false} enum CLIServeRoute: Equatable {    case health    case usage(provider: String?)    case cost(provider: String?)    case dashboardSnapshot} enum CLIServeRouteError: Error, Equatable {    case methodNotAllowed    case notFound} enum CLIServeRouter {    static func route(method: String, path: String, queryItems: [String: String]) throws -> CLIServeRoute {        guard method.uppercased() == "GET" else {            throw CLIServeRouteError.methodNotAllowed        }         let provider = queryItems["provider"]?.trimmingCharacters(in: .whitespacesAndNewlines)        let normalizedProvider = provider?.isEmpty == false ? provider : nil         switch path {        case "/health":            return .health        case "/usage":            return .usage(provider: normalizedProvider)        case "/cost":            return .cost(provider: normalizedProvider)        case "/dashboard/v1/snapshot":            return .dashboardSnapshot        default:            throw CLIServeRouteError.notFound        }    }} private struct ServeErrorPayload: Encodable {    let error: String} private struct ServeHealthPayload: Encodable {    let status: String    let version: String?     private enum CodingKeys: String, CodingKey {        case status        case version    }     func encode(to encoder: any Encoder) throws {        var container = encoder.container(keyedBy: CodingKeys.self)        try container.encode(self.status, forKey: .status)        try container.encodeIfPresent(self.version, forKey: .version)    }} struct CLIServeConfigSnapshot {    let config: CodexBarConfig    let cacheToken: String} struct ServeRuntime {    let configStore: CodexBarConfigStore    let cache: CLIServeResponseCache    let providerOperations: CLIServeOperationCoordinator<UsageCommandOutput>    let costOperations: CLIServeOperationCoordinator<CostPayload>    let refreshInterval: TimeInterval    let requestTimeout: TimeInterval    let healthVersion: String?    let dashboardAuth: CLIServeDashboardAuth    /// True for non-loopback binds: every data route (`/usage`, `/cost`,    /// `/dashboard/v1/snapshot`) then requires the bearer token, so account data    /// is never exposed to the network unauthenticated. `/health` stays open.    /// Resolved once at startup from the bind host.    let dataRoutesRequireAuth: Bool     init(        configStore: CodexBarConfigStore,        cache: CLIServeResponseCache,        providerOperations: CLIServeOperationCoordinator<UsageCommandOutput>,        costOperations: CLIServeOperationCoordinator<CostPayload>,        refreshInterval: TimeInterval,        requestTimeout: TimeInterval,        healthVersion: String?,        dashboardAuth: CLIServeDashboardAuth,        bindHost: String)    {        self.configStore = configStore        self.cache = cache        self.providerOperations = providerOperations        self.costOperations = costOperations        self.refreshInterval = refreshInterval        self.requestTimeout = requestTimeout        self.healthVersion = healthVersion        self.dashboardAuth = dashboardAuth        self.dataRoutesRequireAuth = !CLIServeSecurity.isLoopbackHost(bindHost)    }} private struct ServeResponseRequest: Sendable {    let key: String    let configFingerprint: String    let refreshInterval: TimeInterval    let deadline: ContinuousClock.Instant?} struct CLIServeCoordinatedResponse: Sendable {    let response: CLILocalHTTPResponse    let isCommitted: Bool} private struct ServeUsageContext: Sendable {    let config: CodexBarConfig    let configFingerprint: String    let refreshInterval: TimeInterval    let providerTimeout: TimeInterval?    let providerDeadline: ContinuousClock.Instant?    let providerOperations: CLIServeOperationCoordinator<UsageCommandOutput>    let includeAllCodexAccounts: Bool     init(        config: CodexBarConfig,        configFingerprint: String,        refreshInterval: TimeInterval,        providerTimeout: TimeInterval?,        providerDeadline: ContinuousClock.Instant?,        providerOperations: CLIServeOperationCoordinator<UsageCommandOutput>,        includeAllCodexAccounts: Bool = true)    {        self.config = config        self.configFingerprint = configFingerprint        self.refreshInterval = refreshInterval        self.providerTimeout = providerTimeout        self.providerDeadline = providerDeadline        self.providerOperations = providerOperations        self.includeAllCodexAccounts = includeAllCodexAccounts    }} private struct ServeDashboardContext: Sendable {    let config: CodexBarConfig    let usage: ServeUsageContext    let costCollection: ServeCostCollectionContext    let codexBarVersion: String?} private struct ServeCostContext: Sendable {    let config: CodexBarConfig    let collection: ServeCostCollectionContext} struct ServeCostCollectionContext: Sendable {    let configFingerprint: String    let providerTimeout: TimeInterval?    let requestDeadline: ContinuousClock.Instant?    let now: @Sendable () -> ContinuousClock.Instant    let providerOperations: CLIServeOperationCoordinator<CostPayload>} actor CLIServeResponseCache {    static let maximumStaleTTL: TimeInterval = 3600    nonisolated let operations: CLIServeOperationCoordinator<CLIServeCoordinatedResponse>     init(operations: CLIServeOperationCoordinator<CLIServeCoordinatedResponse> = CLIServeOperationCoordinator()) {        self.operations = operations    }     private struct Entry {        let expiresAt: Date        let response: CLILocalHTTPResponse    }     struct CachePolicy {        /// How long a successful response stays fresh.        let ttl: TimeInterval        /// How long a last-good response may stand in for a failed refresh.        let staleTTL: TimeInterval    }     private struct LastGoodEntry {        let recordedAt: Date        let response: CLILocalHTTPResponse    }     private struct UsageItemKey: Hashable {        let provider: String        let accountID: String    }     private struct LastGoodUsageItem {        let recordedAt: Date        let data: Data    }     private struct UsageMergeResult {        let response: CLILocalHTTPResponse    }     private struct LastGoodCostItem {        let recordedAt: Date        let data: Data    }     private struct CostMergeResult {        let response: CLILocalHTTPResponse    }     private var entries: [String: Entry] = [:]    private var lastGood: [String: LastGoodEntry] = [:]    private var lastGoodUsageItems: [String: [UsageItemKey: LastGoodUsageItem]] = [:]    private var lastGoodCostItems: [String: [String: LastGoodCostItem]] = [:]    private var lastGoodCostOrder: [String: [String]] = [:]     private func pruneExpiredEntries(now: Date) {        self.entries = self.entries.filter { $0.value.expiresAt > now }        self.lastGood = self.lastGood.filter {            now.timeIntervalSince($0.value.recordedAt) <= Self.maximumStaleTTL        }        self.lastGoodUsageItems = self.lastGoodUsageItems.compactMapValues { items in            let retained = items.filter {                now.timeIntervalSince($0.value.recordedAt) <= Self.maximumStaleTTL            }            return retained.isEmpty ? nil : retained        }        self.lastGoodCostItems = self.lastGoodCostItems.compactMapValues { items in            let retained = items.filter {                now.timeIntervalSince($0.value.recordedAt) <= Self.maximumStaleTTL            }            return retained.isEmpty ? nil : retained        }        self.lastGoodCostOrder = self.lastGoodCostOrder.filter { self.lastGoodCostItems[$0.key] != nil }    }     private func response(for key: String) -> CLILocalHTTPResponse? {        guard let entry = self.entries[key] else { return nil }        return entry.response    }     func cachedResponse(for key: String, now: Date) -> CLILocalHTTPResponse? {        self.pruneExpiredEntries(now: now)        return self.response(for: key)    }     /// Transforms a fetched response through the cache's stale policy. Successful    /// responses are cached normally. Failed non-usage    /// fetches may use a whole-response fallback within `staleTTL`; usage    /// responses only replace keyed error rows from the same identified account.    func completeFetch(        _ response: CLILocalHTTPResponse,        for key: String,        policy: CachePolicy,        now: Date,        shouldCache: Bool) -> CLILocalHTTPResponse    {        let delivered: CLILocalHTTPResponse        let staleResponse = self.staleResponse(for: key, staleTTL: policy.staleTTL, now: now)        let usageMerge = self.mergeLastGoodUsageItems(            into: response,            for: key,            staleTTL: policy.staleTTL,            now: now,            replaceCachedItems: shouldCache)        let costMerge = self.mergeLastGoodCostItems(            into: response,            for: key,            staleTTL: policy.staleTTL,            now: now,            replaceCachedItems: shouldCache)        if shouldCache {            self.store(response, for: key, ttl: policy.ttl, now: now)            if key.hasPrefix("usage:") || key.hasPrefix("cost:") {                self.lastGood[key] = nil            } else {                self.lastGood[key] = LastGoodEntry(recordedAt: now, response: response)            }            delivered = response        } else if let usageMerge {            delivered = usageMerge.response            self.lastGood[key] = nil        } else if let costMerge {            delivered = costMerge.response            self.lastGood[key] = nil        } else {            delivered = staleResponse ?? response        }        return delivered    }     private func staleResponse(        for key: String,        staleTTL: TimeInterval,        now: Date) -> CLILocalHTTPResponse?    {        guard staleTTL > 0 else { return nil }        // A timeout cannot prove which usage account is currently active.        if key.hasPrefix("usage:") {            return nil        }        if key.hasPrefix("cost:") {            return self.staleCostResponse(for: key, staleTTL: staleTTL, now: now)        }        if let entry = self.lastGood[key],           now.timeIntervalSince(entry.recordedAt) <= staleTTL        {            return entry.response        }        if self.lastGood[key] != nil {            self.lastGood[key] = nil        }        return nil    }     private func mergeLastGoodUsageItems(        into response: CLILocalHTTPResponse,        for key: String,        staleTTL: TimeInterval,        now: Date,        replaceCachedItems: Bool) -> UsageMergeResult?    {        guard key.hasPrefix("usage:"),              response.status == .ok,              staleTTL > 0,              var items = try? JSONSerialization.jsonObject(with: response.body) as? [[String: Any]]        else {            return nil        }         var cachedItems = replaceCachedItems ? [:] : self.lastGoodUsageItems[key] ?? [:]        if !replaceCachedItems {            cachedItems = cachedItems.filter { now.timeIntervalSince($0.value.recordedAt) <= staleTTL }        }        let itemKeys = items.indices.compactMap { index in            Self.usageItemKey(                items[index],                accountID: Self.cacheAccountKey(at: index, in: response.usageCacheKeys))        }        let duplicateKeys = Set(            Dictionary(grouping: itemKeys, by: { $0 })                .filter { $0.value.count > 1 }                .map(\.key))        for duplicateKey in duplicateKeys {            cachedItems[duplicateKey] = nil        }        var replacedError = false         for index in items.indices {            let item = items[index]            guard let itemKey = Self.usageItemKey(                item,                accountID: Self.cacheAccountKey(at: index, in: response.usageCacheKeys))            else {                continue            }            guard !duplicateKeys.contains(itemKey) else {                continue            }             if Self.hasError(item) {                if let cached = cachedItems[itemKey],                   let cachedItem = try? JSONSerialization.jsonObject(with: cached.data) as? [String: Any]                {                    items[index] = cachedItem                    replacedError = true                }            } else {                if let data = try? JSONSerialization.data(withJSONObject: item, options: [.sortedKeys]) {                    cachedItems[itemKey] = LastGoodUsageItem(recordedAt: now, data: data)                }            }        }        self.lastGoodUsageItems[key] = cachedItems         guard replacedError,              let body = try? JSONSerialization.data(withJSONObject: items, options: [.sortedKeys])        else {            return UsageMergeResult(response: response)        }        return UsageMergeResult(            response: CLILocalHTTPResponse(                status: response.status,                body: body,                contentType: response.contentType,                usageCacheKeys: response.usageCacheKeys))    }     private static func usageItemKey(_ item: [String: Any], accountID: String?) -> UsageItemKey? {        guard let provider = item["provider"] as? String,              !provider.isEmpty,              let accountID,              !accountID.isEmpty        else {            return nil        }        return UsageItemKey(provider: provider, accountID: accountID)    }     private static func cacheAccountKey(at index: Int, in keys: [String?]?) -> String? {        guard let keys, keys.indices.contains(index) else { return nil }        return keys[index]    }     private static func hasError(_ item: [String: Any]) -> Bool {        guard let error = item["error"] else { return false }        return !(error is NSNull)    }     private func mergeLastGoodCostItems(        into response: CLILocalHTTPResponse,        for key: String,        staleTTL: TimeInterval,        now: Date,        replaceCachedItems: Bool) -> CostMergeResult?    {        guard key.hasPrefix("cost:"),              response.status == .ok,              staleTTL > 0,              var items = try? JSONSerialization.jsonObject(with: response.body) as? [[String: Any]]        else {            return nil        }         let providers = items.compactMap { item -> String? in            guard let provider = item["provider"] as? String, !provider.isEmpty else { return nil }            return provider        }        guard providers.count == items.count, Set(providers).count == providers.count else {            return CostMergeResult(response: response)        }         var cachedItems = replaceCachedItems ? [:] : self.lastGoodCostItems[key] ?? [:]        if !replaceCachedItems {            cachedItems = cachedItems.filter { now.timeIntervalSince($0.value.recordedAt) <= staleTTL }        }        for index in items.indices {            let provider = providers[index]            if Self.hasError(items[index]) {                if let cached = cachedItems[provider],                   let cachedItem = try? JSONSerialization.jsonObject(with: cached.data) as? [String: Any]                {                    items[index] = cachedItem                }            } else if let data = try? JSONSerialization.data(withJSONObject: items[index], options: [.sortedKeys]) {                cachedItems[provider] = LastGoodCostItem(recordedAt: now, data: data)            }        }        self.lastGoodCostItems[key] = cachedItems        self.lastGoodCostOrder[key] = providers         guard let body = try? JSONSerialization.data(withJSONObject: items, options: [.sortedKeys]) else {            return CostMergeResult(response: response)        }        return CostMergeResult(response: CLILocalHTTPResponse(            status: response.status,            body: body,            contentType: response.contentType,            usageCacheKeys: response.usageCacheKeys))    }     private func staleCostResponse(        for key: String,        staleTTL: TimeInterval,

Lines 1–500 of 1460.