ralph/

lsp_client.rs

/// JSON-RPC over stdio LSP client.
///
/// Lazy — the server process is spawned on the first tool call, not at
/// construction time.  This avoids startup latency when LSP tools are never
/// used in a session.
use crate::errors::{RalphError, Result};
use serde_json::{json, Value};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::ChildStdin;

// ── Public types ──────────────────────────────────────────────────────────────

pub struct LspClient {
    language: String,
    server_cmd: String,
    workspace: PathBuf,
    inner: Option<RunningLsp>,
}

struct RunningLsp {
    /// Keep the child alive for the session.
    _child: tokio::process::Child,
    stdin: ChildStdin,
    reader: BufReader<tokio::process::ChildStdout>,
    next_id: u64,
    /// Files already sent via textDocument/didOpen.
    opened: HashSet<PathBuf>,
}

// ── LspClient impl ────────────────────────────────────────────────────────────

impl LspClient {
    pub fn new(language: &str, workspace: &Path) -> Self {
        Self {
            language: language.to_lowercase(),
            server_cmd: server_cmd_for(language),
            workspace: workspace.to_path_buf(),
            inner: None,
        }
    }

    pub fn language(&self) -> &str {
        &self.language
    }

    /// Returns true if this client handles files with the given extension.
    pub fn handles_extension(&self, ext: &str) -> bool {
        extensions_for_language(&self.language).contains(&ext)
    }

    /// Spawn the server and run the initialize handshake (once).
    async fn ensure_running(&mut self) -> Result<()> {
        if self.inner.is_some() {
            return Ok(());
        }

        let parts: Vec<&str> = self.server_cmd.split_whitespace().collect();
        if parts.is_empty() {
            return Err(RalphError::ToolFailed {
                tool: "lsp".to_string(),
                message: format!(
                    "No server command configured for language '{}'",
                    self.language
                ),
            });
        }

        let mut child = tokio::process::Command::new(parts[0])
            .args(&parts[1..])
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::null())
            .spawn()
            .map_err(|e| RalphError::ToolFailed {
                tool: "lsp".to_string(),
                message: format!(
                    "Failed to start '{}': {}. Is it installed and on PATH?",
                    parts[0], e
                ),
            })?;

        let stdin = child.stdin.take().ok_or_else(|| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: format!("Failed to capture stdin for '{}' LSP server", parts[0]),
        })?;
        let stdout = child.stdout.take().ok_or_else(|| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: format!("Failed to capture stdout for '{}' LSP server", parts[0]),
        })?;
        let reader = BufReader::new(stdout);

        let mut r = RunningLsp {
            _child: child,
            stdin,
            reader,
            next_id: 1,
            opened: HashSet::new(),
        };

        // initialize ↔ InitializeResult
        let workspace_uri = path_to_uri(&self.workspace);
        tokio::time::timeout(
            std::time::Duration::from_secs(30),
            r.request(
                "initialize",
                json!({
                    "processId": std::process::id(),
                    "rootUri": workspace_uri,
                    "rootPath": self.workspace.display().to_string(),
                    "capabilities": {
                        "textDocument": {
                            "definition": {},
                            "references": {},
                            "hover": { "contentFormat": ["plaintext", "markdown"] },
                        }
                    },
                    "workspaceFolders": [{ "uri": workspace_uri, "name": "workspace" }],
                }),
            ),
        )
        .await
        .map_err(|_| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: "LSP server timed out during initialization (30 s)".to_string(),
        })??;

        r.notify("initialized", json!({})).await?;
        self.inner = Some(r);
        Ok(())
    }

    /// Jump to the definition of the symbol at `file:line:col` (1-based).
    pub async fn go_to_definition(&mut self, file: &Path, line: u32, col: u32) -> Result<String> {
        self.ensure_running().await?;
        let r = self.inner.as_mut().unwrap();
        r.open_if_needed(file, &self.workspace).await?;
        let uri = path_to_uri(&self.workspace.join(file));
        let result = r
            .request(
                "textDocument/definition",
                json!({
                    "textDocument": { "uri": uri },
                    "position": lsp_pos(line, col),
                }),
            )
            .await?;
        Ok(format_locations(&result, &self.workspace))
    }

    /// Find all references to the symbol at `file:line:col` (1-based).
    pub async fn find_references(&mut self, file: &Path, line: u32, col: u32) -> Result<String> {
        self.ensure_running().await?;
        let r = self.inner.as_mut().unwrap();
        r.open_if_needed(file, &self.workspace).await?;
        let uri = path_to_uri(&self.workspace.join(file));
        let result = r
            .request(
                "textDocument/references",
                json!({
                    "textDocument": { "uri": uri },
                    "position": lsp_pos(line, col),
                    "context": { "includeDeclaration": false },
                }),
            )
            .await?;
        Ok(format_locations(&result, &self.workspace))
    }

    /// Return hover info (type, docs) for the symbol at `file:line:col` (1-based).
    pub async fn hover(&mut self, file: &Path, line: u32, col: u32) -> Result<String> {
        self.ensure_running().await?;
        let r = self.inner.as_mut().unwrap();
        r.open_if_needed(file, &self.workspace).await?;
        let uri = path_to_uri(&self.workspace.join(file));
        let result = r
            .request(
                "textDocument/hover",
                json!({
                    "textDocument": { "uri": uri },
                    "position": lsp_pos(line, col),
                }),
            )
            .await?;
        Ok(format_hover(&result))
    }

    /// Graceful shutdown.  Called when the session ends.
    pub async fn shutdown(&mut self) {
        if let Some(r) = self.inner.as_mut() {
            let _ = r.request("shutdown", json!(null)).await;
            let _ = r.notify("exit", json!(null)).await;
        }
        self.inner = None;
    }
}

