Check folder trust before allowing add directory (#12652)

packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx

@@ -201,6 +201,7 @@ describe('useSlashCommandProcessor', () => {
           },
           new Map(), // extensionsUpdateState
           true, // isConfigInitialized
+          vi.fn(), // setCustomDialog
         ),
       );
       result = hook.result;

packages/cli/src/ui/hooks/slashCommandProcessor.ts

@@ -77,6 +77,7 @@ export const useSlashCommandProcessor = (
   actions: SlashCommandProcessorActions,
   extensionsUpdateState: Map<string, ExtensionUpdateStatus>,
   isConfigInitialized: boolean,
+  setCustomDialog: (dialog: React.ReactNode | null) => void,
 ) => {
   const session = useSessionStats();
   const [commands, setCommands] = useState<readonly SlashCommand[] | undefined>(
@@ -215,6 +216,7 @@ export const useSlashCommandProcessor = (
         dispatchExtensionStateUpdate: actions.dispatchExtensionStateUpdate,
         addConfirmUpdateExtensionRequest:
           actions.addConfirmUpdateExtensionRequest,
+        removeComponent: () => setCustomDialog(null),
       },
       session: {
         stats: session.stats,
@@ -239,6 +241,7 @@ export const useSlashCommandProcessor = (
       sessionShellAllowlist,
       reloadCommands,
       extensionsUpdateState,
+      setCustomDialog,
     ],
   );
 
@@ -505,6 +508,10 @@ export const useSlashCommandProcessor = (
                     true,
                   );
                 }
+                case 'custom_dialog': {
+                  setCustomDialog(result.component);
+                  return { type: 'handled' };
+                }
                 default: {
                   const unhandled: never = result;
                   throw new Error(
@@ -578,6 +585,7 @@ export const useSlashCommandProcessor = (
       setSessionShellAllowlist,
       setIsProcessing,
       setConfirmationRequest,
+      setCustomDialog,
     ],
   );
 

packages/cli/src/ui/hooks/useIncludeDirsTrust.test.tsx

@@ -0,0 +1,214 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { vi, describe, it, expect, beforeEach } from 'vitest';
+import type { Mock } from 'vitest';
+import { renderHook } from '../../test-utils/render.js';
+import { waitFor } from '../../test-utils/async.js';
+import { useIncludeDirsTrust } from './useIncludeDirsTrust.js';
+import * as trustedFolders from '../../config/trustedFolders.js';
+import type { Config, WorkspaceContext } from '@google/gemini-cli-core';
+import type { UseHistoryManagerReturn } from './useHistoryManager.js';
+import type { LoadedTrustedFolders } from '../../config/trustedFolders.js';
+
+import type { MultiFolderTrustDialogProps } from '../components/MultiFolderTrustDialog.js';
+
+vi.mock('../utils/directoryUtils.js', () => ({
+  expandHomeDir: (p: string) => p, // Simple pass-through for testing
+  loadMemoryFromDirectories: vi.fn().mockResolvedValue({ fileCount: 1 }),
+}));
+
+vi.mock('../components/MultiFolderTrustDialog.js', () => ({
+  MultiFolderTrustDialog: (props: MultiFolderTrustDialogProps) => (
+    <div data-testid="mock-dialog">{JSON.stringify(props.folders)}</div>
+  ),
+}));
+
+describe('useIncludeDirsTrust', () => {
+  let mockConfig: Config;
+  let mockHistoryManager: UseHistoryManagerReturn;
+  let mockSetCustomDialog: Mock;
+  let mockWorkspaceContext: WorkspaceContext;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    mockWorkspaceContext = {
+      addDirectory: vi.fn(),
+      getDirectories: vi.fn().mockReturnValue([]),
+      onDirectoriesChangedListeners: new Set(),
+      onDirectoriesChanged: vi.fn(),
+      notifyDirectoriesChanged: vi.fn(),
+      resolveAndValidateDir: vi.fn(),
+      getInitialDirectories: vi.fn(),
+      setDirectories: vi.fn(),
+      isPathWithinWorkspace: vi.fn(),
+      fullyResolvedPath: vi.fn(),
+      isPathWithinRoot: vi.fn(),
+      isFileSymlink: vi.fn(),
+    } as unknown as ReturnType<typeof mockConfig.getWorkspaceContext>;
+
+    mockConfig = {
+      getPendingIncludeDirectories: vi.fn().mockReturnValue([]),
+      clearPendingIncludeDirectories: vi.fn(),
+      getFolderTrust: vi.fn().mockReturnValue(true),
+      getWorkspaceContext: () => mockWorkspaceContext,
+      getGeminiClient: vi
+        .fn()
+        .mockReturnValue({ addDirectoryContext: vi.fn() }),
+    } as unknown as Config;
+
+    mockHistoryManager = {
+      addItem: vi.fn(),
+      history: [],
+      updateItem: vi.fn(),
+      clearItems: vi.fn(),
+      loadHistory: vi.fn(),
+    };
+    mockSetCustomDialog = vi.fn();
+  });
+
+  const renderTestHook = (isTrustedFolder: boolean | undefined) => {
+    renderHook(() =>
+      useIncludeDirsTrust(
+        mockConfig,
+        isTrustedFolder,
+        mockHistoryManager,
+        mockSetCustomDialog,
+      ),
+    );
+  };
+
+  it('should do nothing if isTrustedFolder is undefined', () => {
+    vi.mocked(mockConfig.getPendingIncludeDirectories).mockReturnValue([
+      '/foo',
+    ]);
+    renderTestHook(undefined);
+    expect(mockConfig.clearPendingIncludeDirectories).not.toHaveBeenCalled();
+  });
+
+  it('should do nothing if there are no pending directories', () => {
+    renderTestHook(true);
+    expect(mockConfig.clearPendingIncludeDirectories).not.toHaveBeenCalled();
+  });
+
+  describe('when folder trust is disabled or workspace is untrusted', () => {
+    it.each([
+      { trustEnabled: false, isTrusted: true, scenario: 'trust is disabled' },
+      {
+        trustEnabled: true,
+        isTrusted: false,
+        scenario: 'workspace is untrusted',
+      },
+    ])(
+      'should add directories directly when $scenario',
+      async ({ trustEnabled, isTrusted }) => {
+        vi.mocked(mockConfig.getFolderTrust).mockReturnValue(trustEnabled);
+        vi.mocked(mockConfig.getPendingIncludeDirectories).mockReturnValue([
+          '/dir1',
+          '/dir2',
+        ]);
+        vi.mocked(mockWorkspaceContext.addDirectory).mockImplementation(
+          (path) => {
+            if (path === '/dir2') {
+              throw new Error('Test error');
+            }
+          },
+        );
+
+        renderTestHook(isTrusted);
+
+        await waitFor(() => {
+          expect(mockWorkspaceContext.addDirectory).toHaveBeenCalledWith(
+            '/dir1',
+          );
+          expect(mockWorkspaceContext.addDirectory).toHaveBeenCalledWith(
+            '/dir2',
+          );
+          expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
+            expect.objectContaining({
+              text: expect.stringContaining("Error adding '/dir2': Test error"),
+            }),
+            expect.any(Number),
+          );
+          expect(
+            mockConfig.clearPendingIncludeDirectories,
+          ).toHaveBeenCalledTimes(1);
+        });
+      },
+    );
+  });
+
+  describe('when folder trust is enabled and workspace is trusted', () => {
+    let mockIsPathTrusted: Mock;
+
+    beforeEach(() => {
+      vi.spyOn(mockConfig, 'getFolderTrust').mockReturnValue(true);
+      mockIsPathTrusted = vi.fn();
+      const mockLoadedFolders = {
+        isPathTrusted: mockIsPathTrusted,
+      } as unknown as LoadedTrustedFolders;
+      vi.spyOn(trustedFolders, 'loadTrustedFolders').mockReturnValue(
+        mockLoadedFolders,
+      );
+    });
+
+    afterEach(() => {
+      vi.restoreAllMocks();
+    });
+
+    it('should add trusted dirs, collect untrusted errors, and open dialog for undefined', async () => {
+      const pendingDirs = ['/trusted', '/untrusted', '/undefined'];
+      vi.mocked(mockConfig.getPendingIncludeDirectories).mockReturnValue(
+        pendingDirs,
+      );
+
+      mockIsPathTrusted.mockImplementation((path: string) => {
+        if (path === '/trusted') return true;
+        if (path === '/untrusted') return false;
+        return undefined;
+      });
+
+      renderTestHook(true);
+
+      // Opens dialog for undefined trust dir
+      expect(mockSetCustomDialog).toHaveBeenCalledTimes(1);
+      const customDialogAction = mockSetCustomDialog.mock.calls[0][0];
+      expect(customDialogAction).toBeDefined();
+      const dialogProps = (
+        customDialogAction as React.ReactElement<MultiFolderTrustDialogProps>
+      ).props;
+      expect(dialogProps.folders).toEqual(['/undefined']);
+      expect(dialogProps.trustedDirs).toEqual(['/trusted']);
+      expect(dialogProps.errors as string[]).toEqual([
+        `The following directories are explicitly untrusted and cannot be added to a trusted workspace:\n- /untrusted\nPlease use the permissions command to modify their trust level.`,
+      ]);
+    });
+
+    it('should only add directories and clear pending if no dialog is needed', async () => {
+      const pendingDirs = ['/trusted1', '/trusted2'];
+      vi.mocked(mockConfig.getPendingIncludeDirectories).mockReturnValue(
+        pendingDirs,
+      );
+      mockIsPathTrusted.mockReturnValue(true);
+
+      renderTestHook(true);
+
+      await waitFor(() => {
+        expect(mockWorkspaceContext.addDirectory).toHaveBeenCalledWith(
+          '/trusted1',
+        );
+        expect(mockWorkspaceContext.addDirectory).toHaveBeenCalledWith(
+          '/trusted2',
+        );
+        expect(mockSetCustomDialog).not.toHaveBeenCalled();
+        expect(mockConfig.clearPendingIncludeDirectories).toHaveBeenCalledTimes(
+          1,
+        );
+      });
+    });
+  });
+});

