aether-lspd 0.3.4

use crate::daemon::Timeouts;
use crate::diagnostics_store::{DiagnosticsStore, Freshness};
use crate::document_lifecycle::{AcquireAction, DocumentLifecycle, ReleaseAction};
use crate::error::DaemonResult;
use crate::file_watcher::{FileWatcherBatch, is_ignored_directory_name};
use crate::language_catalog::{DiagnosticsMode, LanguageId, LspConfig};
use crate::process_transport::{LspNotification, ProcessTransport, TransportError, TransportEvent};
use crate::protocol::extract_document_uri;
use crate::refresh_queue::RefreshQueue;
use ignore::WalkBuilder;
use lsp_types::notification::{
    DidChangeWatchedFiles, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Notification,
};
use lsp_types::request::{DocumentDiagnosticRequest, Request};
use lsp_types::{
    DidChangeWatchedFilesParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams, DidSaveTextDocumentParams,
    DocumentDiagnosticReport, DocumentDiagnosticReportResult, PublishDiagnosticsParams, TextDocumentIdentifier,
    TextDocumentItem, Uri,
};
use serde_json::Value;
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::{Mutex, mpsc};

const LSP_CONTENT_MODIFIED: i32 = -32801;
const LSP_SERVER_CANCELLED: i32 = -32802;
const TRANSIENT_RETRY_LIMIT: u32 = 3;
const TRANSIENT_RETRY_DELAY: Duration = Duration::from_millis(500);

#[derive(Clone)]
pub(crate) struct WorkspaceSession {
    transport: ProcessTransport,
    documents: DocumentLifecycle,
    document_sync: Arc<Mutex<()>>,
    diagnostics: DiagnosticsStore,
    refresh: RefreshQueue,
    diagnostics_mode: DiagnosticsMode,
    diagnostics_timeout: Duration,
    supported_extensions: Arc<HashSet<String>>,
    alive: Arc<AtomicBool>,
    #[cfg(feature = "testing")]
    hooks: crate::testing::TestHooks,
}

impl WorkspaceSession {
    pub(crate) fn spawn(
        workspace_root: &Path,
        config: &LspConfig,
        timeouts: Timeouts,
        #[cfg(feature = "testing")] hooks: crate::testing::TestHooks,
    ) -> DaemonResult<Self> {
        let (transport, event_rx) = ProcessTransport::spawn(workspace_root, config)?;
        let session = Self {
            transport,
            documents: DocumentLifecycle::new(),
            document_sync: Arc::default(),
            diagnostics: DiagnosticsStore::new(timeouts.diagnostics_settle),
            refresh: RefreshQueue::new(),
            diagnostics_mode: config.diagnostics_mode,
            diagnostics_timeout: timeouts.diagnostics,
            supported_extensions: Arc::new(supported_extensions(config)),
            alive: Arc::new(AtomicBool::new(true)),
            #[cfg(feature = "testing")]
            hooks,
        };

        tokio::spawn(session.clone().run_events(event_rx));
        tokio::spawn(session.clone().run_refresh_worker());
        tokio::spawn(session.clone().bootstrap_refresh(workspace_root.to_path_buf()));

        Ok(session)
    }

    pub(crate) async fn wait_until_initialized(&self) -> DaemonResult<()> {
        self.transport.wait_until_initialized().await
    }

    pub(crate) async fn wait_until_ready(&self, bound: Duration) {
        self.transport.wait_until_ready(bound).await;
    }

    pub(crate) async fn request(&self, method: &str, params: &Value) -> Result<Value, TransportError> {
        let Some(uri) = extract_document_uri(method, params) else {
            return self.request_with_retry(method, params).await;
        };

        self.open_document(&uri).await;
        let result = self.request_with_retry(method, params).await;
        self.close_document(&uri).await;
        result
    }

    pub(crate) fn queue_diagnostic_refresh(&self, uri: Uri) {
        self.refresh.enqueue(vec![uri]);
    }

