2 files changed, 17 insertions, 4 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),
     });