reasonkit_web/mcp/

transport.rs

//! MCP Transport Abstraction
//!
//! This module provides the `McpTransport` trait for pluggable MCP transport
//! backends, enabling flexible deployment options and performance optimization.
//!
//! # Supported Backends
//!
//! - **Stdio**: Standard I/O for CLI tool integration
//! - **SSE**: Server-Sent Events for web clients
//! - **WebSocket**: Bidirectional real-time communication
//! - **HTTP**: REST-style request/response
//!
//! # Performance Optimization
//!
//! The transport layer supports pluggable backends for performance tuning:
//! - Default `rmcp` SDK backend
//! - High-performance `pmcp` backend (feature-gated)
//!
//! # Usage
//!
//! ```rust,ignore
//! use reasonkit_web::mcp::transport::{McpTransport, StdioTransport};
//!
//! let transport = StdioTransport::new();
//! let message = transport.receive().await?;
//! transport.send(&response).await?;
//! ```

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use thiserror::Error;

/// MCP Transport errors
#[derive(Error, Debug)]
pub enum TransportError {
    /// Connection establishment failed
    #[error("Connection failed: {0}")]
    ConnectionFailed(String),

    /// Failed to send message
    #[error("Send failed: {0}")]
    SendFailed(String),

    /// Failed to receive message
    #[error("Receive failed: {0}")]
    ReceiveFailed(String),

    /// Operation timed out
    #[error("Timeout after {0:?}")]
    Timeout(Duration),

    /// Transport connection closed
    #[error("Transport closed")]
    Closed,

    /// Invalid message format error
    #[error("Invalid message format: {0}")]
    InvalidFormat(String),

    /// IO operation error
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    /// Serialization/deserialization error
    #[error("Serialization error: {0}")]
    Serialization(String),
}

/// Result type for transport operations
pub type TransportResult<T> = Result<T, TransportError>;

/// Transport configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransportConfig {
    /// Read timeout
    pub read_timeout: Duration,
    /// Write timeout
    pub write_timeout: Duration,
    /// Maximum message size in bytes
    pub max_message_size: usize,
    /// Enable compression
    pub compression: bool,
    /// Buffer size for reading
    pub buffer_size: usize,
}

impl Default for TransportConfig {
    fn default() -> Self {
        Self {
            read_timeout: Duration::from_secs(30),
            write_timeout: Duration::from_secs(30),
            max_message_size: 10 * 1024 * 1024, // 10 MB
            compression: false,
            buffer_size: 8192,
        }
    }
}

/// MCP message envelope
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpMessage {
    /// JSON-RPC version
    pub jsonrpc: String,
    /// Message ID (for request/response correlation)
    pub id: Option<serde_json::Value>,
    /// Method name (for requests)
    pub method: Option<String>,
    /// Parameters (for requests)
    pub params: Option<serde_json::Value>,
    /// Result (for responses)
    pub result: Option<serde_json::Value>,
    /// Error (for error responses)
    pub error: Option<McpError>,
}

impl McpMessage {
    /// Create a new request message
    pub fn request(
        id: impl Into<serde_json::Value>,
        method: &str,
        params: serde_json::Value,
    ) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: Some(id.into()),
            method: Some(method.to_string()),
            params: Some(params),
            result: None,
            error: None,
        }
    }

    /// Create a new response message
    pub fn response(id: impl Into<serde_json::Value>, result: serde_json::Value) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: Some(id.into()),
            method: None,
            params: None,
            result: Some(result),
            error: None,
        }
    }

    /// Create a new error response
    pub fn error_response(id: impl Into<serde_json::Value>, error: McpError) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: Some(id.into()),
            method: None,
            params: None,
            result: None,
            error: Some(error),
        }
    }

    /// Create a notification (no ID)
    pub fn notification(method: &str, params: serde_json::Value) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: None,
            method: Some(method.to_string()),
            params: Some(params),
            result: None,
            error: None,
        }
    }

    /// Check if this is a request
    pub fn is_request(&self) -> bool {
        self.method.is_some() && self.id.is_some()
    }

    /// Check if this is a notification
    pub fn is_notification(&self) -> bool {
        self.method.is_some() && self.id.is_none()
    }

    /// Check if this is a response
    pub fn is_response(&self) -> bool {
        self.id.is_some() && (self.result.is_some() || self.error.is_some())
    }
}

