mcpls_core/lsp/

client.rs

//! LSP client implementation with async request/response handling.

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};

use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use tokio::sync::{Mutex, mpsc, oneshot};
use tokio::task::JoinHandle;
use tokio::time::{Duration, timeout};
use tracing::{debug, error, trace, warn};

use crate::config::LspServerConfig;
use crate::error::{Error, Result};
use crate::lsp::transport::LspTransport;
use crate::lsp::types::{
    InboundMessage, JsonRpcError, JsonRpcRequest, JsonRpcResponse, LspNotification, RequestId,
};

/// JSON-RPC protocol version.
const JSONRPC_VERSION: &str = "2.0";

/// LSP error code returned when the server cancels a request and wants the client to retry.
const SERVER_CANCELLED_CODE: i32 = -32802;

/// Maximum number of retry attempts for server-cancelled requests.
const SERVER_CANCELLED_MAX_RETRIES: u32 = 3;

/// Initial backoff delay for server-cancelled retries (milliseconds).
const SERVER_CANCELLED_INITIAL_DELAY_MS: u64 = 500;

/// Type alias for pending request tracking map.
type PendingRequests = HashMap<RequestId, oneshot::Sender<Result<Value>>>;

/// LSP client with async request/response handling.
///
/// This client manages communication with an LSP server, handling:
/// - Concurrent requests with unique ID tracking
/// - Background message loop for receiving responses
/// - Timeout support for all requests
/// - Graceful shutdown
#[derive(Debug)]
pub struct LspClient {
    /// Configuration for this LSP server.
    config: LspServerConfig,

    /// Current server state.
    state: Arc<Mutex<super::ServerState>>,

    /// Atomic counter for request IDs.
    request_counter: Arc<AtomicI64>,

    /// Command sender for outbound messages.
    command_tx: mpsc::Sender<ClientCommand>,

    /// Background receiver task handle.
    receiver_task: Option<JoinHandle<Result<()>>>,
}

impl Clone for LspClient {
    /// Creates a clone that shares the underlying connection.
    ///
    /// The clone does not own the receiver task and cannot perform shutdown.
    /// All clones share the same command channel for sending requests.
    fn clone(&self) -> Self {
        Self {
            config: self.config.clone(),
            state: Arc::clone(&self.state),
            request_counter: Arc::clone(&self.request_counter),
            command_tx: self.command_tx.clone(),
            receiver_task: None,
        }
    }
}

/// Commands for client control.
enum ClientCommand {
    /// Send a request and wait for response.
    SendRequest {
        request: JsonRpcRequest,
        response_tx: oneshot::Sender<Result<Value>>,
    },
    /// Send a notification (no response expected).
    SendNotification {
        method: String,
        params: Option<Value>,
    },
    /// Shutdown the client.
    Shutdown,
}

impl LspClient {
    /// Create a new LSP client with the given configuration.
    ///
    /// The client starts in an uninitialized state. Call `initialize()` to
    /// start the server and complete the initialization handshake.
    #[must_use]
    pub fn new(config: LspServerConfig) -> Self {
        // Placeholder channel - the receiver is intentionally dropped since
        // the client starts uninitialized. A real channel is created when
        // `from_transport` or `from_transport_with_notifications` is called.
        let (command_tx, _command_rx) = mpsc::channel(1); // Minimal capacity for placeholder

        Self {
            config,
            state: Arc::new(Mutex::new(super::ServerState::Uninitialized)),
            request_counter: Arc::new(AtomicI64::new(1)),
            command_tx,
            receiver_task: None,
        }
    }

    /// Create client from transport (for testing or custom spawning).
    ///
    /// This method initializes the background message loop with the provided transport.
    #[cfg(test)]
    pub(crate) fn from_transport(config: LspServerConfig, transport: LspTransport) -> Self {
        let state = Arc::new(Mutex::new(super::ServerState::Initializing));
        let request_counter = Arc::new(AtomicI64::new(1));
        let pending_requests = Arc::new(Mutex::new(HashMap::new()));

        let (command_tx, command_rx) = mpsc::channel(100);

        let receiver_task = tokio::spawn(Self::message_loop(
            transport,
            command_rx,
            pending_requests,
            None,
        ));

        Self {
            config,
            state,
            request_counter,
            command_tx,
            receiver_task: Some(receiver_task),
        }
    }