    pub(crate) async fn get_diagnostics(&self, uri: Option<&Uri>) -> Vec<PublishDiagnosticsParams> {
        match uri {
            Some(uri) => {
                if self.refresh_uri(uri).await == Freshness::Unconfirmed {
                    self.refresh.wait_for_current_generation(self.diagnostics_timeout).await;
                }
            }
            None => self.refresh.wait_for_current_generation(self.diagnostics_timeout).await,
        }
        self.diagnostics.get(uri)
    }

    pub(crate) async fn shutdown(&self) {
        self.refresh.shutdown();
        self.transport.shutdown().await;
    }

    /// Whether the language server behind this session is still usable.
    pub(crate) fn is_alive(&self) -> bool {
        self.alive.load(Ordering::SeqCst)
    }

    /// Mark this session dead so the registry replaces it on the next request.
    pub(crate) fn mark_dead(&self) {
        self.alive.store(false, Ordering::SeqCst);
    }

    /// Declare the language server wedged: mark the session dead and kill the
    /// server process so blocked pipes unwind and pending requests fail.
    pub(crate) fn declare_wedged(&self) {
        if self.alive.swap(false, Ordering::SeqCst) {
            self.transport.kill_process();
        }
    }

    async fn refresh_uri(&self, uri: &Uri) -> Freshness {
        let sync = self.open_document(uri).await;
        let freshness = self.settle_diagnostics(uri, sync).await;
        self.close_document(uri).await;
        freshness
    }

    async fn open_document(&self, uri: &Uri) -> DocumentSync {
        let sync = self.document_sync.lock();
        #[cfg(feature = "testing")]
        let sync = self.hooks.observe_pending(crate::testing::TestPoint::DocumentSyncWait(uri.clone()), sync);
        let _sync = sync.await;
        let notifications = match self.documents.acquire(uri).await {
            AcquireAction::Open { file_path, content } => open_and_save_notifications(uri, &file_path, content),
            AcquireAction::Reopen { file_path, content } => reopen_notifications(uri, &file_path, content),
            AcquireAction::Unchanged => return DocumentSync::AlreadyOpen,
            AcquireAction::MissingOnDisk => {
                self.documents.forget_uri(uri);
                self.diagnostics.forget_uri(uri);
                return DocumentSync::Missing;
            }
        };

        #[cfg(feature = "testing")]
        self.hooks.pause(crate::testing::TestPoint::DocumentOpen(uri.clone())).await;
        let version_before = self.diagnostics.current_uri_version(uri);
        for notification in notifications {
            self.transport.send_notification(notification).await;
        }
        DocumentSync::Sent { version_before }
    }

    async fn close_document(&self, uri: &Uri) {
        let _sync = self.document_sync.lock().await;
        match self.documents.release(uri) {
            ReleaseAction::Close => self.transport.send_notification(close_notification(uri)).await,
            ReleaseAction::CloseAndRefresh => {
                self.transport.send_notification(close_notification(uri)).await;
                self.refresh.enqueue(vec![uri.clone()]);
            }
            ReleaseAction::Unchanged => {}
        }
    }

    async fn settle_diagnostics(&self, uri: &Uri, sync: DocumentSync) -> Freshness {
        match (self.diagnostics_mode, sync) {
            (_, DocumentSync::Missing) | (DiagnosticsMode::Push, DocumentSync::AlreadyOpen) => Freshness::Unconfirmed,
            (DiagnosticsMode::Push, DocumentSync::Sent { version_before }) => {
                let fresh = self.diagnostics.wait_for_uri_fresh(uri, version_before, self.diagnostics_timeout);
                #[cfg(feature = "testing")]
                let fresh = self.hooks.observe_pending(crate::testing::TestPoint::DiagnosticsWait(uri.clone()), fresh);
                fresh.await
            }
            (DiagnosticsMode::Pull, DocumentSync::Sent { .. } | DocumentSync::AlreadyOpen) => {
                self.pull_diagnostics(uri).await
            }
        }
    }

