openheim 0.3.0

//! Tool abstraction layer: trait definitions, built-in tools, and the runtime
//! executor that routes LLM tool calls to the correct handler.
//!
//! # Built-in tools
//!
//! Three tools are registered by default via [`SystemToolExecutor::register_builtins`]:
//!
//! | Name | Description |
//! |------|-------------|
//! | `execute_command` | Run a shell command (`sh -c` on Unix, `cmd /C` on Windows) |
//! | `read_file` | Read a file from disk |
//! | `write_file` | Write a file to disk, creating parent directories as needed |
//!
//! Additional tools are loaded from MCP servers and registered under the
//! `{server_name}__{tool_name}` namespace.
//!
//! # Implementing a custom tool
//!
//! ```rust,no_run
//! use async_trait::async_trait;
//! use serde_json::json;
//!
//! struct GreetTool;
//!
//! # use openheim::tools::ToolHandler;
//! # use openheim::core::models::{Tool, FunctionDefinition};
//! # use openheim::error::Result;
//! #[async_trait]
//! impl ToolHandler for GreetTool {
//!     fn definition(&self) -> Tool {
//!         Tool {
//!             tool_type: "function".to_string(),
//!             function: FunctionDefinition {
//!                 name: "greet".to_string(),
//!                 description: "Greet someone by name.".to_string(),
//!                 parameters: json!({
//!                     "type": "object",
//!                     "properties": { "name": { "type": "string" } },
//!                     "required": ["name"]
//!                 }),
//!             },
//!         }
//!     }
//!
//!     async fn execute(&self, args: &str) -> Result<String> {
//!         let v: serde_json::Value = serde_json::from_str(args)?;
//!         let name = v["name"].as_str().unwrap_or("world");
//!         Ok(format!("Hello, {name}!"))
//!     }
//! }
//! ```

mod execute_command;
mod read_file;
pub mod sandbox;
mod sandboxed_executor;
mod write_file;

use std::collections::{BTreeMap, HashMap};

use async_trait::async_trait;

use crate::config::McpServerConfig;
use crate::core::models::Tool;
use crate::error::{Error, Result};

pub use sandboxed_executor::SandboxedExecutor;

#[async_trait]
pub trait ToolHandler: Send + Sync {
    /// Returns the tool definition (name, description, JSON-schema parameters).
    fn definition(&self) -> Tool;

    /// Executes the tool with the given JSON-encoded arguments.
    async fn execute(&self, args: &str) -> Result<String>;
}

/// Routes LLM tool-call requests to the correct [`ToolHandler`].
///
/// The production implementation is [`SystemToolExecutor`]. Tests typically use
/// lightweight mock implementations of this trait.
#[async_trait]
pub trait ToolExecutor: Send + Sync {
    /// Returns the list of tools available to the LLM.
    fn list_tools(&self) -> Vec<Tool>;

    /// Dispatches a tool call by name with JSON-encoded arguments.
    ///
    /// Returns the tool output as a string, or an error if the tool is unknown
    /// or its execution fails.
    async fn execute(&self, name: &str, args_json: &str) -> Result<String>;
}

/// The default tool executor used by the agent runtime.
///
/// Maintains a registry of [`ToolHandler`]s keyed by tool name and dispatches
/// LLM tool calls to the appropriate handler. Built-in tools are registered via
/// [`register_builtins`]; MCP tools are added during [`build`].
pub struct SystemToolExecutor {
    handlers: HashMap<String, Box<dyn ToolHandler>>,
}

impl SystemToolExecutor {
    /// Creates an empty executor with no registered tools.
    pub fn new() -> Self {
        Self {
            handlers: HashMap::new(),
        }
    }

    /// Builds a fully-configured executor: registers built-in tools then connects
    /// to all configured MCP servers and registers their tools.
    ///
    /// When `allow_shell` is `false` the `execute_command` tool is omitted so
    /// the LLM never sees it in its tool list.
    ///
    /// Returns the executor alongside [`McpServerStatus`] entries for each server
    /// so callers can inspect which connections succeeded.
    pub async fn build(
        mcp_configs: &BTreeMap<String, McpServerConfig>,
        allow_shell: bool,
    ) -> (Self, Vec<crate::mcp::McpServerStatus>) {
        let mut executor = Self::new();
        executor.register_builtins();
        if !allow_shell {
            executor.handlers.remove("execute_command");
        }
        let (handlers, statuses) = crate::mcp::load_mcp_tools(mcp_configs).await;
        for handler in handlers {
            executor.register(handler);
        }
        (executor, statuses)
    }

    /// Registers the three built-in tools: `execute_command`, `read_file`, `write_file`.
    pub fn register_builtins(&mut self) {
        self.register(Box::new(execute_command::ExecuteCommandTool));
        self.register(Box::new(read_file::ReadFileTool));
        self.register(Box::new(write_file::WriteFileTool));
    }

    /// Registers a single tool handler.
    ///
    /// If a tool with the same name is already registered it is overwritten and a
    /// warning is logged.
    pub fn register(&mut self, handler: Box<dyn ToolHandler>) {
        let name = handler.definition().function.name.clone();
        if self.handlers.contains_key(&name) {
            tracing::warn!(name = %name, "Tool name collision: overwriting existing handler");
        }
        self.handlers.insert(name, handler);
    }
}

impl Default for SystemToolExecutor {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl ToolExecutor for SystemToolExecutor {
    fn list_tools(&self) -> Vec<Tool> {
        self.handlers.values().map(|h| h.definition()).collect()
    }

    async fn execute(&self, name: &str, args_json: &str) -> Result<String> {
        let handler = self
            .handlers
            .get(name)
            .ok_or_else(|| Error::ToolExecutionError(format!("Unknown tool: {}", name)))?;
        handler.execute(args_json).await
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::*;

    #[test]
    fn new_executor_is_empty() {
        let executor = SystemToolExecutor::new();
        assert_eq!(executor.handlers.len(), 0);
    }

    #[test]
    fn register_builtins_adds_three_tools() {
        let mut executor = SystemToolExecutor::new();
        executor.register_builtins();
        assert!(executor.handlers.contains_key("execute_command"));
        assert!(executor.handlers.contains_key("read_file"));
        assert!(executor.handlers.contains_key("write_file"));
        assert_eq!(executor.handlers.len(), 3);
    }

    #[tokio::test]
    async fn executor_returns_error_for_unknown_tool() {
        let executor = SystemToolExecutor::new();
        let result = executor.execute("nonexistent_tool", "{}").await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Unknown tool"));
    }

    #[tokio::test]
    async fn build_without_shell_omits_execute_command() {
        let (executor, _) = SystemToolExecutor::build(&BTreeMap::new(), false).await;
        assert!(!executor.handlers.contains_key("execute_command"));
        assert!(executor.handlers.contains_key("read_file"));
        assert!(executor.handlers.contains_key("write_file"));
    }
}