Skip to main content

headless_lms_mcp/tools/
handler.rs

1//! The tool handler trait and the registry dispatch looks tools up in.
2
3use async_trait::async_trait;
4use indexmap::IndexMap;
5use schemars::JsonSchema;
6use serde::de::DeserializeOwned;
7use serde_json::{Map, Value};
8
9use super::catalog::ToolDefinition;
10use super::{McpPrincipal, RequestContext, ToolError, ToolOutcome};
11
12/// Implements one tool. `S` is the application state the protocol layer passes through, such as
13/// the database pool, config and validator client.
14#[async_trait(?Send)]
15pub trait ToolHandler<S>: Send + Sync {
16    /// The tool's arguments; their JSON Schema is the tool's input schema.
17    type Args: DeserializeOwned + JsonSchema + 'static;
18
19    /// Whether `call` can answer [`ToolOutcome::NeedsConfirmation`]. Only then does dispatch bind
20    /// the call's arguments to a confirmation, which costs a pass over them.
21    const ASKS_FOR_CONFIRMATION: bool = false;
22
23    async fn call(
24        &self,
25        state: &S,
26        ctx: &RequestContext,
27        args: Self::Args,
28    ) -> Result<ToolOutcome, ToolError>;
29}
30
31/// A [`ToolHandler`] taking its arguments as JSON, so that one registry holds every tool.
32#[async_trait(?Send)]
33pub(crate) trait JsonToolHandler<S>: Send + Sync {
34    fn asks_for_confirmation(&self) -> bool;
35
36    async fn call_with_json(
37        &self,
38        state: &S,
39        ctx: &RequestContext,
40        arguments: Map<String, Value>,
41    ) -> Result<ToolOutcome, ToolError>;
42}
43
44#[async_trait(?Send)]
45impl<S, H: ToolHandler<S>> JsonToolHandler<S> for H {
46    fn asks_for_confirmation(&self) -> bool {
47        H::ASKS_FOR_CONFIRMATION
48    }
49
50    async fn call_with_json(
51        &self,
52        state: &S,
53        ctx: &RequestContext,
54        arguments: Map<String, Value>,
55    ) -> Result<ToolOutcome, ToolError> {
56        let args = serde_json::from_value(Value::Object(arguments)).map_err(|e| {
57            ToolError::InvalidArguments {
58                message: format!("invalid arguments: {e}"),
59            }
60        })?;
61        self.call(state, ctx, args).await
62    }
63}
64
65/// A tool with its handler.
66pub struct Tool<S> {
67    pub definition: ToolDefinition,
68    pub(crate) handler: Box<dyn JsonToolHandler<S>>,
69}
70
71/// The tools in their `tools/list` order.
72pub struct ToolRegistry<S> {
73    tools: IndexMap<&'static str, Tool<S>>,
74}
75
76impl<S> ToolRegistry<S> {
77    /// A registry listing `tools` in the given order.
78    pub fn new(tools: Vec<Tool<S>>) -> Self {
79        Self {
80            tools: tools
81                .into_iter()
82                .map(|tool| (tool.definition.name, tool))
83                .collect(),
84        }
85    }
86
87    /// The tool called `name`, without a scope check; [`Self::listings_for`] lists what a principal may call.
88    pub fn get(&self, name: &str) -> Option<&Tool<S>> {
89        self.tools.get(name)
90    }
91
92    #[cfg(test)]
93    pub fn definitions(&self) -> impl Iterator<Item = &ToolDefinition> {
94        self.tools.values().map(|tool| &tool.definition)
95    }
96
97    /// The `tools/list` entries of the tools `principal` may call.
98    pub fn listings_for(&self, principal: &McpPrincipal) -> Vec<rmcp::model::Tool> {
99        self.tools
100            .values()
101            .filter(|tool| principal.has_scope(tool.definition.scope))
102            .map(|tool| tool.definition.listing.clone())
103            .collect()
104    }
105}