use crate::{
cancellation::AgentCancellation,
config::LspSettings,
lsp::{
client::LspClient,
diagnostics::format_injected_diagnostics,
sync::{LanguageRoute, route_for_path},
},
};
use std::{
collections::BTreeMap,
path::{Path, PathBuf},
sync::{Arc, Mutex, mpsc},
thread::{self, JoinHandle},
time::{Duration, Instant},
};
const MAX_UNEXPECTED_FAILURES_PER_SESSION: u8 = 3;
pub(crate) struct LspManager {
settings: LspSettings,
workspace_root: PathBuf,
servers: Arc<Mutex<BTreeMap<String, Arc<Mutex<ManagedServer>>>>>,
watchdog_stop: Mutex<Option<mpsc::Sender<()>>>,
watchdog: Mutex<Option<JoinHandle<()>>>,
}
struct ManagedServer {
route: LanguageRoute,
client: Option<LspClient>,
unexpected_failures: u8,
disabled_for_session: bool,
last_used: Instant,
notice_emitted: bool,
}
pub(crate) struct EditDiagnosticsRequest<'a> {
pub(crate) path: &'a Path,
pub(crate) content: String,
}
struct WatchdogAccess {
servers: Arc<Mutex<BTreeMap<String, Arc<Mutex<ManagedServer>>>>>,
idle_after: Duration,
}
impl WatchdogAccess {
fn shutdown_idle(&self) {
let now = Instant::now();
let clients = self
.servers
.lock()
.ok()
.map(|servers| {
servers
.values()
.filter_map(|entry| {
let mut server = entry.lock().ok()?;
if server.client.is_some()
&& now.duration_since(server.last_used) >= self.idle_after
{
server.client.take()
} else {
None
}
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
for client in clients {
client.shutdown();
}
}
}
impl LspManager {
pub(crate) fn new(settings: LspSettings, workspace_root: PathBuf) -> Self {
let idle_after = Duration::from_secs(settings.idle_shutdown_minutes.saturating_mul(60));
let tick = idle_after
.min(Duration::from_secs(1))
.max(Duration::from_millis(10));
Self::new_with_timing_impl(settings, workspace_root, idle_after, tick)
}
#[cfg(test)]
fn new_with_timing(
settings: LspSettings,
workspace_root: PathBuf,
idle_after: Duration,
tick: Duration,
) -> Self {
Self::new_with_timing_impl(settings, workspace_root, idle_after, tick)
}
fn new_with_timing_impl(
settings: LspSettings,
workspace_root: PathBuf,
idle_after: Duration,
tick: Duration,
) -> Self {
let workspace_root = workspace_root.canonicalize().unwrap_or(workspace_root);
let (stop_tx, stop_rx) = mpsc::channel();
let servers = Arc::new(Mutex::new(BTreeMap::new()));
let access = Arc::new(WatchdogAccess {
servers: Arc::clone(&servers),
idle_after,
});
let handle = thread::spawn(move || {
loop {
match stop_rx.recv_timeout(tick) {
Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => break,
Err(mpsc::RecvTimeoutError::Timeout) => access.shutdown_idle(),
}
}
});
Self {
settings,
workspace_root,
servers,
watchdog_stop: Mutex::new(Some(stop_tx)),
watchdog: Mutex::new(Some(handle)),
}
}
pub(crate) fn is_enabled(&self) -> bool {
self.settings.enabled
}
pub(crate) fn inject_diagnostics_on_edit(&self) -> bool {
self.settings.inject_diagnostics_on_edit
}
pub(crate) fn edit_budget(&self) -> Duration {
Duration::from_millis(self.settings.diagnostics_wait_ms)
}
pub(crate) fn sync_and_wait_diagnostics(
&self,
request: EditDiagnosticsRequest<'_>,
deadline: Instant,
cancellation: &AgentCancellation,
) -> Option<String> {
if !self.is_enabled() || !self.inject_diagnostics_on_edit() || cancellation.is_canceled() {
return None;
}
if !self.is_inside_workspace(request.path) {
return None;
}
let route = route_for_path(request.path, &self.settings)?;
let server = self.server_entry(route)?;
let mut server = server.lock().ok()?;
server.last_used = Instant::now();
let client = self.client_for(&mut server, deadline, cancellation).ok()?;
client.drain_pending_notifications();
let (version, started) = match client.sync_full_document(
request.path,
request.content,
deadline,
cancellation,
) {
Ok(result) => result,
Err(_) => {
let failed = Self::take_failed_client(&mut server);
drop(server);
if let Some(client) = failed {
client.shutdown();
}
return None;
}
};
let diagnostics = client.wait_for_fresh_diagnostics(
request.path,
version,
started,
deadline,
cancellation,
);
if diagnostics.is_none() && client.is_closed() {
let failed = Self::take_failed_client(&mut server);
drop(server);
if let Some(client) = failed {
client.shutdown();
}
return None;
}
let diagnostics = diagnostics?;
server.last_used = Instant::now();
let block =
format_injected_diagnostics(&server.route.server_id, request.path, &diagnostics);
(!block.trim().is_empty()).then_some(block)
}
pub(crate) fn shutdown_all(&self) {
let clients = self
.servers
.lock()
.ok()
.map(|servers| {
servers
.values()
.filter_map(|entry| entry.lock().ok()?.client.take())
.collect::<Vec<_>>()
})
.unwrap_or_default();
for client in clients {
client.shutdown();
}
}
fn server_entry(&self, route: LanguageRoute) -> Option<Arc<Mutex<ManagedServer>>> {
let mut servers = self.servers.lock().ok()?;
Some(
servers
.entry(route.server_id.clone())
.or_insert_with(|| {
Arc::new(Mutex::new(ManagedServer {
route,
client: None,
unexpected_failures: 0,
disabled_for_session: false,
last_used: Instant::now(),
notice_emitted: false,
}))
})
.clone(),
)
}
fn client_for<'a>(
&self,
server: &'a mut ManagedServer,
deadline: Instant,
_cancellation: &AgentCancellation,
) -> anyhow::Result<&'a mut LspClient> {
if server.disabled_for_session {
anyhow::bail!("LSP server disabled for session");
}
if server.client.is_none() {
if server.unexpected_failures >= MAX_UNEXPECTED_FAILURES_PER_SESSION {
server.disabled_for_session = true;
server.notice_emitted = true;
anyhow::bail!("LSP server unexpected failure limit reached");
}
let timeout = deadline.saturating_duration_since(Instant::now());
if timeout.is_zero() {
anyhow::bail!("LSP deadline elapsed before spawn");
}
match LspClient::start_with_timeout(
server.route.server_id.clone(),
&server.route.config,
&self.workspace_root,
server.route.language_id.clone(),
timeout,
) {
Ok(client) => {
server.client = Some(client);
}
Err(error) => {
server.unexpected_failures = server.unexpected_failures.saturating_add(1);
if server.unexpected_failures >= MAX_UNEXPECTED_FAILURES_PER_SESSION {
server.disabled_for_session = true;
server.notice_emitted = true;
}
return Err(error);
}
}
}
server
.client
.as_mut()
.ok_or_else(|| anyhow::anyhow!("LSP client unavailable"))
}
fn take_failed_client(server: &mut ManagedServer) -> Option<LspClient> {
let client = server.client.take();
if client.is_some() {
server.unexpected_failures = server.unexpected_failures.saturating_add(1);
}
if server.unexpected_failures >= MAX_UNEXPECTED_FAILURES_PER_SESSION {
server.disabled_for_session = true;
server.notice_emitted = true;
}
client
}
fn is_inside_workspace(&self, path: &Path) -> bool {
path.canonicalize()
.map(|path| path.starts_with(&self.workspace_root))
.unwrap_or(false)
}
}
impl Drop for LspManager {
fn drop(&mut self) {
if let Ok(mut stop) = self.watchdog_stop.lock() {
stop.take();
}
if let Ok(mut watchdog) = self.watchdog.lock()
&& let Some(handle) = watchdog.take()
{
let _ = handle.join();
}
self.shutdown_all();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::LspServerConfig;
use std::fs;
fn fake_server_script(mode: &str) -> String {
format!(
r#"
import json, sys, time
mode = {mode:?}
root_uri = None
def read_msg():
header = b''
while not header.endswith(b'\r\n\r\n'):
chunk = sys.stdin.buffer.readline()
if not chunk:
return None
header += chunk
length = 0
for line in header.decode().splitlines():
if line.lower().startswith('content-length:'):
length = int(line.split(':', 1)[1].strip())
return json.loads(sys.stdin.buffer.read(length).decode())
def send(value):
body = json.dumps(value, separators=(',', ':')).encode()
sys.stdout.buffer.write(b'Content-Length: ' + str(len(body)).encode() + b'\r\n\r\n' + body)
sys.stdout.buffer.flush()
while True:
msg = read_msg()
if msg is None:
break
method = msg.get('method')
if method == 'initialize':
if mode == 'crash':
sys.exit(7)
if mode == 'crash_after_initialize':
root_uri = msg.get('params', {{}}).get('rootUri')
send({{'jsonrpc':'2.0','id':msg['id'],'result':{{'capabilities':{{'textDocumentSync':1,'referencesProvider':True}}}}}})
continue
root_uri = msg.get('params', {{}}).get('rootUri')
if mode == 'stale_unversioned' and root_uri:
send({{'jsonrpc':'2.0','method':'textDocument/publishDiagnostics','params':{{'uri':root_uri.rstrip('/') + '/lib.rs','diagnostics':[{{'range':{{'start':{{'line':0,'character':0}},'end':{{'line':0,'character':1}}}},'severity':1,'source':'fake','message':'old unversioned'}}]}}}})
send({{'jsonrpc':'2.0','id':msg['id'],'result':{{'capabilities':{{'textDocumentSync':1,'referencesProvider':True}}}}}})
elif method == 'initialized':
if mode == 'crash_after_initialize':
sys.exit(8)
pass
elif method in ('textDocument/didOpen','textDocument/didChange'):
if mode == 'stale_unversioned':
continue
td = msg['params'].get('textDocument', {{}})
uri = td.get('uri')
version = td.get('version')
send({{'jsonrpc':'2.0','method':'textDocument/publishDiagnostics','params':{{'uri':uri,'version':version,'diagnostics':[{{'range':{{'start':{{'line':0,'character':0}},'end':{{'line':0,'character':1}}}},'severity':1,'source':'fake','message':'boom'}}]}}}})
elif method == 'textDocument/references':
uri = msg['params']['textDocument']['uri']
send({{'jsonrpc':'2.0','id':msg['id'],'result':[{{'uri':uri,'range':{{'start':{{'line':1,'character':0}},'end':{{'line':1,'character':1}}}}}}]}})
elif method == 'shutdown':
send({{'jsonrpc':'2.0','id':msg['id'],'result':None}})
elif method == 'exit':
break
"#
)
}
fn settings_with_fake(mode: &str) -> LspSettings {
let mut settings = LspSettings {
enabled: true,
diagnostics_wait_ms: 1_000,
..LspSettings::default()
};
settings.servers.insert(
"rust-analyzer".to_string(),
LspServerConfig {
command: "python3".to_string(),
args: vec!["-u".to_string(), "-c".to_string(), fake_server_script(mode)],
enabled: true,
},
);
settings
}
#[test]
fn disabled_and_unsupported_paths_do_not_spawn() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let disabled = LspManager::new(LspSettings::default(), temp.path().to_path_buf());
assert!(
disabled
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: "fn main() {}\n".to_string(),
},
Instant::now() + Duration::from_millis(100),
&AgentCancellation::default(),
)
.is_none()
);
let enabled = LspManager::new(settings_with_fake("normal"), temp.path().to_path_buf());
let txt = temp.path().join("notes.txt");
fs::write(&txt, "notes\n").unwrap();
assert!(
enabled
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &txt,
content: "notes\n".to_string(),
},
Instant::now() + Duration::from_millis(100),
&AgentCancellation::default(),
)
.is_none()
);
assert!(enabled.servers.lock().unwrap().is_empty());
}
#[test]
fn sync_injects_diagnostics_and_reuses_server() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let manager = LspManager::new(settings_with_fake("normal"), temp.path().to_path_buf());
let first = manager
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: "fn main() {}\n".to_string(),
},
Instant::now() + Duration::from_secs(2),
&AgentCancellation::default(),
)
.unwrap();
assert!(first.contains("DIAGNOSTICS"));
assert!(first.contains("boom"));
let second = manager
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: "fn main() { }\n".to_string(),
},
Instant::now() + Duration::from_secs(2),
&AgentCancellation::default(),
)
.unwrap();
assert!(second.contains("boom"));
let server = manager.servers.lock().unwrap()["rust-analyzer"].clone();
assert_eq!(server.lock().unwrap().unexpected_failures, 0);
}
#[test]
fn watchdog_autonomously_retires_and_restarts_idle_servers() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let manager = LspManager::new_with_timing(
settings_with_fake("normal"),
temp.path().to_path_buf(),
Duration::from_millis(40),
Duration::from_millis(5),
);
for cycle in 0..3 {
let block = manager
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: format!("fn main() {{ /* cycle {cycle} */ }}\n"),
},
Instant::now() + Duration::from_secs(2),
&AgentCancellation::default(),
)
.unwrap();
assert!(block.contains("DIAGNOSTICS"));
let server = manager.servers.lock().unwrap()["rust-analyzer"].clone();
assert_eq!(server.lock().unwrap().unexpected_failures, 0);
let deadline = Instant::now() + Duration::from_secs(2);
loop {
if server.lock().unwrap().client.is_none() {
break;
}
assert!(
Instant::now() < deadline,
"watchdog did not retire idle client"
);
thread::sleep(Duration::from_millis(5));
}
assert!(!server.lock().unwrap().disabled_for_session);
}
}
#[test]
fn dropping_manager_stops_and_joins_watchdog_promptly() {
let temp = tempfile::TempDir::new().unwrap();
let started = Instant::now();
let manager = LspManager::new_with_timing(
settings_with_fake("normal"),
temp.path().to_path_buf(),
Duration::from_secs(60),
Duration::from_secs(60),
);
drop(manager);
assert!(started.elapsed() < Duration::from_secs(2));
}
#[test]
fn stale_unversioned_diagnostics_before_sync_are_not_injected() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let manager = LspManager::new(
settings_with_fake("stale_unversioned"),
temp.path().to_path_buf(),
);
let block = manager.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: "fn main() {}\n".to_string(),
},
Instant::now() + Duration::from_millis(200),
&AgentCancellation::default(),
);
assert!(
block.is_none(),
"stale unversioned diagnostics were injected: {block:?}"
);
}
#[test]
fn outside_workspace_does_not_sync_and_spawn_failures_hit_respawn_cap() {
let temp = tempfile::TempDir::new().unwrap();
let outside = tempfile::NamedTempFile::new().unwrap();
let manager = LspManager::new(settings_with_fake("normal"), temp.path().to_path_buf());
assert!(
manager
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: outside.path(),
content: "fn main() {}\n".to_string(),
},
Instant::now() + Duration::from_millis(100),
&AgentCancellation::default(),
)
.is_none()
);
assert!(manager.servers.lock().unwrap().is_empty());
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let crashing = LspManager::new(settings_with_fake("crash"), temp.path().to_path_buf());
for _ in 0..3 {
assert!(
crashing
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: "fn main() {}\n".to_string(),
},
Instant::now() + Duration::from_secs(1),
&AgentCancellation::default(),
)
.is_none()
);
}
let server = crashing.servers.lock().unwrap()["rust-analyzer"].clone();
let server = server.lock().unwrap();
assert!(server.disabled_for_session);
assert_eq!(
server.unexpected_failures,
MAX_UNEXPECTED_FAILURES_PER_SESSION
);
}
#[test]
fn post_initialize_crash_drops_stale_client_and_disables_after_respawn_cap() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let manager = LspManager::new(
settings_with_fake("crash_after_initialize"),
temp.path().to_path_buf(),
);
for _ in 0..3 {
assert!(
manager
.sync_and_wait_diagnostics(
EditDiagnosticsRequest {
path: &path,
content: "fn main() {}\n".to_string(),
},
Instant::now() + Duration::from_secs(1),
&AgentCancellation::default(),
)
.is_none()
);
}
let server = manager.servers.lock().unwrap()["rust-analyzer"].clone();
let server = server.lock().unwrap();
assert!(server.disabled_for_session);
assert!(server.client.is_none());
assert_eq!(
server.unexpected_failures,
MAX_UNEXPECTED_FAILURES_PER_SESSION
);
}
}