    /// Create client from transport with notification forwarding.
    ///
    /// Notifications received from the LSP server will be parsed and sent
    /// through the provided channel.
    pub(crate) fn from_transport_with_notifications(
        config: LspServerConfig,
        transport: LspTransport,
        notification_tx: mpsc::Sender<LspNotification>,
    ) -> Self {
        let state = Arc::new(Mutex::new(super::ServerState::Initializing));
        let request_counter = Arc::new(AtomicI64::new(1));
        let pending_requests = Arc::new(Mutex::new(HashMap::new()));

        let (command_tx, command_rx) = mpsc::channel(100);

        let receiver_task = tokio::spawn(Self::message_loop(
            transport,
            command_rx,
            pending_requests,
            Some(notification_tx),
        ));

        Self {
            config,
            state,
            request_counter,
            command_tx,
            receiver_task: Some(receiver_task),
        }
    }

    /// Get the language ID for this client.
    #[must_use]
    pub fn language_id(&self) -> &str {
        &self.config.language_id
    }

    /// Get the current server state.
    pub async fn state(&self) -> super::ServerState {
        *self.state.lock().await
    }

    /// Send request and wait for response with timeout.
    ///
    /// Automatically retries up to 3 times when the server returns error code
    /// -32802 (`ServerCancelled`) with `data.retriggerRequest == true`, using
    /// exponential backoff starting at 500 ms.
    ///
    /// # Type Parameters
    ///
    /// * `P` - The type of the request parameters (must be serializable)
    /// * `R` - The type of the response result (must be deserializable)
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Server has shut down
    /// - Request times out
    /// - Response cannot be deserialized
    /// - LSP server returns an error
    pub async fn request<P, R>(
        &self,
        method: &str,
        params: P,
        timeout_duration: Duration,
    ) -> Result<R>
    where
        P: Serialize,
        R: DeserializeOwned,
    {
        let params_value = serde_json::to_value(params)?;
        let mut delay_ms = SERVER_CANCELLED_INITIAL_DELAY_MS;

        for attempt in 0..=SERVER_CANCELLED_MAX_RETRIES {
            if attempt > 0 {
                debug!(
                    "Retrying {} after ServerCancelled (attempt {}/{}), backoff={}ms",
                    method, attempt, SERVER_CANCELLED_MAX_RETRIES, delay_ms
                );
                tokio::time::sleep(Duration::from_millis(delay_ms)).await;
                delay_ms *= 2;
            }

            let id = RequestId::Number(self.request_counter.fetch_add(1, Ordering::SeqCst));
            let (response_tx, response_rx) = oneshot::channel();
            let request = JsonRpcRequest {
                jsonrpc: JSONRPC_VERSION.to_string(),
                id: id.clone(),
                method: method.to_string(),
                params: Some(params_value.clone()),
            };

            debug!("Sending request: {} (id={:?})", method, id);

            self.command_tx
                .send(ClientCommand::SendRequest {
                    request,
                    response_tx,
                })
                .await
                .map_err(|_| Error::ServerTerminated)?;

            let outcome = timeout(timeout_duration, response_rx)
                .await
                .map_err(|_| Error::Timeout(timeout_duration.as_secs()))?
                .map_err(|_| Error::ServerTerminated)?;

            match outcome {
                Ok(result_value) => {
                    return serde_json::from_value(result_value).map_err(|e| {
                        Error::LspProtocolError(format!("Failed to deserialize response: {e}"))
                    });
                }
                Err(Error::LspServerError {
                    code,
                    ref message,
                    ref data,
                }) if code == SERVER_CANCELLED_CODE && Self::should_retrigger(data.as_ref()) => {
                    warn!(
                        "ServerCancelled (-32802) on '{}', will retry: {}",
                        method, message
                    );
                    if attempt == SERVER_CANCELLED_MAX_RETRIES {
                        return Err(Error::LspServerError {
                            code,
                            message: message.clone(),
                            data: data.clone(),
                        });
                    }
                    // continue loop for next attempt
                }
                Err(e) => return Err(e),
            }
        }

        Err(Error::ServerTerminated)
    }

