mcp/

protocol.rs

//! Tool protocol abstraction layer.
//!
//! This module provides a flexible abstraction for connecting agents to various tool protocols.
//! It supports multiple standards including MCP (Model Context Protocol), custom function calling,
//! Memory persistence, and allows users to implement their own tool communication mechanisms.
//!
//! # Architecture
//!
//! **Single Protocol** (traditional):
//! ```text
//! Agent → ToolRegistry → ToolProtocol → Single Tool Source
//! ```
//!
//! **Multi-Protocol** (new in 0.5.0):
//! ```text
//! Agent → ToolRegistry → [Protocol1, Protocol2, Protocol3]
//!         (routing map)     ↓          ↓          ↓
//!                        Local      YouTube    GitHub
//!                        Tools      Server     Server
//! ```
//!
//! # Key Components
//!
//! - **ToolProtocol trait**: Define how tools are executed, discovered, and described
//! - **ToolRegistry**: Single or multi-protocol tool aggregation with transparent routing
//! - **ToolMetadata**: Tool identity, description, parameters
//! - **ToolParameter**: Type-safe parameter definitions with validation
//! - **ToolResult**: Structured tool execution results
//! - **Tool**: Runtime tool instance bound to a protocol
//!
//! # Single Protocol Example
//!
//! ```rust,ignore
//! use std::sync::Arc;
//! use mcp::protocol::{ToolMetadata, ToolProtocol, ToolRegistry, ToolResult};
//!
//! struct MyProtocol;
//!
//! #[async_trait::async_trait]
//! impl ToolProtocol for MyProtocol {
//!     async fn execute(
//!         &self,
//!         _tool_name: &str,
//!         _parameters: serde_json::Value,
//!     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
//!         Ok(ToolResult::success(serde_json::json!({"ok": true})))
//!     }
//!     async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
//!         Ok(vec![ToolMetadata::new("demo", "Demo tool")])
//!     }
//!     async fn get_tool_metadata(&self, _: &str) -> Result<ToolMetadata, Box<dyn std::error::Error + Send + Sync>> {
//!         Ok(ToolMetadata::new("demo", "Demo tool"))
//!     }
//!     fn protocol_name(&self) -> &str { "demo" }
//! }
//!
//! # #[tokio::main]
//! # async fn main() {
//! let protocol = Arc::new(MyProtocol);
//! let mut registry = ToolRegistry::new(protocol);
//! let _ = registry.discover_tools_from_primary().await;
//! # }
//! ```
//!
//! # Multi-Protocol Example
//!
//! ```rust,ignore
//! use mcp::client::McpClientProtocol;
//! use mcp::protocol::ToolRegistry;
//!
//! # #[tokio::main]
//! # async fn main() {
//! let mut registry = ToolRegistry::empty();
//! let _ = registry
//!     .add_protocol("remote", std::sync::Arc::new(McpClientProtocol::new("http://localhost:8080".to_string())))
//!     .await;
//! # }
//! ```

use crate::resources::ResourceMetadata;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::sync::Arc;

/// Provider-agnostic tool definition suitable for LLM native tool calling.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    /// Tool name as it will appear in the API `tools` array.
    pub name: String,
    /// Human-readable description surfaced to the model.
    pub description: String,
    /// JSON Schema object describing accepted parameters.
    pub parameters_schema: serde_json::Value,
}

/// Represents the result of a tool execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
    /// Whether the tool execution was successful
    pub success: bool,
    /// The output data from the tool
    pub output: serde_json::Value,
    /// Optional error message if execution failed
    pub error: Option<String>,
    /// Metadata about the execution (timing, cost, etc.)
    pub metadata: HashMap<String, serde_json::Value>,
}

impl ToolResult {
    /// Convenience constructor for successful tool execution.
    pub fn success(output: serde_json::Value) -> Self {
        Self {
            success: true,
            output,
            error: None,
            metadata: HashMap::new(),
        }
    }

    /// Convenience constructor for failed tool execution.
    pub fn failure(error: String) -> Self {
        Self {
            success: false,
            output: serde_json::Value::Null,
            error: Some(error),
            metadata: HashMap::new(),
        }
    }

    /// Attach protocol or application specific metadata to the result.
    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.metadata.insert(key.into(), value);
        self
    }
}

