feat: implement protocol parity and harden upstream handling (#88)

* feat: implement protocol parity and harden upstream handling

* docs(agents): add lodash-es utility method and named import guideline

src/server/guards/admin.guard.ts

@@ -0,0 +1,25 @@
+import { CanActivate, ExecutionContext, Injectable, UnauthorizedException } from '@nestjs/common';
+import { getServerConfig } from '../server-config';
+import { extractApiKeyToken, hasConfiguredApiKey, RequestHeaders } from './api-key-auth.util';
+
+@Injectable()
+export class AdminGuard implements CanActivate {
+  canActivate(context: ExecutionContext): boolean {
+    const config = getServerConfig();
+    const apiKey = config?.api_key;
+
+    if (!hasConfiguredApiKey(apiKey)) {
+      return true;
+    }
+
+    const request = context.switchToHttp().getRequest();
+    const headers = request.headers as RequestHeaders;
+    const clientToken = extractApiKeyToken(headers);
+
+    if (clientToken && clientToken === apiKey) {
+      return true;
+    }
+
+    throw new UnauthorizedException('API key validation failed');
+  }
+}

src/server/guards/api-key-auth.util.ts

@@ -0,0 +1,49 @@
+export type RequestHeaderValue = string | string[] | undefined;
+
+export type RequestHeaders = Record<string, RequestHeaderValue>;
+
+export function hasConfiguredApiKey(apiKey: string | undefined): apiKey is string {
+  return typeof apiKey === 'string' && apiKey.trim() !== '';
+}
+
+export function extractApiKeyToken(headers: RequestHeaders): string | null {
+  const authorizationHeader = readHeaderValue(headers['authorization']);
+  if (authorizationHeader) {
+    const [scheme, token, ...rest] = authorizationHeader.trim().split(/\s+/);
+    if (rest.length === 0 && token && scheme?.toLowerCase() === 'bearer') {
+      return token;
+    }
+  }
+
+  const xApiKeyHeader = readHeaderValue(headers['x-api-key']);
+  if (xApiKeyHeader) {
+    return xApiKeyHeader;
+  }
+
+  const xGoogApiKeyHeader = readHeaderValue(headers['x-goog-api-key']);
+  if (xGoogApiKeyHeader) {
+    return xGoogApiKeyHeader;
+  }
+
+  return null;
+}
+
+function readHeaderValue(headerValue: RequestHeaderValue): string | null {
+  if (typeof headerValue === 'string') {
+    const trimmedValue = headerValue.trim();
+    return trimmedValue !== '' ? trimmedValue : null;
+  }
+
+  if (Array.isArray(headerValue)) {
+    for (const value of headerValue) {
+      if (typeof value === 'string') {
+        const trimmedValue = value.trim();
+        if (trimmedValue !== '') {
+          return trimmedValue;
+        }
+      }
+    }
+  }
+
+  return null;
+}

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

@@ -1,16 +1,29 @@
 import { Injectable, Logger } from '@nestjs/common';
-import axios from 'axios';
+import axios, { AxiosProxyConfig, AxiosRequestConfig, AxiosResponse } from 'axios';
+import { isNil } from 'lodash-es';
 import { GeminiRequest, GeminiResponse } from '../interfaces/request-interfaces';
 import { GeminiInternalRequest } from '../../../../lib/antigravity/types';
+import { getServerConfig } from '../../../server-config';
 
 @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 = [
+    'https://cloudcode-pa.googleapis.com/v1internal',
+    'https://daily-cloudcode-pa.googleapis.com/v1internal',
+  ];
 
-  async streamGenerate(model: string, content: GeminiRequest, accessToken: string): Promise<any> {
+  async streamGenerate(
+    model: string,
+    content: GeminiRequest,
+    accessToken: string,
+    upstreamProxyUrl?: string,
+  ): Promise<NodeJS.ReadableStream> {
     const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`;
+    const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
 
     try {
       const response = await axios.post(url, content, {
@@ -20,15 +33,15 @@ export class GeminiClient {
         },
         responseType: 'stream',
         timeout: 60000,
+        proxy: axiosProxy,
       });
       return response.data;
     } catch (error) {
       if (axios.isAxiosError(error)) {
-        this.logger.error(`Gemini Stream API Error: ${error.message}`);
+        this.logger.error(`Gemini stream request failed: ${error.message}`);
         throw new Error(error.response?.data?.error?.message || error.message);
       }
-      // Re-throw as clean Error to avoid circular reference issues
-      throw error instanceof Error ? new Error(error.message) : new Error(String(error));
+      this.throwAsCleanError(error);
     }
   }
 
@@ -36,8 +49,10 @@ export class GeminiClient {
     model: string,
     content: GeminiRequest,
     accessToken: string,
+    upstreamProxyUrl?: string,
   ): Promise<GeminiResponse> {
     const url = `${this.baseUrl}/models/${model}:generateContent`;
+    const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
 
     try {
       const response = await axios.post<GeminiResponse>(url, content, {
@@ -46,68 +61,221 @@ export class GeminiClient {
           'Content-Type': 'application/json',
         },
         timeout: 60000, // 60s timeout
+        proxy: axiosProxy,
       });
       return response.data;
     } catch (error) {
       if (axios.isAxiosError(error)) {
         this.logger.error(
-          `Gemini API Error: ${error.message} - ${this.safeStringify(error.response?.data)}`,
+          `Gemini request failed: ${error.message} - ${this.safeStringify(error.response?.data)}`,
         );
         throw new Error(error.response?.data?.error?.message || error.message);
       }
-      // Re-throw as clean Error to avoid circular reference issues
-      throw error instanceof Error ? new Error(error.message) : new Error(String(error));
+      this.throwAsCleanError(error);
     }
   }
 
-  // --- Antigravity Internal API Support ---
+  // --- Internal Gateway API Support ---
 
-  private readonly internalBaseUrl = 'https://cloudcode-pa.googleapis.com/v1internal';
-
-  async streamGenerateInternal(body: GeminiInternalRequest, accessToken: string): Promise<any> {
-    const url = `${this.internalBaseUrl}:streamGenerateContent?alt=sse`;
-    try {
-      const response = await axios.post(url, body, {
-        headers: {
-          Authorization: `Bearer ${accessToken}`,
-          'Content-Type': 'application/json',
-          'User-Agent': 'antigravity/1.11.9 windows/amd64',
-        },
+  async streamGenerateInternal(
+    body: GeminiInternalRequest,
+    accessToken: string,
+    upstreamProxyUrl?: string,
+    extraHeaders?: Record<string, string>,
+  ): Promise<NodeJS.ReadableStream> {
+    const response = await this.requestWithEndpointFailover<NodeJS.ReadableStream>(
+      ':streamGenerateContent?alt=sse',
+      body,
+      accessToken,
+      upstreamProxyUrl,
+      {
         responseType: 'stream',
-        timeout: 60000,
-      });
-      return response.data;
-    } catch (error) {
-      this.handleAxiosError(error, 'StreamInternal');
-    }
+      },
+      'stream-generate',
+      extraHeaders,
+    );
+
+    return response.data;
   }
 
-  async generateInternal(body: GeminiInternalRequest, accessToken: string): Promise<any> {
-    const url = `${this.internalBaseUrl}:generateContent`;
-    try {
-      const response = await axios.post(url, body, {
-        headers: {
-          Authorization: `Bearer ${accessToken}`,
-          'Content-Type': 'application/json',
-          'User-Agent': 'antigravity/1.11.9 windows/amd64',
-        },
-        timeout: 60000,
-      });
-      // v1internal API wraps the response in a 'response' field, unwrap it
-      return response.data.response || response.data;
-    } catch (error) {
-      this.handleAxiosError(error, 'GenerateInternal');
+  async generateInternal(
+    body: GeminiInternalRequest,
+    accessToken: string,
+    upstreamProxyUrl?: string,
+    extraHeaders?: Record<string, string>,
+  ): Promise<GeminiResponse> {
+    const response = await this.requestWithEndpointFailover<
+      GeminiResponse | { response: GeminiResponse }
+    >(
+      ':generateContent',
+      body,
+      accessToken,
+      upstreamProxyUrl,
+      {},
+      'generate-content',
+      extraHeaders,
+    );
+    const payload = response.data;
+    if (payload && typeof payload === 'object' && 'response' in payload) {
+      return payload.response;
     }
+    return payload;
   }
 
-  private handleAxiosError(error: any, context: string) {
+  private getInternalBaseUrls(): string[] {
+    const fromEnv =
+      process.env.PROXY_INTERNAL_BASE_URLS ?? process.env.ANTIGRAVITY_INTERNAL_BASE_URLS;
+    const configuredBaseUrls = fromEnv
+      ?.split(',')
+      .map((item) => item.trim())
+      .filter(Boolean);
+
+    if (configuredBaseUrls && configuredBaseUrls.length > 0) {
+      return configuredBaseUrls.map((url) => url.replace(/\/+$/, ''));
+    }
+
+    return this.defaultInternalBaseUrls.map((url) => url.replace(/\/+$/, ''));
+  }
+
+  private getInternalTimeoutMs(): number {
+    const config = getServerConfig();
+    const timeoutSeconds = config?.request_timeout ?? 60;
+    return Math.max(1, timeoutSeconds) * 1000;
+  }
+
+  private shouldSwitchEndpointOnError(error: unknown): boolean {
+    if (!axios.isAxiosError(error)) {
+      return false;
+    }
+
+    if (!error.response) {
+      return true;
+    }
+
+    const status = error.response.status;
+
+    // Permanent auth errors should fail fast for current token.
+    if (status === 401 || status === 403) {
+      return false;
+    }
+
+    return status === 408 || status === 429 || status >= 500;
+  }
+
+  private async requestWithEndpointFailover<T>(
+    path: string,
+    body: GeminiInternalRequest,
+    accessToken: string,
+    upstreamProxyUrl: string | undefined,
+    config: AxiosRequestConfig,
+    operation: string,
+    extraHeaders?: Record<string, string>,
+  ): Promise<AxiosResponse<T>> {
+    const baseUrls = this.getInternalBaseUrls();
+    const timeout = this.getInternalTimeoutMs();
+    const requestUserAgent = process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent;
+    const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
+    let lastError: unknown = null;
+
+    for (let index = 0; index < baseUrls.length; index++) {
+      const baseUrl = baseUrls[index];
+      const url = `${baseUrl}${path}`;
+
+      try {
+        return await axios.post<T>(url, body, {
+          headers: {
+            Authorization: `Bearer ${accessToken}`,
+            'Content-Type': 'application/json',
+            'User-Agent': requestUserAgent,
+            ...(extraHeaders ?? {}),
+          },
+          timeout,
+          proxy: axiosProxy,
+          ...config,
+        });
+      } catch (error) {
+        lastError = error;
+        const hasNextEndpoint = index < baseUrls.length - 1;
+
+        if (!hasNextEndpoint || !this.shouldSwitchEndpointOnError(error)) {
+          this.handleAxiosError(error, operation);
+        }
+
+        this.logger.warn(
+          `[${operation}] request failed at ${baseUrl}; trying next endpoint (${index + 2}/${
+            baseUrls.length
+          }).`,
+        );
+      }
+    }
+
+    this.handleAxiosError(lastError, operation);
+  }
+
+  private handleAxiosError(error: unknown, operation: string): never {
     if (axios.isAxiosError(error)) {
+      const responseData = error.response?.data;
+      const upstreamMessage = this.extractAxiosErrorMessage(responseData);
       this.logger.error(
-        `Gemini ${context} API Error: ${error.message} - ${this.safeStringify(error.response?.data)}`,
+        `[${operation}] upstream request error: ${error.message} - ${this.safeStringify(responseData)}`,
       );
-      throw new Error(error.response?.data?.error?.message || error.message);
+      throw new Error(upstreamMessage || error.message);
     }
-    // Re-throw as clean Error to avoid circular reference issues
+    this.throwAsCleanError(error);
+  }
+
+  private extractAxiosErrorMessage(responseData: unknown): string | null {
+    if (!responseData || typeof responseData !== 'object') {
+      return null;
+    }
+
+    const errorRecord = (responseData as { error?: unknown }).error;
+    if (!errorRecord || typeof errorRecord !== 'object') {
+      return null;
+    }
+
+    const message = (errorRecord as { message?: unknown }).message;
+    return typeof message === 'string' && message.trim() !== '' ? message : null;
+  }
+
+  private resolveAxiosProxy(upstreamProxyUrl?: string): AxiosProxyConfig | false | undefined {
+    const config = getServerConfig();
+    const configuredProxyUrl =
+      upstreamProxyUrl ||
+      (config?.upstream_proxy?.enabled && config.upstream_proxy.url
+        ? config.upstream_proxy.url
+        : '');
+
+    if (!configuredProxyUrl) {
+      return undefined;
+    }
+
+    try {
+      const parsed = new URL(configuredProxyUrl);
+      const port = parsed.port ? Number(parsed.port) : parsed.protocol === 'https:' ? 443 : 80;
+
+      const proxyConfig: AxiosProxyConfig = {
+        protocol: parsed.protocol.replace(':', ''),
+        host: parsed.hostname,
+        port,
+      };
+
+      if (parsed.username || parsed.password) {
+        proxyConfig.auth = {
+          username: decodeURIComponent(parsed.username),
+          password: decodeURIComponent(parsed.password),
+        };
+      }
+
+      return proxyConfig;
+    } catch {
+      this.logger.warn(`Upstream proxy URL is invalid: ${configuredProxyUrl}`);
+      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));
   }
 
@@ -115,7 +283,7 @@ export class GeminiClient {
    * Safely stringify an object, handling circular references
    */
   private safeStringify(obj: unknown): string {
-    if (obj === undefined || obj === null) {
+    if (isNil(obj)) {
       return String(obj);
     }
     try {

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

@@ -0,0 +1,292 @@
+import { Body, Controller, Get, HttpStatus, Inject, Param, Post, Res, UseGuards } from '@nestjs/common';
+import { FastifyReply } from 'fastify';
+import { isEmpty } from 'lodash-es';
+import { Observable } from 'rxjs';
+
+import { ProxyGuard } from './proxy.guard';
+import { ProxyService } from './proxy.service';
+import { GeminiRequest, GeminiResponse } from './interfaces/request-interfaces';
+import { getServerConfig } from '../../server-config';
+import { getAllDynamicModels } from '../../../lib/antigravity/ModelMapping';
+
+type GeminiModelMetadata = {
+  name: string;
+  displayName: string;
+  description: string;
+  inputTokenLimit: number;
+  outputTokenLimit: number;
+  supportedGenerationMethods: string[];
+  temperature: number;
+  topK: number;
+  topP: number;
+  version: string;
+};
+
+@Controller('v1beta')
+@UseGuards(ProxyGuard)
+export class GeminiController {
+  constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {}
+
+  @Get('models')
+  listModels(@Res() res: FastifyReply) {
+    const models = this.buildGeminiModelList();
+    res.status(HttpStatus.OK).send({
+      models,
+    });
+  }
+
+  @Get('models/:model')
+  getModel(@Param('model') model: string, @Res() res: FastifyReply) {
+    const targetName = model.startsWith('models/') ? model : `models/${model}`;
+    const matched = this.buildGeminiModelList().find((item) => item.name === targetName);
+
+    if (matched) {
+      res.status(HttpStatus.OK).send({
+        name: matched.name,
+        displayName: matched.displayName,
+      });
+      return;
+    }
+
+    res.status(HttpStatus.OK).send({
+      name: targetName,
+      displayName: targetName.replace(/^models\//, ''),
+    });
+  }
+
+  @Post('models/:modelAction')
+  async modelAction(
+    @Param('modelAction') modelAction: string,
+    @Body() body: GeminiRequest,
+    @Res() res: FastifyReply,
+  ) {
+    const parsed = this.parseCombinedModelAction(modelAction);
+    if (!parsed) {
+      res.status(HttpStatus.BAD_REQUEST).send({
+        error: {
+          code: HttpStatus.BAD_REQUEST,
+          message: 'Model action format is invalid',
+          status: 'INVALID_ARGUMENT',
+        },
+      });
+      return;
+    }
+
+    await this.dispatchAction(parsed.model, parsed.action, body, res);
+  }
+
+  @Post('models/:model/countTokens')
+  async countTokens(
+    @Param('model') model: string,
+    @Body() body: GeminiRequest,
+    @Res() res: FastifyReply,
+  ) {
+    await this.dispatchAction(`models/${model}`, 'countTokens', body, res);
+  }
+
+  private async dispatchAction(
+    model: string,
+    action: string,
+    body: GeminiRequest,
+    res: FastifyReply,
+  ): Promise<void> {
+    if (action === 'countTokens') {
+      res.status(HttpStatus.OK).send({
+        totalTokens: 0,
+      });
+      return;
+    }
+
+    try {
+      if (action === 'streamGenerateContent') {
+        const stream = await this.proxyService.handleGeminiStreamGenerateContent(model, body);
+        if (stream instanceof Observable) {
+          this.sendObservableSse(res, stream);
+          return;
+        }
+      }
+
+      if (action === 'generateContent') {
+        const result = await this.proxyService.handleGeminiGenerateContent(model, body);
+        res.status(HttpStatus.OK).send(this.normalizeGeminiGenerateResponse(result));
+        return;
+      }
+
+      res.status(HttpStatus.BAD_REQUEST).send({
+        error: {
+          code: HttpStatus.BAD_REQUEST,
+          message: `Unsupported model action: ${action}`,
+          status: 'INVALID_ARGUMENT',
+        },
+      });
+    } catch (error) {
+      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({
+        error: {
+          code: HttpStatus.INTERNAL_SERVER_ERROR,
+          message,
+          status: 'INTERNAL',
+        },
+      });
+    }
+  }
+
+  private parseCombinedModelAction(
+    modelAction: string,
+  ): {
+    model: string;
+    action: string;
+  } | null {
+    const colonIndex = modelAction.lastIndexOf(':');
+    if (colonIndex <= 0) {
+      return null;
+    }
+
+    const model = modelAction.slice(0, colonIndex).trim();
+    const action = modelAction.slice(colonIndex + 1).trim();
+    if (!model || !action) {
+      return null;
+    }
+
+    const prefixedModel = model.startsWith('models/') ? model : `models/${model}`;
+    return {
+      model: prefixedModel,
+      action,
+    };
+  }
+
+  private buildGeminiModelList(): GeminiModelMetadata[] {
+    const config = getServerConfig();
+    const dynamicModelIds = getAllDynamicModels(config?.custom_mapping ?? {});
+
+    return dynamicModelIds.map((id) => this.toGeminiModelMetadata(`models/${id}`));
+  }
+
+  private toGeminiModelMetadata(modelName: string): GeminiModelMetadata {
+    const displayName = modelName.replace(/^models\//, '');
+    return {
+      name: modelName,
+      displayName,
+      description: '',
+      inputTokenLimit: 128000,
+      outputTokenLimit: 8192,
+      supportedGenerationMethods: ['generateContent', 'countTokens'],
+      temperature: 1,
+      topK: 64,
+      topP: 0.95,
+      version: '001',
+    };
+  }
+
+  private normalizeGeminiGenerateResponse(response: GeminiResponse): GeminiResponse {
+    const candidates = (response.candidates ?? []).map((candidate, index) => ({
+      content: candidate.content,
+      finishReason: candidate.finishReason,
+      index: typeof candidate.index === 'number' ? candidate.index : index,
+    }));
+
+    const normalized: GeminiResponse = {
+      candidates,
+    };
+
+    if (response.usageMetadata) {
+      const usageMetadata = this.normalizeGeminiUsageMetadata(response.usageMetadata);
+      if (!isEmpty(usageMetadata)) {
+        normalized.usageMetadata = usageMetadata;
+      }
+    }
+
+    return normalized;
+  }
+
+  private normalizeGeminiUsageMetadata(
+    usageMetadata: GeminiResponse['usageMetadata'],
+  ): NonNullable<GeminiResponse['usageMetadata']> {
+    const normalized: NonNullable<GeminiResponse['usageMetadata']> = {};
+    if (usageMetadata?.promptTokenCount !== undefined) {
+      normalized.promptTokenCount = usageMetadata.promptTokenCount;
+    }
+    if (usageMetadata?.candidatesTokenCount !== undefined) {
+      normalized.candidatesTokenCount = usageMetadata.candidatesTokenCount;
+    }
+    if (usageMetadata?.totalTokenCount !== undefined) {
+      normalized.totalTokenCount = usageMetadata.totalTokenCount;
+    }
+    if (usageMetadata?.thoughtsTokenCount !== undefined) {
+      normalized.thoughtsTokenCount = usageMetadata.thoughtsTokenCount;
+    }
+    if (usageMetadata?.promptTokensDetails !== undefined) {
+      normalized.promptTokensDetails = usageMetadata.promptTokensDetails;
+    }
+    if (usageMetadata?.candidatesTokensDetails !== undefined) {
+      normalized.candidatesTokensDetails = usageMetadata.candidatesTokensDetails;
+    }
+    if (usageMetadata?.trafficType !== undefined) {
+      normalized.trafficType = usageMetadata.trafficType;
+    }
+    return normalized;
+  }
+
+  private sendObservableSse(res: FastifyReply, stream: Observable<unknown>): void {
+    if (
+      !res.raw ||
+      typeof res.raw.writeHead !== 'function' ||
+      typeof res.raw.write !== 'function'
+    ) {
+      res.header('Content-Type', 'text/event-stream');
+      res.header('Cache-Control', 'no-cache');
+      res.header('Connection', 'keep-alive');
+      res.send(stream);
+      return;
+    }
+
+    if (this.hasHijack(res)) {
+      res.hijack();
+    }
+
+    res.raw.writeHead(HttpStatus.OK, {
+      'Content-Type': 'text/event-stream',
+      'Cache-Control': 'no-cache',
+      Connection: 'keep-alive',
+    });
+
+    const subscription = stream.subscribe({
+      next: (chunk) => {
+        if (res.raw.writableEnded) {
+          return;
+        }
+        const payload = typeof chunk === 'string' ? chunk : String(chunk ?? '');
+        res.raw.write(payload);
+      },
+      error: (error) => {
+        if (res.raw.writableEnded) {
+          return;
+        }
+        const message = error instanceof Error ? error.message : String(error);
+        res.raw.write(
+          `data: ${JSON.stringify({
+            error: {
+              message,
+              type: 'server_error',
+            },
+          })}\n\n`,
+        );
+        res.raw.end();
+      },
+      complete: () => {
+        if (!res.raw.writableEnded) {
+          res.raw.end();
+        }
+      },
+    });
+
+    res.raw.on('close', () => {
+      subscription.unsubscribe();
+    });
+  }
+
+  private hasHijack(reply: FastifyReply): reply is FastifyReply & { hijack: () => void } {
+    return typeof (reply as { hijack?: unknown }).hijack === 'function';
+  }
+
+}

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

@@ -7,13 +7,18 @@ export interface OpenAIChatRequest {
   stream?: boolean;
   size?: string;
   quality?: string;
-  extra?: Record<string, any>;
+  tools?: OpenAITool[];
+  tool_choice?: string | { type: string; function?: { name: string } };
+  response_format?: { type?: string };
+  extra?: Record<string, unknown>;
 }
 
 export interface OpenAIMessage {
   role: string;
   content: string | OpenAIContentPart[];
   name?: string;
+  tool_calls?: OpenAIToolCall[];
+  tool_call_id?: string;
 }
 
 export interface OpenAIContentPart {
@@ -24,12 +29,32 @@ export interface OpenAIContentPart {
   };
 }
 
+export interface OpenAITool {
+  type: 'function';
+  function: {
+    name: string;
+    description?: string;
+    parameters?: Record<string, unknown>;
+  };
+}
+
+export interface OpenAIToolCall {
+  id: string;
+  type: 'function';
+  function: {
+    name: string;
+    arguments: string;
+  };
+}
+
 export interface AnthropicChatRequest {
   model: string;
   messages: AnthropicMessage[];
   system?: string | AnthropicSystemBlock[];
   max_tokens?: number;
-  metadata?: any;
+  tools?: AnthropicTool[];
+  thinking?: AnthropicThinkingConfig;
+  metadata?: Record<string, unknown>;
   stop_sequences?: string[];
   stream?: boolean;
   temperature?: number;
@@ -37,6 +62,18 @@ export interface AnthropicChatRequest {
   top_k?: number;
 }
 
+export interface AnthropicTool {
+  name: string;
+  description?: string;
+  input_schema?: Record<string, unknown>;
+  type?: string;
+}
+
+export interface AnthropicThinkingConfig {
+  type: 'enabled' | string;
+  budget_tokens?: number;
+}
+
 export interface AnthropicMessage {
   role: string;
   content: string | AnthropicContent[];
@@ -50,7 +87,21 @@ export interface AnthropicSystemBlock {
 export type AnthropicContent =
   | { type: 'text'; text: string }
   | { type: 'thinking'; thinking: string; signature?: string }
-  | { type: 'image'; source: AnthropicImageSource };
+  | { type: 'image'; source: AnthropicImageSource }
+  | {
+      type: 'tool_use';
+      id: string;
+      name: string;
+      input: Record<string, unknown>;
+      signature?: string;
+    }
+  | {
+      type: 'tool_result';
+      tool_use_id: string;
+      content: string | AnthropicContent[];
+      is_error?: boolean;
+    }
+  | { type: 'redacted_thinking'; data: string };
 
 export interface AnthropicImageSource {
   type: 'base64';
@@ -102,6 +153,16 @@ export interface GeminiUsageMetadata {
   promptTokenCount?: number;
   candidatesTokenCount?: number;
   totalTokenCount?: number;
+  thoughtsTokenCount?: number;
+  promptTokensDetails?: Array<{
+    modality?: string;
+    tokenCount?: number;
+  }>;
+  candidatesTokensDetails?: Array<{
+    modality?: string;
+    tokenCount?: number;
+  }>;
+  trafficType?: string;
 }
 
 export interface OpenAIChatResponse {
@@ -121,9 +182,11 @@ export interface OpenAIChoice {
   index: number;
   message: {
     role: string;
-    content: string;
+    content: string | null;
+    tool_calls?: OpenAIToolCall[];
+    reasoning_content?: string;
   };
-  finish_reason: string;
+  finish_reason: string | null;
 }
 
 export interface AnthropicChatResponse {
@@ -133,9 +196,11 @@ export interface AnthropicChatResponse {
   model: string;
   content: AnthropicContent[];
   stop_reason: string | null;
-  stop_sequence: string | null;
+  stop_sequence?: string | null;
   usage: {
     input_tokens: number;
     output_tokens: number;
+    cache_creation_input_tokens?: number;
+    cache_read_input_tokens?: number;
   };
 }

src/server/modules/proxy/proxy.guard.ts

@@ -2,56 +2,37 @@ import {
   CanActivate,
   ExecutionContext,
   Injectable,
-  UnauthorizedException,
   Logger,
+  UnauthorizedException,
 } from '@nestjs/common';
 import { getServerConfig } from '../../server-config';
+import { extractApiKeyToken, hasConfiguredApiKey, RequestHeaders } from '../../guards/api-key-auth.util';
 
 @Injectable()
 export class ProxyGuard implements CanActivate {
   private readonly logger = new Logger(ProxyGuard.name);
 
   canActivate(context: ExecutionContext): boolean {
+    const request = context.switchToHttp().getRequest();
+
     const config = getServerConfig();
 
     // 1. Check for API Key in config
     const apiKey = config?.api_key;
 
+    const headers = request.headers as RequestHeaders;
+    const clientToken = extractApiKeyToken(headers);
+
     // 2. Bypass if no api_key set (Open Mode) or config missing
-    if (!apiKey || apiKey.trim() === '') {
+    if (!hasConfiguredApiKey(apiKey)) {
       return true;
     }
 
-    const request = context.switchToHttp().getRequest();
-    const headers = request.headers;
-
-    // 3. Extract Token from Headers
-    let clientToken: string | undefined;
-
-    // A. standard 'authorization: Bearer <token>'
-    if (headers['authorization']) {
-      const parts = headers['authorization'].split(' ');
-      if (parts.length === 2 && parts[0].toLowerCase() === 'bearer') {
-        clientToken = parts[1];
-      }
-    }
-
-    // B. 'x-api-key' (Anthropic)
-    if (!clientToken && headers['x-api-key']) {
-      clientToken = headers['x-api-key'] as string;
-    }
-
-    // C. 'x-goog-api-key' (Gemini)
-    if (!clientToken && headers['x-goog-api-key']) {
-      clientToken = headers['x-goog-api-key'] as string;
-    }
-
-    // 4. Validate
     if (clientToken === apiKey) {
       return true;
     }
+    this.logger.warn(`Rejected unauthorized request from ${request.ip}`);
 
-    this.logger.warn(`Blocked unauthorized access from ${request.ip}`);
-    throw new UnauthorizedException('Invalid API Key');
+    throw new UnauthorizedException('API key validation failed');
   }
 }

src/server/modules/proxy/proxy.module.ts

@@ -3,9 +3,13 @@ import { ProxyController } from './proxy.controller';
 import { ProxyService } from './proxy.service';
 import { TokenManagerService } from './token-manager.service';
 import { GeminiClient } from './clients/gemini.client';
+import { GeminiController } from './gemini.controller';
+import { ProxyGuard } from './proxy.guard';
 
 @Module({
-  controllers: [ProxyController],
-  providers: [ProxyService, TokenManagerService, GeminiClient],
+  imports: [],
+  controllers: [ProxyController, GeminiController],
+  providers: [ProxyService, TokenManagerService, GeminiClient, ProxyGuard],
+  exports: [TokenManagerService],
 })
 export class ProxyModule {}

src/server/modules/proxy/token-manager.service.ts

@@ -8,20 +8,26 @@ interface TokenData {
   account_id: string;
   access_token: string;
   refresh_token: string;
+  token_type: string;
   expires_in: number;
   expiry_timestamp: number;
   project_id?: string;
   session_id?: string;
+  upstream_proxy_url?: string;
 }
 
 @Injectable()
 export class TokenManagerService implements OnModuleInit {
   private readonly logger = new Logger(TokenManagerService.name);
   private currentIndex = 0;
+  private readonly stickySessionTtlMs = 10 * 60 * 1000;
+  private readonly rateLimitCooldownMs = 5 * 60 * 1000;
+  private readonly forbiddenCooldownMs = 30 * 60 * 1000;
   // In-memory cache of tokens with additional data
   private tokens: Map<string, TokenData> = new Map();
   // Cooldown map for rate-limited accounts
   private cooldowns: Map<string, number> = new Map();
+  private sessionBindings: Map<string, { accountId: string; expiresAt: number }> = new Map();
 
   async onModuleInit() {
     // Load accounts on module initialization
@@ -41,10 +47,10 @@ export class TokenManagerService implements OnModuleInit {
         }
       }
 
-      this.logger.log(`Loaded ${count} accounts`);
+      this.logger.log(`Initialized token pool with ${count} accounts`);
       return count;
     } catch (e) {
-      this.logger.error('Failed to load accounts', e);
+      this.logger.error('Unable to load accounts', e);
       return 0;
     }
   }
@@ -57,10 +63,12 @@ export class TokenManagerService implements OnModuleInit {
       email: account.email,
       access_token: account.token.access_token,
       refresh_token: account.token.refresh_token,
+      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,
     };
   }
 
@@ -72,7 +80,10 @@ export class TokenManagerService implements OnModuleInit {
     return (-(min + rand)).toString();
   }
 
-  async getNextToken(): Promise<CloudAccount | null> {
+  async getNextToken(options?: {
+    sessionKey?: string;
+    excludeAccountIds?: string[];
+  }): Promise<CloudAccount | null> {
     try {
       // Reload if empty
       if (this.tokens.size === 0) {
@@ -82,25 +93,81 @@ export class TokenManagerService implements OnModuleInit {
 
       const now = Date.now();
       const nowSeconds = Math.floor(now / 1000);
+      const sessionKey = options?.sessionKey?.trim();
+      const excludedAccountIds = new Set(options?.excludeAccountIds ?? []);
+
+      this.clearExpiredSessionBindings(now);
+
+      const fullTokenPool = Array.from(this.tokens.entries());
+      const excludedFilteredTokens = fullTokenPool.filter(
+        ([accountId]) => !excludedAccountIds.has(accountId),
+      );
+      const allTokens = excludedFilteredTokens.length > 0 ? excludedFilteredTokens : fullTokenPool;
+      if (excludedFilteredTokens.length === 0 && excludedAccountIds.size > 0) {
+        this.logger.warn(
+          'Retry exclusions removed all accounts; temporarily reusing excluded accounts',
+        );
+      }
 
       // Filter out accounts in cooldown
-      const validTokens = Array.from(this.tokens.entries()).filter(([email, _]) => {
-        const cooldownUntil = this.cooldowns.get(email);
+      const validTokens = allTokens.filter(([accountId]) => {
+        if (excludedAccountIds.has(accountId)) {
+          return false;
+        }
+
+        const cooldownUntil = this.cooldowns.get(accountId);
         return !cooldownUntil || cooldownUntil <= now;
       });
 
-      if (validTokens.length === 0) {
-        this.logger.warn('All accounts are in cooldown');
+      const candidateTokens = validTokens.length > 0 ? validTokens : allTokens;
+      if (candidateTokens.length === 0) {
+        this.logger.warn('No account available after applying exclusions');
         return null;
       }
+      if (validTokens.length === 0) {
+        this.logger.warn(
+          'All accounts are currently in cooldown; temporarily bypassing cooldown to keep service available',
+        );
+      }
+
+      if (sessionKey) {
+        const stickyBinding = this.sessionBindings.get(sessionKey);
+        if (stickyBinding && stickyBinding.expiresAt > now) {
+          const stickyMatch = candidateTokens.find(
+            ([accountId]) => accountId === stickyBinding.accountId,
+          );
+          if (stickyMatch) {
+            const [stickyAccountId, stickyTokenData] = stickyMatch;
+            return this.finalizeSelectedToken(
+              stickyAccountId,
+              stickyTokenData,
+              nowSeconds,
+              sessionKey,
+            );
+          }
+        }
+      }
 
       // Round robin selection
-      const [accountId, tokenData] = validTokens[this.currentIndex % validTokens.length];
+      const [accountId, tokenData] = candidateTokens[this.currentIndex % candidateTokens.length];
       this.currentIndex++;
+      return this.finalizeSelectedToken(accountId, tokenData, nowSeconds, sessionKey);
+    } catch (error) {
+      this.logger.error('Unable to acquire token', error);
+      return null;
+    }
+  }
 
+  private async finalizeSelectedToken(
+    accountId: string,
+    tokenData: TokenData,
+    nowSeconds: number,
+    sessionKey?: string,
+  ): Promise<CloudAccount | null> {
+    try {
       // Check if token needs refresh (expires in < 5 minutes)
       if (nowSeconds >= tokenData.expiry_timestamp - 300) {
-        this.logger.log(`Token for ${tokenData.email} expiring soon, refreshing...`);
+        this.logger.log(`Token for ${tokenData.email} is close to expiry; refreshing...`);
         try {
           const newTokens = await GoogleAPIService.refreshAccessToken(tokenData.refresh_token);
 
@@ -113,51 +180,96 @@ export class TokenManagerService implements OnModuleInit {
           await this.saveRefreshedToken(accountId, tokenData);
           this.tokens.set(accountId, tokenData);
 
-          this.logger.log(`Token refreshed for ${tokenData.email}`);
+          this.logger.log(`Refreshed token for ${tokenData.email}`);
         } catch (e) {
-          this.logger.error(`Failed to refresh token for ${tokenData.email}`, e);
+          this.logger.error(`Unable to refresh token for ${tokenData.email}`, e);
         }
       }
 
-      // Resolve project ID if missing (mock for now, like original)
-      if (!tokenData.project_id) {
-        const mockId = `cloud-code-${Math.floor(Math.random() * 100000)}`;
-        tokenData.project_id = mockId;
-        await this.saveProjectId(accountId, mockId);
+      if (
+        typeof tokenData.project_id !== 'string' ||
+        tokenData.project_id.trim() === '' ||
+        /^cloud-code-\d+$/i.test(tokenData.project_id.trim())
+      ) {
+        tokenData.project_id = undefined;
       }
 
-      this.logger.log(`Selected account: ${tokenData.email}`);
+      this.logger.log(`Using account: ${tokenData.email}`);
+      if (sessionKey) {
+        this.sessionBindings.set(sessionKey, {
+          accountId,
+          expiresAt: Date.now() + this.stickySessionTtlMs,
+        });
+      }
 
       // Return in CloudAccount format for compatibility
-      return {
+      const timestamp = Date.now();
+      const cloudAccount: CloudAccount = {
         id: accountId,
+        provider: 'google',
         email: tokenData.email,
         token: {
           access_token: tokenData.access_token,
           refresh_token: tokenData.refresh_token,
+          token_type: tokenData.token_type,
           expires_in: tokenData.expires_in,
           expiry_timestamp: tokenData.expiry_timestamp,
           project_id: tokenData.project_id,
           session_id: tokenData.session_id,
+          upstream_proxy_url: tokenData.upstream_proxy_url,
         },
-      } as CloudAccount;
+        created_at: timestamp,
+        last_used: timestamp,
+      };
+      return cloudAccount;
     } catch (error) {
-      this.logger.error('Failed to get token', error);
+      this.logger.error('Unable to finalize selected token', error);
       return null;
     }
   }
 
-  markAsRateLimited(email: string) {
-    // Cooldown for 5 minutes
-    const until = Date.now() + 5 * 60 * 1000;
-    this.cooldowns.set(email, until);
-    this.logger.warn(
-      `Account ${email} marked as rate limited until ${new Date(until).toISOString()}`,
-    );
+  markAsRateLimited(accountIdOrEmail: string) {
+    this.markInCooldown(accountIdOrEmail, 'rate limited', this.rateLimitCooldownMs);
   }
 
-  resetCooldown(email: string) {
-    this.cooldowns.delete(email);
+  markAsForbidden(accountIdOrEmail: string) {
+    this.markInCooldown(accountIdOrEmail, 'forbidden', this.forbiddenCooldownMs);
+  }
+
+  private resolveAccountId(accountIdOrEmail: string): string | null {
+    if (this.tokens.has(accountIdOrEmail)) {
+      return accountIdOrEmail;
+    }
+
+    for (const [accountId, tokenData] of this.tokens.entries()) {
+      if (tokenData.email === accountIdOrEmail) {
+        return accountId;
+      }
+    }
+
+    return null;
+  }
+
+  private clearExpiredSessionBindings(now: number): void {
+    for (const [sessionKey, binding] of this.sessionBindings.entries()) {
+      if (binding.expiresAt <= now) {
+        this.sessionBindings.delete(sessionKey);
+      }
+    }
+  }
+
+  private markInCooldown(
+    accountIdOrEmail: string,
+    reason: 'rate limited' | 'forbidden',
+    durationMs: number,
+  ): void {
+    const accountId = this.resolveAccountId(accountIdOrEmail) ?? accountIdOrEmail;
+    const cooldownUntil = Date.now() + durationMs;
+
+    this.cooldowns.set(accountId, cooldownUntil);
+    this.logger.warn(
+      `Account ${accountIdOrEmail} (cooldown key: ${accountId}) set to ${reason} until ${new Date(cooldownUntil).toISOString()}`,
+    );
   }
 
   private async saveRefreshedToken(accountId: string, tokenData: TokenData) {
@@ -173,22 +285,7 @@ export class TokenManagerService implements OnModuleInit {
         await CloudAccountRepo.updateToken(accountId, newToken);
       }
     } catch (e) {
-      this.logger.error('Failed to save refreshed token to DB', e);
-    }
-  }
-
-  private async saveProjectId(accountId: string, projectId: string) {
-    try {
-      const acc = await CloudAccountRepo.getAccount(accountId);
-      if (acc && acc.token) {
-        const newToken = {
-          ...acc.token,
-          project_id: projectId,
-        };
-        await CloudAccountRepo.updateToken(accountId, newToken);
-      }
-    } catch (e) {
-      this.logger.error('Failed to save project ID to DB', e);
+      this.logger.error('Unable to persist refreshed token to DB', e);
     }
   }