mcp_host/registry/

tools.rs

//! Tool registry for MCP servers
//!
//! Provides registration and execution of MCP tools with circuit breaker protection

use std::sync::{Arc, RwLock};

use async_trait::async_trait;
use breaker_machines::{CircuitBreaker, Config as BreakerConfig};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use tokio::sync::mpsc;

use crate::content::types::{Content, ImageContent, TextContent};
use crate::server::multiplexer::ClientRequester;
use crate::server::session::Session;
use crate::server::visibility::{ExecutionContext, VisibilityContext};
use crate::transport::traits::JsonRpcNotification;

/// Tool execution errors
#[derive(Debug, Error)]
pub enum ToolError {
    /// Tool not found
    #[error("Tool not found: {0}")]
    NotFound(String),

    /// Invalid arguments
    #[error("Invalid arguments: {0}")]
    InvalidArguments(String),

    /// Execution error
    #[error("Execution error: {0}")]
    Execution(String),

    /// Internal error
    #[error("Internal error: {0}")]
    Internal(String),

    /// Circuit breaker is open
    #[error("Circuit breaker open for tool '{tool}': {message}")]
    CircuitOpen { tool: String, message: String },
}

/// Tool metadata for listing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInfo {
    /// Tool name
    pub name: String,

    /// Tool description
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Input schema (JSON Schema)
    #[serde(rename = "inputSchema")]
    pub input_schema: Value,

    /// Execution metadata (task support, etc.)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub execution: Option<crate::protocol::types::ToolExecution>,
}

/// Tool trait for implementing MCP tools
#[async_trait]
pub trait Tool: Send + Sync {
    /// Get tool name
    fn name(&self) -> &str;

    /// Get tool description
    fn description(&self) -> Option<&str> {
        None
    }

    /// Get input schema (JSON Schema)
    fn input_schema(&self) -> Value;

    /// Get execution metadata (task support, etc.)
    ///
    /// Override this to declare task support for this tool.
    /// Returns None by default (forbidden task execution).
    fn execution(&self) -> Option<crate::protocol::types::ToolExecution> {
        None
    }

    /// Check if this tool should be visible in the given context
    ///
    /// Override this to implement contextual visibility. The default implementation
    /// always returns true (always visible).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// fn is_visible(&self, ctx: &VisibilityContext) -> bool {
    ///     // Only show git commit tool if repo has uncommitted changes
    ///     ctx.environment
    ///         .map(|e| e.has_git_repo() && !e.git_is_clean())
    ///         .unwrap_or(false)
    /// }
    /// ```
    fn is_visible(&self, _ctx: &VisibilityContext) -> bool {
        true
    }

    /// Execute the tool with execution context
    ///
    /// The execution context provides access to:
    /// - `ctx.params`: Input parameters from the tool call
    /// - `ctx.session`: Current session (roles, state, client info)
    /// - `ctx.environment`: Optional environment state
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// async fn execute(&self, ctx: ExecutionContext<'_>) -> Result<Vec<Box<dyn Content>>, ToolError> {
    ///     let name = ctx.params.get("name").and_then(|v| v.as_str());
    ///
    ///     // Check session roles
    ///     if ctx.is_admin() {
    ///         // Admin-only behavior
    ///     }
    ///
    ///     Ok(vec![])
    /// }
    /// ```
    async fn execute(&self, ctx: ExecutionContext<'_>) -> Result<Vec<Box<dyn Content>>, ToolError>;
}

/// Helper methods for tool implementations
///
/// Provides ergonomic shortcuts for creating common content types.
/// Auto-implemented for all Tool trait objects.
pub trait ToolHelpers {
    /// Create text content
    fn text(&self, content: &str) -> Box<dyn Content> {
        Box::new(TextContent::new(content))
    }

