feat(proxy): inject token usage recording into all handler paths

This commit is contained in:
Gemini Nexus Bot committed 2026-04-22 13:30:46 -03:00
1 parent ce21283170
commit 1a1c0acd16
5 files changed
+53 -9

No files matched your search

+40 -3
View File
@@ -2,6 +2,7 @@ import { Injectable, Logger, Inject } from '@nestjs/common';
import { isEmpty, isNil, isNumber, isPlainObject, isString } from 'lodash-es'; import { isEmpty, isNil, isNumber, isPlainObject, isString } from 'lodash-es';
import { TokenManagerService } from './token-manager.service'; import { TokenManagerService } from './token-manager.service';
import { GeminiClient } from './clients/gemini.client'; import { GeminiClient } from './clients/gemini.client';
import { TokenUsageService } from '../usage/token-usage.service';
import { v4 as uuidv4 } from 'uuid'; import { v4 as uuidv4 } from 'uuid';
import { Observable } from 'rxjs'; import { Observable } from 'rxjs';
import { transformClaudeRequestIn } from '../../../lib/geminiNexus/ClaudeRequestMapper'; import { transformClaudeRequestIn } from '../../../lib/geminiNexus/ClaudeRequestMapper';
@@ -42,8 +43,29 @@ export class ProxyService {
constructor( constructor(
@Inject(TokenManagerService) private readonly tokenManager: TokenManagerService, @Inject(TokenManagerService) private readonly tokenManager: TokenManagerService,
@Inject(GeminiClient) private readonly geminiClient: GeminiClient, @Inject(GeminiClient) private readonly geminiClient: GeminiClient,
@Inject(TokenUsageService) private readonly tokenUsageService: TokenUsageService,
) {} ) {}
private recordUsage(
accountId: string,
model: string,
usage: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number } | undefined | null,
requestType: string,
): void {
if (!usage) {
return;
}
this.tokenUsageService.recordUsage({
accountId,
model,
promptTokens: usage.promptTokenCount ?? 0,
completionTokens: usage.candidatesTokenCount ?? 0,
totalTokens: usage.totalTokenCount ?? 0,
timestamp: Date.now(),
requestType,
});
}
// --- Anthropic Handlers --- // --- Anthropic Handlers ---
async handleAnthropicMessages( async handleAnthropicMessages(
@@ -101,7 +123,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
return this.processAnthropicInternalStream(stream, geminiBody.model); return this.processAnthropicInternalStream(stream, geminiBody.model, token.id);
} else { } else {
const response = await this.generateInternalWithStreamFallback( const response = await this.generateInternalWithStreamFallback(
geminiBody, geminiBody,
@@ -109,6 +131,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
return this.toAnthropicChatResponse(transformResponse(response)); return this.toAnthropicChatResponse(transformResponse(response));
} }
} catch (error) { } catch (error) {
@@ -132,7 +155,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
return this.processAnthropicInternalStream(stream, fallbackBody.model); return this.processAnthropicInternalStream(stream, fallbackBody.model, token.id);
} else { } else {
const response = await this.generateInternalWithStreamFallback( const response = await this.generateInternalWithStreamFallback(
fallbackBody, fallbackBody,
@@ -140,6 +163,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
return this.toAnthropicChatResponse(transformResponse(response)); return this.toAnthropicChatResponse(transformResponse(response));
} }
} catch (fallbackErr) { } catch (fallbackErr) {
@@ -170,7 +194,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
return this.processAnthropicInternalStream(stream, downgradedBody.model); return this.processAnthropicInternalStream(stream, downgradedBody.model, token.id);
} else { } else {
const response = await this.generateInternalWithStreamFallback( const response = await this.generateInternalWithStreamFallback(
downgradedBody, downgradedBody,
@@ -178,6 +202,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic');
const transformed = this.toAnthropicChatResponse(transformResponse(response)); const transformed = this.toAnthropicChatResponse(transformResponse(response));
return { return {
...transformed, ...transformed,
@@ -205,6 +230,7 @@ export class ProxyService {
private processAnthropicInternalStream( private processAnthropicInternalStream(
upstreamStream: NodeJS.ReadableStream, upstreamStream: NodeJS.ReadableStream,
_model: string, _model: string,
accountId: string,
): Observable<string> { ): Observable<string> {
return new Observable<string>((subscriber) => { return new Observable<string>((subscriber) => {
const decoder = new TextDecoder(); const decoder = new TextDecoder();
@@ -272,6 +298,12 @@ 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));
this.recordUsage(
accountId,
_model,
lastUsageMetadata as { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number },
'anthropic',
);
subscriber.complete(); subscriber.complete();
}); });
@@ -339,6 +371,7 @@ export class ProxyService {
extraHeaders, extraHeaders,
); );
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)) {
@@ -361,6 +394,7 @@ export class ProxyService {
token.token.upstream_proxy_url, token.token.upstream_proxy_url,
extraHeaders, extraHeaders,
); );
this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini');
return this.normalizeGeminiGenerateResponse(response); return this.normalizeGeminiGenerateResponse(response);
} catch (fallbackErr) { } catch (fallbackErr) {
lastError = fallbackErr; lastError = fallbackErr;
@@ -769,6 +803,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);
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)}`,
); );
@@ -789,6 +824,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)}`,
); );
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(
@@ -826,6 +862,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);
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) {
@@ -19,9 +19,13 @@ const mockGeminiClient: any = {
generateInternal: async () => undefined, generateInternal: async () => undefined,
}; };
const mockTokenUsageService: any = {
recordUsage: () => undefined,
};
class TestableProxyService extends ProxyService { class TestableProxyService extends ProxyService {
constructor() { constructor() {
super(mockTokenManager, mockGeminiClient); super(mockTokenManager, mockGeminiClient, mockTokenUsageService);
} }
public createGeminiInternal( public createGeminiInternal(
@@ -255,7 +259,7 @@ async function validateRuntimeAnthropicRequestFromRealTokenManager(): Promise<vo
}; };
try { try {
const service = new ProxyService(tokenManagerProxy as any, mockGeminiClient as any); const service = new ProxyService(tokenManagerProxy as any, mockGeminiClient as any, mockTokenUsageService as any);
await service.handleAnthropicMessages({ await service.handleAnthropicMessages({
model: 'claude-sonnet-4-5', model: 'claude-sonnet-4-5',
stream: false, stream: false,
+2 -1
View File
@@ -15,10 +15,11 @@ const mockTokenManager = {
recordParityError: vi.fn(), recordParityError: vi.fn(),
}; };
const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
const mockTokenUsageService = { recordUsage: vi.fn() };
class TestableProxyService extends ProxyService { class TestableProxyService extends ProxyService {
constructor() { constructor() {
super(mockTokenManager as any, mockGeminiClient as any); super(mockTokenManager as any, mockGeminiClient as any, mockTokenUsageService as any);
} }
public toAnthropic(request: any): any { public toAnthropic(request: any): any {
+3 -2
View File
@@ -16,16 +16,17 @@ const mockTokenManager = {
resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model), resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model),
}; };
const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
const mockTokenUsageService = { recordUsage: vi.fn() };
// Subclass to access private method // Subclass to access private method
class TestableProxyService extends ProxyService { class TestableProxyService extends ProxyService {
constructor() { constructor() {
super(mockTokenManager as any, mockGeminiClient as any); super(mockTokenManager as any, mockGeminiClient as any, mockTokenUsageService as any);
} }
public testProcessStream(stream: any, model: string = 'model'): Observable<string> { public testProcessStream(stream: any, model: string = 'model'): Observable<string> {
// Access private method using type assertion // Access private method using type assertion
return (this as any).processAnthropicInternalStream(stream, model); return (this as any).processAnthropicInternalStream(stream, model, 'acc-test');
} }
public testPassthroughStream(stream: any): Observable<string> { public testPassthroughStream(stream: any): Observable<string> {
+2 -1
View File
@@ -15,10 +15,11 @@ const mockTokenManager = {
resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model), resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model),
}; };
const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
const mockTokenUsageService = { recordUsage: vi.fn() };
class TestableProxyService extends ProxyService { class TestableProxyService extends ProxyService {
constructor() { constructor() {
super(mockTokenManager as any, mockGeminiClient as any); super(mockTokenManager as any, mockGeminiClient as any, mockTokenUsageService as any);
} }
public testProcessStreamResponse( public testProcessStreamResponse(