From 7dade1f0e2d088160e5bb0dade40c60d2d8138d1 Mon Sep 17 00:00:00 2001 From: shrutip90 Date: Wed, 17 Sep 2025 13:05:40 -0700 Subject: [PATCH] fix(cli): Auto restart CLI inner node process on trust change (#8378) --- packages/cli/src/gemini.test.tsx | 14 +++++ packages/cli/src/gemini.tsx | 57 +++++++++++++++---- .../ui/components/FolderTrustDialog.test.tsx | 24 ++++---- .../src/ui/components/FolderTrustDialog.tsx | 31 +++++----- packages/cli/src/utils/processUtils.test.ts | 22 +++++++ packages/cli/src/utils/processUtils.ts | 20 +++++++ packages/cli/src/utils/sandbox.ts | 27 ++++++--- 7 files changed, 150 insertions(+), 45 deletions(-) create mode 100644 packages/cli/src/utils/processUtils.test.ts create mode 100644 packages/cli/src/utils/processUtils.ts diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 78c0589f79..41dbbb8e19 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -163,11 +163,16 @@ describe('gemini.tsx main function', () => { }); describe('gemini.tsx main function kitty protocol', () => { + let originalEnvNoRelaunch: string | undefined; let setRawModeSpy: MockInstance< (mode: boolean) => NodeJS.ReadStream & { fd: 0 } >; beforeEach(() => { + // Set no relaunch in tests since process spawning causing issues in tests + originalEnvNoRelaunch = process.env['GEMINI_CLI_NO_RELAUNCH']; + process.env['GEMINI_CLI_NO_RELAUNCH'] = 'true'; + // eslint-disable-next-line @typescript-eslint/no-explicit-any if (!(process.stdin as any).setRawMode) { // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -176,6 +181,15 @@ describe('gemini.tsx main function kitty protocol', () => { setRawModeSpy = vi.spyOn(process.stdin, 'setRawMode'); }); + afterEach(() => { + // Restore original env variables + if (originalEnvNoRelaunch !== undefined) { + process.env['GEMINI_CLI_NO_RELAUNCH'] = originalEnvNoRelaunch; + } else { + delete process.env['GEMINI_CLI_NO_RELAUNCH']; + } + }); + it('should call setRawMode and detectAndEnableKittyProtocol when isInteractive is true', async () => { const { detectAndEnableKittyProtocol } = await import( './ui/utils/kittyProtocolDetector.js' diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index bb58022d22..60339ebeea 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -16,6 +16,7 @@ import dns from 'node:dns'; import { spawn } from 'node:child_process'; import { start_sandbox } from './utils/sandbox.js'; import type { DnsResolutionOrder, LoadedSettings } from './config/settings.js'; +import { RELAUNCH_EXIT_CODE } from './utils/processUtils.js'; import { loadSettings, SettingScope } from './config/settings.js'; import { themeManager } from './ui/themes/theme-manager.js'; import { getStartupWarnings } from './utils/startupWarnings.js'; @@ -103,17 +104,50 @@ function getNodeMemoryArgs(config: Config): string[] { return []; } -async function relaunchWithAdditionalArgs(additionalArgs: string[]) { - const nodeArgs = [...additionalArgs, ...process.argv.slice(1)]; - const newEnv = { ...process.env, GEMINI_CLI_NO_RELAUNCH: 'true' }; +async function relaunchOnExitCode(runner: () => Promise) { + while (true) { + try { + const exitCode = await runner(); + + if (exitCode !== RELAUNCH_EXIT_CODE) { + process.exit(exitCode); + } + } catch (error) { + process.stdin.resume(); + console.error('Fatal error: Failed to relaunch the CLI process.', error); + process.exit(1); + } + } +} + +async function relaunchAppInChildProcess(additionalArgs: string[]) { + if (process.env['GEMINI_CLI_NO_RELAUNCH']) { + return; + } + + const runner = () => { + const nodeArgs = [...additionalArgs, ...process.argv.slice(1)]; + const newEnv = { ...process.env, GEMINI_CLI_NO_RELAUNCH: 'true' }; + + // The parent process should not be reading from stdin while the child is running. + process.stdin.pause(); const child = spawn(process.execPath, nodeArgs, { stdio: 'inherit', env: newEnv, }); - await new Promise((resolve) => child.on('close', resolve)); - process.exit(0); + return new Promise((resolve, reject) => { + child.on('error', reject); + child.on('close', (code) => { + // Resume stdin before the parent process exits. + process.stdin.resume(); + resolve(code ?? 1); + }); + }); + }; + + await relaunchOnExitCode(runner); } import { runZedIntegration } from './zed-integration/zedIntegration.js'; @@ -350,15 +384,14 @@ export async function main() { const sandboxArgs = injectStdinIntoArgs(process.argv, stdinData); - await start_sandbox(sandboxConfig, memoryArgs, config, sandboxArgs); + await relaunchOnExitCode(() => + start_sandbox(sandboxConfig, memoryArgs, config, sandboxArgs), + ); process.exit(0); } else { - // Not in a sandbox and not entering one, so relaunch with additional - // arguments to control memory usage if needed. - if (memoryArgs.length > 0) { - await relaunchWithAdditionalArgs(memoryArgs); - process.exit(0); - } + // Relaunch app so we always have a child process that can be internally + // restarted if needed. + await relaunchAppInChildProcess(memoryArgs); } } diff --git a/packages/cli/src/ui/components/FolderTrustDialog.test.tsx b/packages/cli/src/ui/components/FolderTrustDialog.test.tsx index c583617aae..5cf912937d 100644 --- a/packages/cli/src/ui/components/FolderTrustDialog.test.tsx +++ b/packages/cli/src/ui/components/FolderTrustDialog.test.tsx @@ -8,6 +8,11 @@ import { renderWithProviders } from '../../test-utils/render.js'; import { waitFor } from '@testing-library/react'; import { vi } from 'vitest'; import { FolderTrustDialog, FolderTrustChoice } from './FolderTrustDialog.js'; +import * as processUtils from '../../utils/processUtils.js'; + +vi.mock('../../utils/processUtils.js', () => ({ + relaunchApp: vi.fn(), +})); const mockedExit = vi.hoisted(() => vi.fn()); const mockedCwd = vi.hoisted(() => vi.fn()); @@ -69,21 +74,18 @@ describe('FolderTrustDialog', () => { , ); - expect(lastFrame()).toContain( - 'To see changes, Gemini CLI must be restarted', - ); + expect(lastFrame()).toContain(' Gemini CLI is restarting'); }); - it('should call process.exit when "r" is pressed and isRestarting is true', async () => { - const { stdin } = renderWithProviders( + it('should call relaunchApp when isRestarting is true', async () => { + vi.useFakeTimers(); + const relaunchApp = vi.spyOn(processUtils, 'relaunchApp'); + renderWithProviders( , ); - - stdin.write('r'); - - await waitFor(() => { - expect(mockedExit).toHaveBeenCalledWith(0); - }); + await vi.advanceTimersByTimeAsync(1000); + expect(relaunchApp).toHaveBeenCalled(); + vi.useRealTimers(); }); it('should not call process.exit when "r" is pressed and isRestarting is false', async () => { diff --git a/packages/cli/src/ui/components/FolderTrustDialog.tsx b/packages/cli/src/ui/components/FolderTrustDialog.tsx index 0060187925..f5a1c49b64 100644 --- a/packages/cli/src/ui/components/FolderTrustDialog.tsx +++ b/packages/cli/src/ui/components/FolderTrustDialog.tsx @@ -6,12 +6,15 @@ import { Box, Text } from 'ink'; import type React from 'react'; +import { useEffect } from 'react'; +import { theme } from '../semantic-colors.js'; import { Colors } from '../colors.js'; import type { RadioSelectItem } from './shared/RadioButtonSelect.js'; import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; import { useKeypress } from '../hooks/useKeypress.js'; import * as process from 'node:process'; import * as path from 'node:path'; +import { relaunchApp } from '../../utils/processUtils.js'; export enum FolderTrustChoice { TRUST_FOLDER = 'trust_folder', @@ -28,6 +31,17 @@ export const FolderTrustDialog: React.FC = ({ onSelect, isRestarting, }) => { + useEffect(() => { + const doRelaunch = async () => { + if (isRestarting) { + setTimeout(async () => { + await relaunchApp(); + }, 250); + } + }; + doRelaunch(); + }, [isRestarting]); + useKeypress( (key) => { if (key.name === 'escape') { @@ -37,20 +51,12 @@ export const FolderTrustDialog: React.FC = ({ { isActive: !isRestarting }, ); - useKeypress( - (key) => { - if (key.name === 'r') { - process.exit(0); - } - }, - { isActive: !!isRestarting }, - ); - + const dirName = path.basename(process.cwd()); const parentFolder = path.basename(path.dirname(process.cwd())); const options: Array> = [ { - label: 'Trust folder', + label: `Trust folder (${dirName})`, value: FolderTrustChoice.TRUST_FOLDER, }, { @@ -90,9 +96,8 @@ export const FolderTrustDialog: React.FC = ({ {isRestarting && ( - - To see changes, Gemini CLI must be restarted. Press r to exit and - apply changes now. + + Gemini CLI is restarting to apply the trust changes... )} diff --git a/packages/cli/src/utils/processUtils.test.ts b/packages/cli/src/utils/processUtils.test.ts new file mode 100644 index 0000000000..be85a4dbad --- /dev/null +++ b/packages/cli/src/utils/processUtils.test.ts @@ -0,0 +1,22 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi } from 'vitest'; +import { RELAUNCH_EXIT_CODE, relaunchApp } from './processUtils.js'; +import * as cleanup from './cleanup.js'; + +describe('processUtils', () => { + const processExit = vi + .spyOn(process, 'exit') + .mockReturnValue(undefined as never); + const runExitCleanup = vi.spyOn(cleanup, 'runExitCleanup'); + + it('should run cleanup and exit with the relaunch code', async () => { + await relaunchApp(); + expect(runExitCleanup).toHaveBeenCalledTimes(1); + expect(processExit).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE); + }); +}); diff --git a/packages/cli/src/utils/processUtils.ts b/packages/cli/src/utils/processUtils.ts new file mode 100644 index 0000000000..7d3f1bcbf9 --- /dev/null +++ b/packages/cli/src/utils/processUtils.ts @@ -0,0 +1,20 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { runExitCleanup } from './cleanup.js'; + +/** + * Exit code used to signal that the CLI should be relaunched. + */ +export const RELAUNCH_EXIT_CODE = 42; + +/** + * Exits the process with a special code to signal that the parent process should relaunch it. + */ +export async function relaunchApp(): Promise { + await runExitCleanup(); + process.exit(RELAUNCH_EXIT_CODE); +} diff --git a/packages/cli/src/utils/sandbox.ts b/packages/cli/src/utils/sandbox.ts index dc4ef3bd7b..6a1592bbb2 100644 --- a/packages/cli/src/utils/sandbox.ts +++ b/packages/cli/src/utils/sandbox.ts @@ -188,7 +188,7 @@ export async function start_sandbox( nodeArgs: string[] = [], cliConfig?: Config, cliArgs: string[] = [], -) { +): Promise { const patcher = new ConsolePatcher({ debugMode: cliConfig?.getDebugMode() || !!process.env['DEBUG'], stderr: true, @@ -339,11 +339,17 @@ export async function start_sandbox( ); } // spawn child and let it inherit stdio + process.stdin.pause(); sandboxProcess = spawn(config.command, args, { stdio: 'inherit', }); - await new Promise((resolve) => sandboxProcess?.on('close', resolve)); - return; + return new Promise((resolve, reject) => { + sandboxProcess?.on('error', reject); + sandboxProcess?.on('close', (code) => { + process.stdin.resume(); + resolve(code ?? 1); + }); + }); } console.error(`hopping into sandbox (command: ${config.command}) ...`); @@ -790,22 +796,25 @@ export async function start_sandbox( } // spawn child and let it inherit stdio + process.stdin.pause(); sandboxProcess = spawn(config.command, args, { stdio: 'inherit', }); - sandboxProcess.on('error', (err) => { - console.error('Sandbox process error:', err); - }); + return new Promise((resolve, reject) => { + sandboxProcess.on('error', (err) => { + console.error('Sandbox process error:', err); + reject(err); + }); - await new Promise((resolve) => { sandboxProcess?.on('close', (code, signal) => { - if (code !== 0) { + process.stdin.resume(); + if (code !== 0 && code !== null) { console.log( `Sandbox process exited with code: ${code}, signal: ${signal}`, ); } - resolve(); + resolve(code ?? 1); }); }); } finally {