/// Defines the type of a tool parameter
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ToolParameterType {
    /// UTF-8 string value.
    String,
    /// Floating-point number.
    Number,
    /// Integer number.
    Integer,
    /// Boolean value.
    Boolean,
    /// JSON array.
    Array,
    /// JSON object.
    Object,
}

/// Defines a parameter for a tool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolParameter {
    /// Parameter name.
    pub name: String,
    /// Parameter value type.
    #[serde(rename = "type")]
    pub param_type: ToolParameterType,
    /// Human-readable description.
    pub description: Option<String>,
    /// Whether the caller must provide this parameter.
    pub required: bool,
    /// Optional default value.
    pub default: Option<serde_json::Value>,
    /// For array types, specifies the type of items
    pub items: Option<Box<ToolParameterType>>,
    /// For object types, specifies nested properties
    pub properties: Option<HashMap<String, ToolParameter>>,
}

impl ToolParameter {
    /// Define a new tool parameter with the provided name and type.
    pub fn new(name: impl Into<String>, param_type: ToolParameterType) -> Self {
        Self {
            name: name.into(),
            param_type,
            description: None,
            required: false,
            default: None,
            items: None,
            properties: None,
        }
    }

    /// Add a human readable description that will surface in generated schemas.
    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.description = Some(description.into());
        self
    }

    /// Mark the argument as required.
    pub fn required(mut self) -> Self {
        self.required = true;
        self
    }

    /// Provide a default value that will be used when the LLM omits the parameter.
    pub fn with_default(mut self, default: serde_json::Value) -> Self {
        self.default = Some(default);
        self
    }

    /// For array parameters, declare the type of the contained items.
    pub fn with_items(mut self, item_type: ToolParameterType) -> Self {
        self.items = Some(Box::new(item_type));
        self
    }

    /// For object parameters, describe the nested properties.
    pub fn with_properties(mut self, properties: HashMap<String, ToolParameter>) -> Self {
        self.properties = Some(properties);
        self
    }

    /// Converts this parameter to a JSON Schema snippet suitable for native tool-calling.
    ///
    /// The generated schema follows [JSON Schema draft-07] conventions accepted by all major
    /// providers (OpenAI, Anthropic, Grok, Gemini).
    ///
    /// | `ToolParameterType` | Schema produced |
    /// |---|---|
    /// | `String` | `{"type":"string"}` |
    /// | `Number` | `{"type":"number"}` |
    /// | `Integer` | `{"type":"integer"}` |
    /// | `Boolean` | `{"type":"boolean"}` |
    /// | `Array` | `{"type":"array","items":{...}}` |
    /// | `Object` | `{"type":"object","properties":{...},"required":[...]}` |
    ///
    /// # Example
    ///
    /// ```rust
    /// use mcp::{ToolParameter, ToolParameterType};
    ///
    /// let param = ToolParameter::new("query", ToolParameterType::String)
    ///     .with_description("Search query string")
    ///     .required();
    ///
    /// let schema = param.to_json_schema();
    /// assert_eq!(schema["type"], "string");
    /// assert_eq!(schema["description"], "Search query string");
    /// ```
    pub fn to_json_schema(&self) -> serde_json::Value {
        let mut schema = match &self.param_type {
            ToolParameterType::String => serde_json::json!({"type": "string"}),
            ToolParameterType::Number => serde_json::json!({"type": "number"}),
            ToolParameterType::Integer => serde_json::json!({"type": "integer"}),
            ToolParameterType::Boolean => serde_json::json!({"type": "boolean"}),
            ToolParameterType::Array => {
                let items_schema = self
                    .items
                    .as_ref()
                    .map(|t| param_type_to_schema(t))
                    .unwrap_or_else(|| serde_json::json!({}));
                serde_json::json!({"type": "array", "items": items_schema})
            }
            ToolParameterType::Object => {
                if let Some(props) = &self.properties {
                    let properties: serde_json::Map<String, serde_json::Value> = props
                        .iter()
                        .map(|(k, v)| (k.clone(), v.to_json_schema()))
                        .collect();
                    let required: Vec<&str> = props
                        .values()
                        .filter(|p| p.required)
                        .map(|p| p.name.as_str())
                        .collect();
                    serde_json::json!({
                        "type": "object",
                        "properties": properties,
                        "required": required
                    })
                } else {
                    serde_json::json!({"type": "object"})
                }
            }
        };

        // Attach description if present
        if let Some(desc) = &self.description {
            schema["description"] = serde_json::Value::String(desc.clone());
        }

        schema
    }
}

