perspt_agent/

tools.rs

//! Agent Tooling
//!
//! Tools available to agents for interacting with the workspace.
//! Implements: read_file, search_code, apply_patch, run_command

use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::Command as AsyncCommand;

/// Tool result from agent execution
#[derive(Debug, Clone)]
pub struct ToolResult {
    pub tool_name: String,
    pub success: bool,
    pub output: String,
    pub error: Option<String>,
}

impl ToolResult {
    pub fn success(tool_name: &str, output: String) -> Self {
        Self {
            tool_name: tool_name.to_string(),
            success: true,
            output,
            error: None,
        }
    }

    pub fn failure(tool_name: &str, error: String) -> Self {
        Self {
            tool_name: tool_name.to_string(),
            success: false,
            output: String::new(),
            error: Some(error),
        }
    }
}

/// Tool call request from LLM
#[derive(Debug, Clone)]
pub struct ToolCall {
    pub name: String,
    pub arguments: HashMap<String, String>,
}

/// Agent tools for workspace interaction
pub struct AgentTools {
    /// Working directory (sandbox root)
    working_dir: PathBuf,
    /// Whether to require user approval for commands
    require_approval: bool,
    /// Event sender for streaming output
    event_sender: Option<perspt_core::events::channel::EventSender>,
}

impl AgentTools {
    /// Create new agent tools instance
    pub fn new(working_dir: PathBuf, require_approval: bool) -> Self {
        Self {
            working_dir,
            require_approval,
            event_sender: None,
        }
    }

    /// Set event sender for streaming output
    pub fn set_event_sender(&mut self, sender: perspt_core::events::channel::EventSender) {
        self.event_sender = Some(sender);
    }

    /// Execute a tool call
    pub async fn execute(&self, call: &ToolCall) -> ToolResult {
        match call.name.as_str() {
            "read_file" => self.read_file(call),
            "search_code" => self.search_code(call),
            "apply_patch" => self.apply_patch(call),
            "run_command" => self.run_command(call).await,
            "list_files" => self.list_files(call),
            "write_file" => self.write_file(call),
            // Power Tools (OS-level)
            "sed_replace" => self.sed_replace(call),
            "awk_filter" => self.awk_filter(call),
            "diff_files" => self.diff_files(call),
            _ => ToolResult::failure(&call.name, format!("Unknown tool: {}", call.name)),
        }
    }

    /// Read a file's contents
    fn read_file(&self, call: &ToolCall) -> ToolResult {
        let path = match call.arguments.get("path") {
            Some(p) => self.resolve_path(p),
            None => return ToolResult::failure("read_file", "Missing 'path' argument".to_string()),
        };

        match fs::read_to_string(&path) {
            Ok(content) => ToolResult::success("read_file", content),
            Err(e) => ToolResult::failure("read_file", format!("Failed to read {:?}: {}", path, e)),
        }
    }

    /// Search for code patterns using grep
    fn search_code(&self, call: &ToolCall) -> ToolResult {
        let query = match call.arguments.get("query") {
            Some(q) => q,
            None => {
                return ToolResult::failure("search_code", "Missing 'query' argument".to_string())
            }
        };

        let path = call
            .arguments
            .get("path")
            .map(|p| self.resolve_path(p))
            .unwrap_or_else(|| self.working_dir.clone());

        // Use ripgrep if available, fallback to grep
        let output = Command::new("rg")
            .args(["--json", "-n", query])
            .current_dir(&path)
            .output()
            .or_else(|_| {
                Command::new("grep")
                    .args(["-rn", query, "."])
                    .current_dir(&path)
                    .output()
            });

        match output {
            Ok(out) => {
                let stdout = String::from_utf8_lossy(&out.stdout).to_string();
                ToolResult::success("search_code", stdout)
            }
            Err(e) => ToolResult::failure("search_code", format!("Search failed: {}", e)),
        }
    }

