ricecoder_mcp/
storage_integration.rs

1//! Integration with ricecoder-storage framework
2//!
3//! This module provides integration between the MCP tool system and the ricecoder-storage
4//! framework, enabling tool registry persistence and configuration management.
5
6use crate::error::Result;
7use crate::metadata::ToolMetadata;
8use crate::registry::ToolRegistry;
9use std::collections::HashMap;
10use std::path::{Path, PathBuf};
11use std::sync::Arc;
12
13/// Tool registry storage interface
14///
15/// This trait defines how tool registries are persisted and loaded from storage.
16pub trait ToolRegistryStorage: Send + Sync {
17    /// Save a tool registry to storage
18    ///
19    /// # Arguments
20    ///
21    /// * `registry` - The tool registry to save
22    /// * `path` - The path where the registry should be saved
23    ///
24    /// # Returns
25    ///
26    /// Result indicating success or failure
27    fn save_registry(&self, registry: &ToolRegistry, path: &Path) -> Result<()>;
28
29    /// Load a tool registry from storage
30    ///
31    /// # Arguments
32    ///
33    /// * `path` - The path where the registry is stored
34    ///
35    /// # Returns
36    ///
37    /// The loaded tool registry
38    fn load_registry(&self, path: &Path) -> Result<ToolRegistry>;
39
40    /// Save a single tool to storage
41    ///
42    /// # Arguments
43    ///
44    /// * `tool` - The tool to save
45    /// * `path` - The path where the tool should be saved
46    ///
47    /// # Returns
48    ///
49    /// Result indicating success or failure
50    fn save_tool(&self, tool: &ToolMetadata, path: &Path) -> Result<()>;
51
52    /// Load a single tool from storage
53    ///
54    /// # Arguments
55    ///
56    /// * `path` - The path where the tool is stored
57    ///
58    /// # Returns
59    ///
60    /// The loaded tool metadata
61    fn load_tool(&self, path: &Path) -> Result<ToolMetadata>;
62
63    /// List all tools in a directory
64    ///
65    /// # Arguments
66    ///
67    /// * `path` - The directory path
68    ///
69    /// # Returns
70    ///
71    /// A list of tool metadata for all tools in the directory
72    fn list_tools(&self, path: &Path) -> Result<Vec<ToolMetadata>>;
73}
74
75/// JSON-based tool registry storage
76///
77/// This implementation stores tool registries as JSON files.
78pub struct JsonToolRegistryStorage;
79
80impl JsonToolRegistryStorage {
81    /// Creates a new JSON tool registry storage
82    pub fn new() -> Self {
83        Self
84    }
85}
86
87impl Default for JsonToolRegistryStorage {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93impl ToolRegistryStorage for JsonToolRegistryStorage {
94    fn save_registry(&self, registry: &ToolRegistry, path: &Path) -> Result<()> {
95        let tools: Vec<ToolMetadata> = registry
96            .list_tools()
97            .into_iter()
98            .cloned()
99            .collect();
100
101        let json = serde_json::to_string_pretty(&tools)
102            .map_err(|e| crate::error::Error::SerializationError(e))?;
103
104        std::fs::write(path, json)
105            .map_err(|e| crate::error::Error::IoError(e))?;
106
107        Ok(())
108    }
109
110    fn load_registry(&self, path: &Path) -> Result<ToolRegistry> {
111        let content = std::fs::read_to_string(path)
112            .map_err(|e| crate::error::Error::IoError(e))?;
113
114        let tools: Vec<ToolMetadata> = serde_json::from_str(&content)
115            .map_err(|e| crate::error::Error::SerializationError(e))?;
116
117        let mut registry = ToolRegistry::new();
118        for tool in tools {
119            registry.register_tool(tool)?;
120        }
121
122        Ok(registry)
123    }
124
125    fn save_tool(&self, tool: &ToolMetadata, path: &Path) -> Result<()> {
126        let json = serde_json::to_string_pretty(tool)
127            .map_err(|e| crate::error::Error::SerializationError(e))?;
128
129        std::fs::write(path, json)
130            .map_err(|e| crate::error::Error::IoError(e))?;
131
132        Ok(())
133    }
134
135    fn load_tool(&self, path: &Path) -> Result<ToolMetadata> {
136        let content = std::fs::read_to_string(path)
137            .map_err(|e| crate::error::Error::IoError(e))?;
138
139        let tool: ToolMetadata = serde_json::from_str(&content)
140            .map_err(|e| crate::error::Error::SerializationError(e))?;
141
142        Ok(tool)
143    }
144
145    fn list_tools(&self, path: &Path) -> Result<Vec<ToolMetadata>> {
146        if !path.exists() {
147            return Ok(Vec::new());
148        }
149
150        let mut tools = Vec::new();
151
152        for entry in std::fs::read_dir(path)
153            .map_err(|e| crate::error::Error::IoError(e))?
154        {
155            let entry = entry.map_err(|e| crate::error::Error::IoError(e))?;
156            let file_path = entry.path();
157
158            if file_path.extension().map_or(false, |ext| ext == "json") {
159                match self.load_tool(&file_path) {
160                    Ok(tool) => tools.push(tool),
161                    Err(e) => {
162                        tracing::warn!("Failed to load tool from {:?}: {}", file_path, e);
163                    }
164                }
165            }
166        }
167
168        Ok(tools)
169    }
170}
171
172/// Tool registry persistence manager
173///
174/// This struct manages the persistence of tool registries using storage backends.
175pub struct ToolRegistryPersistence {
176    storage: Arc<dyn ToolRegistryStorage>,
177    registry_path: PathBuf,
178    tools_dir: PathBuf,
179}
180
181impl ToolRegistryPersistence {
182    /// Creates a new tool registry persistence manager
183    ///
184    /// # Arguments
185    ///
186    /// * `storage` - The storage backend to use
187    /// * `registry_path` - The path where the registry file is stored
188    /// * `tools_dir` - The directory where individual tool files are stored
189    pub fn new(
190        storage: Arc<dyn ToolRegistryStorage>,
191        registry_path: PathBuf,
192        tools_dir: PathBuf,
193    ) -> Self {
194        Self {
195            storage,
196            registry_path,
197            tools_dir,
198        }
199    }
200
201    /// Saves a tool registry to persistent storage
202    pub fn save_registry(&self, registry: &ToolRegistry) -> Result<()> {
203        // Create parent directories if they don't exist
204        if let Some(parent) = self.registry_path.parent() {
205            std::fs::create_dir_all(parent)
206                .map_err(|e| crate::error::Error::IoError(e))?;
207        }
208
209        self.storage.save_registry(registry, &self.registry_path)?;
210        tracing::info!("Tool registry saved to {:?}", self.registry_path);
211
212        Ok(())
213    }
214
215    /// Loads a tool registry from persistent storage
216    pub fn load_registry(&self) -> Result<ToolRegistry> {
217        if !self.registry_path.exists() {
218            tracing::info!("Registry file not found at {:?}, creating new registry", self.registry_path);
219            return Ok(ToolRegistry::new());
220        }
221
222        let registry = self.storage.load_registry(&self.registry_path)?;
223        tracing::info!("Tool registry loaded from {:?}", self.registry_path);
224
225        Ok(registry)
226    }
227
228    /// Saves a single tool to persistent storage
229    pub fn save_tool(&self, tool: &ToolMetadata) -> Result<()> {
230        // Create tools directory if it doesn't exist
231        std::fs::create_dir_all(&self.tools_dir)
232            .map_err(|e| crate::error::Error::IoError(e))?;
233
234        let tool_path = self.tools_dir.join(format!("{}.json", tool.id));
235        self.storage.save_tool(tool, &tool_path)?;
236        tracing::info!("Tool saved to {:?}", tool_path);
237
238        Ok(())
239    }
240
241    /// Loads a single tool from persistent storage
242    pub fn load_tool(&self, tool_id: &str) -> Result<ToolMetadata> {
243        let tool_path = self.tools_dir.join(format!("{}.json", tool_id));
244        self.storage.load_tool(&tool_path)
245    }
246
247    /// Lists all tools in persistent storage
248    pub fn list_tools(&self) -> Result<Vec<ToolMetadata>> {
249        self.storage.list_tools(&self.tools_dir)
250    }
251
252    /// Exports a registry to a specific format
253    pub fn export_registry(&self, registry: &ToolRegistry, path: &Path, format: &str) -> Result<()> {
254        match format {
255            "json" => {
256                let tools: Vec<ToolMetadata> = registry
257                    .list_tools()
258                    .into_iter()
259                    .cloned()
260                    .collect();
261
262                let json = serde_json::to_string_pretty(&tools)
263                    .map_err(|e| crate::error::Error::SerializationError(e))?;
264
265                std::fs::write(path, json)
266                    .map_err(|e| crate::error::Error::IoError(e))?;
267
268                Ok(())
269            }
270            "yaml" => {
271                let tools: Vec<ToolMetadata> = registry
272                    .list_tools()
273                    .into_iter()
274                    .cloned()
275                    .collect();
276
277                let yaml = serde_yaml::to_string(&tools)
278                    .map_err(|e| crate::error::Error::ConfigError(e.to_string()))?;
279
280                std::fs::write(path, yaml)
281                    .map_err(|e| crate::error::Error::IoError(e))?;
282
283                Ok(())
284            }
285            _ => Err(crate::error::Error::ConfigError(format!(
286                "Unsupported export format: {}",
287                format
288            ))),
289        }
290    }
291
292    /// Imports a registry from a specific format
293    pub fn import_registry(&self, path: &Path, format: &str) -> Result<ToolRegistry> {
294        match format {
295            "json" => {
296                let content = std::fs::read_to_string(path)
297                    .map_err(|e| crate::error::Error::IoError(e))?;
298
299                let tools: Vec<ToolMetadata> = serde_json::from_str(&content)
300                    .map_err(|e| crate::error::Error::SerializationError(e))?;
301
302                let mut registry = ToolRegistry::new();
303                for tool in tools {
304                    registry.register_tool(tool)?;
305                }
306
307                Ok(registry)
308            }
309            "yaml" => {
310                let content = std::fs::read_to_string(path)
311                    .map_err(|e| crate::error::Error::IoError(e))?;
312
313                let tools: Vec<ToolMetadata> = serde_yaml::from_str(&content)
314                    .map_err(|e| crate::error::Error::ConfigError(e.to_string()))?;
315
316                let mut registry = ToolRegistry::new();
317                for tool in tools {
318                    registry.register_tool(tool)?;
319                }
320
321                Ok(registry)
322            }
323            _ => Err(crate::error::Error::ConfigError(format!(
324                "Unsupported import format: {}",
325                format
326            ))),
327        }
328    }
329}
330
331/// Tool registry cache for in-memory caching
332///
333/// This struct provides in-memory caching of tool registries to reduce
334/// disk I/O and improve performance.
335pub struct ToolRegistryCache {
336    cache: std::sync::Mutex<HashMap<String, (ToolRegistry, std::time::Instant)>>,
337    ttl_secs: u64,
338}
339
340impl ToolRegistryCache {
341    /// Creates a new tool registry cache
342    ///
343    /// # Arguments
344    ///
345    /// * `ttl_secs` - Time-to-live for cache entries in seconds
346    pub fn new(ttl_secs: u64) -> Self {
347        Self {
348            cache: std::sync::Mutex::new(HashMap::new()),
349            ttl_secs,
350        }
351    }
352
353    /// Gets a cached registry
354    pub fn get(&self, key: &str) -> Option<ToolRegistry> {
355        let cache = self.cache.lock().unwrap();
356        if let Some((registry, timestamp)) = cache.get(key) {
357            let elapsed = timestamp.elapsed().as_secs();
358            if elapsed < self.ttl_secs {
359                return Some(registry.clone());
360            }
361        }
362        None
363    }
364
365    /// Sets a cached registry
366    pub fn set(&self, key: String, registry: ToolRegistry) {
367        let mut cache = self.cache.lock().unwrap();
368        cache.insert(key, (registry, std::time::Instant::now()));
369    }
370
371    /// Clears the cache
372    pub fn clear(&self) {
373        let mut cache = self.cache.lock().unwrap();
374        cache.clear();
375    }
376
377    /// Removes expired entries from the cache
378    pub fn cleanup_expired(&self) {
379        let mut cache = self.cache.lock().unwrap();
380        let now = std::time::Instant::now();
381        cache.retain(|_, (_, timestamp)| now.duration_since(*timestamp).as_secs() < self.ttl_secs);
382    }
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388    use crate::metadata::ToolSource;
389    use tempfile::TempDir;
390
391    #[test]
392    fn test_json_tool_registry_storage_save_and_load() {
393        let temp_dir = TempDir::new().unwrap();
394        let registry_path = temp_dir.path().join("registry.json");
395
396        let mut registry = ToolRegistry::new();
397        let tool = ToolMetadata::new(
398            "test-tool".to_string(),
399            "Test Tool".to_string(),
400            "A test tool".to_string(),
401            "test".to_string(),
402            "string".to_string(),
403            ToolSource::Custom,
404        );
405        registry.register_tool(tool).unwrap();
406
407        let storage = JsonToolRegistryStorage::new();
408        storage.save_registry(&registry, &registry_path).unwrap();
409
410        assert!(registry_path.exists());
411
412        let loaded_registry = storage.load_registry(&registry_path).unwrap();
413        assert_eq!(loaded_registry.tool_count(), 1);
414    }
415
416    #[test]
417    fn test_tool_registry_persistence_save_and_load() {
418        let temp_dir = TempDir::new().unwrap();
419        let registry_path = temp_dir.path().join("registry.json");
420        let tools_dir = temp_dir.path().join("tools");
421
422        let storage = Arc::new(JsonToolRegistryStorage::new());
423        let persistence = ToolRegistryPersistence::new(storage, registry_path, tools_dir);
424
425        let mut registry = ToolRegistry::new();
426        let tool = ToolMetadata::new(
427            "test-tool".to_string(),
428            "Test Tool".to_string(),
429            "A test tool".to_string(),
430            "test".to_string(),
431            "string".to_string(),
432            ToolSource::Custom,
433        );
434        registry.register_tool(tool).unwrap();
435
436        persistence.save_registry(&registry).unwrap();
437        let loaded_registry = persistence.load_registry().unwrap();
438
439        assert_eq!(loaded_registry.tool_count(), 1);
440    }
441
442    #[test]
443    fn test_tool_registry_cache() {
444        let cache = ToolRegistryCache::new(60);
445        let registry = ToolRegistry::new();
446
447        cache.set("test".to_string(), registry.clone());
448        assert!(cache.get("test").is_some());
449
450        cache.clear();
451        assert!(cache.get("test").is_none());
452    }
453
454    #[test]
455    fn test_tool_registry_cache_expiration() {
456        let cache = ToolRegistryCache::new(0); // 0 second TTL
457        let registry = ToolRegistry::new();
458
459        cache.set("test".to_string(), registry);
460        std::thread::sleep(std::time::Duration::from_millis(10));
461
462        assert!(cache.get("test").is_none());
463    }
464
465    #[test]
466    fn test_tool_registry_persistence_export_json() {
467        let temp_dir = TempDir::new().unwrap();
468        let registry_path = temp_dir.path().join("registry.json");
469        let tools_dir = temp_dir.path().join("tools");
470        let export_path = temp_dir.path().join("export.json");
471
472        let storage = Arc::new(JsonToolRegistryStorage::new());
473        let persistence = ToolRegistryPersistence::new(storage, registry_path, tools_dir);
474
475        let mut registry = ToolRegistry::new();
476        let tool = ToolMetadata::new(
477            "test-tool".to_string(),
478            "Test Tool".to_string(),
479            "A test tool".to_string(),
480            "test".to_string(),
481            "string".to_string(),
482            ToolSource::Custom,
483        );
484        registry.register_tool(tool).unwrap();
485
486        persistence.export_registry(&registry, &export_path, "json").unwrap();
487        assert!(export_path.exists());
488
489        let imported_registry = persistence.import_registry(&export_path, "json").unwrap();
490        assert_eq!(imported_registry.tool_count(), 1);
491    }
492}
ricecoder_mcp/storage_integration.rs

ricecoder_mcp/
storage_integration.rs
