From 0273154c6f7cea13294b5a4fe897456b2e24268f Mon Sep 17 00:00:00 2001 From: Alvin0412 Date: Thu, 11 Jun 2026 13:53:19 +0800 Subject: [PATCH] fix(cli): prompt for folder trust before auth --- .../cli/src/config/folderTrustPrompt.test.ts | 154 ++++++++++++++++++ packages/cli/src/config/folderTrustPrompt.ts | 115 +++++++++++++ packages/cli/src/gemini.test.tsx | 86 ++++++++++ packages/cli/src/gemini.tsx | 3 + 4 files changed, 358 insertions(+) create mode 100644 packages/cli/src/config/folderTrustPrompt.test.ts create mode 100644 packages/cli/src/config/folderTrustPrompt.ts diff --git a/packages/cli/src/config/folderTrustPrompt.test.ts b/packages/cli/src/config/folderTrustPrompt.test.ts new file mode 100644 index 00000000000..1dfe39143f7 --- /dev/null +++ b/packages/cli/src/config/folderTrustPrompt.test.ts @@ -0,0 +1,154 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import prompts from 'prompts'; +import { coreEvents, ExitCodes } from '@google/gemini-cli-core'; +import { relaunchApp } from '../utils/processUtils.js'; +import { + isWorkspaceTrusted, + loadTrustedFolders, + TrustLevel, +} from './trustedFolders.js'; +import { maybePromptForFolderTrust } from './folderTrustPrompt.js'; +import type { MergedSettings } from './settings.js'; + +vi.mock('prompts', () => ({ + default: vi.fn(), +})); + +vi.mock('../utils/processUtils.js', () => ({ + relaunchApp: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = + await importOriginal(); + const mockCoreEvents = Object.create( + actual.coreEvents, + ) as typeof actual.coreEvents; + mockCoreEvents.emitFeedback = vi.fn(); + return { + ...actual, + coreEvents: mockCoreEvents, + }; +}); + +vi.mock('./trustedFolders.js', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + isWorkspaceTrusted: vi.fn(), + loadTrustedFolders: vi.fn(), + }; +}); + +describe('maybePromptForFolderTrust', () => { + const settings = { + security: { + folderTrust: { + enabled: true, + }, + }, + } as MergedSettings; + const argv = { + prompt: undefined, + query: undefined, + isCommand: undefined, + }; + const setValue = vi.fn().mockResolvedValue(undefined); + const stdinIsTTY = Object.getOwnPropertyDescriptor(process.stdin, 'isTTY'); + const stdoutIsTTY = Object.getOwnPropertyDescriptor(process.stdout, 'isTTY'); + + beforeEach(() => { + vi.clearAllMocks(); + Object.defineProperty(process.stdin, 'isTTY', { + configurable: true, + value: true, + }); + Object.defineProperty(process.stdout, 'isTTY', { + configurable: true, + value: true, + }); + vi.mocked(isWorkspaceTrusted).mockReturnValue({ + isTrusted: undefined, + source: undefined, + }); + vi.mocked(loadTrustedFolders).mockReturnValue({ + setValue, + } as unknown as ReturnType); + vi.mocked(prompts).mockResolvedValue({ + trustLevel: TrustLevel.DO_NOT_TRUST, + }); + }); + + afterEach(() => { + Object.defineProperty(process.stdin, 'isTTY', { + configurable: true, + value: stdinIsTTY?.value, + }); + Object.defineProperty(process.stdout, 'isTTY', { + configurable: true, + value: stdoutIsTTY?.value, + }); + }); + + it('does nothing when the workspace trust state is already known', async () => { + vi.mocked(isWorkspaceTrusted).mockReturnValue({ + isTrusted: false, + source: 'file', + }); + + await maybePromptForFolderTrust(settings, argv, '/repo'); + + expect(prompts).not.toHaveBeenCalled(); + expect(setValue).not.toHaveBeenCalled(); + expect(relaunchApp).not.toHaveBeenCalled(); + }); + + it('saves an explicit do-not-trust decision and continues without relaunching', async () => { + await maybePromptForFolderTrust(settings, argv, '/repo'); + + expect(prompts).toHaveBeenCalledExactlyOnceWith( + expect.objectContaining({ + stdin: process.stdin, + stdout: expect.anything(), + }), + ); + expect(setValue).toHaveBeenCalledWith('/repo', TrustLevel.DO_NOT_TRUST); + expect(relaunchApp).not.toHaveBeenCalled(); + }); + + it('saves a trust decision and relaunches before authentication continues', async () => { + vi.mocked(prompts).mockResolvedValue({ + trustLevel: TrustLevel.TRUST_FOLDER, + }); + + await maybePromptForFolderTrust(settings, argv, '/repo'); + + expect(setValue).toHaveBeenCalledWith('/repo', TrustLevel.TRUST_FOLDER); + expect(relaunchApp).toHaveBeenCalledOnce(); + }); + + it('exits with a config error if the trust decision cannot be saved', async () => { + const processExitSpy = vi + .spyOn(process, 'exit') + .mockImplementation((code?: string | number | null) => { + throw new Error(`process.exit:${code}`); + }); + setValue.mockRejectedValueOnce(new Error('disk full')); + + await expect( + maybePromptForFolderTrust(settings, argv, '/repo'), + ).rejects.toThrow(`process.exit:${ExitCodes.FATAL_CONFIG_ERROR}`); + + expect(coreEvents.emitFeedback).toHaveBeenCalledWith( + 'error', + 'Failed to save trust settings. Exiting Gemini CLI.', + ); + processExitSpy.mockRestore(); + }); +}); diff --git a/packages/cli/src/config/folderTrustPrompt.ts b/packages/cli/src/config/folderTrustPrompt.ts new file mode 100644 index 00000000000..96e227dceb1 --- /dev/null +++ b/packages/cli/src/config/folderTrustPrompt.ts @@ -0,0 +1,115 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as path from 'node:path'; +import prompts from 'prompts'; +import { + coreEvents, + createWorkingStdio, + ExitCodes, + isHeadlessMode, +} from '@google/gemini-cli-core'; +import { runExitCleanup } from '../utils/cleanup.js'; +import { relaunchApp } from '../utils/processUtils.js'; +import type { CliArgs } from './config.js'; +import type { MergedSettings } from './settings.js'; +import { + isFolderTrustEnabled, + isTrustLevel, + isWorkspaceTrusted, + loadTrustedFolders, + TrustLevel, +} from './trustedFolders.js'; + +type FolderTrustPromptArgs = Pick; + +async function exitAfterFailedTrustSave(): Promise { + coreEvents.emitFeedback( + 'error', + 'Failed to save trust settings. Exiting Gemini CLI.', + ); + await runExitCleanup(); + process.exit(ExitCodes.FATAL_CONFIG_ERROR); +} + +async function exitAfterCancelledTrustPrompt(): Promise { + await runExitCleanup(); + process.exit(ExitCodes.FATAL_CANCELLATION_ERROR); +} + +/** + * Resolves an unknown workspace trust state before slow startup work like auth. + * + * Trusting a folder requires a relaunch so startup can reload workspace settings, + * hooks, MCP servers, and context files under the new trust decision. + */ +export async function maybePromptForFolderTrust( + settings: MergedSettings, + argv: FolderTrustPromptArgs, + workspaceDir: string = process.cwd(), +): Promise { + if (!isFolderTrustEnabled(settings)) { + return; + } + + if ( + argv.isCommand || + isHeadlessMode({ prompt: argv.prompt, query: argv.query }) + ) { + return; + } + + const { isTrusted } = isWorkspaceTrusted(settings, workspaceDir, { + prompt: argv.prompt, + query: argv.query, + }); + if (isTrusted !== undefined) { + return; + } + + const dirName = path.basename(workspaceDir); + const parentFolder = path.basename(path.dirname(workspaceDir)); + const { stdout } = createWorkingStdio(); + const response = await prompts({ + type: 'select', + name: 'trustLevel', + message: 'Do you want to trust this folder?', + choices: [ + { + title: `Trust folder (${dirName})`, + value: TrustLevel.TRUST_FOLDER, + }, + { + title: `Trust parent folder (${parentFolder})`, + value: TrustLevel.TRUST_PARENT, + }, + { + title: "Don't trust", + value: TrustLevel.DO_NOT_TRUST, + }, + ], + initial: 0, + stdin: process.stdin, + stdout, + }); + + if (!isTrustLevel(response.trustLevel)) { + await exitAfterCancelledTrustPrompt(); + } + + try { + await loadTrustedFolders().setValue(workspaceDir, response.trustLevel); + } catch { + await exitAfterFailedTrustSave(); + } + + if ( + response.trustLevel === TrustLevel.TRUST_FOLDER || + response.trustLevel === TrustLevel.TRUST_PARENT + ) { + await relaunchApp(); + } +} diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 4c7fec75a72..d35639afa2e 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -265,6 +265,11 @@ vi.mock('./deferred.js', () => ({ defer: vi.fn((m) => m), })); +const folderTrustPromptMock = vi.hoisted(() => ({ + maybePromptForFolderTrust: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('./config/folderTrustPrompt.js', () => folderTrustPromptMock); + vi.mock('./ui/utils/mouse.js', () => ({ enableMouseEvents: vi.fn(), disableMouseEvents: vi.fn(), @@ -637,6 +642,87 @@ describe('gemini.tsx main function kitty protocol', () => { resumeSpy.mockRestore(); }); + it('should resolve folder trust before refreshing auth', async () => { + const processExitSpy = vi + .spyOn(process, 'exit') + .mockImplementation((code) => { + throw new MockProcessExitError(code); + }); + const refreshAuth = vi.fn().mockResolvedValue(undefined); + const partialConfig = createMockConfig({ + isInteractive: () => true, + getQuestion: () => '', + getSandbox: () => undefined, + refreshAuth, + }); + const finalConfig = createMockConfig({ + isInteractive: () => false, + getQuestion: () => '', + getSandbox: () => undefined, + getListExtensions: () => true, + getExtensions: () => [], + }); + vi.mocked(loadCliConfig) + .mockResolvedValueOnce(partialConfig) + .mockResolvedValueOnce(finalConfig); + vi.mocked(loadSandboxConfig).mockResolvedValue(undefined); + vi.mocked(loadSettings).mockReturnValue( + createMockSettings({ + merged: { + advanced: {}, + security: { auth: { selectedType: AuthType.LOGIN_WITH_GOOGLE } }, + ui: {}, + }, + }), + ); + vi.mocked(parseArguments).mockResolvedValue({ + model: undefined, + sandbox: undefined, + debug: undefined, + prompt: undefined, + promptInteractive: undefined, + query: undefined, + yolo: undefined, + approvalMode: undefined, + policy: undefined, + adminPolicy: undefined, + allowedMcpServerNames: undefined, + allowedTools: undefined, + experimentalAcp: undefined, + extensions: undefined, + listExtensions: undefined, + includeDirectories: undefined, + screenReader: undefined, + useWriteTodos: undefined, + resume: undefined, + sessionId: undefined, + listSessions: undefined, + deleteSession: undefined, + outputFormat: undefined, + fakeResponses: undefined, + recordResponses: undefined, + rawOutput: undefined, + acceptRawOutputRisk: undefined, + isCommand: undefined, + skipTrust: undefined, + }); + + try { + await main(); + } catch (e) { + if (!(e instanceof MockProcessExitError)) throw e; + } + + expect(folderTrustPromptMock.maybePromptForFolderTrust).toHaveBeenCalled(); + expect(refreshAuth).toHaveBeenCalled(); + expect( + folderTrustPromptMock.maybePromptForFolderTrust.mock + .invocationCallOrder[0], + ).toBeLessThan(refreshAuth.mock.invocationCallOrder[0]); + expect(processExitSpy).toHaveBeenCalledWith(ExitCodes.SUCCESS); + processExitSpy.mockRestore(); + }); + it.each([ { flag: 'listExtensions' }, { flag: 'listSessions' }, diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index fb76a880684..e1922d05f24 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -95,6 +95,7 @@ import { relaunchOnExitCode } from './utils/relaunch.js'; import { loadSandboxConfig } from './config/sandboxConfig.js'; import { deleteSession, listSessions } from './utils/sessions.js'; import { createPolicyUpdater } from './config/policy.js'; +import { maybePromptForFolderTrust } from './config/folderTrustPrompt.js'; import { setupTerminalAndTheme } from './utils/terminalTheme.js'; import { runDeferredCommand } from './deferred.js'; @@ -464,6 +465,8 @@ export async function main() { process.exit(ExitCodes.FATAL_INPUT_ERROR); } + await maybePromptForFolderTrust(settings.merged, argv); + const isDebugMode = cliConfig.isDebugMode(argv); const consolePatcher = new ConsolePatcher({ stderr: argv.isCommand ? false : true,