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/ipc/config/handlers.ts
+2
View File
@@ -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<void> {
// 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);
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
src/services/GoogleAPIService.ts
+198 -51
View File
@@ -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<string, ModelQuotaInfo>;
model_forwarding_rules?: Record<string, string>;
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<string, boolean>;
}
// 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<string, boolean>;
}
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<string, ModelInfoRaw>;
deprecatedModelIds?: Record<string, DeprecatedModelInfoRaw>;
}
interface ProjectContext {
projectId?: string;
subscriptionTier?: string;
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => {
setTimeout(resolve, ms);
});
}
function buildInternalApiHeaders(accessToken: string): Record<string, string> {
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<string, string> | undefined {
if (!deprecatedModelIds || Object.keys(deprecatedModelIds).length === 0) {
return undefined;
}
const forwardingRules: Record<string, string> = {};
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<UserInfo>;
}
/**
* Core logic: Fetches the internal project ID needed for quota checks.
*/
public static async fetchProjectId(accessToken: string): Promise<string | null> {
public static async fetchProjectContext(accessToken: string): Promise<ProjectContext> {
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<string | null> {
const context = await this.fetchProjectContext(accessToken);
return context.projectId ?? null;
}
/**
* Core logic: Fetches detailed model quota information.
*/
static async fetchQuota(accessToken: string): Promise<QuotaData> {
// 1. Get Project ID (Critical)
const projectId = await this.fetchProjectId(accessToken);
const { projectId, subscriptionTier } = await this.fetchProjectContext(accessToken);
// 2. Build Payload
const payload: Record<string, unknown> = {};
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<string, ModelInfoRaw> };
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);
}
}
}
src/tests/unit/gemini-controller.integration.test.ts
+10 -3
View File
@@ -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);
src/tests/unit/parity-rollout-guard.test.ts
+106
View File
@@ -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>): 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: '[email protected]',
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: '[email protected]',
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');
});
});
src/tests/unit/proxy-controller.integration.test.ts
+67
View File
@@ -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({
src/tests/unit/streaming-state.test.ts
+45 -1
View File
@@ -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');
});
});
});
src/tests/unit/token-manager-parity-scheduling.test.ts
+166
View File
@@ -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>): 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: '[email protected]',
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: '[email protected]',
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: '[email protected]',
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');
});
});
src/types/config.ts
+27 -7
View File
@@ -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,