mcp_protocol_sdk/server/

lifecycle.rs

//! Server lifecycle management
//!
//! This module handles the lifecycle events and state management for MCP servers,
//! including initialization, running state, graceful shutdown, and error recovery.

use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{watch, Mutex, RwLock};
use tokio::time::timeout;

use crate::core::error::{McpError, McpResult};
use crate::server::mcp_server::McpServer;

/// Server lifecycle state
#[derive(Debug, Clone, PartialEq)]
pub enum LifecycleState {
    /// Server is created but not yet started
    Created,
    /// Server is starting up
    Starting,
    /// Server is running and ready to handle requests
    Running,
    /// Server is gracefully shutting down
    Stopping,
    /// Server has stopped
    Stopped,
    /// Server encountered an error
    Error(String),
}

/// Server lifecycle manager
#[derive(Clone)]
pub struct LifecycleManager {
    /// Current lifecycle state
    state: Arc<RwLock<LifecycleState>>,
    /// State change broadcaster
    state_tx: Arc<watch::Sender<LifecycleState>>,
    /// State change receiver
    state_rx: watch::Receiver<LifecycleState>,
    /// Server start time
    start_time: Arc<Mutex<Option<Instant>>>,
    /// Shutdown signal
    shutdown_tx: Arc<Mutex<Option<watch::Sender<()>>>>,
}

/// Lifecycle event listener
pub trait LifecycleListener: Send + Sync {
    /// Called when the server state changes
    fn on_state_change(&self, old_state: LifecycleState, new_state: LifecycleState);

    /// Called when the server starts
    fn on_start(&self) {}

    /// Called when the server stops
    fn on_stop(&self) {}

    /// Called when an error occurs
    fn on_error(&self, _error: &McpError) {}
}

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

impl LifecycleManager {
    /// Create a new lifecycle manager
    pub fn new() -> Self {
        let (state_tx, state_rx) = watch::channel(LifecycleState::Created);

        Self {
            state: Arc::new(RwLock::new(LifecycleState::Created)),
            state_tx: Arc::new(state_tx),
            state_rx,
            start_time: Arc::new(Mutex::new(None)),
            shutdown_tx: Arc::new(Mutex::new(None)),
        }
    }

    /// Get the current lifecycle state
    pub async fn state(&self) -> LifecycleState {
        let state = self.state.read().await;
        state.clone()
    }

    /// Subscribe to state changes
    pub fn subscribe(&self) -> watch::Receiver<LifecycleState> {
        self.state_rx.clone()
    }

    /// Transition to a new state
    pub async fn transition_to(&self, new_state: LifecycleState) -> McpResult<()> {
        let _old_state = {
            let mut state = self.state.write().await;
            let old = state.clone();
            *state = new_state.clone();
            old
        };

        // Broadcast the state change
        if self.state_tx.send(new_state.clone()).is_err() {
            // Receiver may have been dropped, which is okay
        }

        // Handle special state transitions
        match new_state {
            LifecycleState::Running => {
                let mut start_time = self.start_time.lock().await;
                *start_time = Some(Instant::now());
            }
            LifecycleState::Stopped => {
                let mut start_time = self.start_time.lock().await;
                *start_time = None;
            }
            _ => {}
        }

        Ok(())
    }

    /// Check if the server is in a running state
    pub async fn is_running(&self) -> bool {
        let state = self.state.read().await;
        matches!(*state, LifecycleState::Running)
    }

    /// Check if the server can be started
    pub async fn can_start(&self) -> bool {
        let state = self.state.read().await;
        matches!(*state, LifecycleState::Created | LifecycleState::Stopped)
    }

    /// Check if the server can be stopped
    pub async fn can_stop(&self) -> bool {
        let state = self.state.read().await;
        matches!(*state, LifecycleState::Running | LifecycleState::Starting)
    }

    /// Get server uptime
    pub async fn uptime(&self) -> Option<Duration> {
        let start_time = self.start_time.lock().await;
        start_time.map(|start| start.elapsed())
    }

    /// Create a shutdown signal
    pub async fn create_shutdown_signal(&self) -> watch::Receiver<()> {
        let (tx, rx) = watch::channel(());
        let mut shutdown_tx = self.shutdown_tx.lock().await;
        *shutdown_tx = Some(tx);
        rx
    }

