787 lines · 27.4 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("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?} enum CLIServeRoute: Equatable {    case health    case usage(provider: String?)    case cost(provider: String?)} 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)        default:            throw CLIServeRouteError.notFound        }    }} private struct ServeErrorPayload: Encodable {    let error: String} private struct ServeHealthPayload: Encodable {    let status: String} struct CLIServeConfigSnapshot {    let config: CodexBarConfig    let cacheToken: String} private final class CLIServeDeadlineState: @unchecked Sendable {    private let lock = NSLock()    private var continuation: CheckedContinuation<CLILocalHTTPResponse, Never>?    private var workTask: Task<Void, Never>?    private var timeoutTask: Task<Void, Never>?     init(continuation: CheckedContinuation<CLILocalHTTPResponse, Never>) {        self.continuation = continuation    }     func setWorkTask(_ task: Task<Void, Never>) {        var shouldCancel = false        self.lock.lock()        if self.continuation == nil {            shouldCancel = true        } else {            self.workTask = task        }        self.lock.unlock()         if shouldCancel {            task.cancel()        }    }     func setTimeoutTask(_ task: Task<Void, Never>) {        var shouldCancel = false        self.lock.lock()        if self.continuation == nil {            shouldCancel = true        } else {            self.timeoutTask = task        }        self.lock.unlock()         if shouldCancel {            task.cancel()        }    }     func finish(_ response: CLILocalHTTPResponse, cancelWork: Bool, cancelTimeout: Bool) {        let continuation: CheckedContinuation<CLILocalHTTPResponse, Never>?        let workTask: Task<Void, Never>?        let timeoutTask: Task<Void, Never>?         self.lock.lock()        continuation = self.continuation        self.continuation = nil        workTask = cancelWork ? self.workTask : nil        timeoutTask = cancelTimeout ? self.timeoutTask : nil        self.workTask = nil        self.timeoutTask = nil        self.lock.unlock()         workTask?.cancel()        timeoutTask?.cancel()        continuation?.resume(returning: response)    }} enum CLIServeCacheLookup {    case response(CLILocalHTTPResponse)    case miss} actor CLIServeResponseCache {    static let maximumStaleTTL: TimeInterval = 3600     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 var entries: [String: Entry] = [:]    private var lastGood: [String: LastGoodEntry] = [:]    private var lastGoodUsageItems: [String: [UsageItemKey: LastGoodUsageItem]] = [:]    private var inFlightKeys: Set<String> = []    private var waiters: [String: [CheckedContinuation<CLIServeCacheLookup, Never>]] = [:]     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        }    }     private func response(for key: String) -> CLILocalHTTPResponse? {        guard let entry = self.entries[key] else { return nil }        return entry.response    }     func responseOrStartFetch(for key: String, now: Date) async -> CLIServeCacheLookup {        self.pruneExpiredEntries(now: now)        if let cached = self.response(for: key) {            return .response(cached)        }         if self.inFlightKeys.contains(key) {            return await withCheckedContinuation { continuation in                self.waiters[key, default: []].append(continuation)            }        }         self.inFlightKeys.insert(key)        return .miss    }     /// Completes an in-flight fetch and returns the response delivered to    /// waiters. 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)        if shouldCache {            self.store(response, for: key, ttl: policy.ttl, now: now)            if key.hasPrefix("usage:") {                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 {            delivered = staleResponse ?? response        }        self.inFlightKeys.remove(key)        let waiters = self.waiters.removeValue(forKey: key) ?? []        for waiter in waiters {            waiter.resume(returning: .response(delivered))        }        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 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 store(_ response: CLILocalHTTPResponse, for key: String, ttl: TimeInterval, now: Date) {        guard ttl > 0, response.status == .ok else { return }        self.entries[key] = Entry(expiresAt: now.addingTimeInterval(ttl), response: response)    }     func cachedEntryCount() -> Int {        self.entries.count    }     func cachedStaleVariantCount() -> Int {        self.lastGood.count + self.lastGoodUsageItems.count    }} private enum CLIServeArgumentError: LocalizedError {    case invalidPort    case invalidRefreshInterval    case invalidRequestTimeout    case invalidProvider(String)     var errorDescription: String? {        switch self {        case .invalidPort:            "--port must be between 1 and 65535."        case .invalidRefreshInterval:            "--refresh-interval must be zero or greater."        case .invalidRequestTimeout:            "--request-timeout must be zero or greater."        case let .invalidProvider(provider):            "Unknown provider '\(provider)'."        }    }} extension CodexBarCLI {    static let defaultServeRequestTimeout: TimeInterval = 30     static func runServe(_ values: ParsedValues) async {        let output = CLIOutputPreferences(format: .json, jsonOnly: true, pretty: false)        let port = Self.decodeServePort(from: values)        let refreshInterval = Self.decodeServeRefreshInterval(from: values)        let requestTimeout = Self.decodeServeRequestTimeout(from: values)         guard let port else {            Self.exit(                code: .failure,                message: CLIServeArgumentError.invalidPort.localizedDescription,                output: output,                kind: .args)        }         guard let refreshInterval else {            Self.exit(                code: .failure,                message: CLIServeArgumentError.invalidRefreshInterval.localizedDescription,                output: output,                kind: .args)        }         guard let requestTimeout else {            Self.exit(                code: .failure,                message: CLIServeArgumentError.invalidRequestTimeout.localizedDescription,                output: output,                kind: .args)        }         let configStore = CodexBarConfigStore()        let cache = CLIServeResponseCache()        let server = CLILocalHTTPServer(host: "127.0.0.1", port: port) { request in            await Self.handleServeRequest(                request,                configStore: configStore,                cache: cache,                refreshInterval: refreshInterval,                requestTimeout: requestTimeout)        }        let signalMonitor = CLITerminationSignalMonitor { _ in            TTYCommandRunner.terminateActiveProcessesForAppShutdown()            server.stop()        }        defer { signalMonitor.cancel() }         do {            try await server.run {                Self.writeStderr("CodexBar server listening on http://127.0.0.1:\(port)\n")            }        } catch {            await Self.shutdownServeSessions()            Self.exit(code: .failure, message: error.localizedDescription, output: output, kind: .runtime)        }        await Self.shutdownServeSessions()    }     private static func shutdownServeSessions() async {        await ProviderCLISessionLifecycle.shutdownPersistentSessions()        TTYCommandRunner.terminateActiveProcessesForAppShutdown()    }     static func decodeServePort(from values: ParsedValues) -> UInt16? {        let raw = values.options["port"]?.last        let parsed: Int        if let raw {            guard let value = Int(raw) else { return nil }            parsed = value        } else {            parsed = 8080        }        guard parsed > 0, parsed <= Int(UInt16.max) else { return nil }        return UInt16(parsed)    }     static func decodeServeRefreshInterval(from values: ParsedValues) -> TimeInterval? {        let raw = values.options["refreshInterval"]?.last        let parsed: Double        if let raw {            guard let value = Double(raw) else { return nil }            parsed = value        } else {            parsed = 60        }        guard parsed.isFinite, parsed >= 0 else { return nil }        return parsed    }     static func decodeServeRequestTimeout(from values: ParsedValues) -> TimeInterval? {        let raw = values.options["requestTimeout"]?.last        let parsed: Double        if let raw {            guard let value = Double(raw) else { return nil }            parsed = value        } else {            parsed = Self.defaultServeRequestTimeout

Lines 1–500 of 787.