/// Convert a bare [`ToolParameterType`] to a minimal JSON Schema value (no description).
///
/// Used internally when building array `items` schemas.
fn param_type_to_schema(t: &ToolParameterType) -> serde_json::Value {
    match t {
        ToolParameterType::String => serde_json::json!({"type": "string"}),
        ToolParameterType::Number => serde_json::json!({"type": "number"}),
        ToolParameterType::Integer => serde_json::json!({"type": "integer"}),
        ToolParameterType::Boolean => serde_json::json!({"type": "boolean"}),
        ToolParameterType::Array => serde_json::json!({"type": "array"}),
        ToolParameterType::Object => serde_json::json!({"type": "object"}),
    }
}

/// Metadata about a tool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolMetadata {
    /// Tool name.
    pub name: String,
    /// Human-readable description.
    pub description: String,
    /// Parameter schema entries.
    pub parameters: Vec<ToolParameter>,
    /// Additional metadata specific to the protocol
    pub protocol_metadata: HashMap<String, serde_json::Value>,
}

impl ToolMetadata {
    /// Create metadata with the supplied identifier and description.
    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            parameters: Vec::new(),
            protocol_metadata: HashMap::new(),
        }
    }

    /// Append a parameter definition to the tool metadata.
    pub fn with_parameter(mut self, param: ToolParameter) -> Self {
        self.parameters.push(param);
        self
    }

    /// Add protocol specific metadata (e.g. MCP capability flags).
    pub fn with_protocol_metadata(
        mut self,
        key: impl Into<String>,
        value: serde_json::Value,
    ) -> Self {
        self.protocol_metadata.insert(key.into(), value);
        self
    }

    /// Builds a [`ToolDefinition`] (JSON Schema) from this metadata.
    ///
    /// The resulting [`ToolDefinition`] is ready to be passed to
    /// [`ClientWrapper::send_message`](crate::client_wrapper::ClientWrapper::send_message)
    /// as part of the `tools` slice.  Parameters marked `required` are collected into the
    /// JSON Schema `"required"` array automatically.
    ///
    /// # Example
    ///
    /// ```rust
    /// use mcp::{ToolMetadata, ToolParameter, ToolParameterType};
    ///
    /// let meta = ToolMetadata::new("calculator", "Evaluates a math expression")
    ///     .with_parameter(
    ///         ToolParameter::new("expression", ToolParameterType::String)
    ///             .with_description("The expression to evaluate")
    ///             .required(),
    ///     );
    ///
    /// let def = meta.to_tool_definition();
    /// assert_eq!(def.name, "calculator");
    /// let schema = &def.parameters_schema;
    /// assert_eq!(schema["type"], "object");
    /// assert!(schema["properties"]["expression"].is_object());
    /// ```
    pub fn to_tool_definition(&self) -> ToolDefinition {
        let mut properties = serde_json::Map::new();
        let mut required: Vec<String> = Vec::new();

        for param in &self.parameters {
            properties.insert(param.name.clone(), param.to_json_schema());
            if param.required {
                required.push(param.name.clone());
            }
        }

        let parameters_schema = serde_json::json!({
            "type": "object",
            "properties": properties,
            "required": required
        });

        ToolDefinition {
            name: self.name.clone(),
            description: self.description.clone(),
            parameters_schema,
        }
    }
}

/// Trait for implementing tool execution protocols
#[async_trait]
pub trait ToolProtocol: Send + Sync {
    /// Execute a tool with the given parameters
    async fn execute(
        &self,
        tool_name: &str,
        parameters: serde_json::Value,
    ) -> Result<ToolResult, Box<dyn Error + Send + Sync>>;

    /// Get metadata about available tools
    async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn Error + Send + Sync>>;

    /// Get metadata about a specific tool
    async fn get_tool_metadata(
        &self,
        tool_name: &str,
    ) -> Result<ToolMetadata, Box<dyn Error + Send + Sync>>;

    /// Protocol identifier (e.g., "mcp", "custom", "openai-functions")
    fn protocol_name(&self) -> &str;

