diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 497f367592..101725d0cc 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -70,6 +70,7 @@ vi.mock('./hooks/useBracketedPaste.js'); vi.mock('./hooks/useKeypress.js'); vi.mock('./hooks/useLoadingIndicator.js'); vi.mock('./hooks/useFolderTrust.js'); +vi.mock('./hooks/useIdeTrustListener.js'); vi.mock('./hooks/useMessageQueue.js'); vi.mock('./hooks/useAutoAcceptIndicator.js'); vi.mock('./hooks/useWorkspaceMigration.js'); @@ -96,6 +97,7 @@ import { useConsoleMessages } from './hooks/useConsoleMessages.js'; import { useGeminiStream } from './hooks/useGeminiStream.js'; import { useVim } from './hooks/vim.js'; import { useFolderTrust } from './hooks/useFolderTrust.js'; +import { useIdeTrustListener } from './hooks/useIdeTrustListener.js'; import { useMessageQueue } from './hooks/useMessageQueue.js'; import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js'; import { useWorkspaceMigration } from './hooks/useWorkspaceMigration.js'; @@ -127,6 +129,7 @@ describe('AppContainer State Management', () => { const mockedUseGeminiStream = useGeminiStream as Mock; const mockedUseVim = useVim as Mock; const mockedUseFolderTrust = useFolderTrust as Mock; + const mockedUseIdeTrustListener = useIdeTrustListener as Mock; const mockedUseMessageQueue = useMessageQueue as Mock; const mockedUseAutoAcceptIndicator = useAutoAcceptIndicator as Mock; const mockedUseWorkspaceMigration = useWorkspaceMigration as Mock; @@ -222,6 +225,10 @@ describe('AppContainer State Management', () => { handleFolderTrustSelect: vi.fn(), isRestarting: false, }); + mockedUseIdeTrustListener.mockReturnValue({ + needsRestart: false, + restartReason: 'NONE', + }); mockedUseMessageQueue.mockReturnValue({ messageQueue: [], addMessage: vi.fn(), diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 55bc26e68a..f2ed95fc2f 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -779,7 +779,10 @@ Logging in with Google... Please restart Gemini CLI to continue. const { isFolderTrustDialogOpen, handleFolderTrustSelect, isRestarting } = useFolderTrust(settings, setIsTrustedFolder); - const { needsRestart: ideNeedsRestart } = useIdeTrustListener(); + const { + needsRestart: ideNeedsRestart, + restartReason: ideTrustRestartReason, + } = useIdeTrustListener(); const isInitialMount = useRef(true); useEffect(() => { @@ -973,14 +976,6 @@ Logging in with Google... Please restart Gemini CLI to continue. ); useKeypress(handleGlobalKeypress, { isActive: true }); - useKeypress( - (key) => { - if (key.name === 'r' || key.name === 'R') { - process.exit(0); - } - }, - { isActive: showIdeRestartPrompt }, - ); // Update terminal title with Gemini CLI status and thoughts useEffect(() => { @@ -1051,6 +1046,7 @@ Logging in with Google... Please restart Gemini CLI to continue. isAuthDialogOpen || isEditorDialogOpen || showPrivacyNotice || + showIdeRestartPrompt || !!proQuotaRequest; const pendingHistoryItems = useMemo( @@ -1133,6 +1129,7 @@ Logging in with Google... Please restart Gemini CLI to continue. currentIDE, updateInfo, showIdeRestartPrompt, + ideTrustRestartReason, isRestarting, extensionsUpdateState, activePtyId, @@ -1209,6 +1206,7 @@ Logging in with Google... Please restart Gemini CLI to continue. currentIDE, updateInfo, showIdeRestartPrompt, + ideTrustRestartReason, isRestarting, currentModel, extensionsUpdateState, diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 9486f54792..23106dc696 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -27,6 +27,7 @@ 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 { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js'; interface DialogManagerProps { addItem: UseHistoryManagerReturn['addItem']; @@ -43,14 +44,7 @@ export const DialogManager = ({ addItem }: DialogManagerProps) => { uiState; if (uiState.showIdeRestartPrompt) { - return ( - - - Workspace trust has changed. Press 'r' to restart Gemini to - apply the changes. - - - ); + return ; } if (uiState.showWorkspaceMigrationDialog) { return ( diff --git a/packages/cli/src/ui/components/IdeTrustChangeDialog.test.tsx b/packages/cli/src/ui/components/IdeTrustChangeDialog.test.tsx new file mode 100644 index 0000000000..4f63095512 --- /dev/null +++ b/packages/cli/src/ui/components/IdeTrustChangeDialog.test.tsx @@ -0,0 +1,91 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import * as processUtils from '../../utils/processUtils.js'; +import { renderWithProviders } from '../../test-utils/render.js'; +import { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js'; + +describe('IdeTrustChangeDialog', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('renders the correct message for CONNECTION_CHANGE', () => { + const { lastFrame } = renderWithProviders( + , + ); + + const frameText = lastFrame(); + expect(frameText).toContain( + 'Workspace trust has changed due to a change in the IDE connection.', + ); + expect(frameText).toContain("Press 'r' to restart Gemini"); + }); + + it('renders the correct message for TRUST_CHANGE', () => { + const { lastFrame } = renderWithProviders( + , + ); + + const frameText = lastFrame(); + expect(frameText).toContain( + 'Workspace trust has changed due to a change in the IDE trust.', + ); + expect(frameText).toContain("Press 'r' to restart Gemini"); + }); + + it('renders a generic message and logs an error for NONE reason', () => { + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + const { lastFrame } = renderWithProviders( + , + ); + + const frameText = lastFrame(); + expect(frameText).toContain('Workspace trust has changed.'); + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'IdeTrustChangeDialog rendered with unexpected reason "NONE"', + ); + }); + + it('calls relaunchApp when "r" is pressed', () => { + const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp'); + const { stdin } = renderWithProviders( + , + ); + + stdin.write('r'); + + expect(relaunchAppSpy).toHaveBeenCalledTimes(1); + }); + + it('calls relaunchApp when "R" is pressed', () => { + const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp'); + const { stdin } = renderWithProviders( + , + ); + + stdin.write('R'); + + expect(relaunchAppSpy).toHaveBeenCalledTimes(1); + }); + + it('does not call relaunchApp when another key is pressed', async () => { + const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp'); + const { stdin } = renderWithProviders( + , + ); + + stdin.write('a'); + + // Give it a moment to ensure no async actions are triggered + await new Promise((resolve) => setTimeout(resolve, 50)); + + expect(relaunchAppSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/cli/src/ui/components/IdeTrustChangeDialog.tsx b/packages/cli/src/ui/components/IdeTrustChangeDialog.tsx new file mode 100644 index 0000000000..ba1288afcd --- /dev/null +++ b/packages/cli/src/ui/components/IdeTrustChangeDialog.tsx @@ -0,0 +1,47 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Box, Text } from 'ink'; +import { theme } from '../semantic-colors.js'; +import { useKeypress } from '../hooks/useKeypress.js'; +import { relaunchApp } from '../../utils/processUtils.js'; +import { type RestartReason } from '../hooks/useIdeTrustListener.js'; + +interface IdeTrustChangeDialogProps { + reason: RestartReason; +} + +export const IdeTrustChangeDialog = ({ reason }: IdeTrustChangeDialogProps) => { + useKeypress( + (key) => { + if (key.name === 'r' || key.name === 'R') { + relaunchApp(); + } + }, + { isActive: true }, + ); + + let message = 'Workspace trust has changed.'; + if (reason === 'NONE') { + // This should not happen, but provides a fallback and a debug log. + console.error( + 'IdeTrustChangeDialog rendered with unexpected reason "NONE"', + ); + } else if (reason === 'CONNECTION_CHANGE') { + message = + 'Workspace trust has changed due to a change in the IDE connection.'; + } else if (reason === 'TRUST_CHANGE') { + message = 'Workspace trust has changed due to a change in the IDE trust.'; + } + + return ( + + + {message} Press 'r' to restart Gemini to apply the changes. + + + ); +}; diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 5e3197ee1f..5d58c0a8da 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -36,6 +36,7 @@ export interface ProQuotaDialogRequest { } import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js'; +import { type RestartReason } from '../hooks/useIdeTrustListener.js'; export interface UIState { history: HistoryItem[]; @@ -112,6 +113,7 @@ export interface UIState { currentIDE: IdeInfo | null; updateInfo: UpdateObject | null; showIdeRestartPrompt: boolean; + ideTrustRestartReason: RestartReason; isRestarting: boolean; extensionsUpdateState: Map; activePtyId: number | undefined; diff --git a/packages/cli/src/ui/hooks/useIdeTrustListener.test.ts b/packages/cli/src/ui/hooks/useIdeTrustListener.test.ts new file mode 100644 index 0000000000..e3d62a218c --- /dev/null +++ b/packages/cli/src/ui/hooks/useIdeTrustListener.test.ts @@ -0,0 +1,226 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** @vitest-environment jsdom */ + +import { renderHook, act } from '@testing-library/react'; +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { + IdeClient, + IDEConnectionStatus, + ideContextStore, + type IDEConnectionState, +} from '@google/gemini-cli-core'; +import { useIdeTrustListener } from './useIdeTrustListener.js'; +import * as trustedFolders from '../../config/trustedFolders.js'; +import { useSettings } from '../contexts/SettingsContext.js'; +import type { LoadedSettings } from '../../config/settings.js'; + +// Mock dependencies +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const original = + await importOriginal(); + const ideClientInstance = { + addTrustChangeListener: vi.fn(), + removeTrustChangeListener: vi.fn(), + addStatusChangeListener: vi.fn(), + removeStatusChangeListener: vi.fn(), + getConnectionStatus: vi.fn(() => ({ + status: IDEConnectionStatus.Disconnected, + })), + }; + return { + ...original, + IdeClient: { + getInstance: vi.fn().mockResolvedValue(ideClientInstance), + }, + ideContextStore: { + get: vi.fn(), + subscribe: vi.fn(), + }, + }; +}); + +vi.mock('../../config/trustedFolders.js'); +vi.mock('../contexts/SettingsContext.js'); + +describe('useIdeTrustListener', () => { + let mockSettings: LoadedSettings; + let mockIdeClient: Awaited>; + let trustChangeCallback: (isTrusted: boolean) => void; + let statusChangeCallback: (state: IDEConnectionState) => void; + + beforeEach(async () => { + vi.clearAllMocks(); + mockIdeClient = await IdeClient.getInstance(); + + mockSettings = { + merged: { + security: { + folderTrust: { + enabled: true, + }, + }, + }, + } as LoadedSettings; + + vi.mocked(useSettings).mockReturnValue(mockSettings); + + vi.mocked(mockIdeClient.addTrustChangeListener).mockImplementation((cb) => { + trustChangeCallback = cb; + }); + vi.mocked(mockIdeClient.addStatusChangeListener).mockImplementation( + (cb) => { + statusChangeCallback = cb; + }, + ); + }); + + it('should initialize correctly with no trust information', () => { + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: undefined, + source: undefined, + }); + + const { result } = renderHook(() => useIdeTrustListener()); + + expect(result.current.isIdeTrusted).toBe(undefined); + expect(result.current.needsRestart).toBe(false); + expect(result.current.restartReason).toBe('NONE'); + }); + + it('should NOT set needsRestart when connecting for the first time', async () => { + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: IDEConnectionStatus.Disconnected, + }); + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: true, + source: 'ide', + }); + const { result } = renderHook(() => useIdeTrustListener()); + + // Manually trigger the initial connection state for the test setup + await act(async () => { + statusChangeCallback({ status: IDEConnectionStatus.Disconnected }); + }); + + expect(result.current.isIdeTrusted).toBe(undefined); + expect(result.current.needsRestart).toBe(false); + + await act(async () => { + vi.mocked(ideContextStore.get).mockReturnValue({ + workspaceState: { isTrusted: true }, + }); + statusChangeCallback({ status: IDEConnectionStatus.Connected }); + }); + + expect(result.current.isIdeTrusted).toBe(true); + expect(result.current.needsRestart).toBe(false); + expect(result.current.restartReason).toBe('CONNECTION_CHANGE'); + }); + + it('should set needsRestart when IDE trust changes', async () => { + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: IDEConnectionStatus.Connected, + }); + vi.mocked(ideContextStore.get).mockReturnValue({ + workspaceState: { isTrusted: true }, + }); + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: true, + source: 'ide', + }); + + const { result } = renderHook(() => useIdeTrustListener()); + + // Manually trigger the initial connection state for the test setup + await act(async () => { + statusChangeCallback({ status: IDEConnectionStatus.Connected }); + }); + + expect(result.current.isIdeTrusted).toBe(true); + expect(result.current.needsRestart).toBe(false); + + await act(async () => { + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: false, + source: 'ide', + }); + vi.mocked(ideContextStore.get).mockReturnValue({ + workspaceState: { isTrusted: false }, + }); + trustChangeCallback(false); + }); + + expect(result.current.isIdeTrusted).toBe(false); + expect(result.current.needsRestart).toBe(true); + expect(result.current.restartReason).toBe('TRUST_CHANGE'); + }); + + it('should set needsRestart when IDE disconnects', async () => { + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: IDEConnectionStatus.Connected, + }); + vi.mocked(ideContextStore.get).mockReturnValue({ + workspaceState: { isTrusted: true }, + }); + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: true, + source: 'ide', + }); + + const { result } = renderHook(() => useIdeTrustListener()); + + // Manually trigger the initial connection state for the test setup + await act(async () => { + statusChangeCallback({ status: IDEConnectionStatus.Connected }); + }); + + expect(result.current.isIdeTrusted).toBe(true); + expect(result.current.needsRestart).toBe(false); + + await act(async () => { + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: undefined, + source: undefined, + }); + vi.mocked(ideContextStore.get).mockReturnValue(undefined); + statusChangeCallback({ status: IDEConnectionStatus.Disconnected }); + }); + + expect(result.current.isIdeTrusted).toBe(undefined); + expect(result.current.needsRestart).toBe(true); + expect(result.current.restartReason).toBe('CONNECTION_CHANGE'); + }); + + it('should NOT set needsRestart if trust value does not change', async () => { + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: IDEConnectionStatus.Connected, + }); + vi.mocked(ideContextStore.get).mockReturnValue({ + workspaceState: { isTrusted: true }, + }); + vi.mocked(trustedFolders.isWorkspaceTrusted).mockReturnValue({ + isTrusted: true, + source: 'ide', + }); + + const { result, rerender } = renderHook(() => useIdeTrustListener()); + + // Manually trigger the initial connection state for the test setup + await act(async () => { + statusChangeCallback({ status: IDEConnectionStatus.Connected }); + }); + + expect(result.current.isIdeTrusted).toBe(true); + expect(result.current.needsRestart).toBe(false); + + rerender(); + + expect(result.current.isIdeTrusted).toBe(true); + expect(result.current.needsRestart).toBe(false); + }); +}); diff --git a/packages/cli/src/ui/hooks/useIdeTrustListener.ts b/packages/cli/src/ui/hooks/useIdeTrustListener.ts index c585100747..f7b4cbb831 100644 --- a/packages/cli/src/ui/hooks/useIdeTrustListener.ts +++ b/packages/cli/src/ui/hooks/useIdeTrustListener.ts @@ -4,44 +4,87 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { useCallback, useEffect, useState, useSyncExternalStore } from 'react'; -import { IdeClient, ideContextStore } from '@google/gemini-cli-core'; +import { + useCallback, + useEffect, + useState, + useSyncExternalStore, + useRef, +} from 'react'; +import { + IdeClient, + IDEConnectionStatus, + ideContextStore, + type IDEConnectionState, +} from '@google/gemini-cli-core'; +import { useSettings } from '../contexts/SettingsContext.js'; +import { isWorkspaceTrusted } from '../../config/trustedFolders.js'; + +export type RestartReason = 'NONE' | 'CONNECTION_CHANGE' | 'TRUST_CHANGE'; /** * This hook listens for trust status updates from the IDE companion extension. - * It provides the current trust status from the IDE and a flag indicating - * if a restart is needed because the trust state has changed. + * It provides the current trust status from the IDE and a reason if a restart + * is needed because the trust state has changed. */ export function useIdeTrustListener() { + const settings = useSettings(); + const [connectionStatus, setConnectionStatus] = useState( + IDEConnectionStatus.Disconnected, + ); + const previousTrust = useRef(undefined); + const [restartReason, setRestartReason] = useState('NONE'); + const [needsRestart, setNeedsRestart] = useState(false); + const subscribe = useCallback((onStoreChange: () => void) => { + const handleStatusChange = (state: IDEConnectionState) => { + setConnectionStatus(state.status); + setRestartReason('CONNECTION_CHANGE'); + // Also notify useSyncExternalStore that the data has changed + onStoreChange(); + }; + + const handleTrustChange = () => { + setRestartReason('TRUST_CHANGE'); + onStoreChange(); + }; + (async () => { const ideClient = await IdeClient.getInstance(); - ideClient.addTrustChangeListener(onStoreChange); + ideClient.addTrustChangeListener(handleTrustChange); + ideClient.addStatusChangeListener(handleStatusChange); + setConnectionStatus(ideClient.getConnectionStatus().status); })(); return () => { (async () => { const ideClient = await IdeClient.getInstance(); - ideClient.removeTrustChangeListener(onStoreChange); + ideClient.removeTrustChangeListener(handleTrustChange); + ideClient.removeStatusChangeListener(handleStatusChange); })(); }; }, []); - const getSnapshot = () => ideContextStore.get()?.workspaceState?.isTrusted; + const getSnapshot = () => { + if (connectionStatus !== IDEConnectionStatus.Connected) { + return undefined; + } + return ideContextStore.get()?.workspaceState?.isTrusted; + }; const isIdeTrusted = useSyncExternalStore(subscribe, getSnapshot); - const [needsRestart, setNeedsRestart] = useState(false); - const [initialTrustValue] = useState(isIdeTrusted); - useEffect(() => { + const currentTrust = isWorkspaceTrusted(settings.merged).isTrusted; + // Trigger a restart if the overall trust status for the CLI has changed, + // but not on the initial trust value. if ( - !needsRestart && - initialTrustValue !== undefined && - initialTrustValue !== isIdeTrusted + previousTrust.current !== undefined && + previousTrust.current !== currentTrust ) { setNeedsRestart(true); } - }, [isIdeTrusted, initialTrustValue, needsRestart]); + previousTrust.current = currentTrust; + }, [isIdeTrusted, settings.merged]); - return { isIdeTrusted, needsRestart }; + return { isIdeTrusted, needsRestart, restartReason }; }