    /// Apply a patch to a file
    fn apply_patch(&self, call: &ToolCall) -> ToolResult {
        let path = match call.arguments.get("path") {
            Some(p) => self.resolve_path(p),
            None => {
                return ToolResult::failure("apply_patch", "Missing 'path' argument".to_string())
            }
        };

        let content = match call.arguments.get("content") {
            Some(c) => c,
            None => {
                return ToolResult::failure("apply_patch", "Missing 'content' argument".to_string())
            }
        };

        // Create parent directories if needed
        if let Some(parent) = path.parent() {
            if let Err(e) = fs::create_dir_all(parent) {
                return ToolResult::failure(
                    "apply_patch",
                    format!("Failed to create directories: {}", e),
                );
            }
        }

        match fs::write(&path, content) {
            Ok(_) => ToolResult::success("apply_patch", format!("Successfully wrote {:?}", path)),
            Err(e) => {
                ToolResult::failure("apply_patch", format!("Failed to write {:?}: {}", path, e))
            }
        }
    }

    /// Run a shell command (requires approval unless auto-approve is set)
    async fn run_command(&self, call: &ToolCall) -> ToolResult {
        let cmd_str = match call.arguments.get("command") {
            Some(c) => c,
            None => {
                return ToolResult::failure("run_command", "Missing 'command' argument".to_string())
            }
        };

        if self.require_approval {
            log::info!("Command requires approval: {}", cmd_str);
        }

        let mut child = match AsyncCommand::new("sh")
            .args(["-c", cmd_str])
            .current_dir(&self.working_dir)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
        {
            Ok(child) => child,
            Err(e) => return ToolResult::failure("run_command", format!("Failed to spawn: {}", e)),
        };

        let stdout = child.stdout.take().expect("Failed to open stdout");
        let stderr = child.stderr.take().expect("Failed to open stderr");
        let sender = self.event_sender.clone();

        let stdout_handle = tokio::spawn(async move {
            let mut reader = BufReader::new(stdout).lines();
            let mut output = String::new();
            while let Ok(Some(line)) = reader.next_line().await {
                if let Some(ref s) = sender {
                    let _ = s.send(perspt_core::AgentEvent::Log(line.clone()));
                }
                output.push_str(&line);
                output.push('\n');
            }
            output
        });

        let sender_err = self.event_sender.clone();
        let stderr_handle = tokio::spawn(async move {
            let mut reader = BufReader::new(stderr).lines();
            let mut output = String::new();
            while let Ok(Some(line)) = reader.next_line().await {
                if let Some(ref s) = sender_err {
                    let _ = s.send(perspt_core::AgentEvent::Log(format!("ERR: {}", line)));
                }
                output.push_str(&line);
                output.push('\n');
            }
            output
        });

        let status = match child.wait().await {
            Ok(s) => s,
            Err(e) => return ToolResult::failure("run_command", format!("Failed to wait: {}", e)),
        };

        let stdout_str = stdout_handle.await.unwrap_or_default();
        let stderr_str = stderr_handle.await.unwrap_or_default();

        if status.success() {
            ToolResult::success("run_command", stdout_str)
        } else {
            ToolResult::failure(
                "run_command",
                format!("Exit code: {:?}\n{}", status.code(), stderr_str),
            )
        }
    }

    /// List files in a directory
    fn list_files(&self, call: &ToolCall) -> ToolResult {
        let path = call
            .arguments
            .get("path")
            .map(|p| self.resolve_path(p))
            .unwrap_or_else(|| self.working_dir.clone());

        match fs::read_dir(&path) {
            Ok(entries) => {
                let files: Vec<String> = entries
                    .filter_map(|e| e.ok())
                    .map(|e| {
                        let name = e.file_name().to_string_lossy().to_string();
                        if e.file_type().map(|t| t.is_dir()).unwrap_or(false) {
                            format!("{}/", name)
                        } else {
                            name
                        }
                    })
                    .collect();
                ToolResult::success("list_files", files.join("\n"))
            }
            Err(e) => {
                ToolResult::failure("list_files", format!("Failed to list {:?}: {}", path, e))
            }
        }
    }

    /// Write content to a file
    fn write_file(&self, call: &ToolCall) -> ToolResult {
        // Alias for apply_patch with different semantics
        self.apply_patch(call)
    }

