154 lines · 5.2 KB · Swift
RawDownload
import Foundation enum CostUsageCacheIO {    private static let compatibleCodexProducerKeys: Set<String> = [        "codex:cu:p3c27f997569eb3c5",    ]     private static func artifactVersion(for provider: UsageProvider) -> Int {        switch provider {        case .codex:            8        case .claude, .vertexai:            4        default:            1        }    }     private static func defaultCacheRoot() -> URL {        let root = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!        return root.appendingPathComponent("CodexBar", isDirectory: true)    }     static func cacheFileURL(provider: UsageProvider, cacheRoot: URL? = nil) -> URL {        let root = cacheRoot ?? self.defaultCacheRoot()        let artifactVersion = self.artifactVersion(for: provider)        return root            .appendingPathComponent("cost-usage", isDirectory: true)            .appendingPathComponent("\(provider.rawValue)-v\(artifactVersion).json", isDirectory: false)    }     static func load(        provider: UsageProvider,        cacheRoot: URL? = nil,        producerKey: String? = nil) -> CostUsageCache    {        let url = self.cacheFileURL(provider: provider, cacheRoot: cacheRoot)        let expectedProducerKey = producerKey ?? self.currentProducerKey(provider: provider)        let compatibleProducerKeys = producerKey == nil && provider == .codex            ? self.compatibleCodexProducerKeys            : []        if let decoded = self.loadCache(            at: url,            expectedProducerKey: expectedProducerKey,            compatibleProducerKeys: compatibleProducerKeys)        {            return decoded        }        return CostUsageCache()    }     private static func loadCache(        at url: URL,        expectedProducerKey: String?,        compatibleProducerKeys: Set<String>) -> CostUsageCache?    {        guard let data = try? Data(contentsOf: url) else { return nil }        guard let decoded = try? JSONDecoder().decode(CostUsageCache.self, from: data)        else { return nil }        guard decoded.version == 1 else { return nil }        if let expectedProducerKey {            guard decoded.producerKey == expectedProducerKey                || decoded.producerKey.map(compatibleProducerKeys.contains) == true            else { return nil }        }        return decoded    }     static func save(        provider: UsageProvider,        cache: CostUsageCache,        cacheRoot: URL? = nil,        producerKey: String? = nil)    {        let url = self.cacheFileURL(provider: provider, cacheRoot: cacheRoot)        let dir = url.deletingLastPathComponent()        try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)         var cache = cache        cache.producerKey = producerKey ?? self.currentProducerKey(provider: provider)         let tmp = dir.appendingPathComponent(".tmp-\(UUID().uuidString).json", isDirectory: false)        let data = (try? JSONEncoder().encode(cache)) ?? Data()        do {            try data.write(to: tmp, options: [.atomic])            if FileManager.default.fileExists(atPath: url.path) {                _ = try FileManager.default.replaceItemAt(url, withItemAt: tmp)            } else {                try FileManager.default.moveItem(at: tmp, to: url)            }        } catch {            try? FileManager.default.removeItem(at: tmp)        }    }     static func currentProducerKey(        provider: UsageProvider,        parserHash: String = CodexParserHash.value) -> String?    {        guard provider == .codex else { return nil }        return "\(provider.rawValue):cu:p\(parserHash)"    }} struct CostUsageCache: Codable {    var version: Int = 1    var producerKey: String?    var lastScanUnixMs: Int64 = 0    var scanSinceKey: String?    var scanUntilKey: String?    var codexPricingKey: String?    var codexPriorityMetadataKey: String?    var codexPriorityTurnKeys: [String: String]?    var codexPriorityTurnIDsByDay: [String: [String]]?     /// filePath -> file usage    var files: [String: CostUsageFileUsage] = [:]     /// dayKey -> model -> packed usage    var days: [String: [String: [Int]]] = [:]     /// rootPath -> mtime (for Claude roots)    var roots: [String: Int64]?} struct CostUsageFileUsage: Codable {    var mtimeUnixMs: Int64    var size: Int64    var days: [String: [String: [Int]]]    var parsedBytes: Int64?    var lastModel: String?    var lastTotals: CostUsageCodexTotals?    var lastCountedTotals: CostUsageCodexTotals?    var lastRawTotalsBaseline: CostUsageCodexTotals?    var hasDivergentTotals: Bool?    var lastCodexTurnID: String?    var sessionId: String?    var forkedFromId: String?    var codexCostNanos: [String: [String: Int64]]?    var codexPrioritySurchargeNanos: [String: [String: Int64]]?    var codexStandardCostNanos: [String: [String: Int64]]?    var codexPriorityCostNanos: [String: [String: Int64]]?    var codexStandardTokens: [String: [String: Int]]?    var codexPriorityTokens: [String: [String: Int]]?    var codexTurnIDs: [String]?    var codexRows: [CostUsageScanner.CodexUsageRow]?    var claudeRows: [CostUsageScanner.ClaudeUsageRow]?} struct CostUsageCodexTotals: Codable {    var input: Int    var cached: Int    var output: Int}