vtcode_mcp/

client.rs

use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use chrono::Utc;
use parking_lot::RwLock;
use rmcp::model::{CallToolResult, ClientCapabilities, InitializeRequestParams, RootsCapabilities};
use rustc_hash::FxHashMap;
use serde_json::{Map, Value, json};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tracing::{info, warn};
use vtcode_commons::fs::{ensure_dir_exists, write_file_with_context};
use vtcode_config::mcp::{McpAllowListConfig, McpClientConfig, McpProviderConfig, McpTransportConfig};

use super::McpSandboxContext;
use super::{
    McpClientStatus, McpElicitationHandler, McpPromptDetail, McpPromptInfo, McpProvider, McpResourceData,
    McpResourceInfo, McpToolExecutor, McpToolInfo, format_tool_markdown, sanitize_filename,
};
use crate::connection_pool::McpConnectionPool;

struct McpClientState {
    providers: FxHashMap<String, Arc<McpProvider>>,
    allowlist: McpAllowListConfig,
    tool_provider_index: FxHashMap<String, String>,
    resource_provider_index: FxHashMap<String, String>,
    prompt_provider_index: FxHashMap<String, String>,
}

/// Stable provider order keeps the aggregated catalog (and LLM prompt-cache prefix) deterministic.
fn providers_sorted_by_name<V: Clone>(providers: &FxHashMap<String, V>) -> Vec<V> {
    let mut entries: Vec<(&String, &V)> = providers.iter().collect();
    entries.sort_unstable_by(|a, b| a.0.cmp(b.0));
    entries.into_iter().map(|(_, provider)| provider.clone()).collect()
}

pub struct McpClient {
    config: McpClientConfig,
    state: RwLock<McpClientState>,
    elicitation_handler: Option<Arc<dyn McpElicitationHandler>>,
    sandbox_context: Option<McpSandboxContext>,
}

impl McpClient {
    /// Create a new MCP client from the configuration.
    pub fn new(config: McpClientConfig) -> Self {
        Self::with_sandbox_context(config, None)
    }

    /// Create an MCP client with an optional inherited stdio sandbox context.
    pub fn with_sandbox_context(config: McpClientConfig, sandbox_context: Option<McpSandboxContext>) -> Self {
        let allowlist = config.allowlist.clone();

        Self {
            config,
            state: RwLock::new(McpClientState {
                providers: FxHashMap::default(),
                allowlist,
                tool_provider_index: FxHashMap::default(),
                resource_provider_index: FxHashMap::default(),
                prompt_provider_index: FxHashMap::default(),
            }),
            elicitation_handler: None,
            sandbox_context,
        }
    }

    /// Register a handler used to satisfy elicitation requests from providers.
    pub fn set_elicitation_handler(&mut self, handler: Arc<dyn McpElicitationHandler>) {
        self.elicitation_handler = Some(handler);
    }

    /// Establish connections to all configured providers and complete the
    /// MCP handshake.
    pub async fn initialize(&mut self) -> Result<()> {
        if !self.config.enabled {
            info!("MCP client is disabled in configuration");
            return Ok(());
        }

        info!("Initializing MCP client with {} configured providers", self.config.providers.len());

        let max_connections = self.config.max_concurrent_connections.max(1);
        let pool = McpConnectionPool::new(max_connections, self.config.request_timeout_seconds);
        let tool_timeout = Some(Duration::from_secs(self.config.request_timeout_seconds));
        let allowlist_snapshot = self.state.read().allowlist.clone();

        let provider_configs: Vec<McpProviderConfig> =
            self.config.providers.iter().filter(|c| c.enabled).cloned().collect();

        let results = pool
            .initialize_providers_parallel(
                provider_configs,
                self.elicitation_handler.clone(),
                self.sandbox_context.clone(),
                tool_timeout,
                &allowlist_snapshot,
            )
            .await
            .map_err(|e| anyhow::anyhow!("MCP connection pool initialization failed: {e}"))?;

        let mut initialized = FxHashMap::default();
        for (name, provider) in results {
            drop(initialized.insert(name, provider));
        }

        self.state.write().providers = initialized;
        info!("MCP client initialization complete. Active providers: {}", self.state.read().providers.len());

        Ok(())
    }

    /// Validate tool arguments based on security configuration
    fn validate_tool_arguments(&self, _tool_name: &str, args: &Value) -> Result<()> {
        // Check argument size
        if self.config.security.validation.max_argument_size > 0 {
            let args_size = serde_json::to_string(args).map_or(0, |s| u32::try_from(s.len()).unwrap_or(u32::MAX));

            if args_size > self.config.security.validation.max_argument_size {
                return Err(anyhow::anyhow!(
                    "Tool arguments exceed maximum size of {} bytes",
                    self.config.security.validation.max_argument_size
                ));
            }
        }

        // Check for path traversal in file-related arguments
        if self.config.security.validation.path_traversal_protection
            && let Some(path) = args.get("path").and_then(|v| v.as_str())
            && (path.contains("../") || path.starts_with("../") || path.contains("..\\") || path.starts_with("..\\"))
        {
            return Err(anyhow::anyhow!("Path traversal detected in arguments"));
        }

        Ok(())
    }

