Skip to content

Commit ec1e3af

Browse files
warjiangDcatfly
authored andcommitted
feat(cli): add explicit proxy option in cli (google-gemini#2526)
Co-authored-by: Dcatfly <[email protected]>
1 parent e7b4854 commit ec1e3af

File tree

13 files changed

+126
-12
lines changed

13 files changed

+126
-12
lines changed

docs/cli/configuration.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,9 @@ Arguments passed directly when running the CLI can override other configurations
370370
- Example: `gemini -e my-extension -e my-other-extension`
371371
- **`--list-extensions`** (**`-l`**):
372372
- Lists all available extensions and exits.
373+
- **`--proxy`**:
374+
- Sets the proxy for the CLI.
375+
- Example: `--proxy http://localhost:7890`.
373376
- **`--version`**:
374377
- Displays the version of the CLI.
375378

package-lock.json

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,73 @@ describe('loadCliConfig', () => {
187187
const config = await loadCliConfig(settings, [], 'test-session', argv);
188188
expect(config.getShowMemoryUsage()).toBe(true);
189189
});
190+
191+
it(`should leave proxy to empty by default`, async () => {
192+
process.argv = ['node', 'script.js'];
193+
const argv = await parseArguments();
194+
const settings: Settings = {};
195+
const config = await loadCliConfig(settings, [], 'test-session', argv);
196+
expect(config.getProxy()).toBeFalsy();
197+
});
198+
199+
const proxy_url = 'http://localhost:7890';
200+
const testCases = [
201+
{
202+
input: {
203+
env_name: 'https_proxy',
204+
proxy_url,
205+
},
206+
expected: proxy_url,
207+
},
208+
{
209+
input: {
210+
env_name: 'http_proxy',
211+
proxy_url,
212+
},
213+
expected: proxy_url,
214+
},
215+
{
216+
input: {
217+
env_name: 'HTTPS_PROXY',
218+
proxy_url,
219+
},
220+
expected: proxy_url,
221+
},
222+
{
223+
input: {
224+
env_name: 'HTTP_PROXY',
225+
proxy_url,
226+
},
227+
expected: proxy_url,
228+
},
229+
];
230+
testCases.forEach(({ input, expected }) => {
231+
it(`should set proxy to ${expected} according to environment variable [${input.env_name}]`, async () => {
232+
process.env[input.env_name] = input.proxy_url;
233+
process.argv = ['node', 'script.js'];
234+
const argv = await parseArguments();
235+
const settings: Settings = {};
236+
const config = await loadCliConfig(settings, [], 'test-session', argv);
237+
expect(config.getProxy()).toBe(expected);
238+
});
239+
});
240+
241+
it('should set proxy when --proxy flag is present', async () => {
242+
process.argv = ['node', 'script.js', '--proxy', 'http://localhost:7890'];
243+
const argv = await parseArguments();
244+
const settings: Settings = {};
245+
const config = await loadCliConfig(settings, [], 'test-session', argv);
246+
expect(config.getProxy()).toBe('http://localhost:7890');
247+
});
248+
249+
it('should prioritize CLI flag over environment variable for proxy (CLI http://localhost:7890, environment variable http://localhost:7891)', async () => {
250+
process.env['http_proxy'] = 'http://localhost:7891';
251+
process.argv = ['node', 'script.js', '--proxy', 'http://localhost:7890'];
252+
const argv = await parseArguments();
253+
const settings: Settings = {};
254+
const config = await loadCliConfig(settings, [], 'test-session', argv);
255+
expect(config.getProxy()).toBe('http://localhost:7890');
256+
});
190257
});
191258