/// MCP error object
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpError {
    /// Error code
    pub code: i32,
    /// Error message
    pub message: String,
    /// Additional error data
    pub data: Option<serde_json::Value>,
}

impl McpError {
    /// Parse error (-32700)
    pub fn parse_error(message: impl Into<String>) -> Self {
        Self {
            code: -32700,
            message: message.into(),
            data: None,
        }
    }

    /// Invalid request (-32600)
    pub fn invalid_request(message: impl Into<String>) -> Self {
        Self {
            code: -32600,
            message: message.into(),
            data: None,
        }
    }

    /// Method not found (-32601)
    pub fn method_not_found(method: &str) -> Self {
        Self {
            code: -32601,
            message: format!("Method not found: {}", method),
            data: None,
        }
    }

    /// Invalid params (-32602)
    pub fn invalid_params(message: impl Into<String>) -> Self {
        Self {
            code: -32602,
            message: message.into(),
            data: None,
        }
    }

    /// Internal error (-32603)
    pub fn internal_error(message: impl Into<String>) -> Self {
        Self {
            code: -32603,
            message: message.into(),
            data: None,
        }
    }
}

/// Transport statistics for monitoring
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TransportStats {
    /// Messages sent
    pub messages_sent: u64,
    /// Messages received
    pub messages_received: u64,
    /// Bytes sent
    pub bytes_sent: u64,
    /// Bytes received
    pub bytes_received: u64,
    /// Send errors
    pub send_errors: u64,
    /// Receive errors
    pub receive_errors: u64,
    /// Average send latency in microseconds
    pub avg_send_latency_us: u64,
    /// Average receive latency in microseconds
    pub avg_receive_latency_us: u64,
}

/// MCP Transport trait for pluggable backends
///
/// Implement this trait to create custom MCP transport backends.
#[async_trait]
pub trait McpTransport: Send + Sync {
    /// Get the transport type name
    fn transport_type(&self) -> &'static str;

    /// Connect the transport
    async fn connect(&mut self) -> TransportResult<()>;

    /// Disconnect the transport
    async fn disconnect(&mut self) -> TransportResult<()>;

    /// Check if connected
    fn is_connected(&self) -> bool;

    /// Send a message
    async fn send(&mut self, message: &McpMessage) -> TransportResult<()>;

    /// Receive a message
    async fn receive(&mut self) -> TransportResult<McpMessage>;

    /// Receive with timeout
    async fn receive_timeout(&mut self, timeout: Duration) -> TransportResult<McpMessage>;

    /// Get transport statistics
    fn stats(&self) -> TransportStats;

    /// Reset statistics
    fn reset_stats(&mut self);
}

/// Transport type enumeration
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TransportType {
    /// Standard I/O
    Stdio,
    /// Server-Sent Events
    Sse,
    /// WebSocket
    WebSocket,
    /// HTTP
    Http,
}

impl std::fmt::Display for TransportType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Stdio => write!(f, "stdio"),
            Self::Sse => write!(f, "sse"),
            Self::WebSocket => write!(f, "websocket"),
            Self::Http => write!(f, "http"),
        }
    }
}

/// Standard I/O transport implementation
pub struct StdioTransport {
    #[allow(dead_code)]
    config: TransportConfig,
    connected: bool,
    stats: TransportStats,
}

impl StdioTransport {
    /// Create a new stdio transport
    pub fn new() -> Self {
        Self::with_config(TransportConfig::default())
    }

    /// Create with configuration
    pub fn with_config(config: TransportConfig) -> Self {
        Self {
            config,
            connected: false,
            stats: TransportStats::default(),
        }
    }
}

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