    /// Execute a tool call after validating arguments.
    ///
    /// Public-facing version that takes ownership of `args` for compatibility
    /// with existing callers. Delegates to the reference-taking implementation
    /// to avoid unnecessary cloning when the caller already has a reference.
    async fn execute_tool_with_validation(&self, tool_name: &str, args: Value) -> Result<Value> {
        self.execute_tool_with_validation_ref(tool_name, &args).await
    }

    // Internal reference-taking implementation to avoid cloning when not necessary.
    async fn execute_tool_with_validation_ref(&self, tool_name: &str, args: &Value) -> Result<Value> {
        if !self.config.enabled {
            return Err(anyhow!("MCP support is disabled in the current configuration"));
        }

        self.validate_tool_arguments(tool_name, args)?;

        let provider = self.resolve_provider_for_tool(tool_name).await?;
        let allowlist_snapshot = self.state.read().allowlist.clone();
        let result = provider
            .call_tool(tool_name, args, self.tool_timeout(), &allowlist_snapshot)
            .await?;

        Self::format_tool_result(&provider.name, tool_name, result)
    }

    /// Refresh the internal allow list at runtime.
    pub fn update_allowlist(&self, allowlist: McpAllowListConfig) {
        let providers: Vec<Arc<McpProvider>> = {
            let mut state = self.state.write();
            state.allowlist = allowlist;
            state.tool_provider_index.clear();
            state.resource_provider_index.clear();
            state.prompt_provider_index.clear();
            state.providers.values().cloned().collect()
        };

        for provider in providers {
            provider.invalidate_caches();
        }
    }

    /// Current allow list snapshot.
    pub fn current_allowlist(&self) -> McpAllowListConfig {
        self.state.read().allowlist.clone()
    }

    /// Return the provider name serving the given tool if previously cached.
    pub fn provider_for_tool(&self, tool_name: &str) -> Option<String> {
        self.state.read().tool_provider_index.get(tool_name).cloned()
    }

    /// Return the provider responsible for the given resource URI if known.
    pub fn provider_for_resource(&self, uri: &str) -> Option<String> {
        self.state.read().resource_provider_index.get(uri).cloned()
    }

    /// Return the provider that exposes the given prompt if known.
    pub fn provider_for_prompt(&self, prompt_name: &str) -> Option<String> {
        self.state.read().prompt_provider_index.get(prompt_name).cloned()
    }

    /// Execute a tool call on the appropriate provider.
    pub async fn execute_tool(&self, tool_name: &str, args: Value) -> Result<Value> {
        self.execute_tool_with_validation(tool_name, args).await
    }

    /// List all tools from all active providers.
    pub async fn list_tools(&self) -> Result<Vec<McpToolInfo>> {
        self.collect_tools(false).await
    }

    /// List all resources exposed by connected MCP providers.
    pub async fn list_resources(&self) -> Result<Vec<McpResourceInfo>> {
        self.collect_resources(false).await
    }

    /// Force refresh and list resources from providers.
    pub async fn refresh_resources(&self) -> Result<Vec<McpResourceInfo>> {
        self.collect_resources(true).await
    }

    /// List all prompts advertised by connected MCP providers.
    pub async fn list_prompts(&self) -> Result<Vec<McpPromptInfo>> {
        self.collect_prompts(false).await
    }

    /// Force refresh and list prompts from providers.
    pub async fn refresh_prompts(&self) -> Result<Vec<McpPromptInfo>> {
        self.collect_prompts(true).await
    }

    /// Read a single resource from its originating provider.
    pub async fn read_resource(&self, uri: &str) -> Result<McpResourceData> {
        let provider = self.resolve_provider_for_resource(uri).await?;
        let provider_name = provider.name.clone();
        let allowlist_snapshot = self.state.read().allowlist.clone();
        let data = provider.read_resource(uri, self.request_timeout(), &allowlist_snapshot).await?;
        drop(self.state.write().resource_provider_index.insert(uri.into(), provider_name));
        Ok(data)
    }

    /// Retrieve a rendered prompt from its originating provider.
    pub async fn get_prompt(
        &self,
        prompt_name: &str,
        arguments: Option<hashbrown::HashMap<String, String>>,
    ) -> Result<McpPromptDetail> {
        let provider = self.resolve_provider_for_prompt(prompt_name).await?;
        let provider_name = provider.name.clone();
        let allowlist_snapshot = self.state.read().allowlist.clone();
        let prompt = provider
            .get_prompt(prompt_name, arguments.unwrap_or_default(), self.request_timeout(), &allowlist_snapshot)
            .await?;
        drop(
            self.state
                .write()
                .prompt_provider_index
                .insert(prompt_name.into(), provider_name),
        );
        Ok(prompt)
    }

    /// Shutdown all active provider connections.
    pub async fn shutdown(&self) -> Result<()> {
        let providers: Vec<Arc<McpProvider>> = {
            let mut state = self.state.write();
            let values: Vec<_> = state.providers.values().cloned().collect();
            state.providers.clear();
            state.tool_provider_index.clear();
            state.resource_provider_index.clear();
            state.prompt_provider_index.clear();
            values
        };

        if providers.is_empty() {
            info!("No active MCP connections to shutdown");
            return Ok(());
        }

        info!("Shutting down {} MCP providers", providers.len());
        for provider in providers {
            if let Err(err) = provider.shutdown().await {
                warn!("Provider '{}' shutdown returned error: {err}", provider.name);
            }
        }
        Ok(())
    }

