2 files changed, 135 insertions, 51 deletions
diff --git a/packages/cli/src/ui/hooks/useFolderTrust.test.ts b/packages/cli/src/ui/hooks/useFolderTrust.test.ts
index 61552af0..e565ab05 100644
--- a/packages/cli/src/ui/hooks/useFolderTrust.test.ts
+++ b/packages/cli/src/ui/hooks/useFolderTrust.test.ts
@@ -4,15 +4,33 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { renderHook, act } from '@testing-library/react';
 import { vi } from 'vitest';
+import { renderHook, act } from '@testing-library/react';
 import { useFolderTrust } from './useFolderTrust.js';
-import { LoadedSettings, SettingScope } from '../../config/settings.js';
+import { type Config } from '@google/gemini-cli-core';
+import { LoadedSettings } from '../../config/settings.js';
 import { FolderTrustChoice } from '../components/FolderTrustDialog.js';
+import {
+  LoadedTrustedFolders,
+  TrustLevel,
+} from '../../config/trustedFolders.js';
+import * as process from 'process';
+
+import * as trustedFolders from '../../config/trustedFolders.js';
+
+vi.mock('process', () => ({
+  cwd: vi.fn(),
+  platform: 'linux',
+}));
 
 describe('useFolderTrust', () => {
-  it('should set isFolderTrustDialogOpen to true when folderTrustFeature is true and folderTrust is undefined', () => {
-    const settings = {
+  let mockSettings: LoadedSettings;
+  let mockConfig: Config;
+  let mockTrustedFolders: LoadedTrustedFolders;
+  let loadTrustedFoldersSpy: vi.SpyInstance;
+
+  beforeEach(() => {
+    mockSettings = {
       merged: {
         folderTrustFeature: true,
         folderTrust: undefined,
@@ -20,59 +38,110 @@ describe('useFolderTrust', () => {
       setValue: vi.fn(),
     } as unknown as LoadedSettings;
 
-    const { result } = renderHook(() => useFolderTrust(settings));
+    mockConfig = {
+      isTrustedFolder: vi.fn().mockReturnValue(undefined),
+    } as unknown as Config;
 
-    expect(result.current.isFolderTrustDialogOpen).toBe(true);
+    mockTrustedFolders = {
+      setValue: vi.fn(),
+    } as unknown as LoadedTrustedFolders;
+
+    loadTrustedFoldersSpy = vi
+      .spyOn(trustedFolders, 'loadTrustedFolders')
+      .mockReturnValue(mockTrustedFolders);
+    (process.cwd as vi.Mock).mockReturnValue('/test/path');
   });
 
-  it('should set isFolderTrustDialogOpen to false when folderTrustFeature is false', () => {
-    const settings = {
-      merged: {
-        folderTrustFeature: false,
-        folderTrust: undefined,
-      },
-      setValue: vi.fn(),
-    } as unknown as LoadedSettings;
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
 
-    const { result } = renderHook(() => useFolderTrust(settings));
+  it('should not open dialog when folder is already trusted', () => {
+    (mockConfig.isTrustedFolder as vi.Mock).mockReturnValue(true);
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
+    expect(result.current.isFolderTrustDialogOpen).toBe(false);
+  });
 
+  it('should not open dialog when folder is already untrusted', () => {
+    (mockConfig.isTrustedFolder as vi.Mock).mockReturnValue(false);
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
     expect(result.current.isFolderTrustDialogOpen).toBe(false);
   });
 
-  it('should set isFolderTrustDialogOpen to false when folderTrust is defined', () => {
-    const settings = {
-      merged: {
-        folderTrustFeature: true,
-        folderTrust: true,
-      },
-      setValue: vi.fn(),
-    } as unknown as LoadedSettings;
+  it('should open dialog when folder trust is undefined', () => {
+    (mockConfig.isTrustedFolder as vi.Mock).mockReturnValue(undefined);
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
+    expect(result.current.isFolderTrustDialogOpen).toBe(true);
+  });
+
+  it('should handle TRUST_FOLDER choice', () => {
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
 
-    const { result } = renderHook(() => useFolderTrust(settings));
+    act(() => {
+      result.current.handleFolderTrustSelect(FolderTrustChoice.TRUST_FOLDER);
+    });
 
+    expect(loadTrustedFoldersSpy).toHaveBeenCalled();
+    expect(mockTrustedFolders.setValue).toHaveBeenCalledWith(
+      '/test/path',
+      TrustLevel.TRUST_FOLDER,
+    );
     expect(result.current.isFolderTrustDialogOpen).toBe(false);
   });
 
-  it('should call setValue and set isFolderTrustDialogOpen to false on handleFolderTrustSelect', () => {
-    const settings = {
-      merged: {
-        folderTrustFeature: true,
-        folderTrust: undefined,
-      },
-      setValue: vi.fn(),
-    } as unknown as LoadedSettings;
+  it('should handle TRUST_PARENT choice', () => {
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
+
+    act(() => {
+      result.current.handleFolderTrustSelect(FolderTrustChoice.TRUST_PARENT);
+    });
 
-    const { result } = renderHook(() => useFolderTrust(settings));
+    expect(mockTrustedFolders.setValue).toHaveBeenCalledWith(
+      '/test/path',
+      TrustLevel.TRUST_PARENT,
+    );
+    expect(result.current.isFolderTrustDialogOpen).toBe(false);
+  });
+
+  it('should handle DO_NOT_TRUST choice', () => {
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
 
     act(() => {
-      result.current.handleFolderTrustSelect(FolderTrustChoice.TRUST_FOLDER);
+      result.current.handleFolderTrustSelect(FolderTrustChoice.DO_NOT_TRUST);
     });
 
-    expect(settings.setValue).toHaveBeenCalledWith(
-      SettingScope.User,
-      'folderTrust',
-      true,
+    expect(mockTrustedFolders.setValue).toHaveBeenCalledWith(
+      '/test/path',
+      TrustLevel.DO_NOT_TRUST,
     );
     expect(result.current.isFolderTrustDialogOpen).toBe(false);
   });
+
+  it('should do nothing for default choice', () => {
+    const { result } = renderHook(() =>
+      useFolderTrust(mockSettings, mockConfig),
+    );
+
+    act(() => {
+      result.current.handleFolderTrustSelect(
+        'invalid_choice' as FolderTrustChoice,
+      );
+    });
+
+    expect(mockTrustedFolders.setValue).not.toHaveBeenCalled();
+    expect(mockSettings.setValue).not.toHaveBeenCalled();
+    expect(result.current.isFolderTrustDialogOpen).toBe(true);
+  });
 });
diff --git a/packages/cli/src/ui/hooks/useFolderTrust.ts b/packages/cli/src/ui/hooks/useFolderTrust.ts
index 90a69132..6458d4aa 100644
--- a/packages/cli/src/ui/hooks/useFolderTrust.ts
+++ b/packages/cli/src/ui/hooks/useFolderTrust.ts
@@ -5,24 +5,39 @@
  */
 
 import { useState, useCallback } from 'react';
-import { LoadedSettings, SettingScope } from '../../config/settings.js';
+import { type Config } from '@google/gemini-cli-core';
+import { LoadedSettings } from '../../config/settings.js';
 import { FolderTrustChoice } from '../components/FolderTrustDialog.js';
+import { loadTrustedFolders, TrustLevel } from '../../config/trustedFolders.js';
+import * as process from 'process';
 
-export const useFolderTrust = (settings: LoadedSettings) => {
+export const useFolderTrust = (settings: LoadedSettings, config: Config) => {
   const [isFolderTrustDialogOpen, setIsFolderTrustDialogOpen] = useState(
-    !!settings.merged.folderTrustFeature &&
-      // TODO: Update to avoid showing dialog for folders that are trusted.
-      settings.merged.folderTrust === undefined,
+    config.isTrustedFolder() === undefined,
   );
 
-  const handleFolderTrustSelect = useCallback(
-    (_choice: FolderTrustChoice) => {
-      // TODO: Store folderPath in the trusted folders config file based on the choice.
-      settings.setValue(SettingScope.User, 'folderTrust', true);
-      setIsFolderTrustDialogOpen(false);
-    },
-    [settings],
-  );
+  const handleFolderTrustSelect = useCallback((choice: FolderTrustChoice) => {
+    const trustedFolders = loadTrustedFolders();
+    const cwd = process.cwd();
+    let trustLevel: TrustLevel;
+
+    switch (choice) {
+      case FolderTrustChoice.TRUST_FOLDER:
+        trustLevel = TrustLevel.TRUST_FOLDER;
+        break;
+      case FolderTrustChoice.TRUST_PARENT:
+        trustLevel = TrustLevel.TRUST_PARENT;
+        break;
+      case FolderTrustChoice.DO_NOT_TRUST:
+        trustLevel = TrustLevel.DO_NOT_TRUST;
+        break;
+      default:
+        return;
+    }
+
+    trustedFolders.setValue(cwd, trustLevel);
+    setIsFolderTrustDialogOpen(false);
+  }, []);
 
   return {
     isFolderTrustDialogOpen,