feat(core): add local experiments override via GEMINI_EXP (#16181)

This commit is contained in:
Kevin Ramdass
2026-01-09 10:47:05 -08:00
committed by GitHub
parent 041463d112
commit ca48661423
4 changed files with 197 additions and 29 deletions

View File

@@ -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<string, Flag>;
@@ -21,13 +23,37 @@ let experimentsPromise: Promise<Experiments> | undefined;
* The experiments are cached so that they are only fetched once.
*/
export async function getExperiments(
server: CodeAssistServer,
server?: CodeAssistServer,
): Promise<Experiments> {
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);

View File

@@ -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();
});
});

View File

@@ -226,6 +226,10 @@ describe('Server Config (config.ts)', () => {
beforeEach(() => {
// Reset mocks if necessary
vi.clearAllMocks();
vi.mocked(getExperiments).mockResolvedValue({
experimentIds: [],
flags: {},
});
});
describe('initialize', () => {

View File

@@ -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 {