use crate::{
cancellation::AgentCancellation,
config::LspServerConfig,
lsp::{
DiagnosticsStore, DocumentUri, DocumentVersion, DocumentVersions, LspDiagnostic,
stdio::LspStdioTransport,
},
};
use serde::Deserialize;
use serde_json::json;
use std::{
path::Path,
time::{Duration, Instant},
};
const SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(500);
#[cfg(test)]
const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10);
pub(crate) struct LspClient {
language_id: String,
transport: LspStdioTransport,
diagnostics: DiagnosticsStore,
versions: DocumentVersions,
}
impl LspClient {
#[cfg(test)]
pub(crate) fn start(
server_id: String,
config: &LspServerConfig,
workspace_root: &Path,
language_id: String,
) -> anyhow::Result<Self> {
Self::start_with_timeout(
server_id,
config,
workspace_root,
language_id,
INITIALIZE_TIMEOUT,
)
}
pub(crate) fn start_with_timeout(
_server_id: String,
config: &LspServerConfig,
workspace_root: &Path,
language_id: String,
initialize_timeout: Duration,
) -> anyhow::Result<Self> {
let transport = LspStdioTransport::spawn(config, workspace_root)?;
let root_uri = DocumentUri::from_path(workspace_root)?;
let params = json!({
"processId": std::process::id(),
"rootUri": root_uri.as_str(),
"capabilities": {
"textDocument": {
"synchronization": {"didSave": false, "willSave": false, "willSaveWaitUntil": false},
"publishDiagnostics": {"versionSupport": true},
},
"workspace": {"configuration": true}
},
"clientInfo": {"name": "magi-code"}
});
transport.request(
"initialize",
params,
initialize_timeout,
&AgentCancellation::default(),
)?;
transport.notify("initialized", json!({}))?;
Ok(Self {
language_id,
transport,
diagnostics: DiagnosticsStore::default(),
versions: DocumentVersions::default(),
})
}
pub(crate) fn sync_full_document(
&mut self,
path: &Path,
text: String,
deadline: Instant,
cancellation: &AgentCancellation,
) -> anyhow::Result<(DocumentVersion, Instant)> {
cancellation.check()?;
let sync_started_at = Instant::now();
if sync_started_at >= deadline {
anyhow::bail!("LSP sync deadline elapsed");
}
let uri = DocumentUri::from_path(path)?;
let version = self.versions.next_for(path);
if version.0 == 1 {
self.transport.notify(
"textDocument/didOpen",
json!({
"textDocument": {
"uri": uri.as_str(),
"languageId": self.language_id,
"version": version.0,
"text": text
}
}),
)?;
} else {
self.transport.notify(
"textDocument/didChange",
json!({
"textDocument": {"uri": uri.as_str(), "version": version.0},
"contentChanges": [{"text": text}]
}),
)?;
}
cancellation.check()?;
Ok((version, sync_started_at))
}
pub(crate) fn wait_for_fresh_diagnostics(
&mut self,
path: &Path,
min_version: DocumentVersion,
sync_started_at: Instant,
deadline: Instant,
cancellation: &AgentCancellation,
) -> Option<Vec<LspDiagnostic>> {
let uri = DocumentUri::from_path(path).ok()?;
loop {
let _ = cancellation.check();
self.drain_notifications();
if let Some(snapshot) = self
.diagnostics
.fresh_for(&uri, min_version, sync_started_at)
{
return Some(snapshot.diagnostics.clone());
}
if Instant::now() >= deadline || cancellation.is_canceled() {
return None;
}
std::thread::sleep(Duration::from_millis(10));
}
}
pub(crate) fn drain_pending_notifications(&mut self) {
self.drain_notifications();
}
pub(crate) fn is_closed(&self) -> bool {
self.transport.is_closed()
}
pub(crate) fn shutdown(mut self) {
self.transport.shutdown(SHUTDOWN_TIMEOUT);
}
fn drain_notifications(&mut self) {
while let Some(notification) = self.transport.try_recv_notification() {
if notification.method == "textDocument/publishDiagnostics" {
let Some(params) = notification.params else {
continue;
};
if let Ok(params) = serde_json::from_value::<PublishDiagnosticsParams>(params) {
self.diagnostics.update(
params.uri,
params.version.map(DocumentVersion),
params.diagnostics,
Instant::now(),
);
}
}
}
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct PublishDiagnosticsParams {
uri: DocumentUri,
#[serde(default)]
version: Option<i32>,
#[serde(default)]
diagnostics: Vec<LspDiagnostic>,
}
#[cfg(test)]
mod tests {
use super::*;
use std::{
fs,
sync::{Arc, atomic::AtomicBool},
};
fn fake_server_script(mode: &str) -> String {
format!(
r#"
import json, os, sys, time
mode = {mode:?}
seen = []
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())
body = sys.stdin.buffer.read(length)
return json.loads(body.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()
def publish(uri, version):
send({{'jsonrpc':'2.0','method':'textDocument/publishDiagnostics','params':{{'uri':uri,'version':version,'diagnostics':[{{'range':{{'start':{{'line':0,'character':1}},'end':{{'line':0,'character':2}}}},'severity':1,'source':'fake','message':'fake diagnostic'}}]}}}})
while True:
msg = read_msg()
if msg is None:
break
method = msg.get('method')
if method == 'initialize':
if mode == 'slow_initialize':
time.sleep(1)
if mode == 'inbound_request':
send({{'jsonrpc':'2.0','id':'server-1','method':'workspace/configuration','params':{{'items':[{{}}]}}}})
read_msg()
send({{'jsonrpc':'2.0','id':msg['id'],'result':{{'capabilities':{{'textDocumentSync':1,'referencesProvider':True}}}}}})
elif method == 'initialized':
pass
elif method == 'textDocument/didOpen':
td = msg['params']['textDocument']; seen.append('open'); publish(td['uri'], td['version'])
elif method == 'textDocument/didChange':
td = msg['params']['textDocument']; seen.append('change'); publish(td['uri'], td['version'])
uri = msg['params']['textDocument']['uri']
send({{'jsonrpc':'2.0','id':msg['id'],'result':[{{'uri':uri,'range':{{'start':{{'line':2,'character':3}},'end':{{'line':2,'character':4}}}}}}]}})
elif method == 'shutdown':
send({{'jsonrpc':'2.0','id':msg['id'],'result':None}})
elif method == 'exit':
break
"#
)
}
fn config_for_script(script: String) -> LspServerConfig {
LspServerConfig {
command: "python3".to_string(),
args: vec!["-u".to_string(), "-c".to_string(), script],
enabled: true,
}
}
fn start_client(mode: &str, root: &Path) -> LspClient {
LspClient::start(
"fake".to_string(),
&config_for_script(fake_server_script(mode)),
root,
"rust".to_string(),
)
.unwrap()
}
#[test]
fn client_initialize_initialized_sync_diagnostics_references_shutdown() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let mut client = start_client("normal", temp.path());
let deadline = Instant::now() + Duration::from_secs(2);
let (version, started) = client
.sync_full_document(
&path,
"fn main() {}\n".to_string(),
deadline,
&AgentCancellation::default(),
)
.unwrap();
let diagnostics = client
.wait_for_fresh_diagnostics(
&path,
version,
started,
deadline,
&AgentCancellation::default(),
)
.unwrap();
assert_eq!(diagnostics[0].message, "fake diagnostic");
let (version, started) = client
.sync_full_document(
&path,
"fn main(){ let x = ; }\n".to_string(),
deadline,
&AgentCancellation::default(),
)
.unwrap();
assert_eq!(version, DocumentVersion(2));
assert!(
client
.wait_for_fresh_diagnostics(
&path,
version,
started,
deadline,
&AgentCancellation::default()
)
.is_some()
);
client.shutdown();
}
#[test]
fn client_handles_inbound_request_during_initialize() {
let temp = tempfile::TempDir::new().unwrap();
let client = start_client("inbound_request", temp.path());
client.shutdown();
}
#[test]
fn slow_initialize_obeys_timeout_or_cancellation() {
let temp = tempfile::TempDir::new().unwrap();
let cancel = Arc::new(AtomicBool::new(true));
let cancellation = AgentCancellation::new(Arc::clone(&cancel));
let config = config_for_script(fake_server_script("slow_initialize"));
let transport = LspStdioTransport::spawn(&config, temp.path()).unwrap();
let error = transport
.request(
"initialize",
json!({}),
Duration::from_secs(5),
&cancellation,
)
.unwrap_err();
assert!(error.to_string().contains("prompt canceled"), "{error}");
}
#[test]
fn diagnostics_for_path_returns_cached_diagnostics() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("lib.rs");
fs::write(&path, "fn main() {}\n").unwrap();
let mut client = start_client("normal", temp.path());
let deadline = Instant::now() + Duration::from_secs(2);
let (version, started) = client
.sync_full_document(
&path,
"fn main() {}\n".to_string(),
deadline,
&AgentCancellation::default(),
)
.unwrap();
let _ = client.wait_for_fresh_diagnostics(
&path,
version,
started,
deadline,
&AgentCancellation::default(),
);
client.shutdown();
}
}