Skip to main content

prism_mcp_rs/core/
tool.rs

1//! Tool system for MCP servers
2//!
3//! This module provides the abstraction for implementing and managing tools in MCP servers.
4//! Tools are functions that can be called by clients to perform specific operations,
5//! enhanced with complete parameter validation, type checking, and metadata support.
6
7use async_trait::async_trait;
8use serde_json::Value;
9use std::collections::HashMap;
10use std::time::Instant;
11
12use crate::core::error::{McpError, McpResult};
13use crate::core::tool_metadata::{
14    CategoryFilter, ImprovedToolMetadata, ToolBehaviorHints, ToolCategory, ToolDeprecation,
15};
16use crate::core::validation::{ParameterValidator, ValidationConfig};
17use crate::protocol::types::{
18    ClientCapabilities, ContentBlock, Icon, Implementation, ToolInfo, ToolInputSchema, ToolResult,
19};
20use crate::protocol::OperationResult;
21
22/// Trait for implementing tool handlers
23#[async_trait]
24pub trait ToolHandler: Send + Sync {
25    /// Execute the tool with the given arguments
26    ///
27    /// # Arguments
28    /// * `arguments` - Tool arguments as key-value pairs
29    ///
30    /// # Returns
31    /// Result containing the tool execution result or an error
32    async fn call(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult>;
33}
34
35/// One invocation of an MCP 2026 multi-round-trip capable tool.
36#[derive(Debug, Clone, PartialEq)]
37pub struct MultiRoundToolCall {
38    /// Tool arguments supplied on the original request and every retry.
39    pub arguments: HashMap<String, Value>,
40    /// Results of input requests issued by the previous server response.
41    pub input_responses: HashMap<String, Value>,
42    /// Opaque state minted by the handler and returned byte-for-byte by the client.
43    pub request_state: Option<String>,
44    /// Self-reported client identity from this request.
45    pub client_info: Option<Implementation>,
46    /// Capabilities declared for this request.
47    pub client_capabilities: ClientCapabilities,
48}
49
50/// Handler for tools that may require MCP 2026 multi-round-trip input.
51///
52/// A handler returns [`OperationResult::Complete`] when execution is finished,
53/// or [`OperationResult::InputRequired`] with client input requests and/or
54/// opaque request state. The SDK rejects an input-required result on a legacy
55/// connection.
56#[async_trait]
57pub trait MultiRoundToolHandler: Send + Sync {
58    async fn call(&self, call: MultiRoundToolCall) -> McpResult<OperationResult<ToolResult>>;
59}
60
61// ============================================================================
62// Flexible tool handler creation support
63// ============================================================================
64
65// Wrapper to convert closures to ToolHandler
66pub struct ClosureWrapper<F>(pub F);
67
68#[async_trait]
69impl<F> ToolHandler for ClosureWrapper<F>
70where
71    F: Fn(&HashMap<String, Value>) -> McpResult<Vec<ContentBlock>> + Send + Sync,
72{
73    async fn call(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
74        let content = (self.0)(&arguments)?;
75        Ok(ToolResult {
76            content,
77            is_error: Some(false),
78            meta: None,
79            structured_content: None,
80        })
81    }
82}
83
84/// Extension to create tools with closures
85pub fn closure_tool<F>(
86    name: String,
87    description: Option<String>,
88    input_schema: Value,
89    handler: F,
90) -> Tool
91where
92    F: Fn(&HashMap<String, Value>) -> McpResult<Vec<ContentBlock>> + Send + Sync + 'static,
93{
94    Tool::new(name, description, input_schema, ClosureWrapper(handler))
95}
96
97/// A registered tool with its handler, validation, and improved metadata
98pub struct Tool {
99    /// Information about the tool
100    pub info: ToolInfo,
101    /// Handler that implements the tool's functionality
102    pub handler: Box<dyn ToolHandler>,
103    /// Whether the tool is currently enabled
104    pub enabled: bool,
105    /// Parameter validator for input validation
106    pub validator: Option<ParameterValidator>,
107    /// improved metadata for tool behavior, categorization, and performance
108    pub improved_metadata: ImprovedToolMetadata,
109}
110
111impl Tool {
112    /// Create a new tool with the given information and handler
113    ///
114    /// # Arguments
115    /// * `name` - Name of the tool
116    /// * `description` - Optional description of the tool
117    /// * `input_schema` - JSON schema describing the input parameters
118    /// * `handler` - Implementation of the tool's functionality
119    pub fn new<H>(
120        name: String,
121        description: Option<String>,
122        input_schema: Value,
123        handler: H,
124    ) -> Self
125    where
126        H: ToolHandler + 'static,
127    {
128        // Create validator from schema
129        let validator = if input_schema.is_object() {
130            Some(ParameterValidator::new(input_schema.clone()))
131        } else {
132            None
133        };
134
135        Self {
136            info: ToolInfo {
137                name,
138                description,
139                input_schema: ToolInputSchema {
140                    schema_type: "object".to_string(),
141                    properties: input_schema
142                        .get("properties")
143                        .and_then(|p| p.as_object())
144                        .map(|obj| obj.iter().map(|(k, v)| (k.clone(), v.clone())).collect()),
145                    required: input_schema
146                        .get("required")
147                        .and_then(|r| r.as_array())
148                        .map(|arr| {
149                            arr.iter()
150                                .filter_map(|v| v.as_str().map(String::from))
151                                .collect()
152                        }),
153                    additional_properties: input_schema
154                        .as_object()
155                        .unwrap_or(&serde_json::Map::new())
156                        .iter()
157                        .filter(|(k, _)| !["type", "properties", "required"].contains(&k.as_str()))
158                        .map(|(k, v)| (k.clone(), v.clone()))
159                        .collect(),
160                },
161                output_schema: None,
162                annotations: None,
163                icons: None,
164                title: None,
165                meta: None,
166            },
167            handler: Box::new(handler),
168            enabled: true,
169            validator,
170            improved_metadata: ImprovedToolMetadata::new(),
171        }
172    }
173
174    /// Create a new tool with custom validation configuration
175    pub fn with_validation<H>(
176        name: String,
177        description: Option<String>,
178        input_schema: Value,
179        handler: H,
180        validation_config: ValidationConfig,
181    ) -> Self
182    where
183        H: ToolHandler + 'static,
184    {
185        let mut tool = Self::new(name, description, input_schema.clone(), handler);
186        if input_schema.is_object() {
187            tool.validator = Some(ParameterValidator::with_config(
188                input_schema,
189                validation_config,
190            ));
191        }
192        tool
193    }
194
195    /// Enable the tool
196    pub fn enable(&mut self) {
197        self.enabled = true;
198    }
199
200    /// Disable the tool
201    pub fn disable(&mut self) {
202        self.enabled = false;
203    }
204
205    /// Check if the tool is enabled
206    pub fn is_enabled(&self) -> bool {
207        self.enabled
208    }
209
210    /// Execute the tool if it's enabled with parameter validation and performance tracking
211    ///
212    /// # Arguments
213    /// * `arguments` - Tool arguments as key-value pairs
214    ///
215    /// # Returns
216    /// Result containing the tool execution result or an error
217    pub async fn call(&self, mut arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
218        if !self.enabled {
219            return Err(McpError::validation(format!(
220                "Tool '{}' is disabled",
221                self.info.name
222            )));
223        }
224
225        // Check for deprecation warning
226        if let Some(warning) = self.improved_metadata.deprecation_warning() {
227            eprintln!("Warning: {warning}");
228        }
229
230        // Validate and coerce parameters if validator is present
231        if let Some(ref validator) = self.validator {
232            validator.validate_and_coerce(&mut arguments).map_err(|e| {
233                McpError::validation(format!(
234                    "Tool '{}' parameter validation failed: {}",
235                    self.info.name, e
236                ))
237            })?;
238        }
239
240        // Track execution time and outcome
241        let start_time = Instant::now();
242        let result = self.handler.call(arguments).await;
243        let execution_time = start_time.elapsed();
244
245        // Update performance metrics using interior mutability
246        match &result {
247            Ok(_) => self.improved_metadata.record_success(execution_time),
248            Err(_) => self.improved_metadata.record_error(execution_time),
249        }
250
251        result
252    }
253
254    /// Execute the tool without validation or performance tracking (for specialized use cases)
255    pub async fn call_unchecked(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
256        if !self.enabled {
257            return Err(McpError::validation(format!(
258                "Tool '{}' is disabled",
259                self.info.name
260            )));
261        }
262
263        self.handler.call(arguments).await
264    }
265
266    /// Validate parameters without executing the tool
267    pub fn validate_parameters(&self, arguments: &mut HashMap<String, Value>) -> McpResult<()> {
268        if let Some(ref validator) = self.validator {
269            validator.validate_and_coerce(arguments).map_err(|e| {
270                McpError::validation(format!(
271                    "Tool '{}' parameter validation failed: {}",
272                    self.info.name, e
273                ))
274            })
275        } else {
276            Ok(())
277        }
278    }
279
280    // improved Metadata Management Methods
281
282    /// Set behavior hints for the tool
283    pub fn set_behavior_hints(&mut self, hints: ToolBehaviorHints) {
284        self.improved_metadata.behavior_hints = hints;
285    }
286
287    /// Get behavior hints for the tool
288    pub fn behavior_hints(&self) -> &ToolBehaviorHints {
289        &self.improved_metadata.behavior_hints
290    }
291
292    /// Set category for the tool
293    pub fn set_category(&mut self, category: ToolCategory) {
294        self.improved_metadata.category = Some(category);
295    }
296
297    /// Get category for the tool
298    pub fn category(&self) -> Option<&ToolCategory> {
299        self.improved_metadata.category.as_ref()
300    }
301
302    /// Set version for the tool
303    pub fn set_version(&mut self, version: String) {
304        self.improved_metadata.version = Some(version);
305    }
306
307    /// Get version of the tool
308    pub fn version(&self) -> Option<&String> {
309        self.improved_metadata.version.as_ref()
310    }
311
312    /// Set author for the tool
313    pub fn set_author(&mut self, author: String) {
314        self.improved_metadata.author = Some(author);
315    }
316
317    /// Get author of the tool
318    pub fn author(&self) -> Option<&String> {
319        self.improved_metadata.author.as_ref()
320    }
321
322    /// Mark tool as deprecated
323    pub fn deprecate(&mut self, deprecation: ToolDeprecation) {
324        self.improved_metadata.deprecation = Some(deprecation);
325    }
326
327    /// Check if tool is deprecated
328    pub fn is_deprecated(&self) -> bool {
329        self.improved_metadata.is_deprecated()
330    }
331
332    /// Get deprecation warning if tool is deprecated
333    pub fn deprecation_warning(&self) -> Option<String> {
334        self.improved_metadata.deprecation_warning()
335    }
336
337    /// Get performance metrics for the tool
338    pub fn performance_metrics(&self) -> crate::core::tool_metadata::ToolPerformanceMetrics {
339        self.improved_metadata.get_performance_snapshot()
340    }
341
342    /// Add custom metadata field
343    pub fn add_custom_metadata(&mut self, key: String, value: serde_json::Value) {
344        self.improved_metadata.custom.insert(key, value);
345    }
346
347    /// Get custom metadata field
348    pub fn get_custom_metadata(&self, key: &str) -> Option<&serde_json::Value> {
349        self.improved_metadata.custom.get(key)
350    }
351
352    /// Check if tool matches a category filter
353    pub fn matches_category_filter(&self, filter: &CategoryFilter) -> bool {
354        if let Some(ref category) = self.improved_metadata.category {
355            category.matches_filter(filter)
356        } else {
357            // If no category set, only match empty filters
358            filter.primary.is_none() && filter.secondary.is_none() && filter.tags.is_empty()
359        }
360    }
361
362    /// Check if tool is suitable for caching based on behavior hints
363    pub fn is_cacheable(&self) -> bool {
364        self.improved_metadata
365            .behavior_hints
366            .cacheable
367            .unwrap_or(false)
368            || (self
369                .improved_metadata
370                .behavior_hints
371                .read_only
372                .unwrap_or(false)
373                && self
374                    .improved_metadata
375                    .behavior_hints
376                    .idempotent
377                    .unwrap_or(false))
378    }
379
380    /// Check if tool is destructive
381    pub fn is_destructive(&self) -> bool {
382        self.improved_metadata
383            .behavior_hints
384            .destructive
385            .unwrap_or(false)
386    }
387
388    /// Check if tool is read-only
389    pub fn is_read_only(&self) -> bool {
390        self.improved_metadata
391            .behavior_hints
392            .read_only
393            .unwrap_or(false)
394    }
395
396    /// Check if tool is idempotent
397    pub fn is_idempotent(&self) -> bool {
398        self.improved_metadata
399            .behavior_hints
400            .idempotent
401            .unwrap_or(false)
402    }
403
404    /// Check if tool requires authentication
405    pub fn requires_auth(&self) -> bool {
406        self.improved_metadata
407            .behavior_hints
408            .requires_auth
409            .unwrap_or(false)
410    }
411}
412
413impl std::fmt::Debug for Tool {
414    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
415        f.debug_struct("Tool")
416            .field("info", &self.info)
417            .field("enabled", &self.enabled)
418            .field("has_validator", &self.validator.is_some())
419            .field("deprecated", &self.is_deprecated())
420            .field("category", &self.improved_metadata.category)
421            .field("version", &self.improved_metadata.version)
422            .field("execution_count", &self.improved_metadata.execution_count())
423            .field("success_rate", &self.improved_metadata.success_rate())
424            .finish()
425    }
426}
427
428/// Helper macro for creating tools with schema validation
429///
430/// # Examples
431/// ```rust
432/// use prism_mcp_rs::{tool, core::tool::ToolHandler};
433/// use serde_json::json;
434///
435/// struct MyHandler;
436/// #[async_trait::async_trait]
437/// impl ToolHandler for MyHandler {
438/// async fn call(&self, _args: std::collections::HashMap<String, serde_json::Value>) -> prism_mcp_rs::McpResult<prism_mcp_rs::protocol::types::ToolResult> {
439/// // Implementation here
440/// todo!()
441/// }
442/// }
443///
444/// let tool = tool!(
445/// "my_tool",
446/// "A sample tool",
447/// json!({
448/// "type": "object",
449/// "properties": {
450/// "input": { "type": "string" }
451/// }
452/// }),
453/// MyHandler
454/// );
455/// ```
456#[macro_export]
457macro_rules! tool {
458    // New 4-parameter signature (aligned with current API)
459    ($name:expr, $description:expr, $schema:expr, $handler:expr) => {
460        $crate::core::tool::Tool::new(
461            $name.to_string(),
462            Some($description.to_string()),
463            $schema,
464            $handler,
465        )
466    };
467    // Legacy 3-parameter signature with None description
468    ($name:expr, $schema:expr, $handler:expr) => {
469        $crate::core::tool::Tool::new($name.to_string(), None, $schema, $handler)
470    };
471}
472
473// Common tool implementations
474
475/// Simple echo tool for testing
476pub struct EchoTool;
477
478#[async_trait]
479impl ToolHandler for EchoTool {
480    async fn call(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
481        let message = arguments
482            .get("message")
483            .and_then(|v| v.as_str())
484            .unwrap_or("Hello, World!");
485
486        Ok(ToolResult {
487            content: vec![ContentBlock::Text {
488                text: message.to_string(),
489                annotations: None,
490                meta: None,
491            }],
492            is_error: None,
493            structured_content: None,
494            meta: None,
495        })
496    }
497}
498
499/// Tool for adding two numbers
500pub struct AdditionTool;
501
502#[async_trait]
503impl ToolHandler for AdditionTool {
504    async fn call(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
505        let a = arguments
506            .get("a")
507            .and_then(|v| v.as_f64())
508            .ok_or_else(|| McpError::validation("Missing or invalid 'a' parameter"))?;
509
510        let b = arguments
511            .get("b")
512            .and_then(|v| v.as_f64())
513            .ok_or_else(|| McpError::validation("Missing or invalid 'b' parameter"))?;
514
515        let result = a + b;
516
517        Ok(ToolResult {
518            content: vec![ContentBlock::Text {
519                text: result.to_string(),
520                annotations: None,
521                meta: None,
522            }],
523            is_error: None,
524            structured_content: None,
525            meta: None,
526        })
527    }
528}
529
530/// Tool for getting current timestamp
531pub struct TimestampTool;
532
533#[async_trait]
534impl ToolHandler for TimestampTool {
535    async fn call(&self, _arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
536        use std::time::{SystemTime, UNIX_EPOCH};
537
538        let timestamp = SystemTime::now()
539            .duration_since(UNIX_EPOCH)
540            .map_err(|e| McpError::internal(e.to_string()))?
541            .as_secs();
542
543        Ok(ToolResult {
544            content: vec![ContentBlock::Text {
545                text: timestamp.to_string(),
546                annotations: None,
547                meta: None,
548            }],
549            is_error: None,
550            structured_content: None,
551            meta: None,
552        })
553    }
554}
555
556/// Builder for creating tools with fluent API, validation, and enhanced metadata
557pub struct ToolBuilder {
558    name: String,
559    description: Option<String>,
560    input_schema: Option<Value>,
561    validation_config: Option<ValidationConfig>,
562    title: Option<String>,
563    icons: Option<Vec<Icon>>,
564    behavior_hints: ToolBehaviorHints,
565    category: Option<ToolCategory>,
566    version: Option<String>,
567    author: Option<String>,
568    deprecation: Option<ToolDeprecation>,
569    custom_metadata: HashMap<String, serde_json::Value>,
570}
571
572impl ToolBuilder {
573    /// Create a new tool builder with the given name
574    pub fn new<S: Into<String>>(name: S) -> Self {
575        Self {
576            name: name.into(),
577            description: None,
578            input_schema: None,
579            validation_config: None,
580            title: None,
581            icons: None,
582            behavior_hints: ToolBehaviorHints::new(),
583            category: None,
584            version: None,
585            author: None,
586            deprecation: None,
587            custom_metadata: HashMap::new(),
588        }
589    }
590
591    /// Set the tool description
592    pub fn description<S: Into<String>>(mut self, description: S) -> Self {
593        self.description = Some(description.into());
594        self
595    }
596
597    /// Set the tool title (for UI display)
598    pub fn title<S: Into<String>>(mut self, title: S) -> Self {
599        self.title = Some(title.into());
600        self
601    }
602
603    /// Set tool icons (for UI display)
604    pub fn icons(mut self, icons: Vec<Icon>) -> Self {
605        self.icons = Some(icons);
606        self
607    }
608
609    /// Add a single tool icon (for UI display)
610    pub fn icon(mut self, icon: Icon) -> Self {
611        self.icons.get_or_insert_with(Vec::new).push(icon);
612        self
613    }
614
615    /// Set the input schema
616    pub fn schema(mut self, schema: Value) -> Self {
617        self.input_schema = Some(schema);
618        self
619    }
620
621    /// Set custom validation configuration
622    pub fn validation_config(mut self, config: ValidationConfig) -> Self {
623        self.validation_config = Some(config);
624        self
625    }
626
627    /// Enable strict validation (no additional properties, strict types)
628    pub fn strict_validation(mut self) -> Self {
629        self.validation_config = Some(ValidationConfig {
630            allow_additional: false,
631            coerce_types: false,
632            detailed_errors: true,
633            max_string_length: Some(1000),
634            max_array_length: Some(100),
635            max_object_properties: Some(50),
636        });
637        self
638    }
639
640    /// Enable permissive validation (allow additional properties, type coercion)
641    pub fn permissive_validation(mut self) -> Self {
642        self.validation_config = Some(ValidationConfig {
643            allow_additional: true,
644            coerce_types: true,
645            detailed_errors: false,
646            max_string_length: None,
647            max_array_length: None,
648            max_object_properties: None,
649        });
650        self
651    }
652
653    // improved Metadata Builder Methods
654
655    /// Set behavior hints for the tool
656    pub fn behavior_hints(mut self, hints: ToolBehaviorHints) -> Self {
657        self.behavior_hints = hints;
658        self
659    }
660
661    /// Mark tool as read-only
662    pub fn read_only(mut self) -> Self {
663        self.behavior_hints = self.behavior_hints.read_only();
664        self
665    }
666
667    /// Mark tool as destructive
668    pub fn destructive(mut self) -> Self {
669        self.behavior_hints = self.behavior_hints.destructive();
670        self
671    }
672
673    /// Mark tool as idempotent
674    pub fn idempotent(mut self) -> Self {
675        self.behavior_hints = self.behavior_hints.idempotent();
676        self
677    }
678
679    /// Mark tool as requiring authentication
680    pub fn requires_auth(mut self) -> Self {
681        self.behavior_hints = self.behavior_hints.requires_auth();
682        self
683    }
684
685    /// Mark tool as potentially long-running
686    pub fn long_running(mut self) -> Self {
687        self.behavior_hints = self.behavior_hints.long_running();
688        self
689    }
690
691    /// Mark tool as resource-intensive
692    pub fn resource_intensive(mut self) -> Self {
693        self.behavior_hints = self.behavior_hints.resource_intensive();
694        self
695    }
696
697    /// Mark tool results as cacheable
698    pub fn cacheable(mut self) -> Self {
699        self.behavior_hints = self.behavior_hints.cacheable();
700        self
701    }
702
703    /// Set tool category
704    pub fn category(mut self, category: ToolCategory) -> Self {
705        self.category = Some(category);
706        self
707    }
708
709    /// Set tool category with primary and secondary classification
710    pub fn category_simple(mut self, primary: String, secondary: Option<String>) -> Self {
711        let mut cat = ToolCategory::new(primary);
712        if let Some(sec) = secondary {
713            cat = cat.with_secondary(sec);
714        }
715        self.category = Some(cat);
716        self
717    }
718
719    /// Add category tag
720    pub fn tag(mut self, tag: String) -> Self {
721        if let Some(ref mut category) = self.category {
722            category.tags.insert(tag);
723        } else {
724            let mut cat = ToolCategory::new("general".to_string());
725            cat.tags.insert(tag);
726            self.category = Some(cat);
727        }
728        self
729    }
730
731    /// Set tool version
732    pub fn version<S: Into<String>>(mut self, version: S) -> Self {
733        self.version = Some(version.into());
734        self
735    }
736
737    /// Set tool author
738    pub fn author<S: Into<String>>(mut self, author: S) -> Self {
739        self.author = Some(author.into());
740        self
741    }
742
743    /// Mark tool as deprecated
744    pub fn deprecated(mut self, deprecation: ToolDeprecation) -> Self {
745        self.deprecation = Some(deprecation);
746        self
747    }
748
749    /// Mark tool as deprecated with simple reason
750    pub fn deprecated_simple<S: Into<String>>(mut self, reason: S) -> Self {
751        self.deprecation = Some(ToolDeprecation::new(reason.into()));
752        self
753    }
754
755    /// Add custom metadata field
756    pub fn custom_metadata<S: Into<String>>(mut self, key: S, value: serde_json::Value) -> Self {
757        self.custom_metadata.insert(key.into(), value);
758        self
759    }
760
761    /// Build the tool with the given handler
762    pub fn build<H>(self, handler: H) -> McpResult<Tool>
763    where
764        H: ToolHandler + 'static,
765    {
766        let schema = self.input_schema.unwrap_or_else(|| {
767            serde_json::json!({
768                "type": "object",
769                "properties": {},
770                "additionalProperties": true
771            })
772        });
773
774        let mut tool = if let Some(config) = self.validation_config {
775            Tool::with_validation(self.name, self.description, schema, handler, config)
776        } else {
777            Tool::new(self.name, self.description, schema, handler)
778        };
779
780        // Set title if provided
781        if let Some(title) = self.title {
782            tool.info.title = Some(title);
783        }
784        if let Some(icons) = self.icons {
785            tool.info.icons = Some(icons);
786        }
787
788        // Apply improved metadata
789        let mut improved_metadata =
790            ImprovedToolMetadata::new().with_behavior_hints(self.behavior_hints);
791
792        if let Some(category) = self.category {
793            improved_metadata = improved_metadata.with_category(category);
794        }
795
796        if let Some(version) = self.version {
797            improved_metadata = improved_metadata.with_version(version);
798        }
799
800        if let Some(author) = self.author {
801            improved_metadata = improved_metadata.with_author(author);
802        }
803
804        if let Some(deprecation) = self.deprecation {
805            improved_metadata = improved_metadata.deprecated(deprecation);
806        }
807
808        // Add custom metadata fields
809        for (key, value) in self.custom_metadata {
810            improved_metadata = improved_metadata.with_custom_field(key, value);
811        }
812
813        tool.improved_metadata = improved_metadata;
814
815        Ok(tool)
816    }
817
818    /// Build the tool with validation chain - allows chaining parameter validation
819    pub fn build_with_validation_chain<H>(
820        self,
821        handler: H,
822        validation_fn: impl Fn(&mut HashMap<String, Value>) -> McpResult<()> + Send + Sync + 'static,
823    ) -> McpResult<ValidationChainTool>
824    where
825        H: ToolHandler + 'static,
826    {
827        let tool = self.build(handler)?;
828        Ok(ValidationChainTool {
829            tool,
830            custom_validator: Box::new(validation_fn),
831        })
832    }
833}
834
835/// Tool wrapper that supports custom validation chains
836/// Type alias for validation function to reduce complexity
837type ValidationFunction = Box<dyn Fn(&mut HashMap<String, Value>) -> McpResult<()> + Send + Sync>;
838
839pub struct ValidationChainTool {
840    tool: Tool,
841    custom_validator: ValidationFunction,
842}
843
844#[async_trait]
845impl ToolHandler for ValidationChainTool {
846    async fn call(&self, mut arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
847        // Run custom validation first
848        (self.custom_validator)(&mut arguments)?;
849
850        // Then run the tool's built-in validation and execution
851        self.tool.call(arguments).await
852    }
853}
854
855// ============================================================================
856// improved Tool Creation Helpers and Macros
857// ============================================================================
858
859/// Create a validated tool with typed parameters
860#[macro_export]
861macro_rules! validated_tool {
862    (
863        name: $name:expr,
864        description: $desc:expr,
865        parameters: {
866            $( $param_name:ident: $param_type:ident $( ( $( $constraint:ident: $value:expr ),* ) )? ),*
867        },
868        handler: $handler:expr
869    ) => {{
870        use $crate::core::validation::{create_tool_schema, param_schema};
871
872        let params = vec![
873            $(
874                {
875                    $(
876                        // With constraints
877                        param_schema!($param_type stringify!($param_name), $( $constraint: $value ),*)
878                    )?
879                    $(
880                        // Without constraints (fallback)
881                        param_schema!($param_type stringify!($param_name))
882                    )?
883                }
884            ),*
885        ];
886
887        let required = vec![ $( stringify!($param_name) ),* ];
888        let schema = create_tool_schema(params, required);
889
890        $crate::core::tool::Tool::new(
891            $name.to_string(),
892            Some($desc.to_string()),
893            schema,
894            $handler
895        )
896    }};
897}
898
899/// Helper function to create a simple string parameter tool
900pub fn create_string_tool<H>(
901    name: &str,
902    description: &str,
903    param_name: &str,
904    param_description: &str,
905    handler: H,
906) -> Tool
907where
908    H: ToolHandler + 'static,
909{
910    use serde_json::json;
911
912    let schema = json!({
913        "type": "object",
914        "properties": {
915            param_name: {
916                "type": "string",
917                "description": param_description
918            }
919        },
920        "required": [param_name]
921    });
922
923    Tool::new(
924        name.to_string(),
925        Some(description.to_string()),
926        schema,
927        handler,
928    )
929}
930
931/// Helper function to create a tool with multiple typed parameters
932pub fn create_typed_tool<H>(
933    name: &str,
934    description: &str,
935    parameters: Vec<(&str, &str, Value)>, // (name, description, schema)
936    required: Vec<&str>,
937    handler: H,
938) -> Tool
939where
940    H: ToolHandler + 'static,
941{
942    use serde_json::{json, Map};
943
944    let mut properties = Map::new();
945    for (param_name, param_desc, param_schema) in parameters {
946        let mut schema_with_desc = param_schema;
947        if let Some(obj) = schema_with_desc.as_object_mut() {
948            obj.insert("description".to_string(), json!(param_desc));
949        }
950        properties.insert(param_name.to_string(), schema_with_desc);
951    }
952
953    let schema = json!({
954        "type": "object",
955        "properties": properties,
956        "required": required
957    });
958
959    Tool::new(
960        name.to_string(),
961        Some(description.to_string()),
962        schema,
963        handler,
964    )
965}
966
967/// Trait for tools that can provide their own parameter validation
968pub trait ValidatedToolHandler: ToolHandler {
969    /// Get the JSON schema for this tool's parameters
970    fn parameter_schema() -> Value;
971
972    /// Get validation configuration for this tool
973    fn validation_config() -> ValidationConfig {
974        ValidationConfig::default()
975    }
976
977    /// Create a tool instance with built-in validation
978    fn create_tool(name: String, description: Option<String>, handler: Self) -> Tool
979    where
980        Self: Sized + 'static,
981    {
982        Tool::with_validation(
983            name,
984            description,
985            Self::parameter_schema(),
986            handler,
987            Self::validation_config(),
988        )
989    }
990}
991
992// ============================================================================
993// improved Built-in Tool Examples
994// ============================================================================
995
996/// Calculator tool with comprehensive validation
997pub struct CalculatorTool;
998
999#[async_trait]
1000impl ToolHandler for CalculatorTool {
1001    async fn call(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
1002        let operation = arguments
1003            .get("operation")
1004            .and_then(|v| v.as_str())
1005            .ok_or_else(|| McpError::validation("Missing 'operation' parameter"))?;
1006
1007        let a = arguments
1008            .get("a")
1009            .and_then(|v| v.as_f64())
1010            .ok_or_else(|| McpError::validation("Missing or invalid 'a' parameter"))?;
1011
1012        let b = arguments
1013            .get("b")
1014            .and_then(|v| v.as_f64())
1015            .ok_or_else(|| McpError::validation("Missing or invalid 'b' parameter"))?;
1016
1017        let result = match operation {
1018            "add" => a + b,
1019            "subtract" => a - b,
1020            "multiply" => a * b,
1021            "divide" => {
1022                if b == 0.0 {
1023                    return Ok(ToolResult {
1024                        content: vec![ContentBlock::Text {
1025                            text: "Error: Division by zero".to_string(),
1026                            annotations: None,
1027                            meta: None,
1028                        }],
1029                        is_error: Some(true),
1030                        structured_content: Some(serde_json::json!({
1031                            "error": "division_by_zero",
1032                            "message": "Cannot divide by zero"
1033                        })),
1034                        meta: None,
1035                    });
1036                }
1037                a / b
1038            }
1039            _ => {
1040                return Err(McpError::validation(format!(
1041                    "Unsupported operation: {operation}"
1042                )));
1043            }
1044        };
1045
1046        Ok(ToolResult {
1047            content: vec![ContentBlock::Text {
1048                text: result.to_string(),
1049                annotations: None,
1050                meta: None,
1051            }],
1052            is_error: None,
1053            structured_content: Some(serde_json::json!({
1054                "operation": operation,
1055                "operands": [a, b],
1056                "result": result
1057            })),
1058            meta: None,
1059        })
1060    }
1061}
1062
1063impl ValidatedToolHandler for CalculatorTool {
1064    fn parameter_schema() -> Value {
1065        use crate::core::validation::create_tool_schema;
1066        use crate::param_schema;
1067
1068        create_tool_schema(
1069            vec![
1070                param_schema!(enum "operation", values: ["add", "subtract", "multiply", "divide"]),
1071                param_schema!(number "a", min: -1000000, max: 1000000),
1072                param_schema!(number "b", min: -1000000, max: 1000000),
1073            ],
1074            vec!["operation", "a", "b"],
1075        )
1076    }
1077
1078    fn validation_config() -> ValidationConfig {
1079        ValidationConfig {
1080            allow_additional: false,
1081            coerce_types: true,
1082            detailed_errors: true,
1083            max_string_length: Some(20),
1084            max_array_length: Some(10),
1085            max_object_properties: Some(10),
1086        }
1087    }
1088}
1089
1090/// Text processing tool with string validation
1091pub struct TextProcessorTool;
1092
1093#[async_trait]
1094impl ToolHandler for TextProcessorTool {
1095    async fn call(&self, arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
1096        let text = arguments
1097            .get("text")
1098            .and_then(|v| v.as_str())
1099            .ok_or_else(|| McpError::validation("Missing 'text' parameter"))?;
1100
1101        let operation = arguments
1102            .get("operation")
1103            .and_then(|v| v.as_str())
1104            .unwrap_or("uppercase");
1105
1106        let result = match operation {
1107            "uppercase" => text.to_uppercase(),
1108            "lowercase" => text.to_lowercase(),
1109            "reverse" => text.chars().rev().collect(),
1110            "word_count" => text.split_whitespace().count().to_string(),
1111            "char_count" => text.len().to_string(),
1112            _ => {
1113                return Err(McpError::validation(format!(
1114                    "Unsupported operation: {operation}"
1115                )));
1116            }
1117        };
1118
1119        Ok(ToolResult {
1120            content: vec![ContentBlock::Text {
1121                text: result.clone(),
1122                annotations: None,
1123                meta: None,
1124            }],
1125            is_error: None,
1126            structured_content: Some(serde_json::json!({
1127                "original_text": text,
1128                "operation": operation,
1129                "result": result,
1130                "length": text.len()
1131            })),
1132            meta: None,
1133        })
1134    }
1135}
1136
1137impl ValidatedToolHandler for TextProcessorTool {
1138    fn parameter_schema() -> Value {
1139        use crate::core::validation::create_tool_schema;
1140        use crate::param_schema;
1141
1142        create_tool_schema(
1143            vec![
1144                param_schema!(string "text", min: 1, max: 10000),
1145                param_schema!(enum "operation", values: ["uppercase", "lowercase", "reverse", "word_count", "char_count"]),
1146            ],
1147            vec!["text"],
1148        )
1149    }
1150}
1151
1152#[cfg(test)]
1153mod tests {
1154    use super::*;
1155
1156    use serde_json::json;
1157
1158    #[tokio::test]
1159    async fn test_echo_tool() {
1160        let tool = EchoTool;
1161        let mut args = HashMap::new();
1162        args.insert("message".to_string(), json!("test message"));
1163
1164        let result = tool.call(args).await.unwrap();
1165        match &result.content[0] {
1166            ContentBlock::Text { text, .. } => assert_eq!(text, "test message"),
1167            _ => panic!("Expected text content"),
1168        }
1169    }
1170
1171    #[tokio::test]
1172    async fn test_addition_tool() {
1173        let tool = AdditionTool;
1174        let mut args = HashMap::new();
1175        args.insert("a".to_string(), json!(5.0));
1176        args.insert("b".to_string(), json!(3.0));
1177
1178        let result = tool.call(args).await.unwrap();
1179        match &result.content[0] {
1180            ContentBlock::Text { text, .. } => assert_eq!(text, "8"),
1181            _ => panic!("Expected text content"),
1182        }
1183    }
1184
1185    #[test]
1186    fn test_tool_creation() {
1187        let tool = Tool::new(
1188            "test_tool".to_string(),
1189            Some("Test tool".to_string()),
1190            json!({"type": "object"}),
1191            EchoTool,
1192        );
1193
1194        assert_eq!(tool.info.name, "test_tool");
1195        assert_eq!(tool.info.description, Some("Test tool".to_string()));
1196        assert!(tool.is_enabled());
1197    }
1198
1199    #[test]
1200    fn test_tool_enable_disable() {
1201        let mut tool = Tool::new(
1202            "test_tool".to_string(),
1203            None,
1204            json!({"type": "object"}),
1205            EchoTool,
1206        );
1207
1208        assert!(tool.is_enabled());
1209
1210        tool.disable();
1211        assert!(!tool.is_enabled());
1212
1213        tool.enable();
1214        assert!(tool.is_enabled());
1215    }
1216
1217    #[tokio::test]
1218    async fn test_disabled_tool() {
1219        let mut tool = Tool::new(
1220            "test_tool".to_string(),
1221            None,
1222            json!({"type": "object"}),
1223            EchoTool,
1224        );
1225
1226        tool.disable();
1227
1228        let result = tool.call(HashMap::new()).await;
1229        assert!(result.is_err());
1230        match result.unwrap_err() {
1231            McpError::Validation(msg) => assert!(msg.contains("disabled")),
1232            _ => panic!("Expected validation error"),
1233        }
1234    }
1235
1236    #[test]
1237    fn test_tool_builder() {
1238        let tool = ToolBuilder::new("test")
1239            .description("A test tool")
1240            .icon(Icon {
1241                src: "https://example.com/tool-icon.svg".to_string(),
1242                mime_type: Some("image/svg+xml".to_string()),
1243                sizes: None,
1244                theme: None,
1245            })
1246            .schema(json!({"type": "object", "properties": {"x": {"type": "number"}}}))
1247            .build(EchoTool)
1248            .unwrap();
1249
1250        assert_eq!(tool.info.name, "test");
1251        assert_eq!(tool.info.description, Some("A test tool".to_string()));
1252        assert_eq!(tool.info.icons.as_ref().map(Vec::len), Some(1));
1253        assert!(tool.validator.is_some());
1254    }
1255
1256    #[test]
1257    fn test_improved_tool_builder() {
1258        let tool = ToolBuilder::new("improved_test")
1259            .title("improved Test Tool")
1260            .description("A test tool with improved features")
1261            .strict_validation()
1262            .schema(json!({
1263                "type": "object",
1264                "properties": {
1265                    "name": {"type": "string", "minLength": 2},
1266                    "age": {"type": "integer", "minimum": 0}
1267                },
1268                "required": ["name"]
1269            }))
1270            .build(EchoTool)
1271            .unwrap();
1272
1273        assert_eq!(tool.info.name, "improved_test");
1274        assert_eq!(tool.info.title, Some("improved Test Tool".to_string()));
1275        assert!(tool.validator.is_some());
1276    }
1277
1278    #[tokio::test]
1279    async fn test_parameter_validation() {
1280        let schema = json!({
1281            "type": "object",
1282            "properties": {
1283                "name": {"type": "string", "minLength": 2},
1284                "age": {"type": "integer", "minimum": 0, "maximum": 150}
1285            },
1286            "required": ["name", "age"]
1287        });
1288
1289        let tool = Tool::new(
1290            "validation_test".to_string(),
1291            Some("Test validation".to_string()),
1292            schema,
1293            EchoTool,
1294        );
1295
1296        // Valid parameters
1297        let mut valid_args = HashMap::new();
1298        valid_args.insert("name".to_string(), json!("Alice"));
1299        valid_args.insert("age".to_string(), json!(25));
1300        assert!(tool.validate_parameters(&mut valid_args).is_ok());
1301
1302        // Missing required parameter
1303        let mut invalid_args = HashMap::new();
1304        invalid_args.insert("name".to_string(), json!("Bob"));
1305        assert!(tool.validate_parameters(&mut invalid_args).is_err());
1306
1307        // Invalid parameter type with coercion
1308        let mut coercible_args = HashMap::new();
1309        coercible_args.insert("name".to_string(), json!("Charlie"));
1310        coercible_args.insert("age".to_string(), json!("30")); // String that can be coerced to number
1311        assert!(tool.validate_parameters(&mut coercible_args).is_ok());
1312        // After validation, should be coerced to number
1313        assert_eq!(coercible_args.get("age").unwrap().as_i64(), Some(30));
1314    }
1315
1316    #[tokio::test]
1317    async fn test_calculator_tool() {
1318        let tool = CalculatorTool::create_tool(
1319            "calculator".to_string(),
1320            Some("complete calculator".to_string()),
1321            CalculatorTool,
1322        );
1323
1324        // Test addition
1325        let mut args = HashMap::new();
1326        args.insert("operation".to_string(), json!("add"));
1327        args.insert("a".to_string(), json!(5));
1328        args.insert("b".to_string(), json!(3));
1329
1330        let result = tool.call(args).await.unwrap();
1331        assert_eq!(
1332            result.content[0],
1333            ContentBlock::Text {
1334                text: "8".to_string(),
1335                annotations: None,
1336                meta: None,
1337            }
1338        );
1339        assert!(result.structured_content.is_some());
1340
1341        // Test division by zero
1342        let mut args = HashMap::new();
1343        args.insert("operation".to_string(), json!("divide"));
1344        args.insert("a".to_string(), json!(10));
1345        args.insert("b".to_string(), json!(0));
1346
1347        let result = tool.call(args).await.unwrap();
1348        assert_eq!(result.is_error, Some(true));
1349        if let ContentBlock::Text { text, .. } = &result.content[0] {
1350            assert!(text.contains("Division by zero"));
1351        } else {
1352            panic!("Expected text content");
1353        }
1354    }
1355
1356    #[tokio::test]
1357    async fn test_text_processor_tool() {
1358        let tool = TextProcessorTool::create_tool(
1359            "text_processor".to_string(),
1360            Some("Text processing utility".to_string()),
1361            TextProcessorTool,
1362        );
1363
1364        // Test uppercase
1365        let mut args = HashMap::new();
1366        args.insert("text".to_string(), json!("hello world"));
1367        args.insert("operation".to_string(), json!("uppercase"));
1368
1369        let result = tool.call(args.clone()).await.unwrap();
1370        assert_eq!(
1371            result.content[0],
1372            ContentBlock::Text {
1373                text: "HELLO WORLD".to_string(),
1374                annotations: None,
1375                meta: None,
1376            }
1377        );
1378
1379        // Test word count
1380        args.insert("operation".to_string(), json!("word_count"));
1381        let result = tool.call(args).await.unwrap();
1382        assert_eq!(
1383            result.content[0],
1384            ContentBlock::Text {
1385                text: "2".to_string(),
1386                annotations: None,
1387                meta: None,
1388            }
1389        );
1390    }
1391
1392    #[test]
1393    fn test_create_typed_tool() {
1394        let tool = create_typed_tool(
1395            "typed_test",
1396            "A typed parameter test tool",
1397            vec![
1398                (
1399                    "username",
1400                    "User's name",
1401                    json!({"type": "string", "minLength": 3}),
1402                ),
1403                (
1404                    "age",
1405                    "User's age",
1406                    json!({"type": "integer", "minimum": 0}),
1407                ),
1408                (
1409                    "active",
1410                    "Whether user is active",
1411                    json!({"type": "boolean"}),
1412                ),
1413            ],
1414            vec!["username", "age"],
1415            EchoTool,
1416        );
1417
1418        assert_eq!(tool.info.name, "typed_test");
1419        assert!(tool.validator.is_some());
1420
1421        // Check that schema was built correctly
1422        let schema = &tool.info.input_schema;
1423        assert!(schema.properties.is_some());
1424        let props = schema.properties.as_ref().unwrap();
1425        assert!(props.contains_key("username"));
1426        assert!(props.contains_key("age"));
1427        assert!(props.contains_key("active"));
1428    }
1429
1430    #[test]
1431    fn test_validation_config_options() {
1432        // Test strict validation
1433        let strict_tool = ToolBuilder::new("strict")
1434            .strict_validation()
1435            .build(EchoTool)
1436            .unwrap();
1437        assert!(strict_tool.validator.is_some());
1438
1439        // Test permissive validation
1440        let permissive_tool = ToolBuilder::new("permissive")
1441            .permissive_validation()
1442            .build(EchoTool)
1443            .unwrap();
1444        assert!(permissive_tool.validator.is_some());
1445    }
1446}
1447
1448// ============================================================================
1449// Extension Trait for Better Ergonomics
1450// ============================================================================
1451
1452/// Extension trait for HashMap to make parameter extraction easier
1453pub trait ParameterExt {
1454    /// Extract a required string parameter
1455    fn get_string(&self, key: &str) -> McpResult<&str>;
1456
1457    /// Extract an optional string parameter
1458    fn get_optional_string(&self, key: &str) -> Option<&str>;
1459
1460    /// Extract a required number parameter
1461    fn get_number(&self, key: &str) -> McpResult<f64>;
1462
1463    /// Extract an optional number parameter
1464    fn get_optional_number(&self, key: &str) -> Option<f64>;
1465
1466    /// Extract a required integer parameter
1467    fn get_integer(&self, key: &str) -> McpResult<i64>;
1468
1469    /// Extract an optional integer parameter
1470    fn get_optional_integer(&self, key: &str) -> Option<i64>;
1471
1472    /// Extract a required boolean parameter
1473    fn get_boolean(&self, key: &str) -> McpResult<bool>;
1474
1475    /// Extract an optional boolean parameter
1476    fn get_optional_boolean(&self, key: &str) -> Option<bool>;
1477}
1478
1479impl ParameterExt for HashMap<String, Value> {
1480    fn get_string(&self, key: &str) -> McpResult<&str> {
1481        self.get(key).and_then(|v| v.as_str()).ok_or_else(|| {
1482            McpError::validation(format!("Missing or invalid string parameter: {key}"))
1483        })
1484    }
1485
1486    fn get_optional_string(&self, key: &str) -> Option<&str> {
1487        self.get(key).and_then(|v| v.as_str())
1488    }
1489
1490    fn get_number(&self, key: &str) -> McpResult<f64> {
1491        self.get(key).and_then(|v| v.as_f64()).ok_or_else(|| {
1492            McpError::validation(format!("Missing or invalid number parameter: {key}"))
1493        })
1494    }
1495
1496    fn get_optional_number(&self, key: &str) -> Option<f64> {
1497        self.get(key).and_then(|v| v.as_f64())
1498    }
1499
1500    fn get_integer(&self, key: &str) -> McpResult<i64> {
1501        self.get(key).and_then(|v| v.as_i64()).ok_or_else(|| {
1502            McpError::validation(format!("Missing or invalid integer parameter: {key}"))
1503        })
1504    }
1505
1506    fn get_optional_integer(&self, key: &str) -> Option<i64> {
1507        self.get(key).and_then(|v| v.as_i64())
1508    }
1509
1510    fn get_boolean(&self, key: &str) -> McpResult<bool> {
1511        self.get(key).and_then(|v| v.as_bool()).ok_or_else(|| {
1512            McpError::validation(format!("Missing or invalid boolean parameter: {key}"))
1513        })
1514    }
1515
1516    fn get_optional_boolean(&self, key: &str) -> Option<bool> {
1517        self.get(key).and_then(|v| v.as_bool())
1518    }
1519}
1520
1521#[cfg(test)]
1522mod improved_tests {
1523    use super::*;
1524    use crate::core::tool_metadata::*;
1525    use crate::prelude::ToolHandler;
1526    use std::time::Duration;
1527    use tokio;
1528
1529    // Test handler for basic tool functionality
1530    struct TestHandler {
1531        result: String,
1532        should_fail: bool,
1533    }
1534
1535    #[async_trait]
1536    impl ToolHandler for TestHandler {
1537        async fn call(&self, _arguments: HashMap<String, Value>) -> McpResult<ToolResult> {
1538            if self.should_fail {
1539                Err(McpError::validation("Test error".to_string()))
1540            } else {
1541                Ok(ToolResult {
1542                    content: vec![ContentBlock::Text {
1543                        text: self.result.clone(),
1544                        annotations: None,
1545                        meta: None,
1546                    }],
1547                    is_error: None,
1548                    structured_content: None,
1549                    meta: None,
1550                })
1551            }
1552        }
1553    }
1554
1555    #[tokio::test]
1556    async fn test_improved_tool_builder() {
1557        let handler = TestHandler {
1558            result: "test result".to_string(),
1559            should_fail: false,
1560        };
1561
1562        let tool = ToolBuilder::new("test_tool")
1563            .description("A test tool")
1564            .title("Test Tool")
1565            .version("1.0.0")
1566            .author("Test Author")
1567            .read_only()
1568            .idempotent()
1569            .cacheable()
1570            .category_simple("data".to_string(), Some("analysis".to_string()))
1571            .tag("testing".to_string())
1572            .tag("utility".to_string())
1573            .custom_metadata("priority".to_string(), serde_json::Value::from("high"))
1574            .build(handler)
1575            .expect("Failed to build tool");
1576
1577        assert_eq!(tool.info.name, "test_tool");
1578        assert_eq!(tool.info.description, Some("A test tool".to_string()));
1579        assert_eq!(tool.info.title, Some("Test Tool".to_string()));
1580        assert_eq!(tool.version(), Some(&"1.0.0".to_string()));
1581        assert_eq!(tool.author(), Some(&"Test Author".to_string()));
1582        assert!(tool.is_read_only());
1583        assert!(tool.is_idempotent());
1584        assert!(tool.is_cacheable());
1585        assert!(!tool.is_destructive());
1586        assert!(!tool.requires_auth());
1587
1588        let category = tool.category().unwrap();
1589        assert_eq!(category.primary, "data");
1590        assert_eq!(category.secondary, Some("analysis".to_string()));
1591        assert!(category.tags.contains("testing"));
1592        assert!(category.tags.contains("utility"));
1593
1594        let custom_priority = tool.get_custom_metadata("priority");
1595        assert_eq!(custom_priority, Some(&serde_json::Value::from("high")));
1596    }
1597
1598    #[tokio::test]
1599    async fn test_performance_tracking() {
1600        let handler = TestHandler {
1601            result: "success".to_string(),
1602            should_fail: false,
1603        };
1604
1605        let tool = ToolBuilder::new("performance_test")
1606            .build(handler)
1607            .expect("Failed to build tool");
1608
1609        // Initial state
1610        let metrics = tool.performance_metrics();
1611        assert_eq!(metrics.execution_count, 0);
1612        assert_eq!(metrics.success_count, 0);
1613        assert_eq!(metrics.error_count, 0);
1614
1615        // Execute tool successfully
1616        let result = tool.call(HashMap::new()).await;
1617        assert!(result.is_ok());
1618
1619        // Check updated metrics
1620        let metrics = tool.performance_metrics();
1621        assert_eq!(metrics.execution_count, 1);
1622        assert_eq!(metrics.success_count, 1);
1623        assert_eq!(metrics.error_count, 0);
1624        assert_eq!(metrics.success_rate, 100.0);
1625        assert!(metrics.average_execution_time > Duration::from_nanos(0));
1626    }
1627
1628    #[tokio::test]
1629    async fn test_performance_tracking_with_errors() {
1630        let handler = TestHandler {
1631            result: "".to_string(),
1632            should_fail: true,
1633        };
1634
1635        let tool = ToolBuilder::new("error_test")
1636            .build(handler)
1637            .expect("Failed to build tool");
1638
1639        // Execute tool with error
1640        let result = tool.call(HashMap::new()).await;
1641        assert!(result.is_err());
1642
1643        // Check error metrics
1644        let metrics = tool.performance_metrics();
1645        assert_eq!(metrics.execution_count, 1);
1646        assert_eq!(metrics.success_count, 0);
1647        assert_eq!(metrics.error_count, 1);
1648        assert_eq!(metrics.success_rate, 0.0);
1649    }
1650
1651    #[tokio::test]
1652    async fn test_deprecation_warning() {
1653        let handler = TestHandler {
1654            result: "deprecated result".to_string(),
1655            should_fail: false,
1656        };
1657
1658        let deprecation = ToolDeprecation::new("This tool is outdated".to_string())
1659            .with_replacement("new_tool".to_string())
1660            .with_severity(DeprecationSeverity::High);
1661
1662        let tool = ToolBuilder::new("deprecated_tool")
1663            .deprecated(deprecation)
1664            .build(handler)
1665            .expect("Failed to build tool");
1666
1667        assert!(tool.is_deprecated());
1668        let warning = tool.deprecation_warning().unwrap();
1669        assert!(warning.contains("deprecated"));
1670        assert!(warning.contains("outdated"));
1671        assert!(warning.contains("new_tool"));
1672    }
1673
1674    #[tokio::test]
1675    async fn test_category_filtering() {
1676        let category = ToolCategory::new("file".to_string())
1677            .with_secondary("read".to_string())
1678            .with_tag("filesystem".to_string())
1679            .with_tag("utility".to_string());
1680
1681        let handler = TestHandler {
1682            result: "filtered result".to_string(),
1683            should_fail: false,
1684        };
1685
1686        let tool = ToolBuilder::new("filterable_tool")
1687            .category(category)
1688            .build(handler)
1689            .expect("Failed to build tool");
1690
1691        // Test primary category filter
1692        let filter = CategoryFilter::new().with_primary("file".to_string());
1693        assert!(tool.matches_category_filter(&filter));
1694
1695        let filter = CategoryFilter::new().with_primary("network".to_string());
1696        assert!(!tool.matches_category_filter(&filter));
1697
1698        // Test tag filter
1699        let filter = CategoryFilter::new().with_tag("filesystem".to_string());
1700        assert!(tool.matches_category_filter(&filter));
1701
1702        let filter = CategoryFilter::new().with_tag("nonexistent".to_string());
1703        assert!(!tool.matches_category_filter(&filter));
1704
1705        // Test secondary category filter
1706        let filter = CategoryFilter::new().with_secondary("read".to_string());
1707        assert!(tool.matches_category_filter(&filter));
1708
1709        let filter = CategoryFilter::new().with_secondary("write".to_string());
1710        assert!(!tool.matches_category_filter(&filter));
1711    }
1712
1713    #[tokio::test]
1714    async fn test_behavior_hints() {
1715        let hints = ToolBehaviorHints::new()
1716            .read_only()
1717            .idempotent()
1718            .cacheable()
1719            .requires_auth()
1720            .long_running()
1721            .resource_intensive();
1722
1723        let handler = TestHandler {
1724            result: "hints result".to_string(),
1725            should_fail: false,
1726        };
1727
1728        let tool = ToolBuilder::new("hints_tool")
1729            .behavior_hints(hints)
1730            .build(handler)
1731            .expect("Failed to build tool");
1732
1733        assert!(tool.is_read_only());
1734        assert!(tool.is_idempotent());
1735        assert!(tool.is_cacheable());
1736        assert!(tool.requires_auth());
1737        assert!(!tool.is_destructive());
1738
1739        let behavior_hints = tool.behavior_hints();
1740        assert_eq!(behavior_hints.read_only, Some(true));
1741        assert_eq!(behavior_hints.idempotent, Some(true));
1742        assert_eq!(behavior_hints.cacheable, Some(true));
1743        assert_eq!(behavior_hints.requires_auth, Some(true));
1744        assert_eq!(behavior_hints.long_running, Some(true));
1745        assert_eq!(behavior_hints.resource_intensive, Some(true));
1746        assert_eq!(behavior_hints.destructive, None);
1747    }
1748
1749    #[tokio::test]
1750    async fn test_tool_enabling_disabling() {
1751        let handler = TestHandler {
1752            result: "enabled result".to_string(),
1753            should_fail: false,
1754        };
1755
1756        let mut tool = ToolBuilder::new("enable_test")
1757            .build(handler)
1758            .expect("Failed to build tool");
1759
1760        assert!(tool.is_enabled());
1761
1762        // Disable tool
1763        tool.disable();
1764        assert!(!tool.is_enabled());
1765
1766        // Try to call disabled tool
1767        let result = tool.call(HashMap::new()).await;
1768        assert!(result.is_err());
1769        assert!(result.unwrap_err().to_string().contains("disabled"));
1770
1771        // Re-enable tool
1772        tool.enable();
1773        assert!(tool.is_enabled());
1774
1775        // Should work again
1776        let result = tool.call(HashMap::new()).await;
1777        assert!(result.is_ok());
1778    }
1779
1780    #[tokio::test]
1781    async fn test_custom_metadata() {
1782        let handler = TestHandler {
1783            result: "metadata result".to_string(),
1784            should_fail: false,
1785        };
1786
1787        let mut tool = ToolBuilder::new("metadata_tool")
1788            .custom_metadata("priority".to_string(), serde_json::Value::from("high"))
1789            .custom_metadata("team".to_string(), serde_json::Value::from("backend"))
1790            .build(handler)
1791            .expect("Failed to build tool");
1792
1793        assert_eq!(
1794            tool.get_custom_metadata("priority"),
1795            Some(&serde_json::Value::from("high"))
1796        );
1797        assert_eq!(
1798            tool.get_custom_metadata("team"),
1799            Some(&serde_json::Value::from("backend"))
1800        );
1801        assert_eq!(tool.get_custom_metadata("nonexistent"), None);
1802
1803        // Add metadata after creation
1804        tool.add_custom_metadata(
1805            "environment".to_string(),
1806            serde_json::Value::from("production"),
1807        );
1808        assert_eq!(
1809            tool.get_custom_metadata("environment"),
1810            Some(&serde_json::Value::from("production"))
1811        );
1812    }
1813
1814    #[test]
1815    fn test_tool_debug_format() {
1816        let handler = TestHandler {
1817            result: "debug result".to_string(),
1818            should_fail: false,
1819        };
1820
1821        let tool = ToolBuilder::new("debug_tool")
1822            .version("2.0.0")
1823            .category_simple("debug".to_string(), None)
1824            .build(handler)
1825            .expect("Failed to build tool");
1826
1827        let debug_str = format!("{tool:?}");
1828        assert!(debug_str.contains("debug_tool"));
1829        assert!(debug_str.contains("enabled"));
1830        assert!(debug_str.contains("execution_count"));
1831        assert!(debug_str.contains("success_rate"));
1832    }
1833}