rustledger-lsp 0.21.0

//! Main LSP server implementation.

use crate::handlers::execute_command::COMMANDS;
use crate::handlers::on_type_formatting::{FIRST_TRIGGER_CHARACTER, MORE_TRIGGER_CHARACTERS};
use crate::handlers::semantic_tokens::get_capabilities as get_semantic_tokens_capabilities;
use crate::handlers::signature_help::TRIGGER_CHARACTERS as SIGNATURE_TRIGGER_CHARACTERS;
use crate::ledger_state::{LspConfig, discover_journal_file};
use crate::main_loop::run_main_loop;
use crate::uri_to_path;
use lsp_server::Connection;
use lsp_types::InitializeParams;

/// The LSP server.
pub struct Server {
    /// Connection to the LSP client.
    connection: Connection,
    /// Initialize parameters from client.
    init_params: InitializeParams,
    /// LSP configuration parsed from init options.
    config: LspConfig,
    /// Position encoding negotiated with the client during
    /// `initialize`. Handler code emitting `Position` values must
    /// consult this so output aligns with what the client expects.
    position_encoding: crate::handlers::utils::PositionEncoding,
}

impl Server {
    /// Create a new LSP server from a connection.
    pub fn new(
        connection: Connection,
        init_params: InitializeParams,
        position_encoding: crate::handlers::utils::PositionEncoding,
    ) -> Self {
        // Parse configuration from initialization options
        let config = LspConfig::from_init_options(init_params.initialization_options.as_ref());

        if let Some(ref journal) = config.journal_file {
            tracing::info!("Journal file configured: {}", journal.display());
        }

        Self {
            connection,
            init_params,
            config,
            position_encoding,
        }
    }

    /// Run the server's main loop. Returns the exit code produced by
    /// the `exit` notification (or 0 if the channel closed without
    /// one). The caller is responsible for draining IO threads
    /// (`io_threads.join()`) before terminating the process — without
    /// that drain, the writer can lose the shutdown response queued
    /// when the loop broke.
    #[must_use]
    pub fn run(self) -> i32 {
        tracing::info!("Starting Beancount Language Server v{}", crate::VERSION);

        // Resolve journal file path relative to workspace root if needed
        let journal_file = self.resolve_journal_path();

        if let Some(ref path) = journal_file {
            tracing::info!("Using journal file: {}", path.display());
        }

        if let Some(folders) = &self.init_params.workspace_folders
            && let Some(folder) = folders.first()
        {
            tracing::info!("Workspace root: {}", folder.uri.as_str());
        }

        // Run the main event loop with the journal file configuration
        // and the negotiated position encoding (so handlers emit
        // positions in the encoding the client expects).
        let (sender, receiver) = (self.connection.sender, self.connection.receiver);
        let code = run_main_loop(receiver, sender, journal_file, self.position_encoding);

        tracing::info!("Server shutdown complete (exit code {code})");
        code
    }

    /// Resolve the journal file path, making it absolute if necessary.
    /// If no explicit journal file is configured, attempts auto-discovery.
    fn resolve_journal_path(&self) -> Option<std::path::PathBuf> {
        // Get workspace root path for resolution and discovery
        let workspace_root = self.get_workspace_root();

        // If explicit config provided, resolve it
        if let Some(journal) = &self.config.journal_file {
            return self.resolve_explicit_journal(journal, workspace_root.as_deref());
        }

        // No explicit config - auto-discover, but only in the right directory:
        // a set workspace folder is authoritative; the process cwd is used ONLY
        // when there is no workspace folder. Otherwise a stray journal in the
        // editor's launch directory silently contaminates an unrelated
        // workspace's state.
        let discovered = Self::discovery_dir(workspace_root, || std::env::current_dir().ok())
            .and_then(|dir| discover_journal_file(&dir));
        if discovered.is_none() {
            tracing::debug!("No journal file configured or discovered");
        }
        discovered
    }

    /// Directory to search for a journal during auto-discovery: the workspace
    /// folder when set, otherwise the process cwd. The cwd is deliberately
    /// *not* consulted when a workspace folder exists, so a journal that
    /// happens to sit in the editor's launch directory cannot leak into an
    /// unrelated workspace. `cwd` is computed lazily so the `current_dir`
    /// syscall is skipped entirely (and its failure ignored) when a workspace
    /// folder is set.
    fn discovery_dir(
        workspace_root: Option<std::path::PathBuf>,
        cwd: impl FnOnce() -> Option<std::path::PathBuf>,
    ) -> Option<std::path::PathBuf> {
        workspace_root.or_else(cwd)
    }

    /// Get the workspace root path from init params.
    fn get_workspace_root(&self) -> Option<std::path::PathBuf> {
        self.init_params
            .workspace_folders
            .as_ref()
            .and_then(|folders| folders.first())
            .and_then(|folder| uri_to_path(&folder.uri))
    }

