Skip to main content

lean_ctx/server/
registry.rs

1use std::collections::HashMap;
2
3use rmcp::model::Tool;
4
5use super::tool_trait::McpTool;
6
7/// Central registry mapping tool names to their trait-based handlers.
8/// Replaces the match-cascade dispatch for migrated tools while
9/// coexisting with the legacy dispatch for tools not yet migrated.
10pub struct ToolRegistry {
11    tools: HashMap<&'static str, Box<dyn McpTool>>,
12}
13
14impl ToolRegistry {
15    pub fn new() -> Self {
16        Self {
17            tools: HashMap::new(),
18        }
19    }
20
21    pub fn register(&mut self, tool: Box<dyn McpTool>) {
22        self.tools.insert(tool.name(), tool);
23    }
24
25    pub fn get(&self, name: &str) -> Option<&dyn McpTool> {
26        self.tools.get(name).map(AsRef::as_ref)
27    }
28
29    pub fn contains(&self, name: &str) -> bool {
30        self.tools.contains_key(name)
31    }
32
33    /// Returns MCP Tool definitions for all registered tools.
34    /// Used by `list_tools` to expose schemas to clients.
35    pub fn tool_defs(&self) -> Vec<Tool> {
36        let mut defs: Vec<Tool> = self.tools.values().map(|t| t.tool_def()).collect();
37        defs.sort_by(|a, b| a.name.as_ref().cmp(b.name.as_ref()));
38        defs
39    }
40
41    pub fn len(&self) -> usize {
42        self.tools.len()
43    }
44
45    pub fn is_empty(&self) -> bool {
46        self.tools.is_empty()
47    }
48
49    pub fn names(&self) -> Vec<&'static str> {
50        let mut names: Vec<_> = self.tools.keys().copied().collect();
51        names.sort_unstable();
52        names
53    }
54}
55
56impl Default for ToolRegistry {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62/// Register all trait-based tools. Called once during server startup.
63/// Tools are added here as they are migrated from the legacy dispatch.
64pub fn build_registry() -> ToolRegistry {
65    let mut registry = ToolRegistry::new();
66
67    use crate::tools::registered;
68    registry.register(Box::new(registered::ctx_tree::CtxTreeTool));
69    registry.register(Box::new(registered::ctx_benchmark::CtxBenchmarkTool));
70    registry.register(Box::new(registered::ctx_analyze::CtxAnalyzeTool));
71    registry.register(Box::new(registered::ctx_discover::CtxDiscoverTool));
72    registry.register(Box::new(registered::ctx_response::CtxResponseTool));
73    registry.register(Box::new(registered::ctx_wrapped::CtxWrappedTool));
74    registry.register(Box::new(registered::ctx_heatmap::CtxHeatmapTool));
75    registry.register(Box::new(registered::ctx_verify::CtxVerifyTool));
76    registry.register(Box::new(registered::ctx_outline::CtxOutlineTool));
77    registry.register(Box::new(registered::ctx_cost::CtxCostTool));
78    registry.register(Box::new(registered::ctx_gain::CtxGainTool));
79    registry.register(Box::new(registered::ctx_expand::CtxExpandTool));
80    registry.register(Box::new(registered::ctx_routes::CtxRoutesTool));
81    registry.register(Box::new(registered::ctx_callers::CtxCallersTool));
82    registry.register(Box::new(registered::ctx_callees::CtxCalleesTool));
83    registry.register(Box::new(registered::ctx_callgraph::CtxCallgraphTool));
84    registry.register(Box::new(registered::ctx_symbol::CtxSymbolTool));
85    registry.register(Box::new(registered::ctx_graph_diagram::CtxGraphDiagramTool));
86    registry.register(Box::new(
87        registered::ctx_discover_tools::CtxDiscoverToolsTool,
88    ));
89    registry.register(Box::new(registered::ctx_review::CtxReviewTool));
90    registry.register(Box::new(registered::ctx_provider::CtxProviderTool));
91    registry.register(Box::new(registered::ctx_impact::CtxImpactTool));
92    registry.register(Box::new(registered::ctx_architecture::CtxArchitectureTool));
93    registry.register(Box::new(registered::ctx_smells::CtxSmellsTool));
94    registry.register(Box::new(registered::ctx_pack::CtxPackTool));
95    registry.register(Box::new(registered::ctx_index::CtxIndexTool));
96    registry.register(Box::new(registered::ctx_artifacts::CtxArtifactsTool));
97    registry.register(Box::new(
98        registered::ctx_compress_memory::CtxCompressMemoryTool,
99    ));
100
101    registry
102}