From 606a7702de6659b9a46bcb9c2cbca84a1a33945d Mon Sep 17 00:00:00 2001 From: warjiang <1096409085@qq.com> Date: Fri, 18 Jul 2025 02:57:37 +0800 Subject: [PATCH] feat(cli): add explicit proxy option in cli (#2526) Co-authored-by: Dcatfly --- docs/cli/configuration.md | 3 + package-lock.json | 1 + packages/cli/src/config/config.test.ts | 67 +++++++++++++++++++ packages/cli/src/config/config.ts | 8 ++- packages/core/package.json | 1 + packages/core/src/code_assist/oauth2.test.ts | 2 + packages/core/src/code_assist/oauth2.ts | 3 + .../core/src/core/contentGenerator.test.ts | 1 + packages/core/src/core/contentGenerator.ts | 9 ++- packages/core/src/core/modelCheck.ts | 5 ++ .../clearcut-logger/clearcut-logger.ts | 32 +++++++-- packages/core/src/tools/web-fetch.test.ts | 1 + packages/core/src/tools/web-fetch.ts | 5 ++ 13 files changed, 126 insertions(+), 12 deletions(-) diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index 8ac4fac9..d1ddf1da 100644 --- a/docs/cli/configuration.md +++ b/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. diff --git a/package-lock.json b/package-lock.json index 05c4e713..9be1bfd0 100644 --- a/package-lock.json +++ b/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", diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index 39259fe1..cc0f112a 100644 --- a/packages/cli/src/config/config.test.ts +++ b/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', () => { diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index c7c23901..f76d6c60 100644 --- a/packages/cli/src/config/config.ts +++ b/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 { @@ -182,7 +183,11 @@ export async function parseArguments(): Promise { 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 || diff --git a/packages/core/package.json b/packages/core/package.json index c34e62fd..240cccfa 100644 --- a/packages/core/package.json +++ b/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", diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts index 4661f49a..7fa98e17 100644 --- a/packages/core/src/code_assist/oauth2.test.ts +++ b/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 = { diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index 48449b5e..3c3f7055 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/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) => { diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index c50678af..78eee386 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/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(() => { diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 109c0ffc..b381de8e 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/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; } diff --git a/packages/core/src/core/modelCheck.ts b/packages/core/src/core/modelCheck.ts index c9078aaf..25d86993 100644 --- a/packages/core/src/core/modelCheck.ts +++ b/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 { 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: { diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts index e42aa677..5addd99e 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts +++ b/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); diff --git a/packages/core/src/tools/web-fetch.test.ts b/packages/core/src/tools/web-fetch.test.ts index f4e3a652..6be9d504 100644 --- a/packages/core/src/tools/web-fetch.test.ts +++ b/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', () => { diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index 0f5be969..ee06880e 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/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 { type: Type.OBJECT, }, ); + const proxy = config.getProxy(); + if (proxy) { + setGlobalDispatcher(new ProxyAgent(proxy as string)); + } } private async executeFallback(