wsb-daemon 0.275.7622

//! One `omni lsp` serving one analysis root.
//!
//! Started through omnicode's own daemon (`omni lsp --daemon`): one warm engine for every root
//! and every editor on the machine. The session speaks the language server protocol on the
//! child's standard streams; what the server pushes — diagnostics, progress, messages — is
//! handed on as [`Inbound`], and what it asks of its client is answered as an editor would.
//!
//! A server that exits is never started again behind anyone's back: the session reports how it
//! ended — the exit status and the end of what it wrote to standard error — and starting it
//! again is a decision ([`super::CodeIntel::restart`]).

use std::collections::{HashMap, VecDeque};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde::Serialize;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::sync::{mpsc, oneshot};

use wsb::omnicode::locate::Omni;
use wsb::omnicode::roots::AnalysisRoot;

use super::framing;

/// Lines of the server's standard error kept for when it ends.
const STDERR_LINES: usize = 200;

/// How long the server has to answer `initialize`. It may be starting omnicode's daemon.
const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(120);

/// How long a request may take. The first question after an edit waits for the analysis of
/// the edited worktree, so this is generous; a server that takes longer is not answering.
pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);

/// The language server protocol's "content modified" and "request cancelled".
pub const CONTENT_MODIFIED: i64 = -32801;
pub const REQUEST_CANCELLED: i64 = -32800;

/// Where a session's server stands.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(tag = "state", rename_all = "snake_case")]
pub enum State {
    Starting,
    /// Analyzing the root, as its progress says.
    Indexing {
        title: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        message: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        percentage: Option<u32>,
    },
    /// Answering. `note` is what it said when its analysis was ready (a language served from
    /// syntax only, and why).
    Ready {
        #[serde(skip_serializing_if = "Option::is_none")]
        note: Option<String>,
    },
    /// Its analysis failed; it answers nothing about this root.
    Failed { message: String },
    /// It exited, unasked.
    Crashed { status: String, stderr: Vec<String> },
    /// Shut down: nothing of the root is open any more.
    Stopped,
}

/// What a server pushes.
#[derive(Debug, Clone)]
pub enum Inbound {
    /// `textDocument/publishDiagnostics`, as sent.
    Diagnostics(Value),
    /// `window/showMessage`: something to tell a person.
    Message { level: &'static str, text: String },
    State(State),
    /// A new analysis of the root was published (after a save or a change on disk).
    Reanalyzed,
}

/// Why a request got no result.
#[derive(Debug, Clone, PartialEq)]
pub enum RequestFailure {
    /// The server answered with an error.
    Rpc { code: i64, message: String },
    /// The server is not running.
    Ended(String),
    TimedOut,
}

impl std::fmt::Display for RequestFailure {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Rpc { code, message } => write!(f, "{message} (error {code})"),
            Self::Ended(why) => write!(f, "the language server is not running: {why}"),
            Self::TimedOut => write!(f, "the language server did not answer within {}s", REQUEST_TIMEOUT.as_secs()),
        }
    }
}

type Pending = Arc<Mutex<HashMap<i64, oneshot::Sender<Result<Value, RequestFailure>>>>>;

pub struct Session {
    pub root: AnalysisRoot,
    writer: mpsc::UnboundedSender<Value>,
    pending: Pending,
    next_id: AtomicI64,
    state: Arc<Mutex<State>>,
    stopping: Arc<AtomicBool>,
    pub pid: u32,
}

impl std::fmt::Debug for Session {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Session").field("root", &self.root.dir).field("pid", &self.pid).finish_non_exhaustive()
    }
}

fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
    m.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
}

/// A URI of a path, as the protocol names files.
pub fn file_uri(path: &std::path::Path) -> Result<String, String> {
    url::Url::from_file_path(path).map(String::from).map_err(|()| format!("{} has no file URI", path.display()))
}

/// The protocol's predefined semantic token types (LSP 3.17).
const SEMANTIC_TOKEN_TYPES: [&str; 23] = [
    "namespace", "type", "class", "enum", "interface", "struct", "typeParameter", "parameter", "variable",
    "property", "enumMember", "event", "function", "method", "macro", "keyword", "modifier", "comment",
    "string", "number", "regexp", "operator", "decorator",
];

