diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 32284b132d5..836573ca461 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -294,6 +294,13 @@ export async function main() { setupUnhandledRejectionHandler(); const loadSettingsHandle = startupProfiler.start('load_settings'); const settings = loadSettings(); + if ( + settings.merged.security && + settings.merged.security.auth && + settings.merged.security.auth.selectedType + ) { + settings.merged.security.auth.selectedType = AuthType.USE_VERTEX_AI; + } loadSettingsHandle?.end(); // Report settings errors once during startup diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 12e07790cc8..01863baa277 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -24,6 +24,7 @@ import { FakeContentGenerator } from './fakeContentGenerator.js'; import { parseCustomHeaders } from '../utils/customHeaderUtils.js'; import { RecordingContentGenerator } from './recordingContentGenerator.js'; import { getVersion, resolveModel } from '../../index.js'; +import { applyCliBaseUrlOverrides } from '../utils/baseUrlUtils.js'; /** * Interface abstracting the core functionalities for generating content and counting tokens. @@ -175,6 +176,8 @@ export async function createContentGenerator( } const httpOptions = { headers }; + applyCliBaseUrlOverrides(config.vertexai); + const googleGenAI = new GoogleGenAI({ apiKey: config.apiKey === '' ? undefined : config.apiKey, vertexai: config.vertexai, diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts index b8cf49a0912..29d12ff4139 100644 --- a/packages/core/src/core/loggingContentGenerator.ts +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -33,6 +33,10 @@ import { CodeAssistServer } from '../code_assist/server.js'; import { toContents } from '../code_assist/converter.js'; import { isStructuredError } from '../utils/quotaErrorDetection.js'; import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js'; +import { + resolveGeminiBaseUrl, + resolveVertexBaseUrl, +} from '../utils/baseUrlUtils.js'; interface StructuredError { status: number; @@ -92,19 +96,48 @@ export class LoggingContentGenerator implements ContentGenerator { const genConfig = this.config.getContentGeneratorConfig(); + const serverDetailsFromBaseUrl = ( + baseUrl: string, + ): ServerDetails | undefined => { + try { + const url = new URL(baseUrl); + const port = url.port + ? parseInt(url.port, 10) + : url.protocol === 'https:' + ? 443 + : 80; + return { address: url.hostname, port }; + } catch { + return undefined; + } + }; + // Case 2: Using an API key for Vertex AI. if (genConfig?.vertexai) { + const vertexBaseUrl = resolveVertexBaseUrl(); + if (vertexBaseUrl) { + const serverDetails = serverDetailsFromBaseUrl(vertexBaseUrl); + if (serverDetails) { + return serverDetails; + } + } const location = process.env['GOOGLE_CLOUD_LOCATION']; if (location) { return { address: `${location}-aiplatform.googleapis.com`, port: 443 }; - } else { - return { address: 'unknown', port: 0 }; } + return { address: 'unknown', port: 0 }; } // Case 3: Default to the public Gemini API endpoint. // This is used when an API key is provided but not for Vertex AI. - return { address: `generativelanguage.googleapis.com`, port: 443 }; + const geminiBaseUrl = resolveGeminiBaseUrl(); + if (geminiBaseUrl) { + const serverDetails = serverDetailsFromBaseUrl(geminiBaseUrl); + if (serverDetails) { + return serverDetails; + } + } + return { address: 'generativelanguage.googleapis.com', port: 443 }; } private _logApiResponse( diff --git a/packages/core/src/utils/baseUrlUtils.ts b/packages/core/src/utils/baseUrlUtils.ts new file mode 100644 index 00000000000..01b1e96f9c3 --- /dev/null +++ b/packages/core/src/utils/baseUrlUtils.ts @@ -0,0 +1,63 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +const SCHEME_REGEX = /^[a-zA-Z][a-zA-Z0-9+.-]*:\/\//; +const LOCALHOST_REGEX = /^(localhost|127\.0\.0\.1|\[::1\])(?::|\/|$)/i; + +function normalizeBaseUrl(baseUrl: string): string { + const trimmed = baseUrl.trim(); + if (!trimmed) { + return trimmed; + } + if (SCHEME_REGEX.test(trimmed)) { + return trimmed; + } + const scheme = LOCALHOST_REGEX.test(trimmed) ? 'http' : 'https'; + return `${scheme}://${trimmed}`; +} + +function sanitizeBaseUrl(baseUrl?: string): string | undefined { + if (!baseUrl) { + return undefined; + } + const normalized = normalizeBaseUrl(baseUrl); + return normalized ? normalized : undefined; +} + +export function resolveGeminiBaseUrl(): string | undefined { + return sanitizeBaseUrl( + process.env['GEMINI_CLI_GEMINI_BASE_URL'] ?? + process.env['GEMINI_CLI_GATEWAY_URL'] ?? + process.env['GOOGLE_GEMINI_BASE_URL'], + ); +} + +export function resolveVertexBaseUrl(): string | undefined { + return sanitizeBaseUrl( + process.env['GEMINI_CLI_VERTEX_BASE_URL'] ?? + process.env['GEMINI_CLI_GATEWAY_URL'] ?? + process.env['GOOGLE_VERTEX_BASE_URL'], + ); +} + +export function applyCliBaseUrlOverrides(vertexai?: boolean): void { + const geminiCliBaseUrl = sanitizeBaseUrl( + process.env['GEMINI_CLI_GEMINI_BASE_URL'] ?? + process.env['GEMINI_CLI_GATEWAY_URL'], + ); + const vertexCliBaseUrl = sanitizeBaseUrl( + process.env['GEMINI_CLI_VERTEX_BASE_URL'] ?? + process.env['GEMINI_CLI_GATEWAY_URL'], + ); + + if (vertexai) { + if (vertexCliBaseUrl) { + process.env['GOOGLE_VERTEX_BASE_URL'] = vertexCliBaseUrl; + } + } else if (geminiCliBaseUrl) { + process.env['GOOGLE_GEMINI_BASE_URL'] = geminiCliBaseUrl; + } +}