diff --git a/packages/api/src/endpoints/google/llm.spec.ts b/packages/api/src/endpoints/google/llm.spec.ts index 3c977c20ec..edfa74f044 100644 --- a/packages/api/src/endpoints/google/llm.spec.ts +++ b/packages/api/src/endpoints/google/llm.spec.ts @@ -308,7 +308,7 @@ describe('getGoogleConfig', () => { }); }); - it('should preserve explicit Vertex AI endpoint overrides', () => { + it('should preserve explicit Google Vertex AI endpoint overrides', () => { process.env.GOOGLE_LOC = 'us'; const credentials = { @@ -323,13 +323,134 @@ describe('getGoogleConfig', () => { }, addParams: { location: 'eu', - endpoint: 'custom-aiplatform.example.com', + endpoint: 'us-central1-aiplatform.googleapis.com', }, }); expect(result.llmConfig).toMatchObject({ location: 'eu', - endpoint: 'custom-aiplatform.example.com', + endpoint: 'us-central1-aiplatform.googleapis.com', + }); + }); + + it('should preserve explicit Google Vertex AI Private Service Connect endpoints', () => { + process.env.GOOGLE_LOC = 'us'; + + const credentials = { + [AuthKeys.GOOGLE_SERVICE_KEY]: { + project_id: 'test-project', + }, + }; + + const result = getGoogleConfig(credentials, { + modelOptions: { + model: 'gemini-3.1-flash-lite-preview', + }, + addParams: { + endpoint: 'aiplatform-genai1.p.googleapis.com', + }, + }); + + expect(result.llmConfig).toMatchObject({ + location: 'us', + endpoint: 'aiplatform-genai1.p.googleapis.com', + }); + }); + + it('should preserve explicit Google Vertex AI restricted Private Service Connect endpoints', () => { + process.env.GOOGLE_LOC = 'us'; + + const credentials = { + [AuthKeys.GOOGLE_SERVICE_KEY]: { + project_id: 'test-project', + }, + }; + + const result = getGoogleConfig(credentials, { + modelOptions: { + model: 'gemini-3.1-flash-lite-preview', + }, + addParams: { + endpoint: 'us-central1-aiplatform-restricted.p.googleapis.com', + }, + }); + + expect(result.llmConfig).toMatchObject({ + location: 'us', + endpoint: 'us-central1-aiplatform-restricted.p.googleapis.com', + }); + }); + + it('should ignore model option Vertex AI endpoint overrides', () => { + process.env.GOOGLE_LOC = 'eu'; + + const credentials = { + [AuthKeys.GOOGLE_SERVICE_KEY]: { + project_id: 'test-project', + }, + }; + + const result = getGoogleConfig(credentials, { + modelOptions: { + model: 'gemini-3.1-flash-lite-preview', + endpoint: 'attacker.example.test', + } as t.GoogleParameters, + }); + + expect(result.llmConfig).toMatchObject({ + location: 'eu', + endpoint: 'aiplatform.eu.rep.googleapis.com', + }); + }); + + it('should ignore model option transport-level overrides', () => { + const credentials = { + [AuthKeys.GOOGLE_SERVICE_KEY]: { + project_id: 'test-project', + client_email: 'test@test-project.iam.gserviceaccount.com', + }, + }; + + const result = getGoogleConfig(credentials, { + modelOptions: { + model: 'gemini-3.1-flash-lite-preview', + apiKey: 'attacker-api-key', + authOptions: { projectId: 'attacker-project' }, + baseUrl: 'https://attacker.example.test', + customHeaders: { Authorization: 'Bearer attacker' }, + } as t.GoogleParameters, + }); + + expect(result.llmConfig).not.toHaveProperty('apiKey', 'attacker-api-key'); + expect(result.llmConfig).not.toHaveProperty('baseUrl'); + expect(result.llmConfig).not.toHaveProperty('customHeaders'); + expect((result.llmConfig as Record).authOptions).toMatchObject({ + projectId: 'test-project', + }); + }); + + it('should ignore non-Google Vertex AI endpoint overrides from additional params', () => { + process.env.GOOGLE_LOC = 'us'; + + const credentials = { + [AuthKeys.GOOGLE_SERVICE_KEY]: { + project_id: 'test-project', + }, + }; + + const result = getGoogleConfig(credentials, { + modelOptions: { + model: 'gemini-3.1-flash-lite-preview', + }, + addParams: { + location: 'eu', + endpoint: 'attacker.example.test', + }, + }); + + expect(result.llmConfig).toMatchObject({ + location: 'eu', + endpoint: 'aiplatform.eu.rep.googleapis.com', }); }); diff --git a/packages/api/src/endpoints/google/llm.ts b/packages/api/src/endpoints/google/llm.ts index e3826b9cdb..460dc74252 100644 --- a/packages/api/src/endpoints/google/llm.ts +++ b/packages/api/src/endpoints/google/llm.ts @@ -25,6 +25,20 @@ const vertexMultiRegionEndpoints = new Map([ ['global', 'aiplatform.googleapis.com'], ]); +const blockedModelOptionParams = [ + 'apiKey', + 'baseUrl', + 'baseURL', + 'endpoint', + 'authOptions', + 'customHeaders', + 'headers', +] as const; + +type BlockedModelOptionParam = (typeof blockedModelOptionParams)[number]; +type GoogleModelOptions = Partial & + Partial>; + /** Known Google/Vertex AI parameters that map directly to the client config */ export const knownGoogleParams = new Set([ 'model', @@ -106,8 +120,57 @@ function getVertexMultiRegionEndpoint(location: string): string | undefined { return vertexMultiRegionEndpoints.get(location); } -function hasStringEndpoint(config: Record): boolean { - return typeof config.endpoint === 'string' && config.endpoint.length > 0; +function sanitizeModelOptions(modelOptions: Partial | undefined) { + const sanitizedOptions: GoogleModelOptions = { ...(modelOptions ?? {}) }; + blockedModelOptionParams.forEach((param) => { + delete sanitizedOptions[param]; + }); + return sanitizedOptions; +} + +function isAllowedVertexEndpoint(endpoint: string): boolean { + if (!/^[a-z0-9][a-z0-9.-]*[a-z0-9]$/.test(endpoint)) { + return false; + } + + if (endpoint === 'aiplatform.googleapis.com') { + return true; + } + + if (/^[a-z0-9-]+-aiplatform\.googleapis\.com$/.test(endpoint)) { + return true; + } + + if (/^aiplatform-[a-z0-9-]+\.p\.googleapis\.com$/.test(endpoint)) { + return true; + } + + if (/^[a-z0-9-]+-aiplatform-restricted\.p\.googleapis\.com$/.test(endpoint)) { + return true; + } + + return /^aiplatform\.[a-z0-9-]+\.rep\.googleapis\.com$/.test(endpoint); +} + +function hasAllowedVertexEndpoint(config: Record): boolean { + const endpoint = config.endpoint; + if (endpoint === undefined) { + return false; + } + + if (typeof endpoint !== 'string' || endpoint.trim() !== endpoint || endpoint.length === 0) { + delete config.endpoint; + return false; + } + + const normalizedEndpoint = endpoint.toLowerCase(); + if (!isAllowedVertexEndpoint(normalizedEndpoint)) { + delete config.endpoint; + return false; + } + + config.endpoint = normalizedEndpoint; + return true; } function applyVertexMultiRegionEndpoint(config: VertexAIClientOptions & { endpoint?: string }) { @@ -205,11 +268,11 @@ export function getGoogleConfig( thinking = googleSettings.thinking.default, thinkingBudget = googleSettings.thinkingBudget.default, ...modelOptions - } = options.modelOptions || {}; + } = sanitizeModelOptions(options.modelOptions); let enableWebSearch = web_search; - const llmConfig: GoogleClientOptions | VertexAIClientOptions = removeNullishValues( + const llmConfig = removeNullishValues( { ...(modelOptions || {}), model: modelOptions?.model ?? '', @@ -220,9 +283,9 @@ export function getGoogleConfig( maxOutputTokens: modelOptions?.maxOutputTokens ?? undefined, }, true, - ); + ) as GoogleClientOptions | VertexAIClientOptions; const initialConfig = llmConfig as Record; - let hasCustomVertexEndpoint = hasStringEndpoint(initialConfig); + let hasCustomVertexEndpoint = hasAllowedVertexEndpoint(initialConfig); let shouldSyncVertexEndpoint = true; /** Used only for Safety Settings */ @@ -341,8 +404,8 @@ export function getGoogleConfig( if (knownGoogleParams.has(key)) { /** Route known Google params to llmConfig only if undefined */ applyDefaultParams(llmConfig as Record, { [key]: value }); - if (key === 'endpoint' && hasStringEndpoint(llmConfig as Record)) { - hasCustomVertexEndpoint = true; + if (key === 'endpoint') { + hasCustomVertexEndpoint = hasAllowedVertexEndpoint(llmConfig as Record); } } /** Leave other params for transform to handle - they might be OpenAI params */ @@ -364,7 +427,7 @@ export function getGoogleConfig( /** Route known Google params to llmConfig */ (llmConfig as Record)[key] = value; if (key === 'endpoint') { - hasCustomVertexEndpoint = hasStringEndpoint(llmConfig as Record); + hasCustomVertexEndpoint = hasAllowedVertexEndpoint(llmConfig as Record); } } /** Leave other params for transform to handle - they might be OpenAI params */