    /// Current status snapshot for UI/debugging purposes.
    pub fn get_status(&self) -> McpClientStatus {
        let state = self.state.read();
        let providers = &state.providers;
        // Use iterator to collect keys directly without intermediate push
        let mut configured_providers: Vec<String> = providers.keys().cloned().collect();
        configured_providers.sort_unstable();
        McpClientStatus {
            enabled: self.config.enabled,
            provider_count: providers.len(),
            active_connections: providers.len(),
            configured_providers,
        }
    }

    /// Return configured MCP servers and their current connection state.
    ///
    /// Connected servers additionally report the protocol version settled
    /// during the handshake (`negotiated_protocol_version`), which may differ
    /// from the configured version when the server negotiates down.
    pub async fn list_servers(&self) -> Vec<Value> {
        let live: Vec<(McpProviderConfig, Option<Arc<McpProvider>>)> = {
            let state = self.state.read();
            self.config
                .providers
                .iter()
                .map(|provider_config| (provider_config.clone(), state.providers.get(&provider_config.name).cloned()))
                .collect()
        };
        let mut servers = Vec::with_capacity(live.len());
        for (provider_config, provider) in live {
            let negotiated = match &provider {
                Some(provider) => provider.negotiated_protocol_version().await,
                None => None,
            };
            let connected = provider.is_some();
            let (transport, target) = match &provider_config.transport {
                McpTransportConfig::Stdio(stdio) => ("stdio", Value::String(stdio.command.clone())),
                McpTransportConfig::Http(http) => ("http", Value::String(http.endpoint.clone())),
            };

            servers.push(json!({
                "name": provider_config.name,
                "enabled": provider_config.enabled,
                "connected": connected,
                "connection_state": if connected { "connected" } else { "disconnected" },
                "transport": transport,
                "target": target,
                "negotiated_protocol_version": negotiated.map(Value::String).unwrap_or(Value::Null),
            }));
        }
        servers
    }

    /// Return whether model-callable lifecycle tools are enabled by config.
    pub fn allow_model_lifecycle_control(&self) -> bool {
        self.config.lifecycle.allow_model_control
    }

    /// Connect one configured MCP server by name.
    pub async fn connect_server(&self, server_name: &str) -> Result<()> {
        if !self.config.enabled {
            bail!("MCP support is disabled in the current configuration");
        }

        if self.state.read().providers.contains_key(server_name) {
            return Ok(());
        }

        let provider_config = self
            .config
            .providers
            .iter()
            .find(|provider| provider.name == server_name)
            .cloned()
            .ok_or_else(|| anyhow!("MCP server '{server_name}' is not configured"))?;

        if !provider_config.enabled {
            bail!("MCP server '{server_name}' is configured but disabled");
        }

        if let Some(reason) = self.requirement_mismatch_reason(&provider_config) {
            bail!("Cannot connect MCP server '{}': {}", provider_config.name, reason);
        }

        if matches!(provider_config.transport, McpTransportConfig::Http(_)) && !self.config.experimental_use_rmcp_client
        {
            bail!(
                "Cannot connect MCP HTTP server '{}' while experimental_use_rmcp_client is disabled",
                provider_config.name
            );
        }

        let allowlist_snapshot = self.state.read().allowlist.clone();
        let tool_timeout = self.tool_timeout();
        let provider = self
            .connect_and_initialize_provider(&provider_config, &allowlist_snapshot, tool_timeout)
            .await?;

        if let Err(err) = provider.cached_tools_or_refresh_shared(&allowlist_snapshot, tool_timeout).await {
            warn!("Connected MCP server '{}' but failed to refresh tools: {err}", server_name);
        } else if let Some(cache) = provider.cached_tools_shared().await {
            self.record_tool_provider(&provider.name, &cache);
        }

        drop(self.state.write().providers.insert(provider.name.clone(), Arc::new(provider)));
        Ok(())
    }

    /// Disconnect one active MCP server by name.
    pub async fn disconnect_server(&self, server_name: &str) -> Result<()> {
        let provider = {
            let mut state = self.state.write();
            let provider = state
                .providers
                .remove(server_name)
                .ok_or_else(|| anyhow!("MCP server '{server_name}' is not connected"))?;
            state
                .tool_provider_index
                .retain(|_, provider_name| provider_name != server_name);
            state
                .resource_provider_index
                .retain(|_, provider_name| provider_name != server_name);
            state
                .prompt_provider_index
                .retain(|_, provider_name| provider_name != server_name);
            provider
        };

        provider.shutdown().await?;
        Ok(())
    }

