feat(core): support custom base URL via env vars (#21561)

Co-authored-by: Spencer <spencertang@google.com>

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',
+    );
+  });
+});

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<string, string>;
       } = { 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.');
+  }
+}