#[async_trait]
impl McpTransport for StdioTransport {
    fn transport_type(&self) -> &'static str {
        "stdio"
    }

    async fn connect(&mut self) -> TransportResult<()> {
        self.connected = true;
        Ok(())
    }

    async fn disconnect(&mut self) -> TransportResult<()> {
        self.connected = false;
        Ok(())
    }

    fn is_connected(&self) -> bool {
        self.connected
    }

    async fn send(&mut self, message: &McpMessage) -> TransportResult<()> {
        let json = serde_json::to_string(message)
            .map_err(|e| TransportError::Serialization(e.to_string()))?;

        // Write to stdout with Content-Length header (LSP style)
        let content = format!("Content-Length: {}\r\n\r\n{}", json.len(), json);

        use std::io::Write;
        let mut stdout = std::io::stdout().lock();
        stdout
            .write_all(content.as_bytes())
            .map_err(TransportError::Io)?;
        stdout.flush().map_err(TransportError::Io)?;

        self.stats.messages_sent += 1;
        self.stats.bytes_sent += content.len() as u64;

        Ok(())
    }

    async fn receive(&mut self) -> TransportResult<McpMessage> {
        use std::io::{BufRead, Read};

        let stdin = std::io::stdin();
        let mut reader = stdin.lock();

        // Read Content-Length header
        let mut header_line = String::new();
        reader
            .read_line(&mut header_line)
            .map_err(TransportError::Io)?;

        let content_length: usize = header_line
            .trim()
            .strip_prefix("Content-Length: ")
            .ok_or_else(|| TransportError::InvalidFormat("Missing Content-Length header".into()))?
            .parse()
            .map_err(|_| TransportError::InvalidFormat("Invalid Content-Length".into()))?;

        // Skip empty line
        let mut empty = String::new();
        reader.read_line(&mut empty).map_err(TransportError::Io)?;

        // Read content
        let mut content = vec![0u8; content_length];
        reader
            .read_exact(&mut content)
            .map_err(TransportError::Io)?;

        let message: McpMessage = serde_json::from_slice(&content)
            .map_err(|e| TransportError::InvalidFormat(e.to_string()))?;

        self.stats.messages_received += 1;
        self.stats.bytes_received += content_length as u64;

        Ok(message)
    }

    async fn receive_timeout(&mut self, _timeout: Duration) -> TransportResult<McpMessage> {
        // For stdio, timeout handling would require threading
        // This is a simplified implementation
        self.receive().await
    }

    fn stats(&self) -> TransportStats {
        self.stats.clone()
    }

    fn reset_stats(&mut self) {
        self.stats = TransportStats::default();
    }
}

/// Transport factory for creating transport instances
pub struct TransportFactory;

impl TransportFactory {
    /// Create a transport of the specified type
    pub fn create(transport_type: TransportType) -> Box<dyn McpTransport> {
        match transport_type {
            TransportType::Stdio => Box::new(StdioTransport::new()),
            // Other transports would be implemented similarly
            TransportType::Sse | TransportType::WebSocket | TransportType::Http => {
                // Placeholder - would need full implementations
                Box::new(StdioTransport::new())
            }
        }
    }

    /// Create a transport with configuration
    pub fn create_with_config(
        transport_type: TransportType,
        config: TransportConfig,
    ) -> Box<dyn McpTransport> {
        match transport_type {
            TransportType::Stdio => Box::new(StdioTransport::with_config(config)),
            _ => Box::new(StdioTransport::with_config(config)),
        }
    }
}

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

    #[test]
    fn test_message_request() {
        let msg = McpMessage::request(1, "tools/list", serde_json::json!({}));
        assert!(msg.is_request());
        assert!(!msg.is_notification());
        assert!(!msg.is_response());
    }

    #[test]
    fn test_message_notification() {
        let msg = McpMessage::notification("progress", serde_json::json!({ "percent": 50 }));
        assert!(!msg.is_request());
        assert!(msg.is_notification());
        assert!(!msg.is_response());
    }

    #[test]
    fn test_message_response() {
        let msg = McpMessage::response(1, serde_json::json!({ "tools": [] }));
        assert!(!msg.is_request());
        assert!(!msg.is_notification());
        assert!(msg.is_response());
    }

    #[test]
    fn test_error_codes() {
        let err = McpError::method_not_found("unknown");
        assert_eq!(err.code, -32601);

        let err = McpError::invalid_params("bad param");
        assert_eq!(err.code, -32602);
    }

    #[test]
    fn test_transport_config_default() {
        let config = TransportConfig::default();
        assert_eq!(config.read_timeout, Duration::from_secs(30));
        assert!(!config.compression);
    }
}