    /// Sync MCP tool descriptions to files for dynamic context discovery
    ///
    /// This implements Cursor-style dynamic context discovery:
    /// - Tool descriptions are written to `.vtcode/mcp/tools/{provider}/{tool}.md`
    /// - Status is written to `.vtcode/mcp/status.json`
    /// - Agents can discover tools via grep/read_file without loading all schemas
    ///
    /// Returns the paths to written files (index path, tool count)
    pub async fn sync_tools_to_files(&self, workspace_root: &Path) -> Result<(PathBuf, usize)> {
        let tools = self.list_tools().await?;
        let mcp_dir = workspace_root.join(".vtcode").join("mcp");
        let tools_dir = mcp_dir.join("tools");

        // Create directories
        ensure_dir_exists(&tools_dir)
            .await
            .with_context(|| format!("Failed to create MCP tools directory: {}", tools_dir.display()))?;

        // Group tools by provider
        let mut by_provider: BTreeMap<String, Vec<&McpToolInfo>> = BTreeMap::new();
        for tool in &tools {
            by_provider.entry(tool.provider.clone()).or_default().push(tool);
        }

        // Write tool files per provider
        for (provider, provider_tools) in &by_provider {
            let provider_dir = tools_dir.join(sanitize_filename(provider));
            ensure_dir_exists(&provider_dir)
                .await
                .with_context(|| format!("Failed to create provider directory: {}", provider_dir.display()))?;

            for tool in provider_tools {
                let tool_content = format_tool_markdown(tool);
                let tool_path = provider_dir.join(format!("{}.md", sanitize_filename(&tool.name)));
                write_file_with_context(&tool_path, &tool_content, "MCP tool file")
                    .await
                    .with_context(|| format!("Failed to write tool file: {}", tool_path.display()))?;
            }
        }

        // Write index file
        let index_content = self.generate_tools_index(&tools, &by_provider);
        let index_path = tools_dir.join("INDEX.md");
        write_file_with_context(&index_path, &index_content, "MCP tools index")
            .await
            .with_context(|| format!("Failed to write MCP tools index: {}", index_path.display()))?;

        // Write status file
        let status = self.generate_status_json();
        let status_path = mcp_dir.join("status.json");
        let status_json = serde_json::to_string_pretty(&status)?;
        write_file_with_context(&status_path, &status_json, "MCP status")
            .await
            .with_context(|| format!("Failed to write MCP status: {}", status_path.display()))?;

        info!(
            tools = tools.len(),
            providers = by_provider.len(),
            index = %index_path.display(),
            "Synced MCP tool descriptions to files"
        );

        Ok((index_path, tools.len()))
    }

    /// Generate INDEX.md content for MCP tools
    fn generate_tools_index(&self, tools: &[McpToolInfo], by_provider: &BTreeMap<String, Vec<&McpToolInfo>>) -> String {
        let mut content = String::new();
        content.push_str("# MCP Tools Index\n\n");
        content.push_str("This file lists all available MCP tools for dynamic discovery.\n");
        content.push_str("Use `read_file` on individual tool files for full schema details.\n\n");

        if tools.is_empty() {
            content.push_str("*No MCP tools available.*\n\n");
            content.push_str("Configure MCP servers in `vtcode.toml` or `.mcp.json`.\n");
        } else {
            content.push_str(&format!("**Total Tools**: {}\n\n", tools.len()));

            // Summary table
            content.push_str("## Quick Reference\n\n");
            content.push_str("| Provider | Tool | Description |\n");
            content.push_str("|----------|------|-------------|\n");

            for tool in tools {
                let desc = tool.description.lines().next().unwrap_or(&tool.description);
                let desc_truncated = vtcode_commons::formatting::truncate_byte_budget(desc, 57, "...");
                content.push_str(&format!(
                    "| {} | `{}` | {} |\n",
                    tool.provider,
                    tool.name,
                    desc_truncated.replace('|', "\\|")
                ));
            }

            // Per-provider sections
            content.push_str("\n## Tools by Provider\n\n");
            for (provider, provider_tools) in by_provider {
                content.push_str(&format!("### {provider}\n\n"));
                for tool in provider_tools {
                    content.push_str(&format!(
                        "- **{}**: {}\n  - Path: `.vtcode/mcp/tools/{}/{}.md`\n",
                        tool.name,
                        tool.description.lines().next().unwrap_or(&tool.description),
                        sanitize_filename(provider),
                        sanitize_filename(&tool.name)
                    ));
                }
                content.push('\n');
            }
        }

        content.push_str("\n---\n");
        content.push_str("*Generated automatically. Do not edit manually.*\n");

        content
    }

    /// Generate status.json content
    fn generate_status_json(&self) -> Value {
        let status = self.get_status();
        json!({
            "enabled": status.enabled,
            "provider_count": status.provider_count,
            "active_connections": status.active_connections,
            "configured_providers": status.configured_providers,
            "last_updated": Utc::now().to_rfc3339(),
        })
    }

    async fn collect_tools(&self, force_refresh: bool) -> Result<Vec<McpToolInfo>> {
        // Collect provider references in one pass
        let (providers, allowlist) = {
            let state = self.state.read();
            (providers_sorted_by_name(&state.providers), state.allowlist.clone())
        };

        if providers.is_empty() {
            return Ok(Vec::new());
        }

        let timeout = self.tool_timeout();
        let mut all_tools = Vec::with_capacity(128);
        let mut index_updates: FxHashMap<String, String> = FxHashMap::with_capacity_and_hasher(128, Default::default());

        for provider in providers {
            let provider_name = provider.name.clone();
            let tools = if force_refresh {
                provider.refresh_tools(&allowlist, timeout).await
            } else {
                provider.list_tools(&allowlist, timeout).await
            };

            match tools {
                Ok(tools) => {
                    for tool in &tools {
                        let _ = index_updates.entry(tool.name.clone()).or_insert_with(|| provider_name.clone());
                    }
                    all_tools.extend(tools);
                }
                Err(err) => {
                    warn!("Failed to list tools for provider '{}': {err}", provider_name);
                }
            }
        }

        if !index_updates.is_empty() || force_refresh {
            let mut state = self.state.write();
            if index_updates.is_empty() {
                state.tool_provider_index.clear();
            } else {
                state.tool_provider_index = index_updates;
            }
        }

        Ok(all_tools)
    }

