aster/mcp/

lifecycle_manager.rs

//! MCP Lifecycle Manager
//!
//! This module implements the lifecycle manager for MCP servers.
//! It manages server processes including starting, stopping, restarting,
//! health monitoring, and dependency management.
//!
//! # Features
//!
//! - Start/stop server processes with configurable timeouts
//! - Automatic restart with exponential backoff on unexpected exits
//! - Health check monitoring
//! - Dependency-based startup ordering
//! - stdout/stderr capture and event emission
//!
//! # Requirements Coverage
//!
//! - 3.1: Server startup with command and arguments
//! - 3.2: Automatic restart with backoff on unexpected exit
//! - 3.3: Server state tracking (stopped, starting, running, stopping, error, crashed)
//! - 3.4: Maximum restart attempts before marking as crashed
//! - 3.5: Graceful shutdown with configurable timeout
//! - 3.6: Force kill after graceful shutdown timeout
//! - 3.7: stdout/stderr capture and event emission
//! - 3.8: Dependency-based startup ordering

use async_trait::async_trait;
use chrono::Utc;
use std::collections::HashMap;
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::{Child, Command};
use tokio::sync::{mpsc, Mutex, RwLock};

use crate::mcp::error::{McpError, McpResult};
use crate::mcp::types::{
    HealthCheckResult, LifecycleOptions, McpServerConfig, ServerProcess, ServerState, TransportType,
};

/// Lifecycle event for monitoring server state changes
#[derive(Debug, Clone)]
pub enum LifecycleEvent {
    /// Server is starting
    Starting { server_name: String },
    /// Server started successfully
    Started {
        server_name: String,
        pid: Option<u32>,
    },
    /// Server is stopping
    Stopping {
        server_name: String,
        reason: Option<String>,
    },
    /// Server stopped
    Stopped { server_name: String },
    /// Server error occurred
    Error { server_name: String, error: String },
    /// Server crashed
    Crashed {
        server_name: String,
        exit_code: Option<i32>,
    },
    /// Server is restarting
    Restarting { server_name: String },
    /// Health check passed
    HealthOk {
        server_name: String,
        result: HealthCheckResult,
    },
    /// Health check failed
    HealthFailed {
        server_name: String,
        result: HealthCheckResult,
    },
    /// stdout output from server
    Stdout { server_name: String, data: String },
    /// stderr output from server
    Stderr { server_name: String, data: String },
}

/// Start options for server startup
#[derive(Debug, Clone, Default)]
pub struct StartOptions {
    /// Force start even if already running
    pub force: bool,
    /// Wait for server to be ready
    pub wait_for_ready: bool,
    /// Dependencies to start first
    pub dependencies: Vec<String>,
}

/// Stop options for server shutdown
#[derive(Debug, Clone, Default)]
pub struct StopOptions {
    /// Force stop (skip graceful shutdown)
    pub force: bool,
    /// Reason for stopping
    pub reason: Option<String>,
}

/// Internal server state with process handle
pub(crate) struct ManagedServer {
    /// Server process info
    process: ServerProcess,
    /// Server configuration
    config: McpServerConfig,
    /// Child process handle (if running)
    child: Option<Child>,
    /// Dependencies (server names)
    dependencies: Vec<String>,
    /// Output capture task handles
    output_handles: Vec<tokio::task::JoinHandle<()>>,
    /// Health check task handle
    health_check_handle: Option<tokio::task::JoinHandle<()>>,
    /// Auto-restart task handle
    restart_handle: Option<tokio::task::JoinHandle<()>>,
}

impl ManagedServer {
    fn new(name: String, config: McpServerConfig) -> Self {
        Self {
            process: ServerProcess::new(name),
            config,
            child: None,
            dependencies: Vec::new(),
            output_handles: Vec::new(),
            health_check_handle: None,
            restart_handle: None,
        }
    }
}

/// Lifecycle manager trait
///
/// Defines the interface for managing MCP server lifecycles.
#[async_trait]
pub trait LifecycleManager: Send + Sync {
    /// Register a server configuration
    fn register_server(&self, name: &str, config: McpServerConfig);

    /// Unregister a server
    async fn unregister_server(&self, name: &str) -> McpResult<()>;

    /// Set server dependencies
    fn set_dependencies(&self, name: &str, dependencies: Vec<String>);

    /// Start a server
    async fn start(&self, server_name: &str, options: Option<StartOptions>) -> McpResult<()>;

