mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-20 02:51:55 -07:00
fix(cli): prevent OOM crash by limiting file search traversal and adding timeout (#16696)
This commit is contained in:
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<string[]> {
|
||||
|
||||
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<string[]> {
|
||||
.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<string[]> {
|
||||
path.posix.join(relativeToCrawlDir, p),
|
||||
);
|
||||
|
||||
if (options.cache) {
|
||||
if (options.cache && !truncated) {
|
||||
const cacheKey = cache.getCacheKey(
|
||||
options.crawlDirectory,
|
||||
options.ignore.getFingerprint(),
|
||||
|
||||
@@ -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<string, string> = {};
|
||||
for (let i = 0; i < 100; i++) {
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user