    /// Initialize/connect to the tool protocol
    async fn initialize(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
        Ok(())
    }

    /// Cleanup/disconnect from the tool protocol
    async fn shutdown(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
        Ok(())
    }

    /// List available resources (MCP Resource support)
    ///
    /// Resources are application-provided contextual data that agents can read.
    /// This method is optional and defaults to returning an empty list.
    async fn list_resources(&self) -> Result<Vec<ResourceMetadata>, Box<dyn Error + Send + Sync>> {
        Ok(Vec::new())
    }

    /// Read the content of a resource by URI (MCP Resource support)
    ///
    /// This method is optional and defaults to returning NotFound.
    async fn read_resource(&self, uri: &str) -> Result<String, Box<dyn Error + Send + Sync>> {
        Err(format!("Resource not found: {}", uri).into())
    }

    /// Check if this protocol supports resources
    fn supports_resources(&self) -> bool {
        false
    }
}

/// Error types for tool operations
#[derive(Debug, Clone)]
pub enum ToolError {
    /// Requested tool is not registered in the current registry/protocol.
    NotFound(String),
    /// Tool execution completed with an application level failure.
    ExecutionFailed(String),
    /// The provided JSON parameters failed validation or deserialization.
    InvalidParameters(String),
    /// A lower level protocol/transport error occurred.
    ProtocolError(String),
}

impl fmt::Display for ToolError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ToolError::NotFound(name) => write!(f, "Tool not found: {}", name),
            ToolError::ExecutionFailed(msg) => write!(f, "Tool execution failed: {}", msg),
            ToolError::InvalidParameters(msg) => write!(f, "Invalid parameters: {}", msg),
            ToolError::ProtocolError(msg) => write!(f, "Protocol error: {}", msg),
        }
    }
}

impl Error for ToolError {}

/// A tool that can be used by agents
pub struct Tool {
    /// Metadata describing the tool interface.
    metadata: ToolMetadata,
    /// Underlying protocol implementation that actually executes the tool.
    protocol: Arc<dyn ToolProtocol>,
}

impl Tool {
    /// Create a new tool bound to the supplied protocol implementation.
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        protocol: Arc<dyn ToolProtocol>,
    ) -> Self {
        Self {
            metadata: ToolMetadata::new(name, description),
            protocol,
        }
    }

    /// Add a parameter definition to the tool builder.
    pub fn with_parameter(mut self, param: ToolParameter) -> Self {
        self.metadata.parameters.push(param);
        self
    }

    /// Attach protocol specific metadata to the tool builder.
    pub fn with_protocol_metadata(
        mut self,
        key: impl Into<String>,
        value: serde_json::Value,
    ) -> Self {
        self.metadata.protocol_metadata.insert(key.into(), value);
        self
    }

    /// Borrow the static metadata for the tool.
    pub fn metadata(&self) -> &ToolMetadata {
        &self.metadata
    }

    /// Execute the tool using the configured protocol.
    pub async fn execute(
        &self,
        parameters: serde_json::Value,
    ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
        self.protocol.execute(&self.metadata.name, parameters).await
    }
}