    async fn collect_resources(&self, force_refresh: bool) -> Result<Vec<McpResourceInfo>> {
        // Collect provider references in one pass
        let (providers, allowlist) = {
            let state = self.state.read();
            (providers_sorted_by_name(&state.providers), state.allowlist.clone())
        };

        if providers.is_empty() {
            self.state.write().resource_provider_index.clear();
            return Ok(Vec::new());
        }

        let timeout = self.request_timeout();
        let mut all_resources = Vec::with_capacity(64);

        for provider in providers {
            let resources = if force_refresh {
                provider.refresh_resources(&allowlist, timeout).await
            } else {
                provider.list_resources(&allowlist, timeout).await
            };

            match resources {
                Ok(resources) => {
                    all_resources.extend(resources);
                }
                Err(err) => {
                    warn!("Failed to list resources for provider '{}': {err}", provider.name);
                }
            }
        }

        let mut state = self.state.write();
        let index = &mut state.resource_provider_index;
        index.clear();
        for resource in &all_resources {
            let _ = index.entry(resource.uri.clone()).or_insert_with(|| resource.provider.clone());
        }

        Ok(all_resources)
    }

    async fn collect_prompts(&self, force_refresh: bool) -> Result<Vec<McpPromptInfo>> {
        // Collect provider references in one pass
        let (providers, allowlist) = {
            let state = self.state.read();
            (providers_sorted_by_name(&state.providers), state.allowlist.clone())
        };

        if providers.is_empty() {
            self.state.write().prompt_provider_index.clear();
            return Ok(Vec::new());
        }

        let timeout = self.request_timeout();
        let mut all_prompts = Vec::with_capacity(32);

        for provider in providers {
            let prompts = if force_refresh {
                provider.refresh_prompts(&allowlist, timeout).await
            } else {
                provider.list_prompts(&allowlist, timeout).await
            };

            match prompts {
                Ok(prompts) => {
                    all_prompts.extend(prompts);
                }
                Err(err) => {
                    warn!("Failed to list prompts for provider '{}': {err}", provider.name);
                }
            }
        }

        let mut state = self.state.write();
        let index = &mut state.prompt_provider_index;
        index.clear();
        for prompt in &all_prompts {
            let _ = index.entry(prompt.name.clone()).or_insert_with(|| prompt.provider.clone());
        }

        Ok(all_prompts)
    }

    async fn resolve_provider_for_tool(&self, tool_name: &str) -> Result<Arc<McpProvider>> {
        if !self.config.enabled {
            return Err(anyhow!("MCP support is disabled in the current configuration"));
        }

        if let Some(provider) = self.provider_for_tool(tool_name)
            && let Some(found) = self.state.read().providers.get(&provider)
        {
            return Ok(found.clone());
        }

        let (allowlist, providers) = {
            let state = self.state.read();
            (state.allowlist.clone(), providers_sorted_by_name(&state.providers))
        };
        let timeout = self.tool_timeout();

        if providers.is_empty() {
            if self.config.providers.is_empty() {
                return Err(anyhow!(
                    "No MCP providers are configured. Use `vtcode mcp add` or update vtcode.toml to register one."
                ));
            }

            return Err(anyhow!(
                "No MCP providers are currently connected. Ensure MCP initialization completed successfully."
            ));
        }

        for provider in providers {
            match provider.has_tool(tool_name, &allowlist, timeout).await {
                Ok(true) => {
                    drop(
                        self.state
                            .write()
                            .tool_provider_index
                            .insert(tool_name.into(), provider.name.clone()),
                    );
                    return Ok(provider);
                }
                Ok(false) => continue,
                Err(err) => {
                    warn!("Error checking tool '{}' on provider '{}': {err}", tool_name, provider.name);
                }
            }
        }

        match self.collect_tools(true).await {
            Ok(_) => {
                if let Some(provider) = self.provider_for_tool(tool_name)
                    && let Some(found) = self.state.read().providers.get(&provider)
                {
                    return Ok(found.clone());
                }
            }
            Err(err) => {
                warn!("Failed to refresh MCP tool caches while resolving '{}': {err}", tool_name);
            }
        }

        Err(anyhow!(
            "Tool '{tool_name}' not found on any MCP provider.\n\n\
            To use this tool:\n\
            1. Install the MCP server: `uv tool install mcp-server-{tool_name}`\n\
            2. Add to vtcode.toml:\n   \
               [[mcp.providers]]\n   \
               name = \"{tool_name}\"\n   \
               command = \"uvx\"\n   \
               args = [\"mcp-server-{tool_name}\"]\n\
            3. Restart VT Code\n\n\
            Or use the built-in alternative if available (e.g., web_fetch instead of mcp_fetch)"
        ))
    }

    async fn resolve_provider_for_resource(&self, uri: &str) -> Result<Arc<McpProvider>> {
        if let Some(provider) = self.provider_for_resource(uri)
            && let Some(found) = self.state.read().providers.get(&provider)
        {
            return Ok(found.clone());
        }

        let (allowlist, providers) = {
            let state = self.state.read();
            (state.allowlist.clone(), providers_sorted_by_name(&state.providers))
        };
        let timeout = self.request_timeout();

        for provider in providers {
            match provider.has_resource(uri, &allowlist, timeout).await {
                Ok(true) => {
                    drop(
                        self.state
                            .write()
                            .resource_provider_index
                            .insert(uri.into(), provider.name.clone()),
                    );
                    return Ok(provider);
                }
                Ok(false) => continue,
                Err(err) => {
                    warn!("Error checking resource '{}' on provider '{}': {err}", uri, provider.name);
                }
            }
        }

        Err(anyhow!("Resource '{uri}' not found on any MCP provider"))
    }

