1use 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#[derive(Debug, Clone)]
17pub struct Parameters<T>(pub T);
18
19impl<T> Parameters<T>
20where
21 T: for<'de> Deserialize<'de>,
22{
23 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#[derive(Debug, Clone, Serialize)]
33pub struct CallToolResult {
34 pub content: Vec<Content>,
35 pub is_error: bool,
36}
37
38#[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#[derive(Debug, Clone, Serialize)]
52pub struct ResourceReference {
53 pub uri: String,
54 pub text: Option<String>,
55}
56
57impl CallToolResult {
58 pub fn success(content: Vec<Content>) -> Self {
60 Self {
61 content,
62 is_error: false,
63 }
64 }
65
66 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 pub fn text(text: String) -> Self {
78 Self::Text { text }
79 }
80
81 pub fn image(data: String, mime_type: String) -> Self {
83 Self::Image { data, mime_type }
84 }
85
86 pub fn resource(uri: String, text: Option<String>) -> Self {
88 Self::Resource {
89 resource: ResourceReference { uri, text },
90 }
91 }
92}
93
94#[async_trait]
96pub trait EnhancedTool: Send + Sync {
97 async fn execute_enhanced(
99 &self,
100 parameters: JsonValue,
101 host: &dyn HostIntegration,
102 ) -> Result<CallToolResult, ToolError>;
103
104 fn name(&self) -> &str;
106
107 fn description(&self) -> &str;
109
110 fn parameter_schema(&self) -> JsonValue;
112
113 fn requires_permission(&self) -> Permission;
115
116 fn clone_enhanced(&self) -> Box<dyn EnhancedTool>;
118}
119
120pub struct ToolRouter {
122 tools: HashMap<String, Box<dyn EnhancedTool>>,
123}
124
125impl ToolRouter {
126 pub fn new() -> Self {
128 Self {
129 tools: HashMap::new(),
130 }
131 }
132
133 pub fn register<T: EnhancedTool + 'static>(&mut self, tool: T) {
135 self.tools.insert(tool.name().to_string(), Box::new(tool));
136 }
137
138 pub fn get(&self, name: &str) -> Option<&dyn EnhancedTool> {
140 self.tools.get(name).map(|t| t.as_ref())
141 }
142
143 pub fn list_tools(&self) -> Vec<&str> {
145 self.tools.keys().map(|s| s.as_str()).collect()
146 }
147
148 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 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#[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
186pub struct EnhancedToolAdapter {
188 enhanced_tool: Box<dyn EnhancedTool>,
189}
190
191impl EnhancedToolAdapter {
192 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 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
259pub mod schema_utils {
261 use super::*;
262 #[cfg(feature = "mcp-server")]
263 use schemars::{JsonSchema, schema_for};
264
265 #[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 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 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}