    /// Resolve a path relative to working directory
    fn resolve_path(&self, path: &str) -> PathBuf {
        let p = Path::new(path);
        if p.is_absolute() {
            p.to_path_buf()
        } else {
            self.working_dir.join(p)
        }
    }

    // =========================================================================
    // Power Tools (OS-level operations)
    // =========================================================================

    /// Replace text in a file using sed-like pattern matching
    fn sed_replace(&self, call: &ToolCall) -> ToolResult {
        let path = match call.arguments.get("path") {
            Some(p) => self.resolve_path(p),
            None => {
                return ToolResult::failure("sed_replace", "Missing 'path' argument".to_string())
            }
        };

        let pattern = match call.arguments.get("pattern") {
            Some(p) => p,
            None => {
                return ToolResult::failure("sed_replace", "Missing 'pattern' argument".to_string())
            }
        };

        let replacement = match call.arguments.get("replacement") {
            Some(r) => r,
            None => {
                return ToolResult::failure(
                    "sed_replace",
                    "Missing 'replacement' argument".to_string(),
                )
            }
        };

        // Read file, perform replacement, write back
        match fs::read_to_string(&path) {
            Ok(content) => {
                let new_content = content.replace(pattern, replacement);
                match fs::write(&path, &new_content) {
                    Ok(_) => ToolResult::success(
                        "sed_replace",
                        format!(
                            "Replaced '{}' with '{}' in {:?}",
                            pattern, replacement, path
                        ),
                    ),
                    Err(e) => ToolResult::failure("sed_replace", format!("Failed to write: {}", e)),
                }
            }
            Err(e) => {
                ToolResult::failure("sed_replace", format!("Failed to read {:?}: {}", path, e))
            }
        }
    }

    /// Filter file content using awk-like field selection
    fn awk_filter(&self, call: &ToolCall) -> ToolResult {
        let path = match call.arguments.get("path") {
            Some(p) => self.resolve_path(p),
            None => {
                return ToolResult::failure("awk_filter", "Missing 'path' argument".to_string())
            }
        };

        let filter = match call.arguments.get("filter") {
            Some(f) => f,
            None => {
                return ToolResult::failure("awk_filter", "Missing 'filter' argument".to_string())
            }
        };

        // Use awk command for filtering
        let output = Command::new("awk").arg(filter).arg(&path).output();

        match output {
            Ok(out) => {
                if out.status.success() {
                    ToolResult::success(
                        "awk_filter",
                        String::from_utf8_lossy(&out.stdout).to_string(),
                    )
                } else {
                    ToolResult::failure(
                        "awk_filter",
                        String::from_utf8_lossy(&out.stderr).to_string(),
                    )
                }
            }
            Err(e) => ToolResult::failure("awk_filter", format!("Failed to run awk: {}", e)),
        }
    }

    /// Show differences between two files
    fn diff_files(&self, call: &ToolCall) -> ToolResult {
        let file1 = match call.arguments.get("file1") {
            Some(p) => self.resolve_path(p),
            None => {
                return ToolResult::failure("diff_files", "Missing 'file1' argument".to_string())
            }
        };

        let file2 = match call.arguments.get("file2") {
            Some(p) => self.resolve_path(p),
            None => {
                return ToolResult::failure("diff_files", "Missing 'file2' argument".to_string())
            }
        };

        // Use diff command
        let output = Command::new("diff")
            .args([
                "--unified",
                &file1.to_string_lossy(),
                &file2.to_string_lossy(),
            ])
            .output();

        match output {
            Ok(out) => {
                // diff exits with 0 if files are same, 1 if different, 2 if error
                let stdout = String::from_utf8_lossy(&out.stdout).to_string();
                if stdout.is_empty() {
                    ToolResult::success("diff_files", "Files are identical".to_string())
                } else {
                    ToolResult::success("diff_files", stdout)
                }
            }
            Err(e) => ToolResult::failure("diff_files", format!("Failed to run diff: {}", e)),
        }
    }
}

