fix(cli): Make IDE trust listener also listen to IDE status changes a… (#9783)

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(),

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,

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 (
-      <Box borderStyle="round" borderColor={theme.status.warning} paddingX={1}>
-        <Text color={theme.status.warning}>
-          Workspace trust has changed. Press &apos;r&apos; to restart Gemini to
-          apply the changes.
-        </Text>
-      </Box>
-    );
+    return <IdeTrustChangeDialog reason={uiState.ideTrustRestartReason} />;
   }
   if (uiState.showWorkspaceMigrationDialog) {
     return (

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(
+      <IdeTrustChangeDialog reason="CONNECTION_CHANGE" />,
+    );
+
+    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(
+      <IdeTrustChangeDialog reason="TRUST_CHANGE" />,
+    );
+
+    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(
+      <IdeTrustChangeDialog reason="NONE" />,
+    );
+
+    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(
+      <IdeTrustChangeDialog reason="NONE" />,
+    );
+
+    stdin.write('r');
+
+    expect(relaunchAppSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('calls relaunchApp when "R" is pressed', () => {
+    const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp');
+    const { stdin } = renderWithProviders(
+      <IdeTrustChangeDialog reason="CONNECTION_CHANGE" />,
+    );
+
+    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(
+      <IdeTrustChangeDialog reason="CONNECTION_CHANGE" />,
+    );
+
+    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();
+  });
+});

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 (
+    <Box borderStyle="round" borderColor={theme.status.warning} paddingX={1}>
+      <Text color={theme.status.warning}>
+        {message} Press &apos;r&apos; to restart Gemini to apply the changes.
+      </Text>
+    </Box>
+  );
+};

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<string, ExtensionUpdateState>;
   activePtyId: number | undefined;

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<typeof import('@google/gemini-cli-core')>();
+  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<ReturnType<typeof IdeClient.getInstance>>;
+  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);
+  });
+});

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>(
+    IDEConnectionStatus.Disconnected,
+  );
+  const previousTrust = useRef<boolean | undefined>(undefined);
+  const [restartReason, setRestartReason] = useState<RestartReason>('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 };
 }