From da3da19844321a23897a62f80313f095a1b72ee0 Mon Sep 17 00:00:00 2001 From: Shreya Keshive Date: Tue, 4 Nov 2025 15:09:53 -0800 Subject: [PATCH] feat(core): Integrate remote experiments configuration (#12539) --- packages/cli/src/config/config.test.ts | 4 +- packages/core/index.ts | 2 + .../experiments/client_metadata.ts | 10 +-- .../code_assist/experiments/experiments.ts | 2 +- .../core/src/code_assist/experiments/types.ts | 38 +++----- packages/core/src/code_assist/server.test.ts | 4 +- packages/core/src/code_assist/server.ts | 4 +- packages/core/src/code_assist/types.ts | 3 +- packages/core/src/config/config.test.ts | 88 +++++++++++++++++++ packages/core/src/config/config.ts | 58 +++++++++++- .../src/services/chatCompressionService.ts | 3 +- 11 files changed, 173 insertions(+), 43 deletions(-) diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index a7d00affdb..08a4eaab4a 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1550,7 +1550,7 @@ describe('loadCliConfig compressionThreshold', () => { }, }; const config = await loadCliConfig(settings, 'test-session', argv); - expect(config.getCompressionThreshold()).toBe(0.5); + expect(await config.getCompressionThreshold()).toBe(0.5); }); it('should have undefined compressionThreshold if not in settings', async () => { @@ -1558,7 +1558,7 @@ describe('loadCliConfig compressionThreshold', () => { const argv = await parseArguments({} as Settings); const settings: Settings = {}; const config = await loadCliConfig(settings, 'test-session', argv); - expect(config.getCompressionThreshold()).toBeUndefined(); + expect(await config.getCompressionThreshold()).toBeUndefined(); }); }); diff --git a/packages/core/index.ts b/packages/core/index.ts index 2369b6b0e2..4b2b0c38ff 100644 --- a/packages/core/index.ts +++ b/packages/core/index.ts @@ -47,3 +47,5 @@ export { logModelSlashCommand } from './src/telemetry/loggers.js'; export { KeychainTokenStorage } from './src/mcp/token-storage/keychain-token-storage.js'; export * from './src/utils/googleQuotaErrors.js'; export type { GoogleApiError } from './src/utils/googleErrors.js'; +export { getCodeAssistServer } from './src/code_assist/codeAssist.js'; +export { getExperiments } from './src/code_assist/experiments/experiments.js'; diff --git a/packages/core/src/code_assist/experiments/client_metadata.ts b/packages/core/src/code_assist/experiments/client_metadata.ts index 953d46aa5b..16108fc736 100644 --- a/packages/core/src/code_assist/experiments/client_metadata.ts +++ b/packages/core/src/code_assist/experiments/client_metadata.ts @@ -5,7 +5,7 @@ */ import { getReleaseChannel } from '../../utils/channel.js'; -import type { ClientMetadata, Platform } from './types.js'; +import type { ClientMetadata, ClientMetadataPlatform } from '../types.js'; import { fileURLToPath } from 'node:url'; import path from 'node:path'; @@ -15,7 +15,7 @@ const __dirname = path.dirname(__filename); // Cache all client metadata. let clientMetadataPromise: Promise | undefined; -function getPlatform(): Platform { +function getPlatform(): ClientMetadataPlatform { const platform = process.platform; const arch = process.arch; @@ -45,10 +45,10 @@ function getPlatform(): Platform { export async function getClientMetadata(): Promise { if (!clientMetadataPromise) { clientMetadataPromise = (async () => ({ - ide_type: 'GEMINI_CLI', - ide_version: process.env['CLI_VERSION'] || process.version, + ideName: 'GEMINI_CLI', + ideVersion: process.env['CLI_VERSION'] || process.version, platform: getPlatform(), - update_channel: await getReleaseChannel(__dirname), + updateChannel: await getReleaseChannel(__dirname), }))(); } return await clientMetadataPromise; diff --git a/packages/core/src/code_assist/experiments/experiments.ts b/packages/core/src/code_assist/experiments/experiments.ts index 7e4f5543b8..f716da9c72 100644 --- a/packages/core/src/code_assist/experiments/experiments.ts +++ b/packages/core/src/code_assist/experiments/experiments.ts @@ -44,6 +44,6 @@ function parseExperiments(response: ListExperimentsResponse): Experiments { } return { flags, - experimentIds: response.experiment_ids ?? [], + experimentIds: response.experimentIds ?? [], }; } diff --git a/packages/core/src/code_assist/experiments/types.ts b/packages/core/src/code_assist/experiments/types.ts index 3c28c71465..924108f347 100644 --- a/packages/core/src/code_assist/experiments/types.ts +++ b/packages/core/src/code_assist/experiments/types.ts @@ -4,26 +4,28 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { ClientMetadata } from '../types.js'; + export interface ListExperimentsRequest { project: string; metadata?: ClientMetadata; } export interface ListExperimentsResponse { - experiment_ids?: number[]; + experimentIds?: number[]; flags?: Flag[]; - filtered_flags?: FilteredFlag[]; - debug_string?: string; + filteredFlags?: FilteredFlag[]; + debugString?: string; } export interface Flag { name?: string; - bool_value?: boolean; - float_value?: number; - int_value?: string; // int64 - string_value?: string; - int32_list_value?: Int32List; - string_list_value?: StringList; + boolValue?: boolean; + floatValue?: number; + intValue?: string; // int64 + stringValue?: string; + int32ListValue?: Int32List; + stringListValue?: StringList; } export interface Int32List { @@ -38,21 +40,3 @@ export interface FilteredFlag { name?: string; reason?: string; } - -export interface ClientMetadata { - ide_type?: IdeType; - ide_version?: string; - platform?: Platform; - update_channel?: 'nightly' | 'preview' | 'stable'; - duet_project?: string; -} - -export type IdeType = 'GEMINI_CLI'; - -export type Platform = - | 'PLATFORM_UNSPECIFIED' - | 'DARWIN_AMD64' - | 'DARWIN_ARM64' - | 'LINUX_AMD64' - | 'LINUX_ARM64' - | 'WINDOWS_AMD64'; diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts index 5ef4bbd352..86c73f601a 100644 --- a/packages/core/src/code_assist/server.test.ts +++ b/packages/core/src/code_assist/server.test.ts @@ -268,13 +268,13 @@ describe('CodeAssistServer', () => { vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse); const metadata = { - ide_version: 'v0.1.0', + ideVersion: 'v0.1.0', }; const response = await server.listExperiments(metadata); expect(server.requestPost).toHaveBeenCalledWith('listExperiments', { project: 'test-project', - metadata: { ide_version: 'v0.1.0', duet_project: 'test-project' }, + metadata: { ideVersion: 'v0.1.0', duetProject: 'test-project' }, }); expect(response).toEqual(mockResponse); }); diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index 6c0b68f448..412d6f46b2 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -13,11 +13,11 @@ import type { LongRunningOperationResponse, OnboardUserRequest, SetCodeAssistGlobalUserSettingRequest, + ClientMetadata, } from './types.js'; import type { ListExperimentsRequest, ListExperimentsResponse, - ClientMetadata, } from './experiments/types.js'; import type { CountTokensParameters, @@ -163,7 +163,7 @@ export class CodeAssistServer implements ContentGenerator { const projectId = this.projectId; const req: ListExperimentsRequest = { project: projectId, - metadata: { ...metadata, duet_project: projectId }, + metadata: { ...metadata, duetProject: projectId }, }; return await this.requestPost( 'listExperiments', diff --git a/packages/core/src/code_assist/types.ts b/packages/core/src/code_assist/types.ts index b79094bbb8..e59dd3ac40 100644 --- a/packages/core/src/code_assist/types.ts +++ b/packages/core/src/code_assist/types.ts @@ -21,7 +21,8 @@ export type ClientMetadataIdeType = | 'INTELLIJ' | 'VSCODE_CLOUD_WORKSTATION' | 'INTELLIJ_CLOUD_WORKSTATION' - | 'CLOUD_SHELL'; + | 'CLOUD_SHELL' + | 'GEMINI_CLI'; export type ClientMetadataPlatform = | 'PLATFORM_UNSPECIFIED' | 'DARWIN_AMD64' diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 1850ffc0e6..6b4e65dc79 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -230,6 +230,49 @@ describe('Server Config (config.ts)', () => { 'Config was already initialized', ); }); + + describe('getCompressionThreshold', () => { + it('should return the local compression threshold if it is set', async () => { + const config = new Config({ + ...baseParams, + compressionThreshold: 0.5, + }); + expect(await config.getCompressionThreshold()).toBe(0.5); + }); + + it('should return the remote experiment threshold if it is a positive number', async () => { + const config = new Config({ + ...baseParams, + experiments: { + flags: { + GeminiCLIContextCompression__threshold_fraction: { + floatValue: 0.8, + }, + }, + }, + } as unknown as ConfigParameters); + expect(await config.getCompressionThreshold()).toBe(0.8); + }); + + it('should return undefined if the remote experiment threshold is 0', async () => { + const config = new Config({ + ...baseParams, + experiments: { + flags: { + GeminiCLIContextCompression__threshold_fraction: { + floatValue: 0.0, + }, + }, + }, + } as unknown as ConfigParameters); + expect(await config.getCompressionThreshold()).toBeUndefined(); + }); + + it('should return undefined if there are no experiments', async () => { + const config = new Config(baseParams); + expect(await config.getCompressionThreshold()).toBeUndefined(); + }); + }); }); describe('refreshAuth', () => { @@ -1312,3 +1355,48 @@ describe('Config getHooks', () => { expect(Object.keys(retrievedHooks!)).toHaveLength(11); // All hook event types }); }); + +describe('Config getExperiments', () => { + const baseParams: ConfigParameters = { + cwd: '/tmp', + targetDir: '/path/to/target', + debugMode: false, + sessionId: 'test-session-id', + model: 'gemini-pro', + usageStatisticsEnabled: false, + }; + + it('should return undefined when no experiments are provided', () => { + const config = new Config(baseParams); + expect(config.getExperiments()).toBeUndefined(); + }); + + it('should return empty object when empty experiments are provided', () => { + const configWithEmptyExps = new Config({ + ...baseParams, + experiments: { flags: {}, experimentIds: [] }, + }); + expect(configWithEmptyExps.getExperiments()).toEqual({ + flags: {}, + experimentIds: [], + }); + }); + + it('should return the experiments configuration when provided', () => { + const mockExps = { + flags: { + testFlag: { boolValue: true }, + }, + experimentIds: [], + }; + + const config = new Config({ + ...baseParams, + experiments: mockExps, + }); + + const retrievedExps = config.getExperiments(); + expect(retrievedExps).toEqual(mockExps); + expect(retrievedExps).toBe(mockExps); // Should return the same reference + }); +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 20cb7cf25c..20fe3b578d 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -75,9 +75,13 @@ import { MessageBus } from '../confirmation-bus/message-bus.js'; import { PolicyEngine } from '../policy/policy-engine.js'; import type { PolicyEngineConfig } from '../policy/types.js'; import type { UserTierId } from '../code_assist/types.js'; +import { getCodeAssistServer } from '../code_assist/codeAssist.js'; +import type { Experiments } from '../code_assist/experiments/experiments.js'; import { AgentRegistry } from '../agents/registry.js'; import { setGlobalProxy } from '../utils/fetch.js'; import { SubagentToolWrapper } from '../agents/subagent-tool-wrapper.js'; +import { getExperiments } from '../code_assist/experiments/experiments.js'; +import { debugLogger } from '../utils/debugLogger.js'; import { ApprovalMode } from '../policy/types.js'; @@ -288,6 +292,7 @@ export interface ConfigParameters { ptyInfo?: string; disableYoloMode?: boolean; enableHooks?: boolean; + experiments?: Experiments; hooks?: { [K in HookEventName]?: HookDefinition[]; }; @@ -396,6 +401,8 @@ export class Config { private readonly hooks: | { [K in HookEventName]?: HookDefinition[] } | undefined; + private experiments: Experiments | undefined; + private experimentsPromise: Promise | undefined; constructor(params: ConfigParameters) { this.sessionId = params.sessionId; @@ -529,6 +536,7 @@ export class Config { this.retryFetchErrors = params.retryFetchErrors ?? false; this.disableYoloMode = params.disableYoloMode ?? false; this.hooks = params.hooks; + this.experiments = params.experiments; if (params.contextFileName) { setGeminiMdFilename(params.contextFileName); @@ -625,6 +633,20 @@ export class Config { // Initialize BaseLlmClient now that the ContentGenerator is available this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this); + const codeAssistServer = getCodeAssistServer(this); + if (codeAssistServer) { + this.experimentsPromise = getExperiments(codeAssistServer) + .then((experiments) => { + this.setExperiments(experiments); + }) + .catch((e) => { + debugLogger.error('Failed to fetch experiments', e); + }); + } else { + this.experiments = undefined; + this.experimentsPromise = undefined; + } + // Reset the session flag since we're explicitly changing auth and using default model this.inFallbackMode = false; } @@ -1064,8 +1086,26 @@ export class Config { this.fileSystemService = fileSystemService; } - getCompressionThreshold(): number | undefined { - return this.compressionThreshold; + async getCompressionThreshold(): Promise { + if (this.compressionThreshold) { + return this.compressionThreshold; + } + + if (this.experimentsPromise) { + try { + await this.experimentsPromise; + } catch (e) { + debugLogger.debug('Failed to fetch experiments', e); + } + } + + const remoteThreshold = + this.experiments?.flags['GeminiCLIContextCompression__threshold_fraction'] + ?.floatValue; + if (remoteThreshold === 0) { + return undefined; + } + return remoteThreshold; } isInteractiveShellEnabled(): boolean { @@ -1318,6 +1358,20 @@ export class Config { getHooks(): { [K in HookEventName]?: HookDefinition[] } | undefined { return this.hooks; } + + /** + * Get experiments configuration + */ + getExperiments(): Experiments | undefined { + return this.experiments; + } + + /** + * Set experiments configuration + */ + setExperiments(experiments: Experiments): void { + this.experiments = experiments; + } } // Export model constants for use in CLI export { DEFAULT_GEMINI_FLASH_MODEL }; diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts index 573b2ae458..985d1c7448 100644 --- a/packages/core/src/services/chatCompressionService.ts +++ b/packages/core/src/services/chatCompressionService.ts @@ -106,7 +106,8 @@ export class ChatCompressionService { // Don't compress if not forced and we are under the limit. if (!force) { const threshold = - config.getCompressionThreshold() ?? DEFAULT_COMPRESSION_TOKEN_THRESHOLD; + (await config.getCompressionThreshold()) ?? + DEFAULT_COMPRESSION_TOKEN_THRESHOLD; if (originalTokenCount < threshold * tokenLimit(model)) { return { newHistory: null,