313 lines
13 KiB
Swift
313 lines
13 KiB
Swift
import Testing
|
|
import Foundation
|
|
import SQLite3
|
|
@testable import TokenLensCore
|
|
|
|
@Test func formatTokensUsesChineseMagnitude() {
|
|
#expect(TokenFormatter.format(9_999) == "9999")
|
|
#expect(TokenFormatter.format(10_000) == "1.000万")
|
|
#expect(TokenFormatter.format(71_400) == "7.140万")
|
|
#expect(TokenFormatter.format(6_600_000) == "660.000万")
|
|
#expect(TokenFormatter.format(12_345_678) == "1234.6万")
|
|
#expect(TokenFormatter.format(196_700_000) == "1.967亿")
|
|
#expect(TokenFormatter.format(1_020_000_000) == "10.200亿")
|
|
#expect(TokenFormatter.format(100_000_000_000) == "1000.0亿")
|
|
}
|
|
|
|
@Test func parseClaudeUsageLineSkipsConversationContent() throws {
|
|
let line = """
|
|
{"timestamp":"2026-06-11T02:20:00.000Z","sessionId":"s1","cwd":"/tmp","message":{"model":"claude-3-5-sonnet","content":"private text","usage":{"input_tokens":120,"output_tokens":30,"cache_read_input_tokens":50,"cache_creation_input_tokens":10}}}
|
|
"""
|
|
|
|
let record = try #require(TokenUsageParser.parseClaudeLine(line, sourcePath: "/tmp/a.jsonl"))
|
|
|
|
#expect(record.provider == .claude)
|
|
#expect(record.model == "claude-3-5-sonnet")
|
|
#expect(record.inputTokens == 120)
|
|
#expect(record.outputTokens == 30)
|
|
#expect(record.cachedTokens == 60)
|
|
#expect(record.totalTokens == 210)
|
|
}
|
|
|
|
@Test func parseCodexUsesLastTokenUsageOnly() throws {
|
|
let line = """
|
|
{"timestamp":"2026-06-11T03:10:00.000Z","type":"event_msg","payload":{"info":{"total_token_usage":{"input_tokens":1000,"output_tokens":300,"cached_input_tokens":200,"reasoning_output_tokens":40,"total_tokens":1300},"last_token_usage":{"input_tokens":100,"output_tokens":30,"cached_input_tokens":20,"reasoning_output_tokens":4,"total_tokens":130}}}}
|
|
"""
|
|
|
|
let record = try #require(TokenUsageParser.parseCodexLine(line, sourcePath: "/tmp/rollout.jsonl"))
|
|
|
|
#expect(record.provider == .codex)
|
|
#expect(record.inputTokens == 100)
|
|
#expect(record.outputTokens == 30)
|
|
#expect(record.cachedTokens == 20)
|
|
#expect(record.reasoningTokens == 4)
|
|
#expect(record.totalTokens == 154)
|
|
}
|
|
|
|
@Test func parseGeminiTokensLine() throws {
|
|
let line = """
|
|
{"timestamp":"2026-06-11T04:00:00.000Z","type":"response","model":"gemini-2.5-pro","tokens":{"input":70,"output":20,"cached":5,"thoughts":3,"tool":2,"total":100}}
|
|
"""
|
|
|
|
let record = try #require(TokenUsageParser.parseGeminiLine(line, sourcePath: "/tmp/session.jsonl"))
|
|
|
|
#expect(record.provider == .gemini)
|
|
#expect(record.model == "gemini-2.5-pro")
|
|
#expect(record.inputTokens == 70)
|
|
#expect(record.outputTokens == 20)
|
|
#expect(record.cachedTokens == 5)
|
|
#expect(record.reasoningTokens == 3)
|
|
#expect(record.toolTokens == 2)
|
|
#expect(record.totalTokens == 100)
|
|
}
|
|
|
|
@Test func summaryAggregatesTodayMonthProvidersAndSources() {
|
|
let calendar = Calendar(identifier: .gregorian)
|
|
let now = ISO8601DateFormatter().date(from: "2026-06-11T04:00:00Z")!
|
|
let records = [
|
|
usage(.claude, "2026-06-11T01:00:00Z", 100, 20, 10),
|
|
usage(.codex, "2026-06-10T01:00:00Z", 80, 20, 0),
|
|
usage(.gemini, "2026-05-31T01:00:00Z", 50, 10, 0)
|
|
]
|
|
|
|
let summary = UsageSummary.make(records: records, now: now, calendar: calendar)
|
|
|
|
#expect(summary.cards.todayTokens == 130)
|
|
#expect(summary.cards.monthTokens == 230)
|
|
#expect(summary.cards.totalTokens == 290)
|
|
#expect(summary.providerTotals[.claude]?.totalTokens == 130)
|
|
#expect(summary.providerTotals[.codex]?.totalTokens == 100)
|
|
#expect(summary.providerTotals[.gemini]?.totalTokens == 60)
|
|
#expect(summary.todayProviderTotals[.claude]?.totalTokens == 130)
|
|
#expect(summary.todayProviderTotals[.codex]?.totalTokens == 0)
|
|
#expect(summary.todayProviderTotals[.gemini]?.totalTokens == 0)
|
|
#expect(summary.toolRows.count == 3)
|
|
}
|
|
|
|
@Test func toolRowsKeepProviderSpecificTrends() {
|
|
let calendar = Calendar(identifier: .gregorian)
|
|
let now = ISO8601DateFormatter().date(from: "2026-06-11T04:00:00Z")!
|
|
let records = [
|
|
usage(.claude, "2026-06-09T01:00:00Z", 10, 0, 0),
|
|
usage(.claude, "2026-06-11T01:00:00Z", 30, 0, 0),
|
|
usage(.codex, "2026-06-10T01:00:00Z", 20, 0, 0),
|
|
usage(.gemini, "2026-06-11T02:00:00Z", 5, 0, 0)
|
|
]
|
|
|
|
let summary = UsageSummary.make(records: records, now: now, calendar: calendar)
|
|
let trends = Dictionary(uniqueKeysWithValues: summary.toolRows.map { ($0.provider, $0.trendTokens) })
|
|
|
|
#expect(trends[.claude] == [10, 0, 30])
|
|
#expect(trends[.codex] == [0, 20, 0])
|
|
#expect(trends[.gemini] == [0, 0, 5])
|
|
}
|
|
|
|
@Test func trendPointsRegroupBySelectedMode() {
|
|
var calendar = Calendar(identifier: .gregorian)
|
|
calendar.timeZone = TimeZone(secondsFromGMT: 0)!
|
|
calendar.firstWeekday = 2
|
|
let now = ISO8601DateFormatter().date(from: "2026-06-11T12:30:00Z")!
|
|
let records = [
|
|
usage(.claude, "2026-06-11T01:00:00Z", 10, 0, 0),
|
|
usage(.codex, "2026-06-11T10:00:00Z", 20, 0, 0),
|
|
usage(.claude, "2026-06-11T01:00:00Z", 30, 0, 0),
|
|
usage(.gemini, "2026-05-11T01:00:00Z", 5, 0, 0),
|
|
usage(.codex, "2026-04-01T01:00:00Z", 7, 0, 0)
|
|
]
|
|
let summary = UsageSummary.make(records: records, now: now, calendar: calendar)
|
|
|
|
let daily = summary.trendPoints(mode: .daily, calendar: calendar)
|
|
#expect(daily.count == 13)
|
|
#expect(daily[1].values[.claude] == 40)
|
|
#expect(daily[10].values[.codex] == 20)
|
|
|
|
let weekly = summary.trendPoints(mode: .weekly, calendar: calendar)
|
|
#expect(weekly.count == 4)
|
|
#expect(weekly.last?.values[.claude] == 40)
|
|
#expect(weekly.last?.values[.codex] == 20)
|
|
|
|
let monthly = summary.trendPoints(mode: .monthly, calendar: calendar)
|
|
#expect(monthly.count == 11)
|
|
#expect(monthly.last?.values[.claude] == 40)
|
|
#expect(monthly.last?.values[.codex] == 20)
|
|
|
|
let yearly = summary.trendPoints(mode: .yearly, calendar: calendar)
|
|
#expect(yearly.count == 6)
|
|
#expect(yearly[3].values[.codex] == 7)
|
|
#expect(yearly[4].values[.gemini] == 5)
|
|
#expect(yearly[5].values[.claude] == 40)
|
|
}
|
|
|
|
@Test func dailyUsageDetailsAggregateProviderTotals() {
|
|
var calendar = Calendar(identifier: .gregorian)
|
|
calendar.timeZone = TimeZone(secondsFromGMT: 0)!
|
|
let now = ISO8601DateFormatter().date(from: "2026-06-11T12:30:00Z")!
|
|
let records = [
|
|
usage(.claude, "2026-06-11T01:00:00Z", 10, 2, 1),
|
|
usage(.codex, "2026-06-11T10:00:00Z", 20, 3, 0),
|
|
usage(.gemini, "2026-06-10T01:00:00Z", 5, 1, 0)
|
|
]
|
|
let summary = UsageSummary.make(records: records, now: now, calendar: calendar)
|
|
let details = summary.dailyUsageDetails(calendar: calendar)
|
|
|
|
#expect(details.count == 2)
|
|
#expect(details[0].totalTokens == 36)
|
|
#expect(details[0].inputTokens == 30)
|
|
#expect(details[0].outputTokens == 5)
|
|
#expect(details[0].cachedTokens == 1)
|
|
#expect(details[0].providerTotals[.claude] == 13)
|
|
#expect(details[0].providerTotals[.codex] == 23)
|
|
#expect(details[1].totalTokens == 6)
|
|
#expect(details[1].providerTotals[.gemini] == 6)
|
|
}
|
|
|
|
@Test func scannerIndexesLogsThroughSQLiteCache() async throws {
|
|
let root = FileManager.default.temporaryDirectory
|
|
.appending(path: "TokenLensTests-\(UUID().uuidString)")
|
|
defer { try? FileManager.default.removeItem(at: root) }
|
|
|
|
let sessions = root.appending(path: ".codex/sessions")
|
|
try FileManager.default.createDirectory(at: sessions, withIntermediateDirectories: true)
|
|
let file = sessions.appending(path: "rollout.jsonl")
|
|
let line = """
|
|
{"timestamp":"2026-06-11T03:10:00.000Z","type":"event_msg","payload":{"info":{"last_token_usage":{"input_tokens":100,"output_tokens":30,"cached_input_tokens":20,"reasoning_output_tokens":4}}}}
|
|
"""
|
|
try line.write(to: file, atomically: true, encoding: .utf8)
|
|
|
|
let scanner = TokenUsageScanner(
|
|
homeDirectory: root,
|
|
cacheURL: root.appending(path: "Library/Application Support/TokenLens/test-cache.sqlite")
|
|
)
|
|
let now = ISO8601DateFormatter().date(from: "2026-06-11T12:30:00Z")!
|
|
|
|
let first = await scanner.scan(now: now)
|
|
let second = await scanner.scan(now: now)
|
|
let appendedLine = """
|
|
{"timestamp":"2026-06-11T03:12:00.000Z","type":"event_msg","payload":{"info":{"last_token_usage":{"input_tokens":200,"output_tokens":60,"cached_input_tokens":40,"reasoning_output_tokens":8}}}}
|
|
"""
|
|
let handle = try FileHandle(forWritingTo: file)
|
|
try handle.seekToEnd()
|
|
try handle.write(contentsOf: Data(("\n" + appendedLine).utf8))
|
|
try handle.close()
|
|
let third = await scanner.scan(now: now)
|
|
|
|
#expect(first.cards.totalTokens == 154)
|
|
#expect(second.cards.totalTokens == 154)
|
|
#expect(second.providerTotals[.codex]?.totalTokens == 154)
|
|
#expect(third.cards.totalTokens == 462)
|
|
#expect(third.providerTotals[.codex]?.totalTokens == 462)
|
|
}
|
|
|
|
@Test func scannerIndexesHermesAndOpenCodeSQLiteSources() async throws {
|
|
let root = FileManager.default.temporaryDirectory
|
|
.appending(path: "TokenLensSQLiteSources-\(UUID().uuidString)")
|
|
defer { try? FileManager.default.removeItem(at: root) }
|
|
|
|
let hermes = root.appending(path: ".hermes")
|
|
try FileManager.default.createDirectory(at: hermes, withIntermediateDirectories: true)
|
|
try makeHermesStateDB(at: hermes.appending(path: "state.db"))
|
|
|
|
let opencode = root.appending(path: ".local/share/opencode")
|
|
try FileManager.default.createDirectory(at: opencode, withIntermediateDirectories: true)
|
|
try makeOpenCodeDB(at: opencode.appending(path: "opencode.db"))
|
|
|
|
let scanner = TokenUsageScanner(
|
|
homeDirectory: root,
|
|
cacheURL: root.appending(path: "Library/Application Support/TokenLens/test-cache.sqlite")
|
|
)
|
|
let now = Date(timeIntervalSince1970: 1_780_938_000)
|
|
let summary = await scanner.scan(now: now)
|
|
|
|
#expect(summary.providerTotals[.hermes]?.totalTokens == 23)
|
|
#expect(summary.providerTotals[.opencode]?.totalTokens == 41)
|
|
#expect(summary.cards.totalTokens == 64)
|
|
#expect(summary.toolRows.map(\.provider).contains(.hermes))
|
|
#expect(summary.toolRows.map(\.provider).contains(.opencode))
|
|
}
|
|
|
|
private func usage(_ provider: Provider, _ timestamp: String, _ input: Int, _ output: Int, _ cached: Int) -> UsageRecord {
|
|
UsageRecord(
|
|
provider: provider,
|
|
source: "\(provider.rawValue) source",
|
|
sourcePath: "/tmp/\(provider.rawValue).jsonl",
|
|
timestamp: ISO8601DateFormatter().date(from: timestamp)!,
|
|
sessionID: "",
|
|
model: provider.rawValue,
|
|
inputTokens: input,
|
|
outputTokens: output,
|
|
cachedTokens: cached,
|
|
reasoningTokens: 0,
|
|
toolTokens: 0
|
|
)
|
|
}
|
|
|
|
private func makeHermesStateDB(at url: URL) throws {
|
|
let statements = [
|
|
"""
|
|
CREATE TABLE sessions (
|
|
id TEXT PRIMARY KEY,
|
|
source TEXT NOT NULL,
|
|
model TEXT,
|
|
started_at REAL NOT NULL,
|
|
input_tokens INTEGER DEFAULT 0,
|
|
output_tokens INTEGER DEFAULT 0,
|
|
cache_read_tokens INTEGER DEFAULT 0,
|
|
cache_write_tokens INTEGER DEFAULT 0,
|
|
reasoning_tokens INTEGER DEFAULT 0
|
|
)
|
|
""",
|
|
"""
|
|
INSERT INTO sessions (
|
|
id, source, model, started_at, input_tokens, output_tokens,
|
|
cache_read_tokens, cache_write_tokens, reasoning_tokens
|
|
) VALUES ('h1', 'cli', 'gpt-5.5', 1780937894.0, 10, 5, 3, 2, 3)
|
|
"""
|
|
]
|
|
try writeSQLiteDatabase(at: url, statements: statements)
|
|
}
|
|
|
|
private func makeOpenCodeDB(at url: URL) throws {
|
|
let statements = [
|
|
"""
|
|
CREATE TABLE session (
|
|
id TEXT PRIMARY KEY,
|
|
time_created INTEGER NOT NULL,
|
|
time_updated INTEGER NOT NULL,
|
|
model TEXT,
|
|
tokens_input INTEGER DEFAULT 0 NOT NULL,
|
|
tokens_output INTEGER DEFAULT 0 NOT NULL,
|
|
tokens_reasoning INTEGER DEFAULT 0 NOT NULL,
|
|
tokens_cache_read INTEGER DEFAULT 0 NOT NULL,
|
|
tokens_cache_write INTEGER DEFAULT 0 NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
INSERT INTO session (
|
|
id, time_created, time_updated, model, tokens_input, tokens_output,
|
|
tokens_reasoning, tokens_cache_read, tokens_cache_write
|
|
) VALUES (
|
|
'o1', 1780937894000, 1780937895000,
|
|
'{"id":"MiniMax-M3","providerID":"minimax-cn"}',
|
|
11, 7, 5, 13, 5
|
|
)
|
|
"""
|
|
]
|
|
try writeSQLiteDatabase(at: url, statements: statements)
|
|
}
|
|
|
|
private func writeSQLiteDatabase(at url: URL, statements: [String]) throws {
|
|
var database: OpaquePointer?
|
|
guard sqlite3_open_v2(url.path, &database, SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE, nil) == SQLITE_OK else {
|
|
sqlite3_close(database)
|
|
struct OpenError: Error {}
|
|
throw OpenError()
|
|
}
|
|
defer { sqlite3_close(database) }
|
|
for statement in statements {
|
|
guard sqlite3_exec(database, statement, nil, nil, nil) == SQLITE_OK else {
|
|
struct ExecuteError: Error {}
|
|
throw ExecuteError()
|
|
}
|
|
}
|
|
}
|