magi_tool/

pegboard.rs

use dashmap::DashMap;
use rmcp::{
    RoleClient,
    model::{CallToolRequestParam, CallToolResult},
    service::{DynService, RunningService},
};
use serde_json::Value;
use std::borrow::Cow;
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use crate::Tool;

/// Internal implementation of PegBoard that manages tool registration and their associated MCP services.
///
/// Use `PegBoard` (which is `Arc<InternalPegBoard>`) instead of using this directly.
struct InternalPegBoard {
    /// All registered tools with prefixed names (e.g., "namespace-tool_name")
    /// These tools have their `name` field modified to include the prefix
    tools: DashMap<String, Tool>,

    /// MCP services with their namespaces: DashMap<service_index, (namespace, service)>
    /// DashMap provides lock-free concurrent access
    services: DashMap<
        usize,
        (
            String,
            RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
        ),
    >,

    /// Next available service index (atomic counter)
    /// Using atomic counter instead of len() ensures unique indices even after service removal
    next_service_idx: AtomicUsize,

    /// Maps prefixed tool name to (service_index, original_tool_name)
    /// Used for routing: when LLM calls "namespace-search", we look up which service
    /// and what the original tool name was
    tool_routing: DashMap<String, ToolRoute>,

    /// Maps namespace to list of prefixed tool names in that namespace
    namespace_tools: DashMap<String, Vec<String>>,
}

/// Routing information for a tool
#[derive(Debug, Clone)]
struct ToolRoute {
    /// Index of the service that provides this tool
    service_index: usize,
    /// Original tool name (before prefixing)
    original_name: String,
}

/// Helper to create a prefixed tool name
fn prefix_tool_name(namespace: &str, tool_name: &str) -> String {
    format!("{}-{}", namespace, tool_name)
}

impl InternalPegBoard {
    /// Creates a new empty InternalPegBoard
    fn new() -> Self {
        Self {
            tools: DashMap::new(),
            services: DashMap::new(),
            next_service_idx: AtomicUsize::new(0),
            tool_routing: DashMap::new(),
            namespace_tools: DashMap::new(),
        }
    }

    /// Registers a service and automatically discovers all its tools.
    ///
    /// This method:
    /// 1. Calls `list_tools()` on the service to discover all available tools
    /// 2. Converts tools from rmcp format to PegBoard's Tool format
    /// 3. If namespace is provided, prefixes each tool name (e.g., "namespace-tool_name")
    /// 4. If namespace is None or empty, uses original tool names (no prefixing)
    /// 5. Adds the service to the registry
    /// 6. Registers all tools for use with LLM
    ///
    /// # Arguments
    /// * `namespace` - Optional namespace prefix. Use None or empty string if no conflicts expected
    /// * `service` - The MCP service to register
    ///
    /// # Returns
    /// * Number of tools discovered and registered
    ///
    /// # Example
    /// ```ignore
    /// // With namespace (prefixing enabled)
    /// pegboard.add_service(Some("web".to_string()), service).await?;
    /// // Tool "search" is now available as "web-search"
    ///
    /// // Without namespace (no prefixing)
    /// pegboard.add_service(None, service).await?;
    /// // Tool "search" keeps its original name "search"
    /// ```
    pub async fn add_service(
        &self,
        namespace: Option<String>,
        service: RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
    ) -> Result<usize, PegBoardError> {
        // Normalize namespace: None or empty string means no namespace
        let namespace_str = namespace.unwrap_or_default();
        let has_namespace = !namespace_str.is_empty();

        // Check if namespace already exists (only if we have one)
        if has_namespace && self.namespace_tools.contains_key(&namespace_str) {
            return Err(PegBoardError::NamespaceAlreadyExists(namespace_str));
        }

        // Get the next service index atomically
        // Using fetch_add ensures unique indices even after service removal
        let service_idx = self.next_service_idx.fetch_add(1, Ordering::SeqCst);

        // List tools from the service (always available in rmcp)
        let tools_response = service
            .list_tools(None) // No pagination params - get all tools
            .await
            .map_err(|e| PegBoardError::ServiceError(format!("Failed to list tools: {:?}", e)))?;

        // Convert rmcp tools to our Tool format
        let tools_list: Vec<Tool> = tools_response
            .tools
            .into_iter()
            .map(|rmcp_tool| Tool {
                name: rmcp_tool.name,
                description: rmcp_tool.description.map(Cow::from),
                input_schema: serde_json::Value::Object((*rmcp_tool.input_schema).clone()),
            })
            .collect();

        // Register the service (store empty string if no namespace)
        self.services
            .insert(service_idx, (namespace_str.clone(), service));

        // Track tool names
        let mut registered_tool_names = Vec::new();

        // Register each tool
        for original_tool in tools_list {
            let original_name = original_tool.name.to_string();

            // Prefix name only if namespace is provided
            let final_name = if has_namespace {
                prefix_tool_name(&namespace_str, &original_name)
            } else {
                original_name.clone()
            };

            // Check if this tool name already exists
            if self.tools.contains_key(&final_name) {
                return Err(PegBoardError::ToolAlreadyExists(final_name));
            }

            // Create a tool with the final name (prefixed or original)
            let mut final_tool = original_tool.clone();
            final_tool.name = Cow::Owned(final_name.clone());

            // Register the tool
            self.tools.insert(final_name.clone(), final_tool);

            // Store routing information (service index + original name)
            self.tool_routing.insert(
                final_name.clone(),
                ToolRoute {
                    service_index: service_idx,
                    original_name,
                },
            );

            registered_tool_names.push(final_name);
        }

        let tool_count = registered_tool_names.len();

        // Only track namespace if we have one
        if has_namespace {
            self.namespace_tools
                .insert(namespace_str, registered_tool_names);
        }

        Ok(tool_count)
    }

