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

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);
 }

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

packages/core/src/core/contentGenerator.ts

@@ -42,18 +42,19 @@ export type ContentGeneratorConfig = {
 export async function createContentGenerator(
   config: ContentGeneratorConfig,
 ): Promise<ContentGenerator> {
-  if (config.codeAssist) {
-    return createCodeAssistContentGenerator();
-  }
   const version = process.env.CLI_VERSION || process.version;
-  const googleGenAI = new GoogleGenAI({
+  const httpOptions = {
-    apiKey: config.apiKey === '' ? undefined : config.apiKey,
-    vertexai: config.vertexai,
-    httpOptions: {
     headers: {
       'User-Agent': `GeminiCLI/${version}/(${process.platform}; ${process.arch})`,
     },
-    },
+  };
+  if (config.codeAssist) {
+    return createCodeAssistContentGenerator(httpOptions);
+  }
+  const googleGenAI = new GoogleGenAI({
+    apiKey: config.apiKey === '' ? undefined : config.apiKey,
+    vertexai: config.vertexai,
+    httpOptions,
   });
   return googleGenAI.models;
 }