    /// Start all registered servers
    async fn start_all(&self) -> McpResult<()>;

    /// Start a server with its dependencies
    async fn start_with_dependencies(&self, server_name: &str) -> McpResult<()>;

    /// Stop a server
    async fn stop(&self, server_name: &str, options: Option<StopOptions>) -> McpResult<()>;

    /// Stop all servers
    async fn stop_all(&self, force: bool) -> McpResult<()>;

    /// Restart a server
    async fn restart(&self, server_name: &str) -> McpResult<()>;

    /// Restart all servers
    async fn restart_all(&self) -> McpResult<()>;

    /// Perform health check on a server
    async fn health_check(&self, server_name: &str) -> HealthCheckResult;

    /// Perform health check on all servers
    async fn health_check_all(&self) -> HashMap<String, HealthCheckResult>;

    /// Get server state
    fn get_state(&self, server_name: &str) -> ServerState;

    /// Get server process info
    fn get_process(&self, server_name: &str) -> Option<ServerProcess>;

    /// Get all server processes
    fn get_all_processes(&self) -> Vec<ServerProcess>;

    /// Check if a server is running
    fn is_running(&self, server_name: &str) -> bool;

    /// Get list of running servers
    fn get_running_servers(&self) -> Vec<String>;

    /// Subscribe to lifecycle events
    fn subscribe(&self) -> mpsc::Receiver<LifecycleEvent>;

    /// Cleanup all resources
    async fn cleanup(&self) -> McpResult<()>;
}

/// Default implementation of the lifecycle manager
pub struct McpLifecycleManager {
    /// Managed servers
    pub(crate) servers: Arc<RwLock<HashMap<String, ManagedServer>>>,
    /// Default lifecycle options
    pub options: LifecycleOptions,
    /// Event channel sender
    event_tx: Arc<Mutex<Option<mpsc::Sender<LifecycleEvent>>>>,
    /// Enable auto-restart
    enable_auto_restart: bool,
    /// Enable health checks
    enable_health_checks: bool,
}

impl McpLifecycleManager {
    /// Create a new lifecycle manager with default options
    pub fn new() -> Self {
        Self::with_options(LifecycleOptions::default())
    }

    /// Create a new lifecycle manager with custom options
    pub fn with_options(options: LifecycleOptions) -> Self {
        Self {
            servers: Arc::new(RwLock::new(HashMap::new())),
            options,
            event_tx: Arc::new(Mutex::new(None)),
            enable_auto_restart: true,
            enable_health_checks: true,
        }
    }

    /// Enable or disable auto-restart
    pub fn set_auto_restart_enabled(&mut self, enabled: bool) {
        self.enable_auto_restart = enabled;
    }

    /// Enable or disable health checks
    pub fn set_health_checks_enabled(&mut self, enabled: bool) {
        self.enable_health_checks = enabled;
    }

    /// Emit a lifecycle event
    async fn emit_event(&self, event: LifecycleEvent) {
        if let Some(tx) = self.event_tx.lock().await.as_ref() {
            let _ = tx.send(event).await;
        }
    }

    /// Calculate restart delay with exponential backoff
    pub fn calculate_restart_delay(&self, attempt: u32) -> Duration {
        let base = self.options.restart_delay.as_millis() as u64;
        // Exponential backoff: base * 2^attempt, capped at 60 seconds
        let max_delay_ms = 60_000u64;
        let delay_ms = base.saturating_mul(1u64 << attempt.min(10));
        Duration::from_millis(delay_ms.min(max_delay_ms))
    }

    /// Start output capture for a child process
    fn start_output_capture(
        &self,
        server_name: String,
        child: &mut Child,
    ) -> Vec<tokio::task::JoinHandle<()>> {
        let mut handles = Vec::new();
        let event_tx = self.event_tx.clone();

        // Capture stdout
        if let Some(stdout) = child.stdout.take() {
            let name = server_name.clone();
            let tx = event_tx.clone();
            let handle = tokio::spawn(async move {
                let reader = BufReader::new(stdout);
                let mut lines = reader.lines();
                while let Ok(Some(line)) = lines.next_line().await {
                    if let Some(sender) = tx.lock().await.as_ref() {
                        let _ = sender
                            .send(LifecycleEvent::Stdout {
                                server_name: name.clone(),
                                data: line,
                            })
                            .await;
                    }
                }
            });
            handles.push(handle);
        }

        // Capture stderr
        if let Some(stderr) = child.stderr.take() {
            let name = server_name;
            let tx = event_tx;
            let handle = tokio::spawn(async move {
                let reader = BufReader::new(stderr);
                let mut lines = reader.lines();
                while let Ok(Some(line)) = lines.next_line().await {
                    if let Some(sender) = tx.lock().await.as_ref() {
                        let _ = sender
                            .send(LifecycleEvent::Stderr {
                                server_name: name.clone(),
                                data: line,
                            })
                            .await;
                    }
                }
            });
            handles.push(handle);
        }

        handles
    }

