diff --git a/src/commands/config-command.ts b/src/commands/config-command.ts index 28bb0dfa..611b553c 100644 --- a/src/commands/config-command.ts +++ b/src/commands/config-command.ts @@ -160,10 +160,10 @@ export async function handleConfigCommand(args: string[]): Promise { try { // Start server - const { server, wss } = await startServer({ port, dev: options.dev }); + const { server, wss, cleanup } = await startServer({ port, dev: options.dev }); // Setup graceful shutdown - setupGracefulShutdown(server, wss); + setupGracefulShutdown(server, wss, cleanup); const url = `http://localhost:${port}`; diff --git a/src/web-server/index.ts b/src/web-server/index.ts index fd64547c..61dd2941 100644 --- a/src/web-server/index.ts +++ b/src/web-server/index.ts @@ -13,6 +13,7 @@ import { WebSocketServer } from 'ws'; import { setupWebSocket } from './websocket'; import { createSessionMiddleware, authMiddleware } from './middleware/auth-middleware'; import { startAutoSyncWatcher, stopAutoSyncWatcher } from '../cliproxy/sync'; +import { shutdownUsageAggregator } from './usage/aggregator'; export interface ServerOptions { port: number; @@ -107,6 +108,7 @@ export async function startServer(options: ServerOptions): Promise { wsCleanup(); stopAutoSyncWatcher().catch(() => {}); + shutdownUsageAggregator(); }; // Start listening diff --git a/src/web-server/usage/aggregator.ts b/src/web-server/usage/aggregator.ts index ca84bc5e..daca52ef 100644 --- a/src/web-server/usage/aggregator.ts +++ b/src/web-server/usage/aggregator.ts @@ -7,13 +7,7 @@ import * as fs from 'fs'; import * as path from 'path'; -import { - loadDailyUsageData, - loadMonthlyUsageData, - loadSessionData, - loadAllUsageData, - loadHourlyUsageData, -} from './data-aggregator'; +import { loadAllUsageData } from './data-aggregator'; import type { DailyUsage, HourlyUsage, MonthlyUsage, SessionUsage } from './types'; import { readDiskCache, @@ -29,6 +23,7 @@ import { loadCachedCliproxyData, startCliproxySync, stopCliproxySync, + syncCliproxyUsage, } from './cliproxy-usage-syncer'; // ============================================================================ @@ -268,6 +263,14 @@ let diskCacheInitialized = false; // Track if background refresh is in progress let isRefreshing = false; +// Coalesced full refresh promise shared across all usage loaders +let pendingFullRefresh: Promise<{ + daily: DailyUsage[]; + hourly: HourlyUsage[]; + monthly: MonthlyUsage[]; + session: SessionUsage[]; +}> | null = null; + /** * Persist cache to disk when we have enough data to be useful. */ @@ -293,6 +296,10 @@ async function refreshFromSource(): Promise<{ monthly: MonthlyUsage[]; session: SessionUsage[]; }> { + // Try to sync CLIProxy snapshot before reading it. + // Non-fatal: syncer handles unavailability and stale fallback. + await syncCliproxyUsage(); + // Load default data (from ~/.claude/projects/ or CLAUDE_CONFIG_DIR) const defaultData = await loadAllUsageData(); @@ -365,10 +372,35 @@ async function refreshFromSource(): Promise<{ return { daily, hourly, monthly, session }; } +async function refreshFromSourceCoalesced(force = false): Promise<{ + daily: DailyUsage[]; + hourly: HourlyUsage[]; + monthly: MonthlyUsage[]; + session: SessionUsage[]; +}> { + if (force) { + return refreshFromSource(); + } + + if (pendingFullRefresh) { + return pendingFullRefresh; + } + + pendingFullRefresh = refreshFromSource().finally(() => { + pendingFullRefresh = null; + }); + + return pendingFullRefresh; +} + /** * Initialize in-memory cache from disk cache (lazy - called on first API request). */ function ensureDiskCacheLoaded(): void { + // Start sync when usage APIs are actually accessed. + // startCliproxySync() is idempotent. + startCliproxySync(); + if (diskCacheInitialized) return; diskCacheInitialized = true; @@ -445,28 +477,28 @@ async function getCachedData(key: string, ttl: number, loader: () => Promise< /** Cached loader for daily usage data */ export async function getCachedDailyData(): Promise { return getCachedData('daily', CACHE_TTL.daily, async () => { - return await loadDailyUsageData(); + return (await refreshFromSourceCoalesced()).daily; }); } /** Cached loader for monthly usage data */ export async function getCachedMonthlyData(): Promise { return getCachedData('monthly', CACHE_TTL.monthly, async () => { - return await loadMonthlyUsageData(); + return (await refreshFromSourceCoalesced()).monthly; }); } /** Cached loader for session data */ export async function getCachedSessionData(): Promise { return getCachedData('session', CACHE_TTL.session, async () => { - return await loadSessionData(); + return (await refreshFromSourceCoalesced()).session; }); } /** Cached loader for hourly usage data */ export async function getCachedHourlyData(): Promise { return getCachedData('hourly', CACHE_TTL.daily, async () => { - return await loadHourlyUsageData(); + return (await refreshFromSourceCoalesced()).hourly; }); } @@ -475,9 +507,12 @@ export async function getCachedHourlyData(): Promise { */ export function clearUsageCache(): void { cache.clear(); + pendingRequests.clear(); + pendingFullRefresh = null; clearDiskCache(); // Reset so next API call will try to reload from disk/source diskCacheInitialized = false; + lastFetchTimestamp = null; } /** @@ -537,7 +572,7 @@ export async function prewarmUsageCache(): Promise<{ // Background refresh if (!isRefreshing) { isRefreshing = true; - refreshFromSource() + refreshFromSourceCoalesced() .then(() => console.log(ok('Background refresh complete'))) .catch((err) => console.error(fail(`Background refresh failed: ${err}`))) .finally(() => { @@ -550,7 +585,7 @@ export async function prewarmUsageCache(): Promise<{ // No usable disk cache - refresh from source (blocking for first startup only) console.log(info('No disk cache, loading from source...')); - await refreshFromSource(); + await refreshFromSourceCoalesced(); const elapsed = Date.now() - start; console.log(ok(`Usage cache ready (${elapsed}ms)`)); @@ -567,3 +602,13 @@ export async function prewarmUsageCache(): Promise<{ export function shutdownUsageAggregator(): void { stopCliproxySync(); } + +/** + * Force refresh usage cache from all sources. + * Used by manual refresh endpoint. + */ +export async function refreshUsageCache(): Promise { + // Ensure periodic sync is running for subsequent updates. + startCliproxySync(); + await refreshFromSourceCoalesced(true); +} diff --git a/src/web-server/usage/cliproxy-usage-syncer.ts b/src/web-server/usage/cliproxy-usage-syncer.ts index 6d894e49..cc36763f 100644 --- a/src/web-server/usage/cliproxy-usage-syncer.ts +++ b/src/web-server/usage/cliproxy-usage-syncer.ts @@ -145,6 +145,10 @@ export async function syncCliproxyUsage(): Promise { * Performs an immediate sync on startup. */ export function startCliproxySync(): void { + if (syncIntervalId !== null) { + return; + } + console.log(info('Starting CLIProxy usage sync (interval: 5 min)')); // Fire-and-forget initial sync diff --git a/src/web-server/usage/handlers.ts b/src/web-server/usage/handlers.ts index fb42867c..e9e34668 100644 --- a/src/web-server/usage/handlers.ts +++ b/src/web-server/usage/handlers.ts @@ -13,8 +13,8 @@ import { getCachedMonthlyData, getCachedSessionData, getCachedHourlyData, - clearUsageCache, getLastFetchTimestamp, + refreshUsageCache, } from './aggregator'; // ============================================================================ @@ -602,9 +602,13 @@ export async function handleMonthly( } } -export function handleRefresh(_req: Request, res: Response): void { - clearUsageCache(); - res.json({ success: true, message: 'Usage cache cleared' }); +export async function handleRefresh(_req: Request, res: Response): Promise { + try { + await refreshUsageCache(); + res.json({ success: true, message: 'Usage cache refreshed' }); + } catch (error) { + errorResponse(res, error, 'Failed to refresh usage cache'); + } } export function handleStatus(_req: Request, res: Response): void { diff --git a/tests/unit/web-server/cliproxy-usage-syncer.test.ts b/tests/unit/web-server/cliproxy-usage-syncer.test.ts new file mode 100644 index 00000000..1d3f165e --- /dev/null +++ b/tests/unit/web-server/cliproxy-usage-syncer.test.ts @@ -0,0 +1,97 @@ +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import type { CliproxyUsageApiResponse } from '../../../src/cliproxy/stats-fetcher'; + +let ccsDir = ''; +let rawResponse: CliproxyUsageApiResponse | null = null; +let fetchCalls = 0; + +mock.module('../../../src/utils/config-manager', () => ({ + getCcsDir: () => ccsDir, +})); + +mock.module('../../../src/cliproxy/stats-fetcher', () => ({ + fetchCliproxyUsageRaw: async () => { + fetchCalls++; + return rawResponse; + }, +})); + +let syncer: typeof import('../../../src/web-server/usage/cliproxy-usage-syncer'); + +beforeAll(async () => { + syncer = await import('../../../src/web-server/usage/cliproxy-usage-syncer'); +}); + +beforeEach(() => { + ccsDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-cliproxy-syncer-')); + fetchCalls = 0; + rawResponse = { + usage: { + apis: { + gemini: { + models: { + 'gemini-2.5-pro': { + details: [ + { + timestamp: '2026-03-02T12:00:00.000Z', + source: 'account-a', + auth_index: 0, + tokens: { + input_tokens: 100, + output_tokens: 20, + reasoning_tokens: 0, + cached_tokens: 10, + total_tokens: 130, + }, + failed: false, + }, + ], + }, + }, + }, + }, + }, + }; + syncer.stopCliproxySync(); +}); + +afterEach(() => { + syncer.stopCliproxySync(); + fs.rmSync(ccsDir, { recursive: true, force: true }); +}); + +afterAll(() => { + mock.restore(); +}); + +describe('cliproxy usage syncer', () => { + it('writes and loads snapshot data', async () => { + await syncer.syncCliproxyUsage(); + + const snapshotPath = path.join(ccsDir, 'cache', 'cliproxy-usage', 'latest.json'); + expect(fs.existsSync(snapshotPath)).toBe(true); + + const cached = await syncer.loadCachedCliproxyData(); + expect(cached.daily).toHaveLength(1); + expect(cached.daily[0].source).toBe('cliproxy'); + expect(cached.daily[0].inputTokens).toBe(100); + expect(cached.hourly).toHaveLength(1); + expect(cached.monthly).toHaveLength(1); + }); + + it('startCliproxySync is idempotent and starts only one interval', () => { + const intervalSpy = spyOn(globalThis, 'setInterval'); + + syncer.startCliproxySync(); + syncer.startCliproxySync(); + + expect(intervalSpy).toHaveBeenCalledTimes(1); + expect(fetchCalls).toBeGreaterThan(0); + + syncer.stopCliproxySync(); + intervalSpy.mockRestore(); + }); +}); diff --git a/tests/unit/web-server/cliproxy-usage-transformer.test.ts b/tests/unit/web-server/cliproxy-usage-transformer.test.ts new file mode 100644 index 00000000..ac2e6eec --- /dev/null +++ b/tests/unit/web-server/cliproxy-usage-transformer.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, it } from 'bun:test'; +import type { CliproxyUsageApiResponse } from '../../../src/cliproxy/stats-fetcher'; +import { + flattenCliproxyDetails, + transformCliproxyToDailyUsage, + transformCliproxyToHourlyUsage, + transformCliproxyToMonthlyUsage, +} from '../../../src/web-server/usage/cliproxy-usage-transformer'; + +const sampleResponse: CliproxyUsageApiResponse = { + usage: { + apis: { + gemini: { + models: { + 'gemini-2.5-pro': { + details: [ + { + timestamp: '2026-03-01T10:15:00.000Z', + source: 'account-a', + auth_index: 0, + tokens: { + input_tokens: 100, + output_tokens: 50, + reasoning_tokens: 0, + cached_tokens: 20, + total_tokens: 170, + }, + failed: false, + }, + { + timestamp: '2026-03-01T11:30:00.000Z', + source: 'account-a', + auth_index: 0, + tokens: { + input_tokens: 40, + output_tokens: 10, + reasoning_tokens: 0, + cached_tokens: 5, + total_tokens: 55, + }, + failed: true, + }, + { + timestamp: '2026-03-01T10:45:00.000Z', + source: 'account-a', + auth_index: 0, + tokens: { + input_tokens: 30, + output_tokens: 20, + reasoning_tokens: 0, + cached_tokens: 10, + total_tokens: 60, + }, + failed: false, + }, + ], + }, + }, + }, + codex: { + models: { + 'gpt-4.1': { + details: [ + { + timestamp: '2026-03-02T01:00:00.000Z', + source: 'account-b', + auth_index: 1, + tokens: { + input_tokens: 70, + output_tokens: 30, + reasoning_tokens: 0, + cached_tokens: 0, + total_tokens: 100, + }, + failed: false, + }, + ], + }, + }, + }, + }, + }, +}; + +describe('cliproxy usage transformer', () => { + it('flattens nested API details and skips failed requests', () => { + const flat = flattenCliproxyDetails(sampleResponse); + expect(flat).toHaveLength(3); + expect(flat.every((entry) => entry.detail.failed === false)).toBe(true); + }); + + it('transforms daily usage with aggregated model totals', () => { + const daily = transformCliproxyToDailyUsage(sampleResponse); + + expect(daily).toHaveLength(2); + expect(daily[0].date).toBe('2026-03-02'); + expect(daily[0].source).toBe('cliproxy'); + expect(daily[1].date).toBe('2026-03-01'); + + const marchFirst = daily.find((d) => d.date === '2026-03-01'); + expect(marchFirst?.inputTokens).toBe(130); + expect(marchFirst?.outputTokens).toBe(70); + expect(marchFirst?.cacheReadTokens).toBe(30); + expect(marchFirst?.modelsUsed).toContain('gemini-2.5-pro'); + }); + + it('transforms hourly usage with hour buckets', () => { + const hourly = transformCliproxyToHourlyUsage(sampleResponse); + + expect(hourly).toHaveLength(2); + expect(hourly[0].hour).toBe('2026-03-02 01:00'); + + const tenAm = hourly.find((h) => h.hour === '2026-03-01 10:00'); + expect(tenAm?.inputTokens).toBe(130); + expect(tenAm?.outputTokens).toBe(70); + }); + + it('transforms monthly usage with cliproxy source', () => { + const monthly = transformCliproxyToMonthlyUsage(sampleResponse); + + expect(monthly).toHaveLength(1); + expect(monthly[0].month).toBe('2026-03'); + expect(monthly[0].source).toBe('cliproxy'); + expect(monthly[0].inputTokens).toBe(200); + expect(monthly[0].outputTokens).toBe(100); + expect(monthly[0].cacheReadTokens).toBe(30); + }); +}); diff --git a/tests/unit/web-server/usage-aggregator-cliproxy-integration.test.ts b/tests/unit/web-server/usage-aggregator-cliproxy-integration.test.ts new file mode 100644 index 00000000..3b5cd6cf --- /dev/null +++ b/tests/unit/web-server/usage-aggregator-cliproxy-integration.test.ts @@ -0,0 +1,196 @@ +import { afterEach, beforeEach, describe, expect, it } from 'bun:test'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; + +let tempRoot = ''; +let ccsDir = ''; +let claudeDir = ''; +let aggregator: typeof import('../../../src/web-server/usage/aggregator'); +let originalCcsDir: string | undefined; +let originalClaudeConfigDir: string | undefined; + +function writeClaudeJsonlFixture(): void { + const projectDir = path.join(claudeDir, 'projects', 'project-one'); + fs.mkdirSync(projectDir, { recursive: true }); + + const line = JSON.stringify({ + type: 'assistant', + sessionId: 'session-a', + timestamp: '2026-03-02T10:00:00.000Z', + version: '1.0.0', + cwd: '/tmp/project', + message: { + model: 'claude-sonnet-4-5', + usage: { + input_tokens: 100, + output_tokens: 40, + }, + }, + }); + + fs.writeFileSync(path.join(projectDir, 'usage.jsonl'), `${line}\n`, 'utf-8'); +} + +function writeCliproxySnapshotFixture(): void { + const snapshotDir = path.join(ccsDir, 'cache', 'cliproxy-usage'); + fs.mkdirSync(snapshotDir, { recursive: true }); + + const snapshot = { + version: 1, + timestamp: Date.now(), + daily: [ + { + date: '2026-03-02', + source: 'cliproxy', + inputTokens: 50, + outputTokens: 10, + cacheCreationTokens: 0, + cacheReadTokens: 5, + cost: 0.2, + totalCost: 0.2, + modelsUsed: ['gemini-2.5-pro'], + modelBreakdowns: [ + { + modelName: 'gemini-2.5-pro', + inputTokens: 50, + outputTokens: 10, + cacheCreationTokens: 0, + cacheReadTokens: 5, + cost: 0.2, + }, + ], + }, + ], + hourly: [ + { + hour: '2026-03-02 10:00', + source: 'cliproxy', + inputTokens: 50, + outputTokens: 10, + cacheCreationTokens: 0, + cacheReadTokens: 5, + cost: 0.2, + totalCost: 0.2, + modelsUsed: ['gemini-2.5-pro'], + modelBreakdowns: [ + { + modelName: 'gemini-2.5-pro', + inputTokens: 50, + outputTokens: 10, + cacheCreationTokens: 0, + cacheReadTokens: 5, + cost: 0.2, + }, + ], + }, + ], + monthly: [ + { + month: '2026-03', + source: 'cliproxy', + inputTokens: 50, + outputTokens: 10, + cacheCreationTokens: 0, + cacheReadTokens: 5, + totalCost: 0.2, + modelsUsed: ['gemini-2.5-pro'], + modelBreakdowns: [ + { + modelName: 'gemini-2.5-pro', + inputTokens: 50, + outputTokens: 10, + cacheCreationTokens: 0, + cacheReadTokens: 5, + cost: 0.2, + }, + ], + }, + ], + }; + + fs.writeFileSync(path.join(snapshotDir, 'latest.json'), JSON.stringify(snapshot), 'utf-8'); +} + +function writeUnifiedConfigFixture(): void { + const yaml = `version: 2 +accounts: {} +profiles: {} +preferences: + theme: system + telemetry: false + auto_update: true +cliproxy: + oauth_accounts: {} + providers: + - gemini + - codex + - agy + variants: {} +cliproxy_server: + local: + port: 65534 +`; + + fs.mkdirSync(ccsDir, { recursive: true }); + fs.writeFileSync(path.join(ccsDir, 'config.yaml'), yaml, 'utf-8'); +} + +beforeEach(async () => { + tempRoot = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-usage-agg-')); + ccsDir = path.join(tempRoot, '.ccs'); + claudeDir = path.join(tempRoot, '.claude'); + + originalCcsDir = process.env.CCS_DIR; + originalClaudeConfigDir = process.env.CLAUDE_CONFIG_DIR; + process.env.CCS_DIR = ccsDir; + process.env.CLAUDE_CONFIG_DIR = claudeDir; + + writeUnifiedConfigFixture(); + writeClaudeJsonlFixture(); + writeCliproxySnapshotFixture(); + + aggregator = await import('../../../src/web-server/usage/aggregator'); + aggregator.clearUsageCache(); +}); + +afterEach(() => { + aggregator.shutdownUsageAggregator(); + aggregator.clearUsageCache(); + + if (originalCcsDir !== undefined) { + process.env.CCS_DIR = originalCcsDir; + } else { + delete process.env.CCS_DIR; + } + + if (originalClaudeConfigDir !== undefined) { + process.env.CLAUDE_CONFIG_DIR = originalClaudeConfigDir; + } else { + delete process.env.CLAUDE_CONFIG_DIR; + } + + fs.rmSync(tempRoot, { recursive: true, force: true }); +}); + +describe('usage aggregator cliproxy integration', () => { + it('merges cliproxy snapshot data into getCachedDailyData', async () => { + const daily = await aggregator.getCachedDailyData(); + + expect(daily).toHaveLength(1); + expect(daily[0].date).toBe('2026-03-02'); + expect(daily[0].inputTokens).toBe(150); + expect(daily[0].outputTokens).toBe(50); + expect(daily[0].cacheReadTokens).toBe(5); + expect(daily[0].modelsUsed).toContain('claude-sonnet-4-5'); + expect(daily[0].modelsUsed).toContain('gemini-2.5-pro'); + }); + + it('clearUsageCache resets last fetch timestamp', async () => { + await aggregator.getCachedDailyData(); + expect(aggregator.getLastFetchTimestamp()).not.toBeNull(); + + aggregator.clearUsageCache(); + expect(aggregator.getLastFetchTimestamp()).toBeNull(); + }); +});