diff --git a/packages/cli/src/config/auth.ts b/packages/cli/src/config/auth.ts index 1ca07f98eb4..21f3965a512 100644 --- a/packages/cli/src/config/auth.ts +++ b/packages/cli/src/config/auth.ts @@ -45,5 +45,9 @@ export async function validateAuthMethod( return null; } + if (authMethod === AuthType.GATEWAY) { + return null; + } + return 'Invalid auth method selected.'; } diff --git a/packages/core/src/agents/codebase-investigator.test.ts b/packages/core/src/agents/codebase-investigator.test.ts index 3637daa9e36..110a0363e6e 100644 --- a/packages/core/src/agents/codebase-investigator.test.ts +++ b/packages/core/src/agents/codebase-investigator.test.ts @@ -11,6 +11,7 @@ import { GREP_TOOL_NAME, LS_TOOL_NAME, READ_FILE_TOOL_NAME, + AST_SEARCH_TOOL_NAME, } from '../tools/tool-names.js'; import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; import { makeFakeConfig } from '../test-utils/config.js'; @@ -50,6 +51,7 @@ describe('CodebaseInvestigatorAgent', () => { READ_FILE_TOOL_NAME, GLOB_TOOL_NAME, GREP_TOOL_NAME, + AST_SEARCH_TOOL_NAME, ]); }); @@ -77,4 +79,15 @@ describe('CodebaseInvestigatorAgent', () => { const agent = CodebaseInvestigatorAgent(config); expect(agent.promptConfig.systemPrompt).toContain('`ls -R`'); }); + + it('should mention ast_search tool in system prompt', () => { + const agent = CodebaseInvestigatorAgent(config); + expect(agent.promptConfig.systemPrompt).toContain('ast_search'); + expect(agent.promptConfig.systemPrompt).toContain('scope'); + }); + + it('should include ast_search in tool config', () => { + const agent = CodebaseInvestigatorAgent(config); + expect(agent.toolConfig?.tools).toContain(AST_SEARCH_TOOL_NAME); + }); }); diff --git a/packages/core/src/agents/codebase-investigator.ts b/packages/core/src/agents/codebase-investigator.ts index 5036bd28239..fea6ad71805 100644 --- a/packages/core/src/agents/codebase-investigator.ts +++ b/packages/core/src/agents/codebase-investigator.ts @@ -10,6 +10,7 @@ import { GREP_TOOL_NAME, LS_TOOL_NAME, READ_FILE_TOOL_NAME, + AST_SEARCH_TOOL_NAME, } from '../tools/tool-names.js'; import { DEFAULT_THINKING_MODE, @@ -121,6 +122,7 @@ export const CodebaseInvestigatorAgent = ( READ_FILE_TOOL_NAME, GLOB_TOOL_NAME, GREP_TOOL_NAME, + AST_SEARCH_TOOL_NAME, ], }, @@ -132,6 +134,7 @@ export const CodebaseInvestigatorAgent = ( systemPrompt: `You are **Codebase Investigator**, a hyper-specialized AI agent and an expert in reverse-engineering complex software projects. You are a sub-agent within a larger development system. Your **SOLE PURPOSE** is to build a complete mental model of the code relevant to a given investigation. You must identify all relevant files, understand their roles, and foresee the direct architectural consequences of potential changes. You are a sub-agent in a larger system. Your only responsibility is to provide deep, actionable context. +- **DO:** Use the \`ast_search\` tool to quickly locate symbol boundaries and get file outlines before reading entire files. For broad exploration, use \`ast_search\` with scope "map" to get a compressed structural overview of the codebase. - **DO:** Find the key modules, classes, and functions that are part of the problem and its solution. - **DO:** Understand *why* the code is written the way it is. Question everything. - **DO:** Foresee the ripple effects of a change. If \`function A\` is modified, you must check its callers. If a data structure is altered, you must identify where its type definitions need to be updated. diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index c53066573d0..c2ef8bdea48 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -41,6 +41,7 @@ import { EditTool } from '../tools/edit.js'; import { ShellTool } from '../tools/shell.js'; import { WriteFileTool } from '../tools/write-file.js'; import { WebFetchTool } from '../tools/web-fetch.js'; +import { ASTSearchTool } from '../tools/ast-search.js'; import { setGeminiMdFilename, getCurrentGeminiMdFilename, @@ -4020,6 +4021,9 @@ export class Config implements McpContext, AgentLoopContext { maybeRegister(ListMcpResourcesTool, () => registry.registerTool(new ListMcpResourcesTool(this, this.messageBus)), ); + maybeRegister(ASTSearchTool, () => + registry.registerTool(new ASTSearchTool(this, this.messageBus)), + ); maybeRegister(ShellTool, () => registry.registerTool(new ShellTool(this, this.messageBus)), ); diff --git a/packages/core/src/prompts/snippets.ts b/packages/core/src/prompts/snippets.ts index d62613a614d..950b9fe8164 100644 --- a/packages/core/src/prompts/snippets.ts +++ b/packages/core/src/prompts/snippets.ts @@ -732,7 +732,7 @@ function workflowStepResearch(options: PrimaryWorkflowsOptions): string { subAgentSearch = ` For **simple, targeted searches** (like finding a specific function name, file path, or variable declaration), use ${toolsStr} directly in parallel.`; } - return `1. **Research:** Systematically map the codebase and validate assumptions. Utilize specialized sub-agents (e.g., \`codebase_investigator\`) as the primary mechanism for initial discovery when the task involves **complex refactoring, codebase exploration or system-wide analysis**.${subAgentSearch} Use ${formatToolName(READ_FILE_TOOL_NAME)} to validate all assumptions. **Prioritize empirical reproduction of reported issues to confirm the failure state.**${suggestion}`; + return `1. **Research:** Systematically map the codebase and validate assumptions. Utilize specialized sub-agents (e.g., \`codebase_investigator\`) as the primary mechanism for initial discovery when the task involves **complex refactoring, codebase exploration or system-wide analysis**.${subAgentSearch} Use \`ast_search\` with scope "map" or "outline" to quickly understand codebase structure before reading files in full. Use ${formatToolName(READ_FILE_TOOL_NAME)} to validate all assumptions. **Prioritize empirical reproduction of reported issues to confirm the failure state.**${suggestion}`; } return `1. **Research:** Systematically map the codebase and validate assumptions.${searchSentence} Use ${formatToolName(READ_FILE_TOOL_NAME)} to validate all assumptions. **Prioritize empirical reproduction of reported issues to confirm the failure state.**${suggestion}`; diff --git a/packages/core/src/services/astAnalysisService.test.ts b/packages/core/src/services/astAnalysisService.test.ts new file mode 100644 index 00000000000..045c5967096 --- /dev/null +++ b/packages/core/src/services/astAnalysisService.test.ts @@ -0,0 +1,295 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import * as fs from 'node:fs/promises'; +import * as path from 'node:path'; +import * as os from 'node:os'; +import { + ASTAnalysisService, + extractSymbols, + findClosingBrace, + findIndentEnd, +} from './astAnalysisService.js'; + +describe('ASTAnalysisService', () => { + let tmpDir: string; + + beforeEach(async () => { + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'ast-svc-')); + }); + + afterEach(async () => { + await fs.rm(tmpDir, { recursive: true, force: true }); + }); + + describe('findClosingBrace', () => { + it('should find the closing brace of a simple block', () => { + const lines = ['function f() {', ' return 1;', '}']; + expect(findClosingBrace(lines, 0)).toBe(2); + }); + + it('should handle nested braces correctly', () => { + const lines = [ + 'class C {', + ' m() {', + ' if (x) {', + ' }', + ' }', + '}', + ]; + expect(findClosingBrace(lines, 0)).toBe(5); + expect(findClosingBrace(lines, 1)).toBe(4); + }); + + it('should ignore braces inside string literals', () => { + const lines = ['function f() {', ' const s = "}{";', '}']; + expect(findClosingBrace(lines, 0)).toBe(2); + }); + }); + + describe('findIndentEnd', () => { + it('should find the end of a Python indentation block', () => { + const lines = ['def f():', ' x = 1', ' return x', 'def g():']; + expect(findIndentEnd(lines, 0)).toBe(2); + }); + + it('should skip blank lines within a block', () => { + const lines = ['def f():', ' x = 1', '', ' y = 2', 'z = 3']; + expect(findIndentEnd(lines, 0)).toBe(3); + }); + }); + + describe('extractSymbols', () => { + it('should extract TypeScript class with methods', () => { + const lines = [ + 'export class MyService {', + ' private val: number;', + ' public process(x: string): void {', + ' console.log(x);', + ' }', + '}', + ]; + const syms = extractSymbols(lines, 'typescript'); + expect(syms).toHaveLength(1); + expect(syms[0].name).toBe('MyService'); + expect(syms[0].kind).toBe('class'); + expect(syms[0].children.length).toBeGreaterThanOrEqual(1); + }); + + it('should extract standalone functions', () => { + const lines = [ + 'export function doWork() {', + ' return 42;', + '}', + '', + 'export async function fetchData() {', + ' return null;', + '}', + ]; + const syms = extractSymbols(lines, 'typescript'); + expect(syms).toHaveLength(2); + expect(syms[0].name).toBe('doWork'); + expect(syms[1].name).toBe('fetchData'); + }); + + it('should extract interfaces and types', () => { + const lines = [ + 'export interface Config {', + ' host: string;', + '}', + 'export type Status = "ok" | "err";', + ]; + const syms = extractSymbols(lines, 'typescript'); + expect( + syms.some((s) => s.name === 'Config' && s.kind === 'interface'), + ).toBe(true); + expect(syms.some((s) => s.name === 'Status' && s.kind === 'type')).toBe( + true, + ); + }); + + it('should extract Python classes and functions', () => { + const lines = [ + 'class Handler:', + ' def run(self):', + ' pass', + 'def util():', + ' pass', + ]; + const syms = extractSymbols(lines, 'python'); + expect(syms).toHaveLength(2); + expect(syms[0].name).toBe('Handler'); + expect(syms[0].kind).toBe('class'); + expect(syms[0].children).toHaveLength(1); + expect(syms[0].children[0].name).toBe('run'); + expect(syms[0].children[0].kind).toBe('method'); + expect(syms[1].name).toBe('util'); + }); + + it('should skip comments and imports', () => { + const lines = [ + '// comment', + 'import { X } from "y";', + 'export function real() {', + ' return 1;', + '}', + ]; + const syms = extractSymbols(lines, 'typescript'); + expect(syms).toHaveLength(1); + expect(syms[0].name).toBe('real'); + }); + + it('should return empty array for empty input', () => { + expect(extractSymbols([], 'typescript')).toHaveLength(0); + }); + + it('should truncate long signatures', () => { + const longLine = 'export function ' + 'a'.repeat(130) + '() {'; + const syms = extractSymbols([longLine, '}'], 'typescript'); + expect(syms).toHaveLength(1); + expect(syms[0].signature.length).toBeLessThanOrEqual(120); + }); + }); + + describe('getFileOutline', () => { + it('should outline a TypeScript file from disk', async () => { + await fs.writeFile( + path.join(tmpDir, 'svc.ts'), + 'export class Svc {\n run() {\n return 1;\n }\n}\nexport function helper() {\n return 2;\n}\n', + ); + const service = new ASTAnalysisService(tmpDir); + const outline = await service.getFileOutline('svc.ts'); + expect(outline).not.toBeNull(); + expect(outline!.language).toBe('typescript'); + expect(outline!.symbols.length).toBeGreaterThanOrEqual(2); + }); + + it('should return null for unsupported extensions', async () => { + await fs.writeFile(path.join(tmpDir, 'data.json'), '{}'); + const service = new ASTAnalysisService(tmpDir); + expect(await service.getFileOutline('data.json')).toBeNull(); + }); + + it('should return null for missing files', async () => { + const service = new ASTAnalysisService(tmpDir); + expect(await service.getFileOutline('nope.ts')).toBeNull(); + }); + }); + + describe('findSymbolBounds', () => { + it('should locate a class precisely', async () => { + const content = [ + 'import { X } from "x";', + '', + 'export class Target {', + ' method() {', + ' return 1;', + ' }', + '}', + '', + 'export function other() {}', + ].join('\n'); + await fs.writeFile(path.join(tmpDir, 'f.ts'), content); + const svc = new ASTAnalysisService(tmpDir); + const bounds = await svc.findSymbolBounds('f.ts', 'Target'); + expect(bounds).not.toBeNull(); + expect(bounds!.startLine).toBe(3); + expect(bounds!.endLine).toBe(7); + }); + + it('should return null for a non-existent symbol', async () => { + await fs.writeFile( + path.join(tmpDir, 'f.ts'), + 'export function real() {}\n', + ); + const svc = new ASTAnalysisService(tmpDir); + expect(await svc.findSymbolBounds('f.ts', 'ghost')).toBeNull(); + }); + }); + + describe('getCodebaseMap', () => { + it('should map multiple source files', async () => { + const src = path.join(tmpDir, 'src'); + await fs.mkdir(src); + await fs.writeFile(path.join(src, 'a.ts'), 'export class A {}\n'); + await fs.writeFile(path.join(src, 'b.ts'), 'export function b() {}\n'); + await fs.writeFile(path.join(src, 'c.json'), '{}'); + + const svc = new ASTAnalysisService(tmpDir); + const map = await svc.getCodebaseMap(); + expect(map).toContain('Codebase Map:'); + expect(map).toContain('class A'); + expect(map).toContain('function b'); + expect(map).not.toContain('.json'); + }); + + it('should skip node_modules', async () => { + const nm = path.join(tmpDir, 'node_modules', 'pkg'); + await fs.mkdir(nm, { recursive: true }); + await fs.writeFile(path.join(nm, 'index.ts'), 'export class X {}\n'); + await fs.writeFile( + path.join(tmpDir, 'main.ts'), + 'export class Main {}\n', + ); + + const svc = new ASTAnalysisService(tmpDir); + const map = await svc.getCodebaseMap(); + expect(map).toContain('Main'); + expect(map).not.toContain('node_modules'); + }); + }); + + describe('shouldIgnore callback', () => { + it('getFileOutline should return null for ignored files', async () => { + await fs.writeFile( + path.join(tmpDir, 'secret.ts'), + 'export class Secret {}\n', + ); + const ignoreFn = (p: string) => p.includes('secret'); + const svc = new ASTAnalysisService(tmpDir, ignoreFn); + expect(await svc.getFileOutline('secret.ts')).toBeNull(); + }); + + it('getFileOutline should still work for non-ignored files', async () => { + await fs.writeFile( + path.join(tmpDir, 'public.ts'), + 'export class Public {}\n', + ); + const ignoreFn = (p: string) => p.includes('secret'); + const svc = new ASTAnalysisService(tmpDir, ignoreFn); + const outline = await svc.getFileOutline('public.ts'); + expect(outline).not.toBeNull(); + expect(outline!.symbols[0].name).toBe('Public'); + }); + + it('getCodebaseMap should exclude ignored files', async () => { + await fs.writeFile( + path.join(tmpDir, 'visible.ts'), + 'export class Visible {}\n', + ); + await fs.writeFile( + path.join(tmpDir, 'hidden.ts'), + 'export class Hidden {}\n', + ); + const ignoreFn = (p: string) => p.includes('hidden'); + const svc = new ASTAnalysisService(tmpDir, ignoreFn); + const map = await svc.getCodebaseMap(); + expect(map).toContain('Visible'); + expect(map).not.toContain('Hidden'); + }); + + it('findSymbolBounds should return null for ignored files', async () => { + await fs.writeFile( + path.join(tmpDir, 'ignored.ts'), + 'export function target() { return 1; }\n', + ); + const ignoreFn = (p: string) => p.includes('ignored'); + const svc = new ASTAnalysisService(tmpDir, ignoreFn); + expect(await svc.findSymbolBounds('ignored.ts', 'target')).toBeNull(); + }); + }); +}); diff --git a/packages/core/src/services/astAnalysisService.ts b/packages/core/src/services/astAnalysisService.ts new file mode 100644 index 00000000000..fceb217414d --- /dev/null +++ b/packages/core/src/services/astAnalysisService.ts @@ -0,0 +1,694 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as path from 'node:path'; +import * as fs from 'node:fs/promises'; +// Debug logging available via: import { debugLogger } from '../utils/debugLogger.js'; + +/** A symbol extracted from source code. */ +export interface ASTSymbol { + name: string; + kind: 'class' | 'function' | 'method' | 'interface' | 'type' | 'enum'; + startLine: number; + endLine: number; + signature: string; + children: ASTSymbol[]; +} + +/** Structural outline of a single file. */ +export interface ASTFileOutline { + filePath: string; + language: string; + symbols: ASTSymbol[]; + totalLines: number; +} + +const LANG_MAP: Record = { + '.ts': 'typescript', + '.tsx': 'typescript', + '.js': 'javascript', + '.jsx': 'javascript', + '.py': 'python', + '.go': 'go', + '.rs': 'rust', + '.java': 'java', +}; + +const SKIP_DIRS = new Set([ + 'node_modules', + '.git', + 'dist', + 'build', + '__pycache__', + 'vendor', + 'target', +]); + +/** + * Extracts structural outlines from source files using regex heuristics. + * + * Designed as a zero-dependency fallback for when ast-grep (sg) is not installed. + * Covers the three capabilities outlined in issue #22745: + * 1. Symbol-boundary detection for precise method-level reads + * 2. Structural search by symbol name + * 3. Compressed codebase mapping (class/function/signature outlines) + */ +export class ASTAnalysisService { + constructor( + private readonly targetDir: string, + private readonly shouldIgnore?: (filePath: string) => boolean, + ) {} + + /** + * Returns the structural outline of a single source file. + */ + async getFileOutline(filePath: string): Promise { + const resolvedTargetDir = path.resolve(this.targetDir); + const absPath = path.resolve(resolvedTargetDir, filePath); + + // Guard against path traversal + const relative = path.relative(resolvedTargetDir, absPath); + if (relative.startsWith('..') || path.isAbsolute(relative)) { + return null; + } + + // Enforce ignore patterns (e.g. .gitignore, .geminiignore) when provided + if ( + this.shouldIgnore && + (this.shouldIgnore(filePath) || this.shouldIgnore(absPath)) + ) { + return null; + } + + const ext = path.extname(absPath).toLowerCase(); + const language = LANG_MAP[ext]; + if (!language) return null; + + // Reject files larger than 2MB to prevent OOM or event loop blocking + // on minified bundles, logs, or database dumps. + const MAX_FILE_SIZE = 2 * 1024 * 1024; + try { + const stat = await fs.stat(absPath); + if (stat.size > MAX_FILE_SIZE) return null; + } catch { + return null; + } + + let content: string; + try { + content = await fs.readFile(absPath, 'utf-8'); + } catch { + return null; + } + + const lines = content.split('\n'); + const symbols = extractSymbols(lines, language); + return { filePath, language, symbols, totalLines: lines.length }; + } + + /** + * Finds the start/end line bounds of a named symbol in a file. + */ + async findSymbolBounds( + filePath: string, + symbolName: string, + ): Promise<{ startLine: number; endLine: number } | null> { + const outline = await this.getFileOutline(filePath); + if (!outline) return null; + + const found = findSymbolRecursive(outline.symbols, symbolName); + if (!found) return null; + + return { startLine: found.startLine, endLine: found.endLine }; + } + + /** + * Generates a compressed codebase map for LLM consumption. + * Walks source files up to `maxFiles`, extracts top-level symbols, + * and returns a text outline with file paths and signatures. + */ + async getCodebaseMap( + subDir?: string, + maxFiles: number = 100, + ): Promise { + const resolvedTargetDir = path.resolve(this.targetDir); + const searchDir = subDir + ? path.resolve(resolvedTargetDir, subDir) + : resolvedTargetDir; + + // Guard against path traversal + const relative = path.relative(resolvedTargetDir, searchDir); + if (relative.startsWith('..') || path.isAbsolute(relative)) { + return 'Error: path traversal detected, directory must be within workspace.'; + } + + const files = await collectSourceFiles( + searchDir, + maxFiles, + this.shouldIgnore, + ); + const sections: string[] = []; + let totalSymbols = 0; + + for (const file of files) { + const relPath = path.relative(resolvedTargetDir, file); + const outline = await this.getFileOutline(relPath); + if (!outline || outline.symbols.length === 0) continue; + + totalSymbols += countSymbols(outline.symbols); + sections.push(formatOutline(outline)); + } + + const header = `Codebase Map: ${sections.length} files, ${totalSymbols} symbols\n${'='.repeat(50)}`; + return header + '\n\n' + sections.join('\n\n'); + } +} + +// ── Pure helpers (exported for unit testing) ─────────────────────────────── + +export function extractSymbols(lines: string[], language: string): ASTSymbol[] { + const symbols: ASTSymbol[] = []; + const patterns = getDeclarationPatterns(language); + + // Pre-process: strip block comments by replacing their content with spaces. + // This handles mid-line comments like `code /* comment */ more_code` and + // multi-line blocks, while preserving line numbers and offsets. + const cleaned = stripBlockComments(lines, language); + + // Pre-strip string/regex literals once for the entire file so that + // findClosingBrace does not re-run the regex on every call (O(N) instead of O(M*N)). + const stripped = language !== 'python' ? preStripLines(cleaned) : undefined; + + for (let i = 0; i < cleaned.length; i++) { + const trimmed = cleaned[i].trim(); + + if ( + trimmed === '' || + trimmed.startsWith('//') || + trimmed.startsWith('#') || + trimmed.startsWith('*') || + trimmed.startsWith('import ') || + trimmed.startsWith('from ') + ) { + continue; + } + + // Only match at top-level indentation (<=8 spaces for brace langs to support + // 4-space indented codebases and namespace/module nesting). + // Use original lines for indent check since cleaned lines may have inflated + // indentation from blanked block comments. + const indent = lines[i].length - lines[i].trimStart().length; + if (language !== 'python' && indent > 8) continue; + if (language === 'python' && indent > 0) continue; + + for (const { regex, kind } of patterns) { + const match = regex.exec(trimmed); + if (!match?.[1]) continue; + + const endLine = + language === 'python' + ? findIndentEnd(cleaned, i) + : findClosingBrace(cleaned, i, stripped); + + const sym: ASTSymbol = { + name: match[1], + kind, + startLine: i + 1, + endLine: endLine + 1, + signature: + trimmed.length > 120 ? trimmed.slice(0, 117) + '...' : trimmed, + children: + kind === 'class' || kind === 'interface' + ? extractMembers(cleaned, i + 1, endLine, language, stripped) + : [], + }; + + symbols.push(sym); + // Skip past the symbol body for all block declarations (class, function, enum, etc.) + // Only 'type' aliases are single-line and should not advance + if (kind !== 'type') i = endLine; + break; + } + } + + return symbols; +} + +/** + * Strips block comments from source lines by replacing comment + * characters with spaces. This preserves line count and character offsets + * so that line-number-based logic (brace counting, indentation) stays correct. + * Handles mid-line comments, multi-line blocks, and lines with code after a comment. + */ +export function stripBlockComments( + lines: string[], + language?: string, +): string[] { + const result: string[] = []; + let inComment = false; + let inString: string | null = null; + const isPython = language === 'python'; + + for (const line of lines) { + let out = ''; + let j = 0; + let inLineComment = false; + + while (j < line.length) { + if (inComment) { + if (j + 1 < line.length && line[j] === '*' && line[j + 1] === '/') { + out += ' '; + j += 2; + inComment = false; + } else { + out += ' '; + j++; + } + } else if (inLineComment) { + out += ' '; + j++; + } else if (inString) { + // Inside string: blank contents but preserve delimiters + if (line[j] === '\\' && j + 1 < line.length) { + out += ' '; + j += 2; + } else if (inString.length === 3 && line.startsWith(inString, j)) { + // Closing triple-quote + out += inString; + inString = null; + j += 3; + } else if (inString.length === 1 && line[j] === inString) { + out += line[j]; + inString = null; + j++; + } else { + out += ' '; + j++; + } + } else { + // Normal code context - language-aware comment detection + if ( + !isPython && + j + 1 < line.length && + line[j] === '/' && + line[j + 1] === '/' + ) { + inLineComment = true; + out += '//'; + j += 2; + } else if (isPython && line[j] === '#') { + inLineComment = true; + out += '#'; + j++; + } else if ( + !isPython && + j + 1 < line.length && + line[j] === '/' && + line[j + 1] === '*' + ) { + out += ' '; + j += 2; + inComment = true; + } else if ( + isPython && + (line.startsWith('"""', j) || line.startsWith("'''", j)) + ) { + // Python triple-quoted string + inString = line.startsWith('"""', j) ? '"""' : "'''"; + out += inString; + j += 3; + } else if (line[j] === "'" || line[j] === '"' || line[j] === '`') { + inString = line[j]; + out += line[j]; + j++; + } else { + out += line[j]; + j++; + } + } + } + // Reset single-line string state at end of line. If a single or double + // quote was not closed (syntax error, unescaped quote, regex), don't let + // it bleed into subsequent lines. Multi-line delimiters (backtick, triple + // quotes) intentionally persist across lines. + if (inString && inString.length === 1 && inString !== '`') { + inString = null; + } + result.push(out); + } + return result; +} + +/** + * Pre-strips string/regex literals and single-line comments from source lines. + * Call once per file and pass the result to findClosingBrace to avoid O(M*N) + * repeated regex replacements when parsing multiple declarations. + */ +export function preStripLines(lines: string[]): string[] { + const re = /'(?:[^'\\]|\\.)*'|"(?:[^"\\]|\\.)*"|`(?:[^`\\]|\\.)*`/g; + return lines.map((line) => { + let stripped = line.replace(re, ''); + const commentIdx = stripped.indexOf('//'); + if (commentIdx >= 0) { + stripped = stripped.slice(0, commentIdx); + } + return stripped; + }); +} + +export function findClosingBrace( + lines: string[], + startLine: number, + strippedLines?: string[], +): number { + if (startLine < 0 || startLine >= lines.length) { + return startLine; + } + // When no pre-stripped lines are provided, compute them on the fly + // so that direct callers (tests, extractMembers without cache) still work. + const effective = strippedLines ?? preStripLines(lines); + let depth = 0; + let parenDepth = 0; + let opened = false; + for (let i = startLine; i < lines.length; i++) { + const stripped = effective[i] ?? ''; + for (const ch of stripped) { + if (ch === '(') parenDepth++; + else if (ch === ')') parenDepth = Math.max(0, parenDepth - 1); + // Ignore braces inside parentheses (inline object types in params) + if (parenDepth > 0) continue; + if (ch === '{') { + depth++; + opened = true; + } else if (ch === '}') { + depth--; + if (opened && depth === 0) return i; + } + } + } + // If brace matching failed, return startLine rather than the end of the file + // to prevent skipping the entire rest of the file during parsing. + return startLine; +} + +export function findIndentEnd(lines: string[], startLine: number): number { + if (startLine < 0 || startLine >= lines.length) { + return startLine; + } + const baseIndent = + lines[startLine].length - lines[startLine].trimStart().length; + let last = startLine; + let tripleQuoteChar: '"""' | "'''" | null = null; + for (let i = startLine + 1; i < lines.length; i++) { + const trimmed = lines[i].trim(); + + // Track Python triple-quoted strings - must match the same quote type + let j = 0; + while (j < trimmed.length) { + if (tripleQuoteChar) { + if (trimmed.startsWith(tripleQuoteChar, j)) { + tripleQuoteChar = null; + j += 3; + } else { + j++; + } + } else { + if (trimmed.startsWith('"""', j)) { + tripleQuoteChar = '"""'; + j += 3; + } else if (trimmed.startsWith("'''", j)) { + tripleQuoteChar = "'''"; + j += 3; + } else { + j++; + } + } + } + if (tripleQuoteChar) { + last = i; + continue; + } + + if (trimmed === '') continue; + // Skip comment lines - they may have arbitrary indentation and should + // not terminate the block (e.g. a top-level # comment inside a function) + if (trimmed.startsWith('#') || trimmed.startsWith('//')) { + last = i; + continue; + } + const indent = lines[i].length - lines[i].trimStart().length; + if (indent <= baseIndent) return last; + last = i; + } + return last; +} + +function extractMembers( + lines: string[], + start: number, + end: number, + language: string, + strippedLines?: string[], +): ASTSymbol[] { + const members: ASTSymbol[] = []; + const patterns = getMemberPatterns(language); + + for (let i = start; i < end; i++) { + const line = lines[i]; + if (line === undefined) continue; + const trimmed = line.trim(); + if (trimmed === '' || trimmed === '{' || trimmed === '}') continue; + for (const { regex, kind } of patterns) { + const match = regex.exec(trimmed); + if (!match?.[1]) continue; + const name = match[1]; + // Skip control flow keywords that look like function calls + if ( + [ + 'if', + 'for', + 'while', + 'switch', + 'catch', + 'synchronized', + 'return', + ].includes(name) + ) { + continue; + } + const memberEnd = + language === 'python' + ? Math.min(findIndentEnd(lines, i), end) + : Math.min(findClosingBrace(lines, i, strippedLines), end); + members.push({ + name, + kind, + startLine: i + 1, + endLine: memberEnd + 1, + signature: + trimmed.length > 100 ? trimmed.slice(0, 97) + '...' : trimmed, + children: [], + }); + if (kind === 'method' || kind === 'function') i = memberEnd; + break; + } + } + + return members; +} + +function findSymbolRecursive( + symbols: ASTSymbol[], + name: string, +): ASTSymbol | null { + for (const s of symbols) { + if (s.name === name) return s; + const found = findSymbolRecursive(s.children, name); + if (found) return found; + } + return null; +} + +function getDeclarationPatterns(lang: string) { + const p: Array<{ regex: RegExp; kind: ASTSymbol['kind'] }> = []; + switch (lang) { + case 'typescript': + case 'javascript': + p.push({ + regex: /(?:export\s+)?(?:abstract\s+)?class\s+(\w+)/, + kind: 'class', + }); + p.push({ regex: /(?:export\s+)?interface\s+(\w+)/, kind: 'interface' }); + p.push({ regex: /(?:export\s+)?type\s+(\w+)/, kind: 'type' }); + p.push({ regex: /(?:export\s+)?enum\s+(\w+)/, kind: 'enum' }); + p.push({ + regex: /(?:export\s+)?(?:async\s+)?function\s+(\w+)/, + kind: 'function', + }); + break; + case 'python': + p.push({ regex: /^class\s+(\w+)/, kind: 'class' }); + p.push({ regex: /^(?:async\s+)?def\s+(\w+)/, kind: 'function' }); + break; + case 'go': + p.push({ regex: /^type\s+(\w+)\s+struct/, kind: 'class' }); + p.push({ regex: /^type\s+(\w+)\s+interface/, kind: 'interface' }); + p.push({ regex: /^func\s+(?:\([^)]*\)\s+)?(\w+)/, kind: 'function' }); + break; + case 'rust': + p.push({ regex: /(?:pub\s+)?struct\s+(\w+)/, kind: 'class' }); + p.push({ + regex: /^impl(?:\s*<[^>]+>)?\s+(?:[\w:]+\s+for\s+)?(\w+)/, + kind: 'class', + }); + p.push({ regex: /(?:pub\s+)?trait\s+(\w+)/, kind: 'interface' }); + p.push({ regex: /(?:pub\s+)?enum\s+(\w+)/, kind: 'enum' }); + p.push({ + regex: /(?:pub\s+)?(?:async\s+)?fn\s+(\w+)/, + kind: 'function', + }); + break; + case 'java': + p.push({ + regex: /\bclass\s+(\w+)/, + kind: 'class', + }); + p.push({ + regex: /\binterface\s+(\w+)/, + kind: 'interface', + }); + p.push({ + regex: /\benum\s+(\w+)/, + kind: 'enum', + }); + p.push({ + regex: /\brecord\s+(\w+)/, + kind: 'class', + }); + break; + default: + break; + } + return p; +} + +function getMemberPatterns(lang: string) { + const p: Array<{ regex: RegExp; kind: ASTSymbol['kind'] }> = []; + switch (lang) { + case 'typescript': + case 'javascript': + p.push({ + regex: + /(?:public|private|protected|static|async|override|get|set)\s+(\w+)\s*[(<]/, + kind: 'method', + }); + p.push({ regex: /^(\w+)\s*\(/, kind: 'method' }); + break; + case 'python': + p.push({ regex: /^(?:async\s+)?def\s+(\w+)/, kind: 'method' }); + break; + case 'go': + p.push({ regex: /func\s+\([^)]+\)\s+(\w+)/, kind: 'method' }); + p.push({ regex: /^(\w+)\s*\(/, kind: 'method' }); + break; + case 'rust': + p.push({ regex: /(?:pub\s+)?(?:async\s+)?fn\s+(\w+)/, kind: 'method' }); + break; + case 'java': + p.push({ + regex: + /^(?:(?:public|private|protected|static|final|synchronized|abstract|default)\s+)*[\w<>[\]]+\s+(\w+)\s*\(/, + kind: 'method', + }); + // Constructors have no return type, match modifier(s) + name + ( + p.push({ + regex: + /^(?:(?:public|private|protected|static|final|synchronized|abstract|default)\s+)*(\w+)\s*\(/, + kind: 'method', + }); + break; + default: + break; + } + return p; +} + +function formatOutline(outline: ASTFileOutline): string { + const header = `## ${outline.filePath} (${outline.language}, ${outline.totalLines} lines)`; + const body = outline.symbols + .map((s) => { + const range = `L${s.startLine}-${s.endLine}`; + let line = ` ${s.kind} ${s.name} [${range}]: ${s.signature}`; + for (const c of s.children) { + line += `\n ${c.kind} ${c.name} [L${c.startLine}-${c.endLine}]: ${c.signature}`; + } + return line; + }) + .join('\n'); + return header + '\n' + body; +} + +function countSymbols(syms: ASTSymbol[]): number { + let n = 0; + for (const s of syms) { + n += 1 + countSymbols(s.children); + } + return n; +} + +/** + * Walks directories to collect source files with known extensions. + * When a `shouldIgnore` callback is provided (backed by FileDiscoveryService), + * .gitignore and .geminiignore patterns are respected during traversal. + * The SKIP_DIRS set provides an additional hard-coded fast path that always + * applies regardless of ignore patterns. + */ +async function collectSourceFiles( + dir: string, + max: number, + shouldIgnore?: (filePath: string) => boolean, +): Promise { + // Collect matching files with a hard cap to prevent OOM in huge repos. + // Entries are sorted at each directory level for deterministic output + // regardless of OS readdir order, then the final list is sorted and sliced. + const HARD_CAP = 10_000; + const files: string[] = []; + async function walk( + d: string, + depth: number, + shouldIgnore?: (filePath: string) => boolean, + ) { + if (depth > 15 || files.length >= HARD_CAP) return; + let entries; + try { + entries = await fs.readdir(d, { withFileTypes: true }); + } catch { + return; + } + // Sort entries alphabetically for deterministic traversal order + entries.sort((a, b) => a.name.localeCompare(b.name)); + for (const e of entries) { + if (files.length >= HARD_CAP) return; + const full = path.join(d, e.name); + if (shouldIgnore && shouldIgnore(full)) { + continue; + } + if ( + e.isDirectory() && + !SKIP_DIRS.has(e.name) && + !e.name.startsWith('.') + ) { + await walk(full, depth + 1, shouldIgnore); + } else if (e.isFile() && LANG_MAP[path.extname(e.name).toLowerCase()]) { + files.push(full); + } + } + } + await walk(dir, 0, shouldIgnore); + // Already mostly sorted due to per-directory sort; final sort ensures + // cross-directory ordering, then slice to requested max. + return files.sort().slice(0, max); +} diff --git a/packages/core/src/tools/ast-search.test.ts b/packages/core/src/tools/ast-search.test.ts new file mode 100644 index 00000000000..4483f193955 --- /dev/null +++ b/packages/core/src/tools/ast-search.test.ts @@ -0,0 +1,224 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import * as fs from 'node:fs/promises'; +import * as path from 'node:path'; +import * as os from 'node:os'; +import { ASTSearchTool } from './ast-search.js'; +import { AST_SEARCH_TOOL_NAME } from './tool-names.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { makeFakeConfig } from '../test-utils/config.js'; + +describe('ASTSearchTool', () => { + let tmpDir: string; + let mockMessageBus: MessageBus; + + beforeEach(async () => { + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'ast-tool-')); + mockMessageBus = { + publish: vi.fn().mockResolvedValue(undefined), + subscribe: vi.fn(), + } as unknown as MessageBus; + }); + + afterEach(async () => { + await fs.rm(tmpDir, { recursive: true, force: true }); + }); + + function makeTool() { + const config = makeFakeConfig({ targetDir: tmpDir }); + return new ASTSearchTool(config, mockMessageBus); + } + + describe('static properties', () => { + it('should have the correct tool name', () => { + expect(ASTSearchTool.Name).toBe(AST_SEARCH_TOOL_NAME); + }); + + it('should produce a valid schema', () => { + const tool = makeTool(); + const schema = tool.getSchema(); + expect(schema.name).toBe(AST_SEARCH_TOOL_NAME); + expect(schema.description).toBeDefined(); + expect(schema.parametersJsonSchema).toBeDefined(); + }); + }); + + describe('symbol scope', () => { + it('should find a class and return its line bounds', async () => { + const content = [ + 'import { X } from "x";', + '', + 'export class TargetClass {', + ' method() {', + ' return 1;', + ' }', + '}', + ].join('\n'); + await fs.writeFile(path.join(tmpDir, 'target.ts'), content); + + const tool = makeTool(); + const invocation = tool.build({ + symbol_name: 'TargetClass', + file_path: 'target.ts', + scope: 'symbol', + }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('TargetClass'); + expect(result.llmContent).toContain('Lines: 3-7'); + expect(result.llmContent).toContain('TIP: Use read_file'); + }); + + it('should find a function', async () => { + const content = [ + 'export function processData(input: string): string {', + ' return input.trim();', + '}', + ].join('\n'); + await fs.writeFile(path.join(tmpDir, 'utils.ts'), content); + + const tool = makeTool(); + const invocation = tool.build({ + symbol_name: 'processData', + file_path: 'utils.ts', + }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('processData'); + expect(result.llmContent).toContain('Lines: 1-3'); + }); + + it('should return a helpful message for missing symbols', async () => { + await fs.writeFile( + path.join(tmpDir, 'empty.ts'), + 'export const x = 1;\n', + ); + + const tool = makeTool(); + const invocation = tool.build({ + symbol_name: 'NonExistent', + file_path: 'empty.ts', + scope: 'symbol', + }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('not found'); + expect(result.llmContent).toContain('grep_search'); + }); + + it('should error when symbol_name is missing', async () => { + await fs.writeFile(path.join(tmpDir, 'f.ts'), 'class A {}\n'); + + const tool = makeTool(); + expect(() => tool.build({ file_path: 'f.ts', scope: 'symbol' })).toThrow( + 'symbol_name', + ); + }); + + it('should error when file_path is missing', async () => { + const tool = makeTool(); + expect(() => tool.build({ symbol_name: 'X', scope: 'symbol' })).toThrow( + 'file_path', + ); + }); + }); + + describe('outline scope', () => { + it('should return a file outline with symbols', async () => { + const content = [ + 'export interface Config {', + ' host: string;', + '}', + '', + 'export class Server {', + ' constructor() {}', + ' public start(): void {', + ' console.log("started");', + ' }', + '}', + '', + 'export function createServer(): Server {', + ' return new Server();', + '}', + ].join('\n'); + await fs.writeFile(path.join(tmpDir, 'server.ts'), content); + + const tool = makeTool(); + const invocation = tool.build({ + file_path: 'server.ts', + scope: 'outline', + }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('server.ts'); + expect(result.llmContent).toContain('interface Config'); + expect(result.llmContent).toContain('class Server'); + expect(result.llmContent).toContain('function createServer'); + }); + + it('should fail gracefully for unsupported file types', async () => { + await fs.writeFile(path.join(tmpDir, 'data.txt'), 'hello'); + + const tool = makeTool(); + const invocation = tool.build({ + file_path: 'data.txt', + scope: 'outline', + }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('Could not outline'); + }); + }); + + describe('map scope', () => { + it('should generate a codebase map', async () => { + const src = path.join(tmpDir, 'src'); + await fs.mkdir(src); + await fs.writeFile(path.join(src, 'a.ts'), 'export class Alpha {}\n'); + await fs.writeFile(path.join(src, 'b.ts'), 'export function beta() {}\n'); + + const tool = makeTool(); + const invocation = tool.build({ scope: 'map' }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('Codebase Map:'); + expect(result.llmContent).toContain('Alpha'); + expect(result.llmContent).toContain('beta'); + }); + + it('should generate a map scoped to a subdirectory', async () => { + const sub = path.join(tmpDir, 'pkg'); + await fs.mkdir(sub); + await fs.writeFile(path.join(sub, 'c.ts'), 'export enum Status { OK }\n'); + await fs.writeFile( + path.join(tmpDir, 'root.ts'), + 'export class Root {}\n', + ); + + const tool = makeTool(); + const invocation = tool.build({ file_path: 'pkg', scope: 'map' }); + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + }); + + expect(result.llmContent).toContain('Status'); + }); + }); +}); diff --git a/packages/core/src/tools/ast-search.ts b/packages/core/src/tools/ast-search.ts new file mode 100644 index 00000000000..6ff73f542d8 --- /dev/null +++ b/packages/core/src/tools/ast-search.ts @@ -0,0 +1,387 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import path from 'node:path'; +import { resolveDefensiveToolPath, resolveToRealPath } from '../utils/paths.js'; +import { ToolErrorType } from './tool-error.js'; +import { + BaseDeclarativeTool, + BaseToolInvocation, + Kind, + type ToolInvocation, + type ToolResult, + type ExecuteOptions, +} from './tools.js'; +import type { Config } from '../config/config.js'; +import { AST_SEARCH_TOOL_NAME, AST_SEARCH_DISPLAY_NAME } from './tool-names.js'; +import { AST_SEARCH_DEFINITION } from './definitions/coreTools.js'; +import { resolveToolDeclaration } from './definitions/resolver.js'; +import { + ASTAnalysisService, + type ASTFileOutline, + type ASTSymbol, +} from '../services/astAnalysisService.js'; +import { debugLogger } from '../utils/debugLogger.js'; +import type { FileDiscoveryService } from '../services/fileDiscoveryService.js'; + +export interface ASTSearchToolParams { + symbol_name?: string; + file_path?: string; + scope?: 'symbol' | 'outline' | 'map'; +} + +class ASTSearchInvocation extends BaseToolInvocation< + ASTSearchToolParams, + ToolResult +> { + constructor( + private config: Config, + params: ASTSearchToolParams, + messageBus: MessageBus, + _toolName?: string, + _toolDisplayName?: string, + ) { + super(params, messageBus, _toolName, _toolDisplayName); + } + + getDescription(): string { + const scope = this.params.scope ?? 'symbol'; + if (scope === 'map') return 'Generating codebase structure map'; + if (scope === 'outline') + return `Outlining ${this.params.file_path ?? '(no file)'}`; + return `Finding symbol "${this.params.symbol_name ?? ''}" in ${this.params.file_path ?? 'workspace'}`; + } + + async execute(_options: ExecuteOptions): Promise { + const scope = this.params.scope ?? 'symbol'; + const targetDir = this.config.getTargetDir(); + + // Use the centralized, cached FileDiscoveryService from config to avoid + // re-reading and re-parsing .gitignore/.geminiignore on every tool call. + const fileFilteringOptions = this.config.getFileFilteringOptions(); + const fileDiscoveryService = this.config.getFileService(); + const shouldIgnore = (filePath: string): boolean => + fileDiscoveryService.shouldIgnoreFile(filePath, fileFilteringOptions); + + const astService = new ASTAnalysisService(targetDir, shouldIgnore); + + try { + // Trim params to prevent whitespace-only values and search mismatches + const filePath = this.params.file_path?.trim(); + const symbolName = this.params.symbol_name?.trim(); + + if (scope === 'map') { + // For map scope, file_path is optional (subdirectory filter) + let safeMapPath: string | undefined; + if (filePath) { + const sanitized = resolveDefensiveToolPath(filePath, targetDir); + let resolved: string; + try { + resolved = resolveToRealPath(path.resolve(targetDir, sanitized)); + } catch { + resolved = path.resolve(targetDir, sanitized); + } + const err = this.config.validatePathAccess(resolved, 'read'); + if (err) { + return { + llmContent: err, + returnDisplay: 'Path not in workspace.', + error: { + message: err, + type: ToolErrorType.PATH_NOT_IN_WORKSPACE, + }, + }; + } + safeMapPath = sanitized; + } + return await this.handleMapScope(astService, safeMapPath); + } + + if (!filePath) { + return { + llmContent: + 'Error: file_path is required for "symbol" and "outline" scopes.', + returnDisplay: 'Missing file_path', + }; + } + + // Validate path stays within workspace boundaries + const sanitizedPath = resolveDefensiveToolPath(filePath, targetDir); + let resolvedPath: string; + try { + resolvedPath = resolveToRealPath( + path.resolve(targetDir, sanitizedPath), + ); + } catch { + resolvedPath = path.resolve(targetDir, sanitizedPath); + } + + const validationError = this.config.validatePathAccess( + resolvedPath, + 'read', + ); + if (validationError) { + return { + llmContent: validationError, + returnDisplay: 'Path not in workspace.', + error: { + message: validationError, + type: ToolErrorType.PATH_NOT_IN_WORKSPACE, + }, + }; + } + + if (scope === 'outline') { + return await this.handleOutlineScope(astService, sanitizedPath); + } + + // Default: symbol scope + if (!symbolName) { + return { + llmContent: 'Error: symbol_name is required for "symbol" scope.', + returnDisplay: 'Missing symbol_name', + }; + } + + return await this.handleSymbolScope( + astService, + sanitizedPath, + symbolName, + ); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + debugLogger.warn('[ASTSearchTool] Error:', msg); + return { + llmContent: `AST search error: ${msg}`, + returnDisplay: 'Error', + }; + } + } + + private async handleSymbolScope( + astService: ASTAnalysisService, + safePath: string, + symbolName: string, + ): Promise { + // Single getFileOutline call to avoid reading and parsing the file twice + // (findSymbolBounds internally calls getFileOutline, so calling both is redundant). + const outline = await astService.getFileOutline(safePath); + + if (!outline) { + return { + llmContent: `Could not search "${safePath}". File may not exist, is ignored, or its language is not supported.`, + returnDisplay: 'File not parsed', + }; + } + + const findSymbolRecursive = ( + symbols: ASTSymbol[], + ): ASTSymbol | undefined => { + for (const s of symbols) { + if (s.name === symbolName) return s; + const found = findSymbolRecursive(s.children); + if (found) return found; + } + return undefined; + }; + const symbol = findSymbolRecursive(outline.symbols); + + if (!symbol) { + return { + llmContent: + `Symbol "${symbolName}" not found in ${safePath}. ` + + 'Try using grep_search for a text-based search, or check the symbol name spelling.', + returnDisplay: 'Symbol not found', + }; + } + + const result = [ + `Found "${symbolName}" in ${safePath}:`, + ` Lines: ${symbol.startLine}-${symbol.endLine} (${symbol.endLine - symbol.startLine + 1} lines)`, + ` Kind: ${symbol.kind}`, + ` Signature: ${symbol.signature}`, + '', + `TIP: Use read_file with start_line=${symbol.startLine} and end_line=${symbol.endLine} to read the exact symbol body.`, + ].join('\n'); + + return { + llmContent: result, + returnDisplay: `${symbolName}: L${symbol.startLine}-${symbol.endLine}`, + display: { + name: AST_SEARCH_DISPLAY_NAME, + description: this.getDescription(), + resultSummary: `L${symbol.startLine}-${symbol.endLine}`, + result: { type: 'text', text: result }, + }, + }; + } + + private async handleOutlineScope( + astService: ASTAnalysisService, + safePath: string, + ): Promise { + const outline = await astService.getFileOutline(safePath); + if (!outline) { + return { + llmContent: `Could not outline "${this.params.file_path}". File may not exist or its language is not supported.`, + returnDisplay: 'Outline failed', + }; + } + + const result = formatOutlineResult(outline); + return { + llmContent: result, + returnDisplay: `${outline.symbols.length} symbols`, + display: { + name: AST_SEARCH_DISPLAY_NAME, + description: this.getDescription(), + resultSummary: `${outline.symbols.length} symbols`, + result: { type: 'text', text: result }, + }, + }; + } + + private async handleMapScope( + astService: ASTAnalysisService, + safePath?: string, + ): Promise { + const map = await astService.getCodebaseMap(safePath); + return { + llmContent: map, + returnDisplay: 'Codebase map generated', + display: { + name: AST_SEARCH_DISPLAY_NAME, + description: this.getDescription(), + result: { type: 'text', text: map.slice(0, 500) + '...' }, + }, + }; + } +} + +function formatOutlineResult(outline: ASTFileOutline): string { + const lines: string[] = []; + lines.push( + `File: ${outline.filePath} (${outline.language}, ${outline.totalLines} lines)`, + ); + lines.push(`Symbols: ${outline.symbols.length} top-level declarations`); + lines.push(''); + + for (const sym of outline.symbols) { + lines.push( + ` ${sym.kind} ${sym.name} [L${sym.startLine}-${sym.endLine}]: ${sym.signature}`, + ); + for (const child of sym.children) { + lines.push( + ` ${child.kind} ${child.name} [L${child.startLine}-${child.endLine}]: ${child.signature}`, + ); + } + } + + return lines.join('\n'); +} + +export class ASTSearchTool extends BaseDeclarativeTool< + ASTSearchToolParams, + ToolResult +> { + static readonly Name = AST_SEARCH_TOOL_NAME; + private readonly fileDiscoveryService: FileDiscoveryService; + + constructor( + private config: Config, + messageBus: MessageBus, + ) { + super( + ASTSearchTool.Name, + AST_SEARCH_DISPLAY_NAME, + AST_SEARCH_DEFINITION.base.description!, + Kind.Search, + AST_SEARCH_DEFINITION.base.parametersJsonSchema, + messageBus, + true, + false, + ); + this.fileDiscoveryService = config.getFileService(); + } + + protected override validateToolParamValues( + params: ASTSearchToolParams, + ): string | null { + const scope = params.scope ?? 'symbol'; + const symbolName = params.symbol_name?.trim(); + const filePath = params.file_path?.trim(); + + if (scope === 'symbol' && (!symbolName || symbolName === '')) { + return "The 'symbol_name' parameter must be non-empty when scope is 'symbol'."; + } + + if ( + (scope === 'symbol' || scope === 'outline') && + (!filePath || filePath === '') + ) { + return "The 'file_path' parameter must be non-empty for 'symbol' and 'outline' scopes."; + } + + if (filePath) { + const sanitizedPath = resolveDefensiveToolPath( + filePath, + this.config.getTargetDir(), + ); + let resolvedPath: string; + try { + resolvedPath = resolveToRealPath( + path.resolve(this.config.getTargetDir(), sanitizedPath), + ); + } catch (err) { + return `Failed to resolve path: ${err instanceof Error ? err.message : String(err)}`; + } + const validationError = this.config.validatePathAccess( + resolvedPath, + 'read', + ); + if (validationError) { + return validationError; + } + + // Enforce .gitignore / .geminiignore patterns, matching ReadFileTool + const fileFilteringOptions = this.config.getFileFilteringOptions(); + if ( + this.fileDiscoveryService.shouldIgnoreFile( + sanitizedPath, + fileFilteringOptions, + ) || + this.fileDiscoveryService.shouldIgnoreFile( + resolvedPath, + fileFilteringOptions, + ) + ) { + return `File path '${resolvedPath}' is ignored by configured ignore patterns.`; + } + } + + return null; + } + + protected createInvocation( + params: ASTSearchToolParams, + messageBus: MessageBus, + _toolName?: string, + _toolDisplayName?: string, + ): ToolInvocation { + return new ASTSearchInvocation( + this.config, + params, + messageBus, + _toolName, + _toolDisplayName, + ); + } + + override getSchema(modelId?: string) { + return resolveToolDeclaration(AST_SEARCH_DEFINITION, modelId); + } +} diff --git a/packages/core/src/tools/definitions/base-declarations.ts b/packages/core/src/tools/definitions/base-declarations.ts index 6c5d45869d8..ff5cb11b2fc 100644 --- a/packages/core/src/tools/definitions/base-declarations.ts +++ b/packages/core/src/tools/definitions/base-declarations.ts @@ -136,3 +136,9 @@ export const COMPLETE_TASK_DISPLAY_NAME = 'Complete Task'; // -- MCP Resources -- export const READ_MCP_RESOURCE_TOOL_NAME = 'read_mcp_resource'; export const LIST_MCP_RESOURCES_TOOL_NAME = 'list_mcp_resources'; + +// -- ast_search (AST-aware structural search) -- +export const AST_SEARCH_TOOL_NAME = 'ast_search'; +export const AST_SEARCH_PARAM_SYMBOL_NAME = 'symbol_name'; +export const AST_SEARCH_PARAM_FILE_PATH = 'file_path'; +export const AST_SEARCH_PARAM_SCOPE = 'scope'; diff --git a/packages/core/src/tools/definitions/coreTools.ts b/packages/core/src/tools/definitions/coreTools.ts index 2e5c0312881..5b0367f8936 100644 --- a/packages/core/src/tools/definitions/coreTools.ts +++ b/packages/core/src/tools/definitions/coreTools.ts @@ -99,6 +99,10 @@ export { TOPIC_PARAM_TITLE, TOPIC_PARAM_SUMMARY, TOPIC_PARAM_STRATEGIC_INTENT, + AST_SEARCH_TOOL_NAME, + AST_SEARCH_PARAM_SYMBOL_NAME, + AST_SEARCH_PARAM_FILE_PATH, + AST_SEARCH_PARAM_SCOPE, } from './base-declarations.js'; // Re-export sets for compatibility @@ -287,3 +291,10 @@ export const LIST_MCP_RESOURCES_DEFINITION: ToolDefinition = { }, overrides: (modelId) => getToolSet(modelId).list_mcp_resources, }; + +export const AST_SEARCH_DEFINITION: ToolDefinition = { + get base() { + return DEFAULT_LEGACY_SET.ast_search; + }, + overrides: (modelId) => getToolSet(modelId).ast_search, +}; diff --git a/packages/core/src/tools/definitions/model-family-sets/default-legacy.ts b/packages/core/src/tools/definitions/model-family-sets/default-legacy.ts index 3dfe8dd40e3..d1a2726c328 100644 --- a/packages/core/src/tools/definitions/model-family-sets/default-legacy.ts +++ b/packages/core/src/tools/definitions/model-family-sets/default-legacy.ts @@ -73,6 +73,10 @@ import { ASK_USER_OPTION_PARAM_LABEL, ASK_USER_OPTION_PARAM_DESCRIPTION, PLAN_MODE_PARAM_REASON, + AST_SEARCH_TOOL_NAME, + AST_SEARCH_PARAM_SYMBOL_NAME, + AST_SEARCH_PARAM_FILE_PATH, + AST_SEARCH_PARAM_SCOPE, } from '../base-declarations.js'; import { getShellDeclaration, @@ -761,4 +765,39 @@ The agent did not use the todo list because this task could be completed by a ti required: [], }, }, + + ast_search: { + name: AST_SEARCH_TOOL_NAME, + description: + 'Searches for a named code symbol (class, function, method, interface, type, enum) ' + + 'and returns its precise line boundaries and structural signature. ' + + 'Use this to read a specific function or class without guessing line numbers. ' + + 'When scope is "outline", returns the full structural outline of the file instead. ' + + 'When scope is "map", returns a compressed map of the codebase showing all top-level symbols.', + parametersJsonSchema: { + type: 'object', + properties: { + [AST_SEARCH_PARAM_SYMBOL_NAME]: { + description: + 'The exact name of the symbol to locate (e.g. "MyClass", "processData"). ' + + 'Required when scope is "symbol" (the default). Ignored when scope is "outline" or "map".', + type: 'string', + }, + [AST_SEARCH_PARAM_FILE_PATH]: { + description: + 'The path to the file to search in. Required for "symbol" and "outline" scopes. ' + + 'Optional for "map" scope (defaults to working directory).', + type: 'string', + }, + [AST_SEARCH_PARAM_SCOPE]: { + description: + 'The type of AST query: "symbol" (default) to find one symbol\'s bounds, ' + + '"outline" to get a file\'s structural skeleton, or "map" to get a compressed codebase overview.', + type: 'string', + enum: ['symbol', 'outline', 'map'], + }, + }, + required: [], + }, + }, }; diff --git a/packages/core/src/tools/definitions/model-family-sets/gemini-3.ts b/packages/core/src/tools/definitions/model-family-sets/gemini-3.ts index 57a897f9eef..65b401cdd6d 100644 --- a/packages/core/src/tools/definitions/model-family-sets/gemini-3.ts +++ b/packages/core/src/tools/definitions/model-family-sets/gemini-3.ts @@ -73,6 +73,10 @@ import { ASK_USER_OPTION_PARAM_LABEL, ASK_USER_OPTION_PARAM_DESCRIPTION, PLAN_MODE_PARAM_REASON, + AST_SEARCH_TOOL_NAME, + AST_SEARCH_PARAM_SYMBOL_NAME, + AST_SEARCH_PARAM_FILE_PATH, + AST_SEARCH_PARAM_SCOPE, } from '../base-declarations.js'; import { getShellDeclaration, @@ -745,4 +749,29 @@ The agent did not use the todo list because this task could be completed by a ti required: [], }, }, + + ast_search: { + name: AST_SEARCH_TOOL_NAME, + description: + 'AST-aware structural search. Finds symbol bounds, file outlines, or codebase maps.', + parametersJsonSchema: { + type: 'object', + properties: { + [AST_SEARCH_PARAM_SYMBOL_NAME]: { + description: 'Symbol name to locate.', + type: 'string', + }, + [AST_SEARCH_PARAM_FILE_PATH]: { + description: 'File to search in.', + type: 'string', + }, + [AST_SEARCH_PARAM_SCOPE]: { + description: '"symbol", "outline", or "map".', + type: 'string', + enum: ['symbol', 'outline', 'map'], + }, + }, + required: [], + }, + }, }; diff --git a/packages/core/src/tools/definitions/types.ts b/packages/core/src/tools/definitions/types.ts index d6f0a723a15..effb9249a00 100644 --- a/packages/core/src/tools/definitions/types.ts +++ b/packages/core/src/tools/definitions/types.ts @@ -52,4 +52,5 @@ export interface CoreToolSet { read_mcp_resource: FunctionDeclaration; list_mcp_resources: FunctionDeclaration; update_topic?: FunctionDeclaration; + ast_search: FunctionDeclaration; } diff --git a/packages/core/src/tools/read-file.test.ts b/packages/core/src/tools/read-file.test.ts index df0bd171c79..0463ed7889f 100644 --- a/packages/core/src/tools/read-file.test.ts +++ b/packages/core/src/tools/read-file.test.ts @@ -338,6 +338,7 @@ describe('ReadFileTool', () => { 'IMPORTANT: The file content has been truncated', ); expect(result.llmContent).toContain('--- FILE CONTENT (truncated) ---'); + expect(result.llmContent).toContain('ast_search'); expect(result.returnDisplay).toContain('some lines were shortened'); }); diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts index 29ed50c61bb..b275ae18425 100644 --- a/packages/core/src/tools/read-file.ts +++ b/packages/core/src/tools/read-file.ts @@ -164,6 +164,7 @@ class ReadFileToolInvocation extends BaseToolInvocation< IMPORTANT: The file content has been truncated. Status: Showing lines ${start}-${end} of ${total} total lines. Action: To read more of the file, you can use the 'start_line' and 'end_line' parameters in a subsequent 'read_file' call. For example, to read the next section of the file, use start_line: ${end + 1}. +TIP: Use 'ast_search' with scope "outline" to see the file's structural skeleton, or with scope "symbol" to jump directly to a specific function or class. --- FILE CONTENT (truncated) --- ${result.llmContent}`; diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 0987f9f3dd0..537bc25bc21 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -82,6 +82,10 @@ import { TOPIC_PARAM_TITLE, TOPIC_PARAM_SUMMARY, TOPIC_PARAM_STRATEGIC_INTENT, + AST_SEARCH_TOOL_NAME, + AST_SEARCH_PARAM_SYMBOL_NAME, + AST_SEARCH_PARAM_FILE_PATH, + AST_SEARCH_PARAM_SCOPE, } from './definitions/coreTools.js'; export { @@ -162,8 +166,15 @@ export { TOPIC_PARAM_TITLE, TOPIC_PARAM_SUMMARY, TOPIC_PARAM_STRATEGIC_INTENT, + AST_SEARCH_TOOL_NAME, + AST_SEARCH_PARAM_SYMBOL_NAME, + AST_SEARCH_PARAM_FILE_PATH, + AST_SEARCH_PARAM_SCOPE, }; +// Tool Display Names (AST) +export const AST_SEARCH_DISPLAY_NAME = 'ASTSearch'; + export const EDIT_TOOL_NAMES = new Set([EDIT_TOOL_NAME, WRITE_FILE_TOOL_NAME]); /** @@ -273,6 +284,7 @@ export const ALL_BUILTIN_TOOL_NAMES = [ AGENT_TOOL_NAME, READ_MCP_RESOURCE_TOOL_NAME, LIST_MCP_RESOURCES_TOOL_NAME, + AST_SEARCH_TOOL_NAME, ] as const; /** @@ -294,6 +306,7 @@ export const PLAN_MODE_TOOLS = [ 'cli_help', READ_MCP_RESOURCE_TOOL_NAME, LIST_MCP_RESOURCES_TOOL_NAME, + AST_SEARCH_TOOL_NAME, ] as const; /**