    async fn resolve_provider_for_prompt(&self, prompt_name: &str) -> Result<Arc<McpProvider>> {
        if let Some(provider) = self.provider_for_prompt(prompt_name)
            && let Some(found) = self.state.read().providers.get(&provider)
        {
            return Ok(found.clone());
        }

        let (allowlist, providers) = {
            let state = self.state.read();
            (state.allowlist.clone(), providers_sorted_by_name(&state.providers))
        };
        let timeout = self.request_timeout();

        for provider in providers {
            match provider.has_prompt(prompt_name, &allowlist, timeout).await {
                Ok(true) => {
                    drop(
                        self.state
                            .write()
                            .prompt_provider_index
                            .insert(prompt_name.into(), provider.name.clone()),
                    );
                    return Ok(provider);
                }
                Ok(false) => continue,
                Err(err) => {
                    warn!("Error checking prompt '{}' on provider '{}': {err}", prompt_name, provider.name);
                }
            }
        }

        Err(anyhow!("Prompt '{prompt_name}' not found on any MCP provider"))
    }

    fn record_tool_provider(&self, provider: &str, tools: &[McpToolInfo]) {
        let mut state = self.state.write();
        let index = &mut state.tool_provider_index;
        for tool in tools {
            let _ = index
                .entry(tool.name.clone())
                .and_modify(|owner| {
                    if provider < owner.as_str() {
                        *owner = provider.to_string();
                    }
                })
                .or_insert_with(|| provider.to_string());
        }
    }

    async fn connect_and_initialize_provider(
        &self,
        provider_config: &McpProviderConfig,
        allowlist_snapshot: &McpAllowListConfig,
        tool_timeout: Option<Duration>,
    ) -> Result<McpProvider> {
        let total_attempts = self.provider_retry_attempts();
        let mut last_error: Option<anyhow::Error> = None;

        for attempt_idx in 0..total_attempts {
            let attempt_number = attempt_idx + 1;
            match self
                .connect_and_initialize_provider_once(provider_config, allowlist_snapshot, tool_timeout)
                .await
            {
                Ok(provider) => return Ok(provider),
                Err(err) => {
                    if attempt_number == total_attempts {
                        return Err(err);
                    }

                    let retries_remaining = total_attempts - attempt_number;
                    warn!(
                        provider = provider_config.name.as_str(),
                        attempt = attempt_number,
                        retries_remaining,
                        error = %err,
                        "MCP provider initialization failed; retrying"
                    );
                    last_error = Some(err);
                    tokio::time::sleep(Self::provider_retry_delay(attempt_idx)).await;
                }
            }
        }

        Err(last_error.unwrap_or_else(|| anyhow!("Failed to initialize MCP provider '{}'", provider_config.name)))
    }

    async fn connect_and_initialize_provider_once(
        &self,
        provider_config: &McpProviderConfig,
        allowlist_snapshot: &McpAllowListConfig,
        tool_timeout: Option<Duration>,
    ) -> Result<McpProvider> {
        let provider = McpProvider::connect(
            provider_config.clone(),
            self.elicitation_handler.clone(),
            self.sandbox_context.clone(),
        )
        .await
        .with_context(|| format!("Failed to connect to MCP provider '{}'", provider_config.name))?;
        let provider_startup_timeout = self.resolve_startup_timeout(provider_config);
        provider
            .initialize(
                self.build_initialize_params(&provider),
                provider_startup_timeout,
                tool_timeout,
                allowlist_snapshot,
            )
            .await
            .with_context(|| format!("Failed to initialize MCP provider '{}'", provider_config.name))?;
        Ok(provider)
    }

    fn startup_timeout(&self) -> Option<Duration> {
        match self.config.startup_timeout_seconds {
            Some(0) => None,
            Some(value) => Some(Duration::from_secs(value)),
            None => self.request_timeout(),
        }
    }

    fn requirement_mismatch_reason(&self, provider_config: &McpProviderConfig) -> Option<String> {
        let requirements = &self.config.requirements;
        if !requirements.enforce {
            return None;
        }

        match &provider_config.transport {
            McpTransportConfig::Stdio(stdio) => {
                if requirements
                    .allowed_stdio_commands
                    .iter()
                    .any(|allowed| allowed == &stdio.command)
                {
                    None
                } else {
                    Some(format!("stdio command '{}' is not allowlisted", stdio.command))
                }
            }
            McpTransportConfig::Http(http) => {
                if requirements
                    .allowed_http_endpoints
                    .iter()
                    .any(|allowed| allowed == &http.endpoint)
                {
                    None
                } else {
                    Some(format!("HTTP endpoint '{}' is not allowlisted", http.endpoint))
                }
            }
        }
    }

    fn resolve_startup_timeout(&self, provider_config: &McpProviderConfig) -> Option<Duration> {
        if let Some(timeout_ms) = provider_config.startup_timeout_ms {
            if timeout_ms == 0 {
                None
            } else {
                Some(Duration::from_millis(timeout_ms))
            }
        } else {
            self.startup_timeout()
        }
    }

