diff --git a/src/web-server/usage/codex-native-usage-collector.ts b/src/web-server/usage/codex-native-usage-collector.ts index 9d4c8f4b..160fcfb7 100644 --- a/src/web-server/usage/codex-native-usage-collector.ts +++ b/src/web-server/usage/codex-native-usage-collector.ts @@ -1,6 +1,7 @@ import * as fs from 'fs'; import * as path from 'path'; import * as readline from 'readline'; +import { getCcsDir } from '../../utils/config-manager'; import type { RawUsageEntry } from '../jsonl-parser'; import { resolveCodexConfigPaths } from '../services/codex-dashboard-service'; @@ -8,6 +9,24 @@ interface CodexNativeUsageCollectorOptions { env?: NodeJS.ProcessEnv; homeDir?: string; includeCliproxySessions?: boolean; + cacheDir?: string; + disableCache?: boolean; +} + +const CODEX_NATIVE_USAGE_CACHE_VERSION = 1; + +interface CachedRolloutFile { + path: string; + size: number; + mtimeMs: number; + entries: RawUsageEntry[]; +} + +interface CodexNativeUsageCache { + version: typeof CODEX_NATIVE_USAGE_CACHE_VERSION; + includeCliproxySessions: boolean; + generatedAt: number; + files: Record; } interface CodexTokenSnapshot { @@ -82,6 +101,85 @@ async function collectRolloutFiles(dir: string): Promise { return files.sort(); } +function getDefaultCacheDir(): string { + return path.join(getCcsDir(), 'cache'); +} + +function getCacheFilePath(cacheDir: string, includeCliproxySessions: boolean): string { + return path.join( + cacheDir, + includeCliproxySessions + ? 'codex-native-usage-with-cliproxy-v1.json' + : 'codex-native-usage-v1.json' + ); +} + +function isCachedRolloutFile(value: unknown): value is CachedRolloutFile { + if (!isObject(value)) return false; + return ( + typeof value.path === 'string' && + typeof value.size === 'number' && + typeof value.mtimeMs === 'number' && + Array.isArray(value.entries) + ); +} + +function isCodexNativeUsageCache(value: unknown): value is CodexNativeUsageCache { + if (!isObject(value)) return false; + if (value.version !== CODEX_NATIVE_USAGE_CACHE_VERSION) return false; + if (typeof value.includeCliproxySessions !== 'boolean') return false; + if (typeof value.generatedAt !== 'number') return false; + if (!isObject(value.files)) return false; + return Object.values(value.files).every(isCachedRolloutFile); +} + +function readUsageCache( + cacheDir: string, + includeCliproxySessions: boolean +): CodexNativeUsageCache | null { + try { + const cachePath = getCacheFilePath(cacheDir, includeCliproxySessions); + if (!fs.existsSync(cachePath)) return null; + + const parsed = JSON.parse(fs.readFileSync(cachePath, 'utf8')) as unknown; + if (!isCodexNativeUsageCache(parsed)) return null; + if (parsed.includeCliproxySessions !== includeCliproxySessions) return null; + + return parsed; + } catch { + return null; + } +} + +function writeUsageCache(cacheDir: string, cache: CodexNativeUsageCache): void { + try { + fs.mkdirSync(cacheDir, { recursive: true }); + const cachePath = getCacheFilePath(cacheDir, cache.includeCliproxySessions); + const tempPath = `${cachePath}.${process.pid}.tmp`; + fs.writeFileSync(tempPath, JSON.stringify(cache), 'utf8'); + fs.renameSync(tempPath, cachePath); + } catch { + // Best-effort only. + } +} + +function getMtimeMsFingerprint(stats: fs.Stats): number { + return stats.mtimeMs; +} + +function hasMatchingFingerprint( + cached: CachedRolloutFile | undefined, + filePath: string, + stats: fs.Stats +): cached is CachedRolloutFile { + return ( + !!cached && + cached.path === filePath && + cached.size === stats.size && + cached.mtimeMs === getMtimeMsFingerprint(stats) + ); +} + async function parseRolloutFile( filePath: string, includeCliproxySessions: boolean @@ -183,6 +281,7 @@ async function parseRolloutFile( export async function scanCodexNativeUsageEntries( options: CodexNativeUsageCollectorOptions = {} ): Promise { + const includeCliproxySessions = options.includeCliproxySessions === true; const { baseDir } = resolveCodexConfigPaths({ env: options.env, homeDir: options.homeDir, @@ -190,9 +289,42 @@ export async function scanCodexNativeUsageEntries( const rolloutFiles = await collectRolloutFiles(path.join(baseDir, 'sessions')); const entries: RawUsageEntry[] = []; - for (const filePath of rolloutFiles) { - entries.push(...(await parseRolloutFile(filePath, options.includeCliproxySessions === true))); + if (options.disableCache === true) { + for (const filePath of rolloutFiles) { + entries.push(...(await parseRolloutFile(filePath, includeCliproxySessions))); + } + return entries; } + const cacheDir = options.cacheDir ?? getDefaultCacheDir(); + const previousCache = readUsageCache(cacheDir, includeCliproxySessions); + const nextCache: CodexNativeUsageCache = { + version: CODEX_NATIVE_USAGE_CACHE_VERSION, + includeCliproxySessions, + generatedAt: Date.now(), + files: {}, + }; + + for (const filePath of rolloutFiles) { + const stats = await fs.promises.stat(filePath); + const cached = previousCache?.files[filePath]; + + if (hasMatchingFingerprint(cached, filePath, stats)) { + entries.push(...cached.entries); + nextCache.files[filePath] = cached; + continue; + } + + const fileEntries = await parseRolloutFile(filePath, includeCliproxySessions); + entries.push(...fileEntries); + nextCache.files[filePath] = { + path: filePath, + size: stats.size, + mtimeMs: getMtimeMsFingerprint(stats), + entries: fileEntries, + }; + } + + writeUsageCache(cacheDir, nextCache); return entries; } diff --git a/tests/unit/web-server/codex-native-usage-collector.test.ts b/tests/unit/web-server/codex-native-usage-collector.test.ts index aa32db7d..50d68991 100644 --- a/tests/unit/web-server/codex-native-usage-collector.test.ts +++ b/tests/unit/web-server/codex-native-usage-collector.test.ts @@ -2,8 +2,18 @@ import { afterEach, beforeEach, describe, expect, it } from 'bun:test'; import * as fs from 'fs'; import * as os from 'os'; import * as path from 'path'; +import { runWithScopedCcsHome } from '../../../src/utils/config-manager'; import { scanCodexNativeUsageEntries } from '../../../src/web-server/usage/codex-native-usage-collector'; +type TestUsageCache = { + includeCliproxySessions: boolean; + files: Record; +}; + +function readTestUsageCache(cachePath: string): TestUsageCache { + return JSON.parse(fs.readFileSync(cachePath, 'utf8')) as TestUsageCache; +} + function writeCodexRollout( baseDir: string, options: { @@ -12,7 +22,7 @@ function writeCodexRollout( model?: string; cwd?: string; } = {} -): void { +): string { const sessionId = options.sessionId ?? 'codex-session-1'; const rolloutDir = path.join(baseDir, 'sessions', '2026', '03', '02'); const rolloutPath = path.join(rolloutDir, `rollout-${sessionId}.jsonl`); @@ -124,6 +134,38 @@ function writeCodexRollout( fs.mkdirSync(rolloutDir, { recursive: true }); fs.writeFileSync(rolloutPath, `${lines.join('\n')}\n`, 'utf8'); + return rolloutPath; +} + +function appendCodexTokenCount(rolloutPath: string): void { + fs.appendFileSync( + rolloutPath, + `${JSON.stringify({ + timestamp: '2026-03-02T10:15:00.000Z', + type: 'event_msg', + payload: { + type: 'token_count', + info: { + total_token_usage: { + input_tokens: 200, + cached_input_tokens: 40, + output_tokens: 20, + reasoning_output_tokens: 5, + total_tokens: 265, + }, + last_token_usage: { + input_tokens: 50, + cached_input_tokens: 10, + output_tokens: 10, + reasoning_output_tokens: 2, + total_tokens: 72, + }, + model_context_window: 200000, + }, + }, + })}\n`, + 'utf8' + ); } describe('codex native usage collector', () => { @@ -137,13 +179,21 @@ describe('codex native usage collector', () => { fs.rmSync(tempRoot, { recursive: true, force: true }); }); + function getCacheDir(): string { + const cacheDir = path.join(tempRoot, 'ccs-cache'); + fs.mkdirSync(cacheDir, { recursive: true }); + return cacheDir; + } + it('parses token_count events into raw usage entries and suppresses duplicates', async () => { writeCodexRollout(tempRoot); - const entries = await scanCodexNativeUsageEntries({ - env: { CODEX_HOME: tempRoot }, - homeDir: tempRoot, - }); + const entries = await runWithScopedCcsHome(tempRoot, () => + scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + }) + ); expect(entries).toHaveLength(2); expect(entries[0]).toMatchObject({ @@ -166,10 +216,12 @@ describe('codex native usage collector', () => { it('skips cliproxy-backed codex sessions by default to avoid double counting', async () => { writeCodexRollout(tempRoot, { modelProvider: 'cliproxy' }); - const entries = await scanCodexNativeUsageEntries({ - env: { CODEX_HOME: tempRoot }, - homeDir: tempRoot, - }); + const entries = await runWithScopedCcsHome(tempRoot, () => + scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + }) + ); expect(entries).toHaveLength(0); }); @@ -177,11 +229,121 @@ describe('codex native usage collector', () => { it('also skips ccs_runtime-backed codex bridge sessions by default', async () => { writeCodexRollout(tempRoot, { modelProvider: 'ccs_runtime' }); - const entries = await scanCodexNativeUsageEntries({ - env: { CODEX_HOME: tempRoot }, - homeDir: tempRoot, - }); + const entries = await runWithScopedCcsHome(tempRoot, () => + scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + }) + ); expect(entries).toHaveLength(0); }); + + it('reuses cached rollout entries when file size and mtime are unchanged', async () => { + const rolloutPath = writeCodexRollout(tempRoot); + const cacheDir = getCacheDir(); + const stableMtime = new Date('2026-03-02T10:20:00.000Z'); + fs.utimesSync(rolloutPath, stableMtime, stableMtime); + + const firstEntries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + cacheDir, + }); + + const cachePath = path.join(cacheDir, 'codex-native-usage-v1.json'); + const originalStats = fs.statSync(rolloutPath); + const originalCache = readTestUsageCache(cachePath); + const originalCachedRollout = originalCache.files[rolloutPath]; + expect(originalCachedRollout?.mtimeMs).toBe(originalStats.mtimeMs); + + const originalContent = fs.readFileSync(rolloutPath, 'utf8'); + fs.writeFileSync(rolloutPath, '#'.repeat(originalContent.length), 'utf8'); + fs.utimesSync(rolloutPath, stableMtime, stableMtime); + + const secondEntries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + cacheDir, + }); + + expect(secondEntries).toEqual(firstEntries); + }); + + it('reparses a rollout file after its size changes', async () => { + const rolloutPath = writeCodexRollout(tempRoot); + const cacheDir = getCacheDir(); + + const firstEntries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + cacheDir, + }); + + appendCodexTokenCount(rolloutPath); + + const secondEntries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + cacheDir, + }); + + expect(firstEntries).toHaveLength(2); + expect(secondEntries).toHaveLength(3); + expect(secondEntries[2]).toMatchObject({ + inputTokens: 50, + cacheReadTokens: 10, + outputTokens: 12, + }); + }); + + it('falls back to parsing rollout files when the cache file is corrupt', async () => { + writeCodexRollout(tempRoot); + const cacheDir = getCacheDir(); + fs.writeFileSync(path.join(cacheDir, 'codex-native-usage-v1.json'), '{not-json', 'utf8'); + + const entries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + cacheDir, + }); + + expect(entries).toHaveLength(2); + }); + + it('keeps default and include-cliproxy cache entries separate', async () => { + writeCodexRollout(tempRoot, { modelProvider: 'cliproxy' }); + const cacheDir = getCacheDir(); + + const defaultEntries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + cacheDir, + }); + + const includedEntries = await scanCodexNativeUsageEntries({ + env: { CODEX_HOME: tempRoot }, + homeDir: tempRoot, + includeCliproxySessions: true, + cacheDir, + }); + + expect(defaultEntries).toHaveLength(0); + expect(includedEntries).toHaveLength(2); + + const defaultCachePath = path.join(cacheDir, 'codex-native-usage-v1.json'); + const includeCachePath = path.join(cacheDir, 'codex-native-usage-with-cliproxy-v1.json'); + expect(fs.existsSync(defaultCachePath)).toBe(true); + expect(fs.existsSync(includeCachePath)).toBe(true); + + const defaultCache = readTestUsageCache(defaultCachePath); + const includeCache = readTestUsageCache(includeCachePath); + const defaultCachedRollout = defaultCache.files[Object.keys(defaultCache.files)[0] ?? '']; + const includeCachedRollout = includeCache.files[Object.keys(includeCache.files)[0] ?? '']; + + expect(defaultCache.includeCliproxySessions).toBe(false); + expect(defaultCachedRollout?.entries).toHaveLength(0); + expect(includeCache.includeCliproxySessions).toBe(true); + expect(includeCachedRollout?.entries).toHaveLength(2); + }); });