gemini-cli/packages/cli/src/ui/hooks/useAtCompletion.test.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
import { act, useState } from 'react';
import * as path from 'node:path';
import { renderHook } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { useAtCompletion } from './useAtCompletion.js';
import type { Config, FileSearch } from '@google/gemini-cli-core';
import {
  FileSearchFactory,
  FileDiscoveryService,
  escapePath,
} from '@google/gemini-cli-core';
import type { FileSystemStructure } from '@google/gemini-cli-test-utils';
import { createTmpDir, cleanupTmpDir } from '@google/gemini-cli-test-utils';
import type { 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,
      })),
      getEnableRecursiveFileSearch: () => true,
      getFileFilteringEnableFuzzySearch: () => true,
      getResourceRegistry: vi.fn().mockReturnValue({
        getAllResources: () => [],
      }),
    } 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(5);
      });

      expect(result.current.suggestions.length).toBeGreaterThan(0);
      expect(result.current.suggestions.map((s) => s.value)).toEqual([
        'src/',
        'src/components/',
        'file.txt',
        `${escapePath('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/index.js',
        'src/components/',
        'src/components/Button.tsx',
      ]);
    });

    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',
      ]);
    });

    it('should perform a case-insensitive search by lowercasing the pattern', async () => {
      testRootDir = await createTmpDir({ 'cRaZycAsE.txt': '' });

      const fileSearch = FileSearchFactory.create({
        projectRoot: testRootDir,
        ignoreDirs: [],
        fileDiscoveryService: new FileDiscoveryService(testRootDir, {
          respectGitIgnore: false,
          respectGeminiIgnore: false,
        }),
        cache: false,
        cacheTtl: 0,
        enableRecursiveFileSearch: true,
        enableFuzzySearch: true,
      });
      await fileSearch.initialize();

      vi.spyOn(FileSearchFactory, 'create').mockReturnValue(fileSearch);

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(
          true,
          'CrAzYCaSe',
          mockConfig,
          testRootDir,
        ),
      );

      // The hook should find 'cRaZycAsE.txt' even though the pattern is 'CrAzYCaSe'.
      await waitFor(() => {
        expect(result.current.suggestions.map((s) => s.value)).toEqual([
          'cRaZycAsE.txt',
        ]);
      });
    });
  });

  describe('MCP resource suggestions', () => {
    it('should include MCP resources in the suggestion list using fuzzy matching', async () => {
      mockConfig.getResourceRegistry = vi.fn().mockReturnValue({
        getAllResources: () => [
          {
            serverName: 'server-1',
            uri: 'file:///tmp/server-1/logs.txt',
            name: 'logs',
            discoveredAt: Date.now(),
          },
        ],
      });

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(true, 'logs', mockConfig, '/tmp'),
      );

      await waitFor(() => {
        expect(
          result.current.suggestions.some(
            (suggestion) =>
              suggestion.value === 'server-1:file:///tmp/server-1/logs.txt',
          ),
        ).toBe(true);
      });
    });
  });

  describe('UI State and Loading Behavior', () => {
    it('should be in a loading state during initial file system crawl', async () => {
      testRootDir = await createTmpDir({});

      // Mock FileSearch to be slow to catch the loading state
      const mockFileSearch = {
        initialize: vi.fn().mockImplementation(async () => {
          await new Promise((resolve) => setTimeout(resolve, 50));
        }),
        search: vi.fn().mockResolvedValue([]),
      };
      vi.spyOn(FileSearchFactory, 'create').mockReturnValue(
        mockFileSearch as unknown as FileSearch,
      );

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(true, '', mockConfig, testRootDir),
      );

      // It's initially true because the effect runs synchronously.
      await waitFor(() => {
        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 200ms', 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 200ms', async () => {
      const structure: FileSystemStructure = { 'a.txt': '', 'b.txt': '' };
      testRootDir = await createTmpDir(structure);

      const realFileSearch = FileSearchFactory.create({
        projectRoot: testRootDir,
        ignoreDirs: [],
        fileDiscoveryService: new FileDiscoveryService(testRootDir, {
          respectGitIgnore: true,
          respectGeminiIgnore: true,
        }),
        cache: false,
        cacheTtl: 0,
        enableRecursiveFileSearch: true,
        enableFuzzySearch: true,
      });
      await realFileSearch.initialize();

      // Mock that returns results immediately but we'll control timing with fake timers
      const mockFileSearch: FileSearch = {
        initialize: vi.fn().mockResolvedValue(undefined),
        search: vi
          .fn()
          .mockImplementation(async (pattern, options) =>
            realFileSearch.search(pattern, options),
          ),
      };
      vi.spyOn(FileSearchFactory, 'create').mockReturnValue(mockFileSearch);

      const { result, rerender } = renderHook(
        ({ pattern }) =>
          useTestHarnessForAtCompletion(true, pattern, mockConfig, testRootDir),
        { initialProps: { pattern: 'a' } },
      );

      // Wait for the initial search to complete (using real timers)
      await waitFor(() => {
        expect(result.current.suggestions.map((s) => s.value)).toEqual([
          'a.txt',
        ]);
      });

      // Now switch to fake timers for precise control of the loading behavior
      vi.useFakeTimers();

      // Trigger the second search
      act(() => {
        rerender({ pattern: 'b' });
      });

      // Initially, loading should be false (before 200ms timer)
      expect(result.current.isLoadingSuggestions).toBe(false);

      // Advance time by exactly 200ms to trigger the loading state
      act(() => {
        vi.advanceTimersByTime(200);
      });

      // Now loading should be true and suggestions should be cleared
      expect(result.current.isLoadingSuggestions).toBe(true);
      expect(result.current.suggestions).toEqual([]);

      // Switch back to real timers for the final waitFor
      vi.useRealTimers();

      // Wait for the search results to be processed
      await waitFor(() => {
        expect(result.current.suggestions.map((s) => s.value)).toEqual([
          'b.txt',
        ]);
      });

      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 mockFileSearch: FileSearch = {
        initialize: vi.fn().mockResolvedValue(undefined),
        search: vi.fn().mockImplementation(async (pattern: string) => {
          const delay = pattern === 'a' ? 500 : 50;
          await new Promise((resolve) => setTimeout(resolve, delay));
          return [pattern];
        }),
      };
      vi.spyOn(FileSearchFactory, 'create').mockReturnValue(mockFileSearch);

      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(mockFileSearch.search).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(mockFileSearch.search).toHaveBeenCalledWith(
        'b',
        expect.any(Object),
      );
    });
  });

  describe('State Management', () => {
    it('should reset the state when disabled after being in a READY state', async () => {
      const structure: FileSystemStructure = { 'a.txt': '' };
      testRootDir = await createTmpDir(structure);

      const { result, rerender } = renderHook(
        ({ enabled }) =>
          useTestHarnessForAtCompletion(enabled, 'a', mockConfig, testRootDir),
        { initialProps: { enabled: true } },
      );

      // Wait for the hook to be ready and have suggestions
      await waitFor(() => {
        expect(result.current.suggestions.map((s) => s.value)).toEqual([
          'a.txt',
        ]);
      });

      // Now, disable the hook
      rerender({ enabled: false });

      // The suggestions should be cleared immediately because of the RESET action
      expect(result.current.suggestions).toEqual([]);
    });

    it('should reset the state when disabled after being in an ERROR state', async () => {
      testRootDir = await createTmpDir({});

      // Force an error during initialization
      const mockFileSearch: FileSearch = {
        initialize: vi
          .fn()
          .mockRejectedValue(new Error('Initialization failed')),
        search: vi.fn(),
      };
      vi.spyOn(FileSearchFactory, 'create').mockReturnValue(mockFileSearch);

      const { result, rerender } = renderHook(
        ({ enabled }) =>
          useTestHarnessForAtCompletion(enabled, '', mockConfig, testRootDir),
        { initialProps: { enabled: true } },
      );

      // Wait for the hook to enter the error state
      await waitFor(() => {
        expect(result.current.isLoadingSuggestions).toBe(false);
      });
      expect(result.current.suggestions).toEqual([]); // No suggestions on error

      // Now, disable the hook
      rerender({ enabled: false });

      // The state should still be reset (though visually it's the same)
      // We can't directly inspect the internal state, but we can ensure it doesn't crash
      // and the suggestions remain empty.
      expect(result.current.suggestions).toEqual([]);
    });
  });

  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);
    });

    it('should perform a non-recursive search when enableRecursiveFileSearch is false', async () => {
      const structure: FileSystemStructure = {
        'file.txt': '',
        src: {
          'index.js': '',
        },
      };
      testRootDir = await createTmpDir(structure);

      const nonRecursiveConfig = {
        getEnableRecursiveFileSearch: () => false,
        getFileFilteringOptions: vi.fn(() => ({
          respectGitIgnore: true,
          respectGeminiIgnore: true,
        })),
        getFileFilteringEnableFuzzySearch: () => true,
      } as unknown as Config;

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(
          true,
          '',
          nonRecursiveConfig,
          testRootDir,
        ),
      );

      await waitFor(() => {
        expect(result.current.suggestions.length).toBeGreaterThan(0);
      });

      // Should only contain top-level items
      expect(result.current.suggestions.map((s) => s.value)).toEqual([
        'src/',
        'file.txt',
      ]);
    });
  });

  describe('Multi-directory workspace support', () => {
    const multiDirTmpDirs: string[] = [];

    afterEach(async () => {
      await Promise.all(multiDirTmpDirs.map((dir) => cleanupTmpDir(dir)));
      multiDirTmpDirs.length = 0;
    });

    it('should include files from workspace directories beyond cwd', async () => {
      const cwdStructure: FileSystemStructure = { 'main.txt': '' };
      const addedDirStructure: FileSystemStructure = { 'added-file.txt': '' };
      const cwdDir = await createTmpDir(cwdStructure);
      multiDirTmpDirs.push(cwdDir);
      const addedDir = await createTmpDir(addedDirStructure);
      multiDirTmpDirs.push(addedDir);

      const multiDirConfig = {
        ...mockConfig,
        getWorkspaceContext: vi.fn().mockReturnValue({
          getDirectories: () => [cwdDir, addedDir],
          onDirectoriesChanged: vi.fn(() => () => {}),
        }),
      } as unknown as Config;

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(true, '', multiDirConfig, cwdDir),
      );

      await waitFor(() => {
        const values = result.current.suggestions.map((s) => s.value);
        expect(values).toContain('main.txt');
        expect(values).toContain(
          escapePath(path.join(addedDir, 'added-file.txt')),
        );
      });
    });

    it('should pick up newly added directories via onDirectoriesChanged', async () => {
      const cwdStructure: FileSystemStructure = { 'original.txt': '' };
      const addedStructure: FileSystemStructure = { 'new-file.txt': '' };
      const cwdDir = await createTmpDir(cwdStructure);
      multiDirTmpDirs.push(cwdDir);
      const addedDir = await createTmpDir(addedStructure);
      multiDirTmpDirs.push(addedDir);

      let dirChangeListener: (() => void) | null = null;
      const directories = [cwdDir];

      const dynamicConfig = {
        ...mockConfig,
        getWorkspaceContext: vi.fn().mockReturnValue({
          getDirectories: () => [...directories],
          onDirectoriesChanged: vi.fn((listener: () => void) => {
            dirChangeListener = listener;
            return () => {
              dirChangeListener = null;
            };
          }),
        }),
      } as unknown as Config;

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(true, '', dynamicConfig, cwdDir),
      );

      await waitFor(() => {
        const values = result.current.suggestions.map((s) => s.value);
        expect(values).toContain('original.txt');
        expect(values.every((v) => !v.includes('new-file.txt'))).toBe(true);
      });

      directories.push(addedDir);
      act(() => {
        dirChangeListener?.();
      });

      await waitFor(() => {
        const values = result.current.suggestions.map((s) => s.value);
        expect(values).toContain(
          escapePath(path.join(addedDir, 'new-file.txt')),
        );
      });
    });

    it('should show same-named files from different directories without false deduplication', async () => {
      const dir1Structure: FileSystemStructure = { 'readme.md': '' };
      const dir2Structure: FileSystemStructure = { 'readme.md': '' };
      const dir1 = await createTmpDir(dir1Structure);
      multiDirTmpDirs.push(dir1);
      const dir2 = await createTmpDir(dir2Structure);
      multiDirTmpDirs.push(dir2);

      const multiDirConfig = {
        ...mockConfig,
        getWorkspaceContext: vi.fn().mockReturnValue({
          getDirectories: () => [dir1, dir2],
          onDirectoriesChanged: vi.fn(() => () => {}),
        }),
      } as unknown as Config;

      const { result } = renderHook(() =>
        useTestHarnessForAtCompletion(true, 'readme', multiDirConfig, dir1),
      );

      await waitFor(() => {
        const values = result.current.suggestions.map((s) => s.value);
        const readmeEntries = values.filter((v) => v.includes('readme.md'));
        expect(readmeEntries.length).toBe(2);
        expect(readmeEntries).toContain('readme.md');
        expect(readmeEntries).toContain(
          escapePath(path.join(dir2, 'readme.md')),
        );
      });
    });
  });
});