Skip to main content

hematite/agent/lsp/
manager.rs

1use crate::agent::lsp::client::LspClient;
2use serde_json::json;
3use std::collections::{BTreeSet, HashMap};
4use std::path::PathBuf;
5use std::sync::Arc;
6
7/// Orchestrates Language Servers for the agent.
8pub struct LspManager {
9    pub clients: HashMap<String, Arc<LspClient>>,
10    pub workspace_root: PathBuf,
11    pub opened_files: BTreeSet<PathBuf>,
12}
13
14impl LspManager {
15    pub fn new(workspace_root: PathBuf) -> Self {
16        Self {
17            clients: HashMap::new(),
18            workspace_root,
19            opened_files: BTreeSet::new(),
20        }
21    }
22
23    /// Discovers and starts necessary language servers.
24    pub async fn start_servers(&mut self) -> Result<(), String> {
25        // Start rust-analyzer if a Cargo.toml exists
26        if self.workspace_root.join("Cargo.toml").exists() {
27            self.start_server("rust", "rust-analyzer", &[]).await?;
28        }
29
30        // ── Stabilization ──
31        // Give the Language Server a moment to index before the first tool call fires.
32        tokio::time::sleep(std::time::Duration::from_millis(1500)).await;
33
34        Ok(())
35    }
36
37    pub async fn start_server(
38        &mut self,
39        lang: &str,
40        command: &str,
41        args: &[String],
42    ) -> Result<(), String> {
43        let client =
44            LspClient::spawn(command, args).map_err(|e| format!("LSP Spawn Fail: {}", e))?;
45        let arc_client = Arc::new(client);
46
47        // --- LAYER 1: LSP Handshake ---
48        let params = json!({
49            "processId": std::process::id(),
50            "rootUri": format!("file:///{}", self.workspace_root.to_str().unwrap_or_default().replace("\\", "/")),
51            "capabilities": {
52                "textDocument": {
53                    "definition": { "dynamicRegistration": false },
54                    "references": { "dynamicRegistration": false },
55                    "hover": { "dynamicRegistration": false },
56                    "symbol": { "dynamicRegistration": false },
57                    "rename": { "dynamicRegistration": false },
58                    "publishDiagnostics": { "relatedInformation": true }
59                },
60                "workspace": {
61                    "symbol": { "dynamicRegistration": false }
62                }
63            },
64            "initializationOptions": null
65        });
66
67        match arc_client.call("initialize", params).await {
68            Ok(_) => {
69                let _ = arc_client.notify("initialized", json!({})).await;
70                self.clients.insert(lang.to_string(), arc_client);
71                Ok(())
72            }
73            Err(e) => Err(format!("LSP Handshake Fail ({}): {}", lang, e)),
74        }
75    }
76
77    pub fn get_client(&self, lang: &str) -> Option<Arc<LspClient>> {
78        self.clients.get(lang).cloned()
79    }
80
81    /// Helper to find the client for a file extension.
82    pub fn get_client_for_path(&self, path: &str) -> Option<Arc<LspClient>> {
83        let ext = std::path::Path::new(path).extension()?.to_str()?;
84        match ext {
85            "rs" => self.get_client("rust"),
86            "ts" | "js" | "tsx" | "jsx" => self.get_client("typescript"),
87            "py" => self.get_client("python"),
88            _ => None,
89        }
90    }
91
92    pub fn resolve_uri(&self, path: &str) -> String {
93        let abs_path = if std::path::Path::new(path).is_absolute() {
94            std::path::PathBuf::from(path)
95        } else {
96            self.workspace_root.join(path)
97        };
98        format!(
99            "file:///{}",
100            abs_path.to_str().unwrap_or_default().replace("\\", "/")
101        )
102    }
103
104    pub async fn ensure_opened(&mut self, path: &str) -> Result<(), String> {
105        let path_obj = if std::path::Path::new(path).is_absolute() {
106            std::path::PathBuf::from(path)
107        } else {
108            self.workspace_root.join(path)
109        };
110
111        if self.opened_files.contains(&path_obj) {
112            return Ok(());
113        }
114
115        let client = self
116            .get_client_for_path(path)
117            .ok_or_else(|| format!("No LSP client for {}", path))?;
118
119        let content =
120            std::fs::read_to_string(&path_obj).map_err(|e| format!("Read Fail: {}", e))?;
121
122        let lang_id = match path_obj.extension().and_then(|e| e.to_str()) {
123            Some("rs") => "rust",
124            Some("py") => "python",
125            Some("ts") | Some("js") => "typescript",
126            _ => "text",
127        };
128
129        let params = json!({
130            "textDocument": {
131                "uri": self.resolve_uri(path),
132                "languageId": lang_id,
133                "version": 1,
134                "text": content
135            }
136        });
137
138        client.notify("textDocument/didOpen", params).await?;
139        self.opened_files.insert(path_obj);
140        Ok(())
141    }
142}