From e77d7b2e1eb8094d264ffc233a4d22777be068da Mon Sep 17 00:00:00 2001 From: Gal Zahavi <38544478+galz10@users.noreply.github.com> Date: Thu, 15 Jan 2026 12:04:22 -0800 Subject: [PATCH] fix(cli): prevent OOM crash by limiting file search traversal and adding timeout (#16696) --- packages/cli/src/ui/hooks/useAtCompletion.ts | 22 ++++++++++++ packages/core/src/config/config.ts | 14 ++++++++ packages/core/src/config/constants.ts | 6 ++++ .../core/src/utils/filesearch/crawler.test.ts | 35 +++++++++++++++++-- packages/core/src/utils/filesearch/crawler.ts | 21 ++++++++++- .../src/utils/filesearch/fileSearch.test.ts | 28 +++++++++++++++ .../core/src/utils/filesearch/fileSearch.ts | 3 ++ 7 files changed, 126 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/ui/hooks/useAtCompletion.ts b/packages/cli/src/ui/hooks/useAtCompletion.ts index dcb6dfa478..c5f9ddb888 100644 --- a/packages/cli/src/ui/hooks/useAtCompletion.ts +++ b/packages/cli/src/ui/hooks/useAtCompletion.ts @@ -5,6 +5,7 @@ */ import { useEffect, useReducer, useRef } from 'react'; +import { setTimeout as setTimeoutPromise } from 'node:timers/promises'; import type { Config, FileSearch } from '@google/gemini-cli-core'; import { FileSearchFactory, escapePath } from '@google/gemini-cli-core'; import type { Suggestion } from '../components/SuggestionsDisplay.js'; @@ -12,6 +13,8 @@ import { MAX_SUGGESTIONS_TO_SHOW } from '../components/SuggestionsDisplay.js'; import { CommandKind } from '../commands/types.js'; import { AsyncFzf } from 'fzf'; +const DEFAULT_SEARCH_TIMEOUT_MS = 5000; + export enum AtCompletionStatus { IDLE = 'idle', INITIALIZING = 'initializing', @@ -257,6 +260,7 @@ export function useAtCompletion(props: UseAtCompletionProps): void { config?.getEnableRecursiveFileSearch() ?? true, disableFuzzySearch: config?.getFileFilteringDisableFuzzySearch() ?? false, + maxFiles: config?.getFileFilteringOptions()?.maxFileCount, }); await searcher.initialize(); fileSearch.current = searcher; @@ -285,6 +289,22 @@ export function useAtCompletion(props: UseAtCompletionProps): void { dispatch({ type: 'SET_LOADING', payload: true }); }, 200); + const timeoutMs = + config?.getFileFilteringOptions()?.searchTimeout ?? + DEFAULT_SEARCH_TIMEOUT_MS; + + // eslint-disable-next-line @typescript-eslint/no-floating-promises + (async () => { + try { + await setTimeoutPromise(timeoutMs, undefined, { + signal: controller.signal, + }); + controller.abort(); + } catch { + // ignore + } + })(); + try { const results = await fileSearch.current.search(state.pattern, { signal: controller.signal, @@ -332,6 +352,8 @@ export function useAtCompletion(props: UseAtCompletionProps): void { if (!(error instanceof Error && error.name === 'AbortError')) { dispatch({ type: 'ERROR' }); } + } finally { + controller.abort(); } }; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 08bd18216d..19c6f46fc3 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -310,6 +310,8 @@ export interface ConfigParameters { respectGeminiIgnore?: boolean; enableRecursiveFileSearch?: boolean; disableFuzzySearch?: boolean; + maxFileCount?: number; + searchTimeout?: number; }; checkpointing?: boolean; proxy?: string; @@ -441,6 +443,8 @@ export class Config { respectGeminiIgnore: boolean; enableRecursiveFileSearch: boolean; disableFuzzySearch: boolean; + maxFileCount: number; + searchTimeout: number; }; private fileDiscoveryService: FileDiscoveryService | null = null; private gitService: GitService | undefined = undefined; @@ -593,6 +597,14 @@ export class Config { enableRecursiveFileSearch: params.fileFiltering?.enableRecursiveFileSearch ?? true, disableFuzzySearch: params.fileFiltering?.disableFuzzySearch ?? false, + maxFileCount: + params.fileFiltering?.maxFileCount ?? + DEFAULT_FILE_FILTERING_OPTIONS.maxFileCount ?? + 20000, + searchTimeout: + params.fileFiltering?.searchTimeout ?? + DEFAULT_FILE_FILTERING_OPTIONS.searchTimeout ?? + 5000, }; this.checkpointing = params.checkpointing ?? false; this.proxy = params.proxy; @@ -1385,6 +1397,8 @@ export class Config { return { respectGitIgnore: this.fileFiltering.respectGitIgnore, respectGeminiIgnore: this.fileFiltering.respectGeminiIgnore, + maxFileCount: this.fileFiltering.maxFileCount, + searchTimeout: this.fileFiltering.searchTimeout, }; } diff --git a/packages/core/src/config/constants.ts b/packages/core/src/config/constants.ts index 75ef3e4979..9f3047a84f 100644 --- a/packages/core/src/config/constants.ts +++ b/packages/core/src/config/constants.ts @@ -7,16 +7,22 @@ export interface FileFilteringOptions { respectGitIgnore: boolean; respectGeminiIgnore: boolean; + maxFileCount?: number; + searchTimeout?: number; } // For memory files export const DEFAULT_MEMORY_FILE_FILTERING_OPTIONS: FileFilteringOptions = { respectGitIgnore: false, respectGeminiIgnore: true, + maxFileCount: 20000, + searchTimeout: 5000, }; // For all other files export const DEFAULT_FILE_FILTERING_OPTIONS: FileFilteringOptions = { respectGitIgnore: true, respectGeminiIgnore: true, + maxFileCount: 20000, + searchTimeout: 5000, }; diff --git a/packages/core/src/utils/filesearch/crawler.test.ts b/packages/core/src/utils/filesearch/crawler.test.ts index c2865f4401..bf1ccea209 100644 --- a/packages/core/src/utils/filesearch/crawler.test.ts +++ b/packages/core/src/utils/filesearch/crawler.test.ts @@ -503,14 +503,14 @@ describe('crawler', () => { }); }); - const getCrawlResults = (maxDepth?: number) => { + const getCrawlResults = async (maxDepth?: number) => { const ignore = loadIgnoreRules({ projectRoot: tmpDir, useGitignore: false, useGeminiignore: false, ignoreDirs: [], }); - return crawl({ + const paths = await crawl({ crawlDirectory: tmpDir, cwd: tmpDir, ignore, @@ -518,6 +518,7 @@ describe('crawler', () => { cacheTtl: 0, maxDepth, }); + return paths; }; it('should only crawl top-level files when maxDepth is 0', async () => { @@ -571,4 +572,34 @@ describe('crawler', () => { ); }); }); + + it('should detect truncation when maxFiles is hit', async () => { + tmpDir = await createTmpDir({ + 'file1.js': '', + 'file2.js': '', + 'file3.js': '', + }); + + const ignore = loadIgnoreRules({ + projectRoot: tmpDir, + useGitignore: false, + useGeminiignore: false, + ignoreDirs: [], + }); + + const paths = await crawl({ + crawlDirectory: tmpDir, + cwd: tmpDir, + ignore, + cache: false, + cacheTtl: 0, + maxFiles: 2, + }); + + // fdir returns files and directories. + // In our filter, we only increment fileCount for files. + // So we should have 2 files + some directories. + const files = paths.filter((p) => p !== '.' && !p.endsWith('/')); + expect(files.length).toBe(2); + }); }); diff --git a/packages/core/src/utils/filesearch/crawler.ts b/packages/core/src/utils/filesearch/crawler.ts index 9184ba3286..6eb174b968 100644 --- a/packages/core/src/utils/filesearch/crawler.ts +++ b/packages/core/src/utils/filesearch/crawler.ts @@ -16,6 +16,8 @@ export interface CrawlOptions { cwd: string; // The fdir maxDepth option. maxDepth?: number; + // Maximum number of files to return. + maxFiles?: number; // A pre-configured Ignore instance. ignore: Ignore; // Caching options. @@ -43,6 +45,9 @@ export async function crawl(options: CrawlOptions): Promise { const posixCwd = toPosixPath(options.cwd); const posixCrawlDirectory = toPosixPath(options.crawlDirectory); + const maxFiles = options.maxFiles ?? Infinity; + let fileCount = 0; + let truncated = false; let results: string[]; try { @@ -51,7 +56,21 @@ export async function crawl(options: CrawlOptions): Promise { .withRelativePaths() .withDirs() .withPathSeparator('/') // Always use unix style paths + .filter((path, isDirectory) => { + if (!isDirectory) { + fileCount++; + if (fileCount > maxFiles) { + truncated = true; + return false; + } + } + return true; + }) .exclude((_, dirPath) => { + if (fileCount > maxFiles) { + truncated = true; + return true; + } const relativePath = path.posix.relative(posixCrawlDirectory, dirPath); return dirFilter(`${relativePath}/`); }); @@ -72,7 +91,7 @@ export async function crawl(options: CrawlOptions): Promise { path.posix.join(relativeToCrawlDir, p), ); - if (options.cache) { + if (options.cache && !truncated) { const cacheKey = cache.getCacheKey( options.crawlDirectory, options.ignore.getFingerprint(), diff --git a/packages/core/src/utils/filesearch/fileSearch.test.ts b/packages/core/src/utils/filesearch/fileSearch.test.ts index 1b61cdf3f7..ee1af9761f 100644 --- a/packages/core/src/utils/filesearch/fileSearch.test.ts +++ b/packages/core/src/utils/filesearch/fileSearch.test.ts @@ -7,6 +7,7 @@ import { describe, it, expect, afterEach, vi } from 'vitest'; import { FileSearchFactory, AbortError, filter } from './fileSearch.js'; import { createTmpDir, cleanupTmpDir } from '@google/gemini-cli-test-utils'; +import * as crawler from './crawler.js'; describe('FileSearch', () => { let tmpDir: string; @@ -481,6 +482,33 @@ describe('FileSearch', () => { expect(results).toEqual(['src/', 'src/main.js']); }); + it('should respect default maxFiles budget of 20000 in RecursiveFileSearch', async () => { + const crawlSpy = vi.spyOn(crawler, 'crawl'); + + tmpDir = await createTmpDir({ + 'file1.js': '', + }); + + const fileSearch = FileSearchFactory.create({ + projectRoot: tmpDir, + useGitignore: false, + useGeminiignore: false, + ignoreDirs: [], + cache: false, + cacheTtl: 0, + enableRecursiveFileSearch: true, + disableFuzzySearch: false, + }); + + await fileSearch.initialize(); + + expect(crawlSpy).toHaveBeenCalledWith( + expect.objectContaining({ + maxFiles: 20000, + }), + ); + }); + it('should be cancellable via AbortSignal', async () => { const largeDir: Record = {}; for (let i = 0; i < 100; i++) { diff --git a/packages/core/src/utils/filesearch/fileSearch.ts b/packages/core/src/utils/filesearch/fileSearch.ts index 6cd82b75a3..e18fdf3111 100644 --- a/packages/core/src/utils/filesearch/fileSearch.ts +++ b/packages/core/src/utils/filesearch/fileSearch.ts @@ -24,6 +24,7 @@ export interface FileSearchOptions { enableRecursiveFileSearch: boolean; disableFuzzySearch: boolean; maxDepth?: number; + maxFiles?: number; } export class AbortError extends Error { @@ -109,7 +110,9 @@ class RecursiveFileSearch implements FileSearch { cache: this.options.cache, cacheTtl: this.options.cacheTtl, maxDepth: this.options.maxDepth, + maxFiles: this.options.maxFiles ?? 20000, }); + this.buildResultCache(); }