summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--docs/cli/configuration.md3
-rw-r--r--package-lock.json1
-rw-r--r--packages/cli/src/config/config.test.ts67
-rw-r--r--packages/cli/src/config/config.ts8
-rw-r--r--packages/core/package.json1
-rw-r--r--packages/core/src/code_assist/oauth2.test.ts2
-rw-r--r--packages/core/src/code_assist/oauth2.ts3
-rw-r--r--packages/core/src/core/contentGenerator.test.ts1
-rw-r--r--packages/core/src/core/contentGenerator.ts9
-rw-r--r--packages/core/src/core/modelCheck.ts5
-rw-r--r--packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts32
-rw-r--r--packages/core/src/tools/web-fetch.test.ts1
-rw-r--r--packages/core/src/tools/web-fetch.ts5
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<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 ||
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<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: {
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<WebFetchToolParams, ToolResult> {
type: Type.OBJECT,
},
);
+ const proxy = config.getProxy();
+ if (proxy) {
+ setGlobalDispatcher(new ProxyAgent(proxy as string));
+ }
}
private async executeFallback(