magi-code 0.96.2

use crate::{
    cancellation::AgentCancellation,
    config::LspServerConfig,
    lsp::{
        DiagnosticsStore, DocumentUri, DocumentVersion, DocumentVersions, LspDiagnostic,
        stdio::LspStdioTransport,
    },
};
use serde::Deserialize;
use serde_json::json;
use std::{
    path::Path,
    time::{Duration, Instant},
};

const SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(500);

pub(crate) struct LspClient {
    language_id: String,
    transport: LspStdioTransport,
    diagnostics: DiagnosticsStore,
    versions: DocumentVersions,
}

impl LspClient {
    pub(crate) fn start_with_timeout(
        _server_id: String,
        config: &LspServerConfig,
        workspace_root: &Path,
        language_id: String,
        initialize_timeout: Duration,
    ) -> anyhow::Result<Self> {
        let transport = LspStdioTransport::spawn(config, workspace_root)?;
        let root_uri = DocumentUri::from_path(workspace_root)?;
        let params = json!({
            "processId": std::process::id(),
            "rootUri": root_uri.as_str(),
            "capabilities": {
                "textDocument": {
                    "synchronization": {"didSave": false, "willSave": false, "willSaveWaitUntil": false},
                    "publishDiagnostics": {"versionSupport": true},
                },
                "workspace": {"configuration": true}
            },
            "clientInfo": {"name": "magi-code"}
        });
        transport.request(
            "initialize",
            params,
            initialize_timeout,
            &AgentCancellation::default(),
        )?;
        transport.notify("initialized", json!({}))?;
        Ok(Self {
            language_id,
            transport,
            diagnostics: DiagnosticsStore::default(),
            versions: DocumentVersions::default(),
        })
    }

    pub(crate) fn sync_full_document(
        &mut self,
        path: &Path,
        text: String,
        deadline: Instant,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<(DocumentVersion, Instant)> {
        cancellation.check()?;
        let sync_started_at = Instant::now();
        if sync_started_at >= deadline {
            anyhow::bail!("LSP sync deadline elapsed");
        }
        let uri = DocumentUri::from_path(path)?;
        let version = self.versions.next_for(path);
        if version.0 == 1 {
            self.transport.notify(
                "textDocument/didOpen",
                json!({
                    "textDocument": {
                        "uri": uri.as_str(),
                        "languageId": self.language_id,
                        "version": version.0,
                        "text": text
                    }
                }),
            )?;
        } else {
            self.transport.notify(
                "textDocument/didChange",
                json!({
                    "textDocument": {"uri": uri.as_str(), "version": version.0},
                    "contentChanges": [{"text": text}]
                }),
            )?;
        }
        cancellation.check()?;
        Ok((version, sync_started_at))
    }

    pub(crate) fn wait_for_fresh_diagnostics(
        &mut self,
        path: &Path,
        min_version: DocumentVersion,
        sync_started_at: Instant,
        deadline: Instant,
        cancellation: &AgentCancellation,
    ) -> Option<Vec<LspDiagnostic>> {
        let uri = DocumentUri::from_path(path).ok()?;
        loop {
            let _ = cancellation.check();
            self.drain_notifications();
            if let Some(snapshot) = self
                .diagnostics
                .fresh_for(&uri, min_version, sync_started_at)
            {
                return Some(snapshot.diagnostics.clone());
            }
            if Instant::now() >= deadline || cancellation.is_canceled() {
                return None;
            }
            std::thread::sleep(Duration::from_millis(10));
        }
    }

    pub(crate) fn drain_pending_notifications(&mut self) {
        self.drain_notifications();
    }

    pub(crate) fn is_closed(&self) -> bool {
        self.transport.is_closed()
    }

    pub(crate) fn shutdown(mut self) {
        self.transport.shutdown(SHUTDOWN_TIMEOUT);
    }

    fn drain_notifications(&mut self) {
        while let Some(notification) = self.transport.try_recv_notification() {
            if notification.method == "textDocument/publishDiagnostics" {
                let Some(params) = notification.params else {
                    continue;
                };
                if let Ok(params) = serde_json::from_value::<PublishDiagnosticsParams>(params) {
                    self.diagnostics.update(
                        params.uri,
                        params.version.map(DocumentVersion),
                        params.diagnostics,
                        Instant::now(),
                    );
                }
            }
        }
    }
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct PublishDiagnosticsParams {
    uri: DocumentUri,
    #[serde(default)]
    version: Option<i32>,
    #[serde(default)]
    diagnostics: Vec<LspDiagnostic>,
}