metona-ai-desktop/electron/harness/tools/built-in/code-search.ts

/**
 * 代码搜索工具（1 个）
 *
 * code_search — 基于 ripgrep 的高速代码搜索
 *
 * 相比 search_files 的纯 JS 实现，code_search 调用 ripgrep 子进程，
 * 性能提升 10-100 倍，支持正则、文件类型过滤、上下文行展示。
 * 适合大型代码库的精准搜索。
 *
 * @see standard/开发规范.md — 优先使用第三方成熟库
 */

import { execFile } from 'child_process';
import { promisify } from 'util';
import log from 'electron-log';
import type { IMetonaTool, ToolExecutionContext } from '../../types/metona-tool';
import type { MetonaToolDef } from '../../../harness/types';
import { MetonaToolCategory, MetonaRiskLevel } from '../../../harness/types';
// F1-2: 统一用 safeResolvePath（含 isPathWithinWorkspace + MEMORY.md 拦截）
// F4-2: 引入 extractErrorMessage 统一错误处理
import { safeResolvePath, extractErrorMessage } from './file-guard';

const execFileAsync = promisify(execFile);

export class CodeSearchTool implements IMetonaTool {
  /** ripgrep 可用性缓存（实例级，便于测试重置） */
  private rgAvailable: boolean | null = null;

  /** 检测系统是否安装了 ripgrep */
  private async checkRipgrep(): Promise<boolean> {
    if (this.rgAvailable !== null) return this.rgAvailable;
    try {
      await execFileAsync('rg', ['--version'], { timeout: 3_000 });
      this.rgAvailable = true;
    } catch {
      this.rgAvailable = false;
    }
    return this.rgAvailable;
  }

  readonly definition: MetonaToolDef = {
    name: 'code_search',
    description: 'Search code using ripgrep. Supports regex patterns, file type filtering, and context lines. Much faster than search_files for large codebases. Falls back to JS implementation if ripgrep is not installed.',
    parameters: {
      type: 'object',
      properties: {
        pattern: { type: 'string', description: 'Regex pattern to search for' },
        path: { type: 'string', description: 'Search directory (default: workspace root)' },
        file_glob: { type: 'string', description: 'File name glob filter (e.g., "*.ts", "*.py")' },
        case_sensitive: { type: 'boolean', description: 'Case sensitive search (default false)' },
        context_before: { type: 'number', description: 'Lines of context before match (default 0, max 5)' },
        context_after: { type: 'number', description: 'Lines of context after match (default 0, max 5)' },
        max_results: { type: 'number', description: 'Maximum results (default 50, max 200)' },
      },
      required: ['pattern'],
    },
    category: MetonaToolCategory.SEARCH,
    riskLevel: MetonaRiskLevel.SAFE,
    requiresPermission: false,
    timeoutMs: 30_000,
  };

  async execute(args: Record<string, unknown>, context: ToolExecutionContext): Promise<unknown> {
    // F4-2: 外层 try-catch 防止 safeResolvePath 抛出异常向上传播
    // （read_file/write_file/list_directory 均有外层 try-catch，code_search 此前缺失）
    try {
      const pattern = args.pattern as string;
      if (typeof pattern !== 'string' || !pattern) {
        return { results: [], count: 0, error: 'Pattern is required and must be a string' };
      }
      if (pattern.length > 500) {
        return { results: [], count: 0, error: 'Pattern too long (max 500 chars)' };
      }

      const searchPath = args.path
        ? safeResolvePath(args.path as string, context.workspacePath)
        : context.workspacePath;
      const fileGlob = args.file_glob as string | undefined;
      const caseSensitive = (args.case_sensitive as boolean) ?? false;
      const contextBefore = Math.min(5, Math.max(0, (args.context_before as number) ?? 0));
      const contextAfter = Math.min(5, Math.max(0, (args.context_after as number) ?? 0));
      const maxResults = Math.min(200, Math.max(1, (args.max_results as number) ?? 50));

      const opts = { fileGlob, caseSensitive, contextBefore, contextAfter, maxResults };

      // 优先使用 ripgrep，回退到 JS 实现
      const hasRg = await this.checkRipgrep();
      if (hasRg) {
        return this.searchWithRipgrep(pattern, searchPath, opts, context);
      }
      return this.searchWithJs(pattern, searchPath, opts, context);
    } catch (error) {
      return { results: [], count: 0, error: extractErrorMessage(error), success: false };
    }
  }

