feat(cli): add explicit proxy option in cli (#2526)

Co-authored-by: Dcatfly <dcatfly@gmail.com>

docs/cli/configuration.md

@@ -370,6 +370,9 @@ Arguments passed directly when running the CLI can override other configurations
   - Example: `gemini -e my-extension -e my-other-extension`
 - **`--list-extensions`** (**`-l`**):
   - Lists all available extensions and exits.
+- **`--proxy`**:
+  - Sets the proxy for the CLI.
+  - Example: `--proxy http://localhost:7890`.
 - **`--version`**:
   - Displays the version of the CLI.
 

package-lock.json

@@ -11940,6 +11940,7 @@
         "glob": "^10.4.5",
         "google-auth-library": "^9.11.0",
         "html-to-text": "^9.0.5",
+        "https-proxy-agent": "^7.0.6",
         "ignore": "^7.0.0",
         "micromatch": "^4.0.8",
         "open": "^10.1.2",

packages/cli/src/config/config.test.ts

@@ -187,6 +187,73 @@ describe('loadCliConfig', () => {
     const config = await loadCliConfig(settings, [], 'test-session', argv);
     expect(config.getShowMemoryUsage()).toBe(true);
   });
+
+  it(`should leave proxy to empty by default`, async () => {
+    process.argv = ['node', 'script.js'];
+    const argv = await parseArguments();
+    const settings: Settings = {};
+    const config = await loadCliConfig(settings, [], 'test-session', argv);
+    expect(config.getProxy()).toBeFalsy();
+  });
+
+  const proxy_url = 'http://localhost:7890';
+  const testCases = [
+    {
+      input: {
+        env_name: 'https_proxy',
+        proxy_url,
+      },
+      expected: proxy_url,
+    },
+    {
+      input: {
+        env_name: 'http_proxy',
+        proxy_url,
+      },
+      expected: proxy_url,
+    },
+    {
+      input: {
+        env_name: 'HTTPS_PROXY',
+        proxy_url,
+      },
+      expected: proxy_url,
+    },
+    {
+      input: {
+        env_name: 'HTTP_PROXY',
+        proxy_url,
+      },
+      expected: proxy_url,
+    },
+  ];
+  testCases.forEach(({ input, expected }) => {
+    it(`should set proxy to ${expected} according to environment variable [${input.env_name}]`, async () => {
+      process.env[input.env_name] = input.proxy_url;
+      process.argv = ['node', 'script.js'];
+      const argv = await parseArguments();
+      const settings: Settings = {};
+      const config = await loadCliConfig(settings, [], 'test-session', argv);
+      expect(config.getProxy()).toBe(expected);
+    });
+  });
+
+  it('should set proxy when --proxy flag is present', async () => {
+    process.argv = ['node', 'script.js', '--proxy', 'http://localhost:7890'];
+    const argv = await parseArguments();
+    const settings: Settings = {};
+    const config = await loadCliConfig(settings, [], 'test-session', argv);
+    expect(config.getProxy()).toBe('http://localhost:7890');
+  });
+
+  it('should prioritize CLI flag over environment variable for proxy (CLI http://localhost:7890, environment variable http://localhost:7891)', async () => {
+    process.env['http_proxy'] = 'http://localhost:7891';
+    process.argv = ['node', 'script.js', '--proxy', 'http://localhost:7890'];
+    const argv = await parseArguments();
+    const settings: Settings = {};
+    const config = await loadCliConfig(settings, [], 'test-session', argv);
+    expect(config.getProxy()).toBe('http://localhost:7890');
+  });
 });
 
 describe('loadCliConfig telemetry', () => {

packages/cli/src/config/config.ts

@@ -57,6 +57,7 @@ export interface CliArgs {
   extensions: string[] | undefined;
   listExtensions: boolean | undefined;
   ideMode: boolean | undefined;
+  proxy: string | undefined;
 }
 
 export async function parseArguments(): Promise<CliArgs> {
@@ -182,7 +183,11 @@ export async function parseArguments(): Promise<CliArgs> {
       type: 'boolean',
       description: 'Run in IDE mode?',
     })
-
+    .option('proxy', {
+      type: 'string',
+      description:
+        'Proxy for gemini client, like schema://user:password@host:port',
+    })
     .version(await getCliVersion()) // This will enable the --version flag based on package.json
     .alias('v', 'version')
     .help()
@@ -380,6 +385,7 @@ export async function loadCliConfig(
     },
     checkpointing: argv.checkpointing || settings.checkpointing?.enabled,
     proxy:
+      argv.proxy ||
       process.env.HTTPS_PROXY ||
       process.env.https_proxy ||
       process.env.HTTP_PROXY ||

packages/core/package.json

@@ -37,6 +37,7 @@
     "glob": "^10.4.5",
     "google-auth-library": "^9.11.0",
     "html-to-text": "^9.0.5",
+    "https-proxy-agent": "^7.0.6",
     "ignore": "^7.0.0",
     "micromatch": "^4.0.8",
     "open": "^10.1.2",

packages/core/src/code_assist/oauth2.test.ts

@@ -34,6 +34,7 @@ vi.mock('node:readline');
 
 const mockConfig = {
   getNoBrowser: () => false,
+  getProxy: () => 'http://test.proxy.com:8080',
 } as unknown as Config;
 
 // Mock fetch globally
@@ -175,6 +176,7 @@ describe('oauth2', () => {
   it('should perform login with user code', async () => {
     const mockConfigWithNoBrowser = {
       getNoBrowser: () => true,
+      getProxy: () => 'http://test.proxy.com:8080',
     } as unknown as Config;
 
     const mockCodeVerifier = {

packages/core/src/code_assist/oauth2.ts

@@ -73,6 +73,9 @@ export async function getOauthClient(
   const client = new OAuth2Client({
     clientId: OAUTH_CLIENT_ID,
     clientSecret: OAUTH_CLIENT_SECRET,
+    transporterOptions: {
+      proxy: config.getProxy(),
+    },
   });
 
   client.on('tokens', async (tokens: Credentials) => {

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

@@ -68,6 +68,7 @@ describe('createContentGeneratorConfig', () => {
     getModel: vi.fn().mockReturnValue('gemini-pro'),
     setModel: vi.fn(),
     flashFallbackHandler: vi.fn(),
+    getProxy: vi.fn(),
   } as unknown as Config;
 
   beforeEach(() => {

packages/core/src/core/contentGenerator.ts

@@ -50,6 +50,7 @@ export type ContentGeneratorConfig = {
   apiKey?: string;
   vertexai?: boolean;
   authType?: AuthType | undefined;
+  proxy?: string | undefined;
 };
 
 export function createContentGeneratorConfig(
@@ -67,6 +68,7 @@ export function createContentGeneratorConfig(
   const contentGeneratorConfig: ContentGeneratorConfig = {
     model: effectiveModel,
     authType,
+    proxy: config?.getProxy(),
   };
 
   // If we are using Google auth or we are in Cloud Shell, there is nothing else to validate for now
@@ -83,11 +85,8 @@ export function createContentGeneratorConfig(
     getEffectiveModel(
       contentGeneratorConfig.apiKey,
       contentGeneratorConfig.model,
-    ).then((newModel) => {
-      if (newModel !== contentGeneratorConfig.model) {
-        config.flashFallbackHandler?.(contentGeneratorConfig.model, newModel);
-      }
-    });
+      contentGeneratorConfig.proxy,
+    );
 
     return contentGeneratorConfig;
   }

packages/core/src/core/modelCheck.ts

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { setGlobalDispatcher, ProxyAgent } from 'undici';
 import {
   DEFAULT_GEMINI_MODEL,
   DEFAULT_GEMINI_FLASH_MODEL,
@@ -20,6 +21,7 @@ import {
 export async function getEffectiveModel(
   apiKey: string,
   currentConfiguredModel: string,
+  proxy?: string,
 ): Promise<string> {
   if (currentConfiguredModel !== DEFAULT_GEMINI_MODEL) {
     // Only check if the user is trying to use the specific pro model we want to fallback from.
@@ -43,6 +45,9 @@ export async function getEffectiveModel(
   const timeoutId = setTimeout(() => controller.abort(), 2000); // 500ms timeout for the request
 
   try {
+    if (proxy) {
+      setGlobalDispatcher(new ProxyAgent(proxy));
+    }
     const response = await fetch(endpoint, {
       method: 'POST',
       headers: {

packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts

@@ -6,6 +6,8 @@
 
 import { Buffer } from 'buffer';
 import * as https from 'https';
+import { HttpsProxyAgent } from 'https-proxy-agent';
+
 import {
   StartSessionEvent,
   EndSessionEvent,
@@ -132,12 +134,18 @@ export class ClearcutLogger {
         headers: { 'Content-Length': Buffer.byteLength(body) },
       };
       const bufs: Buffer[] = [];
-      const req = https.request(options, (res) => {
-        res.on('data', (buf) => bufs.push(buf));
-        res.on('end', () => {
-          resolve(Buffer.concat(bufs));
-        });
-      });
+      const req = https.request(
+        {
+          ...options,
+          agent: this.getProxyAgent(),
+        },
+        (res) => {
+          res.on('data', (buf) => bufs.push(buf));
+          res.on('end', () => {
+            resolve(Buffer.concat(bufs));
+          });
+        },
+      );
       req.on('error', (e) => {
         if (this.config?.getDebugMode()) {
           console.log('Clearcut POST request error: ', e);
@@ -499,6 +507,18 @@ export class ClearcutLogger {
     });
   }
 
+  getProxyAgent() {
+    const proxyUrl = this.config?.getProxy();
+    if (!proxyUrl) return undefined;
+    // undici which is widely used in the repo can only support http & https proxy protocol,
+    // https://github.com/nodejs/undici/issues/2224
+    if (proxyUrl.startsWith('http')) {
+      return new HttpsProxyAgent(proxyUrl);
+    } else {
+      throw new Error('Unsupported proxy type');
+    }
+  }
+
   shutdown() {
     const event = new EndSessionEvent(this.config);
     this.logEndSessionEvent(event);

packages/core/src/tools/web-fetch.test.ts

@@ -13,6 +13,7 @@ describe('WebFetchTool', () => {
   const mockConfig = {
     getApprovalMode: vi.fn(),
     setApprovalMode: vi.fn(),
+    getProxy: vi.fn(),
   } as unknown as Config;
 
   describe('shouldConfirmExecute', () => {

packages/core/src/tools/web-fetch.ts

@@ -17,6 +17,7 @@ import { Config, ApprovalMode } from '../config/config.js';
 import { getResponseText } from '../utils/generateContentResponseUtilities.js';
 import { fetchWithTimeout, isPrivateIp } from '../utils/fetch.js';
 import { convert } from 'html-to-text';
+import { ProxyAgent, setGlobalDispatcher } from 'undici';
 
 const URL_FETCH_TIMEOUT_MS = 10000;
 const MAX_CONTENT_LENGTH = 100000;
@@ -81,6 +82,10 @@ export class WebFetchTool extends BaseTool<WebFetchToolParams, ToolResult> {
         type: Type.OBJECT,
       },
     );
+    const proxy = config.getProxy();
+    if (proxy) {
+      setGlobalDispatcher(new ProxyAgent(proxy as string));
+    }
   }
 
   private async executeFallback(