    async fn pull_diagnostics(&self, uri: &Uri) -> Freshness {
        // Built by hand: lsp-types serializes `identifier` and `previousResultId` as explicit nulls,
        // which TypeScript 7 rejects.
        let params = serde_json::json!({ "textDocument": TextDocumentIdentifier { uri: uri.clone() } });
        let request = self.request_with_retry(DocumentDiagnosticRequest::METHOD, &params);
        let report = tokio::time::timeout(self.diagnostics_timeout, request)
            .await
            .ok()
            .and_then(Result::ok)
            .and_then(|value| serde_json::from_value(value).ok());

        let Some(DocumentDiagnosticReportResult::Report(DocumentDiagnosticReport::Full(report))) = report else {
            tracing::warn!(uri = uri.as_str(), "Language server returned no full diagnostic report");
            return Freshness::Unconfirmed;
        };
        self.diagnostics.publish(PublishDiagnosticsParams {
            uri: uri.clone(),
            diagnostics: report.full_document_diagnostic_report.items,
            version: None,
        });
        Freshness::Fresh
    }

    async fn request_with_retry(&self, method: &str, params: &Value) -> Result<Value, TransportError> {
        let mut retries = 0;
        loop {
            match self.transport.request_raw(method, params.clone()).await {
                Err(TransportError::Lsp(err))
                    if matches!(err.code, LSP_CONTENT_MODIFIED | LSP_SERVER_CANCELLED)
                        && retries < TRANSIENT_RETRY_LIMIT =>
                {
                    retries += 1;
                    tokio::time::sleep(TRANSIENT_RETRY_DELAY).await;
                }
                result => return result,
            }
        }
    }

    async fn run_events(self, mut event_rx: mpsc::Receiver<TransportEvent>) {
        while let Some(event) = event_rx.recv().await {
            match event {
                #[cfg(feature = "testing")]
                TransportEvent::RequestReceived(method) => {
                    self.hooks.pause(crate::testing::TestPoint::RequestReceived(method)).await;
                }
                TransportEvent::PublishedDiagnostics(params) => self.diagnostics.publish(params),
                TransportEvent::DiagnosticsRefreshRequested => {
                    self.refresh.enqueue(self.filter_supported(self.diagnostics.uris()));
                }
                TransportEvent::FileWatcherBatch(batch) => self.forward_watcher_batch(batch).await,
                TransportEvent::Closed => break,
            }
        }

        self.mark_dead();
        self.refresh.shutdown();
    }

    async fn forward_watcher_batch(&self, batch: FileWatcherBatch) {
        let changes = self.documents.filter_watcher_changes(batch.forwarded_changes);
        let mut refresh_uris = self.filter_supported(changes.iter().map(|change| change.uri.clone()).collect());
        refresh_uris.extend(self.filter_supported(batch.discovered_uris));
        self.refresh.enqueue(refresh_uris);

        if changes.is_empty() {
            return;
        }

        self.transport
            .send_notification(notification::<DidChangeWatchedFiles>(DidChangeWatchedFilesParams { changes }))
            .await;
    }

    async fn run_refresh_worker(self) {
        while let Some(uri) = self.refresh.recv().await {
            self.refresh_uri(&uri).await;
        }
    }

    async fn bootstrap_refresh(self, workspace_root: PathBuf) {
        let supported_extensions = Arc::clone(&self.supported_extensions);
        let uris = tokio::task::spawn_blocking(move || discover_supported_uris(&workspace_root, &supported_extensions))
            .await
            .unwrap_or_default();

        self.refresh.enqueue(uris);
        self.refresh.complete_bootstrap();
    }

    fn filter_supported(&self, uris: Vec<Uri>) -> Vec<Uri> {
        uris.into_iter().filter(|uri| uri_is_supported(uri, &self.supported_extensions)).collect()
    }
}