    /// Manually registers a tool with an optional namespace and service index.
    /// If namespace is provided, the tool name will be prefixed.
    /// If namespace is None or empty, the original tool name is used.
    /// Prefer `add_service()` for automatic tool discovery.
    pub fn register_tool(
        &self,
        namespace: Option<&str>,
        tool: Tool,
        service_idx: usize,
    ) -> Result<(), PegBoardError> {
        // Check if service index is valid
        if !self.services.contains_key(&service_idx) {
            return Err(PegBoardError::InvalidServiceIndex {
                index: service_idx,
                max: self.services.len(),
            });
        }

        let original_name = tool.name.to_string();
        let namespace_str = namespace.unwrap_or("");
        let has_namespace = !namespace_str.is_empty();

        // Prefix name only if namespace is provided
        let final_name = if has_namespace {
            prefix_tool_name(namespace_str, &original_name)
        } else {
            original_name.clone()
        };

        // Check if tool already exists
        if self.tools.contains_key(&final_name) {
            return Err(PegBoardError::ToolAlreadyExists(final_name));
        }

        // Create tool with final name (prefixed or original)
        let mut final_tool = tool;
        final_tool.name = Cow::Owned(final_name.clone());

        // Register tool and routing
        self.tools.insert(final_name.clone(), final_tool);
        self.tool_routing.insert(
            final_name.clone(),
            ToolRoute {
                service_index: service_idx,
                original_name,
            },
        );

        // Update namespace tracking (only if we have a namespace)
        if has_namespace {
            self.namespace_tools
                .entry(namespace_str.to_string())
                .or_default()
                .push(final_name);
        }

        Ok(())
    }

    /// Gets a tool by its name (prefixed if registered with namespace, original otherwise)
    /// This is the name that the LLM sees and uses
    pub fn get_tool(&self, tool_name: &str) -> Option<Tool> {
        self.tools.get(tool_name).map(|entry| entry.value().clone())
    }

