From 9b37bdb2d955e57083587de2222742f2cd040af8 Mon Sep 17 00:00:00 2001 From: Gemini Nexus Bot Date: Wed, 22 Apr 2026 14:34:39 -0300 Subject: [PATCH] fix(usage): address code review findings - SQL injection, stream usage recording, DB singleton --- package-lock.json | 7 + package.json | 1 + src/ipc/database/cloudHandler.ts | 3 +- src/ipc/database/usageHandler.ts | 117 +++++++---- src/localization/i18n.ts | 27 ++- src/localization/vi.ts | 9 +- src/routes/usage.tsx | 91 +++++++-- src/server/modules/proxy/proxy.service.ts | 101 ++++++--- .../modules/usage/token-usage.service.ts | 2 +- src/tests/unit/usage-handler.test.ts | 193 ++++++++++++++++++ src/tests/unit/usage-router.test.ts | 87 ++++++++ 11 files changed, 538 insertions(+), 100 deletions(-) create mode 100644 src/tests/unit/usage-handler.test.ts create mode 100644 src/tests/unit/usage-router.test.ts diff --git a/package-lock.json b/package-lock.json index 96dc473..cfde7e0 100644 --- a/package-lock.json +++ b/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", diff --git a/package.json b/package.json index 2420cc4..2dc7e6c 100644 --- a/package.json +++ b/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", diff --git a/src/ipc/database/cloudHandler.ts b/src/ipc/database/cloudHandler.ts index fce86d9..87a09e1 100644 --- a/src/ipc/database/cloudHandler.ts +++ b/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 diff --git a/src/ipc/database/usageHandler.ts b/src/ipc/database/usageHandler.ts index 7403603..6c90060 100644 --- a/src/ipc/database/usageHandler.ts +++ b/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 } | 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(); } } } diff --git a/src/localization/i18n.ts b/src/localization/i18n.ts index 7b8d184..3076ee3 100644 --- a/src/localization/i18n.ts +++ b/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: 'Все аккаунты', }, }, }, diff --git a/src/localization/vi.ts b/src/localization/vi.ts index 7a318bb..269fea4 100644 --- a/src/localization/vi.ts +++ b/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', }, }; diff --git a/src/routes/usage.tsx b/src/routes/usage.tsx index f1e3d12..65e8be5 100644 --- a/src/routes/usage.tsx +++ b/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(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 ( +
+

{t('usage.loadFailed')}

+ +
+ ); + } + return (
@@ -114,13 +156,27 @@ function UsagePage() {

{t('usage.title')}

{t('usage.description')}

- setRange(v as '24h' | '7d' | '30d')}> - - {t('usage.range24h')} - {t('usage.range7d')} - {t('usage.range30d')} - - +
+ + setRange(v as '24h' | '7d' | '30d')}> + + {t('usage.range24h')} + {t('usage.range7d')} + {t('usage.range30d')} + + +
{/* Summary Cards */} @@ -228,6 +284,9 @@ function UsagePage() { + + + diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index 368b46e..0a48fdc 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/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 { 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 { + function extractUsageFromMetadata(meta: unknown): { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number } | undefined { + if (!meta || typeof meta !== 'object') return undefined; + const m = meta as Record; + 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((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 { + private passthroughSseStream( + upstreamStream: NodeJS.ReadableStream, + accountId: string, + model: string, + ): Observable { return new Observable((subscriber) => { const decoder = new TextDecoder(); + let buffer = ''; let receivedData = false; + let lastUsageMetadata: Record | 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 { return new Observable((subscriber) => { const decoder = new TextDecoder(); @@ -1010,6 +1054,7 @@ export class ProxyService { let hasSentRoleChunk = false; let hasSentDone = false; let toolCallIndex = 0; + let lastUsageMetadata: Record | 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(); }); diff --git a/src/server/modules/usage/token-usage.service.ts b/src/server/modules/usage/token-usage.service.ts index 42c6637..27bb112 100644 --- a/src/server/modules/usage/token-usage.service.ts +++ b/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 { TokenUsageRepo.recordUsage(params); } diff --git a/src/tests/unit/usage-handler.test.ts b/src/tests/unit/usage-handler.test.ts new file mode 100644 index 0000000..1bc487a --- /dev/null +++ b/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'); + }); + }); +}); diff --git a/src/tests/unit/usage-router.test.ts b/src/tests/unit/usage-router.test.ts new file mode 100644 index 0000000..76cc56f --- /dev/null +++ b/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(); + }); +});