Skip to main content

coderlib/tools/
router.rs

1//! Enhanced tool router system for CoderLib
2//!
3//! This module provides enhanced tool patterns inspired by rust-sdk,
4//! including automatic tool discovery, structured parameter handling,
5//! and JSON Schema support for better developer experience.
6
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9use serde_json::Value as JsonValue;
10use std::collections::HashMap;
11
12use crate::integration::HostIntegration;
13use super::{Tool, ToolResponse, ToolError, Permission};
14
15/// Wrapper for structured parameter extraction with JSON Schema support
16#[derive(Debug, Clone)]
17pub struct Parameters<T>(pub T);
18
19impl<T> Parameters<T>
20where
21    T: for<'de> Deserialize<'de>,
22{
23    /// Extract parameters from JSON value with validation
24    pub fn from_json(value: JsonValue) -> Result<Self, ToolError> {
25        let params: T = serde_json::from_value(value)
26            .map_err(|e| ToolError::InvalidParameters(format!("Parameter validation failed: {}", e)))?;
27        Ok(Parameters(params))
28    }
29}
30
31/// Result type for tool execution that matches MCP patterns
32#[derive(Debug, Clone, Serialize)]
33pub struct CallToolResult {
34    pub content: Vec<Content>,
35    pub is_error: bool,
36}
37
38/// Content type for tool responses
39#[derive(Debug, Clone, Serialize)]
40#[serde(tag = "type")]
41pub enum Content {
42    #[serde(rename = "text")]
43    Text { text: String },
44    #[serde(rename = "image")]
45    Image { data: String, mime_type: String },
46    #[serde(rename = "resource")]
47    Resource { resource: ResourceReference },
48}
49
50/// Resource reference for content
51#[derive(Debug, Clone, Serialize)]
52pub struct ResourceReference {
53    pub uri: String,
54    pub text: Option<String>,
55}
56
57impl CallToolResult {
58    /// Create a successful result with text content
59    pub fn success(content: Vec<Content>) -> Self {
60        Self {
61            content,
62            is_error: false,
63        }
64    }
65
66    /// Create an error result
67    pub fn error(message: String) -> Self {
68        Self {
69            content: vec![Content::text(message)],
70            is_error: true,
71        }
72    }
73}
74
75impl Content {
76    /// Create text content
77    pub fn text(text: String) -> Self {
78        Self::Text { text }
79    }
80
81    /// Create image content
82    pub fn image(data: String, mime_type: String) -> Self {
83        Self::Image { data, mime_type }
84    }
85
86    /// Create resource content
87    pub fn resource(uri: String, text: Option<String>) -> Self {
88        Self::Resource {
89            resource: ResourceReference { uri, text },
90        }
91    }
92}
93
94/// Enhanced tool trait that supports structured parameters and JSON Schema
95#[async_trait]
96pub trait EnhancedTool: Send + Sync {
97    /// Execute the tool with structured parameters
98    async fn execute_enhanced(
99        &self,
100        parameters: JsonValue,
101        host: &dyn HostIntegration,
102    ) -> Result<CallToolResult, ToolError>;
103
104    /// Get the tool's name/identifier
105    fn name(&self) -> &str;
106
107    /// Get a description of what this tool does
108    fn description(&self) -> &str;
109
110    /// Get the JSON schema for the tool's parameters
111    fn parameter_schema(&self) -> JsonValue;
112
113    /// Get the permission required to execute this tool
114    fn requires_permission(&self) -> Permission;
115
116    /// Clone the tool (for use in collections)
117    fn clone_enhanced(&self) -> Box<dyn EnhancedTool>;
118}
119
120/// Tool router for managing enhanced tools with automatic discovery
121pub struct ToolRouter {
122    tools: HashMap<String, Box<dyn EnhancedTool>>,
123}
124
125impl ToolRouter {
126    /// Create a new tool router
127    pub fn new() -> Self {
128        Self {
129            tools: HashMap::new(),
130        }
131    }
132
133    /// Register an enhanced tool
134    pub fn register<T: EnhancedTool + 'static>(&mut self, tool: T) {
135        self.tools.insert(tool.name().to_string(), Box::new(tool));
136    }
137
138    /// Get a tool by name
139    pub fn get(&self, name: &str) -> Option<&dyn EnhancedTool> {
140        self.tools.get(name).map(|t| t.as_ref())
141    }
142
143    /// List all registered tools
144    pub fn list_tools(&self) -> Vec<&str> {
145        self.tools.keys().map(|s| s.as_str()).collect()
146    }
147
148    /// Get tool schemas for all registered tools
149    pub fn get_tool_schemas(&self) -> HashMap<String, ToolSchema> {
150        self.tools
151            .iter()
152            .map(|(name, tool)| {
153                let schema = ToolSchema {
154                    name: name.clone(),
155                    description: tool.description().to_string(),
156                    input_schema: tool.parameter_schema(),
157                };
158                (name.clone(), schema)
159            })
160            .collect()
161    }
162
163    /// Execute a tool by name with parameters
164    pub async fn execute_tool(
165        &self,
166        name: &str,
167        parameters: JsonValue,
168        host: &dyn HostIntegration,
169    ) -> Result<CallToolResult, ToolError> {
170        let tool = self.get(name)
171            .ok_or_else(|| ToolError::NotFound(name.to_string()))?;
172        
173        tool.execute_enhanced(parameters, host).await
174    }
175}
176
177/// Tool schema for MCP compatibility
178#[derive(Debug, Clone, Serialize)]
179pub struct ToolSchema {
180    pub name: String,
181    pub description: String,
182    #[serde(rename = "inputSchema")]
183    pub input_schema: JsonValue,
184}
185
186/// Adapter to bridge enhanced tools with the legacy tool interface
187pub struct EnhancedToolAdapter {
188    enhanced_tool: Box<dyn EnhancedTool>,
189}
190
191impl EnhancedToolAdapter {
192    /// Create a new adapter for an enhanced tool
193    pub fn new<T: EnhancedTool + 'static>(tool: T) -> Self {
194        Self {
195            enhanced_tool: Box::new(tool),
196        }
197    }
198}
199
200#[async_trait]
201impl Tool for EnhancedToolAdapter {
202    async fn execute(
203        &self,
204        parameters: JsonValue,
205        host: &dyn HostIntegration,
206    ) -> Result<ToolResponse, ToolError> {
207        let result = self.enhanced_tool.execute_enhanced(parameters, host).await?;
208        
209        // Convert CallToolResult to ToolResponse
210        let content = result.content
211            .into_iter()
212            .map(|c| match c {
213                Content::Text { text } => text,
214                Content::Image { data, mime_type } => format!("Image: {} ({})", data, mime_type),
215                Content::Resource { resource } => {
216                    format!("Resource: {} ({})", resource.uri, resource.text.unwrap_or_default())
217                }
218            })
219            .collect::<Vec<_>>()
220            .join("\n");
221
222        Ok(ToolResponse {
223            content,
224            success: !result.is_error,
225            metadata: JsonValue::Null,
226            affected_files: Vec::new(),
227        })
228    }
229
230    fn requires_permission(&self) -> Permission {
231        self.enhanced_tool.requires_permission()
232    }
233
234    fn description(&self) -> &str {
235        self.enhanced_tool.description()
236    }
237
238    fn name(&self) -> &str {
239        self.enhanced_tool.name()
240    }
241
242    fn parameter_schema(&self) -> JsonValue {
243        self.enhanced_tool.parameter_schema()
244    }
245
246    fn clone_box(&self) -> Box<dyn Tool> {
247        Box::new(EnhancedToolAdapter {
248            enhanced_tool: self.enhanced_tool.clone_enhanced(),
249        })
250    }
251}
252
253impl Default for ToolRouter {
254    fn default() -> Self {
255        Self::new()
256    }
257}
258
259/// Utility functions for creating JSON schemas
260pub mod schema_utils {
261    use super::*;
262    #[cfg(feature = "mcp-server")]
263    use schemars::{JsonSchema, schema_for};
264
265    /// Generate JSON schema for a type that implements JsonSchema
266    #[cfg(feature = "mcp-server")]
267    pub fn generate_schema<T: JsonSchema>() -> JsonValue {
268        let schema = schema_for!(T);
269        serde_json::to_value(schema).unwrap_or(JsonValue::Null)
270    }
271
272    #[cfg(not(feature = "mcp-server"))]
273    pub fn generate_schema<T>() -> JsonValue {
274        serde_json::json!({})
275    }
276
277    /// Create a simple string parameter schema
278    pub fn string_param(description: &str, required: bool) -> JsonValue {
279        serde_json::json!({
280            "type": "object",
281            "properties": {
282                "value": {
283                    "type": "string",
284                    "description": description
285                }
286            },
287            "required": if required { vec!["value"] } else { vec![] }
288        })
289    }
290
291    /// Create a simple number parameter schema
292    pub fn number_param(description: &str, required: bool) -> JsonValue {
293        serde_json::json!({
294            "type": "object",
295            "properties": {
296                "value": {
297                    "type": "number",
298                    "description": description
299                }
300            },
301            "required": if required { vec!["value"] } else { vec![] }
302        })
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    #[cfg(feature = "schemars")]
310    use schemars::JsonSchema;
311
312    #[derive(Deserialize)]
313    #[cfg_attr(feature = "schemars", derive(JsonSchema))]
314    struct TestParams {
315        message: String,
316        count: Option<i32>,
317    }
318
319    struct TestTool;
320
321    #[async_trait]
322    impl EnhancedTool for TestTool {
323        async fn execute_enhanced(
324            &self,
325            parameters: JsonValue,
326            _host: &dyn HostIntegration,
327        ) -> Result<CallToolResult, ToolError> {
328            let Parameters(params): Parameters<TestParams> = Parameters::from_json(parameters)?;
329            let response = format!("Message: {}, Count: {:?}", params.message, params.count);
330            Ok(CallToolResult::success(vec![Content::text(response)]))
331        }
332
333        fn name(&self) -> &str {
334            "test_tool"
335        }
336
337        fn description(&self) -> &str {
338            "A test tool for demonstration"
339        }
340
341        fn parameter_schema(&self) -> JsonValue {
342            schema_utils::generate_schema::<TestParams>()
343        }
344
345        fn requires_permission(&self) -> Permission {
346            Permission::None
347        }
348
349        fn clone_enhanced(&self) -> Box<dyn EnhancedTool> {
350            Box::new(TestTool)
351        }
352    }
353
354    #[tokio::test]
355    async fn test_enhanced_tool_router() {
356        let mut router = ToolRouter::new();
357        router.register(TestTool);
358
359        assert_eq!(router.list_tools(), vec!["test_tool"]);
360        assert!(router.get("test_tool").is_some());
361        assert!(router.get("nonexistent").is_none());
362    }
363
364    #[test]
365    fn test_parameters_extraction() {
366        let json = serde_json::json!({
367            "message": "hello",
368            "count": 42
369        });
370
371        let params: Parameters<TestParams> = Parameters::from_json(json).unwrap();
372        assert_eq!(params.0.message, "hello");
373        assert_eq!(params.0.count, Some(42));
374    }
375
376    #[test]
377    fn test_call_tool_result() {
378        let result = CallToolResult::success(vec![Content::text("success".to_string())]);
379        assert!(!result.is_error);
380        assert_eq!(result.content.len(), 1);
381
382        let error_result = CallToolResult::error("error message".to_string());
383        assert!(error_result.is_error);
384    }
385}