    /// Selects multiple tools by their names.
    /// Returns `Some(Vec<Tool>)` if ALL requested tools are found.
    /// Returns `None` if ANY tool is missing.
    ///
    /// # Arguments
    /// * `tool_names` - A slice of tool names (prefixed if registered with namespace)
    ///
    /// # Returns
    /// * `Some(Vec<Tool>)` if all tools exist
    /// * `None` if any tool is missing
    ///
    /// # Example
    /// ```ignore
    /// let tools = pegboard.select_tools(&["web-search", "file-read"]);
    /// if let Some(tools) = tools {
    ///     // All tools found, can use them
    /// } else {
    ///     // One or more tools not found
    /// }
    /// ```
    pub fn select_tools(&self, tool_names: &[&str]) -> Option<Vec<Tool>> {
        let mut result = Vec::with_capacity(tool_names.len());

        for &tool_name in tool_names {
            let tool = self.get_tool(tool_name)?;
            result.push(tool);
        }

        Some(result)
    }

    /// Gets routing information for a tool by its name
    /// Returns (service_index, original_tool_name) for routing the call
    pub fn get_tool_route(&self, tool_name: &str) -> Option<(usize, String)> {
        self.tool_routing.get(tool_name).map(|entry| {
            let route = entry.value();
            (route.service_index, route.original_name.clone())
        })
    }

    /// Gets all tool names in a namespace (prefixed if namespace was used)
    pub fn list_tools_in_namespace(&self, namespace: &str) -> Vec<String> {
        self.namespace_tools
            .get(namespace)
            .map(|entry| entry.value().clone())
            .unwrap_or_default()
    }

    /// Gets all Tool objects in a namespace (with names as they appear to LLM)
    pub fn get_tools_in_namespace(&self, namespace: &str) -> Vec<Tool> {
        self.list_tools_in_namespace(namespace)
            .iter()
            .filter_map(|tool_name| self.get_tool(tool_name))
            .collect()
    }

    /// Gets all registered tool names across all namespaces
    /// These are the names that should be sent to the LLM
    pub fn list_all_tools(&self) -> Vec<String> {
        self.tools.iter().map(|entry| entry.key().clone()).collect()
    }

    /// Gets all tools as a Vec
    /// These are the tools that should be sent to the LLM
    pub fn get_all_tools(&self) -> Vec<Tool> {
        self.tools
            .iter()
            .map(|entry| entry.value().clone())
            .collect()
    }

    /// Gets all registered namespaces
    pub fn list_namespaces(&self) -> Vec<String> {
        self.namespace_tools
            .iter()
            .map(|entry| entry.key().clone())
            .collect()
    }

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

        self.tool_routing.remove(prefixed_name);

        // Remove from namespace tracking
        // We need to find which namespace this tool belongs to
        for mut namespace_entry in self.namespace_tools.iter_mut() {
            namespace_entry.value_mut().retain(|n| n != prefixed_name);
        }

