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

feat(proxy): improve quota fallback and proxy parity handling

Browse Source
This commit is contained in:
Draculabo committed 2026-04-11 16:47:57 +08:00
1 parent 388fbcb2a9
commit 4fa0a3c63d
9 files changed
+350 -66
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/server/modules/proxy/proxy.service.ts
+39 -24
View File
@@ -15,10 +15,7 @@ import {
} from '../../../lib/antigravity/types';
import { calculateRetryDelay, sleep } from '../../../lib/antigravity/retry-utils';
import { normalizeObjectJsonSchema } from '../../../lib/antigravity/JsonSchemaUtils';
import {
classifyStreamError,
formatErrorForSSE,
} from '../../../lib/antigravity/stream-error-utils';
import { classifyStreamError } from '../../../lib/antigravity/stream-error-utils';
import {
OpenAIChatRequest,
AnthropicChatRequest,
@@ -80,6 +77,10 @@ export class ProxyService {
throw new Error('No available accounts');
}
attemptedAccountIds.add(token.id);
const effectiveTargetModel = this.tokenManager.resolveDynamicModelForAccount(
token.id,
targetModel,
);
try {
const projectId = token.token.project_id ?? '';
@@ -89,7 +90,8 @@ export class ProxyService {
projectId,
requestUserAgent,
);
this.applyInternalGenerationConstraints(geminiBody, targetModel, token.id);
geminiBody.model = effectiveTargetModel;
this.applyInternalGenerationConstraints(geminiBody, effectiveTargetModel, token.id);
if (request.stream) {
const stream = await this.geminiClient.streamGenerateInternal(
@@ -120,7 +122,8 @@ export class ProxyService {
'',
requestUserAgent,
);
this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
fallbackBody.model = effectiveTargetModel;
this.applyInternalGenerationConstraints(fallbackBody, effectiveTargetModel, token.id);
if (request.stream) {
const stream = await this.geminiClient.streamGenerateInternal(
fallbackBody,
@@ -186,7 +189,7 @@ export class ProxyService {
}
lastError = error;
await this.applyUpstreamPenalty(token.id, targetModel, error);
await this.applyUpstreamPenalty(token.id, effectiveTargetModel, error);
}
}
throw lastError || new Error('Request failed after retries');
@@ -267,12 +270,9 @@ export class ProxyService {
upstreamStream.on('error', (err: unknown) => {
const cleanError = err instanceof Error ? err : new Error(String(err));
const { type, message } = classifyStreamError(cleanError);
const { type } = classifyStreamError(cleanError);
this.logger.error(`Stream error: ${type} - ${cleanError.message}`);
// Send SSE error event before closing
subscriber.next(formatErrorForSSE(type, message));
subscriber.error(cleanError);
});
});
@@ -309,17 +309,21 @@ export class ProxyService {
throw new Error('No available accounts (all exhausted or rate limited)');
}
attemptedAccountIds.add(token.id);
const effectiveTargetModel = this.tokenManager.resolveDynamicModelForAccount(
token.id,
targetModel,
);
try {
const requestUserAgent = await resolveRequestUserAgent();
const internalBody = this.createGeminiInternalRequest(
targetModel,
effectiveTargetModel,
request,
token.token.project_id ?? '',
'generate-content',
requestUserAgent,
);
this.applyInternalGenerationConstraints(internalBody, targetModel, token.id);
this.applyInternalGenerationConstraints(internalBody, effectiveTargetModel, token.id);
const response = await this.generateInternalWithStreamFallback(
internalBody,
@@ -337,13 +341,13 @@ export class ProxyService {
try {
const requestUserAgent = await resolveRequestUserAgent();
const fallbackBody = this.createGeminiInternalRequest(
targetModel,
effectiveTargetModel,
request,
'',
'generate-content',
requestUserAgent,
);
this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
this.applyInternalGenerationConstraints(fallbackBody, effectiveTargetModel, token.id);
const response = await this.generateInternalWithStreamFallback(
fallbackBody,
token.token.access_token,
@@ -358,7 +362,7 @@ export class ProxyService {
lastError = err;
}
await this.applyUpstreamPenalty(token.id, targetModel, lastError);
await this.applyUpstreamPenalty(token.id, effectiveTargetModel, lastError);
}
}
@@ -395,17 +399,21 @@ export class ProxyService {
throw new Error('No available accounts (all exhausted or rate limited)');
}
attemptedAccountIds.add(token.id);
const effectiveTargetModel = this.tokenManager.resolveDynamicModelForAccount(
token.id,
targetModel,
);
try {
const requestUserAgent = await resolveRequestUserAgent();
const internalBody = this.createGeminiInternalRequest(
targetModel,
effectiveTargetModel,
request,
token.token.project_id ?? '',
'generate-content',
requestUserAgent,
);
this.applyInternalGenerationConstraints(internalBody, targetModel, token.id);
this.applyInternalGenerationConstraints(internalBody, effectiveTargetModel, token.id);
const stream = await this.geminiClient.streamGenerateInternal(
internalBody,
@@ -422,13 +430,13 @@ export class ProxyService {
try {
const requestUserAgent = await resolveRequestUserAgent();
const fallbackBody = this.createGeminiInternalRequest(
targetModel,
effectiveTargetModel,
request,
'',
'generate-content',
requestUserAgent,
);
this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
this.applyInternalGenerationConstraints(fallbackBody, effectiveTargetModel, token.id);
const stream = await this.geminiClient.streamGenerateInternal(
fallbackBody,
token.token.access_token,
@@ -443,7 +451,7 @@ export class ProxyService {
lastError = err;
}
await this.applyUpstreamPenalty(token.id, targetModel, lastError);
await this.applyUpstreamPenalty(token.id, effectiveTargetModel, lastError);
}
}
@@ -698,13 +706,18 @@ export class ProxyService {
throw new Error('No available accounts (all exhausted or rate limited)');
}
attemptedAccountIds.add(token.id);
const effectiveTargetModel = this.tokenManager.resolveDynamicModelForAccount(
token.id,
targetModel,
);
try {
const claudeRequest = this.convertOpenAIToClaude(request);
const projectId = token.token.project_id ?? '';
const requestUserAgent = await resolveRequestUserAgent();
const geminiBody = transformClaudeRequestIn(claudeRequest, projectId, requestUserAgent);
this.applyInternalGenerationConstraints(geminiBody, targetModel, token.id);
geminiBody.model = effectiveTargetModel;
this.applyInternalGenerationConstraints(geminiBody, effectiveTargetModel, token.id);
// Use v1internal API (same as Anthropic handler)
if (request.stream) {
@@ -765,7 +778,8 @@ export class ProxyService {
const claudeRequest = this.convertOpenAIToClaude(request);
const requestUserAgent = await resolveRequestUserAgent();
const fallbackBody = transformClaudeRequestIn(claudeRequest, '', requestUserAgent);
this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
fallbackBody.model = effectiveTargetModel;
this.applyInternalGenerationConstraints(fallbackBody, effectiveTargetModel, token.id);
if (request.stream) {
const stream = await this.geminiClient.streamGenerateInternal(
fallbackBody,
@@ -791,7 +805,7 @@ export class ProxyService {
lastError = err;
}
await this.applyUpstreamPenalty(token.id, targetModel, lastError);
await this.applyUpstreamPenalty(token.id, effectiveTargetModel, lastError);
}
}
throw lastError || new Error('Request failed after retries');
@@ -1087,6 +1101,7 @@ export class ProxyService {
upstreamStream.on('error', (err: unknown) => {
// Convert to clean Error to avoid circular reference issues (socket objects)
const cleanError = err instanceof Error ? new Error(err.message) : new Error(String(err));
this.logger.error(`OpenAI-compatible stream error: ${cleanError.message}`);
subscriber.error(cleanError);
});
});
src/server/modules/proxy/request-user-agent.ts
+1 -1
View File
@@ -4,7 +4,7 @@ import { logger } from '../../../utils/logger';
const REMOTE_VERSION_URL = 'https://antigravity-auto-updater-974169037036.us-central1.run.app';
const CHANGELOG_URL = 'https://antigravity.google/changelog';
export const FALLBACK_VERSION = '1.19.6';
export const FALLBACK_VERSION = '1.22.2';
const DEFAULT_REMOTE_TIMEOUT_MS = 2500;
const VERSION_REGEX = /\d+\.\d+\.\d+/g;
src/server/modules/proxy/token-manager.service.ts
+128 -2
View File
@@ -11,6 +11,7 @@ interface TokenData {
account_id: string;
access_token: string;
refresh_token: string;
oauth_client_key?: string;
token_type: string;
expires_in: number;
expiry_timestamp: number;
@@ -59,6 +60,14 @@ function normalizeModelId(modelId: string | null | undefined): string | undefine
return normalized !== '' ? normalized : undefined;
}
function normalizeClientKey(clientKey: string | undefined): string | undefined {
if (typeof clientKey !== 'string') {
return undefined;
}
const normalized = clientKey.trim().toLowerCase();
return normalized !== '' ? normalized : undefined;
}
function delay(ms: number): Promise<void> {
return new Promise((resolve) => {
setTimeout(resolve, ms);
@@ -755,12 +764,13 @@ export class TokenManagerService implements OnModuleInit {
email: account.email,
access_token: account.token.access_token,
refresh_token: account.token.refresh_token,
oauth_client_key: normalizeClientKey(account.token.oauth_client_key),
token_type: account.token.token_type || 'Bearer',
expires_in: account.token.expires_in,
expiry_timestamp: account.token.expiry_timestamp,
project_id: account.token.project_id || undefined,
session_id: account.token.session_id || this.generateSessionId(),
upstream_proxy_url: account.token.upstream_proxy_url || undefined,
upstream_proxy_url: account.token.upstream_proxy_url || account.proxy_url || undefined,
quota,
model_quotas: extractedState.modelQuotas,
model_limits: extractedState.modelLimits,
@@ -789,10 +799,18 @@ export class TokenManagerService implements OnModuleInit {
if (nowSeconds >= tokenData.expiry_timestamp - 300) {
this.logger.log(`Access token near expiry for ${tokenData.email}; refreshing`);
try {
const newTokens = await GoogleAPIService.refreshAccessToken(tokenData.refresh_token);
const newTokens = await GoogleAPIService.refreshAccessToken(
tokenData.refresh_token,
tokenData.upstream_proxy_url,
tokenData.oauth_client_key,
);
tokenData.access_token = newTokens.access_token;
tokenData.expires_in = newTokens.expires_in;
tokenData.expiry_timestamp = nowSeconds + newTokens.expires_in;
tokenData.oauth_client_key = this.normalizeRefreshedOauthClientKey(
tokenData,
newTokens.oauth_client_key,
);
await this.persistTokenState(accountId, tokenData);
this.tokens.set(accountId, tokenData);
this.logger.log(`Access token refreshed for ${tokenData.email}`);
@@ -855,6 +873,7 @@ export class TokenManagerService implements OnModuleInit {
expires_in: tokenData.expires_in,
expiry_timestamp: tokenData.expiry_timestamp,
project_id: effectiveProjectId,
oauth_client_key: tokenData.oauth_client_key,
session_id: tokenData.session_id,
upstream_proxy_url: tokenData.upstream_proxy_url,
},
@@ -913,6 +932,10 @@ export class TokenManagerService implements OnModuleInit {
expires_in: tokenData.expires_in,
expiry_timestamp: tokenData.expiry_timestamp,
project_id: tokenData.project_id ?? acc.token.project_id,
oauth_client_key:
tokenData.oauth_client_key ??
normalizeClientKey(acc.token.oauth_client_key) ??
acc.token.oauth_client_key,
session_id: tokenData.session_id ?? acc.token.session_id,
upstream_proxy_url: tokenData.upstream_proxy_url ?? acc.token.upstream_proxy_url,
};
@@ -927,6 +950,13 @@ export class TokenManagerService implements OnModuleInit {
return this.tokens.size;
}
private normalizeRefreshedOauthClientKey(
currentToken: { oauth_client_key?: string; project_id?: string },
refreshedClientKey?: string,
): string | undefined {
return GoogleAPIService.normalizeRefreshedOAuthClientKey(currentToken, refreshedClientKey);
}
getAllCollectedModels(): Set<string> {
const allModels = new Set<string>();
for (const tokenData of this.tokens.values()) {
@@ -937,6 +967,102 @@ export class TokenManagerService implements OnModuleInit {
return allModels;
}
private getAvailableModelsFromToken(tokenData: TokenData): Set<string> {
const availableModels = new Set<string>();
for (const modelId of Object.keys(tokenData.model_quotas ?? {})) {
const normalized = normalizeModelId(modelId)?.toLowerCase();
if (normalized) {
availableModels.add(normalized);
}
}
for (const modelId of Object.keys(tokenData.quota?.models ?? {})) {
const normalized = normalizeModelId(modelId)?.toLowerCase();
if (normalized) {
availableModels.add(normalized);
}
}
return availableModels;
}
private buildDynamicModelCandidates(modelName: string): string[] | null {
const normalizedModel = normalizeModelId(modelName)?.toLowerCase();
if (!normalizedModel) {
return null;
}
const proFamily = new Set([
'gemini-3-pro',
'gemini-3-pro-preview',
'gemini-3-pro-high',
'gemini-3-pro-low',
'gemini-3.1-pro',
'gemini-3.1-pro-preview',
'gemini-3.1-pro-high',
'gemini-3.1-pro-low',
]);
if (!proFamily.has(normalizedModel)) {
return null;
}
const candidates: string[] = [];
const seen = new Set<string>();
const pushCandidate = (candidate: string) => {
if (!seen.has(candidate)) {
seen.add(candidate);
candidates.push(candidate);
}
};
// Keep requested model first, then fallback across the same family.
pushCandidate(normalizedModel);
pushCandidate('gemini-3.1-pro-preview');
pushCandidate('gemini-3-pro-preview');
pushCandidate('gemini-3.1-pro-high');
pushCandidate('gemini-3-pro-high');
pushCandidate('gemini-3.1-pro-low');
pushCandidate('gemini-3-pro-low');
return candidates;
}
resolveDynamicModelForAccount(accountId: string, mappedModel: string): string {
const candidates = this.buildDynamicModelCandidates(mappedModel);
if (!candidates) {
return mappedModel;
}
const tokenData = this.tokens.get(accountId);
if (!tokenData) {
return mappedModel;
}
const availableModels = this.getAvailableModelsFromToken(tokenData);
if (availableModels.size === 0) {
return mappedModel;
}
const normalizedMappedModel = normalizeModelId(mappedModel)?.toLowerCase() ?? mappedModel;
for (const candidate of candidates) {
if (!availableModels.has(candidate)) {
continue;
}
if (candidate !== normalizedMappedModel) {
this.logger.log(
`[Dynamic-Model-Rewrite] account=${accountId} ${mappedModel} -> ${candidate}`,
);
}
return candidate;
}
return mappedModel;
}
getModelOutputLimitForAccount(accountId: string, modelName: string): number | undefined {
const tokenData = this.tokens.get(accountId);
const normalizedModel = normalizeModelId(modelName);