magi-tool 0.0.7

use dashmap::{DashMap, mapref::entry::Entry};
use serde_json::Value;
use std::{future::Future, pin::Pin, sync::Arc};
use thiserror::Error;

use crate::Tool as ToolDefinition;

type ToolCallFuture = Pin<Box<dyn Future<Output = Result<String, String>> + Send>>;
type ToolHandler = Arc<dyn Fn(Value) -> ToolCallFuture + Send + Sync>;

#[derive(Clone)]
struct Tool {
    definition: ToolDefinition,
    handler: ToolHandler,
}

/// Error type for local tool registration and execution.
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum ToolCallError {
    #[error("Tool '{0}' already exists")]
    ToolAlreadyExists(String),
    #[error("Tool '{0}' not found")]
    ToolNotFound(String),
    #[error("Tool '{tool_name}' call failed: {message}")]
    CallFailed { tool_name: String, message: String },
}

/// Local tool registry backed by tool-name-to-closure mappings.
#[derive(Clone)]
pub struct ToolKit {
    tools: Arc<DashMap<String, Tool>>,
}

impl ToolKit {
    /// Creates a new empty ToolKit.
    pub fn new() -> Self {
        Self {
            tools: Arc::new(DashMap::new()),
        }
    }

    /// Creates a new ToolKit containing one registered local tool.
    pub fn with_tool<F, Fut>(tool: ToolDefinition, handler: F) -> Result<Self, ToolCallError>
    where
        F: Fn(Value) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<String, String>> + Send + 'static,
    {
        let tool_kit = Self::new();
        tool_kit.register_tool(tool, handler)?;
        Ok(tool_kit)
    }

    /// Registers a local tool and the closure that executes it.
    pub fn register_tool<F, Fut>(
        &self,
        tool: ToolDefinition,
        handler: F,
    ) -> Result<(), ToolCallError>
    where
        F: Fn(Value) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<String, String>> + Send + 'static,
    {
        let tool_name = tool.name.to_string();
        let handler = Arc::new(move |arguments| Box::pin(handler(arguments)) as ToolCallFuture);

        match self.tools.entry(tool_name) {
            Entry::Occupied(entry) => Err(ToolCallError::ToolAlreadyExists(entry.key().clone())),
            Entry::Vacant(entry) => {
                entry.insert(Tool {
                    definition: tool,
                    handler,
                });
                Ok(())
            }
        }
    }

    /// Gets a tool by its local name.
    pub fn get_tool(&self, tool_name: &str) -> Option<ToolDefinition> {
        self.tools
            .get(tool_name)
            .map(|tool| tool.value().definition.clone())
    }

    /// Looks up multiple tools by local name, preserving missing entries as `None`.
    pub fn lookup_tools(&self, tool_names: &[&str]) -> Vec<Option<ToolDefinition>> {
        tool_names
            .iter()
            .map(|tool_name| self.get_tool(tool_name))
            .collect()
    }

    /// Gets all registered local tool names.
    pub fn tool_names(&self) -> Vec<String> {
        self.tools.iter().map(|entry| entry.key().clone()).collect()
    }

    /// Gets all registered local tool definitions.
    pub fn tools(&self) -> Vec<ToolDefinition> {
        self.tools
            .iter()
            .map(|entry| entry.value().definition.clone())
            .collect()
    }

    /// Removes a local tool by name.
    pub fn unregister_tool(&self, tool_name: &str) -> Result<(), ToolCallError> {
        if self.tools.remove(tool_name).is_none() {
            return Err(ToolCallError::ToolNotFound(tool_name.to_string()));
        }

        Ok(())
    }

    /// Returns the number of registered local tools.
    pub fn tool_count(&self) -> usize {
        self.tools.len()
    }

    /// Calls a local tool by name and returns the text result for the LLM.
    pub async fn call_tool(
        &self,
        tool_name: &str,
        arguments: Value,
    ) -> Result<String, ToolCallError> {
        let handler = {
            self.tools
                .get(tool_name)
                .map(|tool| Arc::clone(&tool.value().handler))
                .ok_or_else(|| ToolCallError::ToolNotFound(tool_name.to_string()))?
        };

        handler(arguments)
            .await
            .map_err(|message| ToolCallError::CallFailed {
                tool_name: tool_name.to_string(),
                message,
            })
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use schemars::JsonSchema;
    use serde_json::json;

    #[derive(JsonSchema)]
    #[allow(dead_code)]
    struct EchoParams {
        value: String,
    }

    fn echo_tool() -> ToolDefinition {
        crate::get_tool::<EchoParams, _, _>("echo", Some("Echo a value")).unwrap()
    }

    #[test]
    fn test_with_tool_registers_one_tool() {
        let tool_kit =
            ToolKit::with_tool(echo_tool(), |_arguments| async { Ok(String::new()) }).unwrap();

        assert_eq!(tool_kit.tool_count(), 1);
        assert!(tool_kit.get_tool("echo").is_some());
    }

    #[tokio::test]
    async fn test_register_and_call_tool() {
        let tool_kit = ToolKit::new();
        tool_kit
            .register_tool(echo_tool(), |arguments| async move {
                let value = arguments
                    .get("value")
                    .and_then(Value::as_str)
                    .unwrap_or_default();
                Ok(value.to_string())
            })
            .unwrap();

        assert_eq!(tool_kit.tool_count(), 1);
        assert!(tool_kit.get_tool("echo").is_some());

        let result = tool_kit.call_tool("echo", json!({"value": "hello"})).await;
        assert_eq!(result.unwrap(), "hello");
    }

    #[test]
    fn test_duplicate_tool_fails() {
        let tool_kit = ToolKit::new();
        tool_kit
            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
            .unwrap();

        let error = tool_kit
            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
            .unwrap_err();

        assert_eq!(error, ToolCallError::ToolAlreadyExists("echo".to_string()));
    }

    #[tokio::test]
    async fn test_missing_tool_call_fails() {
        let tool_kit = ToolKit::new();
        let error = tool_kit
            .call_tool("missing", Value::Null)
            .await
            .unwrap_err();

        assert_eq!(error, ToolCallError::ToolNotFound("missing".to_string()));
    }

    #[tokio::test]
    async fn test_handler_error_includes_called_tool_name() {
        let tool_kit = ToolKit::new();
        tool_kit
            .register_tool(echo_tool(), |_arguments| async {
                Err("missing value".to_string())
            })
            .unwrap();

        let error = tool_kit.call_tool("echo", Value::Null).await.unwrap_err();

        assert_eq!(
            error,
            ToolCallError::CallFailed {
                tool_name: "echo".to_string(),
                message: "missing value".to_string(),
            }
        );
    }

    #[test]
    fn test_lookup_tools_preserves_missing_entries() {
        let tool_kit = ToolKit::new();
        tool_kit
            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
            .unwrap();

        let tools = tool_kit.lookup_tools(&["echo", "missing"]);

        assert_eq!(tools.len(), 2);
        assert!(tools[0].is_some());
        assert!(tools[1].is_none());
    }

    #[test]
    fn test_unregister_tool() {
        let tool_kit = ToolKit::new();
        tool_kit
            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
            .unwrap();

        assert_eq!(tool_kit.tool_names(), vec!["echo".to_string()]);
        assert_eq!(tool_kit.tools().len(), 1);

        tool_kit.unregister_tool("echo").unwrap();

        assert_eq!(tool_kit.tool_count(), 0);
        assert!(tool_kit.get_tool("echo").is_none());
    }
}