        Ok(())
    }

    /// Removes all tools in a namespace and the associated service
    pub fn unregister_namespace(&self, namespace: &str) -> Result<usize, PegBoardError> {
        let prefixed_names = self.list_tools_in_namespace(namespace);
        let count = prefixed_names.len();

        // Remove all tools in this namespace
        for prefixed_name in prefixed_names {
            self.tools.remove(&prefixed_name);
            self.tool_routing.remove(&prefixed_name);
        }

        // Find and remove the service with this namespace
        let service_idx_to_remove = self
            .services
            .iter()
            .find(|entry| entry.value().0 == namespace)
            .map(|entry| *entry.key());

        if let Some(idx) = service_idx_to_remove {
            self.services.remove(&idx);
        }

        self.namespace_tools.remove(namespace);
        Ok(count)
    }

    /// Removes a service by its index and all associated tools
    pub fn unregister_service(&self, service_idx: usize) -> Result<usize, PegBoardError> {
        // Get the namespace for this service
        let namespace = self
            .services
            .get(&service_idx)
            .map(|entry| entry.value().0.clone())
            .ok_or(PegBoardError::InvalidServiceIndex {
                index: service_idx,
                max: self.next_service_idx.load(Ordering::SeqCst),
            })?;

        // Find and remove all tools associated with this service index
        let tools_to_remove: Vec<String> = self
            .tool_routing
            .iter()
            .filter(|entry| entry.value().service_index == service_idx)
            .map(|entry| entry.key().clone())
            .collect();

        let count = tools_to_remove.len();
        for tool_name in tools_to_remove {
            self.tools.remove(&tool_name);
            self.tool_routing.remove(&tool_name);
        }

        // Remove the service
        self.services.remove(&service_idx);

        // If it has a namespace, also clean up namespace tracking
        if !namespace.is_empty() {
            self.namespace_tools.remove(&namespace);
        }

        Ok(count)
    }

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

    /// Returns the number of registered services
    pub fn service_count(&self) -> usize {
        self.services.len()
    }

    /// Returns the number of registered namespaces
    pub fn namespace_count(&self) -> usize {
        self.namespace_tools.len()
    }

    /// Calls a tool by its name (as seen by the LLM) with the given arguments.
    ///
    /// This method:
    /// 1. Looks up the routing information using the tool name
    /// 2. Finds the service that provides this tool
    /// 3. Calls the service's `call_tool` method with the original tool name
    ///
    /// # Arguments
    /// * `tool_name` - The tool name as seen by the LLM (prefixed if namespace was used)
    /// * `arguments` - The arguments to pass to the tool as a JSON value
    ///
    /// # Returns
    /// * `CallToolResult` containing the tool's response
    ///
    /// # Errors
    /// * `PegBoardError::ToolNotFound` - If the tool name is not registered
    /// * `PegBoardError::ServiceError` - If the service call fails
    ///
    /// # Example
    /// ```ignore
    /// // LLM calls "web-search"
    /// let result = pegboard.call_tool(
    ///     "web-search",
    ///     serde_json::json!({"query": "rust programming"}),
    /// ).await?;
    /// ```
    pub async fn call_tool(
        &self,
        tool_name: &str,
        arguments: Value,
    ) -> Result<CallToolResult, PegBoardError> {
        // Get routing information
        let (service_idx, original_name) = self
            .get_tool_route(tool_name)
            .ok_or_else(|| PegBoardError::ToolNotFound(tool_name.to_string()))?;

        // Get the service
        let service_entry =
            self.services
                .get(&service_idx)
                .ok_or(PegBoardError::InvalidServiceIndex {
                    index: service_idx,
                    max: self.services.len(),
                })?;
        let (_namespace, service) = service_entry.value();

        // Convert arguments to JsonObject if it's an object, otherwise use None
        let arguments_obj = match arguments {
            Value::Object(obj) => Some(obj),
            Value::Null => None,
            _ => {
                return Err(PegBoardError::ServiceError(
                    "Tool arguments must be a JSON object or null".to_string(),
                ));
            }
        };

        // Call the service's call_tool method with the original tool name
        let param = CallToolRequestParam {
            name: Cow::from(original_name),
            arguments: arguments_obj,
        };

        service
            .call_tool(param)
            .await
            .map_err(|e| PegBoardError::ServiceError(format!("Tool call failed: {:?}", e)))
    }
}

/// PegBoard manages tool registration and their associated MCP services with namespace support.
///
/// ## Namespace and Tool Name Prefixing
///
/// When MCP services are registered with a namespace, their tool names are automatically
/// prefixed to avoid conflicts. For example:
///
/// - Service "web_search" with tool "search" becomes "web_search-search"
/// - Service "file_search" with tool "search" becomes "file_search-search"
///
/// The Tool's `name` field is modified to include the prefix, so when tools are sent
/// to the LLM, they have unique names. The PegBoard maintains the mapping between
/// prefixed names and original names for routing tool calls back to the correct service.
///
/// ## Thread Safety
///
/// PegBoard is designed for concurrent access and internally uses `Arc` for cheap cloning.
/// It uses `DashMap` for all internal storage, providing lock-free concurrent access.
/// All methods use `&self` and can be called concurrently from multiple threads/tasks.
/// Simply clone the PegBoard to share it across async tasks.
///
/// ## Example
///
/// ```ignore
/// let pegboard = PegBoard::new();
/// let pegboard_clone = pegboard.clone(); // Cheap Arc clone
/// ```
#[derive(Clone)]
pub struct PegBoard {
    inner: Arc<InternalPegBoard>,
}

