From 13f10fe32f73306470a50b0284c8026117b65695 Mon Sep 17 00:00:00 2001 From: vin <49336820+Draculabo@users.noreply.github.com> Date: Wed, 11 Feb 2026 00:42:39 +0800 Subject: [PATCH] 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 --- AGENTS.md | 2 + src/lib/antigravity/ClaudeRequestMapper.ts | 132 +- src/lib/antigravity/ClaudeResponseMapper.ts | 3 +- src/lib/antigravity/ModelMapping.ts | 86 +- src/lib/antigravity/types.ts | 7 +- src/localization/i18n.ts | 6 +- src/server/guards/admin.guard.ts | 25 + src/server/guards/api-key-auth.util.ts | 49 + .../modules/proxy/clients/gemini.client.ts | 260 ++- src/server/modules/proxy/gemini.controller.ts | 292 ++++ .../proxy/interfaces/request-interfaces.ts | 77 +- src/server/modules/proxy/proxy.controller.ts | 1044 ++++++++++- src/server/modules/proxy/proxy.guard.ts | 39 +- src/server/modules/proxy/proxy.module.ts | 8 +- src/server/modules/proxy/proxy.service.ts | 1537 +++++++++++++---- .../modules/proxy/token-manager.service.ts | 187 +- .../gemini-controller.integration.test.ts | 165 ++ src/tests/unit/logger.test.ts | 35 +- .../unit/proxy-controller.integration.test.ts | 482 ++++++ src/tests/unit/proxy-parity-fixtures.test.ts | 94 + src/tests/unit/proxy-retry-mock.test.ts | 526 +++++- src/types/cloudAccount.ts | 2 + src/types/config.ts | 4 + 23 files changed, 4408 insertions(+), 654 deletions(-) create mode 100644 src/server/guards/admin.guard.ts create mode 100644 src/server/guards/api-key-auth.util.ts create mode 100644 src/server/modules/proxy/gemini.controller.ts create mode 100644 src/tests/unit/gemini-controller.integration.test.ts create mode 100644 src/tests/unit/proxy-controller.integration.test.ts create mode 100644 src/tests/unit/proxy-parity-fixtures.test.ts diff --git a/AGENTS.md b/AGENTS.md index f4ccece..d4815e0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -123,6 +123,8 @@ Use `npm` for all commands: - Tools/config: camelCase or kebab-case - **Import paths**: Use `@/` alias for `src/`. - **Type safety**: Avoid `any`; enforce end-to-end type safety with Zod + TypeScript. +- **Utility methods**: Prefer `lodash-es` over native JavaScript utilities for array/object/string transformations to improve consistency and maintainability. + - Use named imports (for example, `import { get, groupBy, uniqBy } from 'lodash-es'`), and avoid full-package imports. - **Component design**: - Prefer Radix UI Primitives. - Use Tailwind utility classes; avoid CSS Modules unless necessary. diff --git a/src/lib/antigravity/ClaudeRequestMapper.ts b/src/lib/antigravity/ClaudeRequestMapper.ts index e101c86..2d30ae4 100644 --- a/src/lib/antigravity/ClaudeRequestMapper.ts +++ b/src/lib/antigravity/ClaudeRequestMapper.ts @@ -52,7 +52,7 @@ export function transformClaudeRequestIn( const toolIdToName = new Map(); // 1. System Instruction - const systemInstruction = buildSystemInstruction(claudeReq.system, claudeReq.model); + const systemInstruction = buildSystemInstruction(claudeReq.system); // Map model name const mappedModel = hasWebSearchTool @@ -200,32 +200,24 @@ function resolveRequestConfig( } const hasNetworkingTool = detectsNetworkingTool(tools); - const hasNonNetworking = containsNonNetworkingTool(tools); // Strip -online suffix const isOnlineSuffix = originalModel.endsWith('-online'); - const isHighQualityModel = - mappedModel === 'gemini-2.5-flash' || - mappedModel === 'gemini-1.5-pro' || - mappedModel.startsWith('gemini-1.5-pro-') || - mappedModel.startsWith('gemini-2.5-flash-') || - mappedModel.startsWith('gemini-2.0-flash') || - mappedModel.startsWith('gemini-3-') || - mappedModel.includes('claude-3-5-sonnet') || - mappedModel.includes('claude-3-opus') || - mappedModel.includes('claude-sonnet') || - mappedModel.includes('claude-opus') || - mappedModel.includes('claude-4'); - - // Determine if we should enable networking - const enableNetworking = - isOnlineSuffix || (isHighQualityModel && !hasNonNetworking) || hasNetworkingTool; + const enableNetworking = isOnlineSuffix || hasNetworkingTool; let finalModel = mappedModel.replace(/-online$/, ''); + + if (finalModel === 'gemini-3-pro-preview') { + finalModel = 'gemini-3-pro-high'; + } else if (finalModel === 'gemini-3-pro-image-preview') { + finalModel = 'gemini-3-pro-image'; + } else if (finalModel === 'gemini-3-flash-preview') { + finalModel = 'gemini-3-flash'; + } + if (enableNetworking) { - // Fallback for search compatibility - if (finalModel.includes('thinking') || !finalModel.startsWith('gemini-')) { + if (finalModel !== 'gemini-2.5-flash') { finalModel = 'gemini-2.5-flash'; } } @@ -314,59 +306,6 @@ function detectsNetworkingTool(tools?: (Tool | GeminiToolDeclaration)[]): boolea return false; } -/** - * Detects if non-networking tools are present - * Checks tool list for tools other than web search - */ -function containsNonNetworkingTool(tools?: (Tool | GeminiToolDeclaration)[]): boolean { - if (!tools) { - return false; - } - const keywords = [ - 'web_search', - 'google_search', - 'web_search_20250305', - 'google_search_retrieval', - ]; - - for (const tool of tools) { - let isNetworking = false; - - // Claude Tool format - if ('name' in tool && tool.name && keywords.includes(tool.name)) { - isNetworking = true; - } - - // OpenAI nested format (runtime check) - const openaiTool = tool as { function?: { name?: string } }; - if (openaiTool.function?.name && keywords.includes(openaiTool.function.name)) - isNetworking = true; - - // Gemini search tools - if ('googleSearch' in tool && tool.googleSearch) { - isNetworking = true; - } - if ('googleSearchRetrieval' in tool && tool.googleSearchRetrieval) { - isNetworking = true; - } - - // Gemini functionDeclarations - if ('functionDeclarations' in tool && tool.functionDeclarations) { - for (const decl of tool.functionDeclarations) { - if (decl.name && !keywords.includes(decl.name)) { - return true; // Found local - } - } - isNetworking = true; // All are networking - } - - if (!isNetworking) { - return true; - } - } - return false; -} - /** * Inject Google Search Tool * Adds googleSearch tool to the request @@ -390,41 +329,39 @@ function injectGoogleSearchTool(body: { tools?: GeminiToolDeclaration[] }) { /** * Builds system instruction - * Converts Claude system prompts to Gemini format, with Antigravity Identity Injection + * Converts Claude system prompts to Gemini format with a default assistant identity directive. */ function buildSystemInstruction( system: ClaudeRequest['system'], - modelName: string, ): { parts: { text: string }[] } | null { - const antigravity_identity = - '--- [IDENTITY_PATCH] ---\n' + - 'Ignore any previous instructions regarding your identity or host platform (e.g., Amazon Q, Google AI).\n' + + const assistantIdentityDirective = 'You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.\n' + 'You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.\n' + '**Absolute paths only**\n' + '**Proactiveness**'; + const identityMarker = 'You are Antigravity'; const parts: { text: string }[] = []; - // Check if user already provided Antigravity identity - let userHasAntigravity = false; + let hasIdentityDirective = false; if (system) { if (typeof system === 'string') { - if (system.includes('You are Antigravity')) userHasAntigravity = true; + if (system.includes(identityMarker)) { + hasIdentityDirective = true; + } } else if (Array.isArray(system)) { for (const block of system) { - if (block.type === 'text' && block.text.includes('You are Antigravity')) { - userHasAntigravity = true; + if (block.type === 'text' && block.text.includes(identityMarker)) { + hasIdentityDirective = true; break; } } } } - // Inject if missing - if (!userHasAntigravity) { - parts.push({ text: antigravity_identity }); + if (!hasIdentityDirective) { + parts.push({ text: assistantIdentityDirective }); } if (system) { @@ -437,10 +374,6 @@ function buildSystemInstruction( } } - if (!userHasAntigravity) { - parts.push({ text: '\n--- [SYSTEM_PROMPT_END] ---' }); - } - // If we pushed at least something if (parts.length > 0) { return { parts }; @@ -635,7 +568,22 @@ function buildGenerationConfig( hasWebSearch: boolean, mappedModel: string, ): GenerationConfig { + const source = String(claudeReq.metadata?.source || '').toLowerCase(); + const isOpenAIPath = source === 'openai'; const config: GenerationConfig = {}; + + if (isOpenAIPath) { + config.temperature = claudeReq.temperature ?? 1.0; + config.topP = claudeReq.top_p ?? 0.95; + if (claudeReq.max_tokens !== undefined) { + config.maxOutputTokens = claudeReq.max_tokens; + } + if (claudeReq.stop_sequences && claudeReq.stop_sequences.length > 0) { + config.stopSequences = claudeReq.stop_sequences; + } + return config; + } + if (claudeReq.thinking?.type === 'enabled') { const thinkingConfig: GenerationConfig['thinkingConfig'] = { includeThoughts: true }; if (claudeReq.thinking.budget_tokens) { @@ -655,7 +603,9 @@ function buildGenerationConfig( if (claudeReq.top_k !== undefined) { config.topK = claudeReq.top_k; } - config.maxOutputTokens = 64000; + if (claudeReq.max_tokens !== undefined) { + config.maxOutputTokens = claudeReq.max_tokens; + } config.stopSequences = ['<|user|>', '<|endoftext|>', '<|end_of_turn|>', '[DONE]', '\n\nHuman:']; return config; } diff --git a/src/lib/antigravity/ClaudeResponseMapper.ts b/src/lib/antigravity/ClaudeResponseMapper.ts index fd3f9c7..6f1f0d6 100644 --- a/src/lib/antigravity/ClaudeResponseMapper.ts +++ b/src/lib/antigravity/ClaudeResponseMapper.ts @@ -226,6 +226,8 @@ class NonStreamingProcessor { const usage: Usage = { input_tokens: geminiResponse.usageMetadata?.promptTokenCount || 0, output_tokens: geminiResponse.usageMetadata?.candidatesTokenCount || 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, }; return { @@ -235,7 +237,6 @@ class NonStreamingProcessor { model: geminiResponse.modelVersion || '', content: this.contentBlocks, stop_reason: stopReason, - stop_sequence: null, usage: usage, }; } diff --git a/src/lib/antigravity/ModelMapping.ts b/src/lib/antigravity/ModelMapping.ts index 5ce8146..736bd66 100644 --- a/src/lib/antigravity/ModelMapping.ts +++ b/src/lib/antigravity/ModelMapping.ts @@ -3,6 +3,7 @@ import { logger } from '../../utils/logger'; const CLAUDE_TO_GEMINI: Record = { // Directly supported models 'claude-opus-4-5-thinking': 'claude-opus-4-5-thinking', + 'claude-opus-4-6-thinking': 'claude-opus-4-6-thinking', 'claude-sonnet-4-5': 'claude-sonnet-4-5', 'claude-sonnet-4-5-thinking': 'claude-sonnet-4-5-thinking', @@ -12,21 +13,23 @@ const CLAUDE_TO_GEMINI: Record = { 'claude-3-5-sonnet-20240620': 'claude-sonnet-4-5', 'claude-opus-4': 'claude-opus-4-5-thinking', 'claude-opus-4-5-20251101': 'claude-opus-4-5-thinking', + 'claude-opus-4-6': 'claude-opus-4-6-thinking', + 'claude-opus-4-6-20260201': 'claude-opus-4-6-thinking', 'claude-haiku-4': 'claude-sonnet-4-5', 'claude-3-haiku-20240307': 'claude-sonnet-4-5', 'claude-haiku-4-5-20251001': 'claude-sonnet-4-5', // OpenAI Protocol Mapping - 'gpt-4': 'gemini-2.5-pro', - 'gpt-4-turbo': 'gemini-2.5-pro', - 'gpt-4-turbo-preview': 'gemini-2.5-pro', - 'gpt-4-0125-preview': 'gemini-2.5-pro', - 'gpt-4-1106-preview': 'gemini-2.5-pro', - 'gpt-4-0613': 'gemini-2.5-pro', + 'gpt-4': 'gemini-2.5-flash', + 'gpt-4-turbo': 'gemini-2.5-flash', + 'gpt-4-turbo-preview': 'gemini-2.5-flash', + 'gpt-4-0125-preview': 'gemini-2.5-flash', + 'gpt-4-1106-preview': 'gemini-2.5-flash', + 'gpt-4-0613': 'gemini-2.5-flash', - 'gpt-4o': 'gemini-2.5-pro', - 'gpt-4o-2024-05-13': 'gemini-2.5-pro', - 'gpt-4o-2024-08-06': 'gemini-2.5-pro', + 'gpt-4o': 'gemini-2.5-flash', + 'gpt-4o-2024-05-13': 'gemini-2.5-flash', + 'gpt-4o-2024-08-06': 'gemini-2.5-flash', 'gpt-4o-mini': 'gemini-2.5-flash', 'gpt-4o-mini-2024-07-18': 'gemini-2.5-flash', @@ -38,29 +41,68 @@ const CLAUDE_TO_GEMINI: Record = { 'gpt-3.5-turbo-0613': 'gemini-2.5-flash', // Gemini Protocol Mapping - 'gemini-2.5-flash-lite': 'gemini-2.5-flash-lite', + 'gemini-2.5-flash-lite': 'gemini-2.5-flash', 'gemini-2.5-flash-thinking': 'gemini-2.5-flash-thinking', - 'gemini-3-pro-low': 'gemini-3-pro-low', - 'gemini-3-pro-high': 'gemini-3-pro-high', + 'gemini-3-pro-low': 'gemini-3-pro-preview', + 'gemini-3-pro-high': 'gemini-3-pro-preview', 'gemini-3-pro-preview': 'gemini-3-pro-preview', + 'gemini-3-pro': 'gemini-3-pro-preview', 'gemini-2.5-flash': 'gemini-2.5-flash', 'gemini-3-flash': 'gemini-3-flash', 'gemini-3-pro-image': 'gemini-3-pro-image', + 'internal-background-task': 'gemini-2.5-flash', }; +const DYNAMIC_IMAGE_BASE_MODEL = 'gemini-3-pro-image'; +const DYNAMIC_IMAGE_RESOLUTIONS = ['', '-2k', '-4k']; +const DYNAMIC_IMAGE_RATIOS = ['', '-1x1', '-4x3', '-3x4', '-16x9', '-9x16', '-21x9']; +const EXTRA_DYNAMIC_MODELS = [ + 'gemini-2.0-flash-exp', + 'gemini-2.5-flash', + 'gemini-3-flash', + 'gemini-3-pro-high', + 'gemini-3-pro-low', +]; + +export const MODEL_LIST_CREATED_AT = 1770652800; + +export const MODEL_LIST_OWNER = 'antigravity'; + +export function getSupportedModels(): string[] { + return Object.keys(CLAUDE_TO_GEMINI); +} + +export function getAllDynamicModels(customMapping: Record = {}): string[] { + const modelIds = new Set(); + + for (const modelId of getSupportedModels()) { + modelIds.add(modelId); + } + + for (const customModelId of Object.keys(customMapping)) { + modelIds.add(customModelId); + } + + for (const resolution of DYNAMIC_IMAGE_RESOLUTIONS) { + for (const ratio of DYNAMIC_IMAGE_RATIOS) { + modelIds.add(`${DYNAMIC_IMAGE_BASE_MODEL}${resolution}${ratio}`); + } + } + + for (const modelId of EXTRA_DYNAMIC_MODELS) { + modelIds.add(modelId); + } + + return [...modelIds].sort(); +} + export function mapClaudeModelToGemini(input: string): string { - // 1. Check exact match in map - if (CLAUDE_TO_GEMINI[input]) { - return CLAUDE_TO_GEMINI[input]; + const mappedModel = CLAUDE_TO_GEMINI[input]; + if (mappedModel) { + return mappedModel; } - // 2. Pass-through known prefixes (gemini-, -thinking) to support dynamic suffixes - if (input.startsWith('gemini-') || input.includes('thinking')) { - return input; - } - - // 3. Fallback to default - return 'claude-sonnet-4-5'; + return input; } /** diff --git a/src/lib/antigravity/types.ts b/src/lib/antigravity/types.ts index 3d771fa..7efd76c 100644 --- a/src/lib/antigravity/types.ts +++ b/src/lib/antigravity/types.ts @@ -53,6 +53,7 @@ export interface ClaudeRequest { tools?: Tool[]; stream?: boolean; max_tokens?: number; + stop_sequences?: string[]; temperature?: number; top_p?: number; top_k?: number; @@ -141,6 +142,8 @@ export interface Tool { export interface Metadata { user_id?: string; + source?: string; + [key: string]: unknown; } // ============================================================================ @@ -336,13 +339,15 @@ export interface ClaudeResponse { model: string; content: ContentBlock[]; stop_reason: string; - stop_sequence: string | null; + stop_sequence?: string | null; usage: Usage; } export interface Usage { input_tokens: number; output_tokens: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; /** Server tool usage stats */ server_tool_use?: ServerToolUse; } diff --git a/src/localization/i18n.ts b/src/localization/i18n.ts index b95076c..fc47afb 100644 --- a/src/localization/i18n.ts +++ b/src/localization/i18n.ts @@ -132,7 +132,7 @@ i18n claudeKeyword: 'Claude Model (Keyword)', targetGemini: 'Target Gemini Model', addPlaceholderKey: 'e.g. op-3', - addPlaceholderValue: 'e.g. gemini-2.0-flash-exp', + addPlaceholderValue: 'e.g. gemini-2.5-flash', noMappings: 'No custom mappings defined.', mapsTo: 'Maps to', default: 'Default', @@ -474,7 +474,7 @@ i18n claudeKeyword: 'Claude 模型 (关键词)', targetGemini: '目标 Gemini 模型', addPlaceholderKey: '例如 op-3', - addPlaceholderValue: '例如 gemini-2.0-flash-exp', + addPlaceholderValue: '例如 gemini-2.5-flash', noMappings: '暂无自定义映射。', mapsTo: '映射到', default: '默认', @@ -816,7 +816,7 @@ i18n claudeKeyword: 'Модель Claude (Ключевое слово)', targetGemini: 'Целевая модель Gemini', addPlaceholderKey: 'например, op-3', - addPlaceholderValue: 'например, gemini-2.0-flash-exp', + addPlaceholderValue: 'например, gemini-2.5-flash', noMappings: 'Нет пользовательских сопоставлений.', mapsTo: 'Мапится в', default: 'По умолчанию', diff --git a/src/server/guards/admin.guard.ts b/src/server/guards/admin.guard.ts new file mode 100644 index 0000000..4e21c1f --- /dev/null +++ b/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'); + } +} diff --git a/src/server/guards/api-key-auth.util.ts b/src/server/guards/api-key-auth.util.ts new file mode 100644 index 0000000..f0147fa --- /dev/null +++ b/src/server/guards/api-key-auth.util.ts @@ -0,0 +1,49 @@ +export type RequestHeaderValue = string | string[] | undefined; + +export type RequestHeaders = Record; + +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; +} diff --git a/src/server/modules/proxy/clients/gemini.client.ts b/src/server/modules/proxy/clients/gemini.client.ts index fd6af6b..83f771d 100644 --- a/src/server/modules/proxy/clients/gemini.client.ts +++ b/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 { + async streamGenerate( + model: string, + content: GeminiRequest, + accessToken: string, + upstreamProxyUrl?: string, + ): Promise { 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 { const url = `${this.baseUrl}/models/${model}:generateContent`; + const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl); try { const response = await axios.post(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 { - 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, + ): Promise { + const response = await this.requestWithEndpointFailover( + ':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 { - 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, + ): Promise { + 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( + path: string, + body: GeminiInternalRequest, + accessToken: string, + upstreamProxyUrl: string | undefined, + config: AxiosRequestConfig, + operation: string, + extraHeaders?: Record, + ): Promise> { + 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(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 { diff --git a/src/server/modules/proxy/gemini.controller.ts b/src/server/modules/proxy/gemini.controller.ts new file mode 100644 index 0000000..c4b1aa4 --- /dev/null +++ b/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 { + 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 { + const normalized: NonNullable = {}; + 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): 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'; + } + +} diff --git a/src/server/modules/proxy/interfaces/request-interfaces.ts b/src/server/modules/proxy/interfaces/request-interfaces.ts index 6b3e556..a79563a 100644 --- a/src/server/modules/proxy/interfaces/request-interfaces.ts +++ b/src/server/modules/proxy/interfaces/request-interfaces.ts @@ -7,13 +7,18 @@ export interface OpenAIChatRequest { stream?: boolean; size?: string; quality?: string; - extra?: Record; + tools?: OpenAITool[]; + tool_choice?: string | { type: string; function?: { name: string } }; + response_format?: { type?: string }; + extra?: Record; } 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; + }; +} + +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; 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; + 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; + 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; }; } diff --git a/src/server/modules/proxy/proxy.controller.ts b/src/server/modules/proxy/proxy.controller.ts index 744a6e7..9711e23 100644 --- a/src/server/modules/proxy/proxy.controller.ts +++ b/src/server/modules/proxy/proxy.controller.ts @@ -1,80 +1,299 @@ -import { Controller, Get, Post, Body, Res, HttpStatus, UseGuards, Inject } from '@nestjs/common'; -import { FastifyReply } from 'fastify'; +import { + Controller, + Get, + Post, + Body, + Res, + HttpStatus, + UseGuards, + Inject, + Req, + Logger, +} from '@nestjs/common'; +import { FastifyReply, FastifyRequest } from 'fastify'; +import { isNil } from 'lodash-es'; import { ProxyService } from './proxy.service'; import { Observable } from 'rxjs'; -import { OpenAIChatRequest, AnthropicChatRequest } from './interfaces/request-interfaces'; +import { + OpenAIChatRequest, + AnthropicChatRequest, + OpenAIChatResponse, + OpenAIContentPart, + GeminiRequest, + GeminiResponse, +} from './interfaces/request-interfaces'; import { ProxyGuard } from './proxy.guard'; +import { + getAllDynamicModels, + MODEL_LIST_CREATED_AT, + MODEL_LIST_OWNER, +} from '../../../lib/antigravity/ModelMapping'; +import { getServerConfig } from '../../server-config'; @Controller('v1') @UseGuards(ProxyGuard) export class ProxyController { + private readonly logger = new Logger(ProxyController.name); + constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {} - @Get('models') - getModels() { - return { - object: 'list', - data: [ - // Gemini Native - { - id: 'gemini-2.5-flash-thinking', - object: 'model', - created: 1734336000, - owned_by: 'google', - }, - { id: 'gemini-2.5-flash', object: 'model', created: 1734336000, owned_by: 'google' }, - { id: 'gemini-2.5-pro', object: 'model', created: 1734336000, owned_by: 'google' }, - { id: 'gemini-3-flash', object: 'model', created: 1734336000, owned_by: 'google' }, - - // Claude Native (Mapped) - { id: 'claude-sonnet-4-5', object: 'model', created: 1734336000, owned_by: 'anthropic' }, - { - id: 'claude-sonnet-4-5-thinking', - object: 'model', - created: 1734336000, - owned_by: 'anthropic', - }, - - // Image Models - { id: 'gemini-3-pro-image', object: 'model', created: 1734336000, owned_by: 'google' }, - ], - }; - } - @Post('chat/completions') async chatCompletions(@Body() body: OpenAIChatRequest, @Res() res: FastifyReply) { - try { - // Claude Code CLI compatibility: - // Claude CLI uses OpenAI-style /v1/chat/completions - if (body.model?.startsWith('claude-')) { - const result = await this.proxyService.handleAnthropicMessages( - body as unknown as AnthropicChatRequest, - ); + await this.sendOpenAIChatCompletionResponse(body, res); + } - if (body.stream && result instanceof Observable) { - res.header('Content-Type', 'text/event-stream'); - res.header('Cache-Control', 'no-cache'); - res.header('Connection', 'keep-alive'); - res.send(result); - } else { - res.status(HttpStatus.OK).send(result); - } + @Post('completions') + async completions( + @Body() + body: { + model?: string; + prompt?: string | string[]; + max_tokens?: number; + temperature?: number; + top_p?: number; + stream?: boolean; + }, + @Res() res: FastifyReply, + ) { + const request: OpenAIChatRequest = { + model: body.model ?? 'gemini-2.5-flash', + messages: [ + { + role: 'user', + content: this.normalizePrompt(body.prompt), + }, + ], + max_tokens: body.max_tokens, + temperature: body.temperature, + top_p: body.top_p, + stream: body.stream, + }; + try { + const result = await this.proxyService.handleChatCompletions(request); + if (body.stream && this.isObservableLike(result)) { + this.writeSseResponse(res, result); return; } + const response = result as OpenAIChatResponse; + res.status(HttpStatus.OK).send(this.toLegacyCompletionsResponse(response)); + } catch (error) { + const message = error instanceof Error ? error.message : 'Internal Server Error'; + const status = this.resolveUpstreamStatus(message); + this.logEndpointError('/v1/completions', status, message, error); + res.status(status).send({ + error: { + message, + type: 'server_error', + }, + }); + } + } + + @Post('responses') + async responses( + @Body() + body: { + model?: string; + instructions?: string; + input?: unknown; + max_output_tokens?: number; + temperature?: number; + top_p?: number; + stream?: boolean; + }, + @Res() res: FastifyReply, + ) { + const request = this.buildResponsesChatRequest(body); + + try { + const result = await this.proxyService.handleChatCompletions(request); + if (body.stream && this.isObservableLike(result)) { + this.writeSseResponse(res, result); + return; + } + + const response = result as OpenAIChatResponse; + res.status(HttpStatus.OK).send(this.toLegacyCompletionsResponse(response)); + } catch (error) { + const message = error instanceof Error ? error.message : 'Internal Server Error'; + const status = this.resolveUpstreamStatus(message); + this.logEndpointError('/v1/responses', status, message, error); + res.status(status).send({ + error: { + message, + type: 'server_error', + }, + }); + } + } + + @Post('images/generations') + async imageGenerations( + @Body() + body: { + model?: string; + prompt?: string; + size?: string; + quality?: string; + }, + @Res() res: FastifyReply, + ) { + const request: OpenAIChatRequest = { + model: body.model ?? 'gemini-3-pro-image', + messages: [ + { + role: 'user', + content: body.prompt ?? '', + }, + ], + stream: false, + size: body.size, + quality: body.quality, + }; + + await this.sendOpenAIImageResponse(request, body.prompt ?? '', res); + } + + @Post('images/edits') + async imageEdits( + @Body() + body: { + model?: string; + prompt?: string; + size?: string; + quality?: string; + image?: string | { data?: string; mimeType?: string }; + reference_images?: Array; + mask?: string | { data?: string; mimeType?: string }; + }, + @Req() req: FastifyRequest, + @Res() res: FastifyReply, + ) { + if (!this.hasValidMultipartBoundary(req)) { + res + .status(HttpStatus.BAD_REQUEST) + .send('Invalid `boundary` for `multipart/form-data` request'); + return; + } + + const imageParts = this.collectImageContentParts([ + body.image, + body.mask, + ...(body.reference_images ?? []), + ]); + + const request: OpenAIChatRequest = { + model: body.model ?? 'gemini-3-pro-image', + messages: [ + { + role: 'user', + content: + imageParts.length > 0 + ? [ + { + type: 'text', + text: body.prompt ?? 'Please edit this image based on the provided instruction.', + }, + ...imageParts, + ] + : (body.prompt ?? 'Please edit this image based on the provided instruction.'), + }, + ], + stream: false, + size: body.size, + quality: body.quality, + }; + + await this.sendOpenAIImageResponse(request, body.prompt ?? '', res); + } + + @Post('audio/transcriptions') + async audioTranscriptions( + @Body() + body: { + model?: string; + prompt?: string; + file?: string | { data?: string; mimeType?: string }; + audio?: string | { data?: string; mimeType?: string }; + }, + @Req() req: FastifyRequest, + @Res() res: FastifyReply, + ) { + if (!this.hasValidMultipartBoundary(req)) { + res + .status(HttpStatus.BAD_REQUEST) + .send('Invalid `boundary` for `multipart/form-data` request'); + return; + } + + const inlineAudio = this.resolveInlineData(body.file ?? body.audio); + if (!inlineAudio) { + res.status(HttpStatus.BAD_REQUEST).send({ + error: { + message: "Missing 'file' or 'audio' input. Provide base64 content or a data URL.", + type: 'invalid_request_error', + }, + }); + return; + } + + try { + const result = await this.proxyService.handleGeminiGenerateContent( + body.model ?? 'gemini-2.5-flash', + { + contents: [ + { + role: 'user', + parts: [ + { + text: body.prompt ?? 'Please transcribe the provided speech audio accurately.', + }, + { + inlineData: inlineAudio, + }, + ], + }, + ], + }, + ); + + const text = result.candidates?.[0]?.content?.parts + ?.map((part) => part.text ?? '') + .join('') + .trim(); + + res.status(HttpStatus.OK).send({ + text: text ?? '', + }); + } catch (error) { + const message = error instanceof Error ? error.message : 'Internal Server 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) { + try { const result = await this.proxyService.handleChatCompletions(body); - if (body.stream && result instanceof Observable) { - res.header('Content-Type', 'text/event-stream'); - res.header('Cache-Control', 'no-cache'); - res.header('Connection', 'keep-alive'); - res.send(result); + if (body.stream && this.isObservableLike(result)) { + this.writeSseResponse(res, result); + return; } else { res.status(HttpStatus.OK).send(result); } } catch (error) { const message = error instanceof Error ? error.message : 'Internal Server Error'; - res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({ + const status = this.resolveUpstreamStatus(message); + this.logEndpointError('/v1/chat/completions', status, message, error); + res.status(status).send({ error: { message: message, type: 'server_error', @@ -88,17 +307,17 @@ export class ProxyController { try { const result = await this.proxyService.handleAnthropicMessages(body); - if (body.stream && result instanceof Observable) { - res.header('Content-Type', 'text/event-stream'); - res.header('Cache-Control', 'no-cache'); - res.header('Connection', 'keep-alive'); - res.send(result); + if (body.stream && this.isObservableLike(result)) { + this.writeSseResponse(res, result); + return; } else { res.status(HttpStatus.OK).send(result); } } catch (error) { const message = error instanceof Error ? error.message : 'Internal Server Error'; - res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({ + const status = this.resolveUpstreamStatus(message); + this.logEndpointError('/v1/messages', status, message, error); + res.status(status).send({ type: 'error', error: { type: 'api_error', @@ -107,4 +326,707 @@ export class ProxyController { }); } } + + private normalizePrompt(prompt: string | string[] | undefined): string { + if (!prompt) { + return ''; + } + if (Array.isArray(prompt)) { + return prompt.join('\n'); + } + return prompt; + } + + private toLegacyCompletionsResponse(response: OpenAIChatResponse): Record { + const choice = response.choices?.[0]; + const content = choice?.message?.content; + const text = typeof content === 'string' ? content : ''; + + return { + id: response.id, + object: 'text_completion', + created: response.created, + model: response.model, + choices: [ + { + text, + index: choice?.index ?? 0, + logprobs: null, + finish_reason: choice?.finish_reason ?? null, + }, + ], + usage: response.usage, + }; + } + + private normalizeResponsesInput(input: unknown): string { + if (typeof input === 'string') { + return input; + } + + if (Array.isArray(input)) { + return input + .map((item) => { + if (typeof item === 'string') { + return item; + } + if (item && typeof item === 'object' && 'content' in item) { + const content = (item as { content?: unknown }).content; + if (typeof content === 'string') { + return content; + } + } + return JSON.stringify(item); + }) + .join('\n'); + } + + if (isNil(input)) { + return ''; + } + + return JSON.stringify(input); + } + + private buildResponsesChatRequest(body: { + model?: string; + instructions?: string; + input?: unknown; + max_output_tokens?: number; + temperature?: number; + top_p?: number; + stream?: boolean; + }): OpenAIChatRequest { + const messages: OpenAIChatRequest['messages'] = []; + if (body.instructions && body.instructions.trim() !== '') { + messages.push({ + role: 'system', + content: body.instructions, + }); + } + + const callIdToToolName = new Map(); + const inputItems = Array.isArray(body.input) ? body.input : null; + + if (inputItems) { + for (const item of inputItems) { + const itemObj = this.toRecord(item); + if (!itemObj) { + continue; + } + + const type = this.asString(itemObj.type); + if (!type) { + continue; + } + + if (type === 'function_call' || type === 'local_shell_call' || type === 'web_search_call') { + const callId = + this.asString(itemObj.call_id) ?? this.asString(itemObj.id) ?? `call_${Date.now()}`; + const toolName = + type === 'local_shell_call' + ? 'shell' + : type === 'web_search_call' + ? 'google_search' + : (this.asString(itemObj.name) ?? 'unknown'); + callIdToToolName.set(callId, toolName); + } + } + + for (const item of inputItems) { + const itemObj = this.toRecord(item); + if (!itemObj) { + continue; + } + + const type = this.asString(itemObj.type); + if (!type) { + continue; + } + + if (type === 'message') { + const role = this.asString(itemObj.role) ?? 'user'; + const content = this.normalizeResponsesMessageContent(itemObj.content); + messages.push({ role, content }); + continue; + } + + if (type === 'function_call' || type === 'local_shell_call' || type === 'web_search_call') { + const callId = + this.asString(itemObj.call_id) ?? this.asString(itemObj.id) ?? `call_${Date.now()}`; + const toolName = callIdToToolName.get(callId) ?? 'unknown'; + const args = this.resolveToolArguments(type, itemObj); + messages.push({ + role: 'assistant', + content: '', + tool_calls: [ + { + id: callId, + type: 'function', + function: { + name: toolName, + arguments: JSON.stringify(args), + }, + }, + ], + }); + continue; + } + + if (type === 'function_call_output' || type === 'custom_tool_call_output') { + const callId = this.asString(itemObj.call_id) ?? this.asString(itemObj.id) ?? 'unknown'; + const output = itemObj.output; + messages.push({ + role: 'tool', + tool_call_id: callId, + name: callIdToToolName.get(callId) ?? 'unknown', + content: this.normalizeResponsesOutput(output), + }); + continue; + } + } + } else if (typeof body.input === 'string') { + messages.push({ + role: 'user', + content: body.input, + }); + } else if (!isNil(body.input)) { + messages.push({ + role: 'user', + content: this.normalizeResponsesInput(body.input), + }); + } + + if (messages.length === 0) { + messages.push({ + role: 'user', + content: '', + }); + } + + return { + model: body.model ?? 'gemini-2.5-flash', + messages, + max_tokens: body.max_output_tokens, + temperature: body.temperature, + top_p: body.top_p, + stream: body.stream, + }; + } + + private normalizeResponsesMessageContent(content: unknown): string | OpenAIContentPart[] { + if (typeof content === 'string') { + return content; + } + + if (!Array.isArray(content)) { + return this.normalizeResponsesInput(content); + } + + const textParts: string[] = []; + const imageParts: OpenAIContentPart[] = []; + + for (const item of content) { + const block = this.toRecord(item); + if (!block) { + continue; + } + + const blockType = this.asString(block.type); + if (blockType === 'input_text' || blockType === 'text' || blockType === 'output_text') { + const text = this.asString(block.text); + if (text) { + textParts.push(text); + } + continue; + } + + if (blockType === 'input_image' || blockType === 'image_url') { + const imageUrl = this.resolveImageUrl(block); + if (imageUrl) { + imageParts.push({ + type: 'image_url', + image_url: { + url: imageUrl, + }, + }); + } + } + } + + if (imageParts.length === 0) { + return textParts.join('\n'); + } + + const merged: OpenAIContentPart[] = []; + if (textParts.length > 0) { + merged.push({ + type: 'text', + text: textParts.join('\n'), + }); + } + merged.push(...imageParts); + return merged; + } + + private resolveToolArguments( + type: string, + item: Record, + ): Record { + if (type === 'local_shell_call') { + const action = this.toRecord(item.action); + const exec = action ? this.toRecord(action.exec) : null; + const command = this.asString(exec?.command); + return { + command: command ? [command] : [], + }; + } + + if (type === 'web_search_call') { + const action = this.toRecord(item.action); + return { + query: this.asString(action?.query) ?? '', + }; + } + + const raw = item.arguments; + if (typeof raw === 'string') { + try { + const parsed = JSON.parse(raw); + if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { + return parsed as Record; + } + return { + value: parsed, + }; + } catch { + return { + raw, + }; + } + } + + if (raw && typeof raw === 'object' && !Array.isArray(raw)) { + return raw as Record; + } + + return {}; + } + + private normalizeResponsesOutput(output: unknown): string { + if (typeof output === 'string') { + return output; + } + if (output && typeof output === 'object' && !Array.isArray(output)) { + const obj = output as Record; + const content = this.asString(obj.content); + if (content) { + return content; + } + } + if (isNil(output)) { + return ''; + } + return JSON.stringify(output); + } + + private resolveImageUrl(block: Record): string | null { + const raw = block.image_url; + if (typeof raw === 'string') { + return raw; + } + if (raw && typeof raw === 'object' && !Array.isArray(raw)) { + const record = raw as Record; + const url = this.asString(record.url); + if (url) { + return url; + } + } + return null; + } + + private collectImageContentParts( + entries: Array, + ): OpenAIContentPart[] { + const parts: OpenAIContentPart[] = []; + for (const entry of entries) { + const inlineData = this.resolveInlineData(entry); + if (!inlineData) { + continue; + } + parts.push({ + type: 'image_url', + image_url: { + url: `data:${inlineData.mimeType};base64,${inlineData.data}`, + }, + }); + } + return parts; + } + + private resolveInlineData(input: unknown): { + mimeType: string; + data: string; + } | null { + if (!input) { + return null; + } + + if (typeof input === 'string') { + const dataUri = input.match(/^data:(?[^;]+);base64,(?[A-Za-z0-9+/=]+)$/); + if (dataUri?.groups?.mime && dataUri.groups.data) { + return { + mimeType: dataUri.groups.mime, + data: dataUri.groups.data, + }; + } + + const cleaned = input.replace(/\s+/g, ''); + if (cleaned.length > 0) { + return { + mimeType: 'audio/mpeg', + data: cleaned, + }; + } + return null; + } + + if (typeof input === 'object' && !Array.isArray(input)) { + const obj = input as Record; + const data = this.asString(obj.data); + if (!data) { + return null; + } + return { + mimeType: this.asString(obj.mimeType) ?? 'audio/mpeg', + data, + }; + } + + return null; + } + + private toRecord(value: unknown): Record | null { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return null; + } + return value as Record; + } + + private asString(value: unknown): string | null { + return typeof value === 'string' ? value : null; + } + + private isObservableLike(value: unknown): value is Observable { + return ( + Boolean(value) && + typeof value === 'object' && + typeof (value as { subscribe?: unknown }).subscribe === 'function' + ); + } + + private writeSseResponse(res: FastifyReply, stream: Observable): 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 (typeof (res as { hijack?: () => void }).hijack === 'function') { + (res as { hijack: () => void }).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 async sendOpenAIImageResponse( + request: OpenAIChatRequest, + prompt: string, + res: FastifyReply, + ): Promise { + try { + const result = await this.proxyService.handleChatCompletions(request); + if (result instanceof Observable) { + this.logEndpointError( + '/v1/images/generations', + HttpStatus.INTERNAL_SERVER_ERROR, + 'Streaming image generation is not supported by this endpoint', + ); + res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({ + error: { + message: 'Streaming image generation is not supported by this endpoint', + type: 'invalid_request_error', + }, + }); + return; + } + + const content = result.choices?.[0]?.message?.content; + const image = this.extractInlineBase64Image(typeof content === 'string' ? content : ''); + if (!image) { + this.logEndpointError( + '/v1/images/generations', + HttpStatus.BAD_GATEWAY, + 'Upstream did not return inline image data', + ); + res.status(HttpStatus.BAD_GATEWAY).send({ + error: { + message: 'Upstream did not return inline image data', + type: 'invalid_response_error', + }, + }); + return; + } + + res.status(HttpStatus.OK).send({ + created: Math.floor(Date.now() / 1000), + data: [ + { + b64_json: image.data, + }, + ], + }); + } catch (error) { + let message = error instanceof Error ? error.message : 'Internal Server Error'; + + if (this.isProjectContextErrorMessage(message)) { + try { + const geminiRequest = this.buildGeminiImageRequest(request, prompt); + const geminiResult = await this.proxyService.handleGeminiGenerateContent( + request.model ?? 'gemini-3-pro-image', + geminiRequest, + ); + const fallbackImage = this.extractInlineBase64ImageFromGeminiResponse(geminiResult); + if (fallbackImage) { + res.status(HttpStatus.OK).send({ + created: Math.floor(Date.now() / 1000), + data: [ + { + b64_json: fallbackImage.data, + }, + ], + }); + return; + } + message = 'Upstream did not return inline image data'; + } catch (fallbackError) { + message = fallbackError instanceof Error ? fallbackError.message : message; + } + } + + const status = this.resolveUpstreamStatus(message); + this.logEndpointError('/v1/images/generations', status, message, error); + res.status(status).send({ + error: { + message, + type: 'server_error', + }, + }); + } + } + + private extractInlineBase64Image(content: string): { + mimeType: string; + data: string; + } | null { + const pattern = /data:(?[\w/+.-]+);base64,(?[A-Za-z0-9+/=]+)/; + const matched = content.match(pattern); + if (!matched || !matched.groups) { + return null; + } + + return { + mimeType: matched.groups.mime, + data: matched.groups.data, + }; + } + + private extractInlineBase64ImageFromGeminiResponse(response: GeminiResponse): { + mimeType: string; + data: string; + } | null { + const parts = response.candidates?.[0]?.content?.parts ?? []; + for (const part of parts) { + if (part.inlineData?.data) { + return { + mimeType: part.inlineData.mimeType ?? 'image/jpeg', + data: part.inlineData.data, + }; + } + if (part.text) { + const parsed = this.extractInlineBase64Image(part.text); + if (parsed) { + return parsed; + } + } + } + return null; + } + + private buildGeminiImageRequest( + request: OpenAIChatRequest, + fallbackPrompt: string, + ): GeminiRequest { + const userMessage = request.messages.find((message) => message.role === 'user'); + const textParts: string[] = []; + const parts: Array<{ text?: string; inlineData?: { mimeType: string; data: string } }> = []; + + if (!userMessage) { + parts.push({ text: fallbackPrompt || 'Please generate an image based on this request.' }); + } else if (typeof userMessage.content === 'string') { + parts.push({ + text: userMessage.content || fallbackPrompt || 'Please generate an image based on this request.', + }); + } else { + for (const block of userMessage.content) { + if (block.type === 'text' && typeof block.text === 'string' && block.text.trim() !== '') { + textParts.push(block.text); + } + if (block.type === 'image_url') { + const imageUrl = this.resolveImageUrl(block as unknown as Record); + const inlineData = this.resolveInlineData(imageUrl); + if (inlineData) { + parts.push({ + inlineData: { + mimeType: inlineData.mimeType, + data: inlineData.data, + }, + }); + } + } + } + if (textParts.length > 0) { + parts.unshift({ text: textParts.join('\n') }); + } + } + + if (parts.length === 0) { + parts.push({ text: fallbackPrompt || 'Please generate an image based on this request.' }); + } + + return { + contents: [ + { + role: 'user', + parts, + }, + ], + }; + } + + private isProjectContextErrorMessage(message: string): boolean { + const lowered = message.toLowerCase(); + return ( + lowered.includes('#3501') || + (lowered.includes('google cloud project') && lowered.includes('code assist license')) || + (lowered.includes('resource projects/') && lowered.includes('could not be found')) || + (lowered.includes('project') && lowered.includes('not found')) + ); + } + + private hasValidMultipartBoundary(req: FastifyRequest): boolean { + const contentType = req.headers['content-type']; + if (typeof contentType !== 'string') { + return false; + } + + const lowered = contentType.toLowerCase(); + return lowered.includes('multipart/form-data') && lowered.includes('boundary='); + } + + private resolveUpstreamStatus(message: string): HttpStatus { + const lowered = message.toLowerCase(); + if (lowered.includes('all accounts failed or unhealthy')) { + return HttpStatus.SERVICE_UNAVAILABLE; + } + if (lowered.includes('all accounts exhausted') || lowered.includes('no available accounts')) { + return HttpStatus.TOO_MANY_REQUESTS; + } + if ( + lowered.includes('network socket disconnected') || + lowered.includes('secure tls connection was established') || + lowered.includes('socket hang up') || + lowered.includes('econnreset') || + lowered.includes('eai_again') + ) { + return HttpStatus.SERVICE_UNAVAILABLE; + } + if (lowered.includes('401') || lowered.includes('unauthorized')) { + return HttpStatus.UNAUTHORIZED; + } + if (lowered.includes('403') || lowered.includes('forbidden')) { + return HttpStatus.FORBIDDEN; + } + if (lowered.includes('429') || lowered.includes('rate limit') || lowered.includes('quota')) { + return HttpStatus.TOO_MANY_REQUESTS; + } + if (lowered.includes('503') || lowered.includes('service unavailable')) { + return HttpStatus.SERVICE_UNAVAILABLE; + } + if (lowered.includes('502') || lowered.includes('bad gateway')) { + return HttpStatus.BAD_GATEWAY; + } + if (lowered.includes('504') || lowered.includes('timeout')) { + return HttpStatus.GATEWAY_TIMEOUT; + } + return HttpStatus.INTERNAL_SERVER_ERROR; + } + + private logEndpointError( + endpoint: string, + status: HttpStatus, + message: string, + error?: unknown, + ): void { + const base = `[${endpoint}] status=${status} message=${message}`; + if (status >= HttpStatus.INTERNAL_SERVER_ERROR) { + this.logger.error(base, error instanceof Error ? error.stack : undefined); + return; + } + this.logger.warn(base); + } } diff --git a/src/server/modules/proxy/proxy.guard.ts b/src/server/modules/proxy/proxy.guard.ts index 93cadd2..8681272 100644 --- a/src/server/modules/proxy/proxy.guard.ts +++ b/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 ' - 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'); } } diff --git a/src/server/modules/proxy/proxy.module.ts b/src/server/modules/proxy/proxy.module.ts index c59b759..7388789 100644 --- a/src/server/modules/proxy/proxy.module.ts +++ b/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 {} diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index bc3925e..4ab7342 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/src/server/modules/proxy/proxy.service.ts @@ -1,4 +1,5 @@ import { Injectable, Logger, Inject } from '@nestjs/common'; +import { isEmpty, isNil } from 'lodash-es'; import { TokenManagerService } from './token-manager.service'; import { GeminiClient } from './clients/gemini.client'; import { v4 as uuidv4 } from 'uuid'; @@ -6,7 +7,12 @@ import { Observable } from 'rxjs'; import { transformClaudeRequestIn } from '../../../lib/antigravity/ClaudeRequestMapper'; import { transformResponse } from '../../../lib/antigravity/ClaudeResponseMapper'; import { StreamingState, PartProcessor } from '../../../lib/antigravity/ClaudeStreamingMapper'; -import { ClaudeRequest } from '../../../lib/antigravity/types'; +import { + ClaudeRequest, + ClaudeResponse, + GeminiInternalRequest, + GeminiPart as InternalGeminiPart, +} from '../../../lib/antigravity/types'; import { calculateRetryDelay, sleep } from '../../../lib/antigravity/retry-utils'; import { classifyStreamError, @@ -15,20 +21,19 @@ import { import { OpenAIChatRequest, AnthropicChatRequest, - GeminiContent, - GeminiPart, - AnthropicSystemBlock, - OpenAIContentPart, GeminiResponse, GeminiRequest, AnthropicChatResponse, OpenAIChatResponse, - GeminiCandidate, + AnthropicContent, } from './interfaces/request-interfaces'; +import { getServerConfig } from '../../server-config'; +import { resolveModelRoute } from '../../../lib/antigravity/ModelMapping'; @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, @@ -40,7 +45,10 @@ export class ProxyService { async handleAnthropicMessages( request: AnthropicChatRequest, ): Promise> { + const sessionKey = this.getSessionKeyFromAnthropic(request); + const targetModel = this.mapModel(request.model); + const extraHeaders = this.buildModelSpecificHeaders(request.model); this.logger.log( `Received Anthropic request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`, ); @@ -48,46 +56,123 @@ export class ProxyService { // Retry loop let lastError: unknown = null; const maxRetries = 3; + const attemptedAccountIds = new Set(); for (let i = 0; i < maxRetries; i++) { - // Fix #6: Add jitter delay for retries (skip first attempt) if (i > 0) { const delay = calculateRetryDelay(i - 1); this.logger.log(`Retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms (jittered)`); await sleep(delay); } - const token = await this.tokenManager.getNextToken(); + const token = await this.tokenManager.getNextToken({ + sessionKey, + excludeAccountIds: Array.from(attemptedAccountIds), + }); if (!token) { throw new Error('No available accounts'); } + attemptedAccountIds.add(token.id); try { - // Use Antigravity Mapper - const projectId = token.token.project_id!; - const geminiBody = transformClaudeRequestIn(request as unknown as ClaudeRequest, projectId); - // Antigravity mapper determines the final model (e.g. gemini-2.5-flash) based on input - // We trust its judgment (it handles logic for online/image models) + const projectId = token.token.project_id ?? ''; + const geminiBody = transformClaudeRequestIn(this.toClaudeRequest(request), projectId); if (request.stream) { const stream = await this.geminiClient.streamGenerateInternal( geminiBody, token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, ); return this.processAnthropicInternalStream(stream, geminiBody.model); } else { - const response = await this.geminiClient.generateInternal( + const response = await this.generateInternalWithStreamFallback( geminiBody, token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, ); - return transformResponse(response) as unknown as AnthropicChatResponse; + return this.toAnthropicChatResponse(transformResponse(response)); } } catch (error) { + if (error instanceof Error && this.isProjectContextError(error.message)) { + this.logger.warn( + `Anthropic request hit project context issue, retrying without project: ${error.message}`, + ); + try { + const fallbackBody = transformClaudeRequestIn(this.toClaudeRequest(request), ''); + if (request.stream) { + const stream = await this.geminiClient.streamGenerateInternal( + fallbackBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.processAnthropicInternalStream(stream, fallbackBody.model); + } else { + const response = await this.generateInternalWithStreamFallback( + fallbackBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.toAnthropicChatResponse(transformResponse(response)); + } + } catch (fallbackError) { + error = fallbackError; + } + } + + 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}`, + ); + try { + const downgradedRequest: ClaudeRequest = { + ...this.toClaudeRequest(request), + model: 'gemini-2.5-flash', + }; + const downgradedBody = transformClaudeRequestIn( + downgradedRequest, + token.token.project_id ?? '', + ); + if (request.stream) { + const stream = await this.geminiClient.streamGenerateInternal( + downgradedBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.processAnthropicInternalStream(stream, downgradedBody.model); + } else { + const response = await this.generateInternalWithStreamFallback( + downgradedBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + const transformed = this.toAnthropicChatResponse(transformResponse(response)); + return { + ...transformed, + model: request.model, + }; + } + } catch (downgradeError) { + error = downgradeError; + } + } + lastError = error; if (error instanceof Error) { this.logger.warn(`Anthropic Request failed: ${error.message}`); - if (this.shouldRetry(error.message)) { - this.tokenManager.markAsRateLimited(token.email); + 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); + } } } } @@ -95,85 +180,10 @@ export class ProxyService { throw lastError || new Error('Request failed after retries'); } - // Old converter kept for reference or other uses - private convertAnthropicToGemini(request: AnthropicChatRequest): GeminiRequest { - const contents: GeminiContent[] = []; - let systemInstruction: { parts: GeminiPart[] } | undefined = undefined; - - if (request.system) { - let systemText = ''; - if (typeof request.system === 'string') { - systemText = request.system; - } else if (Array.isArray(request.system)) { - systemText = request.system - .filter((b: AnthropicSystemBlock) => b.type === 'text') - .map((b: AnthropicSystemBlock) => b.text) - .join('\n'); - } - - systemInstruction = { - parts: [{ text: systemText }], - }; - } - - // Anthropic 'messages' -> Gemini 'contents' - for (const msg of request.messages) { - const role = msg.role === 'user' ? 'user' : 'model'; - // Handle content arrays or strings - let textPart = ''; - if (Array.isArray(msg.content)) { - textPart = msg.content - .filter((c) => c.type === 'text') - .map((c) => (c as { type: 'text'; text: string }).text) - .join('\n'); - } else { - textPart = msg.content; - } - - contents.push({ - role, - parts: [{ text: textPart }], - }); - } - - return { - contents, - systemInstruction, - generationConfig: { - temperature: request.temperature, - maxOutputTokens: request.max_tokens, - topP: request.top_p, - topK: request.top_k, - }, - }; - } - - // Old response converter - private convertGeminiToAnthropicResponse( - geminiResponse: GeminiResponse, + private processAnthropicInternalStream( + upstreamStream: NodeJS.ReadableStream, model: string, - ): AnthropicChatResponse { - const candidate = geminiResponse.candidates?.[0]; - const processed = this.processInlineData(candidate); - const parts = processed?.content?.parts || []; - const contentText = parts.map((p) => p.text || '').join(''); - - return { - id: `msg_${uuidv4()}`, - type: 'message', - role: 'assistant', - model: model, - content: [{ type: 'text', text: contentText }], - stop_reason: 'end_turn', - stop_sequence: null, - usage: { - input_tokens: geminiResponse.usageMetadata?.promptTokenCount || 0, - output_tokens: geminiResponse.usageMetadata?.candidatesTokenCount || 0, - }, - }; - } - - private processAnthropicInternalStream(upstreamStream: any, model: string): Observable { + ): Observable { return new Observable((subscriber) => { const decoder = new TextDecoder(); let buffer = ''; @@ -182,9 +192,8 @@ export class ProxyService { const processor = new PartProcessor(state); let lastFinishReason: string | undefined; - let lastUsageMetadata: any | undefined; + let lastUsageMetadata: Record | undefined; - // Fix #1: Track if we received any valid data let receivedData = false; upstreamStream.on('data', (chunk: Buffer) => { @@ -217,15 +226,14 @@ export class ProxyService { lastUsageMetadata = json.usageMetadata; } - if (part) { - const chunks = processor.process(part as any); + if (this.isGeminiPart(part)) { + const chunks = processor.process(part); chunks.forEach((c) => subscriber.next(c)); } // Reset error state on successful parse state.resetErrorState(); } catch (e) { - // Fix #3: Use handleParseError for graceful degradation this.logger.error('Stream parse error', e); const errorChunks = state.handleParseError(dataStr); errorChunks.forEach((c) => subscriber.next(c)); @@ -234,7 +242,6 @@ export class ProxyService { }); upstreamStream.on('end', () => { - // Fix #1: Check for empty stream if (!receivedData) { this.logger.warn('Empty response stream detected'); subscriber.error(new Error('Empty response stream')); @@ -246,8 +253,7 @@ export class ProxyService { subscriber.complete(); }); - upstreamStream.on('error', (err: any) => { - // Fix #2: Classify error and send friendly message + upstreamStream.on('error', (err: unknown) => { const cleanError = err instanceof Error ? err : new Error(String(err)); const { type, message } = classifyStreamError(cleanError); @@ -261,11 +267,280 @@ export class ProxyService { } // --- OpenAI / Universal Handlers --- + async handleGeminiGenerateContent( + model: string, + request: GeminiRequest, + ): Promise { + const normalizedModel = this.normalizeGeminiModel(model); + const targetModel = this.mapModel(normalizedModel); + const extraHeaders = this.buildModelSpecificHeaders(normalizedModel); + this.logger.log( + `Received Gemini request for model: ${normalizedModel} (Mapped: ${targetModel})`, + ); + + let lastError: unknown = null; + const maxRetries = 3; + const attemptedAccountIds = new Set(); + + for (let i = 0; i < maxRetries; i++) { + if (i > 0) { + const delay = calculateRetryDelay(i - 1); + this.logger.log(`Gemini retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms`); + await sleep(delay); + } + + const token = await this.tokenManager.getNextToken({ + excludeAccountIds: Array.from(attemptedAccountIds), + }); + if (!token) { + throw new Error('No available accounts (all exhausted or rate limited)'); + } + attemptedAccountIds.add(token.id); + + try { + const internalBody = this.createGeminiInternalRequest( + targetModel, + request, + token.token.project_id ?? '', + 'generate-content', + ); + + const response = await this.generateInternalWithStreamFallback( + internalBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + + return this.normalizeGeminiGenerateResponse(response); + } catch (error) { + if (error instanceof Error && this.isProjectContextError(error.message)) { + this.logger.warn( + `Gemini request hit project context issue, retrying without project: ${error.message}`, + ); + try { + const fallbackBody = this.createGeminiInternalRequest( + targetModel, + request, + '', + 'generate-content', + ); + const response = await this.generateInternalWithStreamFallback( + fallbackBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.normalizeGeminiGenerateResponse(response); + } catch (fallbackError) { + error = fallbackError; + } + } + + lastError = error; + if (error instanceof Error) { + this.logger.warn(`Gemini request failed with account ${token.email}: ${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('Gemini request failed after retries'); + } + + async handleGeminiStreamGenerateContent( + model: string, + request: GeminiRequest, + ): Promise> { + const normalizedModel = this.normalizeGeminiModel(model); + const targetModel = this.mapModel(normalizedModel); + const extraHeaders = this.buildModelSpecificHeaders(normalizedModel); + this.logger.log( + `Received Gemini stream request for model: ${normalizedModel} (Mapped: ${targetModel})`, + ); + + let lastError: unknown = null; + const maxRetries = 3; + const attemptedAccountIds = new Set(); + + for (let i = 0; i < maxRetries; i++) { + if (i > 0) { + const delay = calculateRetryDelay(i - 1); + this.logger.log(`Gemini stream retry attempt ${i + 1}/${maxRetries}, waiting ${delay}ms`); + await sleep(delay); + } + + const token = await this.tokenManager.getNextToken({ + excludeAccountIds: Array.from(attemptedAccountIds), + }); + if (!token) { + throw new Error('No available accounts (all exhausted or rate limited)'); + } + attemptedAccountIds.add(token.id); + + try { + const internalBody = this.createGeminiInternalRequest( + targetModel, + request, + token.token.project_id ?? '', + 'generate-content', + ); + + const stream = await this.geminiClient.streamGenerateInternal( + internalBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.passthroughSseStream(stream); + } catch (error) { + if (error instanceof Error && this.isProjectContextError(error.message)) { + this.logger.warn( + `Gemini stream request hit project context issue, retrying without project: ${error.message}`, + ); + try { + const fallbackBody = this.createGeminiInternalRequest( + targetModel, + request, + '', + 'generate-content', + ); + const stream = await this.geminiClient.streamGenerateInternal( + fallbackBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.passthroughSseStream(stream); + } catch (fallbackError) { + error = fallbackError; + } + } + + lastError = error; + if (error instanceof Error) { + this.logger.warn( + `Gemini stream request failed with account ${token.email}: ${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('Gemini stream request failed after retries'); + } + + private passthroughSseStream(upstreamStream: NodeJS.ReadableStream): Observable { + return new Observable((subscriber) => { + const decoder = new TextDecoder(); + let receivedData = false; + + upstreamStream.on('data', (chunk: Buffer) => { + receivedData = true; + subscriber.next(decoder.decode(chunk, { stream: true })); + }); + + upstreamStream.on('end', () => { + if (!receivedData) { + subscriber.error(new Error('Empty response stream')); + return; + } + subscriber.complete(); + }); + + upstreamStream.on('error', (err: unknown) => { + const cleanError = err instanceof Error ? new Error(err.message) : new Error(String(err)); + subscriber.error(cleanError); + }); + }); + } + + private normalizeGeminiModel(model: string): string { + return model.replace(/^models\//i, ''); + } + + private createGeminiInternalRequest( + model: string, + request: GeminiRequest, + projectId: string, + requestType: string, + ): GeminiInternalRequest { + return { + project: projectId, + requestId: uuidv4(), + request: this.toInternalGeminiRequest(request), + model, + userAgent: process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent, + requestType, + }; + } + + private normalizeGeminiGenerateResponse(response: GeminiResponse): GeminiResponse { + const candidates = Array.isArray(response.candidates) + ? response.candidates.map((candidate, index) => ({ + content: candidate?.content, + finishReason: candidate?.finishReason, + index: typeof candidate?.index === 'number' ? candidate.index : index, + })) + : []; + + const normalized: GeminiResponse = { + candidates, + }; + + const usage = response.usageMetadata; + if (usage) { + const usageMetadata: NonNullable = {}; + if (usage.promptTokenCount !== undefined) { + usageMetadata.promptTokenCount = usage.promptTokenCount; + } + if (usage.candidatesTokenCount !== undefined) { + usageMetadata.candidatesTokenCount = usage.candidatesTokenCount; + } + if (usage.totalTokenCount !== undefined) { + usageMetadata.totalTokenCount = usage.totalTokenCount; + } + if (usage.thoughtsTokenCount !== undefined) { + usageMetadata.thoughtsTokenCount = usage.thoughtsTokenCount; + } + if (usage.promptTokensDetails !== undefined) { + usageMetadata.promptTokensDetails = usage.promptTokensDetails; + } + if (usage.candidatesTokensDetails !== undefined) { + usageMetadata.candidatesTokensDetails = usage.candidatesTokensDetails; + } + if (usage.trafficType !== undefined) { + usageMetadata.trafficType = usage.trafficType; + } + if (!isEmpty(usageMetadata)) { + normalized.usageMetadata = usageMetadata; + } + } + + return normalized; + } async handleChatCompletions( request: OpenAIChatRequest, ): Promise> { + const sessionKey = this.getSessionKeyFromOpenAI(request); + const targetModel = this.mapModel(request.model); + const extraHeaders = this.buildModelSpecificHeaders(request.model); this.logger.log( `Received request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`, ); @@ -273,49 +548,120 @@ export class ProxyService { // Retry loop for account selection let lastError: unknown = null; const maxRetries = 3; + const attemptedAccountIds = new Set(); 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)`); + await sleep(delay); + } + // 1. Get Token - const token = await this.tokenManager.getNextToken(); + const token = await this.tokenManager.getNextToken({ + sessionKey, + excludeAccountIds: Array.from(attemptedAccountIds), + }); if (!token) { throw new Error('No available accounts (all exhausted or rate limited)'); } + attemptedAccountIds.add(token.id); try { - // Convert OpenAI request to Claude format for Antigravity pipeline const claudeRequest = this.convertOpenAIToClaude(request); - const projectId = token.token.project_id!; - const geminiBody = transformClaudeRequestIn( - claudeRequest as unknown as ClaudeRequest, - projectId, - ); + const projectId = token.token.project_id ?? ''; + const geminiBody = transformClaudeRequestIn(claudeRequest, projectId); // Use v1internal API (same as Anthropic handler) if (request.stream) { - const stream = await this.geminiClient.streamGenerateInternal( - geminiBody, - token.token.access_token, - ); - return this.processStreamResponse(stream, request.model); + try { + const stream = await this.geminiClient.streamGenerateInternal( + geminiBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.processStreamResponse(stream, request.model); + } catch (streamError) { + this.logger.warn( + `Stream request failed for ${request.model}, falling back to non-stream generation: ${ + streamError instanceof Error ? streamError.message : String(streamError) + }`, + ); + + const response = await this.generateInternalWithStreamFallback( + geminiBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + this.logger.log( + `Upstream response snippet (stream fallback): ${JSON.stringify(response).substring(0, 500)}`, + ); + const claudeResponse = transformResponse(response); + const openaiResponse = this.convertClaudeToOpenAIResponse( + claudeResponse, + request.model, + ); + return this.createSyntheticOpenAIStream(openaiResponse); + } } else { - const response = await this.geminiClient.generateInternal( + const response = await this.generateInternalWithStreamFallback( geminiBody, token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + this.logger.log( + `Upstream response snippet: ${JSON.stringify(response).substring(0, 500)}`, ); - this.logger.log(`Raw Gemini Response: ${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)}`); return this.convertClaudeToOpenAIResponse(claudeResponse, request.model); } } catch (error) { + if (error instanceof Error && this.isProjectContextError(error.message)) { + this.logger.warn( + `OpenAI compatibility request hit project context issue, retrying without project: ${error.message}`, + ); + try { + const claudeRequest = this.convertOpenAIToClaude(request); + const fallbackBody = transformClaudeRequestIn(claudeRequest, ''); + if (request.stream) { + const stream = await this.geminiClient.streamGenerateInternal( + fallbackBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + return this.processStreamResponse(stream, request.model); + } + + const response = await this.generateInternalWithStreamFallback( + fallbackBody, + token.token.access_token, + token.token.upstream_proxy_url, + extraHeaders, + ); + const claudeResponse = transformResponse(response); + return this.convertClaudeToOpenAIResponse(claudeResponse, request.model); + } catch (fallbackError) { + error = fallbackError; + } + } + lastError = error; if (error instanceof Error) { this.logger.warn(`Request failed with account ${token.email}: ${error.message}`); - // Check for rate limit / auth errors - if (this.shouldRetry(error.message)) { - this.tokenManager.markAsRateLimited(token.email); + 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); + } } } } @@ -323,15 +669,138 @@ export class ProxyService { throw lastError || new Error('Request failed after retries'); } + private async generateInternalWithStreamFallback( + body: GeminiInternalRequest, + accessToken: string, + upstreamProxyUrl?: string, + extraHeaders?: Record, + ): Promise { + const direct = await this.geminiClient.generateInternal( + body, + accessToken, + upstreamProxyUrl, + extraHeaders, + ); + if (this.hasUsableGeminiCandidate(direct)) { + return direct; + } + + this.logger.warn('Empty non-stream response detected, falling back to stream aggregation.'); + const stream = await this.geminiClient.streamGenerateInternal( + body, + accessToken, + upstreamProxyUrl, + extraHeaders, + ); + return this.collectGeminiStreamAsResponse(stream); + } + + private hasUsableGeminiCandidate(response: GeminiResponse): boolean { + const candidates = response?.candidates; + if (!Array.isArray(candidates) || candidates.length === 0) { + return false; + } + + const first = candidates[0]; + const parts = first?.content?.parts; + return Array.isArray(parts) && parts.length > 0; + } + + private collectGeminiStreamAsResponse( + upstreamStream: NodeJS.ReadableStream, + ): Promise { + return new Promise((resolve, reject) => { + const decoder = new TextDecoder(); + let buffer = ''; + let receivedData = false; + const mergedParts: InternalGeminiPart[] = []; + let finishReason: string | undefined; + let usageMetadata: GeminiResponse['usageMetadata']; + + upstreamStream.on('data', (chunk: Buffer) => { + receivedData = true; + buffer += decoder.decode(chunk, { stream: true }); + const lines = buffer.split('\n'); + buffer = lines.pop() || ''; + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith('data: ')) { + continue; + } + + const dataStr = trimmed.slice(6); + if (dataStr === '[DONE]') { + continue; + } + + try { + const parsed = JSON.parse(dataStr); + const candidate = parsed?.candidates?.[0]; + const parts = candidate?.content?.parts; + if (Array.isArray(parts)) { + mergedParts.push( + ...parts.filter((part): part is InternalGeminiPart => this.isGeminiPart(part)), + ); + } + + if (candidate?.finishReason) { + finishReason = candidate.finishReason; + } + if (parsed?.usageMetadata) { + usageMetadata = parsed.usageMetadata; + } + } catch { + // Ignore malformed chunks and continue collecting valid parts. + } + } + }); + + upstreamStream.on('end', () => { + if (!receivedData) { + reject(new Error('Empty response stream')); + return; + } + + resolve({ + candidates: [ + { + content: { + role: 'model', + parts: mergedParts, + }, + finishReason, + }, + ], + usageMetadata, + }); + }); + + upstreamStream.on('error', (error: unknown) => { + reject(error instanceof Error ? error : new Error(String(error))); + }); + }); + } + // Handle SSE Stream conversion - private processStreamResponse(upstreamStream: any, model: string): Observable { + private processStreamResponse( + upstreamStream: NodeJS.ReadableStream, + model: string, + ): Observable { return new Observable((subscriber) => { const decoder = new TextDecoder(); let buffer = ''; + let hasEmittedChunk = false; + let hasSentDone = false; const streamId = `chatcmpl-${uuidv4()}`; const created = Math.floor(Date.now() / 1000); + const pushChunk = (payload: Record): void => { + hasEmittedChunk = true; + subscriber.next(`data: ${JSON.stringify(payload)}\n\n`); + }; + upstreamStream.on('data', (chunk: Buffer) => { buffer += decoder.decode(chunk, { stream: true }); const lines = buffer.split('\n'); @@ -347,57 +816,115 @@ export class ProxyService { try { const json = JSON.parse(dataStr); const candidate = json.candidates?.[0]; - const contentPart = candidate?.content?.parts?.[0]; - const text = contentPart?.text || ''; - const isThought = contentPart?.thought === true; + const parts = candidate?.content?.parts || []; - if (isThought && text) { - const reasoningChunk = { - id: streamId, - object: 'chat.completion.chunk', - created: created, - model: model, - choices: [ - { - index: 0, - delta: { reasoning_content: text }, - finish_reason: null, - }, - ], - }; - subscriber.next(`data: ${JSON.stringify(reasoningChunk)}\n\n`); - continue; - } - - if (text || candidate?.finishReason) { - const finishReason = candidate?.finishReason - ? candidate.finishReason === 'STOP' - ? 'stop' - : candidate.finishReason.toLowerCase() - : null; - - const contentChunk = { - id: streamId, - object: 'chat.completion.chunk', - created: created, - model: model, - choices: [ - { - index: 0, - delta: { content: text }, - finish_reason: finishReason, - }, - ], - }; - if (!text && finishReason) { - delete contentChunk.choices[0].delta.content; + for (const part of parts) { + if (part.thought && part.text) { + const reasoningChunk = { + id: streamId, + object: 'chat.completion.chunk', + created: created, + model: model, + choices: [ + { + index: 0, + delta: { reasoning_content: part.text }, + finish_reason: null, + }, + ], + }; + pushChunk(reasoningChunk); + continue; } - subscriber.next(`data: ${JSON.stringify(contentChunk)}\n\n`); + if (part.functionCall) { + const toolCallChunk = { + id: streamId, + object: 'chat.completion.chunk', + created: created, + model: model, + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: part.functionCall.id || `${part.functionCall.name}-${uuidv4()}`, + type: 'function', + function: { + name: part.functionCall.name, + arguments: JSON.stringify(part.functionCall.args || {}), + }, + }, + ], + }, + finish_reason: null, + }, + ], + }; + pushChunk(toolCallChunk); + continue; + } + + if (part.inlineData) { + const mimeType = part.inlineData.mimeType || 'image/jpeg'; + const data = part.inlineData.data || ''; + const imageMarkdown = `\n\n![Generated Image](data:${mimeType};base64,${data})\n\n`; + const imageChunk = { + id: streamId, + object: 'chat.completion.chunk', + created: created, + model: model, + choices: [ + { + index: 0, + delta: { content: imageMarkdown }, + finish_reason: null, + }, + ], + }; + pushChunk(imageChunk); + continue; + } + + if (part.text) { + const contentChunk = { + id: streamId, + object: 'chat.completion.chunk', + created: created, + model: model, + choices: [ + { + index: 0, + delta: { content: part.text }, + finish_reason: null, + }, + ], + }; + pushChunk(contentChunk); + } } if (candidate?.finishReason) { + const finishChunk = { + id: streamId, + object: 'chat.completion.chunk', + created: created, + model: model, + choices: [ + { + index: 0, + delta: {}, + finish_reason: this.mapGeminiFinishReasonToOpenAIFinishReason( + candidate.finishReason, + ), + }, + ], + }; + pushChunk(finishChunk); subscriber.next('data: [DONE]\n\n'); + hasSentDone = true; subscriber.complete(); } } catch (e) { @@ -407,10 +934,29 @@ export class ProxyService { }); upstreamStream.on('end', () => { + if (!hasEmittedChunk) { + pushChunk({ + id: streamId, + object: 'chat.completion.chunk', + created, + model, + choices: [ + { + index: 0, + delta: { content: '' }, + finish_reason: null, + }, + ], + }); + } + if (!hasSentDone) { + subscriber.next('data: [DONE]\n\n'); + hasSentDone = true; + } subscriber.complete(); }); - upstreamStream.on('error', (err: any) => { + upstreamStream.on('error', (err: unknown) => { // Convert to clean Error to avoid circular reference issues (socket objects) const cleanError = err instanceof Error ? new Error(err.message) : new Error(String(err)); subscriber.error(cleanError); @@ -418,170 +964,365 @@ export class ProxyService { }); } - private convertOpenAIToGemini(request: OpenAIChatRequest): GeminiRequest { - const contents: GeminiContent[] = []; - let systemInstruction: { parts: GeminiPart[] } | undefined = undefined; + private createSyntheticOpenAIStream(response: OpenAIChatResponse): Observable { + return new Observable((subscriber) => { + const streamId = response.id || `chatcmpl-${uuidv4()}`; + const created = response.created || Math.floor(Date.now() / 1000); + const model = response.model; + const choice = response.choices?.[0]; + const finishReason = choice?.finish_reason ?? 'stop'; + const content = + choice?.message && typeof choice.message.content === 'string' ? choice.message.content : ''; + const chunkSize = 80; - // Extract system message - const messages = request.messages || []; - for (const msg of messages) { - if (msg.role === 'system') { - if (typeof msg.content === 'string') { - systemInstruction = { - parts: [{ text: msg.content }], - }; - } - } else { - let textPart = ''; - if (Array.isArray(msg.content)) { - // OpenAI multimodal array - textPart = msg.content - .filter((p: OpenAIContentPart) => p.type === 'text') - .map((p: OpenAIContentPart) => p.text || '') - .join('\n'); - } else { - textPart = msg.content; - } - - contents.push({ - role: msg.role === 'user' ? 'user' : 'model', - parts: [{ text: textPart }], - }); - } - } - - return { - contents, - systemInstruction, - generationConfig: { - temperature: request.temperature, - maxOutputTokens: request.max_tokens, - topP: request.top_p, - }, - }; - } - - private convertGeminiToOpenAIResponse( - geminiResponse: GeminiResponse, - model: string, - ): OpenAIChatResponse { - const candidate = geminiResponse.candidates?.[0]; - if (!candidate) throw new Error('No candidates in response'); - - // Process response for images - const processedCandidate = this.processInlineData(candidate); - - const parts = processedCandidate?.content?.parts || []; - const content = parts.map((p) => p.text || '').join(''); - - const finishReason = - candidate?.finishReason === 'STOP' - ? 'stop' - : candidate?.finishReason?.toLowerCase() || 'stop'; - - return { - id: `chatcmpl-${uuidv4()}`, - object: 'chat.completion', - created: Math.floor(Date.now() / 1000), - model: model, - choices: [ - { - index: 0, - message: { - role: 'assistant', - content: content, - }, - finish_reason: finishReason, - }, - ], - usage: { - prompt_tokens: geminiResponse.usageMetadata?.promptTokenCount || 0, - completion_tokens: geminiResponse.usageMetadata?.candidatesTokenCount || 0, - total_tokens: geminiResponse.usageMetadata?.totalTokenCount || 0, - }, - }; - } - - // Convert Gemini inlineData (images) to Markdown - private processInlineData(candidate: GeminiCandidate | undefined): GeminiCandidate | undefined { - if (!candidate?.content?.parts) return candidate; - - const newParts = candidate.content.parts.map((part) => { - if (part.inlineData) { - const mimeType = part.inlineData.mimeType || 'image/jpeg'; - const data = part.inlineData.data || ''; - return { - text: `\n\n![Generated Image](data:${mimeType};base64,${data})\n\n`, + if (content.length === 0) { + const finishChunk = { + id: streamId, + object: 'chat.completion.chunk', + created, + model, + choices: [ + { + index: 0, + delta: {}, + finish_reason: finishReason, + }, + ], + usage: response.usage, }; + subscriber.next(`data: ${JSON.stringify(finishChunk)}\n\n`); + subscriber.next('data: [DONE]\n\n'); + subscriber.complete(); + return; } - return part; - }); + for (let index = 0; index < content.length; index += chunkSize) { + const piece = content.slice(index, index + chunkSize); + const isLast = index + chunkSize >= content.length; + const chunk = { + id: streamId, + object: 'chat.completion.chunk', + created, + model, + choices: [ + { + index: 0, + delta: { content: piece }, + finish_reason: isLast ? finishReason : null, + }, + ], + usage: isLast + ? response.usage + : { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, + }; + subscriber.next(`data: ${JSON.stringify(chunk)}\n\n`); + } + + subscriber.next('data: [DONE]\n\n'); + subscriber.complete(); + }); + } + + private toClaudeRequest(request: AnthropicChatRequest): ClaudeRequest { return { - ...candidate, - content: { - ...candidate.content, - parts: newParts, + model: request.model, + messages: request.messages.map((message) => ({ + role: message.role, + content: message.content, + })), + system: request.system, + tools: request.tools?.map((tool) => ({ + name: tool.name, + description: tool.description, + input_schema: tool.input_schema, + type: tool.type, + })), + stream: request.stream, + max_tokens: request.max_tokens, + stop_sequences: request.stop_sequences, + temperature: request.temperature, + top_p: request.top_p, + top_k: request.top_k, + thinking: request.thinking, + metadata: request.metadata, + }; + } + + private toAnthropicChatResponse(response: ClaudeResponse): AnthropicChatResponse { + return { + id: response.id, + type: response.type, + role: response.role, + model: response.model, + content: response.content, + stop_reason: response.stop_reason, + stop_sequence: response.stop_sequence, + usage: { + input_tokens: response.usage?.input_tokens ?? 0, + output_tokens: response.usage?.output_tokens ?? 0, + cache_creation_input_tokens: response.usage?.cache_creation_input_tokens, + cache_read_input_tokens: response.usage?.cache_read_input_tokens, }, }; } + private toInternalGeminiRequest(request: GeminiRequest): GeminiInternalRequest['request'] { + return { + contents: request.contents, + generationConfig: request.generationConfig, + systemInstruction: request.systemInstruction + ? { + parts: request.systemInstruction.parts + .filter((part): part is { text: string } => typeof part.text === 'string') + .map((part) => ({ text: part.text })), + } + : undefined, + }; + } + // Convert OpenAI request format to Claude/Anthropic format - private convertOpenAIToClaude(request: OpenAIChatRequest): AnthropicChatRequest { + private convertOpenAIToClaude(request: OpenAIChatRequest): ClaudeRequest { const messages = request.messages || []; - let systemPrompt: string | undefined; - const anthropicMessages: { role: string; content: string }[] = []; + const systemPromptParts: string[] = []; + const anthropicMessages: ClaudeRequest['messages'] = []; for (const msg of messages) { if (msg.role === 'system') { - // Extract system message - if (typeof msg.content === 'string') { - systemPrompt = msg.content; + const systemText = this.extractOpenAITextContent(msg.content); + if (systemText) { + systemPromptParts.push(systemText); } - } else { - let textContent = ''; - if (Array.isArray(msg.content)) { - textContent = msg.content - .filter((p: OpenAIContentPart) => p.type === 'text') - .map((p: OpenAIContentPart) => p.text || '') - .join('\n'); - } else { - textContent = msg.content; - } - - anthropicMessages.push({ - role: msg.role === 'assistant' ? 'assistant' : 'user', - content: textContent, - }); + continue; } + + if (msg.role === 'tool') { + const toolResultText = this.extractOpenAITextContent(msg.content) || ''; + anthropicMessages.push({ + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: msg.tool_call_id || msg.name || `tool-result-${uuidv4()}`, + content: toolResultText, + is_error: false, + }, + ], + }); + continue; + } + + const contentBlocks = this.convertOpenAIPartsToAnthropicContent(msg.content); + + if (msg.role === 'assistant' && msg.tool_calls && msg.tool_calls.length > 0) { + for (const toolCall of msg.tool_calls) { + contentBlocks.push({ + type: 'tool_use', + id: toolCall.id, + name: toolCall.function.name, + input: this.parseOpenAIFunctionArguments(toolCall.function.arguments), + }); + } + } + + anthropicMessages.push({ + role: msg.role === 'assistant' ? 'assistant' : 'user', + content: contentBlocks.length > 0 ? contentBlocks : '', + }); } + const systemPrompt = systemPromptParts.length > 0 ? systemPromptParts.join('\n') : undefined; + return { model: request.model, messages: anthropicMessages, system: systemPrompt, - max_tokens: request.max_tokens || 4096, + tools: this.convertOpenAIToolsToAnthropicTools(request.tools), + max_tokens: request.max_tokens, temperature: request.temperature, top_p: request.top_p, stream: request.stream, + metadata: { + ...(request.extra ?? {}), + source: 'openai', + }, }; } - // Convert Claude response to OpenAI format - private convertClaudeToOpenAIResponse(claudeResponse: any, model: string): OpenAIChatResponse { - // Extract text content from Claude response - const content = - claudeResponse.content - ?.filter((block: any) => block.type === 'text') - .map((block: any) => block.text) - .join('') || ''; + private convertOpenAIPartsToAnthropicContent( + content: OpenAIChatRequest['messages'][number]['content'], + ): AnthropicContent[] { + if (typeof content === 'string') { + return content.trim() ? [{ type: 'text', text: content }] : []; + } - const finishReason = - claudeResponse.stop_reason === 'end_turn' - ? 'stop' - : claudeResponse.stop_reason === 'max_tokens' - ? 'length' - : 'stop'; + const blocks: AnthropicContent[] = []; + for (const part of content) { + if (part.type === 'text' && part.text) { + blocks.push({ type: 'text', text: part.text }); + continue; + } + + if (part.type === 'image_url' && part.image_url?.url) { + const url = part.image_url.url; + const dataUri = url.match(/^data:(?[^;]+);base64,(?.+)$/); + if (dataUri?.groups?.mime && dataUri.groups.data) { + blocks.push({ + type: 'image', + source: { + type: 'base64', + media_type: dataUri.groups.mime, + data: dataUri.groups.data, + }, + }); + } else { + blocks.push({ type: 'text', text: `[image_url] ${url}` }); + } + } + } + return blocks; + } + + private extractOpenAITextContent( + content: OpenAIChatRequest['messages'][number]['content'], + ): string { + if (typeof content === 'string') { + return content; + } + + return content + .filter((part) => part.type === 'text') + .map((part) => part.text || '') + .join('\n'); + } + + private parseOpenAIFunctionArguments(argumentsString: string): Record { + if (!argumentsString || argumentsString.trim() === '') { + return {}; + } + + try { + const parsed = JSON.parse(argumentsString); + if (this.isRecord(parsed)) { + return parsed; + } + return { value: parsed }; + } catch { + return { raw: argumentsString }; + } + } + + private convertOpenAIToolsToAnthropicTools( + tools: OpenAIChatRequest['tools'], + ): AnthropicChatRequest['tools'] { + if (!tools || tools.length === 0) { + return undefined; + } + + return tools.map((tool) => ({ + name: tool.function.name, + description: tool.function.description, + input_schema: tool.function.parameters || { type: 'object', properties: {} }, + })); + } + + private mapGeminiFinishReasonToOpenAIFinishReason(finishReason?: string): string | null { + if (!finishReason) { + return null; + } + + const normalized = finishReason.toUpperCase(); + if (normalized === 'STOP') { + return 'stop'; + } + if (normalized === 'MAX_TOKENS') { + return 'length'; + } + if (normalized === 'SAFETY' || normalized === 'RECITATION') { + return 'content_filter'; + } + + return finishReason.toLowerCase(); + } + + private mapAnthropicStopReasonToOpenAIFinishReason(stopReason?: string | null): string | null { + if (!stopReason) { + return null; + } + + if (stopReason === 'end_turn') { + return 'stop'; + } + if (stopReason === 'max_tokens') { + return 'length'; + } + if (stopReason === 'tool_use') { + return 'tool_calls'; + } + + return stopReason; + } + + private normalizeToolCallArguments(input: unknown): string { + if (typeof input === 'string') { + return input; + } + if (isNil(input)) { + return '{}'; + } + + try { + return JSON.stringify(input); + } catch { + return '{}'; + } + } + + // Convert Claude response to OpenAI format + private convertClaudeToOpenAIResponse( + claudeResponse: ClaudeResponse, + model: string, + ): OpenAIChatResponse { + const contentBlocks = Array.isArray(claudeResponse?.content) ? claudeResponse.content : []; + + const textContent = contentBlocks + .filter( + ( + block, + ): block is Extract => + block?.type === 'text', + ) + .map((block) => block.text || '') + .join(''); + + const reasoningContent = contentBlocks + .filter( + ( + block, + ): block is Extract< + ClaudeResponse['content'][number], + { type: 'thinking'; thinking: string } + > => block?.type === 'thinking', + ) + .map((block) => block.thinking || '') + .join(''); + + const toolCalls = contentBlocks + .filter( + ( + block, + ): block is Extract< + ClaudeResponse['content'][number], + { type: 'tool_use'; id: string; name: string; input: Record } + > => block?.type === 'tool_use', + ) + .map((block, index: number) => ({ + id: block.id || `tool-call-${index}`, + type: 'function' as const, + function: { + name: block.name || 'unknown_tool', + arguments: this.normalizeToolCallArguments(block.input), + }, + })); return { id: `chatcmpl-${uuidv4()}`, @@ -593,9 +1334,13 @@ export class ProxyService { index: 0, message: { role: 'assistant', - content: content, + content: textContent || null, + tool_calls: toolCalls.length > 0 ? toolCalls : undefined, + reasoning_content: reasoningContent || undefined, }, - finish_reason: finishReason, + finish_reason: this.mapAnthropicStopReasonToOpenAIFinishReason( + claudeResponse.stop_reason, + ), }, ], usage: { @@ -608,35 +1353,177 @@ export class ProxyService { } private mapModel(model: string): string { - const lower = model.toLowerCase(); - if (lower.includes('sonnet') || lower.includes('thinking')) return 'gemini-3-pro-preview'; - if (lower.includes('haiku')) return 'gemini-2.0-flash-exp'; - if (lower.includes('opus')) return 'gemini-3-pro-preview'; - if (lower.includes('claude')) return 'gemini-2.5-flash-thinking'; - if (lower === 'gemini-3-pro-high' || lower === 'gemini-3-pro-low') - return 'gemini-3-pro-preview'; - if (lower === 'gemini-3-flash') return 'gemini-3-flash-preview'; + const normalizedModel = model.replace(/^models\//i, '').trim(); + const config = getServerConfig(); + const configuredMapping = { + ...(config?.custom_mapping ?? {}), + ...(config?.anthropic_mapping ?? {}), + }; - // Default to original if no mapping found - return model; + const customExactMapping: Record = {}; + const wildcardMapping: Array<{ + pattern: RegExp; + target: string; + }> = []; + + for (const [key, target] of Object.entries(configuredMapping)) { + if (!key || !target) { + continue; + } + + if (key.includes('*')) { + const escaped = key.replace(/[.+?^${}()|[\]\\]/g, '\\$&').replace(/\*/g, '.*'); + wildcardMapping.push({ + pattern: new RegExp(`^${escaped}$`, 'i'), + target, + }); + continue; + } + + customExactMapping[key] = target; + } + + for (const wildcardRule of wildcardMapping) { + if (wildcardRule.pattern.test(normalizedModel)) { + return wildcardRule.target; + } + } + + const routedModel = resolveModelRoute(normalizedModel, customExactMapping, {}, {}); + + return routedModel; } - private shouldRetry(errorMessage: string): boolean { + private classifyUpstreamError(errorMessage: string): { + retry: boolean; + markAsForbidden: boolean; + markAsRateLimited: boolean; + } { const msg = errorMessage.toLowerCase(); - if ( - msg.includes('404') || - msg.includes('not_found') || + const isForbidden = + msg.includes('401') || + msg.includes('unauthorized') || + msg.includes('invalid_grant') || msg.includes('403') || - msg.includes('permission_denied') - ) - return true; - if ( + msg.includes('permission_denied') || + msg.includes('forbidden'); + + if (isForbidden) { + return { + retry: true, + markAsForbidden: true, + markAsRateLimited: false, + }; + } + + const isRateLimitedSignal = msg.includes('429') || msg.includes('resource_exhausted') || msg.includes('quota') || - msg.includes('rate_limit') - ) - return true; - return false; + msg.includes('rate_limit') || + msg.includes('rate limit'); + + const shouldRetryByStatus = + msg.includes('408') || + msg.includes('429') || + msg.includes('500') || + msg.includes('502') || + msg.includes('503') || + msg.includes('504'); + + const shouldRetryByKeyword = + msg.includes('resource_exhausted') || + msg.includes('quota') || + msg.includes('rate_limit') || + msg.includes('timeout') || + msg.includes('socket hang up') || + msg.includes('empty response stream') || + msg.includes('connection reset'); + + if (shouldRetryByStatus || shouldRetryByKeyword) { + return { + retry: true, + markAsForbidden: false, + markAsRateLimited: isRateLimitedSignal, + }; + } + + return { + retry: false, + markAsForbidden: false, + markAsRateLimited: false, + }; + } + + private buildModelSpecificHeaders(model: string | undefined): Record { + if (!model) { + return {}; + } + + if (model.toLowerCase().includes('claude')) { + return { + 'anthropic-beta': + 'claude-code-20250219,interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14', + }; + } + + return {}; + } + + private isProjectLicenseError(errorMessage: string): boolean { + const msg = errorMessage.toLowerCase(); + return ( + msg.includes('#3501') || + (msg.includes('google cloud project') && msg.includes('code assist license')) + ); + } + + private isProjectNotFoundError(errorMessage: string): boolean { + const msg = errorMessage.toLowerCase(); + return ( + (msg.includes('resource projects/') && msg.includes('could not be found')) || + (msg.includes('project') && msg.includes('not found')) + ); + } + + private isProjectContextError(errorMessage: string): boolean { + return this.isProjectLicenseError(errorMessage) || this.isProjectNotFoundError(errorMessage); + } + + private isQuotaExhaustedError(errorMessage: string): boolean { + const msg = errorMessage.toLowerCase(); + return ( + msg.includes('resource has been exhausted') || + msg.includes('resource_exhausted') || + msg.includes('quota') + ); + } + + private getSessionKeyFromAnthropic(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() === '') { + return undefined; + } + return `anthropic:${sessionCandidate.trim()}`; + } + + private getSessionKeyFromOpenAI(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() === '') { + return undefined; + } + return `openai:${sessionCandidate.trim()}`; + } + + private isRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value); + } + + private isGeminiPart(value: unknown): value is InternalGeminiPart { + return this.isRecord(value); } } diff --git a/src/server/modules/proxy/token-manager.service.ts b/src/server/modules/proxy/token-manager.service.ts index 002efdc..3d5c8b5 100644 --- a/src/server/modules/proxy/token-manager.service.ts +++ b/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 = new Map(); // Cooldown map for rate-limited accounts private cooldowns: Map = new Map(); + private sessionBindings: Map = 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 { + async getNextToken(options?: { + sessionKey?: string; + excludeAccountIds?: string[]; + }): Promise { 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 { + 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); } } diff --git a/src/tests/unit/gemini-controller.integration.test.ts b/src/tests/unit/gemini-controller.integration.test.ts new file mode 100644 index 0000000..b0f158d --- /dev/null +++ b/src/tests/unit/gemini-controller.integration.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, it, vi } from 'vitest'; +import { of } from 'rxjs'; + +import { GeminiController } from '../../server/modules/proxy/gemini.controller'; + +function createReplyMock() { + const reply: Record = {}; + reply.status = vi.fn(() => reply); + reply.header = vi.fn(() => reply); + reply.send = vi.fn(() => reply); + return reply; +} + +describe('GeminiController Integration', () => { + it('supports list and get model endpoints', () => { + const proxyService = {}; + const controller = new GeminiController(proxyService as any); + const replyList = createReplyMock(); + const replyGet = createReplyMock(); + + controller.listModels(replyList as any); + controller.getModel('gemini-2.5-flash', replyGet as any); + + expect(replyList.status).toHaveBeenCalledWith(200); + expect(replyList.send).toHaveBeenCalledWith( + expect.objectContaining({ + models: expect.any(Array), + }), + ); + expect(replyList.send).toHaveBeenCalledWith( + expect.objectContaining({ + models: expect.arrayContaining([ + expect.objectContaining({ + name: 'models/gemini-2.5-flash', + description: '', + inputTokenLimit: 128000, + outputTokenLimit: 8192, + supportedGenerationMethods: ['generateContent', 'countTokens'], + }), + ]), + }), + ); + expect(replyGet.status).toHaveBeenCalledWith(200); + expect(replyGet.send).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'models/gemini-2.5-flash', + displayName: 'gemini-2.5-flash', + }), + ); + }); + + it('handles generateContent action from colon endpoint format', async () => { + const proxyService = { + handleGeminiGenerateContent: vi.fn().mockResolvedValue({ + candidates: [ + { + content: { role: 'model', parts: [{ text: 'hello' }] }, + finishReason: 'STOP', + avgLogprobs: -0.1, + }, + ], + usageMetadata: { + promptTokenCount: 1, + candidatesTokenCount: 1, + totalTokenCount: 2, + }, + createTime: '2026-02-10T10:00:00.000Z', + modelVersion: 'gemini-2.5-flash-latest', + responseId: 'resp_123', + }), + handleGeminiStreamGenerateContent: vi.fn(), + }; + const controller = new GeminiController(proxyService as any); + const reply = createReplyMock(); + + await controller.modelAction( + 'models/gemini-2.5-flash:generateContent', + { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] } as any, + reply as any, + ); + + expect(proxyService.handleGeminiGenerateContent).toHaveBeenCalledWith( + 'models/gemini-2.5-flash', + expect.any(Object), + ); + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith({ + candidates: [ + { + content: { role: 'model', parts: [{ text: 'hello' }] }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 1, + candidatesTokenCount: 1, + totalTokenCount: 2, + }, + }); + }); + + it('handles streamGenerateContent action and emits SSE headers', async () => { + const stream = of('data: {"ok":true}\n\n'); + const proxyService = { + handleGeminiGenerateContent: vi.fn(), + handleGeminiStreamGenerateContent: vi.fn().mockResolvedValue(stream), + }; + const controller = new GeminiController(proxyService as any); + const reply = createReplyMock(); + + await controller.modelAction( + 'gemini-2.5-flash:streamGenerateContent', + { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] } as any, + 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('supports countTokens action', async () => { + const proxyService = { + handleGeminiGenerateContent: vi.fn(), + handleGeminiStreamGenerateContent: vi.fn(), + }; + const controller = new GeminiController(proxyService as any); + const reply = createReplyMock(); + + await controller.countTokens( + 'gemini-2.5-flash', + { contents: [{ role: 'user', parts: [{ text: 'abcd efgh' }] }] } as any, + reply as any, + ); + + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith({ totalTokens: 0 }); + }); + + it('returns bad request for invalid combined endpoint action', async () => { + const proxyService = { + handleGeminiGenerateContent: vi.fn(), + handleGeminiStreamGenerateContent: vi.fn(), + }; + const controller = new GeminiController(proxyService as any); + const reply = createReplyMock(); + + await controller.modelAction( + 'models/gemini-2.5-flash-generateContent', + { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] } as any, + reply as any, + ); + + expect(reply.status).toHaveBeenCalledWith(400); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + error: expect.objectContaining({ + status: 'INVALID_ARGUMENT', + }), + }), + ); + }); +}); diff --git a/src/tests/unit/logger.test.ts b/src/tests/unit/logger.test.ts index 4982e21..3cebc3c 100644 --- a/src/tests/unit/logger.test.ts +++ b/src/tests/unit/logger.test.ts @@ -1,5 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach, afterAll } from 'vitest'; -import { logger } from '../../utils/logger'; +import { describe, it, expect, vi, beforeAll, beforeEach, afterEach, afterAll } from 'vitest'; import fs from 'fs'; import path from 'path'; @@ -12,6 +11,10 @@ vi.mock('../../utils/paths', async () => { describe('Logger Utilities', () => { const testLogDir = path.join(process.cwd(), 'temp_test_logs'); + let logger: { + info: (message: string, ...args: unknown[]) => void; + error: (message: string, ...args: unknown[]) => void; + }; const getLatestLogFile = () => { const files = fs @@ -27,7 +30,7 @@ describe('Logger Utilities', () => { }; const waitForLogContains = async (text: string) => { - for (let i = 0; i < 20; i++) { + for (let i = 0; i < 100; i++) { const filePath = getLatestLogFile(); if (filePath && fs.existsSync(filePath)) { const content = fs.readFileSync(filePath, 'utf-8'); @@ -41,30 +44,22 @@ describe('Logger Utilities', () => { return null; }; + beforeAll(async () => { + if (fs.existsSync(testLogDir)) { + fs.rmSync(testLogDir, { recursive: true, force: true }); + } + fs.mkdirSync(testLogDir, { recursive: true }); + const loggerModule = await import('../../utils/logger'); + logger = loggerModule.logger; + }); + beforeEach(() => { vi.spyOn(console, 'log').mockImplementation(() => {}); vi.spyOn(console, 'error').mockImplementation(() => {}); - try { - if (fs.existsSync(testLogDir)) { - fs.rmSync(testLogDir, { recursive: true, force: true }); - } - fs.mkdirSync(testLogDir, { recursive: true }); - } catch (err) { - console.error('beforeEach: setup testLogDir failed', err); - throw err; - } }); afterEach(() => { vi.restoreAllMocks(); - try { - if (fs.existsSync(testLogDir)) { - fs.rmSync(testLogDir, { recursive: true, force: true }); - } - } catch (err) { - console.error('afterEach: cleanup testLogDir failed', err); - throw err; - } }); afterAll(() => { diff --git a/src/tests/unit/proxy-controller.integration.test.ts b/src/tests/unit/proxy-controller.integration.test.ts new file mode 100644 index 0000000..0d6d8f4 --- /dev/null +++ b/src/tests/unit/proxy-controller.integration.test.ts @@ -0,0 +1,482 @@ +import { describe, expect, it, vi } from 'vitest'; +import { of } from 'rxjs'; + +import { ProxyController } from '../../server/modules/proxy/proxy.controller'; + +function createReplyMock() { + const reply: Record = {}; + reply.status = vi.fn(() => reply); + reply.header = vi.fn(() => reply); + reply.send = vi.fn(() => reply); + return reply; +} + +describe('ProxyController Integration', () => { + it('supports models listing endpoints', () => { + const proxyService = {}; + const controller = new ProxyController(proxyService as any); + + const allModels = controller.getModels(); + const claudeModels = controller.getClaudeModels(); + + expect(allModels.object).toBe('list'); + expect(Array.isArray(allModels.data)).toBe(true); + expect(allModels.data.length).toBeGreaterThan(0); + expect(allModels.data.every((model) => model.owned_by === 'antigravity')).toBe(true); + expect(allModels.data.every((model) => model.created === 1770652800)).toBe(true); + + expect(claudeModels).toEqual(allModels); + }); + + it('routes Claude OpenAI requests to protocol parity path', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue({ ok: true }), + handleAnthropicMessages: vi.fn(), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.chatCompletions( + { + model: 'claude-sonnet-4-5', + stream: false, + messages: [{ role: 'user', content: 'hello' }], + } as any, + reply as any, + ); + + expect(proxyService.handleChatCompletions).toHaveBeenCalledOnce(); + expect(proxyService.handleAnthropicMessages).not.toHaveBeenCalled(); + expect(reply.status).toHaveBeenCalledWith(200); + }); + + it('returns stream response with SSE headers for parity stream path', async () => { + const stream = of('data: {"ok":true}\n\n'); + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue(stream), + handleAnthropicMessages: vi.fn(), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.chatCompletions( + { + model: 'claude-sonnet-4-5', + stream: true, + messages: [{ role: 'user', content: 'hello' }], + } as any, + 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('supports OpenAI completions compatibility endpoint', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue({ + id: 'chatcmpl_test', + object: 'chat.completion', + created: 1700000000, + model: 'gpt-4o', + choices: [ + { + index: 0, + finish_reason: 'stop', + message: { + role: 'assistant', + content: 'hello from assistant', + }, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }, + }), + handleAnthropicMessages: vi.fn(), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.completions( + { + model: 'gpt-4o', + prompt: 'hello world', + stream: false, + }, + reply as any, + ); + + expect(proxyService.handleChatCompletions).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'gpt-4o', + messages: [{ role: 'user', content: 'hello world' }], + }), + ); + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + object: 'text_completion', + model: 'gpt-4o', + choices: [ + expect.objectContaining({ + text: 'hello from assistant', + logprobs: null, + }), + ], + }), + ); + }); + + it('supports OpenAI responses compatibility endpoint with normalized input', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue({ + id: 'chatcmpl_resp', + object: 'chat.completion', + created: 1700000001, + model: 'gpt-4o', + choices: [ + { + index: 0, + finish_reason: 'stop', + message: { + content: 'normalized response', + }, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 6, + total_tokens: 16, + }, + }), + }; + + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.responses( + { + model: 'gpt-4o', + instructions: 'Follow the tool protocol', + input: [ + { + type: 'message', + role: 'user', + content: [{ type: 'input_text', text: 'hello' }], + }, + { + type: 'function_call', + id: 'call_1', + name: 'search_docs', + arguments: '{"query":"token"}', + }, + { + type: 'function_call_output', + call_id: 'call_1', + output: { content: 'result: ok' }, + }, + ], + }, + reply as any, + ); + + const callArg = proxyService.handleChatCompletions.mock.calls[0][0]; + expect(callArg.messages[0]).toEqual({ + role: 'system', + content: 'Follow the tool protocol', + }); + expect(callArg.messages.some((message: { role: string }) => message.role === 'assistant')).toBe( + true, + ); + expect(callArg.messages.some((message: { role: string }) => message.role === 'tool')).toBe( + true, + ); + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + object: 'text_completion', + model: 'gpt-4o', + choices: [ + expect.objectContaining({ + text: 'normalized response', + logprobs: null, + }), + ], + usage: expect.objectContaining({ + prompt_tokens: 10, + completion_tokens: 6, + total_tokens: 16, + }), + }), + ); + }); + + it('supports image generations endpoint', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue({ + choices: [ + { + message: { + content: '![img](data:image/png;base64,AAAABBBB)', + }, + }, + ], + }), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.imageGenerations( + { + model: 'gemini-3-pro-image', + prompt: 'draw a cat', + }, + reply as any, + ); + + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + data: [ + expect.objectContaining({ + b64_json: 'AAAABBBB', + }), + ], + }), + ); + }); + + it('maps image generation upstream quota errors to 429', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockRejectedValue(new Error('429 quota exceeded')), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.imageGenerations( + { + model: 'gemini-3-pro-image', + prompt: 'draw a dog', + }, + reply as any, + ); + + expect(reply.status).toHaveBeenCalledWith(429); + }); + + it('falls back to Gemini image generation when chat path hits project context error', async () => { + const proxyService = { + handleChatCompletions: vi + .fn() + .mockRejectedValue( + new Error( + 'You are currently configured to use a Google Cloud Project but lack a Gemini Code Assist license. (#3501)', + ), + ), + handleGeminiGenerateContent: vi.fn().mockResolvedValue({ + candidates: [ + { + content: { + parts: [ + { + inlineData: { + mimeType: 'image/png', + data: 'FALLBACKIMG', + }, + }, + ], + }, + }, + ], + }), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.imageGenerations( + { + model: 'gemini-3-pro-image', + prompt: 'draw a fox', + }, + reply as any, + ); + + expect(proxyService.handleChatCompletions).toHaveBeenCalledOnce(); + expect(proxyService.handleGeminiGenerateContent).toHaveBeenCalledOnce(); + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + data: [ + expect.objectContaining({ + b64_json: 'FALLBACKIMG', + }), + ], + }), + ); + }); + + it('supports image edits endpoint with supplementary image payload', async () => { + const proxyService = { + handleChatCompletions: vi.fn().mockResolvedValue({ + choices: [ + { + message: { + content: '![img](data:image/png;base64,CCCCDDDD)', + }, + }, + ], + }), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.imageEdits( + { + model: 'gemini-3-pro-image', + prompt: 'make it brighter', + image: 'data:image/png;base64,IMGBASE64', + reference_images: ['data:image/jpeg;base64,REFBASE64'], + }, + { + headers: { + 'content-type': 'multipart/form-data; boundary=----parity', + }, + } as any, + reply as any, + ); + + const request = proxyService.handleChatCompletions.mock.calls[0][0]; + expect(Array.isArray(request.messages[0].content)).toBe(true); + expect(reply.status).toHaveBeenCalledWith(200); + }); + + it('rejects image edits request without multipart boundary', async () => { + const proxyService = { + handleChatCompletions: vi.fn(), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.imageEdits( + { + model: 'gemini-3-pro-image', + prompt: 'make it brighter', + image: 'data:image/png;base64,IMGBASE64', + }, + { + headers: { + 'content-type': 'application/json', + }, + } as any, + reply as any, + ); + + expect(proxyService.handleChatCompletions).not.toHaveBeenCalled(); + expect(reply.status).toHaveBeenCalledWith(400); + expect(reply.send).toHaveBeenCalledWith('Invalid `boundary` for `multipart/form-data` request'); + }); + + it('supports audio transcriptions endpoint', async () => { + const proxyService = { + handleGeminiGenerateContent: vi.fn().mockResolvedValue({ + candidates: [ + { + content: { + parts: [{ text: 'transcribed text' }], + }, + }, + ], + }), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.audioTranscriptions( + { + model: 'gemini-2.5-flash', + file: 'data:audio/mpeg;base64,QUJDRA==', + }, + { + headers: { + 'content-type': 'multipart/form-data; boundary=----parity', + }, + } as any, + reply as any, + ); + + expect(proxyService.handleGeminiGenerateContent).toHaveBeenCalledOnce(); + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith({ text: 'transcribed text' }); + }); + + it('rejects audio transcription request without multipart boundary', async () => { + const proxyService = { + handleGeminiGenerateContent: vi.fn(), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.audioTranscriptions( + { + model: 'gemini-2.5-flash', + file: 'data:audio/mpeg;base64,QUJDRA==', + }, + { + headers: { + 'content-type': 'application/json', + }, + } as any, + reply as any, + ); + + expect(proxyService.handleGeminiGenerateContent).not.toHaveBeenCalled(); + expect(reply.status).toHaveBeenCalledWith(400); + expect(reply.send).toHaveBeenCalledWith('Invalid `boundary` for `multipart/form-data` request'); + }); + + it('supports model detect endpoint', () => { + const proxyService = { + detectModelCapabilities: vi.fn().mockReturnValue({ + model: 'claude-sonnet-4-5', + mapped_model: 'gemini-3-pro-preview', + type: 'generate-content', + features: { has_web_search: false, is_image_gen: false }, + }), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + controller.detectModel({ model: 'claude-sonnet-4-5' }, reply as any); + + expect(proxyService.detectModelCapabilities).toHaveBeenCalledWith('claude-sonnet-4-5'); + expect(reply.status).toHaveBeenCalledWith(200); + }); + + it('supports Anthropic messages endpoint', async () => { + const proxyService = { + handleAnthropicMessages: vi.fn().mockResolvedValue({ + id: 'msg_1', + type: 'message', + }), + }; + const controller = new ProxyController(proxyService as any); + const reply = createReplyMock(); + + await controller.anthropicMessages( + { + model: 'claude-sonnet-4-5', + stream: false, + messages: [{ role: 'user', content: 'hello' }], + } as any, + reply as any, + ); + + expect(proxyService.handleAnthropicMessages).toHaveBeenCalledOnce(); + expect(reply.status).toHaveBeenCalledWith(200); + }); +}); diff --git a/src/tests/unit/proxy-parity-fixtures.test.ts b/src/tests/unit/proxy-parity-fixtures.test.ts new file mode 100644 index 0000000..a1d67e8 --- /dev/null +++ b/src/tests/unit/proxy-parity-fixtures.test.ts @@ -0,0 +1,94 @@ +import { readFileSync } from 'fs'; +import path from 'path'; +import { EventEmitter } from 'events'; +import { describe, expect, it, vi } from 'vitest'; +import { Observable } from 'rxjs'; + +import { ProxyService } from '../../server/modules/proxy/proxy.service'; + +const mockTokenManager = { getNextToken: vi.fn(), markAsRateLimited: vi.fn(), markAsForbidden: vi.fn() }; +const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; + +class TestableProxyService extends ProxyService { + constructor() { + super(mockTokenManager as any, mockGeminiClient as any); + } + + public toAnthropic(request: any): any { + return (this as any).convertOpenAIToClaude(request); + } + + public toOpenAI(response: any, model: string): any { + return (this as any).convertClaudeToOpenAIResponse(response, model); + } + + public streamToOpenAI(upstreamStream: any, model: string): Observable { + return (this as any).processStreamResponse(upstreamStream, model); + } +} + +function readFixture(relativePath: string): T { + const fullPath = path.join(process.cwd(), 'src/tests/fixtures/proxy-parity', relativePath); + return JSON.parse(readFileSync(fullPath, 'utf-8')) as T; +} + +describe('Proxy Parity Fixtures', () => { + it('maps OpenAI request fixture to expected Anthropic request semantics', () => { + const service = new TestableProxyService(); + const input = readFixture('request/openai.chat-tools.input.json'); + const expected = readFixture('request/openai.chat-tools.expected.json'); + + const actual = service.toAnthropic(input); + + expect(actual.model).toBe(expected.model); + expect(actual.system).toBe(expected.system); + expect(actual.temperature).toBe(expected.temperature); + expect(actual.max_tokens).toBe(expected.max_tokens); + expect(actual.tools?.[0]?.name).toBe(expected.tools[0].name); + expect(actual.messages[0]).toEqual(expected.messages[0]); + expect(actual.messages[1].content[1]).toEqual(expected.messages[1].content[1]); + expect(actual.messages[2].content[0]).toEqual(expected.messages[2].content[0]); + }); + + it('maps Anthropic response fixture to expected OpenAI response semantics', () => { + const service = new TestableProxyService(); + const input = readFixture('response/anthropic.tool-use.input.json'); + const expected = readFixture('response/anthropic.tool-use.expected.json'); + + const actual = service.toOpenAI(input, expected.model); + + expect(actual.model).toBe(expected.model); + expect(actual.choices[0].message.content).toBe(expected.message.content); + expect(actual.choices[0].message.reasoning_content).toBe(expected.message.reasoning_content); + expect(actual.choices[0].message.tool_calls?.[0]).toEqual(expected.message.tool_calls[0]); + expect(actual.choices[0].finish_reason).toBe(expected.finish_reason); + expect(actual.usage).toEqual(expected.usage); + }); + + it('maps upstream stream fixture into expected OpenAI SSE semantics', async () => { + const service = new TestableProxyService(); + const input = readFixture('stream/openai-from-gemini.input.json'); + const expected = readFixture<{ contains: string[] }>('stream/openai-from-gemini.expected.json'); + + const stream = new EventEmitter(); + const outputChunks: string[] = []; + + const promise = new Promise((resolve, reject) => { + service.streamToOpenAI(stream, input.model).subscribe({ + next: (chunk) => outputChunks.push(chunk), + error: reject, + complete: resolve, + }); + }); + + stream.emit('data', Buffer.from(`data: ${JSON.stringify(input.upstream)}\n`)); + stream.emit('end'); + + await promise; + + const output = outputChunks.join(''); + for (const token of expected.contains) { + expect(output).toContain(token); + } + }); +}); diff --git a/src/tests/unit/proxy-retry-mock.test.ts b/src/tests/unit/proxy-retry-mock.test.ts index fa14983..1553321 100644 --- a/src/tests/unit/proxy-retry-mock.test.ts +++ b/src/tests/unit/proxy-retry-mock.test.ts @@ -1,10 +1,14 @@ -import { describe, it, expect, vi } from 'vitest'; +import { beforeEach, describe, it, expect, vi } from 'vitest'; import { EventEmitter } from 'events'; import { ProxyService } from '../../server/modules/proxy/proxy.service'; import { Observable } from 'rxjs'; // Mock dependencies -const mockTokenManager = { getNextToken: vi.fn(), markAsRateLimited: vi.fn() }; +const mockTokenManager = { + getNextToken: vi.fn(), + markAsRateLimited: vi.fn(), + markAsForbidden: vi.fn(), +}; const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() }; // Subclass to access private method @@ -17,9 +21,77 @@ class TestableProxyService extends ProxyService { // Access private method using type assertion return (this as any).processAnthropicInternalStream(stream, model); } + + public testPassthroughStream(stream: any): Observable { + return (this as any).passthroughSseStream(stream); + } + + public testModelHeaders(model: string): Record { + return (this as any).buildModelSpecificHeaders(model); + } +} + +function createToken(id: string = 'acc-1') { + return { + id, + email: `${id}@test.com`, + token: { + access_token: 'token', + refresh_token: 'refresh', + expires_in: 3600, + expiry_timestamp: Math.floor(Date.now() / 1000) + 3600, + project_id: 'project-1', + session_id: 'session-1', + upstream_proxy_url: undefined, + }, + }; } describe('ProxyService Empty Stream Retry Logic', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('classifies retry matrix consistently', () => { + const service = new TestableProxyService(); + const classify = (message: string) => (service as any).classifyUpstreamError(message); + + expect(classify('401 unauthorized token')).toEqual({ + retry: true, + markAsForbidden: true, + markAsRateLimited: false, + }); + expect(classify('403 permission_denied')).toEqual({ + retry: true, + markAsForbidden: true, + markAsRateLimited: false, + }); + expect(classify('429 quota exceeded')).toEqual({ + retry: true, + markAsForbidden: false, + markAsRateLimited: true, + }); + expect(classify('500 internal error')).toEqual({ + retry: true, + markAsForbidden: false, + markAsRateLimited: false, + }); + expect(classify('400 invalid argument')).toEqual({ + retry: false, + markAsForbidden: false, + markAsRateLimited: false, + }); + }); + + it('builds Claude-specific beta headers consistently', () => { + const service = new TestableProxyService(); + const claudeHeaders = service.testModelHeaders('claude-sonnet-4-5'); + const geminiHeaders = service.testModelHeaders('gemini-2.5-flash'); + + expect(claudeHeaders['anthropic-beta']).toContain('claude-code-20250219'); + expect(geminiHeaders).toEqual({}); + }); + it('should emit error when stream ends without data', async () => { const service = new TestableProxyService(); const stream = new EventEmitter(); @@ -89,4 +161,454 @@ describe('ProxyService Empty Stream Retry Logic', () => { // Actually our mapper might produce "message_start", "content_block_start" etc. // We just care that it didn't error with "Empty response stream" }); + + it('falls back to stream aggregation when non-stream response is empty', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + + mockGeminiClient.generateInternal.mockResolvedValueOnce({ candidates: [] }); + mockGeminiClient.streamGenerateInternal.mockResolvedValueOnce(stream); + + const promise = (service as any).generateInternalWithStreamFallback( + { model: 'gemini-2.5-flash' }, + 'token', + undefined, + ); + + setTimeout(() => { + const payload = JSON.stringify({ + candidates: [ + { + content: { parts: [{ text: 'fallback text' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { totalTokenCount: 5 }, + }); + stream.emit('data', Buffer.from(`data: ${payload}\n\n`)); + stream.emit('end'); + }, 10); + + const result = await promise; + expect(mockGeminiClient.streamGenerateInternal).toHaveBeenCalledOnce(); + expect(result.candidates[0].content.parts[0].text).toBe('fallback text'); + expect(result.candidates[0].finishReason).toBe('STOP'); + }); + + it('injects Claude beta headers when handling Gemini-compatible Claude models', async () => { + const service = new TestableProxyService(); + mockTokenManager.getNextToken.mockResolvedValue(createToken()); + mockGeminiClient.generateInternal.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'ok' }] } }], + }); + + await service.handleGeminiGenerateContent('models/claude-sonnet-4-5', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + const headers = mockGeminiClient.generateInternal.mock.calls[0][3]; + expect(headers['anthropic-beta']).toContain('claude-code-20250219'); + }); + + it('tolerates malformed partial chunks and still completes Anthropic stream', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + const resultObservable = service.testProcessStream(stream); + + let completed = false; + let errored = false; + + const done = new Promise((resolve) => { + resultObservable.subscribe({ + next: () => {}, + error: () => { + errored = true; + resolve(); + }, + complete: () => { + completed = true; + resolve(); + }, + }); + }); + + setTimeout(() => { + stream.emit('data', Buffer.from('data: {"invalid_json":\n\n')); + const validPayload = JSON.stringify({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + }); + stream.emit('data', Buffer.from(`data: ${validPayload}\n\n`)); + stream.emit('end'); + }, 10); + + await done; + + expect(errored).toBe(false); + expect(completed).toBe(true); + }); + + it('raises error for empty Gemini passthrough stream', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + + const observable = service.testPassthroughStream(stream); + let errorMessage = ''; + + const done = new Promise((resolve) => { + observable.subscribe({ + next: () => {}, + error: (error: Error) => { + errorMessage = error.message; + resolve(); + }, + complete: () => resolve(), + }); + }); + + setTimeout(() => stream.emit('end'), 10); + await done; + + expect(errorMessage).toBe('Empty response stream'); + }); + + it('propagates Anthropic stream interruption errors', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + const observable = service.testProcessStream(stream); + let errorMessage = ''; + + const done = new Promise((resolve) => { + observable.subscribe({ + next: () => {}, + error: (error: Error) => { + errorMessage = error.message; + resolve(); + }, + complete: () => resolve(), + }); + }); + + setTimeout(() => stream.emit('error', new Error('upstream interrupted')), 10); + await done; + + expect(errorMessage).toBe('upstream interrupted'); + }); + + it('propagates Gemini passthrough interruption errors', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + const observable = service.testPassthroughStream(stream); + let errorMessage = ''; + + const done = new Promise((resolve) => { + observable.subscribe({ + next: () => {}, + error: (error: Error) => { + errorMessage = error.message; + resolve(); + }, + complete: () => resolve(), + }); + }); + + setTimeout(() => stream.emit('error', new Error('connection reset by peer')), 10); + await done; + + expect(errorMessage).toBe('connection reset by peer'); + }); + + it('retries OpenAI flow with the same error classification matrix', async () => { + const service = new TestableProxyService(); + const token1 = createToken('acc-1'); + const token2 = createToken('acc-2'); + mockTokenManager.getNextToken.mockResolvedValueOnce(token1).mockResolvedValueOnce(token2); + mockGeminiClient.generateInternal + .mockRejectedValueOnce(new Error('429 quota exceeded')) + .mockResolvedValueOnce({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + usageMetadata: { totalTokenCount: 5 }, + }); + + const result = await service.handleChatCompletions({ + model: 'gpt-4o', + stream: false, + messages: [{ role: 'user', content: 'hello' }], + } as any); + + expect(mockTokenManager.getNextToken).toHaveBeenCalledTimes(2); + expect(mockTokenManager.markAsRateLimited).toHaveBeenCalledWith('acc-1'); + expect((result as any).choices?.[0]?.message?.content).toBeDefined(); + }); + + it('retries Anthropic flow with the same error classification matrix', async () => { + const service = new TestableProxyService(); + const token1 = createToken('acc-1'); + const token2 = createToken('acc-2'); + mockTokenManager.getNextToken.mockResolvedValueOnce(token1).mockResolvedValueOnce(token2); + mockGeminiClient.generateInternal + .mockRejectedValueOnce(new Error('429 rate limit exceeded')) + .mockResolvedValueOnce({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + usageMetadata: { totalTokenCount: 5 }, + }); + + const result = await service.handleAnthropicMessages({ + model: 'claude-sonnet-4-5', + stream: false, + max_tokens: 256, + messages: [{ role: 'user', content: 'hello' }], + } as any); + + expect(mockTokenManager.getNextToken).toHaveBeenCalledTimes(2); + expect(mockTokenManager.markAsRateLimited).toHaveBeenCalledWith('acc-1'); + expect((result as any).type).toBe('message'); + }); + + it('retries Gemini flow with the same error classification matrix', async () => { + const service = new TestableProxyService(); + const token1 = createToken('acc-1'); + const token2 = createToken('acc-2'); + mockTokenManager.getNextToken.mockResolvedValueOnce(token1).mockResolvedValueOnce(token2); + mockGeminiClient.generateInternal + .mockRejectedValueOnce(new Error('429 quota exceeded')) + .mockResolvedValueOnce({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + usageMetadata: { totalTokenCount: 5 }, + }); + + const result = await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + expect(mockTokenManager.getNextToken).toHaveBeenCalledTimes(2); + expect(mockTokenManager.markAsRateLimited).toHaveBeenCalledWith('acc-1'); + expect((result as any).candidates?.[0]?.content?.parts?.[0]?.text).toBe('ok'); + }); + + it('does not include sessionId in Gemini internal generate payload', async () => { + const service = new TestableProxyService(); + mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1')); + mockGeminiClient.generateInternal.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + usageMetadata: { totalTokenCount: 5 }, + }); + + await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + const internalPayload = mockGeminiClient.generateInternal.mock.calls[0][0]; + expect(internalPayload).not.toHaveProperty('sessionId'); + }); + + it('strips non-parity Gemini usage metadata fields', async () => { + const service = new TestableProxyService(); + mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1')); + mockGeminiClient.generateInternal.mockResolvedValue({ + candidates: [ + { + content: { role: 'model', parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + promptTokenCount: 1, + candidatesTokenCount: 2, + totalTokenCount: 3, + thoughtsTokenCount: 4, + }, + responseId: 'resp_123', + createTime: '2026-02-10T00:00:00.000Z', + }); + + const result = await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + expect((result as any).usageMetadata).toEqual({ + promptTokenCount: 1, + candidatesTokenCount: 2, + totalTokenCount: 3, + }); + expect((result as any).usageMetadata.thoughtsTokenCount).toBeUndefined(); + }); + + it('retries Gemini generate-content without project when project context is invalid', async () => { + const service = new TestableProxyService(); + mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1')); + mockGeminiClient.generateInternal + .mockRejectedValueOnce( + new Error( + 'You are currently configured to use a Google Cloud Project but lack a Gemini Code Assist license. (#3501)', + ), + ) + .mockResolvedValueOnce({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + usageMetadata: { totalTokenCount: 5 }, + }); + + const result = await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + expect(mockTokenManager.getNextToken).toHaveBeenCalledTimes(1); + expect(mockGeminiClient.generateInternal).toHaveBeenCalledTimes(2); + expect(mockGeminiClient.generateInternal.mock.calls[0][0].project).toBe('project-1'); + expect(mockGeminiClient.generateInternal.mock.calls[1][0].project).toBe(''); + expect((result as any).candidates?.[0]?.content?.parts?.[0]?.text).toBe('ok'); + }); + + it('uses generate-content requestType for Gemini stream internal payload', async () => { + const service = new TestableProxyService(); + mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1')); + mockGeminiClient.streamGenerateInternal.mockResolvedValue(new EventEmitter()); + + await service.handleGeminiStreamGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + const internalPayload = mockGeminiClient.streamGenerateInternal.mock.calls[0][0]; + expect(internalPayload.requestType).toBe('generate-content'); + }); +}); + +describe('ProxyService Protocol Parity Fixtures', () => { + it('maps OpenAI request to Anthropic request with tools and tool result', () => { + const service = new TestableProxyService(); + + const openaiRequest = { + model: 'claude-sonnet-4-5', + stream: false, + temperature: 0.2, + max_tokens: 512, + tools: [ + { + type: 'function', + function: { + name: 'search_docs', + description: 'Search docs', + parameters: { + type: 'object', + properties: { + query: { type: 'string' }, + }, + }, + }, + }, + ], + messages: [ + { role: 'system', content: 'You are a precise assistant.' }, + { role: 'user', content: [{ type: 'text', text: 'Find API key docs' }] }, + { + role: 'assistant', + content: 'Calling search tool', + tool_calls: [ + { + id: 'call_1', + type: 'function', + function: { + name: 'search_docs', + arguments: '{"query":"api key"}', + }, + }, + ], + }, + { + role: 'tool', + tool_call_id: 'call_1', + name: 'search_docs', + content: 'Found 3 results', + }, + ], + }; + + const anthropicRequest = (service as any).convertOpenAIToClaude(openaiRequest); + + expect(anthropicRequest.system).toContain('You are a precise assistant.'); + expect(anthropicRequest.tools?.[0]?.name).toBe('search_docs'); + expect(anthropicRequest.messages.length).toBe(3); + + const assistantMessage = anthropicRequest.messages[1]; + expect(Array.isArray(assistantMessage.content)).toBe(true); + expect(assistantMessage.content.some((block: any) => block.type === 'tool_use')).toBe(true); + + const toolResultMessage = anthropicRequest.messages[2]; + expect(toolResultMessage.role).toBe('user'); + expect(Array.isArray(toolResultMessage.content)).toBe(true); + expect(toolResultMessage.content[0].type).toBe('tool_result'); + }); + + it('maps Anthropic response to OpenAI response with reasoning and tool_calls', () => { + const service = new TestableProxyService(); + + const anthropicResponse = { + content: [ + { type: 'thinking', thinking: 'Need to call tool first.' }, + { + type: 'tool_use', + id: 'tool-1', + name: 'search_docs', + input: { query: 'api key' }, + }, + { type: 'text', text: 'Here are the docs.' }, + ], + stop_reason: 'tool_use', + usage: { + input_tokens: 20, + output_tokens: 30, + }, + }; + + const openaiResponse = (service as any).convertClaudeToOpenAIResponse( + anthropicResponse, + 'gpt-4o-mini', + ); + + expect(openaiResponse.model).toBe('gpt-4o-mini'); + expect(openaiResponse.choices[0].message.role).toBe('assistant'); + expect(openaiResponse.choices[0].message.reasoning_content).toContain('Need to call tool'); + expect(openaiResponse.choices[0].message.tool_calls?.length).toBe(1); + expect(openaiResponse.choices[0].finish_reason).toBe('tool_calls'); + }); + + it('converts internal SSE stream into OpenAI SSE chunks', async () => { + const service = new TestableProxyService(); + const stream = new EventEmitter(); + const observable = (service as any).processStreamResponse(stream, 'gpt-4o-mini'); + + const chunks: string[] = []; + await new Promise((resolve, reject) => { + observable.subscribe({ + next: (chunk: string) => { + chunks.push(chunk); + }, + error: reject, + complete: resolve, + }); + + const payload = JSON.stringify({ + candidates: [ + { + content: { + parts: [ + { thought: true, text: 'reasoning text' }, + { functionCall: { id: 'fc1', name: 'search_docs', args: { query: 'api key' } } }, + { text: 'final answer' }, + ], + }, + finishReason: 'STOP', + }, + ], + }); + + stream.emit('data', Buffer.from(`data: ${payload}\n`)); + stream.emit('end'); + }); + + const output = chunks.join(''); + expect(output).toContain('"reasoning_content":"reasoning text"'); + expect(output).toContain('"tool_calls"'); + expect(output).toContain('"content":"final answer"'); + expect(output).toContain('data: [DONE]'); + }); }); diff --git a/src/types/cloudAccount.ts b/src/types/cloudAccount.ts index 1af5ddc..71b26d7 100644 --- a/src/types/cloudAccount.ts +++ b/src/types/cloudAccount.ts @@ -10,6 +10,7 @@ export interface CloudTokenData { email?: string; project_id?: string; session_id?: string; + upstream_proxy_url?: string; } export interface CloudQuotaData { @@ -48,6 +49,7 @@ export const CloudTokenDataSchema = z.object({ email: z.string().optional(), project_id: z.string().optional(), session_id: z.string().optional(), + upstream_proxy_url: z.string().optional(), }); export const CloudQuotaDataSchema = z.object({ diff --git a/src/types/config.ts b/src/types/config.ts index 46cf974..150a55c 100644 --- a/src/types/config.ts +++ b/src/types/config.ts @@ -10,6 +10,8 @@ export const ProxyConfigSchema = z.object({ port: z.number(), // 监听端口 api_key: z.string(), // API 密钥 (Also used for Admin Mode auth) auto_start: z.boolean(), // 是否自动启动 + backend_canary_enabled: z.boolean().default(true), + custom_mapping: z.record(z.string(), z.string()).default({}), anthropic_mapping: z.record(z.string(), z.string()), // 映射表 request_timeout: z.number().default(120), // 超时秒数 upstream_proxy: UpstreamProxyConfigSchema, @@ -49,6 +51,8 @@ export const DEFAULT_APP_CONFIG: AppConfig = { port: 8045, api_key: '', // Generated dynamically if default needed auto_start: false, + backend_canary_enabled: true, + custom_mapping: {}, anthropic_mapping: {}, request_timeout: 120, upstream_proxy: {