mcp_host/registry/

tools.rs

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

use std::collections::HashSet;
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 state_machines::state_machine;
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 LIFECYCLE - Enabled/Disabled/Defective per tool
// =============================================================================

/// Runtime lifecycle state for a tool.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolLifecycleState {
    /// Tool is enabled and can be listed/called.
    Enabled,
    /// Tool is disabled (hidden and not callable).
    Disabled,
    /// Tool is defective (disabled by circuit breaker).
    Defective,
}

impl ToolLifecycleState {
    fn from_state_name(name: &str) -> Option<Self> {
        match name {
            "Enabled" => Some(Self::Enabled),
            "Disabled" => Some(Self::Disabled),
            "Defective" => Some(Self::Defective),
            _ => None,
        }
    }
}

/// Policy configuration for a tool's lifecycle.
#[derive(Debug, Clone)]
pub struct ToolPolicy {
    /// Desired initial state when the tool is registered.
    pub initial_state: ToolLifecycleState,
    /// Exclusive group name (only one tool in the group can be enabled).
    pub exclusive_group: Option<String>,
    /// Tools to activate immediately when this tool becomes enabled.
    pub activates: Vec<String>,
}

impl Default for ToolPolicy {
    fn default() -> Self {
        Self {
            initial_state: ToolLifecycleState::Enabled,
            exclusive_group: None,
            activates: Vec::new(),
        }
    }
}

impl ToolPolicy {
    /// Create a new policy with defaults (enabled, no group, no activations).
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the initial lifecycle state.
    pub fn initial_state(mut self, state: ToolLifecycleState) -> Self {
        self.initial_state = state;
        self
    }

    /// Set an exclusive group for this tool.
    pub fn exclusive_group(mut self, group: impl Into<String>) -> Self {
        self.exclusive_group = Some(group.into());
        self
    }

    /// Set tools to activate when this tool becomes enabled.
    pub fn activates<I, S>(mut self, tools: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.activates = tools.into_iter().map(Into::into).collect();
        self
    }
}

#[derive(Debug, Default)]
struct ToolLifecycleContext;

impl ToolLifecycleContext {
    fn new() -> Self {
        Self
    }
}

state_machine! {
    name: ToolLifecycle,
    context: ToolLifecycleContext,
    dynamic: true,
    initial: Enabled,
    states: [Enabled, Disabled, Defective],
    events {
        enable {
            transition: { from: [Disabled, Defective], to: Enabled }
        }
        disable {
            transition: { from: [Enabled, Defective], to: Disabled }
        }
        mark_defective {
            transition: { from: [Enabled, Disabled], to: Defective }
        }
        recover {
            transition: { from: Defective, to: Enabled }
        }
    }
}

// =============================================================================
// TOOL OUTPUT - Exclusive content OR structured output
// =============================================================================

/// Tool execution output - either content OR structured, never both.
///
/// MCP spec allows tools to return either:
/// - `content`: Array of content items (text, images, etc.) for display
/// - `structuredContent`: JSON object conforming to the tool's `outputSchema`
///
/// This enum enforces exclusivity at the type level.
pub enum ToolOutput {
    /// Traditional content array (text, images, etc.)
    Content(Vec<Box<dyn Content>>),
    /// Structured JSON output (must conform to outputSchema if defined)
    Structured(Value),
}

impl std::fmt::Debug for ToolOutput {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Content(items) => f
                .debug_struct("Content")
                .field("count", &items.len())
                .finish(),
            Self::Structured(value) => f.debug_tuple("Structured").field(value).finish(),
        }
    }
}

impl ToolOutput {
    /// Create content output with a single text item
    pub fn text(s: impl Into<String>) -> Self {
        Self::Content(vec![Box::new(TextContent::new(s))])
    }

    /// Create content output with multiple text items
    pub fn texts(items: &[&str]) -> Self {
        Self::Content(
            items
                .iter()
                .map(|s| Box::new(TextContent::new(*s)) as Box<dyn Content>)
                .collect(),
        )
    }

