tori/src/tools.rs

use std::collections::HashMap;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};

use crate::llm::{Tool, ToolFunction};

struct ExternalTool {
    path: PathBuf,
    schema: Tool,
}

pub struct ExternalToolManager {
    tools: HashMap<String, ExternalTool>,
    venv_bin: Option<PathBuf>,
}

impl ExternalToolManager {
    /// Build the PATH string with venv bin prepended (if present).
    fn env_path(&self) -> Option<String> {
        self.venv_bin.as_ref().map(|venv_bin| {
            format!(
                "{}:{}",
                venv_bin.display(),
                std::env::var("PATH").unwrap_or_default()
            )
        })
    }

    /// Scan a template directory for external tools and Python dependencies.
    ///
    /// - Discovers executables in `template_dir/tools/`
    /// - If `template_dir/pyproject.toml` exists, runs `uv sync` to create a venv
    pub async fn discover(template_dir: &Path) -> Self {
        let mut tools = HashMap::new();

        // --- Python venv setup ---
        let venv_bin = if template_dir.join("pyproject.toml").is_file() {
            tracing::info!("Found pyproject.toml in {}, running uv sync", template_dir.display());
            let output = tokio::process::Command::new("uv")
                .args(["sync", "--project", &template_dir.to_string_lossy(), "--quiet"])
                .output()
                .await;
            match output {
                Ok(o) if o.status.success() => {
                    let bin = template_dir.join(".venv/bin");
                    tracing::info!("uv sync succeeded, venv bin: {}", bin.display());
                    Some(bin)
                }
                Ok(o) => {
                    tracing::warn!(
                        "uv sync failed: {}",
                        String::from_utf8_lossy(&o.stderr)
                    );
                    None
                }
                Err(e) => {
                    tracing::warn!("Failed to run uv sync: {}", e);
                    None
                }
            }
        } else {
            None
        };

        // --- Tool discovery ---
        let tools_dir = template_dir.join("tools");
        let mut entries = match tokio::fs::read_dir(&tools_dir).await {
            Ok(e) => e,
            Err(_) => return Self { tools, venv_bin },
        };

        // Build PATH with venv for --print-schema calls
        let env_path = venv_bin.as_ref().map(|bin| {
            format!(
                "{}:{}",
                bin.display(),
                std::env::var("PATH").unwrap_or_default()
            )
        });

        while let Ok(Some(entry)) = entries.next_entry().await {
            let path = entry.path();

            // Skip non-files
            let meta = match tokio::fs::metadata(&path).await {
                Ok(m) => m,
                Err(_) => continue,
            };
            if !meta.is_file() {
                continue;
            }

            // Check executable bit
            if meta.permissions().mode() & 0o111 == 0 {
                tracing::debug!("Skipping non-executable: {}", path.display());
                continue;
            }

            // Call --print-schema (with venv PATH if available)
            let mut cmd = tokio::process::Command::new(&path);
            cmd.arg("--print-schema");
            if let Some(ref p) = env_path {
                cmd.env("PATH", p);
            }
            if let Some(ref bin) = venv_bin {
                cmd.env("VIRTUAL_ENV", bin.parent().unwrap().display().to_string());
            }
            let output = match cmd.output().await {
                Ok(o) => o,
                Err(e) => {
                    tracing::warn!("Failed to run --print-schema on {}: {}", path.display(), e);
                    continue;
                }
            };

            if !output.status.success() {
                tracing::warn!(
                    "--print-schema failed for {}: {}",
                    path.display(),
                    String::from_utf8_lossy(&output.stderr)
                );
                continue;
            }

            let schema: serde_json::Value = match serde_json::from_slice(&output.stdout) {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!("Invalid schema JSON from {}: {}", path.display(), e);
                    continue;
                }
            };

            let name = schema["name"]
                .as_str()
                .unwrap_or_else(|| {
                    path.file_name()
                        .and_then(|n| n.to_str())
                        .unwrap_or("unknown")
                })
                .to_string();

            let description = schema["description"].as_str().unwrap_or("").to_string();
            let parameters = schema["parameters"].clone();

            let tool = Tool {
                tool_type: "function".into(),
                function: ToolFunction {
                    name: name.clone(),
                    description,
                    parameters,
                },
            };

            tracing::info!("Discovered external tool: {}", name);
            tools.insert(
                name,
                ExternalTool {
                    path: path.clone(),
                    schema: tool,
                },
            );
        }

        Self { tools, venv_bin }
    }

    /// Return all discovered Tool definitions for LLM API calls.
    pub fn tool_definitions(&self) -> Vec<Tool> {
        self.tools.values().map(|t| t.schema.clone()).collect()
    }

    /// Invoke an external tool by name, passing JSON args as the first argv.
    pub async fn invoke(
        &self,
        name: &str,
        args_json: &str,
        workdir: &str,
    ) -> anyhow::Result<String> {
        let tool = self
            .tools
            .get(name)
            .ok_or_else(|| anyhow::anyhow!("External tool not found: {}", name))?;

        let mut cmd = tokio::process::Command::new(&tool.path);
        cmd.arg(args_json).current_dir(workdir);
        if let Some(ref p) = self.env_path() {
            cmd.env("PATH", p);
        }
        if let Some(ref venv_bin) = self.venv_bin {
            cmd.env("VIRTUAL_ENV", venv_bin.parent().unwrap().display().to_string());
        }
        let output = cmd.output().await?;

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

        if output.status.success() {
            Ok(stdout)
        } else {
            let mut result = stdout;
            if !stderr.is_empty() {
                result.push_str("\nSTDERR: ");
                result.push_str(&stderr);
            }
            result.push_str(&format!(
                "\n[exit code: {}]",
                output.status.code().unwrap_or(-1)
            ));
            Ok(result)
        }
    }

    /// Check if a tool with the given name exists.
    pub fn has_tool(&self, name: &str) -> bool {
        self.tools.contains_key(name)
    }

    /// Number of discovered tools.
    pub fn len(&self) -> usize {
        self.tools.len()
    }
}