dravr_tronc/mcp/
tool.rs

1// ABOUTME: Generic McpTool trait + ToolRegistry, with a per-call ToolContext and capability gating
2// ABOUTME: Parameterized over state type S so each project provides its own ServerState
3//
4// SPDX-License-Identifier: MIT OR Apache-2.0
5// Copyright (c) 2026 dravr.ai
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use bitflags::bitflags;
12use serde_json::Value;
13
14use crate::mcp::schema::{Tool, ToolResponse};
15
16bitflags! {
17    /// Host-agnostic classification flags a tool declares for discovery + gating.
18    ///
19    /// These are the generic capabilities the registry and transports reason
20    /// about (auth/tenant/provider requirements, read vs. write, admin gating).
21    /// Domain taxonomy (e.g. fitness "goals"/"recipes" groupings) belongs in the
22    /// registry's string categories, not here.
23    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
24    pub struct ToolCapabilities: u16 {
25        /// The tool requires an authenticated caller.
26        const REQUIRES_AUTH = 0b0000_0001;
27        /// The tool requires a resolved tenant context.
28        const REQUIRES_TENANT = 0b0000_0010;
29        /// The tool requires a connected upstream provider.
30        const REQUIRES_PROVIDER = 0b0000_0100;
31        /// The tool only reads data (no side effects).
32        const READS_DATA = 0b0000_1000;
33        /// The tool writes or mutates data.
34        const WRITES_DATA = 0b0001_0000;
35        /// The tool may only be invoked by an admin caller.
36        const ADMIN_ONLY = 0b0010_0000;
37    }
38}
39
40/// Per-call context threaded into a tool's `execute`.
41///
42/// Carries the request-scoped identity the host resolved (caller, tenant, how
43/// they authenticated, the request id) plus the precomputed admin flag the
44/// registry uses to gate `ADMIN_ONLY` tools. All identity fields are optional
45/// and host-agnostic (ids as strings) so a server without users/tenants — or a
46/// stdio server with no auth — can pass [`ToolContext::default`].
47#[derive(Debug, Clone, Default)]
48pub struct ToolContext {
49    /// Authenticated caller id, if any (host-defined; e.g. a UUID string).
50    pub user_id: Option<String>,
51    /// Resolved tenant id, if any.
52    pub tenant_id: Option<String>,
53    /// How the caller authenticated (host-defined label, e.g. `"jwt_bearer"`).
54    pub auth_method: Option<String>,
55    /// Correlation id for tracing/logging.
56    pub request_id: Option<Value>,
57    /// Whether the caller holds admin privileges (resolved by the host).
58    pub is_admin: bool,
59}
60
61impl ToolContext {
62    /// An empty context — no identity, not admin. Equivalent to [`Self::default`].
63    #[must_use]
64    pub fn new() -> Self {
65        Self::default()
66    }
67
68    /// Set the authenticated caller id.
69    #[must_use]
70    pub fn with_user(mut self, user_id: impl Into<String>) -> Self {
71        self.user_id = Some(user_id.into());
72        self
73    }
74
75    /// Set the resolved tenant id.
76    #[must_use]
77    pub fn with_tenant(mut self, tenant_id: impl Into<String>) -> Self {
78        self.tenant_id = Some(tenant_id.into());
79        self
80    }
81
82    /// Set the authentication method label.
83    #[must_use]
84    pub fn with_auth_method(mut self, auth_method: impl Into<String>) -> Self {
85        self.auth_method = Some(auth_method.into());
86        self
87    }
88
89    /// Set the request correlation id.
90    #[must_use]
91    pub fn with_request_id(mut self, request_id: Value) -> Self {
92        self.request_id = Some(request_id);
93        self
94    }
95
96    /// Mark whether the caller holds admin privileges.
97    #[must_use]
98    pub const fn as_admin(mut self, is_admin: bool) -> Self {
99        self.is_admin = is_admin;
100        self
101    }
102}
103
104/// Trait implemented by each MCP tool exposed by a server
105///
106/// Generic over `S` — the project-specific server state type, shared as
107/// `Arc<S>`. `S` is `?Sized`, so a host may parameterize it with an unsized
108/// trait object (e.g. a resource façade `dyn HostRuntime`) rather than a
109/// concrete struct. The shared state is handed to `execute` immutably; a host
110/// that needs interior mutability parameterizes `S` with it (e.g.
111/// `S = RwLock<Inner>`, yielding `Arc<RwLock<Inner>>`).
112#[async_trait]
113pub trait McpTool<S: Send + Sync + ?Sized>: Send + Sync {
114    /// Return the tool's MCP definition (name, description, input schema)
115    fn definition(&self) -> Tool;
116
117    /// Declare the tool's host-agnostic capabilities (auth/tenant/admin/...).
118    ///
119    /// Defaults to no capabilities. The registry uses [`ToolCapabilities::ADMIN_ONLY`]
120    /// to gate execution; transports may use the rest for discovery filtering.
121    fn capabilities(&self) -> ToolCapabilities {
122        ToolCapabilities::empty()
123    }
124
125    /// Execute the tool against the shared server state and per-call context
126    async fn execute(&self, state: &Arc<S>, ctx: &ToolContext, arguments: Value) -> ToolResponse;
127}
128
129/// Registry mapping tool names to their handler implementations
130///
131/// Tools are registered at server startup and looked up by name
132/// when `tools/call` requests arrive from the MCP client.
133pub struct ToolRegistry<S: Send + Sync + ?Sized> {
134    tools: HashMap<String, Box<dyn McpTool<S>>>,
135    categories: HashMap<String, Vec<String>>,
136}
137
138impl<S: Send + Sync + ?Sized> Default for ToolRegistry<S> {
139    fn default() -> Self {
140        Self::new()
141    }
142}
143
144impl<S: Send + Sync + ?Sized> ToolRegistry<S> {
145    /// Create an empty registry
146    pub fn new() -> Self {
147        Self {
148            tools: HashMap::new(),
149            categories: HashMap::new(),
150        }
151    }
152
153    /// Register a tool handler, keyed by its definition name
154    pub fn register(&mut self, tool: Box<dyn McpTool<S>>) {
155        let name = tool.definition().name;
156        self.tools.insert(name, tool);
157    }
158
159    /// Register a tool handler and record it under the given category
160    pub fn register_with_category(&mut self, tool: Box<dyn McpTool<S>>, category: &str) {
161        let name = tool.definition().name;
162        self.categories
163            .entry(category.to_owned())
164            .or_default()
165            .push(name.clone());
166        self.tools.insert(name, tool);
167    }
168
169    /// Return the number of registered tools
170    pub fn len(&self) -> usize {
171        self.tools.len()
172    }
173
174    /// Return true if no tools are registered
175    pub fn is_empty(&self) -> bool {
176        self.tools.is_empty()
177    }
178
179    /// List all registered tool definitions for `tools/list` responses
180    pub fn list_definitions(&self) -> Vec<Tool> {
181        self.tools.values().map(|t| t.definition()).collect()
182    }
183
184    /// List tool definitions visible to a non-admin caller (excludes
185    /// `ADMIN_ONLY` tools). Pass `is_admin = true` to include everything.
186    pub fn list_definitions_for(&self, is_admin: bool) -> Vec<Tool> {
187        self.tools
188            .values()
189            .filter(|t| is_admin || !t.capabilities().contains(ToolCapabilities::ADMIN_ONLY))
190            .map(|t| t.definition())
191            .collect()
192    }
193
194    /// Look up a registered tool's declared capabilities.
195    pub fn capabilities_of(&self, name: &str) -> Option<ToolCapabilities> {
196        self.tools.get(name).map(|t| t.capabilities())
197    }
198
199    /// Names of the categories tools have been registered under.
200    pub fn categories(&self) -> Vec<&str> {
201        self.categories.keys().map(String::as_str).collect()
202    }
203
204    /// Tool names registered under the given category.
205    pub fn tools_in_category(&self, category: &str) -> Vec<&str> {
206        self.categories
207            .get(category)
208            .map(|names| names.iter().map(String::as_str).collect())
209            .unwrap_or_default()
210    }
211
212    /// Dispatch a `tools/call` to the named tool handler
213    ///
214    /// Gates `ADMIN_ONLY` tools on `ctx.is_admin` before dispatching.
215    pub async fn execute(
216        &self,
217        name: &str,
218        state: &Arc<S>,
219        ctx: &ToolContext,
220        arguments: Value,
221    ) -> ToolResponse {
222        match self.tools.get(name) {
223            Some(tool) => {
224                if tool.capabilities().contains(ToolCapabilities::ADMIN_ONLY) && !ctx.is_admin {
225                    return ToolResponse::error(format!("Tool '{name}' requires admin privileges"));
226                }
227                tool.execute(state, ctx, arguments).await
228            }
229            None => ToolResponse::error(format!("Unknown tool: {name}")),
230        }
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237    use serde_json::json;
238
239    struct DummyState {
240        counter: i32,
241    }
242
243    struct EchoTool;
244
245    #[async_trait]
246    impl McpTool<DummyState> for EchoTool {
247        fn definition(&self) -> Tool {
248            Tool {
249                name: "echo".to_owned(),
250                description: "Echoes the input".to_owned(),
251                input_schema: json!({
252                    "type": "object",
253                    "properties": {
254                        "message": { "type": "string" }
255                    }
256                }),
257                annotations: None,
258            }
259        }
260
261        fn capabilities(&self) -> ToolCapabilities {
262            ToolCapabilities::READS_DATA
263        }
264
265        async fn execute(
266            &self,
267            _state: &Arc<DummyState>,
268            _ctx: &ToolContext,
269            arguments: Value,
270        ) -> ToolResponse {
271            let msg = arguments
272                .get("message")
273                .and_then(|v| v.as_str())
274                .unwrap_or("(empty)");
275            ToolResponse::text(format!("echo: {msg}"))
276        }
277    }
278
279    struct CounterTool;
280
281    #[async_trait]
282    impl McpTool<DummyState> for CounterTool {
283        fn definition(&self) -> Tool {
284            Tool {
285                name: "counter".to_owned(),
286                description: "Returns the counter value".to_owned(),
287                input_schema: json!({"type": "object"}),
288                annotations: None,
289            }
290        }
291
292        async fn execute(
293            &self,
294            state: &Arc<DummyState>,
295            _ctx: &ToolContext,
296            _arguments: Value,
297        ) -> ToolResponse {
298            ToolResponse::text(format!("counter: {}", state.counter))
299        }
300    }
301
302    struct AdminTool;
303
304    #[async_trait]
305    impl McpTool<DummyState> for AdminTool {
306        fn definition(&self) -> Tool {
307            Tool {
308                name: "admin_reset".to_owned(),
309                description: "Admin-only reset".to_owned(),
310                input_schema: json!({"type": "object"}),
311                annotations: None,
312            }
313        }
314
315        fn capabilities(&self) -> ToolCapabilities {
316            ToolCapabilities::ADMIN_ONLY | ToolCapabilities::WRITES_DATA
317        }
318
319        async fn execute(
320            &self,
321            _state: &Arc<DummyState>,
322            _ctx: &ToolContext,
323            _arguments: Value,
324        ) -> ToolResponse {
325            ToolResponse::text("reset".to_owned())
326        }
327    }
328
329    fn make_state() -> Arc<DummyState> {
330        Arc::new(DummyState { counter: 42 })
331    }
332
333    #[test]
334    fn empty_registry() {
335        let registry = ToolRegistry::<DummyState>::new();
336        assert!(registry.is_empty());
337        assert_eq!(registry.len(), 0);
338        assert!(registry.list_definitions().is_empty());
339    }
340
341    #[test]
342    fn default_is_empty() {
343        let registry = ToolRegistry::<DummyState>::default();
344        assert!(registry.is_empty());
345    }
346
347    #[test]
348    fn register_and_list() {
349        let mut registry = ToolRegistry::<DummyState>::new();
350        registry.register(Box::new(EchoTool));
351        registry.register(Box::new(CounterTool));
352
353        assert_eq!(registry.len(), 2);
354        assert!(!registry.is_empty());
355
356        let defs = registry.list_definitions();
357        assert_eq!(defs.len(), 2);
358
359        let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
360        assert!(names.contains(&"echo"));
361        assert!(names.contains(&"counter"));
362    }
363
364    #[test]
365    fn register_replaces_duplicate_name() {
366        let mut registry = ToolRegistry::<DummyState>::new();
367        registry.register(Box::new(EchoTool));
368        registry.register(Box::new(EchoTool));
369        assert_eq!(registry.len(), 1);
370    }
371
372    #[test]
373    fn register_with_category_tracks_membership() {
374        let mut registry = ToolRegistry::<DummyState>::new();
375        registry.register_with_category(Box::new(EchoTool), "data");
376        registry.register_with_category(Box::new(CounterTool), "data");
377
378        assert!(registry.categories().contains(&"data"));
379        let mut in_data = registry.tools_in_category("data");
380        in_data.sort_unstable();
381        assert_eq!(in_data, vec!["counter", "echo"]);
382        assert!(registry.tools_in_category("missing").is_empty());
383    }
384
385    #[test]
386    fn capabilities_are_reported() {
387        let mut registry = ToolRegistry::<DummyState>::new();
388        registry.register(Box::new(EchoTool));
389        assert_eq!(
390            registry.capabilities_of("echo"),
391            Some(ToolCapabilities::READS_DATA)
392        );
393        assert!(registry.capabilities_of("missing").is_none());
394    }
395
396    #[test]
397    fn admin_only_tools_hidden_from_non_admins() {
398        let mut registry = ToolRegistry::<DummyState>::new();
399        registry.register(Box::new(EchoTool));
400        registry.register(Box::new(AdminTool));
401
402        let user_visible = registry.list_definitions_for(false);
403        assert_eq!(user_visible.len(), 1);
404        assert_eq!(user_visible[0].name, "echo");
405
406        let admin_visible = registry.list_definitions_for(true);
407        assert_eq!(admin_visible.len(), 2);
408    }
409
410    #[tokio::test]
411    async fn execute_known_tool() {
412        let mut registry = ToolRegistry::<DummyState>::new();
413        registry.register(Box::new(EchoTool));
414
415        let state = make_state();
416        let ctx = ToolContext::new();
417        let result = registry
418            .execute("echo", &state, &ctx, json!({"message": "hello"}))
419            .await;
420        assert!(!result.is_error);
421        assert_eq!(result.content[0].as_text(), Some("echo: hello"));
422    }
423
424    #[tokio::test]
425    async fn execute_reads_state() {
426        let mut registry = ToolRegistry::<DummyState>::new();
427        registry.register(Box::new(CounterTool));
428
429        let state = make_state();
430        let ctx = ToolContext::new();
431        let result = registry.execute("counter", &state, &ctx, json!({})).await;
432        assert_eq!(result.content[0].as_text(), Some("counter: 42"));
433    }
434
435    #[tokio::test]
436    async fn execute_unknown_tool_returns_error() {
437        let registry = ToolRegistry::<DummyState>::new();
438        let state = make_state();
439        let ctx = ToolContext::new();
440        let result = registry
441            .execute("nonexistent", &state, &ctx, json!({}))
442            .await;
443        assert!(result.is_error);
444        assert!(result.content[0]
445            .as_text()
446            .expect("text") // Safe: test assertion
447            .contains("Unknown tool"));
448    }
449
450    #[tokio::test]
451    async fn admin_only_tool_rejects_non_admin() {
452        let mut registry = ToolRegistry::<DummyState>::new();
453        registry.register(Box::new(AdminTool));
454        let state = make_state();
455
456        let non_admin = ToolContext::new();
457        let denied = registry
458            .execute("admin_reset", &state, &non_admin, json!({}))
459            .await;
460        assert!(denied.is_error);
461        assert!(denied.content[0]
462            .as_text()
463            .expect("text") // Safe: test assertion
464            .contains("admin"));
465
466        let admin = ToolContext::new().as_admin(true);
467        let allowed = registry
468            .execute("admin_reset", &state, &admin, json!({}))
469            .await;
470        assert!(!allowed.is_error);
471        assert_eq!(allowed.content[0].as_text(), Some("reset"));
472    }
473
474    #[test]
475    fn tool_definitions_have_required_fields() {
476        let tool = EchoTool;
477        let def = tool.definition();
478        assert!(!def.name.is_empty());
479        assert!(!def.description.is_empty());
480        assert!(def.input_schema.is_object());
481    }
482}
dravr_tronc/mcp/tool.rs

dravr_tronc/mcp/
tool.rs