    /// Create content output from a vec of content items
    pub fn content(items: Vec<Box<dyn Content>>) -> Self {
        Self::Content(items)
    }

    /// Create structured output from a serializable value
    pub fn structured<T: Serialize>(value: T) -> Result<Self, serde_json::Error> {
        Ok(Self::Structured(serde_json::to_value(value)?))
    }

    /// Create structured output from a raw JSON value
    pub fn json(value: Value) -> Self {
        Self::Structured(value)
    }

    /// Check if this is content output
    pub fn is_content(&self) -> bool {
        matches!(self, Self::Content(_))
    }

    /// Check if this is structured output
    pub fn is_structured(&self) -> bool {
        matches!(self, Self::Structured(_))
    }
}

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

    /// Human-readable title for UI display
    #[serde(skip_serializing_if = "Option::is_none")]
    pub title: Option<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,

    /// Output schema (JSON Schema) for structured output
    #[serde(rename = "outputSchema", skip_serializing_if = "Option::is_none")]
    pub output_schema: Option<Value>,

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

    /// Tool annotations (behavioral hints for LLMs)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub annotations: Option<crate::protocol::types::ToolAnnotations>,
}

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

    /// Get human-readable title for UI display
    ///
    /// If not provided, clients should use `name` as fallback.
    fn title(&self) -> Option<&str> {
        None
    }

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

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

    /// Get output schema (JSON Schema) for structured output
    ///
    /// When this returns Some, the tool MUST return `ToolOutput::Structured`
    /// that conforms to this schema. When None, the tool returns `ToolOutput::Content`.
    fn output_schema(&self) -> Option<Value> {
        None
    }

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

    /// Get tool annotations (behavioral hints for LLMs)
    ///
    /// Override this to provide hints about tool behavior (read-only, destructive, etc.)
    fn annotations(&self) -> Option<crate::protocol::types::ToolAnnotations> {
        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
    ///
    /// # Returns
    ///
    /// - `ToolOutput::Content` for traditional content (text, images)
    /// - `ToolOutput::Structured` for JSON output (requires `output_schema`)
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// // Content output
    /// async fn execute(&self, ctx: ExecutionContext<'_>) -> Result<ToolOutput, ToolError> {
    ///     Ok(ToolOutput::text("Hello, world!"))
    /// }
    ///
    /// // Structured output
    /// async fn execute(&self, ctx: ExecutionContext<'_>) -> Result<ToolOutput, ToolError> {
    ///     ToolOutput::structured(MyOutput { value: 42 })
    ///         .map_err(|e| ToolError::Internal(e.to_string()))
    /// }
    /// ```
    async fn execute(&self, ctx: ExecutionContext<'_>) -> Result<ToolOutput, 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>>,
    tool_states: Arc<DashMap<String, Arc<RwLock<DynamicToolLifecycle>>>>,
    tool_policies: Arc<DashMap<String, ToolPolicy>>,
}

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,
            tool_states: Arc::new(DashMap::new()),
            tool_policies: Arc::new(DashMap::new()),
        }
    }

    /// 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),
            tool_states: Arc::new(DashMap::new()),
            tool_policies: Arc::new(DashMap::new()),
        }
    }

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

    /// Configure lifecycle policy for a tool.
    ///
    /// This applies the policy immediately, including the desired state.
    pub fn set_tool_policy(&self, name: impl Into<String>, policy: ToolPolicy) {
        let name = name.into();
        self.tool_policies.insert(name.clone(), policy.clone());
        self.ensure_tool_state(&name);

        let mut visited = HashSet::new();
        let changed = self.set_tool_state_internal(&name, policy.initial_state, &mut visited);
        if !changed && self.tool_state(&name) == Some(ToolLifecycleState::Enabled) {
            visited.insert(name.clone());
            self.apply_enable_policy(&name, &policy, &mut visited);
        }
    }

    /// Get lifecycle state for a tool if tracked.
    pub fn tool_state(&self, name: &str) -> Option<ToolLifecycleState> {
        self.tool_states.get(name).and_then(|entry| {
            entry
                .read()
                .ok()
                .and_then(|machine| ToolLifecycleState::from_state_name(machine.current_state()))
        })
    }

    /// Check if a tool is enabled.
    ///
    /// If the tool has no lifecycle tracking yet, it is treated as enabled.
    pub fn is_tool_enabled(&self, name: &str) -> bool {
        match self.tool_state(name) {
            Some(ToolLifecycleState::Enabled) => true,
            Some(_) => false,
            None => true,
        }
    }

    /// Enable a tool (and apply activation/exclusive policies).
    pub fn enable_tool(&self, name: &str) -> bool {
        let mut visited = HashSet::new();
        self.set_tool_state_internal(name, ToolLifecycleState::Enabled, &mut visited)
    }

    /// Disable a tool.
    pub fn disable_tool(&self, name: &str) -> bool {
        let mut visited = HashSet::new();
        self.set_tool_state_internal(name, ToolLifecycleState::Disabled, &mut visited)
    }

    /// Mark a tool as defective (breaker-open state).
    pub fn mark_tool_defective(&self, name: &str) -> bool {
        let mut visited = HashSet::new();
        self.set_tool_state_internal(name, ToolLifecycleState::Defective, &mut visited)
    }

    /// Recover a tool from defective state (no-op if not defective).
    pub fn recover_tool(&self, name: &str) -> bool {
        if self.tool_state(name) != Some(ToolLifecycleState::Defective) {
            return false;
        }
        let mut visited = HashSet::new();
        self.set_tool_state_internal(name, ToolLifecycleState::Enabled, &mut visited)
    }

    /// Register a tool with circuit breaker protection
    pub fn register<T: Tool + 'static>(&self, tool: T) {
        self.register_boxed(Arc::new(tool));
    }

    /// Register a tool with lifecycle policy.
    pub fn register_with_policy<T: Tool + 'static>(&self, tool: T, policy: ToolPolicy) {
        self.register_boxed_with_policy(Arc::new(tool), policy);
    }

    /// Register a boxed tool with circuit breaker
    pub fn register_boxed(&self, tool: Arc<dyn Tool>) {
        let name = tool.name().to_string();
        let policy = self
            .tool_policies
            .get(&name)
            .map(|p| p.clone())
            .unwrap_or_default();
        self.register_boxed_with_policy(tool, policy);
    }

    /// Register a boxed tool with lifecycle policy.
    pub fn register_boxed_with_policy(&self, tool: Arc<dyn Tool>, policy: ToolPolicy) {
        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.clone(), tool);

        // Store policy (overrides any existing policy)
        self.tool_policies.insert(name.clone(), policy.clone());

        // Initialize lifecycle state only if we just created it.
        let (_, created) = self.ensure_tool_state(&name);
        if created && policy.initial_state != ToolLifecycleState::Enabled {
            self.set_tool_state_initial(&name, policy.initial_state);
        }
        if self.tool_state(&name) == Some(ToolLifecycleState::Enabled) {
            let mut visited = HashSet::new();
            visited.insert(name.clone());
            self.apply_enable_policy(&name, &policy, &mut visited);
        }
    }

    /// 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(),
                    title: tool.title().map(|s| s.to_string()),
                    description: tool.description().map(|s| s.to_string()),
                    input_schema: tool.input_schema(),
                    output_schema: tool.output_schema(),
                    execution: tool.execution(),
                    annotations: tool.annotations(),
                }
            })
            .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()) && self.is_tool_enabled(entry.key())
            })
            .map(|entry| {
                let tool = entry.value();
                ToolInfo {
                    name: tool.name().to_string(),
                    title: tool.title().map(|s| s.to_string()),
                    description: tool.description().map(|s| s.to_string()),
                    input_schema: tool.input_schema(),
                    output_schema: tool.output_schema(),
                    execution: tool.execution(),
                    annotations: tool.annotations(),
                }
            })
            .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
            })),
        );
    }

    fn ensure_tool_state(&self, name: &str) -> (Arc<RwLock<DynamicToolLifecycle>>, bool) {
        if let Some(entry) = self.tool_states.get(name) {
            return (Arc::clone(entry.value()), false);
        }

        if !self.tool_policies.contains_key(name) {
            self.tool_policies
                .insert(name.to_string(), ToolPolicy::default());
        }

        let machine = DynamicToolLifecycle::new(ToolLifecycleContext::new());
        let entry = Arc::new(RwLock::new(machine));
        self.tool_states.insert(name.to_string(), Arc::clone(&entry));
        (entry, true)
    }

    fn set_tool_state_initial(&self, name: &str, target: ToolLifecycleState) {
        let event = match target {
            ToolLifecycleState::Enabled => return,
            ToolLifecycleState::Disabled => ToolLifecycleEvent::Disable,
            ToolLifecycleState::Defective => ToolLifecycleEvent::MarkDefective,
        };
        let _ = self.transition_tool_state(name, event, false);
    }

    fn transition_tool_state(
        &self,
        name: &str,
        event: ToolLifecycleEvent,
        notify: bool,
    ) -> Option<(ToolLifecycleState, ToolLifecycleState)> {
        let (entry, _) = self.ensure_tool_state(name);
        let mut guard = entry.write().ok()?;
        let before = ToolLifecycleState::from_state_name(guard.current_state())?;
        let _ = guard.handle(event);
        let after = ToolLifecycleState::from_state_name(guard.current_state())?;
        drop(guard);

        if notify
            && before != after
            && (before == ToolLifecycleState::Enabled
                || after == ToolLifecycleState::Enabled)
        {
            self.notify_tools_changed();
        }

        Some((before, after))
    }

    fn set_tool_state_internal(
        &self,
        name: &str,
        target: ToolLifecycleState,
        visited: &mut HashSet<String>,
    ) -> bool {
        let _ = self.ensure_tool_state(name);
        let current = self
            .tool_state(name)
            .unwrap_or(ToolLifecycleState::Enabled);

        if current == target {
            return false;
        }

        let event = match (current, target) {
            (ToolLifecycleState::Enabled, ToolLifecycleState::Disabled) => {
                ToolLifecycleEvent::Disable
            }
            (ToolLifecycleState::Enabled, ToolLifecycleState::Defective) => {
                ToolLifecycleEvent::MarkDefective
            }
            (ToolLifecycleState::Disabled, ToolLifecycleState::Enabled) => {
                ToolLifecycleEvent::Enable
            }
            (ToolLifecycleState::Disabled, ToolLifecycleState::Defective) => {
                ToolLifecycleEvent::MarkDefective
            }
            (ToolLifecycleState::Defective, ToolLifecycleState::Enabled) => {
                ToolLifecycleEvent::Recover
            }
            (ToolLifecycleState::Defective, ToolLifecycleState::Disabled) => {
                ToolLifecycleEvent::Disable
            }
            _ => return false,
        };

        let change = self.transition_tool_state(name, event, true);
        let changed = matches!(change, Some((prev, next)) if prev != next);
        let became_enabled = matches!(
            change,
            Some((prev, ToolLifecycleState::Enabled))
                if prev != ToolLifecycleState::Enabled
        );

        if became_enabled {
            if !visited.insert(name.to_string()) {
                return true;
            }
            if let Some(policy) = self.tool_policies.get(name).map(|p| p.clone()) {
                self.apply_enable_policy(name, &policy, visited);
            }
        }

        changed
    }

    fn apply_enable_policy(
        &self,
        name: &str,
        policy: &ToolPolicy,
        visited: &mut HashSet<String>,
    ) {
        if let Some(group) = policy.exclusive_group.as_deref() {
            let others: Vec<String> = self
                .tool_policies
                .iter()
                .filter(|entry| {
                    entry.key() != name
                        && entry.value().exclusive_group.as_deref() == Some(group)
                })
                .map(|entry| entry.key().clone())
                .collect();

            for other in others {
                if self.tool_state(&other) == Some(ToolLifecycleState::Enabled) {
                    self.set_tool_state_internal(&other, ToolLifecycleState::Disabled, visited);
                }
            }
        }

        for target in &policy.activates {
            self.set_tool_state_internal(target, ToolLifecycleState::Enabled, visited);
        }
    }

    /// 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<ToolOutput, ToolError> {
        let tool = self
            .get(name)
            .ok_or_else(|| ToolError::NotFound(name.to_string()))?;

        if !self.is_tool_enabled(name) {
            return Err(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() {
                    let _ = self.recover_tool(name);
                    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() {
                    let _ = self.mark_tool_defective(name);
                    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)
                && self.is_tool_enabled(&name)
            {
                let tool = entry.value();
                if tool.is_visible(ctx) {
                    tools.insert(
                        name,
                        ToolInfo {
                            name: tool.name().to_string(),
                            title: tool.title().map(|s| s.to_string()),
                            description: tool.description().map(|s| s.to_string()),
                            input_schema: tool.input_schema(),
                            output_schema: tool.output_schema(),
                            execution: tool.execution(),
                            annotations: tool.annotations(),
                        },
                    );
                }
            }
        }

        // 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) && self.is_tool_enabled(&name) {
                tools.insert(
                    name,
                    ToolInfo {
                        name: tool.name().to_string(),
                        title: tool.title().map(|s| s.to_string()),
                        description: tool.description().map(|s| s.to_string()),
                        input_schema: tool.input_schema(),
                        output_schema: tool.output_schema(),
                        execution: tool.execution(),
                        annotations: tool.annotations(),
                    },
                );
            }
        }

        // 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) && self.is_tool_enabled(&name) {
                tools.insert(
                    name,
                    ToolInfo {
                        name: tool.name().to_string(),
                        title: tool.title().map(|s| s.to_string()),
                        description: tool.description().map(|s| s.to_string()),
                        input_schema: tool.input_schema(),
                        output_schema: tool.output_schema(),
                        execution: tool.execution(),
                        annotations: tool.annotations(),
                    },
                );
            }
        }

        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<ToolOutput, ToolError> {
        // 1. Resolve alias
        let resolved_name = session.resolve_tool_alias(name);
        let resolved = resolved_name.as_ref();

        if !self.is_tool_enabled(resolved) {
            return Err(ToolError::NotFound(name.to_string()));
        }

        // 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::*;

    // 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<ToolOutput, ToolError> {
            let message = ctx
                .params
                .get("message")
                .and_then(|v| v.as_str())
                .ok_or_else(|| {
                    ToolError::InvalidArguments("Missing 'message' field".to_string())
                })?;

            Ok(ToolOutput::text(format!("Echo: {}", message)))
        }
    }

    struct ToolA;
    struct ToolB;

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

        fn input_schema(&self) -> Value {
            serde_json::json!({ "type": "object" })
        }

        async fn execute(
            &self,
            _ctx: ExecutionContext<'_>,
        ) -> Result<ToolOutput, ToolError> {
            Ok(ToolOutput::text("A"))
        }
    }

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

        fn input_schema(&self) -> Value {
            serde_json::json!({ "type": "object" })
        }

        async fn execute(
            &self,
            _ctx: ExecutionContext<'_>,
        ) -> Result<ToolOutput, ToolError> {
            Ok(ToolOutput::text("B"))
        }
    }

    #[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!(result.is_content());
    }

    #[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(_))));
    }

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

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

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

    #[test]
    fn test_exclusive_group_disables_other() {
        let registry = ToolRegistry::new();
        registry.register_with_policy(
            ToolA,
            ToolPolicy::new().exclusive_group("exclusive"),
        );
        registry.register_with_policy(
            ToolB,
            ToolPolicy::new().exclusive_group("exclusive"),
        );

        registry.enable_tool("tool_a");
        registry.enable_tool("tool_b");

        assert_eq!(
            registry.tool_state("tool_a"),
            Some(ToolLifecycleState::Disabled)
        );
        assert_eq!(
            registry.tool_state("tool_b"),
            Some(ToolLifecycleState::Enabled)
        );
    }
}