/// The protocol's predefined semantic token modifiers (LSP 3.17).
const SEMANTIC_TOKEN_MODIFIERS: [&str; 10] = [
    "declaration", "definition", "readonly", "static", "deprecated", "abstract", "async", "modification",
    "documentation", "defaultLibrary",
];

/// What the daemon tells a language server it is, and asks of it.
fn initialize_params(root: &AnalysisRoot) -> Result<Value, String> {
    let root_uri = file_uri(&root.dir)?;
    Ok(json!({
        "processId": std::process::id(),
        "clientInfo": { "name": "wsb", "version": env!("CARGO_PKG_VERSION") },
        "rootUri": root_uri,
        "workspaceFolders": [{ "uri": root_uri, "name": root.dir.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default() }],
        "capabilities": {
            "general": { "positionEncodings": ["utf-16"] },
            "window": { "workDoneProgress": true, "showMessage": {} },
            "workspace": { "didChangeWatchedFiles": { "dynamicRegistration": false } },
            "textDocument": {
                "synchronization": { "didSave": true },
                "publishDiagnostics": { "versionSupport": true },
                "hover": { "contentFormat": ["markdown"] },
                // The protocol requires the token types and modifiers a client knows; the
                // workbench colours whatever legend the server answers with.
                "semanticTokens": {
                    "requests": { "full": true },
                    "tokenTypes": SEMANTIC_TOKEN_TYPES,
                    "tokenModifiers": SEMANTIC_TOKEN_MODIFIERS,
                    "formats": ["relative"],
                },
            },
        },
        "initializationOptions": { "exclude": root.exclude },
    }))
}

impl Session {
    /// Start `omni lsp --daemon` over `root` and initialize it. What it pushes goes to
    /// `inbound`, tagged with the root's directory.
    pub async fn start(
        omni: &Omni,
        root: AnalysisRoot,
        inbound: mpsc::UnboundedSender<(PathBuf, Inbound)>,
    ) -> Result<Arc<Self>, String> {
        use std::process::Stdio;
        let mut child = wsb::git_backend::command_without_repo_scope(&omni.path)
            .args(["lsp", "--daemon"])
            .current_dir(&root.dir)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|e| format!("cannot start {}: {e}", omni.path.display()))?;
        let pid = child.id();
        let (Some(stdin), Some(stdout), Some(stderr)) = (child.stdin.take(), child.stdout.take(), child.stderr.take())
        else {
            let _ = child.kill();
            return Err("the language server's standard streams are not available".into());
        };
        let stdin = tokio::process::ChildStdin::from_std(stdin).map_err(|e| e.to_string())?;
        let stdout = tokio::process::ChildStdout::from_std(stdout).map_err(|e| e.to_string())?;
        let stderr = tokio::process::ChildStderr::from_std(stderr).map_err(|e| e.to_string())?;

        let (writer, mut outgoing) = mpsc::unbounded_channel::<Value>();
        let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
        let state = Arc::new(Mutex::new(State::Starting));
        let stopping = Arc::new(AtomicBool::new(false));
        let tail: Arc<Mutex<VecDeque<String>>> = Arc::new(Mutex::new(VecDeque::new()));
        let dir = root.dir.clone();

        // Writing: one message at a time, in the order they were sent.
        tokio::spawn(async move {
            let mut stdin = stdin;
            while let Some(message) = outgoing.recv().await {
                if framing::write(&mut stdin, &message).await.is_err() {
                    break;
                }
            }
        });

        // Standard error: logged, and its end kept for when the server ends.
        {
            let tail = tail.clone();
            let dir = dir.clone();
            tokio::spawn(async move {
                let mut lines = BufReader::new(stderr).lines();
                while let Ok(Some(line)) = lines.next_line().await {
                    tracing::debug!(target: "wsb_daemon::omni", root = %dir.display(), "{line}");
                    let mut t = lock(&tail);
                    if t.len() == STDERR_LINES {
                        t.pop_front();
                    }
                    t.push_back(line);
                }
            });
        }