    fn provider_retry_attempts(&self) -> usize {
        self.config.retry_attempts.try_into().unwrap_or(usize::MAX).saturating_add(1)
    }

    fn provider_retry_delay(attempt_idx: usize) -> Duration {
        let base_ms = 250u64;
        let max_ms = 5000u64;
        let exp = base_ms.saturating_mul(2u64.saturating_pow(u32::try_from(attempt_idx).unwrap_or(u32::MAX)));
        let delay = exp.min(max_ms);
        Duration::from_millis(delay)
    }

    fn tool_timeout(&self) -> Option<Duration> {
        match self.config.tool_timeout_seconds {
            Some(0) => None,
            Some(value) => Some(Duration::from_secs(value)),
            None => self.request_timeout(),
        }
    }

    fn request_timeout(&self) -> Option<Duration> {
        if self.config.request_timeout_seconds == 0 {
            None
        } else {
            Some(Duration::from_secs(self.config.request_timeout_seconds))
        }
    }

    fn build_initialize_params(&self, _provider: &McpProvider) -> InitializeRequestParams {
        let mut capabilities = ClientCapabilities::default();
        {
            let mut roots_cap = RootsCapabilities::default();
            roots_cap.list_changed = Some(true);
            capabilities.roots = Some(roots_cap);
        }

        if self.elicitation_handler.is_some() {
            // Elicitation is now a first-class capability in rmcp
            capabilities.elicitation = Some(
                rmcp::model::ElicitationCapability::new()
                    .with_form(rmcp::model::FormElicitationCapability::new().with_schema_validation(true)),
            );
        }

        InitializeRequestParams::new(capabilities, super::utils::build_client_implementation())
            .with_protocol_version(super::rmcp_client::latest_protocol_version())
    }

    pub(super) fn normalize_arguments(args: &Value) -> Map<String, Value> {
        match args {
            Value::Null => Map::new(),
            Value::Object(map) => map.clone(),
            other => {
                let mut map = Map::new();
                drop(map.insert("value".to_owned(), other.clone()));
                map
            }
        }
    }

    fn format_tool_result(provider_name: &str, tool_name: &str, result: CallToolResult) -> Result<Value> {
        // Convert result to JSON to access fields flexibly
        let result_json = serde_json::to_value(&result)?;
        let result_obj = result_json.as_object();

        // Check for error - handle both rmcp's is_error field and meta message
        let is_error = result_obj
            .and_then(|o| o.get("isError"))
            .or_else(|| result_obj.and_then(|o| o.get("is_error")))
            .and_then(Value::as_bool)
            .unwrap_or(false);

        if is_error {
            let mut message = result_obj
                .and_then(|o| o.get("_meta"))
                .or_else(|| result_obj.and_then(|o| o.get("meta")))
                .and_then(|m| m.get("message"))
                .and_then(Value::as_str)
                .map(str::to_owned);

            // Try to find text content in the content array
            if message.is_none()
                && let Some(content) = result_obj.and_then(|o| o.get("content")).and_then(Value::as_array)
            {
                message = content
                    .iter()
                    .find_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
            }

            let message = message.unwrap_or_else(|| "Unknown MCP tool error".to_owned());
            return Err(anyhow!("MCP tool '{tool_name}' on provider '{provider_name}' reported an error: {message}"));
        }

        let mut payload = Map::new();
        drop(payload.insert("provider".into(), Value::String(provider_name.to_string())));
        drop(payload.insert("tool".into(), Value::String(tool_name.to_string())));

        // Add meta if present
        if let Some(meta) = result_obj
            .and_then(|o| o.get("_meta"))
            .or_else(|| result_obj.and_then(|o| o.get("meta")))
            .and_then(Value::as_object)
            && !meta.is_empty()
        {
            drop(payload.insert("meta".into(), Value::Object(meta.clone())));
        }

        // Add content if present
        if let Some(content) = result_obj.and_then(|o| o.get("content"))
            && !content.is_null()
            && !content.as_array().map(|a| a.is_empty()).unwrap_or(true)
        {
            drop(payload.insert("content".into(), content.clone()));
        }

        Ok(Value::Object(payload))
    }
}

#[async_trait]
impl McpToolExecutor for McpClient {
    async fn execute_mcp_tool(&self, tool_name: &str, args: &Value) -> Result<Value> {
        self.execute_tool_with_validation_ref(tool_name, args).await
    }

    async fn list_mcp_tools(&self) -> Result<Vec<McpToolInfo>> {
        self.collect_tools(false).await
    }

    async fn has_mcp_tool(&self, tool_name: &str) -> Result<bool> {
        if !self.config.enabled {
            return Ok(false);
        }

        if self.provider_for_tool(tool_name).is_some() {
            return Ok(true);
        }

        if self.state.read().providers.is_empty() {
            if self.config.providers.is_empty() {
                return Ok(false);
            }

            bail!("No MCP providers are currently connected. Ensure MCP initialization completed successfully.");
        }

        let tools = self.collect_tools(false).await?;
        Ok(tools.iter().any(|tool| tool.name == tool_name))
    }

    fn get_status(&self) -> McpClientStatus {
        self.get_status()
    }
}

