feat(proxy): add upstream error and rate-limit tracking support
This commit is contained in:
1 parent
c95f494569
commit
20fffff29c
6 files changed
+803
-2
No files matched your search
@@ -0,0 +1,23 @@
|
||||
export interface UpstreamErrorHeaders {
|
||||
retryAfter?: string;
|
||||
}
|
||||
|
||||
export class UpstreamRequestError extends Error {
|
||||
readonly status?: number;
|
||||
readonly headers?: UpstreamErrorHeaders;
|
||||
readonly body?: string;
|
||||
|
||||
constructor(params: {
|
||||
message: string;
|
||||
status?: number;
|
||||
headers?: UpstreamErrorHeaders;
|
||||
body?: string;
|
||||
}) {
|
||||
super(params.message);
|
||||
this.name = 'UpstreamRequestError';
|
||||
this.status = params.status;
|
||||
this.headers = params.headers;
|
||||
this.body = params.body;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,426 @@
|
||||
export enum RateLimitReason {
|
||||
QuotaExhausted = 'quota_exhausted',
|
||||
RateLimitExceeded = 'rate_limit_exceeded',
|
||||
ModelCapacityExhausted = 'model_capacity_exhausted',
|
||||
ServerError = 'server_error',
|
||||
Unknown = 'unknown',
|
||||
}
|
||||
|
||||
interface RateLimitInfo {
|
||||
resetTimeMs: number;
|
||||
retryAfterSec: number;
|
||||
reason: RateLimitReason;
|
||||
model?: string;
|
||||
}
|
||||
|
||||
type FailureCountEntry = {
|
||||
count: number;
|
||||
lastFailureMs: number;
|
||||
};
|
||||
|
||||
interface GoogleErrorDetail {
|
||||
reason?: string;
|
||||
retryDelay?: string;
|
||||
metadata?: {
|
||||
quotaResetDelay?: string;
|
||||
retryDelay?: string;
|
||||
};
|
||||
}
|
||||
|
||||
interface ParsedGoogleErrorBody {
|
||||
error?: {
|
||||
status?: string;
|
||||
details?: GoogleErrorDetail[];
|
||||
retry_after?: number;
|
||||
};
|
||||
}
|
||||
|
||||
const FAILURE_COUNT_EXPIRY_MS = 60 * 60 * 1000;
|
||||
|
||||
function toLowerText(value: string | undefined): string {
|
||||
return (value ?? '').toLowerCase();
|
||||
}
|
||||
|
||||
function parseDurationToSeconds(text: string): number | null {
|
||||
const match = text.match(/(?:(\d+)h)?(?:(\d+)m)?(?:(\d+(?:\.\d+)?)s)?(?:(\d+(?:\.\d+)?)ms)?/i);
|
||||
if (!match) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const hours = Number(match[1] ?? 0);
|
||||
const minutes = Number(match[2] ?? 0);
|
||||
const seconds = Number(match[3] ?? 0);
|
||||
const milliseconds = Number(match[4] ?? 0);
|
||||
|
||||
const totalSeconds =
|
||||
hours * 3600 + minutes * 60 + Math.ceil(seconds) + Math.ceil(milliseconds / 1000);
|
||||
if (totalSeconds <= 0) {
|
||||
return null;
|
||||
}
|
||||
return totalSeconds;
|
||||
}
|
||||
|
||||
function tryParseGoogleErrorBody(body: string | undefined): ParsedGoogleErrorBody | null {
|
||||
const trimmed = (body ?? '').trim();
|
||||
if (!trimmed.startsWith('{') && !trimmed.startsWith('[')) {
|
||||
return null;
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed = JSON.parse(trimmed) as ParsedGoogleErrorBody;
|
||||
if (!parsed || typeof parsed !== 'object') {
|
||||
return null;
|
||||
}
|
||||
return parsed;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function mapGoogleReasonToTrackerReason(reason: string): RateLimitReason | null {
|
||||
const normalizedReason = reason.trim().toUpperCase();
|
||||
if (normalizedReason === 'QUOTA_EXHAUSTED') {
|
||||
return RateLimitReason.QuotaExhausted;
|
||||
}
|
||||
if (normalizedReason === 'RATE_LIMIT_EXCEEDED') {
|
||||
return RateLimitReason.RateLimitExceeded;
|
||||
}
|
||||
if (normalizedReason === 'MODEL_CAPACITY_EXHAUSTED') {
|
||||
return RateLimitReason.ModelCapacityExhausted;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
export class RateLimitTracker {
|
||||
private readonly lockoutByKey = new Map<string, RateLimitInfo>();
|
||||
private readonly failureCounts = new Map<string, FailureCountEntry>();
|
||||
|
||||
private buildLockoutKey(accountId: string, model?: string): string {
|
||||
if (model && model.trim() !== '') {
|
||||
return `${accountId}:${model}`;
|
||||
}
|
||||
return accountId;
|
||||
}
|
||||
|
||||
getRemainingWaitSeconds(accountId: string, model?: string): number {
|
||||
const now = Date.now();
|
||||
|
||||
const globalLock = this.lockoutByKey.get(accountId);
|
||||
if (globalLock && globalLock.resetTimeMs > now) {
|
||||
return Math.max(0, Math.ceil((globalLock.resetTimeMs - now) / 1000));
|
||||
}
|
||||
|
||||
if (model) {
|
||||
const modelKey = this.buildLockoutKey(accountId, model);
|
||||
const modelLock = this.lockoutByKey.get(modelKey);
|
||||
if (modelLock && modelLock.resetTimeMs > now) {
|
||||
return Math.max(0, Math.ceil((modelLock.resetTimeMs - now) / 1000));
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
getRemainingWaitSec(accountId: string, model?: string): number {
|
||||
return this.getRemainingWaitSeconds(accountId, model);
|
||||
}
|
||||
|
||||
isRateLimited(accountId: string, model?: string): boolean {
|
||||
return this.getRemainingWaitSeconds(accountId, model) > 0;
|
||||
}
|
||||
|
||||
setLockoutUntilIso(
|
||||
accountId: string,
|
||||
resetTimeIso: string,
|
||||
reason: RateLimitReason,
|
||||
model?: string,
|
||||
): boolean {
|
||||
const timestamp = Date.parse(resetTimeIso);
|
||||
if (!Number.isFinite(timestamp)) {
|
||||
return false;
|
||||
}
|
||||
this.setLockoutUntil(accountId, timestamp, reason, model);
|
||||
return true;
|
||||
}
|
||||
|
||||
private setLockoutUntil(
|
||||
accountId: string,
|
||||
resetTimeMs: number,
|
||||
reason: RateLimitReason,
|
||||
model?: string,
|
||||
): void {
|
||||
const retryAfterSec = Math.max(2, Math.ceil((resetTimeMs - Date.now()) / 1000));
|
||||
const key = model && model.trim() !== '' ? this.buildLockoutKey(accountId, model) : accountId;
|
||||
this.lockoutByKey.set(key, {
|
||||
resetTimeMs,
|
||||
retryAfterSec,
|
||||
reason,
|
||||
model,
|
||||
});
|
||||
}
|
||||
|
||||
clear(accountId: string): boolean {
|
||||
return this.lockoutByKey.delete(accountId);
|
||||
}
|
||||
|
||||
clearAll(): number {
|
||||
const size = this.lockoutByKey.size;
|
||||
this.lockoutByKey.clear();
|
||||
return size;
|
||||
}
|
||||
|
||||
cleanupExpired(): number {
|
||||
const now = Date.now();
|
||||
let deleted = 0;
|
||||
for (const [key, info] of this.lockoutByKey.entries()) {
|
||||
if (info.resetTimeMs <= now) {
|
||||
this.lockoutByKey.delete(key);
|
||||
deleted++;
|
||||
}
|
||||
}
|
||||
return deleted;
|
||||
}
|
||||
|
||||
markSuccess(accountId: string): void {
|
||||
this.failureCounts.delete(accountId);
|
||||
this.lockoutByKey.delete(accountId);
|
||||
}
|
||||
|
||||
trackFromUpstreamError(params: {
|
||||
accountId: string;
|
||||
status?: number;
|
||||
retryAfter?: string;
|
||||
body?: string;
|
||||
model?: string;
|
||||
backoffSteps: number[];
|
||||
}): RateLimitInfo | null {
|
||||
const status = params.status ?? 0;
|
||||
if (![404, 429, 500, 503, 529].includes(status)) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const reason = this.detectRateLimitReason(status, params.body);
|
||||
const retryAfterSec = this.computeRetryAfterSeconds({
|
||||
reason,
|
||||
status,
|
||||
retryAfter: params.retryAfter,
|
||||
body: params.body,
|
||||
accountId: params.accountId,
|
||||
backoffSteps: params.backoffSteps,
|
||||
});
|
||||
|
||||
const info: RateLimitInfo = {
|
||||
reason,
|
||||
retryAfterSec,
|
||||
resetTimeMs: Date.now() + retryAfterSec * 1000,
|
||||
model: params.model,
|
||||
};
|
||||
|
||||
const useModelKey = reason === RateLimitReason.QuotaExhausted && Boolean(params.model);
|
||||
const key = useModelKey
|
||||
? this.buildLockoutKey(params.accountId, params.model)
|
||||
: params.accountId;
|
||||
this.lockoutByKey.set(key, info);
|
||||
|
||||
return info;
|
||||
}
|
||||
|
||||
parseAndMarkFromError(params: {
|
||||
accountId: string;
|
||||
status?: number;
|
||||
retryAfter?: string;
|
||||
body?: string;
|
||||
model?: string;
|
||||
backoffSteps: number[];
|
||||
}): RateLimitInfo | null {
|
||||
return this.trackFromUpstreamError(params);
|
||||
}
|
||||
|
||||
private detectRateLimitReason(status: number, body: string | undefined): RateLimitReason {
|
||||
const parsedBody = tryParseGoogleErrorBody(body);
|
||||
const details = Array.isArray(parsedBody?.error?.details) ? parsedBody?.error?.details : [];
|
||||
for (const detail of details) {
|
||||
if (typeof detail.reason !== 'string' || detail.reason.trim() === '') {
|
||||
continue;
|
||||
}
|
||||
const mappedReason = mapGoogleReasonToTrackerReason(detail.reason);
|
||||
if (mappedReason) {
|
||||
return mappedReason;
|
||||
}
|
||||
}
|
||||
|
||||
const statusFromBody = parsedBody?.error?.status?.trim().toUpperCase();
|
||||
if (statusFromBody === 'RESOURCE_EXHAUSTED') {
|
||||
return RateLimitReason.QuotaExhausted;
|
||||
}
|
||||
if (statusFromBody === 'UNAVAILABLE') {
|
||||
const loweredBody = toLowerText(body);
|
||||
if (loweredBody.includes('no capacity available') || loweredBody.includes('model capacity')) {
|
||||
return RateLimitReason.ModelCapacityExhausted;
|
||||
}
|
||||
}
|
||||
|
||||
if (status !== 429) {
|
||||
return RateLimitReason.ServerError;
|
||||
}
|
||||
|
||||
const lowerBody = toLowerText(body);
|
||||
if (
|
||||
lowerBody.includes('per minute') ||
|
||||
lowerBody.includes('rate limit') ||
|
||||
lowerBody.includes('rate_limit') ||
|
||||
lowerBody.includes('too many requests')
|
||||
) {
|
||||
return RateLimitReason.RateLimitExceeded;
|
||||
}
|
||||
if (
|
||||
lowerBody.includes('model capacity exhausted') ||
|
||||
lowerBody.includes('no capacity available')
|
||||
) {
|
||||
return RateLimitReason.ModelCapacityExhausted;
|
||||
}
|
||||
if (lowerBody.includes('quota') || lowerBody.includes('exhausted')) {
|
||||
return RateLimitReason.QuotaExhausted;
|
||||
}
|
||||
|
||||
return RateLimitReason.Unknown;
|
||||
}
|
||||
|
||||
private computeRetryAfterSeconds(params: {
|
||||
reason: RateLimitReason;
|
||||
status: number;
|
||||
retryAfter?: string;
|
||||
body?: string;
|
||||
accountId: string;
|
||||
backoffSteps: number[];
|
||||
}): number {
|
||||
const headerRetryRaw = params.retryAfter?.trim() ?? '';
|
||||
if (headerRetryRaw !== '') {
|
||||
const headerRetry = Number(headerRetryRaw);
|
||||
if (!Number.isNaN(headerRetry) && Number.isFinite(headerRetry) && headerRetry > 0) {
|
||||
return Math.max(2, Math.ceil(headerRetry));
|
||||
}
|
||||
|
||||
const headerRetryAt = Date.parse(headerRetryRaw);
|
||||
if (Number.isFinite(headerRetryAt)) {
|
||||
const retryAfterSec = Math.ceil((headerRetryAt - Date.now()) / 1000);
|
||||
if (retryAfterSec > 0) {
|
||||
return Math.max(2, retryAfterSec);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const bodyRetry = this.parseRetryAfterSecondsFromBody(params.body);
|
||||
if (bodyRetry !== null) {
|
||||
return Math.max(2, Math.ceil(bodyRetry));
|
||||
}
|
||||
|
||||
const failureCount =
|
||||
params.reason !== RateLimitReason.ServerError
|
||||
? this.incrementFailureCount(params.accountId)
|
||||
: 1;
|
||||
|
||||
if (params.reason === RateLimitReason.QuotaExhausted) {
|
||||
const steps = params.backoffSteps.length > 0 ? params.backoffSteps : [60, 300, 1800, 7200];
|
||||
const index = Math.max(0, failureCount - 1);
|
||||
return steps[Math.min(index, steps.length - 1)];
|
||||
}
|
||||
|
||||
if (params.reason === RateLimitReason.RateLimitExceeded) {
|
||||
return 5;
|
||||
}
|
||||
|
||||
if (params.reason === RateLimitReason.ModelCapacityExhausted) {
|
||||
if (failureCount <= 1) {
|
||||
return 5;
|
||||
}
|
||||
if (failureCount === 2) {
|
||||
return 10;
|
||||
}
|
||||
return 15;
|
||||
}
|
||||
|
||||
if (params.reason === RateLimitReason.ServerError) {
|
||||
return params.status === 404 ? 5 : 8;
|
||||
}
|
||||
|
||||
return 60;
|
||||
}
|
||||
|
||||
private parseRetryAfterSecondsFromBody(body: string | undefined): number | null {
|
||||
if (!body) {
|
||||
return null;
|
||||
}
|
||||
const parsedBody = tryParseGoogleErrorBody(body);
|
||||
const details = Array.isArray(parsedBody?.error?.details) ? parsedBody.error.details : [];
|
||||
for (const detail of details) {
|
||||
if (typeof detail.retryDelay === 'string' && detail.retryDelay.trim() !== '') {
|
||||
const parsedDelay = parseDurationToSeconds(detail.retryDelay);
|
||||
if (parsedDelay !== null) {
|
||||
return parsedDelay;
|
||||
}
|
||||
}
|
||||
|
||||
if (
|
||||
typeof detail.metadata?.retryDelay === 'string' &&
|
||||
detail.metadata.retryDelay.trim() !== ''
|
||||
) {
|
||||
const parsedDelay = parseDurationToSeconds(detail.metadata.retryDelay);
|
||||
if (parsedDelay !== null) {
|
||||
return parsedDelay;
|
||||
}
|
||||
}
|
||||
|
||||
if (
|
||||
typeof detail.metadata?.quotaResetDelay === 'string' &&
|
||||
detail.metadata.quotaResetDelay.trim() !== ''
|
||||
) {
|
||||
const parsedDelay = parseDurationToSeconds(detail.metadata.quotaResetDelay);
|
||||
if (parsedDelay !== null) {
|
||||
return parsedDelay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const retryAfter = parsedBody?.error?.retry_after;
|
||||
if (typeof retryAfter === 'number' && retryAfter > 0) {
|
||||
return Math.ceil(retryAfter);
|
||||
}
|
||||
|
||||
const lowerBody = toLowerText(body);
|
||||
const minSecMatch = lowerBody.match(/try again in\s+(\d+)m\s*(\d+)s/);
|
||||
if (minSecMatch) {
|
||||
return Number(minSecMatch[1]) * 60 + Number(minSecMatch[2]);
|
||||
}
|
||||
|
||||
const secMatch = lowerBody.match(/(?:try again in|backoff for|wait)\s*(\d+)s/);
|
||||
if (secMatch) {
|
||||
return Number(secMatch[1]);
|
||||
}
|
||||
|
||||
const resetMatch = lowerBody.match(/quota will reset in\s*(\d+)\s*second/);
|
||||
if (resetMatch) {
|
||||
return Number(resetMatch[1]);
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private incrementFailureCount(accountId: string): number {
|
||||
const now = Date.now();
|
||||
const entry = this.failureCounts.get(accountId);
|
||||
if (!entry) {
|
||||
this.failureCounts.set(accountId, { count: 1, lastFailureMs: now });
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (now - entry.lastFailureMs > FAILURE_COUNT_EXPIRY_MS) {
|
||||
this.failureCounts.set(accountId, { count: 1, lastFailureMs: now });
|
||||
return 1;
|
||||
}
|
||||
|
||||
const nextCount = entry.count + 1;
|
||||
this.failureCounts.set(accountId, { count: nextCount, lastFailureMs: now });
|
||||
return nextCount;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
import axios from 'axios';
|
||||
import { getAntigravityVersion } from '../../../utils/antigravityVersion';
|
||||
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';
|
||||
const DEFAULT_REMOTE_TIMEOUT_MS = 2500;
|
||||
const VERSION_REGEX = /\d+\.\d+\.\d+/g;
|
||||
|
||||
type UserAgentSource = 'local' | 'remote' | 'changelog' | 'fallback';
|
||||
|
||||
interface UserAgentResolution {
|
||||
source: UserAgentSource;
|
||||
userAgent: string;
|
||||
version: string;
|
||||
}
|
||||
|
||||
let cachedUserAgentResolution: UserAgentResolution | null = null;
|
||||
let pendingUserAgentResolution: Promise<UserAgentResolution> | null = null;
|
||||
|
||||
function normalizeNonEmptyString(value: string | null | undefined): string | null {
|
||||
if (typeof value !== 'string') {
|
||||
return null;
|
||||
}
|
||||
|
||||
const trimmed = value.trim();
|
||||
if (!trimmed) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return trimmed;
|
||||
}
|
||||
|
||||
function getPlatformTag(): string {
|
||||
switch (process.platform) {
|
||||
case 'win32':
|
||||
return 'windows';
|
||||
case 'darwin':
|
||||
return 'darwin';
|
||||
default:
|
||||
return 'linux';
|
||||
}
|
||||
}
|
||||
|
||||
function getArchTag(): string {
|
||||
if (process.arch === 'x64') {
|
||||
return 'amd64';
|
||||
}
|
||||
|
||||
if (process.arch === 'arm64') {
|
||||
return 'arm64';
|
||||
}
|
||||
|
||||
return process.arch;
|
||||
}
|
||||
|
||||
export function buildUserAgent(version: string): string {
|
||||
return `antigravity/${version} ${getPlatformTag()}/${getArchTag()}`;
|
||||
}
|
||||
|
||||
function compareSemverVersions(left: string, right: string): number {
|
||||
const leftParts = left.split('.').map((part) => Number.parseInt(part, 10));
|
||||
const rightParts = right.split('.').map((part) => Number.parseInt(part, 10));
|
||||
const maxLength = Math.max(leftParts.length, rightParts.length);
|
||||
|
||||
for (let index = 0; index < maxLength; index += 1) {
|
||||
const leftValue = leftParts[index] ?? 0;
|
||||
const rightValue = rightParts[index] ?? 0;
|
||||
|
||||
if (leftValue > rightValue) {
|
||||
return 1;
|
||||
}
|
||||
if (leftValue < rightValue) {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
function extractHighestSemver(text: string): string | null {
|
||||
const matches = text.match(VERSION_REGEX);
|
||||
if (!matches || matches.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
let best = matches[0];
|
||||
for (const candidate of matches) {
|
||||
if (compareSemverVersions(candidate, best) > 0) {
|
||||
best = candidate;
|
||||
}
|
||||
}
|
||||
|
||||
return best;
|
||||
}
|
||||
|
||||
function shouldSkipRemoteVersionLookup(): boolean {
|
||||
return process.env.NODE_ENV === 'test' || process.env.VITEST === 'true';
|
||||
}
|
||||
|
||||
export function resolveLocalInstalledVersion(): string | null {
|
||||
let localVersionString: string;
|
||||
try {
|
||||
localVersionString = getAntigravityVersion().shortVersion;
|
||||
const rawVersion = normalizeNonEmptyString(localVersionString);
|
||||
if (!rawVersion) {
|
||||
return null;
|
||||
}
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
|
||||
return extractHighestSemver(localVersionString);
|
||||
}
|
||||
|
||||
async function fetchTextPayload(url: string): Promise<string | null> {
|
||||
try {
|
||||
const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION;
|
||||
const response = await axios.get<string>(url, {
|
||||
timeout: DEFAULT_REMOTE_TIMEOUT_MS,
|
||||
responseType: 'text',
|
||||
headers: {
|
||||
'User-Agent': buildUserAgent(discoveryVersion),
|
||||
},
|
||||
});
|
||||
|
||||
if (typeof response.data === 'string') {
|
||||
return response.data;
|
||||
}
|
||||
|
||||
return JSON.stringify(response.data);
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async function resolveRemoteVersion(): Promise<{
|
||||
source: Extract<UserAgentSource, 'remote' | 'changelog'>;
|
||||
version: string;
|
||||
} | null> {
|
||||
if (shouldSkipRemoteVersionLookup()) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const remotePayload = await fetchTextPayload(REMOTE_VERSION_URL);
|
||||
if (remotePayload) {
|
||||
const parsed = extractHighestSemver(remotePayload);
|
||||
if (parsed) {
|
||||
return {
|
||||
source: 'remote',
|
||||
version: parsed,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
const changelogPayload = await fetchTextPayload(CHANGELOG_URL);
|
||||
if (!changelogPayload) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const parsed = extractHighestSemver(changelogPayload);
|
||||
if (!parsed) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
source: 'changelog',
|
||||
version: parsed,
|
||||
};
|
||||
}
|
||||
|
||||
async function resolveDefaultUserAgentResolution(): Promise<UserAgentResolution> {
|
||||
let bestVersion = FALLBACK_VERSION;
|
||||
let bestSource: UserAgentSource = 'fallback';
|
||||
|
||||
const localVersion = resolveLocalInstalledVersion();
|
||||
if (localVersion && compareSemverVersions(localVersion, bestVersion) > 0) {
|
||||
bestVersion = localVersion;
|
||||
bestSource = 'local';
|
||||
}
|
||||
|
||||
const remoteVersion = await resolveRemoteVersion();
|
||||
if (remoteVersion && compareSemverVersions(remoteVersion.version, bestVersion) > 0) {
|
||||
bestVersion = remoteVersion.version;
|
||||
bestSource = remoteVersion.source;
|
||||
}
|
||||
|
||||
return {
|
||||
source: bestSource,
|
||||
version: bestVersion,
|
||||
userAgent: buildUserAgent(bestVersion),
|
||||
};
|
||||
}
|
||||
|
||||
async function getResolvedDefaultRequestUserAgent(): Promise<string> {
|
||||
if (cachedUserAgentResolution) {
|
||||
return cachedUserAgentResolution.userAgent;
|
||||
}
|
||||
|
||||
if (pendingUserAgentResolution) {
|
||||
return (await pendingUserAgentResolution).userAgent;
|
||||
}
|
||||
|
||||
pendingUserAgentResolution = resolveDefaultUserAgentResolution()
|
||||
.then((resolved) => {
|
||||
cachedUserAgentResolution = resolved;
|
||||
logger.info(
|
||||
`Default request User-Agent resolved (source=${resolved.source}, version=${resolved.version}): ${resolved.userAgent}`,
|
||||
);
|
||||
return resolved;
|
||||
})
|
||||
.finally(() => {
|
||||
pendingUserAgentResolution = null;
|
||||
});
|
||||
|
||||
return (await pendingUserAgentResolution).userAgent;
|
||||
}
|
||||
|
||||
export async function resolveRequestUserAgent(): Promise<string> {
|
||||
return await getResolvedDefaultRequestUserAgent();
|
||||
}
|
||||
@@ -10,6 +10,8 @@ const mockTokenManager = {
|
||||
getNextToken: vi.fn(),
|
||||
markAsRateLimited: vi.fn(),
|
||||
markAsForbidden: vi.fn(),
|
||||
markFromUpstreamError: vi.fn(),
|
||||
recordParityError: vi.fn(),
|
||||
};
|
||||
const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
|
||||
|
||||
|
||||
@@ -8,6 +8,10 @@ const mockTokenManager = {
|
||||
getNextToken: vi.fn(),
|
||||
markAsRateLimited: vi.fn(),
|
||||
markAsForbidden: vi.fn(),
|
||||
markFromUpstreamError: vi.fn(),
|
||||
recordParityError: vi.fn(),
|
||||
getModelOutputLimitForAccount: vi.fn(),
|
||||
getModelThinkingBudgetForAccount: vi.fn(),
|
||||
};
|
||||
const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
|
||||
|
||||
@@ -27,7 +31,7 @@ class TestableProxyService extends ProxyService {
|
||||
}
|
||||
|
||||
public testModelHeaders(model: string): Record<string, string> {
|
||||
return (this as any).buildModelSpecificHeaders(model);
|
||||
return (this as any).createModelSpecificHeaders(model);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,7 +58,7 @@ describe('ProxyService Empty Stream Retry Logic', () => {
|
||||
|
||||
it('classifies retry matrix consistently', () => {
|
||||
const service = new TestableProxyService();
|
||||
const classify = (message: string) => (service as any).classifyUpstreamError(message);
|
||||
const classify = (message: string) => (service as any).classifyUpstreamFailure(message);
|
||||
|
||||
expect(classify('401 unauthorized token')).toEqual({
|
||||
retry: true,
|
||||
@@ -401,6 +405,22 @@ describe('ProxyService Empty Stream Retry Logic', () => {
|
||||
expect(internalPayload).not.toHaveProperty('sessionId');
|
||||
});
|
||||
|
||||
it('normalizes Gemini 3.1 preview alias to Gemini 3.1 Pro High for upstream', async () => {
|
||||
const service = new TestableProxyService();
|
||||
mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1'));
|
||||
mockGeminiClient.generateInternal.mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }],
|
||||
usageMetadata: { totalTokenCount: 5 },
|
||||
});
|
||||
|
||||
await service.handleGeminiGenerateContent('models/gemini-3.1-pro-preview', {
|
||||
contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
|
||||
} as any);
|
||||
|
||||
const internalPayload = mockGeminiClient.generateInternal.mock.calls[0][0];
|
||||
expect(internalPayload.model).toBe('gemini-3.1-pro-high');
|
||||
});
|
||||
|
||||
it('strips non-parity Gemini usage metadata fields', async () => {
|
||||
const service = new TestableProxyService();
|
||||
mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1'));
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { RateLimitReason, RateLimitTracker } from '../../server/modules/proxy/rate-limit-tracker';
|
||||
|
||||
describe('RateLimitTracker parity replay', () => {
|
||||
it('uses Retry-After header before body/default', () => {
|
||||
const tracker = new RateLimitTracker();
|
||||
const info = tracker.parseAndMarkFromError({
|
||||
accountId: 'acc-1',
|
||||
status: 429,
|
||||
retryAfter: '30',
|
||||
body: JSON.stringify({
|
||||
error: {
|
||||
details: [{ reason: 'RATE_LIMIT_EXCEEDED' }],
|
||||
retry_after: 1,
|
||||
},
|
||||
}),
|
||||
model: 'gemini-2.5-pro',
|
||||
backoffSteps: [60, 300, 1800, 7200],
|
||||
});
|
||||
|
||||
expect(info).not.toBeNull();
|
||||
expect(info?.reason).toBe(RateLimitReason.RateLimitExceeded);
|
||||
expect(info?.retryAfterSec).toBe(30);
|
||||
expect(tracker.isRateLimited('acc-1')).toBe(true);
|
||||
});
|
||||
|
||||
it('uses model-level key for quota exhausted', () => {
|
||||
const tracker = new RateLimitTracker();
|
||||
tracker.parseAndMarkFromError({
|
||||
accountId: 'acc-2',
|
||||
status: 429,
|
||||
body: JSON.stringify({
|
||||
error: {
|
||||
details: [
|
||||
{
|
||||
reason: 'QUOTA_EXHAUSTED',
|
||||
metadata: { quotaResetDelay: '42s' },
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
model: 'gemini-2.5-flash',
|
||||
backoffSteps: [60, 300, 1800, 7200],
|
||||
});
|
||||
|
||||
expect(tracker.isRateLimited('acc-2', 'gemini-2.5-flash')).toBe(true);
|
||||
expect(tracker.isRateLimited('acc-2', 'gemini-2.5-pro')).toBe(false);
|
||||
});
|
||||
|
||||
it('uses backoff steps when no header/body retry hint', () => {
|
||||
const tracker = new RateLimitTracker();
|
||||
const steps = [60, 300, 1800, 7200];
|
||||
|
||||
const first = tracker.parseAndMarkFromError({
|
||||
accountId: 'acc-3',
|
||||
status: 429,
|
||||
body: JSON.stringify({
|
||||
error: { details: [{ reason: 'QUOTA_EXHAUSTED' }] },
|
||||
}),
|
||||
backoffSteps: steps,
|
||||
});
|
||||
const second = tracker.parseAndMarkFromError({
|
||||
accountId: 'acc-3',
|
||||
status: 429,
|
||||
body: JSON.stringify({
|
||||
error: { details: [{ reason: 'QUOTA_EXHAUSTED' }] },
|
||||
}),
|
||||
backoffSteps: steps,
|
||||
});
|
||||
|
||||
expect(first?.retryAfterSec).toBe(60);
|
||||
expect(second?.retryAfterSec).toBe(300);
|
||||
});
|
||||
|
||||
it('parses MODEL_CAPACITY_EXHAUSTED and RetryInfo.retryDelay from 503 payload', () => {
|
||||
const tracker = new RateLimitTracker();
|
||||
const info = tracker.parseAndMarkFromError({
|
||||
accountId: 'acc-4',
|
||||
status: 503,
|
||||
model: 'gemini-3.1-pro-high',
|
||||
body: JSON.stringify({
|
||||
error: {
|
||||
code: 503,
|
||||
message: 'No capacity available for model gemini-3.1-pro-high on the server',
|
||||
status: 'UNAVAILABLE',
|
||||
details: [
|
||||
{
|
||||
'@type': 'type.googleapis.com/google.rpc.ErrorInfo',
|
||||
reason: 'MODEL_CAPACITY_EXHAUSTED',
|
||||
domain: 'cloudcode-pa.googleapis.com',
|
||||
metadata: { model: 'gemini-3.1-pro-high' },
|
||||
},
|
||||
{
|
||||
'@type': 'type.googleapis.com/google.rpc.RetryInfo',
|
||||
retryDelay: '30s',
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
backoffSteps: [60, 300, 1800, 7200],
|
||||
});
|
||||
|
||||
expect(info).not.toBeNull();
|
||||
expect(info?.reason).toBe(RateLimitReason.ModelCapacityExhausted);
|
||||
expect(info?.retryAfterSec).toBe(30);
|
||||
expect(tracker.isRateLimited('acc-4')).toBe(true);
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user