// ── RunningLsp impl ───────────────────────────────────────────────────────────

impl RunningLsp {
    async fn send_raw(&mut self, body: &str) -> Result<()> {
        let header = format!("Content-Length: {}\r\n\r\n", body.len());
        self.stdin
            .write_all(header.as_bytes())
            .await
            .map_err(io_err)?;
        self.stdin
            .write_all(body.as_bytes())
            .await
            .map_err(io_err)?;
        self.stdin.flush().await.map_err(io_err)?;
        Ok(())
    }

    async fn read_message(&mut self) -> Result<Value> {
        // Read HTTP-style headers until blank line
        let mut content_length: Option<usize> = None;
        loop {
            let mut line = String::new();
            let n = self.reader.read_line(&mut line).await.map_err(io_err)?;
            if n == 0 {
                return Err(RalphError::ToolFailed {
                    tool: "lsp".to_string(),
                    message: "LSP server closed its stdout unexpectedly".to_string(),
                });
            }
            let trimmed = line.trim_end_matches(|c| c == '\r' || c == '\n');
            if trimmed.is_empty() {
                break; // blank line — end of headers
            }
            if let Some(val) = trimmed.strip_prefix("Content-Length: ") {
                content_length = val.trim().parse().ok();
            }
        }

        let len = content_length.ok_or_else(|| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: "LSP message missing Content-Length header".to_string(),
        })?;

        let mut buf = vec![0u8; len];
        self.reader.read_exact(&mut buf).await.map_err(io_err)?;

        serde_json::from_slice(&buf).map_err(|e| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: format!("LSP JSON parse error: {}", e),
        })
    }

    /// Send a request and return its result, discarding notifications and
    /// replying with `null` to any server-initiated requests in between.
    async fn request(&mut self, method: &str, params: Value) -> Result<Value> {
        let id = self.next_id;
        self.next_id += 1;

        let body = serde_json::to_string(&json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params,
        }))
        .map_err(|e| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: format!("JSON encode error: {}", e),
        })?;
        self.send_raw(&body).await?;

        loop {
            let msg = self.read_message().await?;

            match msg.get("id") {
                // Notification (no id) — discard
                None => continue,

                Some(msg_id) => {
                    // Server-to-client request (has both "id" and "method") —
                    // respond with null so the server doesn't block.
                    if msg.get("method").is_some() {
                        let resp = serde_json::to_string(&json!({
                            "jsonrpc": "2.0",
                            "id": msg_id,
                            "result": null,
                        }))
                        .map_err(|e| RalphError::ToolFailed {
                            tool: "lsp".to_string(),
                            message: format!("JSON encode error: {}", e),
                        })?;
                        self.send_raw(&resp).await?;
                        continue;
                    }

                    // Our response
                    if *msg_id == json!(id) {
                        if let Some(err) = msg.get("error") {
                            return Err(RalphError::ToolFailed {
                                tool: "lsp".to_string(),
                                message: format!("LSP error response: {}", err),
                            });
                        }
                        return Ok(msg.get("result").cloned().unwrap_or(Value::Null));
                    }
                    // Different id — skip (shouldn't happen in sequential usage)
                }
            }
        }
    }

    async fn notify(&mut self, method: &str, params: Value) -> Result<()> {
        let body = serde_json::to_string(&json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        }))
        .map_err(|e| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: format!("JSON encode error: {}", e),
        })?;
        self.send_raw(&body).await
    }

    /// Send textDocument/didOpen for `file` (relative to `workspace`) if not
    /// already opened this session.
    async fn open_if_needed(&mut self, file: &Path, workspace: &Path) -> Result<()> {
        let full = workspace.join(file);
        if self.opened.contains(&full) {
            return Ok(());
        }
        let content = std::fs::read_to_string(&full).map_err(|e| RalphError::ToolFailed {
            tool: "lsp".to_string(),
            message: format!("Cannot read '{}' for LSP: {}", full.display(), e),
        })?;
        let lang_id = ext_to_language_id(file.extension().and_then(|e| e.to_str()).unwrap_or(""));
        self.notify(
            "textDocument/didOpen",
            json!({
                "textDocument": {
                    "uri": path_to_uri(&full),
                    "languageId": lang_id,
                    "version": 1,
                    "text": content,
                }
            }),
        )
        .await?;
        self.opened.insert(full);
        Ok(())
    }
}

// ── Public helpers ────────────────────────────────────────────────────────────

