codebook_lsp/
lsp.rs

1use std::collections::HashMap;
2use std::path::Path;
3use std::str::FromStr as _;
4use std::sync::Arc;
5
6use codebook::parser::TextRange;
7use codebook::parser::WordLocation;
8use codebook::parser::get_word_from_string;
9use codebook::queries::LanguageType;
10
11use log::LevelFilter;
12use log::error;
13use serde_json::Value;
14use tokio::task;
15use tower_lsp::jsonrpc::Result as RpcResult;
16use tower_lsp::lsp_types::*;
17use tower_lsp::{Client, LanguageServer};
18
19use codebook::Codebook;
20use codebook_config::CodebookConfig;
21use log::{debug, info};
22
23use crate::file_cache::TextDocumentCache;
24use crate::lsp_logger;
25
26const SOURCE_NAME: &str = "Codebook";
27
28pub struct Backend {
29    pub client: Client,
30    // Wrap every call to codebook in spawn_blocking, it's not async
31    pub codebook: Arc<Codebook>,
32    pub config: Arc<CodebookConfig>,
33    pub document_cache: TextDocumentCache,
34}
35
36enum CodebookCommand {
37    AddWord,
38    AddWordGlobal,
39    Unknown,
40}
41
42impl From<&str> for CodebookCommand {
43    fn from(command: &str) -> Self {
44        match command {
45            "codebook.addWord" => CodebookCommand::AddWord,
46            "codebook.addWordGlobal" => CodebookCommand::AddWordGlobal,
47            _ => CodebookCommand::Unknown,
48        }
49    }
50}
51
52impl From<CodebookCommand> for String {
53    fn from(command: CodebookCommand) -> Self {
54        match command {
55            CodebookCommand::AddWord => "codebook.addWord".to_string(),
56            CodebookCommand::AddWordGlobal => "codebook.addWordGlobal".to_string(),
57            CodebookCommand::Unknown => "codebook.unknown".to_string(),
58        }
59    }
60}
61
62#[tower_lsp::async_trait]
63impl LanguageServer for Backend {
64    async fn initialize(&self, params: InitializeParams) -> RpcResult<InitializeResult> {
65        // Get log level from initialization options
66        let log_level = params
67            .initialization_options
68            .as_ref()
69            .and_then(|options| options.get("logLevel"))
70            .and_then(|level| level.as_str())
71            .map(|level| {
72                if level == "debug" {
73                    LevelFilter::Debug
74                } else {
75                    LevelFilter::Info
76                }
77            })
78            .unwrap_or(LevelFilter::Info);
79
80        // Attach the LSP client to the logger and flush buffered logs
81        lsp_logger::LspLogger::attach_client(self.client.clone(), log_level);
82        info!(
83            "LSP logger attached to client with log level: {}",
84            log_level
85        );
86        Ok(InitializeResult {
87            capabilities: ServerCapabilities {
88                text_document_sync: Some(TextDocumentSyncCapability::Kind(
89                    TextDocumentSyncKind::FULL,
90                )),
91                execute_command_provider: Some(ExecuteCommandOptions {
92                    commands: vec![
93                        CodebookCommand::AddWord.into(),
94                        CodebookCommand::AddWordGlobal.into(),
95                    ],
96                    work_done_progress_options: Default::default(),
97                }),
98                code_action_provider: Some(CodeActionProviderCapability::Options(
99                    CodeActionOptions {
100                        code_action_kinds: Some(vec![CodeActionKind::QUICKFIX]),
101                        resolve_provider: None,
102                        work_done_progress_options: WorkDoneProgressOptions {
103                            work_done_progress: None,
104                        },
105                    },
106                )),
107                ..ServerCapabilities::default()
108            },
109            server_info: Some(ServerInfo {
110                name: format!("{SOURCE_NAME} Language Server"),
111                version: Some(env!("CARGO_PKG_VERSION").to_string()),
112            }),
113        })
114    }
115
116    async fn initialized(&self, _: InitializedParams) {
117        info!("Server ready!");
118        info!(
119            "Project config: {}",
120            self.config.project_config_path().unwrap().display()
121        );
122        info!(
123            "Global config: {}",
124            self.config
125                .global_config_path()
126                .unwrap_or_default()
127                .display()
128        );
129    }
130
131    async fn shutdown(&self) -> RpcResult<()> {
132        info!("Server shutting down");
133        Ok(())
134    }
135
136    async fn did_open(&self, params: DidOpenTextDocumentParams) {
137        debug!(
138            "Opened document: uri {:?}, language: {}, version: {}",
139            params.text_document.uri,
140            params.text_document.language_id,
141            params.text_document.version
142        );
143        self.document_cache.insert(&params.text_document);
144        self.spell_check(&params.text_document.uri).await;
145    }
146
147    async fn did_close(&self, params: DidCloseTextDocumentParams) {
148        self.document_cache.remove(&params.text_document.uri);
149        // Clear diagnostics when a file is closed.
150        self.client
151            .publish_diagnostics(params.text_document.uri, vec![], None)
152            .await;
153    }
154
155    async fn did_save(&self, params: DidSaveTextDocumentParams) {
156        debug!("Saved document: {}", params.text_document.uri);
157        if let Some(text) = params.text {
158            self.document_cache.update(&params.text_document.uri, &text);
159            self.spell_check(&params.text_document.uri).await;
160        }
161    }
162
163    async fn did_change(&self, params: DidChangeTextDocumentParams) {
164        debug!(
165            "Changed document: uri={}, version={}",
166            params.text_document.uri, params.text_document.version
167        );
168        let uri = params.text_document.uri;
169        if let Some(change) = params.content_changes.first() {
170            self.document_cache.update(&uri, &change.text);
171            self.spell_check(&uri).await;
172        }
173    }
174
175    async fn code_action(&self, params: CodeActionParams) -> RpcResult<Option<CodeActionResponse>> {
176        let mut actions: Vec<CodeActionOrCommand> = vec![];
177        let doc = match self.document_cache.get(params.text_document.uri.as_ref()) {
178            Some(doc) => doc,
179            None => return Ok(None),
180        };
181
182        for diag in params.context.diagnostics {
183            // Only process our own diagnostics
184            if diag.source.as_deref() != Some(SOURCE_NAME) {
185                continue;
186            }
187            let line = doc
188                .text
189                .lines()
190                .nth(diag.range.start.line as usize)
191                .unwrap_or_default();
192            let start_char = diag.range.start.character as usize;
193            let end_char = diag.range.end.character as usize;
194            let word = get_word_from_string(start_char, end_char, line);
195            if word.is_empty() || word.contains(" ") {
196                continue;
197            }
198            let cb = self.codebook.clone();
199            let inner_word = word.clone();
200            let suggestions = task::spawn_blocking(move || cb.get_suggestions(&inner_word)).await;
201
202            let suggestions = match suggestions {
203                Ok(suggestions) => suggestions,
204                Err(e) => {
205                    error!(
206                        "Error getting suggestions for word '{}' in file '{}'\n Error: {}",
207                        word,
208                        doc.uri.path(),
209                        e
210                    );
211                    continue;
212                }
213            };
214
215            if suggestions.is_none() {
216                continue;
217            }
218
219            suggestions.unwrap().iter().for_each(|suggestion| {
220                actions.push(CodeActionOrCommand::CodeAction(self.make_suggestion(
221                    suggestion,
222                    &diag.range,
223                    &params.text_document.uri,
224                )));
225            });
226            actions.push(CodeActionOrCommand::CodeAction(CodeAction {
227                title: format!("Add '{word}' to dictionary"),
228                kind: Some(CodeActionKind::QUICKFIX),
229                diagnostics: None,
230                edit: None,
231                command: Some(Command {
232                    title: format!("Add '{word}' to dictionary"),
233                    command: CodebookCommand::AddWord.into(),
234                    arguments: Some(vec![word.to_string().into()]),
235                }),
236                is_preferred: None,
237                disabled: None,
238                data: None,
239            }));
240            actions.push(CodeActionOrCommand::CodeAction(CodeAction {
241                title: format!("Add '{word}' to global dictionary"),
242                kind: Some(CodeActionKind::QUICKFIX),
243                diagnostics: None,
244                edit: None,
245                command: Some(Command {
246                    title: format!("Add '{word}' to global dictionary"),
247                    command: CodebookCommand::AddWordGlobal.into(),
248                    arguments: Some(vec![word.to_string().into()]),
249                }),
250                is_preferred: None,
251                disabled: None,
252                data: None,
253            }));
254        }
255        match actions.is_empty() {
256            true => Ok(None),
257            false => Ok(Some(actions)),
258        }
259    }
260
261    async fn execute_command(&self, params: ExecuteCommandParams) -> RpcResult<Option<Value>> {
262        match CodebookCommand::from(params.command.as_str()) {
263            CodebookCommand::AddWord => {
264                let words = params
265                    .arguments
266                    .iter()
267                    .filter_map(|arg| arg.as_str().map(|s| s.to_string()));
268                info!(
269                    "Adding words to dictionary {}",
270                    words.clone().collect::<Vec<String>>().join(", ")
271                );
272                let updated = self.add_words(words);
273                if updated {
274                    let _ = self.config.save();
275                    self.recheck_all().await;
276                }
277                Ok(None)
278            }
279            CodebookCommand::AddWordGlobal => {
280                let words = params
281                    .arguments
282                    .iter()
283                    .filter_map(|arg| arg.as_str().map(|s| s.to_string()));
284                let updated = self.add_words_global(words);
285                if updated {
286                    let _ = self.config.save_global();
287                    self.recheck_all().await;
288                }
289                Ok(None)
290            }
291            CodebookCommand::Unknown => Ok(None),
292        }
293    }
294}
295
296impl Backend {
297    pub fn new(client: Client, workspace_dir: &Path) -> Self {
298        let config = CodebookConfig::load(Some(workspace_dir)).expect("Unable to make config.");
299        let config_arc = Arc::new(config);
300        let cb_config = Arc::clone(&config_arc);
301        let codebook = Arc::new(Codebook::new(cb_config).expect("Unable to make codebook"));
302
303        Self {
304            client,
305            codebook,
306            config: Arc::clone(&config_arc),
307            document_cache: TextDocumentCache::default(),
308        }
309    }
310    fn make_diagnostic(&self, result: &WordLocation, range: &TextRange) -> Diagnostic {
311        let message = format!("Possible spelling issue '{word}'.", word = result.word);
312        Diagnostic {
313            range: Range {
314                start: Position {
315                    line: range.line,
316                    character: range.start_char,
317                },
318                end: Position {
319                    line: range.line,
320                    character: range.end_char,
321                },
322            },
323            severity: Some(DiagnosticSeverity::INFORMATION),
324            code: None,
325            code_description: None,
326            source: Some(SOURCE_NAME.to_string()),
327            message,
328            related_information: None,
329            tags: None,
330            data: None,
331        }
332    }
333
334    fn add_words(&self, words: impl Iterator<Item = String>) -> bool {
335        let mut should_save = false;
336        for word in words {
337            match self.config.add_word(&word) {
338                Ok(true) => {
339                    should_save = true;
340                }
341                Ok(false) => {
342                    info!("Word '{word}' already exists in dictionary.");
343                }
344                Err(e) => {
345                    error!("Failed to add word: {e}");
346                }
347            }
348        }
349        should_save
350    }
351    fn add_words_global(&self, words: impl Iterator<Item = String>) -> bool {
352        let mut should_save = false;
353        for word in words {
354            match self.config.add_word_global(&word) {
355                Ok(true) => {
356                    should_save = true;
357                }
358                Ok(false) => {
359                    info!("Word '{word}' already exists in global dictionary.");
360                }
361                Err(e) => {
362                    error!("Failed to add word: {e}");
363                }
364            }
365        }
366        should_save
367    }
368
369    fn make_suggestion(&self, suggestion: &str, range: &Range, uri: &Url) -> CodeAction {
370        let title = format!("Replace with '{suggestion}'");
371        let mut map = HashMap::new();
372        map.insert(
373            uri.clone(),
374            vec![TextEdit {
375                range: *range,
376                new_text: suggestion.to_string(),
377            }],
378        );
379        let edit = Some(WorkspaceEdit {
380            changes: Some(map),
381            document_changes: None,
382            change_annotations: None,
383        });
384        CodeAction {
385            title: title.to_string(),
386            kind: Some(CodeActionKind::QUICKFIX),
387            diagnostics: None,
388            edit,
389            command: None,
390            is_preferred: None,
391            disabled: None,
392            data: None,
393        }
394    }
395
396    async fn recheck_all(&self) {
397        let urls = self.document_cache.cached_urls();
398        debug!("Rechecking documents: {urls:?}");
399        for url in urls {
400            self.publish_spellcheck_diagnostics(&url).await;
401        }
402    }
403
404    async fn spell_check(&self, uri: &Url) {
405        let did_reload = match self.config.reload() {
406            Ok(did_reload) => did_reload,
407            Err(e) => {
408                error!("Failed to reload config: {e}");
409                false
410            }
411        };
412
413        if did_reload {
414            self.recheck_all().await;
415        } else {
416            self.publish_spellcheck_diagnostics(uri).await;
417        }
418    }
419
420    /// Helper method to publish diagnostics for spell-checking.
421    async fn publish_spellcheck_diagnostics(&self, uri: &Url) {
422        let doc = match self.document_cache.get(uri.as_ref()) {
423            Some(doc) => doc,
424            None => return,
425        };
426        // Convert the file URI to a local file path.
427        let file_path = doc.uri.to_file_path().unwrap_or_default();
428        debug!("Spell-checking file: {file_path:?}");
429        // 1) Perform spell-check.
430        let lang_type = doc
431            .language_id
432            .as_deref()
433            .and_then(|lang| LanguageType::from_str(lang).ok());
434
435        let cb = self.codebook.clone();
436        let fp = file_path.clone();
437        let spell_results = task::spawn_blocking(move || {
438            cb.spell_check(&doc.text, lang_type, Some(fp.to_str().unwrap_or_default()))
439        })
440        .await;
441
442        let spell_results = match spell_results {
443            Ok(results) => results,
444            Err(err) => {
445                error!("Spell-checking failed for file '{file_path:?}' \n Error: {err}");
446                return;
447            }
448        };
449
450        // 2) Convert the results to LSP diagnostics.
451        let diagnostics: Vec<Diagnostic> = spell_results
452            .into_iter()
453            .flat_map(|res| {
454                // For each misspelling, create a diagnostic for each location.
455                let mut new_locations = vec![];
456                for loc in &res.locations {
457                    let diagnostic = self.make_diagnostic(&res, loc);
458                    new_locations.push(diagnostic);
459                }
460                new_locations
461            })
462            .collect();
463
464        // debug!("Diagnostics: {:?}", diagnostics);
465        // 3) Send the diagnostics to the client.
466        self.client
467            .publish_diagnostics(doc.uri, diagnostics, None)
468            .await;
469        // debug!("Published diagnostics for: {:?}", file_path);
470    }
471}
codebook_lsp/lsp.rs

codebook_lsp/
lsp.rs
