A2a admin setting (#17868)

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<typeof import('@google/gemini-cli-core')>();
   return {
     ...actual,
-    Config: vi.fn().mockImplementation((params) => ({
+    Config: vi.fn().mockImplementation((params) => {
+      const mockConfig = {
+        ...params,
         initialize: vi.fn(),
         refreshAuth: vi.fn(),
-      ...params, // Expose params for assertion
+        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;

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']) {
+  await refreshAuthentication(config, adcFilePath, 'Config');
-    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);
-  }
 
   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);
+  }
+}

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

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<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.
  */