From 05e4ea80eed76be095af0bbdbdd63c16fc738f6b Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Fri, 20 Mar 2026 15:31:01 -0400 Subject: [PATCH] feat(core): refine User-Agent for VS Code traffic (unified format) (#23256) --- .../core/src/core/contentGenerator.test.ts | 121 +++++++++++++++++- packages/core/src/core/contentGenerator.ts | 39 +++++- packages/core/src/utils/surface.ts | 7 +- 3 files changed, 156 insertions(+), 11 deletions(-) diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 4bacd1b488..a264b2fb6c 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -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( diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index ff1739c04b..c901562eb7 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -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 = { + 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 = { - ...customHeadersMap, 'User-Agent': userAgent, + ...customHeadersMap, }; if ( diff --git a/packages/core/src/utils/surface.ts b/packages/core/src/utils/surface.ts index e4b1241d84..7c6bd4da6b 100644 --- a/packages/core/src/utils/surface.ts +++ b/packages/core/src/utils/surface.ts @@ -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; }