consciousness/src/agent/tools/lsp.rs

// 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()
}