brokecrew/src/tool.rs

use serde::{Deserialize, Serialize};
use std::path::Path;
use std::time::Duration;
use tokio::process::Command;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolOutput {
    pub command: String,
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
    pub duration_ms: u64,
}

impl ToolOutput {
    pub fn is_success(&self) -> bool {
        self.exit_code == 0
    }

    pub fn format_for_model(&self) -> String {
        let mut out = format!(
            "[Command: `{}` exit code: {}]\n",
            self.command, self.exit_code
        );
        if !self.stdout.is_empty() {
            out.push_str("stdout:\n");
            out.push_str(&self.stdout);
            if !self.stdout.ends_with('\n') {
                out.push('\n');
            }
        }
        if !self.stderr.is_empty() {
            out.push_str("stderr:\n");
            out.push_str(&self.stderr);
            if !self.stderr.ends_with('\n') {
                out.push('\n');
            }
        }
        if self.stdout.is_empty() && self.stderr.is_empty() {
            out.push_str("(No output produced)\n");
        }
        out
    }
}

pub struct ToolEngine;

impl ToolEngine {
    /// Executes a bash command locally with timeout and output capture
    pub async fn execute_bash(
        command: &str,
        cwd: Option<&Path>,
        timeout_duration: Duration,
    ) -> Result<ToolOutput, String> {
        let start = std::time::Instant::now();

        let mut cmd = Command::new("bash");
        cmd.arg("-c").arg(command);

        if let Some(dir) = cwd {
            cmd.current_dir(dir);
        }

        let child = cmd.output();

        match tokio::time::timeout(timeout_duration, child).await {
            Ok(Ok(output)) => {
                let duration_ms = start.elapsed().as_millis() as u64;
                let exit_code = output.status.code().unwrap_or(-1);

                let stdout_raw = String::from_utf8_lossy(&output.stdout);
                let stderr_raw = String::from_utf8_lossy(&output.stderr);

                // Truncate to maximum 32KB to protect token context
                let stdout = truncate_output(&stdout_raw, 32 * 1024);
                let stderr = truncate_output(&stderr_raw, 16 * 1024);

                Ok(ToolOutput {
                    command: command.to_string(),
                    exit_code,
                    stdout,
                    stderr,
                    duration_ms,
                })
            }
            Ok(Err(e)) => Err(format!("Process spawn error: {e}")),
            Err(_) => Err(format!(
                "Command timed out after {}s",
                timeout_duration.as_secs()
            )),
        }
    }

    /// Extracts any command execution requests formatted as:
    /// `<run>cmd</run>`, `<bash>cmd</bash>`, or markdown ```bash blocks tagged for execution
    pub fn extract_command_calls(text: &str) -> Vec<String> {
        let mut commands = Vec::new();

        // 1. Check for <run>...</run> tags
        let mut rest = text;
        while let Some(start_idx) = rest.find("<run>") {
            let after_start = &rest[start_idx + 5..];
            if let Some(end_idx) = after_start.find("</run>") {
                let cmd = after_start[..end_idx].trim();
                if !cmd.is_empty() {
                    commands.push(cmd.to_string());
                }
                rest = &after_start[end_idx + 6..];
            } else {
                break;
            }
        }

        // 2. Check for <bash>...</bash> tags
        let mut rest_bash = text;
        while let Some(start_idx) = rest_bash.find("<bash>") {
            let after_start = &rest_bash[start_idx + 6..];
            if let Some(end_idx) = after_start.find("</bash>") {
                let cmd = after_start[..end_idx].trim();
                if !cmd.is_empty() && !commands.contains(&cmd.to_string()) {
                    commands.push(cmd.to_string());
                }
                rest_bash = &after_start[end_idx + 7..];
            } else {
                break;
            }
        }

        commands
    }

    /// Tool instructions appended to persona system prompts
    pub fn system_tool_prompt() -> &'static str {
        "\n\n[TOOL CAPABILITY: LOCAL BASH]\nYou can execute local shell commands on the user's system when needed (e.g. running tests, reading files, inspecting directories, running git/cargo/python/npm). To run a command, output it enclosed in <run>...</run> tags, for example:\n<run>cargo check</run>\nor\n<run>git status</run>\nThe command output will be executed and returned to you."
    }
}

fn truncate_output(s: &str, max_bytes: usize) -> String {
    if s.len() <= max_bytes {
        s.to_string()
    } else {
        // Walk back to a char boundary so multibyte output never panics.
        let mut end = max_bytes.min(s.len());
        while end > 0 && !s.is_char_boundary(end) {
            end -= 1;
        }
        format!("{}\n...[output truncated by BrokeCrew]", &s[..end])
    }
}

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

    #[test]
    fn test_extract_command_calls() {
        let text = "Let me check the repository status:\n<run>git status</run>\nand also list files:\n<bash>ls -la</bash>";
        let calls = ToolEngine::extract_command_calls(text);
        assert_eq!(calls, vec!["git status", "ls -la"]);
    }

    #[tokio::test]
    async fn test_execute_echo_command() {
        let res =
            ToolEngine::execute_bash("echo 'hello from brokecrew'", None, Duration::from_secs(5))
                .await
                .expect("execution succeeds");
        assert_eq!(res.exit_code, 0);
        assert!(res.stdout.contains("hello from brokecrew"));
    }

    #[test]
    fn test_truncate_output_respects_char_boundary() {
        // 32KB+ of multibyte output must truncate without panicking.
        let big = "日本語テスト出力 ".repeat(4000);
        let out = truncate_output(&big, 32 * 1024);
        assert!(out.ends_with("...[output truncated by BrokeCrew]"));
        assert!(out.len() <= 32 * 1024 + 64);
        // Short output passes through untouched.
        assert_eq!(truncate_output("ok", 1024), "ok");
    }
}