diff --git a/Sources/CodexBarCore/CostUsageFetcher.swift b/Sources/CodexBarCore/CostUsageFetcher.swift index cc58477c525..473dd35923d 100644 --- a/Sources/CodexBarCore/CostUsageFetcher.swift +++ b/Sources/CodexBarCore/CostUsageFetcher.swift @@ -1,3 +1,4 @@ +// swiftlint:disable file_length import Foundation public enum CostUsageError: LocalizedError, Sendable { @@ -425,14 +426,43 @@ public struct CostUsageFetcher: Sendable { let clampedHistoryDays = max(1, min(365, historyDays)) - if let remoteSnapshot = try await self.loadRemoteTokenSnapshot( + var remoteSnapshot: CostUsageTokenSnapshot? + var remoteError: Error? + // Provider-specific by design: Cursor may fall back to local CSV when its remote dashboard is unavailable. + do { + remoteSnapshot = try await self.loadRemoteTokenSnapshot( + provider: provider, + environment: environment, + now: now, + historyDays: clampedHistoryDays, + cursorCookieHeaderOverride: cursorCookieHeaderOverride) + } catch { + if provider != .cursor { + throw error + } + remoteError = error + } + if let remoteSnapshot { + return remoteSnapshot + } + + // Provider-specific by design: Cursor and Antigravity local readers backfill providers without remote history. + let fallbackCalendar = Self.resolvedScannerOptions( + overrideScannerOptions, provider: provider, - environment: environment, - now: now, - historyDays: clampedHistoryDays, - cursorCookieHeaderOverride: cursorCookieHeaderOverride) + codexHomePath: codexHomePath).calendar + if provider == .cursor, let local = await self.loadCursorLocalSnapshot( + now: now, historyDays: clampedHistoryDays, calendar: fallbackCalendar) { - return remoteSnapshot + return local + } + if let remoteError { + throw remoteError + } + if provider == .antigravity, let local = await self.loadAntigravityLocalSnapshot( + now: now, historyDays: clampedHistoryDays, calendar: fallbackCalendar) + { + return local } var options = Self.resolvedScannerOptions( @@ -1120,6 +1150,74 @@ public struct CostUsageFetcher: Sendable { } #endif + private static func loadCursorLocalSnapshot( + now: Date, + historyDays: Int, + calendar: Calendar = .current) async -> CostUsageTokenSnapshot? + { + let paths = CursorLocalCSVReader.cachedCSVPaths() + guard !paths.isEmpty else { return nil } + var allRows: [CursorLocalCSVReader.Row] = [] + for url in paths { + allRows.append(contentsOf: CursorLocalCSVReader.parseFile(at: url)) + } + guard !allRows.isEmpty else { return nil } + let full = CursorLocalCSVReader.makeDailyReport(from: allRows, calendar: calendar, now: now) + let cal = calendar + let since = cal.date(byAdding: .day, value: -(historyDays - 1), to: cal.startOfDay(for: now)) ?? now + let sinceKey = CostUsageLocalDay.key(from: since, calendar: cal) + let nowKey = CostUsageLocalDay.key(from: now, calendar: cal) + let filtered = full.data.filter { $0.date >= sinceKey && $0.date <= nowKey } + guard !filtered.isEmpty else { return nil } + let costValues = filtered.compactMap(\.costUSD) + let totalCost: Double? = costValues.isEmpty ? nil : costValues.reduce(0, +) + let totalTokens = filtered.compactMap(\.totalTokens).reduce(0, +) + let filteredSummary: CostUsageDailyReport.Summary = .init( + totalInputTokens: nil, + totalOutputTokens: nil, + totalTokens: totalTokens, + totalCostUSD: totalCost) + let daily = CostUsageDailyReport(data: filtered, summary: filteredSummary) + return Self.tokenSnapshot( + from: daily, + now: now, + historyDays: historyDays, + useCurrentLocalDayForSession: true, + calendar: cal, + costProvenance: .listPriceEstimate) + } + + private static func loadAntigravityLocalSnapshot( + now: Date, + historyDays: Int, + calendar: Calendar = .current) async -> CostUsageTokenSnapshot? + { + let cal = calendar + let report = AntigravityLocalReader.makeDailyReport(calendar: cal) + guard !report.data.isEmpty else { return nil } + let since = cal.date(byAdding: .day, value: -(historyDays - 1), to: cal.startOfDay(for: now)) ?? now + let sinceKey = CostUsageLocalDay.key(from: since, calendar: cal) + let nowKey = CostUsageLocalDay.key(from: now, calendar: cal) + let filtered = report.data.filter { $0.date >= sinceKey && $0.date <= nowKey } + guard !filtered.isEmpty else { return nil } + let costValues = filtered.compactMap(\.costUSD) + let totalCost: Double? = costValues.isEmpty ? nil : costValues.reduce(0, +) + let totalTokens = filtered.compactMap(\.totalTokens).reduce(0, +) + let filteredSummary: CostUsageDailyReport.Summary = .init( + totalInputTokens: nil, + totalOutputTokens: nil, + totalTokens: totalTokens, + totalCostUSD: totalCost) + let daily = CostUsageDailyReport(data: filtered, summary: filteredSummary) + return Self.tokenSnapshot( + from: daily, + now: now, + historyDays: historyDays, + useCurrentLocalDayForSession: true, + calendar: cal, + costProvenance: .unknown) + } + static func tokenSnapshot( from daily: CostUsageDailyReport, now: Date, diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalReader.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalReader.swift new file mode 100644 index 00000000000..57f1fdb8961 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalReader.swift @@ -0,0 +1,208 @@ +import Foundation + +// MARK: - Antigravity Local Reader (tokscale compatible) + +enum AntigravityLocalReader { + static func tokiCachePaths(home: URL? = nil) -> [URL] { + let base: URL = if let home { + home.appendingPathComponent(".config/tokscale/antigravity-cache/sessions", isDirectory: true) + } else if let dir = ProcessInfo.processInfo.environment["TOKSCALE_CONFIG_DIR"] { + URL(fileURLWithPath: dir).appendingPathComponent("antigravity-cache/sessions", isDirectory: true) + } else { + FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".config/tokscale/antigravity-cache/sessions", isDirectory: true) + } + guard let contents = try? FileManager.default.contentsOfDirectory( + at: base, includingPropertiesForKeys: nil, options: [.skipsHiddenFiles]) else { return [] } + return contents.filter { $0.pathExtension == "jsonl" }.sorted { $0.lastPathComponent < $1.lastPathComponent } + } + + static func cliDBPaths(home: URL? = nil) -> [URL] { + let base: URL = if let env = ProcessInfo.processInfo.environment["GEMINI_CLI_HOME"], !env.isEmpty { + URL(fileURLWithPath: env).appendingPathComponent("antigravity-cli/conversations", isDirectory: true) + } else if let home { + home.appendingPathComponent(".gemini/antigravity-cli/conversations", isDirectory: true) + } else { + FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".gemini/antigravity-cli/conversations", isDirectory: true) + } + guard let c = try? FileManager.default.contentsOfDirectory( + at: base, includingPropertiesForKeys: nil, options: [.skipsHiddenFiles]) else { return [] } + return c + .filter { + $0.pathExtension == "db" && !$0.lastPathComponent.hasSuffix("-wal") && !$0.lastPathComponent + .hasSuffix("-shm") + } + .sorted { $0.lastPathComponent < $1.lastPathComponent } + } + + static func parseJSONLCache(paths: [URL]? = nil, calendar: Calendar = .current) -> [CostUsageDailyReport.Entry] { + var entries: [CostUsageDailyReport.Entry] = [] + var seenResponseIds = Set() + for url in paths ?? self.tokiCachePaths() { + guard let data = try? Data(contentsOf: url), + let text = String(data: data, encoding: .utf8) else { continue } + var sessionModel: String? + for line in text.components(separatedBy: .newlines) { + let t = line.trimmingCharacters(in: .whitespacesAndNewlines) + guard !t.isEmpty, let d = t.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: d) as? [String: Any] else { continue } + let type = json["type"] as? String + if type == "session_meta" { + sessionModel = json["modelId"] as? String ?? json["model_id"] as? String + continue + } + if type == "usage" || json["input"] != nil { + if let rid = (json["responseId"] as? String ?? json["response_id"] as? String), !rid.isEmpty { + if seenResponseIds.contains(rid) { continue } + seenResponseIds.insert(rid) + } + let modelId = + (json["modelId"] as? String ?? json["model_id"] as? String ?? sessionModel ?? "unknown") + let input = (json["input"] as? Int) ?? 0 + let output = (json["output"] as? Int) ?? 0 + let read = (json["cacheRead"] as? Int) ?? (json["cache_read"] as? Int) ?? 0 + let write = (json["cacheWrite"] as? Int) ?? (json["cache_write"] as? Int) ?? 0 + let reason = (json["reasoning"] as? Int) ?? 0 + let ts = (json["timestamp"] as? Int64) ?? (json["timestamp"] as? Int).map(Int64.init) ?? 0 + let date = self.timestampToDayKey(ts, calendar: calendar) + let total = input + output + read + write + if total == 0 { continue } + if let idx = entries.firstIndex(where: { $0.date == date }) { + var e = entries[idx] + let ne = CostUsageDailyReport.Entry( + date: e.date, + inputTokens: (e.inputTokens ?? 0) + input, + outputTokens: (e.outputTokens ?? 0) + output, + cacheReadTokens: (e.cacheReadTokens ?? 0) + read, + cacheCreationTokens: (e.cacheCreationTokens ?? 0) + write, + reasoningTokens: (e.reasoningTokens ?? 0) + reason, + totalTokens: (e.totalTokens ?? 0) + total, + requestCount: (e.requestCount ?? 0) + 1, + costUSD: e.costUSD, + modelsUsed: nil, + modelBreakdowns: self.mergeBreakdown(e.modelBreakdowns, model: modelId, tokens: total)) + entries[idx] = ne + } else { + let e = CostUsageDailyReport.Entry( + date: date, + inputTokens: input, + outputTokens: output, + cacheReadTokens: read, + cacheCreationTokens: write, + reasoningTokens: reason, + totalTokens: total, + requestCount: 1, + costUSD: nil, + modelsUsed: nil, + modelBreakdowns: [ + CostUsageDailyReport.ModelBreakdown( + modelName: modelId, + costUSD: nil, + totalTokens: total, + requestCount: 1), + ]) + entries.append(e) + } + } + } + } + return entries.sorted { $0.date < $1.date } + } + + static func parseCLIDBs() -> [CostUsageDailyReport.Entry] { + [] + } + + static func makeDailyReport(calendar: Calendar = .current) -> CostUsageDailyReport { + var merged: [String: CostUsageDailyReport.Entry] = [:] + for e in self.parseJSONLCache(calendar: calendar) + self.parseCLIDBs() { + if var ex = merged[e.date] { + let mergedCost: Double? = { + if ex.costUSD == nil, e.costUSD == nil { return nil } + return (ex.costUSD ?? 0) + (e.costUSD ?? 0) + }() + let ne = CostUsageDailyReport.Entry( + date: ex.date, + inputTokens: (ex.inputTokens ?? 0) + (e.inputTokens ?? 0), + outputTokens: (ex.outputTokens ?? 0) + (e.outputTokens ?? 0), + cacheReadTokens: (ex.cacheReadTokens ?? 0) + (e.cacheReadTokens ?? 0), + cacheCreationTokens: (ex.cacheCreationTokens ?? 0) + (e.cacheCreationTokens ?? 0), + reasoningTokens: (ex.reasoningTokens ?? 0) + (e.reasoningTokens ?? 0), + totalTokens: (ex.totalTokens ?? 0) + (e.totalTokens ?? 0), + requestCount: (ex.requestCount ?? 0) + (e.requestCount ?? 0), + costUSD: mergedCost, + modelsUsed: nil, + modelBreakdowns: self.mergeBreakdowns(ex.modelBreakdowns, e.modelBreakdowns)) + merged[e.date] = ne + } else { + merged[e.date] = e + } + } + let sorted = merged.values.sorted { $0.date < $1.date } + let costValues = sorted.compactMap(\.costUSD) + let totalCost: Double? = costValues.isEmpty ? nil : costValues.reduce(0, +) + let totalTokens = sorted.compactMap(\.totalTokens).reduce(0, +) + let summary: CostUsageDailyReport.Summary? = sorted.isEmpty ? nil : .init( + totalInputTokens: nil, + totalOutputTokens: nil, + totalTokens: totalTokens, + totalCostUSD: totalCost) + return CostUsageDailyReport(data: sorted, summary: summary) + } + + private static func timestampToDayKey(_ ms: Int64, calendar: Calendar = .current) -> String { + let sec = Double(ms) / 1000.0 + let date = Date(timeIntervalSince1970: sec) + let c = calendar + let comps = c.dateComponents([.year, .month, .day], from: date) + return String(format: "%04d-%02d-%02d", comps.year ?? 1970, comps.month ?? 1, comps.day ?? 1) + } + + private static func mergeBreakdown( + _ ex: [CostUsageDailyReport.ModelBreakdown]?, + model: String, + tokens: Int) -> [CostUsageDailyReport.ModelBreakdown] + { + var arr = ex ?? [] + if let i = arr.firstIndex(where: { $0.modelName == model }) { + let b = arr[i] + arr[i] = CostUsageDailyReport.ModelBreakdown( + modelName: b.modelName, + costUSD: b.costUSD, + totalTokens: (b.totalTokens ?? 0) + tokens, + requestCount: (b.requestCount ?? 0) + 1) + } else { + arr.append(CostUsageDailyReport.ModelBreakdown( + modelName: model, + costUSD: nil, + totalTokens: tokens, + requestCount: 1)) + } + return arr + } + + private static func mergeBreakdowns( + _ a: [CostUsageDailyReport.ModelBreakdown]?, + _ b: [CostUsageDailyReport.ModelBreakdown]?) -> [CostUsageDailyReport.ModelBreakdown]? + { + var d: [String: CostUsageDailyReport.ModelBreakdown] = [:] + for m in (a ?? []) + (b ?? []) { + if var ex = d[m.modelName] { + let mergedCost: Double? = { + if ex.costUSD == nil, m.costUSD == nil { return nil } + return (ex.costUSD ?? 0) + (m.costUSD ?? 0) + }() + ex = CostUsageDailyReport.ModelBreakdown( + modelName: ex.modelName, + costUSD: mergedCost, + totalTokens: (ex.totalTokens ?? 0) + (m.totalTokens ?? 0), + requestCount: (ex.requestCount ?? 0) + (m.requestCount ?? 0)) + d[m.modelName] = ex + } else { + d[m.modelName] = m + } + } + return d.isEmpty ? nil : Array(d.values) + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift index e748ce0c25c..2608312381d 100644 --- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift @@ -49,8 +49,9 @@ public enum AntigravityProviderDescriptor { ProviderColor(hex: 0xFBBC04), ]), tokenCost: ProviderTokenCostConfig( - supportsTokenCost: false, - noDataMessage: { "Antigravity cost summary is not supported." }), + supportsTokenCost: true, + noDataMessage: { "Antigravity cost summary is not supported." }, + supportsTokenSnapshot: true), pace: ProviderPaceCapability( sessionPaceWindowRule: .custom { window, _ in window.windowMinutes == nil || window.windowMinutes == 300 diff --git a/Sources/CodexBarCore/Providers/Cursor/CursorLocalCSVReader.swift b/Sources/CodexBarCore/Providers/Cursor/CursorLocalCSVReader.swift new file mode 100644 index 00000000000..c7b73ddb0de --- /dev/null +++ b/Sources/CodexBarCore/Providers/Cursor/CursorLocalCSVReader.swift @@ -0,0 +1,204 @@ +import Foundation + +// MARK: - Cursor Local CSV Reader (tokscale compatible) + +enum CursorLocalCSVReader { + static func cachedCSVPaths(home: URL? = nil) -> [URL] { + let base: URL = if let home { + home.appendingPathComponent(".config/tokscale/cursor-cache", isDirectory: true) + } else if let dir = ProcessInfo.processInfo.environment["TOKSCALE_CONFIG_DIR"] { + URL(fileURLWithPath: dir).appendingPathComponent("cursor-cache", isDirectory: true) + } else { + FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".config/tokscale/cursor-cache", isDirectory: true) + } + guard let contents = try? FileManager.default.contentsOfDirectory( + at: base, includingPropertiesForKeys: nil, options: [.skipsHiddenFiles]) + else { return [] } + return contents.filter { + let n = $0.lastPathComponent + return n.hasPrefix("usage") && n.hasSuffix(".csv") && !$0.path.contains("/archive/") && !n + .hasPrefix("usage.backup") + }.sorted { $0.lastPathComponent < $1.lastPathComponent } + } + + struct Row: Sendable { + let date: Date + let model: String + let input: Int + let cacheRead: Int + let cacheWrite: Int + let output: Int + let totalTokens: Int? + let cost: Double + } + + static func parseFile(at url: URL, calendar: Calendar = .current) -> [Row] { + guard let data = try? Data(contentsOf: url), + let text = String(data: data, encoding: .utf8) else { return [] } + let lines = text.components(separatedBy: .newlines).filter { !$0.trimmingCharacters(in: .whitespaces).isEmpty } + guard let header = lines.first else { return [] } + let cols = self.parseCSVLine(header) + let hasKind = cols.contains { $0.lowercased() == "kind" } + struct Idx { let model: Int; let inW: Int; let inWo: Int; let read: Int; let out: Int; let cost: Int } + let idx = if hasKind, cols.count >= 11 { + Idx(model: 4, inW: 6, inWo: 7, read: 8, out: 9, cost: 11) + } else if hasKind { + Idx(model: 2, inW: 4, inWo: 5, read: 6, out: 7, cost: 9) + } else { + Idx(model: 1, inW: 2, inWo: 3, read: 4, out: 5, cost: 7) + } + var rows: [Row] = [] + for line in lines.dropFirst() { + let c = self.parseCSVLine(line) + guard c.count > max(idx.model, idx.cost) else { continue } + let model = c[idx.model].trimmingCharacters(in: .whitespacesAndNewlines) + guard !model.isEmpty else { continue } + let inW = self.parseInt(c[idx.inW]) + let inWo = self.parseInt(c[idx.inWo]) + let read = self.parseInt(c[idx.read]) + let out = self.parseInt(c[idx.out]) + let cost = self.parseCost(c[idx.cost]) + // Parse the authoritative Total Tokens column if present. + let totalIdx = idx.cost - 1 + let totalTokens = totalIdx < cols.count && totalIdx < c.count + && cols[totalIdx].lowercased().contains("total") ? self.parseInt(c[totalIdx]) : nil + guard let date = parseDate(c[0], calendar: calendar) else { continue } + let write = max(0, inW - inWo) + let input = inWo + if input == 0, out == 0, read == 0, write == 0 { continue } + rows.append(Row( + date: date, + model: model, + input: input, + cacheRead: read, + cacheWrite: write, + output: out, + totalTokens: totalTokens, + cost: cost)) + } + return rows + } + + static func makeDailyReport( + from rows: [Row], + calendar: Calendar = .current, + now: Date = Date()) -> CostUsageDailyReport + { + var byDay: [String: CostUsageDailyReport.Entry] = [:] + var totalCost: Double = 0 + var totalTokens = 0 + for row in rows { + let key = CostUsageLocalDay.key(from: row.date, calendar: calendar) + let existing = byDay[key] + var e = existing ?? CostUsageDailyReport.Entry( + date: key, + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + totalTokens: 0, + requestCount: 0, + costUSD: 0, + modelsUsed: nil, + modelBreakdowns: []) + // Honor the CSV's authoritative Total Tokens column when present; + // otherwise derive from component columns. + let tot = row.totalTokens ?? (row.input + row.output + row.cacheRead + row.cacheWrite) + var bds = e.modelBreakdowns ?? [] + if let i = bds.firstIndex(where: { $0.modelName == row.model }) { + let b = bds[i] + bds[i] = CostUsageDailyReport.ModelBreakdown( + modelName: b.modelName, + costUSD: (b.costUSD ?? 0) + row.cost, + totalTokens: (b.totalTokens ?? 0) + tot, + requestCount: (b.requestCount ?? 0) + 1) + } else { + bds.append(CostUsageDailyReport.ModelBreakdown( + modelName: row.model, + costUSD: row.cost, + totalTokens: tot, + requestCount: 1)) + } + let ne = CostUsageDailyReport.Entry( + date: key, + inputTokens: (e.inputTokens ?? 0) + row.input, + outputTokens: (e.outputTokens ?? 0) + row.output, + cacheReadTokens: (e.cacheReadTokens ?? 0) + row.cacheRead, + cacheCreationTokens: (e.cacheCreationTokens ?? 0) + row.cacheWrite, + totalTokens: (e.totalTokens ?? 0) + tot, + requestCount: (e.requestCount ?? 0) + 1, + costUSD: (e.costUSD ?? 0) + row.cost, + modelsUsed: e.modelsUsed, + modelBreakdowns: bds) + byDay[key] = ne + totalCost += row.cost + totalTokens += tot + } + let sorted = byDay.values.sorted { $0.date < $1.date } + let summary: CostUsageDailyReport.Summary? = sorted.isEmpty ? nil : .init( + totalInputTokens: nil, + totalOutputTokens: nil, + totalTokens: totalTokens, + totalCostUSD: totalCost) + return CostUsageDailyReport(data: sorted, summary: summary) + } + + static func parseCSVLine(_ line: String) -> [String] { + var cols: [String] = [] + var cur = "" + var q = false + for ch in line { + if ch == "\"" { q.toggle() } else if ch == ",", !q { + cols.append(cur); cur = "" + } else { cur.append(ch) } + } + cols.append(cur) + return cols + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines).replacingOccurrences(of: "\"\"", with: "\"") } + } + + static func parseInt(_ s: String) -> Int { + Int(s.trimmingCharacters(in: .whitespacesAndNewlines).replacingOccurrences(of: ",", with: "")) ?? 0 + } + + static func parseCost(_ s: String) -> Double { + let t = s.trimmingCharacters(in: .whitespacesAndNewlines) + if t.isEmpty || t == "-" || t.lowercased() == "included" || t.lowercased() == "nan" { return 0 } + return Double(t.replacingOccurrences(of: "$", with: "").replacingOccurrences(of: ",", with: "")) ?? 0 + } + + static func parseDate(_ s: String, calendar: Calendar) -> Date? { + let t = s.trimmingCharacters(in: .whitespacesAndNewlines) + let iso = ISO8601DateFormatter() + iso.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let d = iso.date(from: t) { return d } + iso.formatOptions = [.withInternetDateTime] + if let d = iso.date(from: t) { return d } + let df = DateFormatter() + df.locale = Locale(identifier: "en_US_POSIX") + df.timeZone = TimeZone(secondsFromGMT: 0) + for fmt in ["yyyy-MM-dd'T'HH:mm:ss.SSS", "yyyy-MM-dd'T'HH:mm:ss", "yyyy-MM-dd HH:mm:ss", "yyyy-MM-dd"] { + df.dateFormat = fmt + if let d = df.date(from: t) { + if fmt == "yyyy-MM-dd" { + // Keep date-only rows in the configured calendar's day (noon in that calendar), + // so UTC+13/14 users don't see the row shifted to the next local day. + var utcCal = Calendar(identifier: .gregorian) + utcCal.timeZone = TimeZone(secondsFromGMT: 0)! + let utcComps = utcCal.dateComponents([.year, .month, .day], from: d) + var targetComps = DateComponents() + targetComps.year = utcComps.year + targetComps.month = utcComps.month + targetComps.day = utcComps.day + targetComps.hour = 12 + targetComps.minute = 0 + targetComps.second = 0 + return calendar.date(from: targetComps) ?? d + } + return d + } + } + return nil + } +} diff --git a/Tests/CodexBarTests/CursorAntigravityLocalReaderTests.swift b/Tests/CodexBarTests/CursorAntigravityLocalReaderTests.swift new file mode 100644 index 00000000000..ffbce281746 --- /dev/null +++ b/Tests/CodexBarTests/CursorAntigravityLocalReaderTests.swift @@ -0,0 +1,138 @@ +import Foundation +import Testing +@testable import CodexBarCore + +/// Tokscale-compatible local readers: Cursor CSV caches +/// (`~/.config/tokscale/cursor-cache/usage*.csv`) and Antigravity session caches +/// (`~/.config/tokscale/antigravity-cache/sessions/*.jsonl`). +/// Header layouts mirror `tokscale/crates/tokscale-core/src/sessions/cursor.rs`: +/// - v1: Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to +/// you +/// - v2: Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total +/// Tokens,Cost +/// - v3: Date,Cloud Agent ID,Automation ID,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache +/// Read,Output Tokens,Total Tokens,Cost +struct CursorAntigravityLocalReaderTests { + @Test + func `cursor v1 csv parses token buckets and cost`() throws { + let csv = """ + Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you + 2026-08-20T10:00:00.000Z,claude-sonnet-4-5,1200,800,400,300,1500,0.012,0 + """ + let url = try Self.writeTemporary(csv, extension: "csv") + defer { try? FileManager.default.removeItem(at: url) } + let rows = CursorLocalCSVReader.parseFile(at: url) + #expect(rows.count == 1) + let row = try #require(rows.first) + #expect(row.model == "claude-sonnet-4-5") + #expect(row.input == 800) + #expect(row.cacheRead == 400) + #expect(row.cacheWrite == 400) + #expect(row.output == 300) + #expect(row.cost == 0.012) + } + + @Test + func `cursor v2 csv parses kind rows`() throws { + let csv = [ + "Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write)," + + "Cache Read,Output Tokens,Total Tokens,Cost", + "2026-08-20T11:00:00Z,chat,gpt-5.4,on,500,500,0,50,550,0.001", + ].joined(separator: "\n") + "\n" + let url = try Self.writeTemporary(csv, extension: "csv") + defer { try? FileManager.default.removeItem(at: url) } + let rows = CursorLocalCSVReader.parseFile(at: url) + let row = try #require(rows.first) + #expect(row.model == "gpt-5.4") + #expect(row.input == 500) + #expect(row.cacheRead == 0) + #expect(row.cacheWrite == 0) + #expect(row.output == 50) + } + + @Test + func `cursor v3 csv parses cloud agent rows`() throws { + let csv = [ + "Date,Cloud Agent ID,Automation ID,Kind,Model,Max Mode," + + "Input (w/ Cache Write),Input (w/o Cache Write)," + + "Cache Read,Output Tokens,Total Tokens,Cost", + "2026-08-20,agent-123,,agent,claude-opus-4-6,off,900,700,200,120,1020,0.05", + ].joined(separator: "\n") + "\n" + let url = try Self.writeTemporary(csv, extension: "csv") + defer { try? FileManager.default.removeItem(at: url) } + let rows = CursorLocalCSVReader.parseFile(at: url) + let row = try #require(rows.first) + #expect(row.model == "claude-opus-4-6") + #expect(row.input == 700) + #expect(row.cacheRead == 200) + #expect(row.cacheWrite == 200) + #expect(row.output == 120) + #expect(row.cost == 0.05) + } + + @Test + func `cursor csv aggregation groups rows into day entries`() throws { + let csv = """ + Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you + 2026-08-20T10:00:00.000Z,claude-sonnet-4-5,1200,800,400,300,1500,0.012,0 + 2026-08-20T12:00:00.000Z,claude-sonnet-4-5,600,500,100,100,700,0.004,0 + 2026-08-21T09:00:00.000Z,gpt-5.4,0,0,0,40,40,0.001,0 + """ + let url = try Self.writeTemporary(csv, extension: "csv") + defer { try? FileManager.default.removeItem(at: url) } + let report = CursorLocalCSVReader.makeDailyReport( + from: CursorLocalCSVReader.parseFile(at: url)) + #expect(report.data.count == 2) + let first = try #require(report.data.first) + #expect(first.inputTokens == 1300) + #expect(first.cacheReadTokens == 500) + #expect(first.cacheCreationTokens == 500) + #expect(first.outputTokens == 400) + #expect(first.requestCount == 2) + #expect(first.costUSD ?? 0 == 0.016) + let breakdowns = try #require(first.modelBreakdowns) + #expect(breakdowns.count == 1) + #expect(breakdowns.first?.modelName == "claude-sonnet-4-5") + let summary = try #require(report.summary) + // The CSV's own Total Tokens column is authoritative when present + // (1500 + 700 + 40), rather than the recomputed bucket sum. + #expect(summary.totalTokens == 2240) + } + + @Test + func `antigravity cache jsonl aggregates usage with session model fallback`() throws { + // Derive timestamps from local noon so both events land on the same local day in any TZ. + let calendar = Calendar.current + let noon = calendar.date( + byAdding: .hour, value: 12, to: calendar.startOfDay(for: Date())) ?? Date() + let firstStamp = Int64(noon.timeIntervalSince1970 * 1000) + let secondStamp = firstStamp + 3_600_000 + let jsonl = [ + #"{"type":"session_meta","modelId":"test-model-a"}"#, + #"{"type":"usage","modelId":"test-model-a","input":100,"output":30,"cacheRead":50,"cacheWrite":10,"# + + #""timestamp":\#(firstStamp)}"#, + #"{"type":"usage","input":40,"output":10,"cacheRead":0,"cacheWrite":0,"timestamp":\#(secondStamp)}"#, + ].joined(separator: "\n") + "\n" + let url = try Self.writeTemporary(jsonl, extension: "jsonl") + defer { try? FileManager.default.removeItem(at: url) } + let entries = AntigravityLocalReader.parseJSONLCache(paths: [url]) + #expect(entries.count == 1) + let entry = try #require(entries.first) + #expect(entry.inputTokens == 140) + #expect(entry.outputTokens == 40) + #expect(entry.cacheReadTokens == 50) + #expect(entry.cacheCreationTokens == 10) + #expect(entry.requestCount == 2) + // The second event omits `modelId` and falls back to the session_meta model. + let breakdowns = try #require(entry.modelBreakdowns) + #expect(breakdowns.map(\.modelName) == ["test-model-a"]) + #expect(breakdowns.first?.requestCount == 2) + } + + private static func writeTemporary(_ contents: String, extension fileExtension: String) throws -> URL { + let url = FileManager.default.temporaryDirectory + .appendingPathComponent("cursor-antigravity-reader-\(UUID().uuidString).\(fileExtension)") + try contents.write(to: url, atomically: true, encoding: .utf8) + return url + } +} diff --git a/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift b/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift index de1832210b2..cbdc54201ec 100644 --- a/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift +++ b/Tests/CodexBarTests/ProviderArchitectureGatekeeperTests.swift @@ -200,12 +200,13 @@ struct ProviderArchitectureGatekeeperTests { .warp, .kilo, .mistral, .deepseek, .deepinfra, .qoder, .crof, .chutes, ]) #if os(macOS) + // Antigravity joined via the tokscale-compatible local usage readers. #expect(Set(descriptors.filter(\.tokenCost.supportsTokenSnapshot).map(\.id)) == [ - .codex, .claude, .cursor, .vertexai, .bedrock, + .codex, .claude, .cursor, .vertexai, .bedrock, .antigravity, ]) #else #expect(Set(descriptors.filter(\.tokenCost.supportsTokenSnapshot).map(\.id)) == [ - .codex, .claude, .vertexai, .bedrock, + .codex, .claude, .vertexai, .bedrock, .antigravity, ]) #endif #expect(Set(descriptors.filter { $0.cli.binaryLocator != nil }.map(\.id)) == [ @@ -1353,31 +1354,31 @@ struct ProviderArchitectureGatekeeperTests { reason: "This provider-owned integration passes its fixed identity to a shared helper."), SuppressedProviderReference( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 291, + line: 292, anchor: "provider: .codex,", expectedProviderIDs: ["codex"], reason: "This provider-specific core branch passes its already-selected identity to a shared helper."), SuppressedProviderReference( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 307, + line: 308, anchor: "provider: .codex,", expectedProviderIDs: ["codex"], reason: "This provider-specific core branch passes its already-selected identity to a shared helper."), SuppressedProviderReference( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 779, + line: 809, anchor: "provider: .codex,", expectedProviderIDs: ["codex"], reason: "This provider-specific core branch passes its already-selected identity to a shared helper."), SuppressedProviderReference( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 859, + line: 889, anchor: "provider: .codex,", expectedProviderIDs: ["codex"], reason: "This provider-specific core branch passes its already-selected identity to a shared helper."), SuppressedProviderReference( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 950, + line: 980, anchor: "provider: .codex,", expectedProviderIDs: ["codex"], reason: "This provider-specific core branch passes its already-selected identity to a shared helper."), @@ -3474,7 +3475,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact shared construct dispatches a provider-owned capability at the generic integration boundary."), AllowedProviderConstruct( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 584, + line: 614, anchor: "if provider == .codex {", expectedProviderIDs: ["codex"], expectedReferenceCount: 1, @@ -3482,7 +3483,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact cost scanner dispatch selects a provider-owned transcript, cache, or pricing format."), AllowedProviderConstruct( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 612, + line: 642, anchor: "provider == .claude || (provider == .codex && options.shouldMergePiUsage)", expectedProviderIDs: ["claude", "codex"], expectedReferenceCount: 5, @@ -3490,7 +3491,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact cost scanner dispatch selects a provider-owned transcript, cache, or pricing format."), AllowedProviderConstruct( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 659, + line: 689, anchor: "options.provider == .codex || options.provider == .claude", expectedProviderIDs: ["claude", "codex"], expectedReferenceCount: 2, @@ -3498,7 +3499,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact cost scanner dispatch selects a provider-owned transcript, cache, or pricing format."), AllowedProviderConstruct( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 690, + line: 720, anchor: "guard provider == .codex || provider == .claude else { return nil }", expectedProviderIDs: ["claude", "codex"], expectedReferenceCount: 3, @@ -3506,7 +3507,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact cost scanner dispatch selects a provider-owned transcript, cache, or pricing format."), AllowedProviderConstruct( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 1237, + line: 1335, anchor: "if provider == .vertexai {", expectedProviderIDs: ["claude", "vertexai"], expectedReferenceCount: 2, @@ -3514,7 +3515,7 @@ struct ProviderArchitectureGatekeeperTests { reason: "This exact cost scanner dispatch selects a provider-owned transcript, cache, or pricing format."), AllowedProviderConstruct( path: "Sources/CodexBarCore/CostUsageFetcher.swift", - line: 1593, + line: 1691, anchor: "if provider == .cursor {", expectedProviderIDs: ["cursor"], expectedReferenceCount: 1, diff --git a/Tests/CodexBarTests/SpendDashboardModelTests.swift b/Tests/CodexBarTests/SpendDashboardModelTests.swift index 6056331e6fc..69b0986d75b 100644 --- a/Tests/CodexBarTests/SpendDashboardModelTests.swift +++ b/Tests/CodexBarTests/SpendDashboardModelTests.swift @@ -81,6 +81,8 @@ struct SpendDashboardModelTests { .opencodego, .openrouter, .xai, + // Antigravity joined via the tokscale-compatible local usage readers. + .antigravity, ]) }