fix(usage): apply code review fixes - SSE passthrough type, per-chart errors, Select component

This commit is contained in:
Gemini Nexus Bot committed 2026-04-22 15:31:54 -03:00
1 parent 9b37bdb2d9
commit 425b0009b8
8 files changed
+303 -157

No files matched your search

+1
View File
@@ -37,4 +37,5 @@ export const tokenUsage = sqliteTable('token_usage', {
totalTokens: integer('total_tokens').notNull().default(0), totalTokens: integer('total_tokens').notNull().default(0),
timestamp: integer('timestamp').notNull(), timestamp: integer('timestamp').notNull(),
requestType: text('request_type'), requestType: text('request_type'),
isEstimated: integer('is_estimated'),
}); });
+2
View File
@@ -21,6 +21,7 @@ export interface RecordUsageParams {
totalTokens: number; totalTokens: number;
timestamp: number; timestamp: number;
requestType?: string; requestType?: string;
isEstimated?: boolean;
} }
let usageDbConnection: { raw: Database.Database; orm: BetterSQLite3Database<typeof import('./schema')> } | null = null; let usageDbConnection: { raw: Database.Database; orm: BetterSQLite3Database<typeof import('./schema')> } | null = null;
@@ -75,6 +76,7 @@ export class TokenUsageRepo {
totalTokens: params.totalTokens, totalTokens: params.totalTokens,
timestamp: params.timestamp, timestamp: params.timestamp,
requestType: params.requestType ?? null, requestType: params.requestType ?? null,
isEstimated: params.isEstimated ? 1 : null,
}) })
.run(); .run();
return true; return true;
+3
View File
@@ -510,6 +510,7 @@ i18n
loadFailed: 'Failed to load usage data.', loadFailed: 'Failed to load usage data.',
retry: 'Retry', retry: 'Retry',
allAccounts: 'All Accounts', allAccounts: 'All Accounts',
noData: 'No data available',
}, },
}, },
}, },
@@ -1003,6 +1004,7 @@ i18n
loadFailed: '加载用量数据失败。', loadFailed: '加载用量数据失败。',
retry: '重试', retry: '重试',
allAccounts: '所有账号', allAccounts: '所有账号',
noData: '暂无数据',
}, },
}, },
}, },
@@ -1502,6 +1504,7 @@ i18n
loadFailed: 'Не удалось загрузить данные об использовании.', loadFailed: 'Не удалось загрузить данные об использовании.',
retry: 'Повторить', retry: 'Повторить',
allAccounts: 'Все аккаунты', allAccounts: 'Все аккаунты',
noData: 'Нет данных',
}, },
}, },
}, },
+1
View File
@@ -489,6 +489,7 @@ const vi = {
loadFailed: 'Không thể tải dữ liệu sử dụng.', loadFailed: 'Không thể tải dữ liệu sử dụng.',
retry: 'Thử lại', retry: 'Thử lại',
allAccounts: 'Tất cả tài khoản', allAccounts: 'Tất cả tài khoản',
noData: 'Không có dữ liệu',
}, },
}; };
+203 -129
View File
@@ -2,6 +2,13 @@ import { createFileRoute } from '@tanstack/react-router';
import { useQuery } from '@tanstack/react-query'; import { useQuery } from '@tanstack/react-query';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'; import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { Tabs, TabsList, TabsTrigger } from '@/components/ui/tabs'; import { Tabs, TabsList, TabsTrigger } from '@/components/ui/tabs';
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from '@/components/ui/select';
import { useTranslation } from 'react-i18next'; import { useTranslation } from 'react-i18next';
import { Loader2, TrendingUp, TrendingDown, Activity, Layers } from 'lucide-react'; import { Loader2, TrendingUp, TrendingDown, Activity, Layers } from 'lucide-react';
import { ipc } from '@/ipc/manager'; import { ipc } from '@/ipc/manager';
@@ -27,7 +34,7 @@ function UsagePage() {
const { t } = useTranslation(); const { t } = useTranslation();
const [range, setRange] = useState<'24h' | '7d' | '30d'>('7d'); const [range, setRange] = useState<'24h' | '7d' | '30d'>('7d');
const [selectedAccountId, setSelectedAccountId] = useState<string | undefined>(undefined); const [selectedAccountId, setSelectedAccountId] = useState<string | undefined>(undefined);
const { data: accounts } = useCloudAccounts(); const { data: accounts, isLoading: accountsLoading } = useCloudAccounts();
function getTimeRange(r: '24h' | '7d' | '30d') { function getTimeRange(r: '24h' | '7d' | '30d') {
const now = Date.now(); const now = Date.now();
@@ -52,6 +59,8 @@ function UsagePage() {
return ipc.client.usage.getUsageByDay({ accountId: selectedAccountId, start, end }); return ipc.client.usage.getUsageByDay({ accountId: selectedAccountId, start, end });
}, },
refetchInterval: 30000, refetchInterval: 30000,
refetchIntervalInBackground: false,
staleTime: 60_000,
}); });
const { const {
@@ -66,6 +75,8 @@ function UsagePage() {
return ipc.client.usage.getUsageByModel({ accountId: selectedAccountId, start, end }); return ipc.client.usage.getUsageByModel({ accountId: selectedAccountId, start, end });
}, },
refetchInterval: 30000, refetchInterval: 30000,
refetchIntervalInBackground: false,
staleTime: 60_000,
}); });
const { const {
@@ -80,10 +91,11 @@ function UsagePage() {
return ipc.client.usage.getUsageByHour({ accountId: selectedAccountId, start, end }); return ipc.client.usage.getUsageByHour({ accountId: selectedAccountId, start, end });
}, },
refetchInterval: 30000, refetchInterval: 30000,
refetchIntervalInBackground: false,
staleTime: 60_000,
}); });
const isLoading = dailyLoading || modelLoading || hourlyLoading; const isAllLoading = dailyLoading && modelLoading && hourlyLoading;
const anyError = dailyError || modelError || hourlyError;
const totals = useMemo(() => { const totals = useMemo(() => {
const source = dailyData || []; const source = dailyData || [];
@@ -123,7 +135,7 @@ function UsagePage() {
})); }));
}, [hourlyData]); }, [hourlyData]);
if (isLoading) { if (isAllLoading) {
return ( return (
<div className="flex h-full items-center justify-center"> <div className="flex h-full items-center justify-center">
<Loader2 className="animate-spin" /> <Loader2 className="animate-spin" />
@@ -131,23 +143,65 @@ function UsagePage() {
); );
} }
if (anyError) { const dailyPlaceholder = dailyError ? (
return ( <Card>
<div className="flex h-full flex-col items-center justify-center gap-4"> <CardContent className="flex flex-col items-center justify-center gap-2 py-12">
<p className="text-destructive">{t('usage.loadFailed')}</p> <p className="text-sm text-destructive">{t('usage.loadFailed')}</p>
<button <button
onClick={() => { onClick={() => refetchDaily()}
refetchDaily(); className="text-primary text-sm underline"
refetchModel();
refetchHourly();
}}
className="text-primary underline"
> >
{t('usage.retry')} {t('usage.retry')}
</button> </button>
</div> </CardContent>
); </Card>
} ) : !dailyData ? (
<Card>
<CardContent className="flex items-center justify-center py-12">
<p className="text-muted-foreground text-sm">{t('usage.noData')}</p>
</CardContent>
</Card>
) : null;
const modelPlaceholder = modelError ? (
<Card>
<CardContent className="flex flex-col items-center justify-center gap-2 py-12">
<p className="text-sm text-destructive">{t('usage.loadFailed')}</p>
<button
onClick={() => refetchModel()}
className="text-primary text-sm underline"
>
{t('usage.retry')}
</button>
</CardContent>
</Card>
) : !modelData ? (
<Card>
<CardContent className="flex items-center justify-center py-12">
<p className="text-muted-foreground text-sm">{t('usage.noData')}</p>
</CardContent>
</Card>
) : null;
const hourlyPlaceholder = hourlyError ? (
<Card>
<CardContent className="flex flex-col items-center justify-center gap-2 py-12">
<p className="text-sm text-destructive">{t('usage.loadFailed')}</p>
<button
onClick={() => refetchHourly()}
className="text-primary text-sm underline"
>
{t('usage.retry')}
</button>
</CardContent>
</Card>
) : !hourlyData ? (
<Card>
<CardContent className="flex items-center justify-center py-12">
<p className="text-muted-foreground text-sm">{t('usage.noData')}</p>
</CardContent>
</Card>
) : null;
return ( return (
<div className="container mx-auto max-w-6xl space-y-6 p-6"> <div className="container mx-auto max-w-6xl space-y-6 p-6">
@@ -157,18 +211,30 @@ function UsagePage() {
<p className="text-muted-foreground mt-1">{t('usage.description')}</p> <p className="text-muted-foreground mt-1">{t('usage.description')}</p>
</div> </div>
<div className="flex items-center gap-4"> <div className="flex items-center gap-4">
<select <Select
value={selectedAccountId || ''} value={selectedAccountId || ''}
onChange={(e) => setSelectedAccountId(e.target.value || undefined)} onValueChange={(value) => setSelectedAccountId(value || undefined)}
className="border rounded px-2 py-1" disabled={accountsLoading}
> >
<option value="">{t('usage.allAccounts')}</option> <SelectTrigger className="w-[200px]">
{accounts?.map((acc) => ( {accountsLoading ? (
<option key={acc.id} value={acc.id}> <span className="flex items-center gap-2">
{acc.email} <Loader2 className="h-4 w-4 animate-spin" />
</option> {t('common.loading')}
))} </span>
</select> ) : (
<SelectValue placeholder={t('usage.allAccounts')} />
)}
</SelectTrigger>
<SelectContent>
<SelectItem value="">{t('usage.allAccounts')}</SelectItem>
{accounts?.map((acc) => (
<SelectItem key={acc.id} value={acc.id}>
{acc.email}
</SelectItem>
))}
</SelectContent>
</Select>
<Tabs value={range} onValueChange={(v) => setRange(v as '24h' | '7d' | '30d')}> <Tabs value={range} onValueChange={(v) => setRange(v as '24h' | '7d' | '30d')}>
<TabsList> <TabsList>
<TabsTrigger value="24h">{t('usage.range24h')}</TabsTrigger> <TabsTrigger value="24h">{t('usage.range24h')}</TabsTrigger>
@@ -180,118 +246,126 @@ function UsagePage() {
</div> </div>
{/* Summary Cards */} {/* Summary Cards */}
<div className="grid grid-cols-1 gap-4 sm:grid-cols-2 lg:grid-cols-4"> {dailyPlaceholder || (
<Card> <div className="grid grid-cols-1 gap-4 sm:grid-cols-2 lg:grid-cols-4">
<CardHeader className="flex flex-row items-center justify-between pb-2"> <Card>
<CardTitle className="text-sm font-medium">{t('usage.totalTokens')}</CardTitle> <CardHeader className="flex flex-row items-center justify-between pb-2">
<Activity className="text-muted-foreground h-4 w-4" /> <CardTitle className="text-sm font-medium">{t('usage.totalTokens')}</CardTitle>
</CardHeader> <Activity className="text-muted-foreground h-4 w-4" />
<CardContent> </CardHeader>
<div className="text-2xl font-bold">{totals.totalTokens.toLocaleString()}</div> <CardContent>
</CardContent> <div className="text-2xl font-bold">{totals.totalTokens.toLocaleString()}</div>
</Card> </CardContent>
</Card>
<Card> <Card>
<CardHeader className="flex flex-row items-center justify-between pb-2"> <CardHeader className="flex flex-row items-center justify-between pb-2">
<CardTitle className="text-sm font-medium">{t('usage.promptTokens')}</CardTitle> <CardTitle className="text-sm font-medium">{t('usage.promptTokens')}</CardTitle>
<TrendingUp className="text-muted-foreground h-4 w-4" /> <TrendingUp className="text-muted-foreground h-4 w-4" />
</CardHeader> </CardHeader>
<CardContent> <CardContent>
<div className="text-2xl font-bold">{totals.promptTokens.toLocaleString()}</div> <div className="text-2xl font-bold">{totals.promptTokens.toLocaleString()}</div>
</CardContent> </CardContent>
</Card> </Card>
<Card> <Card>
<CardHeader className="flex flex-row items-center justify-between pb-2"> <CardHeader className="flex flex-row items-center justify-between pb-2">
<CardTitle className="text-sm font-medium">{t('usage.completionTokens')}</CardTitle> <CardTitle className="text-sm font-medium">{t('usage.completionTokens')}</CardTitle>
<TrendingDown className="text-muted-foreground h-4 w-4" /> <TrendingDown className="text-muted-foreground h-4 w-4" />
</CardHeader> </CardHeader>
<CardContent> <CardContent>
<div className="text-2xl font-bold">{totals.completionTokens.toLocaleString()}</div> <div className="text-2xl font-bold">{totals.completionTokens.toLocaleString()}</div>
</CardContent> </CardContent>
</Card> </Card>
<Card> <Card>
<CardHeader className="flex flex-row items-center justify-between pb-2"> <CardHeader className="flex flex-row items-center justify-between pb-2">
<CardTitle className="text-sm font-medium">{t('usage.totalRequests')}</CardTitle> <CardTitle className="text-sm font-medium">{t('usage.totalRequests')}</CardTitle>
<Layers className="text-muted-foreground h-4 w-4" /> <Layers className="text-muted-foreground h-4 w-4" />
</CardHeader> </CardHeader>
<CardContent> <CardContent>
<div className="text-2xl font-bold">{totals.requests.toLocaleString()}</div> <div className="text-2xl font-bold">{totals.requests.toLocaleString()}</div>
</CardContent> </CardContent>
</Card> </Card>
</div> </div>
)}
{/* Charts */} {/* Charts */}
<div className="grid grid-cols-1 gap-6 lg:grid-cols-2"> <div className="grid grid-cols-1 gap-6 lg:grid-cols-2">
<Card> {dailyPlaceholder || (
<CardHeader> <Card>
<CardTitle>{t('usage.dailyUsage')}</CardTitle> <CardHeader>
</CardHeader> <CardTitle>{t('usage.dailyUsage')}</CardTitle>
<CardContent> </CardHeader>
<ResponsiveContainer width="100%" height={300}> <CardContent>
<BarChart data={chartData}> <ResponsiveContainer width="100%" height={300}>
<CartesianGrid strokeDasharray="3 3" /> <BarChart data={chartData}>
<XAxis dataKey="name" /> <CartesianGrid strokeDasharray="3 3" />
<YAxis /> <XAxis dataKey="name" />
<Tooltip /> <YAxis />
<Legend /> <Tooltip />
<Bar dataKey="Prompt" fill="#3b82f6" /> <Legend />
<Bar dataKey="Completion" fill="#10b981" /> <Bar dataKey="Prompt" fill="#3b82f6" />
</BarChart> <Bar dataKey="Completion" fill="#10b981" />
</ResponsiveContainer> </BarChart>
</CardContent> </ResponsiveContainer>
</Card> </CardContent>
</Card>
)}
<Card> {modelPlaceholder || (
<CardHeader> <Card>
<CardTitle>{t('usage.byModel')}</CardTitle> <CardHeader>
</CardHeader> <CardTitle>{t('usage.byModel')}</CardTitle>
<CardContent> </CardHeader>
<ResponsiveContainer width="100%" height={300}> <CardContent>
<PieChart> <ResponsiveContainer width="100%" height={300}>
<Pie <PieChart>
data={pieData} <Pie
cx="50%" data={pieData}
cy="50%" cx="50%"
innerRadius={60} cy="50%"
outerRadius={100} innerRadius={60}
paddingAngle={5} outerRadius={100}
dataKey="value" paddingAngle={5}
label dataKey="value"
> label
{pieData.map((_, index) => ( >
<Cell key={`cell-${index}`} fill={COLORS[index % COLORS.length]} /> {pieData.map((_, index) => (
))} <Cell key={`cell-${index}`} fill={COLORS[index % COLORS.length]} />
</Pie> ))}
<Tooltip /> </Pie>
<Legend /> <Tooltip />
</PieChart> <Legend />
</ResponsiveContainer> </PieChart>
</CardContent> </ResponsiveContainer>
</Card> </CardContent>
</Card>
)}
</div> </div>
{range === '24h' && ( {range === '24h' && (
<Card> hourlyPlaceholder || (
<CardHeader> <Card>
<CardTitle>{t('usage.hourlyUsage')}</CardTitle> <CardHeader>
</CardHeader> <CardTitle>{t('usage.hourlyUsage')}</CardTitle>
<CardContent> </CardHeader>
<ResponsiveContainer width="100%" height={300}> <CardContent>
<BarChart data={hourlyChartData}> <ResponsiveContainer width="100%" height={300}>
<CartesianGrid strokeDasharray="3 3" /> <BarChart data={hourlyChartData}>
<XAxis dataKey="name" /> <CartesianGrid strokeDasharray="3 3" />
<YAxis /> <XAxis dataKey="name" />
<Tooltip /> <YAxis />
<Legend /> <Tooltip />
<Bar dataKey="Prompt" fill="#3b82f6" /> <Legend />
<Bar dataKey="Completion" fill="#10b981" /> <Bar dataKey="Prompt" fill="#3b82f6" />
<Bar dataKey="Total" fill="#8b5cf6" /> <Bar dataKey="Completion" fill="#10b981" />
</BarChart> <Bar dataKey="Total" fill="#8b5cf6" />
</ResponsiveContainer> </BarChart>
</CardContent> </ResponsiveContainer>
</Card> </CardContent>
</Card>
)
)} )}
</div> </div>
); );
+77 -26
View File
@@ -52,6 +52,7 @@ export class ProxyService {
model: string, model: string,
usage: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number } | undefined | null, usage: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number } | undefined | null,
requestType: string, requestType: string,
isEstimated?: boolean,
): Promise<void> { ): Promise<void> {
if (!usage) { if (!usage) {
return; return;
@@ -64,6 +65,7 @@ export class ProxyService {
totalTokens: usage.totalTokenCount ?? 0, totalTokens: usage.totalTokenCount ?? 0,
timestamp: Date.now(), timestamp: Date.now(),
requestType, requestType,
isEstimated,
}); });
} }
@@ -75,6 +77,18 @@ export class ProxyService {
} }
} }
private 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,
};
}
// --- Anthropic Handlers --- // --- Anthropic Handlers ---
async handleAnthropicMessages( async handleAnthropicMessages(
@@ -140,7 +154,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
return this.toAnthropicChatResponse(transformResponse(response)); return this.toAnthropicChatResponse(transformResponse(response));
} }
} catch (error) { } catch (error) {
@@ -172,7 +186,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
return this.toAnthropicChatResponse(transformResponse(response)); return this.toAnthropicChatResponse(transformResponse(response));
} }
} catch (fallbackErr) { } catch (fallbackErr) {
@@ -211,7 +225,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
await this.recordUsage(token.id, 'gemini-3-flash', response.usageMetadata, 'anthropic'); this.recordUsage(token.id, 'gemini-3-flash', response.usageMetadata, 'anthropic');
const transformed = this.toAnthropicChatResponse(transformResponse(response)); const transformed = this.toAnthropicChatResponse(transformResponse(response));
return { return {
...transformed, ...transformed,
@@ -241,16 +255,6 @@ export class ProxyService {
_model: string, _model: string,
accountId: string, accountId: string,
): Observable<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) => { return new Observable<string>((subscriber) => {
const decoder = new TextDecoder(); const decoder = new TextDecoder();
let buffer = ''; let buffer = '';
@@ -260,6 +264,7 @@ export class ProxyService {
let lastFinishReason: string | undefined; let lastFinishReason: string | undefined;
let lastUsageMetadata: Record<string, unknown> | undefined; let lastUsageMetadata: Record<string, unknown> | undefined;
let accumulatedText = '';
let receivedData = false; let receivedData = false;
@@ -296,6 +301,9 @@ export class ProxyService {
if (this.isGeminiPart(part)) { if (this.isGeminiPart(part)) {
const chunks = processor.process(part); const chunks = processor.process(part);
chunks.forEach((c) => subscriber.next(c)); chunks.forEach((c) => subscriber.next(c));
if (part.text) {
accumulatedText += part.text;
}
} }
// Reset error state on successful parse // Reset error state on successful parse
@@ -317,9 +325,20 @@ export class ProxyService {
const finishChunks = state.emitFinish(lastFinishReason, lastUsageMetadata); const finishChunks = state.emitFinish(lastFinishReason, lastUsageMetadata);
finishChunks.forEach((c) => subscriber.next(c)); finishChunks.forEach((c) => subscriber.next(c));
const usage = extractUsageFromMetadata(lastUsageMetadata); const usage = this.extractUsageFromMetadata(lastUsageMetadata);
if (usage) { if (usage) {
void this.recordUsage(accountId, _model, usage, 'anthropic'); this.recordUsage(accountId, _model, usage, 'anthropic').catch((err) =>
this.logger.error('Usage recording failed', err),
);
} else if (accumulatedText) {
const estimated = this.estimateTokens(accumulatedText);
this.recordUsage(accountId, _model, {
promptTokenCount: 0,
candidatesTokenCount: estimated,
totalTokenCount: estimated,
}, 'anthropic', true).catch((err) =>
this.logger.error('Usage recording failed', err),
);
} }
subscriber.complete(); subscriber.complete();
}); });
@@ -388,7 +407,7 @@ export class ProxyService {
extraHeaders, extraHeaders,
); );
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
return this.normalizeGeminiGenerateResponse(response); return this.normalizeGeminiGenerateResponse(response);
} catch (err) { } catch (err) {
if (err instanceof Error && this.isProjectContextError(err.message)) { if (err instanceof Error && this.isProjectContextError(err.message)) {
@@ -411,7 +430,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
return this.normalizeGeminiGenerateResponse(response); return this.normalizeGeminiGenerateResponse(response);
} catch (fallbackErr) { } catch (fallbackErr) {
lastError = fallbackErr; lastError = fallbackErr;
@@ -538,6 +557,7 @@ export class ProxyService {
let buffer = ''; let buffer = '';
let receivedData = false; let receivedData = false;
let lastUsageMetadata: Record<string, unknown> | undefined; let lastUsageMetadata: Record<string, unknown> | undefined;
let accumulatedText = '';
upstreamStream.on('data', (chunk: Buffer) => { upstreamStream.on('data', (chunk: Buffer) => {
receivedData = true; receivedData = true;
@@ -549,15 +569,20 @@ export class ProxyService {
const trimmed = line.trim(); const trimmed = line.trim();
if (!trimmed.startsWith('data: ')) continue; if (!trimmed.startsWith('data: ')) continue;
const dataStr = trimmed.slice(6); const dataStr = trimmed.slice(6);
subscriber.next(`data: ${dataStr}\n\n`);
if (dataStr === '[DONE]') continue; if (dataStr === '[DONE]') continue;
try { try {
const parsed = JSON.parse(dataStr); const parsed = JSON.parse(dataStr);
if (parsed?.usageMetadata) { if (parsed?.usageMetadata) {
lastUsageMetadata = parsed.usageMetadata; lastUsageMetadata = parsed.usageMetadata;
} }
const text = parsed?.candidates?.[0]?.content?.parts?.[0]?.text;
if (typeof text === 'string') {
accumulatedText += text;
}
} catch { /* ignore */ } } catch { /* ignore */ }
} }
subscriber.next(chunk.toString());
}); });
upstreamStream.on('end', () => { upstreamStream.on('end', () => {
@@ -565,8 +590,20 @@ export class ProxyService {
subscriber.error(new Error('Empty response stream')); subscriber.error(new Error('Empty response stream'));
return; return;
} }
if (lastUsageMetadata) { const usage = this.extractUsageFromMetadata(lastUsageMetadata);
void this.recordUsage(accountId, model, lastUsageMetadata as any, 'gemini'); if (usage) {
this.recordUsage(accountId, model, usage, 'gemini').catch((err) =>
this.logger.error('Usage recording failed', err),
);
} else if (accumulatedText) {
const estimated = this.estimateTokens(accumulatedText);
this.recordUsage(accountId, model, {
promptTokenCount: 0,
candidatesTokenCount: estimated,
totalTokenCount: estimated,
}, 'gemini', true).catch((err) =>
this.logger.error('Usage recording failed', err),
);
} }
subscriber.complete(); subscriber.complete();
}); });
@@ -845,7 +882,7 @@ export class ProxyService {
extraHeaders, extraHeaders,
); );
trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0); trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0);
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
this.logger.log( this.logger.log(
`Upstream response snippet after stream fallback: ${JSON.stringify(response).substring(0, 500)}`, `Upstream response snippet after stream fallback: ${JSON.stringify(response).substring(0, 500)}`,
); );
@@ -866,7 +903,7 @@ export class ProxyService {
this.logger.log( this.logger.log(
`Upstream response snippet (non-stream): ${JSON.stringify(response).substring(0, 500)}`, `Upstream response snippet (non-stream): ${JSON.stringify(response).substring(0, 500)}`,
); );
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
// Transform Gemini response to OpenAI format // Transform Gemini response to OpenAI format
const claudeResponse = transformResponse(response); const claudeResponse = transformResponse(response);
this.logger.log( this.logger.log(
@@ -904,7 +941,7 @@ export class ProxyService {
extraHeaders, extraHeaders,
); );
trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0); trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0);
await this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai'); this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai');
const claudeResponse = transformResponse(response); const claudeResponse = transformResponse(response);
return this.convertClaudeToOpenAIResponse(claudeResponse, request.model); return this.convertClaudeToOpenAIResponse(claudeResponse, request.model);
} catch (fallbackErr) { } catch (fallbackErr) {
@@ -1055,6 +1092,7 @@ export class ProxyService {
let hasSentDone = false; let hasSentDone = false;
let toolCallIndex = 0; let toolCallIndex = 0;
let lastUsageMetadata: Record<string, unknown> | undefined; let lastUsageMetadata: Record<string, unknown> | undefined;
let accumulatedText = '';
const streamId = `chatcmpl-${uuidv4()}`; const streamId = `chatcmpl-${uuidv4()}`;
const created = Math.floor(Date.now() / 1000); const created = Math.floor(Date.now() / 1000);
@@ -1108,6 +1146,7 @@ export class ProxyService {
for (const part of parts) { for (const part of parts) {
if (part.thought && part.text) { if (part.thought && part.text) {
accumulatedText += part.text;
// Reasoning/thinking parts: emit via reasoning_content delta field // Reasoning/thinking parts: emit via reasoning_content delta field
// so clients that support it can display thinking separately. // so clients that support it can display thinking separately.
const reasoningChunk = { const reasoningChunk = {
@@ -1183,6 +1222,7 @@ export class ProxyService {
} }
if (part.text) { if (part.text) {
accumulatedText += part.text;
pushRoleChunk(); pushRoleChunk();
const contentChunk = { const contentChunk = {
id: streamId, id: streamId,
@@ -1253,9 +1293,20 @@ export class ProxyService {
subscriber.next('data: [DONE]\n\n'); subscriber.next('data: [DONE]\n\n');
hasSentDone = true; hasSentDone = true;
} }
if (lastUsageMetadata) { const usage = this.extractUsageFromMetadata(lastUsageMetadata);
// TODO: integrate local tokenizer fallback for missing usageMetadata if (usage) {
void this.recordUsage(accountId, effectiveTargetModel, lastUsageMetadata as any, 'openai'); this.recordUsage(accountId, effectiveTargetModel, usage, 'openai').catch((err) =>
this.logger.error('Usage recording failed', err),
);
} else if (accumulatedText) {
const estimated = this.estimateTokens(accumulatedText);
this.recordUsage(accountId, effectiveTargetModel, {
promptTokenCount: 0,
candidatesTokenCount: estimated,
totalTokenCount: estimated,
}, 'openai', true).catch((err) =>
this.logger.error('Usage recording failed', err),
);
} }
subscriber.complete(); subscriber.complete();
}); });
@@ -1,10 +1,23 @@
import { Injectable } from '@nestjs/common'; import { Injectable } from '@nestjs/common';
import { z } from 'zod';
import { TokenUsageRepo, RecordUsageParams } from '../../../ipc/database/usageHandler'; import { TokenUsageRepo, RecordUsageParams } from '../../../ipc/database/usageHandler';
const RecordUsageInputSchema = z.object({
accountId: z.string().min(1),
model: z.string().min(1),
promptTokens: z.number().min(0),
completionTokens: z.number().min(0),
totalTokens: z.number().min(0),
timestamp: z.number(),
requestType: z.enum(['openai', 'anthropic', 'gemini']).optional(),
isEstimated: z.boolean().optional(),
});
@Injectable() @Injectable()
export class TokenUsageService { export class TokenUsageService {
async recordUsage(params: RecordUsageParams): Promise<void> { recordUsage(params: RecordUsageParams): boolean {
TokenUsageRepo.recordUsage(params); RecordUsageInputSchema.parse(params);
return TokenUsageRepo.recordUsage(params);
} }
getUsageByHour(accountId?: string, start?: number, end?: number) { getUsageByHour(accountId?: string, start?: number, end?: number) {
+1
View File
@@ -80,6 +80,7 @@ export const TokenUsageRowSchema = z.object({
total_tokens: z.number(), total_tokens: z.number(),
timestamp: z.number(), timestamp: z.number(),
request_type: z.string().nullable(), request_type: z.string().nullable(),
is_estimated: z.number().nullable(),
}); });
export type TokenUsageRow = z.infer<typeof TokenUsageRowSchema>; export type TokenUsageRow = z.infer<typeof TokenUsageRowSchema>;