fix: add httpOptions with headers field to CCPA client and set User-Agent header (#1103)

3 files changed, 25 insertions, 11 deletions

diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts
index 5922cb41..5df1502b 100644
--- a/packages/core/src/code_assist/codeAssist.ts
+++ b/packages/core/src/code_assist/codeAssist.ts
@@ -7,13 +7,15 @@
 import { ContentGenerator } from '../core/contentGenerator.js';
 import { getOauthClient } from './oauth2.js';
 import { setupUser } from './setup.js';
-import { CodeAssistServer } from './server.js';
+import { CodeAssistServer, HttpOptions } from './server.js';
 
-export async function createCodeAssistContentGenerator(): Promise<ContentGenerator> {
+export async function createCodeAssistContentGenerator(
+  httpOptions: HttpOptions,
+): Promise<ContentGenerator> {
   const oauth2Client = await getOauthClient();
   const projectId = await setupUser(
     oauth2Client,
     process.env.GOOGLE_CLOUD_PROJECT,
   );
-  return new CodeAssistServer(oauth2Client, projectId);
+  return new CodeAssistServer(oauth2Client, projectId, httpOptions);
 }
diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts
index aa6219fa..fc3a0ce2 100644
--- a/packages/core/src/code_assist/server.ts
+++ b/packages/core/src/code_assist/server.ts
@@ -28,6 +28,12 @@ import {
 } from './converter.js';
 import { PassThrough } from 'node:stream';
 
+/** HTTP options to be used in each of the requests. */
+export interface HttpOptions {
+  /** Additional HTTP headers to be sent with the request. */
+  headers?: Record<string, string>;
+}
+
 // TODO: Use production endpoint once it supports our methods.
 export const CODE_ASSIST_ENDPOINT =
   process.env.CODE_ASSIST_ENDPOINT ??
@@ -38,6 +44,7 @@ export class CodeAssistServer implements ContentGenerator {
   constructor(
     readonly auth: OAuth2Client,
     readonly projectId?: string,
+    readonly httpOptions: HttpOptions = {},
   ) {}
 
   async generateContentStream(
@@ -98,6 +105,7 @@ export class CodeAssistServer implements ContentGenerator {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',
+        ...this.httpOptions.headers,
       },
       responseType: 'json',
       body: JSON.stringify(req),
@@ -115,7 +123,10 @@ export class CodeAssistServer implements ContentGenerator {
       params: {
         alt: 'sse',
       },
-      headers: { 'Content-Type': 'application/json' },
+      headers: {
+        'Content-Type': 'application/json',
+        ...this.httpOptions.headers,
+      },
       responseType: 'stream',
       body: JSON.stringify(req),
     });
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index c8d1866a..3b276738 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -42,18 +42,19 @@ export type ContentGeneratorConfig = {
 export async function createContentGenerator(
   config: ContentGeneratorConfig,
 ): Promise<ContentGenerator> {
+  const version = process.env.CLI_VERSION || process.version;
+  const httpOptions = {
+    headers: {
+      'User-Agent': `GeminiCLI/${version}/(${process.platform}; ${process.arch})`,
+    },
+  };
   if (config.codeAssist) {
-    return createCodeAssistContentGenerator();
+    return createCodeAssistContentGenerator(httpOptions);
   }
-  const version = process.env.CLI_VERSION || process.version;
   const googleGenAI = new GoogleGenAI({
     apiKey: config.apiKey === '' ? undefined : config.apiKey,
     vertexai: config.vertexai,
-    httpOptions: {
-      headers: {
-        'User-Agent': `GeminiCLI/${version}/(${process.platform}; ${process.arch})`,
-      },
-    },
+    httpOptions,
   });
   return googleGenAI.models;
 }