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
This commit is contained in:
23 files changed
+4408
-654
No files matched your search
@@ -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.
|
||||
|
||||
@@ -52,7 +52,7 @@ export function transformClaudeRequestIn(
|
||||
const toolIdToName = new Map<string, string>();
|
||||
|
||||
// 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;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ import { logger } from '../../utils/logger';
|
||||
const CLAUDE_TO_GEMINI: Record<string, string> = {
|
||||
// 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<string, string> = {
|
||||
'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<string, string> = {
|
||||
'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, string> = {}): string[] {
|
||||
const modelIds = new Set<string>();
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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: 'По умолчанию',
|
||||
|
||||
@@ -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');
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
export type RequestHeaderValue = string | string[] | undefined;
|
||||
|
||||
export type RequestHeaders = Record<string, RequestHeaderValue>;
|
||||
|
||||
export function hasConfiguredApiKey(apiKey: string | undefined): apiKey is string {
|
||||
return typeof apiKey === 'string' && apiKey.trim() !== '';
|
||||
}
|
||||
|
||||
export function extractApiKeyToken(headers: RequestHeaders): string | null {
|
||||
const authorizationHeader = readHeaderValue(headers['authorization']);
|
||||
if (authorizationHeader) {
|
||||
const [scheme, token, ...rest] = authorizationHeader.trim().split(/\s+/);
|
||||
if (rest.length === 0 && token && scheme?.toLowerCase() === 'bearer') {
|
||||
return token;
|
||||
}
|
||||
}
|
||||
|
||||
const xApiKeyHeader = readHeaderValue(headers['x-api-key']);
|
||||
if (xApiKeyHeader) {
|
||||
return xApiKeyHeader;
|
||||
}
|
||||
|
||||
const xGoogApiKeyHeader = readHeaderValue(headers['x-goog-api-key']);
|
||||
if (xGoogApiKeyHeader) {
|
||||
return xGoogApiKeyHeader;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
function readHeaderValue(headerValue: RequestHeaderValue): string | null {
|
||||
if (typeof headerValue === 'string') {
|
||||
const trimmedValue = headerValue.trim();
|
||||
return trimmedValue !== '' ? trimmedValue : null;
|
||||
}
|
||||
|
||||
if (Array.isArray(headerValue)) {
|
||||
for (const value of headerValue) {
|
||||
if (typeof value === 'string') {
|
||||
const trimmedValue = value.trim();
|
||||
if (trimmedValue !== '') {
|
||||
return trimmedValue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
@@ -1,16 +1,29 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import axios from 'axios';
|
||||
import axios, { AxiosProxyConfig, AxiosRequestConfig, AxiosResponse } from 'axios';
|
||||
import { isNil } from 'lodash-es';
|
||||
import { GeminiRequest, GeminiResponse } from '../interfaces/request-interfaces';
|
||||
import { GeminiInternalRequest } from '../../../../lib/antigravity/types';
|
||||
import { getServerConfig } from '../../../server-config';
|
||||
|
||||
@Injectable()
|
||||
export class GeminiClient {
|
||||
private readonly logger = new Logger(GeminiClient.name);
|
||||
private readonly defaultRequestUserAgent = 'antigravity/1.11.9 windows/amd64';
|
||||
// Default to v1beta for most features
|
||||
private readonly baseUrl = 'https://generativelanguage.googleapis.com/v1beta';
|
||||
private readonly defaultInternalBaseUrls = [
|
||||
'https://cloudcode-pa.googleapis.com/v1internal',
|
||||
'https://daily-cloudcode-pa.googleapis.com/v1internal',
|
||||
];
|
||||
|
||||
async streamGenerate(model: string, content: GeminiRequest, accessToken: string): Promise<any> {
|
||||
async streamGenerate(
|
||||
model: string,
|
||||
content: GeminiRequest,
|
||||
accessToken: string,
|
||||
upstreamProxyUrl?: string,
|
||||
): Promise<NodeJS.ReadableStream> {
|
||||
const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`;
|
||||
const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
|
||||
|
||||
try {
|
||||
const response = await axios.post(url, content, {
|
||||
@@ -20,15 +33,15 @@ export class GeminiClient {
|
||||
},
|
||||
responseType: 'stream',
|
||||
timeout: 60000,
|
||||
proxy: axiosProxy,
|
||||
});
|
||||
return response.data;
|
||||
} catch (error) {
|
||||
if (axios.isAxiosError(error)) {
|
||||
this.logger.error(`Gemini Stream API Error: ${error.message}`);
|
||||
this.logger.error(`Gemini stream request failed: ${error.message}`);
|
||||
throw new Error(error.response?.data?.error?.message || error.message);
|
||||
}
|
||||
// Re-throw as clean Error to avoid circular reference issues
|
||||
throw error instanceof Error ? new Error(error.message) : new Error(String(error));
|
||||
this.throwAsCleanError(error);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,8 +49,10 @@ export class GeminiClient {
|
||||
model: string,
|
||||
content: GeminiRequest,
|
||||
accessToken: string,
|
||||
upstreamProxyUrl?: string,
|
||||
): Promise<GeminiResponse> {
|
||||
const url = `${this.baseUrl}/models/${model}:generateContent`;
|
||||
const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
|
||||
|
||||
try {
|
||||
const response = await axios.post<GeminiResponse>(url, content, {
|
||||
@@ -46,68 +61,221 @@ export class GeminiClient {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
timeout: 60000, // 60s timeout
|
||||
proxy: axiosProxy,
|
||||
});
|
||||
return response.data;
|
||||
} catch (error) {
|
||||
if (axios.isAxiosError(error)) {
|
||||
this.logger.error(
|
||||
`Gemini API Error: ${error.message} - ${this.safeStringify(error.response?.data)}`,
|
||||
`Gemini request failed: ${error.message} - ${this.safeStringify(error.response?.data)}`,
|
||||
);
|
||||
throw new Error(error.response?.data?.error?.message || error.message);
|
||||
}
|
||||
// Re-throw as clean Error to avoid circular reference issues
|
||||
throw error instanceof Error ? new Error(error.message) : new Error(String(error));
|
||||
this.throwAsCleanError(error);
|
||||
}
|
||||
}
|
||||
|
||||
// --- Antigravity Internal API Support ---
|
||||
// --- Internal Gateway API Support ---
|
||||
|
||||
private readonly internalBaseUrl = 'https://cloudcode-pa.googleapis.com/v1internal';
|
||||
|
||||
async streamGenerateInternal(body: GeminiInternalRequest, accessToken: string): Promise<any> {
|
||||
const url = `${this.internalBaseUrl}:streamGenerateContent?alt=sse`;
|
||||
try {
|
||||
const response = await axios.post(url, body, {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
'Content-Type': 'application/json',
|
||||
'User-Agent': 'antigravity/1.11.9 windows/amd64',
|
||||
},
|
||||
async streamGenerateInternal(
|
||||
body: GeminiInternalRequest,
|
||||
accessToken: string,
|
||||
upstreamProxyUrl?: string,
|
||||
extraHeaders?: Record<string, string>,
|
||||
): Promise<NodeJS.ReadableStream> {
|
||||
const response = await this.requestWithEndpointFailover<NodeJS.ReadableStream>(
|
||||
':streamGenerateContent?alt=sse',
|
||||
body,
|
||||
accessToken,
|
||||
upstreamProxyUrl,
|
||||
{
|
||||
responseType: 'stream',
|
||||
timeout: 60000,
|
||||
});
|
||||
return response.data;
|
||||
} catch (error) {
|
||||
this.handleAxiosError(error, 'StreamInternal');
|
||||
}
|
||||
},
|
||||
'stream-generate',
|
||||
extraHeaders,
|
||||
);
|
||||
|
||||
return response.data;
|
||||
}
|
||||
|
||||
async generateInternal(body: GeminiInternalRequest, accessToken: string): Promise<any> {
|
||||
const url = `${this.internalBaseUrl}:generateContent`;
|
||||
try {
|
||||
const response = await axios.post(url, body, {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
'Content-Type': 'application/json',
|
||||
'User-Agent': 'antigravity/1.11.9 windows/amd64',
|
||||
},
|
||||
timeout: 60000,
|
||||
});
|
||||
// v1internal API wraps the response in a 'response' field, unwrap it
|
||||
return response.data.response || response.data;
|
||||
} catch (error) {
|
||||
this.handleAxiosError(error, 'GenerateInternal');
|
||||
async generateInternal(
|
||||
body: GeminiInternalRequest,
|
||||
accessToken: string,
|
||||
upstreamProxyUrl?: string,
|
||||
extraHeaders?: Record<string, string>,
|
||||
): Promise<GeminiResponse> {
|
||||
const response = await this.requestWithEndpointFailover<
|
||||
GeminiResponse | { response: GeminiResponse }
|
||||
>(
|
||||
':generateContent',
|
||||
body,
|
||||
accessToken,
|
||||
upstreamProxyUrl,
|
||||
{},
|
||||
'generate-content',
|
||||
extraHeaders,
|
||||
);
|
||||
const payload = response.data;
|
||||
if (payload && typeof payload === 'object' && 'response' in payload) {
|
||||
return payload.response;
|
||||
}
|
||||
return payload;
|
||||
}
|
||||
|
||||
private handleAxiosError(error: any, context: string) {
|
||||
private getInternalBaseUrls(): string[] {
|
||||
const fromEnv =
|
||||
process.env.PROXY_INTERNAL_BASE_URLS ?? process.env.ANTIGRAVITY_INTERNAL_BASE_URLS;
|
||||
const configuredBaseUrls = fromEnv
|
||||
?.split(',')
|
||||
.map((item) => item.trim())
|
||||
.filter(Boolean);
|
||||
|
||||
if (configuredBaseUrls && configuredBaseUrls.length > 0) {
|
||||
return configuredBaseUrls.map((url) => url.replace(/\/+$/, ''));
|
||||
}
|
||||
|
||||
return this.defaultInternalBaseUrls.map((url) => url.replace(/\/+$/, ''));
|
||||
}
|
||||
|
||||
private getInternalTimeoutMs(): number {
|
||||
const config = getServerConfig();
|
||||
const timeoutSeconds = config?.request_timeout ?? 60;
|
||||
return Math.max(1, timeoutSeconds) * 1000;
|
||||
}
|
||||
|
||||
private shouldSwitchEndpointOnError(error: unknown): boolean {
|
||||
if (!axios.isAxiosError(error)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!error.response) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const status = error.response.status;
|
||||
|
||||
// Permanent auth errors should fail fast for current token.
|
||||
if (status === 401 || status === 403) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return status === 408 || status === 429 || status >= 500;
|
||||
}
|
||||
|
||||
private async requestWithEndpointFailover<T>(
|
||||
path: string,
|
||||
body: GeminiInternalRequest,
|
||||
accessToken: string,
|
||||
upstreamProxyUrl: string | undefined,
|
||||
config: AxiosRequestConfig,
|
||||
operation: string,
|
||||
extraHeaders?: Record<string, string>,
|
||||
): Promise<AxiosResponse<T>> {
|
||||
const baseUrls = this.getInternalBaseUrls();
|
||||
const timeout = this.getInternalTimeoutMs();
|
||||
const requestUserAgent = process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent;
|
||||
const axiosProxy = this.resolveAxiosProxy(upstreamProxyUrl);
|
||||
let lastError: unknown = null;
|
||||
|
||||
for (let index = 0; index < baseUrls.length; index++) {
|
||||
const baseUrl = baseUrls[index];
|
||||
const url = `${baseUrl}${path}`;
|
||||
|
||||
try {
|
||||
return await axios.post<T>(url, body, {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
'Content-Type': 'application/json',
|
||||
'User-Agent': requestUserAgent,
|
||||
...(extraHeaders ?? {}),
|
||||
},
|
||||
timeout,
|
||||
proxy: axiosProxy,
|
||||
...config,
|
||||
});
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
const hasNextEndpoint = index < baseUrls.length - 1;
|
||||
|
||||
if (!hasNextEndpoint || !this.shouldSwitchEndpointOnError(error)) {
|
||||
this.handleAxiosError(error, operation);
|
||||
}
|
||||
|
||||
this.logger.warn(
|
||||
`[${operation}] request failed at ${baseUrl}; trying next endpoint (${index + 2}/${
|
||||
baseUrls.length
|
||||
}).`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
this.handleAxiosError(lastError, operation);
|
||||
}
|
||||
|
||||
private handleAxiosError(error: unknown, operation: string): never {
|
||||
if (axios.isAxiosError(error)) {
|
||||
const responseData = error.response?.data;
|
||||
const upstreamMessage = this.extractAxiosErrorMessage(responseData);
|
||||
this.logger.error(
|
||||
`Gemini ${context} API Error: ${error.message} - ${this.safeStringify(error.response?.data)}`,
|
||||
`[${operation}] upstream request error: ${error.message} - ${this.safeStringify(responseData)}`,
|
||||
);
|
||||
throw new Error(error.response?.data?.error?.message || error.message);
|
||||
throw new Error(upstreamMessage || error.message);
|
||||
}
|
||||
// Re-throw as clean Error to avoid circular reference issues
|
||||
this.throwAsCleanError(error);
|
||||
}
|
||||
|
||||
private extractAxiosErrorMessage(responseData: unknown): string | null {
|
||||
if (!responseData || typeof responseData !== 'object') {
|
||||
return null;
|
||||
}
|
||||
|
||||
const errorRecord = (responseData as { error?: unknown }).error;
|
||||
if (!errorRecord || typeof errorRecord !== 'object') {
|
||||
return null;
|
||||
}
|
||||
|
||||
const message = (errorRecord as { message?: unknown }).message;
|
||||
return typeof message === 'string' && message.trim() !== '' ? message : null;
|
||||
}
|
||||
|
||||
private resolveAxiosProxy(upstreamProxyUrl?: string): AxiosProxyConfig | false | undefined {
|
||||
const config = getServerConfig();
|
||||
const configuredProxyUrl =
|
||||
upstreamProxyUrl ||
|
||||
(config?.upstream_proxy?.enabled && config.upstream_proxy.url
|
||||
? config.upstream_proxy.url
|
||||
: '');
|
||||
|
||||
if (!configuredProxyUrl) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed = new URL(configuredProxyUrl);
|
||||
const port = parsed.port ? Number(parsed.port) : parsed.protocol === 'https:' ? 443 : 80;
|
||||
|
||||
const proxyConfig: AxiosProxyConfig = {
|
||||
protocol: parsed.protocol.replace(':', ''),
|
||||
host: parsed.hostname,
|
||||
port,
|
||||
};
|
||||
|
||||
if (parsed.username || parsed.password) {
|
||||
proxyConfig.auth = {
|
||||
username: decodeURIComponent(parsed.username),
|
||||
password: decodeURIComponent(parsed.password),
|
||||
};
|
||||
}
|
||||
|
||||
return proxyConfig;
|
||||
} catch {
|
||||
this.logger.warn(`Upstream proxy URL is invalid: ${configuredProxyUrl}`);
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private throwAsCleanError(error: unknown): never {
|
||||
// Re-throw as clean Error to avoid circular reference issues.
|
||||
throw error instanceof Error ? new Error(error.message) : new Error(String(error));
|
||||
}
|
||||
|
||||
@@ -115,7 +283,7 @@ export class GeminiClient {
|
||||
* Safely stringify an object, handling circular references
|
||||
*/
|
||||
private safeStringify(obj: unknown): string {
|
||||
if (obj === undefined || obj === null) {
|
||||
if (isNil(obj)) {
|
||||
return String(obj);
|
||||
}
|
||||
try {
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
import { Body, Controller, Get, HttpStatus, Inject, Param, Post, Res, UseGuards } from '@nestjs/common';
|
||||
import { FastifyReply } from 'fastify';
|
||||
import { isEmpty } from 'lodash-es';
|
||||
import { Observable } from 'rxjs';
|
||||
|
||||
import { ProxyGuard } from './proxy.guard';
|
||||
import { ProxyService } from './proxy.service';
|
||||
import { GeminiRequest, GeminiResponse } from './interfaces/request-interfaces';
|
||||
import { getServerConfig } from '../../server-config';
|
||||
import { getAllDynamicModels } from '../../../lib/antigravity/ModelMapping';
|
||||
|
||||
type GeminiModelMetadata = {
|
||||
name: string;
|
||||
displayName: string;
|
||||
description: string;
|
||||
inputTokenLimit: number;
|
||||
outputTokenLimit: number;
|
||||
supportedGenerationMethods: string[];
|
||||
temperature: number;
|
||||
topK: number;
|
||||
topP: number;
|
||||
version: string;
|
||||
};
|
||||
|
||||
@Controller('v1beta')
|
||||
@UseGuards(ProxyGuard)
|
||||
export class GeminiController {
|
||||
constructor(@Inject(ProxyService) private readonly proxyService: ProxyService) {}
|
||||
|
||||
@Get('models')
|
||||
listModels(@Res() res: FastifyReply) {
|
||||
const models = this.buildGeminiModelList();
|
||||
res.status(HttpStatus.OK).send({
|
||||
models,
|
||||
});
|
||||
}
|
||||
|
||||
@Get('models/:model')
|
||||
getModel(@Param('model') model: string, @Res() res: FastifyReply) {
|
||||
const targetName = model.startsWith('models/') ? model : `models/${model}`;
|
||||
const matched = this.buildGeminiModelList().find((item) => item.name === targetName);
|
||||
|
||||
if (matched) {
|
||||
res.status(HttpStatus.OK).send({
|
||||
name: matched.name,
|
||||
displayName: matched.displayName,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
res.status(HttpStatus.OK).send({
|
||||
name: targetName,
|
||||
displayName: targetName.replace(/^models\//, ''),
|
||||
});
|
||||
}
|
||||
|
||||
@Post('models/:modelAction')
|
||||
async modelAction(
|
||||
@Param('modelAction') modelAction: string,
|
||||
@Body() body: GeminiRequest,
|
||||
@Res() res: FastifyReply,
|
||||
) {
|
||||
const parsed = this.parseCombinedModelAction(modelAction);
|
||||
if (!parsed) {
|
||||
res.status(HttpStatus.BAD_REQUEST).send({
|
||||
error: {
|
||||
code: HttpStatus.BAD_REQUEST,
|
||||
message: 'Model action format is invalid',
|
||||
status: 'INVALID_ARGUMENT',
|
||||
},
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this.dispatchAction(parsed.model, parsed.action, body, res);
|
||||
}
|
||||
|
||||
@Post('models/:model/countTokens')
|
||||
async countTokens(
|
||||
@Param('model') model: string,
|
||||
@Body() body: GeminiRequest,
|
||||
@Res() res: FastifyReply,
|
||||
) {
|
||||
await this.dispatchAction(`models/${model}`, 'countTokens', body, res);
|
||||
}
|
||||
|
||||
private async dispatchAction(
|
||||
model: string,
|
||||
action: string,
|
||||
body: GeminiRequest,
|
||||
res: FastifyReply,
|
||||
): Promise<void> {
|
||||
if (action === 'countTokens') {
|
||||
res.status(HttpStatus.OK).send({
|
||||
totalTokens: 0,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
if (action === 'streamGenerateContent') {
|
||||
const stream = await this.proxyService.handleGeminiStreamGenerateContent(model, body);
|
||||
if (stream instanceof Observable) {
|
||||
this.sendObservableSse(res, stream);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (action === 'generateContent') {
|
||||
const result = await this.proxyService.handleGeminiGenerateContent(model, body);
|
||||
res.status(HttpStatus.OK).send(this.normalizeGeminiGenerateResponse(result));
|
||||
return;
|
||||
}
|
||||
|
||||
res.status(HttpStatus.BAD_REQUEST).send({
|
||||
error: {
|
||||
code: HttpStatus.BAD_REQUEST,
|
||||
message: `Unsupported model action: ${action}`,
|
||||
status: 'INVALID_ARGUMENT',
|
||||
},
|
||||
});
|
||||
} catch (error) {
|
||||
const message = error instanceof Error ? error.message : 'Internal Server Error';
|
||||
res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({
|
||||
error: {
|
||||
code: HttpStatus.INTERNAL_SERVER_ERROR,
|
||||
message,
|
||||
status: 'INTERNAL',
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private parseCombinedModelAction(
|
||||
modelAction: string,
|
||||
): {
|
||||
model: string;
|
||||
action: string;
|
||||
} | null {
|
||||
const colonIndex = modelAction.lastIndexOf(':');
|
||||
if (colonIndex <= 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const model = modelAction.slice(0, colonIndex).trim();
|
||||
const action = modelAction.slice(colonIndex + 1).trim();
|
||||
if (!model || !action) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const prefixedModel = model.startsWith('models/') ? model : `models/${model}`;
|
||||
return {
|
||||
model: prefixedModel,
|
||||
action,
|
||||
};
|
||||
}
|
||||
|
||||
private buildGeminiModelList(): GeminiModelMetadata[] {
|
||||
const config = getServerConfig();
|
||||
const dynamicModelIds = getAllDynamicModels(config?.custom_mapping ?? {});
|
||||
|
||||
return dynamicModelIds.map((id) => this.toGeminiModelMetadata(`models/${id}`));
|
||||
}
|
||||
|
||||
private toGeminiModelMetadata(modelName: string): GeminiModelMetadata {
|
||||
const displayName = modelName.replace(/^models\//, '');
|
||||
return {
|
||||
name: modelName,
|
||||
displayName,
|
||||
description: '',
|
||||
inputTokenLimit: 128000,
|
||||
outputTokenLimit: 8192,
|
||||
supportedGenerationMethods: ['generateContent', 'countTokens'],
|
||||
temperature: 1,
|
||||
topK: 64,
|
||||
topP: 0.95,
|
||||
version: '001',
|
||||
};
|
||||
}
|
||||
|
||||
private normalizeGeminiGenerateResponse(response: GeminiResponse): GeminiResponse {
|
||||
const candidates = (response.candidates ?? []).map((candidate, index) => ({
|
||||
content: candidate.content,
|
||||
finishReason: candidate.finishReason,
|
||||
index: typeof candidate.index === 'number' ? candidate.index : index,
|
||||
}));
|
||||
|
||||
const normalized: GeminiResponse = {
|
||||
candidates,
|
||||
};
|
||||
|
||||
if (response.usageMetadata) {
|
||||
const usageMetadata = this.normalizeGeminiUsageMetadata(response.usageMetadata);
|
||||
if (!isEmpty(usageMetadata)) {
|
||||
normalized.usageMetadata = usageMetadata;
|
||||
}
|
||||
}
|
||||
|
||||
return normalized;
|
||||
}
|
||||
|
||||
private normalizeGeminiUsageMetadata(
|
||||
usageMetadata: GeminiResponse['usageMetadata'],
|
||||
): NonNullable<GeminiResponse['usageMetadata']> {
|
||||
const normalized: NonNullable<GeminiResponse['usageMetadata']> = {};
|
||||
if (usageMetadata?.promptTokenCount !== undefined) {
|
||||
normalized.promptTokenCount = usageMetadata.promptTokenCount;
|
||||
}
|
||||
if (usageMetadata?.candidatesTokenCount !== undefined) {
|
||||
normalized.candidatesTokenCount = usageMetadata.candidatesTokenCount;
|
||||
}
|
||||
if (usageMetadata?.totalTokenCount !== undefined) {
|
||||
normalized.totalTokenCount = usageMetadata.totalTokenCount;
|
||||
}
|
||||
if (usageMetadata?.thoughtsTokenCount !== undefined) {
|
||||
normalized.thoughtsTokenCount = usageMetadata.thoughtsTokenCount;
|
||||
}
|
||||
if (usageMetadata?.promptTokensDetails !== undefined) {
|
||||
normalized.promptTokensDetails = usageMetadata.promptTokensDetails;
|
||||
}
|
||||
if (usageMetadata?.candidatesTokensDetails !== undefined) {
|
||||
normalized.candidatesTokensDetails = usageMetadata.candidatesTokensDetails;
|
||||
}
|
||||
if (usageMetadata?.trafficType !== undefined) {
|
||||
normalized.trafficType = usageMetadata.trafficType;
|
||||
}
|
||||
return normalized;
|
||||
}
|
||||
|
||||
private sendObservableSse(res: FastifyReply, stream: Observable<unknown>): void {
|
||||
if (
|
||||
!res.raw ||
|
||||
typeof res.raw.writeHead !== 'function' ||
|
||||
typeof res.raw.write !== 'function'
|
||||
) {
|
||||
res.header('Content-Type', 'text/event-stream');
|
||||
res.header('Cache-Control', 'no-cache');
|
||||
res.header('Connection', 'keep-alive');
|
||||
res.send(stream);
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.hasHijack(res)) {
|
||||
res.hijack();
|
||||
}
|
||||
|
||||
res.raw.writeHead(HttpStatus.OK, {
|
||||
'Content-Type': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
Connection: 'keep-alive',
|
||||
});
|
||||
|
||||
const subscription = stream.subscribe({
|
||||
next: (chunk) => {
|
||||
if (res.raw.writableEnded) {
|
||||
return;
|
||||
}
|
||||
const payload = typeof chunk === 'string' ? chunk : String(chunk ?? '');
|
||||
res.raw.write(payload);
|
||||
},
|
||||
error: (error) => {
|
||||
if (res.raw.writableEnded) {
|
||||
return;
|
||||
}
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
res.raw.write(
|
||||
`data: ${JSON.stringify({
|
||||
error: {
|
||||
message,
|
||||
type: 'server_error',
|
||||
},
|
||||
})}\n\n`,
|
||||
);
|
||||
res.raw.end();
|
||||
},
|
||||
complete: () => {
|
||||
if (!res.raw.writableEnded) {
|
||||
res.raw.end();
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
res.raw.on('close', () => {
|
||||
subscription.unsubscribe();
|
||||
});
|
||||
}
|
||||
|
||||
private hasHijack(reply: FastifyReply): reply is FastifyReply & { hijack: () => void } {
|
||||
return typeof (reply as { hijack?: unknown }).hijack === 'function';
|
||||
}
|
||||
|
||||
}
|
||||
@@ -7,13 +7,18 @@ export interface OpenAIChatRequest {
|
||||
stream?: boolean;
|
||||
size?: string;
|
||||
quality?: string;
|
||||
extra?: Record<string, any>;
|
||||
tools?: OpenAITool[];
|
||||
tool_choice?: string | { type: string; function?: { name: string } };
|
||||
response_format?: { type?: string };
|
||||
extra?: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface OpenAIMessage {
|
||||
role: string;
|
||||
content: string | OpenAIContentPart[];
|
||||
name?: string;
|
||||
tool_calls?: OpenAIToolCall[];
|
||||
tool_call_id?: string;
|
||||
}
|
||||
|
||||
export interface OpenAIContentPart {
|
||||
@@ -24,12 +29,32 @@ export interface OpenAIContentPart {
|
||||
};
|
||||
}
|
||||
|
||||
export interface OpenAITool {
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
description?: string;
|
||||
parameters?: Record<string, unknown>;
|
||||
};
|
||||
}
|
||||
|
||||
export interface OpenAIToolCall {
|
||||
id: string;
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
arguments: string;
|
||||
};
|
||||
}
|
||||
|
||||
export interface AnthropicChatRequest {
|
||||
model: string;
|
||||
messages: AnthropicMessage[];
|
||||
system?: string | AnthropicSystemBlock[];
|
||||
max_tokens?: number;
|
||||
metadata?: any;
|
||||
tools?: AnthropicTool[];
|
||||
thinking?: AnthropicThinkingConfig;
|
||||
metadata?: Record<string, unknown>;
|
||||
stop_sequences?: string[];
|
||||
stream?: boolean;
|
||||
temperature?: number;
|
||||
@@ -37,6 +62,18 @@ export interface AnthropicChatRequest {
|
||||
top_k?: number;
|
||||
}
|
||||
|
||||
export interface AnthropicTool {
|
||||
name: string;
|
||||
description?: string;
|
||||
input_schema?: Record<string, unknown>;
|
||||
type?: string;
|
||||
}
|
||||
|
||||
export interface AnthropicThinkingConfig {
|
||||
type: 'enabled' | string;
|
||||
budget_tokens?: number;
|
||||
}
|
||||
|
||||
export interface AnthropicMessage {
|
||||
role: string;
|
||||
content: string | AnthropicContent[];
|
||||
@@ -50,7 +87,21 @@ export interface AnthropicSystemBlock {
|
||||
export type AnthropicContent =
|
||||
| { type: 'text'; text: string }
|
||||
| { type: 'thinking'; thinking: string; signature?: string }
|
||||
| { type: 'image'; source: AnthropicImageSource };
|
||||
| { type: 'image'; source: AnthropicImageSource }
|
||||
| {
|
||||
type: 'tool_use';
|
||||
id: string;
|
||||
name: string;
|
||||
input: Record<string, unknown>;
|
||||
signature?: string;
|
||||
}
|
||||
| {
|
||||
type: 'tool_result';
|
||||
tool_use_id: string;
|
||||
content: string | AnthropicContent[];
|
||||
is_error?: boolean;
|
||||
}
|
||||
| { type: 'redacted_thinking'; data: string };
|
||||
|
||||
export interface AnthropicImageSource {
|
||||
type: 'base64';
|
||||
@@ -102,6 +153,16 @@ export interface GeminiUsageMetadata {
|
||||
promptTokenCount?: number;
|
||||
candidatesTokenCount?: number;
|
||||
totalTokenCount?: number;
|
||||
thoughtsTokenCount?: number;
|
||||
promptTokensDetails?: Array<{
|
||||
modality?: string;
|
||||
tokenCount?: number;
|
||||
}>;
|
||||
candidatesTokensDetails?: Array<{
|
||||
modality?: string;
|
||||
tokenCount?: number;
|
||||
}>;
|
||||
trafficType?: string;
|
||||
}
|
||||
|
||||
export interface OpenAIChatResponse {
|
||||
@@ -121,9 +182,11 @@ export interface OpenAIChoice {
|
||||
index: number;
|
||||
message: {
|
||||
role: string;
|
||||
content: string;
|
||||
content: string | null;
|
||||
tool_calls?: OpenAIToolCall[];
|
||||
reasoning_content?: string;
|
||||
};
|
||||
finish_reason: string;
|
||||
finish_reason: string | null;
|
||||
}
|
||||
|
||||
export interface AnthropicChatResponse {
|
||||
@@ -133,9 +196,11 @@ export interface AnthropicChatResponse {
|
||||
model: string;
|
||||
content: AnthropicContent[];
|
||||
stop_reason: string | null;
|
||||
stop_sequence: string | null;
|
||||
stop_sequence?: string | null;
|
||||
usage: {
|
||||
input_tokens: number;
|
||||
output_tokens: number;
|
||||
cache_creation_input_tokens?: number;
|
||||
cache_read_input_tokens?: number;
|
||||
};
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -2,56 +2,37 @@ import {
|
||||
CanActivate,
|
||||
ExecutionContext,
|
||||
Injectable,
|
||||
UnauthorizedException,
|
||||
Logger,
|
||||
UnauthorizedException,
|
||||
} from '@nestjs/common';
|
||||
import { getServerConfig } from '../../server-config';
|
||||
import { extractApiKeyToken, hasConfiguredApiKey, RequestHeaders } from '../../guards/api-key-auth.util';
|
||||
|
||||
@Injectable()
|
||||
export class ProxyGuard implements CanActivate {
|
||||
private readonly logger = new Logger(ProxyGuard.name);
|
||||
|
||||
canActivate(context: ExecutionContext): boolean {
|
||||
const request = context.switchToHttp().getRequest();
|
||||
|
||||
const config = getServerConfig();
|
||||
|
||||
// 1. Check for API Key in config
|
||||
const apiKey = config?.api_key;
|
||||
|
||||
const headers = request.headers as RequestHeaders;
|
||||
const clientToken = extractApiKeyToken(headers);
|
||||
|
||||
// 2. Bypass if no api_key set (Open Mode) or config missing
|
||||
if (!apiKey || apiKey.trim() === '') {
|
||||
if (!hasConfiguredApiKey(apiKey)) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const request = context.switchToHttp().getRequest();
|
||||
const headers = request.headers;
|
||||
|
||||
// 3. Extract Token from Headers
|
||||
let clientToken: string | undefined;
|
||||
|
||||
// A. standard 'authorization: Bearer <token>'
|
||||
if (headers['authorization']) {
|
||||
const parts = headers['authorization'].split(' ');
|
||||
if (parts.length === 2 && parts[0].toLowerCase() === 'bearer') {
|
||||
clientToken = parts[1];
|
||||
}
|
||||
}
|
||||
|
||||
// B. 'x-api-key' (Anthropic)
|
||||
if (!clientToken && headers['x-api-key']) {
|
||||
clientToken = headers['x-api-key'] as string;
|
||||
}
|
||||
|
||||
// C. 'x-goog-api-key' (Gemini)
|
||||
if (!clientToken && headers['x-goog-api-key']) {
|
||||
clientToken = headers['x-goog-api-key'] as string;
|
||||
}
|
||||
|
||||
// 4. Validate
|
||||
if (clientToken === apiKey) {
|
||||
return true;
|
||||
}
|
||||
this.logger.warn(`Rejected unauthorized request from ${request.ip}`);
|
||||
|
||||
this.logger.warn(`Blocked unauthorized access from ${request.ip}`);
|
||||
throw new UnauthorizedException('Invalid API Key');
|
||||
throw new UnauthorizedException('API key validation failed');
|
||||
}
|
||||
}
|
||||
@@ -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 {}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -8,20 +8,26 @@ interface TokenData {
|
||||
account_id: string;
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
token_type: string;
|
||||
expires_in: number;
|
||||
expiry_timestamp: number;
|
||||
project_id?: string;
|
||||
session_id?: string;
|
||||
upstream_proxy_url?: string;
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class TokenManagerService implements OnModuleInit {
|
||||
private readonly logger = new Logger(TokenManagerService.name);
|
||||
private currentIndex = 0;
|
||||
private readonly stickySessionTtlMs = 10 * 60 * 1000;
|
||||
private readonly rateLimitCooldownMs = 5 * 60 * 1000;
|
||||
private readonly forbiddenCooldownMs = 30 * 60 * 1000;
|
||||
// In-memory cache of tokens with additional data
|
||||
private tokens: Map<string, TokenData> = new Map();
|
||||
// Cooldown map for rate-limited accounts
|
||||
private cooldowns: Map<string, number> = new Map();
|
||||
private sessionBindings: Map<string, { accountId: string; expiresAt: number }> = new Map();
|
||||
|
||||
async onModuleInit() {
|
||||
// Load accounts on module initialization
|
||||
@@ -41,10 +47,10 @@ export class TokenManagerService implements OnModuleInit {
|
||||
}
|
||||
}
|
||||
|
||||
this.logger.log(`Loaded ${count} accounts`);
|
||||
this.logger.log(`Initialized token pool with ${count} accounts`);
|
||||
return count;
|
||||
} catch (e) {
|
||||
this.logger.error('Failed to load accounts', e);
|
||||
this.logger.error('Unable to load accounts', e);
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
@@ -57,10 +63,12 @@ export class TokenManagerService implements OnModuleInit {
|
||||
email: account.email,
|
||||
access_token: account.token.access_token,
|
||||
refresh_token: account.token.refresh_token,
|
||||
token_type: account.token.token_type || 'Bearer',
|
||||
expires_in: account.token.expires_in,
|
||||
expiry_timestamp: account.token.expiry_timestamp,
|
||||
project_id: account.token.project_id || undefined,
|
||||
session_id: account.token.session_id || this.generateSessionId(),
|
||||
upstream_proxy_url: account.token.upstream_proxy_url || undefined,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -72,7 +80,10 @@ export class TokenManagerService implements OnModuleInit {
|
||||
return (-(min + rand)).toString();
|
||||
}
|
||||
|
||||
async getNextToken(): Promise<CloudAccount | null> {
|
||||
async getNextToken(options?: {
|
||||
sessionKey?: string;
|
||||
excludeAccountIds?: string[];
|
||||
}): Promise<CloudAccount | null> {
|
||||
try {
|
||||
// Reload if empty
|
||||
if (this.tokens.size === 0) {
|
||||
@@ -82,25 +93,81 @@ export class TokenManagerService implements OnModuleInit {
|
||||
|
||||
const now = Date.now();
|
||||
const nowSeconds = Math.floor(now / 1000);
|
||||
const sessionKey = options?.sessionKey?.trim();
|
||||
const excludedAccountIds = new Set(options?.excludeAccountIds ?? []);
|
||||
|
||||
this.clearExpiredSessionBindings(now);
|
||||
|
||||
const fullTokenPool = Array.from(this.tokens.entries());
|
||||
const excludedFilteredTokens = fullTokenPool.filter(
|
||||
([accountId]) => !excludedAccountIds.has(accountId),
|
||||
);
|
||||
const allTokens = excludedFilteredTokens.length > 0 ? excludedFilteredTokens : fullTokenPool;
|
||||
if (excludedFilteredTokens.length === 0 && excludedAccountIds.size > 0) {
|
||||
this.logger.warn(
|
||||
'Retry exclusions removed all accounts; temporarily reusing excluded accounts',
|
||||
);
|
||||
}
|
||||
|
||||
// Filter out accounts in cooldown
|
||||
const validTokens = Array.from(this.tokens.entries()).filter(([email, _]) => {
|
||||
const cooldownUntil = this.cooldowns.get(email);
|
||||
const validTokens = allTokens.filter(([accountId]) => {
|
||||
if (excludedAccountIds.has(accountId)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const cooldownUntil = this.cooldowns.get(accountId);
|
||||
return !cooldownUntil || cooldownUntil <= now;
|
||||
});
|
||||
|
||||
if (validTokens.length === 0) {
|
||||
this.logger.warn('All accounts are in cooldown');
|
||||
const candidateTokens = validTokens.length > 0 ? validTokens : allTokens;
|
||||
if (candidateTokens.length === 0) {
|
||||
this.logger.warn('No account available after applying exclusions');
|
||||
return null;
|
||||
}
|
||||
if (validTokens.length === 0) {
|
||||
this.logger.warn(
|
||||
'All accounts are currently in cooldown; temporarily bypassing cooldown to keep service available',
|
||||
);
|
||||
}
|
||||
|
||||
if (sessionKey) {
|
||||
const stickyBinding = this.sessionBindings.get(sessionKey);
|
||||
if (stickyBinding && stickyBinding.expiresAt > now) {
|
||||
const stickyMatch = candidateTokens.find(
|
||||
([accountId]) => accountId === stickyBinding.accountId,
|
||||
);
|
||||
if (stickyMatch) {
|
||||
const [stickyAccountId, stickyTokenData] = stickyMatch;
|
||||
return this.finalizeSelectedToken(
|
||||
stickyAccountId,
|
||||
stickyTokenData,
|
||||
nowSeconds,
|
||||
sessionKey,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Round robin selection
|
||||
const [accountId, tokenData] = validTokens[this.currentIndex % validTokens.length];
|
||||
const [accountId, tokenData] = candidateTokens[this.currentIndex % candidateTokens.length];
|
||||
this.currentIndex++;
|
||||
return this.finalizeSelectedToken(accountId, tokenData, nowSeconds, sessionKey);
|
||||
} catch (error) {
|
||||
this.logger.error('Unable to acquire token', error);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private async finalizeSelectedToken(
|
||||
accountId: string,
|
||||
tokenData: TokenData,
|
||||
nowSeconds: number,
|
||||
sessionKey?: string,
|
||||
): Promise<CloudAccount | null> {
|
||||
try {
|
||||
// Check if token needs refresh (expires in < 5 minutes)
|
||||
if (nowSeconds >= tokenData.expiry_timestamp - 300) {
|
||||
this.logger.log(`Token for ${tokenData.email} expiring soon, refreshing...`);
|
||||
this.logger.log(`Token for ${tokenData.email} is close to expiry; refreshing...`);
|
||||
try {
|
||||
const newTokens = await GoogleAPIService.refreshAccessToken(tokenData.refresh_token);
|
||||
|
||||
@@ -113,51 +180,96 @@ export class TokenManagerService implements OnModuleInit {
|
||||
await this.saveRefreshedToken(accountId, tokenData);
|
||||
this.tokens.set(accountId, tokenData);
|
||||
|
||||
this.logger.log(`Token refreshed for ${tokenData.email}`);
|
||||
this.logger.log(`Refreshed token for ${tokenData.email}`);
|
||||
} catch (e) {
|
||||
this.logger.error(`Failed to refresh token for ${tokenData.email}`, e);
|
||||
this.logger.error(`Unable to refresh token for ${tokenData.email}`, e);
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve project ID if missing (mock for now, like original)
|
||||
if (!tokenData.project_id) {
|
||||
const mockId = `cloud-code-${Math.floor(Math.random() * 100000)}`;
|
||||
tokenData.project_id = mockId;
|
||||
await this.saveProjectId(accountId, mockId);
|
||||
if (
|
||||
typeof tokenData.project_id !== 'string' ||
|
||||
tokenData.project_id.trim() === '' ||
|
||||
/^cloud-code-\d+$/i.test(tokenData.project_id.trim())
|
||||
) {
|
||||
tokenData.project_id = undefined;
|
||||
}
|
||||
|
||||
this.logger.log(`Selected account: ${tokenData.email}`);
|
||||
this.logger.log(`Using account: ${tokenData.email}`);
|
||||
if (sessionKey) {
|
||||
this.sessionBindings.set(sessionKey, {
|
||||
accountId,
|
||||
expiresAt: Date.now() + this.stickySessionTtlMs,
|
||||
});
|
||||
}
|
||||
|
||||
// Return in CloudAccount format for compatibility
|
||||
return {
|
||||
const timestamp = Date.now();
|
||||
const cloudAccount: CloudAccount = {
|
||||
id: accountId,
|
||||
provider: 'google',
|
||||
email: tokenData.email,
|
||||
token: {
|
||||
access_token: tokenData.access_token,
|
||||
refresh_token: tokenData.refresh_token,
|
||||
token_type: tokenData.token_type,
|
||||
expires_in: tokenData.expires_in,
|
||||
expiry_timestamp: tokenData.expiry_timestamp,
|
||||
project_id: tokenData.project_id,
|
||||
session_id: tokenData.session_id,
|
||||
upstream_proxy_url: tokenData.upstream_proxy_url,
|
||||
},
|
||||
} as CloudAccount;
|
||||
created_at: timestamp,
|
||||
last_used: timestamp,
|
||||
};
|
||||
return cloudAccount;
|
||||
} catch (error) {
|
||||
this.logger.error('Failed to get token', error);
|
||||
this.logger.error('Unable to finalize selected token', error);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
markAsRateLimited(email: string) {
|
||||
// Cooldown for 5 minutes
|
||||
const until = Date.now() + 5 * 60 * 1000;
|
||||
this.cooldowns.set(email, until);
|
||||
this.logger.warn(
|
||||
`Account ${email} marked as rate limited until ${new Date(until).toISOString()}`,
|
||||
);
|
||||
markAsRateLimited(accountIdOrEmail: string) {
|
||||
this.markInCooldown(accountIdOrEmail, 'rate limited', this.rateLimitCooldownMs);
|
||||
}
|
||||
|
||||
resetCooldown(email: string) {
|
||||
this.cooldowns.delete(email);
|
||||
markAsForbidden(accountIdOrEmail: string) {
|
||||
this.markInCooldown(accountIdOrEmail, 'forbidden', this.forbiddenCooldownMs);
|
||||
}
|
||||
|
||||
private resolveAccountId(accountIdOrEmail: string): string | null {
|
||||
if (this.tokens.has(accountIdOrEmail)) {
|
||||
return accountIdOrEmail;
|
||||
}
|
||||
|
||||
for (const [accountId, tokenData] of this.tokens.entries()) {
|
||||
if (tokenData.email === accountIdOrEmail) {
|
||||
return accountId;
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private clearExpiredSessionBindings(now: number): void {
|
||||
for (const [sessionKey, binding] of this.sessionBindings.entries()) {
|
||||
if (binding.expiresAt <= now) {
|
||||
this.sessionBindings.delete(sessionKey);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private markInCooldown(
|
||||
accountIdOrEmail: string,
|
||||
reason: 'rate limited' | 'forbidden',
|
||||
durationMs: number,
|
||||
): void {
|
||||
const accountId = this.resolveAccountId(accountIdOrEmail) ?? accountIdOrEmail;
|
||||
const cooldownUntil = Date.now() + durationMs;
|
||||
|
||||
this.cooldowns.set(accountId, cooldownUntil);
|
||||
this.logger.warn(
|
||||
`Account ${accountIdOrEmail} (cooldown key: ${accountId}) set to ${reason} until ${new Date(cooldownUntil).toISOString()}`,
|
||||
);
|
||||
}
|
||||
|
||||
private async saveRefreshedToken(accountId: string, tokenData: TokenData) {
|
||||
@@ -173,22 +285,7 @@ export class TokenManagerService implements OnModuleInit {
|
||||
await CloudAccountRepo.updateToken(accountId, newToken);
|
||||
}
|
||||
} catch (e) {
|
||||
this.logger.error('Failed to save refreshed token to DB', e);
|
||||
}
|
||||
}
|
||||
|
||||
private async saveProjectId(accountId: string, projectId: string) {
|
||||
try {
|
||||
const acc = await CloudAccountRepo.getAccount(accountId);
|
||||
if (acc && acc.token) {
|
||||
const newToken = {
|
||||
...acc.token,
|
||||
project_id: projectId,
|
||||
};
|
||||
await CloudAccountRepo.updateToken(accountId, newToken);
|
||||
}
|
||||
} catch (e) {
|
||||
this.logger.error('Failed to save project ID to DB', e);
|
||||
this.logger.error('Unable to persist refreshed token to DB', e);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<string, any> = {};
|
||||
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',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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(() => {
|
||||
|
||||
@@ -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<string, any> = {};
|
||||
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: '',
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
};
|
||||
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: '',
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
};
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -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<string> {
|
||||
return (this as any).processStreamResponse(upstreamStream, model);
|
||||
}
|
||||
}
|
||||
|
||||
function readFixture<T>(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<any>('request/openai.chat-tools.input.json');
|
||||
const expected = readFixture<any>('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<any>('response/anthropic.tool-use.input.json');
|
||||
const expected = readFixture<any>('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<any>('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<void>((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);
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -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<string> {
|
||||
return (this as any).passthroughSseStream(stream);
|
||||
}
|
||||
|
||||
public testModelHeaders(model: string): Record<string, string> {
|
||||
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<void>((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<void>((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<void>((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<void>((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<void>((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]');
|
||||
});
|
||||
});
|
||||
@@ -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({
|
||||
|
||||
@@ -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: {
|
||||
|
||||
Reference in new issue
Block a user