mod protocol;
mod servers;
use crate::{
core::{
extensions::{
Extension, ExtensionMode, ExtensionTool, ForkProjectInfo, ToolAnnotations, ToolExecCtx,
ToolExposure,
},
settings_manager::agent_dir,
tools::{ToolError, ToolResult},
},
extensions::{EXTENSION_FACTORIES, ExtensionFactory, PRIORITY_LSP},
utils::{
http::{path_to_uri, uri_to_path},
paths::resolve_path,
},
};
use futures_util::future::BoxFuture;
use protocol::{DiagEntry, LspClient};
use serde_json::{Value, json};
use servers::ServerSpec;
use std::{
collections::HashMap,
path::{Path, PathBuf},
sync::{Arc, Mutex, OnceLock},
time::{Duration, Instant},
};
use strum_macros::{Display, EnumString, IntoStaticStr};
type ClientCache = HashMap<(String, PathBuf), Arc<LspClient>>;
const EXT: &str = "lsp";
const TOOL_LSP: &str = "lsp";
const REQUEST_TIMEOUT: Duration = Duration::from_secs(45);
const DIAGNOSTICS_WAIT: Duration = Duration::from_secs(20);
const MAX_FILE_BYTES: u64 = 5 * 1024 * 1024;
const MAX_DIAGNOSTICS: usize = 50;
#[linkme::distributed_slice(EXTENSION_FACTORIES)]
static LSP_FACTORY: ExtensionFactory = ExtensionFactory {
priority: PRIORITY_LSP,
make: || -> Arc<dyn Extension> { Arc::new(Lsp) },
};
pub struct Lsp;
impl Extension for Lsp {
fn name(&self) -> &str {
EXT
}
fn description(&self) -> &str {
"Query a language server for definitions, references, symbols and diagnostics (lsp)."
}
fn fork_project(&self) -> Option<ForkProjectInfo> {
Some(ForkProjectInfo {
plugin_name: "oh-my-pi".to_string(),
plugin_version: "18.12.2".to_string(),
url: "https://github.com/can1357/oh-my-pi".to_string(),
})
}
fn default_enabled(&self) -> bool {
false
}
fn modes(&self) -> Vec<ExtensionMode> {
vec![ExtensionMode::Dev]
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![lsp_tool()]
}
fn on_enabled_changed(&self, enabled: bool) {
if !enabled {
shutdown_clients();
}
}
fn execute_tool_async(
&self,
name: String,
args: Value,
ctx: ToolExecCtx,
) -> BoxFuture<'static, Result<ToolResult, ToolError>> {
Box::pin(async move {
match name.as_str() {
TOOL_LSP => run_lsp(&args, &ctx).await,
other => Err(ToolError(format!("extension {EXT}: unknown tool: {other}"))),
}
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, IntoStaticStr, Display)]
#[strum(serialize_all = "lowercase")]
enum Operation {
Diagnostics,
Definition,
References,
Symbols,
}
fn clients() -> &'static Mutex<ClientCache> {
static S: OnceLock<Mutex<ClientCache>> = OnceLock::new();
S.get_or_init(|| Mutex::new(HashMap::new()))
}
async fn client_for(spec: &'static ServerSpec, root: &Path) -> Result<Arc<LspClient>, String> {
let key = (spec.command.to_string(), root.to_path_buf());
if let Some(client) = clients()
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&key)
{
return Ok(Arc::clone(client));
}
let log_path = agent_dir()
.join("logs")
.join(format!("lsp-{}.log", spec.command));
let client = LspClient::spawn(spec, root, &log_path).await?;
let mut map = clients().lock().unwrap_or_else(|e| e.into_inner());
if let Some(existing) = map.get(&key) {
return Ok(Arc::clone(existing));
}
map.insert(key, Arc::clone(&client));
Ok(client)
}
fn shutdown_clients() {
clients().lock().unwrap_or_else(|e| e.into_inner()).clear();
}
fn lsp_tool() -> ExtensionTool {
ExtensionTool {
exposure: ToolExposure::Direct,
namespace: None,
annotations: ToolAnnotations {
read_only: true,
destructive: false,
idempotent: true,
},
output_schema: None,
name: TOOL_LSP.to_string(),
description: "Ask a language server about a file: diagnostics, definition, references, or symbols. \
Unlike grep, this understands types, trait implementations, overloads and macro expansion."
.to_string(),
label: Some("Language server".to_string()),
parameters: json!({
"type": "object",
"properties": {
"operation": {
"type": "string",
"enum": ["diagnostics", "definition", "references", "symbols"],
"description": "What to ask the language server."
},
"path": {
"type": "string",
"description": "File path, resolved against the working directory. Required for diagnostics/definition/references and for file symbols."
},
"line": {
"type": "integer",
"minimum": 1,
"description": "1-based line of the cursor position (definition/references)."
},
"character": {
"type": "integer",
"minimum": 1,
"description": "1-based character offset within the line, counted in characters (definition/references)."
},
"query": {
"type": "string",
"description": "Workspace symbol query (symbols without path)."
}
},
"required": ["operation"]
}),
snippet: "Use to get compiler diagnostics or to resolve definitions, references and symbols across a codebase."
.to_string(),
prompt_guidelines: vec![
"Run `diagnostics` after editing code instead of shelling out to a compiler: it also covers unsaved edits you have written to disk.".to_string(),
"`line` and `character` are 1-based and counted in characters; pass the position of the identifier, not the start of the line.".to_string(),
"The first call in a project may block while the server indexes; if it times out, retry once before falling back to `grep`.".to_string(),
],
constrained_sampling: false,
render_shell: None,
execution_mode: None,
prepare_arguments: None,
grammar_sampling: None,
}
}
async fn run_lsp(args: &Value, ctx: &ToolExecCtx) -> Result<ToolResult, ToolError> {
let operation: Operation = args
.get("operation")
.and_then(Value::as_str)
.ok_or_else(|| ToolError("missing required parameter: operation".to_string()))?
.parse()
.map_err(|_| {
ToolError(
"`operation` must be one of: diagnostics, definition, references, symbols"
.to_string(),
)
})?;
let path_arg = args.get("path").and_then(Value::as_str).map(str::trim);
let query = args.get("query").and_then(Value::as_str).map(str::trim);
match (operation, path_arg, query) {
(Operation::Symbols, None, Some(query)) if !query.is_empty() => {
run_workspace_symbols(query, ctx).await
}
(Operation::Symbols, None, _) => Err(ToolError(
"symbols needs either `path` (file outline) or `query` (workspace search)".to_string(),
)),
(_, None, _) => Err(ToolError(format!(
"missing required parameter: path (required for {operation})"
))),
(_, Some(""), _) => Err(ToolError("`path` must not be empty.".to_string())),
(operation, Some(raw), _) => run_file_query(operation, raw, args, ctx).await,
}
}
async fn run_file_query(
operation: Operation,
raw_path: &str,
args: &Value,
ctx: &ToolExecCtx,
) -> Result<ToolResult, ToolError> {
let path = resolve_path(raw_path, &ctx.cwd);
if !path.is_file() {
return Err(ToolError(format!("not a file: {}", path.display())));
}
let language_id = servers::language_id_for(&path).ok_or_else(|| {
ToolError(format!(
"unsupported file type for the language server: {} (supported: rust, c/cpp, typescript, \
javascript, python, go, java, ruby, lua, csharp, php)",
path.display()
))
})?;
let spec = servers::server_for(language_id)
.ok_or_else(|| ToolError(format!("no language server configured for `{language_id}`")))?;
let metadata = tokio::fs::metadata(&path)
.await
.map_err(|e| ToolError(format!("cannot stat {}: {e}", path.display())))?;
if metadata.len() > MAX_FILE_BYTES {
return Err(ToolError(format!(
"{} is larger than {} MiB; use `read` or `grep` instead of the language server",
path.display(),
MAX_FILE_BYTES / (1024 * 1024)
)));
}
let text = tokio::fs::read_to_string(&path)
.await
.map_err(|e| ToolError(format!("cannot read {}: {e}", path.display())))?;
let absolute = path.canonicalize().unwrap_or_else(|_| path.clone());
let root = servers::workspace_root(&absolute, spec);
let client = client_for(spec, &root).await.map_err(ToolError)?;
let uri = path_to_uri(&absolute);
if uri.is_empty() {
return Err(ToolError(format!(
"cannot convert {} to a file URI",
absolute.display()
)));
}
let synced_at = Instant::now();
let version = client
.sync_document(&uri, language_id, &text)
.await
.map_err(ToolError)?;
let body = match operation {
Operation::Diagnostics => {
let entry = client
.wait_diagnostics(&uri, synced_at, version, DIAGNOSTICS_WAIT)
.await;
format_diagnostics(&absolute, entry)
}
Operation::Symbols => {
let params = json!({ "textDocument": { "uri": uri } });
let result = client
.request("textDocument/documentSymbol", params, REQUEST_TIMEOUT)
.await
.map_err(ToolError)?;
format_document_symbols(&absolute, &result)
}
Operation::Definition | Operation::References => {
let (line, character) = cursor_position(args, &text)?;
let position = json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character }
});
if operation == Operation::Definition {
let result = client
.request("textDocument/definition", position, REQUEST_TIMEOUT)
.await
.map_err(ToolError)?;
format_locations(&absolute, line, character, "definition", &result)
} else {
let params = json!({
"textDocument": { "uri": uri },
"position": { "line": line, "character": character },
"context": { "includeDeclaration": true }
});
let result = client
.request("textDocument/references", params, REQUEST_TIMEOUT)
.await
.map_err(ToolError)?;
format_locations(&absolute, line, character, "reference", &result)
}
}
};
Ok(ToolResult::text(body))
}
async fn run_workspace_symbols(query: &str, ctx: &ToolExecCtx) -> Result<ToolResult, ToolError> {
let (spec, root) = servers::detect_workspace(Path::new(&ctx.cwd)).ok_or_else(|| {
ToolError(format!(
"cannot tell which language this workspace is in; pass `path` to a file whose \
language has a server (cwd: {})",
ctx.cwd
))
})?;
let client = client_for(spec, &root).await.map_err(ToolError)?;
let result = client
.request(
"workspace/symbol",
json!({ "query": query }),
REQUEST_TIMEOUT,
)
.await
.map_err(ToolError)?;
Ok(ToolResult::text(format_workspace_symbols(query, &result)))
}
fn cursor_position(args: &Value, text: &str) -> Result<(u32, u32), ToolError> {
let line = args
.get("line")
.and_then(Value::as_u64)
.ok_or_else(|| ToolError("missing required parameter: line".to_string()))?;
let character = args
.get("character")
.and_then(Value::as_u64)
.ok_or_else(|| ToolError("missing required parameter: character".to_string()))?;
if line == 0 || character == 0 {
return Err(ToolError(
"`line` and `character` are 1-based and must be >= 1".to_string(),
));
}
let line_text = text
.lines()
.nth(line as usize - 1)
.ok_or_else(|| ToolError(format!("line {line} is past the end of the file")))?;
let char_index = character as usize - 1;
if char_index > line_text.chars().count() {
return Err(ToolError(format!(
"character {character} is past the end of line {line}"
)));
}
Ok(((line - 1) as u32, utf16_column(line_text, char_index)))
}
fn utf16_column(line: &str, char_index: usize) -> u32 {
line.chars()
.take(char_index)
.map(|c| c.len_utf16() as u32)
.sum()
}
fn char_column(line: &str, utf16: u64) -> usize {
let mut units = 0u64;
for (index, c) in line.chars().enumerate() {
if units >= utf16 {
return index;
}
units += c.len_utf16() as u64;
}
line.chars().count()
}
fn format_diagnostics(path: &Path, entry: Option<DiagEntry>) -> String {
let Some(entry) = entry else {
return format!(
"{}: the language server did not report diagnostics in time (it may still be indexing). \
Retry in a moment.",
path.display()
);
};
if entry.items.is_empty() {
return format!("{}: no diagnostics.", path.display());
}
let source = std::fs::read_to_string(path).unwrap_or_default();
let source_lines: Vec<&str> = source.lines().collect();
let mut lines = vec![format!(
"{} — {} diagnostic{}:",
path.display(),
entry.items.len(),
if entry.items.len() == 1 { "" } else { "s" }
)];
for item in entry.items.iter().take(MAX_DIAGNOSTICS) {
lines.push(format_diagnostic(&source_lines, item));
}
if entry.items.len() > MAX_DIAGNOSTICS {
lines.push(format!(
" … and {} more",
entry.items.len() - MAX_DIAGNOSTICS
));
}
lines.join("\n")
}
fn format_diagnostic(source_lines: &[&str], item: &Value) -> String {
let start = item
.get("range")
.and_then(|r| r.get("start"))
.cloned()
.unwrap_or(Value::Null);
let line = start.get("line").and_then(Value::as_u64).unwrap_or(0) + 1;
let character = start
.get("character")
.and_then(Value::as_u64)
.map(|c| {
source_lines
.get(line as usize - 1)
.map_or(1, |line_text| char_column(line_text, c) + 1)
})
.unwrap_or(1);
let severity = match item.get("severity").and_then(Value::as_u64) {
Some(1) => "error",
Some(2) => "warning",
Some(3) => "info",
Some(4) => "hint",
_ => "diagnostic",
};
let message = item
.get("message")
.and_then(Value::as_str)
.unwrap_or("")
.trim()
.replace('\n', " ");
let mut suffix = String::new();
if let Some(code) = item.get("code") {
let code = match code {
Value::String(s) => s.clone(),
other => other.to_string(),
};
if !code.is_empty() {
suffix.push_str(&format!(" ({code})"));
}
}
if let Some(source) = item.get("source").and_then(Value::as_str)
&& !source.is_empty()
{
suffix.push_str(&format!(" [{source}]"));
}
format!(" {line}:{character} {severity}: {message}{suffix}")
}
fn format_locations(path: &Path, line: u32, character: u32, noun: &str, result: &Value) -> String {
let locations = collect_locations(result);
if locations.is_empty() {
return format!(
"No {noun} found at {}:{}:{}.",
path.display(),
line + 1,
character + 1
);
}
let mut lines = vec![format!(
"{} {noun}{}:",
locations.len(),
if locations.len() == 1 { "" } else { "s" }
)];
for (path, line, character) in locations {
lines.push(format!(" {}:{line}:{character}", path.display()));
}
lines.join("\n")
}
fn collect_locations(value: &Value) -> Vec<(PathBuf, u64, u64)> {
let mut out = Vec::new();
let items: Vec<&Value> = match value {
Value::Null => Vec::new(),
Value::Array(items) => items.iter().collect(),
other => vec![other],
};
for item in items {
let uri = item
.get("uri")
.or_else(|| item.get("targetUri"))
.and_then(Value::as_str);
let range = item
.get("range")
.or_else(|| item.get("targetSelectionRange"))
.or_else(|| item.get("targetRange"));
let (Some(uri), Some(range)) = (uri, range) else {
continue;
};
let Some(path) = uri_to_path(uri) else {
continue;
};
let start = range.get("start").cloned().unwrap_or(Value::Null);
let line = start.get("line").and_then(Value::as_u64).unwrap_or(0) + 1;
let character = start.get("character").and_then(Value::as_u64).unwrap_or(0) + 1;
out.push((path, line, character));
}
out
}
fn format_document_symbols(path: &Path, result: &Value) -> String {
let Some(items) = result.as_array() else {
return format!("{}: no symbols found.", path.display());
};
if items.is_empty() {
return format!("{}: no symbols found.", path.display());
}
let mut lines = vec![format!("{} — {} symbol(s):", path.display(), items.len())];
let mut count = 0usize;
for item in items {
push_document_symbol(item, 0, &mut lines, &mut count);
}
lines.join("\n")
}
fn push_document_symbol(item: &Value, depth: usize, lines: &mut Vec<String>, count: &mut usize) {
let Some(name) = item.get("name").and_then(Value::as_str) else {
return;
};
let kind = symbol_kind(item.get("kind").and_then(Value::as_u64));
let start = item
.get("selectionRange")
.or_else(|| item.get("range"))
.or_else(|| {
item.get("location")
.and_then(|location| location.get("range"))
})
.and_then(|range| range.get("start"));
let line = start
.and_then(|s| s.get("line"))
.and_then(Value::as_u64)
.unwrap_or(0)
+ 1;
let character = start
.and_then(|s| s.get("character"))
.and_then(Value::as_u64)
.unwrap_or(0)
+ 1;
*count += 1;
lines.push(format!(
" {}{line}:{character} {kind} {name}",
" ".repeat(depth)
));
if let Some(children) = item.get("children").and_then(Value::as_array) {
for child in children {
push_document_symbol(child, depth + 1, lines, count);
}
}
}
fn format_workspace_symbols(query: &str, result: &Value) -> String {
let Some(items) = result.as_array() else {
return format!("No symbols matching `{query}`.");
};
if items.is_empty() {
return format!("No symbols matching `{query}`.");
}
let mut lines = vec![format!("{} symbol(s) matching `{query}`:", items.len())];
for item in items {
let name = item.get("name").and_then(Value::as_str).unwrap_or("?");
let kind = symbol_kind(item.get("kind").and_then(Value::as_u64));
let location = item.get("location").cloned().unwrap_or(Value::Null);
let path = location
.get("uri")
.and_then(Value::as_str)
.and_then(uri_to_path)
.map(|p| p.display().to_string())
.unwrap_or_else(|| "?".to_string());
let line = location
.get("range")
.and_then(|r| r.get("start"))
.and_then(|s| s.get("line"))
.and_then(Value::as_u64)
.unwrap_or(0)
+ 1;
let character = location
.get("range")
.and_then(|r| r.get("start"))
.and_then(|s| s.get("character"))
.and_then(Value::as_u64)
.unwrap_or(0)
+ 1;
lines.push(format!(" {path}:{line}:{character} {kind} {name}"));
}
lines.join("\n")
}
fn symbol_kind(kind: Option<u64>) -> &'static str {
match kind {
Some(1) => "file",
Some(2) => "module",
Some(3) => "namespace",
Some(4) => "package",
Some(5) => "class",
Some(6) => "method",
Some(7) => "property",
Some(8) => "field",
Some(9) => "constructor",
Some(10) => "enum",
Some(11) => "interface",
Some(12) => "function",
Some(13) => "variable",
Some(14) => "const",
Some(15) => "string",
Some(16) => "number",
Some(17) => "boolean",
Some(18) => "array",
Some(19) => "object",
Some(20) => "key",
Some(21) => "null",
Some(22) => "enum member",
Some(23) => "struct",
Some(24) => "event",
Some(25) => "operator",
Some(26) => "type parameter",
_ => "symbol",
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::AgentDirGuard;
#[test]
fn utf16_columns_round_trip_for_ascii() {
let line = "let value = compute();";
assert_eq!(utf16_column(line, 4), 4);
assert_eq!(char_column(line, 4), 4);
}
#[test]
fn utf16_columns_account_for_non_bmp() {
let line = "a😀b";
assert_eq!(utf16_column(line, 2), 3);
assert_eq!(char_column(line, 3), 2);
assert_eq!(char_column(line, 4), 3);
}
#[test]
fn cursor_position_rejects_out_of_range() {
let text = "fn main() {}\n";
let args = json!({ "line": 2, "character": 1 });
assert!(cursor_position(&args, text).is_err());
let args = json!({ "line": 1, "character": 99 });
assert!(cursor_position(&args, text).is_err());
let args = json!({ "line": 1, "character": 4 });
assert_eq!(cursor_position(&args, text).unwrap(), (0, 3));
}
#[test]
fn collect_locations_handles_all_three_shapes() {
let single = json!({
"uri": "file:///tmp/a.rs",
"range": { "start": { "line": 3, "character": 4 }, "end": { "line": 3, "character": 5 } }
});
assert_eq!(collect_locations(&single).len(), 1);
let links = json!([{
"targetUri": "file:///tmp/b.rs",
"targetRange": { "start": { "line": 0, "character": 0 }, "end": { "line": 0, "character": 1 } },
"targetSelectionRange": { "start": { "line": 1, "character": 2 }, "end": { "line": 1, "character": 3 } }
}]);
let locations = collect_locations(&links);
assert_eq!(locations.len(), 1);
assert_eq!(locations[0].1, 2);
assert_eq!(locations[0].2, 3);
assert!(collect_locations(&Value::Null).is_empty());
}
#[test]
fn formats_diagnostics_with_severity_and_code() {
let _guard = AgentDirGuard::temp();
let path = std::env::temp_dir().join("prux-lsp-diag.rs");
std::fs::write(&path, "fn main() {}\n").unwrap();
let entry = DiagEntry {
items: vec![json!({
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 0, "character": 1 } },
"severity": 1,
"message": "mismatched types",
"code": "E0308",
"source": "rust-analyzer"
})],
version: Some(1),
received: Instant::now(),
};
let text = format_diagnostics(&path, Some(entry));
assert!(text.contains("1 diagnostic:"), "{text}");
assert!(
text.contains("1:1 error: mismatched types (E0308) [rust-analyzer]"),
"{text}"
);
let _ = std::fs::remove_file(&path);
}
#[test]
fn document_symbols_are_indented_by_depth() {
let result = json!([{
"name": "Foo",
"kind": 23,
"range": { "start": { "line": 0, "character": 0 }, "end": { "line": 9, "character": 0 } },
"selectionRange": { "start": { "line": 0, "character": 7 }, "end": { "line": 0, "character": 10 } },
"children": [{
"name": "new",
"kind": 6,
"range": { "start": { "line": 2, "character": 0 }, "end": { "line": 4, "character": 0 } },
"selectionRange": { "start": { "line": 2, "character": 7 }, "end": { "line": 2, "character": 10 } }
}]
}]);
let text = format_document_symbols(Path::new("/tmp/lib.rs"), &result);
assert!(text.contains("1:8 struct Foo"), "{text}");
assert!(text.contains(" 3:8 method new"), "{text}");
}
#[test]
fn operation_parses_from_string() {
assert_eq!(
"diagnostics".parse::<Operation>().unwrap(),
Operation::Diagnostics
);
assert!("bogus".parse::<Operation>().is_err());
}
}