diff --git a/integration-tests/grep-total-max-matches.test.ts b/integration-tests/grep-total-max-matches.test.ts new file mode 100644 index 0000000000..445717b2fa --- /dev/null +++ b/integration-tests/grep-total-max-matches.test.ts @@ -0,0 +1,92 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeAll, afterAll } from 'vitest'; +import * as path from 'node:path'; +import * as fs from 'node:fs/promises'; +import * as os from 'node:os'; +import { GrepTool } from '../packages/core/src/tools/grep.js'; +import { Config } from '../packages/core/src/config/config.js'; +import { WorkspaceContext } from '../packages/core/src/utils/workspaceContext.js'; + +// Mock Config to provide necessary context +class MockConfig { + constructor(private targetDir: string) {} + + getTargetDir() { + return this.targetDir; + } + + getWorkspaceContext() { + return new WorkspaceContext(this.targetDir, [this.targetDir]); + } + + getDebugMode() { + return true; + } + + getFileFilteringOptions() { + return { + respectGitIgnore: true, + respectGeminiIgnore: true, + customIgnoreFilePaths: [], + }; + } + + getFileExclusions() { + return { + getGlobExcludes: () => [], + }; + } + + validatePathAccess() { + return null; + } +} + +describe('grep-total-max-matches', () => { + let tempDir: string; + let tool: GrepTool; + + beforeAll(async () => { + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'grep-total-max-test-')); + + // Create a test file with multiple matches + const content = ` + match 1 + match 2 + match 3 + match 4 + match 5 + `; + await fs.writeFile(path.join(tempDir, 'many_matches.txt'), content); + + const config = new MockConfig(tempDir) as unknown as Config; + tool = new GrepTool(config, null!); + }); + + afterAll(async () => { + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + it('should limit total matches when total_max_matches is set', async () => { + const invocation = tool.build({ + pattern: 'match', + total_max_matches: 3, + }); + const result = await invocation.execute(new AbortController().signal); + + expect(result.llmContent).toContain('Found 3 matches'); + expect(result.llmContent).toContain('match 1'); + expect(result.llmContent).toContain('match 2'); + expect(result.llmContent).toContain('match 3'); + expect(result.llmContent).not.toContain('match 4'); + expect(result.llmContent).not.toContain('match 5'); + expect(result.llmContent).toContain( + '(results limited to 3 matches for performance)', + ); + }); +}); diff --git a/integration-tests/ripgrep-max-matches.test.ts b/integration-tests/ripgrep-max-matches.test.ts index 780a3d93c6..d25eaa0a08 100644 --- a/integration-tests/ripgrep-max-matches.test.ts +++ b/integration-tests/ripgrep-max-matches.test.ts @@ -97,4 +97,12 @@ describe('ripgrep-max-matches', () => { expect(result.llmContent).toContain('match 3'); expect(result.llmContent).toContain('match 4'); }); + + it('should return context lines when requested', async () => { + const invocation = tool.build({ pattern: 'match 1', context: 1 }); + const result = await invocation.execute(new AbortController().signal); + + expect(result.llmContent).toContain('match 1'); + expect(result.llmContent).toContain('filler'); // Context line + }); }); diff --git a/integration-tests/ripgrep-total-max-matches.test.ts b/integration-tests/ripgrep-total-max-matches.test.ts new file mode 100644 index 0000000000..0c87bd1b29 --- /dev/null +++ b/integration-tests/ripgrep-total-max-matches.test.ts @@ -0,0 +1,100 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeAll, afterAll } from 'vitest'; +import * as path from 'node:path'; +import * as fs from 'node:fs/promises'; +import * as os from 'node:os'; +import { RipGrepTool } from '../packages/core/src/tools/ripGrep.js'; +import { Config } from '../packages/core/src/config/config.js'; +import { WorkspaceContext } from '../packages/core/src/utils/workspaceContext.js'; + +// Mock Config to provide necessary context +class MockConfig { + constructor(private targetDir: string) {} + + getTargetDir() { + return this.targetDir; + } + + getWorkspaceContext() { + return new WorkspaceContext(this.targetDir, [this.targetDir]); + } + + getDebugMode() { + return true; + } + + getFileFilteringRespectGeminiIgnore() { + return true; + } + + getFileFilteringOptions() { + return { + respectGitIgnore: true, + respectGeminiIgnore: true, + customIgnoreFilePaths: [], + }; + } + + validatePathAccess() { + return null; + } +} + +describe('ripgrep-total-max-matches', () => { + let tempDir: string; + let tool: RipGrepTool; + + beforeAll(async () => { + tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'ripgrep-total-max-test-'), + ); + + // Create a test file with multiple matches + const content = ` + match 1 + match 2 + match 3 + match 4 + match 5 + `; + await fs.writeFile(path.join(tempDir, 'many_matches.txt'), content); + + const config = new MockConfig(tempDir) as unknown as Config; + tool = new RipGrepTool(config); + }); + + afterAll(async () => { + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + it('should limit total matches when total_max_matches is set', async () => { + const invocation = tool.build({ + pattern: 'match', + total_max_matches: 3, + }); + const result = await invocation.execute(new AbortController().signal); + + expect(result.llmContent).toContain('Found 3 matches'); + expect(result.llmContent).toContain('match 1'); + expect(result.llmContent).toContain('match 2'); + expect(result.llmContent).toContain('match 3'); + expect(result.llmContent).not.toContain('match 4'); + expect(result.llmContent).not.toContain('match 5'); + expect(result.llmContent).toContain( + '(results limited to 3 matches for performance)', + ); + }); + + it('should use default limit when total_max_matches is not set', async () => { + // We can't easily test the default 100 without 100 matches, but we can verify it doesn't fail. + const invocation = tool.build({ pattern: 'match' }); + const result = await invocation.execute(new AbortController().signal); + + expect(result.llmContent).toContain('Found 5 matches'); + }); +}); diff --git a/packages/core/src/tools/constants.ts b/packages/core/src/tools/constants.ts index 81765ba628..8251dd1b17 100644 --- a/packages/core/src/tools/constants.ts +++ b/packages/core/src/tools/constants.ts @@ -3,5 +3,5 @@ * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -export const DEFAULT_TOTAL_MAX_MATCHES = 100; +export const DEFAULT_TOTAL_MAX_MATCHES = 2000; export const DEFAULT_SEARCH_TIMEOUT_MS = 30000; diff --git a/packages/core/src/tools/grep.ts b/packages/core/src/tools/grep.ts index 56f96bf892..1012340170 100644 --- a/packages/core/src/tools/grep.ts +++ b/packages/core/src/tools/grep.ts @@ -56,6 +56,11 @@ export interface GrepToolParams { * Optional: Maximum number of matches to return per file. */ max_matches_per_file?: number; + + /** + * Optional: Maximum number of total matches to return. + */ + total_max_matches?: number; } /** @@ -194,7 +199,8 @@ class GrepToolInvocation extends BaseToolInvocation< // Collect matches from all search directories let allMatches: GrepMatch[] = []; - const totalMaxMatches = DEFAULT_TOTAL_MAX_MATCHES; + const totalMaxMatches = + this.params.total_max_matches ?? DEFAULT_TOTAL_MAX_MATCHES; // Create a timeout controller to prevent indefinitely hanging searches const timeoutController = new AbortController(); @@ -797,6 +803,12 @@ export class GrepTool extends BaseDeclarativeTool { type: 'integer', minimum: 1, }, + total_max_matches: { + description: + 'Optional: Maximum number of total matches to return. Use this to limit the overall size of the response. Defaults to 100 if omitted.', + type: 'integer', + minimum: 1, + }, }, required: ['pattern'], type: 'object', diff --git a/packages/core/src/tools/ripGrep.ts b/packages/core/src/tools/ripGrep.ts index 6d999157b3..128c591223 100644 --- a/packages/core/src/tools/ripGrep.ts +++ b/packages/core/src/tools/ripGrep.ts @@ -141,6 +141,11 @@ export interface RipGrepToolParams { * Optional: Maximum number of matches to return per file. */ max_matches_per_file?: number; + + /** + * Optional: Maximum number of total matches to return. + */ + total_max_matches?: number; } /** @@ -214,7 +219,8 @@ class GrepToolInvocation extends BaseToolInvocation< const searchDirDisplay = pathParam; - const totalMaxMatches = DEFAULT_TOTAL_MAX_MATCHES; + const totalMaxMatches = + this.params.total_max_matches ?? DEFAULT_TOTAL_MAX_MATCHES; if (this.config.getDebugMode()) { debugLogger.log(`[GrepTool] Total result limit: ${totalMaxMatches}`); } @@ -530,7 +536,7 @@ class GrepToolInvocation extends BaseToolInvocation< ): GrepMatch | null { try { const json = JSON.parse(line); - if (json.type === 'match') { + if (json.type === 'match' || json.type === 'context') { const match = json.data; // Defensive check: ensure text properties exist (skips binary/invalid encoding) if (match.path?.text && match.lines?.text) { @@ -667,6 +673,12 @@ export class RipGrepTool extends BaseDeclarativeTool< type: 'integer', minimum: 1, }, + total_max_matches: { + description: + 'Optional: Maximum number of total matches to return. Use this to limit the overall size of the response. Defaults to 2000 if omitted.', + type: 'integer', + minimum: 1, + }, }, required: ['pattern'], type: 'object',