    /// Start monitoring a process for unexpected exit
    fn start_exit_monitor(&self, server_name: String) {
        if !self.enable_auto_restart {
            return;
        }

        let servers = self.servers.clone();
        let event_tx = self.event_tx.clone();
        let options = self.options.clone();
        let enable_auto_restart = self.enable_auto_restart;

        tokio::spawn(async move {
            loop {
                // Check if process is still running
                let should_restart = {
                    let mut servers_guard = servers.write().await;
                    if let Some(server) = servers_guard.get_mut(&server_name) {
                        if let Some(ref mut child) = server.child {
                            match child.try_wait() {
                                Ok(Some(status)) => {
                                    // Process exited
                                    let exit_code = status.code();
                                    server.process.state = ServerState::Crashed;
                                    server.process.stopped_at = Some(Utc::now());
                                    server.process.consecutive_failures += 1;
                                    server.child = None;

                                    // Emit crashed event
                                    if let Some(tx) = event_tx.lock().await.as_ref() {
                                        let _ = tx
                                            .send(LifecycleEvent::Crashed {
                                                server_name: server_name.clone(),
                                                exit_code,
                                            })
                                            .await;
                                    }

                                    // Check if we should restart
                                    if enable_auto_restart
                                        && server.process.restart_count < options.max_restarts
                                    {
                                        true
                                    } else {
                                        server.process.state = ServerState::Crashed;
                                        server.process.last_error = Some(format!(
                                            "Process exited with code {:?}, max restarts exceeded",
                                            exit_code
                                        ));
                                        false
                                    }
                                }
                                Ok(None) => {
                                    // Process still running
                                    false
                                }
                                Err(e) => {
                                    // Error checking process
                                    server.process.last_error = Some(e.to_string());
                                    false
                                }
                            }
                        } else {
                            // No child process, stop monitoring
                            break;
                        }
                    } else {
                        // Server not found, stop monitoring
                        break;
                    }
                };

                if should_restart {
                    // Calculate delay based on restart count
                    let restart_count = {
                        let servers_guard = servers.read().await;
                        servers_guard
                            .get(&server_name)
                            .map(|s| s.process.restart_count)
                            .unwrap_or(0)
                    };

                    let base = options.restart_delay.as_millis() as u64;
                    let delay_ms = base.saturating_mul(1u64 << restart_count.min(10));
                    let delay = Duration::from_millis(delay_ms.min(60_000));

                    // Emit restarting event
                    if let Some(tx) = event_tx.lock().await.as_ref() {
                        let _ = tx
                            .send(LifecycleEvent::Restarting {
                                server_name: server_name.clone(),
                            })
                            .await;
                    }

                    tokio::time::sleep(delay).await;

                    // Attempt restart
                    let mut servers_guard = servers.write().await;
                    if let Some(server) = servers_guard.get_mut(&server_name) {
                        server.process.restart_count += 1;
                        // The actual restart will be handled by the start method
                        // For now, just update state
                        server.process.state = ServerState::Starting;
                    }
                }

                // Sleep before next check
                tokio::time::sleep(Duration::from_millis(500)).await;
            }
        });
    }

    /// Start health check monitoring for a server
    fn start_health_check_monitor(&self, server_name: String) -> tokio::task::JoinHandle<()> {
        let servers = self.servers.clone();
        let event_tx = self.event_tx.clone();
        let interval = self.options.health_check_interval;

        tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);

