756 lines · 29.3 KB · Swift
RawDownload
import Foundation enum PiSessionCostScanner {    struct Options {        var piSessionsRoot: URL?        var cacheRoot: URL?        var refreshMinIntervalSeconds: TimeInterval = 60        var forceRescan: Bool = false         init(            piSessionsRoot: URL? = nil,            cacheRoot: URL? = nil,            refreshMinIntervalSeconds: TimeInterval = 60,            forceRescan: Bool = false)        {            self.piSessionsRoot = piSessionsRoot            self.cacheRoot = cacheRoot            self.refreshMinIntervalSeconds = refreshMinIntervalSeconds            self.forceRescan = forceRescan        }    }     private struct ParseResult {        let contributions: [String: [String: [String: PiPackedUsage]]]        let parsedBytes: Int64        let lastModelContext: PiModelContext?    }     private struct AssistantIdentity {        let provider: UsageProvider        let modelName: String    }     private static let costScale = 1_000_000_000.0    private static let maxLineBytes = 16 * 1024 * 1024    private static let maxSafeRoundedInt = Double(Int.max) - 1     static func loadDailyReport(        provider: UsageProvider,        since: Date,        until: Date,        now: Date = Date(),        options: Options = Options()) -> CostUsageDailyReport    {        guard provider == .codex || provider == .claude else {            return CostUsageDailyReport(data: [], summary: nil)        }         let range = CostUsageScanner.CostUsageDayRange(since: since, until: until)        var cache = PiSessionCostCacheIO.load(cacheRoot: options.cacheRoot)        let nowMs = Int64(now.timeIntervalSince1970 * 1000)        let refreshMs = Int64(max(0, options.refreshMinIntervalSeconds) * 1000)        let windowExpanded = self.requestedWindowExpandsCache(range: range, cache: cache)        let shouldRefresh = options.forceRescan            || windowExpanded            || refreshMs == 0            || cache.lastScanUnixMs == 0            || nowMs - cache.lastScanUnixMs > refreshMs         if shouldRefresh {            let root = self.defaultPiSessionsRoot(options: options)            let startCutoff = self.dateFromDayKey(range.scanSinceKey) ?? since            let files = self.listPiSessionFiles(root: root, startCutoffLocal: startCutoff)            let filePathsInScan = Set(files.map(\.path))             for fileURL in files {                self.scanPiSessionFile(                    fileURL: fileURL,                    range: range,                    forceRescan: options.forceRescan || windowExpanded,                    cache: &cache)            }             for key in cache.files.keys where !filePathsInScan.contains(key) {                if let old = cache.files[key] {                    self.applyContributions(                        daysByProvider: &cache.daysByProvider,                        contributions: old.contributions,                        sign: -1)                }                cache.files.removeValue(forKey: key)            }             cache.scanSinceKey = range.scanSinceKey            cache.scanUntilKey = range.scanUntilKey            cache.lastScanUnixMs = nowMs            PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot)        }         return self.buildReport(provider: provider, cache: cache, range: range)    }     private static func requestedWindowExpandsCache(        range: CostUsageScanner.CostUsageDayRange,        cache: PiSessionCostCache) -> Bool    {        guard let cachedSince = cache.scanSinceKey,              let cachedUntil = cache.scanUntilKey        else {            return true        }         if range.scanSinceKey < cachedSince {            return true        }        if range.scanUntilKey > cachedUntil {            return true        }        return false    }     private static func defaultPiSessionsRoot(options: Options) -> URL {        if let override = options.piSessionsRoot { return override }        return FileManager.default.homeDirectoryForCurrentUser            .appendingPathComponent(".pi", isDirectory: true)            .appendingPathComponent("agent", isDirectory: true)            .appendingPathComponent("sessions", isDirectory: true)    }     private static func listPiSessionFiles(root: URL, startCutoffLocal: Date) -> [URL] {        guard FileManager.default.fileExists(atPath: root.path) else { return [] }         let keys: Set<URLResourceKey> = [.isRegularFileKey, .contentModificationDateKey]        guard let enumerator = FileManager.default.enumerator(            at: root,            includingPropertiesForKeys: Array(keys),            options: [.skipsHiddenFiles])        else {            return []        }         var output: [URL] = []        while let item = enumerator.nextObject() as? URL {            guard item.pathExtension.lowercased() == "jsonl" else { continue }            let values = try? item.resourceValues(forKeys: keys)            guard values?.isRegularFile == true else { continue }             let startedAt = self.parseSessionStartFromFilename(item.lastPathComponent)            let modifiedAt = values?.contentModificationDate            if self                .shouldIncludeFile(startedAt: startedAt, modifiedAt: modifiedAt, startCutoffLocal: startCutoffLocal)            {                output.append(item)            }        }         return output.sorted(by: { $0.path < $1.path })    }     private static func shouldIncludeFile(        startedAt: Date?,        modifiedAt: Date?,        startCutoffLocal: Date) -> Bool    {        if let modifiedAt, self.localMidnight(modifiedAt) >= startCutoffLocal {            return true        }        if let startedAt, self.localMidnight(startedAt) >= startCutoffLocal {            return true        }        return false    }     private static func scanPiSessionFile(        fileURL: URL,        range: CostUsageScanner.CostUsageDayRange,        forceRescan: Bool,        cache: inout PiSessionCostCache)    {        let path = fileURL.path        let attrs = (try? FileManager.default.attributesOfItem(atPath: path)) ?? [:]        let mtime = (attrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0        let size = (attrs[.size] as? NSNumber)?.int64Value ?? 0        let mtimeMs = Int64(mtime * 1000)         func storeFileUsage(_ usage: PiSessionFileUsage) {            cache.files[path] = usage        }         let cached = cache.files[path]        if !forceRescan,           let cached,           cached.mtimeUnixMs == mtimeMs,           cached.size == size        {            return        }         if !forceRescan,           let cached,           size > cached.size,           cached.parsedBytes > 0,           cached.parsedBytes <= size        {            let delta = self.parsePiSessionFile(                fileURL: fileURL,                range: range,                startOffset: cached.parsedBytes,                initialModelContext: cached.lastModelContext)            if !delta.contributions.isEmpty {                self.applyContributions(                    daysByProvider: &cache.daysByProvider,                    contributions: delta.contributions,                    sign: 1)            }            let merged = self.mergedContributions(existing: cached.contributions, delta: delta.contributions)            storeFileUsage(PiSessionFileUsage(                mtimeUnixMs: mtimeMs,                size: size,                parsedBytes: delta.parsedBytes,                lastModelContext: delta.lastModelContext,                contributions: merged))            return        }         if let cached {            self.applyContributions(                daysByProvider: &cache.daysByProvider,                contributions: cached.contributions,                sign: -1)        }         let parsed = self.parsePiSessionFile(fileURL: fileURL, range: range)        if !parsed.contributions.isEmpty {            self.applyContributions(daysByProvider: &cache.daysByProvider, contributions: parsed.contributions, sign: 1)        }         storeFileUsage(PiSessionFileUsage(            mtimeUnixMs: mtimeMs,            size: size,            parsedBytes: parsed.parsedBytes,            lastModelContext: parsed.lastModelContext,            contributions: parsed.contributions))    }     private static func parsePiSessionFile(        fileURL: URL,        range: CostUsageScanner.CostUsageDayRange,        startOffset: Int64 = 0,        initialModelContext: PiModelContext? = nil) -> ParseResult    {        var currentModelContext = initialModelContext        var contributions: [String: [String: [String: PiPackedUsage]]] = [:]         func add(provider: UsageProvider, dayKey: String, modelName: String, usage: PiPackedUsage) {            guard !usage.isZero else { return }            guard CostUsageScanner.CostUsageDayRange.isInRange(                dayKey: dayKey,                since: range.scanSinceKey,                until: range.scanUntilKey)            else {                return            }             let providerKey = provider.rawValue            var providerDays = contributions[providerKey] ?? [:]            var dayModels = providerDays[dayKey] ?? [:]            let merged = self.addPacked(a: dayModels[modelName] ?? PiPackedUsage(), b: usage, sign: 1)            if merged.isZero {                dayModels.removeValue(forKey: modelName)            } else {                dayModels[modelName] = merged            }            if dayModels.isEmpty {                providerDays.removeValue(forKey: dayKey)            } else {                providerDays[dayKey] = dayModels            }            if providerDays.isEmpty {                contributions.removeValue(forKey: providerKey)            } else {                contributions[providerKey] = providerDays            }        }         let parsedBytes = (try? CostUsageJsonl.scan(            fileURL: fileURL,            offset: startOffset,            maxLineBytes: Self.maxLineBytes,            prefixBytes: Self.maxLineBytes,            onLine: { line in                guard !line.bytes.isEmpty, !line.wasTruncated else { return }                guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any]                else { return }                guard let type = object["type"] as? String else { return }                 if type == "model_change" {                    currentModelContext = self.modelContext(from: object)                    return                }                 guard type == "message", let message = object["message"] as? [String: Any] else { return }                guard (message["role"] as? String) == "assistant" else { return }                 let identity = self.resolveAssistantIdentity(                    entry: object,                    message: message,                    fallback: currentModelContext)                guard let identity else { return }                guard let date = self.timestampDate(entry: object, message: message) else { return }                let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date)                let usage = self.extractUsage(                    provider: identity.provider,                    modelName: identity.modelName,                    message: message)                add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage)            })) ?? startOffset         return ParseResult(            contributions: contributions,            parsedBytes: parsedBytes,            lastModelContext: currentModelContext)    }     private static func modelContext(from object: [String: Any]) -> PiModelContext? {        guard let providerText = object["provider"] as? String,              let provider = self.mappedProvider(fromPiProvider: providerText)        else {            return nil        }        let rawModelName = (object["modelId"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""        guard let modelName = self.normalizeModelName(rawModelName, provider: provider) else { return nil }        return PiModelContext(providerRawValue: provider.rawValue, modelName: modelName)    }     private static func resolveAssistantIdentity(        entry: [String: Any],        message: [String: Any],        fallback: PiModelContext?) -> AssistantIdentity?    {        let explicitProviderText = self.extractProviderText(entry: entry, message: message)        let explicitProvider = explicitProviderText.flatMap(self.mappedProvider(fromPiProvider:))        let explicitModelText = self.extractModelText(entry: entry, message: message)         if explicitProviderText != nil, explicitProvider == nil {            return nil        }         if let explicitProvider,           let explicitModelText,           let explicitModel = self.normalizeModelName(explicitModelText, provider: explicitProvider)        {            return AssistantIdentity(provider: explicitProvider, modelName: explicitModel)        }         if let explicitProvider,           let fallback,           fallback.providerRawValue == explicitProvider.rawValue        {            return AssistantIdentity(provider: explicitProvider, modelName: fallback.modelName)        }         if explicitProviderText == nil,           let explicitModelText,           let fallbackProvider = fallback.flatMap({ UsageProvider(rawValue: $0.providerRawValue) }),           let explicitModel = self.normalizeModelName(explicitModelText, provider: fallbackProvider)        {            return AssistantIdentity(provider: fallbackProvider, modelName: explicitModel)        }         if explicitProviderText == nil,           let fallback,           let provider = UsageProvider(rawValue: fallback.providerRawValue)        {            return AssistantIdentity(provider: provider, modelName: fallback.modelName)        }         return nil    }     private static func extractProviderText(entry: [String: Any], message: [String: Any]) -> String? {        if let provider = (message["provider"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),           !provider.isEmpty        {            return provider        }        if let provider = (entry["provider"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),           !provider.isEmpty        {            return provider        }        return nil    }     private static func extractModelText(entry: [String: Any], message: [String: Any]) -> String? {        for value in [message["model"], entry["model"], message["modelId"], entry["modelId"]] {            if let model = (value as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), !model.isEmpty {                return model            }        }        return nil    }     private static func normalizeModelName(_ raw: String, provider: UsageProvider) -> String? {        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)        guard !trimmed.isEmpty else { return nil }        return switch provider {        case .codex:            CostUsagePricing.normalizeCodexModel(trimmed)        case .claude:            CostUsagePricing.normalizeClaudeModel(trimmed)        default:            trimmed        }    }     private static func timestampDate(entry: [String: Any], message: [String: Any]) -> Date? {        self.parseTimestampValue(message["timestamp"])            ?? self.parseTimestampValue(entry["timestamp"])    }     private static func parseTimestampValue(_ value: Any?) -> Date? {        if let number = value as? NSNumber {            let raw = number.doubleValue            guard raw.isFinite else { return nil }            if raw > 1_000_000_000_000 {                return Date(timeIntervalSince1970: raw / 1000)            }            return Date(timeIntervalSince1970: raw)        }         if let string = value as? String {            if let numeric = Double(string), numeric.isFinite {                if numeric > 1_000_000_000_000 {                    return Date(timeIntervalSince1970: numeric / 1000)                }                return Date(timeIntervalSince1970: numeric)            }            return self.parseISO(string)        }         return nil    }     private static func extractUsage(        provider: UsageProvider,        modelName: String,        message: [String: Any]) -> PiPackedUsage    {        let usage = (message["usage"] as? [String: Any]) ?? [:]        let input = self.readNonNegativeInt(            usage["input"]                ?? usage["inputTokens"]                ?? usage["input_tokens"]                ?? usage["promptTokens"]                ?? usage["prompt_tokens"])        let cacheRead = self.readNonNegativeInt(            usage["cacheRead"]                ?? usage["cacheReadTokens"]                ?? usage["cache_read"]                ?? usage["cache_read_tokens"]                ?? usage["cacheReadInputTokens"]                ?? usage["cache_read_input_tokens"])        let cacheWrite = self.readNonNegativeInt(            usage["cacheWrite"]                ?? usage["cacheWriteTokens"]                ?? usage["cache_write"]                ?? usage["cache_write_tokens"]                ?? usage["cacheCreationTokens"]                ?? usage["cache_creation_tokens"]                ?? usage["cacheCreationInputTokens"]                ?? usage["cache_creation_input_tokens"])        let output = self.readNonNegativeInt(            usage["output"]                ?? usage["outputTokens"]                ?? usage["output_tokens"]                ?? usage["completionTokens"]                ?? usage["completion_tokens"])         let directTotal = self.readNonNegativeInt(            usage["totalTokens"]                ?? usage["total_tokens"]                ?? usage["tokenCount"]                ?? usage["token_count"]                ?? usage["tokens"])        let derivedTotal = input + cacheRead + cacheWrite + output        let totalTokens = max(directTotal, derivedTotal)         let rawUsage = PiPackedUsage(            inputTokens: input,            cacheReadTokens: cacheRead,            cacheWriteTokens: cacheWrite,            outputTokens: output,            totalTokens: totalTokens)        let costUSD = self.computedCostUSD(provider: provider, modelName: modelName, usage: rawUsage)        let costNanos = costUSD.map { Int64(($0 * self.costScale).rounded()) } ?? 0         return PiPackedUsage(            inputTokens: rawUsage.inputTokens,            cacheReadTokens: rawUsage.cacheReadTokens,            cacheWriteTokens: rawUsage.cacheWriteTokens,            outputTokens: rawUsage.outputTokens,            totalTokens: rawUsage.totalTokens,            costNanos: costNanos,            costSampleCount: costUSD == nil ? 0 : 1)    }     private static func computedCostUSD(        provider: UsageProvider,        modelName: String,

Lines 1–500 of 756.