From ca4866142339601a70db5daa837ab10baf450755 Mon Sep 17 00:00:00 2001 From: Kevin Ramdass Date: Fri, 9 Jan 2026 10:47:05 -0800 Subject: [PATCH] feat(core): add local experiments override via GEMINI_EXP (#16181) --- .../code_assist/experiments/experiments.ts | 28 +++- .../experiments/experiments_local.test.ts | 146 ++++++++++++++++++ packages/core/src/config/config.test.ts | 4 + packages/core/src/config/config.ts | 48 +++--- 4 files changed, 197 insertions(+), 29 deletions(-) create mode 100644 packages/core/src/code_assist/experiments/experiments_local.test.ts diff --git a/packages/core/src/code_assist/experiments/experiments.ts b/packages/core/src/code_assist/experiments/experiments.ts index 90eae36679..ecb98491eb 100644 --- a/packages/core/src/code_assist/experiments/experiments.ts +++ b/packages/core/src/code_assist/experiments/experiments.ts @@ -7,6 +7,8 @@ import type { CodeAssistServer } from '../server.js'; import { getClientMetadata } from './client_metadata.js'; import type { ListExperimentsResponse, Flag } from './types.js'; +import * as fs from 'node:fs'; +import { debugLogger } from '../../utils/debugLogger.js'; export interface Experiments { flags: Record; @@ -21,13 +23,37 @@ let experimentsPromise: Promise | undefined; * The experiments are cached so that they are only fetched once. */ export async function getExperiments( - server: CodeAssistServer, + server?: CodeAssistServer, ): Promise { if (experimentsPromise) { return experimentsPromise; } experimentsPromise = (async () => { + if (process.env['GEMINI_EXP']) { + try { + const expPath = process.env['GEMINI_EXP']; + debugLogger.debug('Reading experiments from', expPath); + const content = await fs.promises.readFile(expPath, 'utf8'); + const response = JSON.parse(content); + if ( + (response.flags && !Array.isArray(response.flags)) || + (response.experimentIds && !Array.isArray(response.experimentIds)) + ) { + throw new Error( + 'Invalid format for experiments file: `flags` and `experimentIds` must be arrays if present.', + ); + } + return parseExperiments(response as ListExperimentsResponse); + } catch (e) { + debugLogger.debug('Failed to read experiments from GEMINI_EXP', e); + } + } + + if (!server) { + return { flags: {}, experimentIds: [] }; + } + const metadata = await getClientMetadata(); const response = await server.listExperiments(metadata); return parseExperiments(response); diff --git a/packages/core/src/code_assist/experiments/experiments_local.test.ts b/packages/core/src/code_assist/experiments/experiments_local.test.ts new file mode 100644 index 0000000000..f7bed37319 --- /dev/null +++ b/packages/core/src/code_assist/experiments/experiments_local.test.ts @@ -0,0 +1,146 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import type { CodeAssistServer } from '../server.js'; +import * as fs from 'node:fs'; +import * as os from 'node:os'; +import type { ListExperimentsResponse } from './types.js'; +import type { ClientMetadata } from '../types.js'; + +// Mock dependencies +vi.mock('node:fs', () => ({ + promises: { + readFile: vi.fn(), + }, + readFileSync: vi.fn(), +})); +vi.mock('node:os'); +vi.mock('../server.js'); +vi.mock('./client_metadata.js', () => ({ + getClientMetadata: vi.fn(), +})); + +describe('experiments with GEMINI_EXP', () => { + let mockServer: CodeAssistServer; + + beforeEach(() => { + vi.resetModules(); + vi.clearAllMocks(); + process.env['GEMINI_EXP'] = ''; // Clear env var + + // Default mocks + vi.mocked(os.homedir).mockReturnValue('/home/user'); + mockServer = { + listExperiments: vi.fn(), + } as unknown as CodeAssistServer; + }); + + afterEach(() => { + delete process.env['GEMINI_EXP']; + }); + + it('should read experiments from local file if GEMINI_EXP is set', async () => { + process.env['GEMINI_EXP'] = '/tmp/experiments.json'; + const mockFileContent = JSON.stringify({ + flags: [{ flagId: 111, boolValue: true }], + experimentIds: [999], + }); + vi.mocked(fs.promises.readFile).mockResolvedValue(mockFileContent); + + const { getExperiments } = await import('./experiments.js'); + const experiments = await getExperiments(mockServer); + + expect(fs.promises.readFile).toHaveBeenCalledWith( + '/tmp/experiments.json', + 'utf8', + ); + expect(experiments.flags[111]).toEqual({ + flagId: 111, + boolValue: true, + }); + expect(experiments.experimentIds).toEqual([999]); + expect(mockServer.listExperiments).not.toHaveBeenCalled(); + }); + + it('should fall back to server if reading file fails', async () => { + process.env['GEMINI_EXP'] = '/tmp/missing.json'; + vi.mocked(fs.promises.readFile).mockRejectedValue( + new Error('File not found'), + ); + + // Mock server response + const mockApiResponse = { + flags: [{ flagId: 222, boolValue: true }], + experimentIds: [111], + }; + vi.mocked(mockServer.listExperiments).mockResolvedValue( + mockApiResponse as ListExperimentsResponse, + ); + const { getClientMetadata } = await import('./client_metadata.js'); + vi.mocked(getClientMetadata).mockResolvedValue( + {} as unknown as ClientMetadata, + ); + + const { getExperiments } = await import('./experiments.js'); + const experiments = await getExperiments(mockServer); + + expect(experiments.flags[222]).toBeDefined(); + expect(mockServer.listExperiments).toHaveBeenCalled(); + }); + + it('should work without server if file read succeeds', async () => { + process.env['GEMINI_EXP'] = '/tmp/experiments.json'; + const mockFileContent = JSON.stringify({ + flags: [{ flagId: 333, boolValue: true }], + experimentIds: [999], + }); + vi.mocked(fs.promises.readFile).mockResolvedValue(mockFileContent); + + const { getExperiments } = await import('./experiments.js'); + const experiments = await getExperiments(undefined); + + expect(experiments.flags[333]).toEqual({ + flagId: 333, + boolValue: true, + }); + }); + + it('should return empty if no server and no GEMINI_EXP', async () => { + const { getExperiments } = await import('./experiments.js'); + const experiments = await getExperiments(undefined); + expect(experiments.flags).toEqual({}); + expect(experiments.experimentIds).toEqual([]); + }); + + it('should fallback to server if file has invalid structure', async () => { + process.env['GEMINI_EXP'] = '/tmp/invalid.json'; + const mockFileContent = JSON.stringify({ + flags: 'invalid-flags-type', // Should be array + experimentIds: 123, // Should be array + }); + vi.mocked(fs.promises.readFile).mockResolvedValue(mockFileContent); + + // Mock server response + const mockApiResponse = { + flags: [{ flagId: 444, boolValue: true }], + experimentIds: [555], + }; + vi.mocked(mockServer.listExperiments).mockResolvedValue( + mockApiResponse as ListExperimentsResponse, + ); + const { getClientMetadata } = await import('./client_metadata.js'); + vi.mocked(getClientMetadata).mockResolvedValue( + {} as unknown as ClientMetadata, + ); + + const { getExperiments } = await import('./experiments.js'); + const experiments = await getExperiments(mockServer); + + expect(experiments.flags[444]).toBeDefined(); + expect(mockServer.listExperiments).toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 265785a891..8e1c9d9b68 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -226,6 +226,10 @@ describe('Server Config (config.ts)', () => { beforeEach(() => { // Reset mocks if necessary vi.clearAllMocks(); + vi.mocked(getExperiments).mockResolvedValue({ + experimentIds: [], + flags: {}, + }); }); describe('initialize', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2a9ae19284..f877a2e797 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -814,32 +814,27 @@ export class Config { this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this); const codeAssistServer = getCodeAssistServer(this); - if (codeAssistServer) { - if (codeAssistServer.projectId) { - await this.refreshUserQuota(); - } - - this.experimentsPromise = getExperiments(codeAssistServer) - .then((experiments) => { - this.setExperiments(experiments); - - // If preview features have not been set and the user authenticated through Google, we enable preview based on remote config only if it's true - if (this.getPreviewFeatures() === undefined) { - const remotePreviewFeatures = - experiments.flags[ExperimentFlags.ENABLE_PREVIEW]?.boolValue; - if (remotePreviewFeatures === true) { - this.setPreviewFeatures(remotePreviewFeatures); - } - } - }) - .catch((e) => { - debugLogger.error('Failed to fetch experiments', e); - }); - } else { - this.experiments = undefined; - this.experimentsPromise = undefined; + if (codeAssistServer?.projectId) { + await this.refreshUserQuota(); } + this.experimentsPromise = getExperiments(codeAssistServer) + .then((experiments) => { + this.setExperiments(experiments); + + // If preview features have not been set and the user authenticated through Google, we enable preview based on remote config only if it's true + if (this.getPreviewFeatures() === undefined) { + const remotePreviewFeatures = + experiments.flags[ExperimentFlags.ENABLE_PREVIEW]?.boolValue; + if (remotePreviewFeatures === true) { + this.setPreviewFeatures(remotePreviewFeatures); + } + } + }) + .catch((e) => { + debugLogger.error('Failed to fetch experiments', e); + }); + const authType = this.contentGeneratorConfig.authType; if ( authType === AuthType.USE_GEMINI || @@ -859,10 +854,7 @@ export class Config { return this.experiments; } const codeAssistServer = getCodeAssistServer(this); - if (codeAssistServer) { - return getExperiments(codeAssistServer); - } - return undefined; + return getExperiments(codeAssistServer); } getUserTier(): UserTierId | undefined {