impl PegBoard {
    /// Creates a new empty PegBoard
    pub fn new() -> Self {
        Self {
            inner: Arc::new(InternalPegBoard::new()),
        }
    }

    /// Registers a service and automatically discovers all its tools.
    ///
    /// See `InternalPegBoard::add_service` for full documentation.
    pub async fn add_service(
        &self,
        namespace: Option<String>,
        service: RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
    ) -> Result<usize, PegBoardError> {
        self.inner.add_service(namespace, service).await
    }

    /// Manually registers a tool with an optional namespace and service index.
    pub fn register_tool(
        &self,
        namespace: Option<&str>,
        tool: Tool,
        service_idx: usize,
    ) -> Result<(), PegBoardError> {
        self.inner.register_tool(namespace, tool, service_idx)
    }

    /// Gets a tool by its name (prefixed if registered with namespace, original otherwise)
    pub fn get_tool(&self, tool_name: &str) -> Option<Tool> {
        self.inner.get_tool(tool_name)
    }

    /// Selects multiple tools by their names.
    pub fn select_tools(&self, tool_names: &[&str]) -> Option<Vec<Tool>> {
        self.inner.select_tools(tool_names)
    }

    /// Gets routing information for a tool by its name
    pub fn get_tool_route(&self, tool_name: &str) -> Option<(usize, String)> {
        self.inner.get_tool_route(tool_name)
    }

    /// Gets all tool names in a namespace (prefixed if namespace was used)
    pub fn list_tools_in_namespace(&self, namespace: &str) -> Vec<String> {
        self.inner.list_tools_in_namespace(namespace)
    }

    /// Gets all Tool objects in a namespace (with names as they appear to LLM)
    pub fn get_tools_in_namespace(&self, namespace: &str) -> Vec<Tool> {
        self.inner.get_tools_in_namespace(namespace)
    }

    /// Gets all registered tool names across all namespaces
    pub fn list_all_tools(&self) -> Vec<String> {
        self.inner.list_all_tools()
    }

    /// Gets all tools as a Vec
    pub fn get_all_tools(&self) -> Vec<Tool> {
        self.inner.get_all_tools()
    }

    /// Gets all registered namespaces
    pub fn list_namespaces(&self) -> Vec<String> {
        self.inner.list_namespaces()
    }

    /// Removes a tool by its prefixed name
    pub fn unregister_tool(&self, prefixed_name: &str) -> Result<(), PegBoardError> {
        self.inner.unregister_tool(prefixed_name)
    }

    /// Removes all tools in a namespace and the associated service
    pub fn unregister_namespace(&self, namespace: &str) -> Result<usize, PegBoardError> {
        self.inner.unregister_namespace(namespace)
    }

    /// Removes a service by its index and all associated tools
    pub fn unregister_service(&self, service_idx: usize) -> Result<usize, PegBoardError> {
        self.inner.unregister_service(service_idx)
    }

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

    /// Returns the number of registered services
    pub fn service_count(&self) -> usize {
        self.inner.service_count()
    }

    /// Returns the number of registered namespaces
    pub fn namespace_count(&self) -> usize {
        self.inner.namespace_count()
    }

    /// Calls a tool by its name (as seen by the LLM) with the given arguments.
    pub async fn call_tool(
        &self,
        tool_name: &str,
        arguments: Value,
    ) -> Result<CallToolResult, PegBoardError> {
        self.inner.call_tool(tool_name, arguments).await
    }
}

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

/// Errors that can occur when working with PegBoard
#[derive(Debug, thiserror::Error)]
pub enum PegBoardError {
    #[error("Tool '{0}' already exists")]
    ToolAlreadyExists(String),

    #[error("Tool '{0}' not found")]
    ToolNotFound(String),

    #[error("Invalid service index {index}, max is {max}")]
    InvalidServiceIndex { index: usize, max: usize },