packages/cli/src/ui/hooks/useIncludeDirsTrust.tsx

@@ -0,0 +1,160 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useEffect } from 'react';
+import type { Config } from '@google/gemini-cli-core';
+import { loadTrustedFolders } from '../../config/trustedFolders.js';
+import { expandHomeDir } from '../utils/directoryUtils.js';
+import { refreshServerHierarchicalMemory } from '@google/gemini-cli-core';
+import { MultiFolderTrustDialog } from '../components/MultiFolderTrustDialog.js';
+import type { UseHistoryManagerReturn } from './useHistoryManager.js';
+import { MessageType, type HistoryItem } from '../types.js';
+
+async function finishAddingDirectories(
+  config: Config,
+  addItem: (itemData: Omit<HistoryItem, 'id'>, baseTimestamp: number) => number,
+  added: string[],
+  errors: string[],
+) {
+  if (!config) {
+    addItem(
+      {
+        type: MessageType.ERROR,
+        text: 'Configuration is not available.',
+      },
+      Date.now(),
+    );
+    return;
+  }
+
+  try {
+    if (config.shouldLoadMemoryFromIncludeDirectories()) {
+      await refreshServerHierarchicalMemory(config);
+    }
+  } catch (error) {
+    errors.push(`Error refreshing memory: ${(error as Error).message}`);
+  }
+
+  if (added.length > 0) {
+    const gemini = config.getGeminiClient();
+    if (gemini) {
+      await gemini.addDirectoryContext();
+    }
+  }
+
+  if (errors.length > 0) {
+    addItem({ type: MessageType.ERROR, text: errors.join('\n') }, Date.now());
+  }
+}
+
+export function useIncludeDirsTrust(
+  config: Config,
+  isTrustedFolder: boolean | undefined,
+  historyManager: UseHistoryManagerReturn,
+  setCustomDialog: (dialog: React.ReactNode | null) => void,
+) {
+  const { addItem } = historyManager;
+
+  useEffect(() => {
+    // Don't run this until the initial trust is determined.
+    if (isTrustedFolder === undefined || !config) {
+      return;
+    }
+
+    const pendingDirs = config.getPendingIncludeDirectories();
+    if (pendingDirs.length === 0) {
+      return;
+    }
+
+    console.log('Inside useIncludeDirsTrust');
+
+    // If folder trust is disabled, isTrustedFolder will be undefined.
+    // In that case, or if the user decided not to trust the main folder,
+    // we can just add the directories without checking them.
+    if (config.getFolderTrust() === false || isTrustedFolder === false) {
+      const added: string[] = [];
+      const errors: string[] = [];
+      const workspaceContext = config.getWorkspaceContext();
+      for (const pathToAdd of pendingDirs) {
+        try {
+          workspaceContext.addDirectory(expandHomeDir(pathToAdd.trim()));
+          added.push(pathToAdd.trim());
+        } catch (e) {
+          const error = e as Error;
+          errors.push(`Error adding '${pathToAdd.trim()}': ${error.message}`);
+        }
+      }
+
+      if (added.length > 0 || errors.length > 0) {
+        finishAddingDirectories(config, addItem, added, errors);
+      }
+      config.clearPendingIncludeDirectories();
+      return;
+    }
+
+    const trustedFolders = loadTrustedFolders();
+    const untrustedDirs: string[] = [];
+    const undefinedTrustDirs: string[] = [];
+    const trustedDirs: string[] = [];
+    const added: string[] = [];
+    const errors: string[] = [];
+
+    for (const pathToAdd of pendingDirs) {
+      const expandedPath = expandHomeDir(pathToAdd.trim());
+      const isTrusted = trustedFolders.isPathTrusted(expandedPath);
+      if (isTrusted === false) {
+        untrustedDirs.push(pathToAdd.trim());
+      } else if (isTrusted === undefined) {
+        undefinedTrustDirs.push(pathToAdd.trim());
+      } else {
+        trustedDirs.push(pathToAdd.trim());
+      }
+    }
+
+    if (untrustedDirs.length > 0) {
+      errors.push(
+        `The following directories are explicitly untrusted and cannot be added to a trusted workspace:\n- ${untrustedDirs.join(
+          '\n- ',
+        )}\nPlease use the permissions command to modify their trust level.`,
+      );
+    }
+
+    const workspaceContext = config.getWorkspaceContext();
+    for (const pathToAdd of trustedDirs) {
+      try {
+        workspaceContext.addDirectory(expandHomeDir(pathToAdd));
+        added.push(pathToAdd);
+      } catch (e) {
+        const error = e as Error;
+        errors.push(`Error adding '${pathToAdd}': ${error.message}`);
+      }
+    }
+
+    if (undefinedTrustDirs.length > 0) {
+      console.log(
+        'Creating custom dialog with undecidedDirs:',
+        undefinedTrustDirs,
+      );
+      setCustomDialog(
+        <MultiFolderTrustDialog
+          folders={undefinedTrustDirs}
+          onComplete={() => {
+            setCustomDialog(null);
+            config.clearPendingIncludeDirectories();
+          }}
+          trustedDirs={added}
+          errors={errors}
+          finishAddingDirectories={finishAddingDirectories}
+          config={config}
+          addItem={addItem}
+        />,
+      );
+    } else if (added.length > 0 || errors.length > 0) {
+      finishAddingDirectories(config, addItem, added, errors);
+      config.clearPendingIncludeDirectories();
+    }
+  }, [isTrustedFolder, config, addItem, setCustomDialog]);
+}