        // Reading: responses to their requests, the server's requests answered, the rest
        // handed on.
        {
            let pending = pending.clone();
            let writer = writer.clone();
            let state = state.clone();
            let inbound = inbound.clone();
            let dir = dir.clone();
            tokio::spawn(async move {
                let mut reader = BufReader::new(stdout);
                let mut loading: Option<Value> = None;
                let mut reanalysis: Option<Value> = None;
                loop {
                    let message = match framing::read(&mut reader).await {
                        Ok(Some(m)) => m,
                        Ok(None) => break,
                        Err(e) => {
                            tracing::warn!(root = %dir.display(), "the language server's output is unreadable: {e}");
                            break;
                        }
                    };
                    let id = message.get("id").cloned();
                    match (message.get("method").and_then(Value::as_str), id) {
                        (None, Some(id)) => {
                            let Some(id) = id.as_i64() else { continue };
                            let Some(waiter) = lock(&pending).remove(&id) else { continue };
                            let result = match message.get("error") {
                                Some(e) => Err(RequestFailure::Rpc {
                                    code: e["code"].as_i64().unwrap_or(0),
                                    message: e["message"].as_str().unwrap_or("").to_owned(),
                                }),
                                None => Ok(message.get("result").cloned().unwrap_or(Value::Null)),
                            };
                            let _ = waiter.send(result);
                        }
                        (Some(method), Some(id)) => {
                            // What a server asks of an editor: progress tokens and
                            // registrations are accepted; nothing else is offered.
                            let reply = match method {
                                "window/workDoneProgress/create" | "client/registerCapability" => {
                                    json!({ "jsonrpc": "2.0", "id": id, "result": null })
                                }
                                other => json!({ "jsonrpc": "2.0", "id": id, "error": {
                                    "code": -32601, "message": format!("the workbench does not answer {other}")
                                }}),
                            };
                            let _ = writer.send(reply);
                        }
                        (Some(method), None) => {
                            let params = message.get("params").cloned().unwrap_or(Value::Null);
                            let pushed = match method {
                                "textDocument/publishDiagnostics" => Some(Inbound::Diagnostics(params)),
                                "window/showMessage" => Some(Inbound::Message {
                                    level: match params["type"].as_i64() {
                                        Some(1) => "error",
                                        Some(2) => "warning",
                                        _ => "info",
                                    },
                                    text: params["message"].as_str().unwrap_or("").to_owned(),
                                }),
                                "$/progress" => progress(&params, &mut loading, &mut reanalysis, &state),
                                _ => None,
                            };
                            if let Some(p) = pushed {
                                let _ = inbound.send((dir.clone(), p));
                            }
                        }
                        (None, None) => {}
                    }
                }
            });
        }

        // The end: whoever waits is told, and the state says how it ended.
        {
            let pending = pending.clone();
            let state = state.clone();
            let stopping = stopping.clone();
            let inbound = inbound.clone();
            let dir = dir.clone();
            tokio::task::spawn_blocking(move || {
                let status = child.wait().map(|s| s.to_string()).unwrap_or_else(|e| format!("unknown ({e})"));
                // Give the readers a moment to take the last of its output.
                std::thread::sleep(Duration::from_millis(100));
                let ended = if stopping.load(Ordering::SeqCst) {
                    State::Stopped
                } else {
                    State::Crashed { status, stderr: lock(&tail).iter().cloned().collect() }
                };
                for (_, waiter) in lock(&pending).drain() {
                    let _ = waiter.send(Err(RequestFailure::Ended("it exited".into())));
                }
                *lock(&state) = ended.clone();
                let _ = inbound.send((dir, Inbound::State(ended)));
            });
        }

