A2a admin setting (#17868)

This commit is contained in:
David Pierce
2026-02-03 20:16:20 +00:00
committed by GitHub
parent 675ca07c8b
commit 75dbf9022c
4 changed files with 317 additions and 30 deletions

View File

@@ -11,6 +11,11 @@ import type { Settings } from './settings.js';
import {
type ExtensionLoader,
FileDiscoveryService,
getCodeAssistServer,
Config,
ExperimentFlags,
fetchAdminControlsOnce,
type FetchAdminControlsResponse,
} from '@google/gemini-cli-core';
// Mock dependencies
@@ -19,11 +24,23 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
await importOriginal<typeof import('@google/gemini-cli-core')>();
return {
...actual,
Config: vi.fn().mockImplementation((params) => ({
initialize: vi.fn(),
refreshAuth: vi.fn(),
...params, // Expose params for assertion
})),
Config: vi.fn().mockImplementation((params) => {
const mockConfig = {
...params,
initialize: vi.fn(),
refreshAuth: vi.fn(),
getExperiments: vi.fn().mockReturnValue({
flags: {
[actual.ExperimentFlags.ENABLE_ADMIN_CONTROLS]: {
boolValue: false,
},
},
}),
getRemoteAdminSettings: vi.fn(),
setRemoteAdminSettings: vi.fn(),
};
return mockConfig;
}),
loadServerHierarchicalMemory: vi
.fn()
.mockResolvedValue({ memoryContent: '', fileCount: 0, filePaths: [] }),
@@ -31,6 +48,11 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
flush: vi.fn(),
},
FileDiscoveryService: vi.fn(),
getCodeAssistServer: vi.fn(),
fetchAdminControlsOnce: vi.fn(),
coreEvents: {
emitAdminSettingsChanged: vi.fn(),
},
};
});
@@ -56,6 +78,121 @@ describe('loadConfig', () => {
delete process.env['GEMINI_API_KEY'];
});
describe('admin settings overrides', () => {
it('should not fetch admin controls if experiment is disabled', async () => {
await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(fetchAdminControlsOnce).not.toHaveBeenCalled();
});
describe('when admin controls experiment is enabled', () => {
beforeEach(() => {
// We need to cast to any here to modify the mock implementation
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(Config as any).mockImplementation((params: unknown) => {
const mockConfig = {
...(params as object),
initialize: vi.fn(),
refreshAuth: vi.fn(),
getExperiments: vi.fn().mockReturnValue({
flags: {
[ExperimentFlags.ENABLE_ADMIN_CONTROLS]: {
boolValue: true,
},
},
}),
getRemoteAdminSettings: vi.fn().mockReturnValue({}),
setRemoteAdminSettings: vi.fn(),
};
return mockConfig;
});
});
it('should fetch admin controls and apply them', async () => {
const mockAdminSettings: FetchAdminControlsResponse = {
mcpSetting: {
mcpEnabled: false,
},
cliFeatureSetting: {
extensionsSetting: {
extensionsEnabled: false,
},
},
strictModeDisabled: false,
};
vi.mocked(fetchAdminControlsOnce).mockResolvedValue(mockAdminSettings);
await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(Config).toHaveBeenCalledWith(
expect.objectContaining({
disableYoloMode: !mockAdminSettings.strictModeDisabled,
mcpEnabled: mockAdminSettings.mcpSetting?.mcpEnabled,
extensionsEnabled:
mockAdminSettings.cliFeatureSetting?.extensionsSetting
?.extensionsEnabled,
}),
);
});
it('should treat unset admin settings as false when admin settings are passed', async () => {
const mockAdminSettings: FetchAdminControlsResponse = {
mcpSetting: {
mcpEnabled: true,
},
};
vi.mocked(fetchAdminControlsOnce).mockResolvedValue(mockAdminSettings);
await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(Config).toHaveBeenCalledWith(
expect.objectContaining({
disableYoloMode: !false,
mcpEnabled: mockAdminSettings.mcpSetting?.mcpEnabled,
extensionsEnabled: false,
}),
);
});
it('should not pass default unset admin settings when no admin settings are present', async () => {
const mockAdminSettings: FetchAdminControlsResponse = {};
vi.mocked(fetchAdminControlsOnce).mockResolvedValue(mockAdminSettings);
await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(Config).toHaveBeenCalledWith(expect.objectContaining({}));
});
it('should fetch admin controls using the code assist server when available', async () => {
const mockAdminSettings: FetchAdminControlsResponse = {
mcpSetting: {
mcpEnabled: true,
},
strictModeDisabled: true,
};
const mockCodeAssistServer = { projectId: 'test-project' };
vi.mocked(getCodeAssistServer).mockReturnValue(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
mockCodeAssistServer as any,
);
vi.mocked(fetchAdminControlsOnce).mockResolvedValue(mockAdminSettings);
await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(fetchAdminControlsOnce).toHaveBeenCalledWith(
mockCodeAssistServer,
true,
);
expect(Config).toHaveBeenCalledWith(
expect.objectContaining({
disableYoloMode: !mockAdminSettings.strictModeDisabled,
mcpEnabled: mockAdminSettings.mcpSetting?.mcpEnabled,
extensionsEnabled: false,
}),
);
});
});
});
it('should set customIgnoreFilePaths when CUSTOM_IGNORE_FILE_PATHS env var is present', async () => {
const testPath = '/tmp/ignore';
process.env['CUSTOM_IGNORE_FILE_PATHS'] = testPath;

View File

@@ -24,6 +24,9 @@ import {
PREVIEW_GEMINI_MODEL,
homedir,
GitService,
fetchAdminControlsOnce,
getCodeAssistServer,
ExperimentFlags,
} from '@google/gemini-cli-core';
import { logger } from '../utils/logger.js';
@@ -124,37 +127,54 @@ export async function loadConfig(
configParams.userMemory = memoryContent;
configParams.geminiMdFileCount = fileCount;
configParams.geminiMdFilePaths = filePaths;
const config = new Config({
// Set an initial config to use to get a code assist server.
// This is needed to fetch admin controls.
const initialConfig = new Config({
...configParams,
});
const codeAssistServer = getCodeAssistServer(initialConfig);
const adminControlsEnabled =
initialConfig.getExperiments()?.flags[ExperimentFlags.ENABLE_ADMIN_CONTROLS]
?.boolValue ?? false;
// Initialize final config parameters to the previous parameters.
// If no admin controls are needed, these will be used as-is for the final
// config.
const finalConfigParams = { ...configParams };
if (adminControlsEnabled) {
const adminSettings = await fetchAdminControlsOnce(
codeAssistServer,
adminControlsEnabled,
);
// Admin settings are able to be undefined if unset, but if any are present,
// we should initialize them all.
// If any are present, undefined settings should be treated as if they were
// set to false.
// If NONE are present, disregard admin settings entirely, and pass the
// final config as is.
if (Object.keys(adminSettings).length !== 0) {
finalConfigParams.disableYoloMode = !(
adminSettings.strictModeDisabled ?? false
);
finalConfigParams.mcpEnabled =
adminSettings.mcpSetting?.mcpEnabled ?? false;
finalConfigParams.extensionsEnabled =
adminSettings.cliFeatureSetting?.extensionsSetting?.extensionsEnabled ??
false;
}
}
const config = new Config(finalConfigParams);
// Needed to initialize ToolRegistry, and git checkpointing if enabled
await config.initialize();
startupProfiler.flush(config);
if (process.env['USE_CCPA']) {
logger.info('[Config] Using CCPA Auth:');
try {
if (adcFilePath) {
path.resolve(adcFilePath);
}
} catch (e) {
logger.error(
`[Config] USE_CCPA env var is true but unable to resolve GOOGLE_APPLICATION_CREDENTIALS file path ${adcFilePath}. Error ${e}`,
);
}
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
logger.info(
`[Config] GOOGLE_CLOUD_PROJECT: ${process.env['GOOGLE_CLOUD_PROJECT']}`,
);
} else if (process.env['GEMINI_API_KEY']) {
logger.info('[Config] Using Gemini API Key');
await config.refreshAuth(AuthType.USE_GEMINI);
} else {
const errorMessage =
'[Config] Unable to set GeneratorConfig. Please provide a GEMINI_API_KEY or set USE_CCPA.';
logger.error(errorMessage);
throw new Error(errorMessage);
}
await refreshAuthentication(config, adcFilePath, 'Config');
return config;
}
@@ -222,3 +242,33 @@ function findEnvFile(startDir: string): string | null {
currentDir = parentDir;
}
}
async function refreshAuthentication(
config: Config,
adcFilePath: string | undefined,
logPrefix: string,
): Promise<void> {
if (process.env['USE_CCPA']) {
logger.info(`[${logPrefix}] Using CCPA Auth:`);
try {
if (adcFilePath) {
path.resolve(adcFilePath);
}
} catch (e) {
logger.error(
`[${logPrefix}] USE_CCPA env var is true but unable to resolve GOOGLE_APPLICATION_CREDENTIALS file path ${adcFilePath}. Error ${e}`,
);
}
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
logger.info(
`[${logPrefix}] GOOGLE_CLOUD_PROJECT: ${process.env['GOOGLE_CLOUD_PROJECT']}`,
);
} else if (process.env['GEMINI_API_KEY']) {
logger.info(`[${logPrefix}] Using Gemini API Key`);
await config.refreshAuth(AuthType.USE_GEMINI);
} else {
const errorMessage = `[${logPrefix}] Unable to set GeneratorConfig. Please provide a GEMINI_API_KEY or set USE_CCPA.`;
logger.error(errorMessage);
throw new Error(errorMessage);
}
}

View File

@@ -15,6 +15,7 @@ import {
} from 'vitest';
import {
fetchAdminControls,
fetchAdminControlsOnce,
sanitizeAdminSettings,
stopAdminControlsPolling,
getAdminErrorMessage,
@@ -248,6 +249,71 @@ describe('Admin Controls', () => {
});
});
describe('fetchAdminControlsOnce', () => {
it('should return empty object if server is missing', async () => {
const result = await fetchAdminControlsOnce(undefined, true);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).not.toHaveBeenCalled();
});
it('should return empty object if project ID is missing', async () => {
mockServer = {
fetchAdminControls: vi.fn(),
} as unknown as CodeAssistServer;
const result = await fetchAdminControlsOnce(mockServer, true);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).not.toHaveBeenCalled();
});
it('should return empty object if admin controls are disabled', async () => {
const result = await fetchAdminControlsOnce(mockServer, false);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).not.toHaveBeenCalled();
});
it('should fetch from server and sanitize the response', async () => {
const serverResponse = {
strictModeDisabled: true,
unknownField: 'should be removed',
};
(mockServer.fetchAdminControls as Mock).mockResolvedValue(serverResponse);
const result = await fetchAdminControlsOnce(mockServer, true);
expect(result).toEqual({ strictModeDisabled: true });
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});
it('should return empty object on 403 fetch error', async () => {
const error403 = new Error('Forbidden');
Object.assign(error403, { status: 403 });
(mockServer.fetchAdminControls as Mock).mockRejectedValue(error403);
const result = await fetchAdminControlsOnce(mockServer, true);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});
it('should return empty object on any other fetch error', async () => {
(mockServer.fetchAdminControls as Mock).mockRejectedValue(
new Error('Network error'),
);
const result = await fetchAdminControlsOnce(mockServer, true);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});
it('should not start or stop any polling timers', async () => {
const setIntervalSpy = vi.spyOn(global, 'setInterval');
const clearIntervalSpy = vi.spyOn(global, 'clearInterval');
(mockServer.fetchAdminControls as Mock).mockResolvedValue({});
await fetchAdminControlsOnce(mockServer, true);
expect(setIntervalSpy).not.toHaveBeenCalled();
expect(clearIntervalSpy).not.toHaveBeenCalled();
});
});
describe('polling', () => {
it('should poll and emit changes', async () => {
// Initial fetch

View File

@@ -89,6 +89,40 @@ export async function fetchAdminControls(
}
}
/**
* Fetches the admin controls from the server a single time.
* This function does not start or stop any polling.
*
* @param server The CodeAssistServer instance.
* @param adminControlsEnabled Whether admin controls are enabled.
* @returns The fetched settings if enabled and successful, otherwise undefined.
*/
export async function fetchAdminControlsOnce(
server: CodeAssistServer | undefined,
adminControlsEnabled: boolean,
): Promise<FetchAdminControlsResponse> {
if (!server || !server.projectId || !adminControlsEnabled) {
return {};
}
try {
const rawSettings = await server.fetchAdminControls({
project: server.projectId,
});
return sanitizeAdminSettings(rawSettings);
} catch (e) {
// Non-enterprise users don't have access to fetch settings.
if (isGaxiosError(e) && e.status === 403) {
return {};
}
debugLogger.error(
'Failed to fetch admin controls: ',
e instanceof Error ? e.message : e,
);
return {};
}
}
/**
* Starts polling for admin controls.
*/