fix(usage): address code review findings - SQL injection, stream usage recording, DB singleton

package-lock.json

@@ -51,6 +51,7 @@
         "electron-squirrel-startup": "^1.0.1",
         "fastify": "^5.8.4",
         "find-process": "^2.0.0",
+        "gpt-tokenizer": "^3.4.0",
         "i18next": "^25.6.3",
         "i18next-browser-languagedetector": "^8.2.0",
         "keytar": "^7.9.0",
@@ -18531,6 +18532,12 @@
         "url": "https://github.com/sindresorhus/got?sponsor=1"
       }
     },
+    "node_modules/gpt-tokenizer": {
+      "version": "3.4.0",
+      "resolved": "https://registry.npmjs.org/gpt-tokenizer/-/gpt-tokenizer-3.4.0.tgz",
+      "integrity": "sha512-wxFLnhIXTDjYebd9A9pGl3e31ZpSypbpIJSOswbgop5jLte/AsZVDvjlbEuVFlsqZixVKqbcoNmRlFDf6pz/UQ==",
+      "license": "MIT"
+    },
     "node_modules/graceful-fs": {
       "version": "4.2.11",
       "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz",

package.json

@@ -136,6 +136,7 @@
     "electron-squirrel-startup": "^1.0.1",
     "fastify": "^5.8.4",
     "find-process": "^2.0.0",
+    "gpt-tokenizer": "^3.4.0",
     "i18next": "^25.6.3",
     "i18next-browser-languagedetector": "^8.2.0",
     "keytar": "^7.9.0",

src/ipc/database/cloudHandler.ts

@@ -59,7 +59,7 @@ function sleepSync(ms: number): void {
  * Ensures that the cloud database file and schema exist.
  * @param dbPath {string} The path to the database file.
  */
-function ensureDatabaseInitialized(dbPath: string): void {
+export function ensureDatabaseInitialized(dbPath: string): void {
   const dir = path.dirname(dbPath);
   if (!fs.existsSync(dir)) {
     fs.mkdirSync(dir, { recursive: true });
@@ -68,6 +68,7 @@ function ensureDatabaseInitialized(dbPath: string): void {
   let db: Database.Database | null = null;
   try {
     db = new Database(dbPath);
+    db.pragma('journal_mode = WAL');
     configureDatabase(db, { busyTimeoutMs: SQLITE_BUSY_TIMEOUT_MS });
 
     // Create accounts table

src/ipc/database/usageHandler.ts

@@ -1,5 +1,9 @@
 import { logger } from '../../utils/logger';
-import { getCloudDb } from './cloudHandler';
+import Database from 'better-sqlite3';
+import type { BetterSQLite3Database } from 'drizzle-orm/better-sqlite3';
+import { openDrizzleConnection } from './dbConnection';
+import { getCloudAccountsDbPath } from '../../utils/paths';
+import { ensureDatabaseInitialized } from './cloudHandler';
 import { tokenUsage } from './schema';
 import {
   HourlyUsageSchema,
@@ -19,24 +23,47 @@ export interface RecordUsageParams {
   requestType?: string;
 }
 
-function buildDateFilter(start?: number, end?: number): string {
-  const conditions: string[] = [];
-  if (start !== undefined) {
-    conditions.push(`timestamp >= ${start}`);
+let usageDbConnection: { raw: Database.Database; orm: BetterSQLite3Database<typeof import('./schema')> } | null = null;
+
+function getUsageDb() {
+  if (!usageDbConnection) {
+    const dbPath = getCloudAccountsDbPath();
+    ensureDatabaseInitialized(dbPath);
+    usageDbConnection = openDrizzleConnection(
+      dbPath,
+      { readonly: false, fileMustExist: false },
+      { busyTimeoutMs: 3000 },
+    );
   }
-  if (end !== undefined) {
-    conditions.push(`timestamp <= ${end}`);
-  }
-  return conditions.length > 0 ? `AND ${conditions.join(' AND ')}` : '';
+  return usageDbConnection;
 }
 
-function buildAccountFilter(accountId?: string): string {
-  return accountId ? `AND account_id = '${accountId}'` : '';
+function buildDateFilter(start?: number, end?: number): { sql: string; params: number[] } {
+  const conditions: string[] = [];
+  const params: number[] = [];
+  if (start !== undefined) {
+    conditions.push(`timestamp >= ?`);
+    params.push(start);
+  }
+  if (end !== undefined) {
+    conditions.push(`timestamp <= ?`);
+    params.push(end);
+  }
+  return {
+    sql: conditions.length > 0 ? `AND ${conditions.join(' AND ')}` : '',
+    params,
+  };
+}
+
+function buildAccountFilter(accountId?: string): { sql: string; params: unknown[] } {
+  return accountId
+    ? { sql: 'AND account_id = ?', params: [accountId] }
+    : { sql: '', params: [] };
 }
 
 export class TokenUsageRepo {
-  static recordUsage(params: RecordUsageParams): void {
-    const { raw, orm } = getCloudDb();
+  static recordUsage(params: RecordUsageParams): boolean {
+    const { orm } = getUsageDb();
     try {
       orm
         .insert(tokenUsage)
@@ -50,10 +77,10 @@ export class TokenUsageRepo {
           requestType: params.requestType ?? null,
         })
         .run();
+      return true;
     } catch (error) {
       logger.error('Failed to record token usage', error);
-    } finally {
-      raw.close();
+      return false;
     }
   }
 
@@ -68,8 +95,10 @@ export class TokenUsageRepo {
     totalTokens: number;
     requests: number;
   }> {
-    const { raw } = getCloudDb();
+    const { raw } = getUsageDb();
     try {
+      const accountFilter = buildAccountFilter(accountId);
+      const dateFilter = buildDateFilter(start, end);
       const sql = `
         SELECT
           strftime('%Y-%m-%d %H:00', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
@@ -79,18 +108,16 @@ export class TokenUsageRepo {
           COUNT(*) as requests
         FROM token_usage
         WHERE 1=1
-        ${buildAccountFilter(accountId)}
-        ${buildDateFilter(start, end)}
+        ${accountFilter.sql}
+        ${dateFilter.sql}
         GROUP BY bucket
         ORDER BY bucket
       `;
-      const rows = raw.prepare(sql).all() as unknown[];
+      const rows = raw.prepare(sql).all(...accountFilter.params, ...dateFilter.params) as unknown[];
       return rows.map((row) => HourlyUsageSchema.parse(row));
     } catch (error) {
       logger.error('Failed to get hourly usage', error);
       return [];
-    } finally {
-      raw.close();
     }
   }
 
@@ -105,8 +132,10 @@ export class TokenUsageRepo {
     totalTokens: number;
     requests: number;
   }> {
-    const { raw } = getCloudDb();
+    const { raw } = getUsageDb();
     try {
+      const accountFilter = buildAccountFilter(accountId);
+      const dateFilter = buildDateFilter(start, end);
       const sql = `
         SELECT
           strftime('%Y-%m-%d', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
@@ -116,18 +145,16 @@ export class TokenUsageRepo {
           COUNT(*) as requests
         FROM token_usage
         WHERE 1=1
-        ${buildAccountFilter(accountId)}
-        ${buildDateFilter(start, end)}
+        ${accountFilter.sql}
+        ${dateFilter.sql}
         GROUP BY bucket
         ORDER BY bucket
       `;
-      const rows = raw.prepare(sql).all() as unknown[];
+      const rows = raw.prepare(sql).all(...accountFilter.params, ...dateFilter.params) as unknown[];
       return rows.map((row) => DailyUsageSchema.parse(row));
     } catch (error) {
       logger.error('Failed to get daily usage', error);
       return [];
-    } finally {
-      raw.close();
     }
   }
 
@@ -142,8 +169,10 @@ export class TokenUsageRepo {
     totalTokens: number;
     requests: number;
   }> {
-    const { raw } = getCloudDb();
+    const { raw } = getUsageDb();
     try {
+      const accountFilter = buildAccountFilter(accountId);
+      const dateFilter = buildDateFilter(start, end);
       const sql = `
         SELECT
           strftime('%Y-%W', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
@@ -153,18 +182,16 @@ export class TokenUsageRepo {
           COUNT(*) as requests
         FROM token_usage
         WHERE 1=1
-        ${buildAccountFilter(accountId)}
-        ${buildDateFilter(start, end)}
+        ${accountFilter.sql}
+        ${dateFilter.sql}
         GROUP BY bucket
         ORDER BY bucket
       `;
-      const rows = raw.prepare(sql).all() as unknown[];
+      const rows = raw.prepare(sql).all(...accountFilter.params, ...dateFilter.params) as unknown[];
       return rows.map((row) => WeeklyUsageSchema.parse(row));
     } catch (error) {
       logger.error('Failed to get weekly usage', error);
       return [];
-    } finally {
-      raw.close();
     }
   }
 
@@ -179,8 +206,10 @@ export class TokenUsageRepo {
     totalTokens: number;
     requests: number;
   }> {
-    const { raw } = getCloudDb();
+    const { raw } = getUsageDb();
     try {
+      const accountFilter = buildAccountFilter(accountId);
+      const dateFilter = buildDateFilter(start, end);
       const sql = `
         SELECT
           strftime('%Y-%m', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
@@ -190,18 +219,16 @@ export class TokenUsageRepo {
           COUNT(*) as requests
         FROM token_usage
         WHERE 1=1
-        ${buildAccountFilter(accountId)}
-        ${buildDateFilter(start, end)}
+        ${accountFilter.sql}
+        ${dateFilter.sql}
         GROUP BY bucket
         ORDER BY bucket
       `;
-      const rows = raw.prepare(sql).all() as unknown[];
+      const rows = raw.prepare(sql).all(...accountFilter.params, ...dateFilter.params) as unknown[];
       return rows.map((row) => MonthlyUsageSchema.parse(row));
     } catch (error) {
       logger.error('Failed to get monthly usage', error);
       return [];
-    } finally {
-      raw.close();
     }
   }
 
@@ -216,8 +243,10 @@ export class TokenUsageRepo {
     totalTokens: number;
     requests: number;
   }> {
-    const { raw } = getCloudDb();
+    const { raw } = getUsageDb();
     try {
+      const accountFilter = buildAccountFilter(accountId);
+      const dateFilter = buildDateFilter(start, end);
       const sql = `
         SELECT
           model,
@@ -227,18 +256,16 @@ export class TokenUsageRepo {
           COUNT(*) as requests
         FROM token_usage
         WHERE 1=1
-        ${buildAccountFilter(accountId)}
-        ${buildDateFilter(start, end)}
+        ${accountFilter.sql}
+        ${dateFilter.sql}
         GROUP BY model
         ORDER BY totalTokens DESC
       `;
-      const rows = raw.prepare(sql).all() as unknown[];
+      const rows = raw.prepare(sql).all(...accountFilter.params, ...dateFilter.params) as unknown[];
       return rows.map((row) => ModelUsageSchema.parse(row));
     } catch (error) {
       logger.error('Failed to get model usage', error);
       return [];
-    } finally {
-      raw.close();
     }
   }
 }

src/localization/i18n.ts

@@ -497,9 +497,9 @@ i18n
           usage: {
             title: 'Token Usage',
             description: 'Monitor your API token consumption over time.',
-            range24h: '24 Hours',
-            range7d: '7 Days',
-            range30d: '30 Days',
+            range24h: 'Last 24 Hours',
+            range7d: 'Last 7 Days',
+            range30d: 'Last 30 Days',
             totalTokens: 'Total Tokens',
             promptTokens: 'Prompt Tokens',
             completionTokens: 'Completion Tokens',
@@ -507,6 +507,9 @@ i18n
             dailyUsage: 'Daily Usage',
             byModel: 'Usage by Model',
             hourlyUsage: 'Hourly Usage',
+            loadFailed: 'Failed to load usage data.',
+            retry: 'Retry',
+            allAccounts: 'All Accounts',
           },
         },
       },
@@ -987,9 +990,9 @@ i18n
           usage: {
             title: 'Token 用量',
             description: '监控您的 API Token 消耗情况。',
-            range24h: '24 小时',
-            range7d: '7 天',
-            range30d: '30 天',
+            range24h: '最近 24 小时',
+            range7d: '最近 7 天',
+            range30d: '最近 30 天',
             totalTokens: '总 Token 数',
             promptTokens: 'Prompt Token 数',
             completionTokens: 'Completion Token 数',
@@ -997,6 +1000,9 @@ i18n
             dailyUsage: '每日用量',
             byModel: '按模型统计',
             hourlyUsage: '每小时用量',
+            loadFailed: '加载用量数据失败。',
+            retry: '重试',
+            allAccounts: '所有账号',
           },
         },
       },
@@ -1483,9 +1489,9 @@ i18n
           usage: {
             title: 'Использование токенов',
             description: 'Мониторинг потребления API токенов.',
-            range24h: '24 часа',
-            range7d: '7 дней',
-            range30d: '30 дней',
+            range24h: 'Последние 24 часа',
+            range7d: 'Последние 7 дней',
+            range30d: 'Последние 30 дней',
             totalTokens: 'Всего токенов',
             promptTokens: 'Prompt токены',
             completionTokens: 'Completion токены',
@@ -1493,6 +1499,9 @@ i18n
             dailyUsage: 'Ежедневное использование',
             byModel: 'По моделям',
             hourlyUsage: 'Почасовое использование',
+            loadFailed: 'Не удалось загрузить данные об использовании.',
+            retry: 'Повторить',
+            allAccounts: 'Все аккаунты',
           },
         },
       },

src/localization/vi.ts

@@ -476,9 +476,9 @@ const vi = {
   usage: {
     title: 'Token Usage',
     description: 'Theo dõi mức tiêu thụ API token của bạn.',
-    range24h: '24 giờ',
-    range7d: '7 ngày',
-    range30d: '30 ngày',
+    range24h: '24 giờ qua',
+    range7d: '7 ngày qua',
+    range30d: '30 ngày qua',
     totalTokens: 'Tổng Token',
     promptTokens: 'Prompt Token',
     completionTokens: 'Completion Token',
@@ -486,6 +486,9 @@ const vi = {
     dailyUsage: 'Sử dụng hàng ngày',
     byModel: 'Theo Model',
     hourlyUsage: 'Sử dụng theo giờ',
+    loadFailed: 'Không thể tải dữ liệu sử dụng.',
+    retry: 'Thử lại',
+    allAccounts: 'Tất cả tài khoản',
   },
 };
 

src/routes/usage.tsx

@@ -19,12 +19,15 @@ import {
   Legend,
 } from 'recharts';
 import { useMemo, useState } from 'react';
+import { useCloudAccounts } from '@/hooks/useCloudAccounts';
 
 const COLORS = ['#3b82f6', '#10b981', '#f59e0b', '#ef4444', '#8b5cf6', '#ec4899', '#06b6d4'];
 
 function UsagePage() {
   const { t } = useTranslation();
   const [range, setRange] = useState<'24h' | '7d' | '30d'>('7d');
+  const [selectedAccountId, setSelectedAccountId] = useState<string | undefined>(undefined);
+  const { data: accounts } = useCloudAccounts();
 
   function getTimeRange(r: '24h' | '7d' | '30d') {
     const now = Date.now();
@@ -37,31 +40,50 @@ function UsagePage() {
     return { start: now - 30 * 24 * 60 * 60 * 1000, end: now };
   }
 
-  const { data: dailyData, isLoading: dailyLoading } = useQuery({
-    queryKey: ['usage', 'day', range],
+  const {
+    data: dailyData,
+    isLoading: dailyLoading,
+    error: dailyError,
+    refetch: refetchDaily,
+  } = useQuery({
+    queryKey: ['usage', 'day', range, selectedAccountId],
     queryFn: () => {
       const { start, end } = getTimeRange(range);
-      return ipc.client.usage.getUsageByDay({ start, end });
+      return ipc.client.usage.getUsageByDay({ accountId: selectedAccountId, start, end });
     },
+    refetchInterval: 30000,
   });
 
-  const { data: modelData, isLoading: modelLoading } = useQuery({
-    queryKey: ['usage', 'model', range],
+  const {
+    data: modelData,
+    isLoading: modelLoading,
+    error: modelError,
+    refetch: refetchModel,
+  } = useQuery({
+    queryKey: ['usage', 'model', range, selectedAccountId],
     queryFn: () => {
       const { start, end } = getTimeRange(range);
-      return ipc.client.usage.getUsageByModel({ start, end });
+      return ipc.client.usage.getUsageByModel({ accountId: selectedAccountId, start, end });
     },
+    refetchInterval: 30000,
   });
 
-  const { data: hourlyData, isLoading: hourlyLoading } = useQuery({
-    queryKey: ['usage', 'hour', range],
+  const {
+    data: hourlyData,
+    isLoading: hourlyLoading,
+    error: hourlyError,
+    refetch: refetchHourly,
+  } = useQuery({
+    queryKey: ['usage', 'hour', range, selectedAccountId],
     queryFn: () => {
       const { start, end } = getTimeRange(range);
-      return ipc.client.usage.getUsageByHour({ start, end });
+      return ipc.client.usage.getUsageByHour({ accountId: selectedAccountId, start, end });
     },
+    refetchInterval: 30000,
   });
 
   const isLoading = dailyLoading || modelLoading || hourlyLoading;
+  const anyError = dailyError || modelError || hourlyError;
 
   const totals = useMemo(() => {
     const source = dailyData || [];
@@ -95,6 +117,8 @@ function UsagePage() {
   const hourlyChartData = useMemo(() => {
     return (hourlyData || []).map((d) => ({
       name: d.bucket,
+      Prompt: d.promptTokens,
+      Completion: d.completionTokens,
       Total: d.totalTokens,
     }));
   }, [hourlyData]);
@@ -107,6 +131,24 @@ function UsagePage() {
     );
   }
 
+  if (anyError) {
+    return (
+      <div className="flex h-full flex-col items-center justify-center gap-4">
+        <p className="text-destructive">{t('usage.loadFailed')}</p>
+        <button
+          onClick={() => {
+            refetchDaily();
+            refetchModel();
+            refetchHourly();
+          }}
+          className="text-primary underline"
+        >
+          {t('usage.retry')}
+        </button>
+      </div>
+    );
+  }
+
   return (
     <div className="container mx-auto max-w-6xl space-y-6 p-6">
       <div className="flex items-center justify-between">
@@ -114,13 +156,27 @@ function UsagePage() {
           <h2 className="text-3xl font-bold tracking-tight">{t('usage.title')}</h2>
           <p className="text-muted-foreground mt-1">{t('usage.description')}</p>
         </div>
-        <Tabs value={range} onValueChange={(v) => setRange(v as '24h' | '7d' | '30d')}>
-          <TabsList>
-            <TabsTrigger value="24h">{t('usage.range24h')}</TabsTrigger>
-            <TabsTrigger value="7d">{t('usage.range7d')}</TabsTrigger>
-            <TabsTrigger value="30d">{t('usage.range30d')}</TabsTrigger>
-          </TabsList>
-        </Tabs>
+        <div className="flex items-center gap-4">
+          <select
+            value={selectedAccountId || ''}
+            onChange={(e) => setSelectedAccountId(e.target.value || undefined)}
+            className="border rounded px-2 py-1"
+          >
+            <option value="">{t('usage.allAccounts')}</option>
+            {accounts?.map((acc) => (
+              <option key={acc.id} value={acc.id}>
+                {acc.email}
+              </option>
+            ))}
+          </select>
+          <Tabs value={range} onValueChange={(v) => setRange(v as '24h' | '7d' | '30d')}>
+            <TabsList>
+              <TabsTrigger value="24h">{t('usage.range24h')}</TabsTrigger>
+              <TabsTrigger value="7d">{t('usage.range7d')}</TabsTrigger>
+              <TabsTrigger value="30d">{t('usage.range30d')}</TabsTrigger>
+            </TabsList>
+          </Tabs>
+        </div>
       </div>
 
       {/* Summary Cards */}
@@ -228,6 +284,9 @@ function UsagePage() {
                 <XAxis dataKey="name" />
                 <YAxis />
                 <Tooltip />
+                <Legend />
+                <Bar dataKey="Prompt" fill="#3b82f6" />
+                <Bar dataKey="Completion" fill="#10b981" />
                 <Bar dataKey="Total" fill="#8b5cf6" />
               </BarChart>
             </ResponsiveContainer>

src/server/modules/proxy/proxy.service.ts

@@ -35,6 +35,7 @@ import { getMaxOutputTokens, getThinkingBudget } from '../../../lib/geminiNexus/
 import { resolveRequestUserAgent } from './request-user-agent';
 import { UpstreamRequestError } from './clients/upstream-error';
 import { trafficLogger } from '../../../utils/traffic-logger';
+import { encode } from 'gpt-tokenizer';
 
 @Injectable()
 export class ProxyService {
@@ -46,16 +47,16 @@ export class ProxyService {
     @Inject(TokenUsageService) private readonly tokenUsageService: TokenUsageService,
   ) {}
 
-  private recordUsage(
+  private async recordUsage(
     accountId: string,
     model: string,
     usage: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number } | undefined | null,
     requestType: string,
-  ): void {
+  ): Promise<void> {
     if (!usage) {
       return;
     }
-    this.tokenUsageService.recordUsage({
+    await this.tokenUsageService.recordUsage({
       accountId,
       model,
       promptTokens: usage.promptTokenCount ?? 0,
@@ -66,6 +67,14 @@ export class ProxyService {
     });
   }
 
+  private estimateTokens(text: string): number {
+    try {
+      return encode(text).length;
+    } catch {
+      return Math.ceil(text.length / 4);
+    }
+  }
+
   // --- Anthropic Handlers ---
 
   async handleAnthropicMessages(
@@ -131,7 +140,7 @@ export class ProxyService {
             token.token.upstream_proxy_url,
             extraHeaders,
           );
-          this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
+          await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
           return this.toAnthropicChatResponse(transformResponse(response));
         }
       } catch (error) {
@@ -163,7 +172,7 @@ export class ProxyService {
                 token.token.upstream_proxy_url,
                 extraHeaders,
               );
-              this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
+              await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
               return this.toAnthropicChatResponse(transformResponse(response));
             }
           } catch (fallbackErr) {
@@ -202,7 +211,7 @@ export class ProxyService {
                 token.token.upstream_proxy_url,
                 extraHeaders,
               );
-              this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
+              await this.recordUsage(token.id, 'gemini-3-flash', response.usageMetadata, 'anthropic');
               const transformed = this.toAnthropicChatResponse(transformResponse(response));
               return {
                 ...transformed,
@@ -232,6 +241,16 @@ export class ProxyService {
     _model: string,
     accountId: string,
   ): Observable<string> {
+    function extractUsageFromMetadata(meta: unknown): { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number } | undefined {
+      if (!meta || typeof meta !== 'object') return undefined;
+      const m = meta as Record<string, unknown>;
+      return {
+        promptTokenCount: typeof m.promptTokenCount === 'number' ? m.promptTokenCount : undefined,
+        candidatesTokenCount: typeof m.candidatesTokenCount === 'number' ? m.candidatesTokenCount : undefined,
+        totalTokenCount: typeof m.totalTokenCount === 'number' ? m.totalTokenCount : undefined,
+      };
+    }
+
     return new Observable<string>((subscriber) => {
       const decoder = new TextDecoder();
       let buffer = '';
@@ -298,12 +317,10 @@ export class ProxyService {
 
         const finishChunks = state.emitFinish(lastFinishReason, lastUsageMetadata);
         finishChunks.forEach((c) => subscriber.next(c));
-        this.recordUsage(
-          accountId,
-          _model,
-          lastUsageMetadata as { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number },
-          'anthropic',
-        );
+        const usage = extractUsageFromMetadata(lastUsageMetadata);
+        if (usage) {
+          void this.recordUsage(accountId, _model, usage, 'anthropic');
+        }
         subscriber.complete();
       });
 
@@ -371,7 +388,7 @@ export class ProxyService {
           extraHeaders,
         );
 
-        this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
+        await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
         return this.normalizeGeminiGenerateResponse(response);
       } catch (err) {
         if (err instanceof Error && this.isProjectContextError(err.message)) {
@@ -394,7 +411,7 @@ export class ProxyService {
               token.token.upstream_proxy_url,
               extraHeaders,
             );
-            this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
+            await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
             return this.normalizeGeminiGenerateResponse(response);
           } catch (fallbackErr) {
             lastError = fallbackErr;
@@ -468,7 +485,7 @@ export class ProxyService {
           token.token.upstream_proxy_url,
           extraHeaders,
         );
-        return this.passthroughSseStream(stream);
+            return this.passthroughSseStream(stream, token.id, effectiveTargetModel);
       } catch (err) {
         if (err instanceof Error && this.isProjectContextError(err.message)) {
           this.logger.warn(
@@ -490,7 +507,7 @@ export class ProxyService {
               token.token.upstream_proxy_url,
               extraHeaders,
             );
-            return this.passthroughSseStream(stream);
+        return this.passthroughSseStream(stream, token.id, effectiveTargetModel);
           } catch (fallbackErr) {
             lastError = fallbackErr;
           }
@@ -511,14 +528,36 @@ export class ProxyService {
     throw lastError || new Error('Gemini stream request failed after retries');
   }
 
-  private passthroughSseStream(upstreamStream: NodeJS.ReadableStream): Observable<string> {
+  private passthroughSseStream(
+    upstreamStream: NodeJS.ReadableStream,
+    accountId: string,
+    model: string,
+  ): Observable<string> {
     return new Observable<string>((subscriber) => {
       const decoder = new TextDecoder();
+      let buffer = '';
       let receivedData = false;
+      let lastUsageMetadata: Record<string, unknown> | undefined;
 
       upstreamStream.on('data', (chunk: Buffer) => {
         receivedData = true;
-        subscriber.next(decoder.decode(chunk, { stream: true }));
+        buffer += decoder.decode(chunk, { stream: true });
+        const lines = buffer.split('\n');
+        buffer = lines.pop() || '';
+
+        for (const line of lines) {
+          const trimmed = line.trim();
+          if (!trimmed.startsWith('data: ')) continue;
+          const dataStr = trimmed.slice(6);
+          subscriber.next(`data: ${dataStr}\n\n`);
+          if (dataStr === '[DONE]') continue;
+          try {
+            const parsed = JSON.parse(dataStr);
+            if (parsed?.usageMetadata) {
+              lastUsageMetadata = parsed.usageMetadata;
+            }
+          } catch { /* ignore */ }
+        }
       });
 
       upstreamStream.on('end', () => {
@@ -526,6 +565,9 @@ export class ProxyService {
           subscriber.error(new Error('Empty response stream'));
           return;
         }
+        if (lastUsageMetadata) {
+          void this.recordUsage(accountId, model, lastUsageMetadata as any, 'gemini');
+        }
         subscriber.complete();
       });
 
@@ -787,7 +829,7 @@ export class ProxyService {
               token.token.upstream_proxy_url,
               extraHeaders,
             );
-            return this.processStreamResponse(stream, request.model, requestId);
+              return this.processStreamResponse(stream, request.model, requestId, token.id, effectiveTargetModel);
           } catch (streamError) {
             this.logger.warn(
               `Stream path failed for model=${request.model}; falling back to non-stream generation: ${
@@ -803,7 +845,7 @@ export class ProxyService {
               extraHeaders,
             );
             trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0);
-            this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
+            await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
             this.logger.log(
               `Upstream response snippet after stream fallback: ${JSON.stringify(response).substring(0, 500)}`,
             );
@@ -824,7 +866,7 @@ export class ProxyService {
           this.logger.log(
             `Upstream response snippet (non-stream): ${JSON.stringify(response).substring(0, 500)}`,
           );
-          this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
+          await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
           // Transform Gemini response to OpenAI format
           const claudeResponse = transformResponse(response);
           this.logger.log(
@@ -851,7 +893,7 @@ export class ProxyService {
                 token.token.upstream_proxy_url,
                 extraHeaders,
               );
-              return this.processStreamResponse(stream, request.model, requestId);
+            return this.processStreamResponse(stream, request.model, requestId, token.id, effectiveTargetModel);
             }
 
             trafficLogger.logUpstream(requestId || 'unknown', '/v1/chat/completions', fallbackBody, extraHeaders, { model: effectiveTargetModel });
@@ -862,7 +904,7 @@ export class ProxyService {
               extraHeaders,
             );
             trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0);
-            this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
+            await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
             const claudeResponse = transformResponse(response);
             return this.convertClaudeToOpenAIResponse(claudeResponse, request.model);
           } catch (fallbackErr) {
@@ -1001,7 +1043,9 @@ export class ProxyService {
   private processStreamResponse(
     upstreamStream: NodeJS.ReadableStream,
     model: string,
-    requestId?: string,
+    requestId: string | undefined,
+    accountId: string,
+    effectiveTargetModel: string,
   ): Observable<string> {
     return new Observable<string>((subscriber) => {
       const decoder = new TextDecoder();
@@ -1010,6 +1054,7 @@ export class ProxyService {
       let hasSentRoleChunk = false;
       let hasSentDone = false;
       let toolCallIndex = 0;
+      let lastUsageMetadata: Record<string, unknown> | undefined;
 
       const streamId = `chatcmpl-${uuidv4()}`;
       const created = Math.floor(Date.now() / 1000);
@@ -1055,6 +1100,9 @@ export class ProxyService {
             trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, json, 0, { rawLine: trimmed });
             // v1internal API wraps response in {"response": {"candidates": [...]}}
             const responsePayload = json.response ?? json;
+            if (responsePayload.usageMetadata) {
+              lastUsageMetadata = responsePayload.usageMetadata;
+            }
             const candidate = responsePayload.candidates?.[0];
             const parts = candidate?.content?.parts || [];
 
@@ -1178,7 +1226,6 @@ export class ProxyService {
               pushChunk(finishChunk);
               subscriber.next('data: [DONE]\n\n');
               hasSentDone = true;
-              subscriber.complete();
             }
           } catch {
             // ignore parse errors on individual SSE lines
@@ -1206,6 +1253,10 @@ export class ProxyService {
           subscriber.next('data: [DONE]\n\n');
           hasSentDone = true;
         }
+        if (lastUsageMetadata) {
+          // TODO: integrate local tokenizer fallback for missing usageMetadata
+          void this.recordUsage(accountId, effectiveTargetModel, lastUsageMetadata as any, 'openai');
+        }
         subscriber.complete();
       });
 

src/server/modules/usage/token-usage.service.ts

@@ -3,7 +3,7 @@ import { TokenUsageRepo, RecordUsageParams } from '../../../ipc/database/usageHa
 
 @Injectable()
 export class TokenUsageService {
-  recordUsage(params: RecordUsageParams): void {
+  async recordUsage(params: RecordUsageParams): Promise<void> {
     TokenUsageRepo.recordUsage(params);
   }
 

src/tests/unit/usage-handler.test.ts

@@ -0,0 +1,193 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { TokenUsageRepo } from '../../ipc/database/usageHandler';
+
+const mockRun = vi.fn();
+const mockAll = vi.fn();
+const mockPrepare = vi.fn(() => ({ all: mockAll, run: mockRun }));
+const mockClose = vi.fn();
+const mockInsert = vi.fn(() => ({
+  values: vi.fn(() => ({ run: mockRun })),
+}));
+
+vi.mock('../../ipc/database/dbConnection', () => ({
+  openDrizzleConnection: vi.fn(() => ({
+    raw: {
+      prepare: mockPrepare,
+      close: mockClose,
+    },
+    orm: {
+      insert: mockInsert,
+    },
+  })),
+}));
+
+vi.mock('../../ipc/database/cloudHandler', () => ({
+  ensureDatabaseInitialized: vi.fn(),
+}));
+
+vi.mock('../../utils/paths', () => ({
+  getCloudAccountsDbPath: vi.fn(() => '/mock/db/path.sqlite'),
+}));
+
+vi.mock('../../utils/logger', () => ({
+  logger: {
+    info: vi.fn(),
+    error: vi.fn(),
+    debug: vi.fn(),
+    warn: vi.fn(),
+  },
+}));
+
+describe('TokenUsageRepo', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('recordUsage', () => {
+    it('should insert usage record into database', () => {
+      const params = {
+        accountId: 'acc-1',
+        model: 'gemini-pro',
+        promptTokens: 10,
+        completionTokens: 20,
+        totalTokens: 30,
+        timestamp: Date.now(),
+        requestType: 'chat',
+      };
+
+      TokenUsageRepo.recordUsage(params);
+
+      expect(mockInsert).toHaveBeenCalled();
+      expect(mockRun).toHaveBeenCalled();
+    });
+
+    it('should handle null requestType', () => {
+      const params = {
+        accountId: 'acc-1',
+        model: 'gemini-pro',
+        promptTokens: 10,
+        completionTokens: 20,
+        totalTokens: 30,
+        timestamp: Date.now(),
+      };
+
+      TokenUsageRepo.recordUsage(params);
+
+      expect(mockInsert).toHaveBeenCalled();
+      expect(mockRun).toHaveBeenCalled();
+    });
+  });
+
+  describe('getUsageByHour', () => {
+    it('should return hourly aggregated data', () => {
+      mockAll.mockReturnValue([
+        {
+          bucket: '2024-01-01 10:00',
+          promptTokens: 100,
+          completionTokens: 200,
+          totalTokens: 300,
+          requests: 5,
+        },
+      ]);
+
+      const result = TokenUsageRepo.getUsageByHour('acc-1', 1704110400000, 1704114000000);
+
+      expect(mockPrepare).toHaveBeenCalled();
+      expect(result).toHaveLength(1);
+      expect(result[0].bucket).toBe('2024-01-01 10:00');
+      expect(result[0].totalTokens).toBe(300);
+    });
+
+    it('should return empty array on error', () => {
+      mockPrepare.mockImplementationOnce(() => {
+        throw new Error('DB error');
+      });
+
+      const result = TokenUsageRepo.getUsageByHour();
+
+      expect(result).toEqual([]);
+    });
+  });
+
+  describe('getUsageByDay', () => {
+    it('should return daily aggregated data', () => {
+      mockAll.mockReturnValue([
+        {
+          bucket: '2024-01-01',
+          promptTokens: 1000,
+          completionTokens: 2000,
+          totalTokens: 3000,
+          requests: 50,
+        },
+      ]);
+
+      const result = TokenUsageRepo.getUsageByDay();
+
+      expect(mockPrepare).toHaveBeenCalled();
+      expect(result).toHaveLength(1);
+      expect(result[0].bucket).toBe('2024-01-01');
+    });
+  });
+
+  describe('getUsageByWeek', () => {
+    it('should return weekly aggregated data', () => {
+      mockAll.mockReturnValue([
+        {
+          bucket: '2024-01',
+          promptTokens: 5000,
+          completionTokens: 10000,
+          totalTokens: 15000,
+          requests: 200,
+        },
+      ]);
+
+      const result = TokenUsageRepo.getUsageByWeek();
+
+      expect(result).toHaveLength(1);
+    });
+  });
+
+  describe('getUsageByMonth', () => {
+    it('should return monthly aggregated data', () => {
+      mockAll.mockReturnValue([
+        {
+          bucket: '2024-01',
+          promptTokens: 20000,
+          completionTokens: 40000,
+          totalTokens: 60000,
+          requests: 1000,
+        },
+      ]);
+
+      const result = TokenUsageRepo.getUsageByMonth();
+
+      expect(result).toHaveLength(1);
+    });
+  });
+
+  describe('getUsageByModel', () => {
+    it('should return model aggregated data', () => {
+      mockAll.mockReturnValue([
+        {
+          model: 'gemini-pro',
+          promptTokens: 10000,
+          completionTokens: 20000,
+          totalTokens: 30000,
+          requests: 500,
+        },
+        {
+          model: 'claude-sonnet',
+          promptTokens: 5000,
+          completionTokens: 10000,
+          totalTokens: 15000,
+          requests: 250,
+        },
+      ]);
+
+      const result = TokenUsageRepo.getUsageByModel();
+
+      expect(result).toHaveLength(2);
+      expect(result[0].model).toBe('gemini-pro');
+    });
+  });
+});

src/tests/unit/usage-router.test.ts

@@ -0,0 +1,87 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { TokenUsageRepo } from '../../ipc/database/usageHandler';
+import { usageRouter } from '../../ipc/usage/router';
+
+describe('usageRouter', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('should expose all usage endpoints', () => {
+    expect(usageRouter).toHaveProperty('recordUsage');
+    expect(usageRouter).toHaveProperty('getUsageByHour');
+    expect(usageRouter).toHaveProperty('getUsageByDay');
+    expect(usageRouter).toHaveProperty('getUsageByWeek');
+    expect(usageRouter).toHaveProperty('getUsageByMonth');
+    expect(usageRouter).toHaveProperty('getUsageByModel');
+  });
+
+  it('recordUsage should delegate to TokenUsageRepo.recordUsage', async () => {
+    const spy = vi.spyOn(TokenUsageRepo, 'recordUsage').mockImplementation(() => true);
+    const input = {
+      accountId: 'acc-1',
+      model: 'gemini-pro',
+      promptTokens: 10,
+      completionTokens: 20,
+      totalTokens: 30,
+      timestamp: Date.now(),
+      requestType: 'chat',
+    };
+
+    const handler = (usageRouter.recordUsage as any)['~orpc'].handler;
+    await handler({ input });
+
+    expect(spy).toHaveBeenCalledWith(input);
+    spy.mockRestore();
+  });
+
+  it('getUsageByHour should delegate to TokenUsageRepo.getUsageByHour', async () => {
+    const spy = vi.spyOn(TokenUsageRepo, 'getUsageByHour').mockReturnValue([]);
+    const input = { accountId: 'acc-1', start: 1000, end: 2000 };
+    const handler = (usageRouter.getUsageByHour as any)['~orpc'].handler;
+    await handler({ input });
+
+    expect(spy).toHaveBeenCalledWith('acc-1', 1000, 2000);
+    spy.mockRestore();
+  });
+
+  it('getUsageByDay should delegate to TokenUsageRepo.getUsageByDay', async () => {
+    const spy = vi.spyOn(TokenUsageRepo, 'getUsageByDay').mockReturnValue([]);
+    const input = { start: 1000, end: 2000 };
+    const handler = (usageRouter.getUsageByDay as any)['~orpc'].handler;
+    await handler({ input });
+
+    expect(spy).toHaveBeenCalledWith(undefined, 1000, 2000);
+    spy.mockRestore();
+  });
+
+  it('getUsageByWeek should delegate to TokenUsageRepo.getUsageByWeek', async () => {
+    const spy = vi.spyOn(TokenUsageRepo, 'getUsageByWeek').mockReturnValue([]);
+    const input = {};
+    const handler = (usageRouter.getUsageByWeek as any)['~orpc'].handler;
+    await handler({ input });
+
+    expect(spy).toHaveBeenCalledWith(undefined, undefined, undefined);
+    spy.mockRestore();
+  });
+
+  it('getUsageByMonth should delegate to TokenUsageRepo.getUsageByMonth', async () => {
+    const spy = vi.spyOn(TokenUsageRepo, 'getUsageByMonth').mockReturnValue([]);
+    const input = { accountId: 'acc-2' };
+    const handler = (usageRouter.getUsageByMonth as any)['~orpc'].handler;
+    await handler({ input });
+
+    expect(spy).toHaveBeenCalledWith('acc-2', undefined, undefined);
+    spy.mockRestore();
+  });
+
+  it('getUsageByModel should delegate to TokenUsageRepo.getUsageByModel', async () => {
+    const spy = vi.spyOn(TokenUsageRepo, 'getUsageByModel').mockReturnValue([]);
+    const input = { accountId: 'acc-1', start: 1, end: 2 };
+    const handler = (usageRouter.getUsageByModel as any)['~orpc'].handler;
+    await handler({ input });
+
+    expect(spy).toHaveBeenCalledWith('acc-1', 1, 2);
+    spy.mockRestore();
+  });
+});