/// Return the server command string for a language name/alias.
pub fn server_cmd_for(language: &str) -> String {
    match language.to_lowercase().as_str() {
        "rust" | "rs" => "rust-analyzer".to_string(),
        "ts" | "typescript" => "typescript-language-server --stdio".to_string(),
        "js" | "javascript" => "typescript-language-server --stdio".to_string(),
        "python" | "py" => "pylsp".to_string(),
        "go" => "gopls".to_string(),
        "java" => "jdtls".to_string(),
        "c" => "clangd".to_string(),
        "cpp" | "c++" | "cxx" => "clangd".to_string(),
        "kotlin" | "kt" => "kotlin-language-server".to_string(),
        other => other.to_string(),
    }
}

/// Return the file extensions handled by a language.
pub fn extensions_for_language(language: &str) -> &'static [&'static str] {
    match language.to_lowercase().as_str() {
        "rust" | "rs" => &["rs"],
        "ts" | "typescript" => &["ts", "tsx"],
        "js" | "javascript" => &["js", "jsx"],
        "python" | "py" => &["py"],
        "go" => &["go"],
        "java" => &["java"],
        "c" => &["c", "h"],
        "cpp" | "c++" | "cxx" => &["cpp", "cc", "cxx", "hpp", "hxx", "h"],
        "kotlin" | "kt" => &["kt"],
        _ => &[],
    }
}

// ── Private helpers ───────────────────────────────────────────────────────────

fn path_to_uri(path: &Path) -> String {
    let canonical = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
    let s = canonical.display().to_string();
    if s.starts_with('/') {
        format!("file://{}", s)
    } else {
        // Windows: C:\foo → file:///C:/foo
        format!("file:///{}", s.replace(std::path::MAIN_SEPARATOR, "/"))
    }
}

fn uri_to_rel_path(uri: &str, workspace: &Path) -> String {
    // Strip "file://" and reconstruct absolute path
    let without_scheme = uri.strip_prefix("file://").unwrap_or(uri);
    let abs = if without_scheme.starts_with('/') {
        PathBuf::from(without_scheme)
    } else {
        PathBuf::from(format!("/{}", without_scheme))
    };
    let ws = workspace
        .canonicalize()
        .unwrap_or_else(|_| workspace.to_path_buf());
    abs.strip_prefix(&ws)
        .map(|p| p.display().to_string())
        .unwrap_or_else(|_| abs.display().to_string())
}

/// Convert 1-based line/col to an LSP Position (0-based).
fn lsp_pos(line: u32, col: u32) -> Value {
    json!({
        "line": line.saturating_sub(1),
        "character": col.saturating_sub(1),
    })
}

fn format_locations(result: &Value, workspace: &Path) -> String {
    if result.is_null() {
        return "Not found.".to_string();
    }
    let locs: Vec<&Value> = match result.as_array() {
        Some(arr) => arr.iter().collect(),
        None => vec![result],
    };
    if locs.is_empty() {
        return "Not found.".to_string();
    }
    let mut out = String::new();
    for loc in locs {
        // Supports both Location and LocationLink shapes
        let uri = loc
            .get("uri")
            .or_else(|| loc.get("targetUri"))
            .and_then(|v| v.as_str())
            .unwrap_or("?");
        let range = loc
            .get("range")
            .or_else(|| loc.get("targetSelectionRange"))
            .or_else(|| loc.get("targetRange"));
        let line = range
            .and_then(|r| r.get("start"))
            .and_then(|s| s.get("line"))
            .and_then(|l| l.as_u64())
            .map(|l| l + 1) // 0-based → 1-based
            .unwrap_or(0);
        out.push_str(&format!("  {}:{}\n", uri_to_rel_path(uri, workspace), line));
    }
    out.trim_end().to_string()
}

fn format_hover(result: &Value) -> String {
    if result.is_null() {
        return "No hover information available.".to_string();
    }
    match result.get("contents") {
        Some(c) => extract_hover_text(c),
        None => "No hover information available.".to_string(),
    }
}

fn extract_hover_text(v: &Value) -> String {
    match v {
        Value::String(s) => s.clone(),
        Value::Object(obj) => obj
            .get("value")
            .and_then(|v| v.as_str())
            .unwrap_or("(empty)")
            .to_string(),
        Value::Array(arr) => arr
            .iter()
            .map(extract_hover_text)
            .collect::<Vec<_>>()
            .join("\n"),
        _ => v.to_string(),
    }
}

fn ext_to_language_id(ext: &str) -> &'static str {
    match ext {
        "rs" => "rust",
        "ts" | "tsx" => "typescript",
        "js" | "jsx" => "javascript",
        "py" => "python",
        "go" => "go",
        "java" => "java",
        "c" | "h" => "c",
        "cpp" | "cc" | "cxx" | "hpp" | "hxx" => "cpp",
        "kt" => "kotlin",
        _ => "plaintext",
    }
}

fn io_err(e: std::io::Error) -> RalphError {
    RalphError::ToolFailed {
        tool: "lsp".to_string(),
        message: format!("LSP I/O error: {}", e),
    }
}