feat(agent): add AST-aware structural search tool for precise symbol navigation - #29396

export class ASTSearchTool extends BaseDeclarativeTool<
  ASTSearchToolParams,
  ToolResult
> {
  static readonly Name = AST_SEARCH_TOOL_NAME;

  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,
    );
  }

  protected override validateToolParamValues(
    params: ASTSearchToolParams,
  ): string | null {
    if (params.file_path) {
      const sanitizedPath = resolveDefensiveToolPath(
        params.file_path,
        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;
      }
    }
    return null;
  }

  protected createInvocation(
    params: ASTSearchToolParams,
    messageBus: MessageBus,
    _toolName?: string,
    _toolDisplayName?: string,
  ): ToolInvocation<ASTSearchToolParams, ToolResult> {
    return new ASTSearchInvocation(
      this.config,
      params,
      messageBus,
      _toolName,
      _toolDisplayName,
    );
  }

  override getSchema(modelId?: string) {
    return resolveToolDeclaration(AST_SEARCH_DEFINITION, modelId);
  }
}

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,
} from '../services/astAnalysisService.js';
import { debugLogger } from '../utils/debugLogger.js';

export function findClosingBrace(lines: string[], startLine: number): number {
  let depth = 0;
  let parenDepth = 0;
  let opened = false;
  for (let i = startLine; i < lines.length; i++) {
    const lineWithoutComments = lines[i].split('//')[0].split('/*')[0];
    const stripped = lineWithoutComments.replace(/'[^']*'|"[^"]*"|\x60[^\x60]*\x60/g, '');
    for (const ch of stripped) {
      if (ch === '(') {
        parenDepth++;
      } else if (ch === ')') {
        parenDepth = Math.max(0, parenDepth - 1);
      } else if (ch === '{') {
        if (parenDepth === 0) {
          depth++;
          opened = true;
        }
      } else if (ch === '}') {
        if (parenDepth === 0) {
          depth--;
          if (opened && depth === 0) return i;
        }
      }
    }
  }
  return Math.min(startLine + 50, lines.length - 1);
}

