feat(core): refine User-Agent for VS Code traffic (unified format) (#23256)

This commit is contained in:
Sehoon Shon
2026-03-20 15:31:01 -04:00
committed by GitHub
parent cca595971d
commit 05e4ea80ee
3 changed files with 156 additions and 11 deletions

View File

@@ -131,6 +131,10 @@ describe('createContentGenerator', () => {
// Set a fixed version for testing
vi.stubEnv('CLI_VERSION', '1.2.3');
vi.stubEnv('TERM_PROGRAM', 'iTerm.app');
vi.stubEnv('VSCODE_PID', '');
vi.stubEnv('GITHUB_SHA', '');
vi.stubEnv('GEMINI_CLI_SURFACE', '');
const mockGenerator = {
models: {},
@@ -149,7 +153,7 @@ describe('createContentGenerator', () => {
httpOptions: expect.objectContaining({
headers: expect.objectContaining({
'User-Agent': expect.stringMatching(
/GeminiCLI\/1\.2\.3\/gemini-pro \(.*; .*; .*\)/,
/GeminiCLI\/1\.2\.3\/gemini-pro \(.*; .*; terminal\)/,
),
}),
}),
@@ -159,7 +163,7 @@ describe('createContentGenerator', () => {
);
});
it('should include clientName prefix in User-Agent when specified', async () => {
it('should use standard User-Agent for a2a-server running outside VS Code', async () => {
const mockConfig = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue(undefined),
@@ -169,6 +173,10 @@ describe('createContentGenerator', () => {
// Set a fixed version for testing
vi.stubEnv('CLI_VERSION', '1.2.3');
vi.stubEnv('TERM_PROGRAM', 'iTerm.app');
vi.stubEnv('VSCODE_PID', '');
vi.stubEnv('GITHUB_SHA', '');
vi.stubEnv('GEMINI_CLI_SURFACE', '');
const mockGenerator = {
models: {},
@@ -185,7 +193,7 @@ describe('createContentGenerator', () => {
httpOptions: expect.objectContaining({
headers: expect.objectContaining({
'User-Agent': expect.stringMatching(
/GeminiCLI-a2a-server\/.*\/gemini-pro \(.*; .*; .*\)/,
/GeminiCLI-a2a-server\/1\.2\.3\/gemini-pro \(.*; .*; terminal\)/,
),
}),
}),
@@ -193,6 +201,113 @@ describe('createContentGenerator', () => {
);
});
it('should include unified User-Agent for a2a-server (VS Code Agent Mode)', async () => {
const mockConfig = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue(undefined),
getUsageStatisticsEnabled: () => true,
getClientName: vi.fn().mockReturnValue('a2a-server'),
} as unknown as Config;
// Set a fixed version for testing
vi.stubEnv('CLI_VERSION', '1.2.3');
// Mock the environment variable that the VS Code extension host would provide to the a2a-server process
vi.stubEnv('VSCODE_PID', '12345');
vi.stubEnv('TERM_PROGRAM', 'vscode');
vi.stubEnv('TERM_PROGRAM_VERSION', '1.85.0');
const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
await createContentGenerator(
{ apiKey: 'test-api-key', authType: AuthType.USE_GEMINI },
mockConfig,
undefined,
);
expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
httpOptions: expect.objectContaining({
headers: expect.objectContaining({
'User-Agent': expect.stringMatching(
/CloudCodeVSCode\/1\.2\.3 \(aidev_client; os_type=.*; os_version=.*; arch=.*; host_path=VSCode\/1\.85\.0; proxy_client=geminicli\)/,
),
}),
}),
}),
);
});
it('should include clientName prefix in User-Agent when specified (non-VSCode)', async () => {
const mockConfig = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue(undefined),
getUsageStatisticsEnabled: () => true,
getClientName: vi.fn().mockReturnValue('my-client'),
} as unknown as Config;
// Set a fixed version for testing
vi.stubEnv('CLI_VERSION', '1.2.3');
vi.stubEnv('TERM_PROGRAM', 'iTerm.app');
vi.stubEnv('VSCODE_PID', '');
vi.stubEnv('GITHUB_SHA', '');
vi.stubEnv('GEMINI_CLI_SURFACE', '');
const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
await createContentGenerator(
{ apiKey: 'test-api-key', authType: AuthType.USE_GEMINI },
mockConfig,
undefined,
);
expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
httpOptions: expect.objectContaining({
headers: expect.objectContaining({
'User-Agent': expect.stringMatching(
/GeminiCLI-my-client\/1\.2\.3\/gemini-pro \(.*; .*; terminal\)/,
),
}),
}),
}),
);
});
it('should allow custom headers to override User-Agent', async () => {
const mockConfig = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
getProxy: vi.fn().mockReturnValue(undefined),
getUsageStatisticsEnabled: () => true,
getClientName: vi.fn().mockReturnValue(undefined),
} as unknown as Config;
vi.stubEnv('GEMINI_CLI_CUSTOM_HEADERS', 'User-Agent:MyCustomUA');
const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
await createContentGenerator(
{ apiKey: 'test-api-key', authType: AuthType.USE_GEMINI },
mockConfig,
undefined,
);
expect(GoogleGenAI).toHaveBeenCalledWith(
expect.objectContaining({
httpOptions: expect.objectContaining({
headers: expect.objectContaining({
'User-Agent': 'MyCustomUA',
}),
}),
}),
);
});
it('should include custom headers from GEMINI_CLI_CUSTOM_HEADERS for Code Assist requests', async () => {
const mockGenerator = {} as unknown as ContentGenerator;
vi.mocked(createCodeAssistContentGenerator).mockResolvedValue(

View File

@@ -13,7 +13,9 @@ import {
type EmbedContentResponse,
type EmbedContentParameters,
} from '@google/genai';
import * as os from 'node:os';
import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
import { isCloudShell } from '../ide/detect-ide.js';
import type { Config } from '../config/config.js';
import { loadApiKey } from './apiKeyCredentialStorage.js';
@@ -185,19 +187,46 @@ export async function createContentGenerator(
const customHeadersEnv =
process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined;
const clientName = gcConfig.getClientName();
const userAgentPrefix = clientName
? `GeminiCLI-${clientName}`
: 'GeminiCLI';
const surface = determineSurface();
const userAgent = `${userAgentPrefix}/${version}/${model} (${process.platform}; ${process.arch}; ${surface})`;
let userAgent: string;
// Use unified format for VS Code traffic.
// Note: We don't automatically assume a2a-server is VS Code,
// as it could be used by other clients unless the surface explicitly says 'vscode'.
if (clientName === 'acp-vscode' || surface === 'vscode') {
const osTypeMap: Record<string, string> = {
darwin: 'macOS',
win32: 'Windows',
linux: 'Linux',
};
const osType = osTypeMap[process.platform] || process.platform;
const osVersion = os.release();
const arch = process.arch;
const vscodeVersion = process.env['TERM_PROGRAM_VERSION'] || 'unknown';
let hostPath = `VSCode/${vscodeVersion}`;
if (isCloudShell()) {
const cloudShellVersion =
process.env['CLOUD_SHELL_VERSION'] || 'unknown';
hostPath += ` > CloudShell/${cloudShellVersion}`;
}
userAgent = `CloudCodeVSCode/${version} (aidev_client; os_type=${osType}; os_version=${osVersion}; arch=${arch}; host_path=${hostPath}; proxy_client=geminicli)`;
} else {
const userAgentPrefix = clientName
? `GeminiCLI-${clientName}`
: 'GeminiCLI';
userAgent = `${userAgentPrefix}/${version}/${model} (${process.platform}; ${process.arch}; ${surface})`;
}
const customHeadersMap = parseCustomHeaders(customHeadersEnv);
const apiKeyAuthMechanism =
process.env['GEMINI_API_KEY_AUTH_MECHANISM'] || 'x-goog-api-key';
const apiVersionEnv = process.env['GOOGLE_GENAI_API_VERSION'];
const baseHeaders: Record<string, string> = {
...customHeadersMap,
'User-Agent': userAgent,
...customHeadersMap,
};
if (

View File

@@ -37,9 +37,10 @@ export function determineSurface(): string {
return ide.name;
}
// If the detected IDE is 'vscode', we only accept it if TERM_PROGRAM confirms it.
// This prevents generic terminals from being misidentified as VSCode.
if (process.env['TERM_PROGRAM'] === 'vscode') {
// If the detected IDE is 'vscode', we only accept it if TERM_PROGRAM or VSCODE_PID confirms it.
// This prevents generic terminals from being misidentified as VSCode, while still detecting
// background processes spawned by the VS Code extension host (like a2a-server).
if (process.env['TERM_PROGRAM'] === 'vscode' || process.env['VSCODE_PID']) {
return ide.name;
}