Add Databricks auth support and custom header option to gemini cli (#11893)

Co-authored-by: Taylor Mullen <ntaylormullen@google.com>

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 <key>` 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.

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,

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<string, string> = {
+      ...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

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

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<string, string> {
+  const headers: Record<string, string> = {};
+  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;
+}