diff --git a/src/lib/antigravity/ClaudeRequestMapper.ts b/src/lib/antigravity/ClaudeRequestMapper.ts index 880240b..8a1ad63 100644 --- a/src/lib/antigravity/ClaudeRequestMapper.ts +++ b/src/lib/antigravity/ClaudeRequestMapper.ts @@ -370,18 +370,6 @@ function detectsNetworkingTool(tools?: (Tool | GeminiToolDeclaration)[]): boolea return false; } -/** - * Inject Google Search Tool - * Adds googleSearch tool to the request - */ -function supportsMixedTools(mappedModel?: string): boolean { - if (!mappedModel) { - return false; - } - const modelLower = mappedModel.toLowerCase(); - return modelLower.includes('gemini-3'); -} - function injectGoogleSearchTool(body: { tools?: GeminiToolDeclaration[] }, mappedModel?: string) { if (!body.tools) { body.tools = []; @@ -389,9 +377,9 @@ function injectGoogleSearchTool(body: { tools?: GeminiToolDeclaration[] }, mappe const toolsArr = body.tools; const hasFunctions = toolsArr.some((t) => t.functionDeclarations); - if (hasFunctions && !supportsMixedTools(mappedModel)) { + if (hasFunctions) { logger.info( - 'Skipping googleSearch injection due to existing functionDeclarations on old model', + `[Claude-Request] Skipping googleSearch injection for ${mappedModel ?? 'unknown-model'} because functionDeclarations are present (v1internal incompatible)`, ); return; } @@ -631,13 +619,9 @@ function buildTools( if (functionDeclarations.length > 0) { toolList.push({ functionDeclarations }); if (hasGoogleSearch) { - if (supportsMixedTools(mappedModel)) { - toolList.push({ googleSearch: {} }); - } else { - logger.info( - `[Claude-Request] Skipping googleSearch injection for ${mappedModel} due to existing functionDeclarations on old model`, - ); - } + logger.info( + `[Claude-Request] Skipping googleSearch injection for ${mappedModel} because functionDeclarations are present (v1internal incompatible)`, + ); } } else if (hasGoogleSearch) { toolList.push({ googleSearch: {} }); diff --git a/src/lib/antigravity/ModelMapping.ts b/src/lib/antigravity/ModelMapping.ts index 5c2d470..4628b42 100644 --- a/src/lib/antigravity/ModelMapping.ts +++ b/src/lib/antigravity/ModelMapping.ts @@ -59,6 +59,10 @@ const CLAUDE_TO_GEMINI: Record = { 'gemini-3.1-pro-preview': 'gemini-3.1-pro-high', 'gemini-3.1-pro': 'gemini-3.1-pro-high', 'gemini-3.0-pro': 'gemini-3.1-pro-high', + 'gemini-3-pro': 'gemini-3-pro-preview', + 'gemini-3-pro-preview': 'gemini-3-pro-preview', + 'gemini-3-pro-low': 'gemini-3-pro-low', + 'gemini-3-pro-high': 'gemini-3-pro-high', 'gemini-2.5-flash': 'gemini-3-flash', 'gemini-2.5-pro': 'gemini-3.1-pro-high', 'gemini-2.0-flash': 'gemini-3-flash', @@ -79,11 +83,10 @@ export const MODEL_LIST_CREATED_AT = 1770652800; export const MODEL_LIST_OWNER = 'antigravity'; const GEMINI_MODEL_ALIASES: Record = { - 'gemini-3.1-pro': 'gemini-3.1-pro-high', - 'gemini-3.0-pro': 'gemini-3.1-pro-high', - 'gemini-3.1-pro-preview': 'gemini-3.1-pro-high', 'gemini-3-pro-image-preview': 'gemini-3-pro-image', 'gemini-3-flash-preview': 'gemini-3-flash', + 'gemini-3.1-pro-preview': 'gemini-3.1-pro-high', + 'gemini-3.1-pro': 'gemini-3.1-pro-high', }; export function getSupportedModels(): string[] { diff --git a/src/lib/antigravity/QuotaService.ts b/src/lib/antigravity/QuotaService.ts index d9b8bdb..ea42e21 100644 --- a/src/lib/antigravity/QuotaService.ts +++ b/src/lib/antigravity/QuotaService.ts @@ -3,7 +3,11 @@ import { QuotaData, LoadProjectResponse, QuotaApiResponse } from './types'; import { logger } from '../../utils/logger'; // Constants -const QUOTA_API_URL = 'https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels'; +const QUOTA_API_ENDPOINTS = [ + 'https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:fetchAvailableModels', + 'https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels', + 'https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels', +] as const; const CLOUD_CODE_BASE_URL = 'https://cloudcode-pa.googleapis.com'; const USER_AGENT = 'antigravity/1.11.3 Darwin/arm64'; // Keeping the same UA as source @@ -83,15 +87,16 @@ export class QuotaService { const finalProjectId = projectId; const client = this.createClient(); - const payload = { project: finalProjectId }; - const url = QUOTA_API_URL; - const maxRetries = 3; + const payload = finalProjectId ? { project: finalProjectId } : {}; + let lastError: Error | null = null; - logger.info(`Sending quota request to ${url}`); + for (let endpointIndex = 0; endpointIndex < QUOTA_API_ENDPOINTS.length; endpointIndex++) { + const endpoint = QUOTA_API_ENDPOINTS[endpointIndex]; + const hasNextEndpoint = endpointIndex + 1 < QUOTA_API_ENDPOINTS.length; + logger.info(`Sending quota request to ${endpoint}`); - for (let attempt = 1; attempt <= maxRetries; attempt++) { try { - const response = await client.post(url, payload, { + const response = await client.post(endpoint, payload, { headers: { Authorization: `Bearer ${accessToken}`, 'User-Agent': USER_AGENT, @@ -127,8 +132,14 @@ export class QuotaService { } } - return { quotaData, projectId: projectId }; + if (endpointIndex > 0) { + logger.info(`Quota API fallback succeeded at endpoint #${endpointIndex + 1}`); + } + + return { quotaData, projectId }; } catch (error: any) { + let shouldFallback = true; + if (axios.isAxiosError(error)) { const status = error.response?.status; let text = ''; @@ -147,23 +158,36 @@ export class QuotaService { isForbidden: true, subscriptionTier: subscriptionTier, }, - projectId: projectId, + projectId, }; } - logger.warn(`API Error: ${status} - ${text} (Attempt ${attempt}/${maxRetries})`); + if (hasNextEndpoint && (status === 429 || (typeof status === 'number' && status >= 500))) { + logger.warn( + `Quota API ${endpoint} returned ${status}, falling back to next endpoint`, + ); + lastError = new Error(`HTTP ${status} - ${text}`); + await new Promise((r) => setTimeout(r, 1000)); + continue; + } + + logger.warn(`API Error: ${status} - ${text}`); + lastError = new Error(`HTTP ${status} - ${text}`); + shouldFallback = typeof status !== 'number'; } else { - logger.warn(`Request Failed: ${error.message} (Attempt ${attempt}/${maxRetries})`); + logger.warn(`Request Failed at ${endpoint}: ${error.message}`); + lastError = error instanceof Error ? error : new Error(String(error)); } - if (attempt < maxRetries) { + if (hasNextEndpoint && shouldFallback) { + logger.warn(`Quota API request failed at ${endpoint}, falling back to next endpoint`); await new Promise((r) => setTimeout(r, 1000)); } else { - throw new Error(`Quota query failed: ${error.message}`); + throw lastError ?? new Error(`Quota query failed: ${error.message}`); } } } - throw new Error('Unknown error in fetchQuota'); // Should not reach here + throw lastError ?? new Error('Unknown error in fetchQuota'); } } diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index bee81bf..2d50fab 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/src/server/modules/proxy/proxy.service.ts @@ -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); }); }); diff --git a/src/server/modules/proxy/request-user-agent.ts b/src/server/modules/proxy/request-user-agent.ts index f7736c2..f8b77ba 100644 --- a/src/server/modules/proxy/request-user-agent.ts +++ b/src/server/modules/proxy/request-user-agent.ts @@ -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; diff --git a/src/server/modules/proxy/token-manager.service.ts b/src/server/modules/proxy/token-manager.service.ts index 10df5ee..6a35ebe 100644 --- a/src/server/modules/proxy/token-manager.service.ts +++ b/src/server/modules/proxy/token-manager.service.ts @@ -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 { 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 { const allModels = new Set(); for (const tokenData of this.tokens.values()) { @@ -937,6 +967,102 @@ export class TokenManagerService implements OnModuleInit { return allModels; } + private getAvailableModelsFromToken(tokenData: TokenData): Set { + const availableModels = new Set(); + + 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(); + 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); diff --git a/src/tests/unit/proxy-parity-fixtures.test.ts b/src/tests/unit/proxy-parity-fixtures.test.ts index 16c8159..4b2d6d7 100644 --- a/src/tests/unit/proxy-parity-fixtures.test.ts +++ b/src/tests/unit/proxy-parity-fixtures.test.ts @@ -97,4 +97,5 @@ describe('Proxy Parity Fixtures', () => { expect(output).toContain(token); } }); + }); diff --git a/src/tests/unit/proxy-retry-mock.test.ts b/src/tests/unit/proxy-retry-mock.test.ts index 0c0d695..7051404 100644 --- a/src/tests/unit/proxy-retry-mock.test.ts +++ b/src/tests/unit/proxy-retry-mock.test.ts @@ -12,6 +12,7 @@ const mockTokenManager = { recordParityError: vi.fn(), getModelOutputLimitForAccount: vi.fn(), getModelThinkingBudgetForAccount: vi.fn(), + resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model), }; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; @@ -651,4 +652,43 @@ describe('ProxyService Protocol Parity Fixtures', () => { expect(output).toContain('"content":"final answer"'); expect(output).toContain('data: [DONE]'); }); + + it('propagates OpenAI-compatible upstream stream errors instead of completing with [DONE]', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + const observable = (service as any).processStreamResponse(stream, 'gpt-4o-mini'); + const chunks: string[] = []; + + const streamResult = await new Promise<{ error?: Error; completed: boolean }>((resolve) => { + observable.subscribe({ + next: (chunk: string) => { + chunks.push(chunk); + }, + error: (error: unknown) => + resolve({ + completed: false, + error: error instanceof Error ? error : new Error(String(error)), + }), + complete: () => resolve({ completed: true }), + }); + + const payload = JSON.stringify({ + candidates: [ + { + content: { + parts: [{ text: 'partial output' }], + }, + }, + ], + }); + + stream.emit('data', Buffer.from(`data: ${payload}\n`)); + stream.emit('error', new Error('socket hang up')); + }); + + expect(streamResult.completed).not.toBe(true); + expect(streamResult.error?.message).toContain('socket hang up'); + expect(chunks.join('')).toContain('"content":"partial output"'); + expect(chunks.join('')).not.toContain('data: [DONE]'); + }); }); diff --git a/src/tests/unit/token-manager-parity-scheduling.test.ts b/src/tests/unit/token-manager-parity-scheduling.test.ts index 075ae51..29a7a1a 100644 --- a/src/tests/unit/token-manager-parity-scheduling.test.ts +++ b/src/tests/unit/token-manager-parity-scheduling.test.ts @@ -1,7 +1,8 @@ -import { beforeEach, describe, expect, it } from 'vitest'; +import { beforeEach, describe, expect, it, vi } 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'; +import { GoogleAPIService } from '../../services/GoogleAPIService'; function createProxyConfig(overrides: Partial): ProxyConfig { return { @@ -56,6 +57,96 @@ describe('TokenManagerService parity scheduling replay', () => { seedTokens(service); }); + it('rewrites gemini pro model to first available account candidate', () => { + const nowSec = Math.floor(Date.now() / 1000); + (service as any).tokens = new Map([ + [ + 'acc-1', + { + account_id: 'acc-1', + email: 'acc-1@test.dev', + 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', + model_quotas: { + 'gemini-3.1-pro-low': 80, + }, + model_limits: {}, + model_reset_times: {}, + model_forwarding_rules: {}, + }, + ], + ]); + + const resolved = service.resolveDynamicModelForAccount('acc-1', 'gemini-3-pro'); + expect(resolved).toBe('gemini-3.1-pro-low'); + }); + + it('keeps original model when dynamic rewrite is not applicable', () => { + const resolved = service.resolveDynamicModelForAccount('acc-1', 'gemini-3-flash'); + expect(resolved).toBe('gemini-3-flash'); + }); + + it('passes oauth_client_key when refreshing token and persists refreshed key', async () => { + const nowSec = Math.floor(Date.now() / 1000); + const tokenData = { + account_id: 'acc-1', + email: 'acc-1@test.dev', + access_token: 'token-1', + refresh_token: 'refresh-1', + oauth_client_key: 'custom-client', + upstream_proxy_url: 'http://127.0.0.1:8080', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: nowSec - 1, + project_id: 'project-1', + session_id: 'session-1', + model_quotas: {}, + model_limits: {}, + model_reset_times: {}, + model_forwarding_rules: {}, + }; + + (service as any).tokens = new Map([['acc-1', tokenData]]); + + const refreshSpy = vi.spyOn(GoogleAPIService, 'refreshAccessToken').mockResolvedValue({ + access_token: 'token-new', + expires_in: 7200, + token_type: 'Bearer', + oauth_client_key: 'custom-fallback', + }); + const persistSpy = vi.spyOn(service as any, 'persistTokenState').mockResolvedValue(undefined); + + const selected = await (service as any).finalizeSelectedToken('acc-1', tokenData, nowSec); + + expect(refreshSpy).toHaveBeenCalledWith( + 'refresh-1', + 'http://127.0.0.1:8080', + 'custom-client', + ); + expect(selected?.token.oauth_client_key).toBe('custom-fallback'); + expect((service as any).tokens.get('acc-1')?.oauth_client_key).toBe('custom-fallback'); + + refreshSpy.mockRestore(); + persistSpy.mockRestore(); + }); + + it('keeps oauth_client_key unset for legacy account refreshed by enterprise client', () => { + const normalized = (service as any).normalizeRefreshedOauthClientKey( + { + oauth_client_key: undefined, + project_id: undefined, + }, + 'antigravity_enterprise', + ); + + expect(normalized).toBeUndefined(); + }); + it('prioritizes preferred account in parity mode', async () => { setServerConfig( createProxyConfig({