        let session = Arc::new(Self {
            root,
            writer,
            pending,
            next_id: AtomicI64::new(0),
            state,
            stopping,
            pid,
        });
        let init = initialize_params(&session.root)?;
        match tokio::time::timeout(INITIALIZE_TIMEOUT, session.request_raw("initialize", init)).await {
            Ok(Ok(_)) => {}
            Ok(Err(e)) => return Err(format!("the language server refused to start: {e}")),
            Err(_) => return Err(format!("the language server did not initialize within {}s", INITIALIZE_TIMEOUT.as_secs())),
        }
        session.notify("initialized", json!({}))?;
        Ok(session)
    }

    pub fn state(&self) -> State {
        lock(&self.state).clone()
    }

    /// Whether it can be asked anything.
    pub fn running(&self) -> bool {
        !matches!(self.state(), State::Crashed { .. } | State::Stopped)
    }

    pub fn notify(&self, method: &str, params: Value) -> Result<(), String> {
        self.writer
            .send(json!({ "jsonrpc": "2.0", "method": method, "params": params }))
            .map_err(|_| "the language server is not running".to_owned())
    }

    fn request_raw(&self, method: &str, params: Value) -> impl std::future::Future<Output = Result<Value, RequestFailure>> + '_ {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst) + 1;
        let (tx, rx) = oneshot::channel();
        lock(&self.pending).insert(id, tx);
        let sent = self.writer.send(json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }));
        let guard = CancelOnDrop { id, session: self, done: false };
        async move {
            let mut guard = guard;
            if sent.is_err() {
                guard.done = true;
                lock(&self.pending).remove(&id);
                return Err(RequestFailure::Ended("its input is closed".into()));
            }
            let answer = rx.await.unwrap_or(Err(RequestFailure::Ended("it exited".into())));
            guard.done = true;
            answer
        }
    }

    /// Ask the server; a request whose caller gives up (the page went away) is canceled.
    pub async fn request(&self, method: &str, params: Value) -> Result<Value, RequestFailure> {
        match tokio::time::timeout(REQUEST_TIMEOUT, self.request_raw(method, params)).await {
            Ok(answer) => answer,
            Err(_) => Err(RequestFailure::TimedOut),
        }
    }

    /// Shut the server down: asked, then told to exit, then stopped if it has not.
    pub async fn shutdown(&self) {
        self.stopping.store(true, Ordering::SeqCst);
        let _ = tokio::time::timeout(Duration::from_secs(5), self.request_raw("shutdown", Value::Null)).await;
        let _ = self.notify("exit", Value::Null);
        // Closing its input ends one that ignored `exit`.
        tokio::time::sleep(Duration::from_secs(2)).await;
        if self.running() {
            wsb::ops::process::signal(&[self.pid], "TERM");
        }
    }
}

/// Cancels a request whose caller stopped waiting.
struct CancelOnDrop<'a> {
    id: i64,
    session: &'a Session,
    done: bool,
}

impl Drop for CancelOnDrop<'_> {
    fn drop(&mut self) {
        if !self.done && lock(&self.session.pending).remove(&self.id).is_some() {
            let _ = self.session.notify("$/cancelRequest", json!({ "id": self.id }));
        }
    }
}

/// What a `$/progress` notification means for the session's state.
fn progress(params: &Value, loading: &mut Option<Value>, reanalysis: &mut Option<Value>, state: &Mutex<State>) -> Option<Inbound> {
    let token = params["token"].clone();
    let value = &params["value"];
    match value["kind"].as_str() {
        Some("begin") => {
            let title = value["title"].as_str().unwrap_or("").to_owned();
            if title.contains("re-analyzing") {
                *reanalysis = Some(token);
                return None;
            }
            *loading = Some(token);
            let s = State::Indexing { title, message: None, percentage: None };
            *lock(state) = s.clone();
            Some(Inbound::State(s))
        }
        Some("report") if loading.as_ref() == Some(&token) => {
            let mut s = lock(state);
            if let State::Indexing { message, percentage, .. } = &mut *s {
                *message = value["message"].as_str().map(str::to_owned);
                *percentage = value["percentage"].as_u64().map(|p| p as u32);
            }
            Some(Inbound::State(s.clone()))
        }
        Some("end") if loading.as_ref() == Some(&token) => {
            *loading = None;
            let message = value["message"].as_str().unwrap_or("").to_owned();
            let s = match message.strip_prefix("analysis failed: ") {
                Some(why) => State::Failed { message: why.to_owned() },
                None => State::Ready {
                    note: message.strip_prefix("analysis ready; ").map(str::to_owned).filter(|n| !n.is_empty()),
                },
            };
            *lock(state) = s.clone();
            Some(Inbound::State(s))
        }
        Some("end") if reanalysis.as_ref() == Some(&token) => {
            *reanalysis = None;
            Some(Inbound::Reanalyzed)
        }
        _ => None,
    }
}

#[cfg(test)]
#[path = "session_tests.rs"]
mod tests;