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

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: {} });

src/lib/antigravity/ModelMapping.ts

@@ -59,6 +59,10 @@ const CLAUDE_TO_GEMINI: Record<string, string> = {
   '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<string, string> = {
-  '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[] {

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<QuotaApiResponse>(url, payload, {
+        const response = await client.post<QuotaApiResponse>(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');
   }
 }

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);
       });
     });

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;
 

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<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);

src/tests/unit/proxy-parity-fixtures.test.ts

@@ -97,4 +97,5 @@ describe('Proxy Parity Fixtures', () => {
       expect(output).toContain(token);
     }
   });
+
 });

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]');
+  });
 });

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>): 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: '[email protected]',
+          access_token: 'token-1',
+          refresh_token: 'refresh-1',
+          token_type: 'Bearer',
+          expires_in: 3600,
+          expiry_timestamp: nowSec + 3600,
+          project_id: 'project-1',
+          session_id: 'session-1',
+          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: '[email protected]',
+      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({