vtcode 0.137.0

use anyhow::{Context, Result};
use std::fmt::{self, Display};
use std::sync::{Arc, RwLock as StdRwLock};
use tokio::sync::{Mutex, RwLock};
use tokio::time::{Duration, timeout};
use tracing::{error, info, warn};
use vtcode_core::config::mcp::McpClientConfig;
use vtcode_core::exec_policy::{AskForApproval, RejectConfig};
use vtcode_core::mcp::McpClient;

use crate::agent::runloop::mcp_events::McpEvent;

pub(crate) fn approval_policy_from_human_in_the_loop(human_in_the_loop: bool) -> AskForApproval {
    if human_in_the_loop {
        AskForApproval::OnRequest
    } else {
        AskForApproval::Reject(RejectConfig {
            sandbox_approval: true,
            rules: true,
            request_permissions: true,
            mcp_elicitations: true,
        })
    }
}

/// Represents the initialization status of MCP components
#[derive(Clone)]
pub(crate) enum McpInitStatus {
    /// MCP is not enabled
    Disabled,
    /// MCP initialization is in progress
    Initializing { progress: String },
    /// MCP initialization completed successfully
    Ready { client: Arc<McpClient> },
    /// MCP initialization failed
    Error { message: String },
}

impl McpInitStatus {
    #[cfg(test)]
    pub fn is_ready(&self) -> bool {
        matches!(self, McpInitStatus::Ready { .. })
    }

    pub fn is_error(&self) -> bool {
        matches!(self, McpInitStatus::Error { .. })
    }

    pub fn get_error_message(&self) -> Option<&str> {
        match self {
            McpInitStatus::Error { message } => Some(message),
            _ => None,
        }
    }

    #[cfg(test)]
    pub fn is_initializing(&self) -> bool {
        matches!(self, McpInitStatus::Initializing { .. })
    }
}

impl fmt::Debug for McpInitStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            McpInitStatus::Disabled => f.write_str("Disabled"),
            McpInitStatus::Initializing { progress } => {
                f.debug_struct("Initializing").field("progress", progress).finish()
            }
            McpInitStatus::Ready { .. } => f.write_str("Ready { client: <redacted> }"),
            McpInitStatus::Error { message } => f.debug_struct("Error").field("message", message).finish(),
        }
    }
}

impl Display for McpInitStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            McpInitStatus::Disabled => write!(f, "MCP is disabled"),
            McpInitStatus::Initializing { progress } => write!(f, "MCP initializing: {progress}"),
            McpInitStatus::Ready { .. } => write!(f, "MCP ready with active connections"),
            McpInitStatus::Error { message } => write!(f, "MCP error: {message}"),
        }
    }
}

/// Async manager for MCP client initialization and lifecycle
pub(crate) struct AsyncMcpManager {
    /// Configuration for MCP client
    config: Arc<StdRwLock<McpClientConfig>>,
    /// Whether to ring terminal bell for HITL prompts
    hitl_notification_bell: bool,
    /// Approval policy used by MCP elicitation handling.
    approval_policy: Arc<StdRwLock<AskForApproval>>,
    /// Current initialization status
    status: Arc<RwLock<McpInitStatus>>,
    /// Mutex to prevent multiple concurrent initializations
    initialization_mutex: Arc<Mutex<()>>,
    /// Event callback for MCP events
    event_callback: Arc<dyn Fn(McpEvent) + Send + Sync>,
    /// Handle for the background initialization task, aborted on drop.
    init_task: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
}

impl AsyncMcpManager {
    pub(crate) fn new(
        config: McpClientConfig,
        hitl_notification_bell: bool,
        approval_policy: AskForApproval,
        event_callback: Arc<dyn Fn(McpEvent) + Send + Sync>,
    ) -> Self {
        let init_status = if config.enabled {
            McpInitStatus::Initializing { progress: "Initializing MCP client...".to_string() }
        } else {
            McpInitStatus::Disabled
        };

        Self {
            config: Arc::new(StdRwLock::new(config)),
            hitl_notification_bell,
            approval_policy: Arc::new(StdRwLock::new(approval_policy)),
            status: Arc::new(RwLock::new(init_status)),
            initialization_mutex: Arc::new(Mutex::new(())),
            event_callback,
            init_task: std::sync::Mutex::new(None),
        }
    }

