vtcode_core/mcp/

client.rs

use crate::config::mcp::{
    McpAllowListConfig, McpClientConfig, McpProviderConfig, McpTransportConfig,
};
use crate::utils::file_utils::{ensure_dir_exists, write_file_with_context};
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::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, error, info, warn};

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

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

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

impl McpClient {
    /// Create a new MCP client from the configuration.
    pub fn new(config: McpClientConfig) -> 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,
        }
    }

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

        // Sequential initialization
        Box::pin(self.initialize_sequential()).await
    }

    /// Initialize providers sequentially (fallback method)
    async fn initialize_sequential(&mut self) -> Result<()> {
        let tool_timeout = self.tool_timeout();
        let allowlist_snapshot = self.state.read().allowlist.clone();

        let mut initialized = FxHashMap::default();

        for provider_config in &self.config.providers {
            if !provider_config.enabled {
                debug!(
                    "MCP provider '{}' is disabled; skipping",
                    provider_config.name
                );
                continue;
            }

            if let Some(reason) = self.requirement_mismatch_reason(provider_config) {
                warn!(
                    "Skipping MCP provider '{}' due to requirements policy: {}",
                    provider_config.name, reason
                );
                continue;
            }

            if matches!(provider_config.transport, McpTransportConfig::Http(_))
                && !self.config.experimental_use_rmcp_client
            {
                warn!(
                    "Skipping MCP HTTP provider '{}' because experimental_use_rmcp_client is disabled",
                    provider_config.name
                );
                continue;
            }

            match self
                .connect_and_initialize_provider(provider_config, &allowlist_snapshot, tool_timeout)
                .await
            {
                Ok(provider) => {
                    if let Err(err) = provider
                        .cached_tools_or_refresh(&allowlist_snapshot, tool_timeout)
                        .await
                    {
                        error!(
                            "Failed to fetch tools for provider '{}': {err}",
                            provider_config.name
                        );
                    } else if let Some(cache) = provider.cached_tools().await {
                        self.record_tool_provider(&provider.name, &cache);
                    }

                    initialized.insert(provider.name.clone(), Arc::new(provider));
                    info!(
                        "Successfully initialized MCP provider '{}'",
                        provider_config.name
                    );
                }
                Err(err) => {
                    error!(
                        "Failed to connect to MCP provider '{}': {err}",
                        provider_config.name
                    );
                }
            }
        }

        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| s.len()) as u32;

            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.
    pub 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?;
        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?;
        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 configured_providers: Vec<String> = providers.keys().cloned().collect();
        McpClientStatus {
            enabled: self.config.enabled,
            provider_count: providers.len(),
            active_connections: providers.len(),
            configured_providers,
        }
    }

    /// Return configured MCP servers and their current connection state.
    pub fn list_servers(&self) -> Vec<Value> {
        let state = self.state.read();
        self.config
            .providers
            .iter()
            .map(|provider_config| {
                let connected = state.providers.contains_key(&provider_config.name);
                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()))
                    }
                };

                json!({
                    "name": provider_config.name,
                    "enabled": provider_config.enabled,
                    "connected": connected,
                    "connection_state": if connected { "connected" } else { "disconnected" },
                    "transport": transport,
                    "target": target,
                })
            })
            .collect()
    }

    /// 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 '{}' is not configured", server_name))?;

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

        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(&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().await {
            self.record_tool_provider(&provider.name, &cache);
        }

        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 '{}' is not connected", server_name))?;
            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: FxHashMap<String, Vec<&McpToolInfo>> = FxHashMap::default();
        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: &FxHashMap<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!("### {}\n\n", provider));
                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();
            (
                state.providers.values().cloned().collect::<Vec<_>>(),
                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 {
                        index_updates.insert(tool.name.clone(), 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();
            (
                state.providers.values().cloned().collect::<Vec<_>>(),
                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 {
            index.insert(resource.uri.clone(), 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();
            (
                state.providers.values().cloned().collect::<Vec<_>>(),
                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 {
            index.insert(prompt.name.clone(), 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(),
                state.providers.values().cloned().collect::<Vec<_>>(),
            )
        };
        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) => {
                    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 '{}' not found on any MCP provider.\n\n\
            To use this tool:\n\
            1. Install the MCP server: `uv tool install mcp-server-{}`\n\
            2. Add to vtcode.toml:\n   \
               [[mcp.providers]]\n   \
               name = \"{}\"\n   \
               command = \"uvx\"\n   \
               args = [\"mcp-server-{}\"]\n\
            3. Restart VT Code\n\n\
            Or use the built-in alternative if available (e.g., web_fetch instead of mcp_fetch)",
            tool_name,
            tool_name,
            tool_name,
            tool_name
        ))
    }

    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(),
                state.providers.values().cloned().collect::<Vec<_>>(),
            )
        };
        let timeout = self.request_timeout();

        for provider in providers {
            match provider.has_resource(uri, &allowlist, timeout).await {
                Ok(true) => {
                    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 '{}' not found on any MCP provider", uri))
    }

    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(),
                state.providers.values().cloned().collect::<Vec<_>>(),
            )
        };
        let timeout = self.request_timeout();

        for provider in providers {
            match provider.has_prompt(prompt_name, &allowlist, timeout).await {
                Ok(true) => {
                    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 '{}' not found on any MCP provider",
            prompt_name
        ))
    }

    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 {
            index.insert(tool.name.clone(), 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())
                .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 step = u64::try_from(attempt_idx)
            .unwrap_or(u64::MAX)
            .saturating_add(1);
        Duration::from_millis((step * 250).min(1_000))
    }

    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();
        capabilities.roots = Some(RootsCapabilities {
            list_changed: Some(true),
        });

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

        InitializeRequestParams::new(capabilities, super::utils::build_client_implementation())
            .with_protocol_version(rmcp::model::ProtocolVersion::V_2024_11_05)
    }

    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();
                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 '{}' on provider '{}' reported an error: {}",
                tool_name,
                provider_name,
                message
            ));
        }

        let mut payload = Map::new();
        payload.insert("provider".into(), Value::String(provider_name.to_string()));
        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()
        {
            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)
        {
            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 crate::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"))
        );
    }

    #[test]
    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();
        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");
    }

    #[tokio::test]
    async fn connect_server_rejects_unknown_server_name() {
        let client = McpClient::new(base_config());
        let err = 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"));
    }
}