Merge pull request #668 from kaitranntt/kai/wip/post-merge-pr660-260303-211636

fix(usage): harden cliproxy refresh lifecycle

src/commands/config-command.ts

@@ -160,10 +160,10 @@ export async function handleConfigCommand(args: string[]): Promise<void> {
 
   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}`;
 

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<ServerInstanc
   const cleanup = () => {
     wsCleanup();
     stopAutoSyncWatcher().catch(() => {});
+    shutdownUsageAggregator();
   };
 
   // Start listening

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<T>(key: string, ttl: number, loader: () => Promise<
 /** Cached loader for daily usage data */
 export async function getCachedDailyData(): Promise<DailyUsage[]> {
   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<MonthlyUsage[]> {
   return getCachedData('monthly', CACHE_TTL.monthly, async () => {
-    return await loadMonthlyUsageData();
+    return (await refreshFromSourceCoalesced()).monthly;
   });
 }
 
 /** Cached loader for session data */
 export async function getCachedSessionData(): Promise<SessionUsage[]> {
   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<HourlyUsage[]> {
   return getCachedData('hourly', CACHE_TTL.daily, async () => {
-    return await loadHourlyUsageData();
+    return (await refreshFromSourceCoalesced()).hourly;
   });
 }
 
@@ -475,9 +507,12 @@ export async function getCachedHourlyData(): Promise<HourlyUsage[]> {
  */
 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<void> {
+  // Ensure periodic sync is running for subsequent updates.
+  startCliproxySync();
+  await refreshFromSourceCoalesced(true);
+}

src/web-server/usage/cliproxy-usage-syncer.ts

@@ -145,6 +145,10 @@ export async function syncCliproxyUsage(): Promise<void> {
  * 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

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<void> {
+  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 {

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();
+  });
+});

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);
+  });
+});

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();
+  });
+});