laburnum 1.17.1

An LSP framework for building language servers and compilers, powered by an incremental query tree with content-addressed storage, task-based dataflow, and parallel queries.
Documentation
// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

mod ariadne_diagnostics;
mod db_stats;
mod query_records;
mod query_source_cache;
mod query_source_keys;

use {
  crate::{
    Partitions,
    TRACER,
    protocol::{
      lsp::LanguageServer,
      task::RpcTask,
    },
    scheduler::task::TaskContext,
  },
  otel::span,
  serde_json::Value,
};

#[cfg(feature = "testing-commands")]
use crate::protocol::jsonrpc;

async fn handle_custom_command<P: Partitions, T: LanguageServer<P>>(
  command: &str,
  arguments: &[Value],
  ctx: &mut TaskContext<P, T>,
) -> jsonrpc::Result<Option<Value>> {
  match command {
    | "laburnum/dbStats" => db_stats::handle_db_stats(arguments, ctx).await,
    | "laburnum/querySourceCache" => {
      query_source_cache::handle_query_source_cache(arguments, ctx).await
    },
    | "laburnum/queryRecords" => {
      query_records::handle_query_records(arguments, ctx).await
    },
    | "laburnum/querySourceKeys" => {
      query_source_keys::handle_query_source_keys(arguments, ctx).await
    },
    | "laburnum/ariadneDiagnostics" => {
      ariadne_diagnostics::handle_ariadne_diagnostics(arguments, ctx).await
    },
    | _ => Err(jsonrpc::Error::method_not_found()),
  }
}

impl<P: Partitions, T: LanguageServer<P>> RpcTask<P, T> {
  #[cfg(feature = "testing-commands")]
  pub(super) fn try_handle_laburnum_command(
    &mut self,
    method: &str,
    params: &serde_json::Value,
    request_id: jsonrpc::Id,
  ) -> bool {
    use crate::scheduler::lanes::SYNC_LANE;

    if method != "workspace/executeCommand" {
      return false;
    }

    let command = match params.get("command").and_then(|v| v.as_str()) {
      | Some(cmd) if cmd.starts_with("laburnum/") => cmd.to_string(),
      | _ => return false,
    };

    let arguments = params
      .get("arguments")
      .and_then(|v| v.as_array())
      .cloned()
      .unwrap_or_default();

    let cx = span!(
      ^ @TRACER,
      "laburnum.commands.try_handle",
      "command" = command.clone(),
      "method" = method.to_string(),
      "request_id" = request_id.to_string()
    );

    let tx = self.response_tx.clone();
    let request_id_clone = request_id.clone();

    self.ctx.spawn_task(
      move |mut ctx| -> _ {
        use opentelemetry::trace::FutureExt;

        async move {
          let result =
            handle_custom_command(&command, &arguments, &mut ctx).await;

          let response = match result {
            | Ok(value) => {
              jsonrpc::Response::from_ok(
                request_id_clone.clone(),
                value.unwrap_or(serde_json::Value::Null),
              )
            },
            | Err(error) => {
              jsonrpc::Response::from_error(request_id_clone.clone(), error)
            },
          };

          if let Err(_e) = tx.send(response).await {}

          None
        }
        .with_context(cx)
      },
      SYNC_LANE,
    );

    true
  }
}