/// Get tool definitions for LLM function calling
pub fn get_tool_definitions() -> Vec<ToolDefinition> {
    vec![
        ToolDefinition {
            name: "read_file".to_string(),
            description: "Read the contents of a file".to_string(),
            parameters: vec![ToolParameter {
                name: "path".to_string(),
                description: "Path to the file to read".to_string(),
                required: true,
            }],
        },
        ToolDefinition {
            name: "search_code".to_string(),
            description: "Search for code patterns in the workspace using grep/ripgrep".to_string(),
            parameters: vec![
                ToolParameter {
                    name: "query".to_string(),
                    description: "Search pattern (regex supported)".to_string(),
                    required: true,
                },
                ToolParameter {
                    name: "path".to_string(),
                    description: "Directory to search in (default: working directory)".to_string(),
                    required: false,
                },
            ],
        },
        ToolDefinition {
            name: "apply_patch".to_string(),
            description: "Write or replace file contents".to_string(),
            parameters: vec![
                ToolParameter {
                    name: "path".to_string(),
                    description: "Path to the file to write".to_string(),
                    required: true,
                },
                ToolParameter {
                    name: "content".to_string(),
                    description: "New file contents".to_string(),
                    required: true,
                },
            ],
        },
        ToolDefinition {
            name: "run_command".to_string(),
            description: "Execute a shell command in the working directory".to_string(),
            parameters: vec![ToolParameter {
                name: "command".to_string(),
                description: "Shell command to execute".to_string(),
                required: true,
            }],
        },
        ToolDefinition {
            name: "list_files".to_string(),
            description: "List files in a directory".to_string(),
            parameters: vec![ToolParameter {
                name: "path".to_string(),
                description: "Directory path (default: working directory)".to_string(),
                required: false,
            }],
        },
        // Power Tools
        ToolDefinition {
            name: "sed_replace".to_string(),
            description: "Replace text in a file using sed-like pattern matching".to_string(),
            parameters: vec![
                ToolParameter {
                    name: "path".to_string(),
                    description: "Path to the file".to_string(),
                    required: true,
                },
                ToolParameter {
                    name: "pattern".to_string(),
                    description: "Search pattern".to_string(),
                    required: true,
                },
                ToolParameter {
                    name: "replacement".to_string(),
                    description: "Replacement text".to_string(),
                    required: true,
                },
            ],
        },
        ToolDefinition {
            name: "awk_filter".to_string(),
            description: "Filter file content using awk-like field selection".to_string(),
            parameters: vec![
                ToolParameter {
                    name: "path".to_string(),
                    description: "Path to the file".to_string(),
                    required: true,
                },
                ToolParameter {
                    name: "filter".to_string(),
                    description: "Awk filter expression (e.g., '$1 == \"error\"')".to_string(),
                    required: true,
                },
            ],
        },
        ToolDefinition {
            name: "diff_files".to_string(),
            description: "Show differences between two files".to_string(),
            parameters: vec![
                ToolParameter {
                    name: "file1".to_string(),
                    description: "First file path".to_string(),
                    required: true,
                },
                ToolParameter {
                    name: "file2".to_string(),
                    description: "Second file path".to_string(),
                    required: true,
                },
            ],
        },
    ]
}

/// Tool definition for LLM function calling
#[derive(Debug, Clone)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub parameters: Vec<ToolParameter>,
}

/// Tool parameter definition
#[derive(Debug, Clone)]
pub struct ToolParameter {
    pub name: String,
    pub description: String,
    pub required: bool,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::env::temp_dir;

    #[tokio::test]
    async fn test_read_file() {
        let dir = temp_dir();
        let test_file = dir.join("test_read.txt");
        fs::write(&test_file, "Hello, World!").unwrap();

        let tools = AgentTools::new(dir.clone(), false);
        let call = ToolCall {
            name: "read_file".to_string(),
            arguments: [("path".to_string(), test_file.to_string_lossy().to_string())]
                .into_iter()
                .collect(),
        };

        let result = tools.execute(&call).await;
        assert!(result.success);
        assert_eq!(result.output, "Hello, World!");
    }

    #[tokio::test]
    async fn test_list_files() {
        let dir = temp_dir();
        let tools = AgentTools::new(dir.clone(), false);
        let call = ToolCall {
            name: "list_files".to_string(),
            arguments: HashMap::new(),
        };

        let result = tools.execute(&call).await;
        assert!(result.success);
    }
}