/// Registry for managing tools available to agents
///
/// Supports single or multiple tool protocols, enabling agents to transparently
/// access tools from multiple sources (local functions, MCP servers, etc.)
///
/// # Single Protocol
///
/// ```ignore
/// use async_trait::async_trait;
/// use mcp::{ToolMetadata, ToolProtocol, ToolRegistry, ToolResult};
/// use std::sync::Arc;
///
/// struct DemoProtocol;
///
/// #[async_trait]
/// impl ToolProtocol for DemoProtocol {
///     async fn execute(
///         &self,
///         _tool_name: &str,
///         _parameters: serde_json::Value,
///     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
///         Ok(ToolResult::success(serde_json::json!({"ok": true})))
///     }
///
///     async fn list_tools(
///         &self,
///     ) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
///         Ok(vec![])
///     }
/// }
///
/// let protocol = Arc::new(DemoProtocol);
/// let registry = ToolRegistry::new(protocol);
/// ```
///
/// # Multiple Protocols
///
/// ```ignore
/// use async_trait::async_trait;
/// use mcp::{ToolMetadata, ToolProtocol, ToolRegistry, ToolResult};
/// use std::sync::Arc;
///
/// struct LocalProtocol;
/// struct RemoteProtocol;
///
/// #[async_trait]
/// impl ToolProtocol for LocalProtocol {
///     async fn execute(
///         &self,
///         _tool_name: &str,
///         _parameters: serde_json::Value,
///     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
///         Ok(ToolResult::success(serde_json::json!({"source": "local"})))
///     }
///
///     async fn list_tools(
///         &self,
///     ) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
///         Ok(vec![])
///     }
/// }
///
/// #[async_trait]
/// impl ToolProtocol for RemoteProtocol {
///     async fn execute(
///         &self,
///         _tool_name: &str,
///         _parameters: serde_json::Value,
///     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
///         Ok(ToolResult::success(serde_json::json!({"source": "remote"})))
///     }
///
///     async fn list_tools(
///         &self,
///     ) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
///         Ok(vec![])
///     }
/// }
///
/// # async {
/// let mut registry = ToolRegistry::empty();
///
/// // Add local tools
/// registry.add_protocol(
///     "local",
///     Arc::new(LocalProtocol)
/// ).await.ok();
///
/// // Add remote MCP server
/// registry.add_protocol(
///     "youtube",
///     Arc::new(RemoteProtocol)
/// ).await.ok();
///
/// // Agent transparently accesses both
/// # };
/// ```
pub struct ToolRegistry {
    /// All discovered tools from all protocols
    tools: HashMap<String, Tool>,
    /// Mapping of tool_name -> protocol_name for routing
    tool_to_protocol: HashMap<String, String>,
    /// All registered protocols
    protocols: HashMap<String, Arc<dyn ToolProtocol>>,
    /// Primary protocol (for backwards compatibility with single-protocol code)
    primary_protocol: Option<Arc<dyn ToolProtocol>>,
}

impl ToolRegistry {
    /// Build a registry powered by a single protocol implementation.
    ///
    /// This is the traditional single-protocol mode. Use `empty()` and `add_protocol()`
    /// for multi-protocol support.
    pub fn new(protocol: Arc<dyn ToolProtocol>) -> Self {
        Self {
            tools: HashMap::new(),
            tool_to_protocol: HashMap::new(),
            protocols: {
                let mut m = HashMap::new();
                m.insert("primary".to_string(), protocol.clone());
                m
            },
            primary_protocol: Some(protocol),
        }
    }

    /// Create an empty registry ready to accept multiple protocols.
    ///
    /// Use `add_protocol()` to register protocols.
    pub fn empty() -> Self {
        Self {
            tools: HashMap::new(),
            tool_to_protocol: HashMap::new(),
            protocols: HashMap::new(),
            primary_protocol: None,
        }
    }

    /// Register a protocol and discover its tools.
    ///
    /// # Arguments
    ///
    /// * `protocol_name` - Unique identifier for this protocol (e.g., "local", "youtube", "github")
    /// * `protocol` - The ToolProtocol implementation
    ///
    /// # Tool Discovery
    ///
    /// This method calls `protocol.list_tools()` to discover available tools
    /// and automatically registers them in the registry.
    ///
    /// # Conflicts
    ///
    /// If a tool with the same name already exists, it will be replaced.
    /// The new protocol's tool takes precedence.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use async_trait::async_trait;
    /// use mcp::{ToolMetadata, ToolProtocol, ToolRegistry, ToolResult};
    /// use std::sync::Arc;
    ///
    /// struct MemoryServerProtocol;
    ///
    /// #[async_trait]
    /// impl ToolProtocol for MemoryServerProtocol {
    ///     async fn execute(
    ///         &self,
    ///         _tool_name: &str,
    ///         _parameters: serde_json::Value,
    ///     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
    ///         Ok(ToolResult::success(serde_json::json!({"ok": true})))
    ///     }
    ///
    ///     async fn list_tools(
    ///         &self,
    ///     ) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
    ///         Ok(vec![])
    ///     }
    /// }
    ///
    /// # async {
    /// let mut registry = ToolRegistry::empty();
    /// registry.add_protocol(
    ///     "memory_server",
    ///     Arc::new(MemoryServerProtocol)
    /// ).await?;
    /// # Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
    /// # };
    /// ```
    pub async fn add_protocol(
        &mut self,
        protocol_name: &str,
        protocol: Arc<dyn ToolProtocol>,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        // Discover tools from this protocol
        let discovered_tools = protocol.list_tools().await?;

        // Register the protocol
        self.protocols
            .insert(protocol_name.to_string(), protocol.clone());

        // Register discovered tools
        for tool_meta in discovered_tools {
            let tool_name = tool_meta.name.clone();

            // Create a Tool that routes through this protocol
            let tool = Tool::new(
                tool_name.clone(),
                tool_meta.description.clone(),
                protocol.clone(),
            );

            // Copy over any additional parameters and metadata
            let mut tool = tool;
            for param in &tool_meta.parameters {
                tool = tool.with_parameter(param.clone());
            }
            for (key, value) in &tool_meta.protocol_metadata {
                tool = tool.with_protocol_metadata(key.clone(), value.clone());
            }

            // Register the tool and its routing
            self.tools.insert(tool_name.clone(), tool);
            self.tool_to_protocol
                .insert(tool_name, protocol_name.to_string());
        }

        Ok(())
    }