  /** 使用 ripgrep 子进程搜索 */
  private async searchWithRipgrep(
    pattern: string,
    searchPath: string,
    opts: { fileGlob?: string; caseSensitive: boolean; contextBefore: number; contextAfter: number; maxResults: number },
    context: ToolExecutionContext,
  ): Promise<unknown> {
    const rgArgs: string[] = ['--json'];

    if (!opts.caseSensitive) rgArgs.push('-i');
    if (opts.contextBefore > 0) rgArgs.push('-B', String(opts.contextBefore));
    if (opts.contextAfter > 0) rgArgs.push('-A', String(opts.contextAfter));
    rgArgs.push('-g', '!MEMORY.md');
    if (opts.fileGlob) rgArgs.push('-g', opts.fileGlob);

    // #22 修复: 在 pattern 前加 -- 终止选项解析，防止以 - 开头的 pattern 被解释为选项
    // 攻击场景：pattern="--help" 输出帮助而非搜索，pattern="--ignore-file /etc/passwd" 可读取任意文件，
    // pattern="--files" 列出所有文件。execFile 已用数组参数防 shell 注入，但 ripgrep 自身选项解析仍需防护
    rgArgs.push('--', pattern, searchPath);

    try {
      const { stdout } = await execFileAsync('rg', rgArgs, {
        maxBuffer: 10 * 1024 * 1024,
        timeout: 25_000,
      });

      const results = this.parseRipgrepJsonOutput(stdout);
      return { results: results.slice(0, opts.maxResults), count: results.length, engine: 'ripgrep' };
    } catch (error) {
      const err = error as { code?: number; signal?: string; stdout?: string; stderr?: string; killed?: boolean; message?: string };
      // rg 退出码 1 = 无匹配，不是错误
      if (err.code === 1) {
        return { results: [], count: 0, engine: 'ripgrep' };
      }
      // 超时被 kill
      if (err.killed || err.signal === 'SIGTERM') {
        return { results: [], count: 0, error: 'ripgrep search timed out', engine: 'ripgrep' };
      }
      // 其他错误回退到 JS
      log.warn('[CodeSearch] ripgrep failed, falling back to JS:', err.stderr || err.message);
      return this.searchWithJs(pattern, searchPath, opts, context);
    }
  }

  /** 解析 ripgrep --json 输出 */
  /** @visibleForTesting 纯函数，供单元测试直接断言 ripgrep JSON 状态机 */
  parseRipgrepJsonOutput(output: string): Array<{
    path: string;
    line: number;
    column: number;
    match: string;
    before?: string[];
    after?: string[];
  }> {
    const results: Array<{ path: string; line: number; column: number; match: string; before?: string[]; after?: string[] }> = [];
    const lines = output.split('\n').filter((l) => l.trim());

    let currentMatch: { path: string; line: number; column: number; match: string; before?: string[]; after?: string[] } | null = null;
    let beforeBuffer: string[] = [];
    let afterBuffer: string[] = [];

    for (const line of lines) {
      let entry: Record<string, unknown>;
      try {
        entry = JSON.parse(line);
      } catch {
        continue;
      }

      const type = entry.type as string;
      const data = entry.data as Record<string, unknown>;

      if (type === 'context') {
        const text = (data.lines as { text?: string } | undefined)?.text ?? '';

        if (currentMatch) {
          // 当前有 match，这是 after context
          afterBuffer.push(text);
        } else {
          // 当前无 match，这是 before context
          beforeBuffer.push(text);
        }
      } else if (type === 'match') {
        // 新 match：先保存上一个 match 的 after context
        if (currentMatch) {
          if (afterBuffer.length > 0) currentMatch.after = [...afterBuffer];
          results.push(currentMatch);
          afterBuffer = [];
        }

        const text = (data.lines as { text?: string } | undefined)?.text ?? '';
        const submatches = (data.submatches as Array<{ match: { text?: string }; start?: number }> | undefined) ?? [];
        const matchText = submatches[0]?.match?.text ?? text;
        const column = (submatches[0]?.start ?? 0) + 1;

        currentMatch = {
          path: (data.path as { text?: string } | undefined)?.text ?? '',
          line: data.line_number as number,
          column,
          match: matchText,
          before: beforeBuffer.length > 0 ? [...beforeBuffer] : undefined,
        };
        beforeBuffer = [];
        afterBuffer = [];
      }
    }

    // 保存最后一个 match
    if (currentMatch) {
      if (afterBuffer.length > 0) currentMatch.after = [...afterBuffer];
      results.push(currentMatch);
    }

    return results;
  }

  /** JS 回退实现 */
  private async searchWithJs(
    pattern: string,
    searchPath: string,
    opts: { fileGlob?: string; caseSensitive: boolean; contextBefore: number; contextAfter: number; maxResults: number },
    context: ToolExecutionContext,
  ): Promise<unknown> {
    // 动态导入以避免循环依赖
    const { SearchFilesTool } = await import('./filesystem');
    const searchTool = new SearchFilesTool();
    return searchTool.execute({
      pattern,
      target: 'content',
      path: searchPath,
      file_glob: opts.fileGlob,
      limit: opts.maxResults,
    }, context);
  }
}