    /// Create image content
    fn image(&self, data: &str, mime_type: &str) -> Box<dyn Content> {
        Box::new(ImageContent::new(data, mime_type))
    }
}

// Auto-implement ToolHelpers for all Tool trait objects
impl<T: Tool + ?Sized> ToolHelpers for T {}

/// Circuit breaker configuration for tools
#[derive(Debug, Clone)]
pub struct ToolBreakerConfig {
    /// Number of failures before opening circuit
    pub failure_threshold: usize,
    /// Time window in seconds for counting failures
    pub failure_window_secs: f64,
    /// Timeout before transitioning to half-open
    pub half_open_timeout_secs: f64,
    /// Successes needed to close circuit
    pub success_threshold: usize,
}

impl Default for ToolBreakerConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 5,
            failure_window_secs: 60.0,
            half_open_timeout_secs: 30.0,
            success_threshold: 2,
        }
    }
}

impl From<ToolBreakerConfig> for BreakerConfig {
    fn from(cfg: ToolBreakerConfig) -> Self {
        BreakerConfig {
            failure_threshold: Some(cfg.failure_threshold),
            failure_rate_threshold: None,
            minimum_calls: 1,
            failure_window_secs: cfg.failure_window_secs,
            half_open_timeout_secs: cfg.half_open_timeout_secs,
            success_threshold: cfg.success_threshold,
            jitter_factor: 0.1,
        }
    }
}

/// Tool registry for managing available tools with circuit breaker protection
#[derive(Clone)]
pub struct ToolRegistry {
    tools: Arc<DashMap<String, Arc<dyn Tool>>>,
    breakers: Arc<DashMap<String, Arc<RwLock<CircuitBreaker>>>>,
    breaker_config: Arc<RwLock<ToolBreakerConfig>>,
    notification_tx: Option<mpsc::UnboundedSender<JsonRpcNotification>>,
}

impl ToolRegistry {
    /// Create new tool registry
    pub fn new() -> Self {
        Self {
            tools: Arc::new(DashMap::new()),
            breakers: Arc::new(DashMap::new()),
            breaker_config: Arc::new(RwLock::new(ToolBreakerConfig::default())),
            notification_tx: None,
        }
    }

    /// Create tool registry with notification channel
    pub fn with_notifications(notification_tx: mpsc::UnboundedSender<JsonRpcNotification>) -> Self {
        Self {
            tools: Arc::new(DashMap::new()),
            breakers: Arc::new(DashMap::new()),
            breaker_config: Arc::new(RwLock::new(ToolBreakerConfig::default())),
            notification_tx: Some(notification_tx),
        }
    }

    /// Set notification channel
    pub fn set_notification_tx(&mut self, tx: mpsc::UnboundedSender<JsonRpcNotification>) {
        self.notification_tx = Some(tx);
    }

    /// Configure circuit breakers for all tools
    pub fn set_breaker_config(&self, config: ToolBreakerConfig) {
        if let Ok(mut cfg) = self.breaker_config.write() {
            *cfg = config;
        }
    }

    /// Register a tool with circuit breaker protection
    pub fn register<T: Tool + 'static>(&self, tool: T) {
        let name = tool.name().to_string();

        // Create circuit breaker for this tool
        let breaker_config = self
            .breaker_config
            .read()
            .map(|c| c.clone())
            .unwrap_or_default();

        let breaker = CircuitBreaker::builder(&name)
            .failure_threshold(breaker_config.failure_threshold)
            .failure_window_secs(breaker_config.failure_window_secs)
            .half_open_timeout_secs(breaker_config.half_open_timeout_secs)
            .success_threshold(breaker_config.success_threshold)
            .build();