    /// Trigger shutdown
    pub async fn trigger_shutdown(&self) -> McpResult<()> {
        let shutdown_tx = self.shutdown_tx.lock().await;
        if let Some(tx) = shutdown_tx.as_ref() {
            let _ = tx.send(()); // Ignore error if receiver is dropped
        }
        Ok(())
    }
}

/// Server runner that manages the complete lifecycle
pub struct ServerRunner {
    /// The MCP server instance
    server: Arc<Mutex<McpServer>>,
    /// Lifecycle manager
    lifecycle: LifecycleManager,
    /// Lifecycle listeners
    listeners: Arc<RwLock<Vec<Box<dyn LifecycleListener>>>>,
}

impl ServerRunner {
    /// Create a new server runner
    pub fn new(server: McpServer) -> Self {
        Self {
            server: Arc::new(Mutex::new(server)),
            lifecycle: LifecycleManager::new(),
            listeners: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Add a lifecycle listener
    pub async fn add_listener<L>(&self, listener: L)
    where
        L: LifecycleListener + 'static,
    {
        let mut listeners = self.listeners.write().await;
        listeners.push(Box::new(listener));
    }

    /// Get the lifecycle manager
    pub fn lifecycle(&self) -> &LifecycleManager {
        &self.lifecycle
    }

    /// Start the server with a transport
    pub async fn start<T>(&self, transport: T) -> McpResult<()>
    where
        T: crate::transport::traits::ServerTransport + 'static,
    {
        // Check if we can start
        if !self.lifecycle.can_start().await {
            return Err(McpError::Protocol(
                "Server cannot be started in current state".to_string(),
            ));
        }

        // Transition to starting state
        self.lifecycle
            .transition_to(LifecycleState::Starting)
            .await?;

        // Notify listeners
        self.notify_listeners(|listener| listener.on_start()).await;

        // Start the server
        let result = {
            let mut server = self.server.lock().await;
            server.start(transport).await
        };

        match result {
            Ok(()) => {
                // Transition to running state
                self.lifecycle
                    .transition_to(LifecycleState::Running)
                    .await?;
                Ok(())
            }
            Err(err) => {
                // Transition to error state
                let error_msg = err.to_string();
                self.lifecycle
                    .transition_to(LifecycleState::Error(error_msg.clone()))
                    .await?;

                // Notify listeners
                self.notify_listeners(|listener| listener.on_error(&err))
                    .await;

                Err(err)
            }
        }
    }

    /// Stop the server gracefully
    pub async fn stop(&self) -> McpResult<()> {
        // Check if we can stop
        if !self.lifecycle.can_stop().await {
            return Err(McpError::Protocol(
                "Server cannot be stopped in current state".to_string(),
            ));
        }

        // Transition to stopping state
        self.lifecycle
            .transition_to(LifecycleState::Stopping)
            .await?;

        // Stop the server
        let result = {
            let server = self.server.lock().await;
            server.stop().await
        };

        match result {
            Ok(()) => {
                // Transition to stopped state
                self.lifecycle
                    .transition_to(LifecycleState::Stopped)
                    .await?;

                // Notify listeners
                self.notify_listeners(|listener| listener.on_stop()).await;

                Ok(())
            }
            Err(err) => {
                // Transition to error state
                let error_msg = err.to_string();
                self.lifecycle
                    .transition_to(LifecycleState::Error(error_msg.clone()))
                    .await?;

                // Notify listeners
                self.notify_listeners(|listener| listener.on_error(&err))
                    .await;

                Err(err)
            }
        }
    }

    /// Stop the server with a timeout
    pub async fn stop_with_timeout(&self, shutdown_timeout: Duration) -> McpResult<()> {
        match timeout(shutdown_timeout, self.stop()).await {
            Ok(result) => result,
            Err(_) => {
                // Force stop if timeout exceeded
                self.lifecycle
                    .transition_to(LifecycleState::Error(
                        "Shutdown timeout exceeded".to_string(),
                    ))
                    .await?;
                Err(McpError::Protocol(
                    "Server shutdown timeout exceeded".to_string(),
                ))
            }
        }
    }

    /// Run the server until shutdown signal
    pub async fn run_until_shutdown<T>(&self, transport: T) -> McpResult<()>
    where
        T: crate::transport::traits::ServerTransport + 'static,
    {
        // Start the server
        self.start(transport).await?;

        // Wait for shutdown signal
        let mut shutdown_rx = self.lifecycle.create_shutdown_signal().await;
        let _ = shutdown_rx.changed().await;

        // Stop the server
        self.stop().await?;

        Ok(())
    }

    /// Run the server with graceful shutdown on CTRL+C
    pub async fn run_with_signals<T>(&self, transport: T) -> McpResult<()>
    where
        T: crate::transport::traits::ServerTransport + 'static,
    {
        // Start the server
        self.start(transport).await?;

        // Set up signal handling
        let lifecycle = self.lifecycle.clone();
        tokio::spawn(async move {
            tokio::signal::ctrl_c()
                .await
                .expect("Failed to listen for ctrl+c");
            let _ = lifecycle.trigger_shutdown().await;
        });

        // Wait for shutdown signal
        let mut shutdown_rx = self.lifecycle.create_shutdown_signal().await;
        let _ = shutdown_rx.changed().await;

        // Stop the server gracefully
        let config = {
            let server = self.server.lock().await;
            server.config().clone()
        };

        let shutdown_timeout = Duration::from_millis(config.request_timeout_ms * 2);
        self.stop_with_timeout(shutdown_timeout).await?;

        Ok(())
    }

    /// Get the server instance (for advanced usage)
    pub fn server(&self) -> Arc<Mutex<McpServer>> {
        self.server.clone()
    }

    /// Check if the server is running
    pub async fn is_running(&self) -> bool {
        self.lifecycle.is_running().await
    }

    /// Get server uptime
    pub async fn uptime(&self) -> Option<Duration> {
        self.lifecycle.uptime().await
    }

    /// Restart the server
    pub async fn restart<T>(&self, transport: T) -> McpResult<()>
    where
        T: crate::transport::traits::ServerTransport + 'static,
    {
        // Stop if running
        if self.is_running().await {
            self.stop().await?;
        }

        // Start with new transport
        self.start(transport).await?;

        Ok(())
    }

    /// Wait for the server to reach a specific state
    pub async fn wait_for_state(
        &self,
        target_state: LifecycleState,
        timeout_duration: Option<Duration>,
    ) -> McpResult<()> {
        let mut state_rx = self.lifecycle.subscribe();

        // Check current state first
        if *state_rx.borrow() == target_state {
            return Ok(());
        }

        let wait_future = async {
            while state_rx.changed().await.is_ok() {
                if *state_rx.borrow() == target_state {
                    return Ok(());
                }
            }
            Err(McpError::Protocol(
                "State change channel closed".to_string(),
            ))
        };

        match timeout_duration {
            Some(duration) => timeout(duration, wait_future)
                .await
                .map_err(|_| McpError::Protocol("Timeout waiting for state change".to_string()))?,
            None => wait_future.await,
        }
    }

    /// Notify all listeners with a closure
    pub async fn notify_listeners<F>(&self, f: F)
    where
        F: Fn(&dyn LifecycleListener) + Send + Sync,
    {
        let listeners = self.listeners.read().await;
        for listener in listeners.iter() {
            f(listener.as_ref());
        }
    }
}

/// Default lifecycle listener that logs state changes
pub struct LoggingLifecycleListener;

impl LifecycleListener for LoggingLifecycleListener {
    fn on_state_change(&self, old_state: LifecycleState, new_state: LifecycleState) {
        tracing::info!("Server state changed: {:?} -> {:?}", old_state, new_state);
    }

    fn on_start(&self) {
        tracing::info!("Server started");
    }

    fn on_stop(&self) {
        tracing::info!("Server stopped");
    }

    fn on_error(&self, error: &McpError) {
        tracing::error!("Server error: {}", error);
    }
}

/// Health check information
#[derive(Debug, Clone)]
pub struct HealthStatus {
    /// Current lifecycle state
    pub state: LifecycleState,
    /// Server uptime
    pub uptime: Option<Duration>,
    /// Whether the server is healthy
    pub healthy: bool,
    /// Additional health details
    pub details: std::collections::HashMap<String, String>,
}

impl ServerRunner {
    /// Get health status
    pub async fn health_status(&self) -> HealthStatus {
        let state = self.lifecycle.state().await;
        let uptime = self.lifecycle.uptime().await;
        let healthy = matches!(state, LifecycleState::Running);

        let mut details = std::collections::HashMap::new();
        details.insert("state".to_string(), format!("{:?}", state));

        if let Some(uptime) = uptime {
            details.insert("uptime_seconds".to_string(), uptime.as_secs().to_string());
        }

        HealthStatus {
            state,
            uptime,
            healthy,
            details,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::server::McpServer;

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

        // Initial state should be Created
        assert_eq!(manager.state().await, LifecycleState::Created);
        assert!(manager.can_start().await);
        assert!(!manager.can_stop().await);
        assert!(!manager.is_running().await);

        // Transition to Running
        manager
            .transition_to(LifecycleState::Running)
            .await
            .unwrap();
        assert_eq!(manager.state().await, LifecycleState::Running);
        assert!(!manager.can_start().await);
        assert!(manager.can_stop().await);
        assert!(manager.is_running().await);
        assert!(manager.uptime().await.is_some());

        // Transition to Stopped
        manager
            .transition_to(LifecycleState::Stopped)
            .await
            .unwrap();
        assert_eq!(manager.state().await, LifecycleState::Stopped);
        assert!(manager.can_start().await);
        assert!(!manager.can_stop().await);
        assert!(!manager.is_running().await);
        assert!(manager.uptime().await.is_none());
    }

    #[tokio::test]
    async fn test_server_runner() {
        let server = McpServer::new("test-server".to_string(), "1.0.0".to_string());
        let runner = ServerRunner::new(server);

        // Initial state
        assert!(!runner.is_running().await);
        assert_eq!(runner.lifecycle().state().await, LifecycleState::Created);

        // Add a logging listener
        runner.add_listener(LoggingLifecycleListener).await;

        // Test health status
        let health = runner.health_status().await;
        assert_eq!(health.state, LifecycleState::Created);
        assert!(!health.healthy);
    }

    #[tokio::test]
    async fn test_state_subscription() {
        let manager = LifecycleManager::new();
        let mut state_rx = manager.subscribe();

        // Initial state
        assert_eq!(*state_rx.borrow(), LifecycleState::Created);

        // Change state
        manager
            .transition_to(LifecycleState::Running)
            .await
            .unwrap();

        // Wait for change
        state_rx.changed().await.unwrap();
        assert_eq!(*state_rx.borrow(), LifecycleState::Running);
    }

    #[tokio::test]
    async fn test_shutdown_signal() {
        let manager = LifecycleManager::new();
        let mut shutdown_rx = manager.create_shutdown_signal().await;

        // Trigger shutdown
        manager.trigger_shutdown().await.unwrap();

        // Wait for signal
        shutdown_rx.changed().await.unwrap();
    }

    struct TestLifecycleListener {
        events: Arc<Mutex<Vec<String>>>,
    }

    impl TestLifecycleListener {
        fn new() -> (Self, Arc<Mutex<Vec<String>>>) {
            let events = Arc::new(Mutex::new(Vec::new()));
            let listener = Self {
                events: events.clone(),
            };
            (listener, events)
        }
    }

    impl LifecycleListener for TestLifecycleListener {
        fn on_state_change(&self, old_state: LifecycleState, new_state: LifecycleState) {
            // Use blocking approach for test to avoid race conditions
            if let Ok(mut events) = self.events.try_lock() {
                events.push(format!("state_change: {:?} -> {:?}", old_state, new_state));
            }
        }

        fn on_start(&self) {
            if let Ok(mut events) = self.events.try_lock() {
                events.push("start".to_string());
            }
        }

        fn on_stop(&self) {
            if let Ok(mut events) = self.events.try_lock() {
                events.push("stop".to_string());
            }
        }

        fn on_error(&self, error: &McpError) {
            if let Ok(mut events) = self.events.try_lock() {
                events.push(format!("error: {}", error));
            }
        }
    }

    #[tokio::test]
    async fn test_lifecycle_listeners() {
        let server = McpServer::new("test-server".to_string(), "1.0.0".to_string());
        let runner = ServerRunner::new(server);

        let (listener, events) = TestLifecycleListener::new();
        runner.add_listener(listener).await;

        // Test calling listeners directly via the notify method
        // Since transition_to doesn't call listeners, we need to test the actual listener functionality
        runner
            .notify_listeners(|listener| {
                listener.on_state_change(LifecycleState::Created, LifecycleState::Running);
            })
            .await;

        // Check that events were captured
        let events = events.lock().await;
        assert!(
            events.len() > 0,
            "Expected at least one event, but got: {:?}",
            *events
        );

        // Verify the specific event was captured
        let has_state_change = events.iter().any(|event| event.contains("state_change"));
        assert!(
            has_state_change,
            "Expected state_change event, but got: {:?}",
            *events
        );
    }
}