laburnum 1.17.1

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

//! Tool trait and registry for the MCP server.
//!
//! Tools are zero-sized marker types with typed associated `Input` /
//! `Output`, mirroring the [`Request`](crate::connect::lsp::request::Request)
//! pattern used by the LSP client. `Input` and `Output` reuse LSP type
//! definitions wherever possible — see [`crate::mcp::tools`] for the
//! built-in catalog.

use {
  crate::{
    connect::lsp::{
      LspClient,
      errors::LspClientError,
    },
    protocol::mcp::{
      CallToolResult,
      ToolDescriptor,
    },
  },
  std::{
    collections::HashMap,
    sync::Arc,
  },
};

/// Errors a [`Tool`] body can surface as a `CallToolResult` with
/// `isError: true`.
///
/// Distinct from protocol-level errors (unknown tool, bad arguments) which
/// the dispatcher converts to JSON-RPC error responses.
#[derive(Debug)]
pub enum ToolError {
  /// The underlying LSP request failed.
  Lsp(LspClientError),
  /// Tool-specific failure with a human-readable message.
  Other(String),
}

impl From<LspClientError> for ToolError {
  fn from(value: LspClientError) -> Self {
    Self::Lsp(value)
  }
}

impl std::fmt::Display for ToolError {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    match self {
      | Self::Lsp(e) => write!(f, "LSP request failed: {e}"),
      | Self::Other(s) => f.write_str(s),
    }
  }
}

/// An MCP tool the server exposes to agents.
///
/// `Tool` is a marker (no `&self`) so implementations can be zero-sized
/// types, matching [`crate::connect::lsp::request::Request`]. The framework
/// constructs them via [`ToolRegistry::register::<T>()`].
///
/// # Reuse of LSP types
///
/// `Input` and `Output` should be the corresponding LSP types directly
/// (e.g. `lsp::HoverParams` and `Option<lsp::Hover>`) when a tool wraps a
/// single LSP request. The agent-facing JSON shape is whatever the LSP
/// type's serde representation already produces. For tools that chain
/// multiple LSP requests (e.g. call hierarchy), define small composite
/// types using LSP primitives (`Uri`, `Position`, `Range`,
/// `CallHierarchyItem`).
pub trait Tool: Send + Sync + 'static {
  /// Agent-supplied input. Typically an LSP `*Params` type.
  type Input: serde::de::DeserializeOwned + Send + Sync + 'static;

  /// Returned value. Typically the LSP `Request::Result` type.
  type Output: serde::Serialize + Send + Sync + 'static;

  /// MCP-visible tool name. Must be unique within a registry.
  const NAME: &'static str;

  /// Human-readable description shown to the agent in `tools/list`.
  const DESCRIPTION: &'static str;

  /// JSON Schema for [`Self::Input`].
  ///
  /// Built by composing fragments from [`crate::mcp::schema`].
  fn input_schema() -> serde_json::Value;

  /// Execute the tool.
  fn call(
    client: &LspClient,
    input: Self::Input,
  ) -> impl std::future::Future<Output = Result<Self::Output, ToolError>> + Send;
}

/// Type-erased wrapper so heterogeneous `Tool` impls can live in one map.
///
/// The single (de)serialization boundary in the MCP server: agent-supplied
/// `serde_json::Value` arguments are deserialized into `Tool::Input` here,
/// and `Tool::Output` is serialized back to `CallToolResult` content.
pub trait DynTool: Send + Sync {
  fn name(&self) -> &str;
  fn description(&self) -> &str;
  fn input_schema(&self) -> serde_json::Value;
  fn call<'a>(
    &'a self,
    client: &'a LspClient,
    arguments: serde_json::Value,
  ) -> std::pin::Pin<
    Box<dyn std::future::Future<Output = CallToolResult> + Send + 'a>,
  >;
}

/// Zero-sized adapter from `Tool` (typed) to `DynTool` (erased).
struct Adapter<T: Tool>(std::marker::PhantomData<fn() -> T>);

impl<T: Tool> Adapter<T> {
  const fn new() -> Self {
    Self(std::marker::PhantomData)
  }
}

impl<T: Tool> DynTool for Adapter<T> {
  fn name(&self) -> &str {
    T::NAME
  }

  fn description(&self) -> &str {
    T::DESCRIPTION
  }

  fn input_schema(&self) -> serde_json::Value {
    T::input_schema()
  }

  fn call<'a>(
    &'a self,
    client: &'a LspClient,
    arguments: serde_json::Value,
  ) -> std::pin::Pin<
    Box<dyn std::future::Future<Output = CallToolResult> + Send + 'a>,
  > {
    Box::pin(async move {
      let input: T::Input = match serde_json::from_value(arguments) {
        | Ok(v) => v,
        | Err(e) => {
          return CallToolResult::error_text(format!("invalid arguments: {e}"));
        },
      };
      match T::call(client, input).await {
        | Ok(output) => match serde_json::to_value(&output) {
          | Ok(json) => CallToolResult::ok_json(&json),
          | Err(e) => CallToolResult::error_text(format!(
            "serialization failed: {e}"
          )),
        },
        | Err(e) => CallToolResult::error_text(e.to_string()),
      }
    })
  }
}

/// Catalog of tools available on an `McpServer`.
#[derive(Default)]
pub struct ToolRegistry {
  tools: HashMap<String, Arc<dyn DynTool>>,
}

impl ToolRegistry {
  pub fn new() -> Self {
    Self::default()
  }

  /// Register a tool. Panics if a tool with the same `NAME` is already
  /// registered — name collisions are a programmer error, not a runtime
  /// condition.
  pub fn register<T: Tool>(&mut self) {
    self.insert(T::NAME.to_string(), Arc::new(Adapter::<T>::new()));
  }

  /// Register a dynamically-constructed tool. Used by the
  /// `executeCommand` extension hatch where the catalog is supplied at
  /// runtime rather than via static `Tool` impls.
  pub fn register_dyn(&mut self, tool: Arc<dyn DynTool>) {
    let name = tool.name().to_string();
    self.insert(name, tool);
  }

  fn insert(&mut self, name: String, tool: Arc<dyn DynTool>) {
    if self.tools.contains_key(&name) {
      panic!("duplicate MCP tool registration: {name}");
    }
    self.tools.insert(name, tool);
  }

  /// Look up a tool by name.
  pub(crate) fn get(&self, name: &str) -> Option<Arc<dyn DynTool>> {
    self.tools.get(name).cloned()
  }

  /// Produce the descriptor list for `tools/list`, sorted by name for
  /// deterministic output.
  pub fn descriptors(&self) -> Vec<ToolDescriptor> {
    let mut out: Vec<ToolDescriptor> = self
      .tools
      .values()
      .map(|t| ToolDescriptor {
        name:         t.name().to_string(),
        description:  t.description().to_string(),
        input_schema: t.input_schema(),
      })
      .collect();
    out.sort_by(|a, b| a.name.cmp(&b.name));
    out
  }

  pub fn len(&self) -> usize {
    self.tools.len()
  }

  pub fn is_empty(&self) -> bool {
    self.tools.is_empty()
  }
}