evandro/gemini-nexus
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(proxy): improve parity, rate limiting, and token scheduling

Browse Source
This commit is contained in:
Draculabo committed 2026-03-06 21:50:18 +08:00
1 parent 0a60c0eafa
commit 5401e40943
14 files changed
+1878 -456
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/server/modules/proxy/clients/gemini.client.ts
+59 -17
View File
@@ -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<NodeJS.ReadableStream> {
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<GeminiResponse> {
const url = `${this.baseUrl}/models/${model}:generateContent`;
const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
try {
const response = await axios.post<GeminiResponse>(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<string, string>,
): Promise<NodeJS.ReadableStream> {
const response = await this.requestWithEndpointFailover<NodeJS.ReadableStream>(
const response = await this.executeRequestWithEndpointFailover<NodeJS.ReadableStream>(
':streamGenerateContent?alt=sse',
body,
accessToken,
@@ -104,7 +119,7 @@ export class GeminiClient {
upstreamProxyUrl?: string,
extraHeaders?: Record<string, string>,
): Promise<GeminiResponse> {
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<T>(
private async executeRequestWithEndpointFailover<T>(
path: string,
body: GeminiInternalRequest,
accessToken: string,
@@ -173,8 +188,8 @@ export class GeminiClient {
): Promise<AxiosResponse<T>> {
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<never> {
private async throwUpstreamRequestError(error: unknown, operation: string): Promise<never> {
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<string, unknown>)['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));
src/server/modules/proxy/gemini.controller.ts
+21 -13
View File
@@ -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<unknown>): void {
private writeObservableSseResponse(res: FastifyReply, stream: Observable<unknown>): 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';
}
}
src/server/modules/proxy/interfaces/request-interfaces.ts
+3 -2
View File
@@ -30,12 +30,13 @@ export interface OpenAIContentPart {
}
export interface OpenAITool {
type: 'function';
function: {
type: string;
function?: {
name: string;
description?: string;
parameters?: Record<string, unknown>;
};
[key: string]: unknown;
}
export interface OpenAIToolCall {
src/server/modules/proxy/proxy.controller.ts
+111 -112
View File
@@ -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<string, unknown> {
private toLegacyTextCompletionsResponse(response: OpenAIChatResponse): Record<string, unknown> {
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<string, unknown>;
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<string, unknown>;
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<string, unknown>;
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, unknown>): 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<string, unknown>;
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:(?<mime>[^;]+);base64,(?<data>[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<string, unknown>;
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<string, unknown> | null {
if (!value || typeof value !== 'object' || Array.isArray(value)) {
if (!isPlainObject(value)) {
return null;
}
return value as Record<string, unknown>;
}
private asString(value: unknown): string | null {
return typeof value === 'string' ? value : null;
return isString(value) ? value : null;
}
private isObservableLike(value: unknown): value is Observable<unknown> {
@@ -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,
src/server/modules/proxy/proxy.service.ts
+318 -122
View File
@@ -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<AnthropicChatResponse | Observable<string>> {
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<string> {
return new Observable<string>((subscriber) => {
const decoder = new TextDecoder();
@@ -272,10 +281,10 @@ export class ProxyService {
request: GeminiRequest,
): Promise<GeminiResponse> {
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<Observable<string>> {
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<string, unknown>)
| 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<OpenAIChatResponse | Observable<string>> {
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<AnthropicChatRequest['tools']> = [];
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<void> {
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<string, string> {
private createModelSpecificHeaders(model: string | undefined): Record<string, string> {
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<string, unknown> {
return value !== null && typeof value === 'object' && !Array.isArray(value);
return isPlainObject(value);
}
private isGeminiPart(value: unknown): value is InternalGeminiPart {
src/server/modules/proxy/token-manager.service.ts
+745 -128
View File
File diff suppressed because it is too large. Load diff