    /// Returns true when the error data from a `ServerCancelled` (-32802) response
    /// indicates the server wants the client to retrigger the request.
    ///
    /// Per the LSP specification, `data.retriggerRequest == true` is the signal.
    /// When `data` is absent (older servers), we default to retrying anyway because
    /// code -32802 is exclusively used for this purpose.
    fn should_retrigger(data: Option<&Value>) -> bool {
        data.is_none_or(|v| {
            v.get("retriggerRequest")
                .and_then(Value::as_bool)
                .unwrap_or(true)
        })
    }

    /// Send notification (fire-and-forget, no response expected).
    ///
    /// # Errors
    ///
    /// Returns an error if the server has shut down.
    pub async fn notify<P>(&self, method: &str, params: P) -> Result<()>
    where
        P: Serialize,
    {
        let params_value = serde_json::to_value(params)?;

        debug!("Sending notification: {}", method);

        self.command_tx
            .send(ClientCommand::SendNotification {
                method: method.to_string(),
                params: Some(params_value),
            })
            .await
            .map_err(|_| Error::ServerTerminated)?;

        Ok(())
    }

    /// Shutdown client gracefully.
    ///
    /// This sends a shutdown command to the background task and waits for it to complete.
    ///
    /// # Errors
    ///
    /// Returns an error if the background task failed.
    pub async fn shutdown(mut self) -> Result<()> {
        debug!("Shutting down LSP client");

        let _ = self.command_tx.send(ClientCommand::Shutdown).await;

        if let Some(task) = self.receiver_task.take() {
            task.await
                .map_err(|e| Error::Transport(format!("Receiver task failed: {e}")))??;
        }

        *self.state.lock().await = super::ServerState::Shutdown;

        Ok(())
    }

    /// Background task: handle message I/O.
    ///
    /// This task runs in the background, handling:
    /// - Outbound requests and notifications
    /// - Inbound responses and server notifications
    /// - Matching responses to pending requests
    async fn message_loop(
        mut transport: LspTransport,
        mut command_rx: mpsc::Receiver<ClientCommand>,
        pending_requests: Arc<Mutex<PendingRequests>>,
        notification_tx: Option<mpsc::Sender<LspNotification>>,
    ) -> Result<()> {
        debug!("Message loop started");
        let result = Self::message_loop_inner(
            &mut transport,
            &mut command_rx,
            &pending_requests,
            notification_tx.as_ref(),
        )
        .await;
        if let Err(ref e) = result {
            error!("Message loop exiting with error: {}", e);
        } else {
            debug!("Message loop exiting normally");
        }
        result
    }

