435 lines
14 KiB
TypeScript
435 lines
14 KiB
TypeScript
import { feature } from 'bun:bundle'
|
|
import { randomBytes } from 'crypto'
|
|
import { open } from 'fs/promises'
|
|
import { join } from 'path'
|
|
import type { ModelUsage } from '../entrypoints/agentSdkTypes.js'
|
|
import { logForDebugging } from './debug.js'
|
|
import { getClaudeConfigHomeDir } from './envUtils.js'
|
|
import { errorMessage } from './errors.js'
|
|
import { getFsImplementation } from './fsOperations.js'
|
|
import { logError } from './log.js'
|
|
import { jsonParse, jsonStringify } from './slowOperations.js'
|
|
import type { DailyActivity, DailyModelTokens, SessionStats } from './stats.js'
|
|
|
|
export const STATS_CACHE_VERSION = 3
|
|
const MIN_MIGRATABLE_VERSION = 1
|
|
const STATS_CACHE_FILENAME = 'stats-cache.json'
|
|
|
|
/**
|
|
* Simple in-memory lock to prevent concurrent cache operations.
|
|
*/
|
|
let statsCacheLockPromise: Promise<void> | null = null
|
|
|
|
/**
|
|
* Execute a function while holding the stats cache lock.
|
|
* Only one operation can hold the lock at a time.
|
|
*/
|
|
export async function withStatsCacheLock<T>(fn: () => Promise<T>): Promise<T> {
|
|
// Wait for any existing lock to be released
|
|
while (statsCacheLockPromise) {
|
|
await statsCacheLockPromise
|
|
}
|
|
|
|
// Create our lock
|
|
let releaseLock: (() => void) | undefined
|
|
statsCacheLockPromise = new Promise<void>(resolve => {
|
|
releaseLock = resolve
|
|
})
|
|
|
|
try {
|
|
return await fn()
|
|
} finally {
|
|
// Release the lock
|
|
statsCacheLockPromise = null
|
|
releaseLock?.()
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Persisted stats cache stored on disk.
|
|
* Contains aggregated historical stats that won't change.
|
|
* All fields are bounded to prevent unbounded file growth.
|
|
*/
|
|
export type PersistedStatsCache = {
|
|
version: number
|
|
// Last date that was fully computed (YYYY-MM-DD format)
|
|
// Stats up to and including this date are considered complete
|
|
lastComputedDate: string | null
|
|
// Daily aggregates needed for heatmap, streaks, trends (bounded by days)
|
|
dailyActivity: DailyActivity[]
|
|
dailyModelTokens: DailyModelTokens[]
|
|
// Model usage aggregated (bounded by number of models)
|
|
modelUsage: { [modelName: string]: ModelUsage }
|
|
// Session aggregates (replaces unbounded sessionStats array)
|
|
totalSessions: number
|
|
totalMessages: number
|
|
longestSession: SessionStats | null
|
|
// First session date ever recorded
|
|
firstSessionDate: string | null
|
|
// Hour counts for peak hour calculation (bounded to 24 entries)
|
|
hourCounts: { [hour: number]: number }
|
|
// Speculation time saved across all sessions
|
|
totalSpeculationTimeSavedMs: number
|
|
// Shot distribution: map of shot count → number of sessions (ant-only)
|
|
shotDistribution?: { [shotCount: number]: number }
|
|
}
|
|
|
|
export function getStatsCachePath(): string {
|
|
return join(getClaudeConfigHomeDir(), STATS_CACHE_FILENAME)
|
|
}
|
|
|
|
function getEmptyCache(): PersistedStatsCache {
|
|
return {
|
|
version: STATS_CACHE_VERSION,
|
|
lastComputedDate: null,
|
|
dailyActivity: [],
|
|
dailyModelTokens: [],
|
|
modelUsage: {},
|
|
totalSessions: 0,
|
|
totalMessages: 0,
|
|
longestSession: null,
|
|
firstSessionDate: null,
|
|
hourCounts: {},
|
|
totalSpeculationTimeSavedMs: 0,
|
|
shotDistribution: {},
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Migrate an older cache to the current schema.
|
|
* Returns null if the version is unknown or too old to migrate.
|
|
*
|
|
* Preserves historical aggregates that would otherwise be lost when
|
|
* transcript files have already aged out past cleanupPeriodDays.
|
|
* Pre-migration days may undercount (e.g. v2 lacked subagent tokens);
|
|
* we accept that rather than drop the history.
|
|
*/
|
|
function migrateStatsCache(
|
|
parsed: Partial<PersistedStatsCache> & { version: number },
|
|
): PersistedStatsCache | null {
|
|
if (
|
|
typeof parsed.version !== 'number' ||
|
|
parsed.version < MIN_MIGRATABLE_VERSION ||
|
|
parsed.version > STATS_CACHE_VERSION
|
|
) {
|
|
return null
|
|
}
|
|
if (
|
|
!Array.isArray(parsed.dailyActivity) ||
|
|
!Array.isArray(parsed.dailyModelTokens) ||
|
|
typeof parsed.totalSessions !== 'number' ||
|
|
typeof parsed.totalMessages !== 'number'
|
|
) {
|
|
return null
|
|
}
|
|
return {
|
|
version: STATS_CACHE_VERSION,
|
|
lastComputedDate: parsed.lastComputedDate ?? null,
|
|
dailyActivity: parsed.dailyActivity,
|
|
dailyModelTokens: parsed.dailyModelTokens,
|
|
modelUsage: parsed.modelUsage ?? {},
|
|
totalSessions: parsed.totalSessions,
|
|
totalMessages: parsed.totalMessages,
|
|
longestSession: parsed.longestSession ?? null,
|
|
firstSessionDate: parsed.firstSessionDate ?? null,
|
|
hourCounts: parsed.hourCounts ?? {},
|
|
totalSpeculationTimeSavedMs: parsed.totalSpeculationTimeSavedMs ?? 0,
|
|
// Preserve undefined (don't default to {}) so the SHOT_STATS recompute
|
|
// check in loadStatsCache fires for v1/v2 caches that lacked this field.
|
|
shotDistribution: parsed.shotDistribution,
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Load the stats cache from disk.
|
|
* Returns an empty cache if the file doesn't exist or is invalid.
|
|
*/
|
|
export async function loadStatsCache(): Promise<PersistedStatsCache> {
|
|
const fs = getFsImplementation()
|
|
const cachePath = getStatsCachePath()
|
|
|
|
try {
|
|
const content = await fs.readFile(cachePath, { encoding: 'utf-8' })
|
|
const parsed = jsonParse(content) as PersistedStatsCache
|
|
|
|
// Validate version
|
|
if (parsed.version !== STATS_CACHE_VERSION) {
|
|
const migrated = migrateStatsCache(parsed)
|
|
if (!migrated) {
|
|
logForDebugging(
|
|
`Stats cache version ${parsed.version} not migratable (expected ${STATS_CACHE_VERSION}), returning empty cache`,
|
|
)
|
|
return getEmptyCache()
|
|
}
|
|
logForDebugging(
|
|
`Migrated stats cache from v${parsed.version} to v${STATS_CACHE_VERSION}`,
|
|
)
|
|
// Persist migration so we don't re-migrate on every load.
|
|
// aggregateClaudeCodeStats() skips its save when lastComputedDate is
|
|
// already current, so without this the on-disk file stays at the old
|
|
// version indefinitely.
|
|
await saveStatsCache(migrated)
|
|
if (feature('SHOT_STATS') && !migrated.shotDistribution) {
|
|
logForDebugging(
|
|
'Migrated stats cache missing shotDistribution, forcing recomputation',
|
|
)
|
|
return getEmptyCache()
|
|
}
|
|
return migrated
|
|
}
|
|
|
|
// Basic validation
|
|
if (
|
|
!Array.isArray(parsed.dailyActivity) ||
|
|
!Array.isArray(parsed.dailyModelTokens) ||
|
|
typeof parsed.totalSessions !== 'number' ||
|
|
typeof parsed.totalMessages !== 'number'
|
|
) {
|
|
logForDebugging(
|
|
'Stats cache has invalid structure, returning empty cache',
|
|
)
|
|
return getEmptyCache()
|
|
}
|
|
|
|
// If SHOT_STATS is enabled but cache doesn't have shotDistribution,
|
|
// force full recomputation to get historical shot data
|
|
if (feature('SHOT_STATS') && !parsed.shotDistribution) {
|
|
logForDebugging(
|
|
'Stats cache missing shotDistribution, forcing recomputation',
|
|
)
|
|
return getEmptyCache()
|
|
}
|
|
|
|
return parsed
|
|
} catch (error) {
|
|
logForDebugging(`Failed to load stats cache: ${errorMessage(error)}`)
|
|
return getEmptyCache()
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Save the stats cache to disk atomically.
|
|
* Uses a temp file + rename pattern to prevent corruption.
|
|
*/
|
|
export async function saveStatsCache(
|
|
cache: PersistedStatsCache,
|
|
): Promise<void> {
|
|
const fs = getFsImplementation()
|
|
const cachePath = getStatsCachePath()
|
|
const tempPath = `${cachePath}.${randomBytes(8).toString('hex')}.tmp`
|
|
|
|
try {
|
|
// Ensure the directory exists
|
|
const configDir = getClaudeConfigHomeDir()
|
|
try {
|
|
await fs.mkdir(configDir)
|
|
} catch {
|
|
// Directory already exists or other error - proceed
|
|
}
|
|
|
|
// Write to temp file with fsync for atomic write safety
|
|
const content = jsonStringify(cache, null, 2)
|
|
const handle = await open(tempPath, 'w', 0o600)
|
|
try {
|
|
await handle.writeFile(content, { encoding: 'utf-8' })
|
|
await handle.sync()
|
|
} finally {
|
|
await handle.close()
|
|
}
|
|
|
|
// Atomic rename
|
|
await fs.rename(tempPath, cachePath)
|
|
logForDebugging(
|
|
`Stats cache saved successfully (lastComputedDate: ${cache.lastComputedDate})`,
|
|
)
|
|
} catch (error) {
|
|
logError(error)
|
|
// Clean up temp file
|
|
try {
|
|
await fs.unlink(tempPath)
|
|
} catch {
|
|
// Ignore cleanup errors
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Merge new stats into an existing cache.
|
|
* Used when incrementally adding new days to the cache.
|
|
*/
|
|
export function mergeCacheWithNewStats(
|
|
existingCache: PersistedStatsCache,
|
|
newStats: {
|
|
dailyActivity: DailyActivity[]
|
|
dailyModelTokens: DailyModelTokens[]
|
|
modelUsage: { [modelName: string]: ModelUsage }
|
|
sessionStats: SessionStats[]
|
|
hourCounts: { [hour: number]: number }
|
|
totalSpeculationTimeSavedMs: number
|
|
shotDistribution?: { [shotCount: number]: number }
|
|
},
|
|
newLastComputedDate: string,
|
|
): PersistedStatsCache {
|
|
// Merge daily activity - combine by date
|
|
const dailyActivityMap = new Map<string, DailyActivity>()
|
|
for (const day of existingCache.dailyActivity) {
|
|
dailyActivityMap.set(day.date, { ...day })
|
|
}
|
|
for (const day of newStats.dailyActivity) {
|
|
const existing = dailyActivityMap.get(day.date)
|
|
if (existing) {
|
|
existing.messageCount += day.messageCount
|
|
existing.sessionCount += day.sessionCount
|
|
existing.toolCallCount += day.toolCallCount
|
|
} else {
|
|
dailyActivityMap.set(day.date, { ...day })
|
|
}
|
|
}
|
|
|
|
// Merge daily model tokens - combine by date
|
|
const dailyModelTokensMap = new Map<string, { [model: string]: number }>()
|
|
for (const day of existingCache.dailyModelTokens) {
|
|
dailyModelTokensMap.set(day.date, { ...day.tokensByModel })
|
|
}
|
|
for (const day of newStats.dailyModelTokens) {
|
|
const existing = dailyModelTokensMap.get(day.date)
|
|
if (existing) {
|
|
for (const [model, tokens] of Object.entries(day.tokensByModel)) {
|
|
existing[model] = (existing[model] || 0) + tokens
|
|
}
|
|
} else {
|
|
dailyModelTokensMap.set(day.date, { ...day.tokensByModel })
|
|
}
|
|
}
|
|
|
|
// Merge model usage
|
|
const modelUsage = { ...existingCache.modelUsage }
|
|
for (const [model, usage] of Object.entries(newStats.modelUsage)) {
|
|
if (modelUsage[model]) {
|
|
modelUsage[model] = {
|
|
inputTokens: modelUsage[model]!.inputTokens + usage.inputTokens,
|
|
outputTokens: modelUsage[model]!.outputTokens + usage.outputTokens,
|
|
cacheReadInputTokens:
|
|
modelUsage[model]!.cacheReadInputTokens + usage.cacheReadInputTokens,
|
|
cacheCreationInputTokens:
|
|
modelUsage[model]!.cacheCreationInputTokens +
|
|
usage.cacheCreationInputTokens,
|
|
webSearchRequests:
|
|
modelUsage[model]!.webSearchRequests + usage.webSearchRequests,
|
|
costUSD: modelUsage[model]!.costUSD + usage.costUSD,
|
|
contextWindow: Math.max(
|
|
modelUsage[model]!.contextWindow,
|
|
usage.contextWindow,
|
|
),
|
|
maxOutputTokens: Math.max(
|
|
modelUsage[model]!.maxOutputTokens,
|
|
usage.maxOutputTokens,
|
|
),
|
|
}
|
|
} else {
|
|
modelUsage[model] = { ...usage }
|
|
}
|
|
}
|
|
|
|
// Merge hour counts
|
|
const hourCounts = { ...existingCache.hourCounts }
|
|
for (const [hour, count] of Object.entries(newStats.hourCounts)) {
|
|
const hourNum = parseInt(hour, 10)
|
|
hourCounts[hourNum] = (hourCounts[hourNum] || 0) + count
|
|
}
|
|
|
|
// Update session aggregates
|
|
const totalSessions =
|
|
existingCache.totalSessions + newStats.sessionStats.length
|
|
const totalMessages =
|
|
existingCache.totalMessages +
|
|
newStats.sessionStats.reduce((sum, s) => sum + s.messageCount, 0)
|
|
|
|
// Find longest session (compare existing with new)
|
|
let longestSession = existingCache.longestSession
|
|
for (const session of newStats.sessionStats) {
|
|
if (!longestSession || session.duration > longestSession.duration) {
|
|
longestSession = session
|
|
}
|
|
}
|
|
|
|
// Find first session date
|
|
let firstSessionDate = existingCache.firstSessionDate
|
|
for (const session of newStats.sessionStats) {
|
|
if (!firstSessionDate || session.timestamp < firstSessionDate) {
|
|
firstSessionDate = session.timestamp
|
|
}
|
|
}
|
|
|
|
const result: PersistedStatsCache = {
|
|
version: STATS_CACHE_VERSION,
|
|
lastComputedDate: newLastComputedDate,
|
|
dailyActivity: Array.from(dailyActivityMap.values()).sort((a, b) =>
|
|
a.date.localeCompare(b.date),
|
|
),
|
|
dailyModelTokens: Array.from(dailyModelTokensMap.entries())
|
|
.map(([date, tokensByModel]) => ({ date, tokensByModel }))
|
|
.sort((a, b) => a.date.localeCompare(b.date)),
|
|
modelUsage,
|
|
totalSessions,
|
|
totalMessages,
|
|
longestSession,
|
|
firstSessionDate,
|
|
hourCounts,
|
|
totalSpeculationTimeSavedMs:
|
|
existingCache.totalSpeculationTimeSavedMs +
|
|
newStats.totalSpeculationTimeSavedMs,
|
|
}
|
|
|
|
if (feature('SHOT_STATS')) {
|
|
const shotDistribution: { [shotCount: number]: number } = {
|
|
...(existingCache.shotDistribution || {}),
|
|
}
|
|
for (const [count, sessions] of Object.entries(
|
|
newStats.shotDistribution || {},
|
|
)) {
|
|
const key = parseInt(count, 10)
|
|
shotDistribution[key] = (shotDistribution[key] || 0) + sessions
|
|
}
|
|
result.shotDistribution = shotDistribution
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
/**
|
|
* Extract the date portion (YYYY-MM-DD) from a Date object.
|
|
*/
|
|
export function toDateString(date: Date): string {
|
|
const parts = date.toISOString().split('T')
|
|
const dateStr = parts[0]
|
|
if (!dateStr) {
|
|
throw new Error('Invalid ISO date string')
|
|
}
|
|
return dateStr
|
|
}
|
|
|
|
/**
|
|
* Get today's date in YYYY-MM-DD format.
|
|
*/
|
|
export function getTodayDateString(): string {
|
|
return toDateString(new Date())
|
|
}
|
|
|
|
/**
|
|
* Get yesterday's date in YYYY-MM-DD format.
|
|
*/
|
|
export function getYesterdayDateString(): string {
|
|
const yesterday = new Date()
|
|
yesterday.setDate(yesterday.getDate() - 1)
|
|
return toDateString(yesterday)
|
|
}
|
|
|
|
/**
|
|
* Check if a date string is before another date string.
|
|
* Both should be in YYYY-MM-DD format.
|
|
*/
|
|
export function isDateBefore(date1: string, date2: string): boolean {
|
|
return date1 < date2
|
|
}
|