enum DocumentSync {
    Sent { version_before: u64 },
    AlreadyOpen,
    Missing,
}

fn supported_extensions(config: &LspConfig) -> HashSet<String> {
    config
        .languages
        .iter()
        .flat_map(|language| language.metadata().extensions.iter().copied())
        .map(ToOwned::to_owned)
        .collect()
}

fn discover_supported_uris(workspace_root: &Path, supported_extensions: &HashSet<String>) -> Vec<Uri> {
    if supported_extensions.is_empty() {
        return Vec::new();
    }

    let mut builder = WalkBuilder::new(workspace_root);
    builder
        .standard_filters(true)
        .filter_entry(|entry| entry.depth() == 0 || !is_ignored_directory_name(entry.file_name()));

    builder
        .build()
        .filter_map(Result::ok)
        .filter(|entry| entry.file_type().is_some_and(|file_type| file_type.is_file()))
        .filter(|entry| path_is_supported(entry.path(), supported_extensions))
        .filter_map(|entry| crate::path_to_uri(entry.path()).ok())
        .collect()
}

fn uri_is_supported(uri: &Uri, supported_extensions: &HashSet<String>) -> bool {
    let path = crate::uri_to_path(uri);
    path_is_supported(Path::new(&path), supported_extensions)
}

fn path_is_supported(path: &Path, supported_extensions: &HashSet<String>) -> bool {
    path.extension().and_then(|ext| ext.to_str()).is_some_and(|ext| supported_extensions.contains(ext))
}

fn open_and_save_notifications(uri: &Uri, file_path: &str, content: String) -> Vec<LspNotification> {
    vec![open_notification(uri, file_path, 1, content), save_notification(uri)]
}

fn reopen_notifications(uri: &Uri, file_path: &str, content: String) -> Vec<LspNotification> {
    vec![close_notification(uri), open_notification(uri, file_path, 1, content), save_notification(uri)]
}

fn open_notification(uri: &Uri, file_path: &str, version: i32, content: String) -> LspNotification {
    let language_id = LanguageId::from_path(Path::new(file_path)).as_str().to_string();
    notification::<DidOpenTextDocument>(DidOpenTextDocumentParams {
        text_document: TextDocumentItem::new(uri.clone(), language_id, version, content),
    })
}

fn save_notification(uri: &Uri) -> LspNotification {
    notification::<DidSaveTextDocument>(DidSaveTextDocumentParams {
        text_document: TextDocumentIdentifier::new(uri.clone()),
        text: None,
    })
}

fn close_notification(uri: &Uri) -> LspNotification {
    notification::<DidCloseTextDocument>(DidCloseTextDocumentParams {
        text_document: TextDocumentIdentifier::new(uri.clone()),
    })
}

fn notification<N: Notification>(params: N::Params) -> LspNotification {
    LspNotification { method: N::METHOD.to_string(), params: serde_json::to_value(params).unwrap() }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn open_and_save_notifications_emit_open_then_save() {
        let uri: Uri = "file:///workspace/main.rs".parse().unwrap();
        let notifications = open_and_save_notifications(&uri, "/workspace/main.rs", "fn main() {}\n".to_string());

        assert_eq!(notifications.len(), 2);
        assert_eq!(notifications[0].method, DidOpenTextDocument::METHOD);
        assert_eq!(notifications[1].method, DidSaveTextDocument::METHOD);
    }

    #[test]
    fn reopen_notifications_emit_close_open_save() {
        let uri: Uri = "file:///workspace/main.rs".parse().unwrap();
        let notifications = reopen_notifications(&uri, "/workspace/main.rs", "fn main() {}\n".to_string());

        assert_eq!(notifications.len(), 3);
        assert_eq!(notifications[0].method, DidCloseTextDocument::METHOD);
        assert_eq!(notifications[1].method, DidOpenTextDocument::METHOD);
        assert_eq!(notifications[2].method, DidSaveTextDocument::METHOD);
    }
}