            loop {
                interval_timer.tick().await;

                let is_running = {
                    let servers_guard = servers.read().await;
                    servers_guard
                        .get(&server_name)
                        .map(|s| s.process.state == ServerState::Running)
                        .unwrap_or(false)
                };

                if !is_running {
                    break;
                }

                // Perform health check (check if process is still alive)
                let start = std::time::Instant::now();
                let result = {
                    let mut servers_guard = servers.write().await;
                    if let Some(server) = servers_guard.get_mut(&server_name) {
                        if let Some(ref mut child) = server.child {
                            match child.try_wait() {
                                Ok(None) => {
                                    // Process is still running
                                    HealthCheckResult {
                                        healthy: true,
                                        latency: Some(start.elapsed()),
                                        last_check: Utc::now(),
                                        error: None,
                                    }
                                }
                                Ok(Some(_)) => {
                                    // Process has exited
                                    HealthCheckResult {
                                        healthy: false,
                                        latency: Some(start.elapsed()),
                                        last_check: Utc::now(),
                                        error: Some("Process has exited".to_string()),
                                    }
                                }
                                Err(e) => HealthCheckResult {
                                    healthy: false,
                                    latency: Some(start.elapsed()),
                                    last_check: Utc::now(),
                                    error: Some(e.to_string()),
                                },
                            }
                        } else {
                            HealthCheckResult {
                                healthy: false,
                                latency: None,
                                last_check: Utc::now(),
                                error: Some("No child process".to_string()),
                            }
                        }
                    } else {
                        break;
                    }
                };

                // Emit health event
                if let Some(tx) = event_tx.lock().await.as_ref() {
                    let event = if result.healthy {
                        LifecycleEvent::HealthOk {
                            server_name: server_name.clone(),
                            result,
                        }
                    } else {
                        LifecycleEvent::HealthFailed {
                            server_name: server_name.clone(),
                            result,
                        }
                    };
                    let _ = tx.send(event).await;
                }
            }
        })
    }

    /// Get topologically sorted server names based on dependencies
    pub(crate) fn topological_sort(&self, servers: &HashMap<String, ManagedServer>) -> Vec<String> {
        let mut result = Vec::new();
        let mut visited = std::collections::HashSet::new();
        let mut temp_visited = std::collections::HashSet::new();

        fn visit(
            name: &str,
            servers: &HashMap<String, ManagedServer>,
            visited: &mut std::collections::HashSet<String>,
            temp_visited: &mut std::collections::HashSet<String>,
            result: &mut Vec<String>,
        ) {
            if visited.contains(name) {
                return;
            }
            if temp_visited.contains(name) {
                // Cycle detected, skip
                return;
            }

            temp_visited.insert(name.to_string());

            if let Some(server) = servers.get(name) {
                for dep in &server.dependencies {
                    visit(dep, servers, visited, temp_visited, result);
                }
            }

            temp_visited.remove(name);
            visited.insert(name.to_string());
            result.push(name.to_string());
        }

        for name in servers.keys() {
            visit(name, servers, &mut visited, &mut temp_visited, &mut result);
        }

        result
    }
}

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

#[async_trait]
impl LifecycleManager for McpLifecycleManager {
    fn register_server(&self, name: &str, config: McpServerConfig) {
        let servers = self.servers.clone();
        let name = name.to_string();
        tokio::spawn(async move {
            let mut servers_guard = servers.write().await;
            servers_guard.insert(name.clone(), ManagedServer::new(name, config));
        });
    }

    async fn unregister_server(&self, name: &str) -> McpResult<()> {
        // Stop the server first if running
        if self.is_running(name) {
            self.stop(
                name,
                Some(StopOptions {
                    force: true,
                    reason: Some("Unregistering server".to_string()),
                }),
            )
            .await?;
        }

        let mut servers = self.servers.write().await;
        servers.remove(name);
        Ok(())
    }

    fn set_dependencies(&self, name: &str, dependencies: Vec<String>) {
        let servers = self.servers.clone();
        let name = name.to_string();
        tokio::spawn(async move {
            let mut servers_guard = servers.write().await;
            if let Some(server) = servers_guard.get_mut(&name) {
                server.dependencies = dependencies;
            }
        });
    }