    /// Remove a protocol and all its tools from the registry.
    pub fn remove_protocol(&mut self, protocol_name: &str) {
        self.protocols.remove(protocol_name);

        // Collect tool names to remove
        let tools_to_remove: Vec<String> = self
            .tool_to_protocol
            .iter()
            .filter(|(_, pn)| *pn == protocol_name)
            .map(|(tn, _)| tn.clone())
            .collect();

        // Remove the tools
        for tool_name in tools_to_remove {
            self.tools.remove(&tool_name);
            self.tool_to_protocol.remove(&tool_name);
        }
    }

    /// Insert or replace a tool definition (for manual tool registration).
    pub fn add_tool(&mut self, tool: Tool) {
        self.tools.insert(tool.metadata.name.clone(), tool);
    }

    /// Remove a tool by name returning the owned entry if present.
    pub fn remove_tool(&mut self, name: &str) -> Option<Tool> {
        self.tool_to_protocol.remove(name);
        self.tools.remove(name)
    }

    /// Borrow a tool by name.
    pub fn get_tool(&self, name: &str) -> Option<&Tool> {
        self.tools.get(name)
    }

    /// List metadata for registered tools (iteration order follows the underlying map).
    pub fn list_tools(&self) -> Vec<&ToolMetadata> {
        self.tools.values().map(|t| &t.metadata).collect()
    }

    /// Discover tools from the primary protocol (for single-protocol registries).
    ///
    /// This is useful after registering tools with the protocol to populate the registry.
    /// For multi-protocol registries, use `add_protocol()` instead.
    pub async fn discover_tools_from_primary(
        &mut self,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        if let Some(protocol) = &self.primary_protocol {
            let discovered_tools = protocol.list_tools().await?;
            for tool_meta in discovered_tools {
                let tool_name = tool_meta.name.clone();
                let tool = Tool::new(
                    tool_name.clone(),
                    tool_meta.description.clone(),
                    protocol.clone(),
                );

                // Copy over parameters and metadata
                let mut tool = tool;
                for param in &tool_meta.parameters {
                    tool = tool.with_parameter(param.clone());
                }
                for (key, value) in &tool_meta.protocol_metadata {
                    tool = tool.with_protocol_metadata(key.clone(), value.clone());
                }

                self.tools.insert(tool_name.clone(), tool);
                self.tool_to_protocol
                    .insert(tool_name, "primary".to_string());
            }
            Ok(())
        } else {
            Err("No primary protocol available".into())
        }
    }

    /// Get which protocol handles a specific tool.
    pub fn get_tool_protocol(&self, tool_name: &str) -> Option<&str> {
        self.tool_to_protocol.get(tool_name).map(|s| s.as_str())
    }

    /// Get all registered protocol names.
    pub fn list_protocols(&self) -> Vec<&str> {
        self.protocols.keys().map(|s| s.as_str()).collect()
    }

    /// Execute a named tool with serialized parameters.
    pub async fn execute_tool(
        &self,
        tool_name: &str,
        parameters: serde_json::Value,
    ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
        let tool = self
            .tools
            .get(tool_name)
            .ok_or_else(|| ToolError::NotFound(tool_name.to_string()))?;

        tool.execute(parameters).await
    }

