From 5401e40943824d54bc95971229a83934e55ad88f Mon Sep 17 00:00:00 2001 From: Draculabo Date: Fri, 6 Mar 2026 21:50:18 +0800 Subject: [PATCH] feat(proxy): improve parity, rate limiting, and token scheduling --- src/ipc/config/handlers.ts | 2 + .../modules/proxy/clients/gemini.client.ts | 76 +- src/server/modules/proxy/gemini.controller.ts | 34 +- .../proxy/interfaces/request-interfaces.ts | 5 +- src/server/modules/proxy/proxy.controller.ts | 223 +++-- src/server/modules/proxy/proxy.service.ts | 440 ++++++--- .../modules/proxy/token-manager.service.ts | 873 +++++++++++++++--- src/services/GoogleAPIService.ts | 249 ++++- .../gemini-controller.integration.test.ts | 13 +- src/tests/unit/parity-rollout-guard.test.ts | 106 +++ .../unit/proxy-controller.integration.test.ts | 67 ++ src/tests/unit/streaming-state.test.ts | 46 +- .../token-manager-parity-scheduling.test.ts | 166 ++++ src/types/config.ts | 34 +- 14 files changed, 1878 insertions(+), 456 deletions(-) create mode 100644 src/tests/unit/parity-rollout-guard.test.ts create mode 100644 src/tests/unit/token-manager-parity-scheduling.test.ts diff --git a/src/ipc/config/handlers.ts b/src/ipc/config/handlers.ts index c986472..7e98884 100644 --- a/src/ipc/config/handlers.ts +++ b/src/ipc/config/handlers.ts @@ -2,6 +2,7 @@ import { AppConfig } from '../../types/config'; import { ConfigManager } from './manager'; import { syncAutoStart } from '../../utils/autoStart'; import { logger } from '../../utils/logger'; +import { setServerConfig } from '../../server/server-config'; export function loadConfig(): AppConfig { return ConfigManager.loadConfig(); @@ -13,6 +14,7 @@ export async function saveConfig(config: AppConfig): Promise { // For now just save const previous = ConfigManager.getCachedConfig() ?? ConfigManager.loadConfig(); await ConfigManager.saveConfig(config); + setServerConfig(config.proxy); logger.setErrorReportingEnabled(config.error_reporting_enabled); if (previous.auto_startup !== config.auto_startup) { syncAutoStart(config); diff --git a/src/server/modules/proxy/clients/gemini.client.ts b/src/server/modules/proxy/clients/gemini.client.ts index c9a92e4..9a1a233 100644 --- a/src/server/modules/proxy/clients/gemini.client.ts +++ b/src/server/modules/proxy/clients/gemini.client.ts @@ -4,11 +4,12 @@ import { isNil } from 'lodash-es'; import { GeminiRequest, GeminiResponse } from '../interfaces/request-interfaces'; import { GeminiInternalRequest } from '../../../../lib/antigravity/types'; import { getServerConfig } from '../../../server-config'; +import { resolveRequestUserAgent } from '../request-user-agent'; +import { UpstreamRequestError } from './upstream-error'; @Injectable() export class GeminiClient { private readonly logger = new Logger(GeminiClient.name); - private readonly defaultRequestUserAgent = 'antigravity/1.11.9 windows/amd64'; // Default to v1beta for most features private readonly baseUrl = 'https://generativelanguage.googleapis.com/v1beta'; private readonly defaultInternalBaseUrls = [ @@ -23,7 +24,7 @@ export class GeminiClient { upstreamProxyUrl?: string, ): Promise { const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`; - const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl); + const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl); try { const response = await axios.post(url, content, { @@ -39,7 +40,14 @@ export class GeminiClient { } catch (error) { if (axios.isAxiosError(error)) { this.logger.error(`Gemini stream request failed: ${error.message}`); - throw new Error(error.response?.data?.error?.message || error.message); + throw new UpstreamRequestError({ + message: error.response?.data?.error?.message || error.message, + status: error.response?.status, + headers: { + retryAfter: this.extractRetryAfterHeader(error.response?.headers), + }, + body: this.describeAxiosErrorData(error.response?.data), + }); } this.throwAsCleanError(error); } @@ -52,7 +60,7 @@ export class GeminiClient { upstreamProxyUrl?: string, ): Promise { const url = `${this.baseUrl}/models/${model}:generateContent`; - const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl); + const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl); try { const response = await axios.post(url, content, { @@ -69,7 +77,14 @@ export class GeminiClient { this.logger.error( `Gemini request failed: ${error.message} - ${this.safeStringify(error.response?.data)}`, ); - throw new Error(error.response?.data?.error?.message || error.message); + throw new UpstreamRequestError({ + message: error.response?.data?.error?.message || error.message, + status: error.response?.status, + headers: { + retryAfter: this.extractRetryAfterHeader(error.response?.headers), + }, + body: this.describeAxiosErrorData(error.response?.data), + }); } this.throwAsCleanError(error); } @@ -83,7 +98,7 @@ export class GeminiClient { upstreamProxyUrl?: string, extraHeaders?: Record, ): Promise { - const response = await this.requestWithEndpointFailover( + const response = await this.executeRequestWithEndpointFailover( ':streamGenerateContent?alt=sse', body, accessToken, @@ -104,7 +119,7 @@ export class GeminiClient { upstreamProxyUrl?: string, extraHeaders?: Record, ): Promise { - const response = await this.requestWithEndpointFailover< + const response = await this.executeRequestWithEndpointFailover< GeminiResponse | { response: GeminiResponse } >( ':generateContent', @@ -143,7 +158,7 @@ export class GeminiClient { return Math.max(1, timeoutSeconds) * 1000; } - private shouldSwitchEndpointOnError(error: unknown): boolean { + private shouldFailoverToNextEndpoint(error: unknown): boolean { if (!axios.isAxiosError(error)) { return false; } @@ -162,7 +177,7 @@ export class GeminiClient { return status === 408 || status === 429 || status >= 500; } - private async requestWithEndpointFailover( + private async executeRequestWithEndpointFailover( path: string, body: GeminiInternalRequest, accessToken: string, @@ -173,8 +188,8 @@ export class GeminiClient { ): Promise> { const baseUrls = this.getInternalBaseUrls(); const timeout = this.getInternalTimeoutMs(); - const requestUserAgent = process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent; - const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl); + const requestUserAgent = await resolveRequestUserAgent(); + const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl); let lastError: unknown = null; for (let index = 0; index < baseUrls.length; index++) { @@ -197,8 +212,8 @@ export class GeminiClient { lastError = error; const hasNextEndpoint = index < baseUrls.length - 1; - if (!hasNextEndpoint || !this.shouldSwitchEndpointOnError(error)) { - await this.handleAxiosError(error, operation); + if (!hasNextEndpoint || !this.shouldFailoverToNextEndpoint(error)) { + await this.throwUpstreamRequestError(error, operation); } this.logger.warn( @@ -209,11 +224,11 @@ export class GeminiClient { } } - await this.handleAxiosError(lastError, operation); + await this.throwUpstreamRequestError(lastError, operation); throw new Error(`[${operation}] unexpected control flow after upstream error handling`); } - private async handleAxiosError(error: unknown, operation: string): Promise { + private async throwUpstreamRequestError(error: unknown, operation: string): Promise { if (axios.isAxiosError(error)) { const responseData = error.response?.data; const upstreamMessage = await this.extractAxiosErrorMessage(responseData); @@ -222,7 +237,14 @@ export class GeminiClient { responseData, )}`, ); - throw new Error(upstreamMessage || error.message); + throw new UpstreamRequestError({ + message: upstreamMessage || error.message, + status: error.response?.status, + headers: { + retryAfter: this.extractRetryAfterHeader(error.response?.headers), + }, + body: this.describeAxiosErrorData(responseData), + }); } this.throwAsCleanError(error); } @@ -346,7 +368,9 @@ export class GeminiClient { return this.safeStringify(responseData); } - private resolveAxiosProxy(upstreamProxyUrl?: string): AxiosProxyConfig | false | undefined { + private resolveUpstreamAxiosProxy( + upstreamProxyUrl?: string, + ): AxiosProxyConfig | false | undefined { const config = getServerConfig(); const configuredProxyUrl = upstreamProxyUrl || @@ -382,6 +406,24 @@ export class GeminiClient { } } + private extractRetryAfterHeader(headers: unknown): string | undefined { + if (!headers || typeof headers !== 'object') { + return undefined; + } + + const retryAfter = (headers as Record)['retry-after']; + if (typeof retryAfter === 'string' && retryAfter.trim() !== '') { + return retryAfter.trim(); + } + if (Array.isArray(retryAfter) && retryAfter.length > 0) { + const first = retryAfter[0]; + if (typeof first === 'string' && first.trim() !== '') { + return first.trim(); + } + } + return undefined; + } + private throwAsCleanError(error: unknown): never { // Re-throw as clean Error to avoid circular reference issues. throw error instanceof Error ? new Error(error.message) : new Error(String(error)); diff --git a/src/server/modules/proxy/gemini.controller.ts b/src/server/modules/proxy/gemini.controller.ts index e90eb06..a6c289f 100644 --- a/src/server/modules/proxy/gemini.controller.ts +++ b/src/server/modules/proxy/gemini.controller.ts @@ -8,6 +8,7 @@ import { Post, Res, UseGuards, + Optional, } from '@nestjs/common'; import { FastifyReply } from 'fastify'; import { isEmpty } from 'lodash-es'; @@ -18,6 +19,7 @@ import { ProxyService } from './proxy.service'; import { GeminiRequest, GeminiResponse } from './interfaces/request-interfaces'; import { getServerConfig } from '../../server-config'; import { getAllDynamicModels } from '../../../lib/antigravity/ModelMapping'; +import { TokenManagerService } from './token-manager.service'; type GeminiModelMetadata = { name: string; @@ -35,7 +37,10 @@ type GeminiModelMetadata = { @Controller('v1beta') @UseGuards(ProxyGuard) export class GeminiController { - constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {} + constructor( + @Inject(ProxyService) private readonly proxyService: ProxyService, + @Optional() @Inject(TokenManagerService) private readonly tokenManager?: TokenManagerService, + ) {} @Get('models') listModels(@Res() res: FastifyReply) { @@ -70,7 +75,7 @@ export class GeminiController { @Body() body: GeminiRequest, @Res() res: FastifyReply, ) { - const parsed = this.parseCombinedModelAction(modelAction); + const parsed = this.parseModelActionToken(modelAction); if (!parsed) { res.status(HttpStatus.BAD_REQUEST).send({ error: { @@ -82,7 +87,7 @@ export class GeminiController { return; } - await this.dispatchAction(parsed.model, parsed.action, body, res); + await this.handleModelActionDispatch(parsed.model, parsed.action, body, res); } @Post('models/:model/countTokens') @@ -91,10 +96,10 @@ export class GeminiController { @Body() body: GeminiRequest, @Res() res: FastifyReply, ) { - await this.dispatchAction(`models/${model}`, 'countTokens', body, res); + await this.handleModelActionDispatch(`models/${model}`, 'countTokens', body, res); } - private async dispatchAction( + private async handleModelActionDispatch( model: string, action: string, body: GeminiRequest, @@ -111,14 +116,14 @@ export class GeminiController { if (action === 'streamGenerateContent') { const stream = await this.proxyService.handleGeminiStreamGenerateContent(model, body); if (stream instanceof Observable) { - this.sendObservableSse(res, stream); + this.writeObservableSseResponse(res, stream); return; } } if (action === 'generateContent') { const result = await this.proxyService.handleGeminiGenerateContent(model, body); - res.status(HttpStatus.OK).send(this.normalizeGeminiGenerateResponse(result)); + res.status(HttpStatus.OK).send(this.buildNormalizedGeminiGenerateResponse(result)); return; } @@ -141,7 +146,7 @@ export class GeminiController { } } - private parseCombinedModelAction(modelAction: string): { + private parseModelActionToken(modelAction: string): { model: string; action: string; } | null { @@ -165,7 +170,10 @@ export class GeminiController { private buildGeminiModelList(): GeminiModelMetadata[] { const config = getServerConfig(); - const dynamicModelIds = getAllDynamicModels(config?.custom_mapping ?? {}); + const dynamicModelIds = getAllDynamicModels( + config?.custom_mapping ?? {}, + this.tokenManager?.getAllCollectedModels(), + ); return dynamicModelIds.map((id) => this.toGeminiModelMetadata(`models/${id}`)); } @@ -186,7 +194,7 @@ export class GeminiController { }; } - private normalizeGeminiGenerateResponse(response: GeminiResponse): GeminiResponse { + private buildNormalizedGeminiGenerateResponse(response: GeminiResponse): GeminiResponse { const candidates = (response.candidates ?? []).map((candidate, index) => ({ content: candidate.content, finishReason: candidate.finishReason, @@ -235,7 +243,7 @@ export class GeminiController { return normalized; } - private sendObservableSse(res: FastifyReply, stream: Observable): void { + private writeObservableSseResponse(res: FastifyReply, stream: Observable): void { if ( !res.raw || typeof res.raw.writeHead !== 'function' || @@ -248,7 +256,7 @@ export class GeminiController { return; } - if (this.hasHijack(res)) { + if (this.supportsReplyHijack(res)) { res.hijack(); } @@ -293,7 +301,7 @@ export class GeminiController { }); } - private hasHijack(reply: FastifyReply): reply is FastifyReply & { hijack: () => void } { + private supportsReplyHijack(reply: FastifyReply): reply is FastifyReply & { hijack: () => void } { return typeof (reply as { hijack?: unknown }).hijack === 'function'; } } diff --git a/src/server/modules/proxy/interfaces/request-interfaces.ts b/src/server/modules/proxy/interfaces/request-interfaces.ts index a79563a..1807614 100644 --- a/src/server/modules/proxy/interfaces/request-interfaces.ts +++ b/src/server/modules/proxy/interfaces/request-interfaces.ts @@ -30,12 +30,13 @@ export interface OpenAIContentPart { } export interface OpenAITool { - type: 'function'; - function: { + type: string; + function?: { name: string; description?: string; parameters?: Record; }; + [key: string]: unknown; } export interface OpenAIToolCall { diff --git a/src/server/modules/proxy/proxy.controller.ts b/src/server/modules/proxy/proxy.controller.ts index be65835..6aec030 100644 --- a/src/server/modules/proxy/proxy.controller.ts +++ b/src/server/modules/proxy/proxy.controller.ts @@ -9,9 +9,10 @@ import { Inject, Req, Logger, + Optional, } from '@nestjs/common'; import { FastifyReply, FastifyRequest } from 'fastify'; -import { isNil } from 'lodash-es'; +import { isNil, isPlainObject, isString } from 'lodash-es'; import { ProxyService } from './proxy.service'; import { Observable } from 'rxjs'; import { @@ -29,20 +30,27 @@ import { MODEL_LIST_OWNER, } from '../../../lib/antigravity/ModelMapping'; import { getServerConfig } from '../../server-config'; +import { TokenManagerService } from './token-manager.service'; @Controller('v1') @UseGuards(ProxyGuard) export class ProxyController { private readonly logger = new Logger(ProxyController.name); - constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {} + constructor( + @Inject(ProxyService) private readonly proxyService: ProxyService, + @Optional() @Inject(TokenManagerService) private readonly tokenManager?: TokenManagerService, + ) {} @Get('models') listModels(@Res() res: FastifyReply) { try { const config = getServerConfig(); const customMapping = config?.custom_mapping ?? {}; - const modelIds = getAllDynamicModels(customMapping); + const modelIds = getAllDynamicModels( + customMapping, + this.tokenManager?.getAllCollectedModels(), + ); const data = modelIds.map((id) => ({ id, @@ -69,7 +77,7 @@ export class ProxyController { @Post('chat/completions') async chatCompletions(@Body() body: OpenAIChatRequest, @Res() res: FastifyReply) { - await this.sendOpenAIChatCompletionResponse(body, res); + await this.respondOpenAIChatCompletions(body, res); } @Post('completions') @@ -86,11 +94,11 @@ export class ProxyController { @Res() res: FastifyReply, ) { const request: OpenAIChatRequest = { - model: body.model ?? 'gemini-2.5-flash', + model: body.model ?? 'gemini-3-flash', messages: [ { role: 'user', - content: this.normalizePrompt(body.prompt), + content: this.normalizeCompletionPrompt(body.prompt), }, ], max_tokens: body.max_tokens, @@ -106,17 +114,9 @@ export class ProxyController { } const response = result as OpenAIChatResponse; - res.status(HttpStatus.OK).send(this.toLegacyCompletionsResponse(response)); + res.status(HttpStatus.OK).send(this.toLegacyTextCompletionsResponse(response)); } catch (error) { - const message = error instanceof Error ? error.message : 'Internal Server Error'; - const status = this.resolveUpstreamStatus(message); - this.logEndpointError('/v1/completions', status, message, error); - res.status(status).send({ - error: { - message, - type: 'server_error', - }, - }); + this.sendOpenAIErrorResponse(res, '/v1/completions', error); } } @@ -127,6 +127,7 @@ export class ProxyController { model?: string; instructions?: string; input?: unknown; + tools?: OpenAIChatRequest['tools']; max_output_tokens?: number; temperature?: number; top_p?: number; @@ -144,17 +145,9 @@ export class ProxyController { } const response = result as OpenAIChatResponse; - res.status(HttpStatus.OK).send(this.toLegacyCompletionsResponse(response)); + res.status(HttpStatus.OK).send(this.toLegacyTextCompletionsResponse(response)); } catch (error) { - const message = error instanceof Error ? error.message : 'Internal Server Error'; - const status = this.resolveUpstreamStatus(message); - this.logEndpointError('/v1/responses', status, message, error); - res.status(status).send({ - error: { - message, - type: 'server_error', - }, - }); + this.sendOpenAIErrorResponse(res, '/v1/responses', error); } } @@ -182,7 +175,7 @@ export class ProxyController { quality: body.quality, }; - await this.sendOpenAIImageResponse(request, body.prompt ?? '', res); + await this.sendOpenAIImageGenerationResponse(request, body.prompt ?? '', res); } @Post('images/edits') @@ -200,7 +193,7 @@ export class ProxyController { @Req() req: FastifyRequest, @Res() res: FastifyReply, ) { - if (!this.hasValidMultipartBoundary(req)) { + if (!this.hasMultipartBoundary(req)) { res .status(HttpStatus.BAD_REQUEST) .send('Invalid `boundary` for `multipart/form-data` request'); @@ -236,7 +229,7 @@ export class ProxyController { quality: body.quality, }; - await this.sendOpenAIImageResponse(request, body.prompt ?? '', res); + await this.sendOpenAIImageGenerationResponse(request, body.prompt ?? '', res); } @Post('audio/transcriptions') @@ -251,7 +244,7 @@ export class ProxyController { @Req() req: FastifyRequest, @Res() res: FastifyReply, ) { - if (!this.hasValidMultipartBoundary(req)) { + if (!this.hasMultipartBoundary(req)) { res .status(HttpStatus.BAD_REQUEST) .send('Invalid `boundary` for `multipart/form-data` request'); @@ -271,7 +264,7 @@ export class ProxyController { try { const result = await this.proxyService.handleGeminiGenerateContent( - body.model ?? 'gemini-2.5-flash', + body.model ?? 'gemini-3-flash', { contents: [ { @@ -298,19 +291,11 @@ export class ProxyController { text: text ?? '', }); } catch (error) { - const message = error instanceof Error ? error.message : 'Internal Server Error'; - const status = this.resolveUpstreamStatus(message); - this.logEndpointError('/v1/audio/transcriptions', status, message, error); - res.status(status).send({ - error: { - message, - type: 'server_error', - }, - }); + this.sendOpenAIErrorResponse(res, '/v1/audio/transcriptions', error); } } - private async sendOpenAIChatCompletionResponse(body: OpenAIChatRequest, res: FastifyReply) { + private async respondOpenAIChatCompletions(body: OpenAIChatRequest, res: FastifyReply) { try { const result = await this.proxyService.handleChatCompletions(body); @@ -321,15 +306,7 @@ export class ProxyController { res.status(HttpStatus.OK).send(result); } } catch (error) { - const message = error instanceof Error ? error.message : 'Internal Server Error'; - const status = this.resolveUpstreamStatus(message); - this.logEndpointError('/v1/chat/completions', status, message, error); - res.status(status).send({ - error: { - message: message, - type: 'server_error', - }, - }); + this.sendOpenAIErrorResponse(res, '/v1/chat/completions', error); } } @@ -345,20 +322,11 @@ export class ProxyController { res.status(HttpStatus.OK).send(result); } } catch (error) { - const message = error instanceof Error ? error.message : 'Internal Server Error'; - const status = this.resolveUpstreamStatus(message); - this.logEndpointError('/v1/messages', status, message, error); - res.status(status).send({ - type: 'error', - error: { - type: 'api_error', - message: message, - }, - }); + this.sendAnthropicErrorResponse(res, '/v1/messages', error); } } - private normalizePrompt(prompt: string | string[] | undefined): string { + private normalizeCompletionPrompt(prompt: string | string[] | undefined): string { if (!prompt) { return ''; } @@ -368,7 +336,7 @@ export class ProxyController { return prompt; } - private toLegacyCompletionsResponse(response: OpenAIChatResponse): Record { + private toLegacyTextCompletionsResponse(response: OpenAIChatResponse): Record { const choice = response.choices?.[0]; const content = choice?.message?.content; const text = typeof content === 'string' ? content : ''; @@ -391,21 +359,20 @@ export class ProxyController { } private normalizeResponsesInput(input: unknown): string { - if (typeof input === 'string') { + if (isString(input)) { return input; } if (Array.isArray(input)) { return input .map((item) => { - if (typeof item === 'string') { + if (isString(item)) { return item; } - if (item && typeof item === 'object' && 'content' in item) { - const content = (item as { content?: unknown }).content; - if (typeof content === 'string') { - return content; - } + const itemRecord = this.toRecord(item); + const content = this.asString(itemRecord?.content); + if (content) { + return content; } return JSON.stringify(item); }) @@ -423,6 +390,7 @@ export class ProxyController { model?: string; instructions?: string; input?: unknown; + tools?: OpenAIChatRequest['tools']; max_output_tokens?: number; temperature?: number; top_p?: number; @@ -458,7 +426,7 @@ export class ProxyController { type === 'local_shell_call' ? 'shell' : type === 'web_search_call' - ? 'google_search' + ? 'builtin_web_search' : (this.asString(itemObj.name) ?? 'unknown'); callIdToToolName.set(callId, toolName); } @@ -536,8 +504,9 @@ export class ProxyController { } return { - model: body.model ?? 'gemini-2.5-flash', + model: body.model ?? 'gemini-3-flash', messages, + tools: body.tools, max_tokens: body.max_output_tokens, temperature: body.temperature, top_p: body.top_p, @@ -546,7 +515,7 @@ export class ProxyController { } private normalizeResponsesMessageContent(content: unknown): string | OpenAIContentPart[] { - if (typeof content === 'string') { + if (isString(content)) { return content; } @@ -621,11 +590,12 @@ export class ProxyController { } const raw = item.arguments; - if (typeof raw === 'string') { + if (isString(raw)) { try { const parsed = JSON.parse(raw); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { - return parsed as Record; + const parsedRecord = this.toRecord(parsed); + if (parsedRecord) { + return parsedRecord; } return { value: parsed, @@ -637,23 +607,22 @@ export class ProxyController { } } - if (raw && typeof raw === 'object' && !Array.isArray(raw)) { - return raw as Record; + const rawRecord = this.toRecord(raw); + if (rawRecord) { + return rawRecord; } return {}; } private normalizeResponsesOutput(output: unknown): string { - if (typeof output === 'string') { + if (isString(output)) { return output; } - if (output && typeof output === 'object' && !Array.isArray(output)) { - const obj = output as Record; - const content = this.asString(obj.content); - if (content) { - return content; - } + const outputRecord = this.toRecord(output); + const content = this.asString(outputRecord?.content); + if (content) { + return content; } if (isNil(output)) { return ''; @@ -663,15 +632,13 @@ export class ProxyController { private resolveImageUrl(block: Record): string | null { const raw = block.image_url; - if (typeof raw === 'string') { + if (isString(raw)) { return raw; } - if (raw && typeof raw === 'object' && !Array.isArray(raw)) { - const record = raw as Record; - const url = this.asString(record.url); - if (url) { - return url; - } + const rawRecord = this.toRecord(raw); + const url = this.asString(rawRecord?.url); + if (url) { + return url; } return null; } @@ -703,7 +670,7 @@ export class ProxyController { return null; } - if (typeof input === 'string') { + if (isString(input)) { const dataUri = input.match(/^data:(?[^;]+);base64,(?[A-Za-z0-9+/=]+)$/); if (dataUri?.groups?.mime && dataUri.groups.data) { return { @@ -722,14 +689,14 @@ export class ProxyController { return null; } - if (typeof input === 'object' && !Array.isArray(input)) { - const obj = input as Record; - const data = this.asString(obj.data); + const inputRecord = this.toRecord(input); + if (inputRecord) { + const data = this.asString(inputRecord.data); if (!data) { return null; } return { - mimeType: this.asString(obj.mimeType) ?? 'audio/mpeg', + mimeType: this.asString(inputRecord.mimeType) ?? 'audio/mpeg', data, }; } @@ -738,14 +705,14 @@ export class ProxyController { } private toRecord(value: unknown): Record | null { - if (!value || typeof value !== 'object' || Array.isArray(value)) { + if (!isPlainObject(value)) { return null; } return value as Record; } private asString(value: unknown): string | null { - return typeof value === 'string' ? value : null; + return isString(value) ? value : null; } private isObservableLike(value: unknown): value is Observable { @@ -814,7 +781,7 @@ export class ProxyController { }); } - private async sendOpenAIImageResponse( + private async sendOpenAIImageGenerationResponse( request: OpenAIChatRequest, prompt: string, res: FastifyReply, @@ -822,7 +789,7 @@ export class ProxyController { try { const result = await this.proxyService.handleChatCompletions(request); if (result instanceof Observable) { - this.logEndpointError( + this.logProxyEndpointError( '/v1/images/generations', HttpStatus.INTERNAL_SERVER_ERROR, 'Streaming image generation is not supported by this endpoint', @@ -839,7 +806,7 @@ export class ProxyController { const content = result.choices?.[0]?.message?.content; const image = this.extractInlineBase64Image(typeof content === 'string' ? content : ''); if (!image) { - this.logEndpointError( + this.logProxyEndpointError( '/v1/images/generations', HttpStatus.BAD_GATEWAY, 'Upstream did not return inline image data', @@ -889,14 +856,7 @@ export class ProxyController { } } - const status = this.resolveUpstreamStatus(message); - this.logEndpointError('/v1/images/generations', status, message, error); - res.status(status).send({ - error: { - message, - type: 'server_error', - }, - }); + this.sendOpenAIErrorResponse(res, '/v1/images/generations', error, message); } } @@ -1002,7 +962,7 @@ export class ProxyController { ); } - private hasValidMultipartBoundary(req: FastifyRequest): boolean { + private hasMultipartBoundary(req: FastifyRequest): boolean { const contentType = req.headers['content-type']; if (typeof contentType !== 'string') { return false; @@ -1012,7 +972,46 @@ export class ProxyController { return lowered.includes('multipart/form-data') && lowered.includes('boundary='); } - private resolveUpstreamStatus(message: string): HttpStatus { + private resolveErrorMessageText(error: unknown): string { + return error instanceof Error ? error.message : 'Internal Server Error'; + } + + private sendOpenAIErrorResponse( + res: FastifyReply, + endpoint: string, + error: unknown, + overrideMessage?: string, + ): void { + const message = overrideMessage ?? this.resolveErrorMessageText(error); + const status = this.resolveErrorHttpStatus(message); + this.logProxyEndpointError(endpoint, status, message, error); + res.status(status).send({ + error: { + message, + type: 'server_error', + }, + }); + } + + private sendAnthropicErrorResponse( + res: FastifyReply, + endpoint: string, + error: unknown, + overrideMessage?: string, + ): void { + const message = overrideMessage ?? this.resolveErrorMessageText(error); + const status = this.resolveErrorHttpStatus(message); + this.logProxyEndpointError(endpoint, status, message, error); + res.status(status).send({ + type: 'error', + error: { + type: 'api_error', + message, + }, + }); + } + + private resolveErrorHttpStatus(message: string): HttpStatus { const lowered = message.toLowerCase(); if (lowered.includes('all accounts failed or unhealthy')) { return HttpStatus.SERVICE_UNAVAILABLE; @@ -1050,7 +1049,7 @@ export class ProxyController { return HttpStatus.INTERNAL_SERVER_ERROR; } - private logEndpointError( + private logProxyEndpointError( endpoint: string, status: HttpStatus, message: string, diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index 9563bb2..e9fbb1a 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/src/server/modules/proxy/proxy.service.ts @@ -1,5 +1,5 @@ import { Injectable, Logger, Inject } from '@nestjs/common'; -import { isEmpty, isNil } from 'lodash-es'; +import { isEmpty, isNil, isPlainObject, isString } from 'lodash-es'; import { TokenManagerService } from './token-manager.service'; import { GeminiClient } from './clients/gemini.client'; import { v4 as uuidv4 } from 'uuid'; @@ -14,6 +14,7 @@ import { GeminiPart as InternalGeminiPart, } from '../../../lib/antigravity/types'; import { calculateRetryDelay, sleep } from '../../../lib/antigravity/retry-utils'; +import { normalizeObjectJsonSchema } from '../../../lib/antigravity/JsonSchemaUtils'; import { classifyStreamError, formatErrorForSSE, @@ -28,12 +29,14 @@ import { AnthropicContent, } from './interfaces/request-interfaces'; import { getServerConfig } from '../../server-config'; -import { resolveModelRoute } from '../../../lib/antigravity/ModelMapping'; +import { normalizeGeminiModelAlias, resolveModelRoute } from '../../../lib/antigravity/ModelMapping'; +import { getMaxOutputTokens, getThinkingBudget } from '../../../lib/antigravity/ModelSpecs'; +import { resolveRequestUserAgent } from './request-user-agent'; +import { UpstreamRequestError } from './clients/upstream-error'; @Injectable() export class ProxyService { private readonly logger = new Logger(ProxyService.name); - private readonly defaultRequestUserAgent = 'antigravity/1.11.9 windows/amd64'; constructor( @Inject(TokenManagerService) private readonly tokenManager: TokenManagerService, @@ -45,12 +48,12 @@ export class ProxyService { async handleAnthropicMessages( request: AnthropicChatRequest, ): Promise> { - const sessionKey = this.getSessionKeyFromAnthropic(request); + const sessionKey = this.extractAnthropicSessionKey(request); - const targetModel = this.mapModel(request.model); - const extraHeaders = this.buildModelSpecificHeaders(request.model); + const targetModel = this.resolveTargetModel(request.model); + const extraHeaders = this.createModelSpecificHeaders(request.model); this.logger.log( - `Received Anthropic request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`, + `Anthropic request received: model=${request.model}, mappedModel=${targetModel}, stream=${request.stream}`, ); // Retry loop @@ -61,13 +64,14 @@ export class ProxyService { for (let i = 0; i < maxRetries; i++) { if (i > 0) { const delay = calculateRetryDelay(i - 1); - this.logger.log(`Retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms (jittered)`); + this.logger.log(`Anthropic retry ${i + 1}/${maxRetries}, backoff=${delay}ms (jittered)`); await sleep(delay); } const token = await this.tokenManager.getNextToken({ sessionKey, excludeAccountIds: Array.from(attemptedAccountIds), + model: targetModel, }); if (!token) { throw new Error('No available accounts'); @@ -76,7 +80,13 @@ export class ProxyService { try { const projectId = token.token.project_id ?? ''; - const geminiBody = transformClaudeRequestIn(this.toClaudeRequest(request), projectId); + const requestUserAgent = await resolveRequestUserAgent(); + const geminiBody = transformClaudeRequestIn( + this.toClaudeRequest(request), + projectId, + requestUserAgent, + ); + this.applyInternalGenerationConstraints(geminiBody, targetModel, token.id); if (request.stream) { const stream = await this.geminiClient.streamGenerateInternal( @@ -101,7 +111,13 @@ export class ProxyService { `Anthropic request hit project context issue, retrying without project: ${error.message}`, ); try { - const fallbackBody = transformClaudeRequestIn(this.toClaudeRequest(request), ''); + const requestUserAgent = await resolveRequestUserAgent(); + const fallbackBody = transformClaudeRequestIn( + this.toClaudeRequest(request), + '', + requestUserAgent, + ); + this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id); if (request.stream) { const stream = await this.geminiClient.streamGenerateInternal( fallbackBody, @@ -126,17 +142,20 @@ export class ProxyService { if (error instanceof Error && this.isQuotaExhaustedError(error.message)) { this.logger.warn( - `Anthropic request hit quota exhaustion on mapped model, retrying with fallback model gemini-2.5-flash: ${error.message}`, + `Anthropic request hit quota exhaustion on mapped model, retrying with fallback model gemini-3-flash: ${error.message}`, ); try { const downgradedRequest: ClaudeRequest = { ...this.toClaudeRequest(request), - model: 'gemini-2.5-flash', + model: 'gemini-3-flash', }; + const requestUserAgent = await resolveRequestUserAgent(); const downgradedBody = transformClaudeRequestIn( downgradedRequest, token.token.project_id ?? '', + requestUserAgent, ); + this.applyInternalGenerationConstraints(downgradedBody, 'gemini-3-flash', token.id); if (request.stream) { const stream = await this.geminiClient.streamGenerateInternal( downgradedBody, @@ -164,17 +183,7 @@ export class ProxyService { } lastError = error; - if (error instanceof Error) { - this.logger.warn(`Anthropic Request failed: ${error.message}`); - const decision = this.classifyUpstreamError(error.message); - if (decision.retry) { - if (decision.markAsForbidden) { - this.tokenManager.markAsForbidden(token.id); - } else if (decision.markAsRateLimited) { - this.tokenManager.markAsRateLimited(token.id); - } - } - } + await this.applyUpstreamPenalty(token.id, targetModel, error); } } throw lastError || new Error('Request failed after retries'); @@ -182,7 +191,7 @@ export class ProxyService { private processAnthropicInternalStream( upstreamStream: NodeJS.ReadableStream, - model: string, + _model: string, ): Observable { return new Observable((subscriber) => { const decoder = new TextDecoder(); @@ -272,10 +281,10 @@ export class ProxyService { request: GeminiRequest, ): Promise { const normalizedModel = this.normalizeGeminiModel(model); - const targetModel = this.mapModel(normalizedModel); - const extraHeaders = this.buildModelSpecificHeaders(normalizedModel); + const targetModel = this.resolveTargetModel(normalizedModel); + const extraHeaders = this.createModelSpecificHeaders(normalizedModel); this.logger.log( - `Received Gemini request for model: ${normalizedModel} (Mapped: ${targetModel})`, + `Gemini generate request received: model=${normalizedModel}, mappedModel=${targetModel}`, ); let lastError: unknown = null; @@ -291,6 +300,7 @@ export class ProxyService { const token = await this.tokenManager.getNextToken({ excludeAccountIds: Array.from(attemptedAccountIds), + model: targetModel, }); if (!token) { throw new Error('No available accounts (all exhausted or rate limited)'); @@ -298,12 +308,15 @@ export class ProxyService { attemptedAccountIds.add(token.id); try { + const requestUserAgent = await resolveRequestUserAgent(); const internalBody = this.createGeminiInternalRequest( targetModel, request, token.token.project_id ?? '', 'generate-content', + requestUserAgent, ); + this.applyInternalGenerationConstraints(internalBody, targetModel, token.id); const response = await this.generateInternalWithStreamFallback( internalBody, @@ -319,12 +332,15 @@ export class ProxyService { `Gemini request hit project context issue, retrying without project: ${err.message}`, ); try { + const requestUserAgent = await resolveRequestUserAgent(); const fallbackBody = this.createGeminiInternalRequest( targetModel, request, '', 'generate-content', + requestUserAgent, ); + this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id); const response = await this.generateInternalWithStreamFallback( fallbackBody, token.token.access_token, @@ -339,17 +355,7 @@ export class ProxyService { lastError = err; } - if (lastError instanceof Error) { - this.logger.warn(`Gemini request failed with account ${token.email}: ${lastError.message}`); - const decision = this.classifyUpstreamError(lastError.message); - if (decision.retry) { - if (decision.markAsForbidden) { - this.tokenManager.markAsForbidden(token.id); - } else if (decision.markAsRateLimited) { - this.tokenManager.markAsRateLimited(token.id); - } - } - } + await this.applyUpstreamPenalty(token.id, targetModel, lastError); } } @@ -361,10 +367,10 @@ export class ProxyService { request: GeminiRequest, ): Promise> { const normalizedModel = this.normalizeGeminiModel(model); - const targetModel = this.mapModel(normalizedModel); - const extraHeaders = this.buildModelSpecificHeaders(normalizedModel); + const targetModel = this.resolveTargetModel(normalizedModel); + const extraHeaders = this.createModelSpecificHeaders(normalizedModel); this.logger.log( - `Received Gemini stream request for model: ${normalizedModel} (Mapped: ${targetModel})`, + `Gemini stream request received: model=${normalizedModel}, mappedModel=${targetModel}`, ); let lastError: unknown = null; @@ -380,6 +386,7 @@ export class ProxyService { const token = await this.tokenManager.getNextToken({ excludeAccountIds: Array.from(attemptedAccountIds), + model: targetModel, }); if (!token) { throw new Error('No available accounts (all exhausted or rate limited)'); @@ -387,12 +394,15 @@ export class ProxyService { attemptedAccountIds.add(token.id); try { + const requestUserAgent = await resolveRequestUserAgent(); const internalBody = this.createGeminiInternalRequest( targetModel, request, token.token.project_id ?? '', 'generate-content', + requestUserAgent, ); + this.applyInternalGenerationConstraints(internalBody, targetModel, token.id); const stream = await this.geminiClient.streamGenerateInternal( internalBody, @@ -407,12 +417,15 @@ export class ProxyService { `Gemini stream request hit project context issue, retrying without project: ${err.message}`, ); try { + const requestUserAgent = await resolveRequestUserAgent(); const fallbackBody = this.createGeminiInternalRequest( targetModel, request, '', 'generate-content', + requestUserAgent, ); + this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id); const stream = await this.geminiClient.streamGenerateInternal( fallbackBody, token.token.access_token, @@ -427,19 +440,7 @@ export class ProxyService { lastError = err; } - if (lastError instanceof Error) { - this.logger.warn( - `Gemini stream request failed with account ${token.email}: ${lastError.message}`, - ); - const decision = this.classifyUpstreamError(lastError.message); - if (decision.retry) { - if (decision.markAsForbidden) { - this.tokenManager.markAsForbidden(token.id); - } else if (decision.markAsRateLimited) { - this.tokenManager.markAsRateLimited(token.id); - } - } - } + await this.applyUpstreamPenalty(token.id, targetModel, lastError); } } @@ -475,11 +476,130 @@ export class ProxyService { return model.replace(/^models\//i, ''); } + private normalizeModelIdentifier(model: string): string { + return model.replace(/^models\//i, '').trim(); + } + + private resolveThinkingLevelBudget(level: string): number | undefined { + const normalized = level.trim().toUpperCase(); + if (normalized === 'NONE') { + return 0; + } + if (normalized === 'LOW') { + return 4096; + } + if (normalized === 'MEDIUM') { + return 8192; + } + if (normalized === 'HIGH') { + return 24576; + } + return undefined; + } + + private getModelOutputCap(accountId: string, model: string): number { + const normalizedModel = this.normalizeModelIdentifier(model); + const dynamicCap = this.tokenManager.getModelOutputLimitForAccount(accountId, normalizedModel); + if (typeof dynamicCap === 'number' && Number.isFinite(dynamicCap) && dynamicCap > 0) { + return Math.floor(dynamicCap); + } + return getMaxOutputTokens(normalizedModel); + } + + private getModelThinkingBudget(accountId: string, model: string): number { + const normalizedModel = this.normalizeModelIdentifier(model); + const dynamicBudget = this.tokenManager.getModelThinkingBudgetForAccount( + accountId, + normalizedModel, + ); + if (typeof dynamicBudget === 'number' && Number.isFinite(dynamicBudget) && dynamicBudget >= 0) { + return Math.floor(dynamicBudget); + } + return getThinkingBudget(normalizedModel); + } + + private applyInternalGenerationConstraints( + body: GeminiInternalRequest, + model: string, + accountId: string, + ): void { + const generationConfig = body.request.generationConfig; + if (!generationConfig) { + return; + } + + const outputCap = this.getModelOutputCap(accountId, model); + const thinkingBudgetCap = this.getModelThinkingBudget(accountId, model); + const normalizedModel = this.normalizeModelIdentifier(model).toLowerCase(); + const isClaudeModel = normalizedModel.includes('claude'); + const thinkingConfig = generationConfig.thinkingConfig as + | ({ thinkingLevel?: string; thinkingBudget?: number } & Record) + | undefined; + const adaptiveSentinel = + thinkingConfig && + (typeof thinkingConfig.thinkingLevel === 'string' || + thinkingConfig.thinkingBudget === -1 || + thinkingConfig.thinkingBudget === 32768); + + if (thinkingConfig) { + if (!isClaudeModel && typeof thinkingConfig.thinkingLevel === 'string') { + const converted = this.resolveThinkingLevelBudget(thinkingConfig.thinkingLevel); + if (converted !== undefined) { + thinkingConfig.thinkingBudget = converted; + } + delete thinkingConfig.thinkingLevel; + } + + if (typeof thinkingConfig.thinkingBudget === 'number' && thinkingConfig.thinkingBudget < 0) { + thinkingConfig.thinkingBudget = Math.min(thinkingBudgetCap, 24576); + } + + if ( + typeof thinkingConfig.thinkingBudget === 'number' && + Number.isFinite(thinkingConfig.thinkingBudget) + ) { + thinkingConfig.thinkingBudget = Math.min( + Math.floor(thinkingConfig.thinkingBudget), + Math.max(0, outputCap - 1), + thinkingBudgetCap, + ); + + if (adaptiveSentinel) { + if ( + generationConfig.maxOutputTokens === undefined || + generationConfig.maxOutputTokens < 131072 + ) { + generationConfig.maxOutputTokens = 131072; + } + } else if ( + generationConfig.maxOutputTokens === undefined || + generationConfig.maxOutputTokens <= thinkingConfig.thinkingBudget + ) { + const hasExplicitMax = generationConfig.maxOutputTokens !== undefined; + const overhead = hasExplicitMax ? 8192 : 32768; + const minRequired = Math.min(outputCap, thinkingConfig.thinkingBudget + overhead); + generationConfig.maxOutputTokens = minRequired; + } + } + } + + if ( + typeof generationConfig.maxOutputTokens === 'number' && + Number.isFinite(generationConfig.maxOutputTokens) + ) { + generationConfig.maxOutputTokens = Math.min( + Math.floor(generationConfig.maxOutputTokens), + outputCap, + ); + } + } + private createGeminiInternalRequest( model: string, request: GeminiRequest, projectId: string | undefined, requestType: string, + requestUserAgent: string, ): GeminiInternalRequest { const normalizedProjectId = projectId?.trim(); @@ -487,7 +607,7 @@ export class ProxyService { requestId: uuidv4(), request: this.toInternalGeminiRequest(request), model, - userAgent: process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent, + userAgent: requestUserAgent, requestType, }; @@ -523,9 +643,6 @@ export class ProxyService { if (usage.totalTokenCount !== undefined) { usageMetadata.totalTokenCount = usage.totalTokenCount; } - if (usage.thoughtsTokenCount !== undefined) { - usageMetadata.thoughtsTokenCount = usage.thoughtsTokenCount; - } if (usage.promptTokensDetails !== undefined) { usageMetadata.promptTokensDetails = usage.promptTokensDetails; } @@ -546,12 +663,12 @@ export class ProxyService { async handleChatCompletions( request: OpenAIChatRequest, ): Promise> { - const sessionKey = this.getSessionKeyFromOpenAI(request); + const sessionKey = this.extractOpenAISessionKey(request); - const targetModel = this.mapModel(request.model); - const extraHeaders = this.buildModelSpecificHeaders(request.model); + const targetModel = this.resolveTargetModel(request.model); + const extraHeaders = this.createModelSpecificHeaders(request.model); this.logger.log( - `Received request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`, + `OpenAI-compatible request received: model=${request.model}, mappedModel=${targetModel}, stream=${request.stream}`, ); // Retry loop for account selection @@ -562,7 +679,7 @@ export class ProxyService { for (let i = 0; i < maxRetries; i++) { if (i > 0) { const delay = calculateRetryDelay(i - 1); - this.logger.log(`Retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms (jittered)`); + this.logger.log(`OpenAI-compatible retry ${i + 1}/${maxRetries}, backoff=${delay}ms (jittered)`); await sleep(delay); } @@ -570,6 +687,7 @@ export class ProxyService { const token = await this.tokenManager.getNextToken({ sessionKey, excludeAccountIds: Array.from(attemptedAccountIds), + model: targetModel, }); if (!token) { throw new Error('No available accounts (all exhausted or rate limited)'); @@ -579,7 +697,9 @@ export class ProxyService { try { const claudeRequest = this.convertOpenAIToClaude(request); const projectId = token.token.project_id ?? ''; - const geminiBody = transformClaudeRequestIn(claudeRequest, projectId); + const requestUserAgent = await resolveRequestUserAgent(); + const geminiBody = transformClaudeRequestIn(claudeRequest, projectId, requestUserAgent); + this.applyInternalGenerationConstraints(geminiBody, targetModel, token.id); // Use v1internal API (same as Anthropic handler) if (request.stream) { @@ -593,7 +713,7 @@ export class ProxyService { return this.processStreamResponse(stream, request.model); } catch (streamError) { this.logger.warn( - `Stream request failed for ${request.model}, falling back to non-stream generation: ${ + `Stream path failed for model=${request.model}; falling back to non-stream generation: ${ streamError instanceof Error ? streamError.message : String(streamError) }`, ); @@ -605,7 +725,7 @@ export class ProxyService { extraHeaders, ); this.logger.log( - `Upstream response snippet (stream fallback): ${JSON.stringify(response).substring(0, 500)}`, + `Upstream response snippet after stream fallback: ${JSON.stringify(response).substring(0, 500)}`, ); const claudeResponse = transformResponse(response); const openaiResponse = this.convertClaudeToOpenAIResponse( @@ -622,11 +742,13 @@ export class ProxyService { extraHeaders, ); this.logger.log( - `Upstream response snippet: ${JSON.stringify(response).substring(0, 500)}`, + `Upstream response snippet (non-stream): ${JSON.stringify(response).substring(0, 500)}`, ); // Transform Gemini response to OpenAI format const claudeResponse = transformResponse(response); - this.logger.log(`Claude Response: ${JSON.stringify(claudeResponse).substring(0, 500)}`); + this.logger.log( + `Transformed Claude response snippet: ${JSON.stringify(claudeResponse).substring(0, 500)}`, + ); return this.convertClaudeToOpenAIResponse(claudeResponse, request.model); } } catch (err) { @@ -636,7 +758,9 @@ export class ProxyService { ); try { const claudeRequest = this.convertOpenAIToClaude(request); - const fallbackBody = transformClaudeRequestIn(claudeRequest, ''); + const requestUserAgent = await resolveRequestUserAgent(); + const fallbackBody = transformClaudeRequestIn(claudeRequest, '', requestUserAgent); + this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id); if (request.stream) { const stream = await this.geminiClient.streamGenerateInternal( fallbackBody, @@ -662,18 +786,7 @@ export class ProxyService { lastError = err; } - if (lastError instanceof Error) { - this.logger.warn(`Request failed with account ${token.email}: ${lastError.message}`); - - const decision = this.classifyUpstreamError(lastError.message); - if (decision.retry) { - if (decision.markAsForbidden) { - this.tokenManager.markAsForbidden(token.id); - } else if (decision.markAsRateLimited) { - this.tokenManager.markAsRateLimited(token.id); - } - } - } + await this.applyUpstreamPenalty(token.id, targetModel, lastError); } } throw lastError || new Error('Request failed after retries'); @@ -937,7 +1050,7 @@ export class ProxyService { hasSentDone = true; subscriber.complete(); } - } catch (e) { + } catch { // ignore parse errors } } @@ -1228,11 +1341,51 @@ export class ProxyService { return undefined; } - return tools.map((tool) => ({ - name: tool.function.name, - description: tool.function.description, - input_schema: tool.function.parameters || { type: 'object', properties: {} }, - })); + const result: NonNullable = []; + const searchToolTypes = new Set([ + 'web_search_20250305', + 'google_search', + 'google_search_retrieval', + 'builtin_web_search', + ]); + + for (const tool of tools) { + if (!tool) { + continue; + } + + const toolType = isString(tool.type) ? tool.type.toLowerCase() : ''; + const functionName = isString(tool.function?.name) ? tool.function.name : ''; + const normalizedFunctionName = functionName.toLowerCase(); + const isSearchTool = + searchToolTypes.has(toolType) || searchToolTypes.has(normalizedFunctionName); + + if (isSearchTool) { + result.push({ + name: functionName || 'builtin_web_search', + type: 'web_search_20250305', + input_schema: { + type: 'object', + properties: {}, + }, + }); + continue; + } + + if (!tool.function || !functionName) { + continue; + } + + const inputSchema = normalizeObjectJsonSchema(tool.function.parameters); + + result.push({ + name: functionName, + description: tool.function.description, + input_schema: inputSchema, + }); + } + + return result.length > 0 ? result : undefined; } private mapGeminiFinishReasonToOpenAIFinishReason(finishReason?: string): string | null { @@ -1362,7 +1515,7 @@ export class ProxyService { }; } - private mapModel(model: string): string { + private resolveTargetModel(model: string): string { const normalizedModel = model.replace(/^models\//i, '').trim(); const config = getServerConfig(); const configuredMapping = { @@ -1400,23 +1553,66 @@ export class ProxyService { } const routedModel = resolveModelRoute(normalizedModel, customExactMapping, {}, {}); - - return routedModel; + return normalizeGeminiModelAlias(routedModel); } - private classifyUpstreamError(errorMessage: string): { + private async applyUpstreamPenalty( + accountId: string, + model: string, + error: unknown, + ): Promise { + this.tokenManager.recordParityError(); + + if (error instanceof UpstreamRequestError) { + const status = error.status; + if (status === 401 || status === 403) { + this.tokenManager.markAsForbidden(accountId); + return; + } + + await this.tokenManager.markFromUpstreamError({ + accountIdOrEmail: accountId, + status, + retryAfter: error.headers?.retryAfter, + body: error.body, + model, + }); + return; + } + + if (!(error instanceof Error)) { + return; + } + + this.logger.warn(`Upstream request failed for account ${accountId}: ${error.message}`); + const penaltyDecision = this.classifyUpstreamFailure(error.message); + if (!penaltyDecision.retry) { + return; + } + + if (penaltyDecision.markAsForbidden) { + this.tokenManager.markAsForbidden(accountId); + return; + } + + if (penaltyDecision.markAsRateLimited) { + this.tokenManager.markAsRateLimited(accountId); + } + } + + private classifyUpstreamFailure(errorMessage: string): { retry: boolean; markAsForbidden: boolean; markAsRateLimited: boolean; } { - const msg = errorMessage.toLowerCase(); + const normalizedErrorMessage = errorMessage.toLowerCase(); const isForbidden = - msg.includes('401') || - msg.includes('unauthorized') || - msg.includes('invalid_grant') || - msg.includes('403') || - msg.includes('permission_denied') || - msg.includes('forbidden'); + normalizedErrorMessage.includes('401') || + normalizedErrorMessage.includes('unauthorized') || + normalizedErrorMessage.includes('invalid_grant') || + normalizedErrorMessage.includes('403') || + normalizedErrorMessage.includes('permission_denied') || + normalizedErrorMessage.includes('forbidden'); if (isForbidden) { return { @@ -1427,28 +1623,28 @@ export class ProxyService { } const isRateLimitedSignal = - msg.includes('429') || - msg.includes('resource_exhausted') || - msg.includes('quota') || - msg.includes('rate_limit') || - msg.includes('rate limit'); + normalizedErrorMessage.includes('429') || + normalizedErrorMessage.includes('resource_exhausted') || + normalizedErrorMessage.includes('quota') || + normalizedErrorMessage.includes('rate_limit') || + normalizedErrorMessage.includes('rate limit'); const shouldRetryByStatus = - msg.includes('408') || - msg.includes('429') || - msg.includes('500') || - msg.includes('502') || - msg.includes('503') || - msg.includes('504'); + normalizedErrorMessage.includes('408') || + normalizedErrorMessage.includes('429') || + normalizedErrorMessage.includes('500') || + normalizedErrorMessage.includes('502') || + normalizedErrorMessage.includes('503') || + normalizedErrorMessage.includes('504'); const shouldRetryByKeyword = - msg.includes('resource_exhausted') || - msg.includes('quota') || - msg.includes('rate_limit') || - msg.includes('timeout') || - msg.includes('socket hang up') || - msg.includes('empty response stream') || - msg.includes('connection reset'); + normalizedErrorMessage.includes('resource_exhausted') || + normalizedErrorMessage.includes('quota') || + normalizedErrorMessage.includes('rate_limit') || + normalizedErrorMessage.includes('timeout') || + normalizedErrorMessage.includes('socket hang up') || + normalizedErrorMessage.includes('empty response stream') || + normalizedErrorMessage.includes('connection reset'); if (shouldRetryByStatus || shouldRetryByKeyword) { return { @@ -1465,7 +1661,7 @@ export class ProxyService { }; } - private buildModelSpecificHeaders(model: string | undefined): Record { + private createModelSpecificHeaders(model: string | undefined): Record { if (!model) { return {}; } @@ -1510,28 +1706,28 @@ export class ProxyService { ); } - private getSessionKeyFromAnthropic(request: AnthropicChatRequest): string | undefined { + private extractAnthropicSessionKey(request: AnthropicChatRequest): string | undefined { const metadata = request.metadata; const sessionCandidate = metadata?.session_id ?? metadata?.sessionId ?? metadata?.user_id ?? metadata?.userId; - if (typeof sessionCandidate !== 'string' || sessionCandidate.trim() === '') { + if (!isString(sessionCandidate) || sessionCandidate.trim() === '') { return undefined; } return `anthropic:${sessionCandidate.trim()}`; } - private getSessionKeyFromOpenAI(request: OpenAIChatRequest): string | undefined { + private extractOpenAISessionKey(request: OpenAIChatRequest): string | undefined { const extra = request.extra; const sessionCandidate = extra?.session_id ?? extra?.sessionId ?? extra?.user_id ?? extra?.userId; - if (typeof sessionCandidate !== 'string' || sessionCandidate.trim() === '') { + if (!isString(sessionCandidate) || sessionCandidate.trim() === '') { return undefined; } return `openai:${sessionCandidate.trim()}`; } private isRecord(value: unknown): value is Record { - return value !== null && typeof value === 'object' && !Array.isArray(value); + return isPlainObject(value); } private isGeminiPart(value: unknown): value is InternalGeminiPart { diff --git a/src/server/modules/proxy/token-manager.service.ts b/src/server/modules/proxy/token-manager.service.ts index 4024e99..8302cff 100644 --- a/src/server/modules/proxy/token-manager.service.ts +++ b/src/server/modules/proxy/token-manager.service.ts @@ -1,7 +1,10 @@ import { Injectable, Logger, OnModuleInit } from '@nestjs/common'; import { CloudAccountRepo } from '../../../ipc/database/cloudHandler'; -import { CloudAccount } from '../../../types/cloudAccount'; +import { CloudAccount, CloudQuotaData } from '../../../types/cloudAccount'; import { GoogleAPIService } from '../../../services/GoogleAPIService'; +import { getServerConfig } from '../../server-config'; +import { RateLimitReason, RateLimitTracker } from './rate-limit-tracker'; +import { updateDynamicForwardingRules } from '../../../lib/antigravity/ModelMapping'; interface TokenData { email: string; @@ -14,8 +17,23 @@ interface TokenData { project_id?: string; session_id?: string; upstream_proxy_url?: string; + quota?: CloudQuotaData; + model_quotas: Record; + model_limits: Record; + model_reset_times: Record; + model_forwarding_rules: Record; } +type SchedulingMode = 'cache-first' | 'balance' | 'performance-first'; + +interface GetNextTokenOptions { + sessionKey?: string; + excludeAccountIds?: string[]; + model?: string; +} + +type TokenEntry = [string, TokenData]; + function normalizeProjectId(projectId: string | null | undefined): string | undefined { if (typeof projectId !== 'string') { return undefined; @@ -26,8 +44,6 @@ function normalizeProjectId(projectId: string | null | undefined): string | unde return undefined; } - // Legacy upgrades may carry malformed resource names (e.g. "projects/"). - // The internal gateway expects a plain project id, not a resource path. if (/^projects(?:\/.*)?$/i.test(trimmedProjectId)) { return undefined; } @@ -35,6 +51,20 @@ function normalizeProjectId(projectId: string | null | undefined): string | unde return trimmedProjectId; } +function normalizeModelId(modelId: string | null | undefined): string | undefined { + if (typeof modelId !== 'string') { + return undefined; + } + const normalized = modelId.replace(/^models\//i, '').trim(); + return normalized !== '' ? normalized : undefined; +} + +function delay(ms: number): Promise { + return new Promise((resolve) => { + setTimeout(resolve, ms); + }); +} + @Injectable() export class TokenManagerService implements OnModuleInit { private readonly logger = new Logger(TokenManagerService.name); @@ -43,14 +73,20 @@ export class TokenManagerService implements OnModuleInit { private readonly stickySessionTtlMs = 10 * 60 * 1000; private readonly rateLimitCooldownMs = 5 * 60 * 1000; private readonly forbiddenCooldownMs = 30 * 60 * 1000; - // In-memory cache of tokens with additional data + private readonly defaultBackoffSteps = [60, 300, 1800, 7200]; + private tokens: Map = new Map(); - // Cooldown map for rate-limited accounts - private cooldowns: Map = new Map(); + private accountCooldowns: Map = new Map(); private sessionBindings: Map = new Map(); + private rateLimitTracker = new RateLimitTracker(); + + private shadowComparisonCount = 0; + private shadowMismatchCount = 0; + private parityRequestCount = 0; + private parityErrorCount = 0; + private noGoBlocked = false; async onModuleInit() { - // Load accounts on module initialization await this.loadAccounts(); } @@ -59,24 +95,659 @@ export class TokenManagerService implements OnModuleInit { const accounts = await CloudAccountRepo.getAccounts(); let count = 0; + this.tokens.clear(); + for (const account of accounts) { - const tokenData = this.convertAccountToToken(account); + const tokenData = this.mapAccountToTokenData(account); if (tokenData) { this.tokens.set(account.id, tokenData); count++; } } - this.logger.log(`Initialized token pool with ${count} accounts`); + this.logger.log(`Token manager loaded ${count} cloud accounts into cache`); return count; } catch (e) { - this.logger.error('Unable to load accounts', e); + this.logger.error('Failed to load cloud accounts into token cache', e); return 0; } } - private convertAccountToToken(account: CloudAccount): TokenData | null { - if (!account.token) return null; + async reloadAllAccounts(): Promise { + const count = await this.loadAccounts(); + this.clearAllRateLimits(); + this.clearAllSessions(); + return count; + } + + clearAllSessions(): void { + this.sessionBindings.clear(); + } + + clearAllRateLimits(): void { + this.accountCooldowns.clear(); + this.rateLimitTracker.clearAll(); + } + + recordParityError(): void { + if (!this.isParitySchedulingEnabled()) { + return; + } + + this.parityErrorCount++; + const threshold = this.getNoGoErrorRateThreshold(); + const errorRate = this.parityErrorCount / Math.max(1, this.parityRequestCount); + if (errorRate > threshold) { + this.noGoBlocked = true; + this.logger.error( + `Parity no-go triggered by error threshold: rate=${errorRate.toFixed(4)}, requests=${this.parityRequestCount}, errors=${this.parityErrorCount}`, + ); + } + } + + setPreferredAccount(accountId?: string): void { + const config = getServerConfig(); + if (!config) { + return; + } + config.preferred_account_id = accountId ?? ''; + } + + isRateLimited(accountIdOrEmail: string, model?: string): boolean { + const accountId = this.resolveAccountId(accountIdOrEmail) ?? accountIdOrEmail; + const now = Date.now(); + const legacyCooldownUntil = this.accountCooldowns.get(accountId); + if (legacyCooldownUntil && legacyCooldownUntil > now) { + return true; + } + return this.rateLimitTracker.isRateLimited(accountId, model); + } + + markAsRateLimited(accountIdOrEmail: string) { + this.setAccountCooldown(accountIdOrEmail, 'rate limited', this.rateLimitCooldownMs); + } + + markAsForbidden(accountIdOrEmail: string) { + this.setAccountCooldown(accountIdOrEmail, 'forbidden', this.forbiddenCooldownMs); + } + + async markFromUpstreamError(params: { + accountIdOrEmail: string; + status?: number; + retryAfter?: string; + body?: string; + model?: string; + }): Promise { + const accountId = this.resolveAccountId(params.accountIdOrEmail) ?? params.accountIdOrEmail; + const normalizedModel = normalizeModelId(params.model); + const hasExplicitRetryWindow = + Boolean(params.retryAfter && params.retryAfter.trim() !== '') || + Boolean(params.body && params.body.includes('quotaResetDelay')); + + if (!hasExplicitRetryWindow && (params.status ?? 0) === 429) { + const reason = this.detectRateLimitReasonFromBody(params.body); + const shouldAttemptPreciseLockout = + reason === RateLimitReason.QuotaExhausted || reason === RateLimitReason.Unknown; + + if (!shouldAttemptPreciseLockout) { + const parsed = this.rateLimitTracker.trackFromUpstreamError({ + accountId, + status: params.status, + retryAfter: params.retryAfter, + body: params.body, + model: normalizedModel, + backoffSteps: this.getCircuitBreakerBackoffSteps(), + }); + + if (!parsed) { + return; + } + + if ( + parsed.reason !== RateLimitReason.QuotaExhausted || + !parsed.model || + parsed.model.trim() === '' + ) { + this.accountCooldowns.set(accountId, Date.now() + parsed.retryAfterSec * 1000); + } + return; + } + + const isLockedByRealtimeQuota = await this.refreshRealtimeQuotaAndSetPreciseLockout( + accountId, + reason, + normalizedModel, + ); + if (isLockedByRealtimeQuota) { + return; + } + + const isLockedByQuotaCache = this.setPreciseLockoutFromCachedQuota( + accountId, + reason, + normalizedModel, + ); + if (isLockedByQuotaCache) { + return; + } + } + + const parsed = this.rateLimitTracker.trackFromUpstreamError({ + accountId, + status: params.status, + retryAfter: params.retryAfter, + body: params.body, + model: normalizedModel, + backoffSteps: this.getCircuitBreakerBackoffSteps(), + }); + + if (!parsed) { + return; + } + + // Keep legacy account-level cooldown for reasons that affect the full account. + if ( + parsed.reason !== RateLimitReason.QuotaExhausted || + !parsed.model || + parsed.model.trim() === '' + ) { + this.accountCooldowns.set(accountId, Date.now() + parsed.retryAfterSec * 1000); + } + + this.logger.warn( + `Recorded upstream limit for account ${accountId}: reason=${parsed.reason}, wait=${parsed.retryAfterSec}s, model=${parsed.model ?? 'n/a'}`, + ); + } + + async getNextToken(options?: GetNextTokenOptions): Promise { + try { + if (this.tokens.size === 0) { + await this.loadAccounts(); + } + if (this.tokens.size === 0) { + return null; + } + + const now = Date.now(); + const nowSeconds = Math.floor(now / 1000); + const sessionKey = options?.sessionKey?.trim(); + const model = options?.model; + const excludedAccountIds = new Set(options?.excludeAccountIds ?? []); + + this.clearExpiredSessionBindings(now); + this.rateLimitTracker.cleanupExpired(); + + const fullAccountPool = Array.from(this.tokens.entries()); + const filteredAccountPool = fullAccountPool.filter( + ([accountId]) => !excludedAccountIds.has(accountId), + ); + const candidateAccountPool = + filteredAccountPool.length > 0 ? filteredAccountPool : fullAccountPool; + + if (filteredAccountPool.length === 0 && excludedAccountIds.size > 0) { + this.logger.warn( + 'Exclusion filter removed all accounts; retrying with the full account pool', + ); + } + + if (candidateAccountPool.length === 0) { + this.logger.warn('No eligible account found after exclusion filtering'); + return null; + } + + if (this.shouldExecuteShadowComparison()) { + this.executeShadowComparison(candidateAccountPool, sessionKey, model); + } + + const selectedTokenEntry = this.isParitySchedulingEnabled() + ? await this.selectParityTokenCandidate(candidateAccountPool, sessionKey, model, now) + : this.selectLegacyTokenCandidate(candidateAccountPool, sessionKey, now); + + if (!selectedTokenEntry) { + return null; + } + + if (this.isParitySchedulingEnabled()) { + this.parityRequestCount++; + } + + const [accountId, tokenData] = selectedTokenEntry; + return this.finalizeSelectedToken(accountId, tokenData, nowSeconds, sessionKey); + } catch (error) { + this.logger.error('Failed to select the next account token', error); + return null; + } + } + + private shouldExecuteShadowComparison(): boolean { + const config = getServerConfig(); + return ( + Boolean(config?.parity_shadow_enabled) && + !this.isParitySchedulingEnabled() && + !this.noGoBlocked + ); + } + + private isParitySchedulingEnabled(): boolean { + const config = getServerConfig(); + if (!config) { + return false; + } + if (config.parity_kill_switch) { + return false; + } + if (this.noGoBlocked) { + return false; + } + return Boolean(config.parity_enabled); + } + + private getSchedulingMode(): SchedulingMode { + const config = getServerConfig(); + const mode = (config?.scheduling_mode ?? 'balance').toLowerCase(); + if (mode === 'cache-first' || mode === 'performance-first' || mode === 'balance') { + return mode; + } + return 'balance'; + } + + private getMaxWaitDurationMs(): number { + const config = getServerConfig(); + const seconds = config?.max_wait_seconds ?? 60; + return Math.max(0, seconds) * 1000; + } + + private getCircuitBreakerBackoffSteps(): number[] { + const config = getServerConfig(); + const configured = config?.circuit_breaker_backoff_steps ?? this.defaultBackoffSteps; + const normalized = configured + .map((value) => Number(value)) + .filter((value) => Number.isFinite(value) && value > 0) + .map((value) => Math.ceil(value)); + if (normalized.length > 0) { + return normalized; + } + return this.defaultBackoffSteps; + } + + private getPreferredAccountId(): string | undefined { + const config = getServerConfig(); + const preferred = config?.preferred_account_id?.trim(); + return preferred ? preferred : undefined; + } + + private getNoGoMismatchRateThreshold(): number { + const config = getServerConfig(); + const threshold = config?.parity_no_go_mismatch_rate ?? 0.15; + if (!Number.isFinite(threshold)) { + return 0.15; + } + return Math.min(1, Math.max(0, threshold)); + } + + private getNoGoErrorRateThreshold(): number { + const config = getServerConfig(); + const threshold = config?.parity_no_go_error_rate ?? 0.4; + if (!Number.isFinite(threshold)) { + return 0.4; + } + return Math.min(1, Math.max(0, threshold)); + } + + private collectEligibleTokens( + allTokens: TokenEntry[], + model: string | undefined, + now: number, + ): TokenEntry[] { + return allTokens.filter(([accountId]) => { + const cooldownUntil = this.accountCooldowns.get(accountId); + if (cooldownUntil && cooldownUntil > now) { + return false; + } + return !this.rateLimitTracker.isRateLimited(accountId, model); + }); + } + + private getValidSessionBinding( + sessionKey: string | undefined, + now: number, + ): { accountId: string; expiresAt: number } | null { + if (!sessionKey) { + return null; + } + const stickyBinding = this.sessionBindings.get(sessionKey); + if (!stickyBinding || stickyBinding.expiresAt <= now) { + return null; + } + return stickyBinding; + } + + private findStickySessionToken( + candidates: TokenEntry[], + sessionKey: string | undefined, + now: number, + ): TokenEntry | null { + const stickyBinding = this.getValidSessionBinding(sessionKey, now); + if (!stickyBinding) { + return null; + } + + return candidates.find(([accountId]) => accountId === stickyBinding.accountId) ?? null; + } + + private pickRoundRobinEntry(candidates: TokenEntry[]): TokenEntry | null { + if (candidates.length === 0) { + return null; + } + const picked = candidates[this.currentIndex % candidates.length]; + this.currentIndex++; + return picked; + } + + private peekRoundRobinCandidateAccountId(candidates: TokenEntry[]): string | null { + if (candidates.length === 0) { + return null; + } + return candidates[this.currentIndex % candidates.length][0]; + } + + private selectLegacyTokenCandidate( + allTokens: TokenEntry[], + sessionKey: string | undefined, + now: number, + ): TokenEntry | null { + const availableByCooldown = allTokens.filter(([accountId]) => { + const cooldownUntil = this.accountCooldowns.get(accountId); + return !cooldownUntil || cooldownUntil <= now; + }); + + const candidateAccountPool = availableByCooldown.length > 0 ? availableByCooldown : allTokens; + if (candidateAccountPool.length === 0) { + return null; + } + + if (availableByCooldown.length === 0) { + this.logger.warn( + 'All accounts are cooling down; temporarily bypassing cooldown gate to preserve availability', + ); + } + + const stickyToken = this.findStickySessionToken(candidateAccountPool, sessionKey, now); + if (stickyToken) { + return stickyToken; + } + + return this.pickRoundRobinEntry(candidateAccountPool); + } + + private async selectParityTokenCandidate( + allTokens: TokenEntry[], + sessionKey: string | undefined, + model: string | undefined, + now: number, + ): Promise { + const mode = this.getSchedulingMode(); + const availableTokens = this.collectEligibleTokens(allTokens, model, now); + if (availableTokens.length === 0) { + return null; + } + + const preferredAccountId = this.getPreferredAccountId(); + if (preferredAccountId) { + const preferred = availableTokens.find(([accountId]) => accountId === preferredAccountId); + if (preferred) { + return preferred; + } + } + + const stickyToken = this.findStickySessionToken(availableTokens, sessionKey, now); + if (stickyToken) { + return stickyToken; + } + + const stickyBinding = this.getValidSessionBinding(sessionKey, now); + if (stickyBinding && mode === 'cache-first') { + const waitSec = this.rateLimitTracker.getRemainingWaitSeconds( + stickyBinding.accountId, + model, + ); + const waitMs = waitSec * 1000; + const maxWaitMs = this.getMaxWaitDurationMs(); + if (waitMs > 0 && waitMs <= maxWaitMs) { + await delay(waitMs); + const refreshedAvailable = this.collectEligibleTokens(allTokens, model, Date.now()); + const stickyAfterWait = + refreshedAvailable.find(([accountId]) => accountId === stickyBinding.accountId) ?? null; + if (stickyAfterWait) { + return stickyAfterWait; + } + if (refreshedAvailable.length > 0) { + return this.pickRoundRobinEntry(refreshedAvailable); + } + } + } + + return this.pickRoundRobinEntry(availableTokens); + } + + private predictLegacyAccountCandidateId( + allTokens: TokenEntry[], + sessionKey: string | undefined, + now: number, + ): string | null { + const availableByCooldown = allTokens.filter(([accountId]) => { + const cooldownUntil = this.accountCooldowns.get(accountId); + return !cooldownUntil || cooldownUntil <= now; + }); + const candidateAccountPool = availableByCooldown.length > 0 ? availableByCooldown : allTokens; + if (candidateAccountPool.length === 0) { + return null; + } + + const stickyToken = this.findStickySessionToken(candidateAccountPool, sessionKey, now); + if (stickyToken) { + return stickyToken[0]; + } + + return this.peekRoundRobinCandidateAccountId(candidateAccountPool); + } + + private predictParityAccountCandidateId( + allTokens: TokenEntry[], + sessionKey: string | undefined, + model: string | undefined, + now: number, + ): string | null { + const availableTokens = this.collectEligibleTokens(allTokens, model, now); + if (availableTokens.length === 0) { + return null; + } + + const preferredAccountId = this.getPreferredAccountId(); + if (preferredAccountId) { + const preferred = availableTokens.find(([accountId]) => accountId === preferredAccountId); + if (preferred) { + return preferred[0]; + } + } + + const stickyToken = this.findStickySessionToken(availableTokens, sessionKey, now); + if (stickyToken) { + return stickyToken[0]; + } + + return this.peekRoundRobinCandidateAccountId(availableTokens); + } + + private executeShadowComparison( + allTokens: TokenEntry[], + sessionKey: string | undefined, + model: string | undefined, + ): void { + const now = Date.now(); + const legacyAccountId = this.predictLegacyAccountCandidateId(allTokens, sessionKey, now); + const parityAccountId = this.predictParityAccountCandidateId( + allTokens, + sessionKey, + model, + now, + ); + this.shadowComparisonCount++; + + if (legacyAccountId !== parityAccountId) { + this.shadowMismatchCount++; + this.logger.warn( + `Parity shadow mismatch detected: legacy=${legacyAccountId ?? 'n/a'}, parity=${parityAccountId ?? 'n/a'}`, + ); + } + + const mismatchRate = this.shadowMismatchCount / Math.max(1, this.shadowComparisonCount); + if (mismatchRate > this.getNoGoMismatchRateThreshold()) { + this.noGoBlocked = true; + this.logger.error( + `Parity no-go triggered by mismatch threshold: rate=${mismatchRate.toFixed(4)}, comparisons=${this.shadowComparisonCount}`, + ); + } + } + + private detectRateLimitReasonFromBody(body: string | undefined): RateLimitReason { + const lowerBody = (body ?? '').toLowerCase(); + if (lowerBody.includes('model_capacity')) { + return RateLimitReason.ModelCapacityExhausted; + } + if (lowerBody.includes('exhausted') || lowerBody.includes('quota')) { + return RateLimitReason.QuotaExhausted; + } + if ( + lowerBody.includes('per minute') || + lowerBody.includes('rate limit') || + lowerBody.includes('rate_limit') + ) { + return RateLimitReason.RateLimitExceeded; + } + return RateLimitReason.Unknown; + } + + private extractQuotaSnapshot(quota: CloudQuotaData | undefined): { + modelQuotas: Record; + modelLimits: Record; + modelResetTimes: Record; + modelForwardingRules: Record; + } { + const modelQuotas: Record = {}; + const modelLimits: Record = {}; + const modelResetTimes: Record = {}; + const modelForwardingRules: Record = {}; + + for (const [modelName, modelInfo] of Object.entries(quota?.models ?? {})) { + const normalizedModel = normalizeModelId(modelName); + if (!normalizedModel) { + continue; + } + + if (Number.isFinite(modelInfo.percentage)) { + modelQuotas[normalizedModel] = Math.floor(modelInfo.percentage); + } + + const limitCandidate = modelInfo.max_output_tokens ?? modelInfo.max_tokens; + if ( + typeof limitCandidate === 'number' && + Number.isFinite(limitCandidate) && + limitCandidate > 0 + ) { + modelLimits[normalizedModel] = Math.floor(limitCandidate); + } + + if (typeof modelInfo.resetTime === 'string' && modelInfo.resetTime.trim() !== '') { + modelResetTimes[normalizedModel] = modelInfo.resetTime; + } + } + + for (const [oldModel, newModel] of Object.entries(quota?.model_forwarding_rules ?? {})) { + const normalizedOld = normalizeModelId(oldModel); + const normalizedNew = normalizeModelId(newModel); + if (!normalizedOld || !normalizedNew) { + continue; + } + modelForwardingRules[normalizedOld] = normalizedNew; + updateDynamicForwardingRules(normalizedOld, normalizedNew); + } + + return { + modelQuotas, + modelLimits, + modelResetTimes, + modelForwardingRules, + }; + } + + private findEarliestQuotaResetTime(modelResetTimes: Record): string | null { + const validTimes = Object.values(modelResetTimes).filter((value) => value.trim() !== ''); + if (validTimes.length === 0) { + return null; + } + return [...validTimes].sort()[0]; + } + + private setPreciseLockoutFromCachedQuota( + accountId: string, + reason: RateLimitReason, + model?: string, + ): boolean { + const tokenData = this.tokens.get(accountId); + if (!tokenData) { + return false; + } + + const resetTime = this.findEarliestQuotaResetTime(tokenData.model_reset_times); + if (!resetTime) { + return false; + } + + return this.rateLimitTracker.setLockoutUntilIso(accountId, resetTime, reason, model); + } + + private async refreshRealtimeQuotaAndSetPreciseLockout( + accountId: string, + reason: RateLimitReason, + model?: string, + ): Promise { + const tokenData = this.tokens.get(accountId); + if (!tokenData) { + return false; + } + + try { + const latestQuota = await GoogleAPIService.fetchQuota(tokenData.access_token); + const extractedState = this.extractQuotaSnapshot(latestQuota); + + tokenData.quota = latestQuota; + tokenData.model_quotas = extractedState.modelQuotas; + tokenData.model_limits = extractedState.modelLimits; + tokenData.model_reset_times = extractedState.modelResetTimes; + tokenData.model_forwarding_rules = extractedState.modelForwardingRules; + this.tokens.set(accountId, tokenData); + + await CloudAccountRepo.updateQuota(accountId, latestQuota); + + const resetTime = this.findEarliestQuotaResetTime(extractedState.modelResetTimes); + if (!resetTime) { + return false; + } + return this.rateLimitTracker.setLockoutUntilIso(accountId, resetTime, reason, model); + } catch (error) { + this.logger.warn(`Failed to refresh realtime quota for account ${accountId}`, error); + return false; + } + } + + private mapAccountToTokenData(account: CloudAccount): TokenData | null { + if (!account.token) { + return null; + } + + const quota = account.quota; + const extractedState = this.extractQuotaSnapshot(quota); return { account_id: account.id, @@ -89,6 +760,11 @@ export class TokenManagerService implements OnModuleInit { project_id: account.token.project_id || undefined, session_id: account.token.session_id || this.generateSessionId(), upstream_proxy_url: account.token.upstream_proxy_url || undefined, + quota, + model_quotas: extractedState.modelQuotas, + model_limits: extractedState.modelLimits, + model_reset_times: extractedState.modelResetTimes, + model_forwarding_rules: extractedState.modelForwardingRules, }; } @@ -100,84 +776,6 @@ export class TokenManagerService implements OnModuleInit { return (-(min + rand)).toString(); } - async getNextToken(options?: { - sessionKey?: string; - excludeAccountIds?: string[]; - }): Promise { - try { - // Reload if empty - if (this.tokens.size === 0) { - await this.loadAccounts(); - } - if (this.tokens.size === 0) return null; - - const now = Date.now(); - const nowSeconds = Math.floor(now / 1000); - const sessionKey = options?.sessionKey?.trim(); - const excludedAccountIds = new Set(options?.excludeAccountIds ?? []); - - this.clearExpiredSessionBindings(now); - - const fullTokenPool = Array.from(this.tokens.entries()); - const excludedFilteredTokens = fullTokenPool.filter( - ([accountId]) => !excludedAccountIds.has(accountId), - ); - const allTokens = excludedFilteredTokens.length > 0 ? excludedFilteredTokens : fullTokenPool; - if (excludedFilteredTokens.length === 0 && excludedAccountIds.size > 0) { - this.logger.warn( - 'Retry exclusions removed all accounts; temporarily reusing excluded accounts', - ); - } - - // Filter out accounts in cooldown - const validTokens = allTokens.filter(([accountId]) => { - if (excludedAccountIds.has(accountId)) { - return false; - } - - const cooldownUntil = this.cooldowns.get(accountId); - return !cooldownUntil || cooldownUntil <= now; - }); - - const candidateTokens = validTokens.length > 0 ? validTokens : allTokens; - if (candidateTokens.length === 0) { - this.logger.warn('No account available after applying exclusions'); - return null; - } - if (validTokens.length === 0) { - this.logger.warn( - 'All accounts are currently in cooldown; temporarily bypassing cooldown to keep service available', - ); - } - - if (sessionKey) { - const stickyBinding = this.sessionBindings.get(sessionKey); - if (stickyBinding && stickyBinding.expiresAt > now) { - const stickyMatch = candidateTokens.find( - ([accountId]) => accountId === stickyBinding.accountId, - ); - if (stickyMatch) { - const [stickyAccountId, stickyTokenData] = stickyMatch; - return this.finalizeSelectedToken( - stickyAccountId, - stickyTokenData, - nowSeconds, - sessionKey, - ); - } - } - } - - // Round robin selection - const [accountId, tokenData] = candidateTokens[this.currentIndex % candidateTokens.length]; - this.currentIndex++; - return this.finalizeSelectedToken(accountId, tokenData, nowSeconds, sessionKey); - } catch (error) { - this.logger.error('Unable to acquire token', error); - return null; - } - } - private async finalizeSelectedToken( accountId: string, tokenData: TokenData, @@ -187,24 +785,18 @@ export class TokenManagerService implements OnModuleInit { try { let effectiveProjectId: string | undefined; - // Check if token needs refresh (expires in < 5 minutes) if (nowSeconds >= tokenData.expiry_timestamp - 300) { - this.logger.log(`Token for ${tokenData.email} is close to expiry; refreshing...`); + this.logger.log(`Access token near expiry for ${tokenData.email}; refreshing`); try { const newTokens = await GoogleAPIService.refreshAccessToken(tokenData.refresh_token); - - // Update token data tokenData.access_token = newTokens.access_token; tokenData.expires_in = newTokens.expires_in; tokenData.expiry_timestamp = nowSeconds + newTokens.expires_in; - - // Save to DB - await this.saveTokenState(accountId, tokenData); + await this.persistTokenState(accountId, tokenData); this.tokens.set(accountId, tokenData); - - this.logger.log(`Refreshed token for ${tokenData.email}`); + this.logger.log(`Access token refreshed for ${tokenData.email}`); } catch (e) { - this.logger.error(`Unable to refresh token for ${tokenData.email}`, e); + this.logger.error(`Failed to refresh access token for ${tokenData.email}`, e); } } @@ -217,32 +809,32 @@ export class TokenManagerService implements OnModuleInit { try { const fetchedProjectId = await GoogleAPIService.fetchProjectId(tokenData.access_token); const normalizedProjectId = normalizeProjectId(fetchedProjectId); - if (normalizedProjectId) { tokenData.project_id = normalizedProjectId; effectiveProjectId = normalizedProjectId; - await this.saveTokenState(accountId, tokenData); + await this.persistTokenState(accountId, tokenData); this.tokens.set(accountId, tokenData); this.logger.log(`Resolved project ID for ${tokenData.email}: ${normalizedProjectId}`); } else { this.logger.warn( - `Project ID is unavailable for ${tokenData.email}; request will continue without project context`, + `Project ID unavailable for ${tokenData.email}; continuing without project context`, ); } } catch (error) { - this.logger.warn(`Failed to resolve project ID for ${tokenData.email}`, error); + this.logger.warn(`Unable to resolve project ID for ${tokenData.email}`, error); } } if (!effectiveProjectId) { - const fallbackProjectId = this.getFallbackProjectId(); + const fallbackProjectId = this.resolveFallbackProjectId(); effectiveProjectId = fallbackProjectId; this.logger.warn( - `Using fallback project ID for ${tokenData.email} without persistence: ${fallbackProjectId}`, + `Using non-persistent fallback project ID for ${tokenData.email}: ${fallbackProjectId}`, ); } - this.logger.log(`Using account: ${tokenData.email}`); + this.rateLimitTracker.markSuccess(accountId); + if (sessionKey) { this.sessionBindings.set(sessionKey, { accountId, @@ -250,9 +842,8 @@ export class TokenManagerService implements OnModuleInit { }); } - // Return in CloudAccount format for compatibility const timestamp = Date.now(); - const cloudAccount: CloudAccount = { + return { id: accountId, provider: 'google', email: tokenData.email, @@ -269,21 +860,12 @@ export class TokenManagerService implements OnModuleInit { created_at: timestamp, last_used: timestamp, }; - return cloudAccount; } catch (error) { - this.logger.error('Unable to finalize selected token', error); + this.logger.error('Failed to finalize selected account token', error); return null; } } - markAsRateLimited(accountIdOrEmail: string) { - this.markInCooldown(accountIdOrEmail, 'rate limited', this.rateLimitCooldownMs); - } - - markAsForbidden(accountIdOrEmail: string) { - this.markInCooldown(accountIdOrEmail, 'forbidden', this.forbiddenCooldownMs); - } - private resolveAccountId(accountIdOrEmail: string): string | null { if (this.tokens.has(accountIdOrEmail)) { return accountIdOrEmail; @@ -306,7 +888,7 @@ export class TokenManagerService implements OnModuleInit { } } - private markInCooldown( + private setAccountCooldown( accountIdOrEmail: string, reason: 'rate limited' | 'forbidden', durationMs: number, @@ -314,13 +896,13 @@ export class TokenManagerService implements OnModuleInit { const accountId = this.resolveAccountId(accountIdOrEmail) ?? accountIdOrEmail; const cooldownUntil = Date.now() + durationMs; - this.cooldowns.set(accountId, cooldownUntil); + this.accountCooldowns.set(accountId, cooldownUntil); this.logger.warn( - `Account ${accountIdOrEmail} (cooldown key: ${accountId}) set to ${reason} until ${new Date(cooldownUntil).toISOString()}`, + `Applied ${reason} cooldown: source=${accountIdOrEmail}, accountId=${accountId}, until=${new Date(cooldownUntil).toISOString()}`, ); } - private async saveTokenState(accountId: string, tokenData: TokenData) { + private async persistTokenState(accountId: string, tokenData: TokenData) { try { const acc = await CloudAccountRepo.getAccount(accountId); if (acc && acc.token) { @@ -336,18 +918,53 @@ export class TokenManagerService implements OnModuleInit { await CloudAccountRepo.updateToken(accountId, newToken); } } catch (e) { - this.logger.error('Unable to persist token state to DB', e); + this.logger.error('Failed to persist token state to database', e); } } - /** - * Get the number of loaded accounts (for status) - */ getAccountCount(): number { return this.tokens.size; } - private getFallbackProjectId(): string { + getAllCollectedModels(): Set { + const allModels = new Set(); + for (const tokenData of this.tokens.values()) { + for (const modelId of Object.keys(tokenData.model_quotas)) { + allModels.add(modelId); + } + } + return allModels; + } + + getModelOutputLimitForAccount(accountId: string, modelName: string): number | undefined { + const tokenData = this.tokens.get(accountId); + const normalizedModel = normalizeModelId(modelName); + if (!tokenData || !normalizedModel) { + return undefined; + } + return tokenData.model_limits[normalizedModel]; + } + + getModelThinkingBudgetForAccount(accountId: string, modelName: string): number | undefined { + const tokenData = this.tokens.get(accountId); + const normalizedModel = normalizeModelId(modelName); + if (!tokenData || !normalizedModel) { + return undefined; + } + + for (const [quotaModelName, modelInfo] of Object.entries(tokenData.quota?.models ?? {})) { + if (normalizeModelId(quotaModelName) !== normalizedModel) { + continue; + } + const budget = modelInfo?.thinking_budget; + if (typeof budget === 'number' && Number.isFinite(budget) && budget >= 0) { + return Math.floor(budget); + } + } + return undefined; + } + + private resolveFallbackProjectId(): string { const fromEnv = process.env.PROXY_FALLBACK_PROJECT_ID?.trim(); const normalizedFromEnv = normalizeProjectId(fromEnv); if (normalizedFromEnv) { diff --git a/src/services/GoogleAPIService.ts b/src/services/GoogleAPIService.ts index 8ca742f..3d656bb 100644 --- a/src/services/GoogleAPIService.ts +++ b/src/services/GoogleAPIService.ts @@ -1,6 +1,11 @@ import { ConfigManager } from '../ipc/config/manager'; import { ProxyAgent } from 'undici'; import { logger } from '../utils/logger'; +import { + buildUserAgent, + FALLBACK_VERSION, + resolveLocalInstalledVersion, +} from '@/server/modules/proxy/request-user-agent'; // --- Constants & Config --- @@ -52,13 +57,23 @@ export interface UserInfo { } export interface QuotaData { - models: Record< - string, - { - percentage: number; - resetTime: string; - } - >; + models: Record; + model_forwarding_rules?: Record; + subscription_tier?: string; + is_forbidden?: boolean; +} + +export interface ModelQuotaInfo { + percentage: number; + resetTime: string; + display_name?: string; + supports_images?: boolean; + supports_thinking?: boolean; + thinking_budget?: number; + recommended?: boolean; + max_tokens?: number; + max_output_tokens?: number; + supported_mime_types?: Record; } // Internal types for API parsing @@ -67,10 +82,138 @@ interface ModelInfoRaw { remainingFraction?: number; resetTime?: string; }; + displayName?: string; + supportsImages?: boolean; + supportsThinking?: boolean; + thinkingBudget?: number; + recommended?: boolean; + maxTokens?: number; + maxOutputTokens?: number; + supportedMimeTypes?: Record; +} + +interface DeprecatedModelInfoRaw { + newModelId?: string; +} + +interface IneligibleTierRaw { + reasonCode?: string; +} + +interface TierRaw { + is_default?: boolean; + id?: string; + quotaTier?: string; + name?: string; + slug?: string; } interface LoadProjectResponse { cloudaicompanionProject?: string; + currentTier?: TierRaw; + paidTier?: TierRaw; + allowedTiers?: TierRaw[]; + ineligibleTiers?: IneligibleTierRaw[]; +} + +interface FetchModelsResponse { + models?: Record; + deprecatedModelIds?: Record; +} + +interface ProjectContext { + projectId?: string; + subscriptionTier?: string; +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => { + setTimeout(resolve, ms); + }); +} + +function buildInternalApiHeaders(accessToken: string): Record { + const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION; + return { + Authorization: `Bearer ${accessToken}`, + 'User-Agent': buildUserAgent(discoveryVersion), + 'Content-Type': 'application/json', + }; +} + +function resolveSubscriptionTier(payload: LoadProjectResponse): string | undefined { + const paidTier = payload.paidTier; + if (paidTier?.name && paidTier.name.trim() !== '') { + return paidTier.name; + } + if (paidTier?.id && paidTier.id.trim() !== '') { + return paidTier.id; + } + + const ineligible = Array.isArray(payload.ineligibleTiers) && payload.ineligibleTiers.length > 0; + if (!ineligible) { + const currentTier = payload.currentTier; + if (currentTier?.name && currentTier.name.trim() !== '') { + return currentTier.name; + } + if (currentTier?.id && currentTier.id.trim() !== '') { + return currentTier.id; + } + } + + if (Array.isArray(payload.allowedTiers)) { + const preferredAllowedTier = + payload.allowedTiers.find((tier) => tier.is_default === true) ?? payload.allowedTiers[0]; + if (preferredAllowedTier?.name && preferredAllowedTier.name.trim() !== '') { + return ineligible ? `${preferredAllowedTier.name} (Restricted)` : preferredAllowedTier.name; + } + if (preferredAllowedTier?.id && preferredAllowedTier.id.trim() !== '') { + return ineligible ? `${preferredAllowedTier.id} (Restricted)` : preferredAllowedTier.id; + } + } + + return undefined; +} + +function isTrackedModel(modelName: string): boolean { + return /^(gemini|claude|gpt|image|imagen)/i.test(modelName); +} + +function toModelQuotaInfo(modelName: string, info: ModelInfoRaw): ModelQuotaInfo | null { + if (!isTrackedModel(modelName) || !info.quotaInfo) { + return null; + } + + const fraction = info.quotaInfo.remainingFraction ?? 0; + return { + percentage: Math.floor(fraction * 100), + resetTime: info.quotaInfo.resetTime || '', + display_name: info.displayName, + supports_images: info.supportsImages, + supports_thinking: info.supportsThinking, + thinking_budget: info.thinkingBudget, + recommended: info.recommended, + max_tokens: info.maxTokens, + max_output_tokens: info.maxOutputTokens, + supported_mime_types: info.supportedMimeTypes, + }; +} + +function toModelForwardingRules( + deprecatedModelIds: FetchModelsResponse['deprecatedModelIds'], +): Record | undefined { + if (!deprecatedModelIds || Object.keys(deprecatedModelIds).length === 0) { + return undefined; + } + + const forwardingRules: Record = {}; + for (const [oldModelId, deprecatedInfo] of Object.entries(deprecatedModelIds)) { + if (typeof deprecatedInfo.newModelId === 'string' && deprecatedInfo.newModelId !== '') { + forwardingRules[oldModelId] = deprecatedInfo.newModelId; + } + } + + return Object.keys(forwardingRules).length > 0 ? forwardingRules : undefined; } // --- Service Implementation --- @@ -218,70 +361,71 @@ export class GoogleAPIService { return response.json() as Promise; } - /** - * Core logic: Fetches the internal project ID needed for quota checks. - */ - public static async fetchProjectId(accessToken: string): Promise { + public static async fetchProjectContext(accessToken: string): Promise { const body = { metadata: { ideType: 'ANTIGRAVITY' }, }; - // Simple retry logic (2 attempts) + let projectId: string | undefined; + let subscriptionTier: string | undefined; + for (let i = 0; i < 2; i++) { try { const response = await fetch(URLS.LOAD_PROJECT, { method: 'POST', - headers: { - Authorization: `Bearer ${accessToken}`, - 'User-Agent': USER_AGENT, - 'Content-Type': 'application/json', - }, + headers: buildInternalApiHeaders(accessToken), body: JSON.stringify(body), + signal: createTimeoutSignal(REQUEST_TIMEOUT_MS), ...this.getFetchOptions(), }); if (response.ok) { const data = (await response.json()) as LoadProjectResponse; - if (data.cloudaicompanionProject) { - return data.cloudaicompanionProject; + if (typeof data.cloudaicompanionProject === 'string') { + projectId = data.cloudaicompanionProject; } + subscriptionTier = resolveSubscriptionTier(data); + break; } - } catch (e) { - logger.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, e); - await new Promise((r) => setTimeout(r, 500)); // Sleep 500ms + } catch (error) { + logger.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, error); + await sleep(500); } } - return null; + + return { + projectId, + subscriptionTier, + }; + } + + public static async fetchProjectId(accessToken: string): Promise { + const context = await this.fetchProjectContext(accessToken); + return context.projectId ?? null; } /** * Core logic: Fetches detailed model quota information. */ static async fetchQuota(accessToken: string): Promise { - // 1. Get Project ID (Critical) - const projectId = await this.fetchProjectId(accessToken); + const { projectId, subscriptionTier } = await this.fetchProjectContext(accessToken); - // 2. Build Payload const payload: Record = {}; if (projectId) { payload['project'] = projectId; } - // 3. Send Request with Retries const maxRetries = 3; let lastError: Error | null = null; - const fetchOptions = this.getFetchOptions(); // Reuse options + const fetchOptions = this.getFetchOptions(); for (let attempt = 1; attempt <= maxRetries; attempt++) { try { const response = await fetch(URLS.QUOTA, { method: 'POST', - headers: { - Authorization: `Bearer ${accessToken}`, - 'User-Agent': USER_AGENT, - 'Content-Type': 'application/json', - }, + headers: buildInternalApiHeaders(accessToken), body: JSON.stringify(payload), + signal: createTimeoutSignal(REQUEST_TIMEOUT_MS), ...fetchOptions, }); @@ -302,39 +446,42 @@ export class GoogleAPIService { ); if (attempt < maxRetries) { - await new Promise((r) => setTimeout(r, 1000)); + await sleep(1000); continue; } else { throw new Error(errorMsg); } } - const data = (await response.json()) as { models: Record }; - const result: QuotaData = { models: {} }; + const data = (await response.json()) as FetchModelsResponse; + const result: QuotaData = { + models: {}, + subscription_tier: subscriptionTier, + is_forbidden: false, + }; - // Parse relevant models - for (const [name, info] of Object.entries(data.models || {})) { - if (info.quotaInfo) { - const fraction = info.quotaInfo.remainingFraction ?? 0; - const percentage = Math.floor(fraction * 100); - const resetTime = info.quotaInfo.resetTime || ''; - - if (name.toLowerCase().includes('gemini') || name.toLowerCase().includes('claude')) { - result.models[name] = { percentage, resetTime }; - } + for (const [modelName, modelInfoRaw] of Object.entries(data.models || {})) { + const modelQuota = toModelQuotaInfo(modelName, modelInfoRaw); + if (modelQuota) { + result.models[modelName] = modelQuota; } } + const modelForwardingRules = toModelForwardingRules(data.deprecatedModelIds); + if (modelForwardingRules) { + result.model_forwarding_rules = modelForwardingRules; + } + return result; - } catch (e: unknown) { - if (e instanceof Error) { + } catch (error: unknown) { + if (error instanceof Error) { logger.warn( - `[GoogleAPIService] Request failed: ${e.message} (Attempt ${attempt}/${maxRetries})`, + `[GoogleAPIService] Request failed: ${error.message} (Attempt ${attempt}/${maxRetries})`, ); } - lastError = e instanceof Error ? e : new Error(String(e)); + lastError = error instanceof Error ? error : new Error(String(error)); if (attempt < maxRetries) { - await new Promise((r) => setTimeout(r, 1000)); + await sleep(1000); } } } diff --git a/src/tests/unit/gemini-controller.integration.test.ts b/src/tests/unit/gemini-controller.integration.test.ts index b0f158d..abde542 100644 --- a/src/tests/unit/gemini-controller.integration.test.ts +++ b/src/tests/unit/gemini-controller.integration.test.ts @@ -31,7 +31,14 @@ describe('GeminiController Integration', () => { expect.objectContaining({ models: expect.arrayContaining([ expect.objectContaining({ - name: 'models/gemini-2.5-flash', + name: 'models/gemini-3-flash', + description: '', + inputTokenLimit: 128000, + outputTokenLimit: 8192, + supportedGenerationMethods: ['generateContent', 'countTokens'], + }), + expect.objectContaining({ + name: 'models/gemini-3.1-pro-high', description: '', inputTokenLimit: 128000, outputTokenLimit: 8192, @@ -74,13 +81,13 @@ describe('GeminiController Integration', () => { const reply = createReplyMock(); await controller.modelAction( - 'models/gemini-2.5-flash:generateContent', + 'models/gemini-3.1-pro-high:generateContent', { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] } as any, reply as any, ); expect(proxyService.handleGeminiGenerateContent).toHaveBeenCalledWith( - 'models/gemini-2.5-flash', + 'models/gemini-3.1-pro-high', expect.any(Object), ); expect(reply.status).toHaveBeenCalledWith(200); diff --git a/src/tests/unit/parity-rollout-guard.test.ts b/src/tests/unit/parity-rollout-guard.test.ts new file mode 100644 index 0000000..f8aa26b --- /dev/null +++ b/src/tests/unit/parity-rollout-guard.test.ts @@ -0,0 +1,106 @@ +import { beforeEach, describe, expect, it } from 'vitest'; +import { DEFAULT_APP_CONFIG, ProxyConfig } from '../../types/config'; +import { setServerConfig } from '../../server/server-config'; +import { TokenManagerService } from '../../server/modules/proxy/token-manager.service'; + +function createProxyConfig(overrides: Partial): ProxyConfig { + return { + ...DEFAULT_APP_CONFIG.proxy, + ...overrides, + upstream_proxy: { + ...DEFAULT_APP_CONFIG.proxy.upstream_proxy, + ...(overrides.upstream_proxy ?? {}), + }, + }; +} + +function seedTokens(service: TokenManagerService): void { + const nowSec = Math.floor(Date.now() / 1000); + (service as any).tokens = new Map([ + [ + 'acc-1', + { + account_id: 'acc-1', + email: 'acc-1@test.dev', + access_token: 'token-1', + refresh_token: 'refresh-1', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: nowSec + 3600, + project_id: 'project-1', + }, + ], + [ + 'acc-2', + { + account_id: 'acc-2', + email: 'acc-2@test.dev', + access_token: 'token-2', + refresh_token: 'refresh-2', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: nowSec + 3600, + project_id: 'project-2', + }, + ], + ]); +} + +describe('Parity rollout guard replay', () => { + let service: TokenManagerService; + + beforeEach(() => { + service = new TokenManagerService(); + seedTokens(service); + }); + + it('supports kill-switch rollback from parity path to baseline path', async () => { + setServerConfig( + createProxyConfig({ + parity_enabled: true, + parity_kill_switch: false, + preferred_account_id: 'acc-2', + }), + ); + const paritySelected = await service.getNextToken({ model: 'gemini-2.5-flash' }); + expect(paritySelected?.id).toBe('acc-2'); + + setServerConfig( + createProxyConfig({ + parity_enabled: true, + parity_kill_switch: true, + preferred_account_id: 'acc-2', + }), + ); + (service as any).currentIndex = 0; + const rollbackSelected = await service.getNextToken({ model: 'gemini-2.5-flash' }); + expect(rollbackSelected?.id).toBe('acc-1'); + }); + + it('enforces no-go threshold after shadow mismatch and blocks parity enablement', async () => { + setServerConfig( + createProxyConfig({ + parity_enabled: false, + parity_shadow_enabled: true, + parity_no_go_mismatch_rate: 0, + preferred_account_id: 'acc-2', + }), + ); + + const shadowSelection = await service.getNextToken({ model: 'gemini-2.5-flash' }); + expect(shadowSelection?.id).toBe('acc-1'); + + setServerConfig( + createProxyConfig({ + parity_enabled: true, + parity_shadow_enabled: false, + parity_kill_switch: false, + preferred_account_id: 'acc-2', + }), + ); + (service as any).currentIndex = 0; + const blockedSelection = await service.getNextToken({ model: 'gemini-2.5-flash' }); + expect(blockedSelection?.id).toBe('acc-1'); + }); +}); + diff --git a/src/tests/unit/proxy-controller.integration.test.ts b/src/tests/unit/proxy-controller.integration.test.ts index b43a5b4..d466537 100644 --- a/src/tests/unit/proxy-controller.integration.test.ts +++ b/src/tests/unit/proxy-controller.integration.test.ts @@ -200,6 +200,73 @@ describe('ProxyController Integration', () => { ); }); + it('supports OpenAI responses compatibility endpoint in stream mode with SSE headers', async () => { + const stream = of('data: {"id":"chatcmpl_resp_stream"}\n\n'); + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue(stream), + }; + + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.responses( + { + model: 'gpt-4o', + instructions: 'stream output', + input: 'hello', + stream: true, + }, + reply as any, + ); + + expect(reply.header).toHaveBeenCalledWith('Content-Type', 'text/event-stream'); + expect(reply.header).toHaveBeenCalledWith('Cache-Control', 'no-cache'); + expect(reply.header).toHaveBeenCalledWith('Connection', 'keep-alive'); + expect(reply.send).toHaveBeenCalledWith(stream); + }); + + it('normalizes web_search_call in /v1/responses into builtin_web_search tool messages', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue({ + id: 'chatcmpl_resp_search', + object: 'chat.completion', + created: 1700000002, + model: 'gpt-4o', + choices: [{ index: 0, finish_reason: 'stop', message: { content: 'done' } }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }), + }; + + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.responses( + { + model: 'gpt-4o', + input: [ + { + type: 'web_search_call', + call_id: 'call_search_1', + action: { query: 'gemini api' }, + }, + { + type: 'function_call_output', + call_id: 'call_search_1', + output: { content: 'search result' }, + }, + ], + }, + reply as any, + ); + + const callArg = proxyService.handleChatCompletions.mock.calls[0][0]; + const assistantMessage = callArg.messages.find((message: { role: string }) => message.role === 'assistant'); + const toolMessage = callArg.messages.find((message: { role: string }) => message.role === 'tool'); + + expect(assistantMessage?.tool_calls?.[0]?.function?.name).toBe('builtin_web_search'); + expect(toolMessage?.name).toBe('builtin_web_search'); + }); + it('supports image generations endpoint', async () => { const proxyService = { handleChatCompletions: vi.fn().mockResolvedValue({ diff --git a/src/tests/unit/streaming-state.test.ts b/src/tests/unit/streaming-state.test.ts index 5582ff4..917a3e7 100644 --- a/src/tests/unit/streaming-state.test.ts +++ b/src/tests/unit/streaming-state.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, beforeEach } from 'vitest'; -import { StreamingState } from '../../lib/antigravity/ClaudeStreamingMapper'; +import { PartProcessor, StreamingState } from '../../lib/antigravity/ClaudeStreamingMapper'; describe('StreamingState', () => { let state: StreamingState; @@ -58,4 +58,48 @@ describe('StreamingState', () => { expect(state.getErrorCount()).toBe(2); }); }); + + describe('stream aggregation compatibility', () => { + it('emits tool_use stop reason when functionCall appears in stream', () => { + const processor = new PartProcessor(state); + const functionChunks = processor.process({ + functionCall: { + name: 'builtin_web_search', + args: { query: 'gemini docs' }, + id: 'call_stream_1', + }, + }); + const finishChunks = state.emitFinish('STOP', { + promptTokenCount: 2, + candidatesTokenCount: 3, + } as any); + + const output = [...functionChunks, ...finishChunks].join(''); + expect(output).toContain('"type":"tool_use"'); + expect(output).toContain('"stop_reason":"tool_use"'); + expect(output).toContain('"message_stop"'); + }); + + it('aggregates grounding metadata into final text block', () => { + state.webSearchQuery = 'gemini api'; + state.groundingChunks = [ + { + web: { + title: 'Gemini API Docs', + uri: 'https://example.com/gemini', + }, + }, + ]; + + const chunks = state.emitFinish('STOP', { + promptTokenCount: 1, + candidatesTokenCount: 1, + } as any); + const output = chunks.join(''); + + expect(output).toContain('Searched for you'); + expect(output).toContain('Citations'); + expect(output).toContain('https://example.com/gemini'); + }); + }); }); diff --git a/src/tests/unit/token-manager-parity-scheduling.test.ts b/src/tests/unit/token-manager-parity-scheduling.test.ts new file mode 100644 index 0000000..075ae51 --- /dev/null +++ b/src/tests/unit/token-manager-parity-scheduling.test.ts @@ -0,0 +1,166 @@ +import { beforeEach, describe, expect, it } from 'vitest'; +import { DEFAULT_APP_CONFIG, ProxyConfig } from '../../types/config'; +import { setServerConfig } from '../../server/server-config'; +import { TokenManagerService } from '../../server/modules/proxy/token-manager.service'; + +function createProxyConfig(overrides: Partial): ProxyConfig { + return { + ...DEFAULT_APP_CONFIG.proxy, + ...overrides, + upstream_proxy: { + ...DEFAULT_APP_CONFIG.proxy.upstream_proxy, + ...(overrides.upstream_proxy ?? {}), + }, + }; +} + +function seedTokens(service: TokenManagerService): void { + const nowSec = Math.floor(Date.now() / 1000); + (service as any).tokens = new Map([ + [ + 'acc-1', + { + account_id: 'acc-1', + email: 'acc-1@test.dev', + access_token: 'token-1', + refresh_token: 'refresh-1', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: nowSec + 3600, + project_id: 'project-1', + session_id: 'session-1', + }, + ], + [ + 'acc-2', + { + account_id: 'acc-2', + email: 'acc-2@test.dev', + access_token: 'token-2', + refresh_token: 'refresh-2', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: nowSec + 3600, + project_id: 'project-2', + session_id: 'session-2', + }, + ], + ]); +} + +describe('TokenManagerService parity scheduling replay', () => { + let service: TokenManagerService; + + beforeEach(() => { + service = new TokenManagerService(); + seedTokens(service); + }); + + it('prioritizes preferred account in parity mode', async () => { + setServerConfig( + createProxyConfig({ + parity_enabled: true, + parity_kill_switch: false, + scheduling_mode: 'balance', + preferred_account_id: 'acc-2', + }), + ); + + const token = await service.getNextToken({ model: 'gemini-2.5-flash' }); + expect(token?.id).toBe('acc-2'); + }); + + it('rotates sticky account when limited in balance mode', async () => { + setServerConfig( + createProxyConfig({ + parity_enabled: true, + parity_kill_switch: false, + scheduling_mode: 'balance', + preferred_account_id: '', + }), + ); + + const first = await service.getNextToken({ + sessionKey: 'openai:user-1', + model: 'gemini-2.5-flash', + }); + expect(first?.id).toBe('acc-1'); + + await service.markFromUpstreamError({ + accountIdOrEmail: 'acc-1', + status: 429, + model: 'gemini-2.5-flash', + body: JSON.stringify({ + error: { + details: [{ reason: 'RATE_LIMIT_EXCEEDED' }], + }, + }), + }); + + const second = await service.getNextToken({ + sessionKey: 'openai:user-1', + model: 'gemini-2.5-flash', + }); + expect(second?.id).toBe('acc-2'); + }); + + it('applies model-level lock for quota exhausted only on the same model', async () => { + setServerConfig( + createProxyConfig({ + parity_enabled: true, + parity_kill_switch: false, + scheduling_mode: 'performance-first', + }), + ); + + await service.markFromUpstreamError({ + accountIdOrEmail: 'acc-1', + status: 429, + model: 'gemini-2.5-flash', + body: JSON.stringify({ + error: { + details: [{ reason: 'QUOTA_EXHAUSTED', metadata: { quotaResetDelay: '30s' } }], + }, + }), + }); + + const sameModel = await service.getNextToken({ model: 'gemini-2.5-flash' }); + expect(sameModel?.id).toBe('acc-2'); + + const otherModel = await service.getNextToken({ + model: 'gemini-2.5-pro', + excludeAccountIds: ['acc-2'], + }); + expect(otherModel?.id).toBe('acc-1'); + }); + + it('falls back to excluded pool when retry exclusions would empty all candidates', async () => { + setServerConfig( + createProxyConfig({ + parity_enabled: false, + parity_kill_switch: false, + }), + ); + + const nowSec = Math.floor(Date.now() / 1000); + (service as any).tokens = new Map([ + [ + 'acc-1', + { + account_id: 'acc-1', + email: 'acc-1@test.dev', + access_token: 'token-1', + refresh_token: 'refresh-1', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: nowSec + 3600, + project_id: 'project-1', + session_id: 'session-1', + }, + ], + ]); + + const token = await service.getNextToken({ excludeAccountIds: ['acc-1'] }); + expect(token?.id).toBe('acc-1'); + }); +}); diff --git a/src/types/config.ts b/src/types/config.ts index f9d4094..b2a32e8 100644 --- a/src/types/config.ts +++ b/src/types/config.ts @@ -6,14 +6,24 @@ export const UpstreamProxyConfigSchema = z.object({ }); export const ProxyConfigSchema = z.object({ - enabled: z.boolean(), // 是否启用 - port: z.number(), // 监听端口 - api_key: z.string(), // API 密钥 (Also used for Admin Mode auth) - auto_start: z.boolean(), // 是否自动启动 + enabled: z.boolean(), + port: z.number(), + api_key: z.string(), + auto_start: z.boolean(), backend_canary_enabled: z.boolean().default(true), + parity_enabled: z.boolean().default(false), + parity_shadow_enabled: z.boolean().default(false), + parity_kill_switch: z.boolean().default(false), + parity_no_go_mismatch_rate: z.number().default(0.15), + parity_no_go_error_rate: z.number().default(0.4), + scheduling_mode: z.enum(['cache-first', 'balance', 'performance-first']).default('balance'), + max_wait_seconds: z.number().default(60), + preferred_account_id: z.string().default(''), + circuit_breaker_enabled: z.boolean().default(true), + circuit_breaker_backoff_steps: z.array(z.number()).default([60, 300, 1800, 7200]), custom_mapping: z.record(z.string(), z.string()).default({}), - anthropic_mapping: z.record(z.string(), z.string()), // 映射表 - request_timeout: z.number().default(120), // 超时秒数 + anthropic_mapping: z.record(z.string(), z.string()), // Mapping table + request_timeout: z.number().default(120), // Timeout in seconds upstream_proxy: UpstreamProxyConfigSchema, }); @@ -27,7 +37,7 @@ export const AppConfigSchema = z.object({ auto_startup: z.boolean(), error_reporting_enabled: z.boolean(), privacy_consent_asked: z.boolean().optional().default(false), // Optional for backward compatibility - default_export_path: z.string().nullable().optional(), // 导出路径 + default_export_path: z.string().nullable().optional(), // Export path model_visibility: z.record(z.string(), z.boolean()).default({}), // Model visibility preferences provider_groupings_enabled: z.boolean().default(false), // Enable provider groupings UI grid_layout: z.enum(['auto', '2-col', '3-col', 'list']).default('auto'), // Account card grid layout @@ -58,6 +68,16 @@ export const DEFAULT_APP_CONFIG: AppConfig = { api_key: '', // Generated dynamically if default needed auto_start: false, backend_canary_enabled: true, + parity_enabled: false, + parity_shadow_enabled: false, + parity_kill_switch: false, + parity_no_go_mismatch_rate: 0.15, + parity_no_go_error_rate: 0.4, + scheduling_mode: 'balance', + max_wait_seconds: 60, + preferred_account_id: '', + circuit_breaker_enabled: true, + circuit_breaker_backoff_steps: [60, 300, 1800, 7200], custom_mapping: {}, anthropic_mapping: {}, request_timeout: 120,