From 20fffff29c47d1147e9b429884d9e8b613e7da6f Mon Sep 17 00:00:00 2001 From: Draculabo Date: Fri, 6 Mar 2026 21:56:30 +0800 Subject: [PATCH] feat(proxy): add upstream error and rate-limit tracking support --- .../modules/proxy/clients/upstream-error.ts | 23 + .../modules/proxy/rate-limit-tracker.ts | 426 ++++++++++++++++++ .../modules/proxy/request-user-agent.ts | 222 +++++++++ src/tests/unit/proxy-parity-fixtures.test.ts | 2 + src/tests/unit/proxy-retry-mock.test.ts | 24 +- src/tests/unit/rate-limit-tracker.test.ts | 108 +++++ 6 files changed, 803 insertions(+), 2 deletions(-) create mode 100644 src/server/modules/proxy/clients/upstream-error.ts create mode 100644 src/server/modules/proxy/rate-limit-tracker.ts create mode 100644 src/server/modules/proxy/request-user-agent.ts create mode 100644 src/tests/unit/rate-limit-tracker.test.ts diff --git a/src/server/modules/proxy/clients/upstream-error.ts b/src/server/modules/proxy/clients/upstream-error.ts new file mode 100644 index 0000000..e637762 --- /dev/null +++ b/src/server/modules/proxy/clients/upstream-error.ts @@ -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; + } +} + diff --git a/src/server/modules/proxy/rate-limit-tracker.ts b/src/server/modules/proxy/rate-limit-tracker.ts new file mode 100644 index 0000000..9e3a753 --- /dev/null +++ b/src/server/modules/proxy/rate-limit-tracker.ts @@ -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(); + private readonly failureCounts = new Map(); + + 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; + } +} diff --git a/src/server/modules/proxy/request-user-agent.ts b/src/server/modules/proxy/request-user-agent.ts new file mode 100644 index 0000000..f7736c2 --- /dev/null +++ b/src/server/modules/proxy/request-user-agent.ts @@ -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 | 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 { + try { + const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION; + const response = await axios.get(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; + 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 { + 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 { + 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 { + return await getResolvedDefaultRequestUserAgent(); +} diff --git a/src/tests/unit/proxy-parity-fixtures.test.ts b/src/tests/unit/proxy-parity-fixtures.test.ts index b626c6b..16c8159 100644 --- a/src/tests/unit/proxy-parity-fixtures.test.ts +++ b/src/tests/unit/proxy-parity-fixtures.test.ts @@ -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() }; diff --git a/src/tests/unit/proxy-retry-mock.test.ts b/src/tests/unit/proxy-retry-mock.test.ts index 1475b56..0c0d695 100644 --- a/src/tests/unit/proxy-retry-mock.test.ts +++ b/src/tests/unit/proxy-retry-mock.test.ts @@ -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 { - 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')); diff --git a/src/tests/unit/rate-limit-tracker.test.ts b/src/tests/unit/rate-limit-tracker.test.ts new file mode 100644 index 0000000..f083d82 --- /dev/null +++ b/src/tests/unit/rate-limit-tracker.test.ts @@ -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); + }); +});