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

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

4 files changed, 167 insertions, 0 deletions

diff --git a/packages/core/src/utils/filesearch/crawlCache.test.ts b/packages/core/src/utils/filesearch/crawlCache.test.ts
index 2feab61a..c8ca0df2 100644
--- a/packages/core/src/utils/filesearch/crawlCache.test.ts
+++ b/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', () => {
diff --git a/packages/core/src/utils/filesearch/crawlCache.ts b/packages/core/src/utils/filesearch/crawlCache.ts
index 3cc948c6..b905c9df 100644
--- a/packages/core/src/utils/filesearch/crawlCache.ts
+++ b/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');
 };
 
diff --git a/packages/core/src/utils/filesearch/fileSearch.test.ts b/packages/core/src/utils/filesearch/fileSearch.test.ts
index b804d623..a7f59f91 100644
--- a/packages/core/src/utils/filesearch/fileSearch.test.ts
+++ b/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',
+      ]);
+    });
+  });
 });
diff --git a/packages/core/src/utils/filesearch/fileSearch.ts b/packages/core/src/utils/filesearch/fileSearch.ts
index 5915821a..db14bc65 100644
--- a/packages/core/src/utils/filesearch/fileSearch.ts
+++ b/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();
   }