#[cfg(test)]
mod tests {
    use super::McpClient;
    use rustc_hash::FxHashMap;
    use vtcode_config::mcp::{
        McpClientConfig, McpHttpServerConfig, McpProviderConfig, McpRequirementsConfig, McpStdioServerConfig,
        McpTransportConfig,
    };

    fn base_config() -> McpClientConfig {
        McpClientConfig {
            enabled: true,
            requirements: McpRequirementsConfig {
                enforce: true,
                allowed_stdio_commands: vec!["uvx".to_string()],
                allowed_http_endpoints: vec!["https://allowed.example/mcp".to_string()],
            },
            ..McpClientConfig::default()
        }
    }

    #[test]
    fn requirements_allow_matching_stdio_command() {
        let client = McpClient::new(base_config());
        let provider = McpProviderConfig {
            name: "time".to_string(),
            transport: McpTransportConfig::Stdio(McpStdioServerConfig {
                command: "uvx".to_string(),
                args: vec![],
                working_directory: None,
            }),
            ..McpProviderConfig::default()
        };

        assert!(client.requirement_mismatch_reason(&provider).is_none());
    }

    #[test]
    fn requirements_block_unmatched_stdio_command() {
        let client = McpClient::new(base_config());
        let provider = McpProviderConfig {
            name: "time".to_string(),
            transport: McpTransportConfig::Stdio(McpStdioServerConfig {
                command: "npx".to_string(),
                args: vec![],
                working_directory: None,
            }),
            ..McpProviderConfig::default()
        };

        assert!(
            client
                .requirement_mismatch_reason(&provider)
                .is_some_and(|reason| reason.contains("not allowlisted"))
        );
    }

    #[test]
    fn requirements_block_unmatched_http_endpoint() {
        let client = McpClient::new(base_config());
        let provider = McpProviderConfig {
            name: "remote".to_string(),
            transport: McpTransportConfig::Http(McpHttpServerConfig {
                endpoint: "https://blocked.example/mcp".to_string(),
                ..McpHttpServerConfig::default()
            }),
            ..McpProviderConfig::default()
        };

        assert!(
            client
                .requirement_mismatch_reason(&provider)
                .is_some_and(|reason| reason.contains("not allowlisted"))
        );
    }

    #[tokio::test]
    async fn list_servers_includes_configured_provider_metadata() {
        let mut config = base_config();
        config.providers = vec![McpProviderConfig {
            name: "calendar".to_string(),
            transport: McpTransportConfig::Http(McpHttpServerConfig {
                endpoint: "https://calendar.example/mcp".to_string(),
                ..McpHttpServerConfig::default()
            }),
            ..McpProviderConfig::default()
        }];

        let client = McpClient::new(config);
        let servers = client.list_servers().await;
        assert_eq!(servers.len(), 1);
        assert_eq!(servers[0]["name"], "calendar");
        assert_eq!(servers[0]["connected"], false);
        assert_eq!(servers[0]["connection_state"], "disconnected");
        assert_eq!(servers[0]["transport"], "http");
        assert_eq!(servers[0]["target"], "https://calendar.example/mcp");
        assert!(
            servers[0]["negotiated_protocol_version"].is_null(),
            "disconnected servers have no negotiated version"
        );
    }

    #[tokio::test]
    async fn connect_server_rejects_unknown_server_name() {
        let client = McpClient::new(base_config());
        let err = Box::pin(client.connect_server("missing"))
            .await
            .expect_err("missing server should error");
        assert!(err.to_string().contains("not configured"));
    }

    #[tokio::test]
    async fn disconnect_server_rejects_unknown_server_name() {
        let client = McpClient::new(base_config());
        let err = client
            .disconnect_server("missing")
            .await
            .expect_err("missing server should error");
        assert!(err.to_string().contains("not connected"));
    }

    #[test]
    fn providers_sorted_by_name_ignores_insertion_order() {
        let names = ["zeta", "Alpha", "alpha", "beta", "alpha2"];
        let forward: FxHashMap<String, usize> =
            names.iter().enumerate().map(|(idx, name)| ((*name).to_owned(), idx)).collect();
        let reverse: FxHashMap<String, usize> = names
            .iter()
            .enumerate()
            .rev()
            .map(|(idx, name)| ((*name).to_owned(), idx))
            .collect();

        // Values are the original indices of Alpha, alpha, alpha2, beta, zeta.
        let expected = vec![1, 2, 4, 3, 0];
        assert_eq!(super::providers_sorted_by_name(&forward), expected);
        assert_eq!(super::providers_sorted_by_name(&reverse), expected);
    }

    #[test]
    fn record_tool_provider_keeps_first_provider_by_name_regardless_of_connect_order() {
        use super::McpToolInfo;
        use serde_json::json;

        let tool = |provider: &str, name: &str| McpToolInfo {
            name: name.to_owned(),
            description: String::new(),
            provider: provider.to_owned(),
            input_schema: json!({}),
            output_schema: None,
        };

        for order in [["alpha", "beta"], ["beta", "alpha"]] {
            let client = McpClient::new(base_config());
            for provider in order {
                client.record_tool_provider(
                    provider,
                    &[tool(provider, "shared"), tool(provider, &format!("{provider}_only"))],
                );
            }
            assert_eq!(client.provider_for_tool("shared").as_deref(), Some("alpha"), "order {order:?}");
            assert_eq!(client.provider_for_tool("beta_only").as_deref(), Some("beta"));
        }
    }
}