    async fn start(&self, server_name: &str, options: Option<StartOptions>) -> McpResult<()> {
        let options = options.unwrap_or_default();

        // Check if server is registered
        let config = {
            let servers = self.servers.read().await;
            let server = servers.get(server_name).ok_or_else(|| {
                McpError::lifecycle(
                    format!("Server not registered: {}", server_name),
                    Some(server_name.to_string()),
                )
            })?;

            // Check if already running
            if !options.force && server.process.state == ServerState::Running {
                return Ok(());
            }

            // Only stdio servers can be started as processes
            if server.config.transport_type != TransportType::Stdio {
                return Err(McpError::lifecycle(
                    format!(
                        "Only stdio servers can be started as processes, got {:?}",
                        server.config.transport_type
                    ),
                    Some(server_name.to_string()),
                ));
            }

            server.config.clone()
        };

        // Emit starting event
        self.emit_event(LifecycleEvent::Starting {
            server_name: server_name.to_string(),
        })
        .await;

        // Update state to starting
        {
            let mut servers = self.servers.write().await;
            if let Some(server) = servers.get_mut(server_name) {
                server.process.state = ServerState::Starting;
            }
        }

        // Get command and args
        let command = config.command.ok_or_else(|| {
            McpError::lifecycle(
                "Stdio server requires a command".to_string(),
                Some(server_name.to_string()),
            )
        })?;

        let args = config.args.unwrap_or_default();
        let env = config.env.unwrap_or_default();

        // Build command
        let mut cmd = Command::new(&command);
        cmd.args(&args)
            .envs(&env)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true);

        // Spawn process with timeout
        let startup_timeout = self.options.startup_timeout;
        let spawn_result = tokio::time::timeout(startup_timeout, async { cmd.spawn() }).await;

        let mut child = match spawn_result {
            Ok(Ok(child)) => child,
            Ok(Err(e)) => {
                // Update state to error
                {
                    let mut servers = self.servers.write().await;
                    if let Some(server) = servers.get_mut(server_name) {
                        server.process.state = ServerState::Error;
                        server.process.last_error = Some(e.to_string());
                        server.process.consecutive_failures += 1;
                    }
                }

                self.emit_event(LifecycleEvent::Error {
                    server_name: server_name.to_string(),
                    error: e.to_string(),
                })
                .await;

                return Err(McpError::lifecycle(
                    format!("Failed to spawn process: {}", e),
                    Some(server_name.to_string()),
                ));
            }
            Err(_) => {
                // Timeout
                {
                    let mut servers = self.servers.write().await;
                    if let Some(server) = servers.get_mut(server_name) {
                        server.process.state = ServerState::Error;
                        server.process.last_error = Some("Startup timeout".to_string());
                        server.process.consecutive_failures += 1;
                    }
                }

                self.emit_event(LifecycleEvent::Error {
                    server_name: server_name.to_string(),
                    error: "Startup timeout".to_string(),
                })
                .await;

                return Err(McpError::lifecycle(
                    format!("Startup timeout after {:?}", startup_timeout),
                    Some(server_name.to_string()),
                ));
            }
        };

        // Get PID
        let pid = child.id();

        // Start output capture
        let output_handles = self.start_output_capture(server_name.to_string(), &mut child);

        // Update server state
        {
            let mut servers = self.servers.write().await;
            if let Some(server) = servers.get_mut(server_name) {
                server.process.state = ServerState::Running;
                server.process.pid = pid;
                server.process.started_at = Some(Utc::now());
                server.process.stopped_at = None;
                server.process.consecutive_failures = 0;
                server.child = Some(child);
                server.output_handles = output_handles;
            }
        }

        // Start exit monitor for auto-restart
        self.start_exit_monitor(server_name.to_string());

        // Start health check monitor if enabled
        if self.enable_health_checks {
            let handle = self.start_health_check_monitor(server_name.to_string());
            let mut servers = self.servers.write().await;
            if let Some(server) = servers.get_mut(server_name) {
                server.health_check_handle = Some(handle);
            }
        }

        // Emit started event
        self.emit_event(LifecycleEvent::Started {
            server_name: server_name.to_string(),
            pid,
        })
        .await;