export function extractSymbols(lines: string[], language: string): ASTSymbol[] {
  const symbols: ASTSymbol[] = [];
  const patterns = getDeclarationPatterns(language);
  let inBlockComment = false;

  for (let i = 0; i < lines.length; i++) {
    let trimmed = lines[i].trim();

    if (inBlockComment) {
      if (trimmed.includes('*/')) {
        inBlockComment = false;
        trimmed = trimmed.slice(trimmed.indexOf('*/') + 2).trim();
      } else {
        continue;
      }
    }
    if (trimmed.includes('/*')) {
      if (trimmed.includes('*/')) {
        trimmed = trimmed.replace(/\/\*[\s\S]*?\*\//g, '').trim();
      } else {
        inBlockComment = true;
        trimmed = trimmed.slice(0, trimmed.indexOf('/*')).trim();
      }
    }

    if (
      trimmed === '' ||
      trimmed.startsWith('//') ||
      trimmed.startsWith('#') ||
      trimmed.startsWith('import ') ||
      trimmed.startsWith('from ')
    ) {
      continue;
    }

export function findIndentEnd(lines: string[], startLine: number): number {
  const baseIndent = 
    lines[startLine].length - lines[startLine].trimStart().length;
  let last = startLine;
  let inTripleQuotes = false;

  for (let i = startLine + 1; i < lines.length; i++) {
    const trimmed = lines[i].trim();
    if (trimmed === '') continue;

    // Track Python triple quotes
    const tripleQuoteMatches = (trimmed.match(/"""|'''/g) || []).length;
    if (tripleQuoteMatches % 2 !== 0) {
      inTripleQuotes = !inTripleQuotes;
    }

    if (inTripleQuotes) {
      last = i;
      continue;
    }

    const indent = lines[i].length - lines[i].trimStart().length;
    if (indent <= baseIndent) return last;
    last = i;
  }
  return last;
}

  async execute(_options: ExecuteOptions): Promise<ToolResult> {
    const scope = this.params.scope ?? 'symbol';
    const astService = new ASTAnalysisService(this.config.getTargetDir());

    if (this.params.file_path) {
      const resolvedPath = resolveToRealPath(this.config.getTargetDir(), this.params.file_path);
      const validationError = this.config.validatePathAccess(resolvedPath, 'read');
      if (validationError) {
        return {
          llmContent: validationError,
          returnDisplay: 'Path not in workspace.',
        };
      }
    }

  protected override validateToolParamValues(
    params: ASTSearchToolParams,
  ): string | null {
    if (params.file_path !== undefined) {
      const trimmedPath = params.file_path.trim();
      if (trimmedPath === "") {
        return "File path cannot be empty";
      }
      const resolvedPath = resolveToRealPath(this.config.getTargetDir(), trimmedPath);
      const validationError = this.config.validatePathAccess(resolvedPath, "read");
      if (validationError) {
        return validationError;
      }
    }
    return null;
  }

  async execute(_options: ExecuteOptions): Promise<ToolResult> {
    const scope = this.params.scope ?? 'symbol';
    const targetDir = this.config.getTargetDir();
    const astService = new ASTAnalysisService(targetDir);

    try {
      let sanitizedPath = '';
      if (this.params.file_path) {
        sanitizedPath = resolveDefensiveToolPath(
          this.params.file_path,
          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 === 'map') {
        return await this.handleMapScope(astService, sanitizedPath || undefined);
      }

      if (!this.params.file_path) {
        return {
          llmContent:
            'Error: file_path is required for "symbol" and "outline" scopes.',
          returnDisplay: 'Missing file_path',
        };
      }

      if (scope === 'outline') {
        return await this.handleOutlineScope(astService, sanitizedPath);
      }

      // Default: symbol scope
      if (!this.params.symbol_name) {
        return {
          llmContent: 'Error: symbol_name is required for "symbol" scope.',
          returnDisplay: 'Missing symbol_name',
        };
      }

      return await this.handleSymbolScope(astService, sanitizedPath);
    } 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,
    sanitizedPath: string,
  ): Promise<ToolResult> {
    const bounds = await astService.findSymbolBounds(
      sanitizedPath,
      this.params.symbol_name!,
    );

    if (!bounds) {
      return {
        llmContent:
          'Symbol "' + this.params.symbol_name + '" not found in ' + this.params.file_path + '. ' +
          'Try using grep_search for a text-based search, or check the symbol name spelling.',
        returnDisplay: 'Symbol not found',
      };
    }

    const outline = await astService.getFileOutline(sanitizedPath);
    const symbol = outline?.symbols
      .flatMap((s) => [s, ...s.children])
      .find((s) => s.name === this.params.symbol_name);

    const result = [
      'Found "' + this.params.symbol_name + '" in ' + this.params.file_path + ':',
      '  Lines: ' + bounds.startLine + '-' + bounds.endLine + ' (' + (bounds.endLine - bounds.startLine + 1) + ' lines)',
      symbol ? '  Kind: ' + symbol.kind : '',
      symbol ? '  Signature: ' + symbol.signature : '',
      '',
      'TIP: Use read_file with start_line=' + bounds.startLine + ' and end_line=' + bounds.endLine + ' to read the exact symbol body.',
    ]
      .filter(Boolean)
      .join('\n');

    return {
      llmContent: result,
      returnDisplay: this.params.symbol_name + ': L' + bounds.startLine + '-' + bounds.endLine,
      display: {
        name: AST_SEARCH_DISPLAY_NAME,
        description: this.getDescription(),
        resultSummary: 'L' + bounds.startLine + '-' + bounds.endLine,
        result: { type: 'text', text: result },
      },
    };
  }

  private async handleOutlineScope(
    astService: ASTAnalysisService,
    sanitizedPath: string,
  ): Promise<ToolResult> {
    const outline = await astService.getFileOutline(sanitizedPath);
    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,
    sanitizedPath?: string,
  ): Promise<ToolResult> {
    const map = await astService.getCodebaseMap(sanitizedPath);
    return {
      llmContent: map,
      returnDisplay: 'Codebase map generated',
      display: {
        name: AST_SEARCH_DISPLAY_NAME,
        description: this.getDescription(),
        result: { type: 'text', text: map.slice(0, 500) + '...' },
      },
    };
  }

export function findClosingBrace(lines: string[], startLine: number): number {
  let depth = 0;
  let parenDepth = 0;
  let opened = false;
  for (let i = startLine; i < lines.length; i++) {
    // Strip string literals and template literals to avoid false brace matches.
    const backtickRegex = "'[^']*'|\"[^\"]*\"|" + String.fromCharCode(96) + "[^" + String.fromCharCode(96) + "]*" + String.fromCharCode(96);
    let stripped = lines[i].replace(new RegExp(backtickRegex, "g"), "");
    
    // Strip single-line comments
    const commentIndex = stripped.indexOf("//");
    if (commentIndex !== -1) {
      stripped = stripped.slice(0, commentIndex);
    }

    for (const ch of stripped) {
      if (ch === "(") parenDepth++;
      else if (ch === ")") parenDepth--;
      // 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 end of file rather than an arbitrary offset
  return lines.length - 1;
}

export function findIndentEnd(lines: string[], startLine: number): number {
  const baseIndent =
    lines[startLine].length - lines[startLine].trimStart().length;
  let last = startLine;
  let inTripleQuote = false;
  for (let i = startLine + 1; i < lines.length; i++) {
    const trimmed = lines[i].trim();

    // Track Python triple-quoted strings which can have arbitrary indentation
    const tripleCount = (trimmed.match(/"""|'''/g) || []).length;
    if (tripleCount % 2 !== 0) {
      inTripleQuote = !inTripleQuote;
    }
    if (inTripleQuote) {
      last = i;
      continue;
    }

    if (trimmed === "" || trimmed.startsWith("#") || trimmed.startsWith("//")) continue;
    const indent = lines[i].length - lines[i].trimStart().length;
    if (indent <= baseIndent) return last;
    last = i;
  }
  return last;
}