From acf5ed595fe33b2e8760dedf267b0305f634c1d9 Mon Sep 17 00:00:00 2001 From: Aarushi Shah Date: Tue, 25 Nov 2025 16:22:47 -0800 Subject: [PATCH] Add Databricks auth support and custom header option to gemini cli (#11893) Co-authored-by: Taylor Mullen --- docs/cli/configuration.md | 12 ++ .../core/src/core/contentGenerator.test.ts | 154 ++++++++++++++++++ packages/core/src/core/contentGenerator.ts | 17 ++ .../core/src/utils/customHeaderUtils.test.ts | 91 +++++++++++ packages/core/src/utils/customHeaderUtils.ts | 41 +++++ 5 files changed, 315 insertions(+) create mode 100644 packages/core/src/utils/customHeaderUtils.test.ts create mode 100644 packages/core/src/utils/customHeaderUtils.ts diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index 4af6444459..4a58ec8006 100644 --- a/docs/cli/configuration.md +++ b/docs/cli/configuration.md @@ -464,6 +464,18 @@ the `excludedProjectEnvVars` setting in your `settings.json` file. - Specifies the default Gemini model to use. - Overrides the hardcoded default - Example: `export GEMINI_MODEL="gemini-2.5-flash"` +- **`GEMINI_CLI_CUSTOM_HEADERS`**: + - Adds extra HTTP headers to Gemini API and Code Assist requests. + - Accepts a comma-separated list of `Name: value` pairs. + - Example: + `export GEMINI_CLI_CUSTOM_HEADERS="X-My-Header: foo, X-Trace-ID: abc123"`. +- **`GEMINI_API_KEY_AUTH_MECHANISM`**: + - Specifies how the API key should be sent for authentication when using + `AuthType.USE_GEMINI` or `AuthType.USE_VERTEX_AI`. + - Valid values are `x-goog-api-key` (default) or `bearer`. + - If set to `bearer`, the API key will be sent in the + `Authorization: Bearer ` header. + - Example: `export GEMINI_API_KEY_AUTH_MECHANISM="bearer"` - **`GOOGLE_API_KEY`**: - Your Google Cloud API key. - Required for using Vertex AI in express mode. diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index c31585f7a8..daf6e35878 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -30,6 +30,14 @@ vi.mock('./fakeContentGenerator.js'); const mockConfig = {} as unknown as Config; describe('createContentGenerator', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + }); + it('should create a FakeContentGenerator', async () => { const mockGenerator = {} as unknown as ContentGenerator; vi.mocked(FakeContentGenerator.fromFile).mockResolvedValue( @@ -135,6 +143,152 @@ describe('createContentGenerator', () => { ); }); + it('should include custom headers from GEMINI_CLI_CUSTOM_HEADERS for Code Assist requests', async () => { + const mockGenerator = {} as unknown as ContentGenerator; + vi.mocked(createCodeAssistContentGenerator).mockResolvedValue( + mockGenerator as never, + ); + vi.stubEnv( + 'GEMINI_CLI_CUSTOM_HEADERS', + 'X-Test-Header: test-value, Another-Header: another value', + ); + + await createContentGenerator( + { + authType: AuthType.LOGIN_WITH_GOOGLE, + }, + mockConfig, + ); + + expect(createCodeAssistContentGenerator).toHaveBeenCalledWith( + { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + 'X-Test-Header': 'test-value', + 'Another-Header': 'another value', + }), + }, + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + undefined, + ); + }); + + it('should include custom headers from GEMINI_CLI_CUSTOM_HEADERS for GoogleGenAI requests without inferring auth mechanism', async () => { + const mockConfig = { + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv( + 'GEMINI_CLI_CUSTOM_HEADERS', + 'X-Test-Header: test, Another: value', + ); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: undefined, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + 'X-Test-Header': 'test', + Another: 'value', + }), + }, + }); + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.not.objectContaining({ + httpOptions: { + headers: expect.objectContaining({ + Authorization: expect.any(String), + }), + }, + }), + ); + }); + + it('should pass api key as Authorization Header when GEMINI_API_KEY_AUTH_MECHANISM is set to bearer', async () => { + const mockConfig = { + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv('GEMINI_API_KEY_AUTH_MECHANISM', 'bearer'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: undefined, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + Authorization: 'Bearer test-api-key', + }), + }, + }); + }); + + it('should not pass api key as Authorization Header when GEMINI_API_KEY_AUTH_MECHANISM is not set (default behavior)', async () => { + const mockConfig = { + getUsageStatisticsEnabled: () => false, + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + // GEMINI_API_KEY_AUTH_MECHANISM is not stubbed, so it will be undefined, triggering default 'x-goog-api-key' + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: undefined, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + }), + }, + }); + // Explicitly assert that Authorization header is NOT present + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.not.objectContaining({ + httpOptions: { + headers: expect.objectContaining({ + Authorization: expect.any(String), + }), + }, + }), + ); + }); + it('should create a GoogleGenAI content generator with client install id logging disabled', async () => { const mockConfig = { getUsageStatisticsEnabled: () => false, diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 4e37a9c0cc..8d6a8ad4ae 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -21,6 +21,7 @@ import type { UserTierId } from '../code_assist/types.js'; import { LoggingContentGenerator } from './loggingContentGenerator.js'; import { InstallationManager } from '../utils/installationManager.js'; import { FakeContentGenerator } from './fakeContentGenerator.js'; +import { parseCustomHeaders } from '../utils/customHeaderUtils.js'; import { RecordingContentGenerator } from './recordingContentGenerator.js'; /** @@ -115,10 +116,26 @@ export async function createContentGenerator( return FakeContentGenerator.fromFile(gcConfig.fakeResponses); } const version = process.env['CLI_VERSION'] || process.version; + const customHeadersEnv = + process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined; const userAgent = `GeminiCLI/${version} (${process.platform}; ${process.arch})`; + const customHeadersMap = parseCustomHeaders(customHeadersEnv); + const apiKeyAuthMechanism = + process.env['GEMINI_API_KEY_AUTH_MECHANISM'] || 'x-goog-api-key'; + const baseHeaders: Record = { + ...customHeadersMap, 'User-Agent': userAgent, }; + + if ( + apiKeyAuthMechanism === 'bearer' && + (config.authType === AuthType.USE_GEMINI || + config.authType === AuthType.USE_VERTEX_AI) && + config.apiKey + ) { + baseHeaders['Authorization'] = `Bearer ${config.apiKey}`; + } if ( config.authType === AuthType.LOGIN_WITH_GOOGLE || config.authType === AuthType.COMPUTE_ADC diff --git a/packages/core/src/utils/customHeaderUtils.test.ts b/packages/core/src/utils/customHeaderUtils.test.ts new file mode 100644 index 0000000000..7cae6331e8 --- /dev/null +++ b/packages/core/src/utils/customHeaderUtils.test.ts @@ -0,0 +1,91 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { parseCustomHeaders } from './customHeaderUtils.js'; + +describe('parseCustomHeaders', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it('should return an empty object if input is undefined', () => { + expect(parseCustomHeaders(undefined)).toEqual({}); + }); + + it('should return an empty object if input is empty string', () => { + expect(parseCustomHeaders('')).toEqual({}); + }); + + it('should parse a single header correctly', () => { + const input = 'Authorization: Bearer abc123'; + expect(parseCustomHeaders(input)).toEqual({ + Authorization: 'Bearer abc123', + }); + }); + + it('should parse multiple headers separated by commas', () => { + const input = + 'Authorization: Bearer abc123, Content-Type: application/json'; + expect(parseCustomHeaders(input)).toEqual({ + Authorization: 'Bearer abc123', + 'Content-Type': 'application/json', + }); + }); + + it('should ignore entries without colon', () => { + const input = 'Authorization Bearer abc123, Content-Type: application/json'; + expect(parseCustomHeaders(input)).toEqual({ + 'Content-Type': 'application/json', + }); + }); + + it('should trim whitespace around names and values', () => { + const input = + ' Authorization : Bearer abc123 , Content-Type : application/json '; + expect(parseCustomHeaders(input)).toEqual({ + Authorization: 'Bearer abc123', + 'Content-Type': 'application/json', + }); + }); + + it('should handle headers with colons in the value', () => { + const input = 'X-Custom: value:with:colons, Authorization: Bearer xyz'; + expect(parseCustomHeaders(input)).toEqual({ + 'X-Custom': 'value:with:colons', + Authorization: 'Bearer xyz', + }); + }); + + it('should skip headers with empty name', () => { + const input = ': no-name, Authorization: Bearer abc'; + expect(parseCustomHeaders(input)).toEqual({ + Authorization: 'Bearer abc', + }); + }); + + it('should skip completely empty entries', () => { + const input = ', , Authorization: Bearer abc'; + expect(parseCustomHeaders(input)).toEqual({ + Authorization: 'Bearer abc', + }); + }); + + it('should handle Authorization Bearer with different casing', () => { + const input = 'authorization: Bearer token123'; + expect(parseCustomHeaders(input)).toEqual({ + authorization: 'Bearer token123', + }); + }); + + it('should handle values with commas correctly', () => { + const input = 'X-Header: value,with,commas, Authorization: Bearer abc'; + expect(parseCustomHeaders(input)).toEqual({ + 'X-Header': 'value,with,commas', + Authorization: 'Bearer abc', + }); + }); +}); diff --git a/packages/core/src/utils/customHeaderUtils.ts b/packages/core/src/utils/customHeaderUtils.ts new file mode 100644 index 0000000000..d367e2783a --- /dev/null +++ b/packages/core/src/utils/customHeaderUtils.ts @@ -0,0 +1,41 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Parses custom headers and returns a map of key and vallues + */ +export function parseCustomHeaders( + envValue: string | undefined, +): Record { + const headers: Record = {}; + if (!envValue) { + return headers; + } + + // Split the string on commas that are followed by a header key (key:), + // but ignore commas that are part of a header value (including values with colons or commas) + for (const entry of envValue.split(/,(?=\s*[^,:]+:)/)) { + const trimmedEntry = entry.trim(); + if (!trimmedEntry) { + continue; + } + + const separatorIndex = trimmedEntry.indexOf(':'); + if (separatorIndex === -1) { + continue; + } + + const name = trimmedEntry.slice(0, separatorIndex).trim(); + const value = trimmedEntry.slice(separatorIndex + 1).trim(); + if (!name) { + continue; + } + + headers[name] = value; + } + + return headers; +}