        Ok(())
    }

    async fn start_all(&self) -> McpResult<()> {
        let server_names: Vec<String> = {
            let servers = self.servers.read().await;
            self.topological_sort(&servers)
        };

        for name in server_names {
            if let Err(e) = self.start(&name, None).await {
                tracing::warn!("Failed to start server {}: {}", name, e);
            }
        }

        Ok(())
    }

    async fn start_with_dependencies(&self, server_name: &str) -> McpResult<()> {
        // Get dependencies
        let dependencies = {
            let servers = self.servers.read().await;
            servers
                .get(server_name)
                .map(|s| s.dependencies.clone())
                .unwrap_or_default()
        };

        // Start dependencies first (recursively)
        for dep in dependencies {
            self.start_with_dependencies(&dep).await?;
        }

        // Start this server
        self.start(server_name, None).await
    }

    async fn stop(&self, server_name: &str, options: Option<StopOptions>) -> McpResult<()> {
        let options = options.unwrap_or_default();

        // Check if server exists and is running
        let child_exists = {
            let servers = self.servers.read().await;
            servers
                .get(server_name)
                .map(|s| s.child.is_some())
                .unwrap_or(false)
        };

        if !child_exists {
            return Ok(());
        }

        // Emit stopping event
        self.emit_event(LifecycleEvent::Stopping {
            server_name: server_name.to_string(),
            reason: options.reason.clone(),
        })
        .await;

        // Update state to stopping
        {
            let mut servers = self.servers.write().await;
            if let Some(server) = servers.get_mut(server_name) {
                server.process.state = ServerState::Stopping;
            }
        }

        // Get child process
        let mut child = {
            let mut servers = self.servers.write().await;
            servers.get_mut(server_name).and_then(|s| s.child.take())
        };

        if let Some(ref mut child) = child {
            if options.force {
                // Force kill immediately
                let _ = child.kill().await;
            } else {
                // Try graceful shutdown first
                let shutdown_timeout = self.options.shutdown_timeout;

                // On Unix, we try to send SIGTERM first via the child's kill method
                // which sends SIGKILL. For graceful shutdown, we just wait with timeout.
                // The process should handle its own graceful shutdown.

                // Wait for graceful shutdown with timeout
                let wait_result = tokio::time::timeout(shutdown_timeout, child.wait()).await;

                match wait_result {
                    Ok(Ok(_)) => {
                        // Process exited gracefully
                    }
                    Ok(Err(e)) => {
                        tracing::warn!("Error waiting for process: {}", e);
                    }
                    Err(_) => {
                        // Timeout - force kill
                        tracing::warn!(
                            "Graceful shutdown timeout for {}, force killing",
                            server_name
                        );
                        let _ = child.kill().await;
                    }
                }
            }
        }

        // Cancel output capture handles
        {
            let mut servers = self.servers.write().await;
            if let Some(server) = servers.get_mut(server_name) {
                for handle in server.output_handles.drain(..) {
                    handle.abort();
                }
                if let Some(handle) = server.health_check_handle.take() {
                    handle.abort();
                }
                if let Some(handle) = server.restart_handle.take() {
                    handle.abort();
                }
            }
        }

        // Update state
        {
            let mut servers = self.servers.write().await;
            if let Some(server) = servers.get_mut(server_name) {
                server.process.state = ServerState::Stopped;
                server.process.pid = None;
                server.process.stopped_at = Some(Utc::now());
                server.child = None;
            }
        }

        // Emit stopped event
        self.emit_event(LifecycleEvent::Stopped {
            server_name: server_name.to_string(),
        })
        .await;

        Ok(())
    }

    async fn stop_all(&self, force: bool) -> McpResult<()> {
        let server_names: Vec<String> = {
            let servers = self.servers.read().await;
            // Stop in reverse dependency order
            let mut sorted = self.topological_sort(&servers);
            sorted.reverse();
            sorted
        };

        for name in server_names {
            let options = StopOptions {
                force,
                reason: Some("Stopping all servers".to_string()),
            };
            if let Err(e) = self.stop(&name, Some(options)).await {
                tracing::warn!("Failed to stop server {}: {}", name, e);
            }
        }

        Ok(())
    }

    async fn restart(&self, server_name: &str) -> McpResult<()> {
        self.stop(server_name, None).await?;
        self.start(server_name, None).await
    }

    async fn restart_all(&self) -> McpResult<()> {
        self.stop_all(false).await?;
        self.start_all().await
    }

    async fn health_check(&self, server_name: &str) -> HealthCheckResult {
        let start = std::time::Instant::now();

        let mut servers = self.servers.write().await;
        if let Some(server) = servers.get_mut(server_name) {
            if let Some(ref mut child) = server.child {
                match child.try_wait() {
                    Ok(None) => {
                        // Process is still running
                        HealthCheckResult {
                            healthy: true,
                            latency: Some(start.elapsed()),
                            last_check: Utc::now(),
                            error: None,
                        }
                    }
                    Ok(Some(status)) => {
                        // Process has exited
                        HealthCheckResult {
                            healthy: false,
                            latency: Some(start.elapsed()),
                            last_check: Utc::now(),
                            error: Some(format!("Process exited with status: {:?}", status)),
                        }
                    }
                    Err(e) => HealthCheckResult {
                        healthy: false,
                        latency: Some(start.elapsed()),
                        last_check: Utc::now(),
                        error: Some(e.to_string()),
                    },
                }
            } else {
                HealthCheckResult {
                    healthy: false,
                    latency: None,
                    last_check: Utc::now(),
                    error: Some("Server not running".to_string()),
                }
            }
        } else {
            HealthCheckResult {
                healthy: false,
                latency: None,
                last_check: Utc::now(),
                error: Some("Server not found".to_string()),
            }
        }
    }

    async fn health_check_all(&self) -> HashMap<String, HealthCheckResult> {
        let server_names: Vec<String> = {
            let servers = self.servers.read().await;
            servers.keys().cloned().collect()
        };

        let mut results = HashMap::new();
        for name in server_names {
            let result = self.health_check(&name).await;
            results.insert(name, result);
        }
        results
    }

    fn get_state(&self, server_name: &str) -> ServerState {
        self.servers
            .try_read()
            .ok()
            .and_then(|servers| servers.get(server_name).map(|s| s.process.state))
            .unwrap_or(ServerState::Stopped)
    }

    fn get_process(&self, server_name: &str) -> Option<ServerProcess> {
        self.servers
            .try_read()
            .ok()
            .and_then(|servers| servers.get(server_name).map(|s| s.process.clone()))
    }

    fn get_all_processes(&self) -> Vec<ServerProcess> {
        self.servers
            .try_read()
            .map(|servers| servers.values().map(|s| s.process.clone()).collect())
            .unwrap_or_default()
    }

    fn is_running(&self, server_name: &str) -> bool {
        self.get_state(server_name) == ServerState::Running
    }

    fn get_running_servers(&self) -> Vec<String> {
        self.servers
            .try_read()
            .map(|servers| {
                servers
                    .iter()
                    .filter(|(_, s)| s.process.state == ServerState::Running)
                    .map(|(name, _)| name.clone())
                    .collect()
            })
            .unwrap_or_default()
    }

    fn subscribe(&self) -> mpsc::Receiver<LifecycleEvent> {
        let (tx, rx) = mpsc::channel(100);
        let event_tx = self.event_tx.clone();
        tokio::spawn(async move {
            *event_tx.lock().await = Some(tx);
        });
        rx
    }

    async fn cleanup(&self) -> McpResult<()> {
        // Stop all servers
        self.stop_all(true).await?;

        // Clear all servers
        let mut servers = self.servers.write().await;
        servers.clear();

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    fn create_test_config() -> McpServerConfig {
        McpServerConfig {
            transport_type: TransportType::Stdio,
            command: Some("echo".to_string()),
            args: Some(vec!["hello".to_string()]),
            env: None,
            url: None,
            headers: None,
            enabled: true,
            timeout: Duration::from_secs(30),
            retries: 3,
            auto_approve: vec![],
            log_level: Default::default(),
        }
    }

    #[test]
    fn test_lifecycle_manager_new() {
        let manager = McpLifecycleManager::new();
        assert!(manager.get_all_processes().is_empty());
    }

    #[test]
    fn test_lifecycle_manager_with_options() {
        let options = LifecycleOptions {
            startup_timeout: Duration::from_secs(60),
            max_restarts: 5,
            ..Default::default()
        };
        let manager = McpLifecycleManager::with_options(options);
        assert_eq!(manager.options.startup_timeout, Duration::from_secs(60));
        assert_eq!(manager.options.max_restarts, 5);
    }

    #[test]
    fn test_calculate_restart_delay() {
        let manager = McpLifecycleManager::new();

        let delay0 = manager.calculate_restart_delay(0);
        let delay1 = manager.calculate_restart_delay(1);
        let delay2 = manager.calculate_restart_delay(2);

        // Each delay should be roughly double the previous
        assert!(delay1 > delay0);
        assert!(delay2 > delay1);

        // Should not exceed 60 seconds
        let delay_max = manager.calculate_restart_delay(100);
        assert!(delay_max <= Duration::from_secs(60));
    }

    #[test]
    fn test_server_state_default() {
        let process = ServerProcess::new("test".to_string());
        assert_eq!(process.state, ServerState::Stopped);
        assert_eq!(process.restart_count, 0);
        assert!(process.pid.is_none());
    }

    #[test]
    fn test_start_options_default() {
        let options = StartOptions::default();
        assert!(!options.force);
        assert!(!options.wait_for_ready);
        assert!(options.dependencies.is_empty());
    }

    #[test]
    fn test_stop_options_default() {
        let options = StopOptions::default();
        assert!(!options.force);
        assert!(options.reason.is_none());
    }

    #[test]
    fn test_lifecycle_options_default() {
        let options = LifecycleOptions::default();
        assert_eq!(options.startup_timeout, Duration::from_secs(30));
        assert_eq!(options.shutdown_timeout, Duration::from_secs(10));
        assert_eq!(options.max_restarts, 3);
    }

    #[tokio::test]
    async fn test_register_and_get_process() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        manager.register_server("test-server", config);

        // Wait for async registration
        tokio::time::sleep(Duration::from_millis(50)).await;

        let process = manager.get_process("test-server");
        assert!(process.is_some());
        assert_eq!(process.unwrap().name, "test-server");
    }

    #[tokio::test]
    async fn test_get_state_unregistered() {
        let manager = McpLifecycleManager::new();
        let state = manager.get_state("nonexistent");
        assert_eq!(state, ServerState::Stopped);
    }

    #[tokio::test]
    async fn test_is_running_not_started() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        manager.register_server("test-server", config);
        tokio::time::sleep(Duration::from_millis(50)).await;

        assert!(!manager.is_running("test-server"));
    }

    #[tokio::test]
    async fn test_get_running_servers_empty() {
        let manager = McpLifecycleManager::new();
        let running = manager.get_running_servers();
        assert!(running.is_empty());
    }

    #[tokio::test]
    async fn test_set_dependencies() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        manager.register_server("server-a", config.clone());
        manager.register_server("server-b", config);
        tokio::time::sleep(Duration::from_millis(50)).await;

        manager.set_dependencies("server-b", vec!["server-a".to_string()]);
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Verify dependencies are set
        let servers = manager.servers.read().await;
        let server_b = servers.get("server-b").unwrap();
        assert_eq!(server_b.dependencies, vec!["server-a".to_string()]);
    }

    #[tokio::test]
    async fn test_topological_sort() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        // Register servers
        manager.register_server("server-a", config.clone());
        manager.register_server("server-b", config.clone());
        manager.register_server("server-c", config);
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Set dependencies: c depends on b, b depends on a
        manager.set_dependencies("server-c", vec!["server-b".to_string()]);
        manager.set_dependencies("server-b", vec!["server-a".to_string()]);
        tokio::time::sleep(Duration::from_millis(50)).await;

        let servers = manager.servers.read().await;
        let sorted = manager.topological_sort(&servers);

        // a should come before b, b should come before c
        let pos_a = sorted.iter().position(|x| x == "server-a").unwrap();
        let pos_b = sorted.iter().position(|x| x == "server-b").unwrap();
        let pos_c = sorted.iter().position(|x| x == "server-c").unwrap();

        assert!(pos_a < pos_b);
        assert!(pos_b < pos_c);
    }

    #[tokio::test]
    async fn test_unregister_server() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        manager.register_server("test-server", config);
        tokio::time::sleep(Duration::from_millis(50)).await;

        assert!(manager.get_process("test-server").is_some());

        manager.unregister_server("test-server").await.unwrap();

        assert!(manager.get_process("test-server").is_none());
    }

    #[tokio::test]
    async fn test_cleanup() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        manager.register_server("server-1", config.clone());
        manager.register_server("server-2", config);
        tokio::time::sleep(Duration::from_millis(50)).await;

        assert_eq!(manager.get_all_processes().len(), 2);

        manager.cleanup().await.unwrap();

        assert!(manager.get_all_processes().is_empty());
    }

    #[tokio::test]
    async fn test_health_check_not_running() {
        let manager = McpLifecycleManager::new();
        let config = create_test_config();

        manager.register_server("test-server", config);
        tokio::time::sleep(Duration::from_millis(50)).await;

        let result = manager.health_check("test-server").await;
        assert!(!result.healthy);
        assert!(result.error.is_some());
    }

    #[tokio::test]
    async fn test_health_check_nonexistent() {
        let manager = McpLifecycleManager::new();

        let result = manager.health_check("nonexistent").await;
        assert!(!result.healthy);
        assert!(result.error.unwrap().contains("not found"));
    }

    #[tokio::test]
    async fn test_subscribe_events() {
        let manager = McpLifecycleManager::new();
        let _rx = manager.subscribe();

        // Just verify subscription works without panic
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
}