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

codebook_lsp/
lsp.rs
