use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use anyhow::{Context, bail};
use crossbeam_channel::Sender;
use serde_json::{Value, json};
use tokio::io::{AsyncRead, AsyncWrite, BufReader};
use tokio::process::Child;
use tokio::sync::mpsc;
use crate::codec;
use crate::config::ServerConfig;
use crate::event::{LspEvent, RpcError, ServerKey};
type PendingMap = Arc<Mutex<HashMap<i64, i64>>>;
const INITIALIZE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
const SHUTDOWN_GRACE: std::time::Duration = std::time::Duration::from_secs(3);
pub struct Server {
pub key: ServerKey,
pub capabilities: Value,
stdin_tx: mpsc::UnboundedSender<Vec<u8>>,
next_request_id: i64,
pending: PendingMap,
kill_tx: Option<tokio::sync::oneshot::Sender<()>>,
wait_handle: Option<tokio::task::JoinHandle<()>>,
}
impl Server {
pub async fn spawn(
key: ServerKey,
cmd: &ServerConfig,
evt_tx: Sender<LspEvent>,
) -> anyhow::Result<Self> {
cmd.validate(&key.language).map_err(anyhow::Error::msg)?;
let mut child = tokio::process::Command::new(&cmd.command)
.args(&cmd.args)
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()
.with_context(|| format!("failed to spawn LSP server {:?}", cmd.command))?;
let stdin = child.stdin.take().context("no stdin")?;
let stdout = child.stdout.take().context("no stdout")?;
let stderr = child.stderr.take().context("no stderr")?;
let (stdin_tx, stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
tokio::spawn(stdin_task(stdin_rx, stdin));
tokio::spawn(stderr_task(stderr, key.language.clone()));
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let init_options = cmd
.initialization_options
.clone()
.or_else(|| default_init_options(&cmd.command));
let capabilities = match initialize_handshake(
&key,
&stdin_tx,
stdout,
evt_tx.clone(),
pending.clone(),
init_options.as_ref(),
)
.await
{
Ok(caps) => caps,
Err(e) => {
let _ = child.start_kill();
let _ = child.wait().await;
return Err(e);
}
};
let (kill_tx, wait_handle) = spawn_wait_task(child, key.clone(), evt_tx);
Ok(Self {
key,
capabilities,
stdin_tx,
next_request_id: 1,
pending,
kill_tx: Some(kill_tx),
wait_handle: Some(wait_handle),
})
}
pub fn send_notification(&mut self, method: &str, params: Value) {
self.enqueue(&rpc_envelope(None, method, params));
}
pub fn send_request(&mut self, app_id: i64, method: &str, params: Value) {
let id = self.next_request_id;
self.next_request_id += 1;
if let Ok(mut map) = self.pending.lock() {
map.insert(id, app_id);
}
self.enqueue(&rpc_envelope(Some(id), method, params));
}
pub async fn shutdown(mut self) {
self.send_request(-1, "shutdown", Value::Null);
tracing::debug!(key = ?self.key, "sent shutdown request");
self.send_notification("exit", Value::Null);
drop(self.stdin_tx);
if let (Some(mut handle), Some(kill_tx)) = (self.wait_handle.take(), self.kill_tx.take()) {
match tokio::time::timeout(SHUTDOWN_GRACE, &mut handle).await {
Ok(_) => {} Err(_) => {
tracing::warn!(
key = ?self.key,
"LSP server did not exit within {SHUTDOWN_GRACE:?}; force-killing"
);
let _ = kill_tx.send(()); let _ = handle.await; }
}
}
}
fn enqueue(&self, msg: &Value) {
match serde_json::to_vec(msg) {
Ok(bytes) => {
let _ = self.stdin_tx.send(bytes);
}
Err(e) => {
tracing::warn!("failed to serialize JSON-RPC message: {e}");
}
}
}
}
fn rpc_envelope(id: Option<i64>, method: &str, params: Value) -> Value {
let mut msg = serde_json::Map::new();
msg.insert("jsonrpc".to_string(), Value::from("2.0"));
if let Some(id) = id {
msg.insert("id".to_string(), Value::from(id));
}
msg.insert("method".to_string(), Value::from(method));
msg.insert("params".to_string(), params);
Value::Object(msg)
}
fn default_init_options(command: &str) -> Option<Value> {
let stem = std::path::Path::new(command)
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or(command);
if stem == "rust-analyzer" {
Some(json!({
"check": { "command": "clippy" },
"checkOnSave": { "command": "clippy" },
}))
} else {
None
}
}
async fn initialize_handshake(
key: &ServerKey,
stdin_tx: &mpsc::UnboundedSender<Vec<u8>>,
stdout: impl AsyncRead + Unpin + Send + 'static,
evt_tx: Sender<LspEvent>,
pending: PendingMap,
init_options: Option<&Value>,
) -> anyhow::Result<Value> {
let root_uri = crate::uri::from_path(&key.root).map_err(|_| {
anyhow::anyhow!(
"cannot convert workspace root {:?} to file:// URI",
key.root
)
})?;
let mut params = json!({
"processId": std::process::id(),
"clientInfo": { "name": "hjkl", "version": env!("CARGO_PKG_VERSION") },
"rootUri": root_uri.as_str(),
"capabilities": {
"general": {
"positionEncodings": ["utf-8", "utf-16"],
},
"textDocument": {
"synchronization": {
"dynamicRegistration": false,
"willSave": false,
"willSaveWaitUntil": false,
"didSave": true,
}
},
"workspace": {}
},
});
if let (Some(opts), Some(obj)) = (init_options, params.as_object_mut()) {
obj.insert("initializationOptions".to_string(), opts.clone());
}
let init_msg = rpc_envelope(Some(0), "initialize", params);
let bytes = serde_json::to_vec(&init_msg)?;
stdin_tx.send(bytes).ok();
let mut reader = BufReader::with_capacity(256 * 1024, stdout);
let capabilities = tokio::time::timeout(INITIALIZE_TIMEOUT, async {
loop {
let raw = codec::read_message(&mut reader).await?.ok_or_else(|| {
anyhow::anyhow!("server closed stdout before initialize response")
})?;
let val: Value = serde_json::from_slice(&raw)?;
if val.get("id").and_then(Value::as_i64) == Some(0) && val.get("method").is_none() {
if let Some(err) = val.get("error") {
bail!("initialize error: {err}");
}
let caps = val
.get("result")
.and_then(|r| r.get("capabilities"))
.cloned()
.unwrap_or(Value::Null);
break Ok::<Value, anyhow::Error>(caps);
}
tracing::debug!(
key = ?key,
"received server message before initialize response; ignoring"
);
}
})
.await
.map_err(|_| {
anyhow::anyhow!(
"initialize handshake timed out after {}s",
INITIALIZE_TIMEOUT.as_secs()
)
})??;
let init_notif = json!({
"jsonrpc": "2.0",
"method": "initialized",
"params": {},
});
let bytes = serde_json::to_vec(&init_notif)?;
stdin_tx.send(bytes).ok();
tracing::info!(key = ?key, "LSP server initialized");
let _ = evt_tx.send(LspEvent::ServerInitialized {
key: key.clone(),
capabilities: capabilities.clone(),
});
let key_clone = key.clone();
tokio::spawn(stdout_task(
reader,
key_clone,
evt_tx,
pending,
stdin_tx.clone(),
));
Ok(capabilities)
}
pub async fn spawn_from_io<R, W>(
key: ServerKey,
stdin_writer: W,
stdout_reader: R,
evt_tx: Sender<LspEvent>,
) -> anyhow::Result<Server>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let (stdin_tx, stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
tokio::spawn(stdin_task(stdin_rx, stdin_writer));
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let capabilities = initialize_handshake(
&key,
&stdin_tx,
stdout_reader,
evt_tx,
pending.clone(),
None,
)
.await?;
Ok(Server {
key,
capabilities,
stdin_tx,
next_request_id: 1,
pending,
kill_tx: None,
wait_handle: None,
})
}
async fn stdin_task<W: AsyncWrite + Unpin>(mut rx: mpsc::UnboundedReceiver<Vec<u8>>, mut w: W) {
while let Some(bytes) = rx.recv().await {
if let Err(e) = codec::write_message(&mut w, &bytes).await {
tracing::debug!("LSP stdin write error: {e}");
break;
}
}
}
async fn stdout_task<R: AsyncRead + Unpin>(
mut reader: BufReader<R>,
key: ServerKey,
evt_tx: Sender<LspEvent>,
pending: PendingMap,
stdin_tx: mpsc::UnboundedSender<Vec<u8>>,
) {
loop {
let raw = match codec::read_message(&mut reader).await {
Ok(Some(r)) => r,
Ok(None) => {
tracing::debug!(key = ?key, "LSP stdout closed (clean EOF)");
break;
}
Err(e) => {
tracing::warn!(key = ?key, "LSP stdout read error: {e}");
break;
}
};
let val: Value = match serde_json::from_slice(&raw) {
Ok(v) => v,
Err(e) => {
tracing::warn!(key = ?key, "LSP: failed to parse JSON frame: {e}");
continue;
}
};
dispatch_message(&key, &val, &evt_tx, &pending, &stdin_tx);
}
}
fn dispatch_message(
key: &ServerKey,
val: &Value,
evt_tx: &Sender<LspEvent>,
pending: &PendingMap,
stdin_tx: &mpsc::UnboundedSender<Vec<u8>>,
) {
let has_id = val.get("id").is_some();
let has_method = val.get("method").is_some();
if has_id && !has_method {
let Some(jsonrpc_id) = val.get("id").and_then(Value::as_i64) else {
tracing::warn!(key = ?key, "LSP response with non-integer id; ignoring");
return;
};
let Some(app_id) = pending.lock().ok().and_then(|mut m| m.remove(&jsonrpc_id)) else {
tracing::debug!(key = ?key, jsonrpc_id, "LSP response for unknown id; ignoring");
return;
};
let result = if let Some(err) = val.get("error") {
let code = err.get("code").and_then(Value::as_i64).unwrap_or(-1);
let message = err
.get("message")
.and_then(Value::as_str)
.unwrap_or("unknown error")
.to_string();
Err(RpcError { code, message })
} else {
Ok(val.get("result").cloned().unwrap_or(Value::Null))
};
let _ = evt_tx.send(LspEvent::Response {
request_id: app_id,
result,
});
} else if has_method {
if has_id {
let id = val.get("id").cloned().unwrap_or(Value::Null);
let method = val
.get("method")
.and_then(Value::as_str)
.unwrap_or("<unknown>");
match method {
"workspace/configuration" => {
let count = val
.get("params")
.and_then(|p| p.get("items"))
.and_then(Value::as_array)
.map_or(0, |a| a.len());
send_response(stdin_tx, id, Value::Array(vec![Value::Null; count]));
}
"client/registerCapability"
| "client/unregisterCapability"
| "window/workDoneProgress/create" => {
send_response(stdin_tx, id, Value::Null);
}
"workspace/applyEdit" => {
send_response(stdin_tx, id, json!({ "applied": false }));
}
_ => {
send_error_response(stdin_tx, id, -32601, "method not supported");
}
}
tracing::debug!(key = ?key, method, "LSP server-initiated request auto-answered");
} else {
let method = val
.get("method")
.and_then(Value::as_str)
.unwrap_or("<unknown>")
.to_string();
let params = val.get("params").cloned().unwrap_or(Value::Null);
tracing::debug!(key = ?key, method, "LSP notification received");
let _ = evt_tx.send(LspEvent::Notification {
key: key.clone(),
method,
params,
});
}
} else {
tracing::warn!(key = ?key, "LSP: unrecognized message shape; ignoring");
}
}
fn send_response(stdin_tx: &mpsc::UnboundedSender<Vec<u8>>, id: Value, result: Value) {
let mut msg = serde_json::Map::new();
msg.insert("jsonrpc".to_string(), Value::from("2.0"));
msg.insert("id".to_string(), id);
msg.insert("result".to_string(), result);
if let Ok(bytes) = serde_json::to_vec(&Value::Object(msg)) {
let _ = stdin_tx.send(bytes);
}
}
fn send_error_response(
stdin_tx: &mpsc::UnboundedSender<Vec<u8>>,
id: Value,
code: i64,
message: &str,
) {
let mut msg = serde_json::Map::new();
msg.insert("jsonrpc".to_string(), Value::from("2.0"));
msg.insert("id".to_string(), id);
msg.insert(
"error".to_string(),
json!({ "code": code, "message": message }),
);
if let Ok(bytes) = serde_json::to_vec(&Value::Object(msg)) {
let _ = stdin_tx.send(bytes);
}
}
async fn stderr_task<R: tokio::io::AsyncRead + Unpin>(stderr: R, lang: String) {
use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
const MAX_LINE_BYTES: u64 = 8 * 1024;
let mut reader = BufReader::new(stderr);
let mut buf = Vec::new();
loop {
buf.clear();
let n = {
let mut limited = (&mut reader).take(MAX_LINE_BYTES);
match limited.read_until(b'\n', &mut buf).await {
Ok(n) => n,
Err(e) => {
tracing::debug!(lang, "LSP stderr read error: {e}");
break;
}
}
};
if n == 0 {
break;
}
let text = String::from_utf8_lossy(&buf);
let trimmed = text.trim_end();
if !trimmed.is_empty() {
tracing::warn!(lang, "LSP stderr: {trimmed}");
}
}
}
fn spawn_wait_task(
mut child: Child,
key: ServerKey,
evt_tx: Sender<LspEvent>,
) -> (
tokio::sync::oneshot::Sender<()>,
tokio::task::JoinHandle<()>,
) {
let (kill_tx, kill_rx) = tokio::sync::oneshot::channel::<()>();
let handle = tokio::spawn(async move {
tokio::select! {
res = child.wait() => {
match res {
Ok(status) => {
tracing::info!(key = ?key, ?status, "LSP server exited");
let _ = evt_tx.send(LspEvent::ServerExited { key, status });
}
Err(e) => {
tracing::warn!(key = ?key, "error waiting for LSP server: {e}");
}
}
}
_ = kill_rx => {
let _ = child.start_kill();
match child.wait().await {
Ok(status) => {
tracing::info!(key = ?key, ?status, "LSP server force-killed on shutdown");
let _ = evt_tx.send(LspEvent::ServerExited { key, status });
}
Err(e) => {
tracing::warn!(key = ?key, "error waiting for force-killed LSP server: {e}");
}
}
}
}
});
(kill_tx, handle)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rpc_envelope_matches_json_macro_shape() {
let params = json!({
"textDocument": { "uri": "file:///tmp/a.rs", "version": 3 },
"contentChanges": [{ "text": "fn main() {}\n" }],
});
let notif = rpc_envelope(None, "textDocument/didChange", params.clone());
let notif_old = json!({
"jsonrpc": "2.0",
"method": "textDocument/didChange",
"params": params,
});
assert_eq!(notif, notif_old);
assert_eq!(
serde_json::to_vec(¬if).unwrap(),
serde_json::to_vec(¬if_old).unwrap(),
"wire bytes must be identical (key order included)"
);
assert_eq!(
notif["params"]["contentChanges"][0]["text"],
"fn main() {}\n"
);
assert!(notif.get("id").is_none(), "notifications carry no id");
let req = rpc_envelope(Some(7), "textDocument/definition", params.clone());
let req_old = json!({
"jsonrpc": "2.0",
"id": 7,
"method": "textDocument/definition",
"params": params,
});
assert_eq!(req, req_old);
assert_eq!(
serde_json::to_vec(&req).unwrap(),
serde_json::to_vec(&req_old).unwrap()
);
let null_params = rpc_envelope(None, "exit", Value::Null);
assert_eq!(
null_params,
json!({ "jsonrpc": "2.0", "method": "exit", "params": null })
);
}
}