diff --git a/packages/a2a-server/src/config/config.test.ts b/packages/a2a-server/src/config/config.test.ts index 06be9581a5..d3f0bfdbdd 100644 --- a/packages/a2a-server/src/config/config.test.ts +++ b/packages/a2a-server/src/config/config.test.ts @@ -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(); 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; diff --git a/packages/a2a-server/src/config/config.ts b/packages/a2a-server/src/config/config.ts index 12ab87439a..d98ae4fb7c 100644 --- a/packages/a2a-server/src/config/config.ts +++ b/packages/a2a-server/src/config/config.ts @@ -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 { + 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); + } +} diff --git a/packages/core/src/code_assist/admin/admin_controls.test.ts b/packages/core/src/code_assist/admin/admin_controls.test.ts index b36daa3c9b..f30d8b7e58 100644 --- a/packages/core/src/code_assist/admin/admin_controls.test.ts +++ b/packages/core/src/code_assist/admin/admin_controls.test.ts @@ -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 diff --git a/packages/core/src/code_assist/admin/admin_controls.ts b/packages/core/src/code_assist/admin/admin_controls.ts index fce50b60f0..7d9c46861c 100644 --- a/packages/core/src/code_assist/admin/admin_controls.ts +++ b/packages/core/src/code_assist/admin/admin_controls.ts @@ -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 { + 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. */