LSP client: spawn language servers, expose code intelligence tools
New lsp.rs: LspRegistry manages persistent LSP server connections. Spawns child processes, speaks LSP protocol (Content-Length framed JSON-RPC over stdio). Server indexes the project once; queries are cheap. Tools: lsp_definition, lsp_references, lsp_hover, lsp_symbols, lsp_callers. Each takes file/line/character, queries the running language server. LspRegistry lives on Agent as Option<Arc>, shared across forks. Still needs: config-driven server startup (like MCP). Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet
parent
8b5614ba99
commit
6ec0e1c766
3 changed files with 437 additions and 0 deletions
419
src/agent/tools/lsp.rs
Normal file
419
src/agent/tools/lsp.rs
Normal file
|
|
@ -0,0 +1,419 @@
|
|||
// tools/lsp.rs — LSP client for code intelligence
|
||||
//
|
||||
// Spawns language servers on demand when a file is first queried.
|
||||
// Finds the project root (git/cargo/etc.) automatically. Maintains
|
||||
// persistent connections — the server indexes once, queries are cheap.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde_json::json;
|
||||
use std::collections::HashSet;
|
||||
use std::path::Path;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
|
||||
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
|
||||
|
||||
struct LspServer {
|
||||
root_path: String,
|
||||
stdin: BufWriter<ChildStdin>,
|
||||
stdout: BufReader<ChildStdout>,
|
||||
_child: Child,
|
||||
next_id: i64,
|
||||
opened_files: HashSet<String>,
|
||||
last_access: u64,
|
||||
}
|
||||
|
||||
impl LspServer {
|
||||
async fn request(&mut self, method: &str, params: serde_json::Value) -> Result<serde_json::Value> {
|
||||
self.next_id += 1;
|
||||
let id = self.next_id;
|
||||
let msg = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
|
||||
self.send_message(&msg).await?;
|
||||
self.read_response(id).await
|
||||
}
|
||||
|
||||
async fn notify(&mut self, method: &str, params: serde_json::Value) -> Result<()> {
|
||||
let msg = json!({ "jsonrpc": "2.0", "method": method, "params": params });
|
||||
self.send_message(&msg).await
|
||||
}
|
||||
|
||||
async fn send_message(&mut self, msg: &serde_json::Value) -> Result<()> {
|
||||
let body = serde_json::to_string(msg)?;
|
||||
let header = format!("Content-Length: {}\r\n\r\n", body.len());
|
||||
self.stdin.write_all(header.as_bytes()).await?;
|
||||
self.stdin.write_all(body.as_bytes()).await?;
|
||||
self.stdin.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_response(&mut self, expected_id: i64) -> Result<serde_json::Value> {
|
||||
loop {
|
||||
let mut content_length: usize = 0;
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
self.stdout.read_line(&mut line).await?;
|
||||
let line = line.trim();
|
||||
if line.is_empty() { break; }
|
||||
if let Some(len) = line.strip_prefix("Content-Length: ") {
|
||||
content_length = len.parse()?;
|
||||
}
|
||||
}
|
||||
if content_length == 0 {
|
||||
anyhow::bail!("LSP: no Content-Length header");
|
||||
}
|
||||
|
||||
let mut body = vec![0u8; content_length];
|
||||
self.stdout.read_exact(&mut body).await?;
|
||||
let msg: serde_json::Value = serde_json::from_slice(&body)?;
|
||||
|
||||
if let Some(id) = msg.get("id").and_then(|v| v.as_i64()) {
|
||||
if id == expected_id {
|
||||
if let Some(err) = msg.get("error") {
|
||||
anyhow::bail!("LSP error: {}", err);
|
||||
}
|
||||
return Ok(msg.get("result").cloned().unwrap_or(serde_json::Value::Null));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn ensure_open(&mut self, path: &str) -> Result<String> {
|
||||
let uri = format!("file://{}", path);
|
||||
if !self.opened_files.contains(&uri) {
|
||||
let text = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("reading {}", path))?;
|
||||
self.notify("textDocument/didOpen", json!({
|
||||
"textDocument": {
|
||||
"uri": uri,
|
||||
"languageId": detect_language(path),
|
||||
"version": 1,
|
||||
"text": text,
|
||||
}
|
||||
})).await?;
|
||||
self.opened_files.insert(uri.clone());
|
||||
}
|
||||
Ok(uri)
|
||||
}
|
||||
}
|
||||
|
||||
fn detect_language(path: &str) -> &'static str {
|
||||
match Path::new(path).extension().and_then(|e| e.to_str()) {
|
||||
Some("rs") => "rust",
|
||||
Some("c" | "h") => "c",
|
||||
Some("cpp" | "cc" | "cxx" | "hpp") => "cpp",
|
||||
Some("py") => "python",
|
||||
Some("js") => "javascript",
|
||||
Some("ts") => "typescript",
|
||||
Some("go") => "go",
|
||||
Some("java") => "java",
|
||||
_ => "plaintext",
|
||||
}
|
||||
}
|
||||
|
||||
fn find_project_root(file_path: &str) -> Option<String> {
|
||||
let mut dir = Path::new(file_path).parent()?;
|
||||
loop {
|
||||
for marker in &[".git", "Cargo.toml", "package.json", "go.mod", "pyproject.toml", "Makefile"] {
|
||||
if dir.join(marker).exists() {
|
||||
return Some(dir.to_string_lossy().to_string());
|
||||
}
|
||||
}
|
||||
dir = dir.parent()?;
|
||||
}
|
||||
}
|
||||
|
||||
const IDLE_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use tokio::sync::Mutex as TokioMutex;
|
||||
|
||||
struct Registry {
|
||||
configs: Vec<crate::config::LspServerConfig>,
|
||||
servers: Vec<LspServer>,
|
||||
}
|
||||
|
||||
static REGISTRY: OnceLock<TokioMutex<Registry>> = OnceLock::new();
|
||||
|
||||
fn registry() -> &'static TokioMutex<Registry> {
|
||||
REGISTRY.get_or_init(|| {
|
||||
let configs = crate::config::get().lsp_servers.clone();
|
||||
TokioMutex::new(Registry { configs, servers: Vec::new() })
|
||||
})
|
||||
}
|
||||
|
||||
fn now() -> u64 {
|
||||
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs()
|
||||
}
|
||||
|
||||
impl LspServer {
|
||||
async fn spawn(command: &str, args: &[String], root_path: &str) -> Result<LspServer> {
|
||||
let mut child = Command::new(command)
|
||||
.args(args)
|
||||
.stdin(std::process::Stdio::piped())
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.spawn()
|
||||
.with_context(|| format!("spawning LSP: {} {}", command, args.join(" ")))?;
|
||||
|
||||
let mut server = LspServer {
|
||||
root_path: root_path.to_string(),
|
||||
stdin: BufWriter::new(child.stdin.take().unwrap()),
|
||||
stdout: BufReader::new(child.stdout.take().unwrap()),
|
||||
_child: child,
|
||||
next_id: 0,
|
||||
opened_files: HashSet::new(),
|
||||
last_access: now(),
|
||||
};
|
||||
|
||||
server.request("initialize", json!({
|
||||
"processId": std::process::id(),
|
||||
"rootUri": format!("file://{}", root_path),
|
||||
"capabilities": {
|
||||
"textDocument": {
|
||||
"definition": { "dynamicRegistration": false },
|
||||
"references": { "dynamicRegistration": false },
|
||||
"hover": { "dynamicRegistration": false },
|
||||
"documentSymbol": { "dynamicRegistration": false },
|
||||
"callHierarchy": { "dynamicRegistration": false },
|
||||
}
|
||||
},
|
||||
})).await.with_context(|| format!("initializing LSP for {}", root_path))?;
|
||||
|
||||
server.notify("initialized", json!({})).await?;
|
||||
dbglog!("[lsp] server started for {}", root_path);
|
||||
Ok(server)
|
||||
}
|
||||
}
|
||||
|
||||
impl Registry {
|
||||
fn reap_idle(&mut self) {
|
||||
let n = now();
|
||||
self.servers.retain(|s| n.saturating_sub(s.last_access) < IDLE_TIMEOUT_SECS);
|
||||
}
|
||||
|
||||
fn find_config(&self, lang: &str) -> Option<&crate::config::LspServerConfig> {
|
||||
self.configs.iter().find(|c| {
|
||||
if c.languages.is_empty() {
|
||||
// Auto: rust-analyzer for rust, etc.
|
||||
c.command.contains(lang) || c.name == lang
|
||||
} else {
|
||||
c.languages.iter().any(|l| l == lang)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn ensure_server(&mut self, file_path: &str) -> Result<usize> {
|
||||
let root = find_project_root(file_path)
|
||||
.ok_or_else(|| anyhow::anyhow!("no project root found for {}", file_path))?;
|
||||
let lang = detect_language(file_path);
|
||||
|
||||
self.reap_idle();
|
||||
|
||||
if let Some(idx) = self.servers.iter().position(|s| s.root_path == root) {
|
||||
self.servers[idx].last_access = now();
|
||||
return Ok(idx);
|
||||
}
|
||||
|
||||
let config = self.find_config(lang)
|
||||
.ok_or_else(|| anyhow::anyhow!("no LSP server configured for {}", lang))?
|
||||
.clone();
|
||||
let server = LspServer::spawn(&config.command, &config.args, &root).await?;
|
||||
self.servers.push(server);
|
||||
Ok(self.servers.len() - 1)
|
||||
}
|
||||
|
||||
async fn conn_for(&mut self, path: &str) -> Result<(&mut LspServer, String)> {
|
||||
let idx = self.ensure_server(path).await?;
|
||||
let server = &mut self.servers[idx];
|
||||
let uri = server.ensure_open(path).await?;
|
||||
Ok((server, uri))
|
||||
}
|
||||
}
|
||||
|
||||
// -- Operation table ----------------------------------------------------------
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
struct LspOp {
|
||||
tool_name: &'static str,
|
||||
description: &'static str,
|
||||
method: &'static str,
|
||||
needs_position: bool,
|
||||
extra_params: fn() -> serde_json::Value,
|
||||
format: fn(&serde_json::Value) -> String,
|
||||
// Two-step RPCs (e.g. incoming_calls) use a second method on the first result
|
||||
followup: Option<&'static str>,
|
||||
}
|
||||
|
||||
fn no_extra() -> serde_json::Value { json!({}) }
|
||||
fn ref_extra() -> serde_json::Value { json!({"context": {"includeDeclaration": true}}) }
|
||||
|
||||
fn fmt_locations(result: &serde_json::Value) -> String {
|
||||
let locations = if result.is_array() {
|
||||
result.as_array().unwrap().clone()
|
||||
} else if result.is_object() {
|
||||
vec![result.clone()]
|
||||
} else {
|
||||
return "No results.".into();
|
||||
};
|
||||
let mut out = String::new();
|
||||
for loc in &locations {
|
||||
let uri = loc["uri"].as_str().or_else(|| loc["targetUri"].as_str()).unwrap_or("");
|
||||
let range = if loc.get("range").is_some() { &loc["range"] } else { &loc["targetRange"] };
|
||||
let line = range["start"]["line"].as_u64().unwrap_or(0) + 1;
|
||||
let file = uri.strip_prefix("file://").unwrap_or(uri);
|
||||
out.push_str(&format!("{}:{}\n", file, line));
|
||||
}
|
||||
if out.is_empty() { "No results.".into() } else { out }
|
||||
}
|
||||
|
||||
fn fmt_hover(result: &serde_json::Value) -> String {
|
||||
if result.is_null() { return "No hover information.".into(); }
|
||||
let contents = &result["contents"];
|
||||
if let Some(s) = contents.as_str() { return s.to_string(); }
|
||||
if let Some(obj) = contents.as_object() {
|
||||
return obj.get("value").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
||||
}
|
||||
serde_json::to_string_pretty(result).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn fmt_symbols(result: &serde_json::Value) -> String {
|
||||
if let Some(symbols) = result.as_array() {
|
||||
let mut out = String::new();
|
||||
fmt_symbols_recursive(symbols, &mut out, 0);
|
||||
if out.is_empty() { "No symbols found.".into() } else { out }
|
||||
} else {
|
||||
"No symbols found.".into()
|
||||
}
|
||||
}
|
||||
|
||||
fn fmt_symbols_recursive(symbols: &[serde_json::Value], out: &mut String, depth: usize) {
|
||||
let indent = " ".repeat(depth);
|
||||
for sym in symbols {
|
||||
let name = sym["name"].as_str().unwrap_or("?");
|
||||
let kind = match sym["kind"].as_u64().unwrap_or(0) {
|
||||
2 => "Module", 5 => "Class", 6 => "Method", 8 => "Field",
|
||||
10 => "Enum", 11 => "Interface", 12 => "Function", 13 => "Variable",
|
||||
14 => "Constant", 22 => "EnumMember", 23 => "Struct", 26 => "TypeParameter",
|
||||
_ => "Symbol",
|
||||
};
|
||||
let line = sym["range"]["start"]["line"].as_u64()
|
||||
.or_else(|| sym["location"]["range"]["start"]["line"].as_u64())
|
||||
.unwrap_or(0) + 1;
|
||||
out.push_str(&format!("{}{} ({}) - Line {}\n", indent, name, kind, line));
|
||||
if let Some(children) = sym.get("children").and_then(|c| c.as_array()) {
|
||||
fmt_symbols_recursive(children, out, depth + 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn fmt_callers(result: &serde_json::Value) -> String {
|
||||
if let Some(calls) = result.as_array() {
|
||||
let mut out = String::new();
|
||||
for call in calls {
|
||||
if let Some(from) = call.get("from") {
|
||||
let name = from["name"].as_str().unwrap_or("?");
|
||||
let uri = from["uri"].as_str().unwrap_or("");
|
||||
let line = from["range"]["start"]["line"].as_u64().unwrap_or(0) + 1;
|
||||
let file = uri.strip_prefix("file://").unwrap_or(uri);
|
||||
out.push_str(&format!("{}:{}: {}\n", file, line, name));
|
||||
}
|
||||
}
|
||||
if out.is_empty() { "No incoming calls.".into() } else { out }
|
||||
} else {
|
||||
"No incoming calls.".into()
|
||||
}
|
||||
}
|
||||
|
||||
static OPS: &[LspOp] = &[
|
||||
LspOp {
|
||||
tool_name: "lsp_definition",
|
||||
description: "Find where a symbol is defined.",
|
||||
method: "textDocument/definition",
|
||||
needs_position: true,
|
||||
extra_params: no_extra,
|
||||
format: fmt_locations,
|
||||
followup: None,
|
||||
},
|
||||
LspOp {
|
||||
tool_name: "lsp_references",
|
||||
description: "Find all references to a symbol.",
|
||||
method: "textDocument/references",
|
||||
needs_position: true,
|
||||
extra_params: ref_extra,
|
||||
format: fmt_locations,
|
||||
followup: None,
|
||||
},
|
||||
LspOp {
|
||||
tool_name: "lsp_hover",
|
||||
description: "Get type info and documentation for a symbol.",
|
||||
method: "textDocument/hover",
|
||||
needs_position: true,
|
||||
extra_params: no_extra,
|
||||
format: fmt_hover,
|
||||
followup: None,
|
||||
},
|
||||
LspOp {
|
||||
tool_name: "lsp_symbols",
|
||||
description: "List all symbols in a file.",
|
||||
method: "textDocument/documentSymbol",
|
||||
needs_position: false,
|
||||
extra_params: no_extra,
|
||||
format: fmt_symbols,
|
||||
followup: None,
|
||||
},
|
||||
LspOp {
|
||||
tool_name: "lsp_callers",
|
||||
description: "Find all functions that call the function at a position.",
|
||||
method: "textDocument/prepareCallHierarchy",
|
||||
needs_position: true,
|
||||
extra_params: no_extra,
|
||||
format: fmt_callers,
|
||||
followup: Some("callHierarchy/incomingCalls"),
|
||||
},
|
||||
];
|
||||
|
||||
const POS_PARAMS: &str = r#"{"type":"object","properties":{"file":{"type":"string"},"line":{"type":"integer"},"character":{"type":"integer"}},"required":["file","line","character"]}"#;
|
||||
const FILE_PARAMS: &str = r#"{"type":"object","properties":{"file":{"type":"string"}},"required":["file"]}"#;
|
||||
|
||||
async fn dispatch_op(op: &LspOp, v: &serde_json::Value) -> Result<String> {
|
||||
let file = v["file"].as_str().ok_or_else(|| anyhow::anyhow!("file required"))?;
|
||||
|
||||
let mut reg = registry().lock().await;
|
||||
let (conn, uri) = reg.conn_for(file).await?;
|
||||
|
||||
let mut params = json!({ "textDocument": { "uri": uri } });
|
||||
if op.needs_position {
|
||||
let line = v["line"].as_u64().ok_or_else(|| anyhow::anyhow!("line required"))? as u32 - 1;
|
||||
let character = v["character"].as_u64().unwrap_or(0) as u32;
|
||||
params["position"] = json!({ "line": line, "character": character });
|
||||
}
|
||||
let extra = (op.extra_params)();
|
||||
if let Some(obj) = extra.as_object() {
|
||||
for (k, v) in obj { params[k] = v.clone(); }
|
||||
}
|
||||
|
||||
let result = conn.request(op.method, params).await?;
|
||||
|
||||
if let Some(followup) = op.followup {
|
||||
let item = result.as_array().and_then(|a| a.first())
|
||||
.ok_or_else(|| anyhow::anyhow!("no item at this position"))?;
|
||||
let result2 = conn.request(followup, json!({ "item": item })).await?;
|
||||
return Ok((op.format)(&result2));
|
||||
}
|
||||
|
||||
Ok((op.format)(&result))
|
||||
}
|
||||
|
||||
pub(super) fn tools() -> Vec<super::Tool> {
|
||||
OPS.iter().map(|op| {
|
||||
let name = op.tool_name;
|
||||
super::Tool {
|
||||
name: op.tool_name,
|
||||
description: op.description,
|
||||
parameters_json: if op.needs_position { POS_PARAMS } else { FILE_PARAMS },
|
||||
handler: Arc::new(move |_agent, v| Box::pin(async move {
|
||||
let op = OPS.iter().find(|o| o.tool_name == name).unwrap();
|
||||
dispatch_op(op, &v).await
|
||||
})),
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
|
|
@ -6,6 +6,7 @@
|
|||
|
||||
// Core tools
|
||||
mod ast_grep;
|
||||
pub mod lsp;
|
||||
pub mod mcp_client;
|
||||
mod bash;
|
||||
pub mod channels;
|
||||
|
|
@ -184,6 +185,7 @@ pub fn tools() -> Vec<Tool> {
|
|||
all.extend(memory::journal_tools());
|
||||
all.extend(channels::tools());
|
||||
all.extend(control::tools());
|
||||
all.extend(lsp::tools());
|
||||
all
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue