tiny-trae 0.1.0

An AI coding assistant with tool integration
use anyhow::{Context, Result};
use regex::Regex;
use serde_json::{json, Value};
use std::path::Path;
use std::process::Command;
use log::{debug, info};

use crate::config::Config;
use crate::claude::{ToolDefinition, FunctionDefinition};
use super::{Tool, ToolResult, get_file_content, is_ignored_path};

pub struct SearchTextTool;

#[async_trait::async_trait]
impl Tool for SearchTextTool {
    fn get_definition(&self, name: &str) -> ToolDefinition {
        ToolDefinition {
            r#type: "function".to_string(),
            function: FunctionDefinition {
                name: name.to_string(),
                description: "Search for text patterns in files".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "pattern": {
                            "type": "string",
                            "description": "The text pattern to search for"
                        },
                        "path": {
                            "type": "string",
                            "description": "The path to search in (optional, defaults to workspace root)"
                        },
                        "case_sensitive": {
                            "type": "boolean",
                            "description": "Whether the search should be case sensitive"
                        }
                    },
                    "required": ["pattern"]
                }),
            },
        }
    }

    async fn execute(&self, input: Value, config: &Config) -> Result<ToolResult> {
        let pattern = input
            .get("pattern")
            .and_then(|v| v.as_str())
            .context("Missing or invalid 'pattern' parameter")?;

        let path = input
            .get("path")
            .and_then(|v| v.as_str())
            .unwrap_or(".");

        let case_sensitive = input
            .get("case_sensitive")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let search_path = Path::new(&config.workspace.root_path).join(path);
        
        let regex = if case_sensitive {
            Regex::new(pattern)
        } else {
            Regex::new(&format!("(?i){}", pattern))
        };

        let regex = match regex {
            Ok(r) => r,
            Err(e) => {
                return Ok(ToolResult::error(format!("Invalid regex pattern: {}", e)));
            }
        };

        let mut results = Vec::new();
        
        if search_path.is_file() {
            if let Ok(content) = get_file_content(&search_path) {
                let matches = find_matches(&content, &regex, &search_path.to_string_lossy());
                results.extend(matches);
            }
        } else if search_path.is_dir() {
            let files = walkdir::WalkDir::new(&search_path)
                .into_iter()
                .filter_map(|e| e.ok())
                .filter(|e| e.path().is_file())
                .filter(|e| !is_ignored_path(e.path(), &config.workspace.ignore_patterns));

            for entry in files {
                let file_path = entry.path();
                if let Ok(content) = get_file_content(file_path) {
                    let matches = find_matches(&content, &regex, &file_path.to_string_lossy());
                    results.extend(matches);
                }
            }
        }

        let output = if results.is_empty() {
            "No matches found".to_string()
        } else {
            results.join("\n")
        };

        info!("Search completed, found {} matches", results.len());
        Ok(ToolResult::success(output))
    }
}

pub struct RipgrepTool;

#[async_trait::async_trait]
impl Tool for RipgrepTool {
    fn get_definition(&self, name: &str) -> ToolDefinition {
        ToolDefinition {
            r#type: "function".to_string(),
            function: FunctionDefinition {
                name: name.to_string(),
                description: "Use ripgrep for fast text searching".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "pattern": {
                            "type": "string",
                            "description": "The pattern to search for"
                        },
                        "path": {
                            "type": "string",
                            "description": "The path to search in (optional, defaults to workspace root)"
                        },
                        "case_sensitive": {
                            "type": "boolean",
                            "description": "Whether the search should be case sensitive"
                        },
                        "file_type": {
                            "type": "string",
                            "description": "File type to search (e.g., 'rs', 'py', 'js')"
                        }
                    },
                    "required": ["pattern"]
                }),
            },
        }
    }

    async fn execute(&self, input: Value, config: &Config) -> Result<ToolResult> {
        let pattern = input
            .get("pattern")
            .and_then(|v| v.as_str())
            .context("Missing or invalid 'pattern' parameter")?;

        let path = input
            .get("path")
            .and_then(|v| v.as_str())
            .unwrap_or(".");

        let case_sensitive = input
            .get("case_sensitive")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let file_type = input
            .get("file_type")
            .and_then(|v| v.as_str());

        let search_path = Path::new(&config.workspace.root_path).join(path);
        
        let mut cmd = Command::new("rg");
        cmd.arg("--line-number")
           .arg("--with-filename")
           .arg("--no-heading");

        if !case_sensitive {
            cmd.arg("--ignore-case");
        }

        if let Some(ft) = file_type {
            cmd.arg("--type").arg(ft);
        }

        // Add ignore patterns
        for pattern in &config.workspace.ignore_patterns {
            cmd.arg("--glob").arg(format!("!{}", pattern));
        }

        cmd.arg(pattern).arg(&search_path);

        debug!("Running ripgrep command: {:?}", cmd);

        match cmd.output() {
            Ok(output) => {
                let stdout = String::from_utf8_lossy(&output.stdout);
                let stderr = String::from_utf8_lossy(&output.stderr);
                
                if output.status.success() {
                    info!("Ripgrep search completed successfully");
                    Ok(ToolResult::success(stdout.to_string()))
                } else {
                    debug!("Ripgrep failed: {}", stderr);
                    Ok(ToolResult::error(format!("Ripgrep failed: {}", stderr)))
                }
            }
            Err(e) => {
                debug!("Failed to run ripgrep: {}", e);
                Ok(ToolResult::error(format!("Failed to run ripgrep: {}", e)))
            }
        }
    }
}

fn find_matches(content: &str, regex: &Regex, file_path: &str) -> Vec<String> {
    let mut matches = Vec::new();
    
    for (line_num, line) in content.lines().enumerate() {
        if regex.is_match(line) {
            matches.push(format!("{}:{}:{}", file_path, line_num + 1, line));
        }
    }
    
    matches
}