feat(file-search): Add support for non-recursive file search (#5648)

Co-authored-by: Jacob Richman <jacob314@gmail.com>

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

@@ -50,6 +50,7 @@ describe('useAtCompletion', () => {
         respectGitIgnore: true,
         respectGeminiIgnore: true,
       })),
+      getEnableRecursiveFileSearch: () => true,
     } as unknown as Config;
     vi.clearAllMocks();
   });
@@ -431,5 +432,42 @@ describe('useAtCompletion', () => {
       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,
+        })),
+      } 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',
+      ]);
+    });
   });
 });

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

@@ -165,6 +165,9 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
             config?.getFileFilteringOptions()?.respectGeminiIgnore ?? true,
           cache: true,
           cacheTtl: 30, // 30 seconds
+          maxDepth: !(config?.getEnableRecursiveFileSearch() ?? true)
+            ? 0
+            : undefined,
         });
         await searcher.initialize();
         fileSearch.current = searcher;

packages/core/src/utils/filesearch/crawlCache.test.ts

@@ -26,6 +26,17 @@ describe('CrawlCache', () => {
       const key2 = getCacheKey('/foo', 'baz');
       expect(key1).not.toBe(key2);
     });
+
+    it('should generate a different hash for different maxDepth values', () => {
+      const key1 = getCacheKey('/foo', 'bar', 1);
+      const key2 = getCacheKey('/foo', 'bar', 2);
+      const key3 = getCacheKey('/foo', 'bar', undefined);
+      const key4 = getCacheKey('/foo', 'bar');
+      expect(key1).not.toBe(key2);
+      expect(key1).not.toBe(key3);
+      expect(key2).not.toBe(key3);
+      expect(key3).toBe(key4);
+    });
   });
 
   describe('in-memory cache operations', () => {

packages/core/src/utils/filesearch/crawlCache.ts

@@ -17,10 +17,14 @@ const cacheTimers = new Map<string, NodeJS.Timeout>();
 export const getCacheKey = (
   directory: string,
   ignoreContent: string,
+  maxDepth?: number,
 ): string => {
   const hash = crypto.createHash('sha256');
   hash.update(directory);
   hash.update(ignoreContent);
+  if (maxDepth !== undefined) {
+    hash.update(String(maxDepth));
+  }
   return hash.digest('hex');
 };
 

packages/core/src/utils/filesearch/fileSearch.test.ts

@@ -446,6 +446,46 @@ describe('FileSearch', () => {
 
       expect(crawlSpy).toHaveBeenCalledTimes(1);
     });
+
+    it('should miss the cache when maxDepth changes', async () => {
+      tmpDir = await createTmpDir({ 'file1.js': '' });
+      const getOptions = (maxDepth?: number) => ({
+        projectRoot: tmpDir,
+        useGitignore: false,
+        useGeminiignore: false,
+        ignoreDirs: [],
+        cache: true,
+        cacheTtl: 10000,
+        maxDepth,
+      });
+
+      // 1. First search with maxDepth: 1, should trigger a crawl.
+      const fs1 = new FileSearch(getOptions(1));
+      const crawlSpy1 = vi.spyOn(
+        fs1 as FileSearchWithPrivateMethods,
+        'performCrawl',
+      );
+      await fs1.initialize();
+      expect(crawlSpy1).toHaveBeenCalledTimes(1);
+
+      // 2. Second search with maxDepth: 2, should be a cache miss and trigger a crawl.
+      const fs2 = new FileSearch(getOptions(2));
+      const crawlSpy2 = vi.spyOn(
+        fs2 as FileSearchWithPrivateMethods,
+        'performCrawl',
+      );
+      await fs2.initialize();
+      expect(crawlSpy2).toHaveBeenCalledTimes(1);
+
+      // 3. Third search with maxDepth: 1 again, should be a cache hit.
+      const fs3 = new FileSearch(getOptions(1));
+      const crawlSpy3 = vi.spyOn(
+        fs3 as FileSearchWithPrivateMethods,
+        'performCrawl',
+      );
+      await fs3.initialize();
+      expect(crawlSpy3).not.toHaveBeenCalled();
+    });
   });
 
   it('should handle empty or commented-only ignore files', async () => {
@@ -639,4 +679,109 @@ describe('FileSearch', () => {
     // 3. Assert that the maxResults limit was respected, even with a cache hit.
     expect(limitedResults).toEqual(['file1.js', 'file2.js']);
   });
+
+  describe('with maxDepth', () => {
+    beforeEach(async () => {
+      tmpDir = await createTmpDir({
+        'file-root.txt': '',
+        level1: {
+          'file-level1.txt': '',
+          level2: {
+            'file-level2.txt': '',
+            level3: {
+              'file-level3.txt': '',
+            },
+          },
+        },
+      });
+    });
+
+    it('should only search top-level files when maxDepth is 0', async () => {
+      const fileSearch = new FileSearch({
+        projectRoot: tmpDir,
+        useGitignore: false,
+        useGeminiignore: false,
+        ignoreDirs: [],
+        cache: false,
+        cacheTtl: 0,
+        maxDepth: 0,
+      });
+
+      await fileSearch.initialize();
+      const results = await fileSearch.search('');
+
+      expect(results).toEqual(['level1/', 'file-root.txt']);
+    });
+
+    it('should search one level deep when maxDepth is 1', async () => {
+      const fileSearch = new FileSearch({
+        projectRoot: tmpDir,
+        useGitignore: false,
+        useGeminiignore: false,
+        ignoreDirs: [],
+        cache: false,
+        cacheTtl: 0,
+        maxDepth: 1,
+      });
+
+      await fileSearch.initialize();
+      const results = await fileSearch.search('');
+
+      expect(results).toEqual([
+        'level1/',
+        'level1/level2/',
+        'file-root.txt',
+        'level1/file-level1.txt',
+      ]);
+    });
+
+    it('should search two levels deep when maxDepth is 2', async () => {
+      const fileSearch = new FileSearch({
+        projectRoot: tmpDir,
+        useGitignore: false,
+        useGeminiignore: false,
+        ignoreDirs: [],
+        cache: false,
+        cacheTtl: 0,
+        maxDepth: 2,
+      });
+
+      await fileSearch.initialize();
+      const results = await fileSearch.search('');
+
+      expect(results).toEqual([
+        'level1/',
+        'level1/level2/',
+        'level1/level2/level3/',
+        'file-root.txt',
+        'level1/file-level1.txt',
+        'level1/level2/file-level2.txt',
+      ]);
+    });
+
+    it('should perform a full recursive search when maxDepth is undefined', async () => {
+      const fileSearch = new FileSearch({
+        projectRoot: tmpDir,
+        useGitignore: false,
+        useGeminiignore: false,
+        ignoreDirs: [],
+        cache: false,
+        cacheTtl: 0,
+        maxDepth: undefined, // Explicitly undefined
+      });
+
+      await fileSearch.initialize();
+      const results = await fileSearch.search('');
+
+      expect(results).toEqual([
+        'level1/',
+        'level1/level2/',
+        'level1/level2/level3/',
+        'file-root.txt',
+        'level1/file-level1.txt',
+        'level1/level2/file-level2.txt',
+        'level1/level2/level3/file-level3.txt',
+      ]);
+    });
+  });
 });

packages/core/src/utils/filesearch/fileSearch.ts

@@ -19,6 +19,7 @@ export type FileSearchOptions = {
   useGeminiignore: boolean;
   cache: boolean;
   cacheTtl: number;
+  maxDepth?: number;
 };
 
 export class AbortError extends Error {
@@ -215,6 +216,7 @@ export class FileSearch {
       const cacheKey = cache.getCacheKey(
         this.absoluteDir,
         this.ignore.getFingerprint(),
+        this.options.maxDepth,
       );
       const cachedResults = cache.read(cacheKey);
 
@@ -230,6 +232,7 @@ export class FileSearch {
       const cacheKey = cache.getCacheKey(
         this.absoluteDir,
         this.ignore.getFingerprint(),
+        this.options.maxDepth,
       );
       cache.write(cacheKey, this.allFiles, this.options.cacheTtl * 1000);
     }
@@ -257,6 +260,10 @@ export class FileSearch {
         return dirFilter(`${relativePath}/`);
       });
 
+    if (this.options.maxDepth !== undefined) {
+      api.withMaxDepth(this.options.maxDepth);
+    }
+
     return api.crawl(this.absoluteDir).withPromise();
   }