    /// Borrow the primary protocol implementation (for single-protocol mode).
    ///
    /// Returns None if registry was created with `empty()` or has multiple protocols.
    pub fn protocol(&self) -> Option<&Arc<dyn ToolProtocol>> {
        self.primary_protocol.as_ref()
    }

    /// Returns all registered tools as [`ToolDefinition`]s ready to pass to the LLM.
    ///
    /// Iterates over every tool in the registry, calling
    /// [`ToolMetadata::to_tool_definition`] on each, and returns the results as a `Vec`.
    /// An empty `Vec` is returned when no tools have been registered.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use async_trait::async_trait;
    /// use mcp::{Tool, ToolDefinition, ToolMetadata, ToolProtocol, ToolRegistry, ToolResult};
    /// use std::sync::Arc;
    ///
    /// struct DemoProtocol;
    ///
    /// #[async_trait]
    /// impl ToolProtocol for DemoProtocol {
    ///     async fn execute(
    ///         &self,
    ///         _tool_name: &str,
    ///         _parameters: serde_json::Value,
    ///     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
    ///         Ok(ToolResult::success(serde_json::json!({"ok": true})))
    ///     }
    ///
    ///     async fn list_tools(
    ///         &self,
    ///     ) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
    ///         Ok(vec![])
    ///     }
    /// }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let protocol = Arc::new(DemoProtocol);
    /// let mut registry = ToolRegistry::new(protocol.clone());
    ///
    /// let tool = Tool::new("calculator", "Evaluates math", protocol);
    /// registry.add_tool(tool);
    ///
    /// let defs = registry.to_tool_definitions();
    /// assert_eq!(defs.len(), 1);
    /// assert_eq!(defs[0].name, "calculator");
    /// # }
    /// ```
    pub fn to_tool_definitions(&self) -> Vec<ToolDefinition> {
        self.tools
            .values()
            .map(|t| t.metadata.to_tool_definition())
            .collect()
    }
}

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

    struct MockProtocol;

    #[async_trait]
    impl ToolProtocol for MockProtocol {
        async fn execute(
            &self,
            tool_name: &str,
            _parameters: serde_json::Value,
        ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
            Ok(ToolResult::success(serde_json::json!({
                "tool": tool_name,
                "result": "mock_result"
            })))
        }

        async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn Error + Send + Sync>> {
            Ok(vec![])
        }

        async fn get_tool_metadata(
            &self,
            _tool_name: &str,
        ) -> Result<ToolMetadata, Box<dyn Error + Send + Sync>> {
            Ok(ToolMetadata::new("mock_tool", "A mock tool"))
        }

        fn protocol_name(&self) -> &str {
            "mock"
        }
    }

    // ── Private-field tests (must stay inline) ────────────────────────────

    #[tokio::test]
    async fn test_get_tool_protocol() {
        let protocol = Arc::new(MockProtocol);
        let mut registry = ToolRegistry::empty();

        registry
            .add_protocol("local", protocol.clone())
            .await
            .unwrap();

        let tool = Tool::new("calculator", "Performs calculations", protocol.clone());
        registry.add_tool(tool);
        registry
            .tool_to_protocol
            .insert("calculator".to_string(), "local".to_string());

        assert_eq!(registry.get_tool_protocol("calculator"), Some("local"));
        assert_eq!(registry.get_tool_protocol("nonexistent"), None);
    }

    #[tokio::test]
    async fn test_remove_protocol_removes_tools() {
        let protocol = Arc::new(MockProtocol);
        let mut registry = ToolRegistry::empty();

        registry
            .add_protocol("protocol1", protocol.clone())
            .await
            .unwrap();

        let tool1 = Tool::new("tool1", "First tool", protocol.clone());
        registry.add_tool(tool1);
        registry
            .tool_to_protocol
            .insert("tool1".to_string(), "protocol1".to_string());

        let tool2 = Tool::new("tool2", "Second tool", protocol.clone());
        registry.add_tool(tool2);
        registry
            .tool_to_protocol
            .insert("tool2".to_string(), "protocol1".to_string());

        assert_eq!(registry.list_tools().len(), 2);

        registry.remove_protocol("protocol1");

        assert_eq!(registry.list_tools().len(), 0);
        assert_eq!(registry.get_tool_protocol("tool1"), None);
        assert_eq!(registry.get_tool_protocol("tool2"), None);
    }
}