From e802776c96077c414f279ef9acfcec1371d096c8 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 12 Mar 2026 01:10:29 +0500 Subject: [PATCH] feat(core): support custom base URL via env vars (#21561) Co-authored-by: Spencer --- .../core/src/core/contentGenerator.test.ts | 141 ++++++++++++++++++ packages/core/src/core/contentGenerator.ts | 30 +++- 2 files changed, 169 insertions(+), 2 deletions(-) diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index d86eb6f738..c5dcc6e22a 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -10,6 +10,7 @@ import { AuthType, createContentGeneratorConfig, type ContentGenerator, + validateBaseUrl, } from './contentGenerator.js'; import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js'; import { GoogleGenAI } from '@google/genai'; @@ -442,6 +443,116 @@ describe('createContentGenerator', () => { ); }); + it('should pass GOOGLE_GEMINI_BASE_URL as httpOptions.baseUrl for Gemini API', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv('GOOGLE_GEMINI_BASE_URL', 'https://my-gemini-proxy.example.com'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + httpOptions: expect.objectContaining({ + baseUrl: 'https://my-gemini-proxy.example.com', + }), + }), + ); + }); + + it('should pass GOOGLE_VERTEX_BASE_URL as httpOptions.baseUrl for Vertex AI', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv('GOOGLE_VERTEX_BASE_URL', 'https://my-vertex-proxy.example.com'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + vertexai: true, + authType: AuthType.USE_VERTEX_AI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + httpOptions: expect.objectContaining({ + baseUrl: 'https://my-vertex-proxy.example.com', + }), + }), + ); + }); + + it('should not include baseUrl in httpOptions when GOOGLE_GEMINI_BASE_URL is not set', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.not.objectContaining({ + httpOptions: expect.objectContaining({ + baseUrl: expect.any(String), + }), + }), + ); + }); + + it('should reject an insecure GOOGLE_GEMINI_BASE_URL for non-local hosts', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + vi.stubEnv('GOOGLE_GEMINI_BASE_URL', 'http://evil-proxy.example.com'); + + await expect( + createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ), + ).rejects.toThrow('Custom base URL must use HTTPS unless it is localhost.'); + }); + it('should pass apiVersion for Vertex AI when GOOGLE_GENAI_API_VERSION is set', async () => { const mockConfig = { getModel: vi.fn().mockReturnValue('gemini-pro'), @@ -560,3 +671,33 @@ describe('createContentGeneratorConfig', () => { expect(config.vertexai).toBeUndefined(); }); }); + +describe('validateBaseUrl', () => { + it('should accept a valid HTTPS URL', () => { + expect(() => validateBaseUrl('https://my-proxy.example.com')).not.toThrow(); + }); + + it('should accept HTTP for localhost', () => { + expect(() => validateBaseUrl('http://localhost:8080')).not.toThrow(); + }); + + it('should accept HTTP for 127.0.0.1', () => { + expect(() => validateBaseUrl('http://127.0.0.1:3000')).not.toThrow(); + }); + + it('should accept HTTP for ::1', () => { + expect(() => validateBaseUrl('http://[::1]:8080')).not.toThrow(); + }); + + it('should reject HTTP for non-local hosts', () => { + expect(() => validateBaseUrl('http://my-proxy.example.com')).toThrow( + 'Custom base URL must use HTTPS unless it is localhost.', + ); + }); + + it('should reject an invalid URL', () => { + expect(() => validateBaseUrl('not-a-url')).toThrow( + 'Invalid custom base URL: not-a-url', + ); + }); +}); diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 2ce5420335..d7da9fb064 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -225,13 +225,25 @@ export async function createContentGenerator( 'x-gemini-api-privileged-user-id': `${installationId}`, }; } + let baseUrl = config.baseUrl; + if (!baseUrl) { + const envBaseUrl = config.vertexai + ? process.env['GOOGLE_VERTEX_BASE_URL'] + : process.env['GOOGLE_GEMINI_BASE_URL']; + if (envBaseUrl) { + validateBaseUrl(envBaseUrl); + baseUrl = envBaseUrl; + } + } else { + validateBaseUrl(baseUrl); + } const httpOptions: { baseUrl?: string; headers: Record; } = { headers }; - if (config.baseUrl) { - httpOptions.baseUrl = config.baseUrl; + if (baseUrl) { + httpOptions.baseUrl = baseUrl; } const googleGenAI = new GoogleGenAI({ @@ -253,3 +265,17 @@ export async function createContentGenerator( return generator; } + +const LOCAL_HOSTNAMES = ['localhost', '127.0.0.1', '[::1]']; + +export function validateBaseUrl(baseUrl: string): void { + let url: URL; + try { + url = new URL(baseUrl); + } catch { + throw new Error(`Invalid custom base URL: ${baseUrl}`); + } + if (url.protocol !== 'https:' && !LOCAL_HOSTNAMES.includes(url.hostname)) { + throw new Error('Custom base URL must use HTTPS unless it is localhost.'); + } +}