diff --git a/packages/cli/src/config/settings.test.ts b/packages/cli/src/config/settings.test.ts index 3896a7f5de..809b8f48ff 100644 --- a/packages/cli/src/config/settings.test.ts +++ b/packages/cli/src/config/settings.test.ts @@ -82,6 +82,7 @@ import { FatalConfigError, GEMINI_DIR, Storage, + AuthType, type MCPServerConfig, } from '@google/gemini-cli-core'; import { updateSettingsFilePreservingFormat } from '../utils/commentJson.js'; @@ -202,6 +203,7 @@ describe('Settings Loading and Merging', () => { afterEach(() => { vi.restoreAllMocks(); + vi.unstubAllEnvs(); }); describe('loadSettings', () => { @@ -3036,6 +3038,7 @@ describe('Settings Loading and Merging', () => { delete process.env['CLOUD_SHELL']; delete process.env['MALICIOUS_VAR']; delete process.env['FOO']; + delete process.env['_GEMINI_USER_GCP_PROJECT']; vi.resetAllMocks(); vi.mocked(fs.existsSync).mockReturnValue(false); }); @@ -3268,6 +3271,108 @@ MALICIOUS_VAR=allowed-because-trusted expect(process.env['GOOGLE_CLOUD_PROJECT']).toBe('cloudshell-gca'); }); + it('should not override GOOGLE_CLOUD_PROJECT in Cloud Shell when auth type is vertex-ai', () => { + vi.stubEnv('CLOUD_SHELL', 'true'); + vi.stubEnv('GOOGLE_CLOUD_PROJECT', 'my-vertex-project'); + process.argv = ['node', 'gemini', '-s', 'prompt']; + vi.mocked(isWorkspaceTrusted).mockReturnValue({ + isTrusted: false, + source: 'file', + }); + + // No .env file + vi.mocked(fs.existsSync).mockReturnValue(false); + + loadEnvironment( + createMockSettings({ + tools: { sandbox: false }, + security: { auth: { selectedType: AuthType.USE_VERTEX_AI } }, + }).merged, + MOCK_WORKSPACE_DIR, + ); + + expect(process.env['GOOGLE_CLOUD_PROJECT']).toBe('my-vertex-project'); + }); + + it('should clear cloudshell-gca when switching to Vertex AI without an original project', () => { + process.env['CLOUD_SHELL'] = 'true'; + process.argv = ['node', 'gemini', '-s', 'prompt']; + vi.mocked(isWorkspaceTrusted).mockReturnValue({ + isTrusted: false, + source: 'file', + }); + vi.mocked(fs.existsSync).mockReturnValue(false); + + // First call: normal Cloud Shell auth sets cloudshell-gca + loadEnvironment( + createMockSettings({ tools: { sandbox: false } }).merged, + MOCK_WORKSPACE_DIR, + ); + expect(process.env['GOOGLE_CLOUD_PROJECT']).toBe('cloudshell-gca'); + + // Second call: user switched to Vertex AI, should remove cloudshell-gca + loadEnvironment( + createMockSettings({ + tools: { sandbox: false }, + security: { auth: { selectedType: AuthType.USE_VERTEX_AI } }, + }).merged, + MOCK_WORKSPACE_DIR, + ); + expect(process.env['GOOGLE_CLOUD_PROJECT']).toBeUndefined(); + }); + + it('should restore original project when switching to Vertex AI after Cloud Shell override', () => { + process.env['CLOUD_SHELL'] = 'true'; + process.env['GOOGLE_CLOUD_PROJECT'] = 'my-real-project'; + process.argv = ['node', 'gemini', '-s', 'prompt']; + vi.mocked(isWorkspaceTrusted).mockReturnValue({ + isTrusted: false, + source: 'file', + }); + vi.mocked(fs.existsSync).mockReturnValue(false); + + // First call: saves original to _GEMINI_USER_GCP_PROJECT, sets cloudshell-gca + loadEnvironment( + createMockSettings({ tools: { sandbox: false } }).merged, + MOCK_WORKSPACE_DIR, + ); + expect(process.env['GOOGLE_CLOUD_PROJECT']).toBe('cloudshell-gca'); + expect(process.env['_GEMINI_USER_GCP_PROJECT']).toBe('my-real-project'); + + // Second call: switching to Vertex AI should restore the saved value + loadEnvironment( + createMockSettings({ + tools: { sandbox: false }, + security: { auth: { selectedType: AuthType.USE_VERTEX_AI } }, + }).merged, + MOCK_WORKSPACE_DIR, + ); + expect(process.env['GOOGLE_CLOUD_PROJECT']).toBe('my-real-project'); + }); + + it('should restore project after restart when child inherits cloudshell-gca', () => { + // Simulate child process after restart: inherits cloudshell-gca and + // the saved original from the parent process. + process.env['CLOUD_SHELL'] = 'true'; + process.env['GOOGLE_CLOUD_PROJECT'] = 'cloudshell-gca'; + process.env['_GEMINI_USER_GCP_PROJECT'] = 'my-real-project'; + process.argv = ['node', 'gemini', '-s', 'prompt']; + vi.mocked(isWorkspaceTrusted).mockReturnValue({ + isTrusted: false, + source: 'file', + }); + vi.mocked(fs.existsSync).mockReturnValue(false); + + loadEnvironment( + createMockSettings({ + tools: { sandbox: false }, + security: { auth: { selectedType: AuthType.USE_VERTEX_AI } }, + }).merged, + MOCK_WORKSPACE_DIR, + ); + expect(process.env['GOOGLE_CLOUD_PROJECT']).toBe('my-real-project'); + }); + it('should sanitize GOOGLE_CLOUD_PROJECT in Cloud Shell when loaded from .env in untrusted mode', () => { process.env['CLOUD_SHELL'] = 'true'; process.argv = ['node', 'gemini', '-s', 'prompt']; diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index aace550a25..2d94e719b2 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -18,6 +18,7 @@ import { Storage, coreEvents, homedir, + AuthType, type AdminControlsSettings, createCache, } from '@google/gemini-cli-core'; @@ -532,16 +533,45 @@ function findEnvFile(startDir: string, isTrusted: boolean): string | null { } } +// Internal env var used to preserve the user's original GOOGLE_CLOUD_PROJECT +// across process restarts in Cloud Shell. This survives relaunch because child +// processes inherit the parent's environment. +const USER_GCP_PROJECT = '_GEMINI_USER_GCP_PROJECT'; + export function setUpCloudShellEnvironment( envFilePath: string | null, isTrusted: boolean, isSandboxed: boolean, + selectedAuthType?: string, ): void { // Special handling for GOOGLE_CLOUD_PROJECT in Cloud Shell: // Because GOOGLE_CLOUD_PROJECT in Cloud Shell tracks the project // set by the user using "gcloud config set project" we do not want to // use its value. So, unless the user overrides GOOGLE_CLOUD_PROJECT in // one of the .env files, we set the Cloud Shell-specific default here. + // + // However, if the user has explicitly selected Vertex AI auth, they intend + // to use their own GCP project, so we restore the original value and skip + // the Cloud Shell override to respect their .env settings. + if (selectedAuthType === AuthType.USE_VERTEX_AI) { + const saved = process.env[USER_GCP_PROJECT]; + if (saved !== undefined) { + process.env['GOOGLE_CLOUD_PROJECT'] = saved; + } else if (process.env['GOOGLE_CLOUD_PROJECT'] === 'cloudshell-gca') { + delete process.env['GOOGLE_CLOUD_PROJECT']; + } + return; + } + + // Save the user's original value before overwriting, so it can be restored + // if the user later switches to Vertex AI (even after a process restart). + if (!process.env[USER_GCP_PROJECT]) { + const current = process.env['GOOGLE_CLOUD_PROJECT']; + if (current && current !== 'cloudshell-gca') { + process.env[USER_GCP_PROJECT] = current; + } + } + let value = 'cloudshell-gca'; if (envFilePath && fs.existsSync(envFilePath)) { @@ -584,7 +614,13 @@ export function loadEnvironment( // Cloud Shell environment variable handling if (process.env['CLOUD_SHELL'] === 'true') { - setUpCloudShellEnvironment(envFilePath, isTrusted, isSandboxed); + const selectedAuthType = settings.security?.auth?.selectedType; + setUpCloudShellEnvironment( + envFilePath, + isTrusted, + isSandboxed, + selectedAuthType, + ); } if (envFilePath) { diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index 83e69d6663..5f5ae4b8dc 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/packages/cli/src/test-utils/render.tsx @@ -592,6 +592,7 @@ const mockUIActions: UIActions = { setActiveBackgroundTaskPid: vi.fn(), setIsBackgroundTaskListOpen: vi.fn(), setAuthContext: vi.fn(), + dismissLoginRestart: vi.fn(), onHintInput: vi.fn(), onHintBackspace: vi.fn(), onHintClear: vi.fn(), diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index b196f8a05a..a09f477045 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -173,7 +173,6 @@ import { QUEUE_ERROR_DISPLAY_DURATION_MS, EXPAND_HINT_DURATION_MS, } from './constants.js'; -import { LoginWithGoogleRestartDialog } from './auth/LoginWithGoogleRestartDialog.js'; import { NewAgentsChoice } from './components/NewAgentsNotification.js'; import { isSlashCommand } from './utils/commandUtils.js'; import { parseSlashCommand } from '../utils/commands.js'; @@ -756,7 +755,7 @@ export const AppContainer = (props: AppContainerProps) => { useEffect(() => { if (authState === AuthState.Authenticated && authContext.requiresRestart) { - setAuthState(AuthState.AwaitingGoogleLoginRestart); + setAuthState(AuthState.AwaitingLoginRestart); setAuthContext({}); } }, [authState, authContext, setAuthState]); @@ -2185,8 +2184,13 @@ Logging in with Google... Restarting Gemini CLI to continue. const nightly = props.version.includes('nightly'); + const isAwaitingLoginRestart = authState === AuthState.AwaitingLoginRestart; + const loginRestartMessage = + settings.merged.security.auth.selectedType === AuthType.USE_VERTEX_AI + ? 'Authenticating to Vertex AI in Cloud Shell requires a restart to apply project settings.' + : undefined; + const dialogsVisible = - shouldShowIdePrompt || shouldShowIdePrompt || isFolderTrustDialogOpen || isPolicyUpdateDialogOpen || @@ -2214,6 +2218,7 @@ Logging in with Google... Restarting Gemini CLI to continue. !!emptyWalletRequest || isSessionBrowserOpen || authState === AuthState.AwaitingApiKeyInput || + isAwaitingLoginRestart || !!newAgents; const hasPendingToolConfirmation = useMemo( @@ -2447,6 +2452,8 @@ Logging in with Google... Restarting Gemini CLI to continue. accountSuspensionInfo, isAuthDialogOpen, isAwaitingApiKeyInput: authState === AuthState.AwaitingApiKeyInput, + isAwaitingLoginRestart, + loginRestartMessage, apiKeyDefaultValue, editorError, isEditorDialogOpen, @@ -2652,6 +2659,8 @@ Logging in with Google... Restarting Gemini CLI to continue. customDialog, apiKeyDefaultValue, authState, + isAwaitingLoginRestart, + loginRestartMessage, transientMessage, bannerData, bannerVisible, @@ -2726,6 +2735,10 @@ Logging in with Google... Restarting Gemini CLI to continue. setActiveBackgroundTaskPid, setIsBackgroundTaskListOpen, setAuthContext, + dismissLoginRestart: () => { + setAuthContext({}); + setAuthState(AuthState.Updating); + }, onHintInput: () => {}, onHintBackspace: () => {}, onHintClear: () => {}, @@ -2832,18 +2845,6 @@ Logging in with Google... Restarting Gemini CLI to continue. ], ); - if (authState === AuthState.AwaitingGoogleLoginRestart) { - return ( - { - setAuthContext({}); - setAuthState(AuthState.Updating); - }} - config={config} - /> - ); - } - return ( diff --git a/packages/cli/src/ui/auth/AuthDialog.test.tsx b/packages/cli/src/ui/auth/AuthDialog.test.tsx index 69593df076..0c4ec68f93 100644 --- a/packages/cli/src/ui/auth/AuthDialog.test.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.test.tsx @@ -243,6 +243,32 @@ describe('AuthDialog', () => { unmount(); }); + it('sets auth context with requiresRestart: true for USE_VERTEX_AI in Cloud Shell', async () => { + vi.stubEnv('CLOUD_SHELL', 'true'); + mockedValidateAuthMethod.mockReturnValue(null); + const { unmount } = await renderWithProviders(); + const { onSelect: handleAuthSelect } = + mockedRadioButtonSelect.mock.calls[0][0]; + await handleAuthSelect(AuthType.USE_VERTEX_AI); + + expect(props.setAuthContext).toHaveBeenCalledWith({ + requiresRestart: true, + }); + unmount(); + }); + + it('sets auth context with empty object for USE_VERTEX_AI outside Cloud Shell', async () => { + vi.stubEnv('CLOUD_SHELL', ''); + mockedValidateAuthMethod.mockReturnValue(null); + const { unmount } = await renderWithProviders(); + const { onSelect: handleAuthSelect } = + mockedRadioButtonSelect.mock.calls[0][0]; + await handleAuthSelect(AuthType.USE_VERTEX_AI); + + expect(props.setAuthContext).toHaveBeenCalledWith({}); + unmount(); + }); + it('sets auth context with empty object for other auth types', async () => { mockedValidateAuthMethod.mockReturnValue(null); const { unmount } = await renderWithProviders(); diff --git a/packages/cli/src/ui/auth/AuthDialog.tsx b/packages/cli/src/ui/auth/AuthDialog.tsx index e73d380bf3..4c52e29bc5 100644 --- a/packages/cli/src/ui/auth/AuthDialog.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.tsx @@ -119,7 +119,12 @@ export function AuthDialog({ return; } if (authType) { - if (authType === AuthType.LOGIN_WITH_GOOGLE) { + const needsRestart = + authType === AuthType.LOGIN_WITH_GOOGLE || + (authType === AuthType.USE_VERTEX_AI && + process.env['CLOUD_SHELL'] === 'true'); + + if (needsRestart) { setAuthContext({ requiresRestart: true }); } else { setAuthContext({}); diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx b/packages/cli/src/ui/auth/LoginRestartDialog.test.tsx similarity index 76% rename from packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx rename to packages/cli/src/ui/auth/LoginRestartDialog.test.tsx index 4dd13a3334..3c5f22109e 100644 --- a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx +++ b/packages/cli/src/ui/auth/LoginRestartDialog.test.tsx @@ -1,12 +1,12 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ import { render } from '../../test-utils/render.js'; import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; -import { LoginWithGoogleRestartDialog } from './LoginWithGoogleRestartDialog.js'; +import { LoginRestartDialog } from './LoginRestartDialog.js'; import { useKeypress } from '../hooks/useKeypress.js'; import { runExitCleanup } from '../../utils/cleanup.js'; import { @@ -27,7 +27,7 @@ vi.mock('../../utils/cleanup.js', () => ({ const mockedUseKeypress = useKeypress as Mock; const mockedRunExitCleanup = runExitCleanup as Mock; -describe('LoginWithGoogleRestartDialog', () => { +describe('LoginRestartDialog', () => { const onDismiss = vi.fn(); const exitSpy = vi .spyOn(process, 'exit') @@ -44,11 +44,20 @@ describe('LoginWithGoogleRestartDialog', () => { _resetRelaunchStateForTesting(); }); - it('renders correctly', async () => { + it('renders correctly with default message', async () => { const { lastFrame, unmount } = await render( - , + ); + expect(lastFrame()).toMatchSnapshot(); + unmount(); + }); + + it('renders correctly with custom message', async () => { + const { lastFrame, unmount } = await render( + , ); expect(lastFrame()).toMatchSnapshot(); @@ -57,10 +66,7 @@ describe('LoginWithGoogleRestartDialog', () => { it('calls onDismiss when escape is pressed', async () => { const { unmount } = await render( - , + , ); const keypressHandler = mockedUseKeypress.mock.calls[0][0]; @@ -82,10 +88,7 @@ describe('LoginWithGoogleRestartDialog', () => { vi.useFakeTimers(); const { unmount } = await render( - , + , ); const keypressHandler = mockedUseKeypress.mock.calls[0][0]; diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx b/packages/cli/src/ui/auth/LoginRestartDialog.tsx similarity index 70% rename from packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx rename to packages/cli/src/ui/auth/LoginRestartDialog.tsx index a781828d09..324f461ed1 100644 --- a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx +++ b/packages/cli/src/ui/auth/LoginRestartDialog.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -10,15 +10,17 @@ import { theme } from '../semantic-colors.js'; import { useKeypress } from '../hooks/useKeypress.js'; import { relaunchApp } from '../../utils/processUtils.js'; -interface LoginWithGoogleRestartDialogProps { +interface LoginRestartDialogProps { onDismiss: () => void; config: Config; + message?: string; } -export const LoginWithGoogleRestartDialog = ({ +export const LoginRestartDialog = ({ onDismiss, config, -}: LoginWithGoogleRestartDialogProps) => { + message, +}: LoginRestartDialogProps) => { useKeypress( (key) => { if (key.name === 'escape') { @@ -44,14 +46,20 @@ export const LoginWithGoogleRestartDialog = ({ { isActive: true }, ); - const message = + const displayMessage = + message ?? "You've successfully signed in with Google. Gemini CLI needs to be restarted."; return ( - + + {displayMessage} - {message} Press R to restart, or Esc to choose a different - authentication method. + Press R to restart, or Esc to choose a different authentication method. ); diff --git a/packages/cli/src/ui/auth/__snapshots__/LoginRestartDialog.test.tsx.snap b/packages/cli/src/ui/auth/__snapshots__/LoginRestartDialog.test.tsx.snap new file mode 100644 index 0000000000..33beb0e9d5 --- /dev/null +++ b/packages/cli/src/ui/auth/__snapshots__/LoginRestartDialog.test.tsx.snap @@ -0,0 +1,17 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`LoginRestartDialog > renders correctly with custom message 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ Authenticating to Vertex AI in Cloud Shell requires a restart to apply project settings. │ +│ Press R to restart, or Esc to choose a different authentication method. │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ +" +`; + +exports[`LoginRestartDialog > renders correctly with default message 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ You've successfully signed in with Google. Gemini CLI needs to be restarted. │ +│ Press R to restart, or Esc to choose a different authentication method. │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ +" +`; diff --git a/packages/cli/src/ui/auth/__snapshots__/LoginWithGoogleRestartDialog.test.tsx.snap b/packages/cli/src/ui/auth/__snapshots__/LoginWithGoogleRestartDialog.test.tsx.snap deleted file mode 100644 index 7c7a95e24f..0000000000 --- a/packages/cli/src/ui/auth/__snapshots__/LoginWithGoogleRestartDialog.test.tsx.snap +++ /dev/null @@ -1,9 +0,0 @@ -// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html - -exports[`LoginWithGoogleRestartDialog > renders correctly 1`] = ` -"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ -│ You've successfully signed in with Google. Gemini CLI needs to be restarted. Press R to restart, │ -│ or Esc to choose a different authentication method. │ -╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ -" -`; diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 40f0b06138..acd2f3472f 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -39,6 +39,7 @@ import { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js'; import { NewAgentsNotification } from './NewAgentsNotification.js'; import { AgentConfigDialog } from './AgentConfigDialog.js'; import { PolicyUpdateDialog } from './PolicyUpdateDialog.js'; +import { LoginRestartDialog } from '../auth/LoginRestartDialog.js'; interface DialogManagerProps { addItem: UseHistoryManagerReturn['addItem']; @@ -306,6 +307,17 @@ export const DialogManager = ({ ); } + if (uiState.isAwaitingLoginRestart) { + return ( + + + + ); + } if (uiState.isAuthDialogOpen) { return ( diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index fb979e7c17..cb89758300 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -87,6 +87,7 @@ export interface UIActions { setActiveBackgroundTaskPid: (pid: number) => void; setIsBackgroundTaskListOpen: (isOpen: boolean) => void; setAuthContext: (context: { requiresRestart?: boolean }) => void; + dismissLoginRestart: () => void; onHintInput: (char: string) => void; onHintBackspace: () => void; onHintClear: () => void; diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 2150218d44..eb998a9de0 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -101,6 +101,8 @@ export interface UIState { accountSuspensionInfo: AccountSuspensionInfo | null; isAuthDialogOpen: boolean; isAwaitingApiKeyInput: boolean; + isAwaitingLoginRestart: boolean; + loginRestartMessage?: string; apiKeyDefaultValue?: string; editorError: string | null; isEditorDialogOpen: boolean; diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 2808d716b7..cdaf37e342 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -42,8 +42,8 @@ export enum AuthState { AwaitingApiKeyInput = 'awaiting_api_key_input', // Successfully authenticated Authenticated = 'authenticated', - // Waiting for the user to restart after a Google login - AwaitingGoogleLoginRestart = 'awaiting_google_login_restart', + // Waiting for the user to restart after a login + AwaitingLoginRestart = 'awaiting_login_restart', } // Only defining the state enum needed by the UI