feat(core, cli): Introduce high-performance FileSearch engine (#5136)

Co-authored-by: Jacob Richman <[email protected]>

1 files changed, 380 insertions, 0 deletions

diff --git a/packages/cli/src/ui/hooks/useAtCompletion.test.ts b/packages/cli/src/ui/hooks/useAtCompletion.test.ts
new file mode 100644
index 00000000..bf2453f5
--- /dev/null
+++ b/packages/cli/src/ui/hooks/useAtCompletion.test.ts
@@ -0,0 +1,380 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/** @vitest-environment jsdom */
+
+import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
+import { renderHook, waitFor, act } from '@testing-library/react';
+import { useAtCompletion } from './useAtCompletion.js';
+import { Config, FileSearch } from '@google/gemini-cli-core';
+import {
+  createTmpDir,
+  cleanupTmpDir,
+  FileSystemStructure,
+} from '@google/gemini-cli-test-utils';
+import { useState } from 'react';
+import { Suggestion } from '../components/SuggestionsDisplay.js';
+
+// Test harness to capture the state from the hook's callbacks.
+function useTestHarnessForAtCompletion(
+  enabled: boolean,
+  pattern: string,
+  config: Config | undefined,
+  cwd: string,
+) {
+  const [suggestions, setSuggestions] = useState<Suggestion[]>([]);
+  const [isLoadingSuggestions, setIsLoadingSuggestions] = useState(false);
+
+  useAtCompletion({
+    enabled,
+    pattern,
+    config,
+    cwd,
+    setSuggestions,
+    setIsLoadingSuggestions,
+  });
+
+  return { suggestions, isLoadingSuggestions };
+}
+
+describe('useAtCompletion', () => {
+  let testRootDir: string;
+  let mockConfig: Config;
+
+  beforeEach(() => {
+    mockConfig = {
+      getFileFilteringOptions: vi.fn(() => ({
+        respectGitIgnore: true,
+        respectGeminiIgnore: true,
+      })),
+    } as unknown as Config;
+    vi.clearAllMocks();
+  });
+
+  afterEach(async () => {
+    if (testRootDir) {
+      await cleanupTmpDir(testRootDir);
+    }
+    vi.restoreAllMocks();
+  });
+
+  describe('File Search Logic', () => {
+    it('should perform a recursive search for an empty pattern', async () => {
+      const structure: FileSystemStructure = {
+        'file.txt': '',
+        src: {
+          'index.js': '',
+          components: ['Button.tsx', 'Button with spaces.tsx'],
+        },
+      };
+      testRootDir = await createTmpDir(structure);
+
+      const { result } = renderHook(() =>
+        useTestHarnessForAtCompletion(true, '', mockConfig, testRootDir),
+      );
+
+      await waitFor(() => {
+        expect(result.current.suggestions.length).toBeGreaterThan(0);
+      });
+
+      expect(result.current.suggestions.map((s) => s.value)).toEqual([
+        'src/',
+        'src/components/',
+        'file.txt',
+        'src/components/Button\\ with\\ spaces.tsx',
+        'src/components/Button.tsx',
+        'src/index.js',
+      ]);
+    });
+
+    it('should correctly filter the recursive list based on a pattern', async () => {
+      const structure: FileSystemStructure = {
+        'file.txt': '',
+        src: {
+          'index.js': '',
+          components: {
+            'Button.tsx': '',
+          },
+        },
+      };
+      testRootDir = await createTmpDir(structure);
+
+      const { result } = renderHook(() =>
+        useTestHarnessForAtCompletion(true, 'src/', mockConfig, testRootDir),
+      );
+
+      await waitFor(() => {
+        expect(result.current.suggestions.length).toBeGreaterThan(0);
+      });
+
+      expect(result.current.suggestions.map((s) => s.value)).toEqual([
+        'src/',
+        'src/components/',
+        'src/components/Button.tsx',
+        'src/index.js',
+      ]);
+    });
+
+    it('should append a trailing slash to directory paths in suggestions', async () => {
+      const structure: FileSystemStructure = {
+        'file.txt': '',
+        dir: {},
+      };
+      testRootDir = await createTmpDir(structure);
+
+      const { result } = renderHook(() =>
+        useTestHarnessForAtCompletion(true, '', mockConfig, testRootDir),
+      );
+
+      await waitFor(() => {
+        expect(result.current.suggestions.length).toBeGreaterThan(0);
+      });
+
+      expect(result.current.suggestions.map((s) => s.value)).toEqual([
+        'dir/',
+        'file.txt',
+      ]);
+    });
+  });
+
+  describe('UI State and Loading Behavior', () => {
+    it('should be in a loading state during initial file system crawl', async () => {
+      testRootDir = await createTmpDir({});
+      const { result } = renderHook(() =>
+        useTestHarnessForAtCompletion(true, '', mockConfig, testRootDir),
+      );
+
+      // It's initially true because the effect runs synchronously.
+      expect(result.current.isLoadingSuggestions).toBe(true);
+
+      // Wait for the loading to complete.
+      await waitFor(() => {
+        expect(result.current.isLoadingSuggestions).toBe(false);
+      });
+    });
+
+    it('should NOT show a loading indicator for subsequent searches that complete under 100ms', async () => {
+      const structure: FileSystemStructure = { 'a.txt': '', 'b.txt': '' };
+      testRootDir = await createTmpDir(structure);
+
+      const { result, rerender } = renderHook(
+        ({ pattern }) =>
+          useTestHarnessForAtCompletion(true, pattern, mockConfig, testRootDir),
+        { initialProps: { pattern: 'a' } },
+      );
+
+      await waitFor(() => {
+        expect(result.current.suggestions.map((s) => s.value)).toEqual([
+          'a.txt',
+        ]);
+      });
+      expect(result.current.isLoadingSuggestions).toBe(false);
+
+      rerender({ pattern: 'b' });
+
+      // Wait for the final result
+      await waitFor(() => {
+        expect(result.current.suggestions.map((s) => s.value)).toEqual([
+          'b.txt',
+        ]);
+      });
+
+      expect(result.current.isLoadingSuggestions).toBe(false);
+    });
+
+    it('should show a loading indicator and clear old suggestions for subsequent searches that take longer than 100ms', async () => {
+      const structure: FileSystemStructure = { 'a.txt': '', 'b.txt': '' };
+      testRootDir = await createTmpDir(structure);
+
+      // Spy on the search method to introduce an artificial delay
+      const originalSearch = FileSearch.prototype.search;
+      vi.spyOn(FileSearch.prototype, 'search').mockImplementation(
+        async function (...args) {
+          await new Promise((resolve) => setTimeout(resolve, 200));
+          return originalSearch.apply(this, args);
+        },
+      );
+
+      const { result, rerender } = renderHook(
+        ({ pattern }) =>
+          useTestHarnessForAtCompletion(true, pattern, mockConfig, testRootDir),
+        { initialProps: { pattern: 'a' } },
+      );
+
+      // Wait for the initial (slow) search to complete
+      await waitFor(() => {
+        expect(result.current.suggestions.map((s) => s.value)).toEqual([
+          'a.txt',
+        ]);
+      });
+
+      // Now, rerender to trigger the second search
+      rerender({ pattern: 'b' });
+
+      // Wait for the loading indicator to appear
+      await waitFor(() => {
+        expect(result.current.isLoadingSuggestions).toBe(true);
+      });
+
+      // Suggestions should be cleared while loading
+      expect(result.current.suggestions).toEqual([]);
+
+      // Wait for the final (slow) search to complete
+      await waitFor(
+        () => {
+          expect(result.current.suggestions.map((s) => s.value)).toEqual([
+            'b.txt',
+          ]);
+        },
+        { timeout: 1000 },
+      ); // Increase timeout for the slow search
+
+      expect(result.current.isLoadingSuggestions).toBe(false);
+    });
+
+    it('should abort the previous search when a new one starts', async () => {
+      const structure: FileSystemStructure = { 'a.txt': '', 'b.txt': '' };
+      testRootDir = await createTmpDir(structure);
+
+      const abortSpy = vi.spyOn(AbortController.prototype, 'abort');
+      const searchSpy = vi
+        .spyOn(FileSearch.prototype, 'search')
+        .mockImplementation(async (...args) => {
+          const delay = args[0] === 'a' ? 500 : 50;
+          await new Promise((resolve) => setTimeout(resolve, delay));
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          return [args[0] as any];
+        });
+
+      const { result, rerender } = renderHook(
+        ({ pattern }) =>
+          useTestHarnessForAtCompletion(true, pattern, mockConfig, testRootDir),
+        { initialProps: { pattern: 'a' } },
+      );
+
+      // Wait for the hook to be ready (initialization is complete)
+      await waitFor(() => {
+        expect(searchSpy).toHaveBeenCalledWith('a', expect.any(Object));
+      });
+
+      // Now that the first search is in-flight, trigger the second one.
+      act(() => {
+        rerender({ pattern: 'b' });
+      });
+
+      // The abort should have been called for the first search.
+      expect(abortSpy).toHaveBeenCalledTimes(1);
+
+      // Wait for the final result, which should be from the second, faster search.
+      await waitFor(
+        () => {
+          expect(result.current.suggestions.map((s) => s.value)).toEqual(['b']);
+        },
+        { timeout: 1000 },
+      );
+
+      // The search spy should have been called for both patterns.
+      expect(searchSpy).toHaveBeenCalledWith('b', expect.any(Object));
+
+      vi.restoreAllMocks();
+    });
+  });
+
+  describe('Filtering and Configuration', () => {
+    it('should respect .gitignore files', async () => {
+      const gitignoreContent = ['dist/', '*.log'].join('\n');
+      const structure: FileSystemStructure = {
+        '.git': {},
+        '.gitignore': gitignoreContent,
+        dist: {},
+        'test.log': '',
+        src: {},
+      };
+      testRootDir = await createTmpDir(structure);
+
+      const { result } = renderHook(() =>
+        useTestHarnessForAtCompletion(true, '', mockConfig, testRootDir),
+      );
+
+      await waitFor(() => {
+        expect(result.current.suggestions.length).toBeGreaterThan(0);
+      });
+
+      expect(result.current.suggestions.map((s) => s.value)).toEqual([
+        'src/',
+        '.gitignore',
+      ]);
+    });
+
+    it('should work correctly when config is undefined', async () => {
+      const structure: FileSystemStructure = {
+        node_modules: {},
+        src: {},
+      };
+      testRootDir = await createTmpDir(structure);
+
+      const { result } = renderHook(() =>
+        useTestHarnessForAtCompletion(true, '', undefined, testRootDir),
+      );
+
+      await waitFor(() => {
+        expect(result.current.suggestions.length).toBeGreaterThan(0);
+      });
+
+      expect(result.current.suggestions.map((s) => s.value)).toEqual([
+        'node_modules/',
+        'src/',
+      ]);
+    });
+
+    it('should reset and re-initialize when the cwd changes', async () => {
+      const structure1: FileSystemStructure = { 'file1.txt': '' };
+      const rootDir1 = await createTmpDir(structure1);
+      const structure2: FileSystemStructure = { 'file2.txt': '' };
+      const rootDir2 = await createTmpDir(structure2);
+
+      const { result, rerender } = renderHook(
+        ({ cwd, pattern }) =>
+          useTestHarnessForAtCompletion(true, pattern, mockConfig, cwd),
+        {
+          initialProps: {
+            cwd: rootDir1,
+            pattern: 'file',
+          },
+        },
+      );
+
+      // Wait for initial suggestions from the first directory
+      await waitFor(() => {
+        expect(result.current.suggestions.map((s) => s.value)).toEqual([
+          'file1.txt',
+        ]);
+      });
+
+      // Change the CWD
+      act(() => {
+        rerender({ cwd: rootDir2, pattern: 'file' });
+      });
+
+      // After CWD changes, suggestions should be cleared and it should load again.
+      await waitFor(() => {
+        expect(result.current.isLoadingSuggestions).toBe(true);
+        expect(result.current.suggestions).toEqual([]);
+      });
+
+      // Wait for the new suggestions from the second directory
+      await waitFor(() => {
+        expect(result.current.suggestions.map((s) => s.value)).toEqual([
+          'file2.txt',
+        ]);
+      });
+      expect(result.current.isLoadingSuggestions).toBe(false);
+
+      await cleanupTmpDir(rootDir1);
+      await cleanupTmpDir(rootDir2);
+    });
+  });
+});