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;

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<mcp_id, (namespace, service)>
    /// DashMap provides lock-free concurrent access
    services: DashMap<
        String,
        (
            String,
            RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
        ),
    >,

    /// Maps prefixed tool name to (service_id, 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>>,

    /// Tool discovery: maps tool name to list of related tool names
    /// When a tool is used, these related tools can be discovered and added dynamically
    tool_discovery: DashMap<String, Vec<String>>,

    /// MCP discovery: maps mcp_id to list of related mcp_ids
    /// When an MCP service is used, these related MCPs can be discovered
    mcp_discovery: DashMap<String, Vec<String>>,
}

/// Routing information for a tool
#[derive(Debug, Clone)]
struct ToolRoute {
    /// ID of the service that provides this tool
    service_id: String,
    /// 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(),
            tool_routing: DashMap::new(),
            namespace_tools: DashMap::new(),
            tool_discovery: DashMap::new(),
            mcp_discovery: 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
    /// * `mcp_id` - Unique identifier for this MCP service (used for discovery and routing)
    /// * `namespace` - Optional namespace prefix. Use None or empty string if no conflicts expected
    /// * `service` - The MCP service to register
    ///
    /// # Example
    /// ```ignore
    /// // With namespace (prefixing enabled)
    /// pegboard.add_service("web-mcp".to_string(), Some("web".to_string()), service).await?;
    /// // Tool "search" is now available as "web-search"
    ///
    /// // Without namespace (no prefixing)
    /// pegboard.add_service("file-mcp".to_string(), None, service).await?;
    /// // Tool "search" keeps its original name "search"
    /// ```
    pub async fn add_service(
        &self,
        mcp_id: String,
        namespace: Option<String>,
        service: RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
    ) -> Result<(), PegBoardError> {
        // Check if mcp_id already exists
        if self.services.contains_key(&mcp_id) {
            return Err(PegBoardError::DuplicateMcpId(mcp_id));
        }

        // 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));
        }

        // 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(mcp_id.clone(), (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_id + original name)
            self.tool_routing.insert(
                final_name.clone(),
                ToolRoute {
                    service_id: mcp_id.clone(),
                    original_name,
                },
            );