    #[error("Namespace '{0}' already exists")]
    NamespaceAlreadyExists(String),

    #[error("Service error: {0}")]
    ServiceError(String),
}

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

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

    #[test]
    fn test_prefix_tool_name() {
        let prefixed = prefix_tool_name("web_search", "search");
        assert_eq!(prefixed, "web_search-search");

        let prefixed2 = prefix_tool_name("fs", "read_file");
        assert_eq!(prefixed2, "fs-read_file");
    }

    #[test]
    fn test_pegboard_register_and_get() {
        let pegboard = PegBoard::new();
        let tool = crate::get_tool::<TestParams, _, _>("search", Some("Search tool")).unwrap();

        // Initially empty
        assert_eq!(pegboard.tool_count(), 0);
        assert_eq!(pegboard.namespace_count(), 0);

        // Register fails with invalid service index
        assert!(
            pegboard
                .register_tool(Some("web"), tool.clone(), 0)
                .is_err()
        );

        // Without namespace also fails with invalid service index
        assert!(pegboard.register_tool(None, tool.clone(), 0).is_err());
    }

    #[test]
    fn test_pegboard_tool_name_prefixing() {
        // Get tool with original name "search"
        let original_tool =
            crate::get_tool::<TestParams, _, _>("search", Some("Search tool")).unwrap();
        assert_eq!(original_tool.name, "search");

        // After registration with namespace "web", the name should be "web-search"
        // But we can't test this without a valid service_idx
        // The prefixing logic is tested in test_prefix_tool_name
    }

    #[test]
    fn test_pegboard_namespace_operations() {
        let pegboard = PegBoard::new();

        // Initially empty
        assert_eq!(pegboard.list_namespaces().len(), 0);
        assert_eq!(pegboard.list_all_tools().len(), 0);
        assert_eq!(pegboard.get_all_tools().len(), 0);

        // List tools in non-existent namespace returns empty
        assert_eq!(pegboard.list_tools_in_namespace("nonexistent").len(), 0);
        assert_eq!(pegboard.get_tools_in_namespace("nonexistent").len(), 0);
    }

    #[test]
    fn test_pegboard_get_tool_methods() {
        let pegboard = PegBoard::new();

        // Get non-existent tool returns None (using prefixed name)
        assert!(pegboard.get_tool("web-search").is_none());
        assert!(pegboard.get_tool_route("web-search").is_none());
    }

    #[test]
    fn test_pegboard_unregister() {
        let pegboard = PegBoard::new();

        // Unregister non-existent tool should fail (using prefixed name)
        assert!(pegboard.unregister_tool("web-nonexistent").is_err());

        // Unregister non-existent namespace
        let result = pegboard.unregister_namespace("nonexistent");
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 0); // 0 tools removed
    }

    #[test]
    fn test_tool_route_structure() {
        let route = ToolRoute {
            service_index: 0,
            original_name: "search".to_string(),
        };

        assert_eq!(route.service_index, 0);
        assert_eq!(route.original_name, "search");
    }

    #[test]
    fn test_optional_namespace() {
        // Test that None and Some("") both mean no namespace
        let namespace_none: Option<String> = None;
        let namespace_empty = Some(String::new());

        // Both should normalize to empty string
        let ns1 = namespace_none.unwrap_or_default();
        let ns2 = namespace_empty.unwrap_or_default();

        assert_eq!(ns1, "");
        assert_eq!(ns2, "");
        assert!(!ns1.is_empty() == false);
        assert!(!ns2.is_empty() == false);
    }

    #[test]
    fn test_prefix_only_when_namespace_provided() {
        let original_name = "search";

        // With namespace - should prefix
        let with_ns = prefix_tool_name("web", original_name);
        assert_eq!(with_ns, "web-search");

        // Without namespace - would use original (tested via add_service logic)
        // The prefix_tool_name function always prefixes, but add_service checks has_namespace first
    }

    // Note: Integration test for add_service() will be added once we have a proper
    // way to create RunningService instances for testing. For now, the logic is
    // validated through unit tests of individual components.
}