From 115f8bce8b23069b260b1fda7791b611da246e46 Mon Sep 17 00:00:00 2001 From: Dmitry Lyalin Date: Thu, 12 Mar 2026 08:21:43 -0400 Subject: [PATCH] Experiment with CLI startup fast paths --- packages/cli/src/config/config.test.ts | 40 +++++++- packages/cli/src/config/config.ts | 62 ++++++++---- packages/cli/src/config/extension-manager.ts | 24 +++-- packages/cli/src/config/settings.ts | 5 + packages/cli/src/gemini.test.tsx | 97 +++++++++++++++++++ packages/cli/src/gemini.tsx | 82 +++++++++++++--- packages/cli/src/ui/AppContainer.tsx | 17 +--- packages/cli/src/ui/auth/AuthDialog.test.tsx | 23 ++--- packages/cli/src/ui/auth/AuthDialog.tsx | 6 +- .../LoginWithGoogleRestartDialog.test.tsx | 64 ++++++------ .../ui/auth/LoginWithGoogleRestartDialog.tsx | 15 +-- .../cli/src/ui/components/DialogManager.tsx | 7 +- packages/cli/src/utils/cleanup.ts | 11 ++- packages/cli/src/utils/processUtils.test.ts | 33 ++++++- packages/cli/src/utils/processUtils.ts | 22 ++++- packages/core/src/config/config.ts | 20 ++++ packages/core/src/utils/extensionLoader.ts | 13 +++ packages/core/src/utils/getPty.ts | 45 ++++++--- 18 files changed, 436 insertions(+), 150 deletions(-) diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index 75812e4442..57d7f39ba7 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -26,6 +26,7 @@ import { type MergedSettings, createTestMergedSettings, } from './settings.js'; +import * as SettingsModule from './settings.js'; import * as ServerConfig from '@google/gemini-cli-core'; import { isWorkspaceTrusted } from './trustedFolders.js'; @@ -809,7 +810,7 @@ describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => { }); it('should pass extension context file paths to loadServerHierarchicalMemory', async () => { - process.argv = ['node', 'script.js']; + process.argv = ['node', 'script.js', '--prompt', 'test']; const settings = createTestMergedSettings(); vi.spyOn(ExtensionManager.prototype, 'getExtensions').mockReturnValue([ { @@ -859,7 +860,7 @@ describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => { }); it('should pass includeDirectories to loadServerHierarchicalMemory when loadMemoryFromIncludeDirectories is true', async () => { - process.argv = ['node', 'script.js']; + process.argv = ['node', 'script.js', '--prompt', 'test']; const includeDir = path.resolve(path.sep, 'path', 'to', 'include'); const settings = createTestMergedSettings({ context: { @@ -888,7 +889,7 @@ describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => { }); it('should NOT pass includeDirectories to loadServerHierarchicalMemory when loadMemoryFromIncludeDirectories is false', async () => { - process.argv = ['node', 'script.js']; + process.argv = ['node', 'script.js', '--prompt', 'test']; const settings = createTestMergedSettings({ context: { includeDirectories: ['/path/to/include'], @@ -914,6 +915,39 @@ describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => { 200, ); }); + + it('should skip extension, memory, and PTY discovery in bootstrap mode', async () => { + process.argv = ['node', 'script.js']; + const settings = createTestMergedSettings(); + const argv = await parseArguments(settings); + const loadSettingsSpy = vi.spyOn(SettingsModule, 'loadSettings'); + const getPtySpy = vi.spyOn(ServerConfig, 'getPty'); + + await loadCliConfig(settings, 'session-id', argv, { + mode: 'bootstrap', + loadedSettings: { merged: settings } as SettingsModule.LoadedSettings, + }); + + expect(loadSettingsSpy).not.toHaveBeenCalled(); + expect(ExtensionManager.prototype.loadExtensions).not.toHaveBeenCalled(); + expect(ServerConfig.loadServerHierarchicalMemory).not.toHaveBeenCalled(); + expect(getPtySpy).not.toHaveBeenCalled(); + }); + + it('should defer extension and memory discovery during interactive startup', async () => { + process.argv = ['node', 'script.js']; + const settings = createTestMergedSettings(); + const argv = await parseArguments(settings); + const getPtySpy = vi.spyOn(ServerConfig, 'getPty'); + + await loadCliConfig(settings, 'session-id', argv, { + loadedSettings: { merged: settings } as SettingsModule.LoadedSettings, + }); + + expect(ExtensionManager.prototype.loadExtensions).not.toHaveBeenCalled(); + expect(ServerConfig.loadServerHierarchicalMemory).not.toHaveBeenCalled(); + expect(getPtySpy).toHaveBeenCalled(); + }); }); describe('mergeMcpServers', () => { diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 6a4bd09470..1dd96883de 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -40,6 +40,7 @@ import { Config, applyAdminAllowlist, getAdminBlockedMcpServersMessage, + startupProfiler, type HookDefinition, type HookEventName, type OutputFormat, @@ -47,6 +48,7 @@ import { import { type Settings, type MergedSettings, + type LoadedSettings, saveModelChange, loadSettings, } from './settings.js'; @@ -436,6 +438,8 @@ export function isDebugMode(argv: CliArgs): boolean { export interface LoadCliConfigOptions { cwd?: string; + mode?: 'bootstrap' | 'full'; + loadedSettings?: LoadedSettings; projectHooks?: { [K in HookEventName]?: HookDefinition[] } & { disabled?: string[]; }; @@ -447,10 +451,15 @@ export async function loadCliConfig( argv: CliArgs, options: LoadCliConfigOptions = {}, ): Promise { - const { cwd = process.cwd(), projectHooks } = options; + const { + cwd = process.cwd(), + mode = 'full', + loadedSettings, + projectHooks, + } = options; const debugMode = isDebugMode(argv); - - const loadedSettings = loadSettings(cwd); + const resolvedLoadedSettings = loadedSettings ?? loadSettings(cwd); + const isBootstrap = mode === 'bootstrap'; if (argv.sandbox) { process.env['GEMINI_SANDBOX'] = 'true'; @@ -458,6 +467,17 @@ export async function loadCliConfig( const memoryImportFormat = settings.context?.importFormat || 'tree'; const includeDirectoryTree = settings.context?.includeDirectoryTree ?? true; + const interactive = + !!argv.promptInteractive || + !!argv.experimentalAcp || + (!isHeadlessMode({ prompt: argv.prompt, query: argv.query }) && + !argv.isCommand); + const shouldDeferInteractiveStartupWork = + !isBootstrap && + interactive && + !argv.listExtensions && + !argv.listSessions && + !argv.deleteSession; const ideMode = settings.ide?.enabled ?? false; @@ -499,6 +519,7 @@ export async function loadCliConfig( .map(resolvePath) .concat((argv.includeDirectories || []).map(resolvePath)); + const clientVersion = await getVersion(); const extensionManager = new ExtensionManager({ settings, requestConsent: requestConsentNonInteractive, @@ -507,9 +528,15 @@ export async function loadCliConfig( enabledExtensionOverrides: argv.extensions, // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion eventEmitter: coreEvents as EventEmitter, - clientVersion: await getVersion(), + clientVersion, }); - await extensionManager.loadExtensions(); + if (!isBootstrap && !shouldDeferInteractiveStartupWork) { + const loadExtensionsHandle = startupProfiler.start( + 'load_extensions_during_config', + ); + await extensionManager.loadExtensions(); + loadExtensionsHandle?.end(); + } const experimentalJitContext = settings.experimental?.jitContext ?? false; @@ -517,7 +544,12 @@ export async function loadCliConfig( let fileCount = 0; let filePaths: string[] = []; - if (!experimentalJitContext) { + if ( + !experimentalJitContext && + !isBootstrap && + !shouldDeferInteractiveStartupWork + ) { + const loadMemoryHandle = startupProfiler.start('load_memory'); // Call the (now wrapper) loadHierarchicalGeminiMemory which calls the server's version const result = await loadServerHierarchicalMemory( cwd, @@ -535,6 +567,7 @@ export async function loadCliConfig( memoryContent = result.memoryContent; fileCount = result.fileCount; filePaths = result.filePaths; + loadMemoryHandle?.end(); } const question = argv.promptInteractive || argv.prompt || ''; @@ -624,14 +657,6 @@ export async function loadCliConfig( throw err; } - // -p/--prompt forces non-interactive (headless) mode - // -i/--prompt-interactive forces interactive mode with an initial prompt - const interactive = - !!argv.promptInteractive || - !!argv.experimentalAcp || - (!isHeadlessMode({ prompt: argv.prompt, query: argv.query }) && - !argv.isCommand); - const allowedTools = argv.allowedTools || settings.tools?.allowed || []; const allowedToolsSet = new Set(allowedTools); @@ -723,7 +748,7 @@ export async function loadCliConfig( ? argv.screenReader : (settings.ui?.accessibility?.screenReader ?? false); - const ptyInfo = await getPty(); + const ptyInfo = isBootstrap ? null : await getPty(); const mcpEnabled = settings.admin?.mcp?.enabled ?? true; const extensionsEnabled = settings.admin?.extensions?.enabled ?? true; @@ -755,7 +780,7 @@ export async function loadCliConfig( return new Config({ sessionId, - clientVersion: await getVersion(), + clientVersion, embeddingModel: DEFAULT_GEMINI_EMBEDDING_MODEL, sandbox: sandboxConfig, targetDir: cwd, @@ -832,6 +857,8 @@ export async function loadCliConfig( skillsSupport: settings.skills?.enabled ?? true, disabledSkills: settings.skills?.disabled, experimentalJitContext: settings.experimental?.jitContext, + deferInitialMemoryLoad: + shouldDeferInteractiveStartupWork && !experimentalJitContext, modelSteering: settings.experimental?.modelSteering, toolOutputMasking: settings.experimental?.toolOutputMasking, noBrowser: !!process.env['NO_BROWSER'], @@ -871,7 +898,8 @@ export async function loadCliConfig( hooks: settings.hooks || {}, disabledHooks: settings.hooksConfig?.disabled || [], projectHooks: projectHooks || {}, - onModelChange: (model: string) => saveModelChange(loadedSettings, model), + onModelChange: (model: string) => + saveModelChange(resolvedLoadedSettings, model), onReload: async () => { const refreshedSettings = loadSettings(cwd); return { diff --git a/packages/cli/src/config/extension-manager.ts b/packages/cli/src/config/extension-manager.ts index 93ad3f3536..ed356a1ccf 100644 --- a/packages/cli/src/config/extension-manager.ts +++ b/packages/cli/src/config/extension-manager.ts @@ -90,7 +90,7 @@ interface ExtensionManagerParams { /** * Actual implementation of an ExtensionLoader. * - * You must call `loadExtensions` prior to calling other methods on this class. + * Extension metadata is loaded lazily via `loadExtensions`/`ensureLoaded`. */ export class ExtensionManager extends ExtensionLoader { private extensionEnablementManager: ExtensionEnablementManager; @@ -138,12 +138,24 @@ export class ExtensionManager extends ExtensionLoader { } getExtensions(): GeminiCLIExtension[] { - if (!this.loadedExtensions) { - throw new Error( - 'Extensions not yet loaded, must call `loadExtensions` first', - ); + return this.loadedExtensions ?? []; + } + + override isLoaded(): boolean { + return this.loadedExtensions !== undefined; + } + + override async ensureLoaded(): Promise { + if (this.loadedExtensions) { + return; } - return this.loadedExtensions; + + if (this.loadingPromise) { + await this.loadingPromise; + return; + } + + await this.loadExtensions(); } async installOrUpdateExtension( diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index 657968a3b6..7ec3237860 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -291,6 +291,11 @@ export function createTestMergedSettings( ) as MergedSettings; } +export function createDefaultMergedSettings(): MergedSettings { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return getDefaultsFromSchema() as MergedSettings; +} + /** * An immutable snapshot of settings state. * Used with useSyncExternalStore for reactive updates. diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 538fb8ee4e..a3d04b11bd 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -255,6 +255,7 @@ vi.mock('./validateNonInterActiveAuth.js', () => ({ describe('gemini.tsx main function', () => { let originalIsTTY: boolean | undefined; + const originalArgv = [...process.argv]; let initialUnhandledRejectionListeners: NodeJS.UnhandledRejectionListener[] = []; @@ -270,6 +271,7 @@ describe('gemini.tsx main function', () => { originalIsTTY = process.stdin.isTTY; // eslint-disable-next-line @typescript-eslint/no-explicit-any (process.stdin as any).isTTY = true; + process.argv = ['node', 'script.js']; }); afterEach(() => { @@ -282,11 +284,70 @@ describe('gemini.tsx main function', () => { // eslint-disable-next-line @typescript-eslint/no-explicit-any (process.stdin as any).isTTY = originalIsTTY; + process.argv = [...originalArgv]; vi.unstubAllEnvs(); vi.restoreAllMocks(); }); + it('should fast-path --version without loading settings', async () => { + process.argv = ['node', 'script.js', '--version']; + + await main(); + + expect(loadSettings).not.toHaveBeenCalled(); + expect(parseArguments).not.toHaveBeenCalled(); + }); + + it('should fast-path --help without loading settings', async () => { + process.argv = ['node', 'script.js', '--help']; + + await main(); + + expect(loadSettings).not.toHaveBeenCalled(); + expect(parseArguments).toHaveBeenCalledTimes(1); + }); + + it('should not relaunch a child process when sandboxing and memory relaunch are unnecessary', async () => { + vi.mocked(loadSandboxConfig).mockResolvedValue(undefined); + vi.mocked(loadSettings).mockReturnValue( + createMockSettings({ + merged: { advanced: {}, security: { auth: {} }, ui: {} }, + workspace: { settings: {} }, + setValue: vi.fn(), + forScope: () => ({ settings: {}, originalSettings: {}, path: '' }), + }), + ); + vi.mocked(parseArguments).mockResolvedValue({ + prompt: 'test', + } as unknown as CliArgs); + vi.mocked(loadCliConfig).mockResolvedValue( + createMockConfig({ + isInteractive: () => false, + getQuestion: () => 'test', + getSandbox: () => undefined, + }), + ); + const processExitSpy = vi + .spyOn(process, 'exit') + .mockImplementation((code) => { + throw new MockProcessExitError(code); + }); + + try { + await main(); + expect.fail('Should have thrown MockProcessExitError'); + } catch (e) { + expect(e).toBeInstanceOf(MockProcessExitError); + expect((e as MockProcessExitError).code).toBe(0); + } finally { + processExitSpy.mockRestore(); + } + + const { relaunchAppInChildProcess } = await import('./utils/relaunch.js'); + expect(relaunchAppInChildProcess).not.toHaveBeenCalled(); + }); + it('should log unhandled promise rejections and open debug console on first error', async () => { const processExitSpy = vi .spyOn(process, 'exit') @@ -641,6 +702,42 @@ describe('gemini.tsx main function kitty protocol', () => { processExitSpy.mockRestore(); }); + it('should use bootstrap config for the pre-relaunch startup pass', async () => { + const mockSettings = createMockSettings({ + merged: { + advanced: {}, + security: { auth: {} }, + ui: {}, + }, + workspace: { settings: {} }, + setValue: vi.fn(), + forScope: () => ({ settings: {}, originalSettings: {}, path: '' }), + }); + vi.mocked(loadSettings).mockReturnValue(mockSettings); + vi.mocked(parseArguments).mockResolvedValue({ + promptInteractive: false, + } as unknown as CliArgs); + + vi.mocked(loadCliConfig).mockResolvedValue( + createMockConfig({ + isInteractive: () => true, + getQuestion: () => '', + getSandbox: () => undefined, + }), + ); + + await main(); + + expect(loadCliConfig).toHaveBeenCalledTimes(2); + expect(vi.mocked(loadCliConfig).mock.calls[0][3]).toMatchObject({ + mode: 'bootstrap', + loadedSettings: mockSettings, + }); + expect(vi.mocked(loadCliConfig).mock.calls[1][3]).toMatchObject({ + loadedSettings: mockSettings, + }); + }); + it('should log warning when theme is not found', async () => { const { themeManager } = await import('./ui/themes/theme-manager.js'); const debugLoggerWarnSpy = vi diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index aa830c0250..ee2ed80556 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -21,7 +21,11 @@ import { loadTrustedFolders, type TrustedFoldersError, } from './config/trustedFolders.js'; -import { loadSettings, SettingScope } from './config/settings.js'; +import { + createDefaultMergedSettings, + loadSettings, + SettingScope, +} from './config/settings.js'; import { getStartupWarnings } from './utils/startupWarnings.js'; import { getUserStartupWarnings } from './utils/userStartupWarnings.js'; import { ConsolePatcher } from './ui/utils/ConsolePatcher.js'; @@ -110,6 +114,24 @@ import { runDeferredCommand } from './deferred.js'; import { SlashCommandConflictHandler } from './services/SlashCommandConflictHandler.js'; const SLOW_RENDER_MS = 200; +const HELP_FLAGS = new Set(['--help', '-h']); +const VERSION_FLAGS = new Set(['--version', '-v']); + +function getStartupFastPath(args: string[]): 'help' | 'version' | null { + if (args.length === 0) { + return null; + } + + if (args.every((arg) => HELP_FLAGS.has(arg))) { + return 'help'; + } + + if (args.every((arg) => VERSION_FLAGS.has(arg))) { + return 'version'; + } + + return null; +} export function validateDnsResolutionOrder( order: string | undefined, @@ -321,7 +343,36 @@ export async function startInteractiveUI( registerCleanup(() => instance.unmount()); } +async function runStartupCleanup( + config: Config, + settings: LoadedSettings, +): Promise { + const projectTempDir = config.storage.getProjectTempDir(); + try { + await Promise.all([ + cleanupCheckpoints(projectTempDir), + cleanupToolOutputFiles( + settings.merged, + config.getDebugMode(), + projectTempDir, + ), + ]); + } catch (error) { + debugLogger.warn('Deferred startup cleanup failed:', error); + } +} + export async function main() { + const fastPath = getStartupFastPath(process.argv.slice(2)); + if (fastPath === 'version') { + writeToStdout(`${await getVersion()}\n`); + return; + } + if (fastPath === 'help') { + await parseArguments(createDefaultMergedSettings()); + return; + } + const cliStartupHandle = startupProfiler.start('cli_startup'); // Listen for admin controls from parent process (IPC) in non-sandbox mode. In @@ -353,6 +404,10 @@ export async function main() { coreEvents.emitFeedback('warning', error.message); }); + const parseArgsHandle = startupProfiler.start('parse_arguments'); + const argv = await parseArguments(settings.merged); + parseArgsHandle?.end(); + const trustedFolders = loadTrustedFolders(); trustedFolders.errors.forEach((error: TrustedFoldersError) => { coreEvents.emitFeedback( @@ -361,15 +416,6 @@ export async function main() { ); }); - await Promise.all([ - cleanupCheckpoints(), - cleanupToolOutputFiles(settings.merged), - ]); - - const parseArgsHandle = startupProfiler.start('parse_arguments'); - const argv = await parseArguments(settings.merged); - parseArgsHandle?.end(); - if ( (argv.allowedTools && argv.allowedTools.length > 0) || (settings.merged.tools?.allowed && settings.merged.tools.allowed.length > 0) @@ -437,9 +483,13 @@ export async function main() { } } + const bootstrapConfigHandle = startupProfiler.start('load_bootstrap_config'); const partialConfig = await loadCliConfig(settings.merged, sessionId, argv, { + mode: 'bootstrap', + loadedSettings: settings, projectHooks: settings.workspace.settings.hooks, }); + bootstrapConfigHandle?.end(); adminControlsListner.setConfig(partialConfig); // Refresh auth to fetch remote admin settings from CCPA and before entering @@ -549,9 +599,7 @@ export async function main() { ); await runExitCleanup(); process.exit(ExitCodes.SUCCESS); - } else { - // Relaunch app so we always have a child process that can be internally - // restarted if needed. + } else if (memoryArgs.length > 0) { await relaunchAppInChildProcess(memoryArgs, [], remoteAdminSettings); } } @@ -562,6 +610,7 @@ export async function main() { { const loadConfigHandle = startupProfiler.start('load_cli_config'); const config = await loadCliConfig(settings.merged, sessionId, argv, { + loadedSettings: settings, projectHooks: settings.workspace.settings.hooks, }); loadConfigHandle?.end(); @@ -639,6 +688,13 @@ export async function main() { process.exit(ExitCodes.SUCCESS); } + const startupCleanup = runStartupCleanup(config, settings); + if (config.isInteractive()) { + void startupCleanup; + } else { + await startupCleanup; + } + const wasRaw = process.stdin.isRaw; if (config.isInteractive() && !wasRaw && process.stdin.isTTY) { // Set this as early as possible to avoid spurious characters from diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index b89d0b83c0..59714feb69 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -126,7 +126,7 @@ import { appEvents, AppEvent, TransientMessageType } from '../utils/events.js'; import { type UpdateObject } from './utils/updateCheck.js'; import { setUpdateHandler } from '../utils/handleAutoUpdate.js'; import { registerCleanup, runExitCleanup } from '../utils/cleanup.js'; -import { RELAUNCH_EXIT_CODE } from '../utils/processUtils.js'; +import { relaunchApp } from '../utils/processUtils.js'; import type { SessionInfo } from '../utils/sessionUtils.js'; import { useMessageQueue } from './hooks/useMessageQueue.js'; import { useMcpStatus } from './hooks/useMcpStatus.js'; @@ -778,13 +778,12 @@ export const AppContainer = (props: AppContainerProps) => { authType === AuthType.LOGIN_WITH_GOOGLE && config.isBrowserLaunchSuppressed() ) { - await runExitCleanup(); writeToStdout(` ---------------------------------------------------------------- Logging in with Google... Restarting Gemini CLI to continue. ---------------------------------------------------------------- `); - process.exit(RELAUNCH_EXIT_CODE); + await relaunchApp(config.getRemoteAdminSettings()); } } setAuthState(AuthState.Authenticated); @@ -2522,17 +2521,7 @@ Logging in with Google... Restarting Gemini CLI to continue. onHintClear: () => {}, onHintSubmit: () => {}, handleRestart: async () => { - if (process.send) { - const remoteSettings = config.getRemoteAdminSettings(); - if (remoteSettings) { - process.send({ - type: 'admin-settings-update', - settings: remoteSettings, - }); - } - } - await runExitCleanup(); - process.exit(RELAUNCH_EXIT_CODE); + await relaunchApp(config.getRemoteAdminSettings()); }, handleNewAgentsSelect: async (choice: NewAgentsChoice) => { if (newAgents && choice === NewAgentsChoice.ACKNOWLEDGE) { diff --git a/packages/cli/src/ui/auth/AuthDialog.test.tsx b/packages/cli/src/ui/auth/AuthDialog.test.tsx index c157a6a40d..1e150db574 100644 --- a/packages/cli/src/ui/auth/AuthDialog.test.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.test.tsx @@ -22,9 +22,8 @@ import { AuthState } from '../types.js'; import { RadioButtonSelect } from '../components/shared/RadioButtonSelect.js'; import { useKeypress } from '../hooks/useKeypress.js'; import { validateAuthMethodWithSettings } from './useAuth.js'; -import { runExitCleanup } from '../../utils/cleanup.js'; import { Text } from 'ink'; -import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js'; +import { relaunchApp } from '../../utils/processUtils.js'; // Mocks vi.mock('@google/gemini-cli-core', async (importOriginal) => { @@ -36,14 +35,14 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }; }); -vi.mock('../../utils/cleanup.js', () => ({ - runExitCleanup: vi.fn(), -})); - vi.mock('./useAuth.js', () => ({ validateAuthMethodWithSettings: vi.fn(), })); +vi.mock('../../utils/processUtils.js', () => ({ + relaunchApp: vi.fn().mockResolvedValue(undefined), +})); + vi.mock('../hooks/useKeypress.js', () => ({ useKeypress: vi.fn(), })); @@ -64,7 +63,7 @@ vi.mock('../components/shared/RadioButtonSelect.js', () => ({ const mockedUseKeypress = useKeypress as Mock; const mockedRadioButtonSelect = RadioButtonSelect as Mock; const mockedValidateAuthMethod = validateAuthMethodWithSettings as Mock; -const mockedRunExitCleanup = runExitCleanup as Mock; +const mockedRelaunchApp = relaunchApp as Mock; describe('AuthDialog', () => { let props: { @@ -85,6 +84,7 @@ describe('AuthDialog', () => { props = { config: { isBrowserLaunchSuppressed: vi.fn().mockReturnValue(false), + getRemoteAdminSettings: vi.fn().mockReturnValue(undefined), } as unknown as Config, settings: { merged: { @@ -351,11 +351,8 @@ describe('AuthDialog', () => { unmount(); }); - it('exits process for Login with Google when browser is suppressed', async () => { + it('restarts for Login with Google when browser is suppressed', async () => { vi.useFakeTimers(); - const exitSpy = vi - .spyOn(process, 'exit') - .mockImplementation(() => undefined as never); const logSpy = vi.spyOn(debugLogger, 'log').mockImplementation(() => {}); vi.mocked(props.config.isBrowserLaunchSuppressed).mockReturnValue(true); mockedValidateAuthMethod.mockReturnValue(null); @@ -371,10 +368,8 @@ describe('AuthDialog', () => { await vi.runAllTimersAsync(); }); - expect(mockedRunExitCleanup).toHaveBeenCalled(); - expect(exitSpy).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE); + expect(mockedRelaunchApp).toHaveBeenCalledWith(undefined); - exitSpy.mockRestore(); logSpy.mockRestore(); vi.useRealTimers(); unmount(); diff --git a/packages/cli/src/ui/auth/AuthDialog.tsx b/packages/cli/src/ui/auth/AuthDialog.tsx index 33652297b6..4877e55f4a 100644 --- a/packages/cli/src/ui/auth/AuthDialog.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.tsx @@ -21,9 +21,8 @@ import { } from '@google/gemini-cli-core'; import { useKeypress } from '../hooks/useKeypress.js'; import { AuthState } from '../types.js'; -import { runExitCleanup } from '../../utils/cleanup.js'; import { validateAuthMethodWithSettings } from './useAuth.js'; -import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js'; +import { relaunchApp } from '../../utils/processUtils.js'; interface AuthDialogProps { config: Config; @@ -134,8 +133,7 @@ export function AuthDialog({ ) { setExiting(true); setTimeout(async () => { - await runExitCleanup(); - process.exit(RELAUNCH_EXIT_CODE); + await relaunchApp(config.getRemoteAdminSettings()); }, 100); return; } diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx index 9079358348..18429fdb7d 100644 --- a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx +++ b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx @@ -8,27 +8,23 @@ import { render } from '../../test-utils/render.js'; import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; import { LoginWithGoogleRestartDialog } from './LoginWithGoogleRestartDialog.js'; import { useKeypress } from '../hooks/useKeypress.js'; -import { runExitCleanup } from '../../utils/cleanup.js'; -import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js'; import { type Config } from '@google/gemini-cli-core'; +import { relaunchApp } from '../../utils/processUtils.js'; // Mocks vi.mock('../hooks/useKeypress.js', () => ({ useKeypress: vi.fn(), })); -vi.mock('../../utils/cleanup.js', () => ({ - runExitCleanup: vi.fn(), +vi.mock('../../utils/processUtils.js', () => ({ + relaunchApp: vi.fn().mockResolvedValue(undefined), })); const mockedUseKeypress = useKeypress as Mock; -const mockedRunExitCleanup = runExitCleanup as Mock; +const mockedRelaunchApp = relaunchApp as Mock; describe('LoginWithGoogleRestartDialog', () => { const onDismiss = vi.fn(); - const exitSpy = vi - .spyOn(process, 'exit') - .mockImplementation(() => undefined as never); const mockConfig = { getRemoteAdminSettings: vi.fn(), @@ -36,7 +32,6 @@ describe('LoginWithGoogleRestartDialog', () => { beforeEach(() => { vi.clearAllMocks(); - exitSpy.mockClear(); vi.useRealTimers(); }); @@ -74,36 +69,33 @@ describe('LoginWithGoogleRestartDialog', () => { unmount(); }); - it.each(['r', 'R'])( - 'calls runExitCleanup and process.exit when %s is pressed', - async (keyName) => { - vi.useFakeTimers(); + it.each(['r', 'R'])('restarts when %s is pressed', async (keyName) => { + vi.useFakeTimers(); - const { waitUntilReady, unmount } = render( - , - ); - await waitUntilReady(); - const keypressHandler = mockedUseKeypress.mock.calls[0][0]; + const { waitUntilReady, unmount } = render( + , + ); + await waitUntilReady(); + const keypressHandler = mockedUseKeypress.mock.calls[0][0]; - keypressHandler({ - name: keyName, - shift: false, - ctrl: false, - cmd: false, - sequence: keyName, - }); + keypressHandler({ + name: keyName, + shift: false, + ctrl: false, + cmd: false, + sequence: keyName, + }); - // Advance timers to trigger the setTimeout callback - await vi.runAllTimersAsync(); + // Advance timers to trigger the setTimeout callback + await vi.runAllTimersAsync(); - expect(mockedRunExitCleanup).toHaveBeenCalledTimes(1); - expect(exitSpy).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE); + expect(mockedRelaunchApp).toHaveBeenCalledTimes(1); + expect(mockedRelaunchApp).toHaveBeenCalledWith(undefined); - vi.useRealTimers(); - unmount(); - }, - ); + vi.useRealTimers(); + unmount(); + }); }); diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx index 86cd645fee..a9c2055e0d 100644 --- a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx +++ b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx @@ -8,8 +8,7 @@ import { type Config } from '@google/gemini-cli-core'; import { Box, Text } from 'ink'; import { theme } from '../semantic-colors.js'; import { useKeypress } from '../hooks/useKeypress.js'; -import { runExitCleanup } from '../../utils/cleanup.js'; -import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js'; +import { relaunchApp } from '../../utils/processUtils.js'; interface LoginWithGoogleRestartDialogProps { onDismiss: () => void; @@ -27,17 +26,7 @@ export const LoginWithGoogleRestartDialog = ({ return true; } else if (key.name === 'r' || key.name === 'R') { setTimeout(async () => { - if (process.send) { - const remoteSettings = config.getRemoteAdminSettings(); - if (remoteSettings) { - process.send({ - type: 'admin-settings-update', - settings: remoteSettings, - }); - } - } - await runExitCleanup(); - process.exit(RELAUNCH_EXIT_CODE); + await relaunchApp(config.getRemoteAdminSettings()); }, 100); return true; } diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 3d56c68e5b..17a723e6ab 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -18,8 +18,7 @@ import { EditorSettingsDialog } from './EditorSettingsDialog.js'; import { PrivacyNotice } from '../privacy/PrivacyNotice.js'; import { ProQuotaDialog } from './ProQuotaDialog.js'; import { ValidationDialog } from './ValidationDialog.js'; -import { runExitCleanup } from '../../utils/cleanup.js'; -import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js'; +import { relaunchApp } from '../../utils/processUtils.js'; import { SessionBrowser } from './SessionBrowser.js'; import { PermissionsModifyTrustDialog } from './PermissionsModifyTrustDialog.js'; import { ModelDialog } from './ModelDialog.js'; @@ -28,7 +27,6 @@ import { useUIState } from '../contexts/UIStateContext.js'; import { useUIActions } from '../contexts/UIActionsContext.js'; 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'; @@ -258,8 +256,7 @@ export const DialogManager = ({ settings={settings} onSelect={() => uiActions.closeSettingsDialog()} onRestartRequest={async () => { - await runExitCleanup(); - process.exit(RELAUNCH_EXIT_CODE); + await relaunchApp(config.getRemoteAdminSettings()); }} availableTerminalHeight={terminalHeight - staticExtraHeight} config={config} diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts index 3fce73dd44..e8c603d4eb 100644 --- a/packages/cli/src/utils/cleanup.ts +++ b/packages/cli/src/utils/cleanup.ts @@ -100,10 +100,13 @@ async function drainStdin() { await new Promise((resolve) => setTimeout(resolve, 50)); } -export async function cleanupCheckpoints() { - const storage = new Storage(process.cwd()); - await storage.initialize(); - const tempDir = storage.getProjectTempDir(); +export async function cleanupCheckpoints(projectTempDir?: string) { + let tempDir = projectTempDir; + if (!tempDir) { + const storage = new Storage(process.cwd()); + await storage.initialize(); + tempDir = storage.getProjectTempDir(); + } const checkpointsDir = join(tempDir, 'checkpoints'); try { await fs.rm(checkpointsDir, { recursive: true, force: true }); diff --git a/packages/cli/src/utils/processUtils.test.ts b/packages/cli/src/utils/processUtils.test.ts index be85a4dbad..fc885e23c0 100644 --- a/packages/cli/src/utils/processUtils.test.ts +++ b/packages/cli/src/utils/processUtils.test.ts @@ -7,16 +7,45 @@ import { vi } from 'vitest'; import { RELAUNCH_EXIT_CODE, relaunchApp } from './processUtils.js'; import * as cleanup from './cleanup.js'; +import * as relaunch from './relaunch.js'; + +vi.mock('./relaunch.js', () => ({ + relaunchAppInChildProcess: vi.fn().mockResolvedValue(undefined), +})); describe('processUtils', () => { const processExit = vi .spyOn(process, 'exit') .mockReturnValue(undefined as never); const runExitCleanup = vi.spyOn(cleanup, 'runExitCleanup'); + const relaunchAppInChildProcess = vi.spyOn( + relaunch, + 'relaunchAppInChildProcess', + ); - it('should run cleanup and exit with the relaunch code', async () => { - await relaunchApp(); + afterEach(() => { + delete process.env['SANDBOX']; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + delete (process as any).send; + vi.clearAllMocks(); + }); + + it('should run cleanup and exit with the relaunch code when a parent wrapper exists', async () => { + process.env['SANDBOX'] = 'sandbox-exec'; + processExit.mockImplementationOnce(() => { + throw new Error('PROCESS_EXIT_CALLED'); + }); + + await expect(relaunchApp()).rejects.toThrow('PROCESS_EXIT_CALLED'); expect(runExitCleanup).toHaveBeenCalledTimes(1); expect(processExit).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE); + expect(relaunchAppInChildProcess).not.toHaveBeenCalled(); + }); + + it('should spawn a child wrapper on demand when no parent wrapper exists', async () => { + await relaunchApp(); + expect(runExitCleanup).toHaveBeenCalledTimes(1); + expect(relaunchAppInChildProcess).toHaveBeenCalledWith([], [], undefined); + expect(processExit).not.toHaveBeenCalled(); }); }); diff --git a/packages/cli/src/utils/processUtils.ts b/packages/cli/src/utils/processUtils.ts index 1122a2b0dc..1b967291c1 100644 --- a/packages/cli/src/utils/processUtils.ts +++ b/packages/cli/src/utils/processUtils.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { AdminControlsSettings } from '@google/gemini-cli-core'; import { runExitCleanup } from './cleanup.js'; /** @@ -12,9 +13,24 @@ import { runExitCleanup } from './cleanup.js'; export const RELAUNCH_EXIT_CODE = 199; /** - * Exits the process with a special code to signal that the parent process should relaunch it. + * Restarts the CLI, either by signaling an existing parent wrapper or by + * spawning one on demand. */ -export async function relaunchApp(): Promise { +export async function relaunchApp( + remoteAdminSettings?: AdminControlsSettings, +): Promise { await runExitCleanup(); - process.exit(RELAUNCH_EXIT_CODE); + + if (process.send || process.env['SANDBOX']) { + if (process.send && remoteAdminSettings) { + process.send({ + type: 'admin-settings-update', + settings: remoteAdminSettings, + }); + } + process.exit(RELAUNCH_EXIT_CODE); + } + + const { relaunchAppInChildProcess } = await import('./relaunch.js'); + await relaunchAppInChildProcess([], [], remoteAdminSettings); } diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 7297693b8e..2915e979a9 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -534,6 +534,7 @@ export interface ConfigParameters { disabledSkills?: string[]; adminSkillsEnabled?: boolean; experimentalJitContext?: boolean; + deferInitialMemoryLoad?: boolean; toolOutputMasking?: Partial; disableLLMCorrection?: boolean; plan?: boolean; @@ -733,6 +734,7 @@ export class Config { private readonly adminSkillsEnabled: boolean; private readonly experimentalJitContext: boolean; + private readonly deferInitialMemoryLoad: boolean; private readonly disableLLMCorrection: boolean; private readonly planEnabled: boolean; private readonly planModeRoutingEnabled: boolean; @@ -832,6 +834,7 @@ export class Config { this.adminSkillsEnabled = params.adminSkillsEnabled ?? true; this.modelAvailabilityService = new ModelAvailabilityService(); this.experimentalJitContext = params.experimentalJitContext ?? false; + this.deferInitialMemoryLoad = params.deferInitialMemoryLoad ?? false; this.modelSteering = params.modelSteering ?? false; this.userHintService = new UserHintService(() => this.isModelSteeringEnabled(), @@ -1047,6 +1050,23 @@ export class Config { this.workspaceContext.addDirectory(plansDir); } + if (!this.getExtensionLoader().isLoaded()) { + const extensionLoadHandle = startupProfiler.start('load_extensions'); + await this.getExtensionLoader().ensureLoaded(); + extensionLoadHandle?.end(); + } else { + await this.getExtensionLoader().ensureLoaded(); + } + + if (this.deferInitialMemoryLoad) { + const memoryLoadHandle = startupProfiler.start('load_memory'); + const { refreshServerHierarchicalMemory } = await import( + '../utils/memoryDiscovery.js' + ); + await refreshServerHierarchicalMemory(this); + memoryLoadHandle?.end(); + } + // Initialize centralized FileDiscoveryService const discoverToolsHandle = startupProfiler.start('discover_tools'); this.getFileService(); diff --git a/packages/core/src/utils/extensionLoader.ts b/packages/core/src/utils/extensionLoader.ts index 7110ba8615..e10fbe0e19 100644 --- a/packages/core/src/utils/extensionLoader.ts +++ b/packages/core/src/utils/extensionLoader.ts @@ -24,6 +24,19 @@ export abstract class ExtensionLoader { constructor(private readonly eventEmitter?: EventEmitter) {} + isLoaded(): boolean { + return true; + } + + /** + * Ensures extension metadata is available before initialization continues. + * + * Most implementations have nothing to do here, but loaders backed by disk + * discovery can override this to defer filesystem work until after first + * render. + */ + async ensureLoaded(): Promise {} + /** * All currently known extensions, both active and inactive. */ diff --git a/packages/core/src/utils/getPty.ts b/packages/core/src/utils/getPty.ts index b5d53ca473..349b86ab96 100644 --- a/packages/core/src/utils/getPty.ts +++ b/packages/core/src/utils/getPty.ts @@ -17,25 +17,38 @@ export interface PtyProcess { kill(signal?: string): void; } +let ptyImplementationPromise: Promise | undefined; + export const getPty = async (): Promise => { if (process.env['GEMINI_PTY_INFO'] === 'child_process') { return null; } - try { - const lydell = '@lydell/node-pty'; - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - const module = await import(lydell); - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - return { module, name: 'lydell-node-pty' }; - } catch (_e) { - try { - const nodePty = 'node-pty'; - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - const module = await import(nodePty); - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - return { module, name: 'node-pty' }; - } catch (_e2) { - return null; - } + if (!ptyImplementationPromise) { + ptyImplementationPromise = (async () => { + try { + const lydell = '@lydell/node-pty'; + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + const module = await import(lydell); + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + return { module, name: 'lydell-node-pty' }; + } catch (_e) { + try { + const nodePty = 'node-pty'; + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + const module = await import(nodePty); + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + return { module, name: 'node-pty' }; + } catch (_e2) { + return null; + } + } + })(); } + + return ptyImplementationPromise; }; + +/** For testing purposes only. */ +export function resetPtyCache(): void { + ptyImplementationPromise = undefined; +}