192259
describe('loadCliConfig telemetry', () => {

packages/cli/src/config/config.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ export interface CliArgs {
5757
extensions: string[] | undefined;
5858
listExtensions: boolean | undefined;
5959
ideMode: boolean | undefined;
60+
proxy: string | undefined;
6061
}
6162

6263
export async function parseArguments(): Promise<CliArgs> {
@@ -182,7 +183,11 @@ export async function parseArguments(): Promise<CliArgs> {
182183
type: 'boolean',
183184
description: 'Run in IDE mode?',
184185
})
185-
186+
.option('proxy', {
187+
type: 'string',
188+
description:
189+
'Proxy for gemini client, like schema://user:password@host:port',
190+
})
186191
.version(await getCliVersion()) // This will enable the --version flag based on package.json
187192
.alias('v', 'version')
188193
.help()
@@ -380,6 +385,7 @@ export async function loadCliConfig(
380385
},
381386
checkpointing: argv.checkpointing || settings.checkpointing?.enabled,
382387
proxy:
388+
argv.proxy ||
383389
process.env.HTTPS_PROXY ||
384390
process.env.https_proxy ||
385391
process.env.HTTP_PROXY ||

packages/core/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
"glob": "^10.4.5",
3838
"google-auth-library": "^9.11.0",
3939
"html-to-text": "^9.0.5",
40+
"https-proxy-agent": "^7.0.6",
4041
"ignore": "^7.0.0",
4142
"micromatch": "^4.0.8",
4243
"open": "^10.1.2",

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ vi.mock('node:readline');
3434

3535
const mockConfig = {
3636
getNoBrowser: () => false,
37+
getProxy: () => 'http://test.proxy.com:8080',
3738
} as unknown as Config;
3839

3940
// Mock fetch globally
@@ -175,6 +176,7 @@ describe('oauth2', () => {
175176
it('should perform login with user code', async () => {
176177
const mockConfigWithNoBrowser = {
177178
getNoBrowser: () => true,
179+
getProxy: () => 'http://test.proxy.com:8080',
178180
} as unknown as Config;
179181

180182
const mockCodeVerifier = {

packages/core/src/code_assist/oauth2.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,9 @@ export async function getOauthClient(
7373
const client = new OAuth2Client({
7474
clientId: OAUTH_CLIENT_ID,
7575
clientSecret: OAUTH_CLIENT_SECRET,
76+
transporterOptions: {
77+
proxy: config.getProxy(),
78+
},
7679
});
7780

7881
client.on('tokens', async (tokens: Credentials) => {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ describe('createContentGeneratorConfig', () => {
6868
getModel: vi.fn().mockReturnValue('gemini-pro'),
6969
setModel: vi.fn(),
7070
flashFallbackHandler: vi.fn(),
71+
getProxy: vi.fn(),
7172
} as unknown as Config;
7273

7374
beforeEach(() => {

packages/core/src/core/contentGenerator.ts

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ export type ContentGeneratorConfig = {
5050
apiKey?: string;
5151
vertexai?: boolean;
5252
authType?: AuthType | undefined;
53+
proxy?: string | undefined;
5354
};
5455

5556
export function createContentGeneratorConfig(
@@ -67,6 +68,7 @@ export function createContentGeneratorConfig(
6768
const contentGeneratorConfig: ContentGeneratorConfig = {
6869
model: effectiveModel,
6970
authType,
71+
proxy: config?.getProxy(),
7072
};
7173

7274
// 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(
8385
getEffectiveModel(
8486
contentGeneratorConfig.apiKey,
8587
contentGeneratorConfig.model,
86-
).then((newModel) => {
87-
if (newModel !== contentGeneratorConfig.model) {
88-
config.flashFallbackHandler?.(contentGeneratorConfig.model, newModel);
89-
}
90-
});
88+
contentGeneratorConfig.proxy,
89+
);
9190

9291
return contentGeneratorConfig;
9392
}

packages/core/src/core/modelCheck.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7+
import { setGlobalDispatcher, ProxyAgent } from 'undici';
78
import {
89
DEFAULT_GEMINI_MODEL,
910
DEFAULT_GEMINI_FLASH_MODEL,
@@ -20,6 +21,7 @@ import {
2021
export async function getEffectiveModel(
2122
apiKey: string,
2223
currentConfiguredModel: string,
24+
proxy?: string,
2325
): Promise<string> {
2426
if (currentConfiguredModel !== DEFAULT_GEMINI_MODEL) {
2527
// 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(
4345
const timeoutId = setTimeout(() => controller.abort(), 2000); // 500ms timeout for the request
4446

4547
try {
48+
if (proxy) {
49+
setGlobalDispatcher(new ProxyAgent(proxy));
50+
}
4651
const response = await fetch(endpoint, {
4752
method: 'POST',
4853
headers: {

0 commit comments

Comments
 (0)