mcp_host/server/

multiplexer.rs

//! Request Multiplexer for bidirectional JSON-RPC communication
//!
//! Enables server→client requests (roots/list, sampling/createMessage, etc.)
//! by tracking pending requests and routing responses back to callers.
//!
//! # Architecture
//!
//! The multiplexer sits between the server and transport, handling:
//! - Client→Server: Normal requests (tools/call, resources/read, etc.)
//! - Server→Client: Requests initiated by the server (roots/list, sampling/createMessage)
//!
//! For server-initiated requests, the multiplexer:
//! 1. Generates a unique request ID (UUID)
//! 2. Stores a oneshot channel in pending requests map
//! 3. Sends the request via transport
//! 4. When response arrives, routes it to the waiting channel
//!
//! # Example
//!
//! ```rust,ignore
//! // Server requesting roots from client
//! let roots = server.request_roots().await?;
//! for root in roots {
//!     println!("Client root: {} ({})", root.name, root.uri);
//! }
//!
//! // Server requesting sampling from client
//! let result = server.request_sampling(messages, preferences).await?;
//! println!("Model response: {}", result.content);
//! ```

use std::sync::Arc;
use std::time::Duration;

use tokio::sync::mpsc;

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use tokio::sync::oneshot;

use crate::protocol::types::JsonRpcResponse;

/// Errors from server→client requests
#[derive(Debug, Error)]
pub enum MultiplexerError {
    /// Request timed out waiting for response
    #[error("request timed out after {0:?}")]
    Timeout(Duration),

    /// Client returned an error response
    #[error("client error {code}: {message}")]
    ClientError { code: i32, message: String },

    /// Transport error
    #[error("transport error: {0}")]
    Transport(String),

    /// Response channel was closed (internal error)
    #[error("response channel closed")]
    ChannelClosed,

    /// Failed to serialize request
    #[error("serialization error: {0}")]
    Serialization(#[from] serde_json::Error),

    /// Client doesn't support the requested capability
    #[error("client does not support {0}")]
    UnsupportedCapability(String),
}

/// A pending server-initiated request awaiting response
pub struct PendingRequest {
    /// Request ID (UUID string)
    pub id: String,

    /// Method name for logging/debugging
    pub method: String,

    /// Channel to send the response
    pub response_tx: oneshot::Sender<Result<Value, MultiplexerError>>,

    /// When this request was created (for timeout tracking)
    pub created_at: std::time::Instant,
}

/// Request multiplexer for handling bidirectional communication
///
/// Tracks pending server→client requests and routes responses
pub struct RequestMultiplexer {
    /// Pending requests awaiting responses: request_id → PendingRequest
    pending: DashMap<String, PendingRequest>,

    /// Default timeout for requests
    default_timeout: Duration,
}

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

impl RequestMultiplexer {
    /// Create a new request multiplexer
    pub fn new() -> Self {
        Self {
            pending: DashMap::new(),
            default_timeout: Duration::from_secs(30),
        }
    }

    /// Create with custom default timeout
    pub fn with_timeout(timeout: Duration) -> Self {
        Self {
            pending: DashMap::new(),
            default_timeout: timeout,
        }
    }

    /// Get default timeout
    pub fn default_timeout(&self) -> Duration {
        self.default_timeout
    }

    /// Create a new pending request and return the receiver
    ///
    /// Returns (request_id, receiver) - the caller should await the receiver
    pub fn create_pending(&self, method: impl Into<String>) -> (String, oneshot::Receiver<Result<Value, MultiplexerError>>) {
        let id = uuid::Uuid::new_v4().to_string();
        let method = method.into();
        let (tx, rx) = oneshot::channel();

        let pending = PendingRequest {
            id: id.clone(),
            method,
            response_tx: tx,
            created_at: std::time::Instant::now(),
        };

        self.pending.insert(id.clone(), pending);

        (id, rx)
    }

    /// Route an incoming response to its pending request
    ///
    /// Returns true if the response was routed, false if no matching request found
    pub fn route_response(&self, response: &JsonRpcResponse) -> bool {
        // Extract request ID from response
        let id = match &response.id {
            Value::String(s) => s.clone(),
            Value::Number(n) => n.to_string(),
            _ => return false,
        };

        // Find and remove the pending request
        if let Some((_, pending)) = self.pending.remove(&id) {
            // Build result or error
            let result = if let Some(ref error) = response.error {
                Err(MultiplexerError::ClientError {
                    code: error.code,
                    message: error.message.clone(),
                })
            } else if let Some(ref result) = response.result {
                Ok(result.clone())
            } else {
                // Empty response (no result, no error) - treat as empty object
                Ok(Value::Object(serde_json::Map::new()))
            };

            // Send to waiting caller (ignore if receiver dropped)
            let _ = pending.response_tx.send(result);

            true
        } else {
            false
        }
    }

    /// Check if a response ID matches a pending request
    ///
    /// Used to distinguish client responses from client requests in the message loop
    pub fn is_pending_response(&self, id: &Value) -> bool {
        let id_str = match id {
            Value::String(s) => s.clone(),
            Value::Number(n) => n.to_string(),
            _ => return false,
        };

        self.pending.contains_key(&id_str)
    }

    /// Get number of pending requests
    pub fn pending_count(&self) -> usize {
        self.pending.len()
    }

    /// Cancel a pending request
    pub fn cancel(&self, id: &str) {
        if let Some((_, pending)) = self.pending.remove(id) {
            let _ = pending.response_tx.send(Err(MultiplexerError::ChannelClosed));
        }
    }

    /// Cancel all pending requests
    pub fn cancel_all(&self) {
        let ids: Vec<String> = self.pending.iter().map(|e| e.key().clone()).collect();
        for id in ids {
            self.cancel(&id);
        }
    }

    /// Clean up timed-out requests
    ///
    /// Returns the number of requests that were cleaned up
    pub fn cleanup_timed_out(&self, timeout: Duration) -> usize {
        let now = std::time::Instant::now();
        let mut cleaned = 0;

        // Collect IDs to remove (can't modify during iteration with DashMap)
        let timed_out: Vec<String> = self
            .pending
            .iter()
            .filter(|e| now.duration_since(e.created_at) > timeout)
            .map(|e| e.key().clone())
            .collect();

        for id in timed_out {
            if let Some((_, pending)) = self.pending.remove(&id) {
                let _ = pending
                    .response_tx
                    .send(Err(MultiplexerError::Timeout(timeout)));
                cleaned += 1;
            }
        }

        cleaned
    }
}

/// Server→client request (what we send to the client)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcClientRequest {
    /// JSON-RPC version
    pub jsonrpc: String,

    /// Request ID (UUID string for server-initiated requests)
    pub id: String,

    /// Method name
    pub method: String,

    /// Request parameters
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

impl JsonRpcClientRequest {
    /// Create a new client request
    pub fn new(id: impl Into<String>, method: impl Into<String>, params: Option<Value>) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: id.into(),
            method: method.into(),
            params,
        }
    }
}

/// Root entry from roots/list response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Root {
    /// URI of the root (e.g., "file:///workspace")
    pub uri: String,

    /// Human-readable name for the root
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

/// Result of roots/list request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ListRootsResult {
    /// List of workspace roots
    pub roots: Vec<Root>,
}

/// Sampling message for createMessage request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SamplingMessage {
    /// Role: "user" or "assistant"
    pub role: String,

    /// Message content
    pub content: SamplingContent,
}

/// Content in a sampling message
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum SamplingContent {
    /// Text content
    #[serde(rename = "text")]
    Text { text: String },

    /// Image content
    #[serde(rename = "image")]
    Image { data: String, mime_type: String },
}

/// Model preferences for sampling
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelPreferences {
    /// Preferred model hints
    #[serde(skip_serializing_if = "Option::is_none")]
    pub hints: Option<Vec<ModelHint>>,

    /// Cost priority (0.0 = prefer cheap, 1.0 = prefer quality)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cost_priority: Option<f64>,

    /// Speed priority (0.0 = prefer slow, 1.0 = prefer fast)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub speed_priority: Option<f64>,

    /// Intelligence priority (0.0 = prefer simple, 1.0 = prefer capable)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub intelligence_priority: Option<f64>,
}

/// Model hint for sampling
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelHint {
    /// Model name pattern
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

/// Parameters for sampling/createMessage request
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CreateMessageParams {
    /// Messages to send
    pub messages: Vec<SamplingMessage>,

    /// Model preferences
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model_preferences: Option<ModelPreferences>,

    /// System prompt
    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_prompt: Option<String>,

    /// Include context from MCP servers
    #[serde(skip_serializing_if = "Option::is_none")]
    pub include_context: Option<String>,

    /// Temperature
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f64>,

    /// Maximum tokens to generate
    pub max_tokens: i32,

    /// Stop sequences
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop_sequences: Option<Vec<String>>,
}

/// Result of sampling/createMessage request
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CreateMessageResult {
    /// Role of the response
    pub role: String,

    /// Response content
    pub content: SamplingContent,

    /// Model that was used
    pub model: String,

    /// Reason for stopping
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop_reason: Option<String>,
}

/// Client requester for making server→client requests from tools
///
/// This is passed to tools via ExecutionContext, allowing them to
/// make requests to the client (roots/list, sampling/createMessage, etc.)
///
/// # Example
///
/// ```rust,ignore
/// impl Tool for ListRootsTool {
///     async fn execute(&self, ctx: ExecutionContext<'_>) -> Result<Vec<Box<dyn Content>>, ToolError> {
///         if let Some(requester) = ctx.client_requester() {
///             match requester.request_roots(None).await {
///                 Ok(roots) => {
///                     let msg = format!("Found {} roots", roots.len());
///                     Ok(vec![Box::new(TextContent::new(msg))])
///                 }
///                 Err(e) => Err(ToolError::Execution(e.to_string())),
///             }
///         } else {
///             Err(ToolError::Execution("No client requester available".into()))
///         }
///     }
/// }
/// ```
#[derive(Clone)]
pub struct ClientRequester {
    /// Channel to send requests to the transport
    request_tx: mpsc::UnboundedSender<JsonRpcClientRequest>,

    /// Multiplexer for tracking pending requests
    multiplexer: Arc<RequestMultiplexer>,

    /// Whether the client supports roots capability
    supports_roots: bool,

    /// Whether the client supports sampling capability
    supports_sampling: bool,
}

impl ClientRequester {
    /// Create a new client requester
    pub fn new(
        request_tx: mpsc::UnboundedSender<JsonRpcClientRequest>,
        multiplexer: Arc<RequestMultiplexer>,
        supports_roots: bool,
        supports_sampling: bool,
    ) -> Self {
        Self {
            request_tx,
            multiplexer,
            supports_roots,
            supports_sampling,
        }
    }

    /// Check if client supports roots capability
    pub fn supports_roots(&self) -> bool {
        self.supports_roots
    }

    /// Check if client supports sampling capability
    pub fn supports_sampling(&self) -> bool {
        self.supports_sampling
    }

    /// Request workspace roots from the client
    ///
    /// Returns an error if the client doesn't support roots capability.
    pub async fn request_roots(
        &self,
        timeout: Option<Duration>,
    ) -> Result<Vec<Root>, MultiplexerError> {
        if !self.supports_roots {
            return Err(MultiplexerError::UnsupportedCapability("roots".to_string()));
        }

        // Create pending request
        let (id, rx) = self.multiplexer.create_pending("roots/list");

        // Build and send the request
        let request = JsonRpcClientRequest::new(&id, "roots/list", Some(serde_json::json!({})));

        self.request_tx
            .send(request)
            .map_err(|e| MultiplexerError::Transport(e.to_string()))?;

        // Wait for response with timeout
        let timeout = timeout.unwrap_or(self.multiplexer.default_timeout());
        let result = tokio::time::timeout(timeout, rx)
            .await
            .map_err(|_| MultiplexerError::Timeout(timeout))?
            .map_err(|_| MultiplexerError::ChannelClosed)??;

        // Parse the result
        let list_result: ListRootsResult = serde_json::from_value(result)?;
        Ok(list_result.roots)
    }

    /// Request an LLM completion from the client
    ///
    /// Returns an error if the client doesn't support sampling capability.
    pub async fn request_sampling(
        &self,
        params: CreateMessageParams,
        timeout: Option<Duration>,
    ) -> Result<CreateMessageResult, MultiplexerError> {
        if !self.supports_sampling {
            return Err(MultiplexerError::UnsupportedCapability("sampling".to_string()));
        }

        // Create pending request
        let (id, rx) = self.multiplexer.create_pending("sampling/createMessage");

        // Build and send the request
        let params_value = serde_json::to_value(&params)?;
        let request = JsonRpcClientRequest::new(&id, "sampling/createMessage", Some(params_value));

        self.request_tx
            .send(request)
            .map_err(|e| MultiplexerError::Transport(e.to_string()))?;

        // Wait for response with timeout
        let timeout = timeout.unwrap_or(self.multiplexer.default_timeout());
        let result = tokio::time::timeout(timeout, rx)
            .await
            .map_err(|_| MultiplexerError::Timeout(timeout))?
            .map_err(|_| MultiplexerError::ChannelClosed)??;

        // Parse the result
        let create_result: CreateMessageResult = serde_json::from_value(result)?;
        Ok(create_result)
    }

    /// Request user input elicitation from the client
    ///
    /// Sends an `elicitation/create` request to the client and waits for the response.
    /// The client must have the `elicitation` capability advertised.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to show the user
    /// * `schema` - The schema defining the structure of requested input
    /// * `timeout` - Optional timeout for the request
    ///
    /// # Returns
    ///
    /// The user's response (accept/decline/cancel with optional data), or an error if:
    /// - Client doesn't support elicitation capability
    /// - Request times out
    /// - Transport error occurs
    pub async fn request_elicitation(
        &self,
        message: String,
        requested_schema: Value,
        timeout: Option<Duration>,
    ) -> Result<crate::protocol::types::CreateElicitationResult, MultiplexerError> {
        // Note: We'd need to track elicitation capability during initialization
        // For now, just attempt the request

        // Create pending request
        let (id, rx) = self.multiplexer.create_pending("elicitation/create");

        // Build request params
        let params = serde_json::json!({
            "message": message,
            "requestedSchema": requested_schema,
        });

        let request = JsonRpcClientRequest::new(&id, "elicitation/create", Some(params));

        self.request_tx
            .send(request)
            .map_err(|e| MultiplexerError::Transport(e.to_string()))?;

        // Wait for response with timeout
        let timeout = timeout.unwrap_or(self.multiplexer.default_timeout());
        let result = tokio::time::timeout(timeout, rx)
            .await
            .map_err(|_| MultiplexerError::Timeout(timeout))?
            .map_err(|_| MultiplexerError::ChannelClosed)??;

        // Parse the result
        let elicitation_result: crate::protocol::types::CreateElicitationResult =
            serde_json::from_value(result)?;
        Ok(elicitation_result)
    }
}

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

    #[test]
    fn test_multiplexer_create_pending() {
        let mux = RequestMultiplexer::new();

        let (id1, _rx1) = mux.create_pending("roots/list");
        let (id2, _rx2) = mux.create_pending("sampling/createMessage");

        assert_ne!(id1, id2);
        assert_eq!(mux.pending_count(), 2);
    }

    #[tokio::test]
    async fn test_multiplexer_route_response() {
        let mux = RequestMultiplexer::new();

        let (id, rx) = mux.create_pending("test/method");

        // Create a response
        let response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: Value::String(id.clone()),
            result: Some(serde_json::json!({"status": "ok"})),
            error: None,
        };

        // Route it
        assert!(mux.route_response(&response));
        assert_eq!(mux.pending_count(), 0);

        // Receiver should get the result
        let result = rx.await.unwrap().unwrap();
        assert_eq!(result["status"], "ok");
    }

    #[tokio::test]
    async fn test_multiplexer_route_error() {
        let mux = RequestMultiplexer::new();

        let (id, rx) = mux.create_pending("test/method");

        // Create an error response
        let response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: Value::String(id.clone()),
            result: None,
            error: Some(crate::protocol::types::JsonRpcError {
                code: -32600,
                message: "Invalid request".to_string(),
                data: None,
            }),
        };

        // Route it
        assert!(mux.route_response(&response));

        // Receiver should get the error
        let result = rx.await.unwrap();
        assert!(matches!(result, Err(MultiplexerError::ClientError { code: -32600, .. })));
    }

    #[test]
    fn test_multiplexer_is_pending() {
        let mux = RequestMultiplexer::new();

        let (id, _rx) = mux.create_pending("test");

        assert!(mux.is_pending_response(&Value::String(id.clone())));
        assert!(!mux.is_pending_response(&Value::String("unknown".to_string())));
    }

    #[test]
    fn test_multiplexer_cancel() {
        let mux = RequestMultiplexer::new();

        let (id, _rx) = mux.create_pending("test");
        assert_eq!(mux.pending_count(), 1);

        mux.cancel(&id);
        assert_eq!(mux.pending_count(), 0);
    }

    #[test]
    fn test_client_request_serialization() {
        let req = JsonRpcClientRequest::new(
            "abc-123",
            "roots/list",
            Some(serde_json::json!({})),
        );

        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("\"jsonrpc\":\"2.0\""));
        assert!(json.contains("\"id\":\"abc-123\""));
        assert!(json.contains("\"method\":\"roots/list\""));
    }

    #[test]
    fn test_root_deserialization() {
        let json = r#"{"uri": "file:///workspace", "name": "My Project"}"#;
        let root: Root = serde_json::from_str(json).unwrap();

        assert_eq!(root.uri, "file:///workspace");
        assert_eq!(root.name, Some("My Project".to_string()));
    }

    #[test]
    fn test_sampling_content() {
        let content = SamplingContent::Text {
            text: "Hello, world!".to_string(),
        };

        let json = serde_json::to_string(&content).unwrap();
        assert!(json.contains("\"type\":\"text\""));
        assert!(json.contains("\"text\":\"Hello, world!\""));
    }
}