            registered_tool_names.push(final_name);
        }

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

        Ok(())
    }

    /// Manually registers a tool with an optional namespace and service ID.
    /// 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_id: &str,
    ) -> Result<(), PegBoardError> {
        // Check if service ID is valid
        if !self.services.contains_key(service_id) {
            return Err(PegBoardError::InvalidServiceId(service_id.to_string()));
        }

        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_id: service_id.to_string(),
                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_id, original_tool_name) for routing the call
    pub fn get_tool_route(&self, tool_name: &str) -> Option<(String, String)> {
        self.tool_routing.get(tool_name).map(|entry| {
            let route = entry.value();
            (route.service_id.clone(), 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_id_to_remove = self
            .services
            .iter()
            .find(|entry| entry.value().0 == namespace)
            .map(|entry| entry.key().clone());

        if let Some(id) = service_id_to_remove {
            self.services.remove(&id);
        }

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

    /// Removes a service by its ID and all associated tools
    pub fn unregister_service(&self, service_id: &str) -> Result<usize, PegBoardError> {
        // Get the namespace for this service
        let namespace = self
            .services
            .get(service_id)
            .map(|entry| entry.value().0.clone())
            .ok_or(PegBoardError::InvalidServiceId(service_id.to_string()))?;

        // Find and remove all tools associated with this service ID
        let tools_to_remove: Vec<String> = self
            .tool_routing
            .iter()
            .filter(|entry| entry.value().service_id == service_id)
            .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_id);

        // 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_id, 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_id)
            .ok_or(PegBoardError::InvalidServiceId(service_id.clone()))?;
        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)))
    }

    /// Registers a tool discovery relationship.
    /// When `tool_name` is used, the related tools can be discovered via `discover_tool()`.
    /// This replaces any existing discovery relationships for the tool.
    ///
    /// # Arguments
    /// * `tool_name` - The tool name (prefixed if namespace was used during registration)
    /// * `related_tools` - List of related tool names that should be discovered when this tool is used
    ///
    /// # Example
    /// ```ignore
    /// // When "web-search" is used, suggest "web-fetch" and "web-parse"
    /// pegboard.register_tool_discovery(
    ///     "web-search",
    ///     vec!["web-fetch".to_string(), "web-parse".to_string()]
    /// );
    /// ```
    pub fn register_tool_discovery(&self, tool_name: &str, related_tools: Vec<String>) {
        self.tool_discovery
            .insert(tool_name.to_string(), related_tools);
    }

    /// Discovers related tools for a given tool name.
    /// Returns the full Tool objects for all related tools that are registered.
    /// Returns an empty Vec if no discovery relationships exist for this tool.
    ///
    /// # Arguments
    /// * `tool_name` - The tool name to discover related tools for
    ///
    /// # Returns
    /// * `Vec<Tool>` - List of related Tool objects (empty if none registered)
    ///
    /// # Example
    /// ```ignore
    /// let related_tools = pegboard.discover_tool("web-search");
    /// // Returns Tool objects for "web-fetch" and "web-parse" if they're registered
    /// ```
    pub fn discover_tool(&self, tool_name: &str) -> Vec<Tool> {
        self.tool_discovery
            .get(tool_name)
            .map(|entry| {
                entry
                    .value()
                    .iter()
                    .filter_map(|name| self.get_tool(name))
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Registers an MCP discovery relationship.
    /// When `mcp_id` is used, the related MCPs can be discovered via `discover_mcp()`.
    /// This replaces any existing discovery relationships for the MCP.
    ///
    /// # Arguments
    /// * `mcp_id` - The MCP ID to register discovery for
    /// * `related_mcps` - List of related MCP IDs that should be discovered when this MCP is used
    ///
    /// # Example
    /// ```ignore
    /// // When "web-mcp" is used, suggest "html-parser-mcp" and "image-fetcher-mcp"
    /// pegboard.register_mcp_discovery(
    ///     "web-mcp",
    ///     vec!["html-parser-mcp".to_string(), "image-fetcher-mcp".to_string()]
    /// );
    /// ```
    pub fn register_mcp_discovery(&self, mcp_id: &str, related_mcps: Vec<String>) {
        self.mcp_discovery.insert(mcp_id.to_string(), related_mcps);
    }

    /// Discovers related MCP IDs for a given MCP ID.
    /// Returns a list of related MCP IDs.
    /// Returns an empty Vec if no discovery relationships exist for this MCP.
    ///
    /// # Arguments
    /// * `mcp_id` - The MCP ID to discover related MCPs for
    ///
    /// # Returns
    /// * `Vec<String>` - List of related MCP IDs (empty if none registered)
    ///
    /// # Example
    /// ```ignore
    /// let related_mcps = pegboard.discover_mcp("web-mcp");
    /// // Returns ["html-parser-mcp", "image-fetcher-mcp"]
    /// ```
    pub fn discover_mcp(&self, mcp_id: &str) -> Vec<String> {
        self.mcp_discovery
            .get(mcp_id)
            .map(|entry| entry.value().clone())
            .unwrap_or_default()
    }
}

/// 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,
        mcp_id: String,
        namespace: Option<String>,
        service: RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
    ) -> Result<(), PegBoardError> {
        self.inner.add_service(mcp_id, namespace, service).await
    }

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

    /// 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<(String, 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 ID and all associated tools
    pub fn unregister_service(&self, service_id: &str) -> Result<usize, PegBoardError> {
        self.inner.unregister_service(service_id)
    }

    /// 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
    }

    /// Registers a tool discovery relationship.
    pub fn register_tool_discovery(&self, tool_name: &str, related_tools: Vec<String>) {
        self.inner.register_tool_discovery(tool_name, related_tools)
    }

    /// Discovers related tools for a given tool name.
    pub fn discover_tool(&self, tool_name: &str) -> Vec<Tool> {
        self.inner.discover_tool(tool_name)
    }

    /// Registers an MCP discovery relationship.
    pub fn register_mcp_discovery(&self, mcp_id: &str, related_mcps: Vec<String>) {
        self.inner.register_mcp_discovery(mcp_id, related_mcps)
    }

    /// Discovers related MCP IDs for a given MCP ID.
    pub fn discover_mcp(&self, mcp_id: &str) -> Vec<String> {
        self.inner.discover_mcp(mcp_id)
    }
}

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 ID '{0}'")]
    InvalidServiceId(String),

    #[error("MCP ID '{0}' already exists")]
    DuplicateMcpId(String),

    #[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 ID
        assert!(
            pegboard
                .register_tool(Some("web"), tool.clone(), "invalid-service-id")
                .is_err()
        );

        // Without namespace also fails with invalid service ID
        assert!(
            pegboard
                .register_tool(None, tool.clone(), "invalid-service-id")
                .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_id: "test-mcp".to_string(),
            original_name: "search".to_string(),
        };

        assert_eq!(route.service_id, "test-mcp");
        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.

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

        // Initially, discovering a non-existent tool returns empty
        let discovered = pegboard.discover_tool("web-search");
        assert_eq!(discovered.len(), 0);

        // Register discovery relationship
        pegboard.register_tool_discovery(
            "web-search",
            vec!["web-fetch".to_string(), "web-parse".to_string()],
        );

        // Discover should return empty if related tools don't exist
        let discovered = pegboard.discover_tool("web-search");
        assert_eq!(discovered.len(), 0);

        // Replace discovery relationship
        pegboard.register_tool_discovery("web-search", vec!["other-tool".to_string()]);

        // Discovery is replaced (not appended)
        let discovered = pegboard.discover_tool("web-search");
        assert_eq!(discovered.len(), 0); // Still empty since tools don't exist
    }

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

        // Initially, discovering a non-existent MCP returns empty
        let discovered = pegboard.discover_mcp("web-mcp");
        assert_eq!(discovered.len(), 0);

        // Register MCP discovery relationship
        pegboard.register_mcp_discovery(
            "web-mcp",
            vec![
                "html-parser-mcp".to_string(),
                "image-fetcher-mcp".to_string(),
            ],
        );

        // Discover returns the registered related MCPs
        let discovered = pegboard.discover_mcp("web-mcp");
        assert_eq!(discovered.len(), 2);
        assert!(discovered.contains(&"html-parser-mcp".to_string()));
        assert!(discovered.contains(&"image-fetcher-mcp".to_string()));

        // Replace MCP discovery relationship
        pegboard.register_mcp_discovery("web-mcp", vec!["other-mcp".to_string()]);

        // Discovery is replaced (not appended)
        let discovered = pegboard.discover_mcp("web-mcp");
        assert_eq!(discovered.len(), 1);
        assert_eq!(discovered[0], "other-mcp");
    }
}