Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,8 @@
"yargs": "^17.7.2"
},
"dependencies": {
"http-proxy-agent": "^7.0.2",
"https-proxy-agent": "^7.0.6",
Comment thread
scidomino marked this conversation as resolved.
Outdated
Comment thread
sotokisehiro marked this conversation as resolved.
Outdated
"ink": "npm:@jrichman/ink@6.6.9",
"latest-version": "^9.0.0",
"node-fetch-native": "^1.6.7",
Expand Down
1 change: 1 addition & 0 deletions packages/core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
"glob": "^12.0.0",
"google-auth-library": "^9.11.0",
"html-to-text": "^9.0.5",
"http-proxy-agent": "^7.0.2",
"https-proxy-agent": "^7.0.6",
"ignore": "^7.0.0",
"ipaddr.js": "^1.9.1",
Expand Down
170 changes: 170 additions & 0 deletions packages/core/src/core/contentGenerator.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ import {
} from './contentGenerator.js';
import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
import { GoogleGenAI } from '@google/genai';
import { HttpProxyAgent } from 'http-proxy-agent';
import { HttpsProxyAgent } from 'https-proxy-agent';
import type { Config } from '../config/config.js';
import { LoggingContentGenerator } from './loggingContentGenerator.js';
import { loadApiKey } from './apiKeyCredentialStorage.js';
Expand Down Expand Up @@ -424,6 +426,174 @@ describe('createContentGenerator', () => {
);
});

it('should inject HttpsProxyAgent into googleAuthOptions when proxy URL uses https://', async () => {
const mockConfigWithProxy = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue('https://proxy.example.com:8080'),
getUsageStatisticsEnabled: () => false,
getClientName: vi.fn().mockReturnValue(undefined),
} as unknown as Config;

const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator);

await createContentGenerator(
{
apiKey: 'test-api-key',
vertexai: true,
authType: AuthType.USE_VERTEX_AI,
proxy: 'https://proxy.example.com:8080',
},
mockConfigWithProxy,
);

expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
googleAuthOptions: {
clientOptions: {
transporterOptions: {
agent: expect.any(HttpsProxyAgent),
},
},
},
}),
);
});

it('should still use HttpsProxyAgent for HTTPS destinations even when proxy URL uses http://', async () => {
const mockConfigWithProxy = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue('http://proxy.example.com:8080'),
getUsageStatisticsEnabled: () => false,
getClientName: vi.fn().mockReturnValue(undefined),
} as unknown as Config;

const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator);

await createContentGenerator(
{
apiKey: 'test-api-key',
vertexai: true,
authType: AuthType.USE_VERTEX_AI,
proxy: 'http://proxy.example.com:8080',
},
mockConfigWithProxy,
);

expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
googleAuthOptions: {
clientOptions: {
transporterOptions: {
agent: expect.any(HttpsProxyAgent),
},
},
},
}),
);
});

it('should inject HttpProxyAgent when destination baseUrl uses http://', async () => {
const mockConfigWithProxy = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue('http://proxy.example.com:8080'),
getUsageStatisticsEnabled: () => false,
getClientName: vi.fn().mockReturnValue(undefined),
} as unknown as Config;

const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator);

vi.stubEnv('GOOGLE_VERTEX_BASE_URL', 'http://localhost:9999');
Comment thread
sotokisehiro marked this conversation as resolved.

await createContentGenerator(
{
apiKey: 'test-api-key',
vertexai: true,
authType: AuthType.USE_VERTEX_AI,
proxy: 'http://proxy.example.com:8080',
},
mockConfigWithProxy,
);

expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
googleAuthOptions: {
clientOptions: {
transporterOptions: {
agent: expect.any(HttpProxyAgent),
Comment thread
sotokisehiro marked this conversation as resolved.
},
},
},
}),
);
});

it('should trim whitespace from proxy URL before instantiating agent', async () => {
const mockConfigWithProxy = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue(' https://proxy.example.com:8080 '),
getUsageStatisticsEnabled: () => false,
getClientName: vi.fn().mockReturnValue(undefined),
} as unknown as Config;

const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator);

await createContentGenerator(
{
apiKey: 'test-api-key',
vertexai: true,
authType: AuthType.USE_VERTEX_AI,
proxy: ' https://proxy.example.com:8080 ',
},
mockConfigWithProxy,
);

expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
googleAuthOptions: {
clientOptions: {
transporterOptions: {
agent: expect.any(HttpsProxyAgent),
},
},
},
}),
);
});

it('should not include googleAuthOptions when no proxy is configured', async () => {
const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator);

await createContentGenerator(
{
apiKey: 'test-api-key',
vertexai: true,
authType: AuthType.USE_VERTEX_AI,
},
mockConfig,
);

const callArg = vi.mocked(GoogleGenAI).mock.calls[0][0] as Record<
string,
unknown
>;
expect(callArg).not.toHaveProperty('googleAuthOptions');
});

it('should pass api key as Authorization Header when GEMINI_API_KEY_AUTH_MECHANISM is set to bearer', async () => {
const mockConfig = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
Expand Down
16 changes: 16 additions & 0 deletions packages/core/src/core/contentGenerator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ import {
type EmbedContentResponse,
type EmbedContentParameters,
} from '@google/genai';
import { HttpProxyAgent } from 'http-proxy-agent';
import { HttpsProxyAgent } from 'https-proxy-agent';
import * as os from 'node:os';
import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
import { isCloudShell } from '../ide/detect-ide.js';
Expand Down Expand Up @@ -336,11 +338,25 @@ export async function createContentGenerator(
httpOptions.baseUrl = baseUrl;
}

const proxyUrl = config.proxy?.trim();
const proxyAgent = proxyUrl
? baseUrl?.startsWith('http://')
? new HttpProxyAgent(proxyUrl)
Comment thread
sotokisehiro marked this conversation as resolved.
: new HttpsProxyAgent(proxyUrl)
: undefined;
Comment thread
sotokisehiro marked this conversation as resolved.

const googleGenAI = new GoogleGenAI({
apiKey: config.apiKey === '' ? undefined : config.apiKey,
vertexai: config.vertexai ?? config.authType === AuthType.USE_VERTEX_AI,
httpOptions,
...(apiVersionEnv && { apiVersion: apiVersionEnv }),
...(proxyAgent && {
googleAuthOptions: {
clientOptions: {
transporterOptions: { agent: proxyAgent },
},
},
}),
Comment thread
sotokisehiro marked this conversation as resolved.
});
return new LoggingContentGenerator(googleGenAI.models, gcConfig);
}
Expand Down
Loading