update

packages/core/src/config/config.test.ts

@@ -500,6 +500,8 @@ describe('Server Config (config.ts)', () => {
         config,
         authType,
         undefined,
+        undefined,
+        undefined,
       );
       // Verify that contentGeneratorConfig is updated
       expect(config.getContentGeneratorConfig()).toEqual(mockContentConfig);

packages/core/src/config/config.ts

@@ -1126,7 +1126,12 @@ export class Config {
     return this.contentGenerator;
   }
 
-  async refreshAuth(authMethod: AuthType, apiKey?: string, baseUrl?: string) {
+  async refreshAuth(
+    authMethod: AuthType,
+    apiKey?: string,
+    baseUrl?: string,
+    headers?: Record<string, string>,
+  ) {
     // Reset availability service when switching auth
     this.modelAvailabilityService.reset();
 
@@ -1154,6 +1159,7 @@ export class Config {
       authMethod,
       apiKey,
       baseUrl,
+      headers,
     );
     this.contentGenerator = await createContentGenerator(
       newContentGeneratorConfig,

packages/core/src/core/contentGenerator.test.ts

@@ -174,7 +174,6 @@ describe('createContentGenerator', () => {
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
           'X-Test-Header': 'test-value',
-          'Another-Header': 'another value',
         }),
       },
       AuthType.LOGIN_WITH_GOOGLE,
@@ -475,6 +474,64 @@ describe('createContentGenerator', () => {
       apiVersion: 'v1alpha',
     });
   });
+
+  it('should include config.headers for Gemini auth but not for others', 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);
+
+    // Test USE_GEMINI gets headers
+    await createContentGenerator(
+      {
+        apiKey: 'test-api-key',
+        authType: AuthType.USE_GEMINI,
+        headers: { 'X-Custom-Config-Header': 'custom-value' },
+      },
+      mockConfig,
+    );
+
+    expect(GoogleGenAI).toHaveBeenCalledWith(
+      expect.objectContaining({
+        httpOptions: expect.objectContaining({
+          headers: expect.objectContaining({
+            'X-Custom-Config-Header': 'custom-value',
+          }),
+        }),
+      }),
+    );
+
+    // Test LOGIN_WITH_GOOGLE does NOT get headers (unless they were in baseHeaders, which they are not)
+    const mockCodeAssistGenerator = {} as unknown as ContentGenerator;
+    vi.mocked(createCodeAssistContentGenerator).mockResolvedValue(
+      mockCodeAssistGenerator as never,
+    );
+
+    await createContentGenerator(
+      {
+        authType: AuthType.LOGIN_WITH_GOOGLE,
+        headers: { 'X-Should-Not-Be-Here': 'nope' },
+      },
+      mockConfig,
+    );
+
+    expect(createCodeAssistContentGenerator).toHaveBeenCalledWith(
+      expect.objectContaining({
+        headers: expect.not.objectContaining({
+          'X-Should-Not-Be-Here': 'nope',
+        }),
+      }),
+      AuthType.LOGIN_WITH_GOOGLE,
+      mockConfig,
+      undefined,
+    );
+  });
 });
 
 describe('createContentGeneratorConfig', () => {

packages/core/src/core/contentGenerator.ts

@@ -86,6 +86,7 @@ export type ContentGeneratorConfig = {
   vertexai?: boolean;
   authType?: AuthType;
   proxy?: string;
+  headers?: Record<string, string>;
 };
 
 export async function createContentGeneratorConfig(
@@ -93,6 +94,7 @@ export async function createContentGeneratorConfig(
   authType: AuthType | undefined,
   apiKey?: string,
   baseUrl?: string,
+  headers?: Record<string, string>,
 ): Promise<ContentGeneratorConfig> {
   const geminiApiKey =
     apiKey ||
@@ -109,6 +111,7 @@ export async function createContentGeneratorConfig(
   const contentGeneratorConfig: ContentGeneratorConfig = {
     authType,
     proxy: config?.getProxy(),
+    headers,
   };
 
   // If we are using Google auth or we are in Cloud Shell, there is nothing else to validate for now
@@ -200,7 +203,10 @@ export async function createContentGenerator(
       config.authType === AuthType.USE_GEMINI ||
       config.authType === AuthType.USE_VERTEX_AI
     ) {
-      let headers: Record<string, string> = { ...baseHeaders };
+      let headers: Record<string, string> = {
+        ...baseHeaders,
+        ...config.headers,
+      };
       if (gcConfig?.getUsageStatisticsEnabled()) {
         const installationManager = new InstallationManager();
         const installationId = installationManager.getInstallationId();