943 lines · 37.1 KB · Swift
RawDownload
import Foundation private final class PiSessionISO8601FormatterBox: @unchecked Sendable {    let lock = NSLock()    let withFractional: ISO8601DateFormatter = {        let formatter = ISO8601DateFormatter()        formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]        return formatter    }()     let plain: ISO8601DateFormatter = {        let formatter = ISO8601DateFormatter()        formatter.formatOptions = [.withInternetDateTime]        return formatter    }()} 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 struct ModelsDevPricingContext {        let catalog: ModelsDevCatalog?        let cacheRoot: URL?        let pricingKey: String    }     private struct ScanContext {        let range: CostUsageScanner.CostUsageDayRange        let forceRescan: Bool        let pricingContext: ModelsDevPricingContext        let checkCancellation: CostUsageScanner.CancellationCheck?    }     private static let costScale = 1_000_000_000.0    /// Bump for Pi-only cost formula changes not represented by the parser or pricing fingerprints.    private static let costFormulaVersion = 1    private static let maxLineBytes = 16 * 1024 * 1024    private static let maxSafeRoundedInt = Double(Int.max) - 1    private static let sessionStartFilenameRegex = try? NSRegularExpression(        pattern: "^(\\d{4}-\\d{2}-\\d{2})T(\\d{2})-(\\d{2})-(\\d{2})-(\\d{3})Z_")    private static let isoFormatterBox = PiSessionISO8601FormatterBox()     static func loadDailyReport(        provider: UsageProvider,        since: Date,        until: Date,        now: Date = Date(),        options: Options = Options()) -> CostUsageDailyReport    {        (            try? self.loadDailyReportCancellable(                provider: provider,                since: since,                until: until,                now: now,                options: options,                checkCancellation: nil)) ?? CostUsageDailyReport(data: [], summary: nil)    }     static func loadDailyReportCancellable(        provider: UsageProvider,        since: Date,        until: Date,        now: Date = Date(),        options: Options = Options(),        checkCancellation: CostUsageScanner.CancellationCheck?) throws -> 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 pricingContext = self.pricingContext(now: now, cacheRoot: options.cacheRoot)        let windowExpanded = self.requestedWindowExpandsCache(range: range, cache: cache)        let pricingChanged = cache.pricingKey != pricingContext.pricingKey        let shouldRefresh = options.forceRescan            || windowExpanded            || pricingChanged            || refreshMs == 0            || cache.lastScanUnixMs == 0            || nowMs - cache.lastScanUnixMs > refreshMs         if shouldRefresh {            try checkCancellation?()            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 {                try self.scanPiSessionFile(                    fileURL: fileURL,                    cache: &cache,                    context: ScanContext(                        range: range,                        forceRescan: options.forceRescan || windowExpanded || pricingChanged,                        pricingContext: pricingContext,                        checkCancellation: checkCancellation))            }            try checkCancellation?()             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.pricingKey = pricingContext.pricingKey            cache.lastScanUnixMs = nowMs            try checkCancellation?()            PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot)        }         return self.buildReport(            provider: provider,            cache: cache,            range: range,            pricingContext: pricingContext)    }     struct CachedDailyReportResult {        let report: CostUsageDailyReport        let lastScanAt: Date?    }     static func loadCachedDailyReport(        provider: UsageProvider,        since: Date,        until: Date,        now: Date = Date(),        cacheRoot: URL? = nil) -> CostUsageDailyReport?    {        self.loadCachedDailyReportResult(            provider: provider,            since: since,            until: until,            now: now,            cacheRoot: cacheRoot)?.report    }     static func loadCachedDailyReportResult(        provider: UsageProvider,        since: Date,        until: Date,        now: Date = Date(),        cacheRoot: URL? = nil) -> CachedDailyReportResult?    {        guard provider == .codex || provider == .claude else { return nil }         let range = CostUsageScanner.CostUsageDayRange(since: since, until: until)        let cache = PiSessionCostCacheIO.load(cacheRoot: cacheRoot)        guard !cache.daysByProvider.isEmpty else { return nil }        guard !self.requestedWindowExpandsCache(range: range, cache: cache) else { return nil }         let pricingContext = self.pricingContext(now: now, cacheRoot: cacheRoot)        guard cache.pricingKey == pricingContext.pricingKey else { return nil }        let report = self.buildReport(            provider: provider,            cache: cache,            range: range,            pricingContext: pricingContext)        guard !report.data.isEmpty else { return nil }        let lastScanAt = cache.lastScanUnixMs > 0            ? Date(timeIntervalSince1970: TimeInterval(cache.lastScanUnixMs) / 1000)            : nil        return CachedDailyReportResult(report: report, lastScanAt: lastScanAt)    }     private static func pricingContext(now: Date, cacheRoot: URL?) -> ModelsDevPricingContext {        let modelsDevArtifact = ModelsDevCache.load(now: now, cacheRoot: cacheRoot).artifact        return ModelsDevPricingContext(            catalog: modelsDevArtifact?.catalog,            cacheRoot: cacheRoot,            pricingKey: CostUsagePricingKey.codex(                modelsDevArtifact: modelsDevArtifact,                formulaVersion: Self.costFormulaVersion,                parserHash: CodexParserHash.value,                modelsDevProviderIDs: ["anthropic", "openai"]))    }     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,        cache: inout PiSessionCostCache,        context: ScanContext)        throws    {        try context.checkCancellation?()        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 !context.forceRescan,           let cached,           cached.mtimeUnixMs == mtimeMs,           cached.size == size        {            return        }         if !context.forceRescan,           let cached,           size > cached.size,           cached.parsedBytes > 0,           cached.parsedBytes <= size        {            let delta = try self.parsePiSessionFile(                fileURL: fileURL,                range: context.range,                startOffset: cached.parsedBytes,                initialModelContext: cached.lastModelContext,                pricingContext: context.pricingContext,                checkCancellation: context.checkCancellation)            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 = try self.parsePiSessionFile(            fileURL: fileURL,            range: context.range,            pricingContext: context.pricingContext,            checkCancellation: context.checkCancellation)        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,        pricingContext: ModelsDevPricingContext? = nil,        checkCancellation: CostUsageScanner.CancellationCheck? = nil) throws -> 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: Int64        do {            parsedBytes = try CostUsageJsonl.scan(                fileURL: fileURL,                offset: startOffset,                maxLineBytes: Self.maxLineBytes,                prefixBytes: Self.maxLineBytes,                checkCancellation: checkCancellation,                onLine: { line in                    guard !line.bytes.isEmpty, !line.wasTruncated else { return }                    autoreleasepool {                        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,                            pricingDate: date,                            pricingContext: pricingContext)                        add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage)                    }                })        } catch is CancellationError {            throw CancellationError()        } catch {            parsedBytes = 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,

Lines 1–500 of 943.