526 lines
19 KiB
Swift
526 lines
19 KiB
Swift
import Foundation
|
|
import SQLite3
|
|
|
|
public actor TokenUsageScanner {
|
|
public let homeDirectory: URL
|
|
public let cacheURL: URL
|
|
private let chunkSize = 256 * 1024
|
|
private var lastSignature: FileSignature?
|
|
private var lastSummary: UsageSummary?
|
|
|
|
public init(
|
|
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
|
|
cacheURL: URL? = nil
|
|
) {
|
|
self.homeDirectory = homeDirectory
|
|
self.cacheURL = cacheURL ?? homeDirectory
|
|
.appending(path: "Library/Application Support/TokenLens")
|
|
.appending(path: "usage-cache.sqlite")
|
|
}
|
|
|
|
public func scan(now: Date = Date()) async -> UsageSummary {
|
|
defer { MemoryPressure.relieve() }
|
|
|
|
if let summary = scanWithCache(now: now) {
|
|
return summary
|
|
}
|
|
|
|
var accumulator = UsageSummaryAccumulator(now: now)
|
|
scanClaude(into: &accumulator)
|
|
scanCodex(into: &accumulator)
|
|
scanGemini(into: &accumulator)
|
|
scanHermes(into: &accumulator)
|
|
scanOpenCode(into: &accumulator)
|
|
return accumulator.makeSummary()
|
|
}
|
|
|
|
private func scanWithCache(now: Date) -> UsageSummary? {
|
|
do {
|
|
let snapshots = try logFileSnapshots()
|
|
let signature = FileSignature(snapshots: snapshots)
|
|
if signature == lastSignature,
|
|
let lastSummary,
|
|
isSameDisplayWindow(lastSummary.generatedAt, now) {
|
|
let updated = lastSummary.updatingGeneratedAt(now)
|
|
self.lastSummary = updated
|
|
return updated
|
|
}
|
|
|
|
let cache = try SQLiteUsageCache(url: cacheURL)
|
|
let currentPaths = try refreshCache(cache, snapshots: snapshots)
|
|
try cache.pruneMissingFiles(currentPaths: currentPaths)
|
|
let summary = try cache.makeSummary(now: now)
|
|
lastSignature = signature
|
|
lastSummary = summary
|
|
return summary
|
|
} catch {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
private func refreshCache(_ cache: SQLiteUsageCache, snapshots: [LogFileSnapshot]) throws -> Set<String> {
|
|
var currentPaths = Set<String>()
|
|
for snapshot in snapshots {
|
|
currentPaths.insert(snapshot.path)
|
|
let state = try cache.fileState(path: snapshot.path)
|
|
let isCurrent = try cache.isFileCurrent(
|
|
path: snapshot.path,
|
|
size: snapshot.size,
|
|
modifiedAt: snapshot.modifiedAt
|
|
)
|
|
if isCurrent && (snapshot.scanner.supportsZeroRecordCache || (state?.recordCount ?? 0) > 0) {
|
|
continue
|
|
}
|
|
|
|
if let state,
|
|
snapshot.size > state.size,
|
|
state.size == state.scannedOffset,
|
|
state.scannedOffset > 0,
|
|
snapshot.scanner.supportsIncrementalScan {
|
|
try cache.appendFileRecords(
|
|
path: snapshot.path,
|
|
size: snapshot.size,
|
|
modifiedAt: snapshot.modifiedAt,
|
|
startingRecordCount: state.recordCount
|
|
) { insert in
|
|
try scanFile(
|
|
snapshot.url,
|
|
startingAt: UInt64(state.scannedOffset),
|
|
scanner: snapshot.scanner,
|
|
consume: insert
|
|
).scannedOffset
|
|
}
|
|
} else {
|
|
try cache.replaceFileRecords(
|
|
path: snapshot.path,
|
|
size: snapshot.size,
|
|
modifiedAt: snapshot.modifiedAt
|
|
) { insert in
|
|
try scanFile(snapshot.url, scanner: snapshot.scanner, consume: insert).scannedOffset
|
|
}
|
|
}
|
|
}
|
|
return currentPaths
|
|
}
|
|
|
|
private func scanClaude(into accumulator: inout UsageSummaryAccumulator) {
|
|
scanJSONLFiles(
|
|
root: homeDirectory.appending(path: ".claude/projects"),
|
|
include: { $0.pathExtension == "jsonl" },
|
|
parse: TokenUsageParser.parseClaudeLine,
|
|
into: &accumulator
|
|
)
|
|
}
|
|
|
|
private func scanCodex(into accumulator: inout UsageSummaryAccumulator) {
|
|
scanJSONLFiles(
|
|
root: homeDirectory.appending(path: ".codex/sessions"),
|
|
include: { $0.pathExtension == "jsonl" },
|
|
parse: TokenUsageParser.parseCodexLine,
|
|
into: &accumulator
|
|
)
|
|
}
|
|
|
|
private func scanGemini(into accumulator: inout UsageSummaryAccumulator) {
|
|
scanJSONLFiles(
|
|
root: homeDirectory.appending(path: ".gemini"),
|
|
include: { url in
|
|
url.pathExtension == "jsonl" && url.pathComponents.contains("chats")
|
|
},
|
|
parse: TokenUsageParser.parseGeminiLine,
|
|
into: &accumulator
|
|
)
|
|
}
|
|
|
|
private func scanHermes(into accumulator: inout UsageSummaryAccumulator) {
|
|
for file in files(root: homeDirectory.appending(path: ".hermes"), include: { $0.lastPathComponent == "state.db" }) {
|
|
_ = try? scanHermesStateDB(file) { record in
|
|
accumulator.add(record)
|
|
}
|
|
}
|
|
}
|
|
|
|
private func scanOpenCode(into accumulator: inout UsageSummaryAccumulator) {
|
|
for file in files(root: homeDirectory.appending(path: ".local/share/opencode"), include: { $0.lastPathComponent == "opencode.db" }) {
|
|
_ = try? scanOpenCodeDB(file) { record in
|
|
accumulator.add(record)
|
|
}
|
|
}
|
|
}
|
|
|
|
private func scanJSONLFiles(
|
|
root: URL,
|
|
include: (URL) -> Bool,
|
|
parse: (String, String) -> UsageRecord?,
|
|
into accumulator: inout UsageSummaryAccumulator
|
|
) {
|
|
for file in files(root: root, include: include) {
|
|
_ = try? scanJSONLFile(file, parse: parse) { record in
|
|
accumulator.add(record)
|
|
}
|
|
}
|
|
}
|
|
|
|
private func scanFile(
|
|
_ file: URL,
|
|
startingAt offset: UInt64 = 0,
|
|
scanner: SourceScanner,
|
|
consume: (UsageRecord) throws -> Void
|
|
) throws -> ScanProgress {
|
|
switch scanner {
|
|
case let .jsonl(parse):
|
|
return try scanJSONLFile(file, startingAt: offset, parse: parse, consume: consume)
|
|
case .hermesStateDB:
|
|
return try scanHermesStateDB(file, consume: consume)
|
|
case .openCodeDB:
|
|
return try scanOpenCodeDB(file, consume: consume)
|
|
}
|
|
}
|
|
|
|
private func scanJSONLFile(
|
|
_ file: URL,
|
|
startingAt offset: UInt64 = 0,
|
|
parse: (String, String) -> UsageRecord?,
|
|
consume: (UsageRecord) throws -> Void
|
|
) throws -> ScanProgress {
|
|
guard let handle = try? FileHandle(forReadingFrom: file) else {
|
|
return ScanProgress(scannedOffset: Int64(offset), recordCount: 0)
|
|
}
|
|
defer { try? handle.close() }
|
|
try handle.seek(toOffset: offset)
|
|
|
|
var pendingLine = Data()
|
|
var scannedOffset = Int64(offset)
|
|
var recordCount = 0
|
|
|
|
while true {
|
|
let chunk = (try? handle.read(upToCount: chunkSize)) ?? Data()
|
|
if chunk.isEmpty { break }
|
|
|
|
var lineStart = chunk.startIndex
|
|
for index in chunk.indices where chunk[index] == 0x0A {
|
|
let lineData = chunk[lineStart..<index]
|
|
if pendingLine.isEmpty {
|
|
scannedOffset += Int64(lineData.count + 1)
|
|
try autoreleasepool {
|
|
if try appendRecord(from: lineData, sourcePath: file.path, parse: parse, consume: consume) {
|
|
recordCount += 1
|
|
}
|
|
}
|
|
} else {
|
|
pendingLine.append(contentsOf: lineData)
|
|
scannedOffset += Int64(pendingLine.count + 1)
|
|
try autoreleasepool {
|
|
if try appendRecord(from: pendingLine, sourcePath: file.path, parse: parse, consume: consume) {
|
|
recordCount += 1
|
|
}
|
|
}
|
|
pendingLine.removeAll(keepingCapacity: true)
|
|
}
|
|
lineStart = chunk.index(after: index)
|
|
}
|
|
|
|
if lineStart < chunk.endIndex {
|
|
pendingLine.append(contentsOf: chunk[lineStart..<chunk.endIndex])
|
|
}
|
|
}
|
|
|
|
if !pendingLine.isEmpty {
|
|
try autoreleasepool {
|
|
if try appendRecord(from: pendingLine, sourcePath: file.path, parse: parse, consume: consume) {
|
|
recordCount += 1
|
|
scannedOffset += Int64(pendingLine.count)
|
|
} else if isCompleteJSONLine(pendingLine) {
|
|
scannedOffset += Int64(pendingLine.count)
|
|
}
|
|
}
|
|
}
|
|
|
|
return ScanProgress(scannedOffset: scannedOffset, recordCount: recordCount)
|
|
}
|
|
|
|
private func appendRecord<LineData: Collection>(
|
|
from lineData: LineData,
|
|
sourcePath: String,
|
|
parse: (String, String) -> UsageRecord?,
|
|
consume: (UsageRecord) throws -> Void
|
|
) throws -> Bool where LineData.Element == UInt8 {
|
|
guard !lineData.isEmpty,
|
|
let line = String(bytes: lineData, encoding: .utf8),
|
|
let record = parse(line, sourcePath)
|
|
else {
|
|
return false
|
|
}
|
|
try consume(record)
|
|
return true
|
|
}
|
|
|
|
private func isCompleteJSONLine(_ lineData: Data) -> Bool {
|
|
guard !lineData.isEmpty else { return true }
|
|
return (try? JSONSerialization.jsonObject(with: lineData)) != nil
|
|
}
|
|
|
|
private func scanHermesStateDB(
|
|
_ file: URL,
|
|
consume: (UsageRecord) throws -> Void
|
|
) throws -> ScanProgress {
|
|
let sql = """
|
|
SELECT id, started_at, COALESCE(model, ''),
|
|
input_tokens, output_tokens, cache_read_tokens,
|
|
cache_write_tokens, reasoning_tokens
|
|
FROM sessions
|
|
WHERE input_tokens + output_tokens + cache_read_tokens + cache_write_tokens + reasoning_tokens > 0
|
|
"""
|
|
return try scanSQLiteRows(file, sql: sql) { statement in
|
|
UsageRecord(
|
|
provider: .hermes,
|
|
source: "state.db/sessions",
|
|
sourcePath: file.path,
|
|
timestamp: Date(timeIntervalSince1970: sqlite3_column_double(statement, 1)),
|
|
sessionID: textColumn(statement, 0),
|
|
model: textColumn(statement, 2).isEmpty ? "Hermes" : textColumn(statement, 2),
|
|
inputTokens: intColumn(statement, 3),
|
|
outputTokens: intColumn(statement, 4),
|
|
cachedTokens: intColumn(statement, 5) + intColumn(statement, 6),
|
|
reasoningTokens: intColumn(statement, 7),
|
|
toolTokens: 0
|
|
)
|
|
} consume: {
|
|
try consume($0)
|
|
}
|
|
}
|
|
|
|
private func scanOpenCodeDB(
|
|
_ file: URL,
|
|
consume: (UsageRecord) throws -> Void
|
|
) throws -> ScanProgress {
|
|
let sql = """
|
|
SELECT id, COALESCE(NULLIF(time_updated, 0), time_created), COALESCE(model, ''),
|
|
tokens_input, tokens_output, tokens_reasoning,
|
|
tokens_cache_read, tokens_cache_write
|
|
FROM session
|
|
WHERE tokens_input + tokens_output + tokens_reasoning + tokens_cache_read + tokens_cache_write > 0
|
|
"""
|
|
return try scanSQLiteRows(file, sql: sql) { statement in
|
|
UsageRecord(
|
|
provider: .opencode,
|
|
source: "opencode.db/session",
|
|
sourcePath: file.path,
|
|
timestamp: Date(timeIntervalSince1970: sqlite3_column_double(statement, 1) / 1000),
|
|
sessionID: textColumn(statement, 0),
|
|
model: openCodeModelName(textColumn(statement, 2)),
|
|
inputTokens: intColumn(statement, 3),
|
|
outputTokens: intColumn(statement, 4),
|
|
cachedTokens: intColumn(statement, 6) + intColumn(statement, 7),
|
|
reasoningTokens: intColumn(statement, 5),
|
|
toolTokens: 0
|
|
)
|
|
} consume: {
|
|
try consume($0)
|
|
}
|
|
}
|
|
|
|
private func scanSQLiteRows(
|
|
_ file: URL,
|
|
sql: String,
|
|
makeRecord: (OpaquePointer?) -> UsageRecord,
|
|
consume: (UsageRecord) throws -> Void
|
|
) throws -> ScanProgress {
|
|
var database: OpaquePointer?
|
|
let flags = SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX
|
|
guard sqlite3_open_v2(file.path, &database, flags, nil) == SQLITE_OK else {
|
|
sqlite3_close(database)
|
|
return ScanProgress(scannedOffset: fileSize(file), recordCount: 0)
|
|
}
|
|
defer { sqlite3_close(database) }
|
|
sqlite3_busy_timeout(database, 1000)
|
|
|
|
var statement: OpaquePointer?
|
|
guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK else {
|
|
sqlite3_finalize(statement)
|
|
return ScanProgress(scannedOffset: fileSize(file), recordCount: 0)
|
|
}
|
|
defer { sqlite3_finalize(statement) }
|
|
|
|
var recordCount = 0
|
|
while sqlite3_step(statement) == SQLITE_ROW {
|
|
try consume(makeRecord(statement))
|
|
recordCount += 1
|
|
}
|
|
return ScanProgress(scannedOffset: fileSize(file), recordCount: recordCount)
|
|
}
|
|
|
|
private func fileSize(_ file: URL) -> Int64 {
|
|
((try? file.resourceValues(forKeys: [.fileSizeKey]).fileSize).flatMap { Int64($0) }) ?? 0
|
|
}
|
|
|
|
private func textColumn(_ statement: OpaquePointer?, _ index: Int32) -> String {
|
|
guard let text = sqlite3_column_text(statement, index) else { return "" }
|
|
return String(cString: UnsafeRawPointer(text).assumingMemoryBound(to: CChar.self))
|
|
}
|
|
|
|
private func intColumn(_ statement: OpaquePointer?, _ index: Int32) -> Int {
|
|
Int(sqlite3_column_int64(statement, index))
|
|
}
|
|
|
|
private func openCodeModelName(_ rawModel: String) -> String {
|
|
guard !rawModel.isEmpty,
|
|
let data = rawModel.data(using: .utf8),
|
|
let object = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
|
|
else {
|
|
return rawModel.isEmpty ? "OpenCode" : rawModel
|
|
}
|
|
let provider = object["providerID"] as? String
|
|
let model = object["id"] as? String
|
|
var parts: [String] = []
|
|
if let provider, !provider.isEmpty {
|
|
parts.append(provider)
|
|
}
|
|
if let model, !model.isEmpty {
|
|
parts.append(model)
|
|
}
|
|
return parts.isEmpty ? "OpenCode" : parts.joined(separator: "/")
|
|
}
|
|
|
|
private func files(root: URL, include: (URL) -> Bool) -> [URL] {
|
|
guard let enumerator = FileManager.default.enumerator(
|
|
at: root,
|
|
includingPropertiesForKeys: [.isRegularFileKey],
|
|
options: [.skipsHiddenFiles]
|
|
) else {
|
|
return []
|
|
}
|
|
|
|
var result: [URL] = []
|
|
for case let file as URL in enumerator where include(file) {
|
|
result.append(file)
|
|
}
|
|
return result
|
|
}
|
|
|
|
private var logSources: [LogSource] {
|
|
[
|
|
LogSource(
|
|
root: homeDirectory.appending(path: ".claude/projects"),
|
|
include: { $0.pathExtension == "jsonl" },
|
|
scanner: .jsonl(TokenUsageParser.parseClaudeLine)
|
|
),
|
|
LogSource(
|
|
root: homeDirectory.appending(path: ".codex/sessions"),
|
|
include: { $0.pathExtension == "jsonl" },
|
|
scanner: .jsonl(TokenUsageParser.parseCodexLine)
|
|
),
|
|
LogSource(
|
|
root: homeDirectory.appending(path: ".gemini"),
|
|
include: { url in
|
|
url.pathExtension == "jsonl" && url.pathComponents.contains("chats")
|
|
},
|
|
scanner: .jsonl(TokenUsageParser.parseGeminiLine)
|
|
),
|
|
LogSource(
|
|
root: homeDirectory.appending(path: ".hermes"),
|
|
include: {
|
|
$0.lastPathComponent == "state.db" &&
|
|
$0.deletingLastPathComponent().lastPathComponent == ".hermes"
|
|
},
|
|
scanner: .hermesStateDB
|
|
),
|
|
LogSource(
|
|
root: homeDirectory.appending(path: ".local/share/opencode"),
|
|
include: {
|
|
$0.lastPathComponent == "opencode.db" &&
|
|
$0.deletingLastPathComponent().lastPathComponent == "opencode"
|
|
},
|
|
scanner: .openCodeDB
|
|
)
|
|
]
|
|
}
|
|
|
|
private func logFileSnapshots() throws -> [LogFileSnapshot] {
|
|
var snapshots: [LogFileSnapshot] = []
|
|
for source in logSources {
|
|
for file in files(root: source.root, include: source.include) {
|
|
let metadata = try fileMetadata(file)
|
|
snapshots.append(LogFileSnapshot(
|
|
url: file,
|
|
path: file.path,
|
|
size: metadata.size,
|
|
modifiedAt: metadata.modifiedAt,
|
|
scanner: source.scanner
|
|
))
|
|
}
|
|
}
|
|
return snapshots.sorted { $0.path < $1.path }
|
|
}
|
|
|
|
private func fileMetadata(_ file: URL) throws -> FileMetadata {
|
|
let values = try file.resourceValues(forKeys: [.fileSizeKey, .contentModificationDateKey])
|
|
return FileMetadata(
|
|
size: Int64(values.fileSize ?? 0),
|
|
modifiedAt: values.contentModificationDate?.timeIntervalSince1970 ?? 0
|
|
)
|
|
}
|
|
|
|
private func isSameDisplayWindow(_ lhs: Date, _ rhs: Date) -> Bool {
|
|
var calendar = Calendar.current
|
|
calendar.firstWeekday = 2
|
|
return calendar.isDate(lhs, equalTo: rhs, toGranularity: .hour)
|
|
}
|
|
}
|
|
|
|
private struct LogSource {
|
|
let root: URL
|
|
let include: (URL) -> Bool
|
|
let scanner: SourceScanner
|
|
}
|
|
|
|
private enum SourceScanner {
|
|
case jsonl((String, String) -> UsageRecord?)
|
|
case hermesStateDB
|
|
case openCodeDB
|
|
|
|
var supportsIncrementalScan: Bool {
|
|
if case .jsonl = self { return true }
|
|
return false
|
|
}
|
|
|
|
var supportsZeroRecordCache: Bool {
|
|
if case .jsonl = self { return true }
|
|
return false
|
|
}
|
|
}
|
|
|
|
private struct ScanProgress {
|
|
let scannedOffset: Int64
|
|
let recordCount: Int
|
|
}
|
|
|
|
private struct FileMetadata {
|
|
let size: Int64
|
|
let modifiedAt: TimeInterval
|
|
}
|
|
|
|
private struct LogFileSnapshot {
|
|
let url: URL
|
|
let path: String
|
|
let size: Int64
|
|
let modifiedAt: TimeInterval
|
|
let scanner: SourceScanner
|
|
}
|
|
|
|
private struct FileSignature: Equatable {
|
|
let entries: [Entry]
|
|
|
|
init(snapshots: [LogFileSnapshot]) {
|
|
entries = snapshots.map {
|
|
Entry(path: $0.path, size: $0.size, modifiedAt: $0.modifiedAt)
|
|
}
|
|
}
|
|
|
|
struct Entry: Equatable {
|
|
let path: String
|
|
let size: Int64
|
|
let modifiedAt: TimeInterval
|
|
}
|
|
}
|