        self.breakers
            .insert(name.clone(), Arc::new(RwLock::new(breaker)));
        self.tools.insert(name, Arc::new(tool));
    }

    /// Register a boxed tool with circuit breaker
    pub fn register_boxed(&self, tool: Arc<dyn Tool>) {
        let name = tool.name().to_string();

        let breaker_config = self
            .breaker_config
            .read()
            .map(|c| c.clone())
            .unwrap_or_default();

        let breaker = CircuitBreaker::builder(&name)
            .failure_threshold(breaker_config.failure_threshold)
            .failure_window_secs(breaker_config.failure_window_secs)
            .half_open_timeout_secs(breaker_config.half_open_timeout_secs)
            .success_threshold(breaker_config.success_threshold)
            .build();

        self.breakers
            .insert(name.clone(), Arc::new(RwLock::new(breaker)));
        self.tools.insert(name, tool);
    }

    /// Get a tool by name
    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tools.get(name).map(|t| Arc::clone(&t))
    }

    /// Check if a tool's circuit breaker is open
    pub fn is_circuit_open(&self, name: &str) -> bool {
        self.breakers
            .get(name)
            .and_then(|b| b.read().ok().map(|breaker| breaker.is_open()))
            .unwrap_or(false)
    }

    /// List all registered tools (excluding those with open circuits if filtered)
    pub fn list(&self) -> Vec<ToolInfo> {
        self.tools
            .iter()
            .map(|entry| {
                let tool = entry.value();
                ToolInfo {
                    name: tool.name().to_string(),
                    description: tool.description().map(|s| s.to_string()),
                    input_schema: tool.input_schema(),
                    execution: tool.execution(),
                }
            })
            .collect()
    }

    /// List only available tools (circuit closed or half-open)
    pub fn list_available(&self) -> Vec<ToolInfo> {
        self.tools
            .iter()
            .filter(|entry| !self.is_circuit_open(entry.key()))
            .map(|entry| {
                let tool = entry.value();
                ToolInfo {
                    name: tool.name().to_string(),
                    description: tool.description().map(|s| s.to_string()),
                    input_schema: tool.input_schema(),
                    execution: tool.execution(),
                }
            })
            .collect()
    }

    /// Send notification about circuit state change
    fn send_notification(&self, method: &str, params: Option<Value>) {
        if let Some(tx) = &self.notification_tx {
            let notification = JsonRpcNotification::new(method, params);
            let _ = tx.send(notification);
        }
    }

    /// Send tools list changed notification
    fn notify_tools_changed(&self) {
        self.send_notification("notifications/tools/list_changed", None);
    }

    /// Send logging message notification (visible to LLM)
    fn notify_message(&self, level: &str, logger: &str, message: &str) {
        self.send_notification(
            "notifications/message",
            Some(serde_json::json!({
                "level": level,
                "logger": logger,
                "data": message
            })),
        );
    }

    /// Call a tool by name with circuit breaker protection
    ///
    /// # Arguments
    ///
    /// * `name` - Tool name to call
    /// * `params` - Tool parameters
    /// * `session` - Current session
    /// * `client_requester` - Optional client requester for server→client requests
    pub async fn call(
        &self,
        name: &str,
        params: Value,
        session: &Session,
        logger: &crate::logging::McpLogger,
        client_requester: Option<ClientRequester>,
    ) -> Result<Vec<Box<dyn Content>>, ToolError> {
        let tool = self
            .get(name)
            .ok_or_else(|| ToolError::NotFound(name.to_string()))?;

        let breaker = self.breakers.get(name).ok_or_else(|| {
            ToolError::Internal(format!("No circuit breaker for tool '{}'", name))
        })?;

        // Check circuit state before execution
        let was_open = {
            let breaker_guard = breaker
                .read()
                .map_err(|e| ToolError::Internal(format!("Breaker lock error: {}", e)))?;
            breaker_guard.is_open()
        };

        if was_open {
            return Err(ToolError::CircuitOpen {
                tool: name.to_string(),
                message: "Too many recent failures. Service temporarily unavailable.".to_string(),
            });
        }

        // Create execution context with optional client requester
        let ctx = match client_requester {
            Some(cr) => ExecutionContext::new(params, session, logger).with_client_requester(cr),
            None => ExecutionContext::new(params, session, logger),
        };

        // Execute the tool with timing
        let start = std::time::Instant::now();
        let result = tool.execute(ctx).await;
        let duration_secs = start.elapsed().as_secs_f64();

        // Update circuit breaker based on result
        let breaker_guard = breaker
            .write()
            .map_err(|e| ToolError::Internal(format!("Breaker lock error: {}", e)))?;

        let was_closed_before = !breaker_guard.is_open();

        match &result {
            Ok(_) => {
                breaker_guard.record_success(duration_secs);
                // Check if we recovered from open state
                if was_open && !breaker_guard.is_open() {
                    self.notify_tools_changed();
                    self.notify_message(
                        "info",
                        "breaker-machines",
                        &format!("Tool '{}' recovered and available", name),
                    );
                }
            }
            Err(_) => {
                breaker_guard.record_failure(duration_secs);
                // Check if we just opened the circuit
                if was_closed_before && breaker_guard.is_open() {
                    self.notify_tools_changed();
                    self.notify_message(
                        "warning",
                        "breaker-machines",
                        &format!(
                            "Tool '{}' disabled: circuit breaker open after failures",
                            name
                        ),
                    );
                }
            }
        }

        result
    }

    // ==================== Session-Aware Methods ====================

    /// List tools visible to a specific session
    ///
    /// Resolution order:
    /// 1. Session overrides (replace global tools with session-specific implementations)
    /// 2. Session extras (additional tools added to session)
    /// 3. Global tools (filtered by hidden list and visibility predicate)
    pub fn list_for_session(
        &self,
        session: &Session,
        ctx: &VisibilityContext<'_>,
    ) -> Vec<ToolInfo> {
        let mut tools = std::collections::HashMap::new();

        // 1. Add global tools (filtered by hidden and visibility)
        for entry in self.tools.iter() {
            let name = entry.key().clone();
            if !session.is_tool_hidden(&name) && !self.is_circuit_open(&name) {
                let tool = entry.value();
                if tool.is_visible(ctx) {
                    tools.insert(
                        name,
                        ToolInfo {
                            name: tool.name().to_string(),
                            description: tool.description().map(|s| s.to_string()),
                            input_schema: tool.input_schema(),
                            execution: tool.execution(),
                        },
                    );
                }
            }
        }

        // 2. Add session extras (can add new tools or override global)
        for entry in session.tool_extras().iter() {
            let name = entry.key().clone();
            let tool = entry.value();
            if tool.is_visible(ctx) {
                tools.insert(
                    name,
                    ToolInfo {
                        name: tool.name().to_string(),
                        description: tool.description().map(|s| s.to_string()),
                        input_schema: tool.input_schema(),
                        execution: tool.execution(),
                    },
                );
            }
        }

        // 3. Apply session overrides (replace implementations)
        for entry in session.tool_overrides().iter() {
            let name = entry.key().clone();
            let tool = entry.value();
            if tool.is_visible(ctx) {
                tools.insert(
                    name,
                    ToolInfo {
                        name: tool.name().to_string(),
                        description: tool.description().map(|s| s.to_string()),
                        input_schema: tool.input_schema(),
                        execution: tool.execution(),
                    },
                );
            }
        }

        tools.into_values().collect()
    }

    /// Call a tool with session context
    ///
    /// Resolution order:
    /// 1. Resolve alias to actual name
    /// 2. Check session overrides (polymorphic tools)
    /// 3. Check session extras (injected tools)
    /// 4. Check session hidden (filtered out)
    /// 5. Check visibility predicate (contextual)
    /// 6. Fall back to global registry with circuit breaker
    pub async fn call_for_session(
        &self,
        name: &str,
        params: Value,
        session: &Session,
        logger: &crate::logging::McpLogger,
        visibility_ctx: &VisibilityContext<'_>,
        client_requester: Option<ClientRequester>,
    ) -> Result<Vec<Box<dyn Content>>, ToolError> {
        // 1. Resolve alias
        let resolved_name = session.resolve_tool_alias(name);
        let resolved = resolved_name.as_ref();

        // Create execution context (reuse environment from visibility context)
        // Clone client_requester so we can use it both for exec_ctx and for nested call()
        let exec_ctx = match (visibility_ctx.environment, client_requester.as_ref()) {
            (Some(env), Some(cr)) => {
                ExecutionContext::with_environment(params.clone(), session, logger, env)
                    .with_client_requester(cr.clone())
            }
            (Some(env), None) => {
                ExecutionContext::with_environment(params.clone(), session, logger, env)
            }
            (None, Some(cr)) => ExecutionContext::new(params.clone(), session, logger)
                .with_client_requester(cr.clone()),
            (None, None) => ExecutionContext::new(params.clone(), session, logger),
        };

        // 2. Check session override first
        if let Some(tool) = session.get_tool_override(resolved) {
            if !tool.is_visible(visibility_ctx) {
                return Err(ToolError::NotFound(name.to_string()));
            }
            return tool.execute(exec_ctx).await;
        }

        // 3. Check session extras
        if let Some(tool) = session.get_tool_extra(resolved) {
            if !tool.is_visible(visibility_ctx) {
                return Err(ToolError::NotFound(name.to_string()));
            }
            return tool.execute(exec_ctx).await;
        }

        // 4. Check if hidden in session
        if session.is_tool_hidden(resolved) {
            return Err(ToolError::NotFound(name.to_string()));
        }

        // 5. Check global registry with visibility check
        let tool = self
            .get(resolved)
            .ok_or_else(|| ToolError::NotFound(name.to_string()))?;

        if !tool.is_visible(visibility_ctx) {
            return Err(ToolError::NotFound(name.to_string()));
        }

        // 6. Execute with circuit breaker (use resolved name for breaker tracking)
        self.call(resolved, params, session, logger, client_requester)
            .await
    }

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

    /// Check if registry is empty
    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::content::types::TextContent;

    // Example tool for testing
    struct EchoTool;

    #[async_trait]
    impl Tool for EchoTool {
        fn name(&self) -> &str {
            "echo"
        }

        fn description(&self) -> Option<&str> {
            Some("Echoes back the input message")
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({
                "type": "object",
                "properties": {
                    "message": {
                        "type": "string",
                        "description": "Message to echo"
                    }
                },
                "required": ["message"]
            })
        }

        async fn execute(
            &self,
            ctx: ExecutionContext<'_>,
        ) -> Result<Vec<Box<dyn Content>>, ToolError> {
            let message = ctx
                .params
                .get("message")
                .and_then(|v| v.as_str())
                .ok_or_else(|| {
                    ToolError::InvalidArguments("Missing 'message' field".to_string())
                })?;

            let content = TextContent::new(format!("Echo: {}", message));
            Ok(vec![Box::new(content)])
        }
    }

    #[test]
    fn test_registry_creation() {
        let registry = ToolRegistry::new();
        assert!(registry.is_empty());
    }

    #[test]
    fn test_tool_registration() {
        let registry = ToolRegistry::new();
        registry.register(EchoTool);

        assert_eq!(registry.len(), 1);
        assert!(!registry.is_empty());
    }

    #[test]
    fn test_get_tool() {
        let registry = ToolRegistry::new();
        registry.register(EchoTool);

        let tool = registry.get("echo");
        assert!(tool.is_some());
        assert_eq!(tool.unwrap().name(), "echo");

        let missing = registry.get("nonexistent");
        assert!(missing.is_none());
    }

    #[test]
    fn test_list_tools() {
        let registry = ToolRegistry::new();
        registry.register(EchoTool);

        let tools = registry.list();
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0].name, "echo");
        assert_eq!(
            tools[0].description,
            Some("Echoes back the input message".to_string())
        );
    }

    #[tokio::test]
    async fn test_call_tool() {
        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let logger = crate::logging::McpLogger::new(_tx, "test");
        let registry = ToolRegistry::new();
        registry.register(EchoTool);
        let session = Session::new();

        let params = serde_json::json!({
            "message": "Hello, world!"
        });

        let result = registry
            .call("echo", params, &session, &logger, None)
            .await
            .unwrap();
        assert_eq!(result.len(), 1);
    }

    #[tokio::test]
    async fn test_call_missing_tool() {
        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let logger = crate::logging::McpLogger::new(_tx, "test");
        let registry = ToolRegistry::new();
        let session = Session::new();

        let params = serde_json::json!({});
        let result = registry
            .call("nonexistent", params, &session, &logger, None)
            .await;

        assert!(matches!(result, Err(ToolError::NotFound(_))));
    }

    #[tokio::test]
    async fn test_tool_invalid_arguments() {
        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let logger = crate::logging::McpLogger::new(_tx, "test");
        let registry = ToolRegistry::new();
        registry.register(EchoTool);
        let session = Session::new();

        let params = serde_json::json!({}); // Missing required 'message' field

        let result = registry.call("echo", params, &session, &logger, None).await;
        assert!(matches!(result, Err(ToolError::InvalidArguments(_))));
    }
}