From 8c5f2708cce443a6f0fbf73183458b65b8624338 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Fri, 6 Mar 2026 01:27:05 -0500 Subject: [PATCH] perf(cli): optimize startup by skipping redundant parent authentication --- packages/cli/src/deferred.ts | 4 + packages/cli/src/gemini.test.tsx | 43 ++++++++- packages/cli/src/gemini.tsx | 91 +++++++++---------- packages/cli/src/ui/AppContainer.tsx | 9 -- .../ui/auth/LoginWithGoogleRestartDialog.tsx | 12 --- packages/cli/src/utils/relaunch.ts | 20 +--- packages/core/src/code_assist/oauth2.ts | 2 +- 7 files changed, 91 insertions(+), 90 deletions(-) diff --git a/packages/cli/src/deferred.ts b/packages/cli/src/deferred.ts index 1864ec2cb5..8780b3a5a3 100644 --- a/packages/cli/src/deferred.ts +++ b/packages/cli/src/deferred.ts @@ -25,6 +25,10 @@ export function setDeferredCommand(command: DeferredCommand) { deferredCommand = command; } +export function getDeferredCommand(): DeferredCommand | undefined { + return deferredCommand; +} + export async function runDeferredCommand(settings: MergedSettings) { if (!deferredCommand) { return; diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 2784c5694a..5b1ebc75be 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -43,6 +43,7 @@ import { debugLogger, coreEvents, AuthType, + fetchCachedCredentials, } from '@google/gemini-cli-core'; import { act } from 'react'; import { type InitializationResult } from './core/initializer.js'; @@ -130,6 +131,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { off: vi.fn(), drainBacklogs: vi.fn(), }, + fetchCachedCredentials: vi.fn().mockResolvedValue({}), }; }); @@ -911,12 +913,47 @@ describe('gemini.tsx main function exit codes', () => { } }); - it('should exit with 41 for auth failure during sandbox setup', async () => { + it('should skip refreshAuth in parent process when relaunch is expected', async () => { vi.stubEnv('SANDBOX', ''); + const refreshAuthSpy = vi.fn(); + vi.mocked(loadCliConfig).mockResolvedValue( + createMockConfig({ + refreshAuth: refreshAuthSpy, + getRemoteAdminSettings: vi.fn().mockReturnValue(undefined), + isInteractive: vi.fn().mockReturnValue(true), + }), + ); + vi.mocked(loadSettings).mockReturnValue( + createMockSettings({ + merged: { + security: { auth: { selectedType: 'google', useExternal: false } }, + advanced: { autoConfigureMemory: false }, + }, + }), + ); + vi.mocked(parseArguments).mockResolvedValue({} as CliArgs); + + // Initial process (no SANDBOX, no NO_RELAUNCH) + delete process.env['GEMINI_CLI_NO_RELAUNCH']; + + try { + await main(); + } catch (e) { + if (!(e instanceof MockProcessExitError)) throw e; + } + + expect(refreshAuthSpy).not.toHaveBeenCalled(); + }); + + it('should exit with 41 for auth failure during sandbox setup when auth IS required', async () => { + vi.stubEnv('SANDBOX', ''); + // Force mustAuthNow = true by simulating sandbox entry without credentials vi.mocked(loadSandboxConfig).mockResolvedValue({ command: 'docker', image: 'test-image', }); + vi.mocked(fetchCachedCredentials).mockResolvedValue(null); + vi.mocked(loadCliConfig).mockResolvedValue( createMockConfig({ refreshAuth: vi.fn().mockRejectedValue(new Error('Auth failed')), @@ -931,7 +968,9 @@ describe('gemini.tsx main function exit codes', () => { }, }), ); - vi.mocked(parseArguments).mockResolvedValue({} as CliArgs); + vi.mocked(parseArguments).mockResolvedValue({ + sandbox: true, + } as unknown as CliArgs); try { await main(); diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 6071488542..49856d7ece 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -72,7 +72,6 @@ import { getVersion, ValidationCancelledError, ValidationRequiredError, - type AdminControlsSettings, } from '@google/gemini-cli-core'; import { initializeApp, @@ -108,8 +107,9 @@ import { OverflowProvider } from './ui/contexts/OverflowContext.js'; import { setupTerminalAndTheme } from './utils/terminalTheme.js'; import { profiler } from './ui/components/DebugProfiler.js'; -import { runDeferredCommand } from './deferred.js'; +import { runDeferredCommand, getDeferredCommand } from './deferred.js'; import { SlashCommandConflictHandler } from './services/SlashCommandConflictHandler.js'; +import { fetchCachedCredentials } from '@google/gemini-cli-core'; const SLOW_RENDER_MS = 200; @@ -328,13 +328,6 @@ 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. @@ -446,13 +439,50 @@ export async function main() { const partialConfig = await loadCliConfig(settings.merged, sessionId, argv, { projectHooks: settings.workspace.settings.hooks, }); - adminControlsListner.setConfig(partialConfig); + + const isEnteringSandbox = !process.env['SANDBOX'] && !!argv.sandbox; + + // We relaunch if we are not already in a sandbox and not already in a + // relaunched process. relaunchAppInChildProcess ensures we always have a + // child process that can be internally restarted. + const willRelaunch = + !process.env['SANDBOX'] && !process.env['GEMINI_CLI_NO_RELAUNCH']; + + // Determine if we MUST authenticate in the parent process + const hasDeferredCommand = !!getDeferredCommand(); + let needsInteractiveSandboxLogin = false; + + if (isEnteringSandbox && !settings.merged.security.auth.useExternal) { + const authType = + settings.merged.security.auth.selectedType || + (process.env['CLOUD_SHELL'] === 'true' || + process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true' + ? AuthType.COMPUTE_ADC + : AuthType.LOGIN_WITH_GOOGLE); + + if (authType === AuthType.LOGIN_WITH_GOOGLE) { + const existingCredentials = await fetchCachedCredentials(); + if (!existingCredentials) { + // Sandbox environment might block interactive login, so do it here + needsInteractiveSandboxLogin = true; + } + } + } + + const mustAuthNow = hasDeferredCommand || needsInteractiveSandboxLogin; // Refresh auth to fetch remote admin settings from CCPA and before entering // the sandbox because the sandbox will interfere with the Oauth2 web // redirect. + // + // We skip this in the parent process if we are about to relaunch, as the + // child process will perform its own authentication and fetch admin settings + // itself. let initialAuthFailed = false; - if (!settings.merged.security.auth.useExternal) { + if ( + !settings.merged.security.auth.useExternal && + (!willRelaunch || mustAuthNow) + ) { try { if ( partialConfig.isInteractive() && @@ -501,6 +531,7 @@ export async function main() { } // Run deferred command now that we have admin settings. + // Note: if auth was skipped, settings.merged.admin will use defaults. await runDeferredCommand(settings.merged); // hop into sandbox if we are outside and sandboxing is enabled @@ -558,7 +589,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, [], remoteAdminSettings); + await relaunchAppInChildProcess(memoryArgs, []); } } @@ -577,8 +608,6 @@ export async function main() { // access to the project identifier. await config.storage.initialize(); - adminControlsListner.setConfig(config); - if (config.isInteractive() && settings.merged.general.devtools) { const { setupInitialActivityLogger } = await import( './utils/devtoolsService.js' @@ -870,37 +899,3 @@ export function initializeOutputListenersAndFlush() { } coreEvents.drainBacklogs(); } - -function setupAdminControlsListener() { - let pendingSettings: AdminControlsSettings | undefined; - let config: Config | undefined; - - const messageHandler = (msg: unknown) => { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const message = msg as { - type?: string; - settings?: AdminControlsSettings; - }; - 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/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 41cc5dec3d..7142808d39 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -2484,15 +2484,6 @@ 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 relaunchApp(); }, handleNewAgentsSelect: async (choice: NewAgentsChoice) => { diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx index 94ca359b59..9d9960c84f 100644 --- a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx +++ b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx @@ -4,7 +4,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -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'; @@ -12,12 +11,10 @@ import { relaunchApp } from '../../utils/processUtils.js'; interface LoginWithGoogleRestartDialogProps { onDismiss: () => void; - config: Config; } export const LoginWithGoogleRestartDialog = ({ onDismiss, - config, }: LoginWithGoogleRestartDialogProps) => { useKeypress( (key) => { @@ -26,15 +23,6 @@ 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 relaunchApp(); }, 100); return true; diff --git a/packages/cli/src/utils/relaunch.ts b/packages/cli/src/utils/relaunch.ts index 7e287e4565..b5e48e6729 100644 --- a/packages/cli/src/utils/relaunch.ts +++ b/packages/cli/src/utils/relaunch.ts @@ -6,10 +6,7 @@ import { spawn } from 'node:child_process'; import { RELAUNCH_EXIT_CODE } from './processUtils.js'; -import { - writeToStderr, - type AdminControlsSettings, -} from '@google/gemini-cli-core'; +import { writeToStderr } from '@google/gemini-cli-core'; export async function relaunchOnExitCode(runner: () => Promise) { while (true) { @@ -34,14 +31,11 @@ export async function relaunchOnExitCode(runner: () => Promise) { export async function relaunchAppInChildProcess( additionalNodeArgs: string[], additionalScriptArgs: string[], - remoteAdminSettings?: AdminControlsSettings, ) { if (process.env['GEMINI_CLI_NO_RELAUNCH']) { return; } - let latestAdminSettings = remoteAdminSettings; - const runner = () => { // process.argv is [node, script, ...args] // We want to construct [ ...nodeArgs, script, ...scriptArgs] @@ -61,20 +55,10 @@ export async function relaunchAppInChildProcess( process.stdin.pause(); const child = spawn(process.execPath, nodeArgs, { - stdio: ['inherit', 'inherit', 'inherit', 'ipc'], + stdio: ['inherit', 'inherit', 'inherit'], env: newEnv, }); - if (latestAdminSettings) { - child.send({ type: 'admin-settings', settings: latestAdminSettings }); - } - - child.on('message', (msg: { type?: string; settings?: unknown }) => { - if (msg.type === 'admin-settings-update' && msg.settings) { - latestAdminSettings = msg.settings as AdminControlsSettings; - } - }); - return new Promise((resolve, reject) => { child.on('error', reject); child.on('close', (code) => { diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index e238a4a860..e031a5c377 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/packages/core/src/code_assist/oauth2.ts @@ -642,7 +642,7 @@ export function getAvailablePort(): Promise { }); } -async function fetchCachedCredentials(): Promise< +export async function fetchCachedCredentials(): Promise< Credentials | JWTInput | null > { const useEncryptedStorage = getUseEncryptedStorageFlag();