diff --git a/Sources/CodexBar/Providers/Muse/MuseProviderImplementation.swift b/Sources/CodexBar/Providers/Muse/MuseProviderImplementation.swift new file mode 100644 index 0000000000..b992f768b1 --- /dev/null +++ b/Sources/CodexBar/Providers/Muse/MuseProviderImplementation.swift @@ -0,0 +1,6 @@ +import CodexBarCore +import Foundation + +struct MuseProviderImplementation: ProviderImplementation { + let id: UsageProvider = .muse +} diff --git a/Sources/CodexBar/Providers/Shared/ProviderImplementationManifest.swift b/Sources/CodexBar/Providers/Shared/ProviderImplementationManifest.swift index 49ee8bb14b..c83bbe0dd1 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderImplementationManifest.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderImplementationManifest.swift @@ -6,6 +6,7 @@ import Foundation /// `ProviderImplementationRegistry.register(_:)`. enum ProviderImplementationManifest { static let makeImplementations: [@Sendable () -> any ProviderImplementation] = [ + { MuseProviderImplementation() }, { CodexProviderImplementation() }, { OpenAIAPIProviderImplementation() }, { AzureOpenAIProviderImplementation() }, diff --git a/Sources/CodexBar/Resources/ProviderIcon-muse.svg b/Sources/CodexBar/Resources/ProviderIcon-muse.svg new file mode 100644 index 0000000000..8c3050cba2 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-muse.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/Sources/CodexBar/UsageStore+TokenCost.swift b/Sources/CodexBar/UsageStore+TokenCost.swift index 076cec6589..49a5a07997 100644 --- a/Sources/CodexBar/UsageStore+TokenCost.swift +++ b/Sources/CodexBar/UsageStore+TokenCost.swift @@ -518,7 +518,7 @@ extension UsageStore { return snapshot?.costUsage case .xai: return snapshot.flatMap { XAICostUsageMapping.tokenSnapshot(from: $0, historyDays: windowDays) } - case .grok: + case .grok, .muse: return self.grokLocalTokenSnapshot(from: snapshot, historyDays: windowDays) default: return nil @@ -529,7 +529,7 @@ extension UsageStore { // Provider-specific by design: these providers project live usage snapshots into the // shared spend catalog instead of running the local CostUsageFetcher JSONL pipeline. switch provider { - case .grok, .mistral, .openai, .opencodego, .openrouter, .xai: + case .grok, .muse, .mistral, .openai, .opencodego, .openrouter, .xai: true default: false diff --git a/Sources/CodexBarCore/Providers/Muse/MuseLocalSessionScanner.swift b/Sources/CodexBarCore/Providers/Muse/MuseLocalSessionScanner.swift new file mode 100644 index 0000000000..cb2ee95042 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Muse/MuseLocalSessionScanner.swift @@ -0,0 +1,556 @@ +import Foundation + +/// One local-calendar day of Muse session-token activity. +public struct MuseLocalDailyBucket: Sendable, Equatable { + public let date: String + public let totalTokens: Int + public let inputTokens: Int + public let outputTokens: Int + public let cacheReadTokens: Int + public let reasoningTokens: Int + public let requestCount: Int + public let models: [String] + public init(date: String, totalTokens: Int, inputTokens: Int, outputTokens: Int, cacheReadTokens: Int, reasoningTokens: Int, requestCount: Int, models: [String]) { + self.date = date; self.totalTokens = totalTokens; self.inputTokens = inputTokens; self.outputTokens = outputTokens; self.cacheReadTokens = cacheReadTokens; self.reasoningTokens = reasoningTokens; self.requestCount = requestCount; self.models = models + } +} + +public struct MuseLocalSessionSummary: Sendable { + public let fileCount: Int; public let totalTokens: Int; public let totalInputTokens: Int; public let totalOutputTokens: Int; public let totalCacheReadTokens: Int; public let totalReasoningTokens: Int; public let requestCount: Int; public let lastEventAt: Date?; public let primaryModel: String?; public let models: [String]; public let daily: [MuseLocalDailyBucket]; public let scannedAt: Date + public init(fileCount: Int, totalTokens: Int, totalInputTokens: Int, totalOutputTokens: Int, totalCacheReadTokens: Int, totalReasoningTokens: Int, requestCount: Int, lastEventAt: Date?, primaryModel: String?, models: [String], daily: [MuseLocalDailyBucket] = [], scannedAt: Date = .init()) { + self.fileCount = fileCount; self.totalTokens = totalTokens; self.totalInputTokens = totalInputTokens; self.totalOutputTokens = totalOutputTokens; self.totalCacheReadTokens = totalCacheReadTokens; self.totalReasoningTokens = totalReasoningTokens; self.requestCount = requestCount; self.lastEventAt = lastEventAt; self.primaryModel = primaryModel; self.models = models; self.daily = daily; self.scannedAt = scannedAt + } + public func toCostUsageTokenSnapshot(historyDays: Int) -> CostUsageTokenSnapshot? { + let entries = self.daily.map { bucket in CostUsageDailyReport.Entry(date: bucket.date, inputTokens: bucket.inputTokens > 0 ? bucket.inputTokens : nil, outputTokens: bucket.outputTokens > 0 ? bucket.outputTokens : nil, cacheReadTokens: bucket.cacheReadTokens > 0 ? bucket.cacheReadTokens : nil, reasoningTokens: bucket.reasoningTokens > 0 ? bucket.reasoningTokens : nil, totalTokens: bucket.totalTokens, requestCount: bucket.requestCount, costUSD: nil, modelsUsed: bucket.models.isEmpty ? nil : bucket.models, modelBreakdowns: nil) } + guard !entries.isEmpty else { return nil } + let todayKey = MuseLocalSessionScanner.dayKey(for: self.scannedAt, calendar: .current) + let todayTokens = todayKey.flatMap { key in self.daily.first { $0.date == key }?.totalTokens } + let todayRequests = todayKey.flatMap { key in self.daily.first { $0.date == key }?.requestCount } + return CostUsageTokenSnapshot(sessionTokens: todayTokens, sessionCostUSD: nil, sessionRequests: todayRequests, last30DaysTokens: self.totalTokens, last30DaysCostUSD: nil, last30DaysRequests: self.requestCount, historyDays: historyDays, historyCoverageIsEstablished: true, costProvenance: .unknown, daily: entries, updatedAt: self.scannedAt) + } +} + +public enum MuseLocalSessionScanner { + public static let defaultLookbackDays = 30 + nonisolated(unsafe) static var test_fileScanObserver: ((URL) -> Void)? + // MARK: - Public entry points + public static func summarize(env: [String: String] = ProcessInfo.processInfo.environment, fileManager: FileManager = .default, lookbackDays: Int = defaultLookbackDays, now: Date = .init(), fileScanObserver: ((URL) -> Void)? = nil) -> MuseLocalSessionSummary { + // Non-throwing wrapper for tests and sync callers; cancellation is not checked. + do { + return try self.summarizeCancellable(env: env, fileManager: fileManager, lookbackDays: lookbackDays, now: now, fileScanObserver: fileScanObserver, checkCancellation: nil) + } catch { + // Should never throw when checkCancellation is nil; return empty summary as fallback + return MuseLocalSessionSummary(fileCount: 0, totalTokens: 0, totalInputTokens: 0, totalOutputTokens: 0, totalCacheReadTokens: 0, totalReasoningTokens: 0, requestCount: 0, lastEventAt: nil, primaryModel: nil, models: [], scannedAt: now) + } + } + + static func summarizeCancellable(env: [String: String], fileManager: FileManager, lookbackDays: Int, now: Date, fileScanObserver: ((URL) -> Void)?, checkCancellation: (() throws -> Void)?) throws -> MuseLocalSessionSummary { + let cacheRoot = self.cacheRootForTesting(env: env, fileManager: fileManager) + var cache = MuseSessionCostCacheIO.load(cacheRoot: cacheRoot) + let calendar = Calendar.current + if cache.timeZoneIdentifier != nil, cache.timeZoneIdentifier != calendar.timeZone.identifier { cache = MuseSessionCostCache(version: cache.version) } + let lookbackCutoff = calendar.date(byAdding: .day, value: -lookbackDays, to: now) ?? now + let scanSinceKey = self.dayKey(for: lookbackCutoff, calendar: calendar) ?? "" + let scanUntilKey = self.dayKey(for: now, calendar: calendar) ?? "" + let sessionsRoot = self.sessionsRoot(env: env, fileManager: fileManager) + let dateDirs = self.relevantDateDirectories(sessionsRoot: sessionsRoot, calendar: calendar, lookbackCutoff: lookbackCutoff, now: now, fileManager: fileManager) + if dateDirs.isEmpty { + // Prune old data even when no date dirs exist + self.pruneExpired(cache: &cache, scanSinceKey: scanSinceKey, scanUntilKey: scanUntilKey, calendar: calendar) + // Only save if not cancelled + try checkCancellation?() + cache.lastScanUnixMs = Int64(now.timeIntervalSince1970 * 1000) + MuseSessionCostCacheIO.save(cache: cache, cacheRoot: cacheRoot, calendar: calendar) + let summary = self.summaryFromCache(cache: cache, calendar: calendar, sinceKey: scanSinceKey, untilKey: scanUntilKey, now: now, fileCount: 0) + return summary + } + var filePathsInScan: Set = [] + for dateDir in dateDirs { + try checkCancellation?() + guard let enumerator = fileManager.enumerator(at: dateDir, includingPropertiesForKeys: [.isRegularFileKey], options: [.skipsHiddenFiles]) else { continue } + for case let url as URL in enumerator { guard url.lastPathComponent == "session.jsonl" else { continue }; filePathsInScan.insert(url.path) } + } + for path in filePathsInScan.sorted() { + try checkCancellation?() + let url = URL(fileURLWithPath: path) + let attrs = (try? fileManager.attributesOfItem(atPath: path)) ?? [:] + let mtime = (attrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0 + let size = (attrs[.size] as? NSNumber)?.int64Value ?? 0 + let mtimeMs = Int64(mtime * 1000) + let cached = cache.files[path] + if let cached, cached.mtimeUnixMs == mtimeMs, cached.size == size { continue } + fileScanObserver?(url); self.test_fileScanObserver?(url) + // Check for append vs replacement via fingerprint + if let cached, size > cached.size, cached.parsedBytes > 0, cached.parsedBytes <= size { + // Verify prefix fingerprint to distinguish append from larger replacement. + // For small files (<4K) any append changes the first-4K hash, so compare only the prefix that was present at cache time. + let fingerprintMatches: Bool = { + guard let cachedFingerprint = cached.prefixFingerprint else { return false } // old cache without fingerprint => force full reparse + guard let fullData = try? Data(contentsOf: url) else { return false } + let compareLength = min(4096, Int(cached.size)) + let currentFingerprint = Self.prefixFingerprint(for: fullData, prefixLength: compareLength) + return currentFingerprint == cachedFingerprint + }() + if fingerprintMatches { + let delta = self.parseMuseSessionFileDelta(fileURL: url, startOffset: cached.parsedBytes, calendar: calendar, lookbackCutoff: lookbackCutoff) + if !delta.contributions.isEmpty { self.applyContributions(to: &cache.days, contributions: delta.contributions, sign: 1) } + let merged = self.mergedContributions(existing: cached.contributions, delta: delta.contributions) + let newFingerprint = (try? Data(contentsOf: url)).map { Self.prefixFingerprint(for: $0) } + cache.files[path] = MuseSessionFileUsage(mtimeUnixMs: mtimeMs, size: size, parsedBytes: delta.parsedBytes, prefixFingerprint: newFingerprint, contributions: merged, entryCount: cached.entryCount + delta.entryCount) + continue + } + // Fingerprint mismatch => treat as replacement (fall through to full reparse) + } + // Full reparse path (replacement or new file). Parse first to detect incomplete truncated writes. + let parsed = self.parseMuseSessionFileFull(fileURL: url, calendar: calendar, lookbackCutoff: lookbackCutoff) + // If file was truncated to an incomplete fragment (no complete JSON, no newline), preserve old contributions + // until the write completes. This prevents losing already-counted usage when a crash leaves a partial line. + // ponytail: O(1) check; if file is intentionally truncated to a smaller complete file, parsedBytes == size so we replace. + if let cached, size < cached.size, parsed.parsedBytes < size, parsed.contributions.isEmpty, parsed.entryCount == 0 { + // Keep old cache entry and contributions; don't update file metadata yet (will retry next scan) + continue + } + if let cached { self.applyContributions(to: &cache.days, contributions: cached.contributions, sign: -1) } + let fingerprint = (try? Data(contentsOf: url)).map { Self.prefixFingerprint(for: $0) } + if !parsed.contributions.isEmpty { self.applyContributions(to: &cache.days, contributions: parsed.contributions, sign: 1) } + cache.files[path] = MuseSessionFileUsage(mtimeUnixMs: mtimeMs, size: size, parsedBytes: parsed.parsedBytes, prefixFingerprint: fingerprint, contributions: parsed.contributions, entryCount: parsed.entryCount) + } + // Handle deletions and aging-out + var deletedPaths: [String] = [] + for (path, usage) in cache.files { + try checkCancellation?() + if filePathsInScan.contains(path) { continue } + // If file still exists but its date directory is outside current lookback, it will not be in filePathsInScan + // We prune such files regardless of existence to prevent unbounded growth + if let dayKey = self.dayKeyFromPath(path, calendar: calendar) { + if dayKey < scanSinceKey { + // Aged out: remove contributions and metadata (already filtered from summary, but clean cache) + self.applyContributions(to: &cache.days, contributions: usage.contributions, sign: -1) + deletedPaths.append(path) + continue + } + if dayKey >= scanSinceKey, dayKey <= scanUntilKey { + if !fileManager.fileExists(atPath: path) { + self.applyContributions(to: &cache.days, contributions: usage.contributions, sign: -1) + deletedPaths.append(path) + } + } else { + // DayKey outside window but not aged out? (future) Keep for now. + } + } else { + // Path doesn't match expected /sessions/YYYY/MM/DD structure; if not in scan and missing, prune + if !fileManager.fileExists(atPath: path) { + self.applyContributions(to: &cache.days, contributions: usage.contributions, sign: -1) + deletedPaths.append(path) + } + } + } + for path in deletedPaths { cache.files.removeValue(forKey: path) } + // Prune expired day aggregates (outside window) that may remain from aged-out files or timezone shifts + self.pruneExpired(cache: &cache, scanSinceKey: scanSinceKey, scanUntilKey: scanUntilKey, calendar: calendar) + try checkCancellation?() + cache.lastScanUnixMs = Int64(now.timeIntervalSince1970 * 1000) + MuseSessionCostCacheIO.save(cache: cache, cacheRoot: cacheRoot, calendar: calendar) + let summary = self.summaryFromCache(cache: cache, calendar: calendar, sinceKey: scanSinceKey, untilKey: scanUntilKey, now: now, fileCount: filePathsInScan.count) + return summary + } + + private static func pruneExpired(cache: inout MuseSessionCostCache, scanSinceKey: String, scanUntilKey: String, calendar: Calendar) { + // Remove day aggregates outside the current window + for day in Array(cache.days.keys) { + if day < scanSinceKey || day > scanUntilKey { + cache.days.removeValue(forKey: day) + } + } + // Also prune files whose dayKey is outside window and that were not already removed + // (keeps cache bounded as years of sessions accumulate) + for (path, usage) in Array(cache.files) { + if let dayKey = self.dayKeyFromPath(path, calendar: calendar), dayKey < scanSinceKey { + // Contributions already pruned from days; just remove file metadata + // Ensure we don't double-subtract if caller already handled + if !usage.contributions.isEmpty { + // Contributions for this file should already be absent from cache.days after day prune, + // but if some remain (e.g., due to earlier bug), subtract + self.applyContributions(to: &cache.days, contributions: usage.contributions, sign: -1) + } + cache.files.removeValue(forKey: path) + } + } + } + + // Deterministic FNV-1a 64-bit for first 4K prefix + static func prefixFingerprint(for data: Data, prefixLength: Int? = nil) -> String { + let len = prefixLength ?? min(4096, data.count) + let clamped = min(len, data.count) + var hash: UInt64 = 14695981039346656037 + for i in 0.. MuseLocalSessionSummary { + try await CostUsageScanExecutor.run { checkCancellation in + try checkCancellation() + let summary = try Self.summarizeCancellable(env: env, fileManager: .default, lookbackDays: lookbackDays, now: now, fileScanObserver: nil, checkCancellation: checkCancellation) + try checkCancellation() + return summary + } + } + static func sessionsRoot(env: [String: String], fileManager: FileManager) -> URL { + if let override = env["MUSE_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines), !override.isEmpty { return URL(fileURLWithPath: override).appendingPathComponent("sessions", isDirectory: true) } + if let override = env["CODEXBAR_MUSE_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines), !override.isEmpty { return URL(fileURLWithPath: override).appendingPathComponent("sessions", isDirectory: true) } + let home = env["HOME"].map { URL(fileURLWithPath: $0) } ?? fileManager.homeDirectoryForCurrentUser + return home.appendingPathComponent(".local", isDirectory: true).appendingPathComponent("share", isDirectory: true).appendingPathComponent("muse", isDirectory: true).appendingPathComponent("sessions", isDirectory: true) + } + static func cacheRootForTesting(env: [String: String], fileManager: FileManager) -> URL? { + if let museHome = env["MUSE_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines), !museHome.isEmpty { return URL(fileURLWithPath: museHome) } + if let museHome = env["CODEXBAR_MUSE_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines), !museHome.isEmpty { return URL(fileURLWithPath: museHome) } + if let museCache = env["MUSE_CACHE_ROOT"]?.trimmingCharacters(in: .whitespacesAndNewlines), !museCache.isEmpty { return URL(fileURLWithPath: museCache) } + return nil + } + static func relevantDateDirectories(sessionsRoot: URL, calendar: Calendar, lookbackCutoff: Date, now: Date, fileManager: FileManager) -> [URL] { + let startDay = calendar.startOfDay(for: lookbackCutoff); let endDay = calendar.startOfDay(for: now); guard startDay <= endDay else { return [] } + var dirs: [URL] = []; var cursor = startDay + while cursor <= endDay { + let comps = calendar.dateComponents([.year, .month, .day], from: cursor) + guard let y = comps.year, let m = comps.month, let d = comps.day else { break } + let dir = sessionsRoot.appendingPathComponent(String(format: "%04d", y), isDirectory: true).appendingPathComponent(String(format: "%02d", m), isDirectory: true).appendingPathComponent(String(format: "%02d", d), isDirectory: true) + var isDir: ObjCBool = false + if fileManager.fileExists(atPath: dir.path, isDirectory: &isDir), isDir.boolValue { dirs.append(dir) } + guard let next = calendar.date(byAdding: .day, value: 1, to: cursor) else { break } + cursor = next + } + return dirs + } + struct ParsedEvent { let date: Date; let inputTokens: Int; let outputTokens: Int; let cacheReadTokens: Int; let reasoningTokens: Int; let model: String } + static func events(in fileURL: URL, fileManager: FileManager) -> [ParsedEvent] { + var out: [ParsedEvent] = [] + do { + _ = try CostUsageJsonl.scan(fileURL: fileURL, maxLineBytes: 32*1024, prefixBytes: 16*1024, onLine: { line in + guard !line.wasTruncated else { return } + if let event = Self.parseLine(line.bytes) { out.append(event); return } + for innerData in Self.extractRecordJSONDatas(from: line.bytes) { if let event = Self.parseLine(innerData) { out.append(event) } } + }) + } catch { return out } + return out + } + static func extractRecordJSONDatas(from data: Data) -> [Data] { + var results: [Data] = [] + data.withUnsafeBytes { raw in + guard let base = raw.baseAddress else { return } + let bytes = UnsafeBufferPointer(start: base.assumingMemoryBound(to: UInt8.self), count: data.count) + let fullRange = 0.. 8 { break } + } + } + return results + } + private static func rangeOfFieldValue(field: [UInt8], bytes: UnsafeBufferPointer, in range: Range) -> Range? { + var idx = range.lowerBound + while idx < range.upperBound { + guard let keyStart = self.indexOfQuote(from: idx, bytes: bytes, limit: range.upperBound) else { break } + var keyEnd = keyStart + 1 + while keyEnd < range.upperBound { + let b = bytes[keyEnd] + if b == UInt8(ascii: "\\") { keyEnd += 2; continue } + if b == UInt8(ascii: "\"") { break } + keyEnd += 1 + } + guard keyEnd < range.upperBound else { break } + let keyLen = keyEnd - (keyStart + 1) + let matches: Bool = (keyLen == field.count) ? self.bytesEqual(bytes: bytes, from: keyStart + 1, field: field) : false + var afterKey = keyEnd + 1 + self.skipWhitespace(bytes: bytes, idx: &afterKey, limit: range.upperBound) + if afterKey < range.upperBound, bytes[afterKey] == UInt8(ascii: ":") { + afterKey += 1; self.skipWhitespace(bytes: bytes, idx: &afterKey, limit: range.upperBound) + if matches { + if afterKey < range.upperBound, bytes[afterKey] == UInt8(ascii: "\"") { + var end = afterKey + 1 + while end < range.upperBound { + let b = bytes[end] + if b == UInt8(ascii: "\\") { end += 2; continue } + if b == UInt8(ascii: "\"") { return afterKey..<(end+1) } + end += 1 + } + } + return nil + } + } + idx = afterKey + } + return nil + } + static func parseLine(_ data: Data) -> ParsedEvent? { + guard let text = String(data: data, encoding: .utf8) else { return nil } + guard text.contains("\"model_completed\"") else { return nil } + guard text.contains("\"runtime.session\"") else { return nil } + return data.withUnsafeBytes { raw -> ParsedEvent? in + guard let base = raw.baseAddress else { return nil } + let bytes = UnsafeBufferPointer(start: base.assumingMemoryBound(to: UInt8.self), count: data.count) + let fullRange = 0.. 0 else { return nil } + return ParsedEvent(date: date, inputTokens: max(0, input), outputTokens: max(0, output), cacheReadTokens: max(0, cached), reasoningTokens: max(0, reasoning), model: model.trimmingCharacters(in: .whitespacesAndNewlines)) + } + } + private static func extractStringField(_ field: [UInt8], from bytes: UnsafeBufferPointer, in range: Range) -> String? { + self.extractField(field, from: bytes, in: range) { idx in self.parseJSONString(at: &idx, bytes: bytes, limit: range.upperBound) } + } + private static func extractIntField(_ field: [UInt8], from bytes: UnsafeBufferPointer, in range: Range) -> Int? { + self.extractField(field, from: bytes, in: range) { idx in self.parseInt(at: &idx, bytes: bytes, limit: range.upperBound) } + } + private static func extractNestedString(outerField: [UInt8], innerField: [UInt8], from bytes: UnsafeBufferPointer, in range: Range) -> String? { + guard let outerRange = self.objectRange(for: outerField, from: bytes, in: range) else { return nil } + return self.extractStringField(innerField, from: bytes, in: outerRange) + } + private static func extractNestedInt(outerField: [UInt8], innerField: [UInt8], from bytes: UnsafeBufferPointer, in range: Range) -> Int? { + guard let outerRange = self.objectRange(for: outerField, from: bytes, in: range) else { return nil } + return self.extractIntField(innerField, from: bytes, in: outerRange) + } + private static func objectRange(for field: [UInt8], from bytes: UnsafeBufferPointer, in range: Range) -> Range? { + self.extractField(field, from: bytes, in: range) { idx in self.parseObjectRange(at: &idx, bytes: bytes, limit: range.upperBound) } + } + private static func extractField(_ field: [UInt8], from bytes: UnsafeBufferPointer, in range: Range, parse: (inout Int) -> T?) -> T? { + var idx = range.lowerBound + while idx < range.upperBound { + guard let keyStart = self.indexOfQuote(from: idx, bytes: bytes, limit: range.upperBound) else { break } + var keyEnd = keyStart + 1; var hasEscape = false + while keyEnd < range.upperBound { + let b = bytes[keyEnd] + if b == UInt8(ascii: "\\") { hasEscape = true; keyEnd += 2; continue } + if b == UInt8(ascii: "\"") { break } + keyEnd += 1 + } + guard keyEnd < range.upperBound else { break } + let keyLen = keyEnd - (keyStart + 1) + let keyMatches: Bool = if hasEscape { self.decodeString(bytes: bytes, from: keyStart + 1, to: keyEnd) == String(bytes: field, encoding: .utf8) } else if keyLen == field.count { self.bytesEqual(bytes: bytes, from: keyStart + 1, field: field) } else { false } + var afterKey = keyEnd + 1 + self.skipWhitespace(bytes: bytes, idx: &afterKey, limit: range.upperBound) + if afterKey < range.upperBound, bytes[afterKey] == UInt8(ascii: ":") { + afterKey += 1; self.skipWhitespace(bytes: bytes, idx: &afterKey, limit: range.upperBound) + if keyMatches { var valueIdx = afterKey; if let v = parse(&valueIdx) { return v } } + } + idx = afterKey + } + return nil + } + private static func indexOfQuote(from idx: Int, bytes: UnsafeBufferPointer, limit: Int) -> Int? { + var i = idx + while i < limit { if bytes[i] == UInt8(ascii: "\"") { return i }; i += 1 } + return nil + } + private static func skipWhitespace(bytes: UnsafeBufferPointer, idx: inout Int, limit: Int) { + while idx < limit, bytes[idx] == 32 || bytes[idx] == 9 || bytes[idx] == 10 || bytes[idx] == 13 { idx += 1 } + } + private static func bytesEqual(bytes: UnsafeBufferPointer, from start: Int, field: [UInt8]) -> Bool { + for i in 0.., from start: Int, to end: Int) -> String? { + var out: [UInt8] = []; out.reserveCapacity(end - start); var i = start + while i < end { + let b = bytes[i] + if b == UInt8(ascii: "\\"), i + 1 < end { + let n = bytes[i+1] + switch n { + case UInt8(ascii: "\""): out.append(UInt8(ascii: "\"")); i += 2 + case UInt8(ascii: "\\"): out.append(UInt8(ascii: "\\")); i += 2 + case UInt8(ascii: "/"): out.append(UInt8(ascii: "/")); i += 2 + case UInt8(ascii: "n"): out.append(10); i += 2 + case UInt8(ascii: "t"): out.append(9); i += 2 + default: out.append(b); i += 1 + } + } else { out.append(b); i += 1 } + } + return String(bytes: out, encoding: .utf8) + } + private static func parseJSONString(at idx: inout Int, bytes: UnsafeBufferPointer, limit: Int) -> String? { + guard idx < limit, bytes[idx] == UInt8(ascii: "\"") else { return nil } + idx += 1; let start = idx; var hasEscape = false + while idx < limit { + let b = bytes[idx] + if b == UInt8(ascii: "\\") { hasEscape = true; idx += 2; continue } + if b == UInt8(ascii: "\"") { + let end = idx; idx += 1 + if hasEscape { return self.decodeString(bytes: bytes, from: start, to: end) } + return String(bytes: bytes[start.., limit: Int) -> Int? { + self.skipWhitespace(bytes: bytes, idx: &idx, limit: limit) + var sign = 1 + if idx < limit, bytes[idx] == UInt8(ascii: "-") { sign = -1; idx += 1 } + var value = 0; var sawDigit = false + while idx < limit, bytes[idx] >= 48, bytes[idx] <= 57 { + sawDigit = true; let d = Int(bytes[idx] - 48) + let (m, o1) = value.multipliedReportingOverflow(by: 10); if o1 { return nil } + let (a, o2) = m.addingReportingOverflow(d); if o2 { return nil } + value = a; idx += 1 + } + return sawDigit ? sign * value : nil + } + private static func parseObjectRange(at idx: inout Int, bytes: UnsafeBufferPointer, limit: Int) -> Range? { + guard idx < limit, bytes[idx] == UInt8(ascii: "{") else { return nil } + let start = idx; var depth = 0; var inString = false; var escape = false + while idx < limit { + let b = bytes[idx] + if inString { + if escape { escape = false } else if b == UInt8(ascii: "\\") { escape = true } else if b == UInt8(ascii: "\"") { inString = false } + idx += 1; continue + } + if b == UInt8(ascii: "\"") { inString = true; idx += 1; continue } + if b == UInt8(ascii: "{") { depth += 1 } else if b == UInt8(ascii: "}") { depth -= 1; if depth == 0 { idx += 1; return start.. String? { + let comps = calendar.dateComponents([.year, .month, .day], from: date) + guard let y = comps.year, let m = comps.month, let d = comps.day else { return nil } + return String(format: "%04d-%02d-%02d", y, m, d) + } + // Cache helpers + private static func applyContributions(to days: inout [String: [String: MusePackedUsage]], contributions: [String: [String: MusePackedUsage]], sign: Int) { + for (day, modelMap) in contributions { for (model, usage) in modelMap { var dayMap = days[day] ?? [:]; let existing = dayMap[model] ?? MusePackedUsage(); let delta = sign > 0 ? (existing + usage) : (existing - usage); if delta.isZero { dayMap.removeValue(forKey: model) } else { dayMap[model] = delta }; if dayMap.isEmpty { days.removeValue(forKey: day) } else { days[day] = dayMap } } } + } + private static func mergedContributions(existing: [String: [String: MusePackedUsage]], delta: [String: [String: MusePackedUsage]]) -> [String: [String: MusePackedUsage]] { + var result = existing; for (day, modelMap) in delta { for (model, usage) in modelMap { var dayMap = result[day] ?? [:]; let merged = (dayMap[model] ?? MusePackedUsage()) + usage; dayMap[model] = merged; result[day] = dayMap } }; return result + } + private static func summaryFromCache(cache: MuseSessionCostCache, calendar: Calendar, sinceKey: String, untilKey: String, now: Date, fileCount: Int) -> MuseLocalSessionSummary { + var totalTokens = 0; var totalInput = 0; var totalOutput = 0; var totalCacheRead = 0; var totalReasoning = 0; var requestCount = 0; var modelCounts: [String: Int] = [:]; var dailyInput: [String: Int] = [:]; var dailyOutput: [String: Int] = [:]; var dailyCacheRead: [String: Int] = [:]; var dailyReasoning: [String: Int] = [:]; var dailyTokens: [String: Int] = [:]; var dailyRequests: [String: Int] = [:]; var dailyModels: [String: [String: Int]] = [:]; var lastEventAt: Date? + for (day, modelMap) in cache.days { guard day >= sinceKey, day <= untilKey else { continue }; for (model, usage) in modelMap { totalTokens += usage.totalTokens; totalInput += usage.inputTokens; totalOutput += usage.outputTokens; totalCacheRead += usage.cacheReadTokens; totalReasoning += usage.reasoningTokens; requestCount += usage.requestCount; modelCounts[model, default: 0] += usage.requestCount; dailyInput[day, default: 0] += usage.inputTokens; dailyOutput[day, default: 0] += usage.outputTokens; dailyCacheRead[day, default: 0] += usage.cacheReadTokens; dailyReasoning[day, default: 0] += usage.reasoningTokens; dailyTokens[day, default: 0] += usage.totalTokens; dailyRequests[day, default: 0] += usage.requestCount; dailyModels[day, default: [:]][model, default: 0] += usage.requestCount; if let dayDate = Self.dateFromDayKey(day, calendar: calendar), dayDate > (lastEventAt ?? Date.distantPast) { lastEventAt = dayDate } } } + if requestCount == 0 { return MuseLocalSessionSummary(fileCount: fileCount, totalTokens: 0, totalInputTokens: 0, totalOutputTokens: 0, totalCacheReadTokens: 0, totalReasoningTokens: 0, requestCount: 0, lastEventAt: nil, primaryModel: nil, models: [], scannedAt: now) } + let sortedModels = modelCounts.sorted { $0.value > $1.value }.map(\.key) + let daily = dailyTokens.keys.sorted().map { day in let models = (dailyModels[day] ?? [:]).sorted { $0.value > $1.value }.map(\.key); return MuseLocalDailyBucket(date: day, totalTokens: dailyTokens[day] ?? 0, inputTokens: dailyInput[day] ?? 0, outputTokens: dailyOutput[day] ?? 0, cacheReadTokens: dailyCacheRead[day] ?? 0, reasoningTokens: dailyReasoning[day] ?? 0, requestCount: dailyRequests[day] ?? 0, models: models) } + return MuseLocalSessionSummary(fileCount: fileCount, totalTokens: totalTokens, totalInputTokens: totalInput, totalOutputTokens: totalOutput, totalCacheReadTokens: totalCacheRead, totalReasoningTokens: totalReasoning, requestCount: requestCount, lastEventAt: lastEventAt, primaryModel: sortedModels.first, models: sortedModels, daily: daily, scannedAt: now) + } + private static func dayKeyFromPath(_ path: String, calendar: Calendar) -> String? { + guard let range = path.range(of: "/sessions/") else { return nil } + let suffix = String(path[range.upperBound...]); let parts = suffix.split(separator: "/") + guard parts.count >= 3 else { return nil } + let y = String(parts[0]), m = String(parts[1]), d = String(parts[2]) + guard y.count == 4, m.count == 2, d.count == 2 else { return nil } + return "\(y)-\(m)-\(d)" + } + private static func dateFromDayKey(_ key: String, calendar: Calendar) -> Date? { + let parts = key.split(separator: "-") + guard parts.count == 3, let y = Int(parts[0]), let m = Int(parts[1]), let d = Int(parts[2]) else { return nil } + var comps = DateComponents(); comps.year = y; comps.month = m; comps.day = d + return calendar.date(from: comps) + } + private struct MuseParseResult { var contributions: [String: [String: MusePackedUsage]] = [:]; var parsedBytes: Int64 = 0; var entryCount: Int = 0 } + private static func parseMuseSessionFileFull(fileURL: URL, calendar: Calendar, lookbackCutoff: Date) -> MuseParseResult { + self.parseMuseSessionFileDelta(fileURL: fileURL, startOffset: 0, calendar: calendar, lookbackCutoff: lookbackCutoff) + } + private static func isCompleteJSONFragment(_ data: Data) -> Bool { + // Check if data, when trimmed, forms a complete JSON object/array. + // Incomplete trailing writes (e.g., truncated mid-object) will not be valid JSON and will not end with } or ]. + // Complete but irrelevant records (valid JSON) and wrapped records are valid JSON. + // Complete malformed lines that are valid JSON structure but not Muse (e.g., missing fields) are still valid JSON. + let trimmed = data.drop(while: { $0 == 32 || $0 == 9 || $0 == 10 || $0 == 13 }) + guard !trimmed.isEmpty else { return false } + // Quick check: must end with } or ] after trimming trailing whitespace + let reversedTrimmed = Data(trimmed.reversed().drop(while: { $0 == 32 || $0 == 9 || $0 == 10 || $0 == 13 })) + guard let last = reversedTrimmed.first, last == UInt8(ascii: "}") || last == UInt8(ascii: "]") else { + return false + } + // Try JSONSerialization to confirm it's complete valid JSON + do { + _ = try JSONSerialization.jsonObject(with: Data(trimmed), options: [.allowFragments]) + return true + } catch { + // Even if JSONSerialization fails due to inner malformed content but still ends with } , treat as complete malformed line that should be consumed + // e.g., "{ not valid json" does NOT end with }, so already returned false above. + // For cases like '{"a":}' which ends with } but is malformed, we still want to consume it once. + // If it ends with } but failed to parse, consider it a complete malformed line. + return true + } + } + private static func parseMuseSessionFileDelta(fileURL: URL, startOffset: Int64, calendar: Calendar, lookbackCutoff: Date) -> MuseParseResult { + var result = MuseParseResult() + guard let fullData = try? Data(contentsOf: fileURL) else { return result } + let fileSize = Int64(fullData.count) + guard startOffset <= fileSize else { return result } + let sliceData: Data + if startOffset > 0 { + sliceData = fullData.suffix(from: Int(startOffset)) + } else { + sliceData = fullData + } + // Determine parsedBytes using parse validity to distinguish valid final line without newline vs incomplete trailing fragment. + let lastNewlineOffset: Int? = sliceData.lastIndex(of: 10).map { sliceData.distance(from: sliceData.startIndex, to: $0) } + let isEndsWithNewline = !sliceData.isEmpty && sliceData.last == 10 + var effectiveData: Data + var parsedBytes: Int64 + if isEndsWithNewline { + effectiveData = sliceData + parsedBytes = fileSize + } else if let lastNL = lastNewlineOffset { + let trailing = Data(sliceData.suffix(from: lastNL + 1)) + if trailing.isEmpty { + effectiveData = sliceData + parsedBytes = fileSize + } else if Self.isCompleteJSONFragment(trailing) { + // Trailing bytes form a complete valid (or complete malformed) JSON record without newline — count it now + effectiveData = sliceData + parsedBytes = fileSize + } else { + // Incomplete trailing fragment — leave for next scan + effectiveData = sliceData.prefix(lastNL + 1) + parsedBytes = startOffset + Int64(lastNL + 1) + } + } else { + // No newline in slice: either single complete line without newline or incomplete single line + if Self.isCompleteJSONFragment(sliceData) { + effectiveData = sliceData + parsedBytes = fileSize + } else { + // Incomplete single line — leave for retry + effectiveData = Data() + parsedBytes = startOffset + } + } + var didParseAny = false + let lines = effectiveData.split(separator: 10, omittingEmptySubsequences: false) + for lineSlice in lines { + if lineSlice.isEmpty { continue } + let lineData = Data(lineSlice) + var events: [ParsedEvent] = [] + if let ev = self.parseLine(lineData) { events.append(ev) } else { for inner in self.extractRecordJSONDatas(from: lineData) { if let ev = self.parseLine(inner) { events.append(ev) } } } + if !events.isEmpty { didParseAny = true } + for ev in events { + guard ev.date >= lookbackCutoff else { continue } + guard let day = self.dayKey(for: ev.date, calendar: calendar) else { continue } + let total = ev.inputTokens + ev.outputTokens; guard total > 0 else { continue } + let packed = MusePackedUsage(inputTokens: ev.inputTokens, cacheReadTokens: ev.cacheReadTokens, outputTokens: ev.outputTokens, reasoningTokens: ev.reasoningTokens, totalTokens: total, requestCount: 1) + var dayMap = result.contributions[day] ?? [:] + let existing = dayMap[ev.model] ?? MusePackedUsage() + dayMap[ev.model] = existing + packed + result.contributions[day] = dayMap + result.entryCount += 1 + } + } + // For the no-newline single-line case where we left effectiveData empty (incomplete), ensure we don't have didParseAny + // No further adjustment needed; parsedBytes already set correctly. + _ = didParseAny + result.parsedBytes = parsedBytes + return result + } +} diff --git a/Sources/CodexBarCore/Providers/Muse/MuseProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Muse/MuseProviderDescriptor.swift new file mode 100644 index 0000000000..eff2255731 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Muse/MuseProviderDescriptor.swift @@ -0,0 +1,79 @@ +import Foundation + +public enum MuseProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .muse, + metadata: ProviderMetadata( + id: .muse, + displayName: "Muse", + sessionLabel: "Today", + weeklyLabel: "30-day", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Muse usage", + cliName: "muse", + defaultEnabled: false, + widgetSelectable: false, + dashboardURL: nil, + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .init(provider: .muse), + iconResourceName: "ProviderIcon-muse", + color: ProviderColor(red: 114 / 255, green: 96 / 255, blue: 255 / 255), + confettiPalette: [ + ProviderColor(hex: 0x7260FF), + ProviderColor(hex: 0x1A1A1A), + ProviderColor(hex: 0xEDE8FF), + ]), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: true, + noDataMessage: { + "No Muse sessions found in ~/.local/share/muse/sessions." + }, + supportsTokenSnapshot: true, + estimateDisclaimer: "From local Muse session logs; tokens only, no billing."), + pace: .unsupported, + history: .optIn, + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [MuseLocalFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "muse", + versionDetector: { _ in ProviderVersionDetector.museVersion() })) + } +} + +struct MuseLocalFetchStrategy: ProviderFetchStrategy { + let id: String = "muse.local" + let kind: ProviderFetchKind = .localProbe + + func isAvailable(_: ProviderFetchContext) async -> Bool { true } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let summary = try await MuseLocalSessionScanner.summarizeOffMainThread( + env: context.env, + lookbackDays: context.costUsageHistoryDays, + now: Date()) + guard summary.toCostUsageTokenSnapshot(historyDays: context.costUsageHistoryDays) != nil else { + throw MuseLocalError.noUsage + } + let snapshot = MuseUsageSnapshot(summary: summary, updatedAt: summary.scannedAt) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "local") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { false } +} + +private enum MuseLocalError: LocalizedError, Sendable { + case noUsage + var errorDescription: String? { + "No Muse sessions found in ~/.local/share/muse/sessions." + } +} diff --git a/Sources/CodexBarCore/Providers/Muse/MuseSessionCostCache.swift b/Sources/CodexBarCore/Providers/Muse/MuseSessionCostCache.swift new file mode 100644 index 0000000000..c603eb35f8 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Muse/MuseSessionCostCache.swift @@ -0,0 +1,160 @@ +import Foundation + +// MARK: - Cache structures (mirrors PiSessionCostCache pattern, simplified for Muse) + +enum MuseSessionCostCacheIO { + private static let artifactVersion = 1 + + private static func defaultCacheRoot() -> URL { + // Prefer the standard Caches directory, but fall back to a sandbox-allowed temp location + // when running under the `muse.bash` Managed sandbox (which denies writes to ~/Library/Caches). + if let root = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first { + let codexRoot = root.appendingPathComponent("CodexBar", isDirectory: true) + // Probe writability; `isWritableFile` is more reliable than trying to create and catching. + if FileManager.default.isWritableFile(atPath: root.path) || FileManager.default.isWritableFile(atPath: codexRoot.path) || FileManager.default.fileExists(atPath: codexRoot.path) { + return codexRoot + } + // Try to create the directory as a probe; if it succeeds, use it, otherwise fall back. + if (try? FileManager.default.createDirectory(at: codexRoot, withIntermediateDirectories: true)) != nil, + FileManager.default.isWritableFile(atPath: codexRoot.path) { + return codexRoot + } + } + // Fallback for sandboxed shells (e.g., `muse.bash`): use the process temp directory. + return FileManager.default.temporaryDirectory.appendingPathComponent("CodexBar", isDirectory: true) + } + + static func cacheFileURL(cacheRoot: URL? = nil) -> URL { + let root = cacheRoot ?? self.defaultCacheRoot() + return root + .appendingPathComponent("cost-usage", isDirectory: true) + .appendingPathComponent("muse-sessions-v\(Self.artifactVersion).json", isDirectory: false) + } + + static func load(cacheRoot: URL? = nil) -> MuseSessionCostCache { + let urls: [URL] = { + if let cacheRoot { + return [self.cacheFileURL(cacheRoot: cacheRoot)] + } + // Try default, then fallback temp for sandboxed shells + let defaultURL = self.cacheFileURL(cacheRoot: nil) + let fallbackURL = self.cacheFileURL(cacheRoot: FileManager.default.temporaryDirectory.appendingPathComponent("CodexBar", isDirectory: true)) + return [defaultURL, fallbackURL] + }() + for url in urls { + if let data = try? Data(contentsOf: url), + let decoded = try? JSONDecoder().decode(MuseSessionCostCache.self, from: data), + decoded.version == Self.artifactVersion { + return decoded + } + } + return MuseSessionCostCache(version: Self.artifactVersion) + } + + static func save(cache: MuseSessionCostCache, cacheRoot: URL? = nil, calendar: Calendar = .current) { + var cache = cache + cache.timeZoneIdentifier = calendar.timeZone.identifier + let urls: [URL] = { + if let cacheRoot { + return [self.cacheFileURL(cacheRoot: cacheRoot)] + } + let defaultURL = self.cacheFileURL(cacheRoot: nil) + let fallbackURL = self.cacheFileURL(cacheRoot: FileManager.default.temporaryDirectory.appendingPathComponent("CodexBar", isDirectory: true)) + return [defaultURL, fallbackURL] + }() + let data = (try? JSONEncoder().encode(cache)) ?? Data() + var saved = false + for url in urls { + let dir = url.deletingLastPathComponent() + try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + let tmp = dir.appendingPathComponent(".tmp-\(UUID().uuidString).json", isDirectory: false) + 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) + } + saved = true + break + } catch { + try? FileManager.default.removeItem(at: tmp) + continue + } + } + if !saved { + // Last resort: try workspace temp atomically + let fallback = FileManager.default.temporaryDirectory.appendingPathComponent("CodexBar/cost-usage/muse-sessions-v\(Self.artifactVersion).json") + try? FileManager.default.createDirectory(at: fallback.deletingLastPathComponent(), withIntermediateDirectories: true) + let tmpFallback = fallback.deletingLastPathComponent().appendingPathComponent(".tmp-\(UUID().uuidString).json", isDirectory: false) + if let _ = try? data.write(to: tmpFallback, options: [.atomic]) { + if FileManager.default.fileExists(atPath: fallback.path) { + _ = try? FileManager.default.replaceItemAt(fallback, withItemAt: tmpFallback) + } else { + try? FileManager.default.moveItem(at: tmpFallback, to: fallback) + } + } + } + } + + static func clear(cacheRoot: URL? = nil) { + let url = self.cacheFileURL(cacheRoot: cacheRoot) + try? FileManager.default.removeItem(at: url) + } +} + +struct MuseSessionCostCache: Codable { + var version: Int + var lastScanUnixMs: Int64 = 0 + var timeZoneIdentifier: String? + // dayKey -> model -> packed usage + var days: [String: [String: MusePackedUsage]] = [:] + // file path -> per-file usage + var files: [String: MuseSessionFileUsage] = [:] + + init(version: Int = 1) { + self.version = version + } +} + +struct MuseSessionFileUsage: Codable, Equatable { + var mtimeUnixMs: Int64 + var size: Int64 + var parsedBytes: Int64 + var prefixFingerprint: String? // hash of first 4K for same-path replacement detection + var contributions: [String: [String: MusePackedUsage]] // day -> model -> usage + var entryCount: Int +} + +struct MusePackedUsage: Codable, Equatable { + var inputTokens: Int = 0 + var cacheReadTokens: Int = 0 + var outputTokens: Int = 0 + var reasoningTokens: Int = 0 + var totalTokens: Int = 0 + var requestCount: Int = 0 + + var isZero: Bool { + self.totalTokens == 0 && self.requestCount == 0 + } + + static func +(lhs: MusePackedUsage, rhs: MusePackedUsage) -> MusePackedUsage { + MusePackedUsage( + inputTokens: lhs.inputTokens + rhs.inputTokens, + cacheReadTokens: lhs.cacheReadTokens + rhs.cacheReadTokens, + outputTokens: lhs.outputTokens + rhs.outputTokens, + reasoningTokens: lhs.reasoningTokens + rhs.reasoningTokens, + totalTokens: lhs.totalTokens + rhs.totalTokens, + requestCount: lhs.requestCount + rhs.requestCount) + } + + static func -(lhs: MusePackedUsage, rhs: MusePackedUsage) -> MusePackedUsage { + MusePackedUsage( + inputTokens: lhs.inputTokens - rhs.inputTokens, + cacheReadTokens: lhs.cacheReadTokens - rhs.cacheReadTokens, + outputTokens: lhs.outputTokens - rhs.outputTokens, + reasoningTokens: lhs.reasoningTokens - rhs.reasoningTokens, + totalTokens: lhs.totalTokens - rhs.totalTokens, + requestCount: lhs.requestCount - rhs.requestCount) + } +} diff --git a/Sources/CodexBarCore/Providers/Muse/MuseUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Muse/MuseUsageSnapshot.swift new file mode 100644 index 0000000000..1351ee573a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Muse/MuseUsageSnapshot.swift @@ -0,0 +1,28 @@ +import Foundation + +public struct MuseUsageSnapshot: Sendable { + public let summary: MuseLocalSessionSummary? + public let updatedAt: Date + + public init(summary: MuseLocalSessionSummary?, updatedAt: Date = Date()) { + self.summary = summary + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let costUsage = summary?.toCostUsageTokenSnapshot( + historyDays: MuseLocalSessionScanner.defaultLookbackDays) + let identity = ProviderIdentitySnapshot( + providerID: .muse, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "Muse") + // Token-history only: no quota windows, no cost, no pace. + return UsageSnapshot( + primary: nil, + secondary: nil, + costUsage: costUsage, + updatedAt: summary?.scannedAt ?? self.updatedAt, + identity: costUsage != nil || summary?.requestCount ?? 0 > 0 ? identity : nil) + } +} diff --git a/Sources/CodexBarCore/Providers/ProviderInstanceIDAliases.generated.swift b/Sources/CodexBarCore/Providers/ProviderInstanceIDAliases.generated.swift index 1a6682c057..73a756df14 100644 --- a/Sources/CodexBarCore/Providers/ProviderInstanceIDAliases.generated.swift +++ b/Sources/CodexBarCore/Providers/ProviderInstanceIDAliases.generated.swift @@ -3,6 +3,7 @@ // First-party spellings keep call sites concise while their values remain validated instance IDs. // swiftformat:disable sortDeclarations extension ProviderInstanceID { + public static let muse = UsageProvider.muse.instanceID public static let codex = UsageProvider.codex.instanceID public static let openai = UsageProvider.openai.instanceID public static let azureopenai = UsageProvider.azureopenai.instanceID diff --git a/Sources/CodexBarCore/Providers/ProviderManifest.swift b/Sources/CodexBarCore/Providers/ProviderManifest.swift index 47631b6564..e4f06b9ae1 100644 --- a/Sources/CodexBarCore/Providers/ProviderManifest.swift +++ b/Sources/CodexBarCore/Providers/ProviderManifest.swift @@ -5,6 +5,7 @@ import Foundation /// `ProviderDescriptorRegistry.register(_:)`. public enum ProviderManifest { public static let allDescriptors: [ProviderDescriptor] = [ + MuseProviderDescriptor.descriptor, CodexProviderDescriptor.descriptor, OpenAIAPIProviderDescriptor.descriptor, AzureOpenAIProviderDescriptor.descriptor, diff --git a/Sources/CodexBarCore/Providers/ProviderVersionDetector.swift b/Sources/CodexBarCore/Providers/ProviderVersionDetector.swift index 9358f70359..68763ebdca 100644 --- a/Sources/CodexBarCore/Providers/ProviderVersionDetector.swift +++ b/Sources/CodexBarCore/Providers/ProviderVersionDetector.swift @@ -238,6 +238,11 @@ public enum ProviderVersionDetector { return nil } + public static func museVersion() -> String? { + guard let path = TTYCommandRunner.which(MuseProviderDescriptor.descriptor.cli.name) else { return nil } + return Self.run(path: path, args: ["--version"]) + } + public static func geminiVersion() -> String? { let env = ProcessInfo.processInfo.environment guard let path = BinaryLocator.resolveGeminiBinary(env: env, loginPATH: nil) diff --git a/Sources/CodexBarCore/Providers/Providers.swift b/Sources/CodexBarCore/Providers/Providers.swift index fbdbdea033..06191ace7d 100644 --- a/Sources/CodexBarCore/Providers/Providers.swift +++ b/Sources/CodexBarCore/Providers/Providers.swift @@ -19,6 +19,7 @@ public struct ProviderDebugPaneCapabilities: Sendable { // swiftformat:disable sortDeclarations public enum UsageProvider: String, CaseIterable, Sendable, Codable { + case muse case codex case openai case azureopenai diff --git a/Tests/CodexBarTests/MuseIncrementalCacheTests.swift b/Tests/CodexBarTests/MuseIncrementalCacheTests.swift new file mode 100644 index 0000000000..ebddf8f8c0 --- /dev/null +++ b/Tests/CodexBarTests/MuseIncrementalCacheTests.swift @@ -0,0 +1,543 @@ +import Testing +import Foundation +@testable import CodexBarCore + +struct MuseIncrementalCacheTests { + private func makeSessionLine(recordedAtMicroseconds: Int64, inputTokens: Int, outputTokens: Int, model: String = "muse-spark-1.2-contributor") -> String { + let outer: [String: Any] = [ + "recorded_at": recordedAtMicroseconds, + "payload_type": "runtime.session", + "payload": [ + "kind": "run", + "event": [ + "kind": "model_completed", + "usage": [ + "input_tokens": inputTokens, + "output_tokens": outputTokens, + "cached_tokens": 0, + "cache_write_tokens": 0, + "cache_read_tokens": 0, + "reasoning_tokens": 0, + ], + "duration_ms": 1000, + "finish_reason": "tool_calls", + "model": model, + ], + ] as [String: Any], + ] + let data = try! JSONSerialization.data(withJSONObject: outer, options: []) + return String(data: data, encoding: .utf8)! + } + + private func makeWrappedLine(recordedAtMicroseconds: Int64, inputTokens: Int, outputTokens: Int) -> String { + let inner: [String: Any] = [ + "recorded_at": recordedAtMicroseconds, + "payload_type": "runtime.session", + "payload": [ + "kind": "run", + "event": [ + "kind": "model_completed", + "usage": [ + "input_tokens": inputTokens, + "output_tokens": outputTokens, + "cached_tokens": 0, + "cache_write_tokens": 0, + "cache_read_tokens": 0, + "reasoning_tokens": 0, + ], + "duration_ms": 1000, + "finish_reason": "tool_calls", + "model": "muse-spark-1.2-contributor", + ], + ] as [String: Any], + ] + let innerData = try! JSONSerialization.data(withJSONObject: inner, options: []) + let innerString = String(data: innerData, encoding: .utf8)! + let outer: [String: Any] = [ + "recorded_at": recordedAtMicroseconds, + "payload_type": "retained_frame", + "record_json": innerString, + ] + let outerData = try! JSONSerialization.data(withJSONObject: outer, options: []) + return String(data: outerData, encoding: .utf8)! + } + + private func microseconds(for date: Date) -> Int64 { Int64(date.timeIntervalSince1970 * 1_000_000) } + + @Test func unchangedSecondScanOpensZeroFiles() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-unchanged-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let line = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + var firstScanned: [URL] = [] + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now, fileScanObserver: { firstScanned.append($0) }) + #expect(first.requestCount == 1) + #expect(firstScanned.count == 1) + + var secondScanned: [URL] = [] + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now, fileScanObserver: { secondScanned.append($0) }) + #expect(second.requestCount == 1) + #expect(secondScanned.count == 0) // unchanged file not reopened + #expect(second.totalTokens == first.totalTokens) + } + + @Test func appendedFileParsesOnlyNewContent() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-append-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let fileURL = sessionDir.appendingPathComponent("session.jsonl") + let line1 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line1.write(to: fileURL, atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 1) + #expect(first.totalTokens == 110) + + // Append second event + let line2 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(60)), inputTokens: 200, outputTokens: 20) + let handle = try FileHandle(forWritingTo: fileURL) + try handle.seekToEnd() + try handle.write(contentsOf: (line2 + "\n").data(using: .utf8)!) + try handle.close() + + var appendedScanned: [URL] = [] + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now, fileScanObserver: { appendedScanned.append($0) }) + #expect(appendedScanned.count == 1) // file was reopened for delta + #expect(second.requestCount == 2) + #expect(second.totalTokens == 330) + } + + @Test func newFileAddsUsage() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-new-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let dir1 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess1", isDirectory: true) + try FileManager.default.createDirectory(at: dir1, withIntermediateDirectories: true) + let line1 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line1.write(to: dir1.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 1) + + let dir2 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess2", isDirectory: true) + try FileManager.default.createDirectory(at: dir2, withIntermediateDirectories: true) + let line2 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 200, outputTokens: 20) + try line2.write(to: dir2.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + var secondScanned: [URL] = [] + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now, fileScanObserver: { secondScanned.append($0) }) + #expect(secondScanned.count == 1) // only new file opened + #expect(second.requestCount == 2) + #expect(second.totalTokens == 330) + } + + @Test func truncatedFileReparsesCorrectly() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-trunc-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let fileURL = sessionDir.appendingPathComponent("session.jsonl") + let line1 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + let line2 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(60)), inputTokens: 200, outputTokens: 20) + try (line1 + "\n" + line2).write(to: fileURL, atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 2) + #expect(first.totalTokens == 330) + + // Truncate to single event (replace) + let line3 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 999, outputTokens: 99) + try line3.write(to: fileURL, atomically: true, encoding: .utf8) + + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(second.requestCount == 1) + #expect(second.totalTokens == 1098) + } + + @Test func deletedFileRemovesContribution() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-delete-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let dir1 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess1", isDirectory: true) + let dir2 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess2", isDirectory: true) + try FileManager.default.createDirectory(at: dir1, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: dir2, withIntermediateDirectories: true) + let line1 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + let line2 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 200, outputTokens: 20) + try line1.write(to: dir1.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try line2.write(to: dir2.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 2) + + try FileManager.default.removeItem(at: dir2.appendingPathComponent("session.jsonl")) + + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(second.requestCount == 1) + #expect(second.totalTokens == 110) + } + + @Test func lookbackExpiryAgesOut() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-expiry-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let oldDate = localCalendar.date(byAdding: .day, value: -10, to: now)! + let oldKey = MuseLocalSessionScanner.dayKey(for: oldDate, calendar: localCalendar)! + let recentKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let oldDir = root.appendingPathComponent("sessions/\(oldKey.replacingOccurrences(of: "-", with: "/"))/old", isDirectory: true) + let recentDir = root.appendingPathComponent("sessions/\(recentKey.replacingOccurrences(of: "-", with: "/"))/recent", isDirectory: true) + try FileManager.default.createDirectory(at: oldDir, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: recentDir, withIntermediateDirectories: true) + let oldLine = makeSessionLine(recordedAtMicroseconds: microseconds(for: oldDate), inputTokens: 100, outputTokens: 10) + let recentLine = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 200, outputTokens: 20) + try oldLine.write(to: oldDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try recentLine.write(to: recentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let with7 = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 7, now: now) + #expect(with7.requestCount == 1) + #expect(with7.totalTokens == 220) + + let with30 = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(with30.requestCount == 2) + #expect(with30.totalTokens == 330) + } + + @Test func wrappedRecordJsonIsSupported() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-wrapped-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let wrapped = makeWrappedLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 500, outputTokens: 50) + try wrapped.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + let summary = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(summary.requestCount == 1) + #expect(summary.totalTokens == 550) + } + + // MARK: - Audit regression tests + + @Test func largerReplacementAtSamePathReparsesInsteadOfAppending() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-larger-replace-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let fileURL = sessionDir.appendingPathComponent("session.jsonl") + let lineA = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10, model: "model-a") + try lineA.write(to: fileURL, atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 1) + #expect(first.totalTokens == 110) + + // Replace with larger completely different content (2 events, different prefix, larger size) + let lineB = makeSessionLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(10)), inputTokens: 200, outputTokens: 20, model: "model-bxxx-different-prefix-to-ensure-fingerprint-mismatch-1234567890") + let lineC = makeSessionLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(20)), inputTokens: 300, outputTokens: 30, model: "model-c") + let replacement = [lineB, lineC].joined(separator: "\n") + // Ensure larger + #expect(replacement.count > lineA.count) + try replacement.write(to: fileURL, atomically: true, encoding: .utf8) + + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + // Must be 2, not 3 (which would be 1 old + 2 new if mistakenly appended) + #expect(second.requestCount == 2) + #expect(second.totalTokens == 550) + } + + @Test func timezoneChangeRebucketsOrInvalidates() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-tz-\(UUID().uuidString)", isDirectory: true) + // Use a date near midnight to test bucket movement across timezones + var utcCalendar = Calendar(identifier: .gregorian) + utcCalendar.timeZone = TimeZone(identifier: "UTC")! + let nowUTC = utcCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 23, minute: 30))! + let dayKeyUTC = MuseLocalSessionScanner.dayKey(for: nowUTC, calendar: utcCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKeyUTC.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let line = makeSessionLine(recordedAtMicroseconds: microseconds(for: nowUTC), inputTokens: 100, outputTokens: 10) + try line.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + // First scan with UTC calendar (simulate by setting Cache timezone to UTC) + var cache = MuseSessionCostCache(version: 1) + cache.timeZoneIdentifier = "UTC" + MuseSessionCostCacheIO.save(cache: cache, cacheRoot: root, calendar: utcCalendar) + // Now summarize with a different timezone (Asia/Tokyo) — scanner uses Calendar.current, so we simulate by + // directly testing cache invalidation logic: save a cache with UTC identifier, then load with Tokyo calendar + // should invalidate (tested via summarizeCancellable which checks Calendar.current) + // Instead we test the raw cache load path: write cache with UTC, then create a calendar with Tokyo and verify + // that a fresh summarize with Tokyo would reset. + // We simulate by manually invoking the timezone check: cache.timeZoneIdentifier != Calendar.current + // For deterministic test, we directly verify that MuseSessionCostCacheIO load respects version but not timezone, + // and that summarize resets when timezone differs. + + // Create a cache with UTC day + let firstSummary = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: nowUTC) + #expect(firstSummary.requestCount == 1) + + // Now simulate timezone change by writing a cache with stale timezone and verifying next summarize rebuilds + var staleCache = MuseSessionCostCacheIO.load(cacheRoot: root) + staleCache.timeZoneIdentifier = "Etc/GMT-12" // different from current + MuseSessionCostCacheIO.save(cache: staleCache, cacheRoot: root, calendar: utcCalendar) + // Next summarize should detect timezone mismatch and rebuild (not double-count) + let secondSummary = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: nowUTC) + #expect(secondSummary.requestCount == 1) + #expect(secondSummary.totalTokens == 110) + // Ensure cache after second scan has current timezone identifier + let reloaded = MuseSessionCostCacheIO.load(cacheRoot: root) + #expect(reloaded.timeZoneIdentifier == Calendar.current.timeZone.identifier) + } + + @Test func historicalCachePruningRemovesExpiredFilesAndDays() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-prune-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let oldDate = localCalendar.date(byAdding: .day, value: -60, to: now)! + let oldKey = MuseLocalSessionScanner.dayKey(for: oldDate, calendar: localCalendar)! + let recentKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let oldDir = root.appendingPathComponent("sessions/\(oldKey.replacingOccurrences(of: "-", with: "/"))/old", isDirectory: true) + let recentDir = root.appendingPathComponent("sessions/\(recentKey.replacingOccurrences(of: "-", with: "/"))/recent", isDirectory: true) + try FileManager.default.createDirectory(at: oldDir, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: recentDir, withIntermediateDirectories: true) + let oldLine = makeSessionLine(recordedAtMicroseconds: microseconds(for: oldDate), inputTokens: 100, outputTokens: 10) + let recentLine = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 200, outputTokens: 20) + try oldLine.write(to: oldDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try recentLine.write(to: recentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + // Initial 30-day scan includes only recent, but we first do a 90-day scan to populate old entry + let wide = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 90, now: now) + #expect(wide.requestCount == 2) + let cacheBefore = MuseSessionCostCacheIO.load(cacheRoot: root) + #expect(cacheBefore.files.count == 2) + #expect(cacheBefore.days.count >= 2) + let sizeBefore = try Data(contentsOf: MuseSessionCostCacheIO.cacheFileURL(cacheRoot: root)).count + + // Now narrow to 7 days: old should be pruned from persisted cache, not just filtered in summary + let narrow = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 7, now: now) + #expect(narrow.requestCount == 1) + let cacheAfter = MuseSessionCostCacheIO.load(cacheRoot: root) + #expect(cacheAfter.files.count == 1) + // Old day should be physically removed + #expect(cacheAfter.days[oldKey] == nil) + let sizeAfter = try Data(contentsOf: MuseSessionCostCacheIO.cacheFileURL(cacheRoot: root)).count + #expect(sizeAfter < sizeBefore) + #expect(sizeAfter > 0) + } + + @Test func corruptCacheRecoveryDoesNotCrash() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-corrupt-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let line = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let valid = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(valid.requestCount == 1) + + let cacheURL = MuseSessionCostCacheIO.cacheFileURL(cacheRoot: root) + // Malformed JSON + try Data("not json at all".utf8).write(to: cacheURL) + let afterMalformed = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(afterMalformed.requestCount == 1) + + // Truncated JSON + let validData = try Data(contentsOf: cacheURL) + try validData.prefix(validData.count / 2).write(to: cacheURL) + let afterTruncated = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(afterTruncated.requestCount == 1) + + // Incompatible version + var cache = MuseSessionCostCache(version: 999) + cache.days = ["2099-01-01": ["m": MusePackedUsage(inputTokens: 999, outputTokens: 999, totalTokens: 1998, requestCount: 1)]] + let badData = try JSONEncoder().encode(cache) + try badData.write(to: cacheURL) + let afterVersion = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(afterVersion.requestCount == 1) + let reloaded = MuseSessionCostCacheIO.load(cacheRoot: root) + #expect(reloaded.version == 1) + #expect(reloaded.days["2099-01-01"] == nil) + } + + @Test func atomicWritesProduceValidJSON() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-atomic-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let line = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + for _ in 0..<5 { + _ = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + let cacheURL = MuseSessionCostCacheIO.cacheFileURL(cacheRoot: root) + let data = try Data(contentsOf: cacheURL) + // Must be valid JSON and decode + let decoded = try JSONDecoder().decode(MuseSessionCostCache.self, from: data) + #expect(decoded.version == 1) + // No partial writes: file should not be empty or truncated + #expect(data.count > 10) + } + } + + @Test func cancellationDoesNotPersistPartialCache() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-cancel-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let dir1 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/a", isDirectory: true) + let dir2 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/b", isDirectory: true) + try FileManager.default.createDirectory(at: dir1, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: dir2, withIntermediateDirectories: true) + let line1 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + let line2 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 200, outputTokens: 20) + try line1.write(to: dir1.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try line2.write(to: dir2.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + // Prime cache with first file only + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 2) + let cacheBefore = MuseSessionCostCacheIO.load(cacheRoot: root) + let beforeData = try JSONEncoder().encode(cacheBefore) + + // Now append to one file and attempt cancellable scan that throws mid-way + let line3 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(60)), inputTokens: 300, outputTokens: 30) + let file1URL = dir1.appendingPathComponent("session.jsonl") + let handle = try FileHandle(forWritingTo: file1URL) + try handle.seekToEnd() + try handle.write(contentsOf: ("\n" + line3).data(using: .utf8)!) + try handle.close() + + var checkCount = 0 + do { + _ = try MuseLocalSessionScanner.summarizeCancellable(env: ["MUSE_HOME": root.path], fileManager: .default, lookbackDays: 30, now: now, fileScanObserver: nil, checkCancellation: { + checkCount += 1 + if checkCount == 3 { throw CancellationError() } // after first file, before second — tests partial not persisted + }) + Issue.record("Expected cancellation") + } catch is CancellationError { + // Expected + } + + let cacheAfter = MuseSessionCostCacheIO.load(cacheRoot: root) + // Cache must be unchanged (no partial delta persisted) — compare files/days, ignore lastScanUnixMs + #expect(cacheAfter.files == cacheBefore.files) + #expect(cacheAfter.days == cacheBefore.days) + + // A normal scan after cancellation should still correctly apply the delta exactly once + let afterRetry = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(afterRetry.requestCount == 3) + #expect(afterRetry.totalTokens == 660) + } + + @Test func appendBoundaryIncompleteTrailingLineNotSkipped() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-boundary-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let sessionDir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + let fileURL = sessionDir.appendingPathComponent("session.jsonl") + let line1 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + // Write first line with newline + try (line1 + "\n").write(to: fileURL, atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 1) + + // Append incomplete line without newline (simulate crash mid-write) + let line2 = makeSessionLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(10)), inputTokens: 200, outputTokens: 20) + let partial = String(line2.prefix(line2.count / 2)) // truncated JSON + let handle = try FileHandle(forWritingTo: fileURL) + try handle.seekToEnd() + try handle.write(contentsOf: partial.data(using: .utf8)!) + try handle.close() + + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + // Incomplete line must not be counted, and must not be permanently skipped + #expect(second.requestCount == 1) + + // Now complete the line by appending the remainder + newline + let remainder = String(line2.suffix(line2.count - partial.count)) + let handle2 = try FileHandle(forWritingTo: fileURL) + try handle2.seekToEnd() + try handle2.write(contentsOf: (remainder + "\n").data(using: .utf8)!) + try handle2.close() + + let third = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + // Must be exactly 2, not 1 and not 3 (double-count) + #expect(third.requestCount == 2) + #expect(third.totalTokens == 330) + + // Also verify wrapped format follows same boundary + let wrapped = makeWrappedLine(recordedAtMicroseconds: microseconds(for: now.addingTimeInterval(20)), inputTokens: 300, outputTokens: 30) + let partialWrapped = String(wrapped.prefix(wrapped.count / 2)) + try partialWrapped.write(to: fileURL, atomically: true, encoding: .utf8) // overwrite for isolated wrapped test + // Instead test wrapped append correctly: create new file + let dir2 = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/wrapped", isDirectory: true) + try FileManager.default.createDirectory(at: dir2, withIntermediateDirectories: true) + let wrappedFile = dir2.appendingPathComponent("session.jsonl") + let wrappedLine = makeWrappedLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 400, outputTokens: 40) + // Write with newline then incomplete wrapped + try (wrappedLine + "\n" + String(wrapped.prefix(10))).write(to: wrappedFile, atomically: true, encoding: .utf8) + let fourth = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + // Should count the wrappedLine (1) plus previous 2 from first file = 3 + #expect(fourth.requestCount == 3) + } + + @Test func deletedDirectoryRemovesContributions() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-del-dir-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let dayKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let dir = root.appendingPathComponent("sessions/\(dayKey.replacingOccurrences(of: "-", with: "/"))/todelete", isDirectory: true) + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + let line = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line.write(to: dir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 1) + + try FileManager.default.removeItem(at: dir) + + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(second.requestCount == 0) + let cache = MuseSessionCostCacheIO.load(cacheRoot: root) + #expect(cache.files.isEmpty) + #expect(cache.days.isEmpty) + } + + @Test func deletedFileNotDoubleSubtractedOnAgingOut() throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent("muse-inc-del-aging-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let recentKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let recentDir = root.appendingPathComponent("sessions/\(recentKey.replacingOccurrences(of: "-", with: "/"))/recent", isDirectory: true) + try FileManager.default.createDirectory(at: recentDir, withIntermediateDirectories: true) + let line = makeSessionLine(recordedAtMicroseconds: microseconds(for: now), inputTokens: 100, outputTokens: 10) + try line.write(to: recentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + let first = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(first.requestCount == 1) + + // Delete file + try FileManager.default.removeItem(at: recentDir.appendingPathComponent("session.jsonl")) + let second = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: now) + #expect(second.requestCount == 0) + + // Aging out should not cause negative or double subtract + let later = localCalendar.date(byAdding: .day, value: 31, to: now)! + let third = MuseLocalSessionScanner.summarize(env: ["MUSE_HOME": root.path], lookbackDays: 30, now: later) + #expect(third.requestCount == 0) + let cache = MuseSessionCostCacheIO.load(cacheRoot: root) + #expect(cache.files.isEmpty) + } +} diff --git a/Tests/CodexBarTests/MuseLocalSessionScannerTests.swift b/Tests/CodexBarTests/MuseLocalSessionScannerTests.swift new file mode 100644 index 0000000000..121a131fc3 --- /dev/null +++ b/Tests/CodexBarTests/MuseLocalSessionScannerTests.swift @@ -0,0 +1,495 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct MuseLocalSessionScannerTests { + // MARK: - Helpers + + private func makeSessionLine( + recordedAtMicroseconds: Int64, + inputTokens: Int, + outputTokens: Int, + cachedTokens: Int = 0, + cacheReadTokens: Int? = nil, + reasoningTokens: Int = 0, + model: String = "muse-spark-1.2-contributor") -> String + { + let usage: [String: Any] = [ + "input_tokens": inputTokens, + "output_tokens": outputTokens, + "cached_tokens": cachedTokens, + "cache_write_tokens": 0, + "cache_read_tokens": cacheReadTokens ?? cachedTokens, + "reasoning_tokens": reasoningTokens, + ] + let event: [String: Any] = [ + "kind": "model_completed", + "usage": usage, + "duration_ms": 1000, + "finish_reason": "tool_calls", + "model": model, + ] + let payload: [String: Any] = ["kind": "run", "event": event] + let outer: [String: Any] = [ + "recorded_at": recordedAtMicroseconds, + "payload_type": "runtime.session", + "payload": payload, + ] + let data = try! JSONSerialization.data(withJSONObject: outer, options: []) + return String(data: data, encoding: .utf8)! + } + + private func makeGoalAttributionLine( + recordedAtMicroseconds: Int64, + inputTokens: Int, + outputTokens: Int) -> String + { + // Intentionally same token counts — must NOT be counted. + let outer: [String: Any] = [ + "recorded_at": recordedAtMicroseconds, + "payload_type": "runtime.session", + "payload": [ + "kind": "run", + "event": [ + "kind": "goal_usage_attribution", + "record": [ + "usage_family": "provider", + "quantity": [ + "unit": "tokens", + "reported": true, + "input_tokens": inputTokens, + "output_tokens": outputTokens, + ], + ], + ], + ] as [String: Any], + ] + let data = try! JSONSerialization.data(withJSONObject: outer, options: []) + return String(data: data, encoding: .utf8)! + } + + private func microseconds(for date: Date) -> Int64 { + Int64(date.timeIntervalSince1970 * 1_000_000) + } + + // MARK: - Tests + + @Test + func `single model_completed produces expected token usage`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-single-\(UUID().uuidString)", isDirectory: true) + let sessionDir = root.appendingPathComponent("sessions/2026/08/27/session-a", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + + let now = Date(timeIntervalSince1970: 1_787_840_784) // 2026-08-27 14:26:24 UTC, matches recorded_at micro convention + let line = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 20374, + outputTokens: 128, + reasoningTokens: 27) + + try line.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 30, + now: now) + #expect(summary.requestCount == 1) + #expect(summary.totalTokens == 20374 + 128) + #expect(summary.totalInputTokens == 20374) + #expect(summary.totalOutputTokens == 128) + #expect(summary.totalReasoningTokens == 27) + #expect(summary.daily.count == 1) + #expect(summary.daily.first?.totalTokens == 20502) + #expect(summary.daily.first?.requestCount == 1) + + let snapshot = try #require(summary.toCostUsageTokenSnapshot(historyDays: 30)) + #expect(snapshot.last30DaysTokens == 20502) + #expect(snapshot.last30DaysRequests == 1) + #expect(snapshot.sessionTokens == 20502) + #expect(snapshot.daily.first?.totalTokens == 20502) + #expect(snapshot.costProvenance == .unknown) + } + + @Test + func `goal_usage_attribution does not double-count`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-dedup-\(UUID().uuidString)", isDirectory: true) + let sessionDir = root.appendingPathComponent("sessions/2026/08/27/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + + let now = Date(timeIntervalSince1970: 1_787_840_784) + let goal = self.makeGoalAttributionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 20374, + outputTokens: 128) + let model = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now.addingTimeInterval(1)), + inputTokens: 20374, + outputTokens: 128) + + let combined = [goal, model].joined(separator: "\n") + try combined.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 30, + now: now) + // Only the model_completed counts; goal attribution is ignored. + #expect(summary.requestCount == 1) + #expect(summary.totalTokens == 20502) + } + + @Test + func `subagent usage is recursively included`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-subagent-\(UUID().uuidString)", isDirectory: true) + let parentDir = root.appendingPathComponent("sessions/2026/08/27/parent", isDirectory: true) + let subagentDir = parentDir.appendingPathComponent("subagent/child-1", isDirectory: true) + try FileManager.default.createDirectory(at: parentDir, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: subagentDir, withIntermediateDirectories: true) + + let now = Date(timeIntervalSince1970: 1_787_840_784) + let parentLine = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 1000, + outputTokens: 100) + let childLine = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now.addingTimeInterval(10)), + inputTokens: 500, + outputTokens: 50) + + try parentLine.write(to: parentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try childLine.write(to: subagentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let summary = MuseLocalSessionScanner.summarize( + env: ["CODEXBAR_MUSE_HOME": root.path], + lookbackDays: 30, + now: now) + #expect(summary.fileCount == 2) + #expect(summary.requestCount == 2) + #expect(summary.totalTokens == 1650) + } + + @Test + func `multiple events aggregate into daily and history buckets`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-multi-\(UUID().uuidString)", isDirectory: true) + let sessionDir = root.appendingPathComponent("sessions/2026/08/27/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(identifier: "UTC")! + let day1 = Date(timeIntervalSince1970: 1_787_840_784) // 2026-08-27 + let day0 = try #require(calendar.date(byAdding: .day, value: -1, to: day1)) + // Use microsecond wall-clock that matches dayKey via Calendar.current — set TZ to UTC for determinism + // but scanner uses Calendar.current; we force day keys via recorded_at near midnight. + // Instead use explicit dates and derive dayKey via scanner's calendar. + let l1 = self.makeSessionLine(recordedAtMicroseconds: self.microseconds(for: day0), inputTokens: 100, outputTokens: 10) + let l2 = self.makeSessionLine(recordedAtMicroseconds: self.microseconds(for: day0.addingTimeInterval(3600)), inputTokens: 200, outputTokens: 20) + let l3 = self.makeSessionLine(recordedAtMicroseconds: self.microseconds(for: day1), inputTokens: 300, outputTokens: 30) + + try [l1, l2, l3].joined(separator: "\n") + .write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 30, + now: day1) + #expect(summary.requestCount == 3) + #expect(summary.totalTokens == 660) + // daily buckets are at least 2 days + #expect(summary.daily.count >= 2) + let totalDaily = summary.daily.reduce(0) { $0 + $1.totalTokens } + #expect(totalDaily == 660) + + let snapshot = try #require(summary.toCostUsageTokenSnapshot(historyDays: 30)) + let window = snapshot.summary(forLastDays: 30) + #expect(window.totalTokens == 660) + #expect(window.totalRequests == 3) + } + + @Test + func `empty roots do not publish bogus usage`() { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-empty-\(UUID().uuidString)", isDirectory: true) + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 30, + now: Date()) + #expect(summary.requestCount == 0) + #expect(summary.toCostUsageTokenSnapshot(historyDays: 30) == nil) + let snap = MuseUsageSnapshot(summary: summary) + #expect(snap.toUsageSnapshot().costUsage == nil) + } + + @Test + func `malformed and unrelated lines are safely ignored`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-malformed-\(UUID().uuidString)", isDirectory: true) + let sessionDir = root.appendingPathComponent("sessions/2026/08/27/sess", isDirectory: true) + try FileManager.default.createDirectory(at: sessionDir, withIntermediateDirectories: true) + + let now = Date(timeIntervalSince1970: 1_787_840_784) + let good = self.makeSessionLine(recordedAtMicroseconds: self.microseconds(for: now), inputTokens: 100, outputTokens: 10) + let badJSON = "{ not valid json" + let missingKind = "{\"recorded_at\":\(self.microseconds(for: now)),\"payload_type\":\"runtime.session\",\"payload\":{\"kind\":\"run\",\"event\":{\"kind\":\"other\"}}}" + let wrongPayloadType = "{\"recorded_at\":\(self.microseconds(for: now)),\"payload_type\":\"other\",\"payload\":{\"kind\":\"run\",\"event\":{\"kind\":\"model_completed\"}}}" + // Sensitive fields outside usage must not affect parsing — tested via makeSensitiveLine below. + + let combined = [good, badJSON, missingKind, wrongPayloadType].joined(separator: "\n") + let sensitiveExtra = try self.makeSensitiveLine(now: now) + let all = [combined, sensitiveExtra].joined(separator: "\n") + try all.write(to: sessionDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 30, + now: now) + // Only good + sensitiveExtra are valid; malformed lines are ignored. + #expect(summary.requestCount == 2) + #expect(summary.totalTokens == 1400) // good(110) + sensitiveExtra(1290) + } + + @Test + func `sensitive prompt fields do not affect token parsing`() throws { + let now = Date(timeIntervalSince1970: 1_787_840_784) + let line = try self.makeSensitiveLine(now: now) + let parsed = MuseLocalSessionScanner.parseLine(Data(line.utf8)) + #expect(parsed != nil) + #expect(parsed?.inputTokens == 1234) + #expect(parsed?.outputTokens == 56) + } + + @Test + func `truncated lines are discarded`() { + // Simulate a line larger than prefixBytes that would be marked truncated. + // We test parseLine directly: truncated lines are filtered in events(in:) not here, + // but a 32 KiB+ line beyond 8 KiB prefix that still contains model_completed should be dropped. + // This is exercised indirectly via CostUsageJsonl truncation — we assert the scanner + // never crashes and ignores oversized lines. + let hugePrompt = String(repeating: "a", count: 40 * 1024) + let now = Date(timeIntervalSince1970: 1_787_840_784) + let outer: [String: Any] = [ + "recorded_at": self.microseconds(for: now), + "payload_type": "runtime.session", + "payload": [ + "kind": "run", + "event": [ + "kind": "model_completed", + "usage": ["input_tokens": 10, "output_tokens": 5, "cached_tokens": 0, "reasoning_tokens": 0], + "model": "muse-spark-1.2-contributor", + ], + ] as [String: Any], + "extra_prompt": hugePrompt, + ] + let data = try! JSONSerialization.data(withJSONObject: outer, options: []) + // Data is >32 KiB, will be truncated when scanned with prefix 8 KiB; parseLine on full data would succeed, + // but the scanner's CostUsageJsonl path marks it truncated and skips it. We verify full data still parses + // (correctness) and that the file-level path handles truncation gracefully elsewhere. + let parsed = MuseLocalSessionScanner.parseLine(data) + #expect(parsed != nil) + } + + @Test + func `timestamp bucketing uses microsecond wall clock`() throws { + let utc = TimeZone(identifier: "UTC")! + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = utc + let midnight = try #require(calendar.date(from: DateComponents(year: 2026, month: 8, day: 27))) + let beforeMidnight = midnight.addingTimeInterval(-1) // 2026-08-26 23:59:59 UTC + let afterMidnight = midnight.addingTimeInterval(1) + + let lineBefore = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: beforeMidnight), + inputTokens: 100, + outputTokens: 10) + let lineAfter = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: afterMidnight), + inputTokens: 200, + outputTokens: 20) + + let e1 = try #require(MuseLocalSessionScanner.parseLine(Data(lineBefore.utf8))) + let e2 = try #require(MuseLocalSessionScanner.parseLine(Data(lineAfter.utf8))) + // Dates preserve wall-clock order and are near midnight. + #expect(e1.date < e2.date) + // Day keys should differ when crossing UTC midnight (Calendar.current may be non-UTC, so we only assert + // that dayKey formatting is consistent with the parsed date's calendar day). + let k1 = try #require(MuseLocalSessionScanner.dayKey(for: e1.date, calendar: calendar)) + let k2 = try #require(MuseLocalSessionScanner.dayKey(for: e2.date, calendar: calendar)) + #expect(k1 != k2) + #expect(k1 == "2026-08-26") + #expect(k2 == "2026-08-27") + } + + @Test + func `cache and reasoning are extracted but not double-counted`() throws { + let now = Date(timeIntervalSince1970: 1_787_840_784) + let line = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 1000, + outputTokens: 200, + cachedTokens: 800, + reasoningTokens: 150) + let parsed = try #require(MuseLocalSessionScanner.parseLine(Data(line.utf8))) + #expect(parsed.inputTokens == 1000) + #expect(parsed.outputTokens == 200) + #expect(parsed.cacheReadTokens == 800) + #expect(parsed.reasoningTokens == 150) + // total = input + output, not input+output+cache+reasoning + #expect(parsed.inputTokens + parsed.outputTokens == 1200) + } + + // MARK: - Helpers for sensitive payload + + @Test + func `files outside lookback date directories are not scanned`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-prune-\(UUID().uuidString)", isDirectory: true) + // Use local calendar for directory layout so pruning is deterministic regardless of host TZ. + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let oldDate = localCalendar.date(byAdding: .day, value: -60, to: now)! + let recentKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let oldKey = MuseLocalSessionScanner.dayKey(for: oldDate, calendar: localCalendar)! + let recentDir = root.appendingPathComponent("sessions/\(recentKey.replacingOccurrences(of: "-", with: "/"))/recent", isDirectory: true) + let oldDir = root.appendingPathComponent("sessions/\(oldKey.replacingOccurrences(of: "-", with: "/"))/old", isDirectory: true) + try FileManager.default.createDirectory(at: recentDir, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: oldDir, withIntermediateDirectories: true) + + let recentLine = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 100, + outputTokens: 10) + // Old event's recorded_at is recent (so if file were scanned, it would count), + // but its directory date is outside lookback — pruning must prevent opening it. + let oldLine = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 9999, + outputTokens: 9999) + + try recentLine.write(to: recentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try oldLine.write(to: oldDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + var scannedURLs: [URL] = [] + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 7, + now: now, + fileScanObserver: { url in scannedURLs.append(url) }) + + // Only recent file should have been opened; old directory must be pruned. + let recentSlash = recentKey.replacingOccurrences(of: "-", with: "/") + let oldSlash = oldKey.replacingOccurrences(of: "-", with: "/") + #expect(scannedURLs.count == 1) + #expect(scannedURLs.first?.path.contains(recentSlash) == true) + #expect(scannedURLs.allSatisfy { !$0.path.contains(oldSlash) }) + #expect(summary.requestCount == 1) + #expect(summary.totalTokens == 110) + } + + @Test + func `subagent files are included within selected dates but old dates remain pruned`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-prune-subagent-\(UUID().uuidString)", isDirectory: true) + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27, hour: 12))! + let oldDate = localCalendar.date(byAdding: .day, value: -60, to: now)! + let recentKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let oldKey = MuseLocalSessionScanner.dayKey(for: oldDate, calendar: localCalendar)! + let recentDir = root.appendingPathComponent("sessions/\(recentKey.replacingOccurrences(of: "-", with: "/"))/recent", isDirectory: true) + let recentSub = recentDir.appendingPathComponent("subagent/child", isDirectory: true) + let oldDir = root.appendingPathComponent("sessions/\(oldKey.replacingOccurrences(of: "-", with: "/"))/old", isDirectory: true) + let oldSub = oldDir.appendingPathComponent("subagent/child", isDirectory: true) + try FileManager.default.createDirectory(at: recentDir, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: recentSub, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: oldDir, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: oldSub, withIntermediateDirectories: true) + + let parentLine = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 100, + outputTokens: 10) + let childLine = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 200, + outputTokens: 20) + let oldParent = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 999, + outputTokens: 999) + let oldChild = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 999, + outputTokens: 999) + + try parentLine.write(to: recentDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try childLine.write(to: recentSub.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try oldParent.write(to: oldDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + try oldChild.write(to: oldSub.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + var scanned: [URL] = [] + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 7, + now: now, + fileScanObserver: { scanned.append($0) }) + + let recentSlash2 = recentKey.replacingOccurrences(of: "-", with: "/") + let oldSlash2 = oldKey.replacingOccurrences(of: "-", with: "/") + #expect(scanned.count == 2) + #expect(scanned.allSatisfy { $0.path.contains(recentSlash2) }) + #expect(scanned.allSatisfy { !$0.path.contains(oldSlash2) }) + // Both recent parent and child are counted; old files pruned entirely. + #expect(summary.requestCount == 2) + #expect(summary.totalTokens == 330) + #expect(summary.fileCount == 2) + } + + @Test + func `nonexistent date directories are skipped cheaply`() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("muse-scan-missing-\(UUID().uuidString)", isDirectory: true) + // Only create one date directory; the other 29 in the lookback window do not exist and must be skipped. + let localCalendar = Calendar.current + let now = localCalendar.date(from: DateComponents(year: 2026, month: 8, day: 27))! + let onlyKey = MuseLocalSessionScanner.dayKey(for: now, calendar: localCalendar)! + let onlyDir = root.appendingPathComponent("sessions/\(onlyKey.replacingOccurrences(of: "-", with: "/"))/only", isDirectory: true) + try FileManager.default.createDirectory(at: onlyDir, withIntermediateDirectories: true) + + let line = self.makeSessionLine( + recordedAtMicroseconds: self.microseconds(for: now), + inputTokens: 10, + outputTokens: 5) + try line.write(to: onlyDir.appendingPathComponent("session.jsonl"), atomically: true, encoding: .utf8) + + var scanned: [URL] = [] + let summary = MuseLocalSessionScanner.summarize( + env: ["MUSE_HOME": root.path], + lookbackDays: 30, + now: now, + fileScanObserver: { scanned.append($0) }) + #expect(scanned.count == 1) + #expect(summary.requestCount == 1) + } + + private func makeSensitiveLine(now: Date) throws -> String { + let outer: [String: Any] = [ + "recorded_at": self.microseconds(for: now), + "payload_type": "runtime.session", + "payload": [ + "kind": "run", + "event": [ + "kind": "model_completed", + "usage": ["input_tokens": 1234, "output_tokens": 56, "cached_tokens": 0, "reasoning_tokens": 0], + "model": "muse-spark-1.2-contributor", + ], + "prompt": "THIS IS A SECRET PROMPT THAT MUST NOT BE RETAINED", + "tool_output": ["secret": "do not persist"], + ] as [String: Any], + "prompt": "top-level secret", + ] + let data = try JSONSerialization.data(withJSONObject: outer, options: []) + return String(data: data, encoding: .utf8)! + } +} diff --git a/docs/provider-ids.md b/docs/provider-ids.md index 214a88e646..bb1a1fd8fe 100644 --- a/docs/provider-ids.md +++ b/docs/provider-ids.md @@ -2,4 +2,4 @@ # Provider IDs -`codex`, `openai`, `azureopenai`, `claude`, `clinepass`, `cursor`, `opencode`, `opencodego`, `alibaba`, `alibabatokenplan`, `qwencloud`, `factory`, `fireworks`, `gemini`, `antigravity`, `copilot`, `devin`, `zai`, `minimax`, `manus`, `kimi`, `kilo`, `kiro`, `vertexai`, `augment`, `jetbrains`, `moonshot`, `amp`, `t3chat`, `ollama`, `synthetic`, `openrouter`, `elevenlabs`, `warp`, `windsurf`, `zed`, `perplexity`, `mimo`, `doubao`, `sakana`, `abacus`, `mistral`, `deepseek`, `deepinfra`, `codebuff`, `crof`, `venice`, `commandcode`, `qoder`, `stepfun`, `bedrock`, `grok`, `groq`, `llmproxy`, `litellm`, `deepgram`, `poe`, `chutes`, `neuralwatt`, `clawrouter`, `longcat`, `sub2api`, `wayfinder`, `zenmux`, `aiand`, `zoommate`, `xai`, `notion`, `ibmbob`. +`muse`, `codex`, `openai`, `azureopenai`, `claude`, `clinepass`, `cursor`, `opencode`, `opencodego`, `alibaba`, `alibabatokenplan`, `qwencloud`, `factory`, `fireworks`, `gemini`, `antigravity`, `copilot`, `devin`, `zai`, `minimax`, `manus`, `kimi`, `kilo`, `kiro`, `vertexai`, `augment`, `jetbrains`, `moonshot`, `amp`, `t3chat`, `ollama`, `synthetic`, `openrouter`, `elevenlabs`, `warp`, `windsurf`, `zed`, `perplexity`, `mimo`, `doubao`, `sakana`, `abacus`, `mistral`, `deepseek`, `deepinfra`, `codebuff`, `crof`, `venice`, `commandcode`, `qoder`, `stepfun`, `bedrock`, `grok`, `groq`, `llmproxy`, `litellm`, `deepgram`, `poe`, `chutes`, `neuralwatt`, `clawrouter`, `longcat`, `sub2api`, `wayfinder`, `zenmux`, `aiand`, `zoommate`, `xai`, `notion`, `ibmbob`.