    /// Resolve an explicitly configured journal path.
    fn resolve_explicit_journal(
        &self,
        journal: &std::path::Path,
        workspace_root: Option<&std::path::Path>,
    ) -> Option<std::path::PathBuf> {
        // If already absolute, use as-is
        if journal.is_absolute() {
            return Some(journal.to_path_buf());
        }

        // Try to resolve relative to workspace root
        if let Some(root) = workspace_root {
            let resolved = root.join(journal);
            if resolved.exists() {
                return Some(resolved);
            }
        }

        // Fall back to current directory
        let resolved = std::env::current_dir()
            .ok()
            .map(|cwd| cwd.join(journal))
            .filter(|p| p.exists());

        if resolved.is_none() {
            tracing::warn!(
                "Journal file '{}' not found relative to workspace or current directory",
                journal.display()
            );
        }

        resolved
    }
}

/// Start the LSP server using stdio transport.
///
/// Returns the exit code that the `exit` notification supplied (0 for
/// a clean shutdown, 1 for exit-without-prior-shutdown per LSP spec).
/// If the channel closed without an `exit` notification, returns 0.
///
/// Critically, `io_threads.join()` is called BEFORE returning, so the
/// shutdown response queued in the writer thread's channel is fully
/// flushed to stdout. Previously the production exit path called
/// `process::exit(code)` from inside the main loop, which terminated
/// the process before the writer could flush — losing the shutdown
/// response on slow runners (the `stdio_smoke` CI flake).
pub fn start_stdio() -> Result<i32, Box<dyn std::error::Error + Send + Sync>> {
    tracing::info!("Starting LSP server on stdio");

    // Create connection using stdio
    let (connection, io_threads) = Connection::stdio();

    // Wait for initialize request
    let (id, params) = connection.initialize_start()?;
    let init_params: InitializeParams = serde_json::from_value(params)?;

    // Negotiate position encoding. Our handler stack emits LSP
    // positions as UTF-8 byte offsets, so prefer UTF-8 if the
    // client advertises it (LSP 3.17+; VS Code, neovim, helix, and
    // most modern clients do). If the client doesn't advertise
    // UTF-8, the LSP spec requires the server to use UTF-16 (the
    // default), in which case our byte-based positions are wrong
    // for non-ASCII content — a server-wide latent bug tracked
    // separately. The negotiation here at least makes us correct
    // for modern clients without any handler-side conversion.
    let position_encoding = init_params
        .capabilities
        .general
        .as_ref()
        .and_then(|g| g.position_encodings.as_ref())
        .and_then(|encs| {
            encs.contains(&lsp_types::PositionEncodingKind::UTF8)
                .then_some(lsp_types::PositionEncodingKind::UTF8)
        });

    // Derive the handler-facing `PositionEncoding` once, BEFORE
    // `position_encoding` moves into the `ServerCapabilities` field
    // below. The `Option<PositionEncodingKind>` is non-Copy, so we
    // can't borrow it after the move; computing here also makes the
    // negotiated-encoding-vs-handler-encoding mapping explicit at
    // the negotiation site.
    let handler_encoding =
        crate::handlers::utils::PositionEncoding::from_negotiated(position_encoding.as_ref());

    // Build server capabilities
    let capabilities = lsp_types::ServerCapabilities {
        position_encoding,
        text_document_sync: Some(lsp_types::TextDocumentSyncCapability::Kind(
            lsp_types::TextDocumentSyncKind::FULL,
        )),
        completion_provider: Some(lsp_types::CompletionOptions {
            trigger_characters: Some(vec![
                ":".to_string(),  // Account segments
                " ".to_string(),  // After keywords
                "\"".to_string(), // Strings (payees, narrations)
                "#".to_string(),  // Tags
                "^".to_string(),  // Links
            ]),
            resolve_provider: Some(true), // Enable completion resolve for detailed info
            ..Default::default()
        }),
        definition_provider: Some(lsp_types::OneOf::Left(true)),
        references_provider: Some(lsp_types::OneOf::Left(true)),
        hover_provider: Some(lsp_types::HoverProviderCapability::Simple(true)),
        document_symbol_provider: Some(lsp_types::OneOf::Left(true)),
        semantic_tokens_provider: Some(get_semantic_tokens_capabilities()),
        code_action_provider: Some(lsp_types::CodeActionProviderCapability::Options(
            lsp_types::CodeActionOptions {
                code_action_kinds: Some(vec![
                    lsp_types::CodeActionKind::QUICKFIX,
                    lsp_types::CodeActionKind::REFACTOR,
                ]),
                resolve_provider: Some(true), // Enable resolve for lazy-loading edits
                work_done_progress_options: Default::default(),
            },
        )),
        workspace_symbol_provider: Some(lsp_types::OneOf::Left(true)),
        rename_provider: Some(lsp_types::OneOf::Right(lsp_types::RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: Default::default(),
        })),
        document_formatting_provider: Some(lsp_types::OneOf::Left(true)),
        document_range_formatting_provider: Some(lsp_types::OneOf::Left(true)),
        document_link_provider: Some(lsp_types::DocumentLinkOptions {
            resolve_provider: Some(true), // Enable resolve to verify file existence
            work_done_progress_options: Default::default(),
        }),
        inlay_hint_provider: Some(lsp_types::OneOf::Right(
            lsp_types::InlayHintServerCapabilities::Options(lsp_types::InlayHintOptions {
                resolve_provider: Some(true), // Enable resolve for rich tooltips
                work_done_progress_options: Default::default(),
            }),
        )),
        selection_range_provider: Some(lsp_types::SelectionRangeProviderCapability::Simple(true)),
        folding_range_provider: Some(lsp_types::FoldingRangeProviderCapability::Simple(true)),
        document_highlight_provider: Some(lsp_types::OneOf::Left(true)),
        linked_editing_range_provider: Some(
            lsp_types::LinkedEditingRangeServerCapabilities::Simple(true),
        ),
        document_on_type_formatting_provider: Some(lsp_types::DocumentOnTypeFormattingOptions {
            first_trigger_character: FIRST_TRIGGER_CHARACTER.to_string(),
            more_trigger_character: Some(
                MORE_TRIGGER_CHARACTERS
                    .iter()
                    .map(|s| s.to_string())
                    .collect(),
            ),
        }),
        code_lens_provider: Some(lsp_types::CodeLensOptions {
            resolve_provider: Some(true), // Enable resolve for lazy-loading balance verification
        }),
        color_provider: Some(lsp_types::ColorProviderCapability::Simple(true)),
        declaration_provider: Some(lsp_types::DeclarationCapability::Simple(true)),
        call_hierarchy_provider: Some(lsp_types::CallHierarchyServerCapability::Simple(true)),
        signature_help_provider: Some(lsp_types::SignatureHelpOptions {
            trigger_characters: Some(
                SIGNATURE_TRIGGER_CHARACTERS
                    .iter()
                    .map(|s| s.to_string())
                    .collect(),
            ),
            retrigger_characters: None,
            work_done_progress_options: Default::default(),
        }),
        execute_command_provider: Some(lsp_types::ExecuteCommandOptions {
            commands: COMMANDS.iter().map(|s| s.to_string()).collect(),
            work_done_progress_options: Default::default(),
        }),
        // Type hierarchy: advertised via experimental until lsp-types adds native support
        experimental: Some(serde_json::json!({
            "typeHierarchyProvider": true
        })),
        // Workspace capabilities
        workspace: Some(lsp_types::WorkspaceServerCapabilities {
            workspace_folders: Some(lsp_types::WorkspaceFoldersServerCapabilities {
                supported: Some(true),
                change_notifications: Some(lsp_types::OneOf::Left(true)),
            }),
            file_operations: None, // File operations (create/rename/delete) not needed for Beancount
        }),
        ..Default::default()
    };

    let server_info = lsp_types::ServerInfo {
        name: "rledger-lsp".to_string(),
        version: Some(crate::VERSION.to_string()),
    };

    let init_result = lsp_types::InitializeResult {
        capabilities,
        server_info: Some(server_info),
    };

    // Complete initialization handshake
    connection.initialize_finish(id, serde_json::to_value(init_result)?)?;

    tracing::info!("LSP initialized successfully");

    // Create and run server with the handler-facing position encoding
    // (derived above at the negotiation site).
    let server = Server::new(connection, init_params, handler_encoding);
    let exit_code = server.run();

    // Drain the writer thread BEFORE returning. The main loop has
    // already broken (either via the `exit` notification or because
    // the channel closed), but the writer may still be flushing the
    // shutdown response queued just before. Without this drain, a
    // subsequent `process::exit` in `main()` would kill the writer
    // mid-flush; with it, the response reaches stdout before main
    // tears down.
    io_threads.join()?;

    Ok(exit_code)
}

#[cfg(test)]
mod tests {
    use super::Server;
    use std::path::PathBuf;

    #[test]
    fn discovery_dir_prefers_workspace_and_ignores_cwd() {
        let ws = PathBuf::from("/work/space");
        let cwd = PathBuf::from("/tmp/launch");
        // Workspace set → search the workspace, never the cwd (no contamination).
        // The cwd closure must not even be invoked in this case.
        let mut cwd_called = false;
        let got = Server::discovery_dir(Some(ws.clone()), || {
            cwd_called = true;
            Some(cwd.clone())
        });
        assert_eq!(got, Some(ws));
        assert!(
            !cwd_called,
            "cwd must not be consulted when a workspace is set"
        );
        // No workspace → fall back to the cwd.
        assert_eq!(Server::discovery_dir(None, || Some(cwd.clone())), Some(cwd));
        // Neither → nothing to discover.
        assert_eq!(Server::discovery_dir(None, || None), None);
    }
}