Skip to main content

rustledger_lsp/
server.rs

1//! Main LSP server implementation.
2
3use crate::handlers::execute_command::COMMANDS;
4use crate::handlers::on_type_formatting::{FIRST_TRIGGER_CHARACTER, MORE_TRIGGER_CHARACTERS};
5use crate::handlers::semantic_tokens::get_capabilities as get_semantic_tokens_capabilities;
6use crate::handlers::signature_help::TRIGGER_CHARACTERS as SIGNATURE_TRIGGER_CHARACTERS;
7use crate::ledger_state::{LspConfig, discover_journal_file};
8use crate::main_loop::run_main_loop;
9use crate::uri_to_path;
10use lsp_server::Connection;
11use lsp_types::InitializeParams;
12
13/// The LSP server.
14pub struct Server {
15    /// Connection to the LSP client.
16    connection: Connection,
17    /// Initialize parameters from client.
18    init_params: InitializeParams,
19    /// LSP configuration parsed from init options.
20    config: LspConfig,
21    /// Position encoding negotiated with the client during
22    /// `initialize`. Handler code emitting `Position` values must
23    /// consult this so output aligns with what the client expects.
24    position_encoding: crate::handlers::utils::PositionEncoding,
25}
26
27impl Server {
28    /// Create a new LSP server from a connection.
29    pub fn new(
30        connection: Connection,
31        init_params: InitializeParams,
32        position_encoding: crate::handlers::utils::PositionEncoding,
33    ) -> Self {
34        // Parse configuration from initialization options
35        let config = LspConfig::from_init_options(init_params.initialization_options.as_ref());
36
37        if let Some(ref journal) = config.journal_file {
38            tracing::info!("Journal file configured: {}", journal.display());
39        }
40
41        Self {
42            connection,
43            init_params,
44            config,
45            position_encoding,
46        }
47    }
48
49    /// Run the server's main loop. Returns the exit code produced by
50    /// the `exit` notification (or 0 if the channel closed without
51    /// one). The caller is responsible for draining IO threads
52    /// (`io_threads.join()`) before terminating the process — without
53    /// that drain, the writer can lose the shutdown response queued
54    /// when the loop broke.
55    #[must_use]
56    pub fn run(self) -> i32 {
57        tracing::info!("Starting Beancount Language Server v{}", crate::VERSION);
58
59        // Resolve journal file path relative to workspace root if needed
60        let journal_file = self.resolve_journal_path();
61
62        if let Some(ref path) = journal_file {
63            tracing::info!("Using journal file: {}", path.display());
64        }
65
66        if let Some(folders) = &self.init_params.workspace_folders
67            && let Some(folder) = folders.first()
68        {
69            tracing::info!("Workspace root: {}", folder.uri.as_str());
70        }
71
72        // Run the main event loop with the journal file configuration
73        // and the negotiated position encoding (so handlers emit
74        // positions in the encoding the client expects).
75        let (sender, receiver) = (self.connection.sender, self.connection.receiver);
76        let code = run_main_loop(receiver, sender, journal_file, self.position_encoding);
77
78        tracing::info!("Server shutdown complete (exit code {code})");
79        code
80    }
81
82    /// Resolve the journal file path, making it absolute if necessary.
83    /// If no explicit journal file is configured, attempts auto-discovery.
84    fn resolve_journal_path(&self) -> Option<std::path::PathBuf> {
85        // Get workspace root path for resolution and discovery
86        let workspace_root = self.get_workspace_root();
87
88        // If explicit config provided, resolve it
89        if let Some(journal) = &self.config.journal_file {
90            return self.resolve_explicit_journal(journal, workspace_root.as_deref());
91        }
92
93        // No explicit config - auto-discover, but only in the right directory:
94        // a set workspace folder is authoritative; the process cwd is used ONLY
95        // when there is no workspace folder. Otherwise a stray journal in the
96        // editor's launch directory silently contaminates an unrelated
97        // workspace's state.
98        let discovered = Self::discovery_dir(workspace_root, || std::env::current_dir().ok())
99            .and_then(|dir| discover_journal_file(&dir));
100        if discovered.is_none() {
101            tracing::debug!("No journal file configured or discovered");
102        }
103        discovered
104    }
105
106    /// Directory to search for a journal during auto-discovery: the workspace
107    /// folder when set, otherwise the process cwd. The cwd is deliberately
108    /// *not* consulted when a workspace folder exists, so a journal that
109    /// happens to sit in the editor's launch directory cannot leak into an
110    /// unrelated workspace. `cwd` is computed lazily so the `current_dir`
111    /// syscall is skipped entirely (and its failure ignored) when a workspace
112    /// folder is set.
113    fn discovery_dir(
114        workspace_root: Option<std::path::PathBuf>,
115        cwd: impl FnOnce() -> Option<std::path::PathBuf>,
116    ) -> Option<std::path::PathBuf> {
117        workspace_root.or_else(cwd)
118    }
119
120    /// Get the workspace root path from init params.
121    fn get_workspace_root(&self) -> Option<std::path::PathBuf> {
122        self.init_params
123            .workspace_folders
124            .as_ref()
125            .and_then(|folders| folders.first())
126            .and_then(|folder| uri_to_path(&folder.uri))
127    }
128
129    /// Resolve an explicitly configured journal path.
130    fn resolve_explicit_journal(
131        &self,
132        journal: &std::path::Path,
133        workspace_root: Option<&std::path::Path>,
134    ) -> Option<std::path::PathBuf> {
135        // If already absolute, use as-is
136        if journal.is_absolute() {
137            return Some(journal.to_path_buf());
138        }
139
140        // Try to resolve relative to workspace root
141        if let Some(root) = workspace_root {
142            let resolved = root.join(journal);
143            if resolved.exists() {
144                return Some(resolved);
145            }
146        }
147
148        // Fall back to current directory
149        let resolved = std::env::current_dir()
150            .ok()
151            .map(|cwd| cwd.join(journal))
152            .filter(|p| p.exists());
153
154        if resolved.is_none() {
155            tracing::warn!(
156                "Journal file '{}' not found relative to workspace or current directory",
157                journal.display()
158            );
159        }
160
161        resolved
162    }
163}
164
165/// Start the LSP server using stdio transport.
166///
167/// Returns the exit code that the `exit` notification supplied (0 for
168/// a clean shutdown, 1 for exit-without-prior-shutdown per LSP spec).
169/// If the channel closed without an `exit` notification, returns 0.
170///
171/// Critically, `io_threads.join()` is called BEFORE returning, so the
172/// shutdown response queued in the writer thread's channel is fully
173/// flushed to stdout. Previously the production exit path called
174/// `process::exit(code)` from inside the main loop, which terminated
175/// the process before the writer could flush — losing the shutdown
176/// response on slow runners (the `stdio_smoke` CI flake).
177pub fn start_stdio() -> Result<i32, Box<dyn std::error::Error + Send + Sync>> {
178    tracing::info!("Starting LSP server on stdio");
179
180    // Create connection using stdio
181    let (connection, io_threads) = Connection::stdio();
182
183    // Wait for initialize request
184    let (id, params) = connection.initialize_start()?;
185    let init_params: InitializeParams = serde_json::from_value(params)?;
186
187    // Negotiate position encoding. Our handler stack emits LSP
188    // positions as UTF-8 byte offsets, so prefer UTF-8 if the
189    // client advertises it (LSP 3.17+; VS Code, neovim, helix, and
190    // most modern clients do). If the client doesn't advertise
191    // UTF-8, the LSP spec requires the server to use UTF-16 (the
192    // default), in which case our byte-based positions are wrong
193    // for non-ASCII content — a server-wide latent bug tracked
194    // separately. The negotiation here at least makes us correct
195    // for modern clients without any handler-side conversion.
196    let position_encoding = init_params
197        .capabilities
198        .general
199        .as_ref()
200        .and_then(|g| g.position_encodings.as_ref())
201        .and_then(|encs| {
202            encs.contains(&lsp_types::PositionEncodingKind::UTF8)
203                .then_some(lsp_types::PositionEncodingKind::UTF8)
204        });
205
206    // Derive the handler-facing `PositionEncoding` once, BEFORE
207    // `position_encoding` moves into the `ServerCapabilities` field
208    // below. The `Option<PositionEncodingKind>` is non-Copy, so we
209    // can't borrow it after the move; computing here also makes the
210    // negotiated-encoding-vs-handler-encoding mapping explicit at
211    // the negotiation site.
212    let handler_encoding =
213        crate::handlers::utils::PositionEncoding::from_negotiated(position_encoding.as_ref());
214
215    // Build server capabilities
216    let capabilities = lsp_types::ServerCapabilities {
217        position_encoding,
218        text_document_sync: Some(lsp_types::TextDocumentSyncCapability::Kind(
219            lsp_types::TextDocumentSyncKind::FULL,
220        )),
221        completion_provider: Some(lsp_types::CompletionOptions {
222            trigger_characters: Some(vec![
223                ":".to_string(),  // Account segments
224                " ".to_string(),  // After keywords
225                "\"".to_string(), // Strings (payees, narrations)
226                "#".to_string(),  // Tags
227                "^".to_string(),  // Links
228            ]),
229            resolve_provider: Some(true), // Enable completion resolve for detailed info
230            ..Default::default()
231        }),
232        definition_provider: Some(lsp_types::OneOf::Left(true)),
233        references_provider: Some(lsp_types::OneOf::Left(true)),
234        hover_provider: Some(lsp_types::HoverProviderCapability::Simple(true)),
235        document_symbol_provider: Some(lsp_types::OneOf::Left(true)),
236        semantic_tokens_provider: Some(get_semantic_tokens_capabilities()),
237        code_action_provider: Some(lsp_types::CodeActionProviderCapability::Options(
238            lsp_types::CodeActionOptions {
239                code_action_kinds: Some(vec![
240                    lsp_types::CodeActionKind::QUICKFIX,
241                    lsp_types::CodeActionKind::REFACTOR,
242                ]),
243                resolve_provider: Some(true), // Enable resolve for lazy-loading edits
244                work_done_progress_options: Default::default(),
245            },
246        )),
247        workspace_symbol_provider: Some(lsp_types::OneOf::Left(true)),
248        rename_provider: Some(lsp_types::OneOf::Right(lsp_types::RenameOptions {
249            prepare_provider: Some(true),
250            work_done_progress_options: Default::default(),
251        })),
252        document_formatting_provider: Some(lsp_types::OneOf::Left(true)),
253        document_range_formatting_provider: Some(lsp_types::OneOf::Left(true)),
254        document_link_provider: Some(lsp_types::DocumentLinkOptions {
255            resolve_provider: Some(true), // Enable resolve to verify file existence
256            work_done_progress_options: Default::default(),
257        }),
258        inlay_hint_provider: Some(lsp_types::OneOf::Right(
259            lsp_types::InlayHintServerCapabilities::Options(lsp_types::InlayHintOptions {
260                resolve_provider: Some(true), // Enable resolve for rich tooltips
261                work_done_progress_options: Default::default(),
262            }),
263        )),
264        selection_range_provider: Some(lsp_types::SelectionRangeProviderCapability::Simple(true)),
265        folding_range_provider: Some(lsp_types::FoldingRangeProviderCapability::Simple(true)),
266        document_highlight_provider: Some(lsp_types::OneOf::Left(true)),
267        linked_editing_range_provider: Some(
268            lsp_types::LinkedEditingRangeServerCapabilities::Simple(true),
269        ),
270        document_on_type_formatting_provider: Some(lsp_types::DocumentOnTypeFormattingOptions {
271            first_trigger_character: FIRST_TRIGGER_CHARACTER.to_string(),
272            more_trigger_character: Some(
273                MORE_TRIGGER_CHARACTERS
274                    .iter()
275                    .map(|s| s.to_string())
276                    .collect(),
277            ),
278        }),
279        code_lens_provider: Some(lsp_types::CodeLensOptions {
280            resolve_provider: Some(true), // Enable resolve for lazy-loading balance verification
281        }),
282        color_provider: Some(lsp_types::ColorProviderCapability::Simple(true)),
283        declaration_provider: Some(lsp_types::DeclarationCapability::Simple(true)),
284        call_hierarchy_provider: Some(lsp_types::CallHierarchyServerCapability::Simple(true)),
285        signature_help_provider: Some(lsp_types::SignatureHelpOptions {
286            trigger_characters: Some(
287                SIGNATURE_TRIGGER_CHARACTERS
288                    .iter()
289                    .map(|s| s.to_string())
290                    .collect(),
291            ),
292            retrigger_characters: None,
293            work_done_progress_options: Default::default(),
294        }),
295        execute_command_provider: Some(lsp_types::ExecuteCommandOptions {
296            commands: COMMANDS.iter().map(|s| s.to_string()).collect(),
297            work_done_progress_options: Default::default(),
298        }),
299        // Type hierarchy: advertised via experimental until lsp-types adds native support
300        experimental: Some(serde_json::json!({
301            "typeHierarchyProvider": true
302        })),
303        // Workspace capabilities
304        workspace: Some(lsp_types::WorkspaceServerCapabilities {
305            workspace_folders: Some(lsp_types::WorkspaceFoldersServerCapabilities {
306                supported: Some(true),
307                change_notifications: Some(lsp_types::OneOf::Left(true)),
308            }),
309            file_operations: None, // File operations (create/rename/delete) not needed for Beancount
310        }),
311        ..Default::default()
312    };
313
314    let server_info = lsp_types::ServerInfo {
315        name: "rledger-lsp".to_string(),
316        version: Some(crate::VERSION.to_string()),
317    };
318
319    let init_result = lsp_types::InitializeResult {
320        capabilities,
321        server_info: Some(server_info),
322    };
323
324    // Complete initialization handshake
325    connection.initialize_finish(id, serde_json::to_value(init_result)?)?;
326
327    tracing::info!("LSP initialized successfully");
328
329    // Create and run server with the handler-facing position encoding
330    // (derived above at the negotiation site).
331    let server = Server::new(connection, init_params, handler_encoding);
332    let exit_code = server.run();
333
334    // Drain the writer thread BEFORE returning. The main loop has
335    // already broken (either via the `exit` notification or because
336    // the channel closed), but the writer may still be flushing the
337    // shutdown response queued just before. Without this drain, a
338    // subsequent `process::exit` in `main()` would kill the writer
339    // mid-flush; with it, the response reaches stdout before main
340    // tears down.
341    io_threads.join()?;
342
343    Ok(exit_code)
344}
345
346#[cfg(test)]
347mod tests {
348    use super::Server;
349    use std::path::PathBuf;
350
351    #[test]
352    fn discovery_dir_prefers_workspace_and_ignores_cwd() {
353        let ws = PathBuf::from("/work/space");
354        let cwd = PathBuf::from("/tmp/launch");
355        // Workspace set → search the workspace, never the cwd (no contamination).
356        // The cwd closure must not even be invoked in this case.
357        let mut cwd_called = false;
358        let got = Server::discovery_dir(Some(ws.clone()), || {
359            cwd_called = true;
360            Some(cwd.clone())
361        });
362        assert_eq!(got, Some(ws));
363        assert!(
364            !cwd_called,
365            "cwd must not be consulted when a workspace is set"
366        );
367        // No workspace → fall back to the cwd.
368        assert_eq!(Server::discovery_dir(None, || Some(cwd.clone())), Some(cwd));
369        // Neither → nothing to discover.
370        assert_eq!(Server::discovery_dir(None, || None), None);
371    }
372}