feat(proxy): improve parity, rate limiting, and token scheduling

src/ipc/config/handlers.ts

@@ -2,6 +2,7 @@ import { AppConfig } from '../../types/config';
 import { ConfigManager } from './manager';
 import { syncAutoStart } from '../../utils/autoStart';
 import { logger } from '../../utils/logger';
+import { setServerConfig } from '../../server/server-config';
 
 export function loadConfig(): AppConfig {
   return ConfigManager.loadConfig();
@@ -13,6 +14,7 @@ export async function saveConfig(config: AppConfig): Promise<void> {
   // For now just save
   const previous = ConfigManager.getCachedConfig() ?? ConfigManager.loadConfig();
   await ConfigManager.saveConfig(config);
+  setServerConfig(config.proxy);
   logger.setErrorReportingEnabled(config.error_reporting_enabled);
   if (previous.auto_startup !== config.auto_startup) {
     syncAutoStart(config);

src/server/modules/proxy/clients/gemini.client.ts

@@ -4,11 +4,12 @@ import { isNil } from 'lodash-es';
 import { GeminiRequest, GeminiResponse } from '../interfaces/request-interfaces';
 import { GeminiInternalRequest } from '../../../../lib/antigravity/types';
 import { getServerConfig } from '../../../server-config';
+import { resolveRequestUserAgent } from '../request-user-agent';
+import { UpstreamRequestError } from './upstream-error';
 
 @Injectable()
 export class GeminiClient {
   private readonly logger = new Logger(GeminiClient.name);
-  private readonly defaultRequestUserAgent = 'antigravity/1.11.9 windows/amd64';
   // Default to v1beta for most features
   private readonly baseUrl = 'https://generativelanguage.googleapis.com/v1beta';
   private readonly defaultInternalBaseUrls = [
@@ -23,7 +24,7 @@ export class GeminiClient {
     upstreamProxyUrl?: string,
   ): Promise<NodeJS.ReadableStream> {
     const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`;
-    const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
+    const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
 
     try {
       const response = await axios.post(url, content, {
@@ -39,7 +40,14 @@ export class GeminiClient {
     } catch (error) {
       if (axios.isAxiosError(error)) {
         this.logger.error(`Gemini stream request failed: ${error.message}`);
-        throw new Error(error.response?.data?.error?.message || error.message);
+        throw new UpstreamRequestError({
+          message: error.response?.data?.error?.message || error.message,
+          status: error.response?.status,
+          headers: {
+            retryAfter: this.extractRetryAfterHeader(error.response?.headers),
+          },
+          body: this.describeAxiosErrorData(error.response?.data),
+        });
       }
       this.throwAsCleanError(error);
     }
@@ -52,7 +60,7 @@ export class GeminiClient {
     upstreamProxyUrl?: string,
   ): Promise<GeminiResponse> {
     const url = `${this.baseUrl}/models/${model}:generateContent`;
-    const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
+    const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
 
     try {
       const response = await axios.post<GeminiResponse>(url, content, {
@@ -69,7 +77,14 @@ export class GeminiClient {
         this.logger.error(
           `Gemini request failed: ${error.message} - ${this.safeStringify(error.response?.data)}`,
         );
-        throw new Error(error.response?.data?.error?.message || error.message);
+        throw new UpstreamRequestError({
+          message: error.response?.data?.error?.message || error.message,
+          status: error.response?.status,
+          headers: {
+            retryAfter: this.extractRetryAfterHeader(error.response?.headers),
+          },
+          body: this.describeAxiosErrorData(error.response?.data),
+        });
       }
       this.throwAsCleanError(error);
     }
@@ -83,7 +98,7 @@ export class GeminiClient {
     upstreamProxyUrl?: string,
     extraHeaders?: Record<string, string>,
   ): Promise<NodeJS.ReadableStream> {
-    const response = await this.requestWithEndpointFailover<NodeJS.ReadableStream>(
+    const response = await this.executeRequestWithEndpointFailover<NodeJS.ReadableStream>(
       ':streamGenerateContent?alt=sse',
       body,
       accessToken,
@@ -104,7 +119,7 @@ export class GeminiClient {
     upstreamProxyUrl?: string,
     extraHeaders?: Record<string, string>,
   ): Promise<GeminiResponse> {
-    const response = await this.requestWithEndpointFailover<
+    const response = await this.executeRequestWithEndpointFailover<
       GeminiResponse | { response: GeminiResponse }
     >(
       ':generateContent',
@@ -143,7 +158,7 @@ export class GeminiClient {
     return Math.max(1, timeoutSeconds) * 1000;
   }
 
-  private shouldSwitchEndpointOnError(error: unknown): boolean {
+  private shouldFailoverToNextEndpoint(error: unknown): boolean {
     if (!axios.isAxiosError(error)) {
       return false;
     }
@@ -162,7 +177,7 @@ export class GeminiClient {
     return status === 408 || status === 429 || status >= 500;
   }
 
-  private async requestWithEndpointFailover<T>(
+  private async executeRequestWithEndpointFailover<T>(
     path: string,
     body: GeminiInternalRequest,
     accessToken: string,
@@ -173,8 +188,8 @@ export class GeminiClient {
   ): Promise<AxiosResponse<T>> {
     const baseUrls = this.getInternalBaseUrls();
     const timeout = this.getInternalTimeoutMs();
-    const requestUserAgent = process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent;
+    const requestUserAgent = await resolveRequestUserAgent();
-    const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
+    const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
     let lastError: unknown = null;
 
     for (let index = 0; index < baseUrls.length; index++) {
@@ -197,8 +212,8 @@ export class GeminiClient {
         lastError = error;
         const hasNextEndpoint = index < baseUrls.length - 1;
 
-        if (!hasNextEndpoint || !this.shouldSwitchEndpointOnError(error)) {
+        if (!hasNextEndpoint || !this.shouldFailoverToNextEndpoint(error)) {
-          await this.handleAxiosError(error, operation);
+          await this.throwUpstreamRequestError(error, operation);
         }
 
         this.logger.warn(
@@ -209,11 +224,11 @@ export class GeminiClient {
       }
     }
 
-    await this.handleAxiosError(lastError, operation);
+    await this.throwUpstreamRequestError(lastError, operation);
     throw new Error(`[${operation}] unexpected control flow after upstream error handling`);
   }
 
-  private async handleAxiosError(error: unknown, operation: string): Promise<never> {
+  private async throwUpstreamRequestError(error: unknown, operation: string): Promise<never> {
     if (axios.isAxiosError(error)) {
       const responseData = error.response?.data;
       const upstreamMessage = await this.extractAxiosErrorMessage(responseData);
@@ -222,7 +237,14 @@ export class GeminiClient {
           responseData,
         )}`,
       );
-      throw new Error(upstreamMessage || error.message);
+      throw new UpstreamRequestError({
+        message: upstreamMessage || error.message,
+        status: error.response?.status,
+        headers: {
+          retryAfter: this.extractRetryAfterHeader(error.response?.headers),
+        },
+        body: this.describeAxiosErrorData(responseData),
+      });
     }
     this.throwAsCleanError(error);
   }
@@ -346,7 +368,9 @@ export class GeminiClient {
     return this.safeStringify(responseData);
   }
 
-  private resolveAxiosProxy(upstreamProxyUrl?: string): AxiosProxyConfig | false | undefined {
+  private resolveUpstreamAxiosProxy(
+    upstreamProxyUrl?: string,
+  ): AxiosProxyConfig | false | undefined {
     const config = getServerConfig();
     const configuredProxyUrl =
       upstreamProxyUrl ||
@@ -382,6 +406,24 @@ export class GeminiClient {
     }
   }
 
+  private extractRetryAfterHeader(headers: unknown): string | undefined {
+    if (!headers || typeof headers !== 'object') {
+      return undefined;
+    }
+
+    const retryAfter = (headers as Record<string, unknown>)['retry-after'];
+    if (typeof retryAfter === 'string' && retryAfter.trim() !== '') {
+      return retryAfter.trim();
+    }
+    if (Array.isArray(retryAfter) && retryAfter.length > 0) {
+      const first = retryAfter[0];
+      if (typeof first === 'string' && first.trim() !== '') {
+        return first.trim();
+      }
+    }
+    return undefined;
+  }
+
   private throwAsCleanError(error: unknown): never {
     // Re-throw as clean Error to avoid circular reference issues.
     throw error instanceof Error ? new Error(error.message) : new Error(String(error));

src/server/modules/proxy/gemini.controller.ts

@@ -8,6 +8,7 @@ import {
   Post,
   Res,
   UseGuards,
+  Optional,
 } from '@nestjs/common';
 import { FastifyReply } from 'fastify';
 import { isEmpty } from 'lodash-es';
@@ -18,6 +19,7 @@ import { ProxyService } from './proxy.service';
 import { GeminiRequest, GeminiResponse } from './interfaces/request-interfaces';
 import { getServerConfig } from '../../server-config';
 import { getAllDynamicModels } from '../../../lib/antigravity/ModelMapping';
+import { TokenManagerService } from './token-manager.service';
 
 type GeminiModelMetadata = {
   name: string;
@@ -35,7 +37,10 @@ type GeminiModelMetadata = {
 @Controller('v1beta')
 @UseGuards(ProxyGuard)
 export class GeminiController {
-  constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {}
+  constructor(
+    @Inject(ProxyService) private readonly proxyService: ProxyService,
+    @Optional() @Inject(TokenManagerService) private readonly tokenManager?: TokenManagerService,
+  ) {}
 
   @Get('models')
   listModels(@Res() res: FastifyReply) {
@@ -70,7 +75,7 @@ export class GeminiController {
     @Body() body: GeminiRequest,
     @Res() res: FastifyReply,
   ) {
-    const parsed = this.parseCombinedModelAction(modelAction);
+    const parsed = this.parseModelActionToken(modelAction);
     if (!parsed) {
       res.status(HttpStatus.BAD_REQUEST).send({
         error: {
@@ -82,7 +87,7 @@ export class GeminiController {
       return;
     }
 
-    await this.dispatchAction(parsed.model, parsed.action, body, res);
+    await this.handleModelActionDispatch(parsed.model, parsed.action, body, res);
   }
 
   @Post('models/:model/countTokens')
@@ -91,10 +96,10 @@ export class GeminiController {
     @Body() body: GeminiRequest,
     @Res() res: FastifyReply,
   ) {
-    await this.dispatchAction(`models/${model}`, 'countTokens', body, res);
+    await this.handleModelActionDispatch(`models/${model}`, 'countTokens', body, res);
   }
 
-  private async dispatchAction(
+  private async handleModelActionDispatch(
     model: string,
     action: string,
     body: GeminiRequest,
@@ -111,14 +116,14 @@ export class GeminiController {
       if (action === 'streamGenerateContent') {
         const stream = await this.proxyService.handleGeminiStreamGenerateContent(model, body);
         if (stream instanceof Observable) {
-          this.sendObservableSse(res, stream);
+          this.writeObservableSseResponse(res, stream);
           return;
         }
       }
 
       if (action === 'generateContent') {
         const result = await this.proxyService.handleGeminiGenerateContent(model, body);
-        res.status(HttpStatus.OK).send(this.normalizeGeminiGenerateResponse(result));
+        res.status(HttpStatus.OK).send(this.buildNormalizedGeminiGenerateResponse(result));
         return;
       }
 
@@ -141,7 +146,7 @@ export class GeminiController {
     }
   }
 
-  private parseCombinedModelAction(modelAction: string): {
+  private parseModelActionToken(modelAction: string): {
     model: string;
     action: string;
   } | null {
@@ -165,7 +170,10 @@ export class GeminiController {
 
   private buildGeminiModelList(): GeminiModelMetadata[] {
     const config = getServerConfig();
-    const dynamicModelIds = getAllDynamicModels(config?.custom_mapping ?? {});
+    const dynamicModelIds = getAllDynamicModels(
+      config?.custom_mapping ?? {},
+      this.tokenManager?.getAllCollectedModels(),
+    );
 
     return dynamicModelIds.map((id) => this.toGeminiModelMetadata(`models/${id}`));
   }
@@ -186,7 +194,7 @@ export class GeminiController {
     };
   }
 
-  private normalizeGeminiGenerateResponse(response: GeminiResponse): GeminiResponse {
+  private buildNormalizedGeminiGenerateResponse(response: GeminiResponse): GeminiResponse {
     const candidates = (response.candidates ?? []).map((candidate, index) => ({
       content: candidate.content,
       finishReason: candidate.finishReason,
@@ -235,7 +243,7 @@ export class GeminiController {
     return normalized;
   }
 
-  private sendObservableSse(res: FastifyReply, stream: Observable<unknown>): void {
+  private writeObservableSseResponse(res: FastifyReply, stream: Observable<unknown>): void {
     if (
       !res.raw ||
       typeof res.raw.writeHead !== 'function' ||
@@ -248,7 +256,7 @@ export class GeminiController {
       return;
     }
 
-    if (this.hasHijack(res)) {
+    if (this.supportsReplyHijack(res)) {
       res.hijack();
     }
 
@@ -293,7 +301,7 @@ export class GeminiController {
     });
   }
 
-  private hasHijack(reply: FastifyReply): reply is FastifyReply & { hijack: () => void } {
+  private supportsReplyHijack(reply: FastifyReply): reply is FastifyReply & { hijack: () => void } {
     return typeof (reply as { hijack?: unknown }).hijack === 'function';
   }
 }

src/server/modules/proxy/interfaces/request-interfaces.ts

@@ -30,12 +30,13 @@ export interface OpenAIContentPart {
 }
 
 export interface OpenAITool {
-  type: 'function';
+  type: string;
-  function: {
+  function?: {
     name: string;
     description?: string;
     parameters?: Record<string, unknown>;
   };
+  [key: string]: unknown;
 }
 
 export interface OpenAIToolCall {

src/server/modules/proxy/proxy.controller.ts

@@ -9,9 +9,10 @@ import {
   Inject,
   Req,
   Logger,
+  Optional,
 } from '@nestjs/common';
 import { FastifyReply, FastifyRequest } from 'fastify';
-import { isNil } from 'lodash-es';
+import { isNil, isPlainObject, isString } from 'lodash-es';
 import { ProxyService } from './proxy.service';
 import { Observable } from 'rxjs';
 import {
@@ -29,20 +30,27 @@ import {
   MODEL_LIST_OWNER,
 } from '../../../lib/antigravity/ModelMapping';
 import { getServerConfig } from '../../server-config';
+import { TokenManagerService } from './token-manager.service';
 
 @Controller('v1')
 @UseGuards(ProxyGuard)
 export class ProxyController {
   private readonly logger = new Logger(ProxyController.name);
 
-  constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {}
+  constructor(
+    @Inject(ProxyService) private readonly proxyService: ProxyService,
+    @Optional() @Inject(TokenManagerService) private readonly tokenManager?: TokenManagerService,
+  ) {}
 
   @Get('models')
   listModels(@Res() res: FastifyReply) {
     try {
       const config = getServerConfig();
       const customMapping = config?.custom_mapping ?? {};
-      const modelIds = getAllDynamicModels(customMapping);
+      const modelIds = getAllDynamicModels(
+        customMapping,
+        this.tokenManager?.getAllCollectedModels(),
+      );
 
       const data = modelIds.map((id) => ({
         id,
@@ -69,7 +77,7 @@ export class ProxyController {
 
   @Post('chat/completions')
   async chatCompletions(@Body() body: OpenAIChatRequest, @Res() res: FastifyReply) {
-    await this.sendOpenAIChatCompletionResponse(body, res);
+    await this.respondOpenAIChatCompletions(body, res);
   }
 
   @Post('completions')
@@ -86,11 +94,11 @@ export class ProxyController {
     @Res() res: FastifyReply,
   ) {
     const request: OpenAIChatRequest = {
-      model: body.model ?? 'gemini-2.5-flash',
+      model: body.model ?? 'gemini-3-flash',
       messages: [
         {
           role: 'user',
-          content: this.normalizePrompt(body.prompt),
+          content: this.normalizeCompletionPrompt(body.prompt),
         },
       ],
       max_tokens: body.max_tokens,
@@ -106,17 +114,9 @@ export class ProxyController {
       }
 
       const response = result as OpenAIChatResponse;
-      res.status(HttpStatus.OK).send(this.toLegacyCompletionsResponse(response));
+      res.status(HttpStatus.OK).send(this.toLegacyTextCompletionsResponse(response));
     } catch (error) {
-      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      this.sendOpenAIErrorResponse(res, '/v1/completions', error);
-      const status = this.resolveUpstreamStatus(message);
-      this.logEndpointError('/v1/completions', status, message, error);
-      res.status(status).send({
-        error: {
-          message,
-          type: 'server_error',
-        },
-      });
     }
   }
 
@@ -127,6 +127,7 @@ export class ProxyController {
       model?: string;
       instructions?: string;
       input?: unknown;
+      tools?: OpenAIChatRequest['tools'];
       max_output_tokens?: number;
       temperature?: number;
       top_p?: number;
@@ -144,17 +145,9 @@ export class ProxyController {
       }
 
       const response = result as OpenAIChatResponse;
-      res.status(HttpStatus.OK).send(this.toLegacyCompletionsResponse(response));
+      res.status(HttpStatus.OK).send(this.toLegacyTextCompletionsResponse(response));
     } catch (error) {
-      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      this.sendOpenAIErrorResponse(res, '/v1/responses', error);
-      const status = this.resolveUpstreamStatus(message);
-      this.logEndpointError('/v1/responses', status, message, error);
-      res.status(status).send({
-        error: {
-          message,
-          type: 'server_error',
-        },
-      });
     }
   }
 
@@ -182,7 +175,7 @@ export class ProxyController {
       quality: body.quality,
     };
 
-    await this.sendOpenAIImageResponse(request, body.prompt ?? '', res);
+    await this.sendOpenAIImageGenerationResponse(request, body.prompt ?? '', res);
   }
 
   @Post('images/edits')
@@ -200,7 +193,7 @@ export class ProxyController {
     @Req() req: FastifyRequest,
     @Res() res: FastifyReply,
   ) {
-    if (!this.hasValidMultipartBoundary(req)) {
+    if (!this.hasMultipartBoundary(req)) {
       res
         .status(HttpStatus.BAD_REQUEST)
         .send('Invalid `boundary` for `multipart/form-data` request');
@@ -236,7 +229,7 @@ export class ProxyController {
       quality: body.quality,
     };
 
-    await this.sendOpenAIImageResponse(request, body.prompt ?? '', res);
+    await this.sendOpenAIImageGenerationResponse(request, body.prompt ?? '', res);
   }
 
   @Post('audio/transcriptions')
@@ -251,7 +244,7 @@ export class ProxyController {
     @Req() req: FastifyRequest,
     @Res() res: FastifyReply,
   ) {
-    if (!this.hasValidMultipartBoundary(req)) {
+    if (!this.hasMultipartBoundary(req)) {
       res
         .status(HttpStatus.BAD_REQUEST)
         .send('Invalid `boundary` for `multipart/form-data` request');
@@ -271,7 +264,7 @@ export class ProxyController {
 
     try {
       const result = await this.proxyService.handleGeminiGenerateContent(
-        body.model ?? 'gemini-2.5-flash',
+        body.model ?? 'gemini-3-flash',
         {
           contents: [
             {
@@ -298,19 +291,11 @@ export class ProxyController {
         text: text ?? '',
       });
     } catch (error) {
-      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      this.sendOpenAIErrorResponse(res, '/v1/audio/transcriptions', error);
-      const status = this.resolveUpstreamStatus(message);
-      this.logEndpointError('/v1/audio/transcriptions', status, message, error);
-      res.status(status).send({
-        error: {
-          message,
-          type: 'server_error',
-        },
-      });
     }
   }
 
-  private async sendOpenAIChatCompletionResponse(body: OpenAIChatRequest, res: FastifyReply) {
+  private async respondOpenAIChatCompletions(body: OpenAIChatRequest, res: FastifyReply) {
     try {
       const result = await this.proxyService.handleChatCompletions(body);
 
@@ -321,15 +306,7 @@ export class ProxyController {
         res.status(HttpStatus.OK).send(result);
       }
     } catch (error) {
-      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      this.sendOpenAIErrorResponse(res, '/v1/chat/completions', error);
-      const status = this.resolveUpstreamStatus(message);
-      this.logEndpointError('/v1/chat/completions', status, message, error);
-      res.status(status).send({
-        error: {
-          message: message,
-          type: 'server_error',
-        },
-      });
     }
   }
 
@@ -345,20 +322,11 @@ export class ProxyController {
         res.status(HttpStatus.OK).send(result);
       }
     } catch (error) {
-      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      this.sendAnthropicErrorResponse(res, '/v1/messages', error);
-      const status = this.resolveUpstreamStatus(message);
-      this.logEndpointError('/v1/messages', status, message, error);
-      res.status(status).send({
-        type: 'error',
-        error: {
-          type: 'api_error',
-          message: message,
-        },
-      });
     }
   }
 
-  private normalizePrompt(prompt: string | string[] | undefined): string {
+  private normalizeCompletionPrompt(prompt: string | string[] | undefined): string {
     if (!prompt) {
       return '';
     }
@@ -368,7 +336,7 @@ export class ProxyController {
     return prompt;
   }
 
-  private toLegacyCompletionsResponse(response: OpenAIChatResponse): Record<string, unknown> {
+  private toLegacyTextCompletionsResponse(response: OpenAIChatResponse): Record<string, unknown> {
     const choice = response.choices?.[0];
     const content = choice?.message?.content;
     const text = typeof content === 'string' ? content : '';
@@ -391,21 +359,20 @@ export class ProxyController {
   }
 
   private normalizeResponsesInput(input: unknown): string {
-    if (typeof input === 'string') {
+    if (isString(input)) {
       return input;
     }
 
     if (Array.isArray(input)) {
       return input
         .map((item) => {
-          if (typeof item === 'string') {
+          if (isString(item)) {
             return item;
           }
-          if (item && typeof item === 'object' && 'content' in item) {
+          const itemRecord = this.toRecord(item);
-            const content = (item as { content?: unknown }).content;
+          const content = this.asString(itemRecord?.content);
-            if (typeof content === 'string') {
+          if (content) {
-              return content;
+            return content;
-            }
           }
           return JSON.stringify(item);
         })
@@ -423,6 +390,7 @@ export class ProxyController {
     model?: string;
     instructions?: string;
     input?: unknown;
+    tools?: OpenAIChatRequest['tools'];
     max_output_tokens?: number;
     temperature?: number;
     top_p?: number;
@@ -458,7 +426,7 @@ export class ProxyController {
             type === 'local_shell_call'
               ? 'shell'
               : type === 'web_search_call'
-                ? 'google_search'
+                ? 'builtin_web_search'
                 : (this.asString(itemObj.name) ?? 'unknown');
           callIdToToolName.set(callId, toolName);
         }
@@ -536,8 +504,9 @@ export class ProxyController {
     }
 
     return {
-      model: body.model ?? 'gemini-2.5-flash',
+      model: body.model ?? 'gemini-3-flash',
       messages,
+      tools: body.tools,
       max_tokens: body.max_output_tokens,
       temperature: body.temperature,
       top_p: body.top_p,
@@ -546,7 +515,7 @@ export class ProxyController {
   }
 
   private normalizeResponsesMessageContent(content: unknown): string | OpenAIContentPart[] {
-    if (typeof content === 'string') {
+    if (isString(content)) {
       return content;
     }
 
@@ -621,11 +590,12 @@ export class ProxyController {
     }
 
     const raw = item.arguments;
-    if (typeof raw === 'string') {
+    if (isString(raw)) {
       try {
         const parsed = JSON.parse(raw);
-        if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
+        const parsedRecord = this.toRecord(parsed);
-          return parsed as Record<string, unknown>;
+        if (parsedRecord) {
+          return parsedRecord;
         }
         return {
           value: parsed,
@@ -637,23 +607,22 @@ export class ProxyController {
       }
     }
 
-    if (raw && typeof raw === 'object' && !Array.isArray(raw)) {
+    const rawRecord = this.toRecord(raw);
-      return raw as Record<string, unknown>;
+    if (rawRecord) {
+      return rawRecord;
     }
 
     return {};
   }
 
   private normalizeResponsesOutput(output: unknown): string {
-    if (typeof output === 'string') {
+    if (isString(output)) {
       return output;
     }
-    if (output && typeof output === 'object' && !Array.isArray(output)) {
+    const outputRecord = this.toRecord(output);
-      const obj = output as Record<string, unknown>;
+    const content = this.asString(outputRecord?.content);
-      const content = this.asString(obj.content);
+    if (content) {
-      if (content) {
+      return content;
-        return content;
-      }
     }
     if (isNil(output)) {
       return '';
@@ -663,15 +632,13 @@ export class ProxyController {
 
   private resolveImageUrl(block: Record<string, unknown>): string | null {
     const raw = block.image_url;
-    if (typeof raw === 'string') {
+    if (isString(raw)) {
       return raw;
     }
-    if (raw && typeof raw === 'object' && !Array.isArray(raw)) {
+    const rawRecord = this.toRecord(raw);
-      const record = raw as Record<string, unknown>;
+    const url = this.asString(rawRecord?.url);
-      const url = this.asString(record.url);
+    if (url) {
-      if (url) {
+      return url;
-        return url;
-      }
     }
     return null;
   }
@@ -703,7 +670,7 @@ export class ProxyController {
       return null;
     }
 
-    if (typeof input === 'string') {
+    if (isString(input)) {
       const dataUri = input.match(/^data:(?<mime>[^;]+);base64,(?<data>[A-Za-z0-9+/=]+)$/);
       if (dataUri?.groups?.mime && dataUri.groups.data) {
         return {
@@ -722,14 +689,14 @@ export class ProxyController {
       return null;
     }
 
-    if (typeof input === 'object' && !Array.isArray(input)) {
+    const inputRecord = this.toRecord(input);
-      const obj = input as Record<string, unknown>;
+    if (inputRecord) {
-      const data = this.asString(obj.data);
+      const data = this.asString(inputRecord.data);
       if (!data) {
         return null;
       }
       return {
-        mimeType: this.asString(obj.mimeType) ?? 'audio/mpeg',
+        mimeType: this.asString(inputRecord.mimeType) ?? 'audio/mpeg',
         data,
       };
     }
@@ -738,14 +705,14 @@ export class ProxyController {
   }
 
   private toRecord(value: unknown): Record<string, unknown> | null {
-    if (!value || typeof value !== 'object' || Array.isArray(value)) {
+    if (!isPlainObject(value)) {
       return null;
     }
     return value as Record<string, unknown>;
   }
 
   private asString(value: unknown): string | null {
-    return typeof value === 'string' ? value : null;
+    return isString(value) ? value : null;
   }
 
   private isObservableLike(value: unknown): value is Observable<unknown> {
@@ -814,7 +781,7 @@ export class ProxyController {
     });
   }
 
-  private async sendOpenAIImageResponse(
+  private async sendOpenAIImageGenerationResponse(
     request: OpenAIChatRequest,
     prompt: string,
     res: FastifyReply,
@@ -822,7 +789,7 @@ export class ProxyController {
     try {
       const result = await this.proxyService.handleChatCompletions(request);
       if (result instanceof Observable) {
-        this.logEndpointError(
+        this.logProxyEndpointError(
           '/v1/images/generations',
           HttpStatus.INTERNAL_SERVER_ERROR,
           'Streaming image generation is not supported by this endpoint',
@@ -839,7 +806,7 @@ export class ProxyController {
       const content = result.choices?.[0]?.message?.content;
       const image = this.extractInlineBase64Image(typeof content === 'string' ? content : '');
       if (!image) {
-        this.logEndpointError(
+        this.logProxyEndpointError(
           '/v1/images/generations',
           HttpStatus.BAD_GATEWAY,
           'Upstream did not return inline image data',
@@ -889,14 +856,7 @@ export class ProxyController {
         }
       }
 
-      const status = this.resolveUpstreamStatus(message);
+      this.sendOpenAIErrorResponse(res, '/v1/images/generations', error, message);
-      this.logEndpointError('/v1/images/generations', status, message, error);
-      res.status(status).send({
-        error: {
-          message,
-          type: 'server_error',
-        },
-      });
     }
   }
 
@@ -1002,7 +962,7 @@ export class ProxyController {
     );
   }
 
-  private hasValidMultipartBoundary(req: FastifyRequest): boolean {
+  private hasMultipartBoundary(req: FastifyRequest): boolean {
     const contentType = req.headers['content-type'];
     if (typeof contentType !== 'string') {
       return false;
@@ -1012,7 +972,46 @@ export class ProxyController {
     return lowered.includes('multipart/form-data') && lowered.includes('boundary=');
   }
 
-  private resolveUpstreamStatus(message: string): HttpStatus {
+  private resolveErrorMessageText(error: unknown): string {
+    return error instanceof Error ? error.message : 'Internal Server Error';
+  }
+
+  private sendOpenAIErrorResponse(
+    res: FastifyReply,
+    endpoint: string,
+    error: unknown,
+    overrideMessage?: string,
+  ): void {
+    const message = overrideMessage ?? this.resolveErrorMessageText(error);
+    const status = this.resolveErrorHttpStatus(message);
+    this.logProxyEndpointError(endpoint, status, message, error);
+    res.status(status).send({
+      error: {
+        message,
+        type: 'server_error',
+      },
+    });
+  }
+
+  private sendAnthropicErrorResponse(
+    res: FastifyReply,
+    endpoint: string,
+    error: unknown,
+    overrideMessage?: string,
+  ): void {
+    const message = overrideMessage ?? this.resolveErrorMessageText(error);
+    const status = this.resolveErrorHttpStatus(message);
+    this.logProxyEndpointError(endpoint, status, message, error);
+    res.status(status).send({
+      type: 'error',
+      error: {
+        type: 'api_error',
+        message,
+      },
+    });
+  }
+
+  private resolveErrorHttpStatus(message: string): HttpStatus {
     const lowered = message.toLowerCase();
     if (lowered.includes('all accounts failed or unhealthy')) {
       return HttpStatus.SERVICE_UNAVAILABLE;
@@ -1050,7 +1049,7 @@ export class ProxyController {
     return HttpStatus.INTERNAL_SERVER_ERROR;
   }
 
-  private logEndpointError(
+  private logProxyEndpointError(
     endpoint: string,
     status: HttpStatus,
     message: string,

src/server/modules/proxy/proxy.service.ts

@@ -1,5 +1,5 @@
 import { Injectable, Logger, Inject } from '@nestjs/common';
-import { isEmpty, isNil } from 'lodash-es';
+import { isEmpty, isNil, isPlainObject, isString } from 'lodash-es';
 import { TokenManagerService } from './token-manager.service';
 import { GeminiClient } from './clients/gemini.client';
 import { v4 as uuidv4 } from 'uuid';
@@ -14,6 +14,7 @@ import {
   GeminiPart as InternalGeminiPart,
 } from '../../../lib/antigravity/types';
 import { calculateRetryDelay, sleep } from '../../../lib/antigravity/retry-utils';
+import { normalizeObjectJsonSchema } from '../../../lib/antigravity/JsonSchemaUtils';
 import {
   classifyStreamError,
   formatErrorForSSE,
@@ -28,12 +29,14 @@ import {
   AnthropicContent,
 } from './interfaces/request-interfaces';
 import { getServerConfig } from '../../server-config';
-import { resolveModelRoute } from '../../../lib/antigravity/ModelMapping';
+import { normalizeGeminiModelAlias, resolveModelRoute } from '../../../lib/antigravity/ModelMapping';
+import { getMaxOutputTokens, getThinkingBudget } from '../../../lib/antigravity/ModelSpecs';
+import { resolveRequestUserAgent } from './request-user-agent';
+import { UpstreamRequestError } from './clients/upstream-error';
 
 @Injectable()
 export class ProxyService {
   private readonly logger = new Logger(ProxyService.name);
-  private readonly defaultRequestUserAgent = 'antigravity/1.11.9 windows/amd64';
 
   constructor(
     @Inject(TokenManagerService) private readonly tokenManager: TokenManagerService,
@@ -45,12 +48,12 @@ export class ProxyService {
   async handleAnthropicMessages(
     request: AnthropicChatRequest,
   ): Promise<AnthropicChatResponse | Observable<string>> {
-    const sessionKey = this.getSessionKeyFromAnthropic(request);
+    const sessionKey = this.extractAnthropicSessionKey(request);
 
-    const targetModel = this.mapModel(request.model);
+    const targetModel = this.resolveTargetModel(request.model);
-    const extraHeaders = this.buildModelSpecificHeaders(request.model);
+    const extraHeaders = this.createModelSpecificHeaders(request.model);
     this.logger.log(
-      `Received Anthropic request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`,
+      `Anthropic request received: model=${request.model}, mappedModel=${targetModel}, stream=${request.stream}`,
     );
 
     // Retry loop
@@ -61,13 +64,14 @@ export class ProxyService {
     for (let i = 0; i < maxRetries; i++) {
       if (i > 0) {
         const delay = calculateRetryDelay(i - 1);
-        this.logger.log(`Retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms (jittered)`);
+        this.logger.log(`Anthropic retry ${i + 1}/${maxRetries}, backoff=${delay}ms (jittered)`);
         await sleep(delay);
       }
 
       const token = await this.tokenManager.getNextToken({
         sessionKey,
         excludeAccountIds: Array.from(attemptedAccountIds),
+        model: targetModel,
       });
       if (!token) {
         throw new Error('No available accounts');
@@ -76,7 +80,13 @@ export class ProxyService {
 
       try {
         const projectId = token.token.project_id ?? '';
-        const geminiBody = transformClaudeRequestIn(this.toClaudeRequest(request), projectId);
+        const requestUserAgent = await resolveRequestUserAgent();
+        const geminiBody = transformClaudeRequestIn(
+          this.toClaudeRequest(request),
+          projectId,
+          requestUserAgent,
+        );
+        this.applyInternalGenerationConstraints(geminiBody, targetModel, token.id);
 
         if (request.stream) {
           const stream = await this.geminiClient.streamGenerateInternal(
@@ -101,7 +111,13 @@ export class ProxyService {
             `Anthropic request hit project context issue, retrying without project: ${error.message}`,
           );
           try {
-            const fallbackBody = transformClaudeRequestIn(this.toClaudeRequest(request), '');
+            const requestUserAgent = await resolveRequestUserAgent();
+            const fallbackBody = transformClaudeRequestIn(
+              this.toClaudeRequest(request),
+              '',
+              requestUserAgent,
+            );
+            this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
             if (request.stream) {
               const stream = await this.geminiClient.streamGenerateInternal(
                 fallbackBody,
@@ -126,17 +142,20 @@ export class ProxyService {
 
         if (error instanceof Error && this.isQuotaExhaustedError(error.message)) {
           this.logger.warn(
-            `Anthropic request hit quota exhaustion on mapped model, retrying with fallback model gemini-2.5-flash: ${error.message}`,
+            `Anthropic request hit quota exhaustion on mapped model, retrying with fallback model gemini-3-flash: ${error.message}`,
           );
           try {
             const downgradedRequest: ClaudeRequest = {
               ...this.toClaudeRequest(request),
-              model: 'gemini-2.5-flash',
+              model: 'gemini-3-flash',
             };
+            const requestUserAgent = await resolveRequestUserAgent();
             const downgradedBody = transformClaudeRequestIn(
               downgradedRequest,
               token.token.project_id ?? '',
+              requestUserAgent,
             );
+            this.applyInternalGenerationConstraints(downgradedBody, 'gemini-3-flash', token.id);
             if (request.stream) {
               const stream = await this.geminiClient.streamGenerateInternal(
                 downgradedBody,
@@ -164,17 +183,7 @@ export class ProxyService {
         }
 
         lastError = error;
-        if (error instanceof Error) {
+        await this.applyUpstreamPenalty(token.id, targetModel, error);
-          this.logger.warn(`Anthropic Request failed: ${error.message}`);
-          const decision = this.classifyUpstreamError(error.message);
-          if (decision.retry) {
-            if (decision.markAsForbidden) {
-              this.tokenManager.markAsForbidden(token.id);
-            } else if (decision.markAsRateLimited) {
-              this.tokenManager.markAsRateLimited(token.id);
-            }
-          }
-        }
       }
     }
     throw lastError || new Error('Request failed after retries');
@@ -182,7 +191,7 @@ export class ProxyService {
 
   private processAnthropicInternalStream(
     upstreamStream: NodeJS.ReadableStream,
-    model: string,
+    _model: string,
   ): Observable<string> {
     return new Observable<string>((subscriber) => {
       const decoder = new TextDecoder();
@@ -272,10 +281,10 @@ export class ProxyService {
     request: GeminiRequest,
   ): Promise<GeminiResponse> {
     const normalizedModel = this.normalizeGeminiModel(model);
-    const targetModel = this.mapModel(normalizedModel);
+    const targetModel = this.resolveTargetModel(normalizedModel);
-    const extraHeaders = this.buildModelSpecificHeaders(normalizedModel);
+    const extraHeaders = this.createModelSpecificHeaders(normalizedModel);
     this.logger.log(
-      `Received Gemini request for model: ${normalizedModel} (Mapped: ${targetModel})`,
+      `Gemini generate request received: model=${normalizedModel}, mappedModel=${targetModel}`,
     );
 
     let lastError: unknown = null;
@@ -291,6 +300,7 @@ export class ProxyService {
 
       const token = await this.tokenManager.getNextToken({
         excludeAccountIds: Array.from(attemptedAccountIds),
+        model: targetModel,
       });
       if (!token) {
         throw new Error('No available accounts (all exhausted or rate limited)');
@@ -298,12 +308,15 @@ export class ProxyService {
       attemptedAccountIds.add(token.id);
 
       try {
+        const requestUserAgent = await resolveRequestUserAgent();
         const internalBody = this.createGeminiInternalRequest(
           targetModel,
           request,
           token.token.project_id ?? '',
           'generate-content',
+          requestUserAgent,
         );
+        this.applyInternalGenerationConstraints(internalBody, targetModel, token.id);
 
         const response = await this.generateInternalWithStreamFallback(
           internalBody,
@@ -319,12 +332,15 @@ export class ProxyService {
             `Gemini request hit project context issue, retrying without project: ${err.message}`,
           );
           try {
+            const requestUserAgent = await resolveRequestUserAgent();
             const fallbackBody = this.createGeminiInternalRequest(
               targetModel,
               request,
               '',
               'generate-content',
+              requestUserAgent,
             );
+            this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
             const response = await this.generateInternalWithStreamFallback(
               fallbackBody,
               token.token.access_token,
@@ -339,17 +355,7 @@ export class ProxyService {
           lastError = err;
         }
 
-        if (lastError instanceof Error) {
+        await this.applyUpstreamPenalty(token.id, targetModel, lastError);
-          this.logger.warn(`Gemini request failed with account ${token.email}: ${lastError.message}`);
-          const decision = this.classifyUpstreamError(lastError.message);
-          if (decision.retry) {
-            if (decision.markAsForbidden) {
-              this.tokenManager.markAsForbidden(token.id);
-            } else if (decision.markAsRateLimited) {
-              this.tokenManager.markAsRateLimited(token.id);
-            }
-          }
-        }
       }
     }
 
@@ -361,10 +367,10 @@ export class ProxyService {
     request: GeminiRequest,
   ): Promise<Observable<string>> {
     const normalizedModel = this.normalizeGeminiModel(model);
-    const targetModel = this.mapModel(normalizedModel);
+    const targetModel = this.resolveTargetModel(normalizedModel);
-    const extraHeaders = this.buildModelSpecificHeaders(normalizedModel);
+    const extraHeaders = this.createModelSpecificHeaders(normalizedModel);
     this.logger.log(
-      `Received Gemini stream request for model: ${normalizedModel} (Mapped: ${targetModel})`,
+      `Gemini stream request received: model=${normalizedModel}, mappedModel=${targetModel}`,
     );
 
     let lastError: unknown = null;
@@ -380,6 +386,7 @@ export class ProxyService {
 
       const token = await this.tokenManager.getNextToken({
         excludeAccountIds: Array.from(attemptedAccountIds),
+        model: targetModel,
       });
       if (!token) {
         throw new Error('No available accounts (all exhausted or rate limited)');
@@ -387,12 +394,15 @@ export class ProxyService {
       attemptedAccountIds.add(token.id);
 
       try {
+        const requestUserAgent = await resolveRequestUserAgent();
         const internalBody = this.createGeminiInternalRequest(
           targetModel,
           request,
           token.token.project_id ?? '',
           'generate-content',
+          requestUserAgent,
         );
+        this.applyInternalGenerationConstraints(internalBody, targetModel, token.id);
 
         const stream = await this.geminiClient.streamGenerateInternal(
           internalBody,
@@ -407,12 +417,15 @@ export class ProxyService {
             `Gemini stream request hit project context issue, retrying without project: ${err.message}`,
           );
           try {
+            const requestUserAgent = await resolveRequestUserAgent();
             const fallbackBody = this.createGeminiInternalRequest(
               targetModel,
               request,
               '',
               'generate-content',
+              requestUserAgent,
             );
+            this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
             const stream = await this.geminiClient.streamGenerateInternal(
               fallbackBody,
               token.token.access_token,
@@ -427,19 +440,7 @@ export class ProxyService {
           lastError = err;
         }
 
-        if (lastError instanceof Error) {
+        await this.applyUpstreamPenalty(token.id, targetModel, lastError);
-          this.logger.warn(
-            `Gemini stream request failed with account ${token.email}: ${lastError.message}`,
-          );
-          const decision = this.classifyUpstreamError(lastError.message);
-          if (decision.retry) {
-            if (decision.markAsForbidden) {
-              this.tokenManager.markAsForbidden(token.id);
-            } else if (decision.markAsRateLimited) {
-              this.tokenManager.markAsRateLimited(token.id);
-            }
-          }
-        }
       }
     }
 
@@ -475,11 +476,130 @@ export class ProxyService {
     return model.replace(/^models\//i, '');
   }
 
+  private normalizeModelIdentifier(model: string): string {
+    return model.replace(/^models\//i, '').trim();
+  }
+
+  private resolveThinkingLevelBudget(level: string): number | undefined {
+    const normalized = level.trim().toUpperCase();
+    if (normalized === 'NONE') {
+      return 0;
+    }
+    if (normalized === 'LOW') {
+      return 4096;
+    }
+    if (normalized === 'MEDIUM') {
+      return 8192;
+    }
+    if (normalized === 'HIGH') {
+      return 24576;
+    }
+    return undefined;
+  }
+
+  private getModelOutputCap(accountId: string, model: string): number {
+    const normalizedModel = this.normalizeModelIdentifier(model);
+    const dynamicCap = this.tokenManager.getModelOutputLimitForAccount(accountId, normalizedModel);
+    if (typeof dynamicCap === 'number' && Number.isFinite(dynamicCap) && dynamicCap > 0) {
+      return Math.floor(dynamicCap);
+    }
+    return getMaxOutputTokens(normalizedModel);
+  }
+
+  private getModelThinkingBudget(accountId: string, model: string): number {
+    const normalizedModel = this.normalizeModelIdentifier(model);
+    const dynamicBudget = this.tokenManager.getModelThinkingBudgetForAccount(
+      accountId,
+      normalizedModel,
+    );
+    if (typeof dynamicBudget === 'number' && Number.isFinite(dynamicBudget) && dynamicBudget >= 0) {
+      return Math.floor(dynamicBudget);
+    }
+    return getThinkingBudget(normalizedModel);
+  }
+
+  private applyInternalGenerationConstraints(
+    body: GeminiInternalRequest,
+    model: string,
+    accountId: string,
+  ): void {
+    const generationConfig = body.request.generationConfig;
+    if (!generationConfig) {
+      return;
+    }
+
+    const outputCap = this.getModelOutputCap(accountId, model);
+    const thinkingBudgetCap = this.getModelThinkingBudget(accountId, model);
+    const normalizedModel = this.normalizeModelIdentifier(model).toLowerCase();
+    const isClaudeModel = normalizedModel.includes('claude');
+    const thinkingConfig = generationConfig.thinkingConfig as
+      | ({ thinkingLevel?: string; thinkingBudget?: number } & Record<string, unknown>)
+      | undefined;
+    const adaptiveSentinel =
+      thinkingConfig &&
+      (typeof thinkingConfig.thinkingLevel === 'string' ||
+        thinkingConfig.thinkingBudget === -1 ||
+        thinkingConfig.thinkingBudget === 32768);
+
+    if (thinkingConfig) {
+      if (!isClaudeModel && typeof thinkingConfig.thinkingLevel === 'string') {
+        const converted = this.resolveThinkingLevelBudget(thinkingConfig.thinkingLevel);
+        if (converted !== undefined) {
+          thinkingConfig.thinkingBudget = converted;
+        }
+        delete thinkingConfig.thinkingLevel;
+      }
+
+      if (typeof thinkingConfig.thinkingBudget === 'number' && thinkingConfig.thinkingBudget < 0) {
+        thinkingConfig.thinkingBudget = Math.min(thinkingBudgetCap, 24576);
+      }
+
+      if (
+        typeof thinkingConfig.thinkingBudget === 'number' &&
+        Number.isFinite(thinkingConfig.thinkingBudget)
+      ) {
+        thinkingConfig.thinkingBudget = Math.min(
+          Math.floor(thinkingConfig.thinkingBudget),
+          Math.max(0, outputCap - 1),
+          thinkingBudgetCap,
+        );
+
+        if (adaptiveSentinel) {
+          if (
+            generationConfig.maxOutputTokens === undefined ||
+            generationConfig.maxOutputTokens < 131072
+          ) {
+            generationConfig.maxOutputTokens = 131072;
+          }
+        } else if (
+          generationConfig.maxOutputTokens === undefined ||
+          generationConfig.maxOutputTokens <= thinkingConfig.thinkingBudget
+        ) {
+          const hasExplicitMax = generationConfig.maxOutputTokens !== undefined;
+          const overhead = hasExplicitMax ? 8192 : 32768;
+          const minRequired = Math.min(outputCap, thinkingConfig.thinkingBudget + overhead);
+          generationConfig.maxOutputTokens = minRequired;
+        }
+      }
+    }
+
+    if (
+      typeof generationConfig.maxOutputTokens === 'number' &&
+      Number.isFinite(generationConfig.maxOutputTokens)
+    ) {
+      generationConfig.maxOutputTokens = Math.min(
+        Math.floor(generationConfig.maxOutputTokens),
+        outputCap,
+      );
+    }
+  }
+
   private createGeminiInternalRequest(
     model: string,
     request: GeminiRequest,
     projectId: string | undefined,
     requestType: string,
+    requestUserAgent: string,
   ): GeminiInternalRequest {
     const normalizedProjectId = projectId?.trim();
 
@@ -487,7 +607,7 @@ export class ProxyService {
       requestId: uuidv4(),
       request: this.toInternalGeminiRequest(request),
       model,
-      userAgent: process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent,
+      userAgent: requestUserAgent,
       requestType,
     };
 
@@ -523,9 +643,6 @@ export class ProxyService {
       if (usage.totalTokenCount !== undefined) {
         usageMetadata.totalTokenCount = usage.totalTokenCount;
       }
-      if (usage.thoughtsTokenCount !== undefined) {
-        usageMetadata.thoughtsTokenCount = usage.thoughtsTokenCount;
-      }
       if (usage.promptTokensDetails !== undefined) {
         usageMetadata.promptTokensDetails = usage.promptTokensDetails;
       }
@@ -546,12 +663,12 @@ export class ProxyService {
   async handleChatCompletions(
     request: OpenAIChatRequest,
   ): Promise<OpenAIChatResponse | Observable<string>> {
-    const sessionKey = this.getSessionKeyFromOpenAI(request);
+    const sessionKey = this.extractOpenAISessionKey(request);
 
-    const targetModel = this.mapModel(request.model);
+    const targetModel = this.resolveTargetModel(request.model);
-    const extraHeaders = this.buildModelSpecificHeaders(request.model);
+    const extraHeaders = this.createModelSpecificHeaders(request.model);
     this.logger.log(
-      `Received request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`,
+      `OpenAI-compatible request received: model=${request.model}, mappedModel=${targetModel}, stream=${request.stream}`,
     );
 
     // Retry loop for account selection
@@ -562,7 +679,7 @@ export class ProxyService {
     for (let i = 0; i < maxRetries; i++) {
       if (i > 0) {
         const delay = calculateRetryDelay(i - 1);
-        this.logger.log(`Retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms (jittered)`);
+        this.logger.log(`OpenAI-compatible retry ${i + 1}/${maxRetries}, backoff=${delay}ms (jittered)`);
         await sleep(delay);
       }
 
@@ -570,6 +687,7 @@ export class ProxyService {
       const token = await this.tokenManager.getNextToken({
         sessionKey,
         excludeAccountIds: Array.from(attemptedAccountIds),
+        model: targetModel,
       });
       if (!token) {
         throw new Error('No available accounts (all exhausted or rate limited)');
@@ -579,7 +697,9 @@ export class ProxyService {
       try {
         const claudeRequest = this.convertOpenAIToClaude(request);
         const projectId = token.token.project_id ?? '';
-        const geminiBody = transformClaudeRequestIn(claudeRequest, projectId);
+        const requestUserAgent = await resolveRequestUserAgent();
+        const geminiBody = transformClaudeRequestIn(claudeRequest, projectId, requestUserAgent);
+        this.applyInternalGenerationConstraints(geminiBody, targetModel, token.id);
 
         // Use v1internal API (same as Anthropic handler)
         if (request.stream) {
@@ -593,7 +713,7 @@ export class ProxyService {
             return this.processStreamResponse(stream, request.model);
           } catch (streamError) {
             this.logger.warn(
-              `Stream request failed for ${request.model}, falling back to non-stream generation: ${
+              `Stream path failed for model=${request.model}; falling back to non-stream generation: ${
                 streamError instanceof Error ? streamError.message : String(streamError)
               }`,
             );
@@ -605,7 +725,7 @@ export class ProxyService {
               extraHeaders,
             );
             this.logger.log(
-              `Upstream response snippet (stream fallback): ${JSON.stringify(response).substring(0, 500)}`,
+              `Upstream response snippet after stream fallback: ${JSON.stringify(response).substring(0, 500)}`,
             );
             const claudeResponse = transformResponse(response);
             const openaiResponse = this.convertClaudeToOpenAIResponse(
@@ -622,11 +742,13 @@ export class ProxyService {
             extraHeaders,
           );
           this.logger.log(
-            `Upstream response snippet: ${JSON.stringify(response).substring(0, 500)}`,
+            `Upstream response snippet (non-stream): ${JSON.stringify(response).substring(0, 500)}`,
           );
           // Transform Gemini response to OpenAI format
           const claudeResponse = transformResponse(response);
-          this.logger.log(`Claude Response: ${JSON.stringify(claudeResponse).substring(0, 500)}`);
+          this.logger.log(
+            `Transformed Claude response snippet: ${JSON.stringify(claudeResponse).substring(0, 500)}`,
+          );
           return this.convertClaudeToOpenAIResponse(claudeResponse, request.model);
         }
       } catch (err) {
@@ -636,7 +758,9 @@ export class ProxyService {
           );
           try {
             const claudeRequest = this.convertOpenAIToClaude(request);
-            const fallbackBody = transformClaudeRequestIn(claudeRequest, '');
+            const requestUserAgent = await resolveRequestUserAgent();
+            const fallbackBody = transformClaudeRequestIn(claudeRequest, '', requestUserAgent);
+            this.applyInternalGenerationConstraints(fallbackBody, targetModel, token.id);
             if (request.stream) {
               const stream = await this.geminiClient.streamGenerateInternal(
                 fallbackBody,
@@ -662,18 +786,7 @@ export class ProxyService {
           lastError = err;
         }
 
-        if (lastError instanceof Error) {
+        await this.applyUpstreamPenalty(token.id, targetModel, lastError);
-          this.logger.warn(`Request failed with account ${token.email}: ${lastError.message}`);
-
-          const decision = this.classifyUpstreamError(lastError.message);
-          if (decision.retry) {
-            if (decision.markAsForbidden) {
-              this.tokenManager.markAsForbidden(token.id);
-            } else if (decision.markAsRateLimited) {
-              this.tokenManager.markAsRateLimited(token.id);
-            }
-          }
-        }
       }
     }
     throw lastError || new Error('Request failed after retries');
@@ -937,7 +1050,7 @@ export class ProxyService {
               hasSentDone = true;
               subscriber.complete();
             }
-          } catch (e) {
+          } catch {
             // ignore parse errors
           }
         }
@@ -1228,11 +1341,51 @@ export class ProxyService {
       return undefined;
     }
 
-    return tools.map((tool) => ({
+    const result: NonNullable<AnthropicChatRequest['tools']> = [];
-      name: tool.function.name,
+    const searchToolTypes = new Set([
-      description: tool.function.description,
+      'web_search_20250305',
-      input_schema: tool.function.parameters || { type: 'object', properties: {} },
+      'google_search',
-    }));
+      'google_search_retrieval',
+      'builtin_web_search',
+    ]);
+
+    for (const tool of tools) {
+      if (!tool) {
+        continue;
+      }
+
+      const toolType = isString(tool.type) ? tool.type.toLowerCase() : '';
+      const functionName = isString(tool.function?.name) ? tool.function.name : '';
+      const normalizedFunctionName = functionName.toLowerCase();
+      const isSearchTool =
+        searchToolTypes.has(toolType) || searchToolTypes.has(normalizedFunctionName);
+
+      if (isSearchTool) {
+        result.push({
+          name: functionName || 'builtin_web_search',
+          type: 'web_search_20250305',
+          input_schema: {
+            type: 'object',
+            properties: {},
+          },
+        });
+        continue;
+      }
+
+      if (!tool.function || !functionName) {
+        continue;
+      }
+
+      const inputSchema = normalizeObjectJsonSchema(tool.function.parameters);
+
+      result.push({
+        name: functionName,
+        description: tool.function.description,
+        input_schema: inputSchema,
+      });
+    }
+
+    return result.length > 0 ? result : undefined;
   }
 
   private mapGeminiFinishReasonToOpenAIFinishReason(finishReason?: string): string | null {
@@ -1362,7 +1515,7 @@ export class ProxyService {
     };
   }
 
-  private mapModel(model: string): string {
+  private resolveTargetModel(model: string): string {
     const normalizedModel = model.replace(/^models\//i, '').trim();
     const config = getServerConfig();
     const configuredMapping = {
@@ -1400,23 +1553,66 @@ export class ProxyService {
     }
 
     const routedModel = resolveModelRoute(normalizedModel, customExactMapping, {}, {});
-
+    return normalizeGeminiModelAlias(routedModel);
-    return routedModel;
   }
 
-  private classifyUpstreamError(errorMessage: string): {
+  private async applyUpstreamPenalty(
+    accountId: string,
+    model: string,
+    error: unknown,
+  ): Promise<void> {
+    this.tokenManager.recordParityError();
+
+    if (error instanceof UpstreamRequestError) {
+      const status = error.status;
+      if (status === 401 || status === 403) {
+        this.tokenManager.markAsForbidden(accountId);
+        return;
+      }
+
+      await this.tokenManager.markFromUpstreamError({
+        accountIdOrEmail: accountId,
+        status,
+        retryAfter: error.headers?.retryAfter,
+        body: error.body,
+        model,
+      });
+      return;
+    }
+
+    if (!(error instanceof Error)) {
+      return;
+    }
+
+    this.logger.warn(`Upstream request failed for account ${accountId}: ${error.message}`);
+    const penaltyDecision = this.classifyUpstreamFailure(error.message);
+    if (!penaltyDecision.retry) {
+      return;
+    }
+
+    if (penaltyDecision.markAsForbidden) {
+      this.tokenManager.markAsForbidden(accountId);
+      return;
+    }
+
+    if (penaltyDecision.markAsRateLimited) {
+      this.tokenManager.markAsRateLimited(accountId);
+    }
+  }
+
+  private classifyUpstreamFailure(errorMessage: string): {
     retry: boolean;
     markAsForbidden: boolean;
     markAsRateLimited: boolean;
   } {
-    const msg = errorMessage.toLowerCase();
+    const normalizedErrorMessage = errorMessage.toLowerCase();
     const isForbidden =
-      msg.includes('401') ||
+      normalizedErrorMessage.includes('401') ||
-      msg.includes('unauthorized') ||
+      normalizedErrorMessage.includes('unauthorized') ||
-      msg.includes('invalid_grant') ||
+      normalizedErrorMessage.includes('invalid_grant') ||
-      msg.includes('403') ||
+      normalizedErrorMessage.includes('403') ||
-      msg.includes('permission_denied') ||
+      normalizedErrorMessage.includes('permission_denied') ||
-      msg.includes('forbidden');
+      normalizedErrorMessage.includes('forbidden');
 
     if (isForbidden) {
       return {
@@ -1427,28 +1623,28 @@ export class ProxyService {
     }
 
     const isRateLimitedSignal =
-      msg.includes('429') ||
+      normalizedErrorMessage.includes('429') ||
-      msg.includes('resource_exhausted') ||
+      normalizedErrorMessage.includes('resource_exhausted') ||
-      msg.includes('quota') ||
+      normalizedErrorMessage.includes('quota') ||
-      msg.includes('rate_limit') ||
+      normalizedErrorMessage.includes('rate_limit') ||
-      msg.includes('rate limit');
+      normalizedErrorMessage.includes('rate limit');
 
     const shouldRetryByStatus =
-      msg.includes('408') ||
+      normalizedErrorMessage.includes('408') ||
-      msg.includes('429') ||
+      normalizedErrorMessage.includes('429') ||
-      msg.includes('500') ||
+      normalizedErrorMessage.includes('500') ||
-      msg.includes('502') ||
+      normalizedErrorMessage.includes('502') ||
-      msg.includes('503') ||
+      normalizedErrorMessage.includes('503') ||
-      msg.includes('504');
+      normalizedErrorMessage.includes('504');
 
     const shouldRetryByKeyword =
-      msg.includes('resource_exhausted') ||
+      normalizedErrorMessage.includes('resource_exhausted') ||
-      msg.includes('quota') ||
+      normalizedErrorMessage.includes('quota') ||
-      msg.includes('rate_limit') ||
+      normalizedErrorMessage.includes('rate_limit') ||
-      msg.includes('timeout') ||
+      normalizedErrorMessage.includes('timeout') ||
-      msg.includes('socket hang up') ||
+      normalizedErrorMessage.includes('socket hang up') ||
-      msg.includes('empty response stream') ||
+      normalizedErrorMessage.includes('empty response stream') ||
-      msg.includes('connection reset');
+      normalizedErrorMessage.includes('connection reset');
 
     if (shouldRetryByStatus || shouldRetryByKeyword) {
       return {
@@ -1465,7 +1661,7 @@ export class ProxyService {
     };
   }
 
-  private buildModelSpecificHeaders(model: string | undefined): Record<string, string> {
+  private createModelSpecificHeaders(model: string | undefined): Record<string, string> {
     if (!model) {
       return {};
     }
@@ -1510,28 +1706,28 @@ export class ProxyService {
     );
   }
 
-  private getSessionKeyFromAnthropic(request: AnthropicChatRequest): string | undefined {
+  private extractAnthropicSessionKey(request: AnthropicChatRequest): string | undefined {
     const metadata = request.metadata;
     const sessionCandidate =
       metadata?.session_id ?? metadata?.sessionId ?? metadata?.user_id ?? metadata?.userId;
-    if (typeof sessionCandidate !== 'string' || sessionCandidate.trim() === '') {
+    if (!isString(sessionCandidate) || sessionCandidate.trim() === '') {
       return undefined;
     }
     return `anthropic:${sessionCandidate.trim()}`;
   }
 
-  private getSessionKeyFromOpenAI(request: OpenAIChatRequest): string | undefined {
+  private extractOpenAISessionKey(request: OpenAIChatRequest): string | undefined {
     const extra = request.extra;
     const sessionCandidate =
       extra?.session_id ?? extra?.sessionId ?? extra?.user_id ?? extra?.userId;
-    if (typeof sessionCandidate !== 'string' || sessionCandidate.trim() === '') {
+    if (!isString(sessionCandidate) || sessionCandidate.trim() === '') {
       return undefined;
     }
     return `openai:${sessionCandidate.trim()}`;
   }
 
   private isRecord(value: unknown): value is Record<string, unknown> {
-    return value !== null && typeof value === 'object' && !Array.isArray(value);
+    return isPlainObject(value);
   }
 
   private isGeminiPart(value: unknown): value is InternalGeminiPart {

src/services/GoogleAPIService.ts

@@ -1,6 +1,11 @@
 import { ConfigManager } from '../ipc/config/manager';
 import { ProxyAgent } from 'undici';
 import { logger } from '../utils/logger';
+import {
+  buildUserAgent,
+  FALLBACK_VERSION,
+  resolveLocalInstalledVersion,
+} from '@/server/modules/proxy/request-user-agent';
 
 // --- Constants & Config ---
 
@@ -52,13 +57,23 @@ export interface UserInfo {
 }
 
 export interface QuotaData {
-  models: Record<
+  models: Record<string, ModelQuotaInfo>;
-    string,
+  model_forwarding_rules?: Record<string, string>;
-    {
+  subscription_tier?: string;
-      percentage: number;
+  is_forbidden?: boolean;
-      resetTime: string;
+}
-    }
+
-  >;
+export interface ModelQuotaInfo {
+  percentage: number;
+  resetTime: string;
+  display_name?: string;
+  supports_images?: boolean;
+  supports_thinking?: boolean;
+  thinking_budget?: number;
+  recommended?: boolean;
+  max_tokens?: number;
+  max_output_tokens?: number;
+  supported_mime_types?: Record<string, boolean>;
 }
 
 // Internal types for API parsing
@@ -67,10 +82,138 @@ interface ModelInfoRaw {
     remainingFraction?: number;
     resetTime?: string;
   };
+  displayName?: string;
+  supportsImages?: boolean;
+  supportsThinking?: boolean;
+  thinkingBudget?: number;
+  recommended?: boolean;
+  maxTokens?: number;
+  maxOutputTokens?: number;
+  supportedMimeTypes?: Record<string, boolean>;
+}
+
+interface DeprecatedModelInfoRaw {
+  newModelId?: string;
+}
+
+interface IneligibleTierRaw {
+  reasonCode?: string;
+}
+
+interface TierRaw {
+  is_default?: boolean;
+  id?: string;
+  quotaTier?: string;
+  name?: string;
+  slug?: string;
 }
 
 interface LoadProjectResponse {
   cloudaicompanionProject?: string;
+  currentTier?: TierRaw;
+  paidTier?: TierRaw;
+  allowedTiers?: TierRaw[];
+  ineligibleTiers?: IneligibleTierRaw[];
+}
+
+interface FetchModelsResponse {
+  models?: Record<string, ModelInfoRaw>;
+  deprecatedModelIds?: Record<string, DeprecatedModelInfoRaw>;
+}
+
+interface ProjectContext {
+  projectId?: string;
+  subscriptionTier?: string;
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}
+
+function buildInternalApiHeaders(accessToken: string): Record<string, string> {
+  const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION;
+  return {
+    Authorization: `Bearer ${accessToken}`,
+    'User-Agent': buildUserAgent(discoveryVersion),
+    'Content-Type': 'application/json',
+  };
+}
+
+function resolveSubscriptionTier(payload: LoadProjectResponse): string | undefined {
+  const paidTier = payload.paidTier;
+  if (paidTier?.name && paidTier.name.trim() !== '') {
+    return paidTier.name;
+  }
+  if (paidTier?.id && paidTier.id.trim() !== '') {
+    return paidTier.id;
+  }
+
+  const ineligible = Array.isArray(payload.ineligibleTiers) && payload.ineligibleTiers.length > 0;
+  if (!ineligible) {
+    const currentTier = payload.currentTier;
+    if (currentTier?.name && currentTier.name.trim() !== '') {
+      return currentTier.name;
+    }
+    if (currentTier?.id && currentTier.id.trim() !== '') {
+      return currentTier.id;
+    }
+  }
+
+  if (Array.isArray(payload.allowedTiers)) {
+    const preferredAllowedTier =
+      payload.allowedTiers.find((tier) => tier.is_default === true) ?? payload.allowedTiers[0];
+    if (preferredAllowedTier?.name && preferredAllowedTier.name.trim() !== '') {
+      return ineligible ? `${preferredAllowedTier.name} (Restricted)` : preferredAllowedTier.name;
+    }
+    if (preferredAllowedTier?.id && preferredAllowedTier.id.trim() !== '') {
+      return ineligible ? `${preferredAllowedTier.id} (Restricted)` : preferredAllowedTier.id;
+    }
+  }
+
+  return undefined;
+}
+
+function isTrackedModel(modelName: string): boolean {
+  return /^(gemini|claude|gpt|image|imagen)/i.test(modelName);
+}
+
+function toModelQuotaInfo(modelName: string, info: ModelInfoRaw): ModelQuotaInfo | null {
+  if (!isTrackedModel(modelName) || !info.quotaInfo) {
+    return null;
+  }
+
+  const fraction = info.quotaInfo.remainingFraction ?? 0;
+  return {
+    percentage: Math.floor(fraction * 100),
+    resetTime: info.quotaInfo.resetTime || '',
+    display_name: info.displayName,
+    supports_images: info.supportsImages,
+    supports_thinking: info.supportsThinking,
+    thinking_budget: info.thinkingBudget,
+    recommended: info.recommended,
+    max_tokens: info.maxTokens,
+    max_output_tokens: info.maxOutputTokens,
+    supported_mime_types: info.supportedMimeTypes,
+  };
+}
+
+function toModelForwardingRules(
+  deprecatedModelIds: FetchModelsResponse['deprecatedModelIds'],
+): Record<string, string> | undefined {
+  if (!deprecatedModelIds || Object.keys(deprecatedModelIds).length === 0) {
+    return undefined;
+  }
+
+  const forwardingRules: Record<string, string> = {};
+  for (const [oldModelId, deprecatedInfo] of Object.entries(deprecatedModelIds)) {
+    if (typeof deprecatedInfo.newModelId === 'string' && deprecatedInfo.newModelId !== '') {
+      forwardingRules[oldModelId] = deprecatedInfo.newModelId;
+    }
+  }
+
+  return Object.keys(forwardingRules).length > 0 ? forwardingRules : undefined;
 }
 
 // --- Service Implementation ---
@@ -218,70 +361,71 @@ export class GoogleAPIService {
     return response.json() as Promise<UserInfo>;
   }
 
-  /**
+  public static async fetchProjectContext(accessToken: string): Promise<ProjectContext> {
-   * Core logic: Fetches the internal project ID needed for quota checks.
-   */
-  public static async fetchProjectId(accessToken: string): Promise<string | null> {
     const body = {
       metadata: { ideType: 'ANTIGRAVITY' },
     };
 
-    // Simple retry logic (2 attempts)
+    let projectId: string | undefined;
+    let subscriptionTier: string | undefined;
+
     for (let i = 0; i < 2; i++) {
       try {
         const response = await fetch(URLS.LOAD_PROJECT, {
           method: 'POST',
-          headers: {
+          headers: buildInternalApiHeaders(accessToken),
-            Authorization: `Bearer ${accessToken}`,
-            'User-Agent': USER_AGENT,
-            'Content-Type': 'application/json',
-          },
           body: JSON.stringify(body),
+          signal: createTimeoutSignal(REQUEST_TIMEOUT_MS),
           ...this.getFetchOptions(),
         });
 
         if (response.ok) {
           const data = (await response.json()) as LoadProjectResponse;
-          if (data.cloudaicompanionProject) {
+          if (typeof data.cloudaicompanionProject === 'string') {
-            return data.cloudaicompanionProject;
+            projectId = data.cloudaicompanionProject;
           }
+          subscriptionTier = resolveSubscriptionTier(data);
+          break;
         }
-      } catch (e) {
+      } catch (error) {
-        logger.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, e);
+        logger.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, error);
-        await new Promise((r) => setTimeout(r, 500)); // Sleep 500ms
+        await sleep(500);
       }
     }
-    return null;
+
+    return {
+      projectId,
+      subscriptionTier,
+    };
+  }
+
+  public static async fetchProjectId(accessToken: string): Promise<string | null> {
+    const context = await this.fetchProjectContext(accessToken);
+    return context.projectId ?? null;
   }
 
   /**
    * Core logic: Fetches detailed model quota information.
    */
   static async fetchQuota(accessToken: string): Promise<QuotaData> {
-    // 1. Get Project ID (Critical)
+    const { projectId, subscriptionTier } = await this.fetchProjectContext(accessToken);
-    const projectId = await this.fetchProjectId(accessToken);
 
-    // 2. Build Payload
     const payload: Record<string, unknown> = {};
     if (projectId) {
       payload['project'] = projectId;
     }
 
-    // 3. Send Request with Retries
     const maxRetries = 3;
     let lastError: Error | null = null;
-    const fetchOptions = this.getFetchOptions(); // Reuse options
+    const fetchOptions = this.getFetchOptions();
 
     for (let attempt = 1; attempt <= maxRetries; attempt++) {
       try {
         const response = await fetch(URLS.QUOTA, {
           method: 'POST',
-          headers: {
+          headers: buildInternalApiHeaders(accessToken),
-            Authorization: `Bearer ${accessToken}`,
-            'User-Agent': USER_AGENT,
-            'Content-Type': 'application/json',
-          },
           body: JSON.stringify(payload),
+          signal: createTimeoutSignal(REQUEST_TIMEOUT_MS),
           ...fetchOptions,
         });
 
@@ -302,39 +446,42 @@ export class GoogleAPIService {
           );
 
           if (attempt < maxRetries) {
-            await new Promise((r) => setTimeout(r, 1000));
+            await sleep(1000);
             continue;
           } else {
             throw new Error(errorMsg);
           }
         }
 
-        const data = (await response.json()) as { models: Record<string, ModelInfoRaw> };
+        const data = (await response.json()) as FetchModelsResponse;
-        const result: QuotaData = { models: {} };
+        const result: QuotaData = {
+          models: {},
+          subscription_tier: subscriptionTier,
+          is_forbidden: false,
+        };
 
-        // Parse relevant models
+        for (const [modelName, modelInfoRaw] of Object.entries(data.models || {})) {
-        for (const [name, info] of Object.entries(data.models || {})) {
+          const modelQuota = toModelQuotaInfo(modelName, modelInfoRaw);
-          if (info.quotaInfo) {
+          if (modelQuota) {
-            const fraction = info.quotaInfo.remainingFraction ?? 0;
+            result.models[modelName] = modelQuota;
-            const percentage = Math.floor(fraction * 100);
-            const resetTime = info.quotaInfo.resetTime || '';
-
-            if (name.toLowerCase().includes('gemini') || name.toLowerCase().includes('claude')) {
-              result.models[name] = { percentage, resetTime };
-            }
           }
         }
 
+        const modelForwardingRules = toModelForwardingRules(data.deprecatedModelIds);
+        if (modelForwardingRules) {
+          result.model_forwarding_rules = modelForwardingRules;
+        }
+
         return result;
-      } catch (e: unknown) {
+      } catch (error: unknown) {
-        if (e instanceof Error) {
+        if (error instanceof Error) {
           logger.warn(
-            `[GoogleAPIService] Request failed: ${e.message} (Attempt ${attempt}/${maxRetries})`,
+            `[GoogleAPIService] Request failed: ${error.message} (Attempt ${attempt}/${maxRetries})`,
           );
         }
-        lastError = e instanceof Error ? e : new Error(String(e));
+        lastError = error instanceof Error ? error : new Error(String(error));
         if (attempt < maxRetries) {
-          await new Promise((r) => setTimeout(r, 1000));
+          await sleep(1000);
         }
       }
     }

src/tests/unit/gemini-controller.integration.test.ts

@@ -31,7 +31,14 @@ describe('GeminiController Integration', () => {
       expect.objectContaining({
         models: expect.arrayContaining([
           expect.objectContaining({
-            name: 'models/gemini-2.5-flash',
+            name: 'models/gemini-3-flash',
+            description: '',
+            inputTokenLimit: 128000,
+            outputTokenLimit: 8192,
+            supportedGenerationMethods: ['generateContent', 'countTokens'],
+          }),
+          expect.objectContaining({
+            name: 'models/gemini-3.1-pro-high',
             description: '',
             inputTokenLimit: 128000,
             outputTokenLimit: 8192,
@@ -74,13 +81,13 @@ describe('GeminiController Integration', () => {
     const reply = createReplyMock();
 
     await controller.modelAction(
-      'models/gemini-2.5-flash:generateContent',
+      'models/gemini-3.1-pro-high:generateContent',
       { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] } as any,
       reply as any,
     );
 
     expect(proxyService.handleGeminiGenerateContent).toHaveBeenCalledWith(
-      'models/gemini-2.5-flash',
+      'models/gemini-3.1-pro-high',
       expect.any(Object),
     );
     expect(reply.status).toHaveBeenCalledWith(200);

src/tests/unit/parity-rollout-guard.test.ts

@@ -0,0 +1,106 @@
+import { beforeEach, describe, expect, it } 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';
+
+function createProxyConfig(overrides: Partial<ProxyConfig>): ProxyConfig {
+  return {
+    ...DEFAULT_APP_CONFIG.proxy,
+    ...overrides,
+    upstream_proxy: {
+      ...DEFAULT_APP_CONFIG.proxy.upstream_proxy,
+      ...(overrides.upstream_proxy ?? {}),
+    },
+  };
+}
+
+function seedTokens(service: TokenManagerService): void {
+  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',
+      },
+    ],
+    [
+      'acc-2',
+      {
+        account_id: 'acc-2',
+        email: '[email protected]',
+        access_token: 'token-2',
+        refresh_token: 'refresh-2',
+        token_type: 'Bearer',
+        expires_in: 3600,
+        expiry_timestamp: nowSec + 3600,
+        project_id: 'project-2',
+      },
+    ],
+  ]);
+}
+
+describe('Parity rollout guard replay', () => {
+  let service: TokenManagerService;
+
+  beforeEach(() => {
+    service = new TokenManagerService();
+    seedTokens(service);
+  });
+
+  it('supports kill-switch rollback from parity path to baseline path', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+    const paritySelected = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(paritySelected?.id).toBe('acc-2');
+
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: true,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+    (service as any).currentIndex = 0;
+    const rollbackSelected = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(rollbackSelected?.id).toBe('acc-1');
+  });
+
+  it('enforces no-go threshold after shadow mismatch and blocks parity enablement', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: false,
+        parity_shadow_enabled: true,
+        parity_no_go_mismatch_rate: 0,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+
+    const shadowSelection = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(shadowSelection?.id).toBe('acc-1');
+
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_shadow_enabled: false,
+        parity_kill_switch: false,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+    (service as any).currentIndex = 0;
+    const blockedSelection = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(blockedSelection?.id).toBe('acc-1');
+  });
+});
+

src/tests/unit/proxy-controller.integration.test.ts

@@ -200,6 +200,73 @@ describe('ProxyController Integration', () => {
     );
   });
 
+  it('supports OpenAI responses compatibility endpoint in stream mode with SSE headers', async () => {
+    const stream = of('data: {"id":"chatcmpl_resp_stream"}\n\n');
+    const proxyService = {
+      handleChatCompletions: vi.fn().mockResolvedValue(stream),
+    };
+
+    const controller = new ProxyController(proxyService as any);
+    const reply = createReplyMock();
+
+    await controller.responses(
+      {
+        model: 'gpt-4o',
+        instructions: 'stream output',
+        input: 'hello',
+        stream: true,
+      },
+      reply as any,
+    );
+
+    expect(reply.header).toHaveBeenCalledWith('Content-Type', 'text/event-stream');
+    expect(reply.header).toHaveBeenCalledWith('Cache-Control', 'no-cache');
+    expect(reply.header).toHaveBeenCalledWith('Connection', 'keep-alive');
+    expect(reply.send).toHaveBeenCalledWith(stream);
+  });
+
+  it('normalizes web_search_call in /v1/responses into builtin_web_search tool messages', async () => {
+    const proxyService = {
+      handleChatCompletions: vi.fn().mockResolvedValue({
+        id: 'chatcmpl_resp_search',
+        object: 'chat.completion',
+        created: 1700000002,
+        model: 'gpt-4o',
+        choices: [{ index: 0, finish_reason: 'stop', message: { content: 'done' } }],
+        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
+      }),
+    };
+
+    const controller = new ProxyController(proxyService as any);
+    const reply = createReplyMock();
+
+    await controller.responses(
+      {
+        model: 'gpt-4o',
+        input: [
+          {
+            type: 'web_search_call',
+            call_id: 'call_search_1',
+            action: { query: 'gemini api' },
+          },
+          {
+            type: 'function_call_output',
+            call_id: 'call_search_1',
+            output: { content: 'search result' },
+          },
+        ],
+      },
+      reply as any,
+    );
+
+    const callArg = proxyService.handleChatCompletions.mock.calls[0][0];
+    const assistantMessage = callArg.messages.find((message: { role: string }) => message.role === 'assistant');
+    const toolMessage = callArg.messages.find((message: { role: string }) => message.role === 'tool');
+
+    expect(assistantMessage?.tool_calls?.[0]?.function?.name).toBe('builtin_web_search');
+    expect(toolMessage?.name).toBe('builtin_web_search');
+  });
+
   it('supports image generations endpoint', async () => {
     const proxyService = {
       handleChatCompletions: vi.fn().mockResolvedValue({

src/tests/unit/streaming-state.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect, beforeEach } from 'vitest';
-import { StreamingState } from '../../lib/antigravity/ClaudeStreamingMapper';
+import { PartProcessor, StreamingState } from '../../lib/antigravity/ClaudeStreamingMapper';
 
 describe('StreamingState', () => {
   let state: StreamingState;
@@ -58,4 +58,48 @@ describe('StreamingState', () => {
       expect(state.getErrorCount()).toBe(2);
     });
   });
+
+  describe('stream aggregation compatibility', () => {
+    it('emits tool_use stop reason when functionCall appears in stream', () => {
+      const processor = new PartProcessor(state);
+      const functionChunks = processor.process({
+        functionCall: {
+          name: 'builtin_web_search',
+          args: { query: 'gemini docs' },
+          id: 'call_stream_1',
+        },
+      });
+      const finishChunks = state.emitFinish('STOP', {
+        promptTokenCount: 2,
+        candidatesTokenCount: 3,
+      } as any);
+
+      const output = [...functionChunks, ...finishChunks].join('');
+      expect(output).toContain('"type":"tool_use"');
+      expect(output).toContain('"stop_reason":"tool_use"');
+      expect(output).toContain('"message_stop"');
+    });
+
+    it('aggregates grounding metadata into final text block', () => {
+      state.webSearchQuery = 'gemini api';
+      state.groundingChunks = [
+        {
+          web: {
+            title: 'Gemini API Docs',
+            uri: 'https://example.com/gemini',
+          },
+        },
+      ];
+
+      const chunks = state.emitFinish('STOP', {
+        promptTokenCount: 1,
+        candidatesTokenCount: 1,
+      } as any);
+      const output = chunks.join('');
+
+      expect(output).toContain('Searched for you');
+      expect(output).toContain('Citations');
+      expect(output).toContain('https://example.com/gemini');
+    });
+  });
 });

src/tests/unit/token-manager-parity-scheduling.test.ts

@@ -0,0 +1,166 @@
+import { beforeEach, describe, expect, it } 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';
+
+function createProxyConfig(overrides: Partial<ProxyConfig>): ProxyConfig {
+  return {
+    ...DEFAULT_APP_CONFIG.proxy,
+    ...overrides,
+    upstream_proxy: {
+      ...DEFAULT_APP_CONFIG.proxy.upstream_proxy,
+      ...(overrides.upstream_proxy ?? {}),
+    },
+  };
+}
+
+function seedTokens(service: TokenManagerService): void {
+  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',
+      },
+    ],
+    [
+      'acc-2',
+      {
+        account_id: 'acc-2',
+        email: '[email protected]',
+        access_token: 'token-2',
+        refresh_token: 'refresh-2',
+        token_type: 'Bearer',
+        expires_in: 3600,
+        expiry_timestamp: nowSec + 3600,
+        project_id: 'project-2',
+        session_id: 'session-2',
+      },
+    ],
+  ]);
+}
+
+describe('TokenManagerService parity scheduling replay', () => {
+  let service: TokenManagerService;
+
+  beforeEach(() => {
+    service = new TokenManagerService();
+    seedTokens(service);
+  });
+
+  it('prioritizes preferred account in parity mode', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        scheduling_mode: 'balance',
+        preferred_account_id: 'acc-2',
+      }),
+    );
+
+    const token = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(token?.id).toBe('acc-2');
+  });
+
+  it('rotates sticky account when limited in balance mode', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        scheduling_mode: 'balance',
+        preferred_account_id: '',
+      }),
+    );
+
+    const first = await service.getNextToken({
+      sessionKey: 'openai:user-1',
+      model: 'gemini-2.5-flash',
+    });
+    expect(first?.id).toBe('acc-1');
+
+    await service.markFromUpstreamError({
+      accountIdOrEmail: 'acc-1',
+      status: 429,
+      model: 'gemini-2.5-flash',
+      body: JSON.stringify({
+        error: {
+          details: [{ reason: 'RATE_LIMIT_EXCEEDED' }],
+        },
+      }),
+    });
+
+    const second = await service.getNextToken({
+      sessionKey: 'openai:user-1',
+      model: 'gemini-2.5-flash',
+    });
+    expect(second?.id).toBe('acc-2');
+  });
+
+  it('applies model-level lock for quota exhausted only on the same model', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        scheduling_mode: 'performance-first',
+      }),
+    );
+
+    await service.markFromUpstreamError({
+      accountIdOrEmail: 'acc-1',
+      status: 429,
+      model: 'gemini-2.5-flash',
+      body: JSON.stringify({
+        error: {
+          details: [{ reason: 'QUOTA_EXHAUSTED', metadata: { quotaResetDelay: '30s' } }],
+        },
+      }),
+    });
+
+    const sameModel = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(sameModel?.id).toBe('acc-2');
+
+    const otherModel = await service.getNextToken({
+      model: 'gemini-2.5-pro',
+      excludeAccountIds: ['acc-2'],
+    });
+    expect(otherModel?.id).toBe('acc-1');
+  });
+
+  it('falls back to excluded pool when retry exclusions would empty all candidates', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: false,
+        parity_kill_switch: false,
+      }),
+    );
+
+    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',
+        },
+      ],
+    ]);
+
+    const token = await service.getNextToken({ excludeAccountIds: ['acc-1'] });
+    expect(token?.id).toBe('acc-1');
+  });
+});

src/types/config.ts

@@ -6,14 +6,24 @@ export const UpstreamProxyConfigSchema = z.object({
 });
 
 export const ProxyConfigSchema = z.object({
-  enabled: z.boolean(), // 是否启用
+  enabled: z.boolean(),
-  port: z.number(), // 监听端口
+  port: z.number(),
-  api_key: z.string(), // API 密钥 (Also used for Admin Mode auth)
+  api_key: z.string(),
-  auto_start: z.boolean(), // 是否自动启动
+  auto_start: z.boolean(),
   backend_canary_enabled: z.boolean().default(true),
+  parity_enabled: z.boolean().default(false),
+  parity_shadow_enabled: z.boolean().default(false),
+  parity_kill_switch: z.boolean().default(false),
+  parity_no_go_mismatch_rate: z.number().default(0.15),
+  parity_no_go_error_rate: z.number().default(0.4),
+  scheduling_mode: z.enum(['cache-first', 'balance', 'performance-first']).default('balance'),
+  max_wait_seconds: z.number().default(60),
+  preferred_account_id: z.string().default(''),
+  circuit_breaker_enabled: z.boolean().default(true),
+  circuit_breaker_backoff_steps: z.array(z.number()).default([60, 300, 1800, 7200]),
   custom_mapping: z.record(z.string(), z.string()).default({}),
-  anthropic_mapping: z.record(z.string(), z.string()), // 映射表
+  anthropic_mapping: z.record(z.string(), z.string()), // Mapping table
-  request_timeout: z.number().default(120), // 超时秒数
+  request_timeout: z.number().default(120), // Timeout in seconds
   upstream_proxy: UpstreamProxyConfigSchema,
 });
 
@@ -27,7 +37,7 @@ export const AppConfigSchema = z.object({
   auto_startup: z.boolean(),
   error_reporting_enabled: z.boolean(),
   privacy_consent_asked: z.boolean().optional().default(false), // Optional for backward compatibility
-  default_export_path: z.string().nullable().optional(), // 导出路径
+  default_export_path: z.string().nullable().optional(), // Export path
   model_visibility: z.record(z.string(), z.boolean()).default({}), // Model visibility preferences
   provider_groupings_enabled: z.boolean().default(false), // Enable provider groupings UI
   grid_layout: z.enum(['auto', '2-col', '3-col', 'list']).default('auto'), // Account card grid layout
@@ -58,6 +68,16 @@ export const DEFAULT_APP_CONFIG: AppConfig = {
     api_key: '', // Generated dynamically if default needed
     auto_start: false,
     backend_canary_enabled: true,
+    parity_enabled: false,
+    parity_shadow_enabled: false,
+    parity_kill_switch: false,
+    parity_no_go_mismatch_rate: 0.15,
+    parity_no_go_error_rate: 0.4,
+    scheduling_mode: 'balance',
+    max_wait_seconds: 60,
+    preferred_account_id: '',
+    circuit_breaker_enabled: true,
+    circuit_breaker_backoff_steps: [60, 300, 1800, 7200],
     custom_mapping: {},
     anthropic_mapping: {},
     request_timeout: 120,