summaryrefslogtreecommitdiff
path: root/packages/core
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core')
-rw-r--r--packages/core/src/core/contentGenerator.test.ts38
-rw-r--r--packages/core/src/core/contentGenerator.ts20
2 files changed, 54 insertions, 4 deletions
diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts
index b92e8046..5a857121 100644
--- a/packages/core/src/core/contentGenerator.test.ts
+++ b/packages/core/src/core/contentGenerator.test.ts
@@ -41,6 +41,44 @@ describe('createContentGenerator', () => {
});
it('should create a GoogleGenAI content generator', async () => {
+ const mockConfig = {
+ getUsageStatisticsEnabled: () => true,
+ } as unknown as Config;
+
+ const mockGenerator = {
+ models: {},
+ } as unknown as GoogleGenAI;
+ vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
+ const generator = await createContentGenerator(
+ {
+ model: 'test-model',
+ apiKey: 'test-api-key',
+ authType: AuthType.USE_GEMINI,
+ },
+ mockConfig,
+ );
+ expect(GoogleGenAI).toHaveBeenCalledWith({
+ apiKey: 'test-api-key',
+ vertexai: undefined,
+ httpOptions: {
+ headers: {
+ 'User-Agent': expect.any(String),
+ 'x-gemini-api-privileged-user-id': expect.any(String),
+ },
+ },
+ });
+ expect(generator).toEqual(
+ new LoggingContentGenerator(
+ (mockGenerator as GoogleGenAI).models,
+ mockConfig,
+ ),
+ );
+ });
+
+ it('should create a GoogleGenAI content generator with client install id logging disabled', async () => {
+ const mockConfig = {
+ getUsageStatisticsEnabled: () => false,
+ } as unknown as Config;
const mockGenerator = {
models: {},
} as unknown as GoogleGenAI;
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index 19b285c7..0556805e 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -19,6 +19,7 @@ import { Config } from '../config/config.js';
import { UserTierId } from '../code_assist/types.js';
import { LoggingContentGenerator } from './loggingContentGenerator.js';
+import { getInstallationId } from '../utils/user_id.js';
/**
* Interface abstracting the core functionalities for generating content and counting tokens.
@@ -108,15 +109,16 @@ export async function createContentGenerator(
sessionId?: string,
): Promise<ContentGenerator> {
const version = process.env['CLI_VERSION'] || process.version;
- const httpOptions = {
- headers: {
- 'User-Agent': `GeminiCLI/${version} (${process.platform}; ${process.arch})`,
- },
+ const userAgent = `GeminiCLI/${version} (${process.platform}; ${process.arch})`;
+ const baseHeaders: Record<string, string> = {
+ 'User-Agent': userAgent,
};
+
if (
config.authType === AuthType.LOGIN_WITH_GOOGLE ||
config.authType === AuthType.CLOUD_SHELL
) {
+ const httpOptions = { headers: baseHeaders };
return new LoggingContentGenerator(
await createCodeAssistContentGenerator(
httpOptions,
@@ -132,6 +134,16 @@ export async function createContentGenerator(
config.authType === AuthType.USE_GEMINI ||
config.authType === AuthType.USE_VERTEX_AI
) {
+ let headers: Record<string, string> = { ...baseHeaders };
+ if (gcConfig?.getUsageStatisticsEnabled()) {
+ const installationId = getInstallationId();
+ headers = {
+ ...headers,
+ 'x-gemini-api-privileged-user-id': `${installationId}`,
+ };
+ }
+ const httpOptions = { headers };
+
const googleGenAI = new GoogleGenAI({
apiKey: config.apiKey === '' ? undefined : config.apiKey,
vertexai: config.vertexai,