tiny-trae 0.1.0

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

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

pub struct ReadFileTool;

#[async_trait::async_trait]
impl Tool for ReadFileTool {
    fn get_definition(&self, name: &str) -> ToolDefinition {
        ToolDefinition {
            r#type: "function".to_string(),
            function: FunctionDefinition {
                name: name.to_string(),
                description: "Read the contents of a file".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "The path to the file to read"
                        }
                    },
                    "required": ["path"]
                }),
            },
        }
    }

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

        let file_path = Path::new(&config.workspace.root_path).join(path);
        
        match get_file_content(&file_path) {
            Ok(content) => {
                info!("Successfully read file: {}", file_path.display());
                Ok(ToolResult::success(content))
            }
            Err(e) => {
                debug!("Failed to read file: {}", e);
                Ok(ToolResult::error(format!("Failed to read file: {}", e)))
            }
        }
    }
}

pub struct ListFilesTool;

#[async_trait::async_trait]
impl Tool for ListFilesTool {
    fn get_definition(&self, name: &str) -> ToolDefinition {
        ToolDefinition {
            r#type: "function".to_string(),
            function: FunctionDefinition {
                name: name.to_string(),
                description: "List files in a directory".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "The path to the directory to list (optional, defaults to workspace root)"
                        }
                    }
                }),
            },
        }
    }

    async fn execute(&self, input: Value, config: &Config) -> Result<ToolResult> {
        let path = input
            .get("path")
            .and_then(|v| v.as_str())
            .unwrap_or(".");

        let dir_path = Path::new(&config.workspace.root_path).join(path);
        
        match list_directory(&dir_path, &config.workspace.ignore_patterns) {
            Ok(files) => {
                info!("Successfully listed directory: {}", dir_path.display());
                let output = files.join("\n");
                Ok(ToolResult::success(output))
            }
            Err(e) => {
                debug!("Failed to list directory: {}", e);
                Ok(ToolResult::error(format!("Failed to list directory: {}", e)))
            }
        }
    }
}

pub struct EditFileTool;

#[async_trait::async_trait]
impl Tool for EditFileTool {
    fn get_definition(&self, name: &str) -> ToolDefinition {
        ToolDefinition {
            r#type: "function".to_string(),
            function: FunctionDefinition {
                name: name.to_string(),
                description: "Edit a file by replacing its contents".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "The path to the file to edit"
                        },
                        "content": {
                            "type": "string",
                            "description": "The new content for the file"
                        }
                    },
                    "required": ["path", "content"]
                }),
            },
        }
    }

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

        let content = input
            .get("content")
            .and_then(|v| v.as_str())
            .context("Missing or invalid 'content' parameter")?;

        let file_path = Path::new(&config.workspace.root_path).join(path);
        
        // Create parent directories if they don't exist
        if let Some(parent) = file_path.parent() {
            if !parent.exists() {
                std::fs::create_dir_all(parent)
                    .with_context(|| format!("Failed to create directory: {}", parent.display()))?;
            }
        }

        match write_file_content(&file_path, content) {
            Ok(_) => {
                info!("Successfully wrote file: {}", file_path.display());
                Ok(ToolResult::success(format!("File written successfully: {}", file_path.display())))
            }
            Err(e) => {
                debug!("Failed to write file: {}", e);
                Ok(ToolResult::error(format!("Failed to write file: {}", e)))
            }
        }
    }
}