    /// Start async initialization of MCP client.
    ///
    /// # Priority Guarantee
    ///
    /// This method is designed to never block tokio worker threads, ensuring
    /// that signal handling (Ctrl+C) remains responsive. The initialization
    /// is spawned as a background task that does not block the caller.
    pub(crate) fn start_initialization(&self) -> Result<()> {
        let config = self.config();
        if !config.enabled {
            // MCP is disabled. Status is already set to Disabled in the
            // constructor (new()), so no additional write is needed.
            // This avoids blocking a tokio worker thread and ensures
            // signal handling remains responsive.
            return Ok(());
        }

        if let Ok(guard) = self.init_task.lock()
            && let Some(existing_task) = guard.as_ref()
            && !existing_task.is_finished()
        {
            return Ok(());
        }

        // Clone what we need for the async task
        let status = Arc::clone(&self.status);
        let mutex = Arc::clone(&self.initialization_mutex);
        let event_callback = Arc::clone(&self.event_callback);
        let hitl_notification_bell = self.hitl_notification_bell;
        let approval_policy = Arc::clone(&self.approval_policy);

        // Spawn the initialization task. Store the JoinHandle so it can be
        // aborted on drop — prevents an orphan task if the manager is dropped
        // before initialization completes.
        let init_handle = tokio::spawn(async move {
            // Acquire the mutex to prevent concurrent initializations
            let _guard = mutex.lock().await;

            // Check if already initialized while mutex was being acquired
            {
                let current_status = status.read().await;
                if matches!(*current_status, McpInitStatus::Ready { .. }) {
                    return; // Already initialized
                }
            }

            // Update status to initializing
            {
                let mut status_guard = status.write().await;
                *status_guard = McpInitStatus::Initializing {
                    progress: "Connecting to MCP providers...".to_string(),
                };
            }

            // Initialize MCP client
            match Self::initialize_mcp_client(config, hitl_notification_bell, approval_policy, event_callback).await {
                Ok(client) => {
                    let mut status_guard = status.write().await;
                    *status_guard = McpInitStatus::Ready { client: Arc::new(client) };
                    info!("MCP client initialized successfully");
                }
                Err(e) => {
                    let error_msg = e.to_string();
                    let mcp_error = if error_msg.contains("No such process")
                        || error_msg.contains("ESRCH")
                        || error_msg.contains("EPIPE")
                        || error_msg.contains("Broken pipe")
                        || error_msg.contains("write EPIPE")
                    {
                        format!("MCP server startup failed: {e}")
                    } else {
                        format!("MCP initialization error: {e}")
                    };

                    let mut status_guard = status.write().await;
                    *status_guard = McpInitStatus::Error { message: mcp_error };
                    error!("MCP client initialization failed: {}", e);
                }
            }
        });
        if let Ok(mut guard) = self.init_task.lock() {
            *guard = Some(init_handle);
        }

        Ok(())
    }

    pub(crate) fn config(&self) -> McpClientConfig {
        match self.config.read() {
            Ok(config) => config.clone(),
            Err(poisoned) => {
                warn!("MCP config lock was poisoned; continuing with last known value");
                poisoned.into_inner().clone()
            }
        }
    }

