use crate::shell::child_pipe_writer::ChildPipeWriter;
use crate::{
cancellation::AgentCancellation,
config::LspServerConfig,
mcp::jsonrpc::{
ErrorData, JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest,
METHOD_NOT_FOUND, RequestId,
},
process_environment::{self, SubprocessEnvProfile},
tools::process::{
PIPE_READER_JOIN_TIMEOUT, PipeReaderHandle, recv_pipe_reader_with_timeout,
spawn_bounded_pipe_reader, terminate_child_tree_and_wait,
},
};
use anyhow::{Context, bail};
use serde_json::{Value, json};
use std::{
collections::HashMap,
io::{BufRead, BufReader, Write},
path::Path,
process::{Child, Command, Stdio},
sync::{
Arc, Mutex,
atomic::{AtomicBool, AtomicI64, Ordering},
mpsc::{self, Receiver, SyncSender},
},
thread::{self, JoinHandle},
time::{Duration, Instant},
};
const MAX_LSP_HEADER_BYTES: usize = 8192;
const MAX_LSP_BODY_BYTES: usize = 1_048_576;
const MAX_LSP_STDERR_BYTES: usize = 8192;
type PendingMap = Arc<Mutex<HashMap<RequestId, SyncSender<anyhow::Result<Value>>>>>;
pub(crate) struct LspStdioTransport {
writer: Arc<ChildPipeWriter>,
child: Arc<Mutex<Child>>,
reader: Option<JoinHandle<()>>,
reader_done: Option<Receiver<()>>,
stderr_receiver: Option<PipeReaderHandle>,
pending: PendingMap,
notifications: Receiver<JsonRpcNotification>,
next_id: AtomicI64,
closed: Arc<AtomicBool>,
}
#[derive(Clone)]
struct LspCleanupHandle {
child: Arc<Mutex<Child>>,
}
struct PendingRequestCleanup {
pending: PendingMap,
id: RequestId,
}
impl Drop for PendingRequestCleanup {
fn drop(&mut self) {
if let Ok(mut pending) = self.pending.lock() {
pending.remove(&self.id);
}
}
}
impl LspCleanupHandle {
fn terminate(&self) {
if let Ok(mut child) = self.child.lock() {
let _ = terminate_child_tree_and_wait(&mut child);
}
}
}
impl LspStdioTransport {
pub(crate) fn spawn(config: &LspServerConfig, cwd: &Path) -> anyhow::Result<Self> {
if config.command.trim().is_empty() {
bail!("LSP stdio command must not be empty");
}
let mut command = Command::new(&config.command);
command
.args(&config.args)
.current_dir(cwd)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
process_environment::apply_profile(&mut command, SubprocessEnvProfile::Lsp);
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
command.process_group(0);
}
let mut child = command.spawn().context("spawning LSP server")?;
let stdin = child.stdin.take().context("LSP stdin pipe unavailable")?;
let stdout = child.stdout.take().context("LSP stdout pipe unavailable")?;
let stderr = child.stderr.take().context("LSP stderr pipe unavailable")?;
let writer = match ChildPipeWriter::new(stdin) {
Ok(writer) => Arc::new(writer),
Err(error) => {
let _ = terminate_child_tree_and_wait(&mut child);
return Err(error.into());
}
};
let child = Arc::new(Mutex::new(child));
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (notification_sender, notifications) = mpsc::sync_channel(256);
let (reader_done_sender, reader_done) = mpsc::sync_channel(1);
let closed = Arc::new(AtomicBool::new(false));
let cleanup = LspCleanupHandle {
child: Arc::clone(&child),
};
let reader_pending = Arc::clone(&pending);
let reader_writer = Arc::clone(&writer);
let reader_closed = Arc::clone(&closed);
let reader = thread::spawn(move || {
let mut reader = BufReader::new(stdout);
loop {
let value = match read_lsp_message(&mut reader) {
Ok(value) => value,
Err(error) => {
reader_closed.store(true, Ordering::SeqCst);
fail_pending(
&reader_pending,
anyhow::anyhow!("LSP reader failed: {error}"),
);
cleanup.terminate();
break;
}
};
if value.get("method").is_some() && value.get("id").is_some() {
match serde_json::from_value::<JsonRpcRequest>(value) {
Ok(request) => {
if let Err(error) = respond_to_server_request(&reader_writer, request) {
reader_closed.store(true, Ordering::SeqCst);
fail_pending(&reader_pending, error);
cleanup.terminate();
break;
}
continue;
}
Err(error) => {
reader_closed.store(true, Ordering::SeqCst);
fail_pending(
&reader_pending,
anyhow::anyhow!("LSP malformed JSON-RPC request: {error}"),
);
cleanup.terminate();
break;
}
}
}
let message = match serde_json::from_value::<JsonRpcMessage>(value.clone()) {
Ok(message) => message,
Err(error) => {
reader_closed.store(true, Ordering::SeqCst);
fail_pending(
&reader_pending,
anyhow::anyhow!("LSP malformed JSON-RPC: {error}"),
);
cleanup.terminate();
break;
}
};
match message {
JsonRpcMessage::Response(response) => {
if let Ok(mut pending) = reader_pending.lock()
&& let Some(sender) = pending.remove(&response.id)
{
let _ = sender.send(Ok(response.result));
}
}
JsonRpcMessage::Error(error) => {
if let Some(id) = error.id
&& let Ok(mut pending) = reader_pending.lock()
&& let Some(sender) = pending.remove(&id)
{
let ErrorData { code, message, .. } = error.error;
let _ = sender
.send(Err(anyhow::anyhow!("LSP protocol error {code}: {message}")));
}
}
JsonRpcMessage::Notification(notification) => {
let _ = notification_sender.try_send(notification);
}
JsonRpcMessage::Request(request) => {
if let Err(error) = respond_to_server_request(&reader_writer, request) {
reader_closed.store(true, Ordering::SeqCst);
fail_pending(&reader_pending, error);
cleanup.terminate();
break;
}
}
}
}
let _ = reader_done_sender.send(());
});
let stderr_truncated = Arc::new(AtomicBool::new(false));
let stderr_receiver = Some(spawn_bounded_pipe_reader(
stderr,
MAX_LSP_STDERR_BYTES,
Arc::clone(&stderr_truncated),
));
Ok(Self {
writer,
child,
reader: Some(reader),
reader_done: Some(reader_done),
stderr_receiver,
pending,
notifications,
next_id: AtomicI64::new(1),
closed,
})
}
pub(crate) fn request(
&self,
method: &str,
params: Value,
timeout: Duration,
cancellation: &AgentCancellation,
) -> anyhow::Result<Value> {
let started = Instant::now();
if self.closed.load(Ordering::SeqCst) {
bail!("LSP server is closed");
}
let id = RequestId::Number(self.next_id.fetch_add(1, Ordering::SeqCst));
let (sender, receiver) = mpsc::sync_channel(1);
self.pending
.lock()
.map_err(|_| anyhow::anyhow!("LSP pending request map poisoned"))?
.insert(id.clone(), sender);
let _pending_cleanup = PendingRequestCleanup {
pending: Arc::clone(&self.pending),
id: id.clone(),
};
let message = serde_json::to_value(JsonRpcRequest::new(id.clone(), method, Some(params)))?;
if let Err(error) = self.write_value_until(&message, started + timeout, Some(cancellation))
{
self.terminate();
return Err(error);
}
loop {
if let Err(error) = cancellation.check() {
self.terminate();
return Err(error);
}
let remaining = timeout.saturating_sub(started.elapsed());
if remaining.is_zero() {
self.terminate();
bail!("LSP request timed out after {}ms", timeout.as_millis());
}
match receiver.recv_timeout(remaining.min(Duration::from_millis(50))) {
Ok(result) => return result,
Err(mpsc::RecvTimeoutError::Timeout) => continue,
Err(mpsc::RecvTimeoutError::Disconnected) => {
bail!("LSP reader disconnected before response");
}
}
}
}
pub(crate) fn notify(&self, method: &str, params: Value) -> anyhow::Result<()> {
let message = serde_json::to_value(JsonRpcNotification::new(method, Some(params)))?;
let result =
self.write_value_until(&message, Instant::now() + Duration::from_secs(1), None);
if result.is_err() {
self.terminate();
}
result
}
pub(crate) fn try_recv_notification(&self) -> Option<JsonRpcNotification> {
self.notifications.try_recv().ok()
}
pub(crate) fn is_closed(&self) -> bool {
self.closed.load(Ordering::SeqCst)
}
pub(crate) fn shutdown(&mut self, timeout: Duration) {
let cancellation = AgentCancellation::default();
if self
.request("shutdown", Value::Null, timeout, &cancellation)
.is_ok()
{
let message = serde_json::to_value(JsonRpcNotification::new("exit", None));
if let Ok(message) = message {
let _ = self.write_value_until(
&message,
Instant::now() + timeout.min(Duration::from_millis(100)),
None,
);
}
}
self.terminate();
if let Some(handle) = self.reader.take()
&& let Some(done) = self.reader_done.take()
{
match done.recv_timeout(PIPE_READER_JOIN_TIMEOUT) {
Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => {
let _ = handle.join();
}
Err(mpsc::RecvTimeoutError::Timeout) => {}
}
}
if let Some(receiver) = self.stderr_receiver.take() {
let _ = recv_pipe_reader_with_timeout(receiver);
}
}
fn write_value_until(
&self,
value: &Value,
deadline: Instant,
cancellation: Option<&AgentCancellation>,
) -> anyhow::Result<()> {
let mut bytes = Vec::new();
write_lsp_message(&mut bytes, value)?;
self.writer.write(&bytes, deadline, cancellation)
}
fn terminate(&self) {
self.closed.store(true, Ordering::SeqCst);
if let Ok(mut child) = self.child.lock() {
let _ = terminate_child_tree_and_wait(&mut child);
}
}
}
impl Drop for LspStdioTransport {
fn drop(&mut self) {
self.shutdown(Duration::from_millis(100));
}
}
fn respond_to_server_request(
writer: &Arc<ChildPipeWriter>,
request: JsonRpcRequest,
) -> anyhow::Result<()> {
let value = if matches!(
request.method.as_str(),
"workspace/configuration" | "window/workDoneProgress/create"
) {
json!({"jsonrpc":"2.0", "id": request.id, "result": null})
} else {
serde_json::to_value(JsonRpcError {
jsonrpc: "2.0".to_string(),
id: Some(request.id),
error: ErrorData {
code: METHOD_NOT_FOUND,
message: "Method not found".to_string(),
data: None,
},
})?
};
let mut bytes = Vec::new();
write_lsp_message(&mut bytes, &value)?;
writer.write(&bytes, Instant::now() + Duration::from_secs(1), None)
}
fn fail_pending(pending: &PendingMap, error: anyhow::Error) {
if let Ok(mut pending) = pending.lock() {
let message = error.to_string();
for (_, sender) in pending.drain() {
let _ = sender.send(Err(anyhow::anyhow!(message.clone())));
}
}
}
pub(crate) fn write_lsp_message(writer: &mut dyn Write, value: &Value) -> anyhow::Result<()> {
let body = serde_json::to_vec(value)?;
write!(writer, "Content-Length: {}\r\n\r\n", body.len())?;
writer.write_all(&body)?;
writer.flush()?;
Ok(())
}
pub(crate) fn read_lsp_message(reader: &mut dyn BufRead) -> anyhow::Result<Value> {
let header = read_lsp_header(reader)?;
let header = String::from_utf8(header)?;
let content_length = parse_content_length_header(&header)?;
if content_length > MAX_LSP_BODY_BYTES {
bail!("LSP body exceeded {MAX_LSP_BODY_BYTES} bytes");
}
let mut body = vec![0; content_length];
reader.read_exact(&mut body)?;
Ok(serde_json::from_slice(&body)?)
}
fn read_lsp_header(reader: &mut dyn BufRead) -> anyhow::Result<Vec<u8>> {
let mut header = Vec::with_capacity(MAX_LSP_HEADER_BYTES);
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
bail!("LSP stdout closed");
}
let remaining = MAX_LSP_HEADER_BYTES.saturating_sub(header.len());
let mut consumed = 0;
let mut complete = false;
for byte in available.iter().take(remaining) {
header.push(*byte);
consumed += 1;
if header.ends_with(b"\r\n\r\n") || header.ends_with(b"\n\n") {
complete = true;
break;
}
}
reader.consume(consumed);
if complete {
return Ok(header);
}
if header.len() == MAX_LSP_HEADER_BYTES {
if reader.fill_buf()?.is_empty() {
bail!("LSP stdout closed");
}
bail!("LSP header exceeded {MAX_LSP_HEADER_BYTES} bytes");
}
}
}
pub(crate) fn parse_content_length_header(headers: &str) -> anyhow::Result<usize> {
for line in headers.lines() {
let Some((name, value)) = line.split_once(':') else {
continue;
};
if name.eq_ignore_ascii_case("Content-Length") {
let length = value.trim().parse::<usize>()?;
return Ok(length);
}
}
bail!("missing Content-Length header")
}