From d8d4d87e2920764403c3167087b6bffed4654362 Mon Sep 17 00:00:00 2001 From: Shreya Keshive Date: Fri, 16 Jan 2026 15:24:53 -0500 Subject: [PATCH] feat(admin): implement admin controls polling and restart prompt (#16627) --- docs/cli/keyboard-shortcuts.md | 1 + packages/cli/src/config/keyBindings.ts | 4 + packages/cli/src/config/settings.ts | 4 +- packages/cli/src/gemini.tsx | 46 ++- packages/cli/src/test-utils/render.tsx | 1 + packages/cli/src/ui/AppContainer.tsx | 17 + .../AdminSettingsChangedDialog.test.tsx | 51 +++ .../components/AdminSettingsChangedDialog.tsx | 36 ++ .../cli/src/ui/components/DialogManager.tsx | 4 + .../AdminSettingsChangedDialog.test.tsx.snap | 8 + .../cli/src/ui/contexts/UIActionsContext.tsx | 1 + .../cli/src/ui/contexts/UIStateContext.tsx | 1 + packages/cli/src/utils/relaunch.ts | 12 +- .../code_assist/admin/admin_controls.test.ts | 330 ++++++++++++++++++ .../src/code_assist/admin/admin_controls.ts | 113 ++++++ .../src/code_assist/experiments/flagNames.ts | 1 + packages/core/src/code_assist/server.ts | 11 + packages/core/src/code_assist/types.ts | 40 ++- packages/core/src/config/config.ts | 25 +- packages/core/src/utils/events.ts | 9 + 20 files changed, 689 insertions(+), 26 deletions(-) create mode 100644 packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx create mode 100644 packages/cli/src/ui/components/AdminSettingsChangedDialog.tsx create mode 100644 packages/cli/src/ui/components/__snapshots__/AdminSettingsChangedDialog.test.tsx.snap create mode 100644 packages/core/src/code_assist/admin/admin_controls.test.ts create mode 100644 packages/core/src/code_assist/admin/admin_controls.ts diff --git a/docs/cli/keyboard-shortcuts.md b/docs/cli/keyboard-shortcuts.md index 3a6abf14f7..c496b416c5 100644 --- a/docs/cli/keyboard-shortcuts.md +++ b/docs/cli/keyboard-shortcuts.md @@ -106,6 +106,7 @@ available combinations. | Focus the shell input from the gemini input. | `Tab (no Shift)` | | Focus the Gemini input from the shell input. | `Tab` | | Clear the terminal screen and redraw the UI. | `Ctrl + L` | +| Restart the application. | `R` | diff --git a/packages/cli/src/config/keyBindings.ts b/packages/cli/src/config/keyBindings.ts index 47e4d6c400..465225f3b4 100644 --- a/packages/cli/src/config/keyBindings.ts +++ b/packages/cli/src/config/keyBindings.ts @@ -81,6 +81,7 @@ export enum Command { FOCUS_SHELL_INPUT = 'app.focusShellInput', UNFOCUS_SHELL_INPUT = 'app.unfocusShellInput', CLEAR_SCREEN = 'app.clearScreen', + RESTART_APP = 'app.restart', } /** @@ -238,6 +239,7 @@ export const defaultKeyBindings: KeyBindingConfig = { [Command.FOCUS_SHELL_INPUT]: [{ key: 'tab', shift: false }], [Command.UNFOCUS_SHELL_INPUT]: [{ key: 'tab' }], [Command.CLEAR_SCREEN]: [{ key: 'l', ctrl: true }], + [Command.RESTART_APP]: [{ key: 'r' }], }; interface CommandCategory { @@ -343,6 +345,7 @@ export const commandCategories: readonly CommandCategory[] = [ Command.FOCUS_SHELL_INPUT, Command.UNFOCUS_SHELL_INPUT, Command.CLEAR_SCREEN, + Command.RESTART_APP, ], }, ]; @@ -428,4 +431,5 @@ export const commandDescriptions: Readonly> = { [Command.FOCUS_SHELL_INPUT]: 'Focus the shell input from the gemini input.', [Command.UNFOCUS_SHELL_INPUT]: 'Focus the Gemini input from the shell input.', [Command.CLEAR_SCREEN]: 'Clear the terminal screen and redraw the UI.', + [Command.RESTART_APP]: 'Restart the application.', }; diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index a7bbd76ca6..0ad4dabe1e 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -16,7 +16,7 @@ import { Storage, coreEvents, homedir, - type GeminiCodeAssistSetting, + type FetchAdminControlsResponse, } from '@google/gemini-cli-core'; import stripJsonComments from 'strip-json-comments'; import { DefaultLight } from '../ui/themes/default-light.js'; @@ -346,7 +346,7 @@ export class LoadedSettings { coreEvents.emitSettingsChanged(); } - setRemoteAdminSettings(remoteSettings: GeminiCodeAssistSetting): void { + setRemoteAdminSettings(remoteSettings: FetchAdminControlsResponse): void { const admin: Settings['admin'] = {}; const { secureModeEnabled, mcpSetting, cliFeatureSetting } = remoteSettings; diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 36411feae5..da6643349b 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -61,6 +61,7 @@ import { SessionStartSource, SessionEndReason, getVersion, + type FetchAdminControlsResponse, } from '@google/gemini-cli-core'; import { initializeApp, @@ -283,6 +284,14 @@ export async function startInteractiveUI( export async function main() { const cliStartupHandle = startupProfiler.start('cli_startup'); + + // Listen for admin controls from parent process (IPC) in non-sandbox mode. In + // sandbox mode, we re-fetch the admin controls from the server once we enter + // the sandbox. + // TODO: Cache settings in sandbox mode as well. + const adminControlsListner = setupAdminControlsListener(); + registerCleanup(adminControlsListner.cleanup); + const cleanupStdio = patchStdio(); registerSyncCleanup(() => { // This is needed to ensure we don't lose any buffered output. @@ -358,6 +367,7 @@ export async function main() { const partialConfig = await loadCliConfig(settings.merged, sessionId, argv, { projectHooks: settings.workspace.settings.hooks, }); + adminControlsListner.setConfig(partialConfig); // Refresh auth to fetch remote admin settings from CCPA and before entering // the sandbox because the sandbox will interfere with the Oauth2 web @@ -451,7 +461,7 @@ export async function main() { } else { // Relaunch app so we always have a child process that can be internally // restarted if needed. - await relaunchAppInChildProcess(memoryArgs, []); + await relaunchAppInChildProcess(memoryArgs, [], remoteAdminSettings); } } @@ -464,6 +474,7 @@ export async function main() { projectHooks: settings.workspace.settings.hooks, }); loadConfigHandle?.end(); + adminControlsListner.setConfig(config); if (config.isInteractive() && config.storage && config.getDebugMode()) { const { registerActivityLogger } = await import( @@ -756,3 +767,36 @@ export function initializeOutputListenersAndFlush() { } coreEvents.drainBacklogs(); } + +function setupAdminControlsListener() { + let pendingSettings: FetchAdminControlsResponse | undefined; + let config: Config | undefined; + + const messageHandler = (msg: unknown) => { + const message = msg as { + type?: string; + settings?: FetchAdminControlsResponse; + }; + if (message?.type === 'admin-settings' && message.settings) { + if (config) { + config.setRemoteAdminSettings(message.settings); + } else { + pendingSettings = message.settings; + } + } + }; + + process.on('message', messageHandler); + + return { + setConfig: (newConfig: Config) => { + config = newConfig; + if (pendingSettings) { + config.setRemoteAdminSettings(pendingSettings); + } + }, + cleanup: () => { + process.off('message', messageHandler); + }, + }; +} diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index 083b636a2f..8e4ba82328 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/packages/cli/src/test-utils/render.tsx @@ -173,6 +173,7 @@ const mockUIActions: UIActions = { setBannerVisible: vi.fn(), setEmbeddedShellFocused: vi.fn(), setAuthContext: vi.fn(), + handleRestart: vi.fn(), }; export const renderWithProviders = ( diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 46dd1a69c2..5d23e86aed 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -186,6 +186,7 @@ export const AppContainer = (props: AppContainerProps) => { ); const [copyModeEnabled, setCopyModeEnabled] = useState(false); const [pendingRestorePrompt, setPendingRestorePrompt] = useState(false); + const [adminSettingsChanged, setAdminSettingsChanged] = useState(false); const [shellModeActive, setShellModeActive] = useState(false); const [modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError] = @@ -364,9 +365,18 @@ export const AppContainer = (props: AppContainerProps) => { setSettingsNonce((prev) => prev + 1); }; + const handleAdminSettingsChanged = () => { + setAdminSettingsChanged(true); + }; + coreEvents.on(CoreEvent.SettingsChanged, handleSettingsChanged); + coreEvents.on(CoreEvent.AdminSettingsChanged, handleAdminSettingsChanged); return () => { coreEvents.off(CoreEvent.SettingsChanged, handleSettingsChanged); + coreEvents.off( + CoreEvent.AdminSettingsChanged, + handleAdminSettingsChanged, + ); }; }, []); @@ -1452,6 +1462,7 @@ Logging in with Google... Restarting Gemini CLI to continue. const dialogsVisible = shouldShowIdePrompt || isFolderTrustDialogOpen || + adminSettingsChanged || !!shellConfirmationRequest || !!confirmationRequest || !!customDialog || @@ -1616,6 +1627,7 @@ Logging in with Google... Restarting Gemini CLI to continue. bannerVisible, terminalBackgroundColor: config.getTerminalBackground(), settingsNonce, + adminSettingsChanged, }), [ isThemeDialogOpen, @@ -1710,6 +1722,7 @@ Logging in with Google... Restarting Gemini CLI to continue. bannerVisible, config, settingsNonce, + adminSettingsChanged, ], ); @@ -1754,6 +1767,10 @@ Logging in with Google... Restarting Gemini CLI to continue. setBannerVisible, setEmbeddedShellFocused, setAuthContext, + handleRestart: async () => { + await runExitCleanup(); + process.exit(RELAUNCH_EXIT_CODE); + }, }), [ handleThemeSelect, diff --git a/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx b/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx new file mode 100644 index 0000000000..479c6950ff --- /dev/null +++ b/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx @@ -0,0 +1,51 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { renderWithProviders } from '../../test-utils/render.js'; +import { describe, it, expect, vi, afterEach } from 'vitest'; +import { act } from 'react'; +import { AdminSettingsChangedDialog } from './AdminSettingsChangedDialog.js'; + +const handleRestartMock = vi.fn(); + +describe('AdminSettingsChangedDialog', () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it('renders correctly', () => { + const { lastFrame } = renderWithProviders(); + expect(lastFrame()).toMatchSnapshot(); + }); + + it('restarts on "r" key press', async () => { + const { stdin } = renderWithProviders(, { + uiActions: { + handleRestart: handleRestartMock, + }, + }); + + act(() => { + stdin.write('r'); + }); + + expect(handleRestartMock).toHaveBeenCalled(); + }); + + it.each(['r', 'R'])('restarts on "%s" key press', async (key) => { + const { stdin } = renderWithProviders(, { + uiActions: { + handleRestart: handleRestartMock, + }, + }); + + act(() => { + stdin.write(key); + }); + + expect(handleRestartMock).toHaveBeenCalled(); + }); +}); diff --git a/packages/cli/src/ui/components/AdminSettingsChangedDialog.tsx b/packages/cli/src/ui/components/AdminSettingsChangedDialog.tsx new file mode 100644 index 0000000000..09571836c4 --- /dev/null +++ b/packages/cli/src/ui/components/AdminSettingsChangedDialog.tsx @@ -0,0 +1,36 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Box, Text } from 'ink'; +import { theme } from '../semantic-colors.js'; +import { useKeypress } from '../hooks/useKeypress.js'; +import { useUIActions } from '../contexts/UIActionsContext.js'; +import { Command, keyMatchers } from '../keyMatchers.js'; + +export const AdminSettingsChangedDialog = () => { + const { handleRestart } = useUIActions(); + + useKeypress( + (key) => { + if (keyMatchers[Command.RESTART_APP](key)) { + handleRestart(); + } + }, + { isActive: true }, + ); + + const message = + 'Admin settings have changed. Please restart the session to apply new settings.'; + + return ( + + + {message} Press 'r' to restart, or 'Ctrl+C' twice to + exit. + + + ); +}; diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 6a2fc46568..ce666346fd 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -30,6 +30,7 @@ import { useConfig } from '../contexts/ConfigContext.js'; import { useSettings } from '../contexts/SettingsContext.js'; import process from 'node:process'; import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js'; +import { AdminSettingsChangedDialog } from './AdminSettingsChangedDialog.js'; import { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js'; interface DialogManagerProps { @@ -50,6 +51,9 @@ export const DialogManager = ({ const { constrainHeight, terminalHeight, staticExtraHeight, mainAreaWidth } = uiState; + if (uiState.adminSettingsChanged) { + return ; + } if (uiState.showIdeRestartPrompt) { return ; } diff --git a/packages/cli/src/ui/components/__snapshots__/AdminSettingsChangedDialog.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/AdminSettingsChangedDialog.test.tsx.snap new file mode 100644 index 0000000000..9c5561db27 --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/AdminSettingsChangedDialog.test.tsx.snap @@ -0,0 +1,8 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`AdminSettingsChangedDialog > renders correctly 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ Admin settings have changed. Please restart the session to apply new settings. Press 'r' to restart, or 'Ctrl+C' │ +│ twice to exit. │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index 85839829f5..6358c26fa7 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -57,6 +57,7 @@ export interface UIActions { setBannerVisible: (visible: boolean) => void; setEmbeddedShellFocused: (value: boolean) => void; setAuthContext: (context: { requiresRestart?: boolean }) => void; + handleRestart: () => void; } export const UIActionsContext = createContext(null); diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 1175b0743a..236e48563a 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -141,6 +141,7 @@ export interface UIState { customDialog: React.ReactNode | null; terminalBackgroundColor: TerminalBackgroundColor; settingsNonce: number; + adminSettingsChanged: boolean; } export const UIStateContext = createContext(null); diff --git a/packages/cli/src/utils/relaunch.ts b/packages/cli/src/utils/relaunch.ts index d0405c32f1..131caf9bae 100644 --- a/packages/cli/src/utils/relaunch.ts +++ b/packages/cli/src/utils/relaunch.ts @@ -6,7 +6,10 @@ import { spawn } from 'node:child_process'; import { RELAUNCH_EXIT_CODE } from './processUtils.js'; -import { writeToStderr } from '@google/gemini-cli-core'; +import { + writeToStderr, + type FetchAdminControlsResponse, +} from '@google/gemini-cli-core'; export async function relaunchOnExitCode(runner: () => Promise) { while (true) { @@ -31,6 +34,7 @@ export async function relaunchOnExitCode(runner: () => Promise) { export async function relaunchAppInChildProcess( additionalNodeArgs: string[], additionalScriptArgs: string[], + remoteAdminSettings?: FetchAdminControlsResponse, ) { if (process.env['GEMINI_CLI_NO_RELAUNCH']) { return; @@ -55,10 +59,14 @@ export async function relaunchAppInChildProcess( process.stdin.pause(); const child = spawn(process.execPath, nodeArgs, { - stdio: 'inherit', + stdio: ['inherit', 'inherit', 'inherit', 'ipc'], env: newEnv, }); + if (remoteAdminSettings) { + child.send({ type: 'admin-settings', settings: remoteAdminSettings }); + } + return new Promise((resolve, reject) => { child.on('error', reject); child.on('close', (code) => { diff --git a/packages/core/src/code_assist/admin/admin_controls.test.ts b/packages/core/src/code_assist/admin/admin_controls.test.ts new file mode 100644 index 0000000000..22876f10a2 --- /dev/null +++ b/packages/core/src/code_assist/admin/admin_controls.test.ts @@ -0,0 +1,330 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + describe, + it, + expect, + vi, + beforeEach, + afterEach, + type Mock, +} from 'vitest'; +import { + fetchAdminControls, + sanitizeAdminSettings, + stopAdminControlsPolling, +} from './admin_controls.js'; +import type { CodeAssistServer } from '../server.js'; + +describe('Admin Controls', () => { + let mockServer: CodeAssistServer; + let mockOnSettingsChanged: Mock; + + beforeEach(() => { + vi.resetAllMocks(); + vi.useFakeTimers(); + + mockServer = { + projectId: 'test-project', + fetchAdminControls: vi.fn(), + } as unknown as CodeAssistServer; + + mockOnSettingsChanged = vi.fn(); + }); + + afterEach(() => { + stopAdminControlsPolling(); + vi.useRealTimers(); + }); + + describe('sanitizeAdminSettings', () => { + it('should strip unknown fields', () => { + const input = { + secureModeEnabled: true, + extraField: 'should be removed', + mcpSetting: { + mcpEnabled: false, + unknownMcpField: 'remove me', + }, + }; + + const result = sanitizeAdminSettings(input); + + expect(result).toEqual({ + secureModeEnabled: true, + mcpSetting: { + mcpEnabled: false, + }, + }); + // Explicitly check that unknown fields are gone + expect((result as Record)['extraField']).toBeUndefined(); + }); + + it('should preserve valid nested fields', () => { + const input = { + cliFeatureSetting: { + extensionsSetting: { + extensionsEnabled: true, + }, + }, + }; + expect(sanitizeAdminSettings(input)).toEqual(input); + }); + }); + + describe('fetchAdminControls', () => { + it('should return empty object and not poll if server is missing', async () => { + const result = await fetchAdminControls( + undefined, + undefined, + true, + mockOnSettingsChanged, + ); + 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 fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(result).toEqual({}); + expect(mockServer.fetchAdminControls).not.toHaveBeenCalled(); + }); + + it('should use cachedSettings and start polling if provided', async () => { + const cachedSettings = { secureModeEnabled: true }; + const result = await fetchAdminControls( + mockServer, + cachedSettings, + true, + mockOnSettingsChanged, + ); + + expect(result).toEqual(cachedSettings); + expect(mockServer.fetchAdminControls).not.toHaveBeenCalled(); + + // Should still start polling + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: false, + }); + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + }); + + it('should return empty object if admin controls are disabled', async () => { + const result = await fetchAdminControls( + mockServer, + undefined, + false, + mockOnSettingsChanged, + ); + expect(result).toEqual({}); + expect(mockServer.fetchAdminControls).not.toHaveBeenCalled(); + }); + + it('should fetch from server if no cachedSettings provided', async () => { + const serverResponse = { secureModeEnabled: true }; + (mockServer.fetchAdminControls as Mock).mockResolvedValue(serverResponse); + + const result = await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(result).toEqual(serverResponse); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + }); + + it('should return empty object on fetch error and still start polling', async () => { + (mockServer.fetchAdminControls as Mock).mockRejectedValue( + new Error('Network error'), + ); + const result = await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + + expect(result).toEqual({}); + + // Polling should have been started and should retry + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: true, + }); + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2); // Initial + poll + }); + + it('should sanitize server response', async () => { + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: true, + unknownField: 'bad', + }); + + const result = await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(result).toEqual({ secureModeEnabled: true }); + expect( + (result as Record)['unknownField'], + ).toBeUndefined(); + }); + + it('should reset polling interval if called again', async () => { + (mockServer.fetchAdminControls as Mock).mockResolvedValue({}); + + // First call + await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + + // Advance time, but not enough to trigger the poll + await vi.advanceTimersByTimeAsync(2 * 60 * 1000); + + // Second call, should reset the timer + await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2); + + // Advance time by 3 mins. If timer wasn't reset, it would have fired (2+3=5) + await vi.advanceTimersByTimeAsync(3 * 60 * 1000); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2); // No new poll + + // Advance time by another 2 mins. Now it should fire. + await vi.advanceTimersByTimeAsync(2 * 60 * 1000); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(3); // Poll fires + }); + }); + + describe('polling', () => { + it('should poll and emit changes', async () => { + // Initial fetch + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: false, + }); + await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + + // Update for next poll + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: true, + }); + + // Fast forward + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + + expect(mockOnSettingsChanged).toHaveBeenCalledWith({ + secureModeEnabled: true, + }); + }); + + it('should NOT emit if settings are deeply equal but not the same instance', async () => { + const settings = { secureModeEnabled: true }; + (mockServer.fetchAdminControls as Mock).mockResolvedValue(settings); + + await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + mockOnSettingsChanged.mockClear(); + + // Next poll returns a different object with the same values + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: true, + }); + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + + expect(mockOnSettingsChanged).not.toHaveBeenCalled(); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2); + }); + + it('should continue polling after a fetch error', async () => { + // Initial fetch is successful + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: false, + }); + await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + + // Next poll fails + (mockServer.fetchAdminControls as Mock).mockRejectedValue( + new Error('Poll failed'), + ); + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2); + expect(mockOnSettingsChanged).not.toHaveBeenCalled(); // No changes on error + + // Subsequent poll succeeds with new data + (mockServer.fetchAdminControls as Mock).mockResolvedValue({ + secureModeEnabled: true, + }); + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(3); + expect(mockOnSettingsChanged).toHaveBeenCalledWith({ + secureModeEnabled: true, + }); + }); + }); + + describe('stopAdminControlsPolling', () => { + it('should stop polling after it has started', async () => { + (mockServer.fetchAdminControls as Mock).mockResolvedValue({}); + + // Start polling + await fetchAdminControls( + mockServer, + undefined, + true, + mockOnSettingsChanged, + ); + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + + // Stop polling + stopAdminControlsPolling(); + + // Advance timer well beyond the polling interval + await vi.advanceTimersByTimeAsync(10 * 60 * 1000); + + // The poll should not have fired again + expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); + }); + }); +}); diff --git a/packages/core/src/code_assist/admin/admin_controls.ts b/packages/core/src/code_assist/admin/admin_controls.ts new file mode 100644 index 0000000000..93af330ecb --- /dev/null +++ b/packages/core/src/code_assist/admin/admin_controls.ts @@ -0,0 +1,113 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { CodeAssistServer } from '../server.js'; +import { debugLogger } from '../../utils/debugLogger.js'; +import { isDeepStrictEqual } from 'node:util'; +import { + type FetchAdminControlsResponse, + FetchAdminControlsResponseSchema, +} from '../types.js'; + +let pollingInterval: NodeJS.Timeout | undefined; +let currentSettings: FetchAdminControlsResponse | undefined; + +export function sanitizeAdminSettings( + settings: FetchAdminControlsResponse, +): FetchAdminControlsResponse { + const result = FetchAdminControlsResponseSchema.safeParse(settings); + if (!result.success) { + return {}; + } + return result.data; +} + +/** + * Fetches the admin controls from the server if enabled by experiment flag. + * Safely handles polling start/stop based on the flag and server availability. + * + * @param server The CodeAssistServer instance. + * @param cachedSettings The cached settings to use if available. + * @param adminControlsEnabled Whether admin controls are enabled. + * @param onSettingsChanged Callback to invoke when settings change during polling. + * @returns The fetched settings if enabled and successful, otherwise undefined. + */ +export async function fetchAdminControls( + server: CodeAssistServer | undefined, + cachedSettings: FetchAdminControlsResponse | undefined, + adminControlsEnabled: boolean, + onSettingsChanged: (settings: FetchAdminControlsResponse) => void, +): Promise { + if (!server || !server.projectId || !adminControlsEnabled) { + stopAdminControlsPolling(); + currentSettings = undefined; + return {}; + } + + // If we already have settings (e.g. from IPC during relaunch), use them + // to avoid blocking startup with another fetch. We'll still start polling. + if (cachedSettings) { + currentSettings = cachedSettings; + startAdminControlsPolling(server, server.projectId, onSettingsChanged); + return cachedSettings; + } + + try { + const rawSettings = await server.fetchAdminControls({ + project: server.projectId, + }); + const sanitizedSettings = sanitizeAdminSettings(rawSettings); + currentSettings = sanitizedSettings; + startAdminControlsPolling(server, server.projectId, onSettingsChanged); + return sanitizedSettings; + } catch (e) { + debugLogger.error('Failed to fetch admin controls: ', e); + // If initial fetch fails, start polling to retry. + currentSettings = {}; + startAdminControlsPolling(server, server.projectId, onSettingsChanged); + return {}; + } +} + +/** + * Starts polling for admin controls. + */ +function startAdminControlsPolling( + server: CodeAssistServer, + project: string, + onSettingsChanged: (settings: FetchAdminControlsResponse) => void, +) { + stopAdminControlsPolling(); + + pollingInterval = setInterval( + async () => { + try { + const rawSettings = await server.fetchAdminControls({ + project, + }); + const newSettings = sanitizeAdminSettings(rawSettings); + + if (!isDeepStrictEqual(newSettings, currentSettings)) { + currentSettings = newSettings; + onSettingsChanged(newSettings); + } + } catch (e) { + debugLogger.error('Failed to poll admin controls: ', e); + } + }, + 5 * 60 * 1000, + ); // 5 minutes +} + +/** + * Stops polling for admin controls. + */ +export function stopAdminControlsPolling() { + if (pollingInterval) { + clearInterval(pollingInterval); + pollingInterval = undefined; + } +} diff --git a/packages/core/src/code_assist/experiments/flagNames.ts b/packages/core/src/code_assist/experiments/flagNames.ts index 5fa70e7c20..71519dd40a 100644 --- a/packages/core/src/code_assist/experiments/flagNames.ts +++ b/packages/core/src/code_assist/experiments/flagNames.ts @@ -10,6 +10,7 @@ export const ExperimentFlags = { BANNER_TEXT_NO_CAPACITY_ISSUES: 45740199, BANNER_TEXT_CAPACITY_ISSUES: 45740200, ENABLE_PREVIEW: 45740196, + ENABLE_ADMIN_CONTROLS: 45752213, } as const; export type ExperimentFlagName = diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index 2fd45e0fd4..fca17b6d95 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -16,6 +16,8 @@ import type { ClientMetadata, RetrieveUserQuotaRequest, RetrieveUserQuotaResponse, + FetchAdminControlsRequest, + FetchAdminControlsResponse, ConversationOffered, ConversationInteraction, StreamingLatency, @@ -182,6 +184,15 @@ export class CodeAssistServer implements ContentGenerator { } } + async fetchAdminControls( + req: FetchAdminControlsRequest, + ): Promise { + return this.requestPost( + 'fetchAdminControls', + req, + ); + } + async getCodeAssistGlobalUserSetting(): Promise { return this.requestGet( 'getCodeAssistGlobalUserSetting', diff --git a/packages/core/src/code_assist/types.ts b/packages/core/src/code_assist/types.ts index 540ae63325..7f13d85398 100644 --- a/packages/core/src/code_assist/types.ts +++ b/packages/core/src/code_assist/types.ts @@ -4,6 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +import { z } from 'zod'; + export interface ClientMetadata { ideType?: ClientMetadataIdeType; ideVersion?: string; @@ -286,25 +288,29 @@ export interface ConversationInteraction { isAgentic?: boolean; } -export interface GeminiCodeAssistSetting { - secureModeEnabled?: boolean; - mcpSetting?: McpSetting; - cliFeatureSetting?: CliFeatureSetting; +export interface FetchAdminControlsRequest { + project: string; } -export interface McpSetting { - mcpEnabled?: boolean; - allowedMcpConfigs?: McpConfig[]; -} +export type FetchAdminControlsResponse = z.infer< + typeof FetchAdminControlsResponseSchema +>; -export interface McpConfig { - mcpServer?: string; -} +const ExtensionsSettingSchema = z.object({ + extensionsEnabled: z.boolean().optional(), +}); -export interface CliFeatureSetting { - extensionsSetting?: ExtensionsSetting; -} +const CliFeatureSettingSchema = z.object({ + extensionsSetting: ExtensionsSettingSchema.optional(), +}); -export interface ExtensionsSetting { - extensionsEnabled?: boolean; -} +const McpSettingSchema = z.object({ + mcpEnabled: z.boolean().optional(), + overrideMcpConfigJson: z.string().optional(), +}); + +export const FetchAdminControlsResponseSchema = z.object({ + secureModeEnabled: z.boolean().optional(), + mcpSetting: McpSettingSchema.optional(), + cliFeatureSetting: CliFeatureSettingSchema.optional(), +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 34a02a849d..ed6343ea35 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -92,7 +92,7 @@ import { ApprovalMode, type PolicyEngineConfig } from '../policy/types.js'; import { HookSystem } from '../hooks/index.js'; import type { UserTierId } from '../code_assist/types.js'; import type { RetrieveUserQuotaResponse } from '../code_assist/types.js'; -import type { GeminiCodeAssistSetting } from '../code_assist/types.js'; +import type { FetchAdminControlsResponse } from '../code_assist/types.js'; import { getCodeAssistServer } from '../code_assist/codeAssist.js'; import type { Experiments } from '../code_assist/experiments/experiments.js'; import { AgentRegistry } from '../agents/registry.js'; @@ -105,6 +105,7 @@ import { debugLogger } from '../utils/debugLogger.js'; import { SkillManager, type SkillDefinition } from '../skills/skillManager.js'; import { startupProfiler } from '../telemetry/startupProfiler.js'; import type { AgentDefinition } from '../agents/types.js'; +import { fetchAdminControls } from '../code_assist/admin/admin_controls.js'; export interface AccessibilitySettings { disableLoadingPhrases?: boolean; @@ -540,7 +541,7 @@ export class Config { private readonly planEnabled: boolean; private contextManager?: ContextManager; private terminalBackground: string | undefined = undefined; - private remoteAdminSettings: GeminiCodeAssistSetting | undefined; + private remoteAdminSettings: FetchAdminControlsResponse | undefined; private latestApiRequest: GenerateContentParameters | undefined; constructor(params: ConfigParameters) { @@ -909,6 +910,22 @@ export class Config { if (!this.hasAccessToPreviewModel && isPreviewModel(this.model)) { this.setModel(DEFAULT_GEMINI_MODEL_AUTO); } + + // Fetch admin controls + await this.ensureExperimentsLoaded(); + const adminControlsEnabled = + this.experiments?.flags[ExperimentFlags.ENABLE_ADMIN_CONTROLS] + ?.boolValue ?? false; + const adminControls = await fetchAdminControls( + codeAssistServer, + this.getRemoteAdminSettings(), + adminControlsEnabled, + (newSettings: FetchAdminControlsResponse) => { + this.setRemoteAdminSettings(newSettings); + coreEvents.emitAdminSettingsChanged(); + }, + ); + this.setRemoteAdminSettings(adminControls); } async getExperimentsAsync(): Promise { @@ -967,11 +984,11 @@ export class Config { this.latestApiRequest = req; } - getRemoteAdminSettings(): GeminiCodeAssistSetting | undefined { + getRemoteAdminSettings(): FetchAdminControlsResponse | undefined { return this.remoteAdminSettings; } - setRemoteAdminSettings(settings: GeminiCodeAssistSetting): void { + setRemoteAdminSettings(settings: FetchAdminControlsResponse): void { this.remoteAdminSettings = settings; } diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts index e6a15c68ab..79e440e9ad 100644 --- a/packages/core/src/utils/events.ts +++ b/packages/core/src/utils/events.ts @@ -119,6 +119,7 @@ export enum CoreEvent { HookStart = 'hook-start', HookEnd = 'hook-end', AgentsRefreshed = 'agents-refreshed', + AdminSettingsChanged = 'admin-settings-changed', RetryAttempt = 'retry-attempt', } @@ -133,6 +134,7 @@ export interface CoreEvents { [CoreEvent.HookStart]: [HookStartPayload]; [CoreEvent.HookEnd]: [HookEndPayload]; [CoreEvent.AgentsRefreshed]: never[]; + [CoreEvent.AdminSettingsChanged]: never[]; [CoreEvent.RetryAttempt]: [RetryAttemptPayload]; } @@ -242,6 +244,13 @@ export class CoreEventEmitter extends EventEmitter { this.emit(CoreEvent.AgentsRefreshed); } + /** + * Notifies subscribers that admin settings have changed. + */ + emitAdminSettingsChanged(): void { + this.emit(CoreEvent.AdminSettingsChanged); + } + /** * Notifies subscribers that a retry attempt is happening. */