    pub(crate) async fn reconfigure(&self, config: McpClientConfig) -> Result<()> {
        if let Ok(mut guard) = self.init_task.lock()
            && let Some(task) = guard.take()
        {
            task.abort();
        }

        if let McpInitStatus::Ready { client } = self.get_status().await {
            client.shutdown().await?;
        }

        {
            let mut guard = match self.config.write() {
                Ok(config_guard) => config_guard,
                Err(poisoned) => {
                    warn!("MCP config lock was poisoned during reconfigure; recovering");
                    poisoned.into_inner()
                }
            };
            *guard = config.clone();
        }

        let mut status_guard = self.status.write().await;
        *status_guard = if config.enabled {
            McpInitStatus::Initializing {
                progress: "MCP runtime reconfigured; activation pending".to_string(),
            }
        } else {
            McpInitStatus::Disabled
        };
        Ok(())
    }

    pub(crate) async fn reconfigure_active_runtime(&self, config: McpClientConfig) -> Result<bool> {
        let should_restart =
            self.has_running_initialization_task() || matches!(self.get_status().await, McpInitStatus::Ready { .. });

        self.reconfigure(config).await?;

        if should_restart {
            self.start_initialization()?;
        }

        Ok(should_restart)
    }

    fn has_running_initialization_task(&self) -> bool {
        if let Ok(guard) = self.init_task.lock()
            && let Some(task) = guard.as_ref()
        {
            !task.is_finished()
        } else {
            false
        }
    }

    async fn initialize_mcp_client(
        config: McpClientConfig,
        hitl_notification_bell: bool,
        approval_policy: Arc<StdRwLock<AskForApproval>>,
        event_callback: Arc<dyn Fn(McpEvent) + Send + Sync>,
    ) -> Result<McpClient> {
        info!("Initializing MCP client with {} providers", config.providers.len());

        // Validate configuration before initializing
        if let Err(e) = vtcode_core::mcp::validate_mcp_config(&config) {
            warn!("MCP configuration validation error: {e}");
        }

        // Get startup timeout from config, default to 30 seconds
        let startup_timeout_secs = config.startup_timeout_seconds.unwrap_or(30);
        let startup_timeout = Duration::from_secs(startup_timeout_secs);

        let mut client = McpClient::new(config);

        // Set up elicitation handler
        use crate::agent::runloop::mcp_elicitation::InteractiveMcpElicitationHandler;
        client.set_elicitation_handler(Arc::new(InteractiveMcpElicitationHandler::new(
            hitl_notification_bell,
            approval_policy,
        )));

        // Initialize with timeout
        match timeout(startup_timeout, client.initialize()).await {
            Ok(Ok(())) => {
                info!("MCP client initialized successfully");

                // Add event reporting for tools
                if let Ok(mcp_tools) = client.list_tools().await {
                    info!("Found {} MCP tools", mcp_tools.len());
                    for mcp_tool in mcp_tools {
                        let mut event = McpEvent::new(mcp_tool.provider.clone(), mcp_tool.name.clone(), None);
                        event.success(None);
                        (*event_callback)(event);
                    }
                } else {
                    warn!("Failed to discover MCP tools after initialization");
                }

                Ok(client)
            }
            Ok(Err(e)) => Err(e).context("MCP client initialization failed"),
            Err(_) => Err(anyhow::anyhow!("MCP client initialization timed out after {startup_timeout_secs} seconds")),
        }
    }

    /// Get current status
    pub(crate) async fn get_status(&self) -> McpInitStatus {
        self.status.read().await.clone()
    }

    pub(crate) fn approval_policy(&self) -> AskForApproval {
        match self.approval_policy.read() {
            Ok(policy) => *policy,
            Err(poisoned) => {
                warn!("MCP approval policy lock was poisoned; continuing with last known value");
                *poisoned.into_inner()
            }
        }
    }

    pub(crate) fn set_approval_policy(&self, approval_policy: AskForApproval) {
        let mut policy_guard = match self.approval_policy.write() {
            Ok(policy) => policy,
            Err(poisoned) => {
                warn!("MCP approval policy lock was poisoned during update; recovering");
                poisoned.into_inner()
            }
        };
        *policy_guard = approval_policy;
    }

    /// Shutdown MCP client if active
    pub(crate) async fn shutdown(&self) -> Result<()> {
        match self.get_status().await {
            McpInitStatus::Ready { client } => {
                if let Err(e) = client.shutdown().await {
                    let error_msg = e.to_string();
                    if error_msg.contains("EPIPE")
                        || error_msg.contains("Broken pipe")
                        || error_msg.contains("write EPIPE")
                    {
                        info!("MCP client shutdown encountered pipe errors (normal): {}", e);
                    } else {
                        warn!("Failed to shutdown MCP client cleanly: {}", e);
                    }
                }
                Ok(())
            }
            McpInitStatus::Disabled => {
                info!("MCP is disabled, no shutdown needed");
                Ok(())
            }
            _ => {
                info!("MCP not initialized, no shutdown needed");
                Ok(())
            }
        }
    }
}

impl Drop for AsyncMcpManager {
    fn drop(&mut self) {
        // Abort the background init task so it doesn't outlive the manager.
        if let Ok(mut guard) = self.init_task.lock()
            && let Some(task) = guard.take()
        {
            task.abort();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use vtcode_core::config::mcp::McpClientConfig;
    use vtcode_core::exec_policy::RejectConfig;

    #[tokio::test]
    async fn test_async_mcp_manager_creation() {
        let config = McpClientConfig::default();
        let event_callback: Arc<dyn Fn(McpEvent) + Send + Sync> = Arc::new(|_event| {});

        let manager = AsyncMcpManager::new(config, true, AskForApproval::OnRequest, event_callback);
        let status = manager.get_status().await;

        // With default config, MCP should be disabled
        assert!(matches!(status, McpInitStatus::Disabled));
    }

    #[tokio::test]
    async fn test_mcp_init_status_display() {
        let disabled_status = McpInitStatus::Disabled;
        assert_eq!(disabled_status.to_string(), "MCP is disabled");

        let initializing_status = McpInitStatus::Initializing { progress: "Connecting...".to_string() };
        assert_eq!(initializing_status.to_string(), "MCP initializing: Connecting...");

        let error_status = McpInitStatus::Error { message: "Connection failed".to_string() };
        assert_eq!(error_status.to_string(), "MCP error: Connection failed");
    }

    #[tokio::test]
    async fn test_start_initialization_skips_when_task_already_running() {
        let config = McpClientConfig { enabled: true, ..McpClientConfig::default() };
        let event_callback: Arc<dyn Fn(McpEvent) + Send + Sync> = Arc::new(|_event| {});
        let manager = AsyncMcpManager::new(config, true, AskForApproval::OnRequest, event_callback);

        let blocker = tokio::spawn(async {
            tokio::time::sleep(Duration::from_secs(60)).await;
        });
        let blocker_id = blocker.id();
        if let Ok(mut guard) = manager.init_task.lock() {
            *guard = Some(blocker);
        }

        manager.start_initialization().expect("start should succeed");

        if let Ok(mut guard) = manager.init_task.lock() {
            let task = guard.as_ref().expect("running init task should still be present");
            assert_eq!(task.id(), blocker_id, "start_initialization should not replace a running init task");
            assert!(!task.is_finished(), "existing task should still be running after skipped start");
            if let Some(task) = guard.take() {
                task.abort();
            }
        } else {
            panic!("failed to lock init_task");
        }
    }

    #[tokio::test]
    async fn reconfigure_active_runtime_restarts_running_initialization() {
        let config = McpClientConfig { enabled: true, ..McpClientConfig::default() };
        let event_callback: Arc<dyn Fn(McpEvent) + Send + Sync> = Arc::new(|_event| {});
        let manager = AsyncMcpManager::new(config.clone(), true, AskForApproval::OnRequest, event_callback);
        let initialization_guard = manager.initialization_mutex.lock().await;

        let blocker = tokio::spawn(async {
            tokio::time::sleep(Duration::from_secs(60)).await;
        });
        let blocker_id = blocker.id();
        if let Ok(mut guard) = manager.init_task.lock() {
            *guard = Some(blocker);
        }

        let restarted = manager
            .reconfigure_active_runtime(config)
            .await
            .expect("reconfigure should succeed");

        assert!(restarted);
        if let Ok(mut guard) = manager.init_task.lock() {
            let task = guard.as_ref().expect("active initialization should be restarted");
            assert_ne!(task.id(), blocker_id);
            assert!(!task.is_finished());
            if let Some(task) = guard.take() {
                task.abort();
            }
        } else {
            panic!("failed to lock init_task");
        }
        drop(initialization_guard);
    }

    #[tokio::test]
    async fn reconfigure_active_runtime_restarts_ready_client_without_registry_attachment() {
        let config = McpClientConfig { enabled: true, ..McpClientConfig::default() };
        let event_callback: Arc<dyn Fn(McpEvent) + Send + Sync> = Arc::new(|_event| {});
        let manager = AsyncMcpManager::new(config.clone(), true, AskForApproval::OnRequest, event_callback);
        *manager.status.write().await = McpInitStatus::Ready { client: Arc::new(McpClient::new(config.clone())) };
        let initialization_guard = manager.initialization_mutex.lock().await;

        let restarted = manager
            .reconfigure_active_runtime(config)
            .await
            .expect("reconfigure should succeed");

        assert!(restarted);
        if let Ok(mut guard) = manager.init_task.lock() {
            let task = guard
                .as_ref()
                .expect("ready client should restart initialization after reconfigure");
            assert!(!task.is_finished());
            if let Some(task) = guard.take() {
                task.abort();
            }
        } else {
            panic!("failed to lock init_task");
        }
        drop(initialization_guard);
    }

    #[tokio::test]
    async fn reconfigure_active_runtime_preserves_manual_activation_pending() {
        let config = McpClientConfig { enabled: true, ..McpClientConfig::default() };
        let event_callback: Arc<dyn Fn(McpEvent) + Send + Sync> = Arc::new(|_event| {});
        let manager = AsyncMcpManager::new(config.clone(), true, AskForApproval::OnRequest, event_callback);

        let restarted = manager
            .reconfigure_active_runtime(config)
            .await
            .expect("reconfigure should succeed");

        assert!(!restarted);
        assert!(manager.init_task.lock().is_ok_and(|guard| guard.is_none()));
        assert!(manager.get_status().await.is_initializing());
    }

    #[tokio::test]
    async fn test_mcp_init_status_helpers() {
        let disabled_status = McpInitStatus::Disabled;
        assert!(!disabled_status.is_ready());
        assert!(!disabled_status.is_error());
        assert!(disabled_status.get_error_message().is_none());

        let error_status = McpInitStatus::Error { message: "Test error".to_string() };
        assert!(!error_status.is_ready());
        assert!(error_status.is_error());
        assert_eq!(error_status.get_error_message(), Some("Test error"));

        let initializing_status = McpInitStatus::Initializing { progress: "Init...".to_string() };
        assert!(initializing_status.is_initializing());
        assert!(!initializing_status.is_ready());
    }

    #[test]
    fn test_approval_policy_mapping_from_hitl() {
        assert_eq!(approval_policy_from_human_in_the_loop(true), AskForApproval::OnRequest);
        assert_eq!(
            approval_policy_from_human_in_the_loop(false),
            AskForApproval::Reject(RejectConfig {
                sandbox_approval: true,
                rules: true,
                request_permissions: true,
                mcp_elicitations: true,
            })
        );
    }

    #[test]
    fn test_set_approval_policy_updates_policy() {
        let config = McpClientConfig::default();
        let event_callback: Arc<dyn Fn(McpEvent) + Send + Sync> = Arc::new(|_event| {});

        let manager = AsyncMcpManager::new(config, true, AskForApproval::OnRequest, event_callback);
        assert_eq!(manager.approval_policy(), AskForApproval::OnRequest);

        manager.set_approval_policy(AskForApproval::Never);
        assert_eq!(manager.approval_policy(), AskForApproval::Never);
    }
}