    async fn message_loop_inner(
        transport: &mut LspTransport,
        command_rx: &mut mpsc::Receiver<ClientCommand>,
        pending_requests: &Arc<Mutex<PendingRequests>>,
        notification_tx: Option<&mpsc::Sender<LspNotification>>,
    ) -> Result<()> {
        loop {
            tokio::select! {
                Some(command) = command_rx.recv() => {
                    match command {
                        ClientCommand::SendRequest { request, response_tx } => {
                            pending_requests.lock().await.insert(
                                request.id.clone(),
                                response_tx,
                            );

                            let value = serde_json::to_value(&request)?;
                            transport.send(&value).await?;
                        }
                        ClientCommand::SendNotification { method, params } => {
                            let notification = serde_json::json!({
                                "jsonrpc": "2.0",
                                "method": method,
                                "params": params,
                            });
                            transport.send(&notification).await?;
                        }
                        ClientCommand::Shutdown => {
                            debug!("Client shutdown requested");
                            break;
                        }
                    }
                }

                message = transport.receive() => {
                    let message = match message {
                        Ok(m) => m,
                        Err(e) => {
                            error!("Transport receive error: {}", e);
                            return Err(e);
                        }
                    };
                    match message {
                        InboundMessage::Response(response) => {
                            trace!("Received response: id={:?}", response.id);

                            let sender = pending_requests.lock().await.remove(&response.id);

                            if let Some(sender) = sender {
                                if let Some(error) = response.error {
                                    let message = if error.message.len() > 200 {
                                        format!("{}... (truncated)", &error.message[..200])
                                    } else {
                                        error.message.clone()
                                    };
                                    error!("LSP error response: {} (code {})", message, error.code);
                                    let _ = sender.send(Err(Error::LspServerError {
                                        code: error.code,
                                        message: error.message,
                                        data: error.data,
                                    }));
                                } else if let Some(result) = response.result {
                                    let _ = sender.send(Ok(result));
                                } else {
                                    // LSP spec allows null result for some requests (e.g., hover with no info).
                                    // Treat as successful response with null value.
                                    trace!("Response with null result: {:?}", response.id);
                                    let _ = sender.send(Ok(Value::Null));
                                }
                            } else {
                                warn!("Received response for unknown request ID: {:?}", response.id);
                            }
                        }
                        InboundMessage::Request(request) => {
                            debug!(
                                "Received server request: {} (id={:?})",
                                request.method, request.id
                            );
                            let response = Self::server_request_response(request);
                            let value = serde_json::to_value(&response)?;
                            transport.send(&value).await?;
                        }
                        InboundMessage::Notification(notification) => {
                            debug!("Received notification: {}", notification.method);

                            // Parse notification into typed variant
                            let typed = LspNotification::parse(&notification.method, notification.params);

                            // Forward to notification handler if sender is available
                            if let Some(tx) = notification_tx {
                                // Log diagnostics count since it's useful for debugging
                                if let LspNotification::PublishDiagnostics(ref params) = typed {
                                    debug!(
                                        "Forwarding diagnostics for {}: {} items",
                                        params.uri.as_str(),
                                        params.diagnostics.len()
                                    );
                                } else {
                                    trace!("Forwarding notification: {:?}", typed);
                                }

                                // Send the notification with backpressure handling
                                if tx.try_send(typed).is_err() {
                                    warn!("Notification channel full or closed, dropping notification");
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    fn server_request_response(request: JsonRpcRequest) -> JsonRpcResponse {
        match Self::server_request_result(&request.method, request.params.as_ref()) {
            Ok(result) => JsonRpcResponse {
                jsonrpc: JSONRPC_VERSION.to_string(),
                id: request.id,
                result: Some(result),
                error: None,
            },
            Err(error) => JsonRpcResponse {
                jsonrpc: JSONRPC_VERSION.to_string(),
                id: request.id,
                result: None,
                error: Some(error),
            },
        }
    }

    fn server_request_result(
        method: &str,
        params: Option<&Value>,
    ) -> std::result::Result<Value, JsonRpcError> {
        match method {
            "client/registerCapability"
            | "client/unregisterCapability"
            | "workspace/workspaceFolders"
            | "workspace/diagnostic/refresh"
            | "workspace/semanticTokens/refresh"
            | "workspace/inlayHint/refresh"
            | "workspace/codeLens/refresh"
            | "window/showMessageRequest" => Ok(Value::Null),
            "workspace/configuration" => Ok(Self::workspace_configuration_result(params)),
            "workspace/applyEdit" => Ok(serde_json::json!({ "applied": false })),
            _ => Err(JsonRpcError {
                code: -32601,
                message: format!("Unhandled server request: {method}"),
                data: None,
            }),
        }
    }

    fn workspace_configuration_result(params: Option<&Value>) -> Value {
        let item_count = params
            .and_then(|value| value.get("items"))
            .and_then(Value::as_array)
            .map_or(0, Vec::len);

        Value::Array(vec![Value::Null; item_count])
    }
}

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

    #[test]
    fn test_request_id_generation() {
        let counter = AtomicI64::new(1);

        let id1 = counter.fetch_add(1, Ordering::SeqCst);
        let id2 = counter.fetch_add(1, Ordering::SeqCst);
        let id3 = counter.fetch_add(1, Ordering::SeqCst);

        assert_eq!(id1, 1);
        assert_eq!(id2, 2);
        assert_eq!(id3, 3);
    }

    #[test]
    fn test_client_creation() {
        let config = LspServerConfig::rust_analyzer();

        let client = LspClient::new(config);
        assert_eq!(client.language_id(), "rust");
    }

    #[test]
    fn test_client_clone() {
        let config = LspServerConfig::rust_analyzer();
        let client = LspClient::new(config);

        #[allow(clippy::redundant_clone)]
        let cloned = client.clone();
        assert_eq!(cloned.language_id(), "rust");

        assert!(
            cloned.receiver_task.is_none(),
            "Cloned client should not own receiver task"
        );
    }

    #[test]
    fn test_register_capability_request_is_acknowledged() {
        let request = JsonRpcRequest {
            jsonrpc: JSONRPC_VERSION.to_string(),
            id: RequestId::String("ts1".to_string()),
            method: "client/registerCapability".to_string(),
            params: Some(serde_json::json!({ "registrations": [] })),
        };

        let response = LspClient::server_request_response(request);

        assert_eq!(response.id, RequestId::String("ts1".to_string()));
        assert_eq!(response.result, Some(Value::Null));
        assert!(response.error.is_none());
    }

    #[test]
    fn test_workspace_configuration_request_returns_null_per_item() {
        let result = LspClient::workspace_configuration_result(Some(&serde_json::json!({
            "items": [{ "section": "typescript" }, { "section": "editor" }]
        })));

        assert_eq!(result, serde_json::json!([null, null]));
    }

    #[test]
    fn test_unknown_server_request_returns_method_not_found() {
        let request = JsonRpcRequest {
            jsonrpc: JSONRPC_VERSION.to_string(),
            id: RequestId::String("unknown-1".to_string()),
            method: "custom/request".to_string(),
            params: None,
        };

        let response = LspClient::server_request_response(request);

        assert!(response.result.is_none());
        match response.error {
            Some(error) => {
                assert_eq!(error.code, -32601);
                assert_eq!(error.message, "Unhandled server request: custom/request");
            }
            None => panic!("unknown request should return error"),
        }
    }

    #[tokio::test]
    async fn test_null_response_handling() {
        use crate::lsp::types::{JsonRpcResponse, RequestId};

        let pending_requests: Arc<Mutex<PendingRequests>> = Arc::new(Mutex::new(HashMap::new()));

        let (response_tx, response_rx) = oneshot::channel::<Result<Value>>();

        pending_requests
            .lock()
            .await
            .insert(RequestId::Number(1), response_tx);

        let null_response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: RequestId::Number(1),
            result: None,
            error: None,
        };

        let sender = pending_requests.lock().await.remove(&null_response.id);
        if let Some(sender) = sender {
            let _ = sender.send(Ok(Value::Null));
        }

        let timeout_result =
            tokio::time::timeout(tokio::time::Duration::from_millis(100), response_rx).await;

        assert!(timeout_result.is_ok(), "Should not timeout");

        let channel_result = timeout_result.unwrap();
        assert!(
            channel_result.is_ok(),
            "Channel should not be closed: {:?}",
            channel_result.err()
        );

        let response = channel_result.unwrap();
        assert!(
            response.is_ok(),
            "Should receive Ok(Value::Null), not Err: {:?}",
            response.err()
        );

        let value = response.unwrap();
        assert_eq!(value, Value::Null, "Should receive Value::Null");
    }

    #[tokio::test]
    async fn test_error_response_handling() {
        use crate::lsp::types::{JsonRpcError, JsonRpcResponse, RequestId};

        let pending_requests: Arc<Mutex<PendingRequests>> = Arc::new(Mutex::new(HashMap::new()));
        let (response_tx, response_rx) = oneshot::channel::<Result<Value>>();

        pending_requests
            .lock()
            .await
            .insert(RequestId::Number(1), response_tx);

        let error_response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: RequestId::Number(1),
            result: None,
            error: Some(JsonRpcError {
                code: -32601,
                message: "Method not found".to_string(),
                data: None,
            }),
        };

        let sender = pending_requests.lock().await.remove(&error_response.id);
        if let Some(sender) = sender
            && let Some(error) = error_response.error
        {
            let _ = sender.send(Err(Error::LspServerError {
                code: error.code,
                message: error.message,
                data: error.data,
            }));
        }

        let result = response_rx.await.unwrap();
        assert!(result.is_err(), "Should receive error");

        if let Err(Error::LspServerError { code, message, .. }) = result {
            assert_eq!(code, -32601);
            assert_eq!(message, "Method not found");
        } else {
            panic!("Expected LspServerError");
        }
    }

    #[tokio::test]
    async fn test_unknown_request_id() {
        use crate::lsp::types::{JsonRpcResponse, RequestId};

        let pending_requests: Arc<Mutex<PendingRequests>> = Arc::new(Mutex::new(HashMap::new()));

        let response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: RequestId::Number(999),
            result: Some(Value::Null),
            error: None,
        };

        let sender = pending_requests.lock().await.remove(&response.id);
        assert!(sender.is_none(), "Should not find sender for unknown ID");
    }

    #[tokio::test]
    async fn test_long_error_message_truncation() {
        use crate::lsp::types::{JsonRpcError, JsonRpcResponse, RequestId};

        let pending_requests: Arc<Mutex<PendingRequests>> = Arc::new(Mutex::new(HashMap::new()));
        let (response_tx, response_rx) = oneshot::channel::<Result<Value>>();

        pending_requests
            .lock()
            .await
            .insert(RequestId::Number(1), response_tx);

        let long_message = "x".repeat(250);
        let error_response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: RequestId::Number(1),
            result: None,
            error: Some(JsonRpcError {
                code: -32700,
                message: long_message.clone(),
                data: None,
            }),
        };

        let sender = pending_requests.lock().await.remove(&error_response.id);
        if let Some(sender) = sender
            && let Some(error) = error_response.error
        {
            let _ = sender.send(Err(Error::LspServerError {
                code: error.code,
                message: error.message,
                data: error.data,
            }));
        }

        let result = response_rx.await.unwrap();
        assert!(result.is_err());

        if let Err(Error::LspServerError { code, message, .. }) = result {
            assert_eq!(code, -32700);
            assert_eq!(
                message.len(),
                250,
                "Full message should be preserved in Error"
            );
        } else {
            panic!("Expected LspServerError");
        }
    }

    #[tokio::test]
    async fn test_concurrent_request_ids() {
        let counter = Arc::new(AtomicI64::new(1));

        let counter1 = Arc::clone(&counter);
        let counter2 = Arc::clone(&counter);
        let counter3 = Arc::clone(&counter);

        let handles = vec![
            tokio::spawn(async move { counter1.fetch_add(1, Ordering::SeqCst) }),
            tokio::spawn(async move { counter2.fetch_add(1, Ordering::SeqCst) }),
            tokio::spawn(async move { counter3.fetch_add(1, Ordering::SeqCst) }),
        ];

        let mut ids = Vec::new();
        for handle in handles {
            ids.push(handle.await.unwrap());
        }

        ids.sort_unstable();
        assert_eq!(ids, vec![1, 2, 3], "IDs should be unique and sequential");
    }

    #[test]
    fn test_jsonrpc_version_constant() {
        assert_eq!(JSONRPC_VERSION, "2.0");
    }
}