diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index 1457270..368b46e 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/src/server/modules/proxy/proxy.service.ts @@ -2,6 +2,7 @@ import { Injectable, Logger, Inject } from '@nestjs/common'; import { isEmpty, isNil, isNumber, isPlainObject, isString } from 'lodash-es'; import { TokenManagerService } from './token-manager.service'; import { GeminiClient } from './clients/gemini.client'; +import { TokenUsageService } from '../usage/token-usage.service'; import { v4 as uuidv4 } from 'uuid'; import { Observable } from 'rxjs'; import { transformClaudeRequestIn } from '../../../lib/geminiNexus/ClaudeRequestMapper'; @@ -42,8 +43,29 @@ export class ProxyService { constructor( @Inject(TokenManagerService) private readonly tokenManager: TokenManagerService, @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 --- async handleAnthropicMessages( @@ -101,7 +123,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); - return this.processAnthropicInternalStream(stream, geminiBody.model); + return this.processAnthropicInternalStream(stream, geminiBody.model, token.id); } else { const response = await this.generateInternalWithStreamFallback( geminiBody, @@ -109,6 +131,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic'); return this.toAnthropicChatResponse(transformResponse(response)); } } catch (error) { @@ -132,7 +155,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); - return this.processAnthropicInternalStream(stream, fallbackBody.model); + return this.processAnthropicInternalStream(stream, fallbackBody.model, token.id); } else { const response = await this.generateInternalWithStreamFallback( fallbackBody, @@ -140,6 +163,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic'); return this.toAnthropicChatResponse(transformResponse(response)); } } catch (fallbackErr) { @@ -170,7 +194,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); - return this.processAnthropicInternalStream(stream, downgradedBody.model); + return this.processAnthropicInternalStream(stream, downgradedBody.model, token.id); } else { const response = await this.generateInternalWithStreamFallback( downgradedBody, @@ -178,6 +202,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'anthropic'); const transformed = this.toAnthropicChatResponse(transformResponse(response)); return { ...transformed, @@ -205,6 +230,7 @@ export class ProxyService { private processAnthropicInternalStream( upstreamStream: NodeJS.ReadableStream, _model: string, + accountId: string, ): Observable { return new Observable((subscriber) => { const decoder = new TextDecoder(); @@ -272,6 +298,12 @@ 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', + ); subscriber.complete(); }); @@ -339,6 +371,7 @@ export class ProxyService { extraHeaders, ); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini'); return this.normalizeGeminiGenerateResponse(response); } catch (err) { if (err instanceof Error && this.isProjectContextError(err.message)) { @@ -361,6 +394,7 @@ export class ProxyService { token.token.upstream_proxy_url, extraHeaders, ); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'gemini'); return this.normalizeGeminiGenerateResponse(response); } catch (fallbackErr) { lastError = fallbackErr; @@ -769,6 +803,7 @@ export class ProxyService { extraHeaders, ); trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai'); this.logger.log( `Upstream response snippet after stream fallback: ${JSON.stringify(response).substring(0, 500)}`, ); @@ -789,6 +824,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'); // Transform Gemini response to OpenAI format const claudeResponse = transformResponse(response); this.logger.log( @@ -826,6 +862,7 @@ export class ProxyService { extraHeaders, ); trafficLogger.logUpstreamResponse(requestId || 'unknown', '/v1/chat/completions', 200, response, 0); + this.recordUsage(token.id, effectiveTargetModel, response.usageMetadata, 'openai'); const claudeResponse = transformResponse(response); return this.convertClaudeToOpenAIResponse(claudeResponse, request.model); } catch (fallbackErr) { diff --git a/src/tests/scripts/project-id-request-validation.ts b/src/tests/scripts/project-id-request-validation.ts index 60a6030..1af97ac 100644 --- a/src/tests/scripts/project-id-request-validation.ts +++ b/src/tests/scripts/project-id-request-validation.ts @@ -19,9 +19,13 @@ const mockGeminiClient: any = { generateInternal: async () => undefined, }; +const mockTokenUsageService: any = { + recordUsage: () => undefined, +}; + class TestableProxyService extends ProxyService { constructor() { - super(mockTokenManager, mockGeminiClient); + super(mockTokenManager, mockGeminiClient, mockTokenUsageService); } public createGeminiInternal( @@ -255,7 +259,7 @@ async function validateRuntimeAnthropicRequestFromRealTokenManager(): Promise model), }; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; +const mockTokenUsageService = { recordUsage: vi.fn() }; // Subclass to access private method class TestableProxyService extends ProxyService { 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 { // 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 { diff --git a/src/tests/unit/proxy-stream-debug.test.ts b/src/tests/unit/proxy-stream-debug.test.ts index 477b0da..0e04624 100644 --- a/src/tests/unit/proxy-stream-debug.test.ts +++ b/src/tests/unit/proxy-stream-debug.test.ts @@ -15,10 +15,11 @@ const mockTokenManager = { resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model), }; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; +const mockTokenUsageService = { recordUsage: vi.fn() }; class TestableProxyService extends ProxyService { constructor() { - super(mockTokenManager as any, mockGeminiClient as any); + super(mockTokenManager as